text:
__precompile__()

module NLopt

export Opt, NLOPT_VERSION, algorithm, algorithm_name, ForcedStop,
       lower_bounds!, lower_bounds, upper_bounds!, upper_bounds, stopval!, stopval, ftol_rel!, ftol_rel, ftol_abs!, ftol_abs, xtol_rel!, xtol_rel, xtol_abs!, xtol_abs, maxeval!, maxeval, maxtime!, maxtime, force_stop!, force_stop, force_stop!, population!, population, vector_storage!, vector_storage, initial_step!, initial_step, default_initial_step!, local_optimizer!,
       min_objective!, max_objective!, equality_constraint!, inequality_constraint!, remove_constraints!,
       optimize!, optimize

using Compat
import Compat: String

import Base.ndims, Base.copy, Base.convert, Base.show

import MathProgBase.SolverInterface
import MathProgBase.SolverInterface.optimize!

const depfile = joinpath(dirname(@__FILE__),"..","deps","deps.jl")
if isfile(depfile)
    include(depfile)
else
    error("NLopt not properly installed. Please run Pkg.build(\"NLopt\")")
end

############################################################################
# separate initializations that must occur at runtime, for precompilation

function __init__()
    # cfunction returns a raw pointer, so it must be called at runtime
    global const munge_callback_ptr = cfunction(munge_callback, Ptr{Void},
                                                (Ptr{Void}, Ptr{Void}))
    global const nlopt_callback_wrapper_ptr =
        cfunction(nlopt_callback_wrapper,
                  Cdouble, (Cuint, Ptr{Cdouble}, Ptr{Cdouble}, Ptr{Void}))
    global const nlopt_vcallback_wrapper_ptr =
        cfunction(nlopt_vcallback_wrapper, Void,
                  (Cuint, Ptr{Cdouble}, Cuint, Ptr{Cdouble}, Ptr{Cdouble},
                   Ptr{Void}))

    # get the version of NLopt at runtime, not compile time
    global const NLOPT_VERSION = version()
end

############################################################################
# Mirrors of NLopt's C enum constants:

# Problem: the NLopt API uses various enum types for both arguments and
# return values, but sizeof(enum) is compiler- and architecture dependent.
# By default, gcc stores enums as int (or unsigned int) unless -fshort-enums
# is used in the compiler options, but there are some architectures (ARM?)
# where -fshort-enums is the default.  [The MSVC and Intel compilers also
# apparently default to sizeof(enum)==sizeof(int).]  So, if Julia is
# ever ported to ARM, this may need to be fixed (since all of these enums
# will be packed into a single byte on -fshort-enums architectures).
typealias Cenum Cint
cenum(x) = convert(Cenum, x)

# enum nlopt_algorithm
const GN_DIRECT = cenum(0)
const GN_DIRECT_L = cenum(1)
const GN_DIRECT_L_RAND = cenum(2)
const GN_DIRECT_NOSCAL = cenum(3)
const GN_DIRECT_L_NOSCAL = cenum(4)
const GN_DIRECT_L_RAND_NOSCAL = cenum(5)
const GN_ORIG_DIRECT = cenum(6)
const GN_ORIG_DIRECT_L = cenum(7)
const GD_STOGO = cenum(8)
const GD_STOGO_RAND = cenum(9)
const LD_LBFGS_NOCEDAL = cenum(10)
const LD_LBFGS = cenum(11)
const LN_PRAXIS = cenum(12)
const LD_VAR1 = cenum(13)
const LD_VAR2 = cenum(14)
const LD_TNEWTON = cenum(15)
const LD_TNEWTON_RESTART = cenum(16)
const LD_TNEWTON_PRECOND = cenum(17)
const LD_TNEWTON_PRECOND_RESTART = cenum(18)
const GN_CRS2_LM = cenum(19)
const GN_MLSL = cenum(20)
const GD_MLSL = cenum(21)
const GN_MLSL_LDS = cenum(22)
const GD_MLSL_LDS = cenum(23)
const LD_MMA = cenum(24)
const LN_COBYLA = cenum(25)
const LN_NEWUOA = cenum(26)
const LN_NEWUOA_BOUND = cenum(27)
const LN_NELDERMEAD = cenum(28)
const LN_SBPLX = cenum(29)
const LN_AUGLAG = cenum(30)
const LD_AUGLAG = cenum(31)
const LN_AUGLAG_EQ = cenum(32)
const LD_AUGLAG_EQ = cenum(33)
const LN_BOBYQA = cenum(34)
const GN_ISRES = cenum(35)
const AUGLAG = cenum(36)
const AUGLAG_EQ = cenum(37)
const G_MLSL = cenum(38)
const G_MLSL_LDS = cenum(39)
const LD_SLSQP = cenum(40)
const LD_CCSAQ = cenum(41)
const GN_ESCH = cenum(42)
const NUM_ALGORITHMS = 43

const alg2int = Dict{Symbol,Cenum}(:GN_DIRECT=>GN_DIRECT, :GN_DIRECT_L=>GN_DIRECT_L, :GN_DIRECT_L_RAND=>GN_DIRECT_L_RAND, :GN_DIRECT_NOSCAL=>GN_DIRECT_NOSCAL, :GN_DIRECT_L_NOSCAL=>GN_DIRECT_L_NOSCAL, :GN_DIRECT_L_RAND_NOSCAL=>GN_DIRECT_L_RAND_NOSCAL, :GN_ORIG_DIRECT=>GN_ORIG_DIRECT, :GN_ORIG_DIRECT_L=>GN_ORIG_DIRECT_L, :GD_STOGO=>GD_STOGO, :GD_STOGO_RAND=>GD_STOGO_RAND, :LD_LBFGS_NOCEDAL=>LD_LBFGS_NOCEDAL, :LD_LBFGS=>LD_LBFGS, :LN_PRAXIS=>LN_PRAXIS, :LD_VAR1=>LD_VAR1, :LD_VAR2=>LD_VAR2, :LD_TNEWTON=>LD_TNEWTON, :LD_TNEWTON_RESTART=>LD_TNEWTON_RESTART, :LD_TNEWTON_PRECOND=>LD_TNEWTON_PRECOND, :LD_TNEWTON_PRECOND_RESTART=>LD_TNEWTON_PRECOND_RESTART, :GN_CRS2_LM=>GN_CRS2_LM, :GN_MLSL=>GN_MLSL, :GD_MLSL=>GD_MLSL, :GN_MLSL_LDS=>GN_MLSL_LDS, :GD_MLSL_LDS=>GD_MLSL_LDS, :LD_MMA=>LD_MMA, :LN_COBYLA=>LN_COBYLA, :LN_NEWUOA=>LN_NEWUOA, :LN_NEWUOA_BOUND=>LN_NEWUOA_BOUND, :LN_NELDERMEAD=>LN_NELDERMEAD, :LN_SBPLX=>LN_SBPLX, :LN_AUGLAG=>LN_AUGLAG, :LD_AUGLAG=>LD_AUGLAG, :LN_AUGLAG_EQ=>LN_AUGLAG_EQ, :LD_AUGLAG_EQ=>LD_AUGLAG_EQ, :LN_BOBYQA=>LN_BOBYQA, :GN_ISRES=>GN_ISRES, :AUGLAG=>AUGLAG, :AUGLAG_EQ=>AUGLAG_EQ, :G_MLSL=>G_MLSL, :G_MLSL_LDS=>G_MLSL_LDS, :LD_SLSQP=>LD_SLSQP, :LD_CCSAQ=>LD_CCSAQ, :GN_ESCH=>GN_ESCH)
const int2alg = Dict{Cenum,Symbol}()

# in 0.4 we could use int2alg=[alg2int[k]=>k for k in keys], and in 0.5 a generator
# expression, but Compat currently does not support generator expressions.
for k in keys(alg2int)
    int2alg[alg2int[k]] = k
end

# enum nlopt_result
const FAILURE = cenum(-1)
const INVALID_ARGS = cenum(-2)
const OUT_OF_MEMORY = cenum(-3)
const ROUNDOFF_LIMITED = cenum(-4)
const FORCED_STOP = cenum(-5)
const SUCCESS = cenum(1)
const STOPVAL_REACHED = cenum(2)
const FTOL_REACHED = cenum(3)
const XTOL_REACHED = cenum(4)
const MAXEVAL_REACHED = cenum(5)
const MAXTIME_REACHED = cenum(6)

const res2sym = Dict{Cenum,Symbol}(FAILURE=>:FAILURE, INVALID_ARGS=>:INVALID_ARGS, OUT_OF_MEMORY=>:OUT_OF_MEMORY, ROUNDOFF_LIMITED=>:ROUNDOFF_LIMITED, FORCED_STOP=>:FORCED_STOP, SUCCESS=>:SUCCESS, STOPVAL_REACHED=>:STOPVAL_REACHED, FTOL_REACHED=>:FTOL_REACHED, XTOL_REACHED=>:XTOL_REACHED, MAXEVAL_REACHED=>:MAXEVAL_REACHED, MAXTIME_REACHED=>:MAXTIME_REACHED)

############################################################################
# wrapper around nlopt_opt type

typealias _Opt Ptr{Void} # nlopt_opt

# pass both f and o to the callback so that we can handle exceptions
type Callback_Data
    f::Function
    o::Any # should be Opt, but see Julia issue #269
end

type Opt
    opt::_Opt

    # need to store callback data for objective and constraints in
    # Opt so that they aren't garbage-collected.  cb[1] is the objective.
    cb::Vector{Callback_Data}

    function Opt(p::_Opt)
        opt = new(p, Array(Callback_Data,1))
        finalizer(opt, destroy)
        opt
    end
    function Opt(algorithm::Integer, n::Integer)
        if algorithm < 0 || algorithm > NUM_ALGORITHMS
            throw(ArgumentError("invalid algorithm $algorithm"))
        elseif n < 0
            throw(ArgumentError("invalid dimension $n < 0"))
        end
        p = ccall((:nlopt_create,libnlopt), _Opt, (Cenum, Cuint),
                  algorithm, n)
        if p == C_NULL
            error("Error in nlopt_create")
        end
        Opt(p)
    end
    Opt(algorithm::Symbol, n::Integer) = Opt(try alg2int[algorithm]
                                             catch
                         throw(ArgumentError("unknown algorithm $algorithm"))
                                             end, n)
end

if VERSION < v"0.4.0-dev+3710"
    convert(::Type{_Opt}, o::Opt) = o.opt # for passing to ccall
else
    Base.unsafe_convert(::Type{_Opt}, o::Opt) = o.opt # for passing to ccall
end

destroy(o::Opt) = ccall((:nlopt_destroy,libnlopt), Void, (_Opt,), o)

ndims(o::Opt) = Int(ccall((:nlopt_get_dimension,libnlopt), Cuint, (_Opt,), o))
algorithm(o::Opt) = int2alg[ccall((:nlopt_get_algorithm,libnlopt),
                                  Cenum, (_Opt,), o)]

show(io::IO, o::Opt) = print(io, "Opt(:$(algorithm(o)), $(ndims(o)))")

############################################################################
# copying is a little tricky because we have to tell NLopt to use
# new Callback_Data.

# callback wrapper for nlopt_munge_data in NLopt 2.4
function munge_callback(p::Ptr{Void}, f_::Ptr{Void})
    f = unsafe_pointer_to_objref(f_)::Function
    f(p)::Ptr{Void}
end

function copy(o::Opt)
    p = ccall((:nlopt_copy,libnlopt), _Opt, (_Opt,), o)
    if p == C_NULL
        error("Error in nlopt_copy")
    end
    n = Opt(p)

    n.cb = similar(o.cb)
    for i = 1:length(o.cb)
        try
            n.cb[i] = Callback_Data(o.cb[i].f, n)
        catch e
            # if objective has not been set, o.cb[1] will throw
            # an UndefRefError, which is okay.
            if i != 1 || !isa(e, UndefRefError)
                rethrow(e) # some not-okay exception
            end
        end
    end

    try
        # n.o, for each callback, stores a pointer to an element of o.cb,
        # and we need to convert this into a pointer to the corresponding
        # element of n.cb.  nlopt_munge_data allows us to call a function
        # to transform each stored pointer in n.o, and we use the cbi
        # dictionary to convert pointers to indices into o.cb, whence
        # we obtain the corresponding element of n.cb.
        cbi = Dict{Ptr{Void},Int}()
        for i in 1:length(o.cb)
            try
                cbi[pointer_from_objref(o.cb[i])] = i
            end
        end
        ccall((:nlopt_munge_data,libnlopt), Void, (_Opt, Ptr{Void}, Any),
              n, munge_callback_ptr,
              p::Ptr{Void} -> p==C_NULL ? C_NULL :
                              pointer_from_objref(n.cb[cbi[p]]))
    catch e0
        # nlopt_munge_data not available, punt unless there is
        # no callback data
        try
            o.cb[1]
        catch e
            if length(o.cb) == 1 && isa(e, UndefRefError)
                return n
            end
        end
        error("copy(o::Opt) not supported for NLopt version < 2.4")
    end

    return n
