text:
# Utility functions, mostly for unit conversions.
# NMR quantities in Bruker notation [1]:
# BF1: Basic transmitter frequency (MHz)
# SF: Spectrometer frequency (MHz)
# SR: Spectrum reference frequency (Hz)
# SR = (SF - BF1) * 10^6
# SFO1: Transmitter frequency (MHz)
# O1: Transmitter frequency offset (Hz)
# O1 = (SFO1 - BF1) * 10^6
# Ω: Absolute frequency (MHz)
# ω: Frequency relative to SF1 (Hz)
# δ: Chemical shift δ (ppm)
# δ = 10^6 * (Ω - SF) / SF = ω / SF
# 1. Numerical indices indicate different nuclei

import Base: in

limits(s::Spectrum) = limits(s["O1P"], s["SW"], s["SF"], s["BF1"])
limits(o1p, sw, sf, bf) = begin
    shift = 1e6(bf-sf)/bf
    (o1p - sw/2 + shift, o1p + sw/2 + shift)
end

function in(δ, s::Spectrum)
    l,h = limits(s)
    l < δ < h
end

function chemical_shifts(s::Spectrum)
    lo,hi = limits(s)
    range(hi; stop=lo, length=length(s[:]))
end

"""
    ppmtomhz_abs(δ, bf[, sr])
Convert chemical shift δ to absolute frequency Ω in MHz.
bf in MHz and sr in Hz.
"""
function ppmtomhz_abs(δ, bf, sr = 0.0)
    bf = bf * 1e6
    1e-6(bf+sr)δ + bf + sr
end

"""
    ppmtohz(δ, bf[, sr])
Convert chemical shift δ in ppm to relative frequency ω in Hz.
bf in MHz and sr in Hz.
"""
function ppmtohz(δ, bf, sr = 0.0)
    (bf + sr*1e-6)δ
end

"""
    hztoppm(ω, bf[, sr])
Convert relative frequency ω in Hz to chemical shift δ in ppm .
bf in MHz and sr in Hz.
"""
function hztoppm(ω, bf, sr = 0.0)
    ω / (bf + sr*1e-6)
end

"""
    sr(sf, bf)
Return spectral reference (SR in Bruker notation) in Hz.
sf and bf in MHz.
"""
sr(sf, bf) = 1e6(sf - bf)

"""
    ppmtoindex(::Spectrum, δ)
Return the index in the processed spectrum corresponding to
chemical shift δ.
"""
function ppmtoindex(s::Spectrum, δ)
    min_δ ,max_δ = limits(s)
    @. Int(cld(s["SI"]*(max_δ - δ), (max_δ - min_δ)))
end

function ppmtoindex(s::Spectrum, rng::Tuple{Float64,Float64})
    r1,r2 = rng
    if r1>r2
        ppmtoindex(s,r1):ppmtoindex(s,r2)
    else
        ppmtoindex(s,r2):ppmtoindex(s,r1)
    end
end

hztoindex(f, sw, sf, si) = Int(cld(f*si, sw*sf))
hztoindex(s::Spectrum, f) = hztoindex(f, s["SW"], s["SF"], s["SI"])

Base.length(s::Spectrum) = length(s[:])
Base.length(p::ProcessedSpectrum) = length(p[:])

function union_range(ss::AbstractArray{Spectrum})
    lims = [limits(s) for s in ss]
    l = minimum(lim[1] for lim in lims)
    h = maximum(lim[2] for lim in lims)
    h,l
end

function union_shifts(ss::AbstractArray{Spectrum})
    h,l = union_range(ss)
    res = minimum(freq_resolution(s) for s in ss)
    h:-res:l
end

union_range(s::Spectrum) = union_range([s])
union_shifts(s::Spectrum) = union_shifts([s])

freq_resolution(s::Spectrum) = s["SW"] / length(s)

title(s::Spectrum) = s[s.default_proc].title

# Returns a copy of s with all but the given
# ranges zeroed out. The default proc for s will
# have its intrng adjusted to Δs as well.
function extract(s::Spectrum, Δs::AbstractArray{Intrng})
    res = deepcopy(s)
    res[res.default_proc].re_ft = zeros(length(s))
    res[res.default_proc].im_ft = zeros(length(s))
    res[res.default_proc].intrng = Δs
    for Δ in Δs
        rng = ppmtoindex(s, Δ)
        res[res.default_proc].re_ft[rng] .= s[s.default_proc].re_ft[rng]
        res[res.default_proc].im_ft[rng] .= s[s.default_proc].im_ft[rng]
    end
    res
end

extract(s::Spectrum, Δ::Intrng) = extract(s, [Δ])

"""Tunes `param` until `expr` evaluates to zero within δ.
`expr` must be monotonically increasing in terms of `param`."""
macro binary_opt(expr, param, min, max, δ, nattempts=5)
    :(m = $(esc(min));M = $(esc(max));d=$(esc(δ));z=zero(d);n=$(esc(nattempts));
    for i=1:n
        $(esc(param)) = (m+M)/2
        e = $(esc(expr))
        # println("val: $e, param:$((m+M)/2), min: $m, max: $M")
        if norm(e) < d
            break
        elseif e < z
            m = $(esc(param))
        else
            M = $(esc(param))
        end
    end)
end

### Pulse power profile

"""     powerprofile(pulse, dt, sfo1)
Pulse power profile for arbitrary pulse shape.
pulse: Amplitude profile of pulse.
dt: 'Sampling frequency' — unit of time between points in time. (s)
sfo1: Centre frequency of pulse. (Hz)"""
function powerprofile(pulse, dt, sfo1)
    ft = abs2.(fft(pulse))
    l = length(pulse)
    df = 1/(l*dt)
    fn = interpolate(ft[1:div(l,2)], BSpline(Cubic(Natural())), OnGrid())
    max = fn[1]
    f -> fn[abs(f-sfo1)/df+1]/max
end


text:
using JLD, Knet, ArgParse
include("src/header.jl")
const mfile = "src/model_nofeat_dyn.jl" # dynamic-oracle
include(mfile)
@msg mfile

# lm-related constants(fixed!)
# const lmdir  = "/scratch/users/okirnap/ud-treebanks-v2.2/chmodel_converted"

# To run everything on aws
const lmdir  = "../ud-treebanks-v2.2/chmodel_converted"
#const datadir = "/scratch/users/okirnap/ud-treebanks-v2.2"

function main(args=ARGS)
    s = ArgParseSettings()
    s.description = "Koc University dependency parser (c) Omer Kirnap, 2018."
    s.exc_handler = ArgParse.debug_handler
    @add_arg_table s begin
        ("--lmfile"; help="Language model file to load pretrained language mode")
        ("--loadfile"; help="Initialize model from file/ language model or all model")
        ("--datafiles"; nargs='+'; help="Input in conllu format. If provided, use the first for training, last for dev. If single file use both for train and dev.")
        ("--bestfile"; help="To save the best model file")
        ("--output"; help="Output parse of first datafile in conllu format")

        ("--epochs"; arg_type=Int; default=100; help="Epochs of training.")
        ("--optimization"; default="Adam"; help="Optimization algorithm and parameters.")
        ("--batchsize"; arg_type=Int; default=8; help="Number of sequences to train on in parallel.")
        ("--transfer"; action = :store_true; help="Whether to transfer or not")

        ("--dropout"; nargs='+'; arg_type=Float64; default=[0.6, 0.6]; help="Dropout probabilities.")
        ("--gclip"; arg_type=Float64; default=0.0; help="Grad clip")
        ("--savefile"; help="To save the final model file")

        ("--hidden"; nargs='+'; arg_type=Int; default=[2048]; help="MLP dims for final layer")
        ("--lstmhiddens"; nargs='+'; arg_type=Int; default=[256,256,256]; help="lstm dims (stack-1),(buff-2),(act-3)")
        ("--embed"; nargs='+'; arg_type=Int; default=[128, 128, 128];help="embedding sizes for postag(17),xpostag(?),feats(?) default 128,?,?.")
        ("--actembed"; arg_type=Int; default=64; help="action space embeddings")
        ("--wembed"; arg_type=Int; default=512; help="Word Embeddings dimension")
        ("--fembed"; arg_type=Int; default=128; help="Features Embeddings dimension")
        ("--deprel"; arg_type=Int; default=128; help="DepRel Embed")
        ("--treeType"; arg_type=Symbol; default=:tanh; help="Tree embedding function")

    end
    isa(args, AbstractString) && (args=split(args))
    odict = parse_args(args, s; as_symbols=true)
    println(s.description); flush(STDOUT);

    # to use in parsing
    testmode = (odict[:epochs] <= 0)

    odict[:ldrops] = (0.5, 0.5, (0.5, 0.5))
    # Set-Up
    odict[:arctype] = ArcHybridR1; odict[:posembed]=odict[:embed][1];
    odict[:embdim] = odict[:wembed]+odict[:posembed] # no features
    odict[:optimization] = eval(parse(odict[:optimization]))

    odict[:sthidden], odict[:acthidden], odict[:bufhidden] = odict[:lstmhiddens]
    odict[:stembed]=odict[:bufembed]=odict[:embdim]
    # hyper - parameters
    pdrop = odict[:dropout]; batchsize = odict[:batchsize];

    (odict[:lmfile] !=nothing) && (odict[:loadfile] != nothing) && error("Can not provide both files once")
    if odict[:lmfile] != nothing
        language_model = joinpath(lmdir, odict[:lmfile])
        @msg language_model
        d = load(language_model); v = create_vocab(d); wmodel = makewmodel(d);
    end

    odict[:transfer] && (odict[:loadfile]==nothing) && error("You have to provide load file for transfer learning")
    if odict[:loadfile] !=nothing # continue training
        @msg "loading model from $(odict[:loadfile])"
        bundle = load(odict[:loadfile])
        model, optims      = bundle["allmodel"], bundle["optims"]
        wmodel, v          = bundle["wordmodel"], bundle["vocab"]
        featdict, xposdict = bundle["featdict"], bundle["xposdict"]
        if gpu() >= 0 # To convert again in transfer
            wmodel = map2gpu(wmodel)
        end
        # Todo : effective way of using odict
    end

    corpora = []
    if length(odict[:datafiles]) > 1 # train
        f = odict[:datafiles][1] # train
        if !odict[:transfer] # creat featdict
            c, xposdict, featdict = load_conllu3(f, v)
            fnav = createfnav(featdict)
            shift_cfeats2!(c, fnav)
        else # load fnav from previous model file
            @msg :transfer_learning
            c, _ = load_conllu3(f, v, featdict, xposdict)
            fnav = bundle["fnav"]
            shift_cfeats2!(c, fnav)
        end
        push!(corpora, c)

        f = odict[:datafiles][2] # dev
        c, _ = load_conllu3(f, v, featdict, xposdict)
        shift_cfeats2!(c, fnav)
        push!(corpora, c)
    else # test or generate
        f = odict[:datafiles][1]
        c, _ = load_conllu3(f, v, featdict, xposdict)
        fnav = bundle["fnav"]
        shift_cfeats2!(c, fnav)
        push!(corpora, c)
    end

    cc=vcat(corpora...);
    @msg("context embeddings...")
    fillvecs!(wmodel, cc, v)
    @msg("Caching lm vectors...")
    map(cachelmvec!,corpora)

    odict[:featdim] = featdim(featdict)

    if !testmode

        if odict[:loadfile] == nothing # if provided no need to reinit
            @msg("Model initialization...")
            model, optims = initmodel1(odict, corpora[1][1])
        end

        println("opts=",[(k,v) for (k,v) in odict]...)
        @msg("calculating initial accuracies"); flush(STDOUT)
        acc1 = oracletest(model, corpora[1], odict[:arctype], odict[:lstmhiddens], odict[:batchsize], odict[:embdim]; pdrop=(0.0, 0.0))
        acc2 = oracletest(model, corpora[2], odict[:arctype], odict[:lstmhiddens], odict[:batchsize], odict[:embdim]; pdrop=(0.0, 0.0))
        bestlas = acc2; bestepoch=0;
        @msg "Initial tracc $acc1 devacc $acc2"
        nsentdev = length(corpora[2]); nworddev = sum(map(length, corpora[2]));
        @msg "nsentdev/nworddev=$nsentdev/$nworddev"
        sentbatches = minibatch(corpora[1], odict[:batchsize], maxlen=64, minlen=2, shuf=true)
        nsent = sum(map(length,sentbatches)); nsent0 = length(corpora[1])
        nword = sum(map(length,vcat(sentbatches...))); nword0 = sum(map(length,corpora[1]))
    end

    #error()
    # train-loop
    for epoch=1:odict[:epochs]
        shuffle!(sentbatches)
        @msg("nsent=$nsent/$nsent0 nword=$nword/$nword0")
        nwords = StopWatch()
        losses = Any[0,0,0]
        @time for sentences in sentbatches
            grads = oraclegrad(model, sentences, odict[:arctype], odict[:lstmhiddens], odict[:embdim],losses=losses, pdrop=pdrop, ldrops=odict[:ldrops])
            update!(model, grads, optims)
            nw = sum(map(length,sentences))
            if (speed = inc(nwords, nw)) != nothing
                date("$(nwords.ncurr) words $(round(Int,speed)) wps $(losses[3]) avgloss")
            end
        end
        empty_parses!(corpora[1]); empty_parses!(corpora[2]); #gc(); Knet.gc();gc();
        acc1 = oracletest(model, corpora[1], odict[:arctype], odict[:lstmhiddens],odict[:batchsize], odict[:embdim]; pdrop=(0.0, 0.0))
        acc2 = oracletest(model, corpora[2], odict[:arctype], odict[:lstmhiddens], odict[:batchsize], odict[:embdim]; pdrop=(0.0, 0.0))
        println()
        @msg "epoch $epoch train acc $acc1 dev acc $acc2"

        currlas = acc2
        if currlas > bestlas
            bestepoch = epoch
            bestlas = currlas
            if odict[:bestfile] != nothing
                JLD.save(odict[:bestfile],
                         "allmodel", model,
                         "optims", optims,
                         "featdict", featdict,
                         "xposdict", xposdict,
                         "wordmodel", map2cpu(wmodel),
                         "vocab", v,
                         "odict", odict,
                         "fnav", fnav
                         )
            end
        end
        if 9 < bestepoch < epoch-10
            @msg "bestlas $bestlas in $bestepoch"
            break
        end
    end

    # test-related
    if testmode # given for testing
        @assert length(corpora) == 1
        empty_parses!(corpora[1])
        odict2 = bundle["odict"]
        @msg :Testing
        testacc = oracletest(model, corpora[1], odict2[:arctype], odict2[:lstmhiddens], odict2[:batchsize], odict2[:embdim]; pdrop=(0.0, 0.0))
        @msg "Test acc $testacc"
    end

    # parsing
    if odict[:output] != nothing
        @msg :parsing
        writeconllu(corpora[1], odict[:datafiles][1], odict[:output])
    end
    @msg :done
end
!isinteractive() && main(ARGS)

ext: jl


text:
module Evaluation

using DataFrames
using DataStructures: OrderedDict
using MLDataPattern: nobs, getobs
using Statistics

using ..Transforms: ColTrans

export prob2class, probs2class, ids2classes, accuracy, cmatrix

# binary for single output
prob2class(data::AbstractMatrix{<:AbstractFloat}) = reshape(Int.(getobs(data)), :)

probs2class(data::AbstractMatrix{<:AbstractFloat}) = reshape(getindex.(argmax(data; dims=1), 1), :)

