text:
function default_algorithm(prob::DiffEqBase.AbstractODEProblem{uType,tType,inplace};kwargs...) where {uType,tType,inplace}
  o = Dict{Symbol,Any}(kwargs)
  extra_kwargs = Any[]; alg=AutoTsit5(Rosenbrock23(autodiff=false)) # Standard default
  uEltype = eltype(prob.u0)

  alg_hints = get_alg_hints(o)
  tol_level = get_tolerance_level(o)
  callbacks = callbacks_exists(o)
  mm = mass_matrix_exists(prob)

  if :stiff ∈ alg_hints && :nonstiff ∈ alg_hints
    error("The problem must either be designated as stiff or non-stiff")
  end

  # If adaptivity is not set and the tType is not a float, turn off adaptivity
  # Bad interaction with ForwardDiff
  #!(tType <: AbstractFloat) && (:adaptive ∉ keys(o)) && push!(extra_kwargs,:adaptive=>false)

  if typeof(prob.f) <: SplitFunction
    alg = KenCarp4(autodiff=false)
  elseif typeof(prob.f) <: DynamicalODEFunction
    if tol_level == :low_tol || tol_level == :med_tol
      alg = Tsit5()
    elseif tol_level == :high_tol
      alg = Vern7(lazy=!callbacks)
    else
      alg = Vern9(lazy=!callbacks)
    end
  else # Standard ODE
    if :nonstiff ∈ alg_hints || length(prob.u0) > 10000
      # Don't default to implicit here because of memory requirements
      # And because the linear system gets unruly
      if (!(uEltype <: Float64) && !(uEltype <: Float32) && !(uEltype <: Complex)) || tol_level == :extreme_tol
        # Most likely higher precision, so use a higher order method
        alg = Vern9(lazy=!callbacks)
      elseif tol_level == :low_tol
        alg = Vern7(lazy=!callbacks)
      elseif tol_level == :med_tol
        alg = Tsit5()
      else # tol_level == :high_tol
        alg = BS3()
      end
    elseif :stiff ∈ alg_hints || mm # The problem is stiff
      if length(prob.u0) > 2000
        # Use Krylov method when huge!
        if callbacks && !m
            alg = CVODE_BDF(linear_solver=:GMRES)
        elseif !callbacks
            alg = QNDF(autodiff=false,linsolve=LinSolveGMRES())
        else
            alg = Rodas4(autodiff=false)
        end
      elseif length(prob.u0) > 100
        if callbacks && !m
            alg = CVODE_BDF()
        elseif !callbacks
            alg = QNDF(autodiff=false)
        else
            alg = Rodas4(autodiff=false)
        end
      elseif tol_level == :high_tol
        alg = Rosenbrock23(autodiff=false)
      else
        alg = Rodas4(autodiff=false)
      end
    else # :auto ∈ alg_hints
      if (!(uEltype <: Float64) && !(uEltype <: Float32)) || tol_level == :extreme_tol
        # Most likely higher precision, so use a higher order method
        if length(prob.u0) > 100
            alg = AutoVern9(KenCarp47(autodiff=false),lazy=!callbacks)
        else
            alg = AutoVern9(Rodas5(autodiff=false),lazy=!callbacks)
        end
      elseif tol_level == :low_tol
          if length(prob.u0) > 100
            alg = AutoVern7(TRBDF2(autodiff=false),lazy=!callbacks)
          else
            alg = AutoVern7(Rodas4(autodiff=false),lazy=!callbacks)
          end
      else # :med or low
        if length(prob.u0) > 100
          alg = AutoTsit5(TRBDF2(autodiff=false))
        else
          alg = AutoTsit5(Rosenbrock23(autodiff=false))
        end
      end
    end
  end
  alg,extra_kwargs
end


text:
fp = open("relu.cu","w")

function relusrc(; BLK=256, THR=256)
  sprint() do s
    for (T,F,Fback) in [("float","relu_32_1","reluback_32_1"),("double","relu_64_1","reluback_64_1")]
        print(s,
"""
__global__ void _$(F)(int n, $T max_value, $T negative_slope, $T threshold, $T *x_, $T *y_) {
  int i = threadIdx.x + blockIdx.x * blockDim.x;
  while (i < n) {
    $T x = x_[i];
    y_[i] = (x >= max_value ? max_value : x >= threshold ? x : negative_slope == 0 ? 0 : negative_slope * (x-threshold));
    i += blockDim.x * gridDim.x;
  }
}

__global__ void _$(Fback)(int n, $T max_value, $T negative_slope, $T threshold, $T *x_, $T *y_, $T *dy_, $T *dx_) {
  int i = threadIdx.x + blockIdx.x * blockDim.x;
  while (i < n) {
    $T x=x_[i];
    dx_[i] = dy_[i] * (x >= max_value ? 0 : x >= threshold ? 1 : negative_slope);
    i += blockDim.x * gridDim.x;
  }
}

extern "C" {
  $DLLEXPORT void $F(int n, $T max_value, $T negative_slope, $T threshold, $T *x, $T *y) {
    if (n>0) _$F<<<$BLK,$THR>>>(n,max_value,negative_slope,threshold,x,y);
  }
  $DLLEXPORT void $(F)_stream(int n, $T max_value, $T negative_slope, $T threshold, $T *x, $T *y, cudaStream_t STR) {
    if (n>0) _$F<<<$BLK,$THR,0,STR>>>(n,max_value,negative_slope,threshold,x,y);
  }
  $DLLEXPORT void $(Fback)(int n, $T max_value, $T negative_slope, $T threshold, $T *x, $T *y, $T *dy, $T *dx) {
    _$(Fback)<<<$BLK,$THR>>>(n,max_value,negative_slope,threshold,x,y,dy,dx);
  }
  $DLLEXPORT void $(Fback)_stream(int n, $T max_value, $T negative_slope, $T threshold, $T *x, $T *y, $T *dy, $T *dx, cudaStream_t STR) {
    _$(Fback)<<<$BLK,$THR,0,STR>>>(n,max_value,negative_slope,threshold,x,y,dy,dx);
  }
}
""")
    end
  end
end

print(fp, relusrc())
close(fp)

ext: jl


text:
export CUPTIError

struct CUPTIError <: Exception
    code::CUptiResult
end

Base.convert(::Type{CUptiResult}, err::CUPTIError) = err.code

Base.showerror(io::IO, err::CUPTIError) =
    print(io, "CUPTIError: ", description(err), " (code $(reinterpret(Int32, err.code)), $(name(err)))")

function name(err::CUPTIError)
    str_ref = Ref{Cstring}()
    cuptiGetResultString(err, str_ref)
    unsafe_string(str_ref[])
end

## COV_EXCL_START
function description(err)
    if err.code == CUPTI_SUCCESS
        "no error"
    elseif err.code == CUPTI_ERROR_INVALID_PARAMETER
        "one or more of the parameters is invalid"
    elseif err.code == CUPTI_ERROR_INVALID_DEVICE
        "the device does not correspond to a valid CUDA device"
    elseif err.code == CUPTI_ERROR_INVALID_CONTEXT
        "the context is NULL or not valid"
    elseif err.code == CUPTI_ERROR_INVALID_EVENT_DOMAIN_ID
        "the event domain id is invalid"
    elseif err.code == CUPTI_ERROR_INVALID_EVENT_ID
        "the event id is invalid"
    elseif err.code == CUPTI_ERROR_INVALID_EVENT_NAME
        "the event name is invalid"
    elseif err.code == CUPTI_ERROR_INVALID_OPERATION
        "the current operation cannot be performed due to dependency on other factors"
    elseif err.code == CUPTI_ERROR_OUT_OF_MEMORY
        "unable to allocate enough memory to perform the requested operation"
    elseif err.code == CUPTI_ERROR_HARDWARE
        "an error occurred on the performance monitoring hardware"
    elseif err.code == CUPTI_ERROR_PARAMETER_SIZE_NOT_SUFFICIENT
        "the output buffer size is not sufficient to return all requested data"
    elseif err.code == CUPTI_ERROR_API_NOT_IMPLEMENTED
        "aPI is not implemented"
    elseif err.code == CUPTI_ERROR_MAX_LIMIT_REACHED
        "the maximum limit is reached"
    elseif err.code == CUPTI_ERROR_NOT_READY
        "the object is not yet ready to perform the requested operation"
    elseif err.code == CUPTI_ERROR_NOT_COMPATIBLE
        "the current operation is not compatible with the current state of the object"
    elseif err.code == CUPTI_ERROR_NOT_INITIALIZED
        "CUPTI is unable to initialize its connection to the CUDA driver"
    elseif err.code == CUPTI_ERROR_INVALID_METRIC_ID
        "the metric id is invalid"
    elseif err.code == CUPTI_ERROR_INVALID_METRIC_NAME
        "the metric name is invalid"
    elseif err.code == CUPTI_ERROR_QUEUE_EMPTY
        "the queue is empty"
    elseif err.code == CUPTI_ERROR_INVALID_HANDLE
        "invalid handle (internal?)"
    elseif err.code == CUPTI_ERROR_INVALID_STREAM
        "invalid stream"
    elseif err.code == CUPTI_ERROR_INVALID_KIND
        "invalid kind"
    elseif err.code == CUPTI_ERROR_INVALID_EVENT_VALUE
        "invalid event value"
    elseif err.code == CUPTI_ERROR_DISABLED
        "CUPTI is disabled due to conflicts with other enabled profilers"
    elseif err.code == CUPTI_ERROR_INVALID_MODULE
        "invalid module"
    elseif err.code == CUPTI_ERROR_INVALID_METRIC_VALUE
        "invalid metric value"
    elseif err.code == CUPTI_ERROR_HARDWARE_BUSY
        "the performance monitoring hardware is in use by other client"
    elseif err.code == CUPTI_ERROR_NOT_SUPPORTED
        "the attempted operation is not supported on the current system or device"
    elseif err.code == CUPTI_ERROR_UM_PROFILING_NOT_SUPPORTED
        "unified memory profiling is not supported on the system. Potential reason could be unsupported OS or architecture"
    elseif err.code == CUPTI_ERROR_UM_PROFILING_NOT_SUPPORTED_ON_DEVICE
        "unified memory profiling is not supported on the device"
    elseif err.code == CUPTI_ERROR_UM_PROFILING_NOT_SUPPORTED_ON_NON_P2P_DEVICES
        "unified memory profiling is not supported on a multi-GPU configuration without P2P support between any pair of devices"
    elseif err.code == CUPTI_ERROR_UM_PROFILING_NOT_SUPPORTED_WITH_MPS
        "unified memory profiling is not supported under the Multi-Process Service (MPS) environment with CUDA 7.5"
    elseif err.code == CUPTI_ERROR_CDP_TRACING_NOT_SUPPORTED
        "Devices with compute capability 7.0 don't support CDP tracing"
    elseif err.code == CUPTI_ERROR_VIRTUALIZED_DEVICE_NOT_SUPPORTED
        "profiling on virtualized GPU is not supported"
    elseif err.code == CUPTI_ERROR_CUDA_COMPILER_NOT_COMPATIBLE
        "Profiling results might be incorrect for CUDA applications compiled with nvcc version older than 9.0 for devices with compute capability 6.0 and 6.1"
    elseif err.code == CUPTI_ERROR_INSUFFICIENT_PRIVILEGES
        "user doesn't have sufficient privileges which are required to start the profiling session"
    elseif err.code == CUPTI_ERROR_OLD_PROFILER_API_INITIALIZED
        "old profiling api's are not supported with new profiling api's"
    elseif err.code == CUPTI_ERROR_OPENACC_UNDEFINED_ROUTINE
        "missing definition of the OpenACC API routine in the linked OpenACC library"
    elseif err.code == CUPTI_ERROR_LEGACY_PROFILER_NOT_SUPPORTED
        "an unknown internal error has occurred. Legacy CUPTI Profiling is not supported on devices with Compute Capability 7.5 or higher (Turing+)"
    elseif err.code == CUPTI_ERROR_MULTIPLE_SUBSCRIBERS_NOT_SUPPORTED
        "CUPTI doesn't allow multiple callback subscribers. Only a single subscriber can be registered at a time."
    elseif err.code == CUPTI_ERROR_UNKNOWN
        "an unknown error has occurred"
    else
        "unknown status"
    end
end
## COV_EXCL_STOP

## API call wrapper

# outlined functionality to avoid GC frame allocation
@noinline function throw_api_error(res)
    if res == CUPTI_ERROR_OUT_OF_MEMORY
        throw(OutOfGPUMemoryError())
    else
        throw(CUPTIError(res))
    end
end

function initialize_api()
    CUDA.prepare_cuda_state()
end

macro check(ex, errs...)
    check = :(isequal(err, CUPTI_ERROR_OUT_OF_MEMORY))
    for err in errs
        check = :($check || isequal(err, $(esc(err))))
    end

    quote
        res = @retry_reclaim err->$check $(esc(ex))
        if res != CUPTI_SUCCESS
            throw_api_error(res)
        end

        nothing
    end
end


text:
function receive_msg!(io::IOBuffer, comm::MPI.Comm)
    s = MPI.Probe(0, MPI.MPI_ANY_TAG, comm)
    count = MPI.Get_count(s, UInt8)
    Base.ensureroom(io, count)

    tag = MPI.Get_tag(s)
    received_from = MPI.Get_source(s)

    buf = MPI.Buffer(io.data, count, MPI.Datatype(UInt8))
    MPI.Recv!(buf, received_from, tag, comm)
    io.size = count
    seek(io, 0)

    tag
end

function run_worker(comm::MPI.Comm, funcs::Vector{Union{Function, Nothing}})
    io = IOBuffer()

    while true
        tag = receive_msg!(io, comm)

        if tag == 0
            ex = deserialize(io)::Expr
            Main.eval(ex)
            return true
        elseif tag == 1
            id = deserialize(io)::Int64
            f = deserialize(io)::Function

            while size(funcs, 1) < id
                push!(funcs, nothing)
            end
            funcs[id] = f
            return true
        elseif tag == 2
            return false
        else tag == 3
            id = deserialize(io)::Union{Function, Int64}
            tracker_id = deserialize(io)::Int64
            args = deserialize(io)

            f = if isa(id, Function)
              id
            else
              funcs[id]
            end

            r = try
              f(args...)
            catch e
              showerror(stdout, e, catch_backtrace())
              e
            end

            seek(io, 0)
            serialize(io, tracker_id)
            serialize(io, r)

            buf = MPI.Buffer(io.data, io.size, MPI.Datatype(UInt8))
            MPI.Send(buf, 0, 0, comm)
        end
    end
