text:
abstract AbstractODE

"""
Explicitly defined ODE of form dy = F(t,y).

Fields:

- t0, y0: initial conditions
- F!: ODE function `F!(t,y,dy)` which modifies `dy` in-place
- jac!: TODO
"""
immutable ExplicitODE{T,S} <: AbstractODE
    t0  ::T
    y0  ::S
    F!  ::Function
    jac!::Function
end

# TODO: change (t,y,J)->fdjacobian(F!,t,y,J) to fdjacobian!(F!)
ExplicitODE{T,S<:AbstractVector}(t0::T, y0::S, F!::Function;
               jac!::Function = (t,y,J)->fdjacobian!(F!,t,y,J)) =
                   ExplicitODE{T,S}(t0,y0,F!,jac!)

"""

This type is not yet implemented, but will serve as an implicitly
defined ODE (i.e. ODE of the form F(t,y,y')=0.

"""
immutable ImplicitODE{T,S} <: AbstractODE
end

"""

Convert a out-of-place explicitly defined ODE function to an in-place function.

Note, this does not help with memory allocations.

"""
function explicit_ineff(t0, y0::AbstractVector, F::Function;
                        jac = (t,y)->fdjacobian(F,t,y))
    function F!(t,y,dy)
        copy!(dy,F(t,y))
    end
    function jac!(t,y,J)
        copy!(J,jac(t,y))
    end
    return ExplicitODE(t0,y0,F!,jac!)
end

# A temporary solution for handling scalars, should be faster then the
# previous implementation.  Should be used only at the top level
# interface.  This function cheats by converting scalar functions F
# and jac to vector functions F! and jac!.  Still, solving this ODE
# will result in a vector of length one result, so additional external
# conversion is necessary.
function explicit_ineff(t0, y0::Number, F::Function;
                        jac = (t,y)->fdjacobian(F,t,y))
    function F!(t,y,dy)
        dy[1]=F(t,y[1])
    end
    function jac!(t,y,J)
        J[1]=jac(t,y[1])
    end
    return ExplicitODE(t0,[y0],F!,jac!)
end

"""

The abstract type of the actual algorithm to solve an ODE.

"""
abstract AbstractStepper

"""

AbstractState keeps the temporary data (state) for the iterator
Solver{::AbstractStepper}.

"""
abstract AbstractState

# m3:
# - docs
# - maybe use the typevars as defined in make_consistent_types for t,
#   y, dy?  T->Et, S->Ty
#   (or something else consistent throughout, maybe nicer would be all
#   uppercase: ET, EFY, TT, TY).
# - if find `Step` a bit confusing name, in particular combined with
#   AbstractStepper, but not sure what's better.

"""

Holds a value of a function and its derivative at time t.  This is
usually used to store the solution of an ODE at particular times.

"""
type Step{T,S}
    t ::T
    y ::S
    dy::S
end

function show(io::IO, state::Step)
    println("t  =$(state.t)")
    println("y  =$(state.y)")
    println("dy =$(state.dy)")
end

