text:
"Any radial basis function."
abstract type RBF{T}
end

"A smooth radial basis function with a shape parameter `ε`."
abstract type SmoothRBF{T} <: RBF{T}
end

shape_parameter(rbf::SmoothRBF) = rbf.ε

"A piecewise smooth radial basis function."
abstract type PiecewiseSmoothRBF{T} <: RBF{T}
end

degree(rbf::PiecewiseSmoothRBF) = rbf.p

struct Multiquadric{T} <: SmoothRBF{T}
    ε   :: T
end

const MQ{T} = Multiquadric{T}

(rbf::Multiquadric)(r) = sqrt(1 + (rbf.ε*r)^2)

eval_derivative1(rbf::Multiquadric, r) = 1/rbf(r)*rbf.ε*r

similar(r::Multiquadric{T}, ε) where {T} = Multiquadric{T}(ε)

struct InverseMultiquadric{T} <: SmoothRBF{T}
    ε   :: T
end

const IMQ{T} = InverseMultiquadric{T}

(rbf::InverseMultiquadric)(r) = 1/sqrt(1 + (rbf.ε*r)^2)

eval_derivative1(rbf::InverseMultiquadric{T}, r) where {T} = (-one(T)/2)*rbf(r)^3*2*rbf.ε*r

similar(r::InverseMultiquadric{T}, ε) where {T} = InverseMultiquadric{T}(ε)

struct InverseQuadratic{T} <: SmoothRBF{T}
    ε   :: T
end

const IQ{T} = InverseQuadratic{T}

(rbf::InverseQuadratic)(r) = 1/(1 + (rbf.ε*r)^2)

similar(r::InverseQuadratic{T}, ε) where {T} = InverseQuadratic{T}(ε)

struct Gaussian{T} <: SmoothRBF{T}
    ε   :: T
end

const GA{T} = Gaussian{T}

(rbf::Gaussian)(r) = exp(-(rbf.ε*r)^2)

eval_derivative1(rbf::Gaussian, r) = -2*rbf.ε*r * rbf(r)

eval_derivative2(rbf::Gaussian, r) = -2*rbf.ε * rbf(r) - 2*rbf.ε*r*eval_derivative1(rbf,r)

similar(r::Gaussian{T}, ε) where {T} = Gaussian{T}(ε)

approximate_support(rbf::Gaussian{T}, threshold = eps(T)) where {T} = sqrt(-log(threshold))/rbf.ε

struct PolyharmonicSpline{T} <: PiecewiseSmoothRBF{T}
    p   ::  Int
end

const SPH{T} = PolyharmonicSpline{T}

PolyharmonicSpline(p::Int) = PolyharmonicSpline{Float64}(p)

(rbf::PolyharmonicSpline)(r) = r^(2*rbf.p+1)

similar(r::PolyharmonicSpline{T}, p) where {T} = PolyharmonicSpline{T}(p)

struct ThinPlateSpline{T} <: PiecewiseSmoothRBF{T}
    p   ::  Int
end

const TPS{T} = ThinPlateSpline{T}

ThinPlateSpline(p::Int) = ThinPlateSpline{Float64}(p)

(rbf::ThinPlateSpline)(r) = r^(2*rbf.p) * log(r)

similar(r::ThinPlateSpline{T}, p) where {T} = ThinPlateSpline{T}(p)


text:
typealias BigFloats Union{BigFloat,Complex{BigFloat}}

function Base.fft{F<:Fun}(x::Vector{F})
    n,T = length(x),mapreduce(eltype,promote_type,x)
    if ispow2(n) return fft_pow2(x) end
    ks = linspace(zero(real(T)),n-one(real(T)),n)
    Wks = exp(-im*convert(T,π)*ks.^2/n)
    xq,wq = x.*Wks,conj([exp(-im*convert(T,π)*n);reverse(Wks);Wks[2:end]])
    return Wks.*conv(xq,wq)[n+1:2n]
end

Base.ifft{F<:Fun}(x::Vector{F}) = conj(fft(conj(x)))/length(x)
function Base.ifft!{F<:Fun}(x::Vector{F})
    y = conj(fft(conj(x)))/length(x)
    x[:] = y
    return x
end

function Base.conv{F<:Fun}(u::StridedVector{F}, v::StridedVector)
    nu,nv = length(u),length(v)
    n = nu + nv - 1
    np2 = nextpow2(n)
    pad!(u,np2),pad!(v,np2)
    y = ifft_pow2(fft_pow2(u).*fft_pow2(v))
    #TODO This would not handle Dual/ComplexDual numbers correctly
    T = promote_type(mapreduce(eltype,promote_type,u),mapreduce(eltype,promote_type,v))
    y = T<:Real ? real(y[1:n]) : y[1:n]
end

######################################################################
# TO BE DEPRECATED FOR V0.4 UPGRADE
######################################################################

# plan_fft for BigFloats (covers Laurent svfft)

Base.plan_fft{F<:Fun}(x::Vector{F}) = fft
Base.plan_ifft{F<:Fun}(x::Vector{F}) = ifft
Base.plan_ifft!{F<:Fun}(x::Vector{F}) = ifft

# Chebyshev transforms and plans for BigFloats

plan_chebyshevtransform{T<:BigFloats}(x::Vector{T};kwds...) = identity
plan_ichebyshevtransform{T<:BigFloats}(x::Vector{T};kwds...) = identity

#following Chebfun's @Chebtech1/vals2coeffs.m and @Chebtech2/vals2coeffs.m
function chebyshevtransform{T<:BigFloats}(x::Vector{T},plan;kind::Integer=1)
    if kind == 1
        n = length(x)
        if n == 1
            x
        else
            w = [2exp(im*convert(T,π)*k/2n) for k=0:n-1]
            ret = w.*ifft([reverse(x);x])[1:n]
            ret = T<:Real ? real(ret) : ret
            ret[1] /= 2
            ret
        end
    elseif kind == 2
        n = length(x)
        if n == 1
            x
        else
            ret = ifft([reverse(x);x[2:end-1]])[1:n]
            ret = T<:Real ? real(ret) : ret
            ret[2:n-1] *= 2
            ret
        end
    end
end

#following Chebfun's @Chebtech1/vals2coeffs.m and @Chebtech2/vals2coeffs.m
function ichebyshevtransform{T<:BigFloats}(x::Vector{T},plan;kind::Integer=1)
    if kind == 1
        n = length(x)
        if n == 1
            x
        else
            w = [exp(-im*convert(T,π)*k/2n)/2 for k=0:2n-1]
            w[1] *= 2;w[n+1] *= 0;w[n+2:end] *= -1
            ret = fft(w.*[x;one(T);x[end:-1:2]])[n:-1:1]
            ret = T<:Real ? real(ret) : ret
        end
    elseif kind == 2
        n = length(x)
        if n == 1
            x
        else
            ##TODO: make thread safe
            x[1] *= 2;x[end] *= 2
            ret = chebyshevtransform(x;kind=kind)
            x[1] /=2;x[end] /=2
            ret[1] *= 2;ret[end] *= 2
            negateeven!(ret)
            ret *= .5*(n-1)
            reverse!(ret)
        end
    end
end

# Fourier space plans for BigFloat

function plan_transform{T<:BigFloat,D}(::Fourier{D},x::Vector{T})
    function plan(x)
        v = fft(x)
        n = div(length(x),2)+1
        [real(v[1:n]);imag(v[n-1:-1:2])]
    end
    plan
end

function plan_itransform{T<:BigFloat,D}(::Fourier{D},x::Vector{T})
    function plan(x)
        n = div(length(x),2)+1
        v = complex([x[1:n];x[n-1:-1:2]],[0;-x[2n-2:-1:n+1];0;x[n+1:2n-2]])
        real(fft(v))
    end
    plan
end

# SinSpace plans for BigFloat

function plan_transform{T<:BigFloat}(::SinSpace,x::Vector{T})
    function plan(x)
        imag(fft([0;-x;0;reverse(x)]))[2:length(x)+1]
    end
    plan
end

function plan_itransform{T<:BigFloat}(::SinSpace,x::Vector{T})
    function plan(x)
        imag(fft([0;-x;0;reverse(x)]))[2:length(x)+1]
    end
    plan
end

# Fourier space & SinSpace plans for Complex{BigFloat}

for SP in (:Fourier,:SinSpace), pl in (:plan_transform,:plan_itransform)
    @eval begin
        function $pl{T<:Complex{BigFloat},D}(::$SP{D},x::Vector{T})
            function plan(x)
                complex($pl($SP(),real(x))(real(x)),$pl($SP(),imag(x))(imag(x)))
            end
            plan
        end
    end
end

ext: jl


text:
import JLD
using Distributed

using NetCDF

"A type for writing checkpoints."
struct Checkpointer <: OutputWriter
    dir::AbstractString
    filename_prefix::AbstractString
    output_frequency::Int
    padding::Int
end

"A type for writing NetCDF output."
mutable struct NetCDFOutputWriter <: OutputWriter
    dir::AbstractString
    filename_prefix::AbstractString
    output_frequency::Int
    padding::Int
    naming_scheme::Symbol
    compression::Int
    async::Bool
end

"A type for writing Binary output."
mutable struct BinaryOutputWriter <: OutputWriter
    dir::AbstractString
    filename_prefix::AbstractString
    output_frequency::Int
    padding::Int
end

function Checkpointer(; dir=".", prefix="", frequency=1, padding=9)
    Checkpointer(dir, prefix, frequency, padding)
end

function NetCDFOutputWriter(; dir=".", prefix="", frequency=1, padding=9,
                              naming_scheme=:iteration, compression=3, async=false)
    NetCDFOutputWriter(dir, prefix, frequency, padding, naming_scheme, compression, async)
end

"Return the filename extension for the `OutputWriter` filetype."
ext(fw::OutputWriter) = throw("Not implemented.")
ext(fw::NetCDFOutputWriter) = ".nc"
ext(fw::Checkpointer) = ".jld"

filename(fw::Checkpointer, iteration) = fw.filename_prefix * "model_checkpoint_" * lpad(iteration, fw.padding, "0") * ext(fw)

function filename(fw, name, iteration)
    if fw.naming_scheme == :iteration
        fw.filename_prefix * name * lpad(iteration, fw.padding, "0") * ext(fw)
    elseif fw.naming_scheme == :file_number
        file_num = Int(iteration / fw.output_frequency)
        fw.filename_prefix * name * lpad(file_num, fw.padding, "0") * ext(fw)
    else
        throw(ArgumentError("Invalid naming scheme: $(fw.naming_scheme)"))
    end
end

#
# Checkpointing functions
#

function write_output(model::Model{arch}, chk::Checkpointer) where arch <: Architecture
    filepath = joinpath(chk.dir, filename(chk, model.clock.iteration))

    forcing_functions = model.forcing

    # Do not include forcing functions and FFT plans. We want to avoid serializing
    # FFTW and CuFFT plans as serializing functions is not supported by JLD, and
    # seems like a tricky business in general.
    model.forcing = nothing
    model.poisson_solver = nothing

    println("WARNING: Forcing functions will not be serialized!")

    println("[Checkpointer] Serializing model to disk: $filepath")
    f = JLD.jldopen(filepath, "w", compress=true)
    JLD.@write f model
    close(f)

    println("[Checkpointer] Reconstructing FFT plans...")
    model.poisson_solver = init_poisson_solver(arch(), model.grid, model.stepper_tmp.fCC1)

    # Putting back in the forcing functions.
    model.forcing = forcing_functions

    return nothing
end

function restore_from_checkpoint(filepath)
    println("Deserializing model from disk: $filepath")
    f = JLD.jldopen(filepath, "r")
    model = read(f, "model");
    close(f)

    println("Reconstructing FFT plans...")
    model.poisson_solver = init_poisson_solver(arch(model)(), model.grid, model.stepper_tmp.fCC1)

    model.forcing = Forcing(nothing, nothing, nothing, nothing, nothing)
    println("WARNING: Forcing functions have been set to nothing!")

    return model
end

#
# Binary output function
#

function write_output(model::Model, fw::BinaryOutputWriter)
    for (field, field_name) in zip(fw.fields, fw.field_names)
        filepath = joinpath(fw.dir, filename(fw, field_name, model.clock.iteration))

        println("[BinaryOutputWriter] Writing $field_name to disk: $filepath")
        if model.metadata == :CPU
            write(filepath, field.data)
        elseif model.metadata == :GPU
            write(filepath, Array(field.data))
        end
    end
end

function read_output(model::Model, fw::BinaryOutputWriter, field_name, time)
    filepath = joinpath(fw.dir, filename(fw, field_name, time_step))
    arr = zeros(model.metadata.float_type, model.grid.Nx, model.grid.Ny, model.grid.Nz)

    open(filepath, "r") do fio
        read!(fio, arr)
    end

    return arr
end

#
# NetCDF output functions
#
# Eventually, we want to permit the user to flexibly define what is outputted.
# The user-defined function that produces output may involve computations, launching kernels,
# etc; so this API needs to be designed. For now, we simply save u, v, w, and T.