end

function main_worker(comm::MPI.Comm=MPI.COMM_WORLD)
    if !MPI.Initialized()
        MPI.Init()
    end

    @assert !MPI.Finalized()

    comm = if MPI.Comm_get_parent() == MPI.COMM_NULL
        comm
    else
        MPI.Intercomm_merge(MPI.Comm_get_parent(), true)
    end

    funcs = Union{Function, Nothing}[]

    ret = true
    while ret
      ret = Base.invokelatest(run_worker, comm, funcs)
    end

     # free the merged intercomm https://github.com/open-mpi/ompi/issues/8426
    if MPI.Comm_get_parent() != MPI.COMM_NULL
        MPI.free(comm)
    end

    nothing
end


text:
struct SymbolSelection <: AbstractSelection
    s::Symbol
    b::Bool
    SymbolSelection(s::Symbol, b::Bool) = new(s, b)
end
SymbolSelection(s::Symbol) = SymbolSelection(s, true)
SymbolSelection(s::Vector{Symbol}) = SymbolSelection.(s)
(-)(s::SymbolSelection) = SymbolSelection(s.s, !s.b)

struct IntSelection <: AbstractSelection
    s::Int
    b::Bool
    function IntSelection(s::Int, b::Bool)
        if s == 0
            throw(ArgumentError("Zero is not a valid column index."))
        else
            new(abs(s), b)
        end
    end
end
function IntSelection(s::Int)
    if s == 0
        throw(ArgumentError("Zero is not a valid column index."))
    else
        IntSelection(abs(s), s > 0)
    end
end
(-)(s::IntSelection) = IntSelection(s.s, !s.b)

struct BoolSelection <: AbstractSelection
    s::AbstractArray{Bool}
    b::Bool
end
BoolSelection(s::AbstractArray{Bool}) = BoolSelection(s, true)
(-)(s::BoolSelection) = BoolSelection(s.s, !s.b)

struct RangeSelection{T} <: AbstractSelection where T <: Union{Int,Symbol}
    s1::T
    s2::T
    step::Int
    b::Bool
    function RangeSelection(s1::Int, s2::Int, step::Int, b::Bool)
        if !(sign(s1) == sign(s2) == 1 && ((s2 > s1 && sign(step)) == 1 || (s1 > s2 && sign(step)) == -1 || (s1 == s2 && sign(step) != 0)))
            throw(ArgumentError("The range used for selection ($(s1):$(step):$(s2)) must be non-empty and non-crossing zero."))
        end

        sign(s1) == 1 ? new{Int}(s1, s2, step, b) : new{Int}(abs(s1), abs(s2), abs(step), b)
    end
    RangeSelection(s1::Symbol, s2::Symbol, step::Int, b::Bool) = new{Symbol}(s1, s2, step, b)
end
(-)(s::RangeSelection) = RangeSelection(s.s1, s.s2, s.step, !bool(s))
colrange(s1::Symbol, s2::Symbol; by::Int=1) = RangeSelection(s1, s2, by, true)
colrange(s1::Int, s2::Int; by::Int=1) = RangeSelection(s1, s2, (s1 > s2 ? -1 : 1) * abs(by), true)
Base.show(io::IO, s::RangeSelection) = (!bool(s) && print(io, "-"); print(io, "RangeSelection(", s.s1, ":", s.step == 1 ? "" : "$(s.step):", s.s2, ")"))

function RangeSelection(x::AbstractRange)
    RangeSelection(x[1], x[end], step(x), true)
end
function RangeSelection(x::Union{StepRange{Int,Int},UnitRange{Int}})
    lo,hi = extrema(x)
    range_sign = sign(lo)

    if !(lo != 0 && hi != 0 && length(x) > 0 && range_sign == sign(hi))
        throw(ArgumentError("The range used for selection ($(x)) must be non-empty and non-crossing zero."))
    end

    if (range_sign == 1)
        RangeSelection(x[1], x[end], step(x), true)
    else
        RangeSelection(abs(hi), abs(lo), abs(step(x)), false)
    end
end

struct ArraySelection{S} <: AbstractSelection where T <: Union{Int,Symbol}
    s::Vector{S}
    b::Bool
    ArraySelection(s::AbstractVector{S}, b::Bool) where S <: Union{Int,Symbol} = new{S}(unique(s), b)
    ArraySelection(s::Tuple{Vararg{S}}, b::Bool) where S <: Union{Int,Symbol} = new{S}(collect(unique(s)), b)
end
ArraySelection(s) = ArraySelection(s, true)
(-)(x::ArraySelection) = ArraySelection(x.s, !x.b)

struct KeyPredicateSelection{F} <: AbstractSelection where F <: Callable
    f::F
    b::Bool
end
if_keys(f::Callable) = KeyPredicateSelection(f, true)
if_matches(s::S) where S <: Union{Regex,AbstractString,Char} = if_keys(x -> occursin(s, string(x)))
(-)(x::KeyPredicateSelection) = KeyPredicateSelection(x.f, !x.b)

struct PredicateSelection{F} <: AbstractSelection where F <: Callable
    f::F
    b::Bool
end
if_values(f::Callable) = PredicateSelection(f, true)
if_eltype(t::DataType) = if_values(x -> eltype(x) <: t || eltype(x) <: Union{Missing, t})
(-)(x::PredicateSelection) = PredicateSelection(x.f, !x.b)

struct PairPredicateSelection{F} <: AbstractSelection where F <: Callable
    f::F
    b::Bool
end
if_pairs(f::Callable) = PairPredicateSelection(f, true)
(-)(x::PairPredicateSelection) = PairPredicateSelection(x.f, !x.b)

selection(x::Symbol) = SymbolSelection(x)
selection(x::Int) = IntSelection(x)
selection(x::AbstractVector{Bool}) = BoolSelection(x)
selection(x::StepRange{Int,Int}) = RangeSelection(x)
selection(x::UnitRange{Int}) = RangeSelection(x)
selection(x::AbstractVector{S}) where {S<:Union{Int,Symbol}} = ArraySelection(x)
selection(x::AbstractVector) = isempty(x) ? ArraySelection(Symbol[]) : cols(x...)
selection(x::Tuple{Vararg{S}}) where {S<:Union{Int,Symbol}} = ArraySelection(x)
selection(x::Tuple) = isempty(x) ? ArraySelection(Symbol[]) : cols(x...)
selection(x::Regex) = if_matches(x)
selection(x::DataType) = if_eltype(x)
selection(x::AbstractSelection) = x
selection(x::SelectionQuery) = x
selection(x::SelectionResult) = x

col(x) = selection(x)

cols() = throw(MethodError(cols, ()))
cols(s) = selection(s)
cols(s...) = mapfoldl(selection, OrSelection, [s...])

not() = throw(MethodError(not, ()))
not(x::Pair) = -selection(first(x)) => last(x)
not(s) = -selection(s)
not(s...) = mapfoldl(not, AndSelection, [s...])

"""
```
if_keys(f::Callable)
if_matches(s::Union{AbstractString,Regex})
if_values(f::Callable)
if_eltype(t::DataType)
if_pairs(f::Callable)
colrange(s1, s2 [; step::Int=1])
```
Select columns based on their values and/or names within a call to `select(tab, args...)`.
Usage:
* `if_values(f)` -- Select all columns where `f(column) == true`
* `if_keys(f)` -- Select all columns where `f(colname) == true`
* `if_pairs(f)` -- Select all columns where `f(colname, column) == true`
* `if_eltype(T)` -- Select all columns where `eltype(column) <: T || eltype(column) <: Union{T,Missing}`
* `if_matches(s)` -- Select all columns where `occursin(s, keys(tab)) == true`
* `colrange(s1, s2)` -- Select all columns between `s1` and `s2`. `s1` and `s2` can be both `Symbols` or `Int`s.

See also: [`select`](@ref), [`cols`](@ref), [`other_cols`](@ref), [`renaming`](@ref), [`transformation`](@ref)
"""
if_matches, if_keys, if_values, if_pairs, if_eltype, selection

"""
```
cols(args...)
not(args...)
col(s)
```

Wrappers that turn its inputs into `selection`s and reduce them into a single chained `selection`.
`cols()` reduces all conditions with an `|` boolean operator, `not()` with `&` and also negates its inputs.
`col()` simply converts its single input into a `selection`.

Useful for wrapping `Symbol`s, `Int`s, ranges and boolean vectors so that they can be interpreted as `selections`
and combined with `|`, `&` or `-`.

I.e. `select(tab , :a | :b)` is not defined and it would be considered type piracy, to define it.
Instead, you can use `select(tab, col(:a) | col(:b))` or `select(tab, cols(:a, :b))`.

`cols(args...)` is rougly `mapfoldl(selection, OrSelection, [args...])`

`not(args...)` is rougly `mapfoldl((-)∘selection, AndSelection, [args...])`

See also: [`select`](@ref), [`selection`](@ref)
"""
cols, not, col


text:
module DistributionWrapper

using JavaCall

const Class = @jimport java.lang.Class
const EnumeratedDistribution = @jimport org.hipparchus.distribution.EnumeratedDistribution
const IntegerDistribution = @jimport org.hipparchus.distribution.IntegerDistribution
const JString = @jimport java.lang.String
const List = @jimport java.util.List
const MultivariateRealDistribution = @jimport org.hipparchus.distribution.MultivariateRealDistribution
const Object = @jimport java.lang.Object
const RealDistribution = @jimport org.hipparchus.distribution.RealDistribution

export EnumeratedDistribution
export IntegerDistribution
export MultivariateRealDistribution
export RealDistribution
export check_and_normalize
export cumulative_probability
export density
export equals
export get_class
export get_dimension
export get_numerical_mean
export get_numerical_variance
export get_pmf
export get_support_lower_bound
export get_support_upper_bound
export hash_code
export inverse_cumulative_probability
export is_support_connected
export log_density
export log_probability
export notify
export notify_all
export probability
export reseed_random_generator
export sample
export to_string
export wait

include("ContinuousWrapper/ContinuousWrapper.jl")
include("DiscreteWrapper/DiscreteWrapper.jl")
include("MultivariateWrapper/MultivariateWrapper.jl")
include("enumerated_distribution.jl")
include("integer_distribution.jl")
include("multivariate_real_distribution.jl")
include("real_distribution.jl")

end


text:
# Tweaked DataFrames Index for SharedStrings and SheetName
# https://github.com/JuliaData/DataFrames.jl/blob/master/src/other/index.jl
struct Index
    lookup::Dict{String, Int}
    names::Tuple
end

function Index(names::Vector{String})
    if !allunique(names)
        error("SheetName is duplicated, Excel Do not duplicated SheetName")
    end
    lookup = Dict{String, Int}(zip(names, 1:length(names)))
    Index(lookup, Tuple(names))
end
Index() = Index(Dict{String, Int}(), Tuple{}())

Base.length(x::Index) = length(x.names)
Base.names(x::Index) = copy(x.names)

Base.getindex(x::Index, idx::String) = x.lookup[idx]
Base.getindex(x::Index, idx::Real) = x.names[idx]


text:
include("order.jl")
include("pricelevel.jl")

tic
for i in 1:1000
    LimitOrder(0,0,0,0)
end
toc


text:
abstract type ASECalculatorParameters end

function get_potential_energy(atoms::PyObject, parameters::ASECalculatorParameters)
    @assert pyisinstance(atoms, pyimport("ase").Atoms)
    configure_calculator!(atoms, parameters)
    return atoms.get_potential_energy()
end

function get_forces(atoms::PyObject, parameters::ASECalculatorParameters)
    @assert pyisinstance(atoms, pyimport("ase").Atoms)
    configure_calculator!(atoms, parameters)
    return atoms.get_forces()
end

function get_potential_energy(symbol::String, positions::Vector{Vector{Real}}, parameters::ASECalculatorParameters)
    get_potential_energy(atoms(symbol, positions), parameters)
end

function get_potential_energy(symbol::String, positions::Vector{Vector{Real}}, cell::Vector{Real}, parameters::ASECalculatorParameters)
    get_potential_energy(atoms(symbol, positions, cell), parameters)
end

function get_forces(symbol::String, positions::Vector{Vector{Real}}, parameters::ASECalculatorParameters)
    get_forces(atoms(symbol, positions), parameters)
end

function get_forces(symbol::String, positions::Vector{Vector{Real}}, cell::Vector{Real}, parameters::ASECalculatorParameters)
    get_forces(atoms(symbol, positions, cell), parameters)
end

@with_kw struct DFTKCalculatorParameters <: ASECalculatorParameters
    ecut::Union{Real, Missing} = missing
    functionals::Union{AbstractVector{String}, Missing} = missing
    kpts::Union{Tuple, AbstractArray, AbstractFloat, Missing} = missing
    mixing::Union{String, Missing} = missing
    nbands::Union{Integer, Missing} = missing
    pps::Union{String, Missing} = missing
    scftol::Union{AbstractFloat, Missing} = missing
    smearing::Union{Tuple{String, Real}, Tuple{String, Real, Integer}, Missing} = missing
    xc::Union{String, Missing} = missing
    n_mpi::Union{Integer, Missing} = missing
    n_threads::Union{Integer, Missing} = missing
end
DFTKCalculatorParameters(ecut::Quantity, functionals, kpts, mixing, nbands, pps, scftol, smearing, xc, n_mpi, n_threads) = DFTKCalculatorParameters(ustrip(u"eV", ecut), functionals, kpts, mixing, nbands, pps, scftol, smearing, xc, n_mpi, n_threads)

function configure_calculator!(atoms::PyObject, parameters::DFTKCalculatorParameters)
    DFTKCalc = pyimport("asedftk").DFTK
    atoms.calc = DFTKCalc(; (v=>getfield(parameters, v) for v in fieldnames(DFTKCalculatorParameters) if getfield(parameters, v) !== missing)...)