ids2classes(is::Symbol...) = begin pairs = split.(string.(is), "_"); Symbol(pairs[1][1]) => [string(v[2]) for v in pairs] end
ids2classes(i::ColTrans) = ids2classes(outids(i)...)

function accuracy(Y::AbstractVector{S}, T::AbstractVector{S}) where S
    size(Y) == size(T) || error("different dimensions size(Y)=$(size(Y)), size(T)=$(size(T))")
    mean(Y .== T)
end

accuracy(Y::AbstractMatrix{<:AbstractFloat}, T::AbstractMatrix{<:AbstractFloat}) = accuracy(probs2class(Y), probs2class(T))

function cmatrix(Y::AbstractVector{S}, T::AbstractVector{S}, classnames::Union{AbstractVector,Nothing} = nothing) where S
    size(Y) == size(T) || error("different dimensions size(Y)=$(size(Y)), size(T)=$(size(T))")
    if isnothing(classnames)
        classnames = unique(vcat(unique(Y), unique(T)))
        sort!(classnames)
    end
    cndict = OrderedDict((cn => i for (i, cn) in enumerate(classnames))...)

    cm = zeros(Int, length(classnames), length(classnames))
    for (gnd, pred) in zip(T, Y)
         cm[cndict[gnd], cndict[pred]] += 1
    end
    names = [Symbol("truth"), (Symbol("$n (pred)") for n in classnames)...]
    DataFrame([classnames, (col for col in eachcol(cm))...], names)
end

cmatrix(Y::AbstractMatrix{<:AbstractFloat}, T::AbstractMatrix{<:AbstractFloat}, classnames::Union{AbstractVector,Nothing} = nothing) =
    cmatrix(probs2class(Y), probs2class(T), classnames)

end

max_stars_repo_name: dhonza/MLGoodies


text:
# predefined adaptors for working with types from the Julia standard library

adapt_structure(to, xs::Union{Tuple,NamedTuple}) = map(x->adapt(to,x), xs)

## Closures

# two things can be captured: static parameters, and actual values (fields)

@eval function adapt_structure(to, f::F) where {F<:Function}
  isempty(F.parameters) && return f
  nsparams = length(F.parameters) - nfields(f)

  # TODO: we should adapt the static parameters too
  #       (but adapt currently only works with values)
  sparams = ntuple(i->F.parameters[i], nsparams)
  fields = ntuple(i->adapt(to, getfield(f, i)), nfields(f))

  # TODO: this assumes the typevars of the closure matches the sparams + fields.
  #       that may not always be true, and definitely isn't for arbitrary callable objects.
  ftyp = F.name.wrapper{sparams..., map(Core.Typeof, fields)...}
  $(Expr(:splatnew, :ftyp, :fields))
end

## Broadcast

import Base.Broadcast: Broadcasted, Extruded

adapt_structure(to, bc::Broadcasted{Style}) where Style =
  Broadcasted{Style}(adapt(to, bc.f), map(arg->adapt(to, arg), bc.args), bc.axes)

adapt_structure(to, ex::Extruded) =
    Extruded(adapt(to, ex.x), ex.keeps, ex.defaults)


text:
example_dir = joinpath(@__DIR__, "reference_image_tests")

tested_diff_path = joinpath(@__DIR__, "tested_different")
test_record_path = joinpath(@__DIR__, "test_recordings")

isdir(tested_diff_path) && rm(tested_diff_path, force = true, recursive = true)
mkpath(tested_diff_path)

isdir(test_record_path) && rm(test_record_path, force = true, recursive = true)
mkpath(test_record_path)
database = MakieGallery.load_database(joinpath.(example_dir, readdir(example_dir)))

examples = MakieGallery.record_examples(test_record_path);

@test length(examples) == length(database)
#
# # Download test images manually, so we can specify the folder
# # TODO, refactor makiegallery
path = MakieGallery.download_reference("v0.6.3")

MakieGallery.run_comparison(test_record_path, tested_diff_path, joinpath(dirname(path), "test_recordings"))


text:
#!/usr/local/bin/julia
#
const PKGS = ["Plots","Distributions","Loess","LossFunctions","RDatasets","Clustering",
              "DecisionTree","Flux","Metalhead","Images","Knet","Mamba","Gadfly"]

println("Installing required packages ...")

using Pkg
ipks = Pkg.installed()
for p in PKGS
    println("> $p")
    !haskey(ipks,p) && Pkg.add(p)
end
println("Done!")


text:
using LinearAlgebra, DataFrames, CSV, Plots, Optim, Statistics, StatsBase

# ============================================================
# Ler dados
# ============================================================
data = CSV.read("MER_T07_01.csv", DataFrame);
YYYYMM = data[!,"YYYYMM"]
index = mod.(YYYYMM,100) .!= 13
Value = data[!,"Value"]
y = parse.(Float64,Value[index][1:588])
y_insample = y[1:120]
y_outofsample =y[121:132]
n = length(y_insample)

function diff(y, t)
    out = zeros(length(y)-t)
    for i in 1:length(y)-t
        out[i] = y[i+t] - y[i]
    end
    return out
end

p = plot(y_insample,label = "y",legend=:bottomright)
savefig("figs\\y.png")
p = plot(autocov(y_insample),label = "FAC")
savefig("figs\\FAC.png")
p = plot(autocov(diff(y_insample, 1)),label = "FAC y\'")
savefig("figs\\FAC y\'.png")
p = plot(autocov(diff(y_insample, 12)),label = "FAC y\' (12)")
savefig("figs\\FAC y\' (12).png")

max_stars_repo_path: Modelos Estruturais/Kalman/src/2a.jl

max_issues_repo_path: Modelos Estruturais/Kalman/src/2a.jl

max_forks_repo_path: Modelos Estruturais/Kalman/src/2a.jl


text:
# # Single stack with HEVI solver tutorial based on the 3D Burgers + tracer equations

# This tutorial implements the Burgers equations with a tracer field
# in a single element stack. The flow is initialized with a horizontally
# uniform profile of horizontal velocity and uniform initial temperature. The fluid
# is heated from the bottom surface. Gaussian noise is imposed to the horizontal
# velocity field at each node at the start of the simulation. The tutorial demonstrates how to
#
#   * Initialize a [`BalanceLaw`](@ref ClimateMachine.BalanceLaws.BalanceLaw) in a single stack configuration;
#   * Return the horizontal velocity field to a given profile (e.g., large-scale advection);
#   * Remove any horizontal inhomogeneities or noise from the flow.
#   * Use horizontal explicit vertical implicit (HEVI) solver in the Single stack setup
#
# The second and third bullet points are demonstrated imposing Rayleigh friction, horizontal
# diffusion and 2D divergence damping to the horizontal momentum prognostic equation.

# Equations solved in balance law form:

# ```math
# \begin{align}
# \frac{∂ ρ}{∂ t} =& - ∇ ⋅ (ρ\mathbf{u}) \\
# \frac{∂ ρ\mathbf{u}}{∂ t} =& - ∇ ⋅ (-μ ∇\mathbf{u}) - ∇ ⋅ (ρ\mathbf{u} \mathbf{u}') - γ[ (ρ\mathbf{u}-ρ̄\mathbf{ū}) - (ρ\mathbf{u}-ρ̄\mathbf{ū})⋅ẑ ẑ] - ν_d ∇_h (∇_h ⋅ ρ\mathbf{u}) \\
# \frac{∂ ρcT}{∂ t} =& - ∇ ⋅ (-α ∇ρcT) - ∇ ⋅ (\mathbf{u} ρcT)
# \end{align}
# ```

# Boundary conditions:
# ```math
# \begin{align}
# z_{\mathrm{min}}: & ρ = 1 \\
# z_{\mathrm{min}}: & ρ\mathbf{u} = \mathbf{0} \\
# z_{\mathrm{min}}: & ρcT = ρc T_{\mathrm{fixed}} \\
# z_{\mathrm{max}}: & ρ = 1 \\
# z_{\mathrm{max}}: & ρ\mathbf{u} = \mathbf{0} \\
# z_{\mathrm{max}}: & -α∇ρcT = 0
# \end{align}
# ```

# where
#  - ``t`` is time
#  - ``ρ`` is the density
#  - ``\mathbf{u}`` is the velocity (vector)
#  - ``\mathbf{ū}`` is the horizontally averaged velocity (vector)
#  - ``μ`` is the dynamic viscosity tensor
#  - ``γ`` is the Rayleigh friction frequency
#  - ``ν_d`` is the horizontal divergence damping coefficient
#  - ``T`` is the temperature
#  - ``α`` is the thermal diffusivity tensor
#  - ``c`` is the heat capacity
#  - ``ρcT`` is the thermal energy

# Solving these equations is broken down into the following steps:
# 1) Preliminary configuration
# 2) PDEs
# 3) Space discretization
# 4) Time discretization
# 5) Solver hooks / callbacks
# 6) Solve
# 7) Post-processing

# # Preliminary configuration

# ## [Loading code](@id Loading-code-burgers-bjfnk)

# First, we'll load our pre-requisites
#  - load external packages:
using MPI
using Distributions
using OrderedCollections
using Plots
using StaticArrays
using LinearAlgebra: Diagonal, tr

#  - load CLIMAParameters and set up to use it:
using CLIMAParameters
struct EarthParameterSet <: AbstractEarthParameterSet end
const param_set = EarthParameterSet()

#  - load necessary ClimateMachine modules:
using ClimateMachine
using ClimateMachine.Mesh.Topologies
using ClimateMachine.Mesh.Grids
using ClimateMachine.Writers
using ClimateMachine.DGMethods
using ClimateMachine.DGMethods.NumericalFluxes
using ClimateMachine.BalanceLaws:
    BalanceLaw, Prognostic, Auxiliary, Gradient, GradientFlux, parameter_set

using ClimateMachine.Mesh.Geometry: LocalGeometry
using ClimateMachine.MPIStateArrays
using ClimateMachine.GenericCallbacks
using ClimateMachine.SystemSolvers
using ClimateMachine.ODESolvers
using ClimateMachine.VariableTemplates
using ClimateMachine.SingleStackUtils

#  - import necessary ClimateMachine modules: (`import`ing enables us to
#  provide implementations of these structs/methods)
using ClimateMachine.Orientations:
    Orientation,
    FlatOrientation,
    init_aux!,
    vertical_unit_vector,
    projection_tangential

import ClimateMachine.BalanceLaws:
    vars_state,
    source!,
    flux_second_order!,
    flux_first_order!,
    compute_gradient_argument!,
    compute_gradient_flux!,
    init_state_auxiliary!,
    init_state_prognostic!,
    BoundaryCondition,
    boundary_conditions,
    boundary_state!

# ## Initialization

# Define the float type (`Float64` or `Float32`)
const FT = Float64;
# Initialize ClimateMachine for CPU.
ClimateMachine.init(; disable_gpu = true);

const clima_dir = dirname(dirname(pathof(ClimateMachine)));

# Load some helper functions for plotting
include(joinpath(clima_dir, "docs", "plothelpers.jl"));

# # Define the set of Partial Differential Equations (PDEs)

# ## Define the model

# Model parameters can be stored in the particular [`BalanceLaw`](@ref
# ClimateMachine.BalanceLaws.BalanceLaw), in this case, the `BurgersEquation`:

Base.@kwdef struct BurgersEquation{FT, APS, O} <: BalanceLaw
    "Parameters"
    param_set::APS
    "Orientation model"
    orientation::O
    "Heat capacity"
    c::FT = 1
    "Vertical dynamic viscosity"
    μv::FT = 1e-4
    "Horizontal dynamic viscosity"
    μh::FT = 1
    "Vertical thermal diffusivity"
    αv::FT = 1e-2
    "Horizontal thermal diffusivity"
    αh::FT = 1
    "IC Gaussian noise standard deviation"
    σ::FT = 5e-2
    "Rayleigh damping"
    γ::FT = 5
    "Domain height"
    zmax::FT = 1
    "Initial conditions for temperature"
    initialT::FT = 295.15
    "Bottom boundary value for temperature (Dirichlet boundary conditions)"
    T_bottom::FT = 300.0
    "Top flux (α∇ρcT) at top boundary (Neumann boundary conditions)"
    flux_top::FT = 0.0
    "Divergence damping coefficient (horizontal)"
    νd::FT = 1
end

# Create an instance of the `BurgersEquation`:
orientation = FlatOrientation()

m = BurgersEquation{FT, typeof(param_set), typeof(orientation)}(
    param_set = param_set,
    orientation = orientation,
);

# This model dictates the flow control, using [Dynamic Multiple
# Dispatch](https://en.wikipedia.org/wiki/Multiple_dispatch), for which
# kernels are executed.

# ## Define the variables

# All of the methods defined in this section were `import`ed in
# [Loading code](@ref Loading-code-burgers) to let us provide
# implementations for our `BurgersEquation` as they will be used
# by the solver.

# Specify auxiliary variables for `BurgersEquation`
function vars_state(m::BurgersEquation, st::Auxiliary, FT)
    @vars begin
        coord::SVector{3, FT}
        orientation::vars_state(m.orientation, st, FT)
    end
end

# Specify prognostic variables, the variables solved for in the PDEs, for
# `BurgersEquation`
vars_state(::BurgersEquation, ::Prognostic, FT) =
    @vars(ρ::FT, ρu::SVector{3, FT}, ρcT::FT);

# Specify state variables whose gradients are needed for `BurgersEquation`
vars_state(::BurgersEquation, ::Gradient, FT) =
    @vars(u::SVector{3, FT}, ρcT::FT, ρu::SVector{3, FT});

# Specify gradient variables for `BurgersEquation`
vars_state(::BurgersEquation, ::GradientFlux, FT) = @vars(
    μ∇u::SMatrix{3, 3, FT, 9},
    α∇ρcT::SVector{3, FT},
    νd∇D::SMatrix{3, 3, FT, 9}
);

# ## Define the compute kernels

# Specify the initial values in `aux::Vars`, which are available in
# `init_state_prognostic!`. Note that
# - this method is only called at `t=0`.
# - `aux.coord` is available here because we've specified `coord` in `vars_state(m, aux, FT)`.
function nodal_init_state_auxiliary!(
    m::BurgersEquation,
    aux::Vars,
    tmp::Vars,
    geom::LocalGeometry,
)
    aux.coord = geom.coord
end;

# `init_aux!` initializes the auxiliary gravitational potential field needed for vertical projections
function init_state_auxiliary!(
    m::BurgersEquation,
    state_auxiliary::MPIStateArray,
    grid,
    direction,
)
    init_aux!(m, m.orientation, state_auxiliary, grid, direction)

    init_state_auxiliary!(
        m,
        nodal_init_state_auxiliary!,
        state_auxiliary,
        grid,
        direction,
    )
end;

# Specify the initial values in `state::Vars`. Note that
# - this method is only called at `t=0`.
# - `state.ρ`, `state.ρu` and`state.ρcT` are available here because we've specified `ρ`, `ρu` and `ρcT` in `vars_state(m, state, FT)`.
function init_state_prognostic!(
    m::BurgersEquation,
    state::Vars,
    aux::Vars,
    localgeo,
    t::Real,
)
    z = aux.coord[3]
    ε1 = rand(Normal(0, m.σ))
    ε2 = rand(Normal(0, m.σ))
    state.ρ = 1
    ρu = 1 - 4 * (z - m.zmax / 2)^2 + ε1
    ρv = 1 - 4 * (z - m.zmax / 2)^2 + ε2
    ρw = 0
    state.ρu = SVector(ρu, ρv, ρw)

    state.ρcT = state.ρ * m.c * m.initialT
end;