end

############################################################################
# converting error results into exceptions

type ForcedStop <: Exception end

# cache current exception for forced stop
nlopt_exception = nothing

# check result and throw an exception if necessary
function chk(result::Integer)
    if result < 0 && result != ROUNDOFF_LIMITED
        if result == INVALID_ARGS
            throw(ArgumentError("invalid NLopt arguments"))
        elseif result == OUT_OF_MEMORY
            throw(OutOfMemoryError())
        elseif result == FORCED_STOP
            global nlopt_exception
            e = nlopt_exception
            if e != nothing && !isa(e, ForcedStop)
                nlopt_exception = nothing
                rethrow(e)
            end
        else
            error("nlopt failure: $result")
        end
    end
    result
end

chks(result::Integer) = res2sym[chk(result)]
chkn(result::Integer) = begin chk(result); nothing; end

############################################################################
# getting and setting scalar and vector parameters

# make a quoted symbol expression out of the arguments
qsym(args...) = Expr(:quote, Symbol(string(args...)))

# scalar parameters p of type T
macro GETSET(T, p)
    Tg = T == :Cdouble ? :Real : (T == :Cint || T == :Cuint ? :Integer : :Any)
    ps = Symbol(string(p, "!"))
    quote
        $(esc(p))(o::Opt) = ccall(($(qsym("nlopt_get_", p)),libnlopt),
                                  $T, (_Opt,), o)
        $(esc(ps))(o::Opt, val::$Tg) =
          chkn(ccall(($(qsym("nlopt_set_", p)),libnlopt),
                     Cenum, (_Opt, $T), o, val))
    end
end

# Vector{Cdouble} parameters p
macro GETSET_VEC(p)
    ps = Symbol(string(p, "!"))
    quote
        function $(esc(p))(o::Opt, v::Vector{Cdouble})
            if length(v) != ndims(o)
                throw(BoundsError())
            end
            chk(ccall(($(qsym("nlopt_get_", p)),libnlopt),
                      Cenum, (_Opt, Ptr{Cdouble}), o, v))
            v
        end
        $(esc(p))(o::Opt) = $(esc(p))(o, Array(Cdouble, ndims(o)))
        function $(esc(ps))(o::Opt, v::Vector{Cdouble})
            if length(v) != ndims(o)
                throw(BoundsError())
            end
            chkn(ccall(($(qsym("nlopt_set_", p)),libnlopt),
                      Cenum, (_Opt, Ptr{Cdouble}), o, v))
        end
        $(esc(ps)){T<:Real}(o::Opt, v::AbstractVector{T}) =
          $(esc(ps))(o, copy!(Array(Cdouble,length(v)), v))
        $(esc(ps))(o::Opt, val::Real) =
          chkn(ccall(($(qsym("nlopt_set_", p, "1")),libnlopt),
                     Cenum, (_Opt, Cdouble), o, val))
    end
end

############################################################################
# Optimizer parameters

@GETSET_VEC lower_bounds
@GETSET_VEC upper_bounds
@GETSET Cdouble stopval
@GETSET Cdouble ftol_rel
@GETSET Cdouble ftol_abs
@GETSET Cdouble xtol_rel
@GETSET_VEC xtol_abs
@GETSET Cint maxeval
@GETSET Cdouble maxtime
@GETSET Cint force_stop
@GETSET Cuint population
@GETSET Cuint vector_storage

force_stop!(o::Opt) = force_stop!(o, 1)

local_optimizer!(o::Opt, lo::Opt) =
  chkn(ccall((:nlopt_set_local_optimizer,libnlopt),
             Cenum, (_Opt, _Opt), o, lo))

# the initial-stepsize stuff is a bit different than GETSET_VEC,
# since the heuristics depend on the position x.

function default_initial_step!(o::Opt, x::Vector{Cdouble})
    if length(x) != ndims(o)
        throw(BoundsError())
    end
    chkn(ccall((:nlopt_set_default_initial_step,libnlopt),
               Cenum, (_Opt, Ptr{Cdouble}), o, x))
end
default_initial_step!{T<:Real}(o::Opt, x::AbstractVector{T}) =
  default_initial_step!(o, copy!(Array(Cdouble,length(x)), x))

function initial_step!(o::Opt, dx::Vector{Cdouble})
    if length(dx) != ndims(o)
        throw(BoundsError())
    end
    chkn(ccall((:nlopt_set_initial_step,libnlopt),
               Cenum, (_Opt, Ptr{Cdouble}), o, dx))
end
initial_step!{T<:Real}(o::Opt, dx::AbstractVector{T}) =
  initial_step!(o, copy!(Array(Cdouble,length(dx)), dx))
initial_step!(o::Opt, dx::Real) =
  chkn(ccall((:nlopt_set_initial_step1,libnlopt),
             Cenum, (_Opt, Cdouble), o, dx))

function initial_step(o::Opt, x::Vector{Cdouble}, dx::Vector{Cdouble})
    if length(x) != ndims(o) || length(dx) != ndims(o)
        throw(BoundsError())
    end
    chkn(ccall((:nlopt_get_initial_step,libnlopt),
               Cenum, (_Opt, Ptr{Cdouble}, Ptr{Cdouble}), o, x, dx))
    dx
end
initial_step{T<:Real}(o::Opt, x::AbstractVector{T}) =
    initial_step(o, copy!(Array(Cdouble,length(x)), x),
                 Array(Cdouble, ndims(o)))

############################################################################

function algorithm_name(a::Integer)
    s = ccall((:nlopt_algorithm_name,libnlopt), Ptr{Uint8}, (Cenum,), a)
    if s == C_NULL
        throw(ArgumentError("invalid algorithm $a"))
    end
    return String(s)
end

algorithm_name(a::Symbol) = algorithm_name(try alg2int[a]
                                           catch
                             throw(ArgumentError("unknown algorithm $a"))
                                           end)
algorithm_name(o::Opt) = algorithm_name(algorithm(o))

############################################################################

function version()
    v = Array(Cint, 3)
    pv = pointer(v)
    ccall((:nlopt_version,libnlopt), Void, (Ptr{Cint},Ptr{Cint},Ptr{Cint}),
          pv, pv + sizeof(Cint), pv + 2*sizeof(Cint))
    VersionNumber(convert(Int, v[1]),convert(Int, v[2]),convert(Int, v[3]))
end

############################################################################

srand(seed::Integer) = ccall((:nlopt_srand,libnlopt),
                             Void, (Culong,), seed)
srand_time() = ccall((:nlopt_srand_time,libnlopt), Void, ())

############################################################################
# Objective function:

const empty_grad = Cdouble[] # for passing when grad == C_NULL

function nlopt_callback_wrapper(n::Cuint, x::Ptr{Cdouble},
                                grad::Ptr{Cdouble}, d_::Ptr{Void})
    d = unsafe_pointer_to_objref(d_)::Callback_Data
    try
        res = convert(Cdouble,
                      d.f(unsafe_wrap(Array, x, (convert(Int, n),)),
                          grad == C_NULL ? empty_grad
                          : unsafe_wrap(Array, grad, (convert(Int, n),))))
        return res::Cdouble
    catch e
        global nlopt_exception
        nlopt_exception = e
        println("in callback catch")
        force_stop!(d.o::Opt)
        return 0.0 # ignored by nlopt
    end
end

for m in (:min, :max)
    mf = Symbol(string(m,"_objective!"))
    @eval function $mf(o::Opt, f::Function)
        o.cb[1] = Callback_Data(f, o)
        chkn(ccall(($(qsym("nlopt_set_", m, "_objective")),libnlopt),
                   Cenum, (_Opt, Ptr{Void}, Any),
                   o, nlopt_callback_wrapper_ptr,
                   o.cb[1]))
    end
end

############################################################################
# Nonlinear constraints:

for c in (:inequality, :equality)
    cf = Symbol(string(c, "_constraint!"))
    @eval function $cf(o::Opt, f::Function, tol::Real)
        push!(o.cb, Callback_Data(f, o))
        chkn(ccall(($(qsym("nlopt_add_", c, "_constraint")),libnlopt),
                   Cenum, (_Opt, Ptr{Void}, Any, Cdouble),
                   o, nlopt_callback_wrapper_ptr,
                   o.cb[end], tol))
    end
    @eval $cf(o::Opt, f::Function) = $cf(o, f, 0.0)
end

function remove_constraints!(o::Opt)
    resize!(o.cb, 1)
    chkn(ccall((:nlopt_remove_inequality_constraints,libnlopt),
               Cenum, (_Opt,), o))
    chkn(ccall((:nlopt_remove_equality_constraints,libnlopt),
               Cenum, (_Opt,), o))
end

############################################################################
# Vector-valued constraints

const empty_jac = Array(Cdouble,0,0) # for passing when grad == C_NULL

function nlopt_vcallback_wrapper(m::Cuint, res::Ptr{Cdouble},
                                 n::Cuint, x::Ptr{Cdouble},
                                 grad::Ptr{Cdouble}, d_::Ptr{Void})
    d = unsafe_pointer_to_objref(d_)::Callback_Data
    try
        d.f(unsafe_wrap(Array, res, (convert(Int, m),)),
            unsafe_wrap(Array, x, (convert(Int, n),)),
            grad == C_NULL ? empty_jac
            : unsafe_wrap(Array, grad, (convert(Int, n),convert(Int, m))))
    catch e
        global nlopt_exception
        nlopt_exception = e
        force_stop!(d.o::Opt)
    end
    nothing
end

for c in (:inequality, :equality)
    cf = Symbol(string(c, "_constraint!"))
    @eval begin
        function $cf(o::Opt, f::Function, tol::Vector{Cdouble})
            push!(o.cb, Callback_Data(f, o))
            chkn(ccall(($(qsym("nlopt_add_", c, "_mconstraint")),
                        libnlopt),
                       Cenum, (_Opt, Cuint, Ptr{Void}, Any, Ptr{Cdouble}),
                       o, length(tol), nlopt_vcallback_wrapper_ptr,
                       o.cb[end], tol))
        end
        $cf{T<:Real}(o::Opt, f::Function, tol::AbstractVector{T}) =
           $cf(o, f, copy!(Array(Float64,length(tol)), tol))
        $cf(o::Opt, m::Integer, f::Function, tol::Real) =
           $cf(o, f, fill!(Array(Cdouble, m), tol))
        $cf(o::Opt, m::Integer, f::Function)=
           $cf(o, m, f, 0.0)
    end
end

############################################################################
# Perform the optimization:

function optimize!(o::Opt, x::Vector{Cdouble})
    if length(x) != ndims(o)
        throw(BoundsError())
    end
    opt_f = Array(Cdouble, 1)
    ret = ccall((:nlopt_optimize,libnlopt), Cenum, (_Opt, Ptr{Cdouble},
                                                     Ptr{Cdouble}),
                o, x, opt_f)
    return (opt_f[1], x, chks(ret))
end

optimize{T<:Real}(o::Opt, x::AbstractVector{T}) =
  optimize!(o, copy!(Array(Cdouble,length(x)), x))

############################################################################

include("NLoptSolverInterface.jl")

if VERSION < v"0.3-"
    __init__() # automatic call to __init__ was added in Julia 0.3
end

end # module


text:
module Symmetry
using LinearAlgebra
import Base

tol = 1E-5

include("Shapes.jl")
include("Transformations.jl")
include("PointGroupGenerators.jl")
include("ctab_hard.jl")
include("character_tables.jl")

end

ext: jl


text:
# Linear Algebra style
# For single linear regresion y= mx .+ b
julia> m
3.3

julia> b
2-element Array{Int64,1}:
 2
 2

julia> x
2×1 Array{Float64,2}:
 0.18541330251480792
 0.07749677781424502

julia> y = m * x .+ b
2×1 Array{Float64,2}:
 2.611863898298866
 2.2557393667870085