end

@with_kw struct LAMMPSCalculatorParameters <: ASECalculatorParameters
    files::Union{Vector{String}, Missing} = missing
    parameters::Dict{String, Union{String, Vector{String}}} = Dict{String, Union{String, Vector{String}}}()
    specorder::Union{Vector{Integer}, Missing} = missing
    keep_tmp_file::Union{Bool, Missing} = missing
    tmp_dir::Union{String, Missing} = missing
    no_data_file::Union{Bool, Missing} = missing
    keep_alive::Union{Bool, Missing} = missing
    always_triclinic::Union{Bool, Missing} = missing
end

function configure_calculator!(atoms::PyObject, parameters::LAMMPSCalculatorParameters)
    LAMPPSCalc = pyimport("ase.calculators.lammpsrun").LAMMPS
    atoms.calc = LAMPPSCalc(; (v=>getfield(parameters, v) for v in fieldnames(LAMMPSCalculatorParameters) if getfield(parameters, v) !== missing)...)
end

struct SNAPCalculatorParameters <: ASECalculatorParameters
    lammps_parameters::LAMMPSCalculatorParameters
    twojmax::Real
    rcutfac::Real
    rfac0::Real
    rmin0::Real
    radelem1::Real
    radelem2::Real
    wj1::Real
    wj2::Real
    quadratic::Real
    bzero::Real
    switch::Real
end

function configure_calculator!(atoms::PyObject, parameters::SNAPCalculatorParameters)
    lammps_parameters = parameters.lammps_parameters
    # TODO
    configure_calculator!(atoms, lammps_parameters)
end


text:
using Test

include("../src/BinTree.jl")

@testset "Testing the BinTree structure" begin

# generate data
N = 1001
X = rand(Float64,N,3)

n = 50

bt = BinTree(X,n)

# this test should succeed whatever the initial distribution as the BinTree splits so that the number of nodes at a given depth remains balanced
@test maxdepth(bt) == 6
@test_throws DomainError BinTree(rand(Float64,N,2),n)

end

max_stars_repo_name: marcbakry/HierarchicalMatrices

max_issues_repo_name: marcbakry/HierarchicalMatrices

max_forks_repo_name: marcbakry/HierarchicalMatrices


text:
# ========================
# Sampling from a LazySet
# ========================

"""
    AbstractSampler

Abstract type for defining new sampling methods.

### Notes

All subtypes should implement a `sample!(D, X, ::Method)` method where the
first argument is the output (vector of vectors), the second argument is the
set to be sampled, and the third argument is the sampler instance.
"""
abstract type AbstractSampler end

"""
    sample(X::LazySet{N}, num_samples::Int;
           [sampler]=_default_sampler(X),
           [rng]::AbstractRNG=GLOBAL_RNG,
           [seed]::Union{Int, Nothing}=nothing,
           [include_vertices]=false,
           [VN]=Vector{N}) where {N}

Random sampling of an arbitrary set `X`.

### Input

- `X`           -- set to be sampled
- `num_samples` -- number of random samples
- `sampler`     -- (optional, default: `_default_sampler(X)`) the sampler used;
                   falls back to `RejectionSampler`
- `rng`         -- (optional, default: `GLOBAL_RNG`) random number generator
- `seed`        -- (optional, default: `nothing`) seed for reseeding
- `include_vertices` -- (optional, default: `false`) option to include the vertices of `X`
- `VN`          -- (optional, default: `Vector{N}`) vector type of the sampled points

### Output

A vector of `num_samples` vectors.
If `num_samples` is not passed, the result is just one sample (not wrapped in a
vector).

### Algorithm

See the documentation of the respective `Sampler`.

### Notes

If `include_vertices == true`, we include all vertices computed with `vertices`.
Alternatively if a number ``k`` is passed, we plot the first ``k`` vertices
returned by `vertices(X)`.
"""
function sample(X::LazySet{N}, num_samples::Int;
                sampler=_default_sampler(X),
                rng::AbstractRNG=GLOBAL_RNG,
                seed::Union{Int, Nothing}=nothing,
                include_vertices=false,
                VN=Vector{N}) where {N}

    D = Vector{VN}(undef, num_samples) # preallocate output
    sample!(D, X, sampler; rng=rng, seed=seed)

    if include_vertices != false
        k = (include_vertices isa Bool) ? Inf : include_vertices
        for v in vertices(X)
            push!(D, v)
            k -= 1
            if k <= 0
                break
            end
        end
    end

    return D
end

# without argument, returns a single element (instead of a singleton)
function sample(X::LazySet; kwargs...)
    return sample(X, 1; kwargs...)[1]
end

# =====================
# Uniform distribution
# =====================

# represents a uniform distribution over the interval [a, b]
# using `rand` from the Julia standard library
struct DefaultUniform{N}
    a::N
    b::N
end

function Base.rand(rng::AbstractRNG, U::DefaultUniform)
    r = rand(rng)
    Δ = U.b - U.a
    return Δ * r + U.a
end

function Base.rand(rng::AbstractRNG, U::DefaultUniform, n::Int)
    return [rand(rng, U) for i in 1:n]
end

function Base.rand(rng::AbstractRNG, U::AbstractVector{<:DefaultUniform})
    return rand.(Ref(rng), U)
end

function rand!(x, rng::AbstractRNG, U::DefaultUniform)
    @inbounds for i in eachindex(x)
        x[i] = rand(rng, U)
    end
    return x
end

# ===================
# Rejection Sampling
# ===================

"""
    RejectionSampler{D} <: Sampler

Type used for rejection sampling of an arbitrary set `X`.

### Fields

- `distribution` -- (optional, default: `DefaultUniform`) distribution from which
                    the sample is drawn
- `tight`        -- (optional, default: `false`) set to `true` if the support of
                    the distribution is known to coincide with the set `X`

### Algorithm

Draw a sample ``x`` from a given distribution of a box-overapproximation of the
original set ``X`` in all ``n`` dimensions. The function rejects a drawn sample
``x`` and redraws as long as the sample is not contained in the original set
``X``, i.e., ``x ∉ X``.
"""
struct RejectionSampler{D} <: AbstractSampler
    distribution::D
    tight::Bool
end

function RejectionSampler(distr; tight::Bool=false)
    return RejectionSampler(distr, tight)
end

function RejectionSampler(distr::DefaultUniform; tight::Bool=false)
    return RejectionSampler([distr], tight)
end

function RejectionSampler(X::LazySet, distribution=DefaultUniform; tight::Bool=false)
    # define the support of the distribution as the smallest box enclosing X
    n = dim(X)
    B = box_approximation(X)

    # distribution over B
    distr = [distribution(low(B, i), high(B, i)) for i in 1:n]

    return RejectionSampler(distr, tight)
end

# the support of this distribution is always tight wrt X
function RejectionSampler(X::AbstractHyperrectangle)
    n = dim(X)
    distr = [DefaultUniform(low(X, i), high(X, i)) for i in 1:n]
    return RejectionSampler(distr, true)
end

# ambiguity fix
function RejectionSampler(X::LazySet, tight::Bool)
    RejectionSampler(X, DefaultUniform; tight=tight)
end
# default sampling for LazySets
_default_sampler(X::LazySet) = RejectionSampler(X)
_default_sampler(X::LineSegment{N}) where {N} = RejectionSampler(DefaultUniform(zero(N), one(N)), true)

function sample!(D::Vector{VN}, X::LazySet, sampler::RejectionSampler;
                 rng::AbstractRNG=GLOBAL_RNG,
                 seed::Union{Int, Nothing}=nothing) where {N, VN<:AbstractVector{N}}

    U = sampler.distribution
    rng = reseed(rng, seed)
    @inbounds for i in 1:length(D)
        w = rand(rng, U)

        if !(sampler.tight)
            while w ∉ X
                w = rand(rng, U)
            end
        end
        D[i] = w
    end
    return D
end

function sample!(D::Vector{VN}, L::LineSegment, sampler::RejectionSampler{<:DefaultUniform};
                 rng::AbstractRNG=GLOBAL_RNG,
                 seed::Union{Int, Nothing}=nothing) where {N, VN<:AbstractVector{N}}

    rng = reseed(rng, seed)
    U = sampler.distribution
    @assert U.a >= zero(N) && U.b <= one(N)
    p = L.p
    q = L.q

    @inbounds for i in 1:length(D)
        λ = rand(rng, U)
        D[i] = p + λ * (q - p)
    end
    return D
end

"""
    RandomWalkSampler{D} <: Sampler

Type used for sampling of a convex polytope `X` using its vertex representation.

### Fields

- `X` -- convex polytope to be sampled

### Algorithm

Choose a random convex combination of the vertices of `X`.

Let ``V = \\{v_i\\}_i`` denote the set of vertices of `X`.
Then any point ``p \\in \\mathbb{R}^n`` of the convex polytope ``X`` is a convex
combination of its vertices, i.e., ``p = \\sum_{i} v_i α_i`` for some
(non-negative) coefficients ``\\{α_i\\}_i`` that add up to 1.
The algorithm chooses a random convex combination (the ``α_i``).
To produce such combination we apply the finite difference operator on a sorted
uniform sample over ``[0, 1]``; the method can be found in [1] and [2].

### Notes

The sampling is not uniform - points in the center of the polytope are more
likely to be sampled.

### References

[1] *Rubin, Donald B. The bayesian bootstrap. The annals of statistics (1981):
130-134.*

[2] https://cs.stackexchange.com/questions/3227/uniform-sampling-from-a-simplex/3229
"""
struct RandomWalkSampler <: AbstractSampler
    #
end

function sample!(D::Vector{VN}, X::LazySet, sampler::RandomWalkSampler;
                 rng::AbstractRNG=GLOBAL_RNG,
                 seed::Union{Int, Nothing}=nothing) where {N, VN<:AbstractVector{N}}
    rng = reseed(rng, seed)
    U = DefaultUniform(zero(N), one(N))
    vlist = vertices_list(X)
    m = length(vlist)

    # vector used to store the combination coefficients
    r = Vector{N}(undef, m-1)

    @inbounds for i in 1:length(D)
        # get a list of m uniform numbers (https://cs.stackexchange.com/a/3229)
        # and compute the corresponding linear combination in-place
        rand!(r, rng, U)
        sort!(r)
        D[i] = r[1] * vlist[1]  # r[1] - 0 == r[1]
        for j in 2:m-1
            α = r[j] - r[j-1]
            D[i] .+= α * vlist[j]
        end
        D[i] .+= (1 - r[m-1]) * vlist[m]
    end

    return D
end

# =============================
# Code requiring Distributions
# =============================

function load_distributions_samples()
return quote

using .Distributions: Uniform, Normal, Distribution, UnivariateDistribution
import .Distributions

RejectionSampler(distr::UnivariateDistribution; tight::Bool=false) = RejectionSampler([distr], tight=tight)

function Base.rand(rng::AbstractRNG, U::AbstractVector{<:UnivariateDistribution})
    return rand.(Ref(rng), U)
end

# ======================================================
# Sampling from a uniform distribution on balls/spheres
# ======================================================

"""
    _sample_unit_nsphere_muller!(D::Vector{Vector{N}}, n::Int, p::Int;
                                 rng::AbstractRNG=GLOBAL_RNG,
                                 seed::Union{Int, Nothing}=nothing) where {N}

Draw samples from a uniform distribution on an ``n``-dimensional unit
sphere using Muller's method.

### Input

- `D`    -- output, vector of points
- `n`    -- dimension of the sphere
- `p`    -- number of random samples
- `rng`  -- (optional, default: `GLOBAL_RNG`) random number generator
- `seed` -- (optional, default: `nothing`) seed for reseeding

### Output

A vector of `nsamples` vectors.

### Algorithm

This function implements Muller's method of normalised Gaussians [1] to uniformly
sample over the ``n``-dimensional sphere ``S^n`` (which is the bounding surface
of the ``n``-dimensional unit ball).

Given ``n`` canonical Gaussian random variables ``Z₁, Z₂, …, Z_n``, the
distribution of the vectors

```math
\\dfrac{1}{α}\\left(z₁, z₂, …, z_n\\right)^T,
```
where ``α := \\sqrt{z₁² + z₂² + … + z_n²}``, is uniform over ``S^n``.

[1] Muller, Mervin E. *A note on a method for generating points uniformly on
    n-dimensional spheres.* Communications of the ACM 2.4 (1959): 19-20.
"""
function _sample_unit_nsphere_muller!(D::Vector{Vector{N}}, n::Int, p::Int;
                                      rng::AbstractRNG=GLOBAL_RNG,
                                      seed::Union{Int, Nothing}=nothing) where {N}
    rng = reseed(rng, seed)
    Zdims = [Normal() for _ in 1:n] # normal distributions, one for each dimension
    v = Vector{N}(undef, n) # sample direction
    @inbounds for j in 1:p
        α = zero(N)
        for i in 1:n
            v[i] = rand(rng, Zdims[i])
            α += v[i]^2
        end
        D[j] = v ./ sqrt(α)
    end
    return D
end