function write_output(model::Model, fw::NetCDFOutputWriter)
    fields = Dict(
        "xC" => collect(model.grid.xC),
        "yC" => collect(model.grid.yC),
        "zC" => collect(model.grid.zC),
        "xF" => collect(model.grid.xF),
        "yF" => collect(model.grid.yF),
        "zF" => collect(model.grid.zF),
        "u" => Array(model.velocities.u.data),
        "v" => Array(model.velocities.v.data),
        "w" => Array(model.velocities.w.data),
        "T" => Array(model.tracers.T.data),
        "S" => Array(model.tracers.S.data)
    )

    if fw.async
        # Execute asynchronously on worker 2.
        i = model.clock.iteration
        @async remotecall(write_output_netcdf, 2, fw, fields, i)
    else
        write_output_netcdf(fw, fields, model.clock.iteration)
    end

    return nothing
end

function write_output_netcdf(fw::NetCDFOutputWriter, fields, iteration)
    xC, yC, zC = fields["xC"], fields["yC"], fields["zC"]
    xF, yF, zF = fields["xF"], fields["yF"], fields["zF"]

    u, v, w = fields["u"], fields["v"], fields["w"]
    T, S    = fields["T"], fields["S"]

    xC_attr = Dict("longname" => "Locations of the cell centers in the x-direction.", "units" => "m")
    yC_attr = Dict("longname" => "Locations of the cell centers in the y-direction.", "units" => "m")
    zC_attr = Dict("longname" => "Locations of the cell centers in the z-direction.", "units" => "m")

    xF_attr = Dict("longname" => "Locations of the cell faces in the x-direction.", "units" => "m")
    yF_attr = Dict("longname" => "Locations of the cell faces in the y-direction.", "units" => "m")
    zF_attr = Dict("longname" => "Locations of the cell faces in the z-direction.", "units" => "m")

    u_attr = Dict("longname" => "Velocity in the x-direction", "units" => "m/s")
    v_attr = Dict("longname" => "Velocity in the y-direction", "units" => "m/s")
    w_attr = Dict("longname" => "Velocity in the z-direction", "units" => "m/s")
    T_attr = Dict("longname" => "Temperature", "units" => "K")
    S_attr = Dict("longname" => "Salinity", "units" => "g/kg")

    filepath = joinpath(fw.dir, filename(fw, "", iteration))

    if fw.async
        println("[Worker $(Distributed.myid()): NetCDFOutputWriter] Writing fields to disk: $filepath")
    else
        println("[NetCDFOutputWriter] Writing fields to disk: $filepath")
    end

    isfile(filepath) && rm(filepath)

    nccreate(filepath, "u", "xF", xC, xC_attr,
                            "yC", yC, yC_attr,
                            "zC", zC, zC_attr,
                            atts=u_attr, compress=fw.compression)

    nccreate(filepath, "v", "xC", xC, xC_attr,
                            "yF", yC, yC_attr,
                            "zC", zC, zC_attr,
                            atts=v_attr, compress=fw.compression)

    nccreate(filepath, "w", "xC", xC, xC_attr,
                            "yC", yC, yC_attr,
                            "zF", zC, zC_attr,
                            atts=w_attr, compress=fw.compression)

    nccreate(filepath, "T", "xC", xC, xC_attr,
                            "yC", yC, yC_attr,
                            "zC", zC, zC_attr,
                            atts=T_attr, compress=fw.compression)

    nccreate(filepath, "S", "xC", xC, xC_attr,
                            "yC", yC, yC_attr,
                            "zC", zC, zC_attr,
                            atts=S_attr, compress=fw.compression)

    ncwrite(u, filepath, "u")
    ncwrite(v, filepath, "v")
    ncwrite(w, filepath, "w")
    ncwrite(T, filepath, "T")
    ncwrite(S, filepath, "S")

    ncclose(filepath)

    return nothing
end

function read_output(fw::NetCDFOutputWriter, field_name, iter)
    filepath = joinpath(fw.dir, filename(fw, "", iter))
    println("[NetCDFOutputWriter] Reading fields from disk: $filepath")
    field_data = ncread(filepath, field_name)
    ncclose(filepath)
    return field_data
end


text:
"""
    Frames

Stores the actual computed frames and the user input
which can be `:same` or `RFrames(10)`.
The `frames` are computed in [`render`](@ref).
"""
mutable struct Frames{T}
    frames::Union{Nothing,UnitRange}
    user::T
end

Base.convert(::Type{Frames}, x) = Frames(nothing, x)
Base.convert(::Type{Frames}, x::Frames) = x
Base.convert(::Type{Frames}, x::UnitRange) = Frames(x, x)

Base.copy(f::Frames) = Frames(f.frames, f.user)

"""
    set_frames!(parent, elem, last_frames::UnitRange; is_first=false)

Compute the frames based on a.frames and `last_frames`.
Save the result in `a.frames.frames` which can be accessed via [`get_frames`](@ref).

# Arguments
- `parent` is either nothing or the Object for the Action
- `elem` is the Object or Action
- `last_frames` holds the frames of the previous object or action.
- `is_first` defines whether this is the first child of the parent (for actions)
"""
function set_frames!(parent, elem, last_frames::UnitRange; is_first = false)
    frames = elem.frames.user
    elem.frames.frames = get_frames(parent, elem, frames, last_frames; is_first = is_first)
end

"""
    get_frames(elem)

Return `elem.frames.frames` which holds the computed frames for the AbstractObject or AbstractAction `a`.
"""
function get_frames(elem)
    elem.frames.frames
end

"""
    get_frames(parent, elem, frames::Symbol, last_frames::UnitRange; is_first=false)

Get the frames based on a symbol defined in `FRAMES_SYMBOL` and the `last_frames`.
Throw `ArgumentError` if symbol is unknown
"""
function get_frames(parent, elem, frames::Symbol, last_frames::UnitRange; is_first = false)
    if frames === :same
        if elem isa AbstractAction && is_first
            return 1:length(last_frames)
        end
        return last_frames
    elseif frames === :all
        return 1:maximum(CURRENT_VIDEO[1].background_frames)
    else
        backtick_frame_symbol = map(x -> "`:$x`", FRAMES_SYMBOL)
        allowed_frames_str = join(backtick_frame_symbol, ", ", " and ")
        err_msg = "Currently the only symbols supported for defining frames are $allowed_frames_str."
        throw(ArgumentError(err_msg))
    end
end

"""
    get_frames(parent, elem, relative::RFrames, last_frames::UnitRange; is_first=falsee)

Return the frames based on a relative frames [`RFrames`](@ref) object and the `last_frames`.
"""
function get_frames(
    parent,
    elem,
    relative::RFrames,
    last_frames::UnitRange;
    is_first = false,
)
    start_frame = last(last_frames) + first(relative.frames)
    last_frame = last(last_frames) + last(relative.frames)
    return start_frame:last_frame
end

"""
    get_frames(parent, elem, glob::GFrames, last_frames::UnitRange)

Return the frames based on a global frames [`GFrames`](@ref) object and the `last_frames`.
If `is_action` is false this is the same as defining the frames as just a unit range.
Inside an action it's now defined globally though.
"""
function get_frames(parent, elem, glob::GFrames, last_frames::UnitRange; is_first = false)
    if elem isa AbstractAction
        return glob.frames .- first(get_frames(parent)) .+ 1
    end
    return glob.frames
end

"""
    function get_frames(parent, elem, func_frames::Function, last_frames::UnitRange; is_first = false)

Return the frames based on a specified function. The function `func_frames` is simply evaluated
"""
function get_frames(
    parent,
    elem,
    func_frames::Function,
    last_frames::UnitRange;
    is_first = false,
)
    return func_frames()
end

"""
    prev_start()

The start frame of the previous object or for an action the start frame of the parental object.
Can be used to provide frame ranges like:
```
@Frames(prev_start(), 10)
```
"""
function prev_start()
    if CURRENT_OBJECT_ACTION_TYPE[1] == :Object
        PREVIOUS_OBJECT[1].frames.frames[1]
    else
        PREVIOUS_ACTION[1].frames.frames[1]
    end
end

"""
    prev_end()

The end frame of the previous object or for an action the end frame of the parental object.
Can be used to provide frame ranges like:
```
@Frames(prev_end()-10, 10)
```
"""
function prev_end()
    if CURRENT_OBJECT_ACTION_TYPE[1] == :Object
        PREVIOUS_OBJECT[1].frames.frames[end]
    else
        PREVIOUS_ACTION[1].frames.frames[end]
    end
end

startof(oa::Union{AbstractAction,AbstractObject}) = oa.frames.frames[1]
endof(oa::Union{AbstractAction,AbstractObject}) = oa.frames.frames[end]

"""
    @Frames(start, len)
    @Frames(start, stop=)

Can be used to define frames using functions like [`prev_start`](@ref) or [`prev_end`](@ref)

# Example
```julia
red_circ = Object(1:90, (args...)->circ("red"))
blue_circ = Object(@Frames(prev_start()+20, 70), (args...)->circ("blue"))
blue_circ = Object(@Frames(prev_start()+20, stop=90), (args...)->circ("blue"))
```
is the same as
```julia
red_circ = Object(1:90, (args...)->circ("red"))
blue_circ = Object(21:90, (args...)->circ("blue"))
blue_circ = Object(41:90, (args...)->circ("blue"))
```
"""
macro Frames(start, in_args...)
    args = []
    kwargs = Pair{Symbol,Any}[]
    kwarg_symbols = Symbol[]
    for el in in_args
        if Meta.isexpr(el, :(=))
            push!(kwargs, Pair(el.args...))
            push!(kwarg_symbols, el.args[1])
        else
            push!(args, el)
        end
    end
    stop_idx = findfirst(==(:stop), kwarg_symbols)
    if stop_idx !== nothing
        stop = kwargs[stop_idx][2]
        return esc(quote
            Javis.Frames(nothing, () -> ($start):($stop))
        end)
    elseif isempty(kwarg_symbols)
        esc(quote
            Javis.Frames(nothing, () -> ($start):($start + $(args[1]) - 1))
        end)
    end
end


text:
values_type(::T) where {T} = values_type(T)
values_type(::Type{T}) where {T} = T

keys_type(::T) where {T<:AbstractVector} = LinearIndices{1,Tuple{Base.OneTo{Int64}}}
keys_type(x::T) where {T} = typeof(keys(x))

## TODO these should go in base if possible
struct BitAnd{F1<:Function,F2<:Function} <: Function
    f1::F1
    f2::F2
end

(f::BitAnd)(x) = f.f1(x) & f.f2(x)

Base.:&(f1::Function, f2::Function) = BitAnd(f1, f2)

struct BitOr{F1<:Function,F2<:Function} <: Function
    f1::F1
    f2::F2
end

(f::BitOr)(x) = f.f1(x) | f.f2(x)

Base.:|(f1::Function, f2::Function) = BitOr(f1, f2)

F2Lt{T} = Fix2{<:Union{typeof(<),typeof(<=)},T}
F2Gt{T} = Fix2{<:Union{typeof(>),typeof(>=)},T}
F2Eq{T} = Fix2{<:Union{typeof(isequal),typeof(==)},T}


text:
using StaticArrays
using SHIPs
using SHIPs.SphericalHarmonics
SH = SHIPs.SphericalHarmonics
##  TODO: there is something wrong here! => test errors due to sqrt(neg numbers)
@info("Test 3: Old vs New Timing")
L = 15
θ = rand() * π
cos_θ = cos(θ)
φ = (rand()-0.5) * 2*π
r = 0.1+rand()
R = SVector(r*sin(θ)*cos(φ), r*sin(θ)*sin(φ), r*cos(θ))
x, y, z, s = SH.compute_rxz(R)
coeff = SH.compute_coefficients(L)
P = Vector{Float64}(undef, SH.sizeP(L))
SH.compute_p!(L, z, coeff, P)
Y = Vector{ComplexF64}(undef, SH.sizeY(L))
@info("Old Implementation:")
@btime compute_y!($L, $cos_θ, $φ, $P, $Y)
@info("New Implementation:")
@btime cYlm_from_cart!($Y, $L, $r, $x, $z, $s, $P)
@info("Experimental Real Spherical Harmonics:")
Yr = Vector{Float64}(undef, SH.sizeY(L))
@btime SH.rYlm_from_cart!($Y, $L, $r, $x, $z, $s, $P)

max_stars_repo_name: cortner/SHIPs

max_issues_repo_name: cortner/SHIPs

max_forks_repo_name: cortner/SHIPs


text:
using Test
include("RandomHypergraphs.jl")
include("load_data.jl")