julia> yn = y + randn(size(y)) * 0.5
2×1 Array{Float64,2}:
 3.0183983473834095
 1.6640798029816912

julia> X = zeros(2,2)
2×2 Array{Float64,2}:
 0.0  0.0
 0.0  0.0

julia> X[:,1] = x
2×1 Array{Float64,2}:
 0.18541330251480792
 0.07749677781424502

julia> X[:,2] = 1.0
ERROR: MethodError: no method matching setindex_shape_check(::Float64, ::Int64, ::Int64)
Closest candidates are:
  setindex_shape_check(::AbstractArray{#s57,1} where #s57, ::Integer, ::Integer) at indices.jl:196
  setindex_shape_check(::AbstractArray{#s57,2} where #s57, ::Integer, ::Integer) at indices.jl:200
  setindex_shape_check(::AbstractArray, ::Integer...) at indices.jl:154
  ...
Stacktrace:
 [1] macro expansion at .\multidimensional.jl:641 [inlined]
 [2] _unsafe_setindex!(::IndexLinear, ::Array{Float64,2}, ::Float64, ::Base.Slice{Base.OneTo{Int64}}, ::Int64) at .\multidimensional.jl:636
 [3] _setindex! at .\multidimensional.jl:631 [inlined]
 [4] setindex!(::Array{Float64,2}, ::Float64, ::Function, ::Int64) at .\abstractarray.jl:998
 [5] top-level scope at none:0

julia> X[:,2] = [1.0]
ERROR: DimensionMismatch("tried to assign 1-element array to 2×1 destination")
Stacktrace:
 [1] throw_setindex_mismatch(::Array{Float64,1}, ::Tuple{Int64,Int64}) at .\indices.jl:144
 [2] setindex_shape_check at .\indices.jl:196 [inlined]
 [3] macro expansion at .\multidimensional.jl:641 [inlined]
 [4] _unsafe_setindex!(::IndexLinear, ::Array{Float64,2}, ::Array{Float64,1}, ::Base.Slice{Base.OneTo{Int64}}, ::Int64) at .\multidimensional.jl:636
 [5] _setindex! at .\multidimensional.jl:631 [inlined]
 [6] setindex!(::Array{Float64,2}, ::Array{Float64,1}, ::Function, ::Int64) at .\abstractarray.jl:998
 [7] top-level scope at none:0

julia> X[:,2] = [1.0; 1.0]
2-element Array{Float64,1}:
 1.0
 1.0

julia> X
2×2 Array{Float64,2}:
 0.185413   1.0
 0.0774968  1.0

julia> coeff_pred = X\yn
2×1 Array{Float64,2}:
 12.549686418827514
  0.6915195429433669

julia> slope =  round(coeff_pred[1], 2)
ERROR: MethodError: no method matching round(::Float64, ::Int64)
Closest candidates are:
  round(::Float64, ::RoundingMode{:Nearest}) at float.jl:368
  round(::Float64, ::RoundingMode{:Up}) at float.jl:366
  round(::Float64, ::RoundingMode{:Down}) at float.jl:364
  ...
Stacktrace:
 [1] top-level scope at none:0

julia>

slope =  round(coeff_pred[1], 2)
intercept = round(coeff_pred[2], 2)

println("The real slope is $m, and the predicted slope is $slope")
println("The real intercept is $b, and the predicted slope is $intercept")


text:
#!/bin/bash
#=
exec julia --color=yes --startup-file=no --threads=auto "${BASH_SOURCE[0]}" "$@"
=#
# Execute this script via: `bash test/benchmarks/benchmark_potential_calculation.jl overwrite`

recalculate = in("overwrite", ARGS)

using CUDAKernels, CUDA
using SolidStateDetectors
using DataFrames, JLD2
using BenchmarkTools
using KernelAbstractions
using KernelAbstractions: get_device

using SolidStateDetectors:
    ConstructiveSolidGeometry.AbstractCoordinateSystem,
    Cartesian, Cylindrical, PotentialCalculationSetup,
    _guess_optimal_number_of_threads_for_SOR, adapt,
    get_sor_kernel, update!

fn_pot_calc_benchmark_df = joinpath(ENV["HOME"], "SSD_potential_calculation_benchmark.jld2")

N_MAX_THREADS = Base.Threads.nthreads()

T = Float32

sim_cyl = Simulation{T}(SSD_examples[:BEGe])
sim_car = Simulation{T}(SSD_examples[:CGD])

hw_backends = (
    (label = "CPU_1",  use_nthreads =   1, DAT = Array),
    (label = "CPU_2",  use_nthreads =   2, DAT = Array),
    (label = "CPU_4",  use_nthreads =   4, DAT = Array),
    (label = "CPU_8",  use_nthreads =   8, DAT = Array),
    (label = "CPU_16", use_nthreads =  16, DAT = Array),
    (label = "CPU_32", use_nthreads =  32, DAT = Array),
    (label = "GPU",    use_nthreads =   1, DAT = CuArray),
)

refinement_limits = [0.2, 0.1, 0.05, 0.025, 0.01]
max_n_iterations = 10_000

df_pot_calc = if !isfile(fn_pot_calc_benchmark_df) || recalculate
    df = DataFrame()
    for sim in (sim_cyl, sim_car)
        S = SolidStateDetectors.get_coordinate_system(sim)
        for backend in hw_backends
            potential_calc_settings = (
                use_nthreads = backend.use_nthreads,
                convergence_limit = 1e-6,
                refinement_limits = refinement_limits,
                device_array_type = backend.DAT,
                max_n_iterations = max_n_iterations,
                depletion_handling = true,
                verbose = true,
            )
            calculate_electric_potential!(sim; # compile time
                max_n_iterations = 40,
                n_iterations_between_checks = 10,
                potential_calc_settings...
            )
            t = (@elapsed calculate_electric_potential!(sim; potential_calc_settings... ))
            push!(df,
                (
                    backend = backend.label,
                    S = S,
                    time = t,
                    grid_size = size(sim.electric_potential),
                    potential_calc_settings...
                )
            )
            display(df)
        end
    end
    save(fn_pot_calc_benchmark_df, Dict("df_pot_calc" => df))
    df
else
    load(fn_pot_calc_benchmark_df, "df_pot_calc")
end


text:
export DSOSPoly, SDSOSPoly, SOSPoly

function JuMP.value(p::MatPolynomial{JuMP.VariableRef})
    MatPolynomial(map(JuMP.value, p.Q), p.x)
end

for poly in (:DSOSPoly, :SDSOSPoly, :SOSPoly)
    @eval begin
        struct $poly{PB<:PolyJuMP.AbstractPolynomialBasis} <: PolyJuMP.AbstractPoly
            polynomial_basis::PB
        end
        $poly(x::AbstractVector{<:MultivariatePolynomials.AbstractPolynomialLike}) = $poly(MonomialBasis(x))
    end
end

const PosPoly{PB} = Union{DSOSPoly{PB}, SDSOSPoly{PB}, SOSPoly{PB}}

JuMP.variable_type(m::JuMP.Model, p::PosPoly) = PolyJuMP.polytype(m, p, p.polynomial_basis)
PolyJuMP.polytype(m::JuMP.Model, ::PosPoly, basis::PolyJuMP.MonomialBasis{MT, MV}) where {MT<:AbstractMonomial, MV<:AbstractVector{MT}} = MatPolynomial{JuMP.VariableRef, MT, MV}

# Sum-of-Squares polynomial

_polytype(m::JuMP.Model, ::PosPoly, x::MVT) where {MT<:AbstractMonomial, MVT<:AbstractVector{MT}} = MatPolynomial{JuMP.VariableRef, MT, MVT}

"""
    constraint_matpoly!(model::JuMP.AbstractModel, p::MatPolynomial, ::Union{SOSPoly, SDSOSPoly, DSOSPoly})

Constraints matrix of `p` to be

* positive semidefinite if the third argument is of type `SOSPoly`,
* scaled diagonally dominant if the third argument is of type `SDSOSPoly`, or
* diagonally dominant if the third argument is of type `DSOSPoly`.

See Definition 4 of [AM17] for a precise definition of the last two items.

[AM17] Ahmadi, A. A. & Majumdar, A.
*DSOS and SDSOS Optimization: More Tractable Alternatives to Sum of Squares and Semidefinite Optimization*
ArXiv e-prints, 2017
"""
function constraint_matpoly! end

function constraint_matpoly!(m, p::MatPolynomial, ::SOSPoly)
    JuMP.add_constraint(m, JuMP.VectorConstraint(p.Q.Q, MOI.PositiveSemidefiniteConeTriangle(length(p.x))))
end
function constraint_matpoly!(model, p::MatPolynomial, ::SDSOSPoly)
    # `p.Q` is SDD iff it is the sum of psd matrices Mij that are zero except for
    # entries ii, ij and jj [Lemma 9, AM17].
    n = length(p.x)
    T = JuMP.GenericAffExpr{Float64, JuMP.VariableRef}
    # `Q[r, c]` will contain the expression `p.Q[r, c] - sum Mij[r, c]`
    Q = SymMatrix{T}(Vector{T}(p.Q.Q), n)
    for i in 1:n
        for j in i+1:n
            Mii = @variable(model)
            JuMP.add_to_expression!(Q[i, i], -1.0, Mii)
            Mij = @variable(model)
            JuMP.add_to_expression!(Q[i, j], -1.0, Mij)
            Mjj = @variable(model)
            JuMP.add_to_expression!(Q[j, j], -1.0, Mjj)
            # PSD constraints on 2x2 matrices are SOC representable
            @constraint(model, [Mii + Mjj, 2Mij, Mii - Mjj] in MOI.SecondOrderCone(3))
        end
    end
    @constraint(model, Q.Q .== 0)
end
function constraint_matpoly!(model, p::MatPolynomial, ::DSOSPoly)
    n = length(p.x)
    Q = Matrix{JuMP.VariableRef}(undef, n, n)
    for i in 1:n
        for j in i:n
            if i == j
                Q[i, j] = p[i, j]
            else
                Q[j, i] = Q[i, j] = JuMP.VariableRef(model)
                @constraint model Q[i, j] >= p[i, j]
                @constraint model Q[i, j] >= -p[i, j]
            end
        end
    end
    for i in 1:n
        @constraint model 2Q[i, i] >= sum(Q[i, :])
    end
    # If n > 1, this is implied by the constraint but it doesn't hurt to add the variable cone
    # Adding things on varCones makes JuMP think that it is SDP
    # push!(m.varCones, (:NonNeg, map(i -> p[i, i].col, 1:n)))
end
function _matpolynomial(m, x::AbstractVector{<:AbstractMonomial}, binary::Bool, integer::Bool)
    if isempty(x)
        zero(JuMP.variable_type(m, SOSPoly(x)))
    else
        function _newvar(i, j)
            v = JuMP.VariableRef(m)
            if length(x) == 1
                # 1x1 matrix is SDP iff its only entry is nonnegative
                # We handle this case here and do not create any SDP constraint
                set_lower_bound(v, 0)
            end
            if binary
                set_binary(v)
            end
            if integer
                set_integer(v)
            end
            v
        end
        MatPolynomial{JuMP.VariableRef}(_newvar, x)
    end
end
function _createpoly(m::JuMP.Model, set::PosPoly, basis::PolyJuMP.MonomialBasis, binary::Bool, integer::Bool)
    p = _matpolynomial(m, basis.monomials, binary, integer)
    if length(basis.monomials) > 1
        constraint_matpoly!(m, p, set)
    end
    p
end
function PolyJuMP.createpoly(m::JuMP.Model, p::PosPoly, binary::Bool, integer::Bool)
    _createpoly(m, p, p.polynomial_basis, binary, integer)
end

# Defer other methods to defaults in PolyJuMP
createpoly(args...) = PolyJuMP.createpoly(args...)


text:
using Pkg
Pkg.activate(".")
using SymEngine

struct CustomFunc
    name::String
    cfunc::Base.CFunction
    func1
    func2
end

struct CFunctionWrapperClass
    name::Ptr{Cchar}
    create_ptr::Ptr{Cvoid}
    diff_ptr::Ptr{Cvoid}
end

function create_func(name, f)
    function _create_func(a::Basic, args_ptr::Ptr{Cvoid})::Cint
        args = SymEngine.CVecBasic(args_ptr)
        try
            res = f(convert(Vector, args))
        catch e
            res = nothing
        end
        if res == nothing
            return 1
        else
            res = Basic(res)
            ccall((:basic_assign, SymEngine.libsymengine), Nothing, (Ref{Basic}, Ref{Basic}), a, res)
            return 0
        end
    end
    create_ptr = @cfunction($_create_func, Cint, (Ref{Basic}, Ptr{Cvoid}))
    return CustomFunc(name, create_ptr, f, _create_func)
end

function my_real(x::Vector{Basic})
   return real(N(x[1]))
end

function my_norm(x::Vector{Basic})
   a, b = N(x[1]), N(x[2])
   return sqrt(a^2+b^2)
end

@vars x

function (c::CustomFunc)(args...)
    v = convert(SymEngine.CVecBasic, args...)
    a = Basic()
    ccall((:basic_create_function_wrapper, SymEngine.libsymengine), Nothing, (Ref{Basic}, Ptr{Cvoid}, Cstring, Ptr{Cvoid}, Ptr{Cvoid}), a, v.ptr, pointer(c.name), c.cfunc.ptr, C_NULL)
    return a
end

c = create_func("Re", my_real)
c(x)

d = create_func("Norm", my_norm)
d(2, 3)


text:
using FractalGeometry

using StaticArrays
using GeometryBasics
using FileIO
import Makie

begin
    testp = Point3(0.6, 0.3, -0.4)
    const tau = 2pi
    const τ = tau
    const TAU = τ
    ori = Point3(-1.1, -1.1, -1.1)
    wid = Point3(2.2, 2.2, 2.2)
    orisv = SVector(ori)
    widsv = SVector(wid)

    outerbox(v) = FractalGeometry.boxdist(v, Point3(0.8, 0.8, 0.8))
    outerboxsdf = FractalGeometry.FunctionSDF(outerbox)
end

begin
    arc_length_sampling_freq = 0.01
    linecurveγ = FractalGeometry.FunctionCurve(FractalGeometry.linecurve, (-1., 1.))
    linecurvesdf = FractalGeometry.curve_to_kdtree_sdf(linecurveγ, 0.125, arc_length_sampling_freq)
    distancep = linecurvesdf(testp)

    linecurveboxsdf = FractalGeometry.IntersectSDF(linecurvesdf, outerboxsdf)
end

begin
    using FractalGeometry
    arc_length_sampling_freq = 0.01
    helixcurve1(t) = Point3(0.3*cos(tau * t), 0.3*sin(tau * t), t)
    helixcurve2(t) = Point3(-0.3*cos(tau * t), -0.3*sin(tau * t), t)
    helixcurveconn1(t) = helixcurve1(0.4) * t + helixcurve2(0.4) * (1 - t)
    helixcurveconn2(t) = helixcurve1(-0.4) * t + helixcurve2(-0.4) * (1 - t)

    curves = [helixcurve1, helixcurve2, helixcurveconn1, helixcurveconn2]
    tspans = [(-1., 1.), (-1., 1.), (0., 1.), (0., 1.)]
    extrusions = [0.125, 0.125, 0.05, 0.05]

    union_sdf = FractalGeometry.curves_to_union_sdf(curves, tspans, extrusions, arc_length_sampling_freq; boundary_sdf=outerboxsdf)

    distancep = union_sdf(testp)

end

begin
    curdist = union_sdf
    curh0 = 0.03

    result = FractalGeometry.distmesh(curdist, FractalGeometry.HUniform(), curh0, FractalGeometry.DistMeshSetup(); origin=ori, widths=wid, maxiters=100)
    resultouterbox = FractalGeometry.distmesh(outerbox, HUniform(), 0.5, DistMeshSetup(); origin=ori, widths=wid, maxiters=100)
    boundarytris, alltris = FractalGeometry.boundary_triangles(result.tetrahedra)
    boundarytrisouterbox, alltrisouterbox = FractalGeometry.boundary_triangles(resultouterbox.tetrahedra)

    boundarymesh = GeometryBasics.Mesh(result.points, boundarytris)
    boundarymeshouterbox = GeometryBasics.Mesh(resultouterbox.points, boundarytrisouterbox)
end
begin

    show_axis = true

    scene = Makie.Scene()
    Makie.mesh!(scene, boundarymesh, color=[norm(v[1:2]) for v in coordinates(boundarymesh)], show_axis=show_axis)

    Makie.wireframe!(scene, boundarymesh, show_axis=show_axis)
    Makie.wireframe!(scene, boundarymeshouterbox, show_axis=show_axis, color=RGB(0.8, 0.3, 0.8))
end

begin
    savefile = "meshes/generated/doublehelixclipped.stl"
    save(savefile, boundarymesh)

end

begin
    mesh = load(savefile)
    scene = Makie.Scene()
    Makie.mesh!(scene, mesh, color=[norm(v[1:2]) for v in coordinates(mesh)], show_axis=show_axis)

    Makie.wireframe!(scene, mesh, show_axis=show_axis)

end

max_stars_repo_name: zoemcc/geometry

max_issues_repo_name: zoemcc/geometry

max_forks_repo_name: zoemcc/geometry


text:
"""
Variational partitioned Runge-Kutta integrator cache.

### Fields

* `q̃`: initial guess of q
* `p̃`: initial guess of p
* `ṽ`: initial guess of v
* `f̃`: initial guess of f
* `s̃`: holds shift due to periodicity of solution
* `Q`: internal stages of q
* `P`: internal stages of p
* `V`: internal stages of v
* `F`: internal stages of f
* `Y`: integral of vector field of internal stages of q
* `Z`: integral of vector field of internal stages of p
"""
mutable struct IntegratorCacheVPRK{ST,D,S} <: IODEIntegratorCache{ST,D}
    # q::Vector{ST}
    # p::Vector{ST}
    λ::Vector{ST}
    λ̅::Vector{ST}

    q₋::Vector{ST}
    q̅₊::Vector{ST}
    p₋::Vector{ST}
    p̅₊::Vector{ST}

    u::Vector{ST}
    g::Vector{ST}

    q̃::Vector{ST}
    p̃::Vector{ST}
    ṽ::Vector{ST}
    f̃::Vector{ST}
    s̃::Vector{ST}

    y::Vector{ST}
    z::Vector{ST}

    Q::Vector{Vector{ST}}
    P::Vector{Vector{ST}}
    V::Vector{Vector{ST}}
    F::Vector{Vector{ST}}
    Λ::Vector{Vector{ST}}
    Φ::Vector{Vector{ST}}

    Y::Vector{Vector{ST}}
    Z::Vector{Vector{ST}}
    U::Vector{Vector{ST}}
    G::Vector{Vector{ST}}
    R::Vector{Vector{ST}}

    function IntegratorCacheVPRK{ST,D,S}(projection::Bool=false) where {ST,D,S}
        # create temporary vectors
        q̃ = zeros(ST,D)
        p̃ = zeros(ST,D)
        ṽ = zeros(ST,D)
        f̃ = zeros(ST,D)
        s̃ = zeros(ST,D)

        # create update vectors
        y = zeros(ST,D)
        z = zeros(ST,D)

        # create internal stage vectors
        Q = create_internal_stage_vector(ST, D, S)
        P = create_internal_stage_vector(ST, D, S)
        V = create_internal_stage_vector(ST, D, S)
        F = create_internal_stage_vector(ST, D, S)
        Y = create_internal_stage_vector(ST, D, S)
        Z = create_internal_stage_vector(ST, D, S)

        # projection vectors
        if projection
            λ = zeros(ST,D)
            λ̅ = zeros(ST,D)

            q₋= zeros(ST,D)
            q̅₊= zeros(ST,D)
            p₋= zeros(ST,D)
            p̅₊= zeros(ST,D)

            u = zeros(ST,D)
            g = zeros(ST,D)

            Λ = create_internal_stage_vector(ST, D, S)
            Φ = create_internal_stage_vector(ST, D, S)

            U = create_internal_stage_vector(ST, D, 2)
            G = create_internal_stage_vector(ST, D, 2)
            R = create_internal_stage_vector(ST, D, S)
        else
            λ = Vector{ST}()
            λ̅ = Vector{ST}()

            q₋= Vector{ST}()
            q̅₊= Vector{ST}()
            p₋= Vector{ST}()
            p̅₊= Vector{ST}()

            u = Vector{ST}()
            g = Vector{ST}()

            Λ = create_internal_stage_vector(ST, 0, 0)
            Φ = create_internal_stage_vector(ST, 0, 0)

            U = create_internal_stage_vector(ST, 0, 0)
            G = create_internal_stage_vector(ST, 0, 0)
            R = create_internal_stage_vector(ST, 0, 0)
        end

        new(λ, λ̅, q₋, q̅₊, p₋, p̅₊,
            u, g, q̃, p̃, ṽ, f̃, s̃, y, z,
            Q, P, V, F, Λ, Φ, Y, Z, U, G, R)
    end
end

function IntegratorCacheVPRK(ST,D,S)
    IntegratorCacheVPRK{ST,D,S}(false)
end

function IntegratorCacheVPRKwProjection(ST,D,S)
    IntegratorCacheVPRK{ST,D,S}(true)
end

function create_integrator_cache(int::AbstractIntegratorVPRK{DT,TT}) where {DT,TT}
    IntegratorCacheVPRK(DT, TT, ndims(int), nstages(int))
end

function create_integrator_cache(int::AbstractIntegratorVPRKwProjection{DT,TT}) where {DT,TT}
    IntegratorCacheVPRKwProjection(DT, TT, ndims(int), nstages(int))
end


text:
struct DQExperience
    s::Array{Float64}
    a::Int64
    r::Float64
    sp::Array{Float64}
    done::Bool
end

mutable struct ReplayBuffer
    max_size::Int64 # maximum size of the buffer
    batch_size::Int64
    rng::AbstractRNG
    _curr_size::Int64
    _idx::Int64
    _experience::Vector{DQExperience}

    _s_batch::Array{Float64}
    _a_batch::Vector{Int64}
    _r_batch::Vector{Float64}
    _sp_batch::Array{Float64}
    _done_batch::Vector{Bool}

    function ReplayBuffer(env::AbstractEnvironment,
                          max_size::Int64,
                          batch_size::Int64,
                          rng::AbstractRNG = MersenneTwister(0))
        s_dim = obs_dimensions(env)
        experience = Vector{DQExperience}(max_size)
        _s_batch = zeros(batch_size, s_dim...)
        _a_batch = zeros(Int64, batch_size)
        _r_batch = zeros(batch_size)
        _sp_batch = zeros(batch_size, s_dim...)
        _done_batch = zeros(Bool, batch_size)
        return new(max_size, batch_size, rng, 0, 1, experience,
                   _s_batch, _a_batch, _r_batch, _sp_batch, _done_batch)
    end
end

is_full(r::ReplayBuffer) = r._curr_size == r.max_size

max_size(r::ReplayBuffer) = r.max_size

function add_exp!(r::ReplayBuffer, expe::DQExperience)
    r._experience[r._idx] = expe
    r._idx = mod1((r._idx + 1),r.max_size)
    if r._curr_size < r.max_size
        r._curr_size += 1
    end
end

# convert array of DQExp to batches
function _encode_sample!(r::ReplayBuffer, samples::Vector{DQExperience})
    @assert length(samples) == size(r._s_batch[1])
    for (i, sample) in enumerate(samples)
        r._s_batch[i] = sample.s
        r._a_batch[i] = sample.a
        r._r_batch[i] = sample.r
        r._sp_batch[i] = sample.sp
        r._done_batch[i] = sample.done
    end
end

function StatsBase.sample(r::ReplayBuffer)
    @assert r._curr_size >= r.batch_size
    @assert r.max_size >= r.batch_size # could be checked during construction
    sample_indices = sample(r.rng, 1:r._curr_size, r.batch_size, replace=false)
    @assert length(sample_indices) == size(r._s_batch)[1]
    for (i, idx) in enumerate(sample_indices)
        r._s_batch[Base.setindex(indices(r._s_batch), i, 1)...] = r._experience[idx].s
        r._a_batch[i] = r._experience[idx].a
        r._r_batch[i] = r._experience[idx].r
        r._sp_batch[Base.setindex(indices(r._sp_batch), i, 1)...] = r._experience[idx].sp
        r._done_batch[i] = r._experience[idx].done
    end
    return r._s_batch, r._a_batch, r._r_batch, r._sp_batch, r._done_batch
end

function populate_replay_buffer!(replay::ReplayBuffer, env::AbstractEnvironment;
                                 max_pop::Int64=replay.max_size, max_steps::Int64=2000)
    o = reset(env)
    done = false
    step = 0
    println("Start populating replay buffer...")
    for t=1:(max_pop - replay._curr_size)
        if !env.random_populate_replay_buffer
            q_low_fi = get_qvals(env.low_fi_policy, convert_ac_state(env.ego_agent.state))
            action = actions(env)[indmax(q_low_fi)]
        else
            action = sample_action(env)
        end
        ai = action_index(env.problem, action)
        op, rew, _, done, info = step!(env, action)
        exp = DQExperience(o, ai, rew, op, done)
        add_exp!(replay, exp)
        o = op
        # println(o, " ", action, " ", rew, " ", done, " ", info) #TODO verbose?
        step += 1
        if done || step >= max_steps
            o = reset(env)
            done = false
            step = 0
        end
    end
    @assert replay._curr_size >= replay.batch_size
    println("Done populating replay buffer!")
end

max_stars_repo_name: parachutel/DeepCorrectionCAS


text:
module GeCKOReader

using Distances
using DataStructures
using DataFrames
using Libz
using BioSequences
using Formatting
using Gadfly
using Query
using HypothesisTests

include("utils.jl")
include("types.jl")
include("io.jl")
include("processing.jl")
include("plotting.jl")

export read_seq_file,
       gen_barcode_mapping,
       ScreenSample,
       get_rel_freqs,
       get_log2fc_pvals,
       fit_lines,
       calc_log2fc,
       compute_gene_stats

end # module


text:
using DiscreteFiltering
using Documenter

DocMeta.setdocmeta!(
    DiscreteFiltering,
    :DocTestSetup,
    :(using DiscreteFiltering);
    recursive = true,
)

makedocs(;
    modules = [DiscreteFiltering],
    authors = "Syver Døving Agdestein <syverda@icloud.com> and contributors",
    repo = "https://github.com/agdestein/DiscreteFiltering.jl/blob/{commit}{path}#{line}",
    sitename = "DiscreteFiltering.jl",
    format = Documenter.HTML(;
        prettyurls = get(ENV, "CI", "false") == "true",
        canonical = "https://agdestein.github.io/DiscreteFiltering.jl",
        assets = String[],
    ),
    pages = ["Home" => "index.md"],
)

deploydocs(; repo = "github.com/agdestein/DiscreteFiltering.jl", devbranch = "main")


text:
using NCDatasets

using Dates: AbstractTime, now

using Oceananigans.Fields

using Oceananigans.Grids: topology, halo_size, all_x_nodes, all_y_nodes, all_z_nodes, parent_index_range
using Oceananigans.Utils: versioninfo_with_gpu, oceananigans_versioninfo, prettykeys
using Oceananigans.TimeSteppers: float_or_date_time
using Oceananigans.Fields: reduced_dimensions, reduced_location, location, validate_indices

mutable struct NetCDFOutputWriter{D, O, T, A} <: AbstractOutputWriter
    filepath :: String
    dataset :: D
    outputs :: O
    schedule :: T
    overwrite_existing :: Bool
    array_type :: A
    previous :: Float64
    verbose :: Bool
end

ext(::Type{NetCDFOutputWriter}) = ".nc"

dictify(outputs) = outputs
dictify(outputs::NamedTuple) = Dict(string(k) => dictify(v) for (k, v) in zip(keys(outputs), values(outputs)))

xdim(::Type{Face}) = ("xF",)
ydim(::Type{Face}) = ("yF",)
zdim(::Type{Face}) = ("zF",)

xdim(::Type{Center}) = ("xC",)
ydim(::Type{Center}) = ("yC",)
zdim(::Type{Center}) = ("zC",)

xdim(::Type{Nothing}) = ()
ydim(::Type{Nothing}) = ()
zdim(::Type{Nothing}) = ()

netcdf_spatial_dimensions(::AbstractField{LX, LY, LZ}) where {LX, LY, LZ} =
    tuple(xdim(LX)..., ydim(LY)..., zdim(LZ)...)

function default_dimensions(output, grid, indices, with_halos)
    Hx, Hy, Hz = halo_size(grid)
    TX, TY, TZ = topo = topology(grid)

    locs = Dict(
                "xC" => (Center, Center, Center),
                "xF" => (  Face, Center, Center),
                "yC" => (Center, Center, Center),
                "yF" => (Center,   Face, Center),
                "zC" => (Center, Center, Center),
                "zF" => (Center, Center,   Face),
               )

    indices = Dict(name => validate_indices(indices, locs[name], grid) for name in keys(locs))

    if !with_halos
        indices = Dict(name => restrict_to_interior.(indices[name], locs[name], topo, size(grid))
                       for name in keys(locs))
    end

    dims = Dict("xC" => parent(all_x_nodes(Center, grid))[parent_index_range(indices["xC"][1], Center, TX, Hx)],
                "xF" => parent(all_x_nodes(Face,   grid))[parent_index_range(indices["xF"][1],   Face, TX, Hx)],
                "yC" => parent(all_y_nodes(Center, grid))[parent_index_range(indices["yC"][2], Center, TY, Hy)],
                "yF" => parent(all_y_nodes(Face,   grid))[parent_index_range(indices["yF"][2],   Face, TY, Hy)],
                "zC" => parent(all_z_nodes(Center, grid))[parent_index_range(indices["zC"][3], Center, TZ, Hz)],
                "zF" => parent(all_z_nodes(Face,   grid))[parent_index_range(indices["zF"][3],   Face, TZ, Hz)])

    return dims
end

const default_dimension_attributes = Dict(
    "xC"          => Dict("longname" => "Locations of the cell centers in the x-direction.", "units" => "m"),
    "xF"          => Dict("longname" => "Locations of the cell faces in the x-direction.",   "units" => "m"),
    "yC"          => Dict("longname" => "Locations of the cell centers in the y-direction.", "units" => "m"),
    "yF"          => Dict("longname" => "Locations of the cell faces in the y-direction.",   "units" => "m"),
    "zC"          => Dict("longname" => "Locations of the cell centers in the z-direction.", "units" => "m"),
    "zF"          => Dict("longname" => "Locations of the cell faces in the z-direction.",   "units" => "m"),
    "time"        => Dict("longname" => "Time", "units" => "s"),
    "particle_id" => Dict("longname" => "Particle ID")
)

const default_output_attributes = Dict(
    "u" => Dict("longname" => "Velocity in the x-direction", "units" => "m/s"),
    "v" => Dict("longname" => "Velocity in the y-direction", "units" => "m/s"),
    "w" => Dict("longname" => "Velocity in the z-direction", "units" => "m/s"),
    "b" => Dict("longname" => "Buoyancy",                    "units" => "m/s²"),
    "T" => Dict("longname" => "Conservative temperature",    "units" => "°C"),
    "S" => Dict("longname" => "Absolute salinity",           "units" => "g/kg")
)

add_schedule_metadata!(attributes, schedule) = nothing

function add_schedule_metadata!(global_attributes, schedule::IterationInterval)
    global_attributes["schedule"] = "IterationInterval"
    global_attributes["interval"] = schedule.interval
    global_attributes["output iteration interval"] =
        "Output was saved every $(schedule.interval) iteration(s)."

    return nothing
end

function add_schedule_metadata!(global_attributes, schedule::TimeInterval)
    global_attributes["schedule"] = "TimeInterval"
    global_attributes["interval"] = schedule.interval
    global_attributes["output time interval"] =
        "Output was saved every $(prettytime(schedule.interval))."

    return nothing
end

function add_schedule_metadata!(global_attributes, schedule::WallTimeInterval)
    global_attributes["schedule"] = "WallTimeInterval"
    global_attributes["interval"] = schedule.interval
    global_attributes["output time interval"] =
        "Output was saved every $(prettytime(schedule.interval))."

    return nothing
end

function add_schedule_metadata!(global_attributes, schedule::AveragedTimeInterval)
    add_schedule_metadata!(global_attributes, TimeInterval(schedule))

    global_attributes["time_averaging_window"] = schedule.window
    global_attributes["time averaging window"] =
        "Output was time averaged with a window size of $(prettytime(schedule.window))"

    global_attributes["time_averaging_stride"] = schedule.stride
    global_attributes["time averaging stride"] =
        "Output was time averaged with a stride of $(schedule.stride) iteration(s) within the time averaging window."

    return nothing
end

"""
    NetCDFOutputWriter(model, outputs; filename, schedule
                                          dir = ".",
                                   array_type = Array{Float32},
                                      indices = nothing,
                            global_attributes = Dict(),
                            output_attributes = Dict(),
                                   dimensions = Dict(),
                           overwrite_existing = false,
                                  compression = 0,
                                      verbose = false)

Construct a `NetCDFOutputWriter` that writes `(label, output)` pairs in `outputs` (which should
be a `Dict`) to a NetCDF file, where `label` is a string that labels the output and `output` is
either a `Field` (e.g. `model.velocities.u`) or a function `f(model)` that
returns something to be written to disk. Custom output requires the spatial `dimensions` (a
`Dict`) to be manually specified (see examples).

Keyword arguments
=================
- `filename` (required): Descriptive filename. ".nc" is appended to `filename` if ".nc" is not detected.

- `schedule` (required): `AbstractSchedule` that determines when output is saved.

- `dir`: Directory to save output to.

- `array_type`: The array type to which output arrays are converted to prior to saving.
                Default: `Array{Float32}`.

- `indices`: Tuple of indices of the output variables to include. Default is `(:, :, :)`, which
             includes the full fields.

- `with_halos`: Boolean defining whether or not to include halos in the outputs.

- `global_attributes`: Dict of model properties to save with every file. Default: `Dict()`.

- `output_attributes`: Dict of attributes to be saved with each field variable (reasonable
                       defaults are provided for velocities, buoyancy, temperature, and salinity;
                       otherwise `output_attributes` *must* be user-provided).

- `dimensions`: A `Dict` of dimension tuples to apply to outputs (required for function outputs)

- `overwrite_existing`: If false, `NetCDFOutputWriter` will be set to append to `filepath`. If true, `NetCDFOutputWriter`
                        will overwrite `filepath` if it exists or create it if it does not.
                        Default: `false`. See NCDatasets.jl documentation for more information about its `mode` option.

- `compression`: Determines the compression level of data (0-9; default: 0)

Examples
========
Saving the ``u`` velocity field and temperature fields, the full 3D fields and surface 2D slices
to separate NetCDF files:

```jldoctest netcdf1
using Oceananigans

grid = RectilinearGrid(size=(16, 16, 16), extent=(1, 1, 1))

model = NonhydrostaticModel(grid=grid, tracers=:c)

simulation = Simulation(model, Δt=12, stop_time=3600)

fields = Dict("u" => model.velocities.u, "c" => model.tracers.c)

simulation.output_writers[:field_writer] =
    NetCDFOutputWriter(model, fields, filename="fields.nc", schedule=TimeInterval(60))

# output
NetCDFOutputWriter scheduled on TimeInterval(1 minute):
├── filepath: ./fields.nc
├── dimensions: zC(16), zF(17), xC(16), yF(16), xF(16), yC(16), time(0)
├── 2 outputs: (c, u)
└── array type: Array{Float32}
```

```jldoctest netcdf1
simulation.output_writers[:surface_slice_writer] =
    NetCDFOutputWriter(model, fields, filename="surface_xy_slice.nc",
                       schedule=TimeInterval(60), indices=(:, :, grid.Nz))

# output
NetCDFOutputWriter scheduled on TimeInterval(1 minute):
├── filepath: ./surface_xy_slice.nc
├── dimensions: zC(1), zF(1), xC(16), yF(16), xF(16), yC(16), time(0)
├── 2 outputs: (c, u)
└── array type: Array{Float32}
```

```jldoctest netcdf1
simulation.output_writers[:averaged_profile_writer] =
    NetCDFOutputWriter(model, fields,
                       filename = "averaged_z_profile.nc",
                       schedule = AveragedTimeInterval(60, window=20),
                       indices = (1, 1, :))

# output
NetCDFOutputWriter scheduled on TimeInterval(1 minute):
├── filepath: ./averaged_z_profile.nc
├── dimensions: zC(16), zF(17), xC(1), yF(1), xF(1), yC(1), time(0)
├── 2 outputs: (c, u) averaged on AveragedTimeInterval(window=20 seconds, stride=1, interval=1 minute)
└── array type: Array{Float32}
```

`NetCDFOutputWriter` also accepts output functions that write scalars and arrays to disk,
provided that their `dimensions` are provided:

```jldoctest
using Oceananigans

grid = RectilinearGrid(size=(16, 16, 16), extent=(1, 2, 3))

model = NonhydrostaticModel(grid=grid)

simulation = Simulation(model, Δt=1.25, stop_iteration=3)

f(model) = model.clock.time^2; # scalar output

g(model) = model.clock.time .* exp.(znodes(Center, grid)) # vector/profile output

h(model) = model.clock.time .* (   sin.(xnodes(Center, grid, reshape=true)[:, :, 1])
                            .*     cos.(ynodes(Face, grid, reshape=true)[:, :, 1])) # xy slice output

outputs = Dict("scalar" => f, "profile" => g, "slice" => h)

dims = Dict("scalar" => (), "profile" => ("zC",), "slice" => ("xC", "yC"))

output_attributes = Dict(
    "scalar"  => Dict("longname" => "Some scalar", "units" => "bananas"),
    "profile" => Dict("longname" => "Some vertical profile", "units" => "watermelons"),
    "slice"   => Dict("longname" => "Some slice", "units" => "mushrooms")
);

global_attributes = Dict("location" => "Bay of Fundy", "onions" => 7)

simulation.output_writers[:things] =
    NetCDFOutputWriter(model, outputs,
                       schedule=IterationInterval(1), filename="things.nc", dimensions=dims, verbose=true,
                       global_attributes=global_attributes, output_attributes=output_attributes)

# output
NetCDFOutputWriter scheduled on IterationInterval(1):
├── filepath: ./things.nc
├── dimensions: zC(16), zF(17), xC(16), yF(16), xF(16), yC(16), time(0)
├── 3 outputs: (profile, slice, scalar)
└── array type: Array{Float32}
```
"""
function NetCDFOutputWriter(model, outputs; filename, schedule,
                                          dir = ".",
                                   array_type = Array{Float32},
                                      indices = (:, :, :),
                                   with_halos = false,
                            global_attributes = Dict(),
                            output_attributes = Dict(),
                                   dimensions = Dict(),
                           overwrite_existing = nothing,
                                  compression = 0,
                                      verbose = false)

    mkpath(dir)
    filename = auto_extension(filename, ".nc")
    filepath = joinpath(dir, filename)

    if isnothing(overwrite_existing)
        if isfile(filepath)
            overwrite_existing = false
        else
            overwrite_existing = true
        end
    else

        if isfile(filepath) && !overwrite_existing
            @warn "$filepath already exists and `overwrite_existing = false`. Mode will be set to append to existing file. " *
                  "You might experience errors when writing output if the existing file belonged to a different simulation!"

        elseif isfile(filepath) && overwrite_existing
            @warn "Overwriting existing $filepath."

        end
    end

    mode = overwrite_existing ? "c" : "a"

    # TODO: This call to dictify is only necessary because "dictify" is hacked to help
    # with LagrangianParticles output (see the end of the file).
    # We shouldn't support this in the future; we should require users to 'name' LagrangianParticles output.
    outputs = dictify(outputs)
    outputs = Dict(string(name) => construct_output(outputs[name], model.grid, indices, with_halos) for name in keys(outputs))
    output_attributes = dictify(output_attributes)
    global_attributes = dictify(global_attributes)
    dimensions = dictify(dimensions)

    # Ensure we can add any kind of metadata to the global attributes later by converting to Dict{Any, Any}.
    global_attributes = Dict{Any, Any}(global_attributes)

    # Add useful metadata
    global_attributes["date"] = "This file was generated on $(now())."
    global_attributes["Julia"] = "This file was generated using " * versioninfo_with_gpu()
    global_attributes["Oceananigans"] = "This file was generated using " * oceananigans_versioninfo()

    add_schedule_metadata!(global_attributes, schedule)

    # Convert schedule to TimeInterval and each output to WindowedTimeAverage if
    # schedule::AveragedTimeInterval
    schedule, outputs = time_average_outputs(schedule, outputs, model)

    dims = default_dimensions(outputs, model.grid, indices, with_halos)

    # Open the NetCDF dataset file
    dataset = NCDataset(filepath, mode, attrib=global_attributes)

    # Define variables for each dimension and attributes if this is a new file.
    if mode == "c"
        for (dim_name, dim_array) in dims
            defVar(dataset, dim_name, array_type(dim_array), (dim_name,),
                   compression=compression, attrib=default_dimension_attributes[dim_name])
        end

        # DateTime and TimeDate are both <: AbstractTime
        time_attrib = model.clock.time isa AbstractTime ?
            Dict("longname" => "Time", "units" => "seconds since 2000-01-01 00:00:00") :
            Dict("longname" => "Time", "units" => "seconds")

        # Creates an unlimited dimension "time"
        defDim(dataset, "time", Inf)
        defVar(dataset, "time", eltype(model.grid), ("time",), attrib=time_attrib)

        # Use default output attributes for known outputs if the user has not specified any.
        # Unknown outputs get an empty tuple (no output attributes).
        for c in keys(outputs)
            if !haskey(output_attributes, c)
                output_attributes[c] = c in keys(default_output_attributes) ? default_output_attributes[c] : ()
            end
        end

        for (name, output) in outputs
            attributes = try output_attributes[name]; catch; Dict(); end
            define_output_variable!(dataset, output, name, array_type, compression, attributes, dimensions)
        end

        sync(dataset)
    end

    close(dataset)

    return NetCDFOutputWriter(filepath, dataset, outputs, schedule, overwrite_existing, array_type, 0.0, verbose)
end

#####
##### Variable definition
#####

""" Defines empty variables for 'custom' user-supplied `output`. """
function define_output_variable!(dataset, output, name, array_type, compression, output_attributes, dimensions)
    name ∉ keys(dimensions) && error("Custom output $name needs dimensions!")

    defVar(dataset, name, eltype(array_type), (dimensions[name]..., "time"),
           compression=compression, attrib=output_attributes)

    return nothing
end

""" Defines empty field variable. """
define_output_variable!(dataset, output::AbstractField, name, array_type, compression, output_attributes, dimensions) =
    defVar(dataset, name, eltype(array_type),
           (netcdf_spatial_dimensions(output)..., "time"),
           compression=compression, attrib=output_attributes)

""" Defines empty field variable for `WindowedTimeAverage`s over fields. """
define_output_variable!(dataset, output::WindowedTimeAverage{<:AbstractField}, args...) =
    define_output_variable!(dataset, output.operand, args...)

#####
##### Write output
#####

Base.open(nc::NetCDFOutputWriter) = NCDataset(nc.filepath, "a")
Base.close(nc::NetCDFOutputWriter) = close(nc.dataset)

function save_output!(ds, output, model, ow, time_index, name)
    data = fetch_and_convert_output(output, model, ow)
    data = drop_output_dims(output, data)
    colons = Tuple(Colon() for _ in 1:ndims(data))
    ds[name][colons..., time_index] = data
    return nothing
end

function save_output!(ds, output::LagrangianParticles, model, ow, time_index, name)
    data = fetch_and_convert_output(output, model, ow)
    for (particle_field, vals) in pairs(data)
        ds[string(particle_field)][:, time_index] = vals
    end

    return nothing
end

"""
    write_output!(output_writer, model)

Writes output to netcdf file `output_writer.filepath` at specified intervals. Increments the `time` dimension
every time an output is written to the file.
"""
function write_output!(ow::NetCDFOutputWriter, model)
    ow.dataset = open(ow)

    ds, verbose, filepath = ow.dataset, ow.verbose, ow.filepath

    time_index = length(ds["time"]) + 1
    ds["time"][time_index] = float_or_date_time(model.clock.time)

    if verbose
        @info "Writing to NetCDF: $filepath..."
        @info "Computing NetCDF outputs for time index $(time_index): $(keys(ow.outputs))..."

        # Time and file size before computing any outputs.
        t0, sz0 = time_ns(), filesize(filepath)
    end

    for (name, output) in ow.outputs
        # Time before computing this output.
        verbose && (t0′ = time_ns())

        save_output!(ds, output, model, ow, time_index, name)

        if verbose
            # Time after computing this output.
            t1′ = time_ns()
            @info "Computing $name done: time=$(prettytime((t1′-t0′) / 1e9))"
        end
    end

    if verbose
        # Time and file size after computing and writing all outputs.
        t1, sz1 = time_ns(), filesize(filepath)
        verbose && @info begin
            @sprintf("Writing done: time=%s, size=%s, Δsize=%s",
                    prettytime((t1-t0)/1e9), pretty_filesize(sz1), pretty_filesize(sz1-sz0))
        end
    end

    sync(ds)
    close(ow)

    return nothing
end

drop_output_dims(output, data) = data # fallback
drop_output_dims(output::Field, data) = dropdims(data, dims=reduced_dimensions(output))
drop_output_dims(output::WindowedTimeAverage{<:Field}, data) = dropdims(data, dims=reduced_dimensions(output.operand))

#####
##### Show
#####

Base.summary(ow::NetCDFOutputWriter) =
    string("NetCDFOutputWriter writing ", prettykeys(ow.outputs), " to ", ow.filepath, " on ", summary(ow.schedule))

function Base.show(io::IO, ow::NetCDFOutputWriter)
    dims = NCDataset(ow.filepath, "r") do ds
        join([dim * "(" * string(length(ds[dim])) * "), "
              for dim in keys(ds.dim)])[1:end-2]
    end

    averaging_schedule = output_averaging_schedule(ow)
    Noutputs = length(ow.outputs)

    print(io, "NetCDFOutputWriter scheduled on $(summary(ow.schedule)):", '\n',
              "├── filepath: ", ow.filepath, '\n',
              "├── dimensions: $dims", '\n',
              "├── $Noutputs outputs: ", prettykeys(ow.outputs), show_averaging_schedule(averaging_schedule), '\n',
              "└── array type: ", show_array_type(ow.array_type))
end

#####
##### Support / hacks for Lagrangian particles output
#####

""" Defines empty variable for particle trackting. """
function define_output_variable!(dataset, output::LagrangianParticles, name, array_type, compression, output_attributes, dimensions)
    particle_fields = eltype(output.properties) |> fieldnames .|> string
    for particle_field in particle_fields
        defVar(dataset, particle_field, eltype(array_type),
               ("particle_id", "time"), compression=compression)
    end
end

dictify(outputs::LagrangianParticles) = Dict("particles" => outputs)

default_dimensions(outputs::Dict{String,<:LagrangianParticles}, grid, indices, with_halos) =
    Dict("particle_id" => collect(1:length(outputs["particles"])))


text:
"""
applies a ↦ b = Ka
"""
function  apply_K_sep!(b::AbstractMatrix, a, k, tmp)
    mul!(tmp, k[1], a)       # :t[1] = k[1] * :b
    mul!(b, tmp, k[2]')
end

"""
applies a ↦ b = Ka
"""
@inline function  apply_K_sep!(b::AbstractVector, a, K, tmp)
    mul!(b, K, a)
end

"""
wasserstein_distance_separated(p, β, a₀, b₀, d₁₀, d₂₀, k, SP, caches)
p: right marginal as nxn matrix
β: left marginal as nxn matrix
k: e^(c/ε) with cost c as dx(nxn) vector of matrices

Only differentiable by α

TODO: higher dimensions than d=2, number of grid points different in dimensions/marginals
"""

function sinkhorn_dvg_separated(α::AbstractArray{T}, β::AbstractArray{V},
                                a₀, b₀, d₁₀, d₂₀,
                                k, SP, caches
                                ) where {T, V}

    MC = caches.MC

    MC[:α,V] .= ForwardDiff.value.(α)

    MC[:a,V] .= a₀;   MC[:b,V] .= b₀
    if SP.debias
        MC[:d₁,V] .= d₁₀; MC[:d₂,V] .= d₂₀
    end

    for l in 1:SP.L
        if SP.averaged_updates
            apply_K_sep!(MC[:t2,V], MC[:b,V], k, MC[:t1,V])
            MC[:a₊,V] .= MC[:α,V] ./ MC[:t2,V]
            apply_K_sep!(MC[:t2,V], MC[:a,V], k, MC[:t1,V])
            MC[:b₊,V] .= β ./ MC[:t2,V]

            MC[:b,V] .= sqrt.(MC[:b,V] .* MC[:b₊,V])
            MC[:a,V] .= sqrt.(MC[:a,V] .* MC[:a₊,V])
        else
            apply_K_sep!(MC[:t2,V], MC[:b,V], k, MC[:t1,V])
            MC[:a,V] .= MC[:α,V] ./ MC[:t2,V]
            apply_K_sep!(MC[:t2,V], MC[:a,V], k, MC[:t1,V])
            MC[:b,V] .= β ./ MC[:t2,V]
        end

        if SP.debias #&& ( l%2==0 || l==SP.L )
            if SP.averaged_updates
                apply_K_sep!(MC[:t2,V], MC[:d₁,V], k, MC[:t1,V])
                MC[:d₁₊,V] .= sqrt.( MC[:d₁,V] .* MC[:α,V] ./ MC[:t2,V] )
                apply_K_sep!(MC[:t2,V], MC[:d₂,V], k, MC[:t1,V])
                MC[:d₂₊,V] .= sqrt.( MC[:d₂,V] .* β ./ MC[:t2,V] )

                MC[:d₁,V] .= sqrt.(MC[:d₁,V] .* MC[:d₁₊,V])
                MC[:d₂,V] .= sqrt.(MC[:d₂,V] .* MC[:d₂₊,V])
            else
                apply_K_sep!(MC[:t2,V], MC[:d₁,V], k, MC[:t1,V])
                MC[:d₁,V] .= sqrt.( MC[:d₁,V] .* MC[:α,V] ./ MC[:t2,V] )
                apply_K_sep!(MC[:t2,V], MC[:d₂,V], k, MC[:t1,V])
                MC[:d₂,V] .= sqrt.( MC[:d₂,V] .* β ./ MC[:t2,V] )
            end
        end
    end

    if SP.debias
        for i in eachindex(MC[:t2,T])
                MC[:t2,T][i] = α[i] * (_safe_log( MC[:a,V][i] ) - _safe_log( MC[:d₁,V][i] )) +
                               β[i] * (_safe_log( MC[:b,V][i] ) - _safe_log( MC[:d₂,V][i] ))
        end
    else
        for i in eachindex(MC[:t2,T])
                MC[:t2,T][i] = α[i] * _safe_log( MC[:a,V][i] ) + β[i] * _safe_log( MC[:b,V][i] )
        end
    end

    S_ε = SP.ε * sum( MC[:t2,T] )

    if SP.update_potentials
        a₀ .= MC[:a,V]
        b₀ .= MC[:b,V]
        if SP.debias
            d₁₀ .= MC[:d₁,V]
            d₂₀ .= MC[:d₂,V]
        end
    end

    return S_ε
end

function sinkhorn_dvg_separated(α::AbstractArray{T}, β::AbstractArray{V},
                                k, SP, caches
                                ) where {T, V}
    sinkhorn_dvg_separated( α, β,
                            ones(V, size(α)), ones(V, size(β)),
                            ones(V, size(α)), ones(V, size(β)),
                            k, SP, caches
                            )
end

"""
wasserstein_barycenter_separated!(α, β, k, ε, L)
w: un-normalized weights as s vector
α: input marginals as sx(nxn) vector of matrices
k: Gibbs factor as a vector of dx(nxn) matrices
ε: entropic regularization parameter
L: number of Sinkhorn iterations
MC: Separated Cache Dictionary

returns the Wasserstein barycenter μ as an nxn matrix

TODO: higher dimensions than d=2, number of grid points different in dimensions/marginals
"""

function sinkhorn_barycenter_separated(  λ::AbstractVector{T}, α::AbstractVector{AT},
                                            b_₀, d₀,
                                            k, SP, caches,
                                            ) where {T, V, AT <: AbstractArray{V}}

    MC = caches.MC
    VMC = caches.VMC

    for s in 1:VMC.S
        VMC[:b,T][s] .= b_₀[s]
    end
    if SP.debias
        MC[:d,T] .= d₀
    end

    #=
    if SP.update_warmstart
        for s in 1:VMC.S
            VMC[:bL,T][s] .= VMC[:b,T][s]
        end
        if SP.debias
            MC[:dL,T] .= MC[:d,T]
        end
        SP.update_warmstart = false
    end

    if SP.warmstart == "ones"
        for s in 1:VMC.S
            VMC[:b,T][s] .= 1
        end
        if SP.debias
            MC[:d,T] .= 1
        end
    elseif SP.warmstart == "last_update"
        for s in 1:VMC.S
            VMC[:b,T][s] .= VMC[:bL,T][s]
        end
        if SP.debias
            VMC[:d,T] .= VMC[:dL,T]
        end
    elseif SP.warmstart == "latest"
        nothing
    else
        error("No valid warmstart")
    end
    =#

    for l in 1:SP.L # Sinkhorn loop
       for s in 1:VMC.S
            apply_K_sep!(MC[:t2,T], VMC[:b,T][s], k, MC[:t1,T])
            MC[:t2,T] .= α[s] ./ MC[:t2,T]      # this is aˡₛ
            apply_K_sep!(VMC[:φ,T][s], MC[:t2,T], k, MC[:t1,T])
        end

        SP.debias ? MC[:μ,T] .= MC[:d,T] : MC[:μ,T] .= 1

        for s in 1:MC.S
            MC[:μ,T] .*= VMC[:φ,T][s] .^ λ[s]
        end

        for s in 1:MC.S
            VMC[:b,T][s] .= MC[:μ,T] ./ VMC[:φ,T][s]
        end

        if SP.debias # && ( l%2==0 || l==SP.L )
            apply_K_sep!(MC[:t2,T], MC[:d,T], k, MC[:t1,T])
            MC[:d,T] .= sqrt.( MC[:d,T] .* MC[:μ,T] ./ MC[:t2,T] )
        end
    end

    if SP.update_potentials
        for s in 1:VMC.S
            b_₀[s] .= ForwardDiff.value.(VMC[:b,T][s])
        end
        if SP.debias
            d₀ .= ForwardDiff.value.(MC[:d,T])
        end
    end

    return copy(MC[:μ,T])
end

function sinkhorn_barycenter_separated(  λ::AbstractVector{T}, α::AbstractVector{AT},
                                            k, SP, caches,
                                            ) where {T, V, AT <: AbstractArray{V}}
    sinkhorn_barycenter_separated(   λ, α,
                                        [ones(V, size(α[s])) for s in eachindex(α)],
                                        ones(V, size(α[1])),
                                        k, SP, caches)
end

max_stars_repo_name: JuliaRCM/WassersteinDictionaries

max_issues_repo_name: JuliaRCM/WassersteinDictionaries

max_forks_repo_name: JuliaRCM/WassersteinDictionaries


text:
# https://spec.openapis.org/oas/v3.1.0

# https://swagger.io/specification/
# module OpenAPIv3
using StructTypes

export OpenAPI

struct ExternalDocumentationObject
    description::String
    url::String
end

StructTypes.StructType(::Type{ExternalDocumentationObject}) = StructTypes.Struct()

struct TagObject
    name::String
    description::String
    externalDocs::ExternalDocumentationObject
end

StructTypes.StructType(::Type{TagObject}) = StructTypes.Struct()

struct DiscriminatorObject
    propertyName::String
    mapping::Dict{String,String}
end

StructTypes.StructType(::Type{DiscriminatorObject}) = StructTypes.Struct()

struct SchemaObject
    nullable::Bool
    discriminator::DiscriminatorObject
    readOnly::Bool
    writeOnly::Bool
    externalDocs::ExternalDocumentationObject
    example::Any
    deprecated::Bool
end

StructTypes.StructType(::Type{SchemaObject}) = StructTypes.Struct()

struct ServerVariable
    enum::Array{String}
    default::String
    description::String
end

StructTypes.StructType(::Type{ServerVariable}) = StructTypes.Struct()

struct ServerObject
    url::String
    description::String
    variables::Dict{String,ServerVariable}
end

StructTypes.StructType(::Type{ServerObject}) = StructTypes.Struct()

struct HeaderObject
    description::String
    required::Bool
    depecated::Bool
    allowEmptyValue::Bool
    style::String
    explode::Bool
    allowReserved::Bool
end

StructTypes.StructType(::Type{HeaderObject}) = StructTypes.Struct()

struct LinkObject
    operationRef::String
    operationId::String
    parameters::Dict{String,String}
    requestBody::String
    description::String
    server::ServerObject
end

StructTypes.StructType(::Type{LinkObject}) = StructTypes.Struct()

struct Contact
    name::String
    url::String
    email::String
end

StructTypes.StructType(::Type{Contact}) = StructTypes.Struct()

struct ExampleObject
    summary::String
    description::String
    # can have either value or external value not both
    value::Any
    externalValue::String
end

StructTypes.StructType(::Type{ExampleObject}) = StructTypes.Struct()

Base.@kwdef struct EncodingObject
    contentType::String # */* for unknown formats
    headers::Dict{String,HeaderObject}
    styled::Union{String,Nothing} = nothing
    explode::Union{Bool,Nothing} = nothing
    allowReserved::Bool = false
end

StructTypes.StructType(::Type{EncodingObject}) = StructTypes.Struct()
StructTypes.omitempties(::Type{EncodingObject}) = true

Base.@kwdef struct MediaTypeObject
    schema::JSONSchema
    example::Union{Any,Nothing} = nothing
    examples::Union{Dict{String,ExampleObject},Nothing} = nothing
    # only relevant for forms, I think?
    encoding::Union{Dict{String,EncodingObject},Nothing} = nothing
end

StructTypes.StructType(::Type{MediaTypeObject}) = StructTypes.UnorderedStruct()
StructTypes.omitempties(::Type{MediaTypeObject}) = true

MediaTypeObject(t::DataType) = MediaTypeObject(
    schema=json_schema(t)
)

Base.@kwdef struct ParameterObject
    name::String
    in::String # any of - query header path cookie
    description::Union{String,Nothing} = nothing
    required::Bool = true
    depecated::Bool = false
    allowEmptyValue::Bool = false
    style::Union{String,Nothing} = nothing
    explode::Union{Bool,Nothing} = nothing
    schema::Union{JSONSchema,Nothing} = nothing
    example::Union{Any,Nothing} = nothing
    examples::Union{Dict{String,ExampleObject},Nothing} = nothing
    content::Union{Dict{String,MediaTypeObject},Nothing} = nothing

    # only applies to query
    allowReserved::Union{Bool,Nothing} = nothing
end

StructTypes.StructType(::Type{ParameterObject}) = StructTypes.Struct()
StructTypes.omitempties(::Type{ParameterObject}) = true

struct OAuthFlowObject
    authorizationUrl::String
    tokenUrl::String
    refreshUrl::String
    scopes::Dict{String,String}
end

StructTypes.StructType(::Type{OAuthFlowObject}) = StructTypes.Struct()

struct SecurityRequirementObject
    # custom serialization
    requirements::Dict{String,Array{String}}
end

StructTypes.StructType(::Type{SecurityRequirementObject}) = StructTypes.Struct()

struct License
    name::String
    url::String
end

StructTypes.StructType(::Type{License}) = StructTypes.Struct()

Base.@kwdef struct RequestBodyObject
    description::Union{String,Nothing} = nothing
    content::Dict{String,MediaTypeObject}
    required::Bool = false
end

StructTypes.StructType(::Type{RequestBodyObject}) = StructTypes.Struct()
StructTypes.omitempties(::Type{RequestBodyObject}) = true

Base.@kwdef struct ResponseObject
    # keys are the content type
    # application/json => value
    content::Dict{String,MediaTypeObject}
    description::Union{String,Nothing} = nothing
    headers::Union{Dict{String,HeaderObject},Nothing} = nothing
    links::Union{Dict{String,LinkObject},Nothing} = nothing
end

StructTypes.StructType(::Type{ResponseObject}) = StructTypes.Struct()
StructTypes.omitempties(::Type{ResponseObject}) = true

struct ResponsesObject
    default::ResponseObject
    # needs custom serialization
    statusCodes::Dict{String,ResponseObject}
end

StructTypes.StructType(::Type{ResponsesObject}) = StructTypes.Struct()
StructTypes.omitempties(::Type{ResponsesObject}) = true

Base.@kwdef mutable struct OperationObject
    responses::Dict{String,ResponseObject}
    tags::Union{Array{String},Nothing} = nothing
    summary::Union{String,Nothing} = nothing
    description::Union{String,Nothing} = nothing
    externalDocs::Union{ExternalDocumentationObject,Nothing} = nothing
    operationId::String = string(uuid4())
    parameters::Array{ParameterObject} = nothing
    requestBody::Union{RequestBodyObject,Nothing} = nothing
    callbacks::Union{Dict{String,Dict},Nothing} = nothing
    depercated::Bool = false
    security::Union{Array{SecurityRequirementObject},Nothing} = nothing
    servers::Union{Array{ServerObject},Nothing} = nothing
end

StructTypes.StructType(::Type{OperationObject}) = StructTypes.Struct()
StructTypes.omitempties(::Type{OperationObject}) = true

Base.@kwdef mutable struct PathItemObject
    summary::Union{String,Nothing} = nothing
    description::Union{String,Nothing} = nothing
    get::Union{Nothing,OperationObject} = nothing
    put::Union{Nothing,OperationObject} = nothing
    post::Union{Nothing,OperationObject} = nothing
    delete::Union{Nothing,OperationObject} = nothing
    options::Union{Nothing,OperationObject} = nothing
    head::Union{Nothing,OperationObject} = nothing
    patch::Union{Nothing,OperationObject} = nothing
    trace::Union{Nothing,OperationObject} = nothing
    servers::Union{Nothing,Array{ServerObject}} = nothing
    parameters::Union{ParameterObject,Nothing} = nothing
end

StructTypes.StructType(::Type{PathItemObject}) = StructTypes.Struct()
StructTypes.omitempties(::Type{PathItemObject}) = true

struct CallbackObject
    # needs custom serialization
    callbacks::Dict{String,PathItemObject}
end

StructTypes.StructType(::Type{CallbackObject}) = StructTypes.Struct()

struct OAuthFlowsObject
    implict::OAuthFlowObject
    password::OAuthFlowObject
    clientCredntials::OAuthFlowObject
    authorizationCode::OAuthFlowObject
end

StructTypes.StructType(::Type{OAuthFlowsObject}) = StructTypes.Struct()

struct SecuritySchemeObject
    type::String
    description::String
    name::String
    in::String
    scheme::String
    bearerFormatString
    flows::OAuthFlowsObject
    openIdConnectUrl::String
end

StructTypes.StructType(::Type{SecuritySchemeObject}) = StructTypes.Struct()

const PathsObject = Pair{String,PathItemObject}

Base.@kwdef struct Info
    title::String = ""
    description::Union{String,Nothing} = nothing
    termsOfService::Union{String,Nothing} = nothing
    contact::Union{Contact,Nothing} = nothing
    license::Union{License,Nothing} = nothing
    version::String = "3.0.0"
end

StructTypes.StructType(::Type{Info}) = StructTypes.Struct()

struct Components
    schema::Dict{String,Union{SchemaObject}}
    responses::Dict{String,ResponseObject}
    parameters::Dict{String,ResponseObject}
    examples::Dict{String,ExampleObject}
    requestBodies::Dict{String,RequestBodyObject}
    securitySchemes::Dict{String,SecuritySchemeObject}
    links::Dict{String,LinkObject}
    callbacks::Dict{String,CallbackObject}
end

StructTypes.StructType(::Type{Components}) = StructTypes.Struct()

Base.@kwdef mutable struct OpenAPI
    openapi::String = "3.0" # semantic version number
    info::Union{Info,Nothing} = nothing
    servers::Array{ServerObject} = []
    paths::Dict{String,PathItemObject} = Dict()
    components::Union{Components,Nothing} = nothing
    security::Array{SecurityRequirementObject} = []
    tags::Array{TagObject} = []
    externalDocs::Union{ExternalDocumentationObject,Nothing} = nothing
end

StructTypes.StructType(::Type{OpenAPI}) = StructTypes.Struct()
StructTypes.omitempties(::Type{OpenAPI}) = true

function open_api_parameters(leaf::Leaf)
    l = ParameterObject[]
    (; handler) = leaf
    for t in handler_reads(handler)
        push!(l, open_api_parameters(t)...)
    end
    l
end

function open_api_parameters(::Type{A}) where {A<:HttpParameter}

    T = parameter_type(A)

    l = ParameterObject[]

    in = if A <: Query
        "query"
    elseif A <: Params
        "path"
    elseif A <: Headers
        "header"
    else
        error("invalid type $A for ParameterObject")
    end

    StructTypes.foreachfield(T) do i, field, field_type
        schema = json_schema(field_type)
        d = (
            name=string(field),
            in=in,
            required=Missing <: field_type,
            schema=schema,
        )
        push!(l, ParameterObject(; d...))
    end
    l
end

function RequestBodyObject(::Type{Body{T}}) where {T}

    if StructType(T) == NoStructType()
        error("no struct type defined for $T")
    end

    media_type = MediaTypeObject(
        schema=json_schema(T),
    )

    RequestBodyObject(
        content=Dict(
            "application/json" => media_type
        ),
        required=true,
    )
end

function ResponseObject(t::DataType)

    st = StructTypes.StructType(t)
    if st == NoStructType()
        error("unsuppored type $t")
    end

    media_type = MediaTypeObject(
        schema=JSONSchema(; json_schema(t)...),
    )

    ResponseObject(
        content=Dict(
            "application/json" => media_type
        ),
    )
end

OperationObject(h::HttpHandler) = OperationObject(h.fn)
OperationObject(h::Middleware) = OperationObject(h.fn)

function OperationObject(handler)
    writes = handler_writes(handler)
    responses = Dict{String,ResponseObject}()
    # for now we will ignore duplicate response codes
    # and assume everything is a json response
    for (res_type, res_code) in writes
        k = string(Int(res_code))
        responses[k] = ResponseObject(
            content=Dict(
                "application/json" => MediaTypeObject(res_type)
            )
        )
    end

    params = ParameterObject[]
    requestBody = nothing

    for p in handler_reads(handler)
        # @debug "parameters" p=p
        if p <: Body
            requestBody = RequestBodyObject(p)
        else
            push!(params, open_api_parameters(p)...)
        end
    end

    OperationObject(
        responses=responses,
        requestBody=requestBody,
        parameters=params,
    )
end

function OperationObject(leaf::Leaf)
    o = OperationObject(leaf.handler)
    o.operationId = leaf.path
    return o
end

# function PathItemObject(leaves::Array{Leaf})
#     PathItemObject(;d...)
# end

function OpenAPI(app)

    paths = Dict{String, Dict{Symbol, Any}}()

    leaves = []
    for n in PostOrderDFS(app.paths)
        if !isempty(n.methods)
            push!(leaves, n.methods...)
        end
    end

    for leaf in leaves
        (; path, method) = leaf

        if path == app.docs
            @warn "skiping $(path)"
            continue
        end

        # @info "METHOD" method=method path=path
        # @info path

        d = get(paths, path, Dict())
        o = OperationObject(leaf)

        method = Symbol(lowercase(method)) # "GET" -> :get
        docs = app.paths_docs[method]
        description = get(docs, path, nothing)
        o.description = description

        d[method] = o
        paths[path] = d
    end

    paths = Dict(
        k => PathItemObject(; v...) for (k, v) in paths
    )

    OpenAPI(
        paths=paths,
        info=Info()
    )
end


text:
cd(@__DIR__)
using Pkg
Pkg.activate(@__DIR__)

# Set Python executable to current and re-build PyCall if necessary
ENV["PYTHON"] = Sys.which("python")
try
    Pkg.build("PyCall")
catch e
    if !isa(e, Pkg.Types.PkgError)
        throw(e)
    end
end

# Precompile everything
Pkg.API.precompile()

# Use packages
# using SpineInterface # SpineOpt uses SpineInterface anyway hence commented
using SpineOpt
using Dates # For debugging with the time slice tools
using CSV # Because is useful sometimes
using DataFrames
using GEPPR # My GEP model to check the results, can get rid of at end
using JuMP
using Plots
using AxisArrays
using DataStructures
using JSON
using PyCall

# Useful vars
dataPath = abspath(@__DIR__, "data")
spineitemsPath = joinpath(
    @__DIR__, "data", "case_study_b2", ".spinetoolbox", "items"
)
GEPPRDataPath = abspath(dataPath, "GEPPR_input_data")


text:
########
# path #
########

training_data_path() = mkpath(joinpath(datadep"SqState", "training_data"))
model_path() = mkpath(joinpath(datadep"SqState", "model"))

############
# gen data #
############

rand2range(x_range) = x_range[1] + (x_range[2]-x_range[1])*rand()

function rand_arg_sqth_th(r_range, θ_range, n̄_range)
    r = rand2range(r_range)
    θ = rand2range(θ_range)
    n̄ = rand2range(n̄_range)
    n̄0 = rand2range(n̄_range)
    c1 = rand()
    c2 = 1 - c1

    return r, θ, n̄, n̄0, c1, c2
end

function rand_arg_sqth(r_range, θ_range, n̄_range)
    r = rand2range(r_range)
    θ = rand2range(θ_range)
    n̄ = rand2range(n̄_range)

    return r, θ, n̄
end

function construct_state_sqth_th(r, θ, n̄, n̄0, c1, c2, dim)
    state =
        c1 * SqueezedThermalState(ξ(r, θ), n̄, dim=dim) +
        c2 * ThermalState(n̄0, dim=dim)

    return state
end

function construct_state_sqth(r, θ, n̄, dim)
    return SqueezedThermalState(ξ(r, θ), n̄, dim=dim)
end

#########
# model #
#########

to_complex(𝐱1::AbstractArray, 𝐱2::AbstractArray) = 𝐱1 + im.*𝐱2

function ChainRulesCore.rrule(::typeof(to_complex), 𝐱1::AbstractArray, 𝐱2::AbstractArray)
    function to_complex_pullback(𝐲̄)
        return NoTangent(), real.(𝐲̄), imag.(𝐲̄)
    end

    return to_complex(𝐱1, 𝐱2), to_complex_pullback
end

function reshape_cholesky(x)
    dim = Int(sqrt(size(x, 1)))
    𝐱_row = reshape(x, dim, dim, :)
    𝐱_real = cat([reshape(tril(𝐱_row[:, :, i]), dim, dim, 1) for i in axes(𝐱_row, 3)]..., dims=3)
    𝐱_imag = cat([reshape(tril(𝐱_row[:, :, i]', -1), dim, dim, 1) for i in axes(𝐱_row, 3)]..., dims=3)
    𝐱 = to_complex(𝐱_real, 𝐱_imag)

    return 𝐱
end

function cholesky2ρ(x)
    𝐱 = reshape_cholesky(Zygote.hook(real, x))
    𝛒 = Flux.batched_mul(𝐱, Flux.batched_adjoint(𝐱))
    𝛒 = cat([reshape(𝛒[:, :, i]/tr(𝛒[:, :, i]), 1, size(𝛒, 1), size(𝛒, 2), 1) for i in axes(𝛒, 3)]..., dims=4)

    return vcat(real.(𝛒), imag.(𝛒))
end

function res_block(
    ch::NTuple{4, <:Integer},
    conv_kernel_size::NTuple{3, <:Integer},
    conv_pad::NTuple{3, <:Any},
    shortcut_kernel_size::Integer,
    shortcut_pad::Any,
    pool_size::Integer,
    σ=identity
)
    conv_layers = Chain(
        Conv((conv_kernel_size[1], ), ch[1]=>ch[2], σ,  pad=conv_pad[1]),
        Conv((conv_kernel_size[2], ), ch[2]=>ch[3], σ, pad=conv_pad[2]),
        Conv((conv_kernel_size[3], ), ch[3]=>ch[4], σ, pad=conv_pad[3]),
    )
    shortcut = Conv((shortcut_kernel_size, ), ch[1]=>ch[end], σ, pad=shortcut_pad)
    pool = (pool_size > 0) ? MeanPool((pool_size, )) : identity

    return Chain(
        Parallel(+, conv_layers, shortcut),
        x -> σ.(x),
        pool,
        BatchNorm(ch[end], σ)
    )
end

############
# training #
############

function update_model!(model_path::String, model_name::String, model)
    model = cpu(model)
    jldsave(joinpath(model_path, "$model_name.jld2"); model)
    @warn "'$model_name' updated!"
end

function get_device()
    if has_cuda()
        @info "CUDA is on"
        device = gpu
        CUDA.allowscalar(false)
    else
        device = cpu
    end

    return device
end

#############
# inference #
#############

function get_model(model_name::String)
    f = jldopen(joinpath(model_path() , "$model_name.jld2"))
    model = f["model"]
    close(f)

    return model
end