"""
    _sample_unit_nball_muller!(D::Vector{Vector{N}}, n::Int, p::Int;
                               rng::AbstractRNG=GLOBAL_RNG,
                               seed::Union{Int, Nothing}=nothing) where {N}

Draw samples from a uniform distribution on an ``n``-dimensional unit ball
using Muller's method.

### Input

- `D`    -- output, vector of points
- `n`    -- dimension of the ball
- `p`    -- number of random samples
- `rng`  -- (optional, default: `GLOBAL_RNG`) random number generator
- `seed` -- (optional, default: `nothing`) seed for reseeding

### Output

A vector of `nsamples` vectors.

### Algorithm

This function implements Muller's method of normalised Gaussians [1] to
sample from the interior of the ball.

Given ``n`` Gaussian random variables ``Z₁, Z₂, …, Z_n`` and a uniformly
distributed random variable ``r`` with support in ``[0, 1]``, the distribution
of the vectors

```math
\\dfrac{r^{1/n}}{α} \\left(z₁, z₂, …, z_n\\right)^T,
```
where ``α := \\sqrt{z₁² + z₂² + … + z_n²}``, is uniform over the
``n``-dimensional unit ball.

[1] Muller, Mervin E. *A note on a method for generating points uniformly on
    n-dimensional spheres.* Communications of the ACM 2.4 (1959): 19-20.
"""
function _sample_unit_nball_muller!(D::Vector{Vector{N}}, n::Int, p::Int;
                                    rng::AbstractRNG=GLOBAL_RNG,
                                    seed::Union{Int, Nothing}=nothing) where {N}

    rng = reseed(rng, seed)
    Zdims = [Normal() for _ in 1:n] # normal distributions, one for each dimension
    Zrad = Uniform() # distribution to pick random radius
    one_over_n = one(N)/n
    v = Vector{N}(undef, n) # sample direction
    @inbounds for j in 1:p
        α = zero(N)
        for i in 1:n
            v[i] = rand(rng, Zdims[i])
            α += v[i]^2
        end
        r = rand(rng, Zrad)
        β = r^one_over_n / sqrt(α)
        D[j] = v .* β
    end
    return D
end

end # quote
end # function load_distributions_samples()


text:
module Day4

function read_input(filepath)

end

end


text:
@testset "handle" begin
    @jyro function m()
        x ~ Normal()
    end

    handle!(m, TraceHandler())
    @test length(m.handlers_stack) == 1
    @test isa(m.handlers_stack[1], TraceHandler)

    new_m = handle(m, LogJointHandler())
    @test length(m.handlers_stack) == 1
    @test length(new_m.handlers_stack) == 2
    @test isa(new_m.handlers_stack[1], TraceHandler)
    @test isa(new_m.handlers_stack[2], LogJointHandler)
end

@testset "trace" begin
    @jyro function m()
        x ~ Normal()
    end

    handle!(m, TraceHandler())
    t = m()
    @test haskey(t, :msgs)
    @test haskey(t[:msgs], :x)
    @test isa(t[:msgs][:x][:value], Float64)
end

@testset "logjoint" begin
    @jyro function m()
        x ~ Normal()
    end

    trace!(m)
    log_joint!(m)
    t = m()
    val = t[:msgs][:x][:value]

    @test logpdf(Normal(), val) == t[:logjoint]
end

@testset "condition" begin
    @jyro function m()
        x ~ Normal()
    end

    condition!(m, Dict(:x => 0.0))
    log_joint!(m)
    trace!(m)

    t = m()
    val = t[:msgs][:x][:value]

    @test val == 0.0
    @test t[:logjoint] == logpdf(Normal(), 0.0)
    @test t[:msgs][:x][:is_observed]
end

@testset "replay handler" begin
    # First sampling.
    @jyro function model1()
        x ~ Normal()
    end

    handle!(model1, TraceHandler())
    trace1 = model1()

    # Second sampling.
    @jyro function model2()
        x ~ Normal()
        y ~ Normal()
    end

    handle!(model2, ReplayHandler(trace1))
    handle!(model2, TraceHandler())
    trace2 = model2()

    @test trace1[:msgs][:x][:value] == trace2[:msgs][:x][:value]
end

@testset "escape handler" begin
    @jyro function model()
        x ~ Normal()
        {:y => 1} ~ Normal()
        z ~ Normal()
    end

    # Escape when we reach variable y.
    handle!(model, EscapeHandler(x -> x[:name] == (:y => 1)))
    handle!(model, TraceHandler())
    try
        t = model()
    catch e
        if isa(e, EscapeException)
            t = e.trace

            @test e.name == (:y => 1)
            @test haskey(t[:msgs], :x)
            @test isa(t[:msgs][:x][:value], Float64)
            @test t[:msgs][e.name][:value] == nothing
            @test !haskey(t[:msgs], :z)
        else
            throw(e)
        end
    end
end

@testset "queue" begin
    @jyro function model()
        a ~ Bernoulli(0.5)
        b ~ Bernoulli(0.5)
    end

    true_enum = Set([(true,true), (false,true), (true,false), (false,false)])
    empty_trace = Dict()
    empty_trace[:msgs] = Dict()
    q = [empty_trace]

    enum_model = queue(model, q)
    generated_enum = Set()
    while length(q) > 0
        t = enum_model()
        push!(generated_enum, (t[:msgs][:a][:value], t[:msgs][:b][:value]))
    end

    @test true_enum == generated_enum
end

@testset "generated handlers" begin
    @jyro function model()
        a ~ Normal()
    end

    conditioned_model = condition(model, Dict(:a => 1.0))
    trace!(model)
    t = model()
    cond_t = trace(conditioned_model)()

    # Traces from conditioned model should be different.
    @test cond_t[:msgs][:a][:value] == 1.0
    @test cond_t[:msgs][:a][:value] != t[:msgs][:a][:value] # This should technically never happen.
end


text:
function coulumb_to_fock!(F, D, J, b1::BasisSet, b2::BasisSet)

    # Create a merged basis set
    # Note that this basis set has repeated atoms. This is done so that
    # we separate auxiliar basis from regular ones. The first N atoms hold the regular basis functions
    # whereas the final N atoms hold auxiliar fitting functions (for a total of 2N atoms).

    b3 = GaussianBasis.BasisSet(b1, b2)
    Nvals = [GaussianBasis.Libcint.CINTcgtos_spheric(i-1, b3.lc_bas) for i = 1:b3.nshells]
    ao_offset = [sum(Nvals[1:(i-1)]) for i = 1:b3.nshells]
    bufs = [zeros(Cdouble, maximum(Nvals)^3) for _ = 1:Threads.nthreads()]

    # Build intermediate c
    c = zeros(b2.nbas)
    @sync for A in (b1.nshells+1):b3.nshells # A runs through Auxiliar shells only
        Threads.@spawn begin
            buf = bufs[Threads.threadid()]
            Nα = Nvals[A]                  # Number of basis in the A shell
            Aoff = ao_offset[A] - b1.nbas  # Number of basis set before the A shell
            for μ in 1:b1.nshells          # Shell index μ
                Nμ = Nvals[μ]              # Number of basis in the μ shell
                μoff = ao_offset[μ]        # Number of basis set before the μ shell
                for ν in μ:b1.nshells      # Shell index ν
                    Nν = Nvals[ν]          # Number of basis in the ν shell
                    νoff = ao_offset[ν]    # Number of basis set before the ν shell
                    f = μ != ν ? 2.0 : 1.0

                    # Compute integral for the (μ,ν,A) triplet shell
                    GaussianBasis.cint3c2e_sph!(buf, [μ, ν, A], b3)

                    for Am in 1:Nα
                        for νm in 1:Nν
                            _ν = νoff + νm  # Basis set index ν
                            for μm in 1:Nμ
                                _μ = μoff + μm      # Basis set index μ
                                γμν = f*D[_μ,_ν]
                                # C(A) = γ(μν)*(μν|A)
                                c[Aoff+Am] += γμν*buf[μm + Nμ*(νm-1) + Nμ*Nν*(Am-1)]
                            end
                        end
                    end
                end
            end
        end
    end

    # Build intermediate d
    # d(A) = J(AB)*C(B)
    d = J*c

    # Add Coulumb contribution to F
    @sync for μ in 1:b1.nshells          # Shell index μ
        Threads.@spawn begin
            buf = bufs[Threads.threadid()]
            Nμ = Nvals[μ]              # Number of basis in the μ shell
            μoff = ao_offset[μ]        # Number of basis set before the μ shell
            for ν in μ:b1.nshells      # Shell index ν
                    Nν = Nvals[ν]          # Number of basis in the ν shell
                    νoff = ao_offset[ν]    # Number of basis set before the ν shell
                for A in (b1.nshells+1):b3.nshells # A runs through Auxiliar shells only
                    Nα = Nvals[A]                  # Number of basis in the A shell
                    Aoff = ao_offset[A] - b1.nbas  # Number of basis set before the A shell
                    # Compute integral for the (μ,ν,A) triplet shell
                    GaussianBasis.cint3c2e_sph!(buf, Cint.([μ-1, ν-1, A-1]), b3.lc_atoms, b3.natoms, b3.lc_bas, b3.nbas, b3.lc_env)

                    for Am in 1:Nα
                        for νm in 1:Nν
                            _ν = νoff + νm  # Basis set index ν
                            for μm in 1:Nμ
                                _μ = μoff + μm      # Basis set index μ
                                if _μ > _ν
                                    continue
                                end
                                # F += 2J  where J = d(A)*(μν|A)
                                Fμν = 2.0*d[Aoff+Am]*buf[μm + Nμ*(νm-1) + Nμ*Nν*(Am-1)]
                                F[_μ, _ν] += Fμν
                                if _μ != _ν
                                    F[_ν, _μ] += Fμν
                                end
                            end
                        end
                    end
                end
            end
        end
    end
end

function exchange_to_fock!(F, C, J, ndocc, b1::BasisSet, b2::BasisSet)

    # Create a merged basis set
    # Note that this basis set has repeated atoms. This is done so that
    # we separate auxiliar basis from regular ones. The first N atoms hold the regular basis functions
    # whereas the final N atoms hold auxiliar fitting functions (for a total of 2N atoms).

    b3 = GaussianBasis.BasisSet(b1, b2)
    Nvals = [GaussianBasis.Libcint.CINTcgtos_spheric(i-1, b3.lc_bas) for i = 1:b3.nshells]
    ao_offset = [sum(Nvals[1:(i-1)]) for i = 1:b3.nshells]
    bufs = [zeros(Cdouble, maximum(Nvals)^3) for _ = 1:Threads.nthreads()]

    # Build half-transformed (μi|A)
    W = zeros(b1.nbas, ndocc, b2.nbas)
    @sync for A in (b1.nshells+1):b3.nshells # A runs through Auxiliar shells only
        Threads.@spawn begin
            buf = bufs[Threads.threadid()]
            Nα = Nvals[A]                  # Number of basis in the A shell
            Aoff = ao_offset[A] - b1.nbas  # Number of basis set before the A shell
            for μ in 1:b1.nshells          # Shell index μ
                Nμ = Nvals[μ]              # Number of basis in the μ shell
                μoff = ao_offset[μ]        # Number of basis set before the μ shell
                for ν in μ:b1.nshells      # Shell index ν
                    Nν = Nvals[ν]          # Number of basis in the ν shell
                    νoff = ao_offset[ν]    # Number of basis set before the ν shell

                    # Compute integral for the (μ,ν,A) triplet shell
                    GaussianBasis.cint3c2e_sph!(buf, [μ, ν, A], b3)

                    for Am in 1:Nα
                        _A = Aoff + Am
                        for νm in 1:Nν
                            _ν = νoff + νm      # Basis set index ν
                            for μm in 1:Nμ
                                _μ = μoff + μm  # Basis set index μ
                                if _μ > _ν
                                    break
                                end
                                for i = 1:ndocc
                                    W[_μ, i, _A] += C[_ν,i]*buf[μm + Nμ*(νm-1) + Nμ*Nν*(Am-1)]
                                    if _μ != _ν
                                        W[_ν, i, _A] += C[_μ,i]*buf[μm + Nμ*(νm-1) + Nμ*Nν*(Am-1)]
                                    end
                                end
                            end
                        end
                    end
                end
            end
        end #spawn
    end #sync

    @tensoropt F[μ,ν] -=  W[μ,i,A]*J[A,B]*W[ν,i,B]
end

max_stars_repo_name: FermiQC/Fermi

max_issues_repo_name: FermiQC/Fermi

max_forks_repo_name: FermiQC/Fermi


text:
mutable struct Function

    alloc::Bool
    ptr::Ptr{Void}
    net::Network

    function Function(a, ptr, n)
        this = new(a, ptr, n)
        finalizer(this, dealloc)
        this
    end
end

function Function(name::String, weight::Float64, net::Network)

    if name == "generation cost"
        Function(true, ccall((:FUNC_GEN_COST_new, libpfnet), Ptr{Void}, (Float64, Ptr{Void},), weight, net.ptr), net)

    elseif name == "consumption utility"
        Function(true, ccall((:FUNC_LOAD_util_new, libpfnet), Ptr{Void}, (Float64, Ptr{Void},), weight, net.ptr), net)

    elseif name == "net consumption cost"
        Function(true, ccall((:FUNC_NETCON_COST_new, libpfnet), Ptr{Void}, (Float64, Ptr{Void},), weight, net.ptr), net)

    elseif name == "phase shift regularization"
        Function(true, ccall((:FUNC_REG_PHASE_new, libpfnet), Ptr{Void}, (Float64, Ptr{Void},), weight, net.ptr), net)

    elseif name == "generator powers regularization"
        Function(true, ccall((:FUNC_REG_PQ_new, libpfnet), Ptr{Void}, (Float64, Ptr{Void},), weight, net.ptr), net)

    elseif name == "tap ratio regularization"
        Function(true, ccall((:FUNC_REG_RATIO_new, libpfnet), Ptr{Void}, (Float64, Ptr{Void},), weight, net.ptr), net)

    elseif name == "susceptance regularization"
        Function(true, ccall((:FUNC_REG_SUSC_new, libpfnet), Ptr{Void}, (Float64, Ptr{Void},), weight, net.ptr), net)

    elseif name == "voltage angle regularization"
        Function(true, ccall((:FUNC_REG_VANG_new, libpfnet), Ptr{Void}, (Float64, Ptr{Void},), weight, net.ptr), net)

    elseif name == "voltage magnitude regularization"
        Function(true, ccall((:FUNC_REG_VMAG_new, libpfnet), Ptr{Void}, (Float64, Ptr{Void},), weight, net.ptr), net)

    elseif name == "variable regularization"
        Function(true, ccall((:FUNC_REG_VAR_new, libpfnet), Ptr{Void}, (Float64, Ptr{Void},), weight, net.ptr), net)

    elseif name == "soft voltage magnitude limts"
        Function(true, ccall((:FUNC_SLIM_VMAG_new, libpfnet), Ptr{Void}, (Float64, Ptr{Void},), weight, net.ptr), net)

    elseif name == "sparse control penalty"
        Function(true, ccall((:FUNC_SP_CONTROLs_new, libpfnet), Ptr{Void}, (Float64, Ptr{Void},), weight, net.ptr), net)

    else
        throw(ArgumentError("invalid function name"))
    end