@testset "Stub Matching" begin
    # Degree and edge-dimension sequences are correct
    DK(H::Hypergraph) = H.D, H.K
    sequence_match(D::Vector{Int64}, K::Vector{Int64}) = return DK(stub_matching(D,K)) == (D, sort(K, rev=true))

    # Realizable hypergraphs
    @test sequence_match([1,1,1,1], [2,2])
    @test sequence_match([1,2,3,4,5,6],[1,2,3,4,5,6])
    @test sequence_match([6,3,3,3], [4,3,3,1,1,1,1,1])
    @test sequence_match([1,2,3,1,1,1],[3,2,3,1])

    # Unrealizable hypergraphs
    @test_throws ArgumentError sequence_match([0],[5])
    @test_throws ArgumentError sequence_match([5,5,5,5,5],[5])
    @test_throws ArgumentError sequence_match([1,1,1,1],[4,4,4])

end
@testset "MCMC" begin
    @testset "Vertex" begin
        @test_throws AssertionError MCMC(Hypergraph(Array{Array{Int64,1},1}(),4,0))
        @test_throws AssertionError MCMC(Hypergraph([[1,2]],2,1))
        @test typeof(MCMC(Hypergraph([[1,2],[3,4],[1,4]], 4, 3))) <: Hypergraph
        @test typeof(MCMC(stub_matching([1,2,3,1,1,1],[3,2,3,1]); samples=1000)) <: Hypergraph
        @test typeof(MCMC(load_hypergraph("Enron"); samples=10000)) <: Hypergraph
    end

    @testset "Stub" begin
        @test_throws AssertionError typeof(MCMC(Hypergraph(Array{Array{Int64,1},1}(),4,0); type="stub")) <: Hypergraph
        @test typeof(MCMC(Hypergraph([[1,2],[3,4],[1,4]], 4, 3); type="stub")) <: Hypergraph
        @test typeof(MCMC(stub_matching([1,2,3,1,1,1],[3,2,3,1]); type="stub")) <: Hypergraph
        @test typeof(MCMC(load_hypergraph("Enron"); samples=10000, type="stub")) <: Hypergraph
    end
end