# The remaining methods, defined in this section, are called at every
# time-step in the solver by the [`BalanceLaw`](@ref
# ClimateMachine.BalanceLaws.BalanceLaw) framework.

# Since we have second-order fluxes, we must tell `ClimateMachine` to compute
# the gradient of `ρcT`, `u` and `ρu`. Here, we specify how `ρcT`, `u` and `ρu` are computed. Note that
# e.g. `transform.ρcT` is available here because we've specified `ρcT` in `vars_state(m, ::Gradient, FT)`.
function compute_gradient_argument!(
    m::BurgersEquation,
    transform::Vars,
    state::Vars,
    aux::Vars,
    t::Real,
)
    transform.ρcT = state.ρcT
    transform.u = state.ρu / state.ρ
    transform.ρu = state.ρu
end;

# Specify where in `diffusive::Vars` to store the computed gradient from
# `compute_gradient_argument!`. Note that:
#  - `diffusive.μ∇u` is available here because we've specified `μ∇u` in `vars_state(m, ::GradientFlux, FT)`.
#  - `∇transform.u` is available here because we've specified `u` in `vars_state(m, ::Gradient, FT)`.
#  - `diffusive.μ∇u` is built using an anisotropic diffusivity tensor.
#  - The `divergence` may be computed from the trace of tensor `∇ρu`.
function compute_gradient_flux!(
    m::BurgersEquation{FT},
    diffusive::Vars,
    ∇transform::Grad,
    state::Vars,
    aux::Vars,
    t::Real,
) where {FT}
    param_set = parameter_set(m)
    ∇ρu = ∇transform.ρu
    ẑ = vertical_unit_vector(m.orientation, param_set, aux)
    divergence = tr(∇ρu) - ẑ' * ∇ρu * ẑ
    diffusive.α∇ρcT = Diagonal(SVector(m.αh, m.αh, m.αv)) * ∇transform.ρcT
    diffusive.μ∇u = Diagonal(SVector(m.μh, m.μh, m.μv)) * ∇transform.u
    diffusive.νd∇D =
        Diagonal(SVector(m.νd, m.νd, FT(0))) *
        Diagonal(SVector(divergence, divergence, FT(0)))
end;

# Introduce Rayleigh friction towards a target profile as a source.
# Note that:
# - Rayleigh damping is only applied in the horizontal using the `projection_tangential` method.
function source!(
    m::BurgersEquation{FT},
    source::Vars,
    state::Vars,
    diffusive::Vars,
    aux::Vars,
    args...,
) where {FT}
    param_set = parameter_set(m)
    ẑ = vertical_unit_vector(m.orientation, param_set, aux)
    z = aux.coord[3]
    ρ̄ū =
        state.ρ * SVector{3, FT}(
            0.5 - 2 * (z - m.zmax / 2)^2,
            0.5 - 2 * (z - m.zmax / 2)^2,
            0.0,
        )
    ρu_p = state.ρu - ρ̄ū
    source.ρu -=
        m.γ * projection_tangential(m.orientation, param_set, aux, ρu_p)
end;

# Compute advective flux.
# Note that:
# - `state.ρu` is available here because we've specified `ρu` in `vars_state(m, state, FT)`.
function flux_first_order!(
    m::BurgersEquation,
    flux::Grad,
    state::Vars,
    aux::Vars,
    t::Real,
    _...,
)
    flux.ρ = state.ρu

    u = state.ρu / state.ρ
    flux.ρu = state.ρu * u'
    flux.ρcT = u * state.ρcT
end;

# Compute diffusive flux (e.g. ``F(μ, \mathbf{u}, t) = -μ∇\mathbf{u}`` in the original PDE).
# Note that:
# - `diffusive.μ∇u` is available here because we've specified `μ∇u` in `vars_state(m, ::GradientFlux, FT)`.
# - The divergence gradient can be written as a diffusive flux using a divergence diagonal tensor.
function flux_second_order!(
    m::BurgersEquation,
    flux::Grad,
    state::Vars,
    diffusive::Vars,
    hyperdiffusive::Vars,
    aux::Vars,
    t::Real,
)
    flux.ρcT -= diffusive.α∇ρcT
    flux.ρu -= diffusive.μ∇u
    flux.ρu -= diffusive.νd∇D
end;

# ### Boundary conditions

# Second-order terms in our equations, ``∇⋅(G)`` where ``G = μ∇\mathbf{u}``, are
# internally reformulated to first-order unknowns.
# Boundary conditions must be specified for all unknowns, both first-order and
# second-order unknowns which have been reformulated.

struct TopBC <: BoundaryCondition end;
struct BottomBC <: BoundaryCondition end;
boundary_conditions(::BurgersEquation) = (BottomBC(), TopBC());

# The boundary conditions for `ρ`, `ρu` and `ρcT` (first order unknowns)
function boundary_state!(
    nf,
    bc::BottomBC,
    m::BurgersEquation,
    state⁺::Vars,
    aux⁺::Vars,
    n⁻,
    _...,
)
    state⁺.ρ = 1
    state⁺.ρu = SVector(0, 0, 0)
    state⁺.ρcT = state⁺.ρ * m.c * m.T_bottom
end;
function boundary_state!(
    nf,
    bc::TopBC,
    m::BurgersEquation,
    state⁺::Vars,
    aux⁺::Vars,
    n⁻,
    _...,
)
    state⁺.ρ = 1
    state⁺.ρu = SVector(0, 0, 0)
end;

# The boundary conditions for `ρ`, `ρu` and `ρcT` are specified here for
# second-order unknowns

function boundary_state!(
    nf,
    bc::BottomBC,
    m::BurgersEquation,
    state⁺::Vars,
    diff⁺::Vars,
    hyperdiff⁺::Vars,
    aux⁺::Vars,
    n⁻,
    _...,
)
    state⁺.ρ = 1
    state⁺.ρu = SVector(0, 0, 0)
    state⁺.ρcT = state⁺.ρ * m.c * m.T_bottom
end;

function boundary_state!(
    nf,
    bc::TopBC,
    m::BurgersEquation,
    state⁺::Vars,
    diff⁺::Vars,
    hyperdiff⁺::Vars,
    aux⁺::Vars,
    n⁻,
    _...,
)
    state⁺.ρ = 1
    state⁺.ρu = SVector(0, 0, 0)
    diff⁺.α∇ρcT = -n⁻ * m.flux_top
end;

# # Spatial discretization

# Prescribe polynomial order of basis functions in finite elements
N_poly = 5;

# Specify the number of vertical elements
nelem_vert = 10;

# Specify the domain height
zmax = m.zmax;

# # Temporal discretization

# This initializes the ODE
# solver, the horizontal explicit vertical implicit scheme
# with ARK2GiraldoKellyConstantinescu method.

ode_solver_type = ClimateMachine.HEVISolverType(
    FT;
    solver_method = ARK2GiraldoKellyConstantinescu,
    linear_max_subspace_size = Int(30),
    linear_atol = FT(-1.0),
    linear_rtol = FT(1e-5),
    nonlinear_max_iterations = Int(10),
    nonlinear_rtol = FT(1e-4),
    nonlinear_ϵ = FT(1.e-10),
    preconditioner_update_freq = Int(50),
)

# Establish a `ClimateMachine` single stack configuration
driver_config = ClimateMachine.SingleStackConfiguration(
    "BurgersEquation",
    N_poly,
    nelem_vert,
    zmax,
    param_set,
    m,
    numerical_flux_first_order = CentralNumericalFluxFirstOrder(),
);

# # Time discretization

# Specify simulation time (SI units)
t0 = FT(0);
timeend = FT(1);

# We'll define the time-step based on the Fourier
# number and the [Courant number](https://en.wikipedia.org/wiki/Courant–Friedrichs–Lewy_condition)
# of the flow
Δ = min_node_distance(driver_config.grid)

given_Fourier = FT(0.5);
Fourier_bound = given_Fourier * Δ^2 / max(m.αh, m.μh, m.νd);
Courant_bound = FT(0.5) * Δ;

# We define the time step 50 times larger than that defined by CFL law
dt = FT(50.0) * min(Fourier_bound, Courant_bound)

# # Configure a `ClimateMachine` solver.

# This initializes the state vector and allocates memory for the solution in
# space (`dg` has the model `m`, which describes the PDEs as well as the
# function used for initialization).

solver_config = ClimateMachine.SolverConfiguration(
    t0,
    timeend,
    driver_config,
    ode_dt = dt,
    ode_solver_type = ode_solver_type,
);

# ## Inspect the initial conditions for a single nodal stack

# Let's export plots of the initial state
output_dir = @__DIR__;

mkpath(output_dir);

z_scale = 100 # convert from meters to cm
z_key = "z"
z_label = "z [cm]"
z = get_z(driver_config.grid; z_scale = z_scale)
state_vars = get_vars_from_nodal_stack(
    driver_config.grid,
    solver_config.Q,
    vars_state(m, Prognostic(), FT),
);

# Create an array to store the solution:
state_data = Dict[state_vars]  # store initial condition at ``t=0``
time_data = FT[0]                                      # store time data

# Generate plots of initial conditions for the southwest nodal stack
export_plot(
    z,
    time_data,
    state_data,
    ("ρcT",),
    joinpath(output_dir, "initial_condition_T_nodal_bjfnk.png");
    xlabel = "ρcT at southwest node",
    ylabel = z_label,
);
export_plot(
    z,
    time_data,
    state_data,
    ("ρu[1]",),
    joinpath(output_dir, "initial_condition_u_nodal_bjfnk.png");
    xlabel = "ρu at southwest node",
    ylabel = z_label,
);
export_plot(
    z,
    time_data,
    state_data,
    ("ρu[2]",),
    joinpath(output_dir, "initial_condition_v_nodal_bjfnk.png");
    xlabel = "ρv at southwest node",
    ylabel = z_label,
);

# ![](initial_condition_T_nodal_bjfnk.png)
# ![](initial_condition_u_nodal_bjfnk.png)

# ## Inspect the initial conditions for the horizontal averages

# Horizontal statistics of variables

state_vars_var = get_horizontal_variance(
    driver_config.grid,
    solver_config.Q,
    vars_state(m, Prognostic(), FT),
);

state_vars_avg = get_horizontal_mean(
    driver_config.grid,
    solver_config.Q,
    vars_state(m, Prognostic(), FT),
);

data_avg = Dict[state_vars_avg]
data_var = Dict[state_vars_var]

export_plot(
    z,
    time_data,
    data_avg,
    ("ρu[1]",),
    joinpath(output_dir, "initial_condition_avg_u_bjfnk.png");
    xlabel = "Horizontal mean of ρu",
    ylabel = z_label,
);
export_plot(
    z,
    time_data,
    data_var,
    ("ρu[1]",),
    joinpath(output_dir, "initial_condition_variance_u_bjfnk.png");
    xlabel = "Horizontal variance of ρu",
    ylabel = z_label,
);

# ![](initial_condition_avg_u_bjfnk.png)
# ![](initial_condition_variance_u_bjfnk.png)

# # Solver hooks / callbacks

# Define the number of outputs from `t0` to `timeend`
const n_outputs = 5;
const every_x_simulation_time = timeend / n_outputs;

# Create a dictionary for `z` coordinate (and convert to cm) NCDatasets IO:
dims = OrderedDict(z_key => collect(z));

# Create dictionaries to store outputs:
data_var = Dict[Dict([k => Dict() for k in 0:n_outputs]...),]
data_var[1] = state_vars_var

data_avg = Dict[Dict([k => Dict() for k in 0:n_outputs]...),]
data_avg[1] = state_vars_avg

data_nodal = Dict[Dict([k => Dict() for k in 0:n_outputs]...),]
data_nodal[1] = state_vars

# The `ClimateMachine`'s time-steppers provide hooks, or callbacks, which
# allow users to inject code to be executed at specified intervals. In this
# callback, the state variables are collected, combined into a single
# `OrderedDict` and written to a NetCDF file (for each output step `step`).
step = [0];
callback = GenericCallbacks.EveryXSimulationTime(every_x_simulation_time) do
    state_vars_var = get_horizontal_variance(
        driver_config.grid,
        solver_config.Q,
        vars_state(m, Prognostic(), FT),
    )
    state_vars_avg = get_horizontal_mean(
        driver_config.grid,
        solver_config.Q,
        vars_state(m, Prognostic(), FT),
    )
    state_vars = get_vars_from_nodal_stack(
        driver_config.grid,
        solver_config.Q,
        vars_state(m, Prognostic(), FT),
        i = 1,
        j = 1,
    )
    step[1] += 1
    push!(data_var, state_vars_var)
    push!(data_avg, state_vars_avg)
    push!(data_nodal, state_vars)
    push!(time_data, gettime(solver_config.solver))
    nothing
end;

# # Solve

# This is the main `ClimateMachine` solver invocation. While users do not have
# access to the time-stepping loop, code may be injected via `user_callbacks`,
# which is a `Tuple` of [`GenericCallbacks`](@ref ClimateMachine.GenericCallbacks).
ClimateMachine.invoke!(solver_config; user_callbacks = (callback,))

# # Post-processing

# Our solution has now been calculated and exported to NetCDF files in
# `output_dir`.

# Let's plot the horizontal statistics of `ρu` and `ρcT`, as well as the evolution of
# `ρu` for the southwest nodal stack:
export_plot(
    z,
    time_data,
    data_avg,
    ("ρu[1]"),
    joinpath(output_dir, "solution_vs_time_u_avg_bjfnk.png");
    xlabel = "Horizontal mean of ρu",
    ylabel = z_label,
);
export_plot(
    z,
    time_data,
    data_var,
    ("ρu[1]"),
    joinpath(output_dir, "variance_vs_time_u_bjfnk.png");
    xlabel = "Horizontal variance of ρu",
    ylabel = z_label,
);
export_plot(
    z,
    time_data,
    data_avg,
    ("ρcT"),
    joinpath(output_dir, "solution_vs_time_T_avg_bjfnk.png");
    xlabel = "Horizontal mean of ρcT",
    ylabel = z_label,
);
export_plot(
    z,
    time_data,
    data_var,
    ("ρcT"),
    joinpath(output_dir, "variance_vs_time_T_bjfnk.png");
    xlabel = "Horizontal variance of ρcT",
    ylabel = z_label,
);
export_plot(
    z,
    time_data,
    data_nodal,
    ("ρu[1]"),
    joinpath(output_dir, "solution_vs_time_u_nodal_bjfnk.png");
    xlabel = "ρu at southwest node",
    ylabel = z_label,
);
# ![](solution_vs_time_u_avg_bjfnk.png)
# ![](variance_vs_time_u_bjfnk.png)
# ![](solution_vs_time_T_avg_bjfnk.png)
# ![](variance_vs_time_T_bjfnk.png)
# ![](solution_vs_time_u_nodal_bjfnk.png)

# Rayleigh friction returns the horizontal velocity to the objective
# profile on the timescale of the simulation (1 second), since `γ`∼1. The horizontal viscosity
# and 2D divergence damping act to reduce the horizontal variance over the same timescale.
# The initial Gaussian noise is propagated to the temperature field through advection.
# The horizontal diffusivity acts to reduce this `ρcT` variance in time, although in a longer
# timescale.

# To run this file, and
# inspect the solution, include this tutorial in the Julia REPL
# with:

# ```julia
# include(joinpath("tutorials", "Atmos", "burgers_single_stack.jl"))
# ```


text:
import MKL_jll
using Test
using IntelVectorMath

include("common.jl")
include("real.jl")
include("complex.jl")


text:
using Documenter, RipQP