end

function dealloc(func::Function)
    if func.alloc
        ccall((:FUNC_del, libpfnet), Void, (Ptr{Void},), func.ptr)
    end
    func.alloc = false
    func.ptr = C_NULL
end


text:
using ProgressiveHedgingSolvers
using Test
using LinearAlgebra
using JuMP
using StochasticPrograms
using GLPKMathProgInterface
using OSQP

τ = 1e-2
reference_solver = GLPKSolverLP()
osqp = OSQP.OSQPMathProgBaseInterface.OSQPSolver(verbose=0)

penalties = [Fixed(),
             Adaptive(θ = 1.01)]

problems = Vector{Tuple{<:StochasticProgram,String}}()
@info "Loading test problems..."
@info "Loading simple..."
include("simple.jl")
@info "Loading farmer..."
include("farmer.jl")
@info "Loading infeasible..."
include("infeasible.jl")

@testset "Sequential solver" begin
    @testset "$(solverstr(ph)): $name" for ph in [ProgressiveHedgingSolver(osqp,
                                                                           penalty = penalty,
                                                                           τ = 1e-3,
                                                                           log = false)
                                                  for penalty in penalties], (sp,name) in problems
        optimize!(sp, solver=reference_solver)
        x̄ = optimal_decision(sp)
        Q̄ = optimal_value(sp)
        optimize!(sp, solver=ph)
        @test abs(optimal_value(sp) - Q̄)/(1e-10+abs(Q̄)) <= τ
        @test norm(optimal_decision(sp) - x̄)/(1e-10+norm(x̄)) <= sqrt(τ)
    end
end

@info "Starting distributed tests..."

include("/usr/share/julia/test/testenv.jl")
push!(test_exeflags.exec,"--color=yes")
cmd = `$test_exename $test_exeflags run_dtests.jl`

if !success(pipeline(cmd; stdout=stdout, stderr=stderr)) && ccall(:jl_running_on_valgrind,Cint,()) == 0
    error("Distributed test failed, cmd : $cmd")
end


text:
module Regression

using Parameters
import LsqFit
import Dierckx
const lq = LsqFit
const dk = Dierckx

include("Utils.jl")
include("Opt.jl")
using .Opt
using .Utils

export curve_fit, is_monotonous

const DEF_RESOLUTION = 0.01

#=
function curve_fit(model, xdata, ydata, p0, lb, ub)
    fit = lq.curve_fit(model, xdata, ydata, p0, lower=lb, upper=ub)

    if !fit.converged
        @warn "The fitting is not converged with residual(=$(param.resid[end]))"
    end

    param = fit.param

    return param

end

function curve_fit(model, xdata, ydata, p0)
    fit = lq.curve_fit(model, xdata, ydata, p0)

    if !fit.converged
        @warn "The fitting is not converged with residual(=$(param.resid[end]))"
    end

    param = fit.param

    return param

end
=#

"""
Fit model with given x,y data.

# Argument
- model::`function`: function with input and parameters. model(x,p)
- xdata::`Array`:input data
- ydata:: `Array`: output data
- p0::'tuple{1,n}': n = number of inputs
- lb::`Array{1,n}`: lower bound of n inputs
- ub::`Array{1,n}`: upper bound of n inputs
- method::`Symbol`: Method of LsqFit. it can be `:forwarddiff` or `:finiteforward`
"""
function curve_fit(model_functor, xdata, ydata, p0; method=:forwarddiff)
    model_func =Utils.convert_functor_model(model_functor)

    fit =  lq.curve_fit(model_func, xdata, ydata, p0;autodiff=method)

    if !fit.converged
        @warn "The fitting is not converged"
    end

    return fit
end

function curve_fit(model_functor, xdata, ydata, p0, lb, ub; method=:forwarddiff)
    model_func = Utils.convert_functor_model(model_functor)
    fit = lq.curve_fit(model_func, xdata, ydata, p0, lower=lb, upper=ub; autodiff=method)

    if !fit.converged
        @warn "The fitting is not converged"
    end

    return fit
end

"""
Check the monotony of a series.
"""
function is_monotonous(series)
    is_mono_increase = all(series == accumulate(max,series) )
    is_mono_decrease = all(series == accumulate(min,series) )
    is_mono = is_mono_increase | is_mono_decrease
    return is_mono
end

"""
Check the mononous of a function with fixed domain.

# Argument
- f::`function`: the function to be tested for monotony.
- domain::`array of 2 numbers`: domain of function.
- resolution::`float`: resolution of screening
- reset_step:: `float`: reset resolution if resolution is larger than the domain range
"""
function is_monotonous(f, domain; resolution=0.01, reset_step=100.0)

    data = Utils.screen_domain(f, domain; resolution=0.01, reset_step=100.0)

    is_mono = is_monotonous(data.y)

    return is_mono
end

"""
Measure the Effective concentration with given ratio

# Arguement
- arrayX{Vector}: should be monotonous
- arrayY{Vector}: should be monotonous
- EC_ratio{Number}: The effective concentration
"""
function measure_EC( arrayX , arrayY , EC_ratio)

    if !is_monotonous(arrayX) | !is_monotonous(arrayY)
        @warn "The given curve is not monotonous"
    end

    if (EC_ratio < 0.0 )| (EC_ratio >= 1.0)
        @warn "Invalid EC_ratio($EC_ratio). EC_ratio ∈ [0,1)"
    end

    min_v, max_v = minimum(arrayY), maximum(arrayY)

    threshold_v = (max_v-min_v) * EC_ratio

    spl = dk.Spline1D(arrayX, arrayY .- threshold_v .- min_v)

    threshold_inputs = dk.roots(spl)

    return threshold_inputs
end

"""Measure Effective concentration with function input and screening resolution."""
function measure_EC(f, domain, EC_ratio, resolution)
    data = Utils.screen_domain(f, domain; resolution=resolution)

    threshold_inputs = measure_EC(data.x, data.y, EC_ratio)

    return threshold_inputs
end

end


text:
using Test
using QuantumLattices.Prerequisites