text:
@testset "pool" begin
    @testset "GlobalPool" begin
        p = GlobalPool(+)
        n = 10
        chin = 6
        X = rand(Float32, 6, n)
        g = GNNGraph(random_regular_graph(n, 4), ndata=X, graph_type=GRAPH_T)
        u = p(g, X)
        @test u ≈ sum(X, dims=2)

        ng = 3
        g = Flux.batch([GNNGraph(random_regular_graph(n, 4),
                                 ndata=rand(Float32, chin, n),
                                 graph_type=GRAPH_T)
                        for i=1:ng])
        u = p(g, g.ndata.x)
        @test size(u) == (chin, ng)
        @test u[:,[1]] ≈ sum(g.ndata.x[:,1:n], dims=2)
        @test p(g).gdata.u == u

        test_layer(p, g, rtol=1e-5, exclude_grad_fields = [:aggr], outtype=:graph)
    end

    @testset "GlobalAttentionPool" begin
        n = 10
        chin = 6
        chout = 5
        ng = 3

        fgate = Dense(chin, 1)
        ffeat = Dense(chin, chout)
        p = GlobalAttentionPool(fgate, ffeat)
        @test length(Flux.params(p)) == 4

        g = Flux.batch([GNNGraph(random_regular_graph(n, 4),
                                 ndata=rand(Float32, chin, n),
                                 graph_type=GRAPH_T)
                        for i=1:ng])

        test_layer(p, g, rtol=1e-5, outtype=:graph, outsize=(chout, ng))
    end

    @testset "TopKPool" begin
        N = 10
        k, in_channel = 4, 7
        X = rand(in_channel, N)
        for T = [Bool, Float64]
            adj = rand(T, N, N)
            p = TopKPool(adj, k, in_channel)
            @test eltype(p.p) === Float32
            @test size(p.p) == (in_channel,)
            @test eltype(p.Ã) === T
            @test size(p.Ã) == (k, k)
            y = p(X)
            @test size(y) == (in_channel, k)
        end
    end

    @testset "topk_index" begin
        X = [8,7,6,5,4,3,2,1]
        @test topk_index(X, 4) == [1,2,3,4]
        @test topk_index(X', 4) == [1,2,3,4]
    end
end


text:
module InitializePhonons

using Random

using ..Models: HolsteinModel, SSHModel, update_model!
using ..Utilities: get_index, reshaped

export init_phonons_single_site!, init_phonons_half_filled!, sample_qho

function init_phonons_half_filled!(ssh::SSHModel{T1,T2}) where {T1,T2}

    # info about temperature of system
    β  = ssh.β::T1
    Δτ = ssh.Δτ::T1
    Lτ = ssh.Lτ::Int

    # number of phonons
    Nph = ssh.Nph

    # get phonon fields
    x = ssh.x

    # keeps track of the fields
    field = 0

    # iterate over phonons
    for phonon in 1:Nph

        # calculate average phonon position
        bond = ssh.phonon_to_bond[phonon]
        t    = ssh.t[bond]
        α    = ssh.α[phonon]
        ω    = ssh.ω[phonon]
        if iszero(α) || iszero(t)
            x0 = sample_qho(ω,β,ssh.rng)
        else
            x0 = (2.0*rand(ssh.rng)-1.0)/5.0 * (t/α)
        end

        # iterate over imaginary time slices
        for τ in 1:Lτ

            # increment field count
            field += 1

            # set phonon value
            x[field] = x0
        end
    end

    # account for equivalent fields
    @views @. x = x[ssh.primary_field]

    # udpate exponentiated interaction matrix
    update_model!(ssh)

    return nothing
end

function init_phonons_half_filled!(holstein::HolsteinModel{T1,T2}) where {T1,T2}

    # info about temperature of system
    β  = holstein.β::T1
    Δτ = holstein.Δτ::T1
    Lτ = holstein.Lτ::Int

    # iterate over sites in lattice
    for site in 1:holstein.Nsites

        # get parameters for site in lattice
        ω = holstein.ω[site]
        λ = holstein.λ[site]

        # get all phonon fields corresponding to current
        # site in lattice.
        i_start = get_index(1,site,Lτ)
        i_end   = get_index(Lτ,site,Lτ)
        path    = @view holstein.x[i_start:i_end]

        # shift levy path by ammount corresponding having either
        # a density of 0, 1 or 2 on the site
        x0 = λ/ω^2 * rand(holstein.rng,-1:1:1)
        xr = x0 + sample_qho(ω,β,holstein.rng)
        @. path = xr
    end

    # udpate exponentiated interaction matrix
    update_model!(holstein)

    return nothing
end

"""
Samples the position distribution of a QHO with frequency `ω` at inverse temperature `β`.
"""
function sample_qho(ω::T,β::T,rng::AbstractRNG)::T where {T<:Number}

    if ω>0
        σ = 1/sqrt(2*ω*tanh(β*ω/2))
    else
        σ = 1.0
    end
    return σ*randn(rng)
end

end

max_stars_repo_name: cohensbw/Langevin

max_issues_repo_name: cohensbw/ElPhDynamics

max_forks_repo_name: cohensbw/ElPhDynamics


text:
# Define symbol inputs for general measure method types
const sampling = :sampling
const quadrature = :quadrature

# Extend Base.copy for new variable types
Base.copy(v::MeasureRef, new_model::InfiniteModel) = MeasureRef(new_model, v.index)

"""
    JuMP.name(mref::MeasureRef)::String

Extend [`JuMP.name`](@ref) to return the name associated with a measure
reference.
"""
function JuMP.name(mref::MeasureRef)::String
    return JuMP.owner_model(mref).meas_to_name[mref.index]
end

"""
    JuMP.set_name(mref::MeasureRef, name::String)

Extend [`JuMP.set_name`](@ref) to specify the name of a measure reference.
"""
function JuMP.set_name(mref::MeasureRef, name::String)
    JuMP.owner_model(mref).meas_to_name[JuMP.index(mref)] = name
    return
end

"""
    JuMP.is_valid(model::InfiniteModel, mref::MeasureRef)::Bool

Extend [`JuMP.is_valid`](@ref) to return `Bool` whether `mref` is valid.
"""
function JuMP.is_valid(model::InfiniteModel, mref::MeasureRef)::Bool
    return (model === JuMP.owner_model(mref) && JuMP.index(mref) in keys(model.measures))
end

# Set a default weight function
_w(t) = 1

"""
    DiscreteMeasureData(parameter_ref::ParameterRef,
                        coefficients::Vector{<:Number},
                        supports::Vector{<:Number}; name::String = "measure",
                        weight_function::Function = w(t) = 1)::DiscreteMeasureData

Returns a `DiscreteMeasureData` object that can be utilized to define
measures using [`measure`](@ref). This accepts input for a scalar (single)
parameter. Note that `name` is used for printing purposes and a description of
the other arguments is provided in the documentation for
[`DiscreteMeasureData`](@ref). Errors if supports are out bounds or an unequal
number of supports and coefficients are given.

**Example**
```julia-repl
julia> data = DiscreteMeasureData(pref, [0.5, 0.5], [1, 2], name = "example")
DiscreteMeasureData(pref, [0.5, 0.5], [1, 2], "example", InfiniteOpt._w)
```
"""
function DiscreteMeasureData(parameter_ref::ParameterRef,
                             coefficients::Vector{<:Number},
                             supports::Vector{<:Number};
                             name::String = "measure",
                             weight_function::Function = _w
                             )::DiscreteMeasureData
    is_finite_parameter(parameter_ref) && error("Measure parameter cannot be " *
                                                "finite.")
    return DiscreteMeasureData(parameter_ref, coefficients, supports, name,
                               weight_function)
end

"""
    DiscreteMeasureData(parameter_ref::AbstractArray{<:ParameterRef},
                        coefficients::Vector{<:Number},
                        supports::Vector{<:AbstractArray{<:Number}};
                        name::String = "measure",
                        weight_function::Function = w(t) = 1
                        )::MultiDiscreteMeasureData

Returns a `MultiDiscreteMeasureData` object that can be utilized to
define measures using [`measure`](@ref). This accepts input for an array (multi)
parameter. The inner arrays in the supports vector need to match the formatting
of the array used for `parameter_ref`. Note that `name` is used for printing
purposes and a description of the other arguments is provided in the
documentation for [`MultiDiscreteMeasureData`](@ref). Errors if supports are out
bounds, an unequal number of supports and coefficients are given, the array
formats do not match, or the parameters have different group IDs.

**Example**
```julia-repl
julia> data = DiscreteMeasureData(prefs, [0.5, 0.5], [[1, 1], [2, 2]], name = "example");

julia> typeof(data)
MultiDiscreteMeasureData
```
"""
function DiscreteMeasureData(parameter_ref::AbstractArray{<:ParameterRef},
                             coefficients::Vector{<:Number},
                             supports::Vector{<:AbstractArray};
                             name::String = "measure",
                             weight_function::Function = _w
                             )::MultiDiscreteMeasureData
    any(is_finite_parameter.(parameter_ref)) && error("Measure parameter cannot " *
                                                      " be finite.")
    supports = [convert(JuMPC.SparseAxisArray, s) for s in supports]
    parameter_ref = convert(JuMPC.SparseAxisArray, parameter_ref)
    return MultiDiscreteMeasureData(parameter_ref, coefficients, supports, name,
                                    weight_function)
end

"""
    measure_name(data::AbstractMeasureData)::String

Return the measure name string stored in `data`. This is intended as an internal
function to be used with measure addition. User-defined measure data types will
need to extend this function otherwise the measure names default to `"measure"`.
"""
function measure_name(data::AbstractMeasureData)::String
    return "measure"
end

# DiscreteMeasureData and MultiDiscreteMeasureData
function measure_name(data::Union{DiscreteMeasureData, MultiDiscreteMeasureData})::String
    return data.name
end

"""
    parameter_refs(data::AbstractMeasureData)::Union{ParameterRef,
                                                     AbstractArray{<:ParameterRef}}

Return the infinite parameter reference(s) in `data`. This is intended as an
internal function to be used with measure addition. User-defined measure data types
will need to extend this function otherwise an error is thrown.
"""
function parameter_refs(data::AbstractMeasureData)
    error("Function `parameter_refs` not extended for measure data of type $(typeof(data)).")
end

# DiscreteMeasureData and MultiDiscreteMeasureData
function parameter_refs(data::Union{DiscreteMeasureData, MultiDiscreteMeasureData})
    return data.parameter_ref
end

"""
    supports(data::AbstractMeasureData)::Vector

Return the supports stored in `data` associated with its infinite parameters.
This is intended as en internal method for measure creation and ensures any
new supports are added to parameters. User-defined measure data types should
extend this function if appropriate, otherwise an empty vector is returned.
"""
function supports(data::AbstractMeasureData)::Vector
    return Number[]
end

# DiscreteMeasureData and MultiDiscreteMeasureData
function supports(data::Union{DiscreteMeasureData, MultiDiscreteMeasureData})::Vector
    return data.supports
end

# Parse the string for displaying a measure
function _make_meas_name(meas::Measure)::String
    return string(measure_name(meas.data), "(", JuMP.function_string(JuMP.REPLMode,
                                                            meas.func), ")")
end

# Used to update the model.var_to_meas and model.param_tomes fields
# this is needed to update measures if variables are deleted
function _update_var_meas_mapping(vrefs::Vector{<:GeneralVariableRef},
                                  mindex::Int)
    for vref in vrefs
        model = JuMP.owner_model(vref)
        if isa(vref, InfOptVariableRef)
            if haskey(model.var_to_meas, JuMP.index(vref))
                push!(model.var_to_meas[JuMP.index(vref)], mindex)
            else
                model.var_to_meas[JuMP.index(vref)] = [mindex]
            end
        elseif isa(vref, ParameterRef)
            if haskey(model.param_to_meas, JuMP.index(vref))
                push!(model.param_to_meas[JuMP.index(vref)], mindex)
            else
                model.param_to_meas[JuMP.index(vref)] = [mindex]
            end
        elseif isa(vref, MeasureRef)
            if haskey(model.meas_to_meas, JuMP.index(vref))
                push!(model.meas_to_meas[JuMP.index(vref)], mindex)
            else
                model.meas_to_meas[JuMP.index(vref)] = [mindex]
            end
        elseif isa(vref, ReducedInfiniteVariableRef)
            if haskey(model.reduced_to_meas, JuMP.index(vref))
                push!(model.reduced_to_meas[JuMP.index(vref)], mindex)
            else
                model.reduced_to_meas[JuMP.index(vref)] = [mindex]
            end
        end
    end
    return
end

## Used to add the measure index to param_to_meas for parameters that are used
## in the evaluation data
# DiscreteMeasureData
function _update_param_data_mapping(model::InfiniteModel,
                                    pref::ParameterRef,
                                    mindex::Int)
    if haskey(model.param_to_meas, JuMP.index(pref))
        if !(mindex in model.param_to_meas[JuMP.index(pref)])
            push!(model.param_to_meas[JuMP.index(pref)], mindex)
        end
    else
        model.param_to_meas[JuMP.index(pref)] = [mindex]
    end
    return
end

# MultiDiscreteMeasureData
function _update_param_data_mapping(model::InfiniteModel,
                                    prefs::AbstractArray{<:ParameterRef},
                                    mindex::Int)
    for pref in prefs
        if haskey(model.param_to_meas, JuMP.index(pref))
            if !(mindex in model.param_to_meas[JuMP.index(pref)])
                push!(model.param_to_meas[JuMP.index(pref)], mindex)
            end
        else
            model.param_to_meas[JuMP.index(pref)] = [mindex]
        end
    end
    return
end

## Internal functions for adding measure data supports to the parameter supports
# scalar pref
function _add_supports_to_parameters(pref::ParameterRef,
                                     supports::Vector{<:Number})
    add_supports(pref, supports)
    return
end

# array pref
function _add_supports_to_parameters(prefs::AbstractArray{<:ParameterRef},
                                     supports::Array{<:AbstractArray{<:Number}})
    for i = eachindex(supports)
        for key in keys(prefs)
            add_supports(prefs[key], supports[i][key])
        end
    end
    return
end

"""
    add_measure(model::InfiniteModel, meas::Measure)::MeasureRef

Add a measure to `model` and return the corresponding measure reference. This
operates in a manner similar to [`JuMP.add_variable`](@ref). Note this intended
as an internal method.
"""
function add_measure(model::InfiniteModel, meas::Measure)::MeasureRef
    model.next_meas_index += 1
    index = model.next_meas_index
    JuMP.check_belongs_to_model(meas.func, model)
    prefs = parameter_refs(meas.data)
    supps = supports(meas.data)
    all(JuMP.is_valid.(model, prefs)) || error("Invalid parameter dependence in " *
                                               "measure data")
    if length(supps) != 0
        _add_supports_to_parameters(prefs, supps)
    end
    vrefs = _all_function_variables(meas.func)
    _update_var_meas_mapping(vrefs, index)
    _update_param_data_mapping(model, prefs, index)
    mref = MeasureRef(model, model.next_meas_index)
    model.measures[mref.index] = meas
    JuMP.set_name(mref, _make_meas_name(meas))
    model.meas_in_objective[index] = false
    return mref
end

"""
    measure_function(mref::MeasureRef)::JuMP.AbstractJuMPScalar

Return the function associated with `mref`.
"""
function measure_function(mref::MeasureRef)::JuMP.AbstractJuMPScalar
    return JuMP.owner_model(mref).measures[JuMP.index(mref)].func
end

"""
    measure_data(mref::MeasureRef)::AbstractMeasureData

Return the measure data associated with `mref`.
"""
function measure_data(mref::MeasureRef)::AbstractMeasureData
    return JuMP.owner_model(mref).measures[JuMP.index(mref)].data
end

# Check a measure function for a particular parameter and return Bool
function _has_parameter(vrefs::Vector{<:GeneralVariableRef},
                        pref::ParameterRef)::Bool
    if _has_variable(vrefs, pref)
        return true
    end
    model = JuMP.owner_model(pref)
    relavent_ivindices = model.param_to_vars[JuMP.index(pref)]
    relavent_ivrefs = [InfiniteVariableRef(model, vindex) for vindex in relavent_ivindices]
    for ivref in relavent_ivrefs
        if _has_variable(vrefs, ivref)
            return true
        elseif used_by_reduced_variable(ivref)
            for index in model.infinite_to_reduced[JuMP.index(ivref)]
                if _has_variable(vrefs, ReducedInfiniteVariableRef(model, index))
                    return true
                end
            end
        end
    end
    return false
end

## Check if expr contains a parameter directly or via an infinite variable
# scalar pref
function _check_has_parameter(vrefs::Vector{<:GeneralVariableRef},
                              pref::ParameterRef)
    if !_has_parameter(vrefs, pref)
        error("Measure expression is not parameterized by the parameter " *
              "specified in the measure data.")
    end
    return
end

# array pref
function _check_has_parameter(vrefs::Vector{<:GeneralVariableRef},
                              pref::JuMPC.SparseAxisArray{<:ParameterRef})
    for key in keys(pref.data)
        if !_has_parameter(vrefs, pref.data[key])
            error("Measure expression is not parameterized by the parameter " *
                  "specified in the measure data.")
        end
    end
    return
end

# Parse the model pertaining to an expression
function _model_from_expr(vrefs::Vector{<:GeneralVariableRef})
    if length(vrefs) > 0
        return JuMP.owner_model(vrefs[1])
    else
        return
    end
end

"""
    measure_data_in_hold_bounds(data::AbstractMeasureData,
                                bounds::ParameterBounds)::Bool

Return a `Bool` whether the domain of `data` is valid in accordance with
`bounds`. This is intended as an internal method and is used to check hold
variables used in measures. User-defined measure data types will need to
extend this function to enable this error checking, otherwise it is skipped and
a warning is given.
"""
function measure_data_in_hold_bounds(data::AbstractMeasureData,
                                     bounds::ParameterBounds)::Bool
    @warn "Unable to check if hold variables bounds are valid in measure " *
           "with measure data type `$(typeof(data))`. This can be resolved by " *
           "extending `measure_data_in_hold_bounds`."
    return true
end

# DiscreteMeasureData
function measure_data_in_hold_bounds(data::DiscreteMeasureData,
                                     bounds::ParameterBounds)::Bool
    pref = parameter_refs(data)
    supps = supports(data)
    if haskey(bounds.intervals, pref)
        return supports_in_set(supps, bounds.intervals[pref])
    end
    return true
end

# MultiDiscreteMeasureData
function measure_data_in_hold_bounds(data::MultiDiscreteMeasureData,
                                     bounds::ParameterBounds)::Bool
    prefs = parameter_refs(data)
    supps = supports(data)
    for i in eachindex(supps)
        for key in keys(prefs)
            if haskey(bounds.intervals, prefs[key])
                if !supports_in_set(supps[i][key], bounds.intervals[prefs[key]])
                    return false
                end
            end
        end
    end
    return true
end

## Check that variables don't violate the parameter bounds
# GeneralVariableRef
function _check_var_bounds(vref::GeneralVariableRef, data::AbstractMeasureData)
    return
end

# HoldVariableRef (single parameter)
function _check_var_bounds(vref::HoldVariableRef, data::AbstractMeasureData)
    bounds = parameter_bounds(vref)
    if !measure_data_in_hold_bounds(data, bounds)
        error("Measure bounds violate hold variable bounds.")
    end
    return
end

# MeasureRef
function _check_var_bounds(mref::MeasureRef, data::AbstractMeasureData)
    vrefs = _all_function_variables(measure_function(mref))
    for vref in vrefs
        _check_var_bounds(vref, data)
    end
    return
end

"""
    measure(expr::JuMP.AbstractJuMPScalar, data::AbstractMeasureData)::MeasureRef

Return a measure reference that evaluates `expr` using according to `data`. This
is the preferred method for implementing measures which follow the form:
``\\int_{p \\in P} expr(p) w(p) dp`` where ``p`` is an infinite parameter (scalar
or vector) and ``w`` is the weight function. The measure data `data` determines
how the measure is to be evaluated. Typically, the [`DiscreteMeasureData`](@ref)
constructor can be used to for `data`. The variable expression `expr` can contain
`InfiniteOpt` variables, infinite parameters, other measure references (meaning
measures can be nested), and constants. Errors if `expr` does not contain
infinite variables, infinite parameters, or measure references. Also errors if
the measure parameter specified in `data` is not in `expr` and is not in any
the nested measure references. Typically, this is called inside of
[`JuMP.@expression`](@ref), [`JuMP.@objective`](@ref), and
[`JuMP.@constraint`](@ref) in a manner similar to `sum`. Note measures are not
explicitly evaluated until [`build_optimizer_model!`](@ref) is called or unless
they are expanded via [`expand`](@ref) or [`expand_all_measures!`](@ref).

**Example**
```julia-repl
julia> tdata = DiscreteMeasureData(t, [0.5, 0.5], [1, 2], name = "name1");

julia> xdata = DiscreteMeasureData(xs, [0.5, 0.5], [[-1, -1], [1, 1]],
                                   name = "name2");

julia> constr_RHS = @expression(model, measure(g - s + 2, tdata) + s^2)
name1(g(t) - s + 2) + s²

julia> @objective(model, Min, measure(g - 1  + measure(T, xdata), tdata))
name1(g(t) - 1 + name2(T(t, x)))
```
"""
function measure(expr::JuMP.AbstractJuMPScalar,
                 data::AbstractMeasureData)::MeasureRef
    vrefs = _all_function_variables(expr)
    model = _model_from_expr(vrefs)
    if model == nothing
        error("Expression contains no variables or parameters.")
    end
    pref = parameter_refs(data)
    if model.has_hold_bounds
        for vref in vrefs
            _check_var_bounds(vref, data)
        end
    end
    meas = Measure(expr, data)
    return add_measure(model, meas)
end

"""
    integral(expr::JuMP.AbstractJuMPScalar,
             [params::Union{ParameterRef, AbstractArray{<:ParameterRef},
                            Nothing} = nothing,
             lb::Union{Number, AbstractArray{<:Number}, Nothing} = nothing,
             ub::Union{Number, AbstractArray{<:Number}, Nothing} = nothing];
             [eval_method::Symbol = sampling,
             num_supports::Int = 10,
             weight_func::Function = _w,
             name = "integral",
             use_existing_supports::Bool = false,
             kwargs...])::MeasureRef

Returns a measure reference that evaluates the integral of `expr` with respect
to infinite parameter(s) from `lb` to `ub`. This thus considers integrals of the
form: ``\\int_{p \\in P} expr(p) w(p) dp`` where ``p`` is an infinite parameter
(scalar or vector) and ``w`` is the weight function is 1 by default. This
function provides a high-level interface that ultimately constructs a
[`DiscreteMeasureData`](@ref) via `eval_method` that is used to call
[`measure`](@ref).

The arugments are as follows:
- `params`: the integral infinite parameter(s) in ``dp``
- `lb` & `ub`: integral upper and lower bounds (defaults to entire parameter domain)
Note that `params` is required if `expr` contains multiple parameters.

The keyword arguments are as follows:
- `eval_method`: method tha generates the supports and coefficients
    - `sampling` --> dispatch to appropriate sampling method (default)
    - `quadrature` --> dispatch to appropriate quadrature method
    - `mc_sampling` --> Monte Carlo sampling
    - `gauss_hermite` --> Gaussian Hermite quadrature (infinite interval)
    - `gauss_legendre` --> Gaussian Legendre quadrature (finite interval)
    - `gauss_laguerre` --> Gaussian Laguerre quadrature (semi-infinite interval)
    - `trapezoid` --> trapezoidal quadrature (finite interval)
- `num_supports`: The number of supports to be generated
- `weight_func`: ``w(p)`` above with parameter value inputs and scalar output
- `name`: the name used in printing
- `use_existing_supports`: Use all supports currently stored in `params`
Note that `mc_sampling` samples from the parameter interval for [`IntervalSet`](@ref)s,
and from the underlying distribution for [`DistributionSet`](@ref)s. Also,
`use_existing_supports` is useful for subsequent integral calls when using
`mc_sampling` such that all the measures use the same supports.

See [`set_integral_defaults`](@ref) to update the default keyword argument values
for all integral calls.

**Example**
```jldoctest; setup = :(using InfiniteOpt, JuMP; model = InfiniteModel(seed = true))
julia> @infinite_parameter(model, x in [0., 1.])
x

julia> @infinite_variable(model, f(x))
f(x)

julia> int = integral(f, num_supports = 5)
(f(x))

julia> expand(int)
0.2 f(0.8236475079774124) + 0.2 f(0.9103565379264364) + 0.2 f(0.16456579813368521) + 0.2 f(0.17732884646626457) + 0.2 f(0.278880109331201)
```
"""
function integral(expr::JuMP.AbstractJuMPScalar,
                  params::Union{ParameterRef, AbstractArray{<:ParameterRef},
                                                            Nothing} = nothing,
                  lb::Union{Number, AbstractArray{<:Number}, Nothing} = nothing,
                  ub::Union{Number, AbstractArray{<:Number}, Nothing} = nothing;
                  kwargs...)::MeasureRef

    # collect parameters from expression if they are not provided
    if isa(params, Nothing)
        if isa(expr, MeasureRef)
            error("Nested call of intregral must specify parameters.")
        end
        params = _all_parameter_refs(expr)
        if length(params) == 0
            error("No infinite parameters in the expression.")
        elseif length(params) == 1
            params = params[1]
        else
            error("Multiple groups of parameters are in the expression. Need to " *
                  "specify one group of parameters only.")
        end
    end

    # count number of parameters
    if isa(params, ParameterRef)
        num_params = 1
    else
        num_params = length(params)
        if num_params == 0
            # error if empty ParameterRef array is provided
            error("No infinite parameter is provided.")
        elseif num_params == 1
            params = first(params)
        else
            ids = unique(group_id.(params))
            if length(ids) > 1
                error("Multiple groups of parameters are specified.")
            end
            # Use SparseAxisArray for params, lb, ub if multiple parameters
            params = convert(JuMPC.SparseAxisArray, params)
            if isa(lb, AbstractArray)
                lb = convert(JuMPC.SparseAxisArray, lb)
            end
            if isa(ub, AbstractArray)
                ub = convert(JuMPC.SparseAxisArray, ub)
            end
        end
    end

    # collect model of the measure
    if isa(params, ParameterRef)
        model = params.model
    else
        model = first(params).model
    end

    # collect keyword arguments
    kwargs = merge(model.integral_defaults, kwargs)
    eval_method = kwargs[:eval_method]
    num_supports = kwargs[:num_supports]
    name = kwargs[:name]
    weight_func = kwargs[:weight_func]
    use_existing_supports = kwargs[:use_existing_supports]

    if eval_method == sampling
        eval_method = mc_sampling
        kwargs[:eval_method] = mc_sampling
    end

    # delete unneeded keyword arguments
    delete!(kwargs, :use_existing_supports)
    delete!(kwargs, :num_supports)

    # make bounds arrays if needed
    if num_params > 1
        if isa(lb, Number)
            lb = JuMPC.SparseAxisArray(Dict(k => lb for k in keys(params)))
        end
        if isa(ub, Number)
            ub = JuMPC.SparseAxisArray(Dict(k => ub for k in keys(params)))
        end
    end

    # check lower bounds
    bounds_valid = true
    if isa(lb, Number)
        bounds_valid = supports_in_set(lb, _parameter_set(first(params)))
    elseif isa(lb, JuMPC.SparseAxisArray)
        bounds_valid = all(supports_in_set.(lb, _parameter_set.(params)))
    end
    bounds_valid || error("Lower bound(s) violate(s) the infinite set domain.")

    # check upper bounds
    if isa(ub, Number)
        bounds_valid = supports_in_set(ub, _parameter_set(first(params)))
    elseif isa(ub, JuMPC.SparseAxisArray)
        bounds_valid = all(supports_in_set.(ub, _parameter_set.(params)))
    end
    bounds_valid || error("Upper bound(s) violate(s) the infinite set domain.")

    # fill in bounds if needed
    set = _parameter_set(first(params))
    if JuMP.has_lower_bound(set)
        # Fill in lower bounds and upper bounds if not given
        if isa(lb, Nothing) || length(lb) == 0
            lb = JuMP.lower_bound.(params)
        end
        if isa(ub, Nothing) || length(ub) == 0
            ub = JuMP.upper_bound.(params)
        end

        # Check the dimension of lb and ub matches number of parameters
        if length(lb) != num_params || length(ub) != num_params
            error("Number of parameters do not match number of lower bounds or " *
                  "upper bounds.")
        end

        if num_params == 1
            if isa(lb, AbstractArray)
                lb = lb[findfirst(x->isa(x, Number), lb)]
            end
            if isa(ub, AbstractArray)
                ub = ub[findfirst(x->isa(x, Number), ub)]
            end
        end

        # Check the input lower bounds and upper bounds are reasonable
        for i in eachindex(lb)
            if lb[i] >= ub[i]
                error("Lower bound is not less than upper bound for some " *
                      "parameter. Please check the input lower bounds " *
                      "and upper bounds.")
            end
        end
    end

    # construct data and return measure if we use existing supports
    if use_existing_supports
        if eval_method == quadrature || eval_method == gauss_legendre ||
           eval_method == gauss_hermite || eval_method == gauss_laguerre
            @warn("Quadrature method will not be used because " *
                  "use_existing_supports is set as true.")
        end
        support = supports(params)
        if !isa(lb, Nothing) && !isa(ub, Nothing)
            support = [i for i in support if all(i .>= lb) && all(i .<= ub)]
        end

        # TODO: generate reasonable coefficients for given supports
        len = length(support)
        if isa(set, DistributionSet)
            data = DiscreteMeasureData(params, ones(len) ./ len, support,
                                       name = name, weight_function = weight_func)
        elseif isa(set, IntervalSet) && eval_method == mc_sampling
            coeffs = ones(len) / len * prod(ub .- lb)
            data = DiscreteMeasureData(params, coeffs, support, name = name,
                                       weight_function = weight_func)
        else
            data = DiscreteMeasureData(params, ones(len), support, name = name,
                                       weight_function = weight_func)
        end
        return measure(expr, data)
    end

    if eval_method == quadrature || eval_method == gauss_legendre ||
       eval_method == gauss_hermite || eval_method == gauss_laguerre
        if num_params > 1
            error("Quadrature method is not supported for multivariate measures.")
        end
        inf_bound_num = (lb == -Inf) + (ub == Inf)
        if inf_bound_num == 0
            kwargs[:eval_method] = gauss_legendre
        elseif inf_bound_num == 1
            kwargs[:eval_method] = gauss_laguerre
        else
            kwargs[:eval_method] = gauss_hermite
        end
    end

    # construct DiscreteMeasureData as data
    data = generate_measure_data(params, num_supports, lb, ub; kwargs...)

    # call measure function to construct the measure
    return measure(expr, data)
end

"""
    integral_defaults(model::InfiniteModel)

Get the default keyword argument values for defining integrals in `model`.

```jldoctest; setup = :(using InfiniteOpt; model = InfiniteModel())
julia> integral_defaults(model)
Dict{Symbol,Any} with 6 entries:
  :num_supports          => 10
  :eval_method           => nothing
  :name                  => "integral"
  :weight_func           => _w
  :use_existing_supports => false
```
"""
function integral_defaults(model::InfiniteModel)
    return model.integral_defaults
end

"""
    set_integral_defaults(model::InfiniteModel; kwargs...)

Set the default keyword argument settings for integrals of the specified model.
The keyword arguments of this function will be recorded in the default keyword
argument values of the model. If the keyword argument has been defined in the
model default, it will be overwritten with the new keyword argument value.
Otherwise, the default will record the new keyword argument and its value for
measures. The default values will be used by integrals constructed from
[`integral`](@ref) calls.

**Example**
```jldoctest; setup = :(using InfiniteOpt; model = InfiniteModel())
julia> integral_defaults(model)
Dict{Symbol,Any} with 6 entries:
  :num_supports          => 10
  :eval_method           => nothing
  :name                  => "integral"
  :weight_func           => _w
  :use_existing_supports => false

julia> set_integral_default(m, num_supports = 5, eval_method = quadrature, new_kwarg = true)

julia> integral_defaults(m)
Dict{Symbol,Any} with 6 entries:
  :new_kwarg             => true
  :num_supports          => 5
  :eval_method           => :quadrature
  :name                  => "integral"
  :weight_func           => _w
  :use_existing_supports => false
```
"""
function set_integral_defaults(model::InfiniteModel; kwargs...)
    merge!(model.integral_defaults, kwargs)
    return
end

"""
    expect(expr::JuMP.AbstractJuMPScalar,
           [params::Union{ParameterRef, AbstractArray{<:ParameterRef},
                          Nothing} = nothing];
           [num_supports::Int = 10,
            use_existing_supports::Bool = false])::MeasureRef

Creates a measure that represents the expected value of an expression in
a random parameter involved in the expression. Return the [`MeasureRef`](@ref)
of the created measure.

**Example**
```jldoctest; setup = :(using InfiniteOpt, JuMP, Distributions; model = InfiniteModel(seed = true))
julia> @infinite_parameter(model, x in Normal(0., 1.))
x

julia> @infinite_variable(model, f(x))
f(x)

julia> meas = expect(f, num_supports = 2)
expect(f(x))

julia> expand(meas)
0.5 f(0.6791074260357777) + 0.5 f(0.8284134829000359)
```
"""
function expect(expr::JuMP.AbstractJuMPScalar,
                params::Union{ParameterRef, AbstractArray{<:ParameterRef},
                              Nothing} = nothing;
                num_supports::Int = 10,
                use_existing_supports::Bool = false)::MeasureRef
    # expectation measure
    return integral(expr, params, num_supports = num_supports,
                    name = "expect", use_existing_supports = use_existing_supports)
end

"""
    support_sum(expr::JuMP.AbstractJuMPScalar,
                [params::Union{ParameterRef, AbstractArray{<:ParameterRef},
                               Nothing} = nothing])::MeasureRef

Creates a measure that represents the sum of the expression over a parameter
using its existing supports. Return the [`MeasureRef`](@ref) of the created
measure.

**Example**
```jldoctest; setup = :(using InfiniteOpt, JuMP; model = InfiniteModel())
julia> @infinite_parameter(model, x in [0, 1], supports = [0.3, 0.7])
x

julia> @infinite_variable(model, f(x))
f(x)

julia> meas = support_sum(f)
sum(f(x))

julia> expand(meas)
f(0.3) + f(0.7)
```
"""
function support_sum(expr::JuMP.AbstractJuMPScalar,
                     params::Union{ParameterRef, AbstractArray{<:ParameterRef},
                                   Nothing} = nothing)::MeasureRef
    # sum measure
    return integral(expr, params, use_existing_supports = true, name = "sum",
                    eval_method = :not_default)
end

"""
    used_by_constraint(mref::MeasureRef)::Bool

Return a `Bool` indicating if `mref` is used by a constraint.

**Example**
```julia-repl
julia> used_by_constraint(mref)
false
```
"""
function used_by_constraint(mref::MeasureRef)::Bool
    return haskey(JuMP.owner_model(mref).meas_to_constrs, JuMP.index(mref))
end

"""
    used_by_measure(mref::MeasureRef)::Bool

Return a `Bool` indicating if `mref` is used by a measure.

**Example**
```julia-repl
julia> used_by_measure(mref)
true
```
"""
function used_by_measure(mref::MeasureRef)::Bool
    return haskey(JuMP.owner_model(mref).meas_to_meas, JuMP.index(mref))
end

"""
    used_by_objective(vmref::MeasureRef)::Bool

Return a `Bool` indicating if `mref` is used by the objective.

**Example**
```julia-repl
julia> used_by_objective(mref)
true
```
"""
function used_by_objective(mref::MeasureRef)::Bool
    return JuMP.owner_model(mref).meas_in_objective[JuMP.index(mref)]
end

"""
    is_used(mref::MeasureRef)::Bool

Return a `Bool` indicating if `mref` is used in the model.

**Example**
```julia-repl
julia> is_used(mref)
true
```
"""
function is_used(mref::MeasureRef)::Bool
    return used_by_measure(mref) || used_by_constraint(mref) || used_by_objective(mref)
end

"""
    JuMP.delete(model::InfiniteModel, mref::MeasureRef)

Extend [`JuMP.delete`](@ref) to delete measures. Errors if measure is invalid,
meaning it does not belong to the model or it has already been deleted.

**Example**
```julia-repl
julia> print(model)
Min measure(g(t)*t) + z
Subject to
 z >= 0.0
 measure(g(t)) == 0
 g(t) + z >= 42.0
 g(0.5) == 0
 t in [0, 6]

julia> delete(model, meas)

julia> print(model)
Min z
Subject to
 z >= 0.0
 0 == 0
 g(t) + z >= 42.0
 g(0.5) == 0
 t in [0, 6]
```
"""
function JuMP.delete(model::InfiniteModel, mref::MeasureRef)
    @assert JuMP.is_valid(model, mref) "Invalid measure reference."
    # Reset the transcription status
    if is_used(mref)
        set_optimizer_model_ready(model, false)
    end
    # Remove from dependent measures if there are any
    if used_by_measure(mref)
        for mindex in model.meas_to_meas[JuMP.index(mref)]
            if isa(model.measures[mindex].func, MeasureRef)
                data = model.measures[mindex].data
                model.measures[mindex] = Measure(zero(JuMP.AffExpr), data)
            else
                _remove_variable(model.measures[mindex].func, mref)
            end
            JuMP.set_name(MeasureRef(model, mindex),
                          _make_meas_name(model.measures[mindex]))
        end
        delete!(model.meas_to_meas, JuMP.index(mref))
    end
    # Remove from dependent constraints if there are any
    if used_by_constraint(mref)
        for cindex in model.meas_to_constrs[JuMP.index(mref)]
            if isa(model.constrs[cindex].func, MeasureRef)
                model.constrs[cindex] = JuMP.ScalarConstraint(zero(JuMP.AffExpr),
                                                      model.constrs[cindex].set)
            else
                _remove_variable(model.constrs[cindex].func, mref)
            end
        end
        delete!(model.meas_to_constrs, JuMP.index(mref))
    end
    # Remove from objective if used there
    if used_by_objective(mref)
        if isa(model.objective_function, MeasureRef)
            model.objective_function = zero(JuMP.AffExpr)
        else
            _remove_variable(model.objective_function, mref)
        end
    end
    # Update that the parameters from the data are no longer dependent
    prefs = parameter_refs(measure_data(mref))
    if prefs isa ParameterRef
        prefs = [prefs]
    end
    for pref in prefs
        filter!(e -> e != JuMP.index(mref),
                model.param_to_meas[JuMP.index(pref)])
        if length(model.param_to_meas[JuMP.index(pref)]) == 0
            delete!(model.param_to_meas, JuMP.index(pref))
        end
    end
    # Update that the variables used by it are no longer used by it
    vrefs = _all_function_variables(measure_function(mref))
    for vref in vrefs
        if isa(vref, InfOptVariableRef)
            filter!(e -> e != JuMP.index(mref),
                    model.var_to_meas[JuMP.index(vref)])
            if length(model.var_to_meas[JuMP.index(vref)]) == 0
                delete!(model.var_to_meas, JuMP.index(vref))
            end
        elseif isa(vref, ParameterRef)
            filter!(e -> e != JuMP.index(mref),
                    model.param_to_meas[JuMP.index(vref)])
            if length(model.param_to_meas[JuMP.index(vref)]) == 0
                delete!(model.param_to_meas, JuMP.index(vref))
            end
        elseif isa(vref, MeasureRef)
            filter!(e -> e != JuMP.index(mref),
                    model.meas_to_meas[JuMP.index(vref)])
            if length(model.meas_to_meas[JuMP.index(vref)]) == 0
                delete!(model.meas_to_meas, JuMP.index(vref))
            end
        elseif isa(vref, ReducedInfiniteVariableRef)
            filter!(e -> e != JuMP.index(mref),
                    model.reduced_to_meas[JuMP.index(vref)])
            if length(model.reduced_to_meas[JuMP.index(vref)]) == 0
                delete!(model.reduced_to_meas, JuMP.index(vref))
            end
        end
    end
    # delete remaining measure information
    delete!(model.meas_in_objective, JuMP.index(mref))
    delete!(model.measures, JuMP.index(mref))
    delete!(model.meas_to_name, JuMP.index(mref))
    return
end


text:
"""
"""
function autodiff_cell_residual_from_energy(uh_to_cell_energy::Function,uh,cell_ids=_default_cell_ids(uh))
  @assert is_a_fe_function(uh)
  U = get_fe_space(uh)
  cell_u_to_cell_energy = _change_argument_to_cell_u(uh_to_cell_energy,U)
  cell_u = get_cell_values(uh)
  cell_r = autodiff_array_gradient(cell_u_to_cell_energy,cell_u,cell_ids)
  cell_r
end

"""
"""
function autodiff_cell_jacobian_from_energy(uh_to_cell_energy::Function,uh,cell_ids=_default_cell_ids(uh))
  @assert is_a_fe_function(uh)
  U = get_fe_space(uh)
  cell_u_to_cell_energy = _change_argument_to_cell_u(uh_to_cell_energy,U)
  cell_u = get_cell_values(uh)
  cell_j = autodiff_array_hessian(cell_u_to_cell_energy,cell_u,cell_ids)
  cell_j
end

"""
"""
function autodiff_cell_jacobian_from_residual(uh_to_cell_residual::Function,uh,cell_ids=_default_cell_ids(uh))
  @assert is_a_fe_function(uh)
  U = get_fe_space(uh)
  cell_u_to_cell_residual = _change_argument_to_cell_u(uh_to_cell_residual,U)
  cell_u = get_cell_values(uh)
  cell_j = autodiff_array_jacobian(cell_u_to_cell_residual,cell_u,cell_ids)
  cell_j
end

function _default_cell_ids(uh)
  cell_u = get_cell_values(uh)
  cell_ids = IdentityVector(length(cell_u))
end

function _change_argument_to_cell_u(uh_to_cell_energy,U)
  function f(cell_u)
    uh = CellField(U,cell_u)
    cell_e = uh_to_cell_energy(uh)
    cell_e
  end
  f
end


text:
export
    q_pdf,
    q_pdf!

real_tr_mul(𝐚, 𝐛) = sum(real(𝐚[i, :]' * 𝐛[:, i]) for i in 1:size(𝐚, 1))

"""
    q_pdf(state::AbstractState, θ::Real, x::Real; T=Float64)

Quadrature prabability at point (θ, x)
"""
function q_pdf(T::Type{<:Real}, state::AbstractState, θ::Real, x::Real)
    𝛑̂_res = Matrix{Complex{T}}(undef, state.dim, state.dim)

    return q_pdf!(𝛑̂_res, state, θ, x)
end

q_pdf(state::AbstractState, θ::Real, x::Real) = q_pdf(Float64, state, θ, x)

function q_pdf!(𝛑̂_res::AbstractMatrix, state::StateMatrix, θ::Real, x::Real)
    return real_tr_mul(𝛑̂!(𝛑̂_res, θ, x, dim=state.dim), state.𝛒)
end

function q_pdf!(𝛑̂_res::AbstractMatrix, state::StateVector, args...; kwargs...)
    return q_pdf!(𝛑̂_res, StateMatrix(state), args...; kwargs...)
end

"""
    q_pdf(state::AbstractState, θ::Real, x::Real; T=Float64)

Quadrature prabability at points (θs, xs)
"""
function q_pdf(T::Type{<:Real}, state::AbstractState, θs, xs)
    𝛑̂_res_vec = [Matrix{Complex{T}}(undef, state.dim, state.dim) for _ in 1:Threads.nthreads()]
    𝐩 = Matrix{T}(undef, length(θs), length(xs))

    return q_pdf!(𝛑̂_res_vec, 𝐩, state, θs, xs)
end

q_pdf(state::AbstractState, θs, xs) = q_pdf(Float64, state, θs, xs)

function q_pdf!(𝛑̂_res_vec::Vector{Matrix{Complex{T}}}, 𝐩::Matrix{T}, state::StateMatrix, θs, xs) where {T}
    @sync for (j, x) in enumerate(xs)
        for (i, θ) in enumerate(θs)
            Threads.@spawn 𝐩[i, j] = q_pdf!(𝛑̂_res_vec[Threads.threadid()], state, θ, x)
        end
    end

    return 𝐩
end

function q_pdf!(𝛑̂_res_vec::Vector{Matrix{Complex{T}}}, 𝐩::Matrix{T}, state::StateVector, args...; kwargs...) where {T}
    q_pdf!(𝛑̂_res_vec, 𝐩, StateMatrix(state), args...; kwargs...)
end


text:
# MIT License
#
# Copyright (c) 2018 Martin Biel
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

"""
    relative_gap(model::StochasticProgram)

Return the final relative optimality gap after a call to `optimize!(stochasticprogram)`.
"""
function JuMP.relative_gap(stochasticprogram::StochasticProgram)::Float64
    return MOI.get(stochasticprogram, MOI.RelativeGap())
end
"""
    relative_gap(model::StochasticProgram, stage::Integer, scenario_index::Integer)

Return the final relative optimality gap in the node at stage `stage` and
scenario `scenario_index` after a call to `optimize!(stochasticprogram)`.
"""
function JuMP.relative_gap(stochasticprogram::StochasticProgram, stage::Integer, scenario_index::Integer)::Float64
    attr = ScenarioDependentModelAttribute(stage, scenario_index, MOI.RelativeGap())
    return MOI.get(stochasticprogram, attr)
end
"""
    objective_bound(stochasticprogram::StochasticProgram)

Return the best known bound on the optimal objective value after a call to `optimize!(stochasticprogram)`.
"""
function JuMP.objective_bound(stochasticprogram::StochasticProgram)::Float64
    return MOI.get(stochasticprogram, MOI.ObjectiveBound())
end
"""
    objective_bound(stochasticprogram::StochasticProgram, stage::Integer, scenario_index::Integer)

Return the best known bound on the optimal objective value in the node at
stage `stage` and scenario `scenario_index` after a call to `optimize!(stochasticprogram)`.
"""
function JuMP.objective_bound(stochasticprogram::StochasticProgram, stage::Integer, scenario_index::Integer)::Float64
    attr = ScenarioDependentModelAttribute(stage, scenario_index, MOI.ObjectiveBound())
    return MOI.get(stochasticprogram, attr)
end
"""
    objective_value(stochasticprogram::StochasticProgram; result::Int = 1)

Return the objective value associated with result index `result` of the most-recent solution after a call to `optimize!(stochasticprogram)`.
"""
function JuMP.objective_value(stochasticprogram::StochasticProgram; result::Int = 1)::Float64
    # Throw NoOptimizer error if no recognized optimizer has been provided
    check_provided_optimizer(stochasticprogram.optimizer)
    # Switch on termination status
    status = JuMP.termination_status(stochasticprogram)
    if status in AcceptableTermination
        return MOI.get(stochasticprogram, MOI.ObjectiveValue(result))
    else
        if status == MOI.OPTIMIZE_NOT_CALLED
            throw(OptimizeNotCalled())
        elseif status == MOI.INFEASIBLE
            @warn "Stochastic program is infeasible"
            return objective_sense(stochasticprogram) == MOI.MAX_SENSE ? -Inf : Inf
        elseif status == MOI.DUAL_INFEASIBLE
            @warn "Stochastic program is unbounded"
            return objective_sense(stochasticprogram) == MOI.MAX_SENSE ? Inf : -Inf
        else
            @warn("Stochastic program could not be solved, returned status: $status")
            return NaN
        end
    end
end
"""
    objective_value(stochasticprogram::StochasticProgram, stage::Integer, scenario_index::Integer; result::Int = 1)

Return the objective value of the node at stage `stage` and scenario `scenario_index` associated
with result index `result` of the most-recent solution after a call to `optimize!(stochasticprogram)`.
"""
function JuMP.objective_value(stochasticprogram::StochasticProgram, stage::Integer, scenario_index::Integer; result::Int = 1)::Float64
    # Throw NoOptimizer error if no recognized optimizer has been provided
    check_provided_optimizer(stochasticprogram.optimizer)
    # Switch on termination status
    status = JuMP.termination_status(stochasticprogram, stage, scenario_index)
    if status in AcceptableTermination
        attr = ScenarioDependentModelAttribute(stage, scenario_index, MOI.ObjectiveValue(result))
        return MOI.get(stochasticprogram, attr)
    else
        if status == MOI.OPTIMIZE_NOT_CALLED
            throw(OptimizeNotCalled())
        elseif status == MOI.INFEASIBLE
            @warn "Stochastic program is infeasible in node at ($stage,$scenario_index)"
            return objective_sense(stochasticprogram, stage, scenario_index) == MOI.MAX_SENSE ? -Inf : Inf
        elseif status == MOI.DUAL_INFEASIBLE
            @warn "Stochastic program is unbounded in node at ($stage,$scenario_index)"
            return objective_sense(stochasticprogram, stage, scenario_index) == MOI.MAX_SENSE ? Inf : -Inf
        else
            @warn("Stochastic program node at ($stage,$scenario_index) could not be solved, returned status: $status")
            return NaN
        end
    end
end
"""
    objective_value(stochasticprogram::TwoStageStochasticProgram, scenario_index::Integer; result::Int = 1)

Return the objective value of scenario `scenario_index` associated with result index `result`
of the most-recent solution after a call to `optimize!(stochasticprogram)`.
"""
function JuMP.objective_value(stochasticprogram::TwoStageStochasticProgram, scenario_index::Integer; result::Int = 1)::Float64
    return objective_value(stochasticprogram, 2, scenario_index; result = result)
end
"""
    dual_objective_value(stochasticprogram::StochasticProgram; result::Int = 1)

Return the objective value of the dual problem associated with result index `result`
of the most-recent solution after a call to `optimize!(stochasticprogram)`.
"""
function JuMP.dual_objective_value(stochasticprogram::StochasticProgram; result::Int = 1)::Float64
    # Throw NoOptimizer error if no recognized optimizer has been provided
    check_provided_optimizer(stochasticprogram.optimizer)
    # Switch on termination status
    status = JuMP.termination_status(stochasticprogram)
    if status in AcceptableTermination
        return MOI.get(stochasticprogram, MOI.DualObjectiveValue(result))
    else
        if status == MOI.OPTIMIZE_NOT_CALLED
            throw(OptimizeNotCalled())
        elseif status == MOI.INFEASIBLE
            @warn "Stochastic program is infeasible"
            return objective_sense(stochasticprogram) == MOI.MAX_SENSE ? Inf : -Inf
        elseif status == MOI.DUAL_INFEASIBLE
            @warn "Stochastic program is unbounded"
            return objective_sense(stochasticprogram) == MOI.MAX_SENSE ? -Inf : Inf
        else
            @warn("Stochastic program could not be solved, returned status: $status")
            return NaN
        end
    end
end
"""
    dual_objective_value(stochasticprogram::StochasticProgram, stage::Integer, scenario_index::Integer; result::Int = 1)

Return the objective value of the dual problem of the node at stage `stage` and scenario `scenario_index` associated
with result index `result` of the most-recent solution after a call to `optimize!(stochasticprogram)`.
"""
function JuMP.dual_objective_value(stochasticprogram::StochasticProgram, stage::Integer, scenario_index::Integer; result::Int = 1)::Float64
    # Throw NoOptimizer error if no recognized optimizer has been provided
    check_provided_optimizer(stochasticprogram.optimizer)
    # Switch on termination status
    status = JuMP.termination_status(stochasticprogram, stage, scenario_index)
    if status in AcceptableTermination
        attr = ScenarioDependentModelAttribute(stage, scenario_index, MOI.DualObjectiveValue(result))
        return MOI.get(stochasticprogram, attr)
    else
        if status == MOI.OPTIMIZE_NOT_CALLED
            throw(OptimizeNotCalled())
        elseif status == MOI.INFEASIBLE
            @warn "Stochastic program is infeasible in node at ($stage,$scenario_index)"
            return objective_sense(stochasticprogram, stage, scenario_index) == MOI.MAX_SENSE ? Inf : -Inf
        elseif status == MOI.DUAL_INFEASIBLE
            @warn "Stochastic program is unbounded in node at ($stage,$scenario_index)"
            return objective_sense(stochasticprogram, stage, scenario_index) == MOI.MAX_SENSE ? -Inf : Inf
        else
            @warn("Stochastic program node at ($stage,$scenario_index) could not be solved, returned status: $status")
            return NaN
        end
    end
end
"""
    dual_objective_value(stochasticprogram::TwoStageStochasticProgram, scenario_index::Integer; result::Int = 1)

Return the objective value of the dual problem of scenario `scenario_index` associated with result index `result`
of the most-recent solution after a call to `optimize!(stochasticprogram)`.
"""
function JuMP.dual_objective_value(stochasticprogram::TwoStageStochasticProgram, scenario_index::Integer; result::Int = 1)::Float64
    return dual_objective_value(stochasticprogram, 2, scenario_index; result = result)
end
"""
    objective_sense(stochasticprogram::StochasticProgram)::MathOptInterface.OptimizationSense

Return the objective sense of the `stochasticprogram`.
"""
function JuMP.objective_sense(stochasticprogram::StochasticProgram)
    return MOI.get(stochasticprogram, MOI.ObjectiveSense())::MOI.OptimizationSense
end
"""
    objective_sense(stochasticprogram::StochasticProgram, stage::Integer)::MathOptInterface.OptimizationSense

Return the objective sense of the `stochasticprogram` stage `stage`.
"""
function JuMP.objective_sense(stochasticprogram::StochasticProgram{N}, stage::Integer) where N
    1 <= stage <= N || error("Stage $stage not in range 1 to $N.")
    return objective_sense(proxy(stochasticprogram, stage))
end
"""
    objective_sense(stochasticprogram::StochasticProgram, stage::Integer, scenario_index::Integer)::MathOptInterface.OptimizationSense

Return the objective sense in the node at stage `stage` and scenario `scenario_index`.
"""
function JuMP.objective_sense(stochasticprogram::StochasticProgram, stage::Integer, scenario_index::Integer)
    stage == 1 && error("The first-stage objective is not scenario dependent, consider `objective_sense(stochasticprogram, 1)`.")
    attr = ScenarioDependentModelAttribute(stage, scenario_index, MOI.ObjectiveSense())
    return MOI.get(stochasticprogram, attr)::MOI.OptimizationSense
end
"""
    set_objective_sense(stochasticprogram::StochasticProgram, sense::MathOptInterface.OptimizationSense)

Sets the objective sense of the `stochasticprogram` to `sense`.
"""
function JuMP.set_objective_sense(stochasticprogram::StochasticProgram, sense::MOI.OptimizationSense)
    if objective_sense(stochasticprogram) == sense
        # Nothing to do
        return nothing
    end
    # Modify proxy
    set_objective_sense(proxy(stochasticprogram, 1), sense)
    # Modify through structure
    MOI.set(structure(stochasticprogram), MOI.ObjectiveSense(), sense)
    return nothing
end
"""
    set_objective_sense(stochasticprogram::StochasticProgram, stage::Integer, sense::MathOptInterface.OptimizationSense)

Sets the objective sense of the `stochasticprogram` at stage `stage` to `sense`.
"""
function JuMP.set_objective_sense(stochasticprogram::StochasticProgram{N}, stage::Integer, sense::MOI.OptimizationSense) where N
    1 <= stage <= N || error("Stage $stage not in range 1 to $N.")
    if objective_sense(stochasticprogram, stage) == sense
        # Nothing to do
        return nothing
    end
    # Modify proxy
    set_objective_sense(proxy(stochasticprogram, stage), sense)
    # Modify through structure
    set_objective_sense(structure(stochasticprogram), stage, sense)
    return nothing
end
"""
    set_objective_sense(stochasticprogram::StochasticProgram, stage::Integer, scenario_index::Integer, sense::MathOptInterface.OptimizationSense)

Sets the objective sense of the stochasticprogram node at stage `stage`
and scenario `scenario_index` to the given `sense`.
"""
function JuMP.set_objective_sense(stochasticprogram::StochasticProgram, stage::Integer, scenario_index::Integer, sense::MOI.OptimizationSense)
    stage == 1 && error("The first-stage objective is not scenario dependent, consider `set_objective_sense(stochasticprogram, 1, sense)`.")
    attr = ScenarioDependentModelAttribute(stage, scenario_index, MOI.ObjectiveSense())
    MOI.set(structure(stochasticprogram), attr, sense)
    return nothing
end
"""
    objective_function_type(stochasticprogram::Stochasticprogram)::AbstractJuMPScalar

Return the type of the objective function of `stochasticprogram`.
"""
function JuMP.objective_function_type(stochasticprogram::StochasticProgram)
    return objective_function_type(structure(stochasticprogram))
end
"""
    objective_function_type(stochasticprogram::Stochasticprogram, stage::Integer)::AbstractJuMPScalar

Return the type of the objective function at stage `stage` of `stochasticprogram`.
"""
function JuMP.objective_function_type(stochasticprogram::StochasticProgram{N}, stage::Integer) where N
    1 <= stage <= N || error("Stage $stage not in range 1 to $N.")
    return objective_function_type(proxy(stochasticprogram, stage))
end
"""
    objective_function_type(stochasticprogram::Stochasticprogram, stage::Integer, scenario_index::Integer)::AbstractJuMPScalar

Return the type of the objective function in the node at stage `stage` and scenario `scenario_index`.
"""
function JuMP.objective_function_type(stochasticprogram::StochasticProgram, stage::Integer, scenario_index::Integer)
    stage == 1 && error("The first-stage objective is not scenario dependent, consider `objective_function_type(stochasticprogram, 1)`.")
    return objective_function_type(structure(stochasticprogram), stage, scenario_index)
end
"""
    objective_function(stochasticprogram::StochasticProgram,
                       T::Type{<:AbstractJuMPScalar}=objective_function_type(model))

Return an object of type `T` representing the full objective function of the `stochasticprogram`. Error if the objective is not convertible to type `T`.
"""
function JuMP.objective_function(stochasticprogram::StochasticProgram, FunType::Type{<:AbstractJuMPScalar} = objective_function_type(stochasticprogram))
    return objective_function(structure(stochasticprogram), FunType)
end
"""
    objective_function(stochasticprogram::StochasticProgram,
                       stage::Integer,
                       T::Type{<:AbstractJuMPScalar}=objective_function_type(model))

Return an object of type `T` representing the objective function at stage `stage` of the `stochasticprogram`. Error if the objective is not convertible to type `T`.
"""
function JuMP.objective_function(stochasticprogram::StochasticProgram{N},
                                 stage::Integer,
                                 FunType::Type{<:AbstractJuMPScalar} = objective_function_type(stochasticprogram, stage)) where N
    1 <= stage <= N || error("Stage $stage not in range 1 to $N.")
    return objective_function(structure(stochasticprogram), stage, FunType)
end
"""
    objective_function(stochasticprogram::StochasticProgram,
                       stage::Integer,
                       scenario_index::Integer,
                       T::Type{<:AbstractJuMPScalar}=objective_function_type(model))

Return an object of type `T` representing the objective function in the node at
stage `stage` and scenario `scenario_index`. Error if the objective is not convertible to type `T`.
"""
function JuMP.objective_function(stochasticprogram::StochasticProgram,
                                 stage::Integer,
                                 scenario_index::Integer,
                                 FunType::Type{<:AbstractJuMPScalar} = objective_function_type(stochasticprogram, stage, scenario_index))
    stage == 1 && error("The first-stage objective is not scenario dependent, consider `objective_function(stochasticprogram, 1)`.")
    return objective_function(structure(stochasticprogram), stage, scenario_index, FunType)
end
"""
    set_objective_coefficient(stochasticprogram::StochasticProgram, dvar::DecisionVariable, stage::Integer, coefficient::Real)

Set the linear objective coefficient associated with `dvar` to `coefficient` in stage `stage`.
"""
function JuMP.set_objective_coefficient(stochasticprogram::StochasticProgram, dvar::DecisionVariable, stage::Integer, coeff::Real)
    # Modify proxy
    proxy_ = proxy(stochasticprogram, stage)
    dref = DecisionRef(proxy_, index(dvar))
    set_objective_coefficient(proxy_, dref, coeff)
    # Modify objective through structure
    set_objective_coefficient(structure(stochasticprogram), index(dvar), StochasticPrograms.stage(dvar), stage, coeff)
end
"""
    set_objective_coefficient(stochasticprogram::StochasticProgram, dvar::DecisionVariable, stage::Integer, scenario_index::Integer, coefficient::Real)

Set the scenario-dependent linear objective coefficient at `scenario_index` associated with `dvar` to `coefficient` in stage `stage`.
"""
function JuMP.set_objective_coefficient(stochasticprogram::StochasticProgram, dvar::DecisionVariable, stage::Integer, scenario_index::Integer, coeff::Real)
    stage == 1 && error("The first-stage objective is not scenario dependent, consider `set_objective_coefficient(stochasticprogram, dvar, 1, coeff)`.")
    # Modify objective through structure
    set_objective_coefficient(structure(stochasticprogram), index(dvar), StochasticPrograms.stage(dvar), stage, scenario_index, coeff)
end


text:
using Logging

Logging.configure(output=STDOUT, level=Logging.INFO)

# For now just have the build script tell the user to add fox to their path.
# Later it could try adding it to /usr/local/bin or a user directory on the path if they
# exist.

Logging.info(
    "A script for test automation with julia playgrounds is located in deps/usr/bin/fox.\n" *
    "If you would like to make your script callable from anywhere please add \n" *
    "deps/usr/bin to your path or symlink it into a folder already on your path such\n" *
    "a local bin directory."
)


text:
module GroupSlices ## registered version is for Julia v0.5; this has trivial updates

import Base.hash
import Base.Cartesian, Base.Cartesian.@nloops, Base.Cartesian.@nref

export groupslices, groupinds, firstinds, lastinds

struct Prehashed
    hash::UInt
end
hash(x::Prehashed) = x.hash

"""
    groupslices(A, dim)
Returns a vector of integers where each integer element of the returned vector
is a group number corresponding to the unique slices along dimension `dim` as
returned from `unique(A, dim)`, where `A` can be a multidimensional array.

Example usage:

If `C = unique(A, dim)`, `ic = groupslices(A, dim)`, and
`ndims(A) == ndims(C) == 3`, then:
```
if dim == 1
   all(A .== C[ic,:,:])
elseif dim == 2
   all(A .== C[:,ic,:])
elseif dim == 3
   all(A .== C[:,:,ic])
end
```
"""
@generated function groupslices(A::AbstractArray{T,N}, dim::Int) where {T,N}
    quote
        if !(1 <= dim <= $N)
            ArgumentError("Input argument dim must be 1 <= dim <= $N, but is currently $dim")
        end
        hashes = zeros(UInt, size(A, dim))

        # Compute hash for each row
        k = 0
        @nloops $N i A d->(if d == dim; k = i_d; end) begin
            @inbounds hashes[k] = hash(hashes[k], hash((@nref $N A i)))
        end

        # Collect index of first row for each hash
        uniquerow = Array{Int}(undef,size(A, dim))
        firstrow = Dict{Prehashed,Int}()
        for k = 1:size(A, dim)
            uniquerow[k] = get!(firstrow, Prehashed(hashes[k]), k)
        end
        uniquerows = collect(values(firstrow))

        # Check for collisions
        collided = falses(size(A, dim))
        @inbounds begin
            @nloops $N i A d->(if d == dim
                k = i_d
                j_d = uniquerow[k]
            else
                j_d = i_d
            end) begin
                if (@nref $N A j) != (@nref $N A i)
                    collided[k] = true
                end
            end
        end

        if any(collided)
            nowcollided = BitArray(undef, size(A, dim))
            while any(collided)
                # Collect index of first row for each collided hash
                empty!(firstrow)
                for j = 1:size(A, dim)
                    collided[j] || continue
                    uniquerow[j] = get!(firstrow, Prehashed(hashes[j]), j)
                end
                for v in values(firstrow)
                    push!(uniquerows, v)
                end

                # Check for collisions
                fill!(nowcollided, false)
                @nloops $N i A d->begin
                    if d == dim
                        k = i_d
                        j_d = uniquerow[k]
                        (!collided[k] || j_d == k) && continue
                    else
                        j_d = i_d
                    end
                end begin
                    if (@nref $N A j) != (@nref $N A i)
                        nowcollided[k] = true
                    end
                end
                (collided, nowcollided) = (nowcollided, collided)
            end
        end
        ie = unique(uniquerow)
        ic_dict = Dict{Int,Int}()
        for k = 1:length(ie)
            ic_dict[ie[k]] = k
        end

        ic = similar(uniquerow)
        for k = 1:length(ic)
            ic[k] = ie[ic_dict[uniquerow[k]]]
        end
        return ic
    end
end

"""
    groupinds(ic)
Returns a vector of vectors of integers wherein the vector of group slice
index integers as returned from `groupslices(A, dim)` is converted into a
grouped vector of vectors.  Each vector entry in the returned vector of
vectors contains all of the positional indices of slices in the original
input array `A` that correspond to the unique slices along dimension `dim`
that are present in the array `C` as returned from `unique(A, dim)`.
"""
function groupinds(ic::Vector{Int})
    d = Dict{Int, Int}()
    ia = unique(ic)
    n = length(ia)
    for i = 1:n
        d[ia[i]]= i
    end

    ib = Array{Vector{Int}}(undef, n)
    for k = 1:n
        ib[k] = Int[]
    end

    for h = 1:length(ic)
        push!(ib[d[ic[h]]], h)
    end
    return ib
end

"""
    firstinds(ic::Vector{Int})
    firstinds(ib::Vector{Vector{Int}})

Returns a vector of integers containing the first index position of each unique
value in the input integer vector `ic`, or the first index position of each
entry in the input vector of integer vectors `ib`.

When operating on the output returned from `unique(A, dim)`, the returned
vector of integers correspond to the positions of the first of each unique slice
present in the original input multidimensional array `A` along dimension `dim`.

The implementation of `firstinds` accepting a vector of integers operates on the
output returned from `groupslices(A, dim)`.

The implementation of `firstinds` accepting a vector of vector of integers
operates on the output returned from `groupinds(ic::Vector{Int})`.
"""
function firstinds(ic::Vector{Int})
    id = unique(ic)
    n = length(id)
    ia = Array{Int}(undef,n)
    for i = 1:n
        ia[i] = something(findfirst(isequal(id[i]), ic), 0) # findfirst(ic, id[i])
    end
    return ia
end

function firstinds(ib::Vector{Vector{Int}})
    ia = map(first, ib)
end

"""
    lastinds(ic::Vector{Int})
Returns a vector of integers containing the last index position of each unique
value in the input integer vector `ic`.

When operating on the output returned from `groupinds(unique(A, dim))`, the
returned vector of integers correspond to the positions of the last of each
unique slice present in the original input multidimensional array `A` along
dimension `dim`.

The implementation of `firstinds` accepting a vector of vector of integers
operates on the output returned from `groupinds(ic::Vector{Int})`.
"""
function lastinds(ib::Vector{Vector{Int}})
    ia = map(last, ib)
end

end # module


text:
using Zygote, Test
using Zygote: gradient

@testset "Zygote" begin

@testset "Features" begin
  include("features.jl")
end

@testset "Gradients" begin
  include("gradcheck.jl")
end

@testset "Compiler" begin
  include("compiler.jl")
end

end


text:
module EmpiricalBayes

using Distributions
using NLopt

include("model.jl")
include("optim.jl")
include("calculations.jl")

export LikelihoodModel
export npmle, mple, refprior
export OptConfig

end


text:
import Base: show

"""

Engines for Hamiltonian Monte Carlo

$(TYPEDEF)

### Engines defined
```julia
* `Nuts`                               : No-U-Turn sampler
* `Static`                             : Static integration time
```
Not exported.
"""
abstract type Engine end

"""

Sampling algorithms

$(TYPEDEF)

### Sampling algorithms defined
```julia
* `HMC`                                : Hamiltonian Monte Carlo sampler
```
Not exported.
"""
abstract type SamplingAlgorithm end

struct Nuts <: Engine
  max_depth::Int64
end

"""

Nuts constructor

$(SIGNATURES)

### Optional arguments
```julia
* `max_depth::Number`                  : Maximum tree depth
```
Not exported
"""
Nuts(;max_depth::Number=10) = Nuts(max_depth)

struct Static <: Engine
  int_time::Float64
end

"""

Static constructor

$(SIGNATURES)

### Optional arguments
```julia
* `;int_time::Number`                  : Static integration time
```
Not exported
"""
Static(;int_time::Number=2 * pi) = Static(int_time)

"""

Metric types for geometry of base manifold

### Types defined
```julia
* unit_e::Metric                       : Euclidean manifold with unit metric
* dense_e::Metric                      : Euclidean manifold with dense netric
* diag_e::Metric                       : Euclidean manifold with diag netric
```
Not exported.
"""
abstract type Metric end
struct unit_e <: Metric
end
struct dense_e <: Metric
end
struct diag_e <: Metric
end

struct Hmc <: SamplingAlgorithm
  engine::Engine
  metric::Metric
  stepsize::Float64
  stepsize_jitter::Float64
end

Hmc(;engine::Engine=Nuts(), metric::DataType=diag_e,
  stepsize::Number=1.0, stepsize_jitter::Number=1.0) =
    Hmc(engine, metric(), stepsize, stepsize_jitter)

"""

Hmc constructor

$(SIGNATURES)

# Extended help

### Method
```julia
Hmc(;
  engine=StanSample.Nuts(),
  metric=StanSample.diag_e,
  stepsize=1.0,
  stepsize_jitter=1.0
)
```
### Optional arguments
```julia
* `engine::Engine`                     : Engine for Hamiltonian Monte Carlo
* `metric::Metric`                     : Geometry for base manifold
* `stepsize::Float64`                  : Stepsize for discrete evolutions
* `stepsize_jitter::Float64`           : Uniform random jitter of the stepsize [%]
```
Not exported
"""
Hmc(engine::Engine) = Hmc(engine, diag_e(), 1.0, 1.0)

"""

Fixed_param type and constructor

$(TYPEDEF)

### Method
```julia
Fixed_param()
```
Not exported
"""
struct Fixed_param <: SamplingAlgorithm end

struct Adapt
  engaged::Bool
  gamma::Float64
  delta::Float64
  kappa::Float64
  t0::Float64
  init_buffer::Int64
  term_buffer::Int64
  window::Int64
end

"""

# Adapt constructor

$(SIGNATURES)

# Extended help

Settings for adapt=CmdStan.Adapt() in Sample().

### Method
```julia
Adapt(;
  engaged=true,
  gamma=0.05,
  delta=0.8,
  kappa=0.75,
  t0=10.0,
  init_buffer=75,
  term_buffer=50,
  window::25
)
```
### Optional arguments
```julia
* `engaged::Bool`              : Adaptation engaged?
* `gamma::Float64`             : Adaptation regularization scale
* `delta::Float64`             : Adaptation target acceptance statistic
* `kappa::Float64`             : Adaptation relaxation exponent
* `t0::Float64`                : Adaptation iteration offset
* `init_buffer::Int64`         : Width of initial adaptation interval
* `term_buffer::Int64`         : Width of final fast adaptation interval
* `window::Int64`              : Initial width of slow adaptation interval
```

Not exported
"""
Adapt(;engaged::Bool=true, gamma::Number=0.05, delta::Number=0.8,
  kappa::Number=0.75, t0::Number=10.0,
  init_buffer::Number=75, term_buffer::Number=50, window::Number=25) =
    Adapt(engaged, gamma, delta, kappa, t0, init_buffer, term_buffer, window)

struct Sample
  num_samples::Int64
  num_warmup::Int64
  save_warmup::Bool
  thin::Int64
  adapt::Adapt
  algorithm::SamplingAlgorithm
end

"""
Settings for method=Sample() in SampleModel.

$(SIGNATURES)

# Extended help

### Method
```julia
Sample(;
  num_samples=1000,
  num_warmup=1000,
  save_warmup=false,
  thin=1,
  adapt=StanSample.Adapt(),
  algorithm=SamplingAlgorithm()
)
```
### Optional arguments
```julia
* `num_samples::Int64`          : Number of sampling iterations ( >= 0 )
* `num_warmup::Int64`           : Number of warmup iterations ( >= 0 )
* `save_warmup::Bool`           : Include warmup samples in output
* `thin::Int64`                 : Period between saved samples
* `adapt::Adapt`                : Warmup adaptation settings
* `algorithm::SamplingAlgorithm`: Sampling algorithm
```
Not exported
"""
Sample(;num_samples::Number=1000, num_warmup::Number=1000,
  save_warmup::Bool=false, thin::Number=1,
  adapt::Adapt=Adapt(), algorithm::SamplingAlgorithm=Hmc()) =
    Sample(num_samples, num_warmup, save_warmup, thin, adapt, algorithm)


text:
using Flux
using OceanParameterizations
using BSON
using OrdinaryDiffEq, DiffEqSensitivity
using WindMixing

# data in which the neural network is trained on
train_files = ["strong_wind"]

PATH = pwd()
OUTPUT_PATH = joinpath(PATH, "Output")
𝒟train = data(train_files, scale_type=ZeroMeanUnitVarianceScaling, enforce_surface_fluxes=true)

uvT = 𝒟train.uvT_scaled
uw = 𝒟train.uw.scaled
vw = 𝒟train.vw.scaled
wT = 𝒟train.wT.scaled

N_inputs = 96
hidden_units = 400
N_outputs = 31
uw_NN = Chain(Dense(N_inputs, hidden_units, relu), Dense(hidden_units, hidden_units, relu), Dense(hidden_units, N_outputs))
vw_NN = Chain(Dense(N_inputs, hidden_units, relu), Dense(hidden_units, hidden_units, relu), Dense(hidden_units, N_outputs))
wT_NN = Chain(Dense(N_inputs, hidden_units, relu), Dense(hidden_units, hidden_units, relu), Dense(hidden_units, N_outputs))

_, re_uw = Flux.destructure(uw_NN)
_, re_vw = Flux.destructure(vw_NN)
_, re_wT = Flux.destructure(wT_NN)

optimizers_NN = [ADAM(0.01)]
epochs_NN = 10

uw_weights = train_NN(uw_NN, uvT, uw, optimizers_NN, epochs_NN, OUTPUT_PATH, "uw_weights_large_NN")
vw_weights = train_NN(vw_NN, uvT, vw, optimizers_NN, epochs_NN, OUTPUT_PATH, "vw_weights_large_NN")
wT_weights = train_NN(wT_NN, uvT, wT, optimizers_NN, epochs_NN, OUTPUT_PATH, "wT_weights_large_NN")

uw_NN = re_uw(uw_weights)
vw_NN = re_vw(vw_weights)
wT_NN = re_wT(wT_weights)

train_tranges_NDE = [1:5:50, 1:5:100]
optimizers_NDE = [ADAM(0.01)]
epochs_NDE = 100

train_NDE_convective_adjustment(uw_NN, vw_NN, wT_NN, 𝒟train, 1:5:100, ROCK4(), optimizers_NDE, epochs_NDE, OUTPUT_PATH, "weights_SWNH_convective_adjustment_100_large")