makedocs(
  modules = [RipQP],
  doctest = true,
  linkcheck = true,
  strict = true,
  format = Documenter.HTML(
    assets = ["assets/style.css"],
    prettyurls = get(ENV, "CI", nothing) == "true",
  ),
  sitename = "RipQP.jl",
  pages = [
    "Home" => "index.md",
    "API" => "API.md",
    "Tutorial" => "tutorial.md",
    "Reference" => "reference.md",
  ],
)

deploydocs(
  deps = nothing,
  make = nothing,
  repo = "github.com/JuliaSmoothOptimizers/RipQP.jl.git",
  target = "build",
  devbranch = "main",
)


text:
module HiddenPoisson

using ArgCheck
using Distributions
using Parameters
using StatsBase

import Base: show

export
    transitions, observation, # users define methods for these
    HiddenState, next_observation, next_observations # simulation interface

"""
```julia
transitions(model, state)
```

Return a vector of `Pair(λ, next_state)` elements, where the `λ`s are Poisson
rates for shocks, while `next_state` is either a closure that generates a new
state, or a deterministic value.

The user needs to supply a method for a given `model`.
"""
function transitions end

"""
```julia
observation(model, state)
```

Draw a (possibly random) observation of the `model` in `state`.

The user needs to supply a method for a given `model`.
"""
function observation end

"""
```julia
competing_poisson(rates)
```
Draw a random first arrival time for competing Poisson processes with `rates`.

Return `T, i` where

- `T` is the arrival time and
- `i` is the index of the event which happened first.
"""
function competing_poisson(rates::AbstractVector)
    first_time = rand(Exponential(1/sum(rates)))
    index = sample(indices(rates, 1), WeightVec(rates))
    first_time, index
end

"""
T, next_state = random_transition(trans)

Given a the object returned by `transitions(model, state)`, draw a random
stopping time and next state generator, and return these.
"""
function random_transition(trans)
    T, i = competing_poisson(first.(trans))
    T, trans[i].second
end

random_transition(model, state) = random_transition(transitions(model, state))

"""
Representation of a hidden Poisson process.

`model` contains the model parameters, and is used by `rates` and `observation`.

`state` is the current state (an arbitrary object).

`T` is the time left until the next Poisson event, which is simulated using `rates`.

`next_state` is the function that generates the new state at `T`.
"""
immutable HiddenState{Tm, Ts, Tt <: Real, Tn}
    model::Tm
    state::Ts
    T::Tt
    next_state::Tn
    function HiddenState(model, state, T, next_state)
        @argcheck T ≥ 0
        new(model, state, T, next_state)
    end
end

HiddenState{Tm,Ts,Tt,Tn}(model::Tm, state::Ts, T::Tt, next_state::Tn) =
    HiddenState{Tm,Ts,Tt,Tn}(model, state, T, next_state)

function HiddenState(model, state)
    HiddenState(model, state, random_transition(model, state)...)
end

function show(io::IO, hs::HiddenState)
    @unpack model, state, T, next_state = hs
    println(io, "HiddenState with model $(model)")
    println(io, "    state $(state)")
    println(io, "    time to next shock $(T)")
    print(io, "    next state $(next_state)")
end

"""
```julia
o, new_hs = next_observation(hs, t)
```

Return observation `o` after time `t` starting from hidden state `hs`, and the
new hidden state.
"""
function next_observation(hs::HiddenState, t::Real)
    @argcheck t ≥ zero(t)
    @unpack model, state, T, next_state = hs
    while t ≥ T
        t -= T
        state = isa(next_state, Function) ? next_state() : next_state
        T, next_state = random_transition(model, state)
    end
    observation(model, state), HiddenState(model, state, T-t, next_state)
end

"""
```julia
os, new_hs = next_observation(hs, ts)
```

Call `next_observation` repeatedly with the given times `ts`, collect the resulting
observations, and return them and the new hidden state.
"""
function next_observations(hs::HiddenState, ts)
    observations = [((o, hs) = next_observation(hs, t); o) for t in ts]
    observations, hs
end

end # module


text:
# Author: AJ Tyler
# Date: 24/01/22
module AtomicShells # This module labels the shell number of the atoms in a clutser
export ShellLabelling # This function returns the shell numbers of the atoms in a cluster
using Graphs # This module is used to call dijkstra_shortest_paths()
include("CNA.jl") # This file contains the CommonNeighbourAnalysis module
using .CommonNeighbourAnalysis # This module contains the CNA algorithm

"""
This function returns the shell number of the atoms of a cluster. This is done by first identifying an innermost core of atoms and then
determining the smallest number of bonds required to get to each atom from the core.

Input:
configuration: (Matrix{Float64}) 2D matrix contatining the coordinates of the atoms of the cluster
N: (Int64) The number of atoms in the cluster
B :(Bool) Whether configurations were found under a magnetic field of 0.3 au or not.
rCut: (Float64) Cutoff radius in units of equilibrium bond length.
req: (Float64) Equilibrium bond length (angstroms).

Output:
shells: (Vector{Int64}) Vector containing the shell number of each atom, same ordering as configuration.
bondGraph: (SimpleGraph) Graph representation of the configuration.
"""
function ShellLabelling(configuration,N,B,rCut,req)
    core = innerCore(configuration,N,req) # Determine which atoms form the innermost core
    bondGraph = adjacencyGraph(configuration,N,rCut,B,req) # Construct the graph representation of the cluster
    shells = fill(N,N) # Initialise shell numbers
    for innerAtom in core # For each innermost atom
        # Update the shell number of each atom to be the shortest number of bonds from innerAtom if smaller than current shell number.
        shells = min.(shells,(dijkstra_shortest_paths(bondGraph,innerAtom)).dists)
    end

	return shells, bondGraph # Return the shell numbers
end

"""
This function identifies the innermost core of atoms.

Input:
configuration: (Matrix{Float64}) 2D matrix contatining the coordinates of the atoms of the cluster
N: (Int64) The number of atoms in the cluster
req: (Float64) Equilibrium bond length (angstroms)
coreSize: (Float64) A measure of how large the core is, i.e. how many equilibrium bond lengths away from the centre of mass contains
all the innermost atoms. Default value of 0.85 (0.75 is too small, 1.0 is too large). May need further tuning.

Output:
core: (Vector{Int64}) Vector containg the atom numbers which are part of the inner core.
"""

function innerCore(configuration,N,req,coreSize=0.85)
    CoM = centreOfMass(configuration,N) # Compute the CoM of the cluster
    distances = radialDistances(configuration,N,CoM) # Compute the dsitances between all the atoms and the CoM
    # Return a vector of the atoms which are within 'coreSize' equilibrium bond lengths of the CoM.
    return findall(x-> x<coreSize*req,distances)
end

"""
This function computes the radial distances of each atom from the centre of mass.

Inputs:
configuration: (Matrix{Float64}) x,y,z positions of the cluster's atoms.
N: (Int64) Number of points
CoM: (Float64) Centre of mass of the points

Outputs:
radialDistances: (Vector{Float64}) The ith element of the array is the distance the ith atom in configuration is from the centre of mass
"""
function radialDistances(configuration,N,CoM)
    distances = Vector{Float64}(undef,N) # Preallocate the distances vector
    for i in 1:N # For each of the atoms in the cluster
        diff = configuration[i,:]-CoM[:] # calculate the vector between it and the CoM
		distances[i] = sqrt(sum(diff.*diff)) # Compute the length of the vector
    end

    return distances # Return the distacnes vector
end

"""
This function computes the centre of mass of N points (with equal mass)

Input:
configuration: (Matrix{Float64}) x,y,z positions of the cluster's atoms.
N: (Int64) Number of points

Output:
CoM: (Vector{Float64})
"""
function centreOfMass(configuration,N)
    # Initialise the CoM
    CoM = zeros(Float64,3,1)
    for i in 1:N # For each point
        CoM .+= configuration[i] # Add to CoM component wise
    end
    return CoM ./ N # Average over all points
end

end # End of module


text:
# PREAMBLE

# PKG_SETUP

# ## Setup

using ContactImplicitMPC
using LinearAlgebra

# ## Simulation
s_sim = get_simulation("flamingo", "slope_smooth_2D_lc", "slope");
s = get_simulation("flamingo", "flat_2D_lc", "flat");
model = s.model
env = s.env

# ## Reference Trajectory
ref_traj = deepcopy(ContactImplicitMPC.get_trajectory(s.model, s.env,
    joinpath(module_dir(), "src/dynamics/flamingo/gaits/gait_forward_36_4.jld2"),
    load_type = :split_traj_alt));

H = ref_traj.H
h = ref_traj.h

# ## MPC setup
N_sample = 5
H_mpc = 15
h_sim = h / N_sample
H_sim = 1000 #35000
κ_mpc = 2.0e-4

obj = TrackingVelocityObjective(model, env, H_mpc,
    v = [Diagonal(1e-3 * [1e0,1,1e4,1,1,1,1,1e4,1e4]) for t = 1:H_mpc],
    q = [Diagonal(1e-1 * [3e2, 1e-6, 3e2, 1, 1, 1, 1, 0.1, 0.1]) for t = 1:H_mpc],
    u = [Diagonal(3e-1 * [0.1; 0.1; 0.3; 0.3; ones(model.nu-6); 2; 2]) for t = 1:H_mpc],
    γ = [Diagonal(1.0e-100 * ones(model.nc)) for t = 1:H_mpc],
    b = [Diagonal(1.0e-100 * ones(model.nc * friction_dim(env))) for t = 1:H_mpc]);

p = ci_mpc_policy(ref_traj, s, obj,
    H_mpc = H_mpc,
    N_sample = N_sample,
    κ_mpc = κ_mpc,
	mode = :configuration,
	ip_opts = InteriorPointOptions(
					undercut = 5.0,
					κ_tol = κ_mpc,
					r_tol = 1.0e-8,
					diff_sol = true,
					solver = :empty_solver,
					max_time = 1e5,),
    n_opts = NewtonOptions(
        r_tol = 3e-4,
        max_iter = 5),
    mpc_opts = CIMPCOptions(
        altitude_update = true,
        altitude_impact_threshold = 0.02,
        )
    );

# ## Initial conditions
q1_sim, v1_sim = initial_conditions(ref_traj);

# ## Simulator
sim = simulator(s, H_sim, h=h_sim, policy=p);

# ## Simulate
simulate!(sim, q1_sim, v1_sim);

# ## Visualizer
vis = ContactImplicitMPC.Visualizer()
ContactImplicitMPC.render(vis)

# ## Visualize
anim = visualize_meshrobot!(vis, model, sim.traj, h=h_sim * 10, sample=10);

# ## Timing result
# Julia is [JIT-ed](https://en.wikipedia.org/wiki/Just-in-time_compilation) so re-run the MPC setup through Simulate for correct timing results.
process!(sim.stats, N_sample) # Time budget
H_sim * h_sim / sum(sim.stats.policy_time) # Speed ratio


text:
using HopfieldNets
using Base.Test

include(Pkg.dir("HopfieldNets", "demo", "letters.jl"))

patterns = hcat(X, O)

n = size(patterns, 1)

dnet = DiscreteHopfieldNet(n)
cnet = ContinuousHopfieldNet(n)

for net in [dnet, cnet]
    train!(net, patterns)

    e0 = energy(net)
    settle!(net, 1_000, false)
    eFinal = energy(net)
    @assert e0 != eFinal

    Xcorrupt = copy(X)
    for i = 2:7
         Xcorrupt[i] = 1
    end
    Xrestored = associate!(net, Xcorrupt)
    @test norm(Xcorrupt - Xrestored) > 1e-4
    @test norm(X - Xrestored) < 1e-4

    Ocorrupt = copy(O)
    for i = 2:7
         Ocorrupt[i] = -1
    end
    Orestored = associate!(net, Ocorrupt)
    @test norm(Ocorrupt - Orestored) > 1e-4
    @test norm(O - Orestored) < 1e-4
end


text:
@testset "execute" begin
    execute(`ls -al`, TEST_CONFIG, "myproject")
    execute(`julia -v`, TEST_CONFIG, join(TEST_TMP_DIR, "test-playground"))
end


text:
#=
P18 (**) Extract a slice from a list.
Given two indices, I and K, the slice is the list containing the elements between the I'th and K'th element of the original list (both limits included). Start counting the elements with 1.