@testset "decimaltostr" begin
    @test decimaltostr(:a) == ":a"
    @test decimaltostr(1) == "1"
    @test decimaltostr(10^6) == "1000000"
    @test decimaltostr(1//7) == "1//7"
    @test decimaltostr(1.0) == "1.0"
    @test decimaltostr(10.0^6) == "1.0e+06"
    @test decimaltostr(10^-5) == "1.0e-05"
    @test decimaltostr(1/7) == "0.14286"
    @test decimaltostr(1/7, 8) == "0.14285714"
    @test decimaltostr(0im) == "0.0"
    @test decimaltostr(1//7+0im) == "1//7"
    @test decimaltostr(1.0im) == "1.0im"
    @test decimaltostr(-1.0im) == "-1.0im"
    @test decimaltostr(0.1+0.12im) == "0.1+0.12im"
    @test decimaltostr(0.1-0.12im) == "0.1-0.12im"
end

@testset "ordinal" begin
    @test ordinal(1) == "1st"
    @test ordinal(2) == "2nd"
    @test ordinal(3) == "3rd"
    @test ordinal(4) == "4th"
    @test ordinal(5) == "5th"
end

@testset "delta" begin
    @test delta(1, 2) == 0
    @test delta(1, 1) == 1
end

@testset "rawtype" begin
    @test rawtype(Int) == Int
    @test rawtype(Vector{Int}) == Array
    @test rawtype(Vector) == Array
    @test rawtype(Vector{<:Real}) == Array
end


text:
using CVChannel
using LinearAlgebra
using QBase
using Test

@testset "./src/channels.jl" begin

maxEntState = 0.5 * [1 0 0 1 ; 0 0 0 0 ; 0 0 0 0 ; 1 0 0 1]
maxMixState = 0.25 * [1 0 0 0 ; 0 1 0 0 ; 0 0 1 0 ; 0 0 0 1]

@testset "choi" begin
    depolChan(X) = 1/2*[1 0 ; 0 1]
    identChan(X) = X
    @test isapprox(choi(identChan,2,2),2*maxEntState, atol = 1e-6)
    @test isapprox(choi(depolChan,2,2),1/2*[1 0 1 0 ; 0 1 0 1 ; 1 0 1 0 ; 0 1 0 1], atol = 1e-6)

    @testset "kraus operator inputs" begin
        id_kraus_ops = [[1 0;0 1]]
        @test choi(id_kraus_ops) ≈ 2 * maxEntState

        depolarizing_kraus_ops = [σI, σx, σy, σz]/2
        @test sum(k -> k' * k, depolarizing_kraus_ops) ≈ I
        @test choi(depolarizing_kraus_ops) ≈ I/2
    end
end

@testset "is_choi_matrix" begin
    @test is_choi_matrix(Matrix(I, 6, 6), 2, 3)
    @test !is_choi_matrix(Matrix(I, 4, 4), 2, 3)
end

@testset "Choi" begin
    @testset "function instantiation" begin
        depolChan(X) = 1/2*[1 0; 0 1]
        depolChoi = Choi(depolChan, 2, 2)

        @test depolChoi isa Choi{ComplexF64}
        @test depolChoi.JN isa Matrix{ComplexF64}
        @test depolChoi.in_dim == 2
        @test depolChoi.out_dim == 2
    end

    @testset "matrix instantiation" begin
        JN = [1 0 0 1;0 0 0 0;0 0 0 0;1 0 0 1]
        choi_channel = Choi(JN, 2, 2)

        @test choi_channel isa Choi{Int}
        @test choi_channel.JN isa Matrix{Int}
        @test choi_channel.JN == JN
        @test choi_channel.in_dim == 2
        @test choi_channel.out_dim == 2
    end

    @testset "kraus operator instantiation" begin
        kraus_ops = [σI, σx, σy, σz]/2
        choi_channel = Choi(kraus_ops)

        @test choi_channel isa Choi{ComplexF64}
        @test choi_channel.JN isa Matrix{ComplexF64}
        @test choi_channel.JN == I/2
        @test choi_channel.in_dim == 2
        @test choi_channel.out_dim == 2
    end

    @testset "DomainError" begin
        @test_throws DomainError Choi(Matrix(I, 4, 4), 2, 3)
    end
end

@testset "parChoi" begin
    @testset "subsystem system swap" begin
        chan = Choi(2*maxEntState, 2, 2)
        par_chan = parChoi(chan, chan)

        @test par_chan isa Choi{Float64}
        @test par_chan.in_dim == 4
        @test par_chan.out_dim == 4

        @test par_chan.JN == [
            1  0  0  0  0  1  0  0  0  0  1  0  0  0  0  1
            0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0
            0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0
            0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0
            0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0
            1  0  0  0  0  1  0  0  0  0  1  0  0  0  0  1
            0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0
            0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0
            0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0
            0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0
            1  0  0  0  0  1  0  0  0  0  1  0  0  0  0  1
            0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0
            0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0
            0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0
            0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0
            1  0  0  0  0  1  0  0  0  0  1  0  0  0  0  1
        ]
    end

    @testset "subsystem dims" begin
        dephrasure_chan(ρ) = dephrasureChannel(ρ, 0 ,0)
        identity_chan(ρ) = ρ

        deph_chan = Choi(dephrasure_chan, 2,3)
        id_chan = Choi(identity_chan, 2, 2)

        par_choi = parChoi(deph_chan, id_chan)

        @test par_choi isa Choi{Complex{Float64}}
        @test par_choi.in_dim == 4
        @test par_choi.out_dim == 6
        @test size(par_choi.JN) == (24,24)
    end
end

@testset "isometricChannel" begin
    kraus_set1 = Vector{Matrix{Complex{Float64}}}(undef, 2)
    kraus_set1[1] = [1 0 ; 0 0 ; 0 0]
    kraus_set1[2] = 1/sqrt(2)*[0 0 ; 0 1 ; 0 1]
    @test isapprox(isometricChannel(kraus_set1),[1 0 ; 0 0 ; 0 0 ; 0 1/sqrt(2) ; 0 0 ; 0 1/sqrt(2)], atol=1e-6)

    kraus_set2 = Vector{Matrix{Complex{Float64}}}(undef, 1)
    kraus_set2[1] = [1 0 ; 0 -1]
    @test isapprox(isometricChannel(kraus_set2),[1 0 ; 0 -1], atol=1e-6)
end

@testset "complementaryChannel" begin
    @testset "First Map" begin
        kraus_set1 = Vector{Matrix{Complex{Float64}}}(undef, 2)
        kraus_set1[1] = [1. 0 ; 0 0 ; 0 0]
        kraus_set1[2] = 1/sqrt(2)*[0 0 ; 0 1 ; 0 1]
        comp_kraus = complementaryChannel(kraus_set1)
        @test comp_kraus[1] == [1. 0 ; 0 0]
        @test comp_kraus[2] == 1/sqrt(2)*[0 0; 0 1]
        @test comp_kraus[3] == 1/sqrt(2)*[0 0; 0 1]
    end
    @testset "Unitary Map" begin
        kraus_set2 = Vector{Matrix{Complex{Float64}}}(undef, 1)
        kraus_set2[1] = [1 0 ; 0 -1]
        comp_kraus = complementaryChannel(kraus_set2)
        @test comp_kraus[1] == [1 0 ; 0 0]
        @test comp_kraus[2] == [0 1 ; 0 0]
    end
end

@testset "krausAction" begin
    #By linearity this should suffice for checking it
    #works on multiple dimensions
    ρ1 = [1 0 ; 0 0]
    ρ2 = [0 0 ; 0 1]
    ρ3 = [1 0 0 ; 0 0 0 ; 0 0 0]
    ρ4 = [0 0 0 ; 0 1 0 ; 0 0 0]
    ρ5 = [0 0 0 ; 0 0 0 ; 0 0 1]
    @testset "Identity Channel" begin
        kraus_set1 = Vector{Matrix{Complex{Float64}}}(undef, 1)
        kraus_set1[1] = [1 0 ; 0 1]
        @test krausAction(kraus_set1,ρ1) == ρ1
        @test krausAction(kraus_set1,ρ2) == ρ2
        kraus_set1 = Vector{Matrix{Complex{Float64}}}(undef, 1)
        kraus_set1[1] = [1 0 0 ; 0 1 0 ; 0 0 1]
        @test krausAction(kraus_set1,ρ3) == ρ3
        @test krausAction(kraus_set1,ρ4) == ρ4
        @test krausAction(kraus_set1,ρ5) == ρ5
    end
    @testset "Completely Depolarizing Channel" begin
        #This uses that applying the Discrete-Weyl operators uniformly
        #to a qudit turns it into the maximally mixed state
        kraus_set1 = Vector{Matrix{Complex{Float64}}}(undef, 4)
        kraus_set1[1] = 1/sqrt(4)*discreteWeylOperator(0,0,2)
        kraus_set1[2] = 1/sqrt(4)*discreteWeylOperator(0,1,2)
        kraus_set1[3] = 1/sqrt(4)*discreteWeylOperator(1,0,2)
        kraus_set1[4] = 1/sqrt(4)*discreteWeylOperator(1,1,2)
        @test krausAction(kraus_set1,ρ1) == 1/2*[1 0 ; 0 1]
        @test krausAction(kraus_set1,ρ2) == 1/2*[1 0 ; 0 1]
        kraus_set1 = Vector{Matrix{Complex{Float64}}}(undef, 9)
        kraus_ctr = 1
        for i = 0:2;
            for j = 0:2;
                kraus_set1[kraus_ctr] = 1/sqrt(9)*discreteWeylOperator(i,j,3)
                kraus_ctr += 1
            end
        end
        @test isapprox(krausAction(kraus_set1,ρ3),1/3*[1 0 0; 0 1 0 ; 0 0 1], atol = 1e-6)
        @test isapprox(krausAction(kraus_set1,ρ4),1/3*[1 0 0; 0 1 0 ; 0 0 1], atol = 1e-6)
        @test isapprox(krausAction(kraus_set1,ρ5),1/3*[1 0 0; 0 1 0 ; 0 0 1], atol = 1e-6)
    end
end

@testset "depolarizingChannel" begin
    @test isapprox(depolarizingChannel(maxEntState,0),maxEntState, atol=1e-6)
    @test isapprox(depolarizingChannel(maxEntState,1),maxMixState, atol=1e-6)

    @testset "errors" begin
        @test_throws DomainError depolarizingChannel(maxEntState,1.1)
        @test_throws DomainError depolarizingChannel(maxEntState,-.1)
        @test_throws DomainError depolarizingChannel([1 0 0;0 0 0],0.5)
    end
end

@testset "dephrasureChannel" begin
    checkQubit = [2/3 1/4*1im ; -1/4*1im 1/3]
    @test isapprox(dephrasureChannel(checkQubit,0,0),[2/3 1/4*1im 0; -1/4*1im 1/3 0; 0 0 0], atol = 1e-6)
    @test isapprox(dephrasureChannel(checkQubit,1/2,0),[2/3 0 0; 0 1/3 0; 0 0 0], atol = 1e-6)
    @test isapprox(dephrasureChannel(checkQubit,0,1/2),[2/6 1/8*1im 0; -1/8*1im 1/6 0; 0 0 1/2], atol = 1e-6)
    @test isapprox(dephrasureChannel(checkQubit,1/2,1/2),[2/6 0 0; 0 1/6 0; 0 0 1/2], atol = 1e-6)

    @testset "errors" begin
        @test_throws DomainError dephrasureChannel([1 0 0;0 0 0;0 0 0], 1, 1)
        @test_throws DomainError dephrasureChannel([1 0;0 0], 1.1, 1)
        @test_throws DomainError dephrasureChannel([1 0;0 0], -.1, 1)
        @test_throws DomainError dephrasureChannel([1 0;0 0], 1, 1.1)
        @test_throws DomainError dephrasureChannel([1 0;0 0], 1, -.1)
    end
end

@testset "wernerHolevoChannel" begin
    checkQubit = [2/3 1/4*1im ; -1/4*1im 1/3]
    @test isapprox(wernerHolevoChannel(checkQubit,1), (1/3*[5/3 -1/4*1im ; 1/4*1im 4/3]), atol = 1e-6)
    @test isapprox(wernerHolevoChannel(checkQubit,0), [1/3 1/4*1im ; -1/4*1im 2/3], atol = 1e-6)
    @test isapprox(wernerHolevoChannel(checkQubit,1/2), [4/9 1/12*1im ; -1/12*1im 5/9], atol = 1e-6)

    @testset "errors" begin
        @test_throws DomainError wernerHolevoChannel([1 0;0 0;0 0], 1)
        @test_throws DomainError wernerHolevoChannel([1 0;0 0], 1.1)
        @test_throws DomainError wernerHolevoChannel([1 0;0 0],-.1)
    end
end

@testset "siddhuChannel" begin
    @testset "verify channel definition" begin
        for s in [0:0.1:0.5;]
            sidchan(X) = siddhuChannel(X,s)
            testchan = Choi(sidchan,3,3)
            α = 1-s
            γ = sqrt(s)
            β = sqrt(1-s)
            @test isapprox(testchan.JN,
                [s 0 0 0 0 γ 0 0 0;
                0 α 0 0 0 0 0 0 β;
                0 0 0 0 0 0 0 0 0 ;
                0 0 0 0 0 0 0 0 0 ;
                0 0 0 0 0 0 0 0 0 ;
                γ 0 0 0 0 1 0 0 0 ;
                0 0 0 0 0 0 0 0 0 ;
                0 0 0 0 0 0 0 0 0 ;
                0 β 0 0 0 0 0 0 1
                ],
                atol = 1e-6
                )
        end
    end

    @testset "errors" begin
        @test_throws DomainError siddhuChannel([1 0 ; 0 0; 0 0], 0.2)
        @test_throws DomainError siddhuChannel([1 0 ; 0 0], 0.3)
        @test_throws DomainError siddhuChannel([1 0 0 ; 0 0 0 ; 0 0 0], 7)
    end
end

@testset "generalizedSiddhu" begin
    @testset "verify channel definition" begin
        for s in [0:0.1:0.5;]
            for μ in [0:0.1:1;]
                genSidChan(X) = generalizedSiddhu(X,s,μ)
                testchan = Choi(genSidChan,3,3)
                α, β, γ, δ = 1-s, 1-μ, sqrt(s*(1-μ)), sqrt(s*μ)
                ϵ, ζ, η = sqrt(μ*(1-s)), sqrt((1-μ)*(1-s)), sqrt(μ*(1-μ))
                target = [s 0 0 0 γ 0 0 0 δ ;
                          0 α 0 0 0 ϵ ζ 0 0;
                          0 0 0 0 0 0 0 0 0 ;
                          0 0 0 0 0 0 0 0 0 ;
                          γ 0 0 0 β 0 0 0 η ;
                          0 ϵ 0 0 0 μ η 0 0 ;
                          0 ζ 0 0 0 η β 0 0 ;
                          0 0 0 0 0 0 0 0 0 ;
                          δ 0 0 0 η 0 0 0 μ]
                @test isapprox(testchan.JN, target, atol = 1e-6)
            end
        end
    end
    @testset "errors" begin
        @test_throws DomainError generalizedSiddhu([1 0 ; 0 0; 0 0], 0.2,1)
        @test_throws DomainError generalizedSiddhu([1 0 ; 0 0], 0.3,1)
        @test_throws DomainError generalizedSiddhu([1 0 0 ; 0 0 0 ; 0 0 0], 7,1)
        @test_throws DomainError generalizedSiddhu([1 0 0 ; 0 0 0 ; 0 0 0], 0.2,7)
    end
end

@testset "GADChannel" begin
    @testset "Verify Channel Definition" begin
        #By linearity we just need to check the computational basis
        zero_state, one_state = [1 0 ; 0 0], [0 0 ; 0 1]
        scan_range = [0:0.1:1;]
        for p in scan_range
            for n in scan_range
                @test isapprox(GADChannel(zero_state,p,n),[1-p*n 0 ; 0 p*n], atol = 1e-6)
                @test isapprox(GADChannel(one_state,p,n),[(1-n)*p 0 ; 0 1-p+p*n], atol = 1e-6)
            end
        end
    end

    @testset "errors" begin
        @test_throws DomainError GADChannel([1 0 ; 0 0; 0 0], 0.2,0.4)
        @test_throws DomainError GADChannel([1 0 0 ; 0 0 0; 0 0 0], 0.4, 0.2)
        @test_throws DomainError GADChannel([1 0 0 ; 0 0 0 ; 0 0 0], 0.3,1.2)
        @test_throws DomainError GADChannel([1 0 0 ; 0 0 0 ; 0 0 0], 1.5,0.1)
    end
end
end

max_forks_repo_name: ChitambarLab/cv-channel


text:
using SnoopCompile, Random
ENV["FILESYSTEMDATASTRUCTURES_GENERATING_PRECOMPILE"] = "true"

inf_timing = @snoopi tmin=0.001 begin
    using FilesystemDatastructures

    mktempdir() do root
        for target_size in (TargetSizeKeepFree(10*1024), TargetSizeConstant(10*1024)),
            discard_policy in (DiscardLRU(), DiscardLFU())

            scfc = SizeConstrainedFileCache(root, target_size, discard_policy)

            for idx in 1:100
                key = randstring(12)
                filesize = rand(512:1024)
                filedata = rand(UInt8, filesize)
                path = add!(scfc, key, filesize)
                open(path, "w") do io
                    write(path, filedata)
                end

                for k in keys(scfc.entries)
                    if rand() > 0.8
                        hit!(scfc, k)
                    end
                    if rand() > 0.98
                        delete!(scfc, k)
                    end
                end
            end
        end
    end
end

pc = SnoopCompile.parcel(inf_timing)
mktempdir() do dir
    SnoopCompile.write(dir, pc)
    cp(joinpath(dir, "precompile_FilesystemDatastructures.jl"),
       joinpath(@__DIR__, "precompile.jl"); force=true)
end


text:
function get_mesh_2d(i)
  coordinates_min = (-1,-1)
  coordinates_max = ( 1, 1)
  mesh1_2d = TreeMesh(coordinates_min, coordinates_max,
                  initial_refinement_level=5,
                  n_cells_max=10_000
                  )

  coordinates_min = (-1,-1)
  coordinates_max = ( 1, 1)
  mesh2_2d = TreeMesh(coordinates_min, coordinates_max,
                  initial_refinement_level=6,
                  n_cells_max=10_000
                  )
  coordinates_min = (-1,-1)
  coordinates_max = ( 1, 1)
  refinement_patches = (
    (type="box", coordinates_min=(-1.0,-1.0), coordinates_max=(-0.5,-0.5)),
    (type="box", coordinates_min=(-0.5,-0.5), coordinates_max=(0.0,0.0)),
    (type="box", coordinates_min=(0.0,0.0), coordinates_max=(0.5,0.5)),
    (type="box", coordinates_min=(0.5,0.5), coordinates_max=(1.0,1.0)),
  )
  mesh3_2d = TreeMesh(coordinates_min, coordinates_max,
                  initial_refinement_level=4,
                  refinement_patches=refinement_patches,
                  n_cells_max=10_000
                  )

  coordinates_min = (-1,-1)
  coordinates_max = ( 1, 1)
  refinement_patches = (
    (type="box", coordinates_min=(-1.0,-1.0), coordinates_max=(-0.5,-0.5)),
    (type="box", coordinates_min=(-0.5,-0.5), coordinates_max=(0.0,0.0)),
    (type="box", coordinates_min=(0.0,0.0), coordinates_max=(0.5,0.5)),
    (type="box", coordinates_min=(0.5,0.5), coordinates_max=(1.0,1.0)),
    (type = "box", coordinates_min = (-1.0, -1.0), coordinates_max = (-0.75, -0.75)),
    (type = "box", coordinates_min = (-0.75, -0.75), coordinates_max = (-0.5, -0.5)),
    (type = "box", coordinates_min = (-0.5, -0.5), coordinates_max = (-0.25, -0.25)),
    (type = "box", coordinates_min = (-0.25, -0.25), coordinates_max = (0.0, 0.0)),
    (type = "box", coordinates_min = (0.0, 0.0), coordinates_max = (0.25, 0.25)),
    (type = "box", coordinates_min = (0.25, 0.25), coordinates_max = (0.5, 0.5)),
    (type = "box", coordinates_min = (0.5, 0.5), coordinates_max = (0.75, 0.75)),
    (type = "box", coordinates_min = (0.75, 0.75), coordinates_max = (1.0, 1.0)),
  )
  mesh4_2d = TreeMesh(coordinates_min, coordinates_max,
                  initial_refinement_level=5,
                  refinement_patches=refinement_patches,
                  n_cells_max=10_000
                  )

  coordinates_min = (-2,-2)
  coordinates_max = ( 2, 2)
  refinement_patches = (
    (type="box", coordinates_min=(-1.3,-1.3), coordinates_max=(-0.8,-0.8)),
    (type="box", coordinates_min=(-1.1,-1.1), coordinates_max=(-0.9,-0.9)),
    (type="box", coordinates_min=(0.8,0.8), coordinates_max=(1.2,1.2)),
    (type="box", coordinates_min=(1.2,1.2), coordinates_max=(1.8,1.8)),
  )
  mesh5_2d = TreeMesh(coordinates_min, coordinates_max,
                  initial_refinement_level=5,
                  refinement_patches=refinement_patches,
                  n_cells_max=10_000
                  )

  if i==1
    mesh = mesh1_2d
  elseif i==2
    mesh = mesh2_2d
  elseif i==3
    mesh = mesh3_2d
  elseif i==4
    mesh = mesh4_2d
  elseif i==5
    mesh = mesh5_2d
  end
  return mesh
end


text:
using CSV, DataFrames, Dates#, Plots
include("src/Hymod.jl")

obsPath = "../example_data/LA011201_obs.csv"
obs = CSV.read(obsPath)
# obs.isodate = Hymod.Utils.parseDates(obs, format="m/d/Y",dateCol=:date)

forcingPath = "../example_data/LA011201_forcings.csv"
forcings = CSV.read(forcingPath)

forcings.pet = Hymod.hargreaves(forcings,tminCol=:tmin,tmaxCol=:tmax,dtCol=:isodate)

paramSpace = Dict{Symbol,Dict}(
    :cmax => Dict{Symbol,Float64}(:lower => 1.0, :upper => 100),
    :bexp => Dict{Symbol,Float64}(:lower => 0.0, :upper => 2.0),
    :alpha => Dict{Symbol,Float64}(:lower => 0.2, :upper => 0.99),
    :ks => Dict{Symbol,Float64}(:lower => 0.01, :upper => 0.5),
    :kq => Dict{Symbol,Float64}(:lower => 0.5, :upper => 1.2)
)

calStart = Date(1986,1,1)
calEnd = Date(1995,12,31)

calForcings = filter(row -> row[:isodate] >= calStart && row[:isodate] <= calEnd, forcings)
obsSel = filter(row -> row[:datetime] >= calStart && row[:datetime] <= calEnd, obs)

nIters = 5

calQ, calPars, loss = Hymod.calibrate(calForcings,obsSel.discharge,paramSpace,nIters)
print(calPars,loss)

# # do some plotting to show the results
# obsSel[:calibrated] = calQ
# obsSel = filter(row -> row[:datetime] >= Date(1987,1,1), obsSel)

# theme(:bright)

# plot(obsSel[:datetime],[obsSel[:discharge] obsSel[:calibrated]],
#     label=["Observed" "Calibrated"],
#     xlabel="Date",
#     xrotation=40,
#     ylabel="Discharge m³/s",
#     dpi=200
# )
# savefig("calibrated_discharge.png")

max_stars_repo_name: jsta/hymod

max_issues_repo_name: jsta/hymod

max_forks_repo_name: jsta/hymod


text:
# Script to generate figures in Figure 2.
using Plots, SparseArrays, JLD2, LinearAlgebra, MultiscaleGraphSignalTransforms

# First define a utility function "Glevel", which generate a Graph Signal that
# are piecewise constants whose values are actual region indices "k".
function Glevel(G::GraphSig, GP::GraphPart, j::Int64)
    f = zeros(size(G.f))
    for k in 1:size(GP.rs,1)
        a = GP.rs[k,j+1]
        b = GP.rs[k+1,j+1] - 1
        if b == -1
            break
        end
        f[a:b] .= k*1.0
    end
    Gsub = deepcopy(G)
    Gsub.f[GP.ind] = f
    return Gsub
end

# Set up the resolution and display size
gr(dpi=200, size=(800,600))

# Load the Toronto street network (vehicular volume counts as its graph signal)
JLD2.@load "Toronto_new.jld2"
tmp1 = toronto["G"]
G=GraphSig(tmp1["W"],xy=tmp1["xy"],f=tmp1["f"],name =tmp1["name"],
    plotspecs = tmp1["plotspecs"])

# Assign correct edge weights via 1/Euclidean distance
G = Adj2InvEuc(G)
# Perform the hierarchical graph partitioning using the Fiedler vectors
GP = partition_tree_fiedler(G) # Lrw by default
# Compute the expansion coefficients of the full GHWT c2f dictionary
dmatrix = ghwt_analysis!(G, GP=GP)
N = length(G.f)

# Fig. 2a: Level 1 first since Level 0 is not interesting
j = 1
GraphSig_Plot(Glevel(G,GP,j), linewidth = 1., markersize = 4.,
    markercolor = :viridis, markerstrokealpha =0., notitle = true, nocolorbar = true)
plot!(showaxis = false, ticks = false)
savefig("toronto_j1.pdf")
savefig("toronto_j1.png")

# Fig. 2b: Level 2
j = 2
GraphSig_Plot(Glevel(G,GP,j), linewidth = 1., markersize = 4.,
    markercolor = :viridis, markerstrokealpha =0., notitle = true, nocolorbar = true)
plot!(showaxis = false, ticks = false)
savefig("toronto_j2.pdf")
savefig("toronto_j2.png")

# Fig. 2c: Level 3
j = 3
GraphSig_Plot(Glevel(G,GP,j), linewidth = 1., markersize = 4.,
    markercolor = :viridis, markerstrokealpha =0., notitle = true, nocolorbar = true)
plot!(showaxis = false, ticks = false)
savefig("toronto_j3.pdf")
savefig("toronto_j3.png")


text:
struct ROPF_infos
    instance_name::String
    matpower_instance_path::String
    output_instance_path::String
    decomposition::String
    output_decomposition_path::String
end

struct BB_infos
    search_strategy::String
    branch_strategy::String
    seuil_u::Float64
    seuil_l::Float64
end

abstract type AbstractROPFConstraint end

abstract type NoConstraint <: AbstractROPFConstraint end
abstract type MAXkshuntsConstraint <:AbstractROPFConstraint end
abstract type MAXkmovesConstraint <: AbstractROPFConstraint end

include("solve_SDP.jl")
include("solve_minlp.jl")
include("B&B_fixingsome1and0.jl")

function optvalue_bounds(ROPF::ROPF_infos, nb_max_shunts, typeofconstraint)
    LB, status = solve_SDP(ROPF, typeofconstraint, nb_max_shunts)
    UB = solve_minlp(ROPF, [], Dict{String,Float64}(), nb_max_shunts, typeofconstraint)
    return UB, LB
end

function solve_BandB(ROPF::ROPF_infos, max_time::Int64, BB_parameters::BB_infos, nb_max_shunts, typeofconstraint)
    LB, status = solve_SDP(ROPF, typeofconstraint, nb_max_shunts)
    (UB, nb_nodes, open_nodes) = BandB_fixingsome1and0(ROPF, BB_parameters, max_time, nb_max_shunts, typeofconstraint)
    return UB, LB
end


text:
module BlockSparseArrays

using ComputedFieldTypes
using LinearAlgebra
using Random
using SparseArrays
using StaticArrays

include("vector.jl")
include("matrix.jl")

end


text:
function BeidouMeo(arg0::AbsoluteDate, arg1::AbsoluteDate, arg2::ExtendedPVCoordinatesProvider, arg3::Frame)
    return BeidouMeo((AbsoluteDate, AbsoluteDate, ExtendedPVCoordinatesProvider, Frame), arg0, arg1, arg2, arg3)
end


text:
vis = Visualizer()
open(vis)

# Parameters
origin = Origin{Float64}()
pbody = Sphere(0.1, 1.0)
joint0to1 = JointConstraint(Floating(origin, pbody))
bodies = [pbody]
joints = [joint0to1]
contact_p = contact_constraint(pbody, [0.0, 0.0, 1.0],
    friction_coefficient=0.5,
    contact_origin=[0.0, 0.0, 0.0],
    contact_radius=0.1,
    contact_type=:nonlinear)

contacts = [contact_p]
mech = Mechanism(origin, bodies, joints, contacts,
            gravity=-9.81,
            timestep=0.1)
mech.bodies[1].state.x2 = [0.0, 0.0, 0.5]

storage = simulate!(mech, 1.0,
    verbose=true,
    record=true)

visualize(mech, storage, vis=vis)


text:
using Base.Test
using Records


text:
using Test
using LVCyl

@testset "calculations" begin
    v1 = LorentzVectorCyl(1761.65,-2.30322,-2.5127,0.105652)
    v2 = LorentzVectorCyl(115.906,-2.28564,-2.50781,0.105713)
    @test energy(v1) ≈ 8901.870789524375 atol=1e-6
    @test px(v1) ≈ -1424.610065192358 atol=1e-6
    @test py(v1) ≈ -1036.2899616674022 atol=1e-6
    @test pz(v1) ≈ -8725.817601790963 atol=1e-6

    @test isapprox((v1+v2).mass, 8.25741602000877, atol=1e-6)
    @test isapprox(fast_mass(v1,v2), 8.25741602000877, atol=1e-6)

    v1 = LorentzVectorCyl(43.71242f0, 1.4733887f0, 1.6855469f0, 0.10571289f0)
    v2 = LorentzVectorCyl(36.994347f0, 0.38684082f0, -1.3935547f0, 0.10571289f0)
    @test (v1+v2).mass == 92.55651f0
    @test fast_mass(v1,v2) == 92.55651f0
    @test isapprox(deltar(v1,v2), 3.265188f0, atol=1e-6)
    @test isapprox(deltaphi(v1,v2), -3.0791016f0, atol=1e-6)

    v3 = v1*5
    @test v3.pt == 5*v1.pt
    @test v3.mass == 5*v1.mass
    @test v3.eta == v1.eta
    @test v3.phi == v1.phi
end

@testset "summing" begin
    v1 = LorentzVectorCyl(1761.65,-2.30322,-2.5127,0.105652)
    v2 = LorentzVectorCyl(115.906,-2.28564,-2.50781,0.105713)
    v3 = LorentzVectorCyl(43.71242f0, 1.4733887f0, 1.6855469f0, 0.10571289f0)
    v4 = LorentzVectorCyl(36.994347f0, 0.38684082f0, -1.3935547f0, 0.10571289f0)
    vs = [v1, v2, v3, v4]
    @test sum(vs).mass ≈ 2153.511000993
    @test sum(LorentzVectorCyl[]).mass ≈ 0
end

@testset "broadcasting" begin
    pts = [1761.65,115.906,43.712420,36.994347]
    etas = [-2.30322,-2.28564,1.4733887,0.38684082]
    phis = [-2.5127,-2.50781,1.6855469,-1.3935547]
    mass = 0.105652
    vs = LorentzVectorCyl.(pts, etas, phis, mass)
    @test all([v.mass for v in vs] .== mass)
    @test fast_mass.(vs[1], vs[2:end]) == fast_mass.(Ref(vs[1]), vs[2:end])
end

@testset "conversions" begin
    v1 = LorentzVectorCyl(1761.65,-2.30322,-2.5127,0.105652)
    v2 = fromcartesian(tocartesian(v1)...)
    @test v1.pt ≈ v2.pt
    @test v1.eta ≈ v2.eta
    @test v1.phi ≈ v2.phi
    @test v1.mass ≈ v2.mass atol=1e-6
end


text:
# ------  Prob of stepping down one or more grid points

"""
	$(SIGNATURES)

Switches for human capital shocks that step down one or more steps on the grid.
Each college has a probability of a shock. If there is more than one shock level (`maxSteps > 1`), then `stepProbV` gives the probability of stepping down another step. This is common to all colleges.
"""
Base.@kwdef mutable struct HcShockSwitchesDown <: HcShockSwitches
    nColleges :: CollInt
    # calProbs :: Bool = true
    # Max number of steps that can be stepped down
    maxSteps :: UInt8 = 1
    # Prob of stepping down one more step
    stepProbV :: Vector{Double} = [0.0]
    calStepProbs :: Bool = true;
end

mutable struct HcShockSetDown <: HcShockSet
    objId :: ObjectId
    switches :: HcShockSwitchesDown
    pvec :: ParamVector
    shockProbs :: BoundedVector{Double}
    # This is of length `max_steps - 1`
    stepProbV :: Vector{Double}
end

struct HcShockDown <: HcShock
    # Prob of stepping down n steps. Sum is total shock prob.
    probV :: Vector{Double}
end

function Base.show(io :: IO, dh :: HcShockSwitchesDown)
    # Shock probs are always calibrated. Step sizes may not be.
    calStr = "calibrated";
    maxSteps = max_steps(dh);
    print(io, "HcShockSwitchesDown: up to $maxSteps grid points; shock probabilities $calStr");
end

# Set
function Base.show(io :: IO, dh :: HcShockSetDown)
    maxSteps = max_steps(dh);
    print(io, "HcShockSetDown with probabilities  ",
        round.(shock_probs(dh), digits = 2));
    print(io,  "  Max shock size is $maxSteps grid points.")
end

function StructLH.describe(dh :: HcShockSetDown)
    maxSteps = max_steps(dh);
    # probV = round.(shock_probs(dh), digits = 2);

    m = [
        "H shocks"  "step down by $maxSteps grid points";
        "Probs for each college"  "by step size"
    ];
    for ic = 1 : n_colleges(dh)
        probV = shock_probs_one_college(dh :: HcShockSetDown, ic :: Integer);
        probV = round.(probV; digits = 2);
        m = vcat(m, ["  College $ic"  "$probV"]);
    end
    return m
end

function settings_table(dh :: HcShockSwitchesDown)
    maxPoints = max_steps(dh);
    return [
        "Human capital shocks"  " ";
        "Shocks drop h down by up to"  "$maxPoints grid points"
    ]
end

Base.show(io :: IO, dh :: HcShockDown) =
    print(io,  "HcShockDown with probability ",
        round(shock_prob(dh), digits = 2));

## -------------  Access

has_pvector(dh :: HcShockSetDown) = true;

# Shock prob by college
shock_probs(dh :: HcShockSetDown) = ModelParams.values(dh.shockProbs);
shock_probs(dh :: HcShockSetDown, idx) = shock_probs(dh)[idx];

# probs_calibrated(switches :: HcShockSwitchesDown) = switches.calStepProbs;

# Max no of steps down
max_steps(switches :: HcShockSwitchesDown) = switches.maxSteps;
max_steps(dh :: HcShockSetDown) = max_steps(dh.switches);
max_steps(dh :: HcShockDown) = length(dh.probV);

cal_step_probs(switches :: HcShockSwitchesDown) = switches.calStepProbs;

## -------------  Construction

function make_hshock_switches_down(nColleges, maxSteps)
    if maxSteps > 2
        stepProbV = collect(range(0.3, 0.2, length = maxSteps - 1));
        calStepProbs = true;
    else
        stepProbV = [0.3];
        calStepProbs = false;
    end
    switches = HcShockSwitchesDown(
        nColleges = nColleges, maxSteps = maxSteps,
        stepProbV = stepProbV, calStepProbs = calStepProbs);
    @assert validate_hshocks(switches)
    return switches
end

function make_h_shock_set(objId, switches :: HcShockSwitchesDown)
    # objId = make_child_id(model_id(), :HcShockSet,  "H shocks in college");

    pProbId = make_child_id(objId, :shockProbs, "H shock probabilities");
    pProb = init_shock_probs(pProbId, switches);
    pStepProbs = init_step_probs(switches);
    pvec = ParamVector(objId, [pStepProbs]);
    return HcShockSetDown(objId, switches, pvec, pProb, ModelParams.value(pStepProbs))
end

function init_shock_probs(objId :: ObjectId, switches :: HcShockSwitchesDown)
    # objId = make_child_id(parentId, :shockProbs,  "H shock probabilities");
    nc = n_colleges(switches);
    dxV = fill(0.2, nc);
    pProb = BoundedVector{Double}(objId, ParamVector(objId),
        :decreasing, 0.01, 0.5, dxV);
    set_pvector!(pProb; description = ldescription(:hShockGrad),
        symbol = lsymbol(:hShockGrad), isCalibrated = true);
    return pProb
end

function init_step_probs(switches :: HcShockSwitchesDown)
    if max_steps(switches) == 1
        # Irrelevant
        doCal = false;
        probV = [0.0];
    else
        probV = switches.stepProbV;
        @assert length(probV) == max_steps(switches) - 1
        doCal = switches.calStepProbs;
    end
    n = length(probV);
    p = Param(:stepProbV, ldescription(:hShockSteps), lsymbol(:hShockSteps),
        probV, probV, fill(0.2, n), fill(0.9, n), switches.calStepProbs);
    return p
end

function validate_hshocks(switches :: HcShockSwitchesDown)
    isValid = true;
    if max_steps(switches) > 1
        isValid = isValid && (length(switches.stepProbV) == max_steps(switches) - 1);
        isValid = isValid && all(switches.stepProbV .>= 0.0) &&
            all(switches.stepProbV .<= 1.0);
    else
        isValid = isValid  &&  !cal_step_probs(switches);
    end
    return isValid
end

function validate_hshocks(dh :: HcShockSetDown)
    isValid = all_at_least(shock_probs(dh), 0.0)  .&
        all_at_most(shock_probs(dh), 1.0)
    return isValid
end

function validate_hshocks(dh :: HcShockDown)
    return (shock_prob(dh) >= 0.0)  &&  (shock_prob(dh) <= 1.0)
end

function make_h_shock(dh :: HcShockSetDown, iCollege :: Integer)
    return HcShockDown(shock_probs_one_college(dh, iCollege))
end

# ----------------  Individual shock object for one college

# Shock probs, all steps, for one college
function shock_probs_one_college(dh :: HcShockSetDown, ic :: Integer)
    shockProb = shock_probs(dh, ic);
    if max_steps(dh) == 1
        probV = [shockProb];
    else
        probV = shockProb .* cumprod(vcat(1.0, dh.stepProbV));
    end
    return probV
end

# Probability of any shock
shock_prob(dh :: HcShockDown) = sum(dh.probV);

# Probability of a shock of size `nSteps`. Step size 0 is no shock
function shock_prob(dh :: HcShockDown, nSteps :: Integer)
    @assert 0 <= nSteps <= max_steps(dh)
    if nSteps == 0
        p = 1.0 - shock_prob(dh);
    else
        p = dh.probV[nSteps];
    end
    return p
end

"""
	$(SIGNATURES)

Simulate h shocks. Returns `h` next period, given `h` at the end of this period (before shocks are realized).

# Arguments
- `dh`: shock object
- `hEndIdxM`: h indices at end of period
"""
function sim_h_shocks(dh :: HcShockDown, hStartM, hEndIdxM :: Array{T1},
    hGridNextV, rng :: AbstractRNG) where T1 <: Integer

    hOutM = copy(hEndIdxM);
    maxSteps = max_steps(dh);
    for nSteps = one(T1) : T1(maxSteps)
        # Step one down, but not below index 1
        hOutM[(rand(rng, Double, size(hOutM)) .< shock_prob(dh, nSteps)) .& (hEndIdxM .> nSteps)] .-= one(T1);
    end
    return hOutM
end

"""
	$(SIGNATURES)

Prob(h on tomorrow's grid | h grid point at end of period today).
"""
function pr_hprime(dh :: HcShockDown, hEndIdx :: I1,
    hStart :: Real, hEnd :: Real, hGridNextV) where I1 <: Integer
    if hEndIdx == 1
        idxV = hEndIdx;
        prV = [1.0];
    else
        maxSteps = I1(min(max_steps(dh), hEndIdx - one(I1)));
        idxV = collect((hEndIdx - maxSteps) : hEndIdx);
        # idxV = [hEndIdx - one(hEndIdx), hEndIdx];
        prV = [shock_prob(dh, nSteps)  for nSteps in (maxSteps : -1 : 0)];
    end
    return idxV, prV
end

# -----------------


text:
# Parameter learning by Accelerated Expectation-Maximimzation (SQUAREM)
# RAVI VARADHAN & CHRISTOPHE ROLAND, Simple and Globally Convergent Methods for Accelerating the Convergence of Any EM Algorithm, J. Scand J Statist 2008
# Yu Du & Ravi Varadhan, SQUAREM: An R Package for Off-the-Shelf Acceleration of EM, MM and Other EM-Like Monotone Algorithms, J. Statistical Software, 2020.

"""
Learn weights using the Expectation Maximization algorithm.
"""
mutable struct SQUAREM <: ParameterLearner
spn::SumProductNetwork
layers::Vector{Vector{Int}}
cache1::SumProductNetwork
cache2::SumProductNetwork
cache3::SumProductNetwork
cache4::SumProductNetwork
diff::Vector{Float64} # to store derivatives
values::Vector{Float64} # to store logprobabilities
# dataset::AbstractMatrix
score::Float64     # score (loglikelihood)
prevscore::Float64 # for checking convergence
tolerance::Float64 # tolerance for convergence criterion
steps::Integer   # number of learning steps (epochs)
minimumvariance::Float64 # minimum variance for Gaussian leaves
SQUAREM(spn::SumProductNetwork) = new(spn, layers(spn), deepcopy(spn), deepcopy(spn), deepcopy(spn), deepcopy(spn), Array{Float64}(undef,length(spn)), Array{Float64}(undef,length(spn)), NaN, NaN, 1e-3, 0, 0.5)
end

"""
Verifies if algorithm has converged.

Convergence is defined as absolute difference of score.
Requires at least 2 steps of optimization.
"""
converged(learner::SQUAREM) = learner.steps < 2 ? false : abs(learner.score - learner.prevscore) < learner.tolerance

# TODO: take filename os CSV File object as input and iterate over file to decreae memory footprint
"""
Improvement step for learning weights using the Squared Iterative Expectation Maximization algorithm. Returns improvement in negative loglikelihood.

spn[i].weights[j] = spn[i].weights[k] * backpropagate(spn)[i]/sum(spn[i].weights[k] * backpropagate(spn)[i] for k=1:length(spn[i].children))

## Arguments

- `learner`: SQUAREM struct
- `data`: Data Matrix
- `smoothing`: weight smoothing factor (= pseudo expected count) [default: 0.0001]
- `minimumvariance`: minimum variance for Gaussian leaves [default: learner.minimumvariance]
"""
function update(learner::SQUAREM, Data::AbstractMatrix, smoothing::Float64 = 0.0001, minimumvariance::Float64 = learner.minimumvariance)

    numrows, numcols = size(Data)

    θ_0 = learner.spn
    θ_1 = learner.cache1
    θ_2 = learner.cache2
    r = learner.cache3
    v = learner.cache4
    # smooth out estaimtors to avoid degenerate probabilities
    sumnodes = filter(i -> isa(learner.spn[i], SumNode), 1:length(learner.spn))
    diff = learner.diff
    values = learner.values
    # Compute theta1 = EM_Update(theta0)
    for t = 1:numrows
        datum = view(Data, t, :)
        lv = plogpdf!(values, θ_0, learner.layers, datum) # parallelized version
        @assert isfinite(lv) "1. logvalue of datum $t is not finite: $lv"
        backpropagate!(diff, θ_0, values) # backpropagate derivatives
        Threads.@threads for i in sumnodes # update each node in parallel
            @inbounds for (k,j) in enumerate(learner.spn[i].children)
                if isfinite(values[j])
                    δ = θ_0[i].weights[k]*diff[i]*exp(values[j]-lv) # improvement
                    @assert isfinite(δ) "1. improvement to weight ($i,$j):$(θ_0[i].weights[k]) is not finite: $δ, $(diff[i]), $(values[j]), $(exp(values[j]-lv))"
                else
                    δ = 0.0
                end
                θ_1[i].weights[k] = ((t-1)/t)*θ_1[i].weights[k] + δ/t # running average for improved precision
                @assert θ_1[i].weights[k] ≥ 0
            end
        end
    end
    @inbounds Threads.@threads for i in sumnodes
        θ_1[i].weights .+= smoothing/length(θ_1[i].weights) # smoothing factor to prevent degenerate probabilities
        θ_1[i].weights ./= sum(θ_1[i].weights)
        @assert sum(θ_1[i].weights) ≈ 1.0 "1. Unnormalized weight vector at node $i: $(sum(θ_1[i].weights)) | $(θ_1[i].weights)"
    end
    # Compute theta2 = EM_Update(theta1)
    for t = 1:numrows
        datum = view(Data, t, :)
        lv = plogpdf!(values, θ_1, learner.layers, datum) # parallelized version
        @assert isfinite(lv) "2. logvalue of datum $t is not finite: $lv"
        backpropagate!(diff, θ_1, values) # backpropagate derivatives
        Threads.@threads for i in sumnodes # update each node in parallel
            @inbounds for (k,j) in enumerate(learner.spn[i].children)
                if isfinite(values[j])
                    δ = θ_1[i].weights[k]*diff[i]*exp(values[j]-lv) # improvement
                    @assert isfinite(δ) "2. improvement to weight ($i,$j):$(θ_1[i].weights[k]) is not finite: $δ, $(diff[i]), $(values[j]), $(exp(values[j]-lv))"
                else
                    δ = 0.0
                end
                θ_2[i].weights[k] = ((t-1)/t)*θ_2[i].weights[k] + δ/t
                @assert θ_2[i].weights[k] ≥ 0
            end
        end
    end
    @inbounds Threads.@threads for i in sumnodes
        θ_2[i].weights .+= smoothing/length(θ_2[i].weights) # smoothing factor to prevent degenerate probabilities
        θ_2[i].weights ./= sum(θ_2[i].weights)
        @assert sum(θ_2[i].weights) ≈ 1.0 "2. Unnormalized weight vector at node $i: $(sum(θ_2[i].weights)) | $(θ_2[i].weights)"
    end
    # Compute r, v, |r| and |v|
    r_norm, v_norm = 0.0, 0.0
    @inbounds Threads.@threads for i in sumnodes
        # r[i].weights .= θ_1[i].weights .- θ_0[i].weights
        # v[i].weights .= θ_2[i].weights .- θ_1[i].weights .- r[i].weights
        for k in 1:length(r[i].weights)
            r[i].weights[k] = θ_1[i].weights[k] - θ_0[i].weights[k]
            v[i].weights[k] = θ_2[i].weights[k] - θ_1[i].weights[k] - r[i].weights[k]
            r_norm += r[i].weights[k] * r[i].weights[k]
            v_norm += v[i].weights[k] * v[i].weights[k]
        end
        # r_norm += sum(r[i].weights .* r[i].weights)
        # v_norm += sum(v[i].weights .* v[i].weights)
    end
    # steplength
    α = -max(sqrt(r_norm)/sqrt(v_norm),1)
    #println("α: $α")
    # Compute θ' (reuse θ_1 for that matter)
    @inbounds Threads.@threads for i in sumnodes
        # θ' = θ0 - 2αr + α^2v
        θ_1[i].weights .= θ_0[i].weights
        θ_1[i].weights .-= ((2*α).*r[i].weights)
        θ_1[i].weights .+= ((α*α).*v[i].weights)
        θ_1[i].weights .+ smoothing/length(θ_1[i].weights) # add term to prevent negative weights due to numerical imprecision
        θ_1[i].weights ./= sum(θ_1[i].weights)
        @assert sum(θ_1[i].weights) ≈ 1.0 "3. Unnormalized weight vector at node $i: $(sum(θ_1[i].weights)) | $(θ_1[i].weights)"
        for w in θ_1[i].weights
            @assert w ≥ 0 "Negative weight at node $i: $(θ_1[i].weights)"
        end
    end
    # Final EM Update: θ_0 = EM_Update(θ')
    score = 0.0 # data loglikelihood
    for t = 1:numrows
        datum = view(Data, t, :)
        lv = plogpdf!(values, θ_1, learner.layers, datum) # parallelized version
        @assert isfinite(lv) "4. logvalue of datum $t is not finite: $lv"
        score += lv
        backpropagate!(diff, θ_1, values) # backpropagate derivatives
        Threads.@threads for i in sumnodes # update each node in parallel
            @inbounds for (k,j) in enumerate(learner.spn[i].children)
                if isfinite(values[j])
                    δ = θ_1[i].weights[k]*diff[i]*exp(values[j]-lv) # improvement
                    @assert isfinite(δ) "4. improvement to weight ($i,$j):$(θ_1[i].weights[k]) is not finite: $δ, $(diff[i]), $(values[j]), $(exp(values[j]-lv))"
                else
                    δ = 0.0
                end
                θ_0[i].weights[k] = ((t-1)/t)*θ_0[i].weights[k] + δ/t
                @assert θ_0[i].weights[k] ≥ 0
            end
        end
    end
    @inbounds Threads.@threads for i in sumnodes
        θ_0[i].weights .+= smoothing/length(θ_0[i].weights) # smoothing factor to prevent degenerate probabilities
        θ_0[i].weights ./= sum(θ_0[i].weights)
        @assert sum(θ_0[i].weights) ≈ 1.0 "4. Unnormalized weight vector at node $i: $(sum(θ_0[i].weights)) | $(θ_0[i].weights)"
    end
    learner.steps += 1
    learner.prevscore = learner.score
    learner.score = -score/numrows
    return learner.prevscore-learner.score, α
end


text:
module ProjB

using ProjC

const C_VERSION = ProjC.VERSION

end # module