"""

Options for ODE solvers.  This type has a key-word constructor which
will fill the structure with default values.

General:

- initstep ::T  initial step
- tstop    ::T  end integration time
- reltol   ::T  relative tolerance (m3: could this be a vector?)
- abstol   ::T  absolute tolerance (m3: could this be a vector?)
- minstep  ::T  minimal allowed step
- maxstep  ::T  maximal allowed step
- norm           function to calculate the norm in step control
- maxiters ::T  maximum number of steps
- isoutofdomain::Function checks if the solution became non-numeric (NaN or Inf)

Dense output options:

- tspan    ::Vector{T}  output times
- points   ::Symbol which points are returned: `:specified` only the
  ones in tspan or `:all` which includes also the step-points of the solver.
- stopevent  Stop integration at a zero of this function
- roottol    TODO

"""
type Options{T} # m3:  T->Et
    # stepper options
    initstep ::T
    tstop    ::T
    reltol   ::T
    abstol   ::T
    minstep  ::T
    maxstep  ::T
    norm     ::Function
    maxiters ::T

    isoutofdomain::Function

    # dense output options
    tspan    ::Vector{T}
    points   ::Symbol

    # m3: I think this should be an array of functions.  Depending on some
    # flag each one returns, the iteration stops or continues.  Rename it
    # to eventfns.  I like matlabs interface.
    # [value,isterminal,direction] = myEventsFcn(t,y,dy)
    # The value gets stored.
    stopevent::Function
    roottol  ::T

    function Options(;
                     tspan    = T[Inf],
                     tstop    = tspan[end],
                     reltol   = eps(T)^T(1//3)/10,
                     abstol   = eps(T)^T(1//2)/10,
                     minstep  = 10*eps(T),
                     maxstep  = 1/minstep,
                     # TODO: we need a better guess here, possibly
                     # overwrite it in the call to solve()
                     initstep = max(min(reltol,abstol,maxstep),minstep),
                     norm     = Base.norm,
                     maxiters = T(Inf),
                     points   = :all,
                     stopevent = (t,y)->false,
                     roottol  = eps(T)^T(1//3),
                     isoutofdomain = isnan,
                     kargs...)
        if all(points .!= [:specified,:all])
            error("Option points = $points is not supported, use :specified or :all")
        end
        #TODO iterate over fields here?
        new(initstep,tstop,reltol,abstol,minstep,maxstep,norm,maxiters,isoutofdomain,sort(tspan),points,stopevent,roottol)
    end

end

function show{T}(io::IO, opts :: Options{T})
    for name in fieldnames(opts)
        @printf("%-20s = %s\n",name,getfield(opts,name))
    end
end

"""

This is an iterable type, each call to next(...) produces a next step
of a numerical solution to an ODE.

- ode: is the prescrived ode, along with the initial data
- stepper: the algorithm used to produce subsequent steps
- options: options passed to the stepper

"""
type Solver{T<:AbstractStepper} # TODO: immutable?
    ode     :: AbstractODE
    stepper :: T
    options :: Options
end

# TODO: is this the right way to implement the mid level interface?
solve(ode, stepper; kargs...) = solve(ode, stepper, Options(kargs...))

# filter the wrong combinations of ode and stepper
solve{T,S}(ode :: T, stepper :: S, options :: Options) = error("The $S doesn't support $T")

# normally we return the working array, which changes at each step and
# expect the user to copy it if necessary.  In order for collect to
# return the expected result we need to copy the output at each step.
function collect{T,S}(t::Type{Tuple{T,S}}, s::Solver)
    if maximum(s.options.tspan) == T(Inf)
        error("Attempting to collect an infinite list, use tstop or tspan with finite numbers only")
    end
    collect(t, imap(x->deepcopy(x),s))
end

function collect(s::Solver)
    if maximum(s.options.tspan) == Inf
        error("Attempting to collect an infinite list, use tstop or tspan with finite numbers only")
    end
    collect(imap(x->deepcopy(x),s))
end

# some leftovers from the previous implementation

# FIXME: This doesn't really work if x is anything but a Vector or a scalar
function fdjacobian(F, t, x::Number)
    ftx = F(t, x)

    # The 100 below is heuristic
    dx = (x .+ (x==0))./100
    dFdx = (F(t,x+dx)-ftx)./dx

    return dFdx
end

function fdjacobian(F, t, x::Vector)
    ftx = F(t, x)
    lx = max(length(x),1)
    dFdx = zeros(eltype(x), lx, lx)
    for j = 1:lx
        # The 100 below is heuristic
        dx = zeros(eltype(x), lx)
        dx[j] = (x[j] .+ (x[j]==0))./100
        dFdx[:,j] = (F(t,x+dx)-ftx)./dx[j]
    end
    return dFdx
end

function fdjacobian!{T}(F!, t, x::Vector{T}, J::Array{T,2})
    ftx = similar(x)
    ftx2= similar(x)
    dx  = similar(x)
    F!(t,x,ftx)
    lx = max(length(x),1)
    dFdx = zeros(eltype(x), lx, lx)
    for j = 1:lx
        # The 100 below is heuristic
        dx[j] = (x[j] .+ (x[j]==0))./100
        F!(t,x+dx,ftx2)
        J[:,j] = (ftx2-ftx)./dx[j]
    end
end


text:
module MyLearn2Classify

using IJulia, Plots
using Random: seed!, randperm

# these are the functions that will be immediately accessible after `using MyLearn2Classify`
export linear, dlinear, dtanh, sigmoid, dsigmoid, gn, grad_loss, learn2classify_asgd

# Activation functions
linear(z) = z
dlinear(z) = 1.0

# tanh is already defined in Julia
dtanh(z) = 1 - tanh(z)^2

sigmoid(z::Real) = 1.0 / (1.0 + exp(-z))
function dsigmoid(z::Real)
    sigmoid_z = 1.0 / (1.0 + exp(-z))
    return sigmoid_z * (1 - sigmoid_z)
end

function gn(x::AbstractArray, w::AbstractVector, b::Number, f_a::Function)
    return f_a.(x' * w .+ b)
end

function grad_loss(
        f_a::Function,
        df_a::Function,
        x::AbstractMatrix,
        y::AbstractVector,
        w::AbstractVector,
        b::Number,
        normalize::Bool=true
    )
    dw = zeros(length(w))
    db = 0.0
    loss = 0.0
    for j in 1:size(x, 2)
        error = y[j] - f_a.(w' * x[:, j] + b)
        common_term = error .* df_a.(w' * x[:, j] + b)
        dw = dw .- 2.0 * common_term .* x[:, j]
        db = db .- 2.0 * common_term
        loss += error^2
    end

    if normalize
        dw /= length(y)
        db /= length(y)
        loss /= length(y)
    end

    return dw, db, loss
end

function learn2classify_asgd(
        f_a::Function,
        df_a::Function,
        grad_loss::Function,
        x::AbstractMatrix,
        y::AbstractVector,
        mu::Number=1e-3,
        iters::Integer=500,
        batch_size::Integer=10,
        show_loss::Bool=true,
        normalize::Bool=true,
        seed::Integer=1
    )
    n, N = size(x)

    if seed == false
        b = 0.0
        w = zeros(n)
    else
        seed!(seed) # initialize random number generator
        w = randn(n)
        b = rand()
    end

    loss = zeros(iters)
    lambdak = 0
    qk = w
    pk = b
    for i in 1:iters
        batch_idx = randperm(N)
        batch_idx = batch_idx[1:min(batch_size, N)]

        dw, db, loss_i = grad_loss(f_a, df_a, x[:, batch_idx], y[batch_idx], w, b, normalize)
        qkp1 = w - mu * dw
        pkp1 = b - mu * db

        lambdakp1 = (1 + sqrt(1 + 4 * lambdak^2)) / 2
        gammak = (1 - lambdak) / lambdakp1

        w = (1 - gammak) * qkp1 + gammak * qk
        b = (1 - gammak) * pkp1 + gammak * pk

        qk = qkp1
        pk = pkp1
        lambdak = lambdakp1

        loss[i] = convert(Float64, loss_i[1])

        if show_loss && (rem(i, 100) == 0)
            IJulia.clear_output(true)
            loss_plot = scatter(
                [1:50:i], loss[1:50:i], yscale=:log10,
                xlabel="iteration",
                ylabel="training loss",
                title="iteration $i, loss = $loss_i"
            )
            display(loss_plot)
        end
    end

    return w, b, loss
end

end # module

ext: jl

max_stars_repo_name: jevogel/ml

max_issues_repo_name: jevogel/ml

max_forks_repo_name: jevogel/ml


text:
## Assumptions
## http://docs.sympy.org/0.7.2/modules/assumptions/index.html

"""
     refine

Simplify an expression using assumptions; [refine](https://docs.sympy.org/dev/modules/assumptions/refine.html).
"""
refine(ex, assumpts...) = sympy.refine(ex, assumpts...)
export refine

"""
    ask(query)

Returns `true`, `false`, or `nothing`; [ask](https://docs.sympy.org/dev/modules/assumptions/ask.html)

Example:

```jldoctest ask
julia> using SymPy

julia> @vars x y integer=true
(x, y)

julia> ask(𝑄.integer(x*y), And(𝑄.integer(x), 𝑄.integer(y)))
true

julia> ## really slow isprime:
       filter(x -> ask(𝑄.prime(x)), 1:10)
4-element Vector{Int64}:
 2
 3
 5
 7
```

"""
ask(x::Sym, args...) = sympy.ask(x, args...)
ask(x::Bool, args...) = x
ask(x::Nothing, args...) = x
export ask

## should we support & and | for (sym,sym) pairs? Not sure
## dependso on what happens to | (x==0) and ex == x usage for ==
## for now, we can combine terms logically with And, Or, Not...
## these are in logic module

## We make a module Q to hold the assumptions
## this follows this page http://docs.sympy.org/0.7.5/_modules/sympy/assumptions/ask.html
"""
    𝑄
    SymPy.Q

Documentation for the `SymPy.Q` module, exported as `𝑄`.

SymPy allows for
[assumptions](https://docs.sympy.org/latest/modules/assumptions/index.html)
on variables. These may be placed on free sympols at construction.

For example, the following creates a real value variable `x` and a postive, real variable `y`:

```jldoctest 𝑄
julia> using SymPy

julia> @vars x real=true
(x,)

julia> @vars y real=true positive=true
(y,)
```

The `Q` module exposes a means to *q*uery the assumptions on a
variable. For example,

```jldoctest 𝑄
julia> ask(𝑄.positive(y))  # true
true

julia> ask(𝑄.negative(y))  # false
false

julia> ask(SymPy.Q.positive(x))  # `nothing`

julia> ask(SymPy.Q.positive(x^2)) # `nothing` -- might be 0

julia> ask(SymPy.Q.positive(1 + x^2)) # true  -- must be postive now.
true
```

The ask function uses tri-state logic, returning one of 3 values:
`true`; `false`; or `nothing`, when the query is indeterminate.

The construction of predicates is done through `Q` methods. These can
be combined logically. For example, this will be `true`:

```jldoctest 𝑄
julia> ask(𝑄.positive(y) & 𝑄.negative(-x^2 - 1))
true
```

The above use `&` as an infix operation for the binary operator
`And`. Values can also be combined with `Or`, `Not`, `Xor`, `Nand`,
`Nor`, `Implies`, `Equivalent`, and `satisfiable`.

!!! note "Matrix predicates"
    As `SymPy.jl` converts symbolic matrices into Julia's `Array`
type and not as matrices within Python, the predicate functions from SymPy for
matrices are not used, though a replacement is given.
"""
module Q
import SymPy
import PyCall
import LinearAlgebra: det, norm

##http://docs.sympy.org/dev/_modules/sympy/assumptions/ask.html#ask
Q_predicates = (:antihermitian,
                :bounded, :finite, # bounded deprecated
                :commutative,
                :complex,
                :composite,
                :even,
                :extended_real,
                :hermitian,
                :imaginary,
                :infinitesimal,
                :infinity, :infinite, # :infinity deprecated
                :integer,
                :irrational,
                :rational,
                :algebraic,
                :transcendental,
                :negative,
                :nonzero, :zero,
                :positive,
                :prime,
                :real,
                :odd,
                :is_true,
                :nonpositive,
                :nonnegative
#                :symmetric,
#                :invertible,
#                :singular,
#                :orthogonal,
#                :unitary,
#                :normal,
#                :positive_definite,
#                :upper_triangular,
#                :lower_triangular,
#                :diagonal,
#                :triangular,
#                :unit_triangular,
#                :fullrank,
#                :square,
#                :real_elements,
#                :complex_elements,
#                :integer_elements
)

for meth in Q_predicates
   nm = string(meth)
      @eval begin
            ($meth)(x) = PyCall.pycall(SymPy.sympy.Q.$nm, SymPy.Sym, x)::SymPy.Sym
   end
end

symmetric(M::Array{T,2}) where {T <: SymPy.Sym} = SymPy.issymmetric(M)
function invertible(M::Array{T,2}) where {T <: SymPy.Sym}
    d = det(M)
    pos = SymPy.ask(positive(d))
    if pos == nothing
        return nothing
    elseif pos == true
        return true
    end
    neg = SymPy.ask(negative(d))
    if neg == nothing
        return nothing
    end
    z = SymPy.ask(zero(d))
    if z == nothing
        return nothing
    elseif z == true
        return false
    end

    return true
end

function singular(M::Array{T,2}) where {T <: SymPy.Sym}
    !invertible(M)
end

function orthogonal(M::Array{T,2}) where {T <: SymPy.Sym}
    vals = SymPy.simplify.(SymPy.simplify.(M*transpose(M)) .== one(T))
    no_nothing = 0
    for val in vals
        a = SymPy.ask(zero(val))
        if a == nothing
            no_nothing += 1
        elseif a == false
            return false
        end
    end

    no_nothing > 0 && return nothing
    return true
end

function unitary(M::Array{T,2}) where {T <: SymPy.Sym}
    vals = SymPy.simplify.(SymPy.simplify.(M*ctranspose(M)) .== one(T))
    no_nothing = 0
    for val in vals
        a = SymPy.ask(zero(val))
        if a == nothing
            no_nothing += 1
        elseif a == false
            return false
        end
    end

    no_nothing > 0 && return nothing
    return true
end

function normal(M::Array{T,2}) where {T <: SymPy.Sym}
    lhs = ctranspose(M) * M
    rhs = M * ctranspose(M)
    vals = zero.(SymPy.simplify.(lhs - rhs))
    no_nothing = 0
    for val in vals
        a = SymPy.ask(val)
        if a == nothing
            no_nothing += 1
        elseif a == false
            return false
        end
    end

    no_nothing > 0 && return nothing
    return true
end

# Use [Sylvester's](https://en.wikipedia.org/wiki/Sylvester%27s_criterion) Criteria
function positive_definite(M::Array{T,2}) where {T <: SymPy.Sym}
    !SymPy.ask(square(M)) && return false
    !SymPy.ask(symmetric(M))  && return false
    m, n = size(M)
    no_false = 0
    no_nothing = 0
    for i in 1:m
        a = SymPy.ask(Q.positive(det(M[1:i, 1:i])))
        if a == nothing no_nothing += 1 end
        if a == false no_false += 1 end
    end
    if no_false > 0
        return false
    elseif no_nothing > 0
        return nothing
    else
        return true
    end
end

upper_triangular(M::Array{T,2}) where {T <: SymPy.Sym} = SymPy.istriu(M)
lower_triangular(M::Array{T,2}) where {T <: SymPy.Sym} = SymPy.istril(M)
diagonal(M::Array{T,2}) where {T <: SymPy.Sym} = upper_triangular(M) && lower_triangular(M)
triangular(M::Array{T,2}) where {T <: SymPy.Sym} = upper_triangular(M) || lower_triangular(M)

## This is likely not the best way as it is a bit fidgety due
## to the call to rref.
function full_rank(M::Array{T,2}) where {T <: SymPy.Sym}
    m,n = size(M)
    m <= n || return full_rank(transpose(M))

    rr, p = SymPy.rref(M)
    lr = rr[end, :] # is this zero?
    no_nothing = 0
    no_nonzero = 0
    for val in lr
        a = SymPy.ask(nonzero(val))
        if a == nothing
            no_nothing += 1
        end
        if a == true
            no_nonzero += 1
        end
    end
    if no_nothing > 0
        return nothing
    elseif no_nonzero == 0
        return false
    else
        return true
    end

end

function square(M::Array{T,2}) where {T <: SymPy.Sym}
    m,n = SymPy.size(M)
    m == n
end

function real_elements(M::Array{T,2}) where {T <: SymPy.Sym}
    vals = real.(M)
    for val in vals
        a = SymPy.ask(real(val))
        (a == nothing || a == false) && return false
    end
    return true
end

function complex_elements(M::Array{T,2}) where {T <: SymPy.Sym}
    vals = real.(M)
    for val in vals
        a = SymPy.ask(SymPy.sympy."Q".complex(val))
        (a == nothing || a == false) && return false
    end
    return true

end

function integer_elements(M::Array{T,2}) where {T <: SymPy.Sym}
    vals = real.(M)
    for val in vals
        a = SymPy.ask(integer(val))
        (a == nothing || a == false) && return false
    end
    return true
end

end

## Issue  #354; request to *not*  export  Q
## export
#export Q

const 𝑄 = Q

"""
    𝑄

Exported  symbol for  [`SymPy.Q`](@ref), a  Julia  module implementing `sympy.Q`. "Questions" can be asked through the patterns
`𝑄.query(value)` (𝑄 is entered as  [slash]itQ[tab]) or `SymPy.Q.query(value)` *but  not* as `sympy.Q.query(value)`

!!! note
    At one time, the symbol `Q` was exported for this. To avoid namespace clutter, the unicode alternative is now used. Legacy code would need a definition like `const Q = SymPy.Q`  to work.

"""
𝑄
export  𝑄


text:
# Note that this script can accept some limited command-line arguments, run
# `julia build_tarballs.jl --help` to see a usage message.
using BinaryBuilder, Pkg

name = "HiGHS"

version = v"1.2.1"

sources = [
    GitSource(
        "https://github.com/ERGO-Code/HiGHS.git",
        "60ad175979bbf6be309de3992adc9852625d1e7b",
    ),
]

# These are the platforms we will build for by default, unless further
# platforms are passed in on the command line
platforms = supported_platforms()

function build_script(; shared_libs::String)
    return "BUILD_SHARED=$(shared_libs)\n" * raw"""
cd $WORKSPACE/srcdir

mkdir -p HiGHS/build
cd HiGHS/build

# Do fully static build only on Windows
if [[ "${BUILD_SHARED}" == "OFF" ]] && [[ "${target}" == *-mingw* ]]; then
    export CXXFLAGS="-static"
fi
cmake -DCMAKE_INSTALL_PREFIX=${prefix} \
    -DCMAKE_TOOLCHAIN_FILE=${CMAKE_TARGET_TOOLCHAIN} \
    -DCMAKE_BUILD_TYPE=Release \
    -DBUILD_SHARED_LIBS=${BUILD_SHARED} \
    -DFAST_BUILD=ON \
    -DJULIA=ON \
    -DIPX=ON ..

if [[ "${target}" == *-linux-* ]]; then
        make -j ${nproc}
else
    if [[ "${target}" == *-mingw* ]]; then
        cmake --build . --config Release
    else
        cmake --build . --config Release --parallel
    fi
fi
make install

install_license ../LICENSE

if [[ "${BUILD_SHARED}" == "OFF" ]]; then
    # Delete the static library to save space
    rm -r ${prefix}/lib
    if [[ "${target}" == *-mingw* ]]; then
        # The Windows build ships also GCC runtime, add its license as well
        install_license /usr/share/licenses/GPL3
    fi
fi
"""
end

max_stars_repo_name: stemann/Yggdrasil

max_issues_repo_name: stemann/Yggdrasil

max_forks_repo_name: stemann/Yggdrasil


text:
using CompScienceMeshes
using BEAST
using Test
using LinearAlgebra

@warn "CompScienceMeshes being patched..."
import StaticArrays
function CompScienceMeshes.chart(Smesh::CompScienceMeshes.subdMesh,E)
    element = Smesh.elements[E]
    Svertices = Smesh.vertices
    nodes = element.RingNodes
    N = length(nodes)
    verticecoords = Vector{StaticArrays.SVector{3,Float64}}(undef,N)

    for i = 1:N
        verticecoords[i] = Svertices[nodes[i]].Coords
    end
    return chart = CompScienceMeshes.subd_chart(E,N,nodes,verticecoords)
end

G = readmesh(joinpath(dirname(@__FILE__),"assets","sphere872.in"))
 # G0 = readmesh(fn)
# G0 = meshsphere(1.0, 0.5)
#G = readmesh("/Users/Benjamin/Documents/sphere.in")
# G = Loop_subdivision(G0)
X = subdsurface(G)

#els, ad = BEAST.assemblydata(X)

identityop  = Identity()
singlelayer = Helmholtz3D.singlelayer(gamma=1.0)
I = assemble(identityop, X, X)
#S = assemble(singlelayer, X, X)
ncd = cond(Matrix(I))

@test ncd ≈ 64.50401358713235 rtol=1e-8


text:
### Two methods to compute APE budgets and Buoyancy budgets from SAM v3 outputs, modified from Da Yang, and based on the mathematics from
# Yang, Da. “Boundary Layer Diabatic Processes, the Virtual Effect, and Convective Self-Aggregation.” Journal of Advances in Modeling Earth Systems 10, no. 9 (September 2018): 2163–76. https://doi.org/10.1029/2017MS001261.

"""
Returns an apebudget object.
Input is
B (buoyancy)
U,V,W The three dimensional velocities U,V and W
N2 The Brunt Va"isala frequency squared
Fs the Surface fluxes
Diabatic_other other sources of diabatic heating
rho0 the mean density
x,y,z,t the coordinate vectors
dx,dy,dz,dt the steps in each coordinate
z_up the maximum height that will be used

"""

function getapebudget_old_old(B, U,V, W, N2, RAD_b, Fs, Diabatic_other, rho0, x,y, z, t, dx,dy, dz, dt, z_up)

#***********Empty array generation***********#
Udb2dx = similar(U)
Vdb2dy = similar(V)
#*************  APE **************
b2      = B.*B/2
APE_b2  = mean(b2,dims=(1,2))[1,1,:,:]./N2;
# # # figure(4)
# # # contourf(APE_b2)

#************ KE ********************
#KE      = U.*U/2 + V.*V/2
xBar_KE = mean(U.*U/2 + V.*V/2,dims=(1,2))[1,1,:,:]

#************ APE rate ***************
xBar_APE_rate = Array{eltype(B),2}(undef,length(z), length(t))
xBar_APE_rate[:,1:end-1] = (APE_b2[:,2:end] - APE_b2[:,1:end-1])/dt;
xBar_APE_rate[:,end] = xBar_APE_rate[:,end-1]

#*************  UdxB2 **************
b2_ghost= Array{typeof(B)}(undef, length(x)+1,length(y)+1, length(z), length(t))
b2_ghost[1:end-1,1:end-1,:,:] = b2

b2_ghost[end,1:end-1,:,:] = b2[1,:,:,:]
b2_ghost[1:end-1,end,:,:] = b2[:,1,:,:]

@. Udb2dx = U*(b2_ghost[2:end,1:end-1,:,:]-b2_ghost[1:end-1,1:end-1,:,:])/dx
@. Vdb2dy = V*(b2_ghost[1:end-1,2:end,:,:]-b2_ghost[1:end-1,1:end-1,:,:])/dy
xBar_APE_Ub2 = mean(Udb2dx,dims=(1,2))[1,1,:,:]./N2
xBar_APE_Vb2 = mean(Vdb2dy,dims=(1,2))[1,1,:,:]./N2

# static stability WN2
#APE_WN2     = W.*B
xBar_APE_WN2= mean(W.*B,dims=(1,2))[1,1,:,:]

# RAD generation
xBar_APE_RAD     = mean(RAD_b.*B,dims=(1,2))[1,1,:,:]./N2;

# Diabatic_other
xBar_APE_DIA     = mean(Diabatic_other.*B,dims=(1,2))[1,1,:,:]./N2;

# interpolation
k_up              = argmin(abs.(z.-z_up));
z1                = z[1]:dz:z[k_up];
@info z1
rho01 = zeros(length(z1),length(t))
xBar_APE_b21 = similar(rho01)
xBar_APE_RAD1 = similar(rho01)
xBar_APE_DIA1 = similar(rho01)
xBar_APE_WN21 = similar(rho01)
xBar_APE_Ub21 = similar(rho01)
xBar_APE_Vb21 = similar(rho01)
xBar_KE1 = similar(rho01)
xBar_APE_rate1 = similar(rho01)
for time in 1:length(t)
    rho01_itp         = LinearInterpolation(z, rho0[:,time]);
    xBar_APE_b21_itp  = LinearInterpolation(z, APE_b2[:,time]);
    xBar_APE_RAD1_itp = LinearInterpolation(z, xBar_APE_RAD[:,time]);
    xBar_APE_DIA1_itp = LinearInterpolation(z, xBar_APE_DIA[:,time]);
    xBar_APE_WN21_itp = LinearInterpolation(z, xBar_APE_WN2[:,time]);
    xBar_APE_Ub21_itp = LinearInterpolation(z, xBar_APE_Ub2[:,time]);
    xBar_APE_Vb21_itp = LinearInterpolation(z, xBar_APE_Vb2[:,time]);
    xBar_KE1_itp      = LinearInterpolation(z, xBar_KE[:,time]);
    xBar_APE_rate1_itp    = LinearInterpolation(z, xBar_APE_rate[:,time]);

    rho01[:,time]              = [rho01_itp(x) for x in z1]
    xBar_APE_b21[:,time]       = [xBar_APE_b21_itp(x) for x in z1]
    xBar_APE_RAD1[:,time]      = [xBar_APE_RAD1_itp(x) for x in z1]
    xBar_APE_DIA1[:,time]      = [xBar_APE_DIA1_itp(x) for x in z1]
    xBar_APE_WN21[:,time]      = [xBar_APE_WN21_itp(x) for x in z1]
    xBar_APE_Ub21[:,time]      = [xBar_APE_Ub21_itp(x) for x in z1]
    xBar_APE_Vb21[:,time]      = [xBar_APE_Vb21_itp(x) for x in z1]
    xBar_KE1[:,time]           = [xBar_KE1_itp(x) for x in z1]
    xBar_APE_rate1[:,time]     = [xBar_APE_rate1_itp(x) for x in z1]
end

# vertical integration
int_mass     = sum(rho01.*dz,dims=1)[1,:]
int_KE       = sum(rho01.*xBar_KE1.*dz,dims=1)[1,:]
int_APE      = sum(rho01.*xBar_APE_b21.*dz,dims=1)[1,:]
int_APE_RAD  = sum(rho01.*xBar_APE_RAD1.*dz,dims=1)[1,:]
int_APE_DIA  = sum(rho01.*xBar_APE_DIA1.*dz,dims=1)[1,:]
int_APE_WN2  = sum(rho01.*xBar_APE_WN21.*dz,dims=1)[1,:]
int_APE_Ub2  = sum(rho01.*xBar_APE_Ub21.*dz,dims=1)[1,:]
int_APE_Vb2  = sum(rho01.*xBar_APE_Vb21.*dz,dims=1)[1,:]
int_APE_rate = sum(rho01.*xBar_APE_rate1.*dz,dims=1)[1,:]

# surface flux contribution
#N2S         = N2[1,:]
#APE_Fs      = B[:,:,1,:].*Fs

xBar_APE_Fs = mean(B[:,:,1,:].*Fs, dims=(1,2))[1,1,1,:]./N2[1,:];

residual    = int_APE_rate .+ int_APE_Ub2 .+ int_APE_Vb2+int_APE_WN2 .- (int_APE_RAD .+ int_APE_DIA .+ xBar_APE_Fs)

return (int_mass, int_KE, int_APE, int_APE_rate, int_APE_Ub2,int_APE_Vb2, int_APE_WN2, int_APE_RAD, int_APE_DIA, xBar_APE_Fs, residual)
end

"""
-------------Computes the APE budgets budget---------------

"""

function getapebudget_old(B, U,V, W, N2, RAD_b, Fs, Diabatic_other, rho0, x,y, z, t, dx,dy, dz, dt, z_up)
    N2            = reshape(N2,1,1,length(z),length(t))
    #***********Empty array generation***********#
    T             = eltype(B)
    lt            = length(t)
    lz            = length(z)
    lx            = length(y)
    ly            = length(x)
    buf           = Array{T}(undef,lx,ly,lz, lt)
    xBar_KE       = Array{T}(undef,1,1,lz, lt)
    APE_b2        = Array{T}(undef,1,1,lz, lt)
    xBar_APE_rate = Array{T}(undef,lz, lt)
    b2_ghost      = Array{T}(undef, length(x)+1,length(y)+1, lz, lt)
    xBar_APE_Ub2  = Array{T}(undef,1,1,lz, lt)
    xBar_APE_Vb2  = Array{T}(undef,1,1,lz, lt)
    xBar_APE_WN2  = Array{T}(undef,1,1,lz, lt)
    xBar_APE_RAD  = Array{T}(undef,1,1,lz, lt)
    xBar_APE_DIA  = Array{T}(undef,1,1,lz, lt)
    xBar_APE_FS   = Array{T}(undef,1,1,lz, lt)
    #*************  APE **************
    @. buf                           = B*B/2
    mean!(APE_b2,buf);
    @. APE_b2                        = APE_b2/N2
    @. b2_ghost[1:end-1,1:end-1,:,:] = buf
    @. b2_ghost[end,1:end-1,:,:]     = buf[1,:,:,:]
    @. b2_ghost[1:end-1,end,:,:]     = buf[:,1,:,:]

    #************ KE ********************
    #KE      = U.*U/2 + V.*V/2
    @. buf  = U*U/2 + V*V/2
    xBar_KE = mean!(xBar_KE,buf)

    #************ APE rate ***************

    @. xBar_APE_rate[:,1:end-1] = (APE_b2[1,1,:,2:end] - APE_b2[1,1,:,1:end-1])/dt;
    @. xBar_APE_rate[:,end] = xBar_APE_rate[:,end-1]

    #*************  UdxB2 **************

    @. buf = U*(b2_ghost[2:end,1:end-1,:,:]-b2_ghost[1:end-1,1:end-1,:,:])/dx
    mean!(xBar_APE_Ub2,buf)

    @. buf = V*(b2_ghost[1:end-1,2:end,:,:]-b2_ghost[1:end-1,1:end-1,:,:])/dy

    mean!(xBar_APE_Vb2,buf)

    @. xBar_APE_Ub2 = xBar_APE_Ub2/N2
    @. xBar_APE_Vb2 = xBar_APE_Vb2/N2

    # static stability WN2
    #APE_WN2     = W.*B
    @. buf = W*B
    mean!(xBar_APE_WN2,buf)

    # RAD generation
    @. buf = RAD_b.*B
    mean!(xBar_APE_RAD,buf);
    @. xBar_APE_RAD = xBar_APE_RAD/N2

    # Diabatic_other
    @. buf = Diabatic_other.*B
    mean!(xBar_APE_DIA,buf)
    @. xBar_APE_DIA = xBar_APE_DIA/N2;

    # Surface fluxes contribution

    xBar_APE_Fs  = mean(B[:,:,1,:].*Fs, dims=(1,2))[1,1,1,:]./N2[1,1,1,:];

    # interpolation
    k_up              = argmin(abs.(z.-z_up));
    z1                = z[1]:dz:z[k_up];
      #    int_mass      = Array{T}(undef,lt)
    int_mass      = zeros(T,lt)
    int_KE        = zeros(T,lt)
    int_APE       = zeros(T,lt)
    int_APE_RAD   = zeros(T,lt)
    int_APE_DIA   = zeros(T,lt)
    int_APE_WN2   = zeros(T,lt)
    int_APE_Ub2   = zeros(T,lt)
    int_APE_Vb2   = zeros(T,lt)
    int_APE_rate  = zeros(T,lt)
    for timeind in 1:lt
        rho01_itp         = interpolate((z,), rho0[:,timeind],Gridded(Linear()))
        xBar_APE_b21_itp  = interpolate((z,), APE_b2[1,1,:,timeind],Gridded(Linear()))
        xBar_APE_RAD1_itp = interpolate((z,), xBar_APE_RAD[1,1,:,timeind],Gridded(Linear()))
        xBar_APE_DIA1_itp = interpolate((z,), xBar_APE_DIA[1,1,:,timeind],Gridded(Linear()))
        xBar_APE_WN21_itp = interpolate((z,), xBar_APE_WN2[1,1,:,timeind],Gridded(Linear()))
        xBar_APE_Ub21_itp = interpolate((z,), xBar_APE_Ub2[1,1,:,timeind],Gridded(Linear()))
        xBar_APE_Vb21_itp = interpolate((z,), xBar_APE_Vb2[1,1,:,timeind],Gridded(Linear()))
        xBar_KE1_itp      = interpolate((z,), xBar_KE[1,1,:,timeind],Gridded(Linear()))
        xBar_APE_rate1_itp    = interpolate((z,), xBar_APE_rate[:,timeind],Gridded(Linear()))

        for zeta in z1
                mass = dz*rho01_itp(zeta)
                int_mass[timeind]         += mass
                int_APE[timeind]          += mass*xBar_APE_b21_itp(zeta)
                int_APE_RAD[timeind]      += mass*xBar_APE_RAD1_itp(zeta)
                int_APE_DIA[timeind]      += mass*xBar_APE_DIA1_itp(zeta)
                int_APE_WN2[timeind]      += mass*xBar_APE_WN21_itp(zeta)
                int_APE_Ub2[timeind]      += mass*xBar_APE_Ub21_itp(zeta)
                int_APE_Vb2[timeind]      += mass*xBar_APE_Vb21_itp(zeta)
                int_KE[timeind]           += mass*xBar_KE1_itp(zeta)
                int_APE_rate[timeind]     += mass*xBar_APE_rate1_itp(zeta)
        end
    end
    residual    = int_APE_rate .+ int_APE_Ub2 .+ int_APE_Vb2 .+ int_APE_WN2 .- (int_APE_RAD .+ int_APE_DIA .+ xBar_APE_Fs)
    return (int_mass, int_KE, int_APE, int_APE_rate, int_APE_Ub2,int_APE_Vb2, int_APE_WN2, int_APE_RAD, int_APE_DIA, xBar_APE_Fs, residual)
    end

"""
-------------Computes the APE budgets budget---------------

"""

function getapebudget(B, U,V, W, N2, RAD_b, Fs, Diabatic_other, rho0, x,y, z, t, dx,dy, dz, dt, z_up)
    N2            = reshape(N2,1,1,length(z),length(t))
    rho0            = reshape(rho0,1,1,length(z),length(t))
    #***********Empty array generation***********#
    T             = eltype(B)
    lx            = length(x)
    ly            = length(y)
    lt            = length(t)
    lz            = length(z)
    buf           = similar(U)
    buf_2d        = similar(Fs)
    xBar_KE       = Array{T}(undef,1,1,lz, lt)
    APE_b2           = Array{T}(undef,1,1,lz, lt)
    xBar_APE_rate = Array{T}(undef,1,1,lz, lt)
    xBar_APE_Ub2  = Array{T}(undef,1,1,lz, lt)
    xBar_APE_Vb2  = Array{T}(undef,1,1,lz, lt)
    xBar_APE_WN2  = Array{T}(undef,1,1,lz, lt)
    xBar_APE_RAD  = Array{T}(undef,1,1,lz, lt)
    xBar_APE_DIA  = Array{T}(undef,1,1,lz, lt)
    xBar_APE_Fs   = Array{T}(undef,1,1,lt)

    neighborx(indx,lx) = mod1(indx+1,lx)
    onex = CartesianIndex((1,0,0,0))
    oney = CartesianIndex((0,1,0,0))
    onet = CartesianIndex((0,0,0,1))

    #    for ind in CartesianIndices((1:(sx - 1), 2:(sy - 1), 1:sz, 2:(st-1)))

    #*************  APE **************
    @. buf = B*B/2
    mean!(APE_b2,buf./N2)

    #************ KE ********************
    @. buf  = U*U/2 + V*V/2
    xBar_KE = mean!(xBar_KE,buf)

    #************ APE rate ***************

    @inbounds  for it in 1:(lt -1), iz in 1:lz, iy in 1:ly, ix in 1:lx
        ind = CartesianIndex((ix,iy,iz,it))
        buf[ind] =  B[ind]*(B[ind + onet] - B[ind]) / N2[1,1,iz,it] / dt
    end
    @inbounds for ind in CartesianIndices((1:lx,1:ly,1:lz,lt:lt))
        buf[ind] = buf[ind-onet]
    end

    mean!(xBar_APE_rate,buf);

    #*************  Advection **************
    @inbounds  for it in 1:(lt - 1), iz in 1:lz, iy in 1:ly, ix in 1:lx
        ind = CartesianIndex((ix,iy,iz,it))
        buf[ind] =  B[ind]*U[ind]*(B[neighborx(ix,lx),iy,iz,it] - B[ind]) / N2[1,1,iz,it] /dx
    end
    @inbounds for ind in CartesianIndices((1:lx,1:ly,1:lz,lt:lt))
        buf[ind] = buf[ind-onet]
    end
    mean!(xBar_APE_Ub2,buf)

    @inbounds  for it in 1:lt, iz in 1:lz, iy in 1:ly, ix in 1:lx
        ind = CartesianIndex((ix,iy,iz,it))
        buf[ind] = B[ind]*V[ind]*(B[ix,neighborx(iy,ly),iz,it] - B[ind]) / N2[1,1,iz,it] /dy
    end
    @inbounds for ind in CartesianIndices((1:lx,1:ly,1:lz,lt:lt))
        buf[ind] = buf[ind-onet]
    end
    mean!(xBar_APE_Vb2,buf)

    ################################# static stability WN2

    @. buf = W*B
    @inbounds for ind in CartesianIndices((1:lx,1:ly,1:lz,lt:lt))
        buf[ind] = buf[ind-onet]
    end
    mean!(xBar_APE_WN2,buf)
    # RAD generation

    @. buf = RAD_b*B/N2
    @inbounds for ind in CartesianIndices((1:lx,1:ly,1:lz,lt:lt))
        buf[ind] = buf[ind-onet]
    end
    mean!(xBar_APE_RAD,buf);

    ########################## Diabatic_other
    @. buf = Diabatic_other*B/N2
    @inbounds for ind in CartesianIndices((1:lx,1:ly,1:lz,lt:lt))
        buf[ind] = buf[ind-onet]
    end
    mean!(xBar_APE_DIA,buf)
    ##################### Surface fluxes contribution
    @inbounds for tind in 1:lt
        for yind in 1:ly
            for xind in 1:lx
                buf_2d[xind,yind,tind] = B[xind,yind,1,tind]*Fs[xind,yind,tind]/N2[1,1,1,tind]
            end
        end
    end
    @inbounds for ind in CartesianIndices((1:lx,1:ly,lt:lt))
        buf_2d[ind] = buf_2d[ind - CartesianIndex((0,0,1))]
    end
    mean!(xBar_APE_Fs,buf_2d)
    xBar_APE_Fs = dropdims(xBar_APE_Fs, dims=(1,2))
    ##################### interpolation
    k_up              = argmin(abs.(z.-z_up));
    z1                = z[1]:dz:z[k_up];
    int_mass      = zeros(T,lt)
    int_KE        = zeros(T,lt)
    int_APE       = zeros(T,lt)
    int_APE_RAD   = zeros(T,lt)
    int_APE_DIA   = zeros(T,lt)
    int_APE_WN2   = zeros(T,lt)
    int_APE_Ub2   = zeros(T,lt)
    int_APE_Vb2   = zeros(T,lt)
    int_APE_rate  = zeros(T,lt)
    @inbounds for time in 1:lt
        rho01_itp         =  interpolate((z,), rho0[1,1,:,time],Gridded(Linear()))
        xBar_APE_b21_itp  =  interpolate((z,), APE_b2[1,1,:,time],Gridded(Linear()))
        xBar_APE_RAD1_itp =  interpolate((z,), xBar_APE_RAD[1,1,:,time],Gridded(Linear()))
        xBar_APE_DIA1_itp =  interpolate((z,), xBar_APE_DIA[1,1,:,time],Gridded(Linear()))
        xBar_APE_WN21_itp =  interpolate((z,), xBar_APE_WN2[1,1,:,time],Gridded(Linear()))
        xBar_APE_Ub21_itp =  interpolate((z,), xBar_APE_Ub2[1,1,:,time],Gridded(Linear()))
        xBar_APE_Vb21_itp =  interpolate((z,), xBar_APE_Vb2[1,1,:,time],Gridded(Linear()))
        xBar_KE1_itp      =  interpolate((z,), xBar_KE[1,1,:,time],Gridded(Linear()))
        xBar_APE_rate1_itp  =  interpolate((z,), xBar_APE_rate[1,1,:,time],Gridded(Linear()))
        @inbounds for zeta in z1
            mass = dz*rho01_itp(zeta)
            int_mass[time]         += mass
            int_APE[time]          += mass*xBar_APE_b21_itp(zeta)
            int_APE_RAD[time]      += mass*xBar_APE_RAD1_itp(zeta)
            int_APE_DIA[time]      += mass*xBar_APE_DIA1_itp(zeta)
            int_APE_WN2[time]      += mass*xBar_APE_WN21_itp(zeta)
            int_APE_Ub2[time]      += mass*xBar_APE_Ub21_itp(zeta)
            int_APE_Vb2[time]      += mass*xBar_APE_Vb21_itp(zeta)
            int_KE[time]           += mass*xBar_KE1_itp(zeta)
            int_APE_rate[time]     += mass*xBar_APE_rate1_itp(zeta)

        end
    end
    residual    = int_APE_rate .+ int_APE_Ub2 .+ int_APE_Vb2 .+ int_APE_WN2 .- (int_APE_RAD .+ int_APE_DIA .+ xBar_APE_Fs)
    return (int_mass, int_KE, int_APE, int_APE_rate, int_APE_Ub2,int_APE_Vb2, int_APE_WN2, int_APE_RAD, int_APE_DIA, xBar_APE_Fs, residual)
    end

"""
-------------This function computes the buoyancy budget---------------
Inputs:
B (buoyancy)
RAD_b the radiative heating, converted to units of buoyancy
Fs the Surface fluxes
U,V,W The three dimensional velocities
N2 The Brunt Va"isala frequency squared
dx,dy,dz,dt the steps in each coordinate
x,y,z,t the coordinate vectors

"""
function buoyancybudget_old(B, RAD_b, Fs, U,V, W, N2, dx,dy, dz, dt, x,y, z, t)
#************ Array creation **************#
Qs      = zeros(typeof(B[1]),length(x),length(y),length(z),length(t))
B_ghost = Array{typeof(B[1])}(undef, length(x)+1,length(y)+1, length(z), length(t))
WN2     = similar(W)
dBdt    = similar(B)
UdBdx   = similar(U)
VdBdy   = similar(U)
#************ WN2 *****************

 WN2  .= reshape(N2,(1,1,size(N2,1),size(N2,2))).*W

#*************  Advection **************

@. B_ghost[1:end-1,1:end-1,:,:] = B
@. B_ghost[end,1:end-1,:,:] = B[1,:,:,:]
@. B_ghost[1:end-1,end,:,:] = B[:,1,:,:]

@. UdBdx = U*(B_ghost[2:end,1:end-1,:,:]-B_ghost[1:end-1,1:end-1,:,:])/dx
@. VdBdy = V*(B_ghost[1:end-1,2:end,:,:]-B_ghost[1:end-1,1:end-1,:,:])/dy

# @. UdBdx[1:end-1,:,:,:] .= U[1:end-1,:,:,:]*(B[2:end,:,:,:]-B[1:end-1,:,:,:])/dx
# @. UdBdx[end,:,:,:] = U[end,:,:,:]*(B[1,:,:,:]-B[end,:,:,:])/dx

# @. VdBdy[:,1:end-1,:,:] = V[:,1:end-1,:,:].*(B[:,2:end,:,:]-B[:,1:end-1,:,:])/dy
# @. VdBdy[:,end,:,:] = V[:,end,:,:]*(B[:,1,:,:]-B[:,end,:,:])/dy

Qs[:,:,1,:] .= Fs./dz
#************ Time evolution *************#

@. dBdt[:,:,:,1:end-1] = (B[:,:,:,2:end] - B[:,:,:,1:end-1])/dt;
@. dBdt[:,:,:,end]     = dBdt[:,:,:,end-1]/dt
#*************** Return ********************#

Diabatic_other  = dBdt .+ UdBdx .+ VdBdy .+ WN2 .- RAD_b .- Qs

return dBdt, UdBdx,VdBdy, WN2, Qs, Diabatic_other
end

"""
-------------This function computes the buoyancy budget---------------
Inputs:
B (buoyancy)
RAD_b the radiative heating, converted to units of buoyancy
Fs the Surface fluxes
U,V,W The three dimensional velocities
N2 The Brunt Va"isala frequency squared
dx,dy,dz,dt the steps in each coordinate
x,y,z,t the coordinate vectors

"""

function buoyancybudget(B, RAD_b, Fs, U,V, W, N2, dx,dy, dz, dt, x,y, z, t)
    #************ Array creation **************#
    #Qs      = zeros(eltype(B),length(x),length(y),length(z),length(t))
    #B_ghost = Array{typeof(B[1])}(undef, length(x)+1,length(y)+1, length(z), length(t))
    WN2     = similar(W)
    dBdt    = similar(B)
    UdBdx   = similar(U)
    VdBdy   = similar(U)
    #************ WN2 *****************

    WN2  .= reshape(N2,(1,1,size(N2,1),size(N2,2))).*W

    #*************  Advection **************

    @. @views UdBdx[1:end-1,:,:,:] = U[1:end-1,:,:,:]*(B[2:end,:,:,:]-B[1:end-1,:,:,:])/dx
    @. @views UdBdx[end,:,:,:] = U[end,:,:,:]*(B[1,:,:,:]-B[end,:,:,:])/dx

    @. @views VdBdy[:,1:end-1,:,:] = V[:,1:end-1,:,:]*(B[:,2:end,:,:]-B[:,1:end-1,:,:])/dy
    @. @views VdBdy[:,end,:,:] = V[:,end,:,:]*(B[:,1,:,:]-B[:,end,:,:])/dy

    #@. Qs[:,:,1,:] = Fs/dz
#************ Time evolution *************#

    @. @views dBdt[:,:,:,1:end-1] = (B[:,:,:,2:end] - B[:,:,:,1:end-1])/dt
    @. @views dBdt[:,:,:,end]     = dBdt[:,:,:,end-1]
    #*************** Return ********************#

    Diabatic_other  = dBdt .+ UdBdx .+ VdBdy .+ WN2 .- RAD_b
    Diabatic_other[:,:,1,:] .= @views  Diabatic_other[:,:,1,:] -  Fs/dz

#    return dBdt, UdBdx,VdBdy, WN2, Fs/dz, Diabatic_other
    return Diabatic_other
end

function get_diabatic_as_residual_buoyancy(B, RAD_b, Fs, U,V, W, N2, dx,dy, dz, dt)
    #************ Array creation **************#
    #Qs      = zeros(eltype(B),length(x),length(y),length(z),length(t))
    #B_ghost = Array{typeof(B[1])}(undef, length(x)+1,length(y)+1, length(z), length(t))

    Diabatic_other = similar(B)

    onex = CartesianIndex((1,0,0,0))
    oney = CartesianIndex((0,1,0,0))
    onet = CartesianIndex((0,0,0,1))

    sx,sy,sz,st = size(Diabatic_other)
    N2            = reshape(N2,1,1,sz,st)
    neighborx(indx,sx) = mod1(indx+1,sx)

    #    for ind in CartesianIndices((1:(sx - 1), 2:(sy - 1), 1:sz, 2:(st-1)))
    #### derivative
    @inbounds  for it in 1:(st - 1), iz in 1:sz, iy in 1:sy, ix in 1:sx
            ind = CartesianIndex((ix,iy,iz,it))
            Diabatic_other[ind] = (B[ind+onet] - B[ind])/dt + U[ind]*(B[neighborx(ix,sx),iy,iz,it] -
                B[ind])/dx + V[ind]*(B[ix,neighborx(iy,sy),iz,it] - B[ind])/dy +
                W[ind]*N2[1,1,iz,it] - RAD_b[ind]
    end

    @inbounds for it in 1:(st - 1), iy in 1:sy, ix in 1:sx
            Diabatic_other[ix,iy,1,it] = Diabatic_other[ix,iy,1,it] - Fs[ix,iy,it]/dz
    end

#    for it in st:st, iz in 1:sz, iy in 1:(sy - 1), ix in 1:(sx - 1)

    @inbounds for ind in CartesianIndices((1:sx,1:sy,1:sz,st:st))
        Diabatic_other[ind] = Diabatic_other[ind-onet]
    end

    return Diabatic_other
end


text:
using DelayDiffEq, DiffEqProblemLibrary.DDEProblemLibrary
using Test

DDEProblemLibrary.importddeproblems()

@testset "Constant delays" begin
    prob = DDEProblemLibrary.prob_dde_constant_2delays_ip
    prob_scalar = DDEProblemLibrary.prob_dde_constant_2delays_scalar
    # disable reuse
    nlsolve = nlsolve=NLNewton(fast_convergence_cutoff=0)

    algdict = Dict(BS3() => 2.4e-6,
                   Tsit5() => 4.5e-3,
                   RK4() => 1.1e-4,
                   Vern6() => 1.0e-3,
                   SDIRK2(nlsolve=nlsolve) => 2.3e-1,
                   TRBDF2(nlsolve=nlsolve) => 6.2e-2,
                   KenCarp4(nlsolve=nlsolve) => 5.6e-2,
                   Rosenbrock23() => 6.5e-4,
                   Rodas4() => 5.4e-4
                  )

    for (alg, error) in algdict
        ddealg = MethodOfSteps(alg)

        sol = solve(prob, ddealg)
        @test sol.errors[:l∞] < error

        sol_scalar = solve(prob_scalar, ddealg)
        @test sol.t ≈ sol_scalar.t
        @test sol[1, :] ≈ sol_scalar.u
    end
end

function lotka_volterra!(du,u,h,p,t)
    🐰, 🐺 = u
    α,β,γ,δ,τ = p
    🕥🐰 = h(p,t-τ;idxs=1)
    du[1] = d🐰 = α*🕥🐰 - β*🐺*🐰
    du[2] = d🐺 = γ*🐺*🐰 -δ*🐺
    nothing
end
uₒ = [1.0,1.0]
tspan = (0.0, 10.0)
h(p,t) = [1.0,1.0]
h(p,t;idxs = 1) = 1.0
p = [1.5,1.0,3.0,1.0,1.0]
prob = DDEProblem(lotka_volterra!,uₒ,h,tspan,p,constant_lags = (p[end],))
sol = solve(prob,MethodOfSteps(AutoTsit5(Rosenbrock23(autodiff=false))))


text:
# Initialize environment in current directory
@info("Ensuring example environment instantiated...")
import Pkg
Pkg.activate(@__DIR__)
Pkg.instantiate()

@info("Loading Zygote and Flux...")
using Zygote, Flux, Random, Statistics
using Flux.Data.MNIST

# We're going to showcase how to use Zygote with Flux; we'll create a simple
# Multi-Layer Perceptron network to do digit classification upon the MNIST
# dataset.  We start with some setup that is ripped straight from the Flux
# model zoo:

# First, we load the MNIST images and flatten them into a giant matrix:
@info("Loading dataset...")
X = hcat(float.(reshape.(MNIST.images(), :))...)

# Load labels as well, one-hot encoding them
Y = float.(Flux.onehotbatch(MNIST.labels(), 0:9))

# Do the same for the test data/labels:
X_test = hcat(float.(reshape.(MNIST.images(:test), :))...)
Y_test = float.(Flux.onehotbatch(MNIST.labels(:test), 0:9))

@info("Constructing MLP model...")
model = Chain(
    Dense(28^2, 32, relu),
    Dense(32, 10),
    softmax,
)

# Until Flux drops Tracker as its default Automatic Differentiation library,
# strip it out with this line:
model = Flux.mapleaves(Flux.data, model)

# Our loss is the classical multiclass crossentropy loss
loss(model, X, Y) = Flux.crossentropy(model(X), Y)

# Helper function to calculate accuracy of our model
accuracy(model, X, Y) = mean(Flux.onecold(model(X)) .== Flux.onecold(Y))

# Recursive zygote update method, this is the general recursion case:
function zyg_update!(opt, model, updates)
	# If this `model` node has no fields, then just return it
    if nfields(model) == 0
        return model
    end

	# If it does have fields, recurse into them:
    for field_idx in 1:nfields(model)
        zyg_update!(opt, getfield(model, field_idx), getfield(updates, field_idx))
    end

    # In the end, return the `model`
    return model
end
# If the `updates` is set to `Nothing`, then just return `model`; this means
# that there were no changes to be applied to this piece of the model.
zyg_update!(opt, model, updates::Nothing) = model

# If `model` is an `AbstractArray` and `updates` is too, then apply our Flux
# optimizer to the incoming gradients and apply them to the model!
function zyg_update!(opt, model::AbstractArray, updates::AbstractArray)
    # Sub off to Flux's ADAM optimizer
    Flux.Optimise.apply!(opt, model, updates)
    return model .-= updates
end

# We will train for a number of epochs, with minibatches, using the `ADAM`
# optimizer to nudge our weights toward perfection.
opt = ADAM(0.001)
num_epochs = 10
@info("Training for $(num_epochs) epochs...")
for epoch_idx in 1:num_epochs
    # "global" here to dodgescoping issues with for loops at top-level
    global X, Y, model

    # Shuffle the data each epoch:
    perm = shuffle(1:size(X,2))
    X = X[:, perm]
    Y = Y[:, perm]

    # Iterate over batches
    batch_size = 512
    batch_idxs = 1:batch_size:(size(X,2) - batch_size)
    for bidx in batch_idxs
        # Calculate gradients upon the model for this batch
        grads = Zygote.gradient(model) do model
            return loss(model, X[:, bidx:bidx+batch_size],
                               Y[:, bidx:bidx+batch_size])
        end

        # Peel outer Tuple to access gradient of first parameter
        grads = grads[1]

        # Apply recursive update to our model:
        zyg_update!(opt, model, grads)
    end

    # After each epoch, report our accuracy on the test set:
    acc = accuracy(model, X_test, Y_test)
    @info("[$(epoch_idx)] Accuracy: $(round(100*acc; digits=1))%")
end


text:
"""
    apply(data::T, ::Transform; kwargs...)

Applies the [`Transform`](@ref) to the data. New transforms should usually only extend
`_apply` which this method delegates to.

Where necessary, this should be extended for new data types `T`.
"""
function apply end

"""
    apply!(data::T, ::Transform; kwargs...) -> T

Applies the [`Transform`](@ref) mutating the input `data`. This method delegates to
[`apply`](@ref) under the hood so does not need to be defined separately.

If [`Transform`](@ref) does not support mutation, this method will error.
"""
function apply! end

"""
    apply(A::AbstractArray, ::Transform; dims=:, inds=:, kwargs...)

Applies the [`Transform`](@ref) to the elements of `A`.

Provide the `dims` keyword to apply the [`Transform`](@ref) along a certain dimension.
For example, given a `Matrix`, `dims=1` applies to each column, while `dims=2` applies
to each row.

Provide the `inds` keyword to apply the [`Transform`](@ref) to certain indices along the
`dims` specified.

Note: if `dims === :` (all dimensions), then `inds` will be the global indices of the array,
instead of being relative to a certain dimension.
"""
function apply(A::AbstractArray, t::Transform; dims=:, inds=:, kwargs...)

    # TODO: remove this https://github.com/invenia/FeatureTransforms.jl/issues/82
    if t isa LinearCombination && dims === Colon()
        Base.depwarn(
           "The default `dims=1` for `LinearCombination` is deprecated and will be removed " *
           "in a future release. Please set `dims` explicitly.",
           :apply,
        )
        dims=1
    end

    c = cardinality(t)
    if dims === Colon()
        if inds === Colon()
            return _apply(_preformat(c, A, :), t; dims=:, kwargs...)
        else
            return _apply(_preformat(c, A[:][inds], :), t; dims=:, kwargs...)
        end
    end

    input = _preformat(c, selectdim(A, dims, inds), dims)
    return _apply(input, t; dims=dims, kwargs...)
end

"""
    apply!(A::AbstractArray, ::Transform; dims=:, kwargs...)

Applies the [`Transform`](@ref) to each element of `A`, mutating the data.
"""
function apply!(A::AbstractArray, t::Transform; kwargs...)
    A[:] = apply(A, t; kwargs...)
    return A
end

"""
    apply(table, ::Transform; [cols], [header], kwargs...) -> Table

Applies the [`Transform`](@ref) to each of the specified columns in the `table`.
If no `cols` are specified, then the [`Transform`](@ref) is applied to all columns.

Optionally provide a `header` for the output table. If none is provided the default in
`Tables.table` is used.
"""
function apply(table, t::Transform; cols=_get_cols(table), header=nothing, kwargs...)
    Tables.istable(table) || throw(MethodError(apply, (table, t)))

    indices = Tables.columnindex.(Ref(table), _to_vec(cols))
    components = Tables.matrix(table)[:, indices]

    # Passing dims=2 only matters for ManyToOne transforms - otherwise it has no effect.
    input = _preformat(cardinality(t), components, 2)
    result = hcat(_apply(input, t; dims=2, kwargs...))
    return Tables.materializer(table)(_to_table(result, header))
end

# Prevents Vector{NamedTuple} from using the AbstractArray method above
function apply(table::Tables.RowTable, t::Transform; kwargs...)
    return rowtable(apply(Tables.columntable(table), t; kwargs...))
end

"""
    apply!(table::T, ::Transform; [cols])::T where T

Applies the [`Transform`](@ref) to each of the specified columns in the `table`.
If no `cols` are specified, then the [`Transform`](@ref) is applied to all columns.

!!! Note
  `apply!` does not support `RowTable`s since the `NamedTuple` rows are immutable.
"""
function apply!(table::T, t::Transform; cols=_get_cols(table), kwargs...)::T where T
    Tables.istable(table) || throw(MethodError(apply!, (table, t)))

    for cname in _to_vec(cols)
        apply!(getproperty(table, cname), t; kwargs...)
    end

    return table
end

# Prevents Vector{NamedTuple} from using the AbstractArray method above
function apply!(table::Tables.RowTable, t::Transform; kwargs...)
    return throw(MethodError(apply!, (table, t)))
end

"""
    apply_append(A::AbstractArray, ::Transform; append_dim, kwargs...)

Applies the [`Transform`](@ref) to `A` and returns the result in a new array where the output
is appended to `A` along the `append_dim` dimension. The remaining `kwargs` correspond to
the usual [`Transform`](@ref) being invoked.
"""
function apply_append(A::AbstractArray, t; append_dim, kwargs...)::AbstractArray
    result = apply(A, t; kwargs...)
    result = _postformat(cardinality(t), result, A, append_dim)
    # Conver result to Array to avoid clashing axis/key names in concatenated result
    return cat(A, Array(result); dims=append_dim)
end

"""
    apply_append(table, ::Transform; [header], kwargs...)

Applies the [`Transform`](@ref) to the `table` and appends the result in a new table with an
optional `header`. If none is provided the default in `Tables.table` is used. The remaining
`kwargs` correspond to the [`Transform`](@ref) being invoked.
"""
function apply_append(table, t; kwargs...)
    T = Tables.materializer(table)
    result = Tables.columntable(apply(table, t; kwargs...))
    return T(merge(Tables.columntable(table), result))
end

# Prevents Vector{NamedTuple} from using the AbstractArray method above
function apply_append(table::Tables.RowTable, t; kwargs...)
    T = Tables.materializer(table)
    result = apply(table, t; kwargs...)
    return T(merge.(Tables.rowtable(table), result))
end

# These methods format data according to the cardinality of the Transform.
# Most Transforms don't require any formatting, only those that are ManyToOne do.
# Note: we don't yet have a ManyToMany transform, so those might need separate treatment.

# _preformat formats the data before calling _apply. Needed for all apply methods.
# Before applying a ManyToOne Transform we must first slice up the data along the dimension
# we are reducing over.
_preformat(::Cardinality, A, d) = A
_preformat(::ManyToOne, A, d) = eachslice(A; dims=d)

# _postformat formats the data after calling _apply so it will work with apply_append.
# Basically, when we call LinearCombination it always returns a column vector. But if we
# want to append the result as a row we have to reshape to get it to fit.
# In general, after applying a ManyToOne Transform, we have to reshape the reduced dimension
# to 1 if we want to cat the result.
_postformat(::Cardinality, result, A, append_dim) = result
function _postformat(::ManyToOne, result, A, append_dim)
    new_size = collect(size(A))
    setindex!(new_size, 1, dim(A, append_dim))
    return copy(reshape(result, new_size...))  # return a copy to remove the reshape type
end


text:
function pposs_mutation(s, n)
    s_temp = copy(s)
    threshold = 1.0 / n
    flipped = false
    for i = 1:n
        if rand(rng, Float64) <= threshold
            s_temp[i] = (s_temp[i] + 1) % 2
            flipped = true
        end
    end
    if !flipped
        mustflip = rand(rng, 1:n)
        s_temp[mustflip] = (s_temp[mustflip] + 1) % 2
    end
    return s_temp
end

function pposs_opt!(objective::Objective, parameter::Parameter)
    sample_set = RemoteChannel(()->Channel(parameter.evaluation_server_num))
    result_set = RemoteChannel(()->Channel(parameter.evaluation_server_num))
    asyn_result = RemoteChannel(()->Channel(1))
    history = RemoteChannel(()->Channel(1))
    population = []
    # require calculator server
    ip = parameter.control_server_ip
    port = parameter.control_server_port
    cs_send = connect(ip, port)
    println(cs_send, "client: require servers#")
    readline(cs_send)
    println(cs_send, string(parameter.evaluation_server_num, "#"))
    msg = readline(cs_send)

    servers_msg = readline(cs_send)
    if servers_msg == "No available evaluation server"
        zoolog("Error: no available evaluation server")
        return Solution()
    end
    servers = split(servers_msg, " ")
    println("get $(length(servers)) servers")
    parameter.ip_port = RemoteChannel(()->Channel(length(servers)))
    for server in servers
      put!(parameter.ip_port, server)
    end
    n = objective.dim.size
    sol = Solution(x=[0 for i = 1:n])
    ip_port = take!(parameter.ip_port)
    br = pposs_compute_fx(sol, ip_port, parameter)
    put!(parameter.ip_port, ip_port)

    push!(population, sol)

    pposs_init_sample_set!(sample_set, sol, parameter)
    println("Initialization succeeds")
    finish = SharedArray{Bool}(1)
    finish[1] = false
    @spawn pposs_updater!(population, sample_set, result_set, asyn_result, history, parameter, finish)
    br = false
    while true
        if finish[1] == true
            break
        end
        # println("$(finish[1])")
        ip_port = take!(parameter.ip_port)
        sol = take!(sample_set)
        @spawn begin
            try
                br = pposs_compute_fx(sol, ip_port, parameter)
                put!(parameter.ip_port, ip_port)
                put!(result_set, sol)
                # pposs_updater!(population, result_set, asyn_result, parameter, finish)
            catch e
                println("Exception")
                println(e)
                cs_exception = connect(ip, port)
                println(cs_exception, "client: restart#")
                readline(cs_send)
                println(cs_receive, string(servers_msg, "#"))
                return Solution()
            end
        end
    end
    result = take!(asyn_result)
    objective.history = take!(history)
    cs_receive = connect(ip, port)
    println(cs_receive, "client: return servers#")
    println(readline(cs_send))
    println(cs_receive, string(servers_msg, "#"))
    return result
end

function pposs_init_sample_set!(sample_set, sol, parameter)
    evaluation_server_num = parameter.evaluation_server_num
    data_temp = parameter.init_sample
    init_num = 0
    if !isnull(data_temp)
        init_num = length(data_temp)
        for i = 1:init_num
            put!(sample_set, data_temp[i])
        end
    end
    for i = 1:(evaluation_server_num - init_num)
        new_x = pposs_mutation(sol.x, length(sol.x))
        put!(sample_set, Solution(x=new_x))
    end
end

function pposs_updater!(population, sample_set, result_set, asyn_result, history, parameter, finish)
    t = 1
    pop_size = 1
    my_history = []
    budget = parameter.budget
    time_log1 = now()
    interval = 10
    time_sum = interval
    output_file = parameter.output_file
    best_sol = population[1]
    f = Nullable()
    if !isnull(output_file)
        f = open(output_file, "w")
    end
    br = false
    while(t <= budget)
        sol = take!(result_set)
        if sol.value[2] >= 0
            push!(my_history, sol.value[1])
        else
            push!(my_history, best_sol.value[1])
        end
        if sol.value[1] < best_sol.value[1] && sol.value[2] >= 0
            best_sol = sol
        end
        has_better = false
        for i in 1:pop_size
            if parameter.isolationfunc(sol.x) != parameter.isolationfunc(population[i].x)
                continue
            else
                if (population[i].value[1] < sol.value[1] && population[i].value[2] >= sol.value[2]) ||
                    (population[i].value[1] <= sol.value[1] && population[i].value[2] > sol.value[2])
                    has_better = true
                    break
                end
            end
        end
        if !has_better
            Q = []
            for i = 1:pop_size
                if sol.value[1] <= population[i].value[1] && sol.value[2] >= population[i].value[2]
                    continue
                else
                    push!(Q, population[i])
                end
            end
            push!(Q, sol)
            population = Q
            pop_size = length(population)
        end
        t += 1
        time_log2 = now()
        time_pass = Dates.value(time_log2 - time_log1) / 1000
        if time_pass >= time_sum
            time_sum = time_sum + interval
            zoolog("time: $(time_pass) update $(t): best_solution value = $(best_sol.value[1])")
            if parameter.show_x == true
                zoolog("best_solution x = $(best_sol.x)")
            end
            str = "$(floor(time_pass)) $(best_sol.value)\n"
            if parameter.show_x == true
                str = string(str, best_sol.x, "\n")
            end
            if !isnull(f)
                write(f, str)
            end
            if !isnull(parameter.time_limit) && time_pass > parameter.time_limit
                zoolog("Exceed time limit: $(parameter.time_limit)")
                br = true
            end
         end
         index = rand(rng, 1:pop_size)
         new_x = pposs_mutation(population[index].x, length(population[index].x))
         put!(sample_set, Solution(x=new_x))
     end
     min_value = population[1].value[1]
     result_index = 1
     for p = 1:pop_size
         fitness = population[p].value
         if fitness[2] >= 0 && fitness[1] < min_value
             min_value = fitness[1]
             result_index = p
         end
     end
     # zoolog("update finish")
     if !isnull(f)
         close(f)
     end
     put!(asyn_result, population[result_index])
     put!(history, my_history)
     finish[1] = true
     index = rand(rng, 1:pop_size)
     new_x = pposs_mutation(population[index].x, length(population[index].x))
     put!(sample_set, Solution(x=new_x))
     # println("$(finish[1])")
end

function pposs_compute_fx(sol::Solution, ip_port, parameter::Parameter)
    ip, port = get_ip_port(ip_port)
    client = connect(ip, port)

    # send calculate info
    println(client, "client: calculate#")
    msg = readline(client)

    br = false
    if check_exception(msg) == true
        br = true
    end
    println(client, "pposs#")
    msg = readline(client)
    if check_exception(msg) == true
        br = true
    end
    # send objective_file:func
    if br == false
        smsg = string(parameter.objective_file, ":", parameter.func, ":", parameter.constraint, "#")
        println(client, smsg)
        msg = readline(client)
        if check_exception(msg) == true
            br = true
        end
    end

    # send x
    if br == false
        str = list2str(sol.x)
        println(client, str)
        receive = split(readline(client), " ")
        if check_exception(receive) == true
            br = true
        end
    end
    if br == false
        value = [parse(Float64, receive[1]), parse(Float64, receive[2])]
        # zoolog("$(value)")
        sol.value = value
    end
    return br
end


text:
# simulation settings

# kernel function: "CubicSpline", "WendlandC2"
kernel = "CubicSpline"

# kernel interpolation: "Standard", "IntegralApproach"
gradient = "Standard"

# sph formulation: "VanillaIce", "Lagrangian"
formulation = "VanillaIce"

# volume element: "mass", "U"
volume_element = "mass"

# time integrator: "RK2", "VL2", "RK3", "Leapfrog", "ModifiedLeapfrog", "symplectic Euler"
time_integrator = "ModifiedLeapfrog"

# output file name
outputfile = "out/out"

# The Courant-Friedrichs-Levi (CFL) number
cfl = 0.3

# time-dependent viscosity. if set to false, viscosity alpha is always 1.
time_dependent_viscosity = false
alpha_max = 1.5
alpha_min = 0.1

# upper limit of time step
dt_max = 1

# factor for evaluation of smoothing length
# this code uses Gadget's definition of smoothing length (smoothing kernel vanishes at 1h rather than at 2h)
eta_hsml = 4.8

# maximum number of SPH neighbour particles
Nngb = 32

# simulation time and output timing
t_start = 0
t_end = 0.1
output_interval = 0.001

# specific heat ratio
gamma = 1.4

# number of particles
Npart = 1000

# simulation box boundaries
x_min = -1.0
x_max = 1.0
x_center = 0

# problem setup: "standard", "strong", "custom". parameters are set below.
setup = "standard"

# plot density, pressure and velocity: true, false
plot_figure = true
x_min_plot = -1
x_max_plot = 1

# number of sample points for exact Riemann solver
Nsample_riemann = 1e4

# tolerance criterion in iterative process in exact Riemann solver and density estimate
TOL = 1e-6

# size of bin to measure L1 error
L1_binsize = 0.05

# standard shock tube (Saitoh & Makino, 2013; Springel et al., 2020)
if setup == "standard"
rho_left = 1
v_left = 0
P_left = 1
rho_right = 0.25
v_right = 0
P_right = 0.1795

# strong shock tube (Saitoh & Makino, 2013)
elseif setup == "strong"
rho_left = 1
v_left = 0
P_left = 1000
rho_right = 1
v_right = 0
P_right = 0.01

elseif setup == "custom"
# density and pressure in the left side
rho_left = 1
v_left = 0
P_left = 1000
# density and pressure in the right side
rho_right = 1
v_right = 0
P_right = 0.01
else
  @assert 0
end

# calculation

if plot_figure
  using GR
  using Plots
  ENV["GKSwstype"] = "100"
  gr()
end
using Printf

include("./src/begin.jl")
include("./src/evaluate.jl")
include("./src/exact_riemann.jl")
include("./src/finish.jl")
include("./src/force.jl")
include("./src/integrate.jl")
include("./src/kernel.jl")
include("./src/output.jl")
include("./src/run.jl")
include("./src/struct.jl")
include("./src/timestep.jl")

initialize()
run()
finish()


text:
mutable struct zmp_com_observer_state_t <: LCMType
    utime::Int64
    com::SVector{2, Float64}
    comd::SVector{2, Float64}
    ground_plane_height::Float64
end

@lcmtypesetup(zmp_com_observer_state_t)


text:
"""
onesite infinite dmrg
"""
@with_kw struct Idmrg1{} <: Algorithm
    tol_galerkin::Float64 = Defaults.tol
    tol_gauge::Float64 = Defaults.tolgauge
    maxiter::Int = Defaults.maxiter
    verbose::Bool = Defaults.verbose
end

function find_groundstate(ost::InfiniteMPS, ham::Hamiltonian,alg::Idmrg1,oenvs=environments(ost,ham))
    st = copy(ost);
    envs = IDMRGEnv(ost,oenvs);

    delta::Float64 = 2*alg.tol_galerkin;

    for topit in 1:alg.maxiter
        delta = 0.0;

        curc = st.CR[0];

        for pos = 1:length(st)
            (eigvals,vecs) = @closure eigsolve(st.AC[pos],1,:SR,Arnoldi()) do x
                ac_prime(x,pos,st,envs)
            end

            st.AC[pos] = vecs[1]
            (st.AL[pos],st.CR[pos]) = leftorth(vecs[1]);

            setleftenv!(envs,pos+1,transfer_left(leftenv(envs,pos),ham[pos],st.AL[pos],st.AL[pos]));
        end

        for pos in 1:length(st)-1
            setleftenv!(envs,pos+1,transfer_left(leftenv(envs,pos),ham[pos],st.AL[pos],st.AL[pos]));
        end

        for pos = length(st):-1:1

            (eigvals,vecs) = @closure eigsolve(st.AC[pos],1,:SR,Arnoldi()) do x
                ac_prime(x,pos,st,envs)
            end
            st.AC[pos] = vecs[1]
            (st.CR[pos-1],temp) = rightorth(_permute_tail(vecs[1]));
            st.AR[pos] = _permute_front(temp);

            setrightenv!(envs,pos-1,transfer_right(rightenv(envs,pos),ham[pos],st.AR[pos],st.AR[pos]));
        end

        for pos = length(st):-1:2
            setrightenv!(envs,pos-1,transfer_right(rightenv(envs,pos),ham[pos],st.AR[pos],st.AR[pos]));
        end

        delta = norm(curc-st.CR[0]);
        delta<alg.tol_galerkin && break;
        alg.verbose && @info "idmrg iter $(topit) err $(delta)"
    end

    nst = InfiniteMPS(st.AR[1:end],tol=alg.tol_gauge);
    nenvs = environments(nst,ham,tol=oenvs.tol,maxiter=oenvs.maxiter)
    return nst,nenvs,delta;
end

"""
twosite infinite dmrg
"""
@with_kw struct Idmrg2{} <: Algorithm
    tol_galerkin::Float64 = Defaults.tol
    tol_gauge::Float64 = Defaults.tolgauge
    maxiter::Int = Defaults.maxiter
    verbose::Bool = Defaults.verbose
    trscheme = truncerr(1e-6);
end

function find_groundstate(ost::InfiniteMPS, ham::Hamiltonian,alg::Idmrg2,oenvs=environments(ost,ham))
    length(ost) < 2 && throw(ArgumentError("unit cell should be >= 2"))

    st = copy(ost);
    envs = IDMRGEnv(ost,oenvs);

    delta::Float64 = 2*alg.tol_galerkin;

    for topit in 1:alg.maxiter
        delta = 0.0;

        curc = st.CR[0];

        #sweep from left to right
        for pos = 1:length(st)-1
            ac2 = st.AC[pos]*_permute_tail(st.AR[pos+1]);

            (eigvals,vecs) = @closure eigsolve(ac2,1,:SR,Arnoldi()) do x
                ac2_prime(x,pos,st,envs)
            end

            (al,c,ar,ϵ) = tsvd(vecs[1],trunc=alg.trscheme,alg=TensorKit.SVD())
            normalize!(c);

            st.AL[pos] = al
            st.CR[pos] = complex(c);
            st.AR[pos+1] = _permute_front(ar);
            st.AC[pos+1] = _permute_front(c*ar);

            setleftenv!(envs,pos+1,transfer_left(leftenv(envs,pos),ham[pos],st.AL[pos],st.AL[pos]));
            setrightenv!(envs,pos,transfer_right(rightenv(envs,pos+1),ham[pos+1],st.AR[pos+1],st.AR[pos+1]))
        end

        #update the edge
        @tensor ac2[-1 -2;-3 -4] := st.AC[end][-1,-2,1]*inv(st.CR[0])[1,2]*st.AL[1][2,-3,3]*st.CR[1][3,-4]
        (eigvals,vecs) = @closure eigsolve(ac2,1,:SR,Arnoldi()) do x
            ac2_prime(x,0,st,envs)
        end
        (al,c,ar,ϵ) = tsvd(vecs[1],trunc=alg.trscheme,alg=TensorKit.SVD())
        normalize!(c);

        st.AC[end] = al*c;
        st.AL[end] = al;
        st.CR[end] = complex(c);
        st.AR[1] = _permute_front(ar);
        st.AC[1] = _permute_front(c*ar);
        st.AL[1] = _permute_front(c*ar)*inv(st.CR[1]);

        curc = complex(c);

        #update environments
        setleftenv!(envs,1,transfer_left(leftenv(envs,0),ham[0],st.AL[0],st.AL[0]));
        setrightenv!(envs,0,transfer_right(rightenv(envs,1),ham[1],st.AR[1],st.AR[1]));

        #sweep from right to left
        for pos = length(st)-1:-1:1
            ac2 = st.AL[pos]*_permute_tail(st.AC[pos+1]);

            (eigvals,vecs) = @closure eigsolve(ac2,1,:SR,Arnoldi()) do x
                ac2_prime(x,pos,st,envs)
            end

            (al,c,ar,ϵ) = tsvd(vecs[1],trunc=alg.trscheme,alg=TensorKit.SVD())
            normalize!(c);

            st.AL[pos] = al
            st.AC[pos] = al*c
            st.CR[pos] = complex(c);
            st.AR[pos+1] = _permute_front(ar);
            st.AC[pos+1] = _permute_front(c*ar);

            setrightenv!(envs,pos,transfer_right(rightenv(envs,pos+1),ham[pos+1],st.AR[pos+1],st.AR[pos+1]))
            setleftenv!(envs,pos+1,transfer_left(leftenv(envs,pos),ham[pos],st.AL[pos],st.AL[pos]));
        end

        #update the edge
        @tensor ac2[-1 -2;-3 -4] :=  st.CR[end-1][-1,1]*st.AR[end][1,-2,2]*inv(st.CR[end])[2,3]*st.AC[1][3,-3,-4]
        (eigvals,vecs) = @closure eigsolve(ac2,1,:SR,Arnoldi()) do x
            ac2_prime(x,0,st,envs)
        end
        (al,c,ar,ϵ) = tsvd(vecs[1],trunc=alg.trscheme,alg=TensorKit.SVD())
        normalize!(c);

        st.AR[end] = _permute_front(inv(st.CR[end-1])*_permute_tail(al*c))
        st.AL[end] = al;
        st.CR[end] = complex(c);
        st.AR[1] = _permute_front(ar);
        st.AC[1] = _permute_front(c*ar);

        setleftenv!(envs,1,transfer_left(leftenv(envs,0),ham[0],st.AL[0],st.AL[0]));
        setrightenv!(envs,0,transfer_right(rightenv(envs,1),ham[1],st.AR[1],st.AR[1]));

        #update error
        d1 = Diagonal(convert(Array,curc));
        d2 = Diagonal(convert(Array,complex(c)));
        minl = min(length(d1),length(d2));
        delta = norm(d1[1:minl]-d2[1:minl])

        alg.verbose && @info "idmrg iter $(topit) err $(delta)"

        delta<alg.tol_galerkin && break;

    end

    nst = InfiniteMPS(st.AR[1:end],tol=alg.tol_gauge);
    nenvs = environments(nst,ham,tol=oenvs.tol,maxiter=oenvs.maxiter)
    return nst,nenvs,delta;
end


text:
# This file is a part of RadiationDetectorSignals.jl, licensed under the MIT License (MIT).

using Test
using RadiationDetectorSignals
import Documenter

Documenter.DocMeta.setdocmeta!(
    RadiationDetectorSignals,
    :DocTestSetup,
    :(using RadiationDetectorSignals);
    recursive=true,
)
Documenter.doctest(RadiationDetectorSignals)


text:
# Autogenerated wrapper script for Rclone_jll for aarch64-linux-gnu
export rclone

JLLWrappers.@generate_wrapper_header("Rclone")
JLLWrappers.@declare_executable_product(rclone)
function __init__()
    JLLWrappers.@generate_init_header()
    JLLWrappers.@init_executable_product(
        rclone,
        "bin/rclone",
    )

    JLLWrappers.@generate_init_footer()
end  # __init__()


text:
function hp_row_sort!(ds::Dataset, cols = names(ds, Union{Missing, Number}); kwargs...)
    colsidx = index(ds)[cols]
    T = mapreduce(eltype, promote_type, eachcol(ds)[colsidx])
    m = Matrix{T}(ds[!, colsidx])
    Threads.@threads for i in 1:size(m, 1)
        @views sort!(m[i, :]; kwargs...)
    end
    # TODO no parallel is needed here to minimise memory
    for i in 1:length(colsidx)
        _columns(ds)[colsidx[i]] = m[:, i]
    end
    removeformat!(ds, cols)
    any(index(ds).sortedcols .∈ Ref(colsidx)) && _reset_grouping_info!(ds)
    _modified(_attributes(ds))
    ds
end

"""
    row_sort!(ds::AbstractDataset[, cols]; kwargs...)
    sort `cols` in each row.
"""
function hp_row_sort(ds::AbstractDataset, cols = names(ds, Union{Missing, Number}); kwargs...)
    dscopy = copy(ds)
    hp_row_sort!(dscopy, cols; kwargs...)
    dscopy
end

function hp_row_generic(ds::AbstractDataset, f::Function, cols::MultiColumnIndex)
    colsidx = multiple_getindex(index(ds), cols)
    if length(colsidx) == 2
        try
            allowmissing(f.(_columns(ds)[colsidx[1]], _columns(ds)[colsidx[2]]))
        catch e
            if e isa MethodError
                _hp_row_generic(ds, f, colsidx)
            else
                rethrow(e)
            end
        end
    else
        _hp_row_generic(ds, f, colsidx)
    end
end

function _hp_row_generic(ds::AbstractDataset, f::Function, colsidx)
    T = mapreduce(eltype, promote_type, view(_columns(ds),colsidx))
    inmat = Matrix{T}(undef, length(colsidx), min(1000, nrow(ds)))

    all_data = view(_columns(ds), colsidx)
    _fill_matrix!(inmat, all_data, 1:min(1000, nrow(ds)), colsidx)
    res_temp = allowmissing(f.(eachcol(inmat)))
    if !(typeof(res_temp) <:  AbstractVector)
        throw(ArgumentError("output of `f` must be a vector"))
    end

    # if length(res_temp[1]) > 1
    #     throw(ArgumentError("The matrix output is not supported"))
    #     res = similar(res_temp, nrow(ds), size(res_temp,2))
    # elseif length(res_temp[1]) == 1
    res = similar(res_temp, nrow(ds))
    # else
        # throw(ArgumentError("the result cannot be with zero dimension"))

    if nrow(ds)>1000
        if size(res, 2) == 1
            view(res, 1:1000) .= res_temp
            _hp_row_generic_vec!(res, ds, f, colsidx, Val(T))
        else
            view(res, 1:1000, :) .= res_temp
            # not implemented yet
            _hp_row_generic_mat!(res, ds, f, colsidx)
        end
    else
        return res_temp
    end
    return res
end

function _hp_row_generic_vec!(res, ds, f, colsidx, ::Val{T}) where T
    nt = Threads.nthreads()
    loopsize = div(length(res) - 1000, 1000)
    all_data = view(_columns(ds), colsidx)
    if loopsize == 0
        st = 1001
        en = length(res)
        inmat = Matrix{T}(undef, length(colsidx), en - st + 1)
        _fill_matrix!(inmat, all_data, st:en, colsidx)
        view(res, st:en) .= f.(eachcol(inmat))
        return
    end
    max_cz = length(res) - 1000 - (loopsize - 1)*1000
    inmat_all = [Matrix{T}(undef, length(colsidx), max_cz) for i in 1:nt]
    # make sure that the variable inside the loop are not the same as the out of scope one
    Threads.@threads for i in 1:loopsize
        t_st = i*1000 + 1
        i == loopsize ? t_en = length(res) : t_en = (i+1)*1000
        _fill_matrix!(inmat_all[Threads.threadid()], all_data, t_st:t_en, colsidx)
        for k in t_st:t_en
            res[k] = f(view(inmat_all[Threads.threadid()], :, k - t_st + 1))
        end
    end
end


text:
using Dates
using DelimitedFiles

using HTTP

import ..Datasets: dataset, datafile, reader, Mode, Test

diagnostic_init() = register(DataDep(
    "GLUE-Diagnostic",
    """
    GLUE Diagnostics Datasets
    """,
    "https://storage.googleapis.com/mtl-sentence-representations.appspot.com/tsvsWithoutLabels%2FAX.tsv?GoogleAccessId=firebase-adminsdk-0khhl@mtl-sentence-representations.iam.gserviceaccount.com&Expires=2498860800&Signature=DuQ2CSPt2Yfre0C%2BiISrVYrIFaZH1Lc7hBVZDD4ZyR7fZYOMNOUGpi8QxBmTNOrNPjR3z1cggo7WXFfrgECP6FBJSsURv8Ybrue8Ypt%2FTPxbuJ0Xc2FhDi%2BarnecCBFO77RSbfuz%2Bs95hRrYhTnByqu3U%2FYZPaj3tZt5QdfpH2IUROY8LiBXoXS46LE%2FgOQc%2FKN%2BA9SoscRDYsnxHfG0IjXGwHN%2Bf88q6hOmAxeNPx6moDulUF6XMUAaXCSFU%2BnRO2RDL9CapWxj%2BDl7syNyHhB7987hZ80B%2FwFkQ3MEs8auvt5XW1%2Bd4aCU7ytgM69r8JDCwibfhZxpaa4gd50QXQ%3D%3D",
    "0e13510b1bb14436ff7e2ee82338f0efb0133ecf2e73507a697dc210db3f05fd";
    fetch_method=(url, localdir)-> begin
      format_progress(x) = round(x, digits=4)
      format_bytes(x) = !isfinite(x) ? "∞ B" : Base.format_bytes(x)
      format_seconds(x) = "$(round(x; digits=2)) s"
      format_bytes_per_second(x) = format_bytes(x) * "/s"

      filepath = joinpath(localdir, "diagnostic.tsv")

      HTTP.open("GET", url) do stream
        resp = HTTP.startread(stream)

        total_bytes = parse(Float64, HTTP.header(resp, "Content-Length"))
        downloaded_bytes = 0
        start_time = now()
        prev_time = now()

        function report_callback()
            prev_time = now()
            taken_time = (prev_time - start_time).value / 1000 # in seconds
            average_speed = downloaded_bytes / taken_time
            remaining_bytes = total_bytes - downloaded_bytes
            remaining_time = remaining_bytes / average_speed
            completion_progress = downloaded_bytes / total_bytes

            @info("Downloading",
                  source=url,
                  dest = filepath,
                  progress = completion_progress |> format_progress,
                  time_taken = taken_time |> format_seconds,
                  time_remaining = remaining_time |> format_seconds,
                  average_speed = average_speed |> format_bytes_per_second,
                  downloaded = downloaded_bytes |> format_bytes,
                  remaining = remaining_bytes |> format_bytes,
                  total = total_bytes |> format_bytes,
                  )
        end

        Base.open(filepath, "w") do fh
            while(!eof(stream))
                downloaded_bytes += write(fh, readavailable(stream))
                if now() - prev_time > Millisecond(1000*2)
                    report_callback()
                end
            end
        end
        report_callback()
      end
      filepath
    end
))

struct Diagnostic <: Dataset end

function testfile(::Diagnostic)
    sets, header = readlm(datadep"GLUE-Diagnostic/diagnostic.tsv", '\t', String; header=true)
    [selectdim(sets, 2, i) for i = (2,3)]
end


text:
using Celeste.Model: CanonicalParams, ids, NUM_COLOR_COMPONENTS

using Celeste.DeterministicVI.ElboMaximize: elbo_constraints

using Celeste.DeterministicVI.ConstraintTransforms: Constraint, BoxConstraint,
                                    SimplexConstraint, ParameterConstraint, ConstraintBatch,
                                    inv_logit, logit, to_free,
                                    to_free!, to_bound, to_bound!, simplexify, unsimplexify,
                                    enforce, enforce!, allocate_free_params,
                                    TransformDerivatives, differentiate!,
                                    propagate_derivatives!

import Celeste.Transform, Celeste.SensitiveFloats

using Base.Test

srand(1)

# Test helpers

function valid_boxed_parameters!(bound, indices, c::BoxConstraint)
    population = linspace(c.lower, c.upper)
    for src in eachindex(bound), i in indices
        bound[src][i] = rand(population)
    end
    return bound
end

function valid_simplexed_parameters!(bound, indices, c::SimplexConstraint)
    @assert length(indices) == c.n
    for src in eachindex(bound)
        params = view(bound[src], indices)
        for i in eachindex(params)
            params[i] = min(rand() + nextfloat(c.lower), 1.0)
        end
        normalize!(params, 1)
        @assert all(c.lower .< params .< 1.0)
        @assert isapprox(sum(params), 1.0)
    end
    return bound
end

function correct_to_free(bound, bound_range, c::SimplexConstraint)
    log_last = log(unsimplexify(bound[bound_range[end]], c))
    free = similar(bound, length(bound_range) - 1)
    for i in 1:length(free)
        free[i] = c.scale * (log(unsimplexify(bound[bound_range[i]], c)) - log_last)
    end
    return free
end

function star_constraints(bound; loc_width = 1.5e-3, loc_scale = 1.0)
    n_sources = length(bound)
    boxes = Vector{Vector{ParameterConstraint{BoxConstraint}}}(n_sources)
    simplexes = Vector{Vector{ParameterConstraint{SimplexConstraint}}}(n_sources)
    for src in 1:n_sources
        i1, i2 = ids.pos[1], ids.pos[2]
        u1, u2 = bound[src][i1], bound[src][i2]
        boxes[src] = [
            ParameterConstraint(BoxConstraint(u1 - loc_width, u1 + loc_width, loc_scale), i1),
            ParameterConstraint(BoxConstraint(u2 - loc_width, u2 + loc_width, loc_scale), i2),
            ParameterConstraint(BoxConstraint(-1.0, 10.0, 1.0), ids.flux_loc[1]),
            ParameterConstraint(BoxConstraint(-10.0, 10.0, 1.0), ids.color_mean[:, 1]),
            ParameterConstraint(BoxConstraint(1e-4, 1.0, 1.0), ids.color_var[:, 1])
        ]
        simplexes[src] = [
            ParameterConstraint(SimplexConstraint(0.005, 1.0, 2), ids.is_star)
            ParameterConstraint(SimplexConstraint(0.01/NUM_COLOR_COMPONENTS, 1.0, NUM_COLOR_COMPONENTS), ids.k[:, 1])
        ]
    end
    return ConstraintBatch(boxes, simplexes)
end

@testset "constraints" begin

    @testset "constraint types" begin

        @test BoxConstraint <: Constraint
        @test SimplexConstraint <: Constraint
        @test ParameterConstraint <: Constraint

        n = 3
        lower = 0.01 / n
        scale = rand()

        box = BoxConstraint(lower, 2*lower, scale)
        @test box.lower === lower
        @test box.upper === 2*lower
        @test box.scale === scale

        @test_throws AssertionError BoxConstraint(lower, lower/2, scale)
        @test_throws AssertionError BoxConstraint(-Inf, lower, scale)
        @test_throws AssertionError BoxConstraint(lower, Inf, scale)
        @test_throws AssertionError BoxConstraint(lower, 2*lower, 0.0)
        @test_throws AssertionError BoxConstraint(lower, 2*lower, -scale)
        @test_throws AssertionError BoxConstraint(lower, 2*lower, Inf)

        simplex = SimplexConstraint(lower, scale, n)
        @test simplex.lower === lower
        @test simplex.scale === scale
        @test simplex.n === n

        @test_throws AssertionError SimplexConstraint(lower, scale, 0)
        @test_throws AssertionError SimplexConstraint(lower, scale, 1)
        @test_throws AssertionError SimplexConstraint(-lower, scale, 1)
        @test_throws AssertionError SimplexConstraint(2*lower, scale, 1)
        @test_throws AssertionError SimplexConstraint(lower, 0.0, 2)
        @test_throws AssertionError SimplexConstraint(lower, 0.0, n)
        @test_throws AssertionError SimplexConstraint(lower, -scale, n)
        @test_throws AssertionError SimplexConstraint(lower, Inf, n)

        c = ParameterConstraint(box, n)
        @test c.constraint === box
        @test c.indices === n:n

        c = ParameterConstraint(box, 1:n)
        @test c.constraint === box
        @test c.indices === 1:n

        c = ParameterConstraint(simplex, n)
        @test c.constraint === simplex
        @test c.indices === n:n

        c = ParameterConstraint(simplex, 1:n)
        @test c.constraint === simplex
        @test c.indices === 1:n
    end

    @testset "to_bound[!]/to_free[!] BoxConstraint" begin
        x = rand()
        @test inv_logit(x) === -log(1.0 / x - 1)
        @test logit(x) === 1.0 / (1.0 + exp(-x))
        @test isapprox(logit(inv_logit(x)), x)

        lower, scale = rand(2)

        c = BoxConstraint(lower, 2*lower, scale)
        bound = 1.5 * c.lower

        @test_throws AssertionError to_free(c.lower, c)
        @test_throws AssertionError to_free(c.upper, c)
        @test to_free(bound, c) === inv_logit((bound - c.lower) / (c.upper - c.lower)) * c.scale

        free = to_free(bound, c)

        @test_throws AssertionError to_bound(-Inf, c)
        @test_throws AssertionError to_bound(Inf, c)
        @test to_bound(free, c) === bound
    end

    @testset "to_bound[!]/to_free[!] SimplexConstraint" begin
        x = rand()
        c = SimplexConstraint(0.01/5, rand(), 5)

        @test simplexify(x, c) === (1 - c.n * c.lower) * x + c.lower
        @test unsimplexify(x, c) === (x - c.lower) / (1 - c.n * c.lower)
        @test unsimplexify(simplexify(x, c), c) === x

        n_params = 10
        n_sources = 3
        bound_range = 1:5
        free_range = 3:6
        bound = [rand(n_params) for _ in 1:n_sources]
        valid_simplexed_parameters!(bound, bound_range, c)
        old_bound = deepcopy(bound)
        free = [zeros(n_params - 1) for _ in 1:n_sources]

        for src in 1:n_sources
            to_free!(free[src], free_range, bound[src], bound_range, c)
            @test (free[src][free_range] ==
                   correct_to_free(bound[src], bound_range, c))

            for i in setdiff(1:n_sources, src)
                for j in setdiff(1:length(free[src]), free_range)
                    @test free[i][j] === 0.0
                end
            end

            bound[src][bound_range] = 0.0
            to_bound!(bound[src], bound_range, free[src], free_range, c)
            @test isapprox(old_bound[src], bound[src])
            @test all(old_bound[i] == bound[i]
                      for i in setdiff(1:n_sources, src))

            copy!(bound[src], old_bound[src])
            fill!(free[src], 0.0)
        end
    end

    @testset "ParameterConstraint BoxConstraint" begin
        n_params = 10
        n_sources = 3
        lower, scale = rand(2)
        c = ParameterConstraint(BoxConstraint(lower, 2*lower, scale), 3:6)
        bound = collect(rand(n_params) for _ in 1:n_sources)
        valid_boxed_parameters!(bound, c.indices, c.constraint)
        old_bound = deepcopy(bound)
        free = collect(zeros(n_params - 1) for _ in 1:n_sources)

        for src in 1:n_sources
            free_index = to_free!(free[src], bound[src], c, 1)
            @test free_index == (length(c.indices) + 1)
            for i in 1:length(c.indices)
                @test free[src][i] == to_free(bound[src][c.indices[i]], c.constraint)
            end

            free_index = to_bound!(bound[src], free[src], c, 1)
            @test free_index == (length(c.indices) + 1)
            @test isapprox(bound[src], old_bound[src])
            @test all(old_bound[i] == bound[i] for i in setdiff(1:n_sources, src))

            copy!(bound[src], old_bound[src])
            fill!(free[src], 0.0)
        end
    end

    @testset "ParameterConstraint SimplexConstraint" begin
        n_params = 10
        n_sources = 3
        c = ParameterConstraint(SimplexConstraint(0.01/5, rand(), 5), 3:7)
        bound = collect(rand(n_params) for _ in 1:n_sources)
        valid_simplexed_parameters!(bound, c.indices, c.constraint)
        old_bound = deepcopy(bound)
        free = collect(zeros(n_params - 1) for _ in 1:n_sources)

        for src in 1:n_sources
            free_index = to_free!(free[src], bound[src], c, 1)
            @test free_index == length(c.indices)
            @test free[src][1:(free_index-1)] == correct_to_free(bound[src], c.indices, c.constraint)

            free_index = to_bound!(bound[src], free[src], c, 1)
            @test free_index == length(c.indices)
            @test isapprox(bound[src], old_bound[src])
            @test all(old_bound[i] == bound[i] for i in setdiff(1:n_sources, src))

            copy!(bound[src], old_bound[src])
            fill!(free[src], 0.0)
        end
    end

    @testset "enforce[!] BoxConstraint" begin
        lower, scale = rand(2)
        c = BoxConstraint(lower, 2*lower, scale)

        @test enforce(c.lower, c) === nextfloat(c.lower)
        @test enforce(c.upper, c) === prevfloat(c.upper)
        @test enforce(-Inf, c) === nextfloat(c.lower)
        @test enforce(Inf, c) === prevfloat(c.upper)
        @test enforce(1.5 * c.lower, c) === 1.5 * c.lower
    end

    @testset "enforce[!] SimplexConstraint" begin
        n = 10
        lower = 0.01 / n
        c = SimplexConstraint(lower, rand(), n)
        bound = valid_simplexed_parameters!([rand(n)], 1:n, c)[1]
        old_bound = copy(bound)

        enforce!(bound, c)
        @test bound == old_bound

        bound[2] = 1.0
        bound[4] = c.lower
        bound[6] = Inf
        bound[7] = c.lower * 1.5

        enforce!(bound, c)
        @test all(c.lower .< bound .< 1.0)
        @test isapprox(sum(bound), 1.0)

        new_bound = copy(old_bound)
        new_bound[2] = prevfloat(1.0)
        new_bound[4] = nextfloat(c.lower)
        new_bound[6] = prevfloat(1.0)
        new_bound[7] = c.lower * 1.5
        rescale = (1 - c.n * c.lower) / (sum(new_bound) - c.n * c.lower)
        new_bound[:] = [nextfloat(c.lower) + rescale * (i - c.lower) for i in new_bound]
        @test bound == new_bound
    end

# ParameterConstraint #
#---------------------#

# BoxConstraint
    @testset "ParameterConstraint BoxConstraint" begin
        n_params = 10
        n_sources = 3
        lower, scale = rand(2)
        c = BoxConstraint(lower, 2*lower, scale)
        pc = ParameterConstraint(c, 3:7)
        bound = collect(rand(n_params) for _ in 1:n_sources)
        valid_boxed_parameters!(bound, pc.indices, c)
        old_bound = deepcopy(bound)

        for src in 1:n_sources
            enforce!(bound[src], pc)

            for i in 1:n_sources
                @test old_bound[i] == bound[i]
                bound[i][3] = c.lower
                bound[i][4] = c.upper
                bound[i][5] = -Inf
                bound[i][6] = Inf
                bound[i][7] = 1.5 * c.lower
            end

            enforce!(bound[src], pc)

            @test bound[src][3] === nextfloat(c.lower)
            @test bound[src][4] === prevfloat(c.upper)
            @test bound[src][5] === nextfloat(c.lower)
            @test bound[src][6] === prevfloat(c.upper)
            @test bound[src][7] === 1.5 * c.lower

            for other_src in setdiff(1:n_sources, src)
                @test bound[other_src][3] === c.lower
                @test bound[other_src][4] === c.upper
                @test bound[other_src][5] === -Inf
                @test bound[other_src][6] === Inf
                @test bound[other_src][7] === 1.5 * c.lower
            end

            for i in 1:n_sources
                for other_param in (1, 2, 8, 9, 10)
                    @test bound[i][other_param] == old_bound[i][other_param]
                end
                copy!(bound[i], old_bound[i])
            end
        end
    end

    @testset "ParameterConstraint SimplexConstraint" begin
        n_params = 10
        n_sources = 3
        n = 5
        lower = 0.01 / n
        c = SimplexConstraint(lower, rand(), n)
        pc = ParameterConstraint(c, 3:7)
        bound = collect(rand(n_params) for _ in 1:n_sources)
        valid_simplexed_parameters!(bound, pc.indices, c)
        old_bound = deepcopy(bound)

        for src in 1:n_sources
            enforce!(bound[src], pc)

            for i in 1:n_sources
                @test old_bound[i] == bound[i]
                bound[i][3] = 1.0
                bound[i][4] = c.lower
                bound[i][6] = Inf
                bound[i][7] = c.lower * 1.5
            end

            new_bound = deepcopy(bound)
            enforce!(bound[src], pc)

            @test all(c.lower .< bound[src][pc.indices] .< 1.0)
            @test isapprox(sum(bound[src][pc.indices]), 1.0)

            new_bound[src][3] = prevfloat(1.0)
            new_bound[src][4] = nextfloat(c.lower)
            new_bound[src][6] = prevfloat(1.0)
            new_bound[src][7] = c.lower * 1.5
            rescale = (1 - c.n * c.lower) / (sum(new_bound[src][pc.indices]) - c.n * c.lower)
            for i in pc.indices
                new_bound[src][i] = nextfloat(c.lower) + rescale * (new_bound[src][i] - c.lower)
            end

            for i in 1:n_sources
                @test new_bound[i] == bound[i]
                copy!(bound[i], old_bound[i])
            end
        end
    end

    @testset "ConstraintBatch" begin
        loc_width, loc_scale = rand(), rand()
        n_sources = 4
        bound = collect(rand(length(CanonicalParams)) for _ in 1:n_sources)
        constraints = elbo_constraints(bound, loc_width, loc_scale)

        for src in 1:n_sources
            u1_box = constraints.boxes[src][1].constraint
            u2_box = constraints.boxes[src][2].constraint
            u1, u2 = bound[src][ids.pos[1]], bound[src][ids.pos[2]]
            @test u1_box === BoxConstraint(u1 - loc_width, u1 + loc_width, loc_scale)
            @test u2_box === BoxConstraint(u2 - loc_width, u2 + loc_width, loc_scale)
        end

        free = allocate_free_params(bound, constraints)

        @test length(free) == length(bound) == n_sources

        free_length = length(CanonicalParams) - length(constraints.simplexes[1])

        for src in 1:n_sources
            @test length(free[src]) == free_length
        end

        enforce!(bound, constraints)
        old_bound = deepcopy(bound)
        to_free!(free, bound, constraints)
        old_free = deepcopy(free)
        to_bound!(bound, free, constraints)

        @test all(isapprox(old_bound[src], bound[src]) for src in 1:n_sources)

        # This test currently fails, but that's probably inevitable as
        # long as doing a round-trip unconstrain/reconstrain on a
        # parameter set yields parameters that are only appromixately
        # equal to the original parameters (which seems like a
        # tolerable behavior)
        to_free!(free, bound, constraints)
        @test_broken all(isapprox(old_free[src], free[src])
                         for src in 1:n_sources)
    end

    @testset "star-only parameter transformations" begin
        n_sources = 10

        bound = collect(rand(length(CanonicalParams)) for _ in 1:n_sources)

        constraints = star_constraints(bound)

        free = allocate_free_params(bound, constraints)

        @test length(free) == n_sources
        @test all(length(free[src]) == 19 for src in 1:n_sources)

        enforce!(bound, constraints)
        old_bound = deepcopy(bound)
        to_free!(free, bound, constraints)
        old_free = deepcopy(free)
        to_bound!(bound, free, constraints)
        @test all(isapprox(old_bound[src], bound[src]) for src in 1:n_sources)
    end

    @testset "check derivatives against old Transform" begin
        ea, vp, catalog = SampleData.gen_n_body_dataset(1)

        bound = vp
        constraints = elbo_constraints(bound)
        free = allocate_free_params(bound, constraints)

        enforce!(bound, constraints)
        to_free!(free, bound, constraints)

        derivs = TransformDerivatives(bound, free)

        differentiate!((y, x) -> to_bound!(y, x, constraints, 1), derivs, free[1])

        dt = Transform.get_mp_transform(vp, [1])
        td = Transform.get_transform_derivatives(vp, [1], dt.bounds)

        @test isapprox(td.dparam_dfree, derivs.jacobian)

        for i in 1:size(derivs.hessian, 1)
            @test isapprox(td.d2param_dfree2[i], derivs.hessian[i, :, :])
        end
    end

    @testset "check SensitiveFloat propagation" begin
        ea, vp, catalog = SampleData.gen_n_body_dataset(1)

        bound = vp
        constraints = elbo_constraints(bound)
        free = allocate_free_params(bound, constraints)

        enforce!(bound, constraints)
        to_free!(free, bound, constraints)

        derivs = TransformDerivatives(bound, free)
        dt = Transform.get_mp_transform(vp, [1])

        sf_bound1 = SensitiveFloats.SensitiveFloat{Float64}(length(bound[1]), 1, true, true)
        sf_bound1.v[] = rand()
        rand!(sf_bound1.d)
        rand!(sf_bound1.h)
        sf_bound2 = SensitiveFloats.SensitiveFloat{Float64}(length(bound[1]), 1, true, true)
        sf_bound2.v[] = sf_bound1.v[]
        copy!(sf_bound2.d, sf_bound1.d)
        copy!(sf_bound2.h, sf_bound1.h)
        sf_free1 = SensitiveFloats.SensitiveFloat{Float64}(length(free[1]), 1, true, true)
        sf_free2 = SensitiveFloats.SensitiveFloat{Float64}(length(free[1]), 1, true, true)

        propagate_derivatives!(to_bound!, sf_bound1, sf_free1, free, constraints, derivs)
        Transform.transform_sensitive_float!(dt, sf_free2, sf_bound2, vp, [1])

        @test isapprox(sf_free1.v[], sf_free2.v[])
        @test isapprox(sf_free1.d, sf_free2.d)
        @test isapprox(sf_free1.h, sf_free2.h)
    end

end # testset


text:
const DEBUG = true

function make_design(rs::RunSummary, lag::Float64)
    if !rs.unit_time
        rs = unit_run_sum(rs)
    end

    unit_time = rs.time[1]

    if lag < unit_time
        error("The 'lag' must be greater than the unit_time (set to 1.0 by default).")
    end

    n_lag = Int(floor(lag / unit_time))

    inv_lag = 1.0 / lag

    n = length(rs)

    X = zeros(n, count_covs()) # intercept, distance, altitude
    # X[:, 1] .= 1.0

    current_dist = 0.0
    current_alt = 0.0

    for i = 1:minimum([n_lag, n])
        current_dist += rs.dist[i]
        current_alt += rs.alt[i]
        X[i, 1] = current_dist * inv_lag
        X[i, 2] = current_alt * inv_lag
    end

    for i = (n_lag + 1):n
        current_dist += rs.dist[i] - rs.dist[i - n_lag]
        current_alt += rs.alt[i] - rs.alt[i - n_lag]

        X[i, 1] = current_dist * inv_lag
        X[i, 2] = current_alt * inv_lag
    end

    return X

end

function regress(y::Vector{Float64}, X::Matrix{Float64})
    n, p = size(X)
    @assert length(y) == n

    # TODO make more memory efficient if necessary

    xty = X' * y
    xtx = X' * X
    β = xtx \ xty

    res = X * β - y

    ssr = res' * res

    return β, ssr
end

function count_covs()
    return 2 #TODO: need a better way to specify covariates in the model
                # currently uses only speed and altitude
end

function multi_regress(summaries::Vector{RunSummary}, lag::Float64)
    n_cov = count_covs()
    p_base = n_cov
    p_total = p_base + length(summaries)

    XtX = zeros(p_total, p_total)
    Xty = zeros(p_total)

    n = length(summaries)

    t_unit = 0.0
    act_rng = 1:n_cov

    for i = 1:n
        if DEBUG
            println("Starting $i of $n")
        end

        rs = summaries[i]
        if !rs.unit_time
            rs = unit_run_sum(rs)
        end

        if i == 1
            t_unit = rs.time[1]
        else
            @assert rs.time[1] == t_unit
        end

        X = make_design(rs, lag)

        design_ind = p_base + i

        #TODO: make below memory efficient
        XtX[1:p_base, 1:p_base] .+= X' * X
        sum_X = vec(sum(X, dims=1))

        XtX[act_rng, design_ind] .= sum_X
        XtX[design_ind, act_rng] .= sum_X
        XtX[design_ind, design_ind] = Float64(length(rs))

        Xty[1:p_base] .+= X' * rs.hr
        Xty[design_ind] = sum(rs.hr) #TODO: can get this from intercept of Xty?

    end

    return XtX \ Xty

end