Example:
* (slice '(a b c d e f g h i k) 3 7)
(C D E F G)
=#

function my_slice(coll, idx1, idx2)
    len = length(coll)
    if idx1 > len
        []
    elseif idx2 > len
        coll[idx1:end]
    else
        coll[idx1:idx2]
    end
end

@assert my_slice([:a, :b, :c, :d, :e, :f, :g, :h, :i, :k], 3, 7) == [:c, :d, :e, :f, :g]
@assert my_slice([:a, :b, :c], 2, 10) == [:b, :c]
@assert my_slice([:a, :b, :c], 5, 10) == []

println("Tests passed: JL-18.jl")


text:
let mode = :selector
    function switch!(s)
        enter_mode!(s)
        nothing
    end
    function selector_mode_parse(s)
        quote
            modes = $modes
            s = Symbol($s)
            if haskey(modes, s)
                $switch!(modes[s])
            else
                println("avaliable modes:")
                for each in keys(modes)
                    println(each)
                end
            end
        end
    end

    register() do
        modes[mode] = initrepl(
            selector_mode_parse;
            prompt_text="selector> ",
            prompt_color = :light_blue,
            start_key="\\",
            mode_name=string(mode)
        )
    end
end


text:
using BenchmarkTools

include("RegionTutorialBenchmarks.jl")
include("getindex.jl")

const SUITE = BenchmarkGroup()
SUITE["one_region"] = @benchmarkable run_oneregion()
SUITE["multi_regions"] = @benchmarkable run_multiregion()
SUITE["getindex"] = @benchmarkable run_getindex()


text:
# First, define the tight-binding system

syst = kwant.Builder()

# Here, we are only working with square lattices
a = 1
lat = kwant.lattice.square(a)

t = 1.0
W = 10
L = 30

# Define the scattering region

for i in range(0,length=L) ## note 1:W not range(W)
    for j in range(0,length=W) ## note 1:W not range(W)
        # On-site Hamiltonian
        syst[lat(i, j)] = 4 * t

        # Hopping in y-direction
        if j > 0
            syst[lat(i,j),lat(i,j-1)] = -t
        end

        # Hopping in x-direction
        if i > 0
            syst[lat(i, j), lat(i - 1, j)] = -t
        end
    end
end

# Then, define and attach the leads:

# First the lead to the left
# (Note: TranslationalSymmetry takes a real-space vector)
sym_left_lead = kwant.TranslationalSymmetry((-a, 0))
left_lead = kwant.Builder(sym_left_lead)

for j in range(0,length=W)
    left_lead[lat(0, j)] = 4 * t
    if j > 0
        left_lead[lat(0, j), lat(0, j - 1)] = -t
    end
    left_lead[lat(1, j), lat(0, j)] = -t
end
syst.attach_lead(left_lead)

# Then the lead to the right
sym_right_lead = kwant.TranslationalSymmetry((a, 0))
right_lead = kwant.Builder(sym_right_lead)

for j in range(0,length=W) ## 1:W not range(W)
    right_lead[lat(0, j)] = 4 * t
    if j > 0 ## note 1-based indexing
        right_lead[lat(0, j), lat(0, j - 1)] = -t
    end
    right_lead[lat(1, j), lat(0, j)] = -t
end


text:
@testset "Database" begin
  DB = Hecke.lattice_database()
  L = lattice(DB, 1)
  @assert L isa ZLat

  DB = Hecke.quadratic_lattice_database()
  L = lattice(DB, 1)
  @assert L isa QuadLat

  DB = Hecke.hermitian_lattice_database()
  L = lattice(DB, 1)
  @assert L isa HermLat
end


text:
@inline function dynamics(mechanism, body::Body{T}) where T
    state = body.state
    Δt = mechanism.Δt
    graph = mechanism.graph

    ezg = SA{T}[0; 0; -mechanism.g]
    dynT = body.m * ((state.vsol[2] - state.vc) / Δt + ezg) - state.Fk[1]

    J = body.J
    ω1 = state.ωc
    ω2 = state.ωsol[2]
    sq1 = sqrt(4 / Δt^2 - ω1' * ω1)
    sq2 = sqrt(4 / Δt^2 - ω2' * ω2)
    dynR = skewplusdiag(ω2, sq2) * (J * ω2) - skewplusdiag(ω1, sq1) * (J * ω1) - 2 * state.τk[1]

    state.d = [dynT;dynR]

    for childid in connections(graph, body.id)
        GtλTof!(mechanism, body, geteqconstraint(mechanism, childid))
    end

    for childid in ineqchildren(graph, body.id)
        NtγTof!(mechanism, body, getineqconstraint(mechanism, childid))
    end

    return state.d
end


text:
using Currencies
using Accounting
using Base.Test

@usingcurrencies USD, CAD

# make a ledger for testing
ledger = Ledger()

cash = asset!(ledger, "Cash")
ob = equity!(ledger, "Opening Balance")

transfer!(ledger, Date(2015, 11, 29), ob, cash, 100USD)

@test debit(ledger, cash, :USD) == 100USD
@test credit(ledger, ob, :USD) == 100USD
@test debit(ledger, cash, :CAD) == 0CAD

bal = balances(ledger)

@test bal[cash] == 100USD
@test bal[ob] == -100USD

# simulate some activity
bankloan = liability!(ledger, "Bank Loan")
inventory = asset!(ledger, "Inventory")
cogs = expense!(ledger, "Cost of Goods Sold")
sales = revenue!(ledger, "Sales")

transfer!(ledger, Date(2015, 11, 30), bankloan, cash, 5000USD)
transfer!(ledger, Date(2015, 12, 01), cash, inventory, 1000USD)
transfer!(ledger, Date(2015, 12, 02), inventory, cogs, 20USD)
transfer!(ledger, Date(2015, 12, 02), sales, cash, 30CAD)

bal = balances(ledger)

@test bal[cash] == Basket([4100USD, 30CAD])
@test bal[inventory] == 980USD
@test bal[bankloan] == -5000USD
@test bal[cogs] == 20USD
@test bal[sales] == -30CAD

sheet = balancesheet(ledger)

@test total(sheet.assets) == Basket([5080USD, 30CAD])
@test total(sheet.assets) == total(sheet.le)


text:
## Preloads
using Statistics
using FFTW
using Plots
using BenchmarkTools
using Profile
using LinearAlgebra
using Measures
#using StaticArrays
#using HybridArrays

# put the cwd on the path so Julia can find the module
push!(LOAD_PATH, pwd()*"/main")
using DHC_2DUtils

##
function rod_image(xcen, ycen, length, pa, fwhm; nx=256)
    # returns image of a rod with some (x,y) position, length,
    #   position angle, and FWHM in an (nx,nx) image.
    rodimage = zeros(nx,nx)

    x=0
    y=0
    sig = fwhm/2.355
    dtor = π/180
    # -------- define a unit vector in direction of rod at position angle pa
    ux = sin(pa*dtor)   #  0 deg is up
    uy = cos(pa*dtor)   # 90 deg to the right

    for i=1:nx
        for j=1:nx
            x=i-nx/2+xcen
            y=j-nx/2+ycen

            # -------- distance parallel and perpendicular to
            dpara =  ux*x + uy*y
            dperp = -uy*x + ux*y

            if abs(dpara)-length <0
                dpara= 0
            end
            dpara = abs(dpara)
            dpara = min(abs(dpara-length),dpara)

            rodimage[i,j] = exp(-(dperp^2+dpara^2)/(2*sig^2))
        end
    end
    rodimage ./= sqrt(sum(rodimage.^2))
    return rodimage
end

function rod_test_old(filter_hash, Nsam=1; doS20=false)

    Nf      = length(filter_hash["filt_value"])
    Nj      = length(filter_hash["j_value"])
    Nl      = length(filter_hash["theta_value"])
    Npix    = filter_hash["npix"]
    jlind   = filter_hash["psi_index"]
    angmax  = 180.0
    angstep = 2.5
    nang    = Int64(angmax ÷ angstep+1)
    energy     = zeros(nang)
    Siso       = zeros(Nj, Nj, Nl, nang)
    Sall       = zeros(Nl, Nj, Nl, Nj, nang)
    S2all      = zeros(Nf, Nf, nang)
    Ef         = zeros(Nf, nang)  # energy in each filter
    dang       = 180/Nl
    for i = 1:nang
        ang = (i-1)*angstep

        t = zeros(2+Nf+Nf*Nf)
        for k=0:Nsam-1
            rod1 = rod_image(10,10,30,ang+dang*k/4,8, nx=fhash["npix"])
            t   += DHC_compute(rod1, filter_hash, doS2=!doS20,doS20=doS20)
        end
        t   ./= Nsam

        #S20_big = reshape(t[2+Nf+1:2+Nf+Nf*Nf], Nf, Nf)
        #S20 = S20_big[1:Nj*Nl, 1:Nj*Nl]
        off2 = 2+Nf
        S2 = reshape(t[off2+1:off2+Nf*Nf], Nf, Nf)
        S2all[:,:,i] = S2
        # energy[i] = sum(diag(S20_big))
        # Ef[:,i]   = diag(S20_big)
        # if filter_hash["pc"]==1 Ef[1:Nj*Nl,i] .*= 2 end
        # Verify sum of diag(S20) = sum(S1)   (= energy)
        # println(i,"   ",energy[i],"   ",sum(t[3:2+Nf]))

        # what about j1, j2, Delta L
        #S = reshape(S20,Nj,Nl,Nj,Nl)  #  (j1, l1, j2, l2), symmetric in j and l
        S = S2[jlind,jlind]  # amazing this works...
        println("Sum of S2:  ",sum(S2))
        for j1 = 1:Nj
            for j2 = 1:Nj
                for l1 = 1:Nl
                    for l2 = 1:Nl
                        DeltaL = mod(l1-l2, Nl)
                        Siso[j1,j2,DeltaL+1,i] += S[j1,l1,j2,l2]
                        # Sall[j1,j2,l1,l2,i] = S[j1,l1,j2,l2]
                        Sall[l1,j1,l2,j2,i] = S[j1,l1,j2,l2] # flip indices
                        #println(Siso[j1,j2,DeltaL+1,i])
                    end
                end
            end
        end

    end
    println(std(Siso))
    return (Siso, Sall)
end

function make_mp4(fhash; doS20=false)

    clims = (-12,-3)
    if doS20
        clims=(-6,-1)
    end

    prefix = "frame"
    suffix = ".png"
    Sisofoo, Sall = rod_test_old(fhash, doS20=doS20)
    Nj      = length(fhash["j_value"])
    Nl      = length(fhash["theta_value"])
    Nframe  = size(Sall, 5)

    for i=1:Nframe
        p=heatmap(log10.(reshape(Sall[:,:,:,:,i],Nj*Nl,Nj*Nl)),
        size=(512,512),aspect_ratio=1,clims=clims)
        #display(p)
        outname = prefix*lpad(i,4,'0')*suffix
        println(outname)
        savefig(p, outname)
    end
    return 0
end

cmd=`ffmpeg -r 15 -f image2 -i frame%04d.png -vcodec libx264 -crf 25 -pix_fmt yuv420p test.mp4`
run(cmd)

function rod_test(filter_hash, Nsam=1; nx=256, doS20=false, doS1=false)
    # Loop over many rod position angles, return Siso for all
    Nf      = length(filter_hash["filt_value"])
    Nj      = length(filter_hash["j_value"])
    Nl      = length(filter_hash["theta_value"])
    Npix    = filter_hash["npix"]
    isomat  = doS1 ? filter_hash["S1_iso_mat"] : filter_hash["S2_iso_mat"]
    Niso    = size(isomat, 1)

    angmax  = 180.0
    angstep = 2.5
    Nang    = Int64(angmax ÷ angstep+1)

    Siso       = zeros(Niso, Nang)
    dang       = 180/Nl
    for i = 1:Nang
        rod1 = zeros(nx,nx)
        ang = i*angstep

        t = 0.0
        for k=1:Nsam
            rod1 = rod_image(1,1,30*(nx/256),ang+dang*k/4,8,nx=nx)
            rod1 .-= mean(rod1)
            doS2 = (!doS20 & !doS1)
            t   = t.+DHC_compute(rod1, filter_hash, doS2=doS2, doS20=doS20)
        end
        t   ./= Nsam

        off2 = 2+Nf
        Sarr = doS1 ? t[3:end] : t[off2+1:off2+Nf*Nf]

        println(i," of ",Nang,"  Sum of S:  ",sum(Sarr), "  Sum of rod: ",sum(rod1))
        Siso[:,i] = isomat * Sarr

    end
    println(std(Siso))
    return Siso
end

function rodtest_rms(Siso; frac=false)
    # reshape(Sisoin,Nj*Nj*Nl,37)
    sig = std(Siso, dims=2)
    mn  = mean(Siso, dims=2)
    if frac return sig./(mn.+1e-6) end
    return sig
end

function rodtest_plot_rms(Siso1,Siso4,Siso16;outname="test.png",
    label="Nsam=".*["1","4","16"], doS20=false, title=title)

    ylim=10.0 .^(-6,0)
    marg   = 10mm
    coefftype = doS20 ? "S20" : "S2"
    xlabel = coefftype*"_ISO Index"
    p1=plot(rodtest_rms(Siso1,frac=true),ylim=ylim,xlabel=xlabel,
            ylabel="Fractional RMS", yscale=:log10, line=(:gray, :dot),
            label=label[1],size=(700,1000),leftmargin=marg)
    plot!(p1,rodtest_rms(Siso4,frac=true), line=(:gray, :dash),label=label[2])
    plot!(p1,rodtest_rms(Siso16,frac=true), line=(:blue),label=label[3])
    annotate!(p1,10,0.6*ylim[2],text(title,:left,:blue,16))

    p2=plot(rodtest_rms(Siso1),ylim=[1E-10,1E-2],xlabel=xlabel,
            ylabel="RMS [Energy]", yscale=:log10, line=(:gray, :dot),
            label=label[1],size=(700,1000))
    plot!(p2,rodtest_rms(Siso4), line=(:gray, :dash),label=label[2])
    plot!(p2,rodtest_rms(Siso16), line=(:green),label=label[3])

    ps = [p1,p2]
    myplot = plot(ps..., layout=(2,1))
    savefig(myplot, outname)

    return
end

function rodtest_plot_rms_S1(Siso1,Siso4,Siso16;outname="test.png",
    label="Nsam=".*["1","4","16"], title=title)

    ylim=10.0 .^(-8,0)
    marg   = 10mm
    coefftype = "S1"
    xlabel = coefftype*"_ISO Index"
    p1=plot(rodtest_rms(Siso1,frac=true),ylim=ylim,xlabel=xlabel,
            ylabel="Fractional RMS", yscale=:log10, line=(:gray, :dot),
            label=label[1],size=(700,1000),leftmargin=marg)
    plot!(p1,rodtest_rms(Siso4,frac=true), line=(:gray, :dash),label=label[2])
    plot!(p1,rodtest_rms(Siso16,frac=true), line=(:blue),label=label[3])
    annotate!(p1,10,0.6*ylim[2],text(title,:left,:blue,16))

    p2=plot(rodtest_rms(Siso1),ylim=[1E-10,1E-2],xlabel=xlabel,
            ylabel="RMS [Energy]", yscale=:log10, line=(:gray, :dot),
            label=label[1],size=(700,1000))
    plot!(p2,rodtest_rms(Siso4), line=(:gray, :dash),label=label[2])
    plot!(p2,rodtest_rms(Siso16), line=(:green),label=label[3])

    ps = [p1,p2]
    myplot = plot(ps..., layout=(2,1))
    savefig(myplot, outname)

    return
end

function rodtest_plot_coeffs(Siso1,Siso4,Siso16;outname="test.png",
    label="Nsam=".*["1","4","16"], doS20=false, title=title)

    marg   = 4mm
    coefftype = doS20 ? "S20" : "S2"
    xval   = collect(0:2.5:180)
    xlabel = "Angle [deg]"
    ticks  = collect(0:45:180)

    p1=plot(xval,Siso1',yscale=:log10,legend=false,xlabel=xlabel,
             ylabel=coefftype,size=(700,1000),leftmargin=marg,xticks=ticks)
    annotate!(p1,0,.9,text(title,:left,:blue,16))

    p2=plot(xval,Siso16',yscale=:log10,legend=false,xlabel=xlabel,
             ylabel=coefftype,size=(700,1000),leftmargin=marg,xticks=ticks)
    #annotate!(p2,0,.9,text(title,:left,:blue,16))

    p3=plot(xval,Siso1',yscale=:log10,legend=false,xlabel=xlabel,
             ylabel=coefftype,size=(700,1000),leftmargin=marg,
             xticks=ticks,ylims=(1e-3,1e-2))
    annotate!(p3,0,.9,text(title,:left,:blue,16))

    p4=plot(xval,Siso16',yscale=:log10,legend=false,xlabel=xlabel,
             ylabel=coefftype,size=(700,1000),leftmargin=marg,
             xticks=ticks,ylims=(1e-3,1e-2))
    annotate!(p4,0,.9,text(title,:left,:blue,16))

    ps = [p1,p2,p3,p4]
    myplot = plot(ps..., layout=(2,2))
    savefig(myplot, outname)

    return
end

function rod_test_plot(Nsam; L=8, wd=1, nx=128,
                       outname="test.png", doS20=false)
    # run rod_test for 3 values of Nsam and output PNG file
    fhash = fink_filter_hash(1, L, wd=wd, nx=nx, Omega=true)
    Siso_A = rod_test(fhash, Nsam[1], nx=nx, doS20=doS20)
    Siso_B = rod_test(fhash, Nsam[2], nx=nx, doS20=doS20)
    Siso_C = rod_test(fhash, Nsam[3], nx=nx, doS20=doS20)
    mylabels = "Nsam=".*string.(Nsam)
    coefftype = doS20 ? "S20" : "S2"
    title = string(coefftype, "iso, L=", L, ", wd=", wd)
    rodtest_plot_rms(Siso_A, Siso_B, Siso_C, outname=outname,
                     label=mylabels, title=title, doS20=doS20)

    cname = replace(outname, "RMS" => "coeff")
    rodtest_plot_coeffs(Siso_A, Siso_B, Siso_C, outname=cname,
                        label=mylabels, title=title, doS20=doS20)

    return
end

function rod_test_plot_S1(Nsam; L=8, wd=1, nx=128,
                       outname="test.png", doS20=false)
    # run rod_test for 3 values of Nsam and output PNG file
    fhash = fink_filter_hash(1, L, wd=wd, nx=nx, Omega=true)
    println(fhash["wd"])
    Siso_A = rod_test(fhash, Nsam[1], nx=nx, doS1=true)
    Siso_B = rod_test(fhash, Nsam[2], nx=nx, doS1=true)
    Siso_C = rod_test(fhash, Nsam[3], nx=nx, doS1=true)
    mylabels = "Nsam=".*string.(Nsam)
    coefftype = doS20 ? "S20" : "S2"
    title = string(coefftype, "iso, L=", L, ", wd=", wd)
    rodtest_plot_rms_S1(Siso_A, Siso_B, Siso_C, outname=outname,
                     label=mylabels, title=title)

    return
end

# Just run this code to make the plots I sent on Mar 14, 2021 (π-day!)
nx=128
snx = string(nx)
Nsam = [1,4,16]

rod_test_plot(Nsam; L=8, wd=1, nx=nx, outname="S2-RMS-NL8-wd1-"*snx*".png")
rod_test_plot(Nsam; L=8, wd=2, nx=nx, outname="S2-RMS-NL8-wd2-"*snx*".png")
rod_test_plot(Nsam; L=8, wd=4, nx=nx, outname="S2-RMS-NL8-wd4-"*snx*".png")
rod_test_plot([1,4,32]; L=16, wd=4, nx=nx, outname="S2-RMS-NL16-wd4-"*snx*".png")

rod_test_plot(Nsam; L=8, wd=1, nx=nx, outname="S20-RMS-NL8-wd1-"*snx*".png", doS20=true)
rod_test_plot(Nsam; L=8, wd=2, nx=nx, outname="S20-RMS-NL8-wd2-"*snx*".png", doS20=true)
rod_test_plot(Nsam; L=8, wd=4, nx=nx, outname="S20-RMS-NL8-wd4-"*snx*".png", doS20=true)
rod_test_plot([1,4,32]; L=16, wd=4, nx=nx, outname="S20-RMS-NL16-wd4-"*snx*".png", doS20=true)

rod_test_plot_S1(Nsam; L=8, wd=1, nx=nx, outname="S1-RMS-NL8-wd1-"*snx*".png")
rod_test_plot_S1(Nsam; L=8, wd=2, nx=nx, outname="S1-RMS-NL8-wd2-"*snx*".png")
rod_test_plot_S1(Nsam; L=8, wd=4, nx=nx, outname="S1-RMS-NL8-wd4-"*snx*".png")
rod_test_plot_S1([1,4,32]; L=16, wd=4, nx=nx, outname="S1-RMS-NL16-wd4-"*snx*".png")

## AKS export
xval   = collect(0:2.5:180)

fhash = fink_filter_hash(1, 8, wd=2, nx=128, Omega=false)
Siso_A = rod_test(fhash, 1, nx=128, doS20=false)
h5write("../DHC/scratch_AKS/paper_data/rod_stable_L8_wd2_N1.h5", "data", Siso_A, deflate=3)
h5write("../DHC/scratch_AKS/paper_data/rod_stable_L8_wd2_N1.h5", "angles", xval, deflate=3)

fhash = fink_filter_hash(1, 16, wd=2, nx=128, Omega=false)
Siso_A = rod_test(fhash, 1, nx=128, doS20=false)
h5write("../DHC/scratch_AKS/paper_data/rod_stable_L16_wd2_N1.h5", "data", Siso_A, deflate=3)
h5write("../DHC/scratch_AKS/paper_data/rod_stable_L16_wd2_N1.h5", "angles", xval, deflate=3)

fhash = fink_filter_hash(1, 16, wd=4, nx=128, Omega=false)
Siso_A = rod_test(fhash, 1, nx=128, doS20=false)
h5write("../DHC/scratch_AKS/paper_data/rod_stable_L16_wd4_N1.h5", "data", Siso_A, deflate=3)
h5write("../DHC/scratch_AKS/paper_data/rod_stable_L16_wd4_N1.h5", "angles", xval, deflate=3)

fhash = fink_filter_hash(1, 8, wd=4, nx=128, Omega=false)
Siso_A = rod_test(fhash, 1, nx=128, doS20=false)
h5write("../DHC/scratch_AKS/paper_data/rod_stable_L8_wd4_N1.h5", "data", Siso_A, deflate=3)
h5write("../DHC/scratch_AKS/paper_data/rod_stable_L8_wd4_N1.h5", "angles", xval, deflate=3)

fhash = fink_filter_hash(1, 8, wd=1, nx=128, Omega=false)
Siso_A = rod_test(fhash, 1, nx=128, doS20=false)
h5write("../DHC/scratch_AKS/paper_data/rod_stable_L8_wd1_N1.h5", "data", Siso_A, deflate=3)
h5write("../DHC/scratch_AKS/paper_data/rod_stable_L8_wd1_N1.h5", "angles", xval, deflate=3)

fhash = fink_filter_hash(1, 8, wd=2, nx=128, Omega=false)
Siso_A = rod_test(fhash, 4, nx=128, doS20=false)
h5write("../DHC/scratch_AKS/paper_data/rod_stable_L8_wd2_N4.h5", "data", Siso_A, deflate=3)
h5write("../DHC/scratch_AKS/paper_data/rod_stable_L8_wd2_N4.h5", "angles", xval, deflate=3)

fhash = fink_filter_hash(1, 8, wd=2, nx=128, Omega=false)
Siso_A = rod_test(fhash, 16, nx=128, doS20=false)
h5write("../DHC/scratch_AKS/paper_data/rod_stable_L8_wd2_N16.h5", "data", Siso_A, deflate=3)
h5write("../DHC/scratch_AKS/paper_data/rod_stable_L8_wd2_N16.h5", "angles", xval, deflate=3)

test_image = rod_image(1,1,30*(128/256),0,8,nx=128)

heatmap(test_image,xlims = [46,81],ylims = [46,81])

function rod_test_noniso(filter_hash, Nsam=1; nx=256, doS20=false, doS1=true)
    # Loop over many rod position angles, return Siso for all
    Nf      = length(filter_hash["filt_value"])
    Nj      = length(filter_hash["j_value"])
    Nl      = length(filter_hash["theta_value"])
    Npix    = filter_hash["npix"]
    isomat  = doS1 ? filter_hash["S1_iso_mat"] : filter_hash["S2_iso_mat"]
    Niso    = size(isomat, 1)

    angmax  = 180.0
    angstep = 2.5
    Nang    = Int64(angmax ÷ angstep+1)

    Siso       = zeros(filter_hash["num_coeff"], Nang)
    dang       = 180/Nl
    for i = 1:Nang
        rod1 = zeros(nx,nx)
        ang = i*angstep

        t = 0.0
        for k=1:Nsam
            rod1 = rod_image(1,1,30*(nx/256),ang+dang*k/4,8,nx=nx)
            rod1 .-= mean(rod1)
            doS2 = true
            t   = t.+DHC_compute(rod1, filter_hash, doS2=doS2, doS20=doS20)
        end
        t   ./= Nsam

        off2 = 2+Nf
        Sarr = t

        println(i," of ",Nang,"  Sum of S:  ",sum(Sarr), "  Sum of rod: ",sum(rod1))
        Siso[:,i] = Sarr

    end
    println(std(Siso))
    return Siso
end

fhash = fink_filter_hash(1, 8, wd=2, nx=128, Omega=false)
Siso_A = rod_test_noniso(fhash, 1, nx=128, doS20=false)
h5write("../DHC/scratch_AKS/paper_data/rod_stable_L8_wd2_N1_noISO.h5", "data", Siso_A, deflate=3)
h5write("../DHC/scratch_AKS/paper_data/rod_stable_L8_wd2_N1_noISO.h5", "angles", xval, deflate=3)

max_stars_repo_name: andrew-saydjari/DHC

max_issues_repo_name: andrew-saydjari/DHC

max_forks_repo_name: andrew-saydjari/DHC


text:
mutable struct Heartbeat
    interval::Int # Heartbeat interval (in sec)
    last_sent::Float64
    last_ack::Float64
end

Heartbeat() = Heartbeat(0, 0.0, 0.0)
set_interval!(h::Heartbeat, msec::Int) = (h.interval = round(Int, msec / 1000))
update!(h::Heartbeat) = (h.last_sent = time())
set_ack!(h::Heartbeat) = (h.last_ack = time())
is_zombie(h::Heartbeat) = h.last_sent > h.last_ack


text:
@testset "633.sum-of-square-numbers.jl" begin
    for (i, j) in
        zip([1, 2, 3, 4, 5, 483761424, 1547186596, 1749724464, 45290257, 2094715898,
             1836642396, 896826837, 317827104, 1421291910, 1305562843, 1168481452],
            [true, true, false, true, true, false, true, false, true, false, false, false,
             false, false, false, false])
        @test judge_square_sum(i) == j
    end
end


text:
using Test
using MPI
using CLIMA
using CLIMA.Mesh.Topologies
using CLIMA.Mesh.Grids
using CLIMA.DGmethods
using CLIMA.DGmethods.NumericalFluxes
using CLIMA.MPIStateArrays
using LinearAlgebra
using Random

include("advection_diffusion_model.jl")

struct Box{dim} end
struct Sphere end

vertical_unit_vector(::Box{2}, ::SVector{3}) = SVector(0, 1, 0)
vertical_unit_vector(::Box{3}, ::SVector{3}) = SVector(0, 0, 1)
vertical_unit_vector(::Sphere, coord::SVector{3}) = coord / norm(coord)

projection(::EveryDirection, ::SVector{3}) = I
projection(::VerticalDirection, k::SVector{3}) = k * k'
projection(::HorizontalDirection, k::SVector{3}) = I - k * k'

struct Advection end
struct Diffusion end

struct TestProblem{adv, diff, dir, topo} <: AdvectionDiffusionProblem end

function init_velocity_diffusion!(::TestProblem{adv, diff, dir, topo}, aux::Vars,
                                  geom::LocalGeometry) where {adv, diff, dir, topo}
  k = vertical_unit_vector(topo, geom.coord)
  P = projection(dir, k)
  aux.u = !isnothing(adv) ? P * sin.(geom.coord) : zeros(SVector{3})
  aux.D = !isnothing(diff) ? SMatrix{3, 3}(P) / 200 : zeros(SMatrix{3, 3})
end

function initial_condition!(::AdvectionDiffusionProblem, state, aux, x, t)
  # random perturbation to trigger numerical fluxes
  state.ρ = prod(sin.(x)) + rand() / 10
end

function create_topology(::Box{dim}, mpicomm, Ne, FT) where {dim}
  brickrange = ntuple(j->range(FT(0); length=Ne+1, stop=1), dim)
  periodicity = ntuple(j->false, dim)
  bc = ntuple(j->(3, 3), dim)
  StackedBrickTopology(mpicomm, brickrange;
                       periodicity=periodicity,
                       boundary=bc)
end

function create_topology(::Sphere, mpicomm, Ne, FT)
  vert_range = grid1d(FT(1), FT(2), nelem=Ne)
  StackedCubedSphereTopology(mpicomm, Ne, vert_range, boundary=(3, 3))
end

create_dg(model, grid, direction) =
  DGModel(model,
          grid,
          Rusanov(),
          CentralNumericalFluxDiffusive(),
          CentralNumericalFluxGradient(),
          direction=direction)

function run(adv, diff, topo, mpicomm, ArrayType, FT, polynomialorder, Ne, level)
  topology = create_topology(topo(), mpicomm, Ne, FT)
  grid = DiscontinuousSpectralElementGrid(topology,
                                          FloatType = FT,
                                          DeviceArray = ArrayType,
                                          polynomialorder = polynomialorder,
                                          meshwarp = topo == Sphere ?
                                            cubedshellwarp : (x...)->identity(x)
                                         )

  problems = (p=TestProblem{adv, diff, EveryDirection(), topo()}(),
              vp=TestProblem{adv, diff, VerticalDirection(), topo()}(),
              hp=TestProblem{adv, diff, HorizontalDirection(), topo()}())

  models = map(AdvectionDiffusion{3, false}, problems)

  dgmodels = map(models) do m
    (dg=create_dg(m, grid, EveryDirection()),
     vdg=create_dg(m, grid, VerticalDirection()),
     hdg=create_dg(m, grid, HorizontalDirection()))
  end

  Q = init_ode_state(dgmodels.p.dg, FT(0), init_on_cpu=true)

  # evaluate all combinations
  dQ = map(x->map(dg->(dQ=similar(Q); dg(dQ, Q, nothing, FT(0)); dQ), x), dgmodels)

  # set up tolerances
  atolm = 6e-13
  atolv = 5e-5 / 5 ^ (level - 1)
  atolh = 0.0002 / 5 ^ (level - 1)

  @testset "total" begin
    atol = topo <: Box || diff == nothing ? atolm : atolh
    @test isapprox(norm(dQ.p.dg  .- dQ.p.vdg  .- dQ.p.hdg ), 0, atol = atol)
    @test isapprox(norm(dQ.vp.dg .- dQ.vp.vdg .- dQ.vp.hdg), 0, atol = atol)
    @test isapprox(norm(dQ.hp.dg .- dQ.hp.vdg .- dQ.hp.hdg), 0, atol = atol)
  end

  @testset "vertical" begin
    atol = topo <: Box ? atolm : atolv
    @test isapprox(norm(dQ.vp.dg .- dQ.vp.vdg), 0, atol = atol)
    @test isapprox(norm(dQ.vp.hdg), 0, atol = atol)
    @test isapprox(norm(dQ.vp.dg .- dQ.p.vdg), 0, atol = atol)
  end

  @testset "horizontal" begin
    atol = topo <: Box ? atolm : atolh
    @test isapprox(norm(dQ.hp.dg .- dQ.hp.hdg), 0, atol = atol)
    @test isapprox(norm(dQ.hp.vdg), 0, atol = atol)
    @test isapprox(norm(dQ.hp.dg - dQ.p.hdg), 0, atol = atol)
  end
end

let
  Random.seed!(44)
  CLIMA.init()
  ArrayType = CLIMA.array_type()
  mpicomm = MPI.COMM_WORLD
  FT = Float64
  numlevels = 2
  polynomialorder = 4
  base_num_elem = 4

  @testset "$(@__FILE__)" begin
    @testset for topo in (Box{2}, Box{3}, Sphere)
      @testset for (adv, diff) in
          ((Advection, nothing), (nothing, Diffusion), (Advection, Diffusion))
        @testset for level = 1:numlevels
          Ne = 2 ^ (level - 1) * base_num_elem
          run(adv, diff, topo, mpicomm, ArrayType, FT,
              polynomialorder, Ne, level)
        end
      end
    end
  end
end
nothing

max_stars_repo_name: yairchn/CLIMA

max_issues_repo_name: yairchn/CLIMA

max_forks_repo_name: yairchn/CLIMA


text:
include("Derivatives.jl")
include("kernel.jl")
include("transforms.jl")

using Distributions
using Printf
using SpecialFunctions
using Plots
using Polynomials
using DataFrames
using CSV
#using .btgDeriv

function normalize(X)
    for j = 1:size(X, 2)
        col = X[:, j]
        mini = minimum(col)
        maxi = maximum(col)
        if mini == maxi
            col = col/maxi
        else
            col = (col .- mini)/(maxi-mini)
        end
        X[:, j] = col
    end
    return X
end

function squash(X)
    maxi = maximum(X)
    return X ./ maxi
end

if true
    #s = [3.0; 4; 2]; s0 = [1; 2; 3]; X = [3 4; 9 5; 7 13]; X0 = [1 1; 2 1; -1 3]; z = [10; 11; 13]
    s = [3.0 6 1 3; 4 2 3 2; 2 1 1 5; 1 4 2 3;5 6 7 8]
    s0 = [1.0 2 3 2; 2 4 2 1; 3 1 2 6; 1 9 4 2; 2 3 8 6]
    X = [3.0 4 4; 9 3 5; 1 7 13; 4 1 2; 5 6 14]
    X0 = [1.0 1 2 ; 3 2 1; 4 -1 3; 5 5 4; 8 -3 5]
    z = [9.0; 11; 13; 6; 7]

    example = setting(s, s0, X, X0, z)
else
    df = DataFrame(CSV.File("data//abalone.csv"))
    data = convert(Matrix, df[:,2:9]) #length, diameter, height, whole weight, shucked weight, viscera weight, shell weight, age
    #target = convert(Matrix, df[:, 9:9]) #age

    #ind = 200:230
    ind = 220:230
    ind0 = 23:24
    block = vcat(data[ind, :], data[ind0, :]) #block contains training vectors and test vectors (values and labels)
    #s = [3 1; 4 1; 2 1; 5 6]; s0 = [1 1; 2 1; 3 1]; X = [3 4; 9 5; 7 13; 7 8]; X0 = [1 1; 2 1; -1 3]; z = [10; 11; 13;3]
    s = block[1:length(ind), 1:3]; s0 = block[length(ind)+1:end, 1:3]; X = block[1:length(ind), 1:1]; X0 = block[length(ind)+1:end, 1:1]; z = squash(block[1:length(ind), 8])
    example = setting(s, s0, X, X0, z)
end

##Example 1
if false
    f = x -> [x.^2; x.^3]
    df = x-> [2*x;3*x]
    x0 = [1]
    (h, A) = checkDerivative(f, df, x0)
    println(polyfit(h, A, 1))
    println("synthetic")
end

#Example 2: rbf
if false
    x = [1;2;3;4]
    f =  θ -> vec(K(x, x, θ[1], rbf));
    df = θ -> vec(K(x, x, θ[1], rbf_prime));
    θ0 = [1];
    (h, A) = checkDerivative(f, df, θ0);
    println("rbf")
    println(polyfit(h, A, 1))
end

#Example 3: Check derivatives of sub-variables of partial_theta
# Beta_hat, qtilde, Btheta, Etheta, sigmatheta, H, m, D, C, qC_inv
if false
    f = θ -> partial_theta(θ[1], 2, example)[1]
    df = θ -> partial_theta(θ[1], 2, example)[2]
    θ0 = [50]
    (h, A) = checkDerivative(f, df, θ0);
    println("partial_theta function return value")
    println(polyfit(h, A, 1))
end

#Example 4: Check derivative of partial_theta
if true
    #f = θ -> partial_theta(θ[1], 2, example)[1](z0)
    #df = θ -> partial_theta(θ[1], 2, example)[2](z0)
    example2 = getExample(1, 25, 1, 1, 10)
    θ0 = [1.1]
    z0 = [6];
    f = θ -> partial_theta(θ[1], 2.0, example2, boxCoxObj, nothing, "Turan")[2][2](z0)
    df = θ -> partial_theta(θ[1], 2.0, example2, boxCoxObj, nothing, "Turan")[2][3](z0)

    (h, A) = checkDerivative(f, df, θ0, 5, 10, 10)
    plt1 = plot(h, A, title = "Finite Difference Derivative Checker", xlabel = "log of h", ylabel = "log of error",fontfamily=font(48, "Courier") , reuse = false)
    #plot(polyfit(h, A, 1), reuse = true)
    println("partial theta of p(z0|theta, lambda, z)")
    println(polyfit(h, A, 1))
    nums = collect(.1:.1:20)
    g = x -> f(x)[1]
    plt2 = plot(nums, g.(nums),xlabel = "theta", ylabel = "p(z0|theta, lambda, z)", fontfamily=font(48, "Courier") ,title = "theta vs p(z0| theta, lambda, z)")
    display(plot(plt1, plt2, fontfamily=font(48, "CoSurier")))
    gui()
end

#Example 5: Check derivative of Box Cox
if false
    z = 2
    f =  λ -> [boxCox(z, λ[1])]
    df = λ -> [boxCoxPrime_lambda(z, λ[1])]
    λ = [2.2323]
    (h, A) = checkDerivative(f, df, λ)
    println("Box Cox partial lambda")
    f = θ -> partial_theta(θ[1], 2, example)[1]
    df = θ -> partial_theta(θ[1], 2, example)[2]
    θ0 = [.5]
    (h, A) = checkDerivative(f, df, θ0, 4, 15, 10)
    plt1 = plot(h, A, title = "Finite Difference Derivative Checker", xlabel = "log of h", ylabel = "log of error",fontfamily=font(48, "Courier") , reuse = false)
    #plot(polyfit(h, A, 1), reuse = true)
    println("partial theta of p(z0|theta, lambda, z)")
    println(polyfit(h, A, 1))
    println(polyfit(h, A, 1))
end

#Example 5: Check derivative of Jacobian of Box Cox w.r.t lambda
if false
    θ = 1.312
    λ0 = [2.123]
    z0 = rand(size(s0, 1))
    f = λ -> [partial_lambda(θ, λ[1], example)[1](z0)]
    df = λ -> [partial_lambda(θ, λ[1], example)[2](z0)]
    (h, A) = checkDerivative(f, df, λ0)
    println("Jacobian w.r.t lambda")
    println(polyfit(h, A, 1))
end

#Example 6: Check derivative of sub-variables of partial_lambda
if false
    θ = 1.312
    λ0 = [2.123]
    f = λ -> [partial_lambda(θ, λ[1], example)[1]]
    df = λ -> [partial_lambda(θ, λ[1], example)[2]]
    (h, A) = checkDerivative(f, df, λ0)
    println("partial lambda")
    println(polyfit(h, A, 1))
end

#Example 7: Check derivative of sub-functions of partial_lambda
if false
    θ = 1
    λ0 = [.5]
    #z0 = [10;11;13]
    z0 = rand(size(s0, 1))

    f = λ -> [partial_lambda(θ, λ[1], example)[1](z0)]
    df = λ -> partial_lambda(θ, λ[1], example)[2](z0)
    (h, A) = checkDerivative(f, df, λ0)
    println("partial lambda")
    println(polyfit(h, A, 1))
end

#Example 7: Check derivative of sub-variables of posterior_theta. No "sub-functions", because we don't
#have a z0 to deal with
if false
    println("loading some simple functions")
    @time begin
    θ0 = [2.0]
    λ = 3.0;
    pθ = θ -> 1/sqrt(2*pi)*exp.(-(θ .- 1) .^2/2)
    dpθ = θ -> (1 .- θ)/sqrt(2*pi)*exp.(-(1 .- θ) .^2/2)
    dpθ2 = θ -> -1/(sqrt(2*pi))*exp.(-(θ .-1).^2/2) + (1 .-θ)^2 * 1/sqrt(2*pi) * exp.(-(1 .- θ)^2/2)
    pλ = λ -> 1
    end
    println("defining posterior theta and derivative")
    @time begin
    f = θ  -> posterior_theta(θ[1], λ, pθ, dpθ, dpθ2, pλ, example2)[1]
    df = θ -> posterior_theta(θ[1], λ, pθ, dpθ, dpθ2, pλ, example2)[2]
    end
    (rr, tt) = checkDerivative(θ -> [dpθ(θ[1])], θ -> dpθ2(θ[1]), [.5], 5, 12)
    println("check derivative of prior")
    println(polyfit(rr, tt, 1))
    println("Check derivative by eval f and df at 10 points")
    @time begin
    (h, A) = checkDerivative(f, df, θ0, 5, 11)
    end
    println("partial theta of p(theta, lambda|z)")
    pfit = polyfit(h, A, 1)
    println(pfit)
    println("Plotting")
    @time begin
    plt1 = plot(h, A, title = "Finite Difference Derivative Checker", xlabel = "log of h", ylabel = "log of error",
    label = [string(pfit)], fontfamily=font(48, "Courier") , reuse = false)
    end
    nums = collect(.2:.2:2)
    g = x -> f(x)[1]
    println("Evaluating f at about 25 points")
    @time begin
    gnums = g.(nums)
    end
    plt2 = plot(nums, gnums,xlabel = "theta", ylabel = "p(theta, lambda| z)", fontfamily=font(48, "Courier") ,title = "theta vs p(theta, lambda| z)")
    display(plot(plt1, plt2, fontfamily=font(48, "Courier") ))
    #display(plot(plt1))
    gui()

end

#Example 8: Check derivative of posterior_lambda
if false
    θ = 1;
    pλ = λ -> 1/sqrt(2*pi)*exp.(-(λ .- 1) .^2/2)
    dpλ = θ -> (1 .- θ)/sqrt(2*pi)*exp.(-(1 .- θ) .^2/2)
    pθ = θ -> 1;
    f = λ -> [posterior_lambda(θ, λ[1], pθ, pλ, dpλ, example)[1]]
    df = λ -> posterior_lambda(θ, λ[1], pθ, pλ, dpλ, example)[2]
    λ0 = [0.6]
    (h, A) = checkDerivative(f, df, λ0)
    println("partial lambda of p(theta, lambda|z)")
    println(polyfit(h, A, 1))
end

#Example 10: Check various derivatives in partial_z0 function
if false
    g = (x, λ) -> -exp.(λ .*x)
    dg = (x, λ) -> -λ .*exp.(λ .*x)
    dg2 = (x, λ) -> -λ^2 .*exp.(λ .*x)
   (jac, djac) =  partial_z0(2, 1, example, g, dg, dg2)
   (h, A) = checkDerivative(jac, djac, [1;20])
   println("jac and djac w.r.t z0")
   println(polyfit(h, A, 1))
end

#Example 11: Check derivative of rbf second derivative
if false
    x0 = [3, 4]
    f = theta -> rbf_prime(x0[1], x0[2], theta)
    df = theta -> rbf_prime2(x0[1], x0[2], theta)
    (h, A) = checkDerivative(f, df, 3.0)
    println("first and second derivatives of rbf")
    println(polyfit(h, A, 1))
end

#Example 12: Check that derivative of trace is trace of derivative
if false
    s1 = [1;2;3]
    s2 = [5;6;7]
    f = θ -> tr(K(s1, s2, θ, rbf))
    df = θ -> tr(K(s1, s2, θ, rbf_prime))

    (h, A) = checkDerivative(f, df, .5)
    println("derivative of trace")
    println(polyfit(h, A, 1))
end

max_stars_repo_name: dbindel/btg

max_issues_repo_name: dbindel/btg

max_forks_repo_name: dbindel/btg


text:
@doc raw"""
```
Photometric(illumination_direction1::Vector{T} where T <: Real = [Inf, Inf, Inf], illumination_direction2::Vector{T} where T <: Real = [Inf, Inf, Inf], illumination_direction3::Vector{T} where T <: Real = [Inf, Inf, Inf], integration_scheme::AbstractIntegrator = Horn())
```
Creates a `Photometric` functor which attempts to produce the gradients and
integrate them to retrieve a height map using three-point photometric stereo.

This algorithm solves a system of linear equation formed by varying the lighting
conditions on an object while holding the viewing angle constant. Note that
the ilumination direction id defined negative of most other algorithms and as
such inouts may need to be corrected (see example).
# Output
Returns an M by N array (matching dimensions of original image) of Float `Z`
that represents the reconstructed height at the point and the gradients in
M by N arrays of Float `p` and `q`.
# Details
Let ñ₁, ñ₂ and ñ₃ define the ilumination direction let N be a 3 X 3 matrix
whos rows are formed by the row vectors of the illumination direction as per:
```math
N=\begin{bmatrix}&n_{11}&n_{12}&n_{13}\\&n_{21}&n_{22}&n_{23}\\&n_{31}&n_{32}&n_{33}\end{bmatrix}
```
and I₁, I₂ and I₃ be the intensity values at a point (x,y) and let Ĩ=[I₁,I₂,I₃]′
be the column vector formed by these values.

The surface normal at the point (x,y) can now be directly calculated using:
```math
ñ=N^{-1}Ĩ
```
where n is the surface normal.
# Arguments
The function arguments are described in more detail below.
## `illumination_direction1`
A `Vector{T} where T <: Real` that specifies the first illumination vector to be used by the
algorithm. The `illumination_direction` should be a vector of the form [x,y,z]
where x,y,z are int he range [0,1]. Defults to `[Inf,Inf,Inf]` which triggers
`estimate_img_properties` to be called.
## `illumination_direction2`
A `Vector{T} where T <: Real` that specifies the second illumination vector to be used by the
algorithm. The `illumination_direction` should be a vector of the form [x,y,z]
where x,y,z are int he range [0,1]. Defults to `[Inf,Inf,Inf]` which triggers
`estimate_img_properties` to be called.
## `illumination_direction3`
A `Vector{T} where T <: Real` that specifies the third illumination vector to be used by the
algorithm. The `illumination_direction` should be a vector of the form [x,y,z]
where x,y,z are int he range [0,1]. Defults to `[Inf,Inf,Inf]` which triggers
`estimate_img_properties` to be called.
## `integration_scheme`
A keywork argument of type `AbstractIntegrator` which specifies which integration
scheme is to be used to reconstruct the surface from the calculated normals.
Defutls to Horn().
!!! note
    If any `illumination_direction` are not defined (i.e. have defulted to `Inf`) they
    will be calculated at runtime using `estimate_img_properties`.
```julia
using Images, Makie, ShapeFromShading

#generate synthetic images
n₁ = [0,0,1.0]
n₂ = [1.0,0.2,1.0]
n₃ = [0.2,0.9,1.0]
img1, img2, img3 = generate_photometric(SynthSphere(50), 1.0, n₁, n₂, n₃)

#calculate the heightmap
photometric = Photometric(-n₁, -n₂, -n₃, Horn())
Z,p,q = photometric(img1, img2, img3)

#normalize to maximum of 1 (not necessary but makes displaying easier)
Z = Z./maximum(Z)

#display using Makie (Note: Makie can often take several minutes first time)
r = 0.0:0.1:2
surface(r, r, Z)
```
# Reference
1. R. Woodham, "Photometric Method For Determining Surface Orientation From Multiple Images", Optical Engineering, vol. 19, no. 1, 1980. [doi:10.1117/12.7972479](https://doi.org/10.1117/12.7972479)
"""
function (algorithm::Photometric)(img1::AbstractArray, img2::AbstractArray, img3::AbstractArray)

    illumination_direction1 = algorithm.illumination_direction1
    illumination_direction2 = algorithm.illumination_direction2
    illumination_direction3 = algorithm.illumination_direction3
    integration_scheme = algorithm.integration_scheme

    if illumination_direction1 == [Inf, Inf, Inf] || illumination_direction2 == [Inf, Inf, Inf] || illumination_direction3 == [Inf, Inf, Inf]
        ρ, illumination_direction1, σ, τ = estimate_img_properties(img1)
        ρ, illumination_direction2, σ, τ = estimate_img_properties(img2)
        ρ, illumination_direction3, σ, τ = estimate_img_properties(img3)
    end

    #setup illumination matrix
    N = zeros(Float64, 3, 3)
    for i = 1:3
        N[1,i] = illumination_direction1[i]
        N[2,i] = illumination_direction2[i]
        N[3,i] = illumination_direction3[i]
    end
    Ninv = inv(N)

    #setup image and gradient arrays
    I₁ = Float64.(img1)
    I₂ = Float64.(img2)
    I₃ = Float64.(img3)
    p = zeros(Float64, size(I₁))
    q = zeros(Float64, size(I₁))

    #calculate gradients
    for i in CartesianIndices(I₁)
        I = [I₁[i], I₂[i], I₃[i]]
        n = Ninv * I
        if n[3] == 0.0
            n = [0.0,0.0,-1.0]
        end
        p[i] = -n[1] / n[3]
        q[i] = -n[2] / n[3]
    end

    #reconstruct surface
    Z = integration_scheme(p, q)
    return Z, p, q
end


text:
module EulerGravity
  export EulerGravityLaw

  import ..Atum
  using ..Atum: avg, logavg, roe_avg, constants
  using StaticArrays
  using LinearAlgebra: I, norm

  struct EulerGravityLaw{FT, D, S, C} <: Atum.AbstractBalanceLaw{FT, D, S, C}
    function EulerGravityLaw{FT, D}(; γ = 7 // 5,
                                      grav = 981 // 100,
                                      pde_level_balance = false) where {FT, D}
      S = 2 + D
      C = (γ = FT(γ), grav = FT(grav), pde_level_balance = pde_level_balance)
      new{FT, D, S, C}()
    end
  end

  function varsindices(law::EulerGravityLaw)
    S = Atum.numberofstates(law)
    ix_ρ = 1
    ix_ρu⃗ = StaticArrays.SUnitRange(2, S - 1)
    ix_ρe = S
    return ix_ρ, ix_ρu⃗, ix_ρe
  end

  function unpackstate(law::EulerGravityLaw, q)
    ix_ρ, ix_ρu⃗, ix_ρe = varsindices(law)
    @inbounds q[ix_ρ], q[ix_ρu⃗], q[ix_ρe]
  end

  referencestate(law::EulerGravityLaw, x⃗) = SVector{0, eltype(law)}()
  reference_ρ(law::EulerGravityLaw, aux) = @inbounds aux[ndims(law) + 1]
  reference_p(law::EulerGravityLaw, aux) = @inbounds aux[ndims(law) + 2]

  function Atum.auxiliary(law::EulerGravityLaw, x⃗)
    vcat(x⃗, referencestate(law, x⃗))
  end

  function coordinates(law::EulerGravityLaw, aux)
    aux[SOneTo(ndims(law))]
  end

  function geopotential(law, aux)
    x⃗ = coordinates(law, aux)
    z = last(x⃗)
    constants(law).grav * z
  end

  function pressure(law::EulerGravityLaw, ρ, ρu⃗, ρe, Φ)
    γ = constants(law).γ
    (γ - 1) * (ρe - ρu⃗' * ρu⃗ / 2ρ - ρ * Φ)
  end
  function energy(law::EulerGravityLaw, ρ, ρu⃗, p, Φ)
    γ = constants(law).γ
    p / (γ - 1) + ρu⃗' * ρu⃗ / 2ρ + ρ * Φ
  end
  function soundspeed(law::EulerGravityLaw, ρ, p)
    γ = constants(law).γ
    sqrt(γ * p / ρ)
  end
  function soundspeed(law::EulerGravityLaw, ρ, ρu⃗, ρe, Φ)
    soundspeed(law, ρ, pressure(law, ρ, ρu⃗, ρe, Φ))
  end

  function Atum.flux(law::EulerGravityLaw, q, aux)
    ρ, ρu⃗, ρe = unpackstate(law, q)

    Φ = geopotential(law, aux)

    u⃗ = ρu⃗ / ρ
    p = pressure(law, ρ, ρu⃗, ρe, Φ)

    δp = p
    if constants(law).pde_level_balance
      δp -= reference_p(law, aux)
    end

    fρ = ρu⃗
    fρu⃗ = ρu⃗ * u⃗' + δp * I
    fρe = u⃗ * (ρe + p)

    hcat(fρ, fρu⃗, fρe)
  end

  function Atum.nonconservative_term!(law::EulerGravityLaw, dq, q, aux)
    ix_ρ, ix_ρu⃗, _ = varsindices(law)

    @inbounds ρ = q[ix_ρ]

    if constants(law).pde_level_balance
      ρ -= reference_ρ(law, aux)
    end

    @inbounds dq[ix_ρu⃗[end]] -= ρ * constants(law).grav
  end

  function Atum.wavespeed(law::EulerGravityLaw, n⃗, q, aux)
    ρ, ρu⃗, ρe = unpackstate(law, q)

    Φ = geopotential(law, aux)

    u⃗ = ρu⃗ / ρ
    abs(n⃗' * u⃗) + soundspeed(law, ρ, ρu⃗, ρe, Φ)
  end

  function Atum.entropy(law::EulerGravityLaw, q, aux)
    ρ, ρu⃗, ρe = unpackstate(law, q)
    Φ = geopotential(law, aux)
    γ = constants(law).γ
    p = pressure(law, ρ, ρu⃗, ρe, Φ)
    s = log(p / ρ ^ γ)

    -ρ * s / (γ - 1)
  end

  function Atum.entropyvariables(law::EulerGravityLaw, q, aux)
    ρ, ρu⃗, ρe = unpackstate(law, q)
    Φ = geopotential(law, aux)
    γ = constants(law).γ
    p = pressure(law, ρ, ρu⃗, ρe, Φ)
    s = log(p / ρ ^ γ)
    b = ρ / 2p
    u⃗ = ρu⃗ / ρ
    vρ = (γ - s) / (γ - 1) - (u⃗' * u⃗  - 2Φ) * b
    vρu⃗ = 2b * u⃗
    vρe = -2b

    SVector(vρ, vρu⃗..., vρe)
  end

  function Atum.surfaceflux(::Atum.RoeFlux, law::EulerGravityLaw, n⃗, q⁻, aux⁻, q⁺, aux⁺)
    Φ = geopotential(law, aux⁻)

    f⁻ = Atum.flux(law, q⁻, aux⁻)
    f⁺ = Atum.flux(law, q⁺, aux⁺)

    ρ⁻, ρu⃗⁻, ρe⁻ = unpackstate(law, q⁻)
    u⃗⁻ = ρu⃗⁻ / ρ⁻
    e⁻ = ρe⁻ / ρ⁻
    p⁻ = pressure(law, ρ⁻, ρu⃗⁻, ρe⁻, Φ)
    h⁻ = e⁻ + p⁻ / ρ⁻
    c⁻ = soundspeed(law, ρ⁻, p⁻)

    ρ⁺, ρu⃗⁺, ρe⁺ = unpackstate(law, q⁺)
    u⃗⁺ = ρu⃗⁺ / ρ⁺
    e⁺ = ρe⁺ / ρ⁺
    p⁺ = pressure(law, ρ⁺, ρu⃗⁺, ρe⁺, Φ)
    h⁺ = e⁺ + p⁺ / ρ⁺
    c⁺ = soundspeed(law, ρ⁺, p⁺)

    ρ = sqrt(ρ⁻ * ρ⁺)
    u⃗ = roe_avg(ρ⁻, ρ⁺, u⃗⁻, u⃗⁺)
    h = roe_avg(ρ⁻, ρ⁺, h⁻, h⁺)
    c = roe_avg(ρ⁻, ρ⁺, c⁻, c⁺)

    uₙ = u⃗' * n⃗

    Δρ = ρ⁺ - ρ⁻
    Δp = p⁺ - p⁻
    Δu⃗ = u⃗⁺ - u⃗⁻
    Δuₙ = Δu⃗' * n⃗

    c⁻² = 1 / c^2
    w1 = abs(uₙ - c) * (Δp - ρ * c * Δuₙ) * c⁻² / 2
    w2 = abs(uₙ + c) * (Δp + ρ * c * Δuₙ) * c⁻² / 2
    w3 = abs(uₙ) * (Δρ - Δp * c⁻²)
    w4 = abs(uₙ) * ρ

    fp_ρ = (w1 + w2 + w3) / 2
    fp_ρu⃗ = (w1 * (u⃗ - c * n⃗) +
             w2 * (u⃗ + c * n⃗) +
             w3 * u⃗ +
             w4 * (Δu⃗ - Δuₙ * n⃗)) / 2
    fp_ρe = (w1 * (h - c * uₙ) +
             w2 * (h + c * uₙ) +
             w3 * (u⃗' * u⃗ / 2 + Φ) +
             w4 * (u⃗' * Δu⃗ - uₙ * Δuₙ)) / 2

    (f⁻ + f⁺)' * n⃗ / 2 - SVector(fp_ρ, fp_ρu⃗..., fp_ρe)
  end

  function Atum.twopointflux(::Atum.EntropyConservativeFlux,
                             law::EulerGravityLaw,
                             q₁, aux₁, q₂, aux₂)
      FT = eltype(law)
      γ = constants(law).γ
      ρ₁, ρu⃗₁, ρe₁ = unpackstate(law, q₁)
      ρ₂, ρu⃗₂, ρe₂ = unpackstate(law, q₂)

      Φ₁ = geopotential(law, aux₁)
      u⃗₁ = ρu⃗₁ / ρ₁
      p₁ = pressure(law, ρ₁, ρu⃗₁, ρe₁, Φ₁)
      b₁ = ρ₁ / 2p₁

      Φ₂ = geopotential(law, aux₂)
      u⃗₂ = ρu⃗₂ / ρ₂
      p₂ = pressure(law, ρ₂, ρu⃗₂, ρe₂, Φ₂)
      b₂ = ρ₂ / 2p₂

      ρ_avg = avg(ρ₁, ρ₂)
      u⃗_avg = avg(u⃗₁, u⃗₂)
      b_avg = avg(b₁, b₂)
      Φ_avg = avg(Φ₁, Φ₂)

      u²_avg = avg(u⃗₁' * u⃗₁, u⃗₂' * u⃗₂)
      ρ_log = logavg(ρ₁, ρ₂)
      b_log = logavg(b₁, b₂)

      fρ = u⃗_avg * ρ_log
      fρu⃗ = u⃗_avg * fρ' + ρ_avg / 2b_avg * I
      fρe = (1 / (2 * (γ - 1) * b_log) - u²_avg / 2 + Φ_avg) * fρ + fρu⃗ * u⃗_avg

      # fluctuation
      α = b_avg * ρ_log / 2b₁
      fρu⃗ -= α * (Φ₁ - Φ₂) * I

      hcat(fρ, fρu⃗, fρe)
  end

  function Atum.twopointflux(::Atum.KennedyGruberFlux,
                             law::EulerGravityLaw,
                             q₁, aux₁, q₂, aux₂)
      FT = eltype(law)
      ρ₁, ρu⃗₁, ρe₁ = unpackstate(law, q₁)
      ρ₂, ρu⃗₂, ρe₂ = unpackstate(law, q₂)

      Φ₁ = geopotential(law, aux₁)
      u⃗₁ = ρu⃗₁ / ρ₁
      e₁ = ρe₁ / ρ₁
      p₁ = pressure(law, ρ₁, ρu⃗₁, ρe₁, Φ₁)

      Φ₂ = geopotential(law, aux₂)
      u⃗₂ = ρu⃗₂ / ρ₂
      e₂ = ρe₂ / ρ₂
      p₂ = pressure(law, ρ₂, ρu⃗₂, ρe₂, Φ₂)

      ρ_avg = avg(ρ₁, ρ₂)
      u⃗_avg = avg(u⃗₁, u⃗₂)
      e_avg = avg(e₁, e₂)
      p_avg = avg(p₁, p₂)

      fρ = u⃗_avg * ρ_avg
      fρu⃗ = u⃗_avg * fρ' + p_avg * I
      fρe = u⃗_avg * (ρ_avg * e_avg + p_avg)

      # fluctuation
      α = ρ_avg / 2
      fρu⃗ -= α * (Φ₁ - Φ₂) * I

      hcat(fρ, fρu⃗, fρe)
  end

  function Atum.surfaceflux(::Atum.MatrixFlux, law::EulerGravityLaw, n⃗, q⁻, aux⁻, q⁺, aux⁺)
    FT = eltype(law)
    γ = constants(law).γ
    ecflux = Atum.surfaceflux(Atum.EntropyConservativeFlux(), law, n⃗, q⁻, aux⁻, q⁺, aux⁺)

    ρ⁻, ρu⃗⁻, ρe⁻ = unpackstate(law, q⁻)
    Φ⁻ = geopotential(law, aux⁻)
    u⃗⁻ = ρu⃗⁻ / ρ⁻
    e⁻ = ρe⁻ / ρ⁻
    p⁻ = pressure(law, ρ⁻, ρu⃗⁻, ρe⁻, Φ⁻)
    b⁻ = ρ⁻ / 2p⁻

    ρ⁺, ρu⃗⁺, ρe⁺ = unpackstate(law, q⁺)
    Φ⁺ = geopotential(law, aux⁺)
    u⃗⁺ = ρu⃗⁺ / ρ⁺
    e⁺ = ρe⁺ / ρ⁺
    p⁺ = pressure(law, ρ⁺, ρu⃗⁺, ρe⁺, Φ⁺)
    b⁺ = ρ⁺ / 2p⁺

    Φ_avg = avg(Φ⁻, Φ⁺)
    ρ_log = logavg(ρ⁻, ρ⁺)
    b_log = logavg(b⁻, b⁺)
    u⃗_avg = avg(u⃗⁻, u⃗⁺)
    p_avg = avg(ρ⁻, ρ⁺) / 2avg(b⁻, b⁺)
    u²_bar = 2 * norm(u⃗_avg) - avg(norm(u⃗⁻), norm(u⃗⁺))
    h_bar = γ / (2 * b_log * (γ - 1)) + u²_bar / 2 + Φ_avg
    c_bar = sqrt(γ * p_avg / ρ_log)

    u⃗mc = u⃗_avg - c_bar * n⃗
    u⃗pc = u⃗_avg + c_bar * n⃗
    u_avgᵀn = u⃗_avg' * n⃗

    v⁻ = Atum.entropyvariables(law, q⁻, aux⁻)
    v⁺ = Atum.entropyvariables(law, q⁺, aux⁺)
    Δv = v⁺ - v⁻

    λ1 = abs(u_avgᵀn - c_bar) * ρ_log / 2γ
    λ2 = abs(u_avgᵀn) * ρ_log * (γ - 1) / γ
    λ3 = abs(u_avgᵀn + c_bar) * ρ_log / 2γ
    λ4 = abs(u_avgᵀn) * p_avg

    Δv_ρ, Δv_ρu⃗, Δv_ρe = unpackstate(law, Δv)
    u⃗ₜ = u⃗_avg - u_avgᵀn * n⃗

    w1 = λ1 * (Δv_ρ + u⃗mc' * Δv_ρu⃗ + (h_bar - c_bar * u_avgᵀn) * Δv_ρe)
    w2 = λ2 * (Δv_ρ + u⃗_avg' * Δv_ρu⃗ + (u²_bar / 2 + Φ_avg) * Δv_ρe)
    w3 = λ3 * (Δv_ρ + u⃗pc' * Δv_ρu⃗ + (h_bar + c_bar * u_avgᵀn) * Δv_ρe)

    Dρ = w1 + w2 + w3

    Dρu⃗ = (w1 * u⃗mc +
           w2 * u⃗_avg +
           w3 * u⃗pc +
           λ4 * (Δv_ρu⃗ - n⃗' * (Δv_ρu⃗) * n⃗ + Δv_ρe * u⃗ₜ))

    Dρe = (w1 * (h_bar - c_bar * u_avgᵀn) +
           w2 * (u²_bar / 2 + Φ_avg) +
           w3 * (h_bar + c_bar * u_avgᵀn) +
           λ4 * (u⃗ₜ' * Δv_ρu⃗ + Δv_ρe * (u⃗_avg' * u⃗_avg - u_avgᵀn ^ 2)))

    ecflux - SVector(Dρ, Dρu⃗..., Dρe) / 2
  end
end