text:
###
### DynamicHMC backend - https://github.com/tpapp/DynamicHMC.jl
###
struct DynamicNUTS{AD, space} <: Hamiltonian{AD} end

using LogDensityProblems: LogDensityProblems

struct FunctionLogDensity{F}
    dimension::Int
    f::F
end

LogDensityProblems.dimension(ℓ::FunctionLogDensity) = ℓ.dimension

function LogDensityProblems.capabilities(::Type{<:FunctionLogDensity})
    LogDensityProblems.LogDensityOrder{1}()
end

function LogDensityProblems.logdensity(ℓ::FunctionLogDensity, x::AbstractVector)
    first(ℓ.f(x))
end

function LogDensityProblems.logdensity_and_gradient(ℓ::FunctionLogDensity,
                                                    x::AbstractVector)
    ℓ.f(x)
end

"""
    DynamicNUTS()

Dynamic No U-Turn Sampling algorithm provided by the DynamicHMC package. To use it, make
sure you have the DynamicHMC package (version `2.*`) loaded:

```julia
using DynamicHMC
``
"""
DynamicNUTS(args...) = DynamicNUTS{ADBackend()}(args...)
DynamicNUTS{AD}() where AD = DynamicNUTS{AD, ()}()
function DynamicNUTS{AD}(space::Symbol...) where AD
    DynamicNUTS{AD, space}()
end

struct DynamicNUTSState{V<:AbstractVarInfo,D}
    vi::V
    draws::Vector{D}
end

DynamicPPL.getspace(::DynamicNUTS{<:Any, space}) where {space} = space

DynamicPPL.initialsampler(::Sampler{<:DynamicNUTS}) = SampleFromUniform()

function DynamicPPL.initialstep(
    rng::AbstractRNG,
    model::Model,
    spl::Sampler{<:DynamicNUTS},
    vi::AbstractVarInfo;
    N::Int,
    kwargs...
)
    # Set up lp function.
    function _lp(x)
        gradient_logp(x, vi, model, spl)
    end

    link!(vi, spl)
    l, dl = _lp(vi[spl])
    while !isfinite(l) || !isfinite(dl)
        model(vi, SampleFromUniform())
        link!(vi, spl)
        l, dl = _lp(vi[spl])
    end

    if spl.selector.tag == :default && !islinked(vi, spl)
        link!(vi, spl)
        model(vi, spl)
    end

    results = mcmc_with_warmup(
        rng,
        FunctionLogDensity(
            length(vi[spl]),
            _lp
        ),
        N
    )
    draws = results.chain

    # Compute first transition and state.
    draw = popfirst!(draws)
    vi[spl] = draw
    transition = Transition(vi)
    state = DynamicNUTSState(vi, draws)

    return transition, state
end

function AbstractMCMC.step(
    rng::AbstractRNG,
    model::Model,
    spl::Sampler{<:DynamicNUTS},
    state::DynamicNUTSState;
    kwargs...
)
    # Extract VarInfo object.
    vi = state.vi

    # Pop the next draw off the vector.
    draw = popfirst!(state.draws)
    vi[spl] = draw

    # Compute next transition.
    transition = Transition(vi)

    return transition, state
end

# Disable the progress logging for DynamicHMC, since it has its own progress meter.
function AbstractMCMC.sample(
    rng::AbstractRNG,
    model::AbstractModel,
    alg::DynamicNUTS,
    N::Integer;
    chain_type=MCMCChains.Chains,
    resume_from=nothing,
    progress=PROGRESS[],
    kwargs...
)
    if progress
        @warn "[HMC] Progress logging in Turing is disabled since DynamicHMC provides its own progress meter"
    end
    if resume_from === nothing
        return AbstractMCMC.sample(rng, model, Sampler(alg, model), N;
                                   chain_type=chain_type, progress=false, N=N, kwargs...)
    else
        return resume(resume_from, N; chain_type=chain_type, progress=false, N=N, kwargs...)
    end
end

function AbstractMCMC.sample(
    rng::AbstractRNG,
    model::AbstractModel,
    alg::DynamicNUTS,
    parallel::AbstractMCMC.AbstractMCMCParallel,
    N::Integer,
    n_chains::Integer;
    chain_type=MCMCChains.Chains,
    progress=PROGRESS[],
    kwargs...
)
    if progress
        @warn "[HMC] Progress logging in Turing is disabled since DynamicHMC provides its own progress meter"
    end
    return AbstractMCMC.sample(rng, model, Sampler(alg, model), parallel, N, n_chains;
                               chain_type=chain_type, progress=false, N=N, kwargs...)
end


text:
using SeaPearl
using SeaPearlExtras
using ReinforcementLearning
const RL = ReinforcementLearning
using Flux
using GeometricFlux
using Random

include("rewards.jl")
include("features.jl")

# -------------------
# Generator
# -------------------
latin_generator = SeaPearl.LatinGenerator(11,0.6)
# -------------------
# Internal variables
# -------------------

SR = SeaPearl.DefaultStateRepresentation{LatinFeaturization, SeaPearl.DefaultTrajectoryState}
numInFeatures = SeaPearl.feature_length(SR)
numGlobalFeature = SeaPearl.global_feature_length(SR)
# -------------------
# Experience variables
# -------------------
nbEpisodes = 500
evalFreq = 30
nbInstances = 1
nbRandomHeuristics = 0

# -------------------
# Agent definition
# -------------------

include("agents.jl")

# -------------------
# Value Heuristic definition
# -------------------
learnedHeuristic=SeaPearl.LearnedHeuristic{SR, InspectReward, SeaPearl.FixedOutput}(agent)
# Basic value-selection heuristic
selectMin(x::SeaPearl.IntVar; cpmodel=nothing) = SeaPearl.minimum(x.domain)
heuristic_min = SeaPearl.BasicHeuristic(selectMin)

function select_random_value(x::SeaPearl.IntVar; cpmodel=nothing)
    selected_number = rand(1:length(x.domain))
    i = 1
    for value in x.domain
        if i == selected_number
            return value
        end
        i += 1
    end
    @assert false "This should not happen"
end

randomHeuristics = []
for i in 1:nbRandomHeuristics
    push!(randomHeuristics, SeaPearl.BasicHeuristic(select_random_value))
end

valueSelectionArray = [learnedHeuristic,heuristic_min]
append!(valueSelectionArray, randomHeuristics)
# -------------------
# Variable Heuristic definition
# -------------------
variableSelection = SeaPearl.MinDomainVariableSelection{false}()

# -------------------
# -------------------
# Core function
# -------------------
# -------------------

function trytrain(nbEpisodes::Int)

    metricsArray, eval_metricsArray = SeaPearl.train!(
        valueSelectionArray=valueSelectionArray,
        generator=latin_generator,
        nbEpisodes=nbEpisodes,
        strategy=SeaPearl.DFSearch(),
        variableHeuristic=variableSelection,
        out_solver=false,
        verbose = false,
        evaluator=SeaPearl.SameInstancesEvaluator(valueSelectionArray,latin_generator; evalFreq = evalFreq, nbInstances = nbInstances),
        restartPerInstances = 1
    )

    return metricsArray, eval_metricsArray
end

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

metricsArray, eval_metricsArray = trytrain(nbEpisodes)
nothing

ext: jl

max_stars_repo_name: corail-research/SeaPearlZoo

max_issues_repo_name: corail-research/SeaPearlZoo

max_forks_repo_name: corail-research/SeaPearlZoo


text:
#=
Copyright (c) 2015, Intel Corporation

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:

    * Redistributions of source code must retain the above copyright notice,
      this list of conditions and the following disclaimer.
    * Redistributions in binary form must reproduce the above copyright
      notice, this list of conditions and the following disclaimer in the
      documentation and/or other materials provided with the distribution.
    * Neither the name of Intel Corporation nor the names of its contributors
      may be used to endorse or promote products derived from this software
      without specific prior written permission.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
=#

include("../../src/Sparso.jl")
include("../../src/simple-show.jl")
include("./utils.jl")
using Sparso

# The original pcg_symgs
#include("./pcg-symgs.jl")
function pcg_symgs(x, A, b, tol, maxiter)
    total_time = time()

    trsv_time = 0.
    spmv_time = 0.
    blas1_time = 0.

    L = tril(A)
    U  = SparseMatrixCSC{Cdouble, Cint}(spdiagm(1./diag(A)))*triu(A)

    spmv_time -= time()
    #r = b - A * x
    r = b - Sparso.SpMV(A,x)
    spmv_time += time()

    blas1_time -= time()
    #normr0 = norm(r)
    normr0 = Sparso.norm(r)
    z = copy(r)
    blas1_time += time()

    rel_err = 1

    trsv_time -= time()
    Base.SparseMatrix.fwdTriSolve!(L, z)
    Base.SparseMatrix.bwdTriSolve!(U, z)
    trsv_time += time()

    blas1_time -= time()
    p = copy(z) #NOTE: do not write "p=z"! That would make p and z aliased (the same variable)
    #rz = dot(r, z)
    rz = Sparso.dot(r,z)
    blas1_time += time()

    k = 1
    while k <= maxiter
        old_rz = rz

        spmv_time -= time()
        #Ap = A*p
        Ap = Sparso.SpMV(A,p)
        spmv_time += time()

        blas1_time -= time()
        #alpha = old_rz / dot(p, Ap)
        alpha = old_rz / Sparso.dot(p, Ap)
        #x += alpha * p
        Sparso.WAXPBY!(x,1,x,alpha,p)
        #r -= alpha * Ap
        Sparso.WAXPBY!(r,1,r,-alpha,Ap)
        #rel_err = norm(r)/normr0
        rel_err = Sparso.norm(r)/normr0
        blas1_time += time()

        if rel_err < tol
            break
        end

        blas1_time -= time()
        #z = copy(r)
        Sparso.copy!(z, r)
        blas1_time += time()

        trsv_time -= time()
        Base.SparseMatrix.fwdTriSolve!(L, z)
        Base.SparseMatrix.bwdTriSolve!(U, z)
        trsv_time += time()

        blas1_time -= time()
        #rz = dot(r, z)
        rz = Sparso.dot(r,z)
        beta = rz/old_rz
        #p = z + beta * p
        Sparso.WAXPBY!(p,1,z,beta,p)
        blas1_time += time()

        k += 1
    end
    total_time = time() - total_time
    println("total = $(total_time)s trsv_time = $(trsv_time)s ($((12.*(nnz(L) + nnz(U)) + 2.*8*(size(L, 1) + size(L, 2)))*k/trsv_time/1e9) gbps) spmv_time = $(spmv_time)s ($((12.*nnz(A) + 8.*(size(A, 1) + size(A, 2)))*(k + 1)/spmv_time/1e9) gbps) blas1_time = $blas1_time")

    return x, k, rel_err
end

# This code is what we expect to generate by context-sensitive optimizations except for reordering.
# It is used for debugging purpose only
function pcg_symgs_with_context_opt_without_reordering(x, A, b, tol, maxiter)
    total_time = time()

    trsv_time = 0.
    spmv_time = 0.
    blas1_time = 0.

    L = tril(A)
    U = SparseMatrixCSC{Cdouble, Cint}(spdiagm(1./diag(A)))*triu(A)

    spmv_time -= time()
    #r = b - A * x
    r = b - Sparso.SpMV(A,x)
    spmv_time += time()

    blas1_time -= time()
    #normr0 = norm(r)
    normr0 = Sparso.norm(r)
    z = copy(r)
    blas1_time += time()

    rel_err = 1

    trsv_time -= time()
    Base.SparseMatrix.fwdTriSolve!(L, z)
    Base.SparseMatrix.bwdTriSolve!(U, z)
    trsv_time += time()

    blas1_time -= time()
    p = copy(z) #NOTE: do not write "p=z"! That would make p and z aliased (the same variable)
    #rz = dot(r, z)
    rz = Sparso.dot(r,z)
    blas1_time += time()

    k = 1

    __mknobA = (Sparso.new_matrix_knob)(:A, true, true, true, true, false, false)
    __mknobL = (Sparso.new_matrix_knob)(:L, true, true, false, false, false, false) # matrix knob for L
    __mknobU = (Sparso.new_matrix_knob)(:U, true, true, false, false, false, false) # matrix knob for L

    (Sparso.set_derivative)(__mknobL, Sparso.DERIVATIVE_TYPE_SYMMETRIC, __mknobA)
    (Sparso.set_derivative)(__mknobU, Sparso.DERIVATIVE_TYPE_SYMMETRIC, __mknobA)

    fknob_spmv = (Sparso.new_function_knob)()
    (Sparso.add_mknob_to_fknob)(__mknobA,fknob_spmv)
    __fknob_8201 = (Sparso.new_function_knob)()
    (Sparso.add_mknob_to_fknob)(__mknobL,__fknob_8201)
    __fknob_8221 = (Sparso.new_function_knob)()
    (Sparso.add_mknob_to_fknob)(__mknobU,__fknob_8221)

    while k <= maxiter
        old_rz = rz

        spmv_time -= time()
        #Ap = A*p
        Ap = Sparso.SpMV(A,p,fknob_spmv)
        spmv_time += time()

        blas1_time -= time()
        #alpha = old_rz / dot(p, Ap)
        alpha = old_rz / Sparso.dot(p, Ap)
        #x += alpha * p
        Sparso.WAXPBY!(x,1,x,alpha,p)
        #r -= alpha * Ap
        Sparso.WAXPBY!(r,1,r,-alpha,Ap)
        #rel_err = norm(r)/normr0
        rel_err = Sparso.norm(r)/normr0
        blas1_time += time()

        if rel_err < tol
            break
        end

        blas1_time -= time()
        #z = copy(r)
        Sparso.copy!(z, r)
        blas1_time += time()

        trsv_time -= time()
        #Base.SparseMatrix.fwdTriSolve!(L, z)
        Sparso.fwdTriSolve!(L, z, __fknob_8201)
        #Base.SparseMatrix.bwdTriSolve!(U, z)
        Sparso.bwdTriSolve!(U, z, __fknob_8221)
        trsv_time += time()

        blas1_time -= time()
        #rz = dot(r, z)
        rz = Sparso.dot(r,z)
        beta = rz/old_rz
        #p = z + beta * p
        Sparso.WAXPBY!(p,1,z,beta,p)
        blas1_time += time()

        k += 1
    end

    (Sparso.delete_function_knob)(__fknob_8221)
    (Sparso.delete_function_knob)(__fknob_8201)
    (Sparso.delete_matrix_knob)(__mknobL)

    total_time = time() - total_time
    println("total = $(total_time)s trsv_time = $(trsv_time)s ($((12.*(nnz(L) + nnz(U)) + 2.*8*(size(L, 1) + size(L, 2)))*k/trsv_time/1e9) gbps) spmv_time = $(spmv_time)s ($((12.*nnz(A) + 8.*(size(A, 1) + size(A, 2)))*(k + 1)/spmv_time/1e9) gbps) blas1_time = $blas1_time")

    return x, k, rel_err
end

# This code is what we expect to generate by context-sensitive optimizations.
# It is used for debugging purpose only
function pcg_symgs_with_context_opt(x, A, b, tol, maxiter)
    total_time = time()

    trsv_time = 0.
    spmv_time = 0.
    blas1_time = 0.
    reorder_time = 0.

    L = tril(A)
    U  = SparseMatrixCSC{Cdouble, Cint}(spdiagm(1./diag(A)))*triu(A)

    spmv_time -= time()
    #r = b - A * x
    r = b - Sparso.SpMV(A,x)
    spmv_time += time()

    blas1_time -= time()
    #normr0 = norm(r)
    normr0 = Sparso.norm(r)
    z = copy(r)
    blas1_time += time()

    rel_err = 1

    trsv_time -= time()
    Base.SparseMatrix.fwdTriSolve!(L, z)
    Base.SparseMatrix.bwdTriSolve!(U, z)
    trsv_time += time()

    blas1_time -= time()
    p = copy(z) #NOTE: do not write "p=z"! That would make p and z aliased (the same variable)
    #rz = dot(r, z)
    rz = Sparso.dot(r,z)
    blas1_time += time()

    k = 1

    __mknobA = (Sparso.new_matrix_knob)(:A, true, true, true, true, false, false)
    __mknobL = (Sparso.new_matrix_knob)(:L, true, true, false, false, false, false) # matrix knob for L
    __mknobU = (Sparso.new_matrix_knob)(:U, true, true, false, false, false, false) # matrix knob for L

    (Sparso.set_derivative)(__mknobL, Sparso.DERIVATIVE_TYPE_SYMMETRIC, __mknobA)
    (Sparso.set_derivative)(__mknobU, Sparso.DERIVATIVE_TYPE_SYMMETRIC, __mknobA)

    fknob_spmv = (Sparso.new_function_knob)()
    (Sparso.add_mknob_to_fknob)(__mknobA,fknob_spmv)
    __fknob_8201 = (Sparso.new_function_knob)()
    (Sparso.add_mknob_to_fknob)(__mknobL,__fknob_8201)
    __fknob_8221 = (Sparso.new_function_knob)()
    (Sparso.add_mknob_to_fknob)(__mknobU,__fknob_8221)

    # Set up reordering:
    # Let Sparso.fwdTriSolve! decides what permutation/inverse
    # permutation vector should be, and reorders its inputs andoutputs accordingly.
    # Other functions just respect the decision, makes no decision, nor does any
    # reordering.
    (Sparso.set_reordering_decision_maker)(__fknob_8201)

    # Reordering_status is a vector. See the comments in lib-interface.jl:
    # reordering() for the meaning of the elements
    reordering_status = [false, C_NULL, C_NULL, C_NULL, C_NULL, reorder_time]
    while k <= maxiter
        old_rz = rz

        spmv_time -= time()
        #Ap = A*p
        Ap = Sparso.SpMV(A,p,fknob_spmv)
        spmv_time += time()

        blas1_time -= time()
        #alpha = old_rz / dot(p, Ap)
        alpha = old_rz / Sparso.dot(p, Ap)
        #x += alpha * p
        Sparso.WAXPBY!(x,1,x,alpha,p)
        #r -= alpha * Ap
        Sparso.WAXPBY!(r,1,r,-alpha,Ap)
        #rel_err = norm(r)/normr0
        rel_err = Sparso.norm(r)/normr0
        blas1_time += time()

        if rel_err < tol
            break
        end

        blas1_time -= time()
        #z = copy(r)
        Sparso.copy!(z, r)
        blas1_time += time()

        trsv_time -= time()
        #Base.SparseMatrix.fwdTriSolve!(L, z)
        Sparso.fwdTriSolve!(L, z, __fknob_8201)
        trsv_time += time()

        Sparso.reordering(
            __fknob_8201,
            reordering_status,
            A, Sparso.COL_PERM, Sparso.ROW_INV_PERM, __mknobA,
            U, Sparso.COL_PERM, Sparso.ROW_INV_PERM, __mknobU,
            :__delimitor__,
            r, Sparso.ROW_PERM,
            x, Sparso.COL_PERM,
            p, Sparso.COL_PERM
        )

        trsv_time -= time()
        #Base.SparseMatrix.bwdTriSolve!(U, z)
        Sparso.bwdTriSolve!(U, z, __fknob_8221)
        trsv_time += time()

        blas1_time -= time()
        #rz = dot(r, z)
        rz = Sparso.dot(r,z)
        beta = rz/old_rz
        #p = z + beta * p
        Sparso.WAXPBY!(p,1,z,beta,p)
        blas1_time += time()

        k += 1
    end

    Sparso.reverse_reordering(
        reordering_status,
        :__delimitor__,
        x, Sparso.COL_PERM
    )

    (Sparso.delete_function_knob)(__fknob_8221)
    (Sparso.delete_function_knob)(__fknob_8201)
    (Sparso.delete_matrix_knob)(__mknobL)

    total_time = time() - total_time
    println("total = $(total_time)s trsv_time = $(trsv_time)s ($((12.*(nnz(L) + nnz(U)) + 2.*8*(size(L, 1) + size(L, 2)))*k/trsv_time/1e9) gbps) spmv_time = $(spmv_time)s ($((12.*nnz(A) + 8.*(size(A, 1) + size(A, 2)))*(k + 1)/spmv_time/1e9) gbps) blas1_time = $blas1_time reorder_time = $reorder_time")

    return x, k, rel_err
end

print_sum = false
if length(ARGS) == 0
    m = 1000
    A = generate_symmetric_nonzerodiagonal_sparse_matrix(m)
elseif length(ARGS) == 1
    # This is for performance testing's purpose
    A = matrix_market_read(ARGS[1], true, true)
else
    # This is to for regression tests' purpose, which needs sum to be printed.
    print_sum = true
    A         = matrix_market_read(ARGS[2], true, true)
end
m       = size(A, 1)
x       = zeros(Float64, m)
b       = ones(Float64, m)
tol     = 1e-7
maxiter = 20000

x, k, rel_err = pcg_symgs(x, A, b, tol, maxiter)
x       = zeros(Float64, m)
println("Original: ")
x, k, rel_err = pcg_symgs(x, A, b, tol, maxiter)
if print_sum
    println("\tOriginal sum of x=", sum(x))
end
println("\tOriginal k=", k)
println("\tOriginal rel_err=", rel_err)

x, k, rel_err = pcg_symgs_with_context_opt_without_reordering(x, A, b, tol, maxiter)
println("\nWith manual context-sensitive optimization without reordering: ")
x       = zeros(Float64, m)
x, k, rel_err = pcg_symgs_with_context_opt_without_reordering(x, A, b, tol, maxiter)
if print_sum
    println("\tManual_context_no_reorder sum of x=", sum(x))
end
println("\tManual_context_no_reorder k=", k)
println("\tManual_context_no_reorder rel_err=", rel_err)

A2 = copy(A) # workaround that we change A in-place
x, k, rel_err = pcg_symgs_with_context_opt(x, A2, b, tol, maxiter)
println("\nWith manual context-sensitive optimization: ")
x       = zeros(Float64, m)
x, k, rel_err = pcg_symgs_with_context_opt(x, A, b, tol, maxiter)
if print_sum
    println("\tManual_context sum of x=", sum(x))
end
println("\tManual_context k=", k)
println("\tManual_context rel_err=", rel_err)

max_stars_repo_name: IntelLabs/Sparso

max_issues_repo_name: IntelLabs/Sparso

max_forks_repo_name: IntelLabs/Sparso


text:
# Common interface for AbstractNLPEvaluator
#
function Base.show(io::IO, nlp::AbstractNLPEvaluator)
    n = n_variables(nlp)
    m = n_constraints(nlp)
    println(io, "A Evaluator object")
    println(io, "    * #vars: ", n)
    println(io, "    * #cons: ", m)
end

## Generic callbacks
function constraint(nlp::AbstractNLPEvaluator, x)
    cons = similar(x, n_constraints(nlp)) ; fill!(cons, 0)
    constraint!(nlp, cons, x)
    return cons
end

function gradient(nlp::AbstractNLPEvaluator, x)
    ∇f = similar(x) ; fill!(∇f, 0)
    gradient!(nlp, ∇f, x)
    return ∇f
end

function jacobian(nlp::AbstractNLPEvaluator, x)
    n = n_variables(nlp)
    m = n_constraints(nlp)
    J = similar(x, m, n) ; fill!(J, 0)
    jacobian!(nlp, J, x)
    return J
end

# Default implementation of jprod!, using full Jacobian matrix
function jprod!(nlp::AbstractNLPEvaluator, jv, u, v)
    nᵤ = length(u)
    m  = n_constraints(nlp)
    @assert nᵤ == length(v)
    jac = jacobian(nlp, u)
    mul!(jv, jac, v)
    return
end

# Joint Objective Jacobian transpose vector product (default implementation)
function ojtprod!(nlp::AbstractNLPEvaluator, jv, u, σ, v)
    gradient!(nlp, jv, u)
    jv .*= σ  # scale gradient
    jtprod!(nlp, jv, u, v)
    return
end

# Common interface for Hessian
function hessian!(nlp::AbstractNLPEvaluator, hess, x)
    n = n_variables(nlp)
    v = similar(x)
    @inbounds for i in 1:n
        hv = @view hess[:, i]
        fill!(v, 0)
        v[i] = 1.0
        hessprod!(nlp, hv, x, v)
    end
end

function hessian_lagrangian_penalty!(nlp::AbstractNLPEvaluator, hess, x, y, σ, D,)
    n = n_variables(nlp)
    v = similar(x)
    @inbounds for i in 1:n
        hv = @view hess[:, i]
        fill!(v, 0)
        v[i] = 1.0
        hessian_lagrangian_penalty_prod!(nlp, hv, x, y, σ, v, D)
    end
end

function hessian(nlp::AbstractNLPEvaluator, x)
    n = n_variables(nlp)
    H = similar(x, n, n) ; fill!(H, 0)
    hessian!(nlp, H, x)
    return H
end

# Counters
abstract type AbstractCounter end

mutable struct NLPCounter <: AbstractCounter
    objective::Int
    gradient::Int
    hessian::Int
    jacobian::Int
    jtprod::Int
    hprod::Int
end
NLPCounter() = NLPCounter(0, 0, 0, 0, 0, 0)

function Base.empty!(c::NLPCounter)
    for attr in fieldnames(NLPCounter)
        setfield!(c, attr, 0)
    end
end

# Active set utils
function _check(val, val_min, val_max)
    violated_inf = findall(val .< val_min)
    violated_sup = findall(val .> val_max)
    n_inf = length(violated_inf)
    n_sup = length(violated_sup)
    err_inf = norm(val_min[violated_inf] .- val[violated_inf], Inf)
    err_sup = norm(val[violated_sup] .- val_max[violated_sup] , Inf)
    return (n_inf, err_inf, n_sup, err_sup)
end

function _inf_pr(nlp::AbstractNLPEvaluator, cons)
    (n_inf, err_inf, n_sup, err_sup) = _check(cons, nlp.g_min, nlp.g_max)
    return max(err_inf, err_sup)
end

# Scaler utils
abstract type AbstractScaler end

scale_factor(h, tol, η) = max(tol, η / max(1.0, h))

struct MaxScaler{T, VT} <: AbstractScaler
    scale_obj::T
    scale_cons::VT
    g_min::VT
    g_max::VT
end
function MaxScaler(g_min, g_max)
    @assert length(g_min) == length(g_max)
    m = length(g_min)
    sc = similar(g_min) ; fill!(sc, 1.0)
    return MaxScaler{eltype(g_min), typeof(g_min)}(1.0, sc, g_min, g_max)
end
function MaxScaler(nlp::AbstractNLPEvaluator, u0::AbstractVector;
                   η=100.0, tol=1e-8)
    n = n_variables(nlp)
    m = n_constraints(nlp)
    conv = update!(nlp, u0)
    ∇g = similar(u0) ; fill!(∇g, 0)
    gradient!(nlp, ∇g, u0)
    s_obj = scale_factor(norm(∇g, Inf), tol, η)

    VT = typeof(u0)
    ∇c = xzeros(VT, n)
    s_cons = xzeros(VT, m)
    v = xzeros(VT, m)
    for i in eachindex(s_cons)
        fill!(v, 0.0)
        v[i] = 1.0
        jtprod!(nlp, ∇c, u0, v)
        s_cons[i] = scale_factor(norm(∇c, Inf), tol, η)
    end

    g♭, g♯ = bounds(nlp, Constraints())

    return MaxScaler{typeof(s_obj), typeof(s_cons)}(s_obj, s_cons, s_cons .* g♭, s_cons .* g♯)
end


text:
struct LatticeBoltzmannModel{
    Q <: Quadrature,
    T,
    CM <: CollisionModel,
    PM <: ProcessingMethod,
    BCs <: AbstractVector{<:BoundaryCondition},
}
    f_stream::T
    f_collision::T
    quadrature::Q
    collision_model::CM
    boundary_conditions::BCs
    processing_method::PM
end
function LatticeBoltzmannModel(
    problem,
    quadrature;
    collision_model = SRT,
    initialization_strategy = InitializationStrategy(problem),
    process_method,
)
    f_stream = initialize(initialization_strategy, quadrature, problem, collision_model)
    f_collision = copy(f_stream)

    LatticeBoltzmannModel(
        f_stream,
        f_collision,
        quadrature,
        CollisionModel(collision_model, quadrature, problem),
        boundary_conditions(problem),
        process_method,
    )
end
function simulate(
    problem::FluidFlowProblem,
    q::Quadrature;
    process_method = nothing,
    should_process = true,
    initialization_strategy = InitializationStrategy(problem),
    t_end = 1.0,
    collision_model = SRT,
)
    Δt = delta_t(problem)
    n_steps = round(Int, t_end / Δt)

    if isnothing(process_method)
        process_method = ProcessingMethod(problem, should_process, n_steps)
    end

    model = LatticeBoltzmannModel(
        problem,
        q,
        collision_model = collision_model,
        initialization_strategy = initialization_strategy,
        process_method = process_method,
    )

    simulate(model, 0:n_steps)
end
function simulate(model::LatticeBoltzmannModel, time)
    Δt = isdefined(model.processing_method, :problem) ?
        delta_t(model.processing_method.problem) : 0.0

    @inbounds for t in time
        collide!(model, time = t * Δt)
        stream!(model)
        apply_boundary_conditions!(model, time = t * Δt)

        if next!(model, t + 1)
            return model
        end
    end

    next!(model, last(time) + 1)

    model
end

# The following are helper functions which use the old interface where we
# explicitely pass the quadrature, collision model etc.
# This will likely be refactored so that we can write specialized function
# for each specific LatticeBoltzmannModel (once we also introduce ddf models)

function collide!(model::LatticeBoltzmannModel; time)
    collide!(
        model.collision_model,
        model.quadrature,
        f_new = model.f_collision,
        f_old = model.f_stream,
        time = time,
    )
end

function stream!(model::LatticeBoltzmannModel)
    stream!(model.quadrature, f_new = model.f_stream, f_old = model.f_collision)
end

function apply_boundary_conditions!(model::LatticeBoltzmannModel; time = 0.0)
    apply!(
        model.boundary_conditions,
        model.quadrature,
        model.f_stream,
        model.f_collision,
        time = time,
    )
end

function next!(model::LatticeBoltzmannModel, t::Int64)
    next!(model.processing_method, model.quadrature, model.f_stream, t)
end


text:
using BenchmarkTools
using GemmKernels
using GemmKernels.Tiling
using GemmKernels.Layout
using CUDA
using Test
using Statistics
using Printf

# TCCG benchmark 1: D_abc = A_bda * B_dc

# sizes for each dimension
SA = 64
SB = 32
SC = 2048
SD = 2048

A = CuArray(rand(Float16, (SB, SD, SA)))
B = CuArray(rand(Float16, (SD, SC)))

# layout for the A tensor
abstract type LayoutA{T} <: Layout.AlignedColMajor{T} end

@inline function Layout.load(::Type{LayoutA{T}}, workspace, tile::Tile{size}) where {T, size}
    NUMEL = 16 ÷ sizeof(T)

    M = tile.base.M + tile.offset.M
    K = tile.base.K + tile.offset.K

    d = K

    a = M ÷ Base.size(workspace, 1)
    b = M % Base.size(workspace, 1)

    offset = 1 + b + d * Base.size(workspace, 1) + a * Base.size(workspace, 1) * Base.size(workspace, 2)

    Layout.vloada(Layout.Vec{NUMEL, T}, pointer(workspace), offset)
end

# layout for the B tensor
abstract type LayoutB{T} <: Layout.AlignedColMajor{T} end

@inline function Layout.load(::Type{LayoutB{T}}, workspace, tile::Tile{size}) where {T, size}
    NUMEL = 16 ÷ sizeof(T)

    K = tile.base.K + tile.offset.K
    N = tile.base.N + tile.offset.N

    d = K
    c = N

    offset = 1 + d + c * Base.size(workspace, 1)

    Layout.vloada(Layout.Vec{NUMEL, T}, pointer(workspace), offset)
end

# layout for the C tensor
abstract type LayoutC{T} <: Layout.AlignedColMajor{T} end

@inline function Layout.load(::Type{LayoutC{T}}, workspace, tile::Tile{size}) where {T, size}
    N = 16 ÷ sizeof(T)

    ntuple(Val(N)) do i
        VecElement{T}(zero(T))
    end
end

# layout for the D tensor
abstract type LayoutD{T} <: Layout.AlignedColMajor{T} end

@inline function Layout.store!(::Type{LayoutD{T}}, workspace, value, tile::Tile{size}) where {T, size}
    NUMEL = 16 ÷ sizeof(T)

    for i = 1 : NUMEL
        M = tile.base.M + tile.offset.M
        N = tile.base.N + tile.offset.N

        c = N

        a = (M + i - 1) ÷ Base.size(workspace, 2)
        b = (M + i - 1) % Base.size(workspace, 2)

        @inbounds workspace[a + 1, b + 1, c + 1] = value[i].value
    end
end

# implementation using GemmKernels.jl
function gemmkernels_impl(;benchmark = false)
    D = CuArray(zeros(Float16, (SA, SB, SC)))

    M = SA * SB
    N = SC
    K = SD

    conf = GemmKernels.get_config(
                                  gemm_shape = (M = M, N = N, K = K),
                                  operator = Operator.WMMAOp{16, 16, 16, Float16},

                                  global_a_layout = LayoutA{Float16},
                                  global_b_layout = LayoutB{Float16},
                                  global_c_layout = LayoutC{Float16},
                                  global_d_layout = LayoutD{Float16},

                                  shared_a_layout = Layout.Padded{Layout.AlignedColMajor{Float16}, 8},
                                  shared_b_layout = Layout.Padded{Layout.AlignedColMajor{Float16}, 8},
                                  shared_c_layout = Layout.AlignedColMajor{Float16},
                                  shared_d_layout = Layout.AlignedColMajor{Float16},

                                  is_a_col_major = true,
                                  is_b_col_major = true,
                                 )

    if !benchmark
        GemmKernels.matmul(A, B, D, D, conf;
                        kernel = Kernel.matmul_pipelined
                        )
        D
    else
        times = []

        for i = 1 : 10000
            synchronize(context())
            time = CUDA.@elapsed GemmKernels.matmul(A, B, D, D, conf;
                            kernel = Kernel.matmul_pipelined
                        )
            push!(times, time)
        end

        times
    end
end

# cuTENSOR implementation
function cutensor_impl(;algo = CUDA.CUTENSOR.CUTENSOR_ALGO_DEFAULT, benchmark = false)
    D = CuArray(zeros(Float16, (SA, SB, SC)))

    plan = CUDA.CUTENSOR.plan_contraction(A, [ 'b', 'd', 'a' ], CUDA.CUTENSOR.CUTENSOR_OP_IDENTITY,
                                          B, [ 'd', 'c' ], CUDA.CUTENSOR.CUTENSOR_OP_IDENTITY,
                                          D, [ 'a', 'b', 'c' ], CUDA.CUTENSOR.CUTENSOR_OP_IDENTITY,
                                          CUDA.CUTENSOR.CUTENSOR_OP_IDENTITY,
                                          algo = algo,
                                          compute_type = Float32)

    if !benchmark
        CUDA.CUTENSOR.contraction!(1,
                                A, [ 'b', 'd', 'a' ], CUDA.CUTENSOR.CUTENSOR_OP_IDENTITY,
                                B, [ 'd', 'c' ], CUDA.CUTENSOR.CUTENSOR_OP_IDENTITY,
                                0,
                                D, [ 'a', 'b', 'c' ], CUDA.CUTENSOR.CUTENSOR_OP_IDENTITY,
                                CUDA.CUTENSOR.CUTENSOR_OP_IDENTITY,
                                compute_type = Float32,
                                plan = plan)
        D
    else
        times = []

        for i = 1 : 10000
            synchronize(context())
            time = CUDA.@elapsed CUDA.CUTENSOR.contraction!(1,
                                    A, [ 'b', 'd', 'a' ], CUDA.CUTENSOR.CUTENSOR_OP_IDENTITY,
                                    B, [ 'd', 'c' ], CUDA.CUTENSOR.CUTENSOR_OP_IDENTITY,
                                    0,
                                    D, [ 'a', 'b', 'c' ], CUDA.CUTENSOR.CUTENSOR_OP_IDENTITY,
                                    CUDA.CUTENSOR.CUTENSOR_OP_IDENTITY,
                                    compute_type = Float32,
                                    plan = plan)
            push!(times, time)
        end

        times
    end
end

# compare
function test()
    D_reference = cutensor_impl()
    D_gemmkernels = gemmkernels_impl()

    @test all(isapprox.(Array(D_reference), Array(D_gemmkernels); rtol = sqrt(eps(Float16))))
end

# Taken from BenchmarkTools.jl: src/trials.jl
function rmskew(t::Vector)
    st = sort(t)
    return st[1:(BenchmarkTools.skewcutoff(st) - 1)]
end

# Benchmark function
function bench()
    bench_cutensor = cutensor_impl(benchmark = true)

    bench_cutensor_gett = cutensor_impl(algo = CUDA.CUTENSOR.CUTENSOR_ALGO_GETT, benchmark = true)

    bench_gemmkernels = gemmkernels_impl(benchmark = true)

    for (title, bench) in [("cuTENSOR (let heuristic choose algorithm):", bench_cutensor),
                           ("cuTENSOR (force GETT):", bench_cutensor_gett),
                           ("GemmKernels:", bench_gemmkernels),
                           ("rmskew: cuTENSOR (let heuristic choose algorithm):", rmskew(bench_cutensor)),
                           ("rmskew: cuTENSOR (force GETT):", rmskew(bench_cutensor_gett)),
                           ("rmskew: GemmKernels:", rmskew(bench_gemmkernels))
                          ]
        println(title)
        bench .= bench .* 1e6 # convert seconds to us
        @printf("Samples:            %u\n", length(bench))
        @printf("Range (min .. max): %.2f us .. %.2f us\n", minimum(bench), maximum(bench))
        @printf("Median:             %.2f us\n", median(bench))
        @printf("Mean ± std:         %.2f us ± %.2f us\n", mean(bench), std(bench))
        println()
    end

    nothing
end

# Main entry point
function main()
    display(test())
    println()

    bench()
end

!isinteractive() && main()


text:
using QuantumOptics: tensor, CompositeBasis
using .PhysicalConstants: ħ, c

export Trap, get_basis, Efield_from_pi_time, Efield_from_pi_time!,
       transition_frequency, set_gradient!, Efield_from_rabi_frequency,
       Efield_from_rabi_frequency!, global_beam!, get_η

#############################################################################################
# an ion trap with a linear ion chain configuration
#############################################################################################

"""
    Trap(;
            configuration::LinearChain, B::Real=0, Bhat::NamedTuple{(:x,:y,:z)=ẑ, ∇B::Real=0,
             δB::Union{Real,Function}=0, lasers::Vector{Laser}
    )

Information necessary to describe the Hamiltonian for a collection of ions in a linear chain
interacting with laser light.
**user-defined fields**
* `configuration<:LinearChain`
* `B`: A real value describing the mean magnitude of the B-field [Tesla].
* `Bhat::NamedTuple{(:x,:y,:z)}`: Describes the direction of the B-field (defaults to ẑ).
* `∇B`: Magnitude of the B-field gradient. We assume that the gradient always points along the
        z-direction. [Tesla / meter]
* `δB::Function`: Time-dependence of the B-field [Tesla]
* `lasers::Array{<:Laser}`: For each laser in the array, the pointing field should contain
        an array of `Tuple{Int,Real}`. The first element specifies the index of an ion
        in the `ions` field that the laser interacts with. The second element specifies a
        scaling factor for the strength of that interaction (to be used, e.g., for
        modeling cross-talk).
**derived fields**
* `_cnst_δB::Bool`: A Boolean flag signifying whether or not `δB` is a constant function.
* `basis<:CompositeBasis`: The basis for describing the combined system, ions + vibrational
        modes. If constructing the Hamiltonian explictly (with [`hamiltonian`](@ref)), then
        the ordering of the basis is set, by convention, as
        ``ion₁ ⊗ ion₂ ⊗ ... ⊗ ion_N ⊗ mode₁ ⊗ mode₂ ⊗ ... ⊗ mode_N``, where the ion bases are
        ordered according to the order in `T.configuration.ions` and the vibrational modes
        are ordered according to the order in
        `[T.configuration.vibrational_modes.x, T.configuration.vibrational_modes.y,
        T.configuration.vibrational_modes.z]`.
    E.g. for:

    ```
    chain = LinearChain(ions=[C1, C2], com_frequencies=(x=2e6,y=2e6,z=1e6),
    selected_modes=(x=[1, 2], y=[], z=[1]))
    ```

    The ordering of the basis would be

    `C1.basis ⊗ C2.basis ⊗ chain.vibrational_modes.x[1].basis
    ⊗ chain.vibrational_modes.x[2].basis ⊗ chain.vibrational_modes.z[1].basis`

    Otherwise, the ordering is according to the form of the initial state used in the solver.
"""
mutable struct Trap
    configuration::LinearChain
    B::Real
    Bhat::NamedTuple{(:x,:y,:z)}
    ∇B::Real
    δB::Function
    lasers::Array{<:Laser}
    basis::CompositeBasis
    _cnst_δB::Bool
    function Trap(;
            configuration::LinearChain, B=0, Bhat=ẑ, ∇B=0, δB::TδB=0, lasers=Laser[]
        ) where {TδB}
        warn = nothing
        for i in 1:length(lasers)
            if length(lasers[i].pointing) == 0
                for n in eachindex(configuration.ions)
                    push!(lasers[i].pointing, (n, 1.0))
                end
            end
            for j in i+1:length(lasers)
                if lasers[j] ≡ lasers[i]
                    lasers[j] = copy(lasers[i])
                    if isnothing(warn)
                        warn = "Some lasers point to the same thing. Making copies."
                        @warn warn
                    end
                end
            end
        end
        @assert isapprox(norm(Bhat), 1, rtol=1e-6) "!(|$Bhat| = 1)"
        for (li, l) in enumerate(lasers), p in l.pointing
            @assert p[1] <= length(configuration.ions) (
                """lasers[$li] points at configuration.ions[$(p[1])], but there are only
                 $(length(configuration.ions)) ions."""
            )
        end
        if TδB <: Number
            _cnst_δB = true
            δBt(t) = δB
        else
            _cnst_δB = false
            δBt = δB
        end
        basis = tensor(
            configuration.ions...,
            configuration.vibrational_modes.x...,
            configuration.vibrational_modes.y...,
            configuration.vibrational_modes.z...,
        )
        new(configuration, B, Bhat, ∇B, δBt, lasers, basis, _cnst_δB)
    end
end

function Base.setproperty!(T::Trap, s::Symbol, v::Tv) where{Tv}
    if s == :δB
        if Tv <: Number
            _cnst_δB = true
            vt(t) = v
        else
            _cnst_δB = false
            vt = v
        end
        Core.setproperty!(T, s, vt)
        Core.setproperty!(T, :_cnst_δB, _cnst_δB)
    elseif s == :basis || s == :_cnst_δB
        return
    elseif s == :configuration
        Core.setproperty!(T, s, v)
        basis = tensor(
            v.ions..., v.vibrational_modes.x..., v.vibrational_modes.y...,
            v.vibrational_modes.z...,
        )
        Core.setproperty!(T, :basis, basis)
    else
        Core.setproperty!(T, s, v)
    end
end

Base.show(io::IO, T::Trap) = print(io, "Trap")  # suppress long output

#############################################################################################
# general functions
#############################################################################################

"""
    get_basis(T::Trap)
Returns the composite basis describing the Hilbert space for `T`.
"""
function get_basis(T::Trap)::CompositeBasis
    tensor(
            T.configuration.ions...,
            T.configuration.vibrational_modes.x...,
            T.configuration.vibrational_modes.y...,
            T.configuration.vibrational_modes.z...,
        )
end

"""
    global_beam!(T::Trap, laser::Laser)
Set `laser` to shine with full intensity on all ions in `Trap`.
"""
function global_beam!(T::Trap, laser::Laser)
    for n in eachindex(T.configuration.ions)
        push!(laser.pointing, (n, 1.0))
    end
end

"""
    Efield_from_pi_time(
        pi_time::Real, Bhat::NamedTuple{(:x,:y,:z)}, laser::Laser, ion::Ion,
        transition::Union{Tuple{String,String},Vector{<:String}}
    )
Compute the E-field needed to get a certain `pi_time` with a certain `laser`-`ion`
`transition`.

Alternatively, one may use
```
Efield_from_pi_time(
            pi_time::Real, T::Trap, laser_index::Int, ion_index::Int,
            transition::Union{Tuple{String,String},Vector{<:String}}
        )
```
which is the same as
`Efield_from_pi_time(pi_time, T.Bhat, T.lasers[laser_index], T.configuration.ions[ion_index],
transition)`
"""
function Efield_from_pi_time(
    pi_time::Real, Bhat::NamedTuple{(:x,:y,:z)}, laser::Laser, ion::Ion,
    transition::Union{Tuple{String,String},Vector{<:String}}
)
    p = laser.pointing
    (γ, ϕ) = map(x -> rad2deg(acos(ndot(Bhat, x))), [laser.ϵ, laser.k])
    s_indx = findall(x -> x[1] == ion.number, p)
    @assert length(s_indx) > 0 "This laser doesn't shine on this ion"
    s = laser.pointing[s_indx[1]][2]
    Ω = s * ion.selected_matrix_elements[tuple(transition...)](1, γ, ϕ)
    if Ω < 1e-15
        # even when coupling strength is zero, numerical error causes it to be finite
        # (on order 1e-16), this is a band-aid to prevent users from unknowingly setting
        # the E-field to something absurd (like 1e20 V/m)
        return Inf
    end
    1 / (2Ω * pi_time)
end

function Efield_from_pi_time(
        pi_time::Real, T::Trap, laser_index::Int, ion_index::Int,
        transition::Union{Tuple{String,String},Vector{<:String}}
    )
    Efield_from_pi_time(
            pi_time, T.Bhat, T.lasers[laser_index], T.configuration.ions[ion_index],
            transition
        )
end

"""
    Efield_from_pi_time!(
            pi_time::Real, Bhat::NamedTuple{(:x,:y,:z)}, laser::Laser, ion::Ion,
            transition::Union{Tuple{String,String},Vector{<:String}}
        )
Same as [`Efield_from_pi_time`](@ref), but updates `laser[:E]` in-place.
"""
function Efield_from_pi_time!(
        pi_time::Real, Bhat::NamedTuple{(:x,:y,:z)}, laser::Laser, ion::Ion,
        transition::Union{Tuple{String,String},Vector{<:String}}
    )
    Efield::Float64 = Efield_from_pi_time(pi_time, Bhat, laser, ion, transition)
    laser.E = t -> Efield
end

function Efield_from_pi_time!(
        pi_time::Real, T::Trap, laser_index::Int, ion_index::Int,
        transition::Union{Tuple{String,String},Vector{<:String}}
    )
    Efield::Float64 = Efield_from_pi_time(pi_time, T, laser_index, ion_index, transition)
    T.lasers[laser_index].E = t -> Efield
end

"""
    Efield_from_rabi_frequency(
        Ω::Real, Bhat::NamedTuple{(:x,:y,:z)}, laser::Laser, ion::Ion,
        transition::Union{Tuple{String,String},Vector{<:String}}
    )
Compute the Efield needed to get a certain rabi frequency `Ω` with a certain `laser`-`ion`
`transition`.

Alternatively, one may use
```
Efield_from_rabi_frequency(
            Ω::Real, T::Trap, laser_index::Int, ion_index::Int,
            transition::Union{Tuple{String,String},Vector{<:String}}
        )
```
which is the same as
`Efield_from_rabi_frequency(pi_time, T.Bhat, T.lasers[laser_index],
T.configuration.ions[ion_index], transition)`
"""
function Efield_from_rabi_frequency(
    Ω::Real, Bhat::NamedTuple{(:x,:y,:z)}, laser::Laser, ion::Ion,
    transition::Union{Tuple{String,String},Vector{<:String}})
    Efield_from_pi_time(1 / 2Ω, Bhat, laser, ion, transition)
end

function Efield_from_rabi_frequency(
        Ω::Real, T::Trap, laser_index::Int, ion_index::Int,
        transition::Union{Tuple{String,String},Vector{<:String}}
    )
    Efield_from_rabi_frequency(
            Ω, T.Bhat, T.lasers[laser_index], T.configuration.ions[ion_index],
            transition
        )
end

"""
    Efield_from_rabi_frequency!(
        Ω::Real, Bhat::NamedTuple{(:x,:y,:z)}, laser::Laser, ion::Ion,
        transition::Union{Tuple{String,String},Vector{<:String}}
    )
Same as [`Efield_from_rabi_frequency`](@ref), but updates `laser[:E]` in-place.
"""
function Efield_from_rabi_frequency!(
        Ω::Real, Bhat::NamedTuple{(:x,:y,:z)}, laser::Laser, ion::Ion,
        transition::Union{Tuple{String,String},Vector{<:String}}
    )
    Efield::Float64 = Efield_from_rabi_frequency(Ω, Bhat, laser, ion, transition)
    laser.E = t -> Efield
end

function Efield_from_rabi_frequency!(
        Ω::Real, T::Trap, laser_index::Int, ion_index::Int,
        transition::Union{Tuple{String,String},Vector{<:String}}
    )
    Efield::Float64 = Efield_from_rabi_frequency(Ω, T, laser_index, ion_index, transition)
    T.lasers[laser_index].E = t -> Efield
end

"""
    transition_frequency(
        B::Real, ion::Ion, transition::Union{Tuple{String,String},Vector{<:String}}
    )
Compute the transition frequency of the `ion`'s selected transition under Bfield `B`.
Alternatively, one may use:
```
transition_frequency(
        T::Trap, ion::Ion, transition::Union{Tuple{String,String},Vector{<:String}}
    )
```
which is the same as `transition_frequency(T.B, ion, transition)` or
```
transition_frequency(
        T::Trap, ion_index::Int, transition::Union{Tuple{String,String},Vector{<:String}}
    )
```
which is the same as `transition_frequency(T.B, T.configuration.ions[ion_index], transition)`.
"""
function transition_frequency(
        B::Real, ion::Ion, transition::Union{Tuple{String,String},Vector{<:String}}
    )
    diff([map(x -> zeeman_shift(B, ion.selected_level_structure[x]), transition)...])[1]
end

function transition_frequency(
        T::Trap, ion::Ion, transition::Union{Tuple{String,String},Vector{<:String}}
    )
    B = T.B + T.∇B * ion.position
    transition_frequency(B, ion, transition)
end

function transition_frequency(
        T::Trap, ion_index::Int, transition::Union{Tuple{String,String},Vector{<:String}}
    )
    ion = T.configuration.ions[ion_index]
    B = T.B + T.∇B * ion.position
    transition_frequency(B, ion, transition)
end

"""
    set_gradient!(
            T::Trap, ion_indxs::Tuple{Int,Int}, transition::Tuple{String,String}, f::Real
        )
Sets the Bfield gradient in place to achieve a detuning `f` between the `transition` of two
ions, which are assumed to be of the same species. `ion_indxs` refer to the
ordering of the ions in the chain.
"""
function set_gradient!(
        T::Trap, ion_indxs::Tuple{Int,Int}, transition::Tuple{String,String}, f::Real
    )
    ion1 = T.configuration.ions[ion_indxs[1]]
    ion2 = T.configuration.ions[ion_indxs[2]]
    separation = abs(ion1.position - ion2.position)
    E1, E2 = map(x -> zeeman_shift(1, ion1.selected_level_structure[x]), transition)
    T.∇B = f / (abs(E2 - E1) * separation)
end

# In QunatumOptics.jl, this method will return true whenever the shapes of b1 and b2 match,
# but we'd like to distinguish, i.e., between Ion ⊗ mode1 ⊗ mode2 and Ion ⊗ mode2 ⊗ mode1
# when mode1.N == mode2.N but mode1.axis ≠ mode2.axis.
function (QuantumOptics.:(==)(b1::T, b2::T)
    where {T<:CompositeBasis{<:Vector{Int},<:Tuple{Vararg{<:IonSimBasis}}}})
N = length(b1.bases)
if N ≠ length(b2.bases)
    return false
end
for i in 1:N
    if !(b1.bases[i] == b2.bases[i])
        return false
    end
end
return true
end

"""
    get_η(V::VibrationalMode, L::Laser, I::Ion)
The Lamb-Dicke parameter:
``|k|cos(\\theta)\\sqrt{\\frac{\\hbar}{2m\\nu}}``
for a given vibrational mode, ion and laser.
"""
function get_η(V::VibrationalMode, L::Laser, I::Ion; scaled=false)
    @fastmath begin
        k = 2π / L.λ
        scaled ? ν = 1 : ν = V.ν
        x0 = √(ħ / (2 * I.mass * 2π * ν))
        cosθ = ndot(L.k, V.axis)
        k * x0 * cosθ * V.mode_structure[I.number]
    end
end


text:
#------------------------------------------------------------------------------
# Author: Edward Tate <edward.tate@erlang-solutions.com>
#------------------------------------------------------------------------------

export reply, reply_close
export reply_stream, reply_stream_close
export reply_info
export noreply, noreply_close

#------------------------------------------------------------------------------
# Generic Server Functions
#------------------------------------------------------------------------------

function reply (s::TCPSocket, x::Any)
    term, nbytes = encode_message((:reply, x))
    nbytes_written = write_erlang_term(s, term)
    nbytes_written
end

function reply_stream (s::TCPSocket, bytes::Array{Uint8})
    term, nbytes = encode_message((:reply, length(bytes)))
    write_erlang_term(s, term)
    bin, nbytes = encode_message(bytes)
    write_binary_stream(s, bin)
end

function reply_stream_close (s::TCPSocket, bytes::Array{Uint8})
    reply_stream(s, bytes)
    close(s)
end

function reply_close (s::TCPSocket, x::Any)
    reply(s, x)
    close(s)
end

function reply_info (s::TCPSocket, command::Symbol, options)
    term, nbytes = encode_message((:info, command, options))
    write_erlang_term(s, term)
end

function noreply (s::TCPSocket)
    term, nbytes = encode_message(tuple(:noreply))
    write_erlang_term(s, term)
end

function noreply_close (s::TCPSocket)
    noreply(s)
    close(s)
end

max_issues_repo_name: esl/FsServer

max_forks_repo_name: esl/FsServer


text:
function LinearRegression(
    x :: AbstractArray{T, 1},
    y :: AbstractArray{T, 1};
    order :: Integer = 1,
) where T <: AbstractFloat

    if order < 0
        throw(ErrorException("Parameter order must be non-negative."))
    end

    N = length(x)

    bases = order + 1

    ϕ = zeros(N, bases)

    for i = 1:bases

        if i == 1
            ϕ[:, 1] .= 1.0
        else
            ϕ[:, i] .= x.^(i-1)
        end

    end

    # ϕ β = y => β = ϕ \ y

    return ϕ \ y

end

function detrend(
    x :: AbstractArray{T, 1},
    y :: AbstractArray{T, 1};
    order :: Integer = 1,
) where T <: AbstractFloat

    if order < 0
        throw(ErrorException("Parameter order must be non-negative."))
    end

    bases = order + 1

    β = LinearRegression(x, y; order=order)

    result = copy(y)
    total_mean = 0.0
    for i = 1:bases
        result .-= β[i] * x.^(i-1)
        total_mean += β[i]/ i * (x[end]^i - x[1]^i)
    end

    result .+= total_mean / (x[end] - x[1])

    return result

end

max_stars_repo_name: meteorologytoday/CESM-diagnostic

max_issues_repo_name: meteorologytoday/CESM-diagnostic

max_forks_repo_name: meteorologytoday/CESM-diagnostic


text:
const julia_version = v"1.4.2"
include("../common.jl")

max_stars_repo_name: waralex/Yggdrasil

max_issues_repo_name: waralex/Yggdrasil

max_forks_repo_name: waralex/Yggdrasil


text:
show(io::IO, e::AEdge) = print(io, "$(src(e))=>$(dst(e))")

==(e1::AEdge, e2::AEdge) = (src(e1) == src(e2) && dst(e1) == dst(e2))

"""
    struct Edge
        src::Int
        dst::Int
    end

A type representing an edge between two vertices of a graph.
"""
struct Edge{T} <: AEdge
    src::T
    dst::T
end

Edge(::Type{G},u,v) where {G<:ADiGraph} = Edge(u,v)
Edge(::Type{G},u,v) where {G<:AGraph} =  v > u ? Edge(u,v) : Edge(v,u)

function Edge(u::T, v::S) where {T,S}
    V = promote_type(T,S)
    return Edge{V}(promote(u, v)...)
end

src(e::Edge) = e.src
dst(e::Edge) = e.dst
iterate(e::AEdge, i=0) = (i+=1; i>2 ? nothing : (getfield(e,i), i))
reverse(e::Edge) = Edge(dst(e), src(e))

show(io::IO, e::AIndexedEdge) = print(io, "($(src(e))=>$(dst(e)),$(idx(e)))")

"""
    struct IndexedEdge <: AIndexedEdge
        src::Int
        dst::Int
        idx::Int
    end

An indexed edge type

    IndexedEdge(u, v) = IndexedEdge(u,v,-1)

Creates an edge with invalid index.
"""
struct IndexedEdge <: AIndexedEdge
    src::Int
    dst::Int
    idx::Int
end

IndexedEdge(u, v) = IndexedEdge(u,v,-1)

src(e::IndexedEdge) = e.src
dst(e::IndexedEdge) = e.dst
idx(e::IndexedEdge) = e.idx
reverse(e::IndexedEdge) = IndexedEdge(e.dst, e.src, e.idx)

Base.sort(e::AEdge) = src(e) <= dst(e) ? e : reverse(e)


text:
module MaxHelpingHandParallaxFadeOutController

using ..Ahorn, Maple

@mapdef Entity "MaxHelpingHand/ParallaxFadeOutController" ParallaxFadeOutController(x::Integer, y::Integer)

const placements = Ahorn.PlacementDict(
    "Parallax Fade Out Controller (max480's Helping Hand)" => Ahorn.EntityPlacement(
        ParallaxFadeOutController
    )
)

function Ahorn.selection(entity::ParallaxFadeOutController)
    x, y = Ahorn.position(entity)

    return Ahorn.Rectangle(x - 12, y - 12, 24, 24)
end

Ahorn.render(ctx::Ahorn.Cairo.CairoContext, entity::ParallaxFadeOutController, room::Maple.Room) = Ahorn.drawImage(ctx, Ahorn.Assets.northernLights, -12, -12)

end

max_stars_repo_name: Kitty-Cats/MaxHelpingHand

max_issues_repo_name: Kitty-Cats/MaxHelpingHand

max_forks_repo_name: Kitty-Cats/MaxHelpingHand


text:
using AvailablePotentialEnergyFramework, ImageFiltering, Test
@testset "Array filtering" begin

    @testset "Kernel creation" begin

    @test AvailablePotentialEnergyFramework.kernel_4d(false,false) ==
    ImageFiltering.kernelfactors((centered([(1.0)]),centered([(1.0)]),centered([(1.0)]),centered([(1.0)]) ))

    @test AvailablePotentialEnergyFramework.kernel_4d(5,false) ==
    ImageFiltering.kernelfactors((centered(ones(5)./5.0),centered(ones(5)./5.0),centered([(1.0)]),centered([(1.0)]) ))

    end
    @testset "Filter arrays 1 and 2" begin
        @testset "Filter_array!" begin
            ones3d  = ones(20,20,20)
            ones4d  = ones(20,20,20,20)
            zeros3d = zeros(20,20,20)
            zeros4d = zeros(20,20,20,20)
            buf3d   = similar(ones3d)
            buf4d   = similar(ones4d)
            filter_array!(buf3d,zeros3d,5,5,1)
            filter_array!(buf3d,ones3d,5,5,1)
            filter_array!(buf4d,zeros4d,5,5,1)
            filter_array!(buf4d,ones4d,5,5,1)
            @test ones3d   == ones(20,20,20)
            @test zeros3d  == zeros(20,20,20)
            @test ones4d   == ones(20,20,20,20)
            @test zeros4d  == zeros(20,20,20,20)
        end

        @testset "Filter_array_2" begin
            zeros3d = zeros(20,20,20);
            ones3d = ones(20,20,20)
            zeros4d = zeros(20,20,20,20);
            ones4d = ones(20,20,20,20)

            filter_array_2!(ones3d,5,5,1)
            filter_array_2!(zeros3d,5,5,1)
            filter_array_2!(ones4d,5,5,1)
            filter_array_2!(zeros4d,5,5,1)

            @test ones3d == ones(20,20,20)
            @test zeros3d == zeros(20,20,20)
            @test ones4d == ones(20,20,20,20)
            @test zeros4d == zeros(20,20,20,20)

        end

         @testset "Both filters simultaneously" begin
            arr3d = rand(20,20,20)
            arr3d2 = copy(arr3d)
            arr3d3 = copy(arr3d)
            arr4d = rand(20,20,20,20)
            arr4d2 = copy(arr4d)
            arr4d3 = copy(arr4d)
            buf3d   = similar(arr3d)
            buf4d   = similar(arr4d)

            filter_array_2!(arr3d2,5,5,1)
            filter_array_2!(arr4d2,5,5,1)
            filter_array!(buf3d,arr3d3,5,5,1)
            filter_array!(buf4d,arr4d3,5,5,1)
            @test arr3d3 == arr3d2
            @test arr4d3 == arr4d2

        end
    end
end


text:
# Similar to JuMP/test/solvers.jl

function try_import(name::Symbol)
    try
        @eval import $name
        return true
    catch e
        return false
    end
end

mos = try_import(:MosekTools)
csd = try_import(:CSDP)
scs = try_import(:SCS)

# Semidefinite factories
sdp_factories = MOI.OptimizerWithAttributes[]
# Need at least Mosek 8 for sosdemo3 to pass; see:
# https://github.com/JuliaOpt/Mosek.jl/issues/92
# and at least Mosek 8.0.0.41 for sosdemo6 to pass; see:
# https://github.com/JuliaOpt/Mosek.jl/issues/98
mos && push!(sdp_factories, optimizer_with_attributes(MosekTools.Mosek.Optimizer, MOI.Silent() => true))
# Currently, need to create a file param.csdp to avoid printing, see https://github.com/JuliaOpt/CSDP.jl/issues/15
csd && push!(sdp_factories, optimizer_with_attributes(CSDP.Optimizer, MOI.Silent() => true))
iscsdp(factory) = occursin("CSDP", string(factory.optimizer_constructor))
# Need 54000 iterations for sosdemo3 to pass on Linux 64 bits
# With 55000, sosdemo3 passes for every platform except Windows 64 bits on AppVeyor
scs && push!(sdp_factories, optimizer_with_attributes(SCS.Optimizer, "eps_abs" => 1e-6, "max_iters" => 60000, MOI.Silent() => true))
isscs(factory) = occursin("SCS", string(factory.optimizer_constructor))

max_stars_repo_name: blegat/SumOfSquares

max_issues_repo_name: blegat/SumOfSquares

max_forks_repo_name: blegat/SumOfSquares


text:
# Autogenerated wrapper script for Infernal_jll for x86_64-linux-musl
export cmalign, cmbuild, cmcalibrate, cmconvert, cmemit, cmfetch, cmpress, cmscan, cmsearch, cmstat, easel, esl_afetch, esl_alimanip, esl_alimap, esl_alimask, esl_alimerge, esl_alipid, esl_alirev, esl_alistat, esl_compalign, esl_compstruct, esl_construct, esl_histplot, esl_mask, esl_mixdchlet, esl_reformat, esl_selectn, esl_seqrange, esl_seqstat, esl_sfetch, esl_shuffle, esl_ssdraw, esl_translate, esl_weight

using GSL_jll
JLLWrappers.@generate_wrapper_header("Infernal")
JLLWrappers.@declare_executable_product(cmalign)
JLLWrappers.@declare_executable_product(cmbuild)
JLLWrappers.@declare_executable_product(cmcalibrate)
JLLWrappers.@declare_executable_product(cmconvert)
JLLWrappers.@declare_executable_product(cmemit)
JLLWrappers.@declare_executable_product(cmfetch)
JLLWrappers.@declare_executable_product(cmpress)
JLLWrappers.@declare_executable_product(cmscan)
JLLWrappers.@declare_executable_product(cmsearch)
JLLWrappers.@declare_executable_product(cmstat)
JLLWrappers.@declare_executable_product(easel)
JLLWrappers.@declare_executable_product(esl_afetch)
JLLWrappers.@declare_executable_product(esl_alimanip)
JLLWrappers.@declare_executable_product(esl_alimap)
JLLWrappers.@declare_executable_product(esl_alimask)
JLLWrappers.@declare_executable_product(esl_alimerge)
JLLWrappers.@declare_executable_product(esl_alipid)
JLLWrappers.@declare_executable_product(esl_alirev)
JLLWrappers.@declare_executable_product(esl_alistat)
JLLWrappers.@declare_executable_product(esl_compalign)
JLLWrappers.@declare_executable_product(esl_compstruct)
JLLWrappers.@declare_executable_product(esl_construct)
JLLWrappers.@declare_executable_product(esl_histplot)
JLLWrappers.@declare_executable_product(esl_mask)
JLLWrappers.@declare_executable_product(esl_mixdchlet)
JLLWrappers.@declare_executable_product(esl_reformat)
JLLWrappers.@declare_executable_product(esl_selectn)
JLLWrappers.@declare_executable_product(esl_seqrange)
JLLWrappers.@declare_executable_product(esl_seqstat)
JLLWrappers.@declare_executable_product(esl_sfetch)
JLLWrappers.@declare_executable_product(esl_shuffle)
JLLWrappers.@declare_executable_product(esl_ssdraw)
JLLWrappers.@declare_executable_product(esl_translate)
JLLWrappers.@declare_executable_product(esl_weight)
function __init__()
    JLLWrappers.@generate_init_header(GSL_jll)
    JLLWrappers.@init_executable_product(
        cmalign,
        "bin/cmalign",
    )

    JLLWrappers.@init_executable_product(
        cmbuild,
        "bin/cmbuild",
    )

    JLLWrappers.@init_executable_product(
        cmcalibrate,
        "bin/cmcalibrate",
    )

    JLLWrappers.@init_executable_product(
        cmconvert,
        "bin/cmconvert",
    )

    JLLWrappers.@init_executable_product(
        cmemit,
        "bin/cmemit",
    )

    JLLWrappers.@init_executable_product(
        cmfetch,
        "bin/cmfetch",
    )

    JLLWrappers.@init_executable_product(
        cmpress,
        "bin/cmpress",
    )

    JLLWrappers.@init_executable_product(
        cmscan,
        "bin/cmscan",
    )

    JLLWrappers.@init_executable_product(
        cmsearch,
        "bin/cmsearch",
    )

    JLLWrappers.@init_executable_product(
        cmstat,
        "bin/cmstat",
    )

    JLLWrappers.@init_executable_product(
        easel,
        "bin/easel",
    )

    JLLWrappers.@init_executable_product(
        esl_afetch,
        "bin/esl-afetch",
    )

    JLLWrappers.@init_executable_product(
        esl_alimanip,
        "bin/esl-alimanip",
    )

    JLLWrappers.@init_executable_product(
        esl_alimap,
        "bin/esl-alimap",
    )

    JLLWrappers.@init_executable_product(
        esl_alimask,
        "bin/esl-alimask",
    )

    JLLWrappers.@init_executable_product(
        esl_alimerge,
        "bin/esl-alimerge",
    )

    JLLWrappers.@init_executable_product(
        esl_alipid,
        "bin/esl-alipid",
    )

    JLLWrappers.@init_executable_product(
        esl_alirev,
        "bin/esl-alirev",
    )

    JLLWrappers.@init_executable_product(
        esl_alistat,
        "bin/esl-alistat",
    )

    JLLWrappers.@init_executable_product(
        esl_compalign,
        "bin/esl-compalign",
    )

    JLLWrappers.@init_executable_product(
        esl_compstruct,
        "bin/esl-compstruct",
    )

    JLLWrappers.@init_executable_product(
        esl_construct,
        "bin/esl-construct",
    )

    JLLWrappers.@init_executable_product(
        esl_histplot,
        "bin/esl-histplot",
    )

    JLLWrappers.@init_executable_product(
        esl_mask,
        "bin/esl-mask",
    )

    JLLWrappers.@init_executable_product(
        esl_mixdchlet,
        "bin/esl-mixdchlet",
    )

    JLLWrappers.@init_executable_product(
        esl_reformat,
        "bin/esl-reformat",
    )

    JLLWrappers.@init_executable_product(
        esl_selectn,
        "bin/esl-selectn",
    )

    JLLWrappers.@init_executable_product(
        esl_seqrange,
        "bin/esl-seqrange",
    )

    JLLWrappers.@init_executable_product(
        esl_seqstat,
        "bin/esl-seqstat",
    )

    JLLWrappers.@init_executable_product(
        esl_sfetch,
        "bin/esl-sfetch",
    )

    JLLWrappers.@init_executable_product(
        esl_shuffle,
        "bin/esl-shuffle",
    )

    JLLWrappers.@init_executable_product(
        esl_ssdraw,
        "bin/esl-ssdraw",
    )

    JLLWrappers.@init_executable_product(
        esl_translate,
        "bin/esl-translate",
    )

    JLLWrappers.@init_executable_product(
        esl_weight,
        "bin/esl-weight",
    )

    JLLWrappers.@generate_init_footer()
end  # __init__()


text:
##### Run the BEFW with chemical effects on y and B0, varying the slope, asymptote and inflection point

#### Initial conditions ####
Z = 1
I_K = 3
T = 293.15

doses = [0:1:100.0;]

# time span for each simulation
tspan = (0.0, 315360000000.0)

# define extinction threshold
threshold = 10e-12

#### Varying slopes ####
slopes = [4.0, 8.0, 16.0, 32.0, 64.0]
targets = ["y", "B0", "r", "x"]

## empty dataframe to store results
df = DataFrame(
    target = [],
    slope = [],
    dose = [],
    K = [],
    r = [],
    x_2 = [],
    y_21 = [],
    B0_21 = [],
    maxB1 = [],
    maxB2 = [],
    maxB3 = [],
    minB1 = [],
    minB2 = [],
    minB3 = [],
    sp_survival = []
)

for rate in targets
    for sl in slopes
        for dose in doses
            # generate parameters specific to this combo of target rate, DR slope and dose
            p = ModelParameters_CG(param, T, I_K, Z, dose, rate, slope = sl, lower = 0.0, xmid = 50.0)

            # Calculate starting biomasses
            u0 = [p.K[1]/2, p.K[1]/2, p.K[1]/2]

            # Define the problem
            prob = ODEProblem(BEFW, u0, tspan, p)

            # Solve the problem
            sol = solve(prob, maxiters = 1e7)

            ###### Output metrics ######
            # turn vector output of biomasses into matrix
            matrix_u = hcat(sol.u...)' # 56×3 adjoint(::Matrix{Float64}) with eltype Float64:

            ### Biomass extremes
            maxB1 = maximum(matrix_u[2 * Int64(round(length(matrix_u[:,1])/3)):Int64(length(matrix_u[:,1])),1])
            maxB2 = maximum(matrix_u[2 * Int64(round(length(matrix_u[:,1])/3)):Int64(length(matrix_u[:,1])),2])
            maxB3 = maximum(matrix_u[2 * Int64(round(length(matrix_u[:,1])/3)):Int64(length(matrix_u[:,1])),3])

            minB1 = minimum(matrix_u[2 * Int64(round(length(matrix_u[:,1])/3)):Int64(length(matrix_u[:,1])),1])
            minB2 = minimum(matrix_u[2 * Int64(round(length(matrix_u[:,1])/3)):Int64(length(matrix_u[:,1])),2])
            minB3 = minimum(matrix_u[2 * Int64(round(length(matrix_u[:,1])/3)):Int64(length(matrix_u[:,1])),3])

            ### Number of surviving species
            # final biomasses
            final_u = matrix_u[length(sol.t),:]

            # count number of species remaining that are above extinction threshold
            survived = count(final_u .> threshold)

            ### Push outputs to DataFrame
            push!(df, [
                rate,
                sl,
                dose,
                p.K[1],
                p.r[1],
                p.x[2],
                p.y[2,1],
                p.B0[2,1],
                maxB1,
                maxB2,
                maxB3,
                minB1,
                minB2,
                minB3,
                survived
            ])

            # print some stuff - see how the simulation is progressing
            println(("target rate = $rate", "slope = $sl", "dose = $dose"))
        end
    end
end

### Save output
CSV.write("output_dfs/full_slope.csv", df)

#### Varying asymptotes ####
asymptotes = [0.0, 0.1, 0.2, 0.3, 0.4]

## empty dataframe to store results
df2 = DataFrame(
    target = [],
    asymptote = [],
    dose = [],
    K = [],
    r = [],
    x_2 = [],
    y_21 = [],
    B0_21 = [],
    maxB1 = [],
    maxB2 = [],
    maxB3 = [],
    minB1 = [],
    minB2 = [],
    minB3 = [],
    sp_survival = []
)

for rate in targets
    for as in asymptotes
        for dose in doses
            # generate parameters specific to this combo of target rate, DR asymptote and dose
            p = ModelParameters_CG(param, T, I_K, Z, dose, rate, slope = 5.0, lower = as, xmid = 50.0)

            # Calculate starting biomasses
            u0 = [p.K[1]/2, p.K[1]/2, p.K[1]/2]

            # Define the problem
            prob = ODEProblem(BEFW, u0, tspan, p)

            # Solve the problem
            sol = solve(prob, maxiters = 1e7)

            ###### Output metrics ######
            # turn vector output of biomasses into matrix
            matrix_u = hcat(sol.u...)' # 56×3 adjoint(::Matrix{Float64}) with eltype Float64:

            ### Biomass extremes
            maxB1 = maximum(matrix_u[2 * Int64(round(length(matrix_u[:,1])/3)):Int64(length(matrix_u[:,1])),1])
            maxB2 = maximum(matrix_u[2 * Int64(round(length(matrix_u[:,1])/3)):Int64(length(matrix_u[:,1])),2])
            maxB3 = maximum(matrix_u[2 * Int64(round(length(matrix_u[:,1])/3)):Int64(length(matrix_u[:,1])),3])

            minB1 = minimum(matrix_u[2 * Int64(round(length(matrix_u[:,1])/3)):Int64(length(matrix_u[:,1])),1])
            minB2 = minimum(matrix_u[2 * Int64(round(length(matrix_u[:,1])/3)):Int64(length(matrix_u[:,1])),2])
            minB3 = minimum(matrix_u[2 * Int64(round(length(matrix_u[:,1])/3)):Int64(length(matrix_u[:,1])),3])

            ### Number of surviving species
            # final biomasses
            final_u = matrix_u[length(sol.t),:]

            # count number of species remaining that are above extinction threshold
            survived = count(final_u .> threshold)

            ### Push outputs to DataFrame
            push!(df2, [
                rate,
                as,
                dose,
                p.K[1],
                p.r[1],
                p.x[2],
                p.y[2,1],
                p.B0[2,1],
                maxB1,
                maxB2,
                maxB3,
                minB1,
                minB2,
                minB3,
                survived
            ])

            # print some stuff - see how the simulation is progressing
            println(("target rate = $rate", "asymptote = $as", "dose = $dose"))
        end
    end
end

### Save output
CSV.write("output_dfs/full_asymptote.csv", df2)

#### Varying inflection points ####
inflections = [10.0, 20.0, 30.0, 40.0, 50.0]

## empty dataframe to store results
df3 = DataFrame(
    target = [],
    xmid = [],
    dose = [],
    K = [],
    r = [],
    x_2 = [],
    y_21 = [],
    B0_21 = [],
    maxB1 = [],
    maxB2 = [],
    maxB3 = [],
    minB1 = [],
    minB2 = [],
    minB3 = [],
    sp_survival = []
)

for rate in targets
    for in in inflections
        for dose in doses
            # generate parameters specific to this combo of target rate, DR inflection points and dose
            p = ModelParameters_CG(param, T, I_K, Z, dose, rate, slope = 5.0, lower = 0.0, xmid = in)

            # Calculate starting biomasses
            u0 = [p.K[1]/2, p.K[1]/2, p.K[1]/2]

            # Define the problem
            prob = ODEProblem(BEFW, u0, tspan, p)

            # Solve the problem
            sol = solve(prob, maxiters = 1e7)

            ###### Output metrics ######
            # turn vector output of biomasses into matrix
            matrix_u = hcat(sol.u...)' # 56×3 adjoint(::Matrix{Float64}) with eltype Float64:

            ### Biomass extremes
            maxB1 = maximum(matrix_u[2 * Int64(round(length(matrix_u[:,1])/3)):Int64(length(matrix_u[:,1])),1])
            maxB2 = maximum(matrix_u[2 * Int64(round(length(matrix_u[:,1])/3)):Int64(length(matrix_u[:,1])),2])
            maxB3 = maximum(matrix_u[2 * Int64(round(length(matrix_u[:,1])/3)):Int64(length(matrix_u[:,1])),3])

            minB1 = minimum(matrix_u[2 * Int64(round(length(matrix_u[:,1])/3)):Int64(length(matrix_u[:,1])),1])
            minB2 = minimum(matrix_u[2 * Int64(round(length(matrix_u[:,1])/3)):Int64(length(matrix_u[:,1])),2])
            minB3 = minimum(matrix_u[2 * Int64(round(length(matrix_u[:,1])/3)):Int64(length(matrix_u[:,1])),3])

            ### Number of surviving species
            # final biomasses
            final_u = matrix_u[length(sol.t),:]

            # count number of species remaining that are above extinction threshold
            survived = count(final_u .> threshold)

            ### Push outputs to DataFrame
            push!(df3, [
                rate,
                in,
                dose,
                p.K[1],
                p.r[1],
                p.x[2],
                p.y[2,1],
                p.B0[2,1],
                maxB1,
                maxB2,
                maxB3,
                minB1,
                minB2,
                minB3,
                survived
            ])

            # print some stuff - see how the simulation is progressing
            println(("target rate = $rate", "inflection point = $in", "dose = $dose"))
        end
    end
end

### Save output
CSV.write("output_dfs/full_inflection.csv", df3)

##### Run the BEFW with chemical effects on functional response, varying the slope, asymptote and inflection point

#### Varying slopes

## empty dataframe to store results
df4 = DataFrame(
    target = [],
    slope = [],
    dose = [],
    K = [],
    r = [],
    x_2 = [],
    y_21 = [],
    B0_21 = [],
    maxB1 = [],
    maxB2 = [],
    maxB3 = [],
    minB1 = [],
    minB2 = [],
    minB3 = [],
    sp_survival = []
)

for sl in slopes
    for dose in doses
        # generate parameters specific to this combo of target rate, DR slope and dose
        p = ModelParameters_CG(param, T, I_K, Z, dose, "FR", slope = 5.0, lower = 0.0, xmid = 50.0)

        # Calculate starting biomasses
        u0 = [p.K[1]/2, p.K[1]/2, p.K[1]/2]

        # Define the problem
        prob = ODEProblem(BEFW_CG, u0, tspan, p)

        # Solve the problem
        sol = solve(prob, maxiters = 1e5)

        ###### Output metrics ######
        # turn vector output of biomasses into matrix
        matrix_u = hcat(sol.u...)' # 56×3 adjoint(::Matrix{Float64}) with eltype Float64:

        ### Biomass extremes
        maxB1 = maximum(matrix_u[2 * Int64(round(length(matrix_u[:,1])/3)):Int64(length(matrix_u[:,1])),1])
        maxB2 = maximum(matrix_u[2 * Int64(round(length(matrix_u[:,1])/3)):Int64(length(matrix_u[:,1])),2])
        maxB3 = maximum(matrix_u[2 * Int64(round(length(matrix_u[:,1])/3)):Int64(length(matrix_u[:,1])),3])

        minB1 = minimum(matrix_u[2 * Int64(round(length(matrix_u[:,1])/3)):Int64(length(matrix_u[:,1])),1])
        minB2 = minimum(matrix_u[2 * Int64(round(length(matrix_u[:,1])/3)):Int64(length(matrix_u[:,1])),2])
        minB3 = minimum(matrix_u[2 * Int64(round(length(matrix_u[:,1])/3)):Int64(length(matrix_u[:,1])),3])

        ### Number of surviving species
        # final biomasses
        final_u = matrix_u[length(sol.t),:]

        # count number of species remaining that are above extinction threshold
        survived = count(final_u .> threshold)

        ### Push outputs to DataFrame
        push!(df4, [
            "FR",
            sl,
            dose,
            p.K[1],
            p.r[1],
            p.x[2],
            p.y[2,1],
            p.B0[2,1],
            maxB1,
            maxB2,
            maxB3,
            minB1,
            minB2,
            minB3,
            survived
        ])

        # print some stuff - see how the simulation is progressing
        println(("slope = $sl", "dose = $dose"))
    end
end

max_stars_repo_name: hanamayall/simple-feeding-DR

max_issues_repo_name: hanamayall/simple-feeding-DR

max_forks_repo_name: hanamayall/simple-feeding-DR


text:
"""
SigmaGrid(;
        nxs,
        nys,
        nσs,
        σ = collect(range(0.0,length=nσs,stop=1.0)),
        ζ = one(eltype(σ)) .- σ
        quadrature_weights = [0.5;ones(nσs-2);0.5]/(nσs-1)
        η,
        θ,
        Φ,
        glen_b
        )

Construct a WAVI.jl SigmaGrid with size (nxs,nys,nσs)
SigmaGrid stores fields that are defined on the problem's three dimensional Sigma grid.

Keyword arguments
=================
- 'nxs': (required) Number of grid cells in x-direction in SigmaGrid (should be same as grid.nx)
        Note that we store the grid size here, even though it can be easily inferred from grid, to increase transparency in velocity solve.
- 'nys': (required) Number of grid cells in y-direction in SigmaGrid (should be same as grid.ny )
- 'nσs': (required) Number of grid cells in the vertical
- 'σ' : (required) Vertical levels, read in from grid
- 'ζ' : Reverse vertical levels
- 'quadrature_weights' : (required) weights associated with sigma levels, used in computation of integrals over thickness
- 'η' : (required) three dimensional viscosity field
- 'θ' : (required) three dimensional temperature field
- 'Φ' : (required) three dimensional damage field
- 'glen_b': (required) three dimensional field of glen_b values in viscosity calcluations
"""

@with_kw struct SigmaGrid{T <: Real, N <: Integer}
nxs :: N
nys :: N
nσs :: N
σ :: Vector{T}
ζ :: Vector{T} = one(eltype(σ)) .- σ ; @assert length(ζ) == nσs
quadrature_weights :: Vector{T} = [0.5;ones(nσs-2);0.5]/(nσs-1); @assert length(quadrature_weights) == nσs
η :: Array{T,3}; @assert size(η)==(nxs,nys,nσs)
θ :: Array{T,3}; @assert size(θ)==(nxs,nys,nσs)
Φ :: Array{T,3}; @assert size(Φ)==(nxs,nys,nσs)
glen_b :: Array{T,3} = glen_b.(θ,Φ); @assert size(glen_b)==(nxs,nys,nσs)
end


text:
int_rules_7_3_5 = @theory begin

    #= ::Subsection::Closed:: =#
    #= 7.3.5*u*(a+b*arctanh(c+d*x))^p =#
    @apply_utils Antiderivative((~(a') + ~(b') * atanh(~c + ~(d') * ~x)) ^ ~(p'), ~x) => (1 / ~d) * Subst(Antiderivative((~a + ~b * atanh(~x)) ^ ~p, ~x), ~x, ~c + ~d * ~x) <-- FreeQ([~a, ~b, ~c, ~d], ~x) && IGtQ(~p, 0)
    @apply_utils Antiderivative((~(a') + ~(b') * acoth(~c + ~(d') * ~x)) ^ ~(p'), ~x) => (1 / ~d) * Subst(Antiderivative((~a + ~b * acoth(~x)) ^ ~p, ~x), ~x, ~c + ~d * ~x) <-- FreeQ([~a, ~b, ~c, ~d], ~x) && IGtQ(~p, 0)
    @apply_utils Antiderivative((~(a') + ~(b') * atanh(~c + ~(d') * ~x)) ^ ~p, ~x) => Unintegrable((~a + ~b * atanh(~c + ~d * ~x)) ^ ~p, ~x) <-- FreeQ([~a, ~b, ~c, ~d, ~p], ~x) && Not(IGtQ(~p, 0))
    @apply_utils Antiderivative((~(a') + ~(b') * acoth(~c + ~(d') * ~x)) ^ ~p, ~x) => Unintegrable((~a + ~b * acoth(~c + ~d * ~x)) ^ ~p, ~x) <-- FreeQ([~a, ~b, ~c, ~d, ~p], ~x) && Not(IGtQ(~p, 0))
    @apply_utils Antiderivative((~(e') + ~(f') * ~x) ^ ~(m') * (~(a') + ~(b') * atanh(~c + ~(d') * ~x)) ^ ~(p'), ~x) => (1 / ~d) * Subst(Antiderivative(((~f * ~x) / ~d) ^ ~m * (~a + ~b * atanh(~x)) ^ ~p, ~x), ~x, ~c + ~d * ~x) <-- FreeQ([~a, ~b, ~c, ~d, ~e, ~f, ~m], ~x) && (EqQ(~d * ~e - ~c * ~f, 0) && IGtQ(~p, 0))
    @apply_utils Antiderivative((~(e') + ~(f') * ~x) ^ ~(m') * (~(a') + ~(b') * acoth(~c + ~(d') * ~x)) ^ ~(p'), ~x) => (1 / ~d) * Subst(Antiderivative(((~f * ~x) / ~d) ^ ~m * (~a + ~b * acoth(~x)) ^ ~p, ~x), ~x, ~c + ~d * ~x) <-- FreeQ([~a, ~b, ~c, ~d, ~e, ~f, ~m], ~x) && (EqQ(~d * ~e - ~c * ~f, 0) && IGtQ(~p, 0))
    @apply_utils Antiderivative((~(e') + ~(f') * ~x) ^ ~m * (~(a') + ~(b') * atanh(~c + ~(d') * ~x)) ^ ~(p'), ~x) => ((~e + ~f * ~x) ^ (~m + 1) * (~a + ~b * atanh(~c + ~d * ~x)) ^ ~p) / (~f * (~m + 1)) - ((~b * ~d * ~p) / (~f * (~m + 1))) * Antiderivative(((~e + ~f * ~x) ^ (~m + 1) * (~a + ~b * atanh(~c + ~d * ~x)) ^ (~p - 1)) / (1 - (~c + ~d * ~x) ^ 2), ~x) <-- FreeQ([~a, ~b, ~c, ~d, ~e, ~f], ~x) && (IGtQ(~p, 0) && ILtQ(~m, -1))
    @apply_utils Antiderivative((~(e') + ~(f') * ~x) ^ ~m * (~(a') + ~(b') * acoth(~c + ~(d') * ~x)) ^ ~(p'), ~x) => ((~e + ~f * ~x) ^ (~m + 1) * (~a + ~b * acoth(~c + ~d * ~x)) ^ ~p) / (~f * (~m + 1)) - ((~b * ~d * ~p) / (~f * (~m + 1))) * Antiderivative(((~e + ~f * ~x) ^ (~m + 1) * (~a + ~b * acoth(~c + ~d * ~x)) ^ (~p - 1)) / (1 - (~c + ~d * ~x) ^ 2), ~x) <-- FreeQ([~a, ~b, ~c, ~d, ~e, ~f], ~x) && (IGtQ(~p, 0) && ILtQ(~m, -1))
    @apply_utils Antiderivative((~(e') + ~(f') * ~x) ^ ~(m') * (~(a') + ~(b') * atanh(~c + ~(d') * ~x)) ^ ~(p'), ~x) => (1 / ~d) * Subst(Antiderivative(((~d * ~e - ~c * ~f) / ~d + (~f * ~x) / ~d) ^ ~m * (~a + ~b * atanh(~x)) ^ ~p, ~x), ~x, ~c + ~d * ~x) <-- FreeQ([~a, ~b, ~c, ~d, ~e, ~f, ~m], ~x) && IGtQ(~p, 0)
    @apply_utils Antiderivative((~(e') + ~(f') * ~x) ^ ~(m') * (~(a') + ~(b') * acoth(~c + ~(d') * ~x)) ^ ~(p'), ~x) => (1 / ~d) * Subst(Antiderivative(((~d * ~e - ~c * ~f) / ~d + (~f * ~x) / ~d) ^ ~m * (~a + ~b * acoth(~x)) ^ ~p, ~x), ~x, ~c + ~d * ~x) <-- FreeQ([~a, ~b, ~c, ~d, ~e, ~f, ~m], ~x) && IGtQ(~p, 0)
    @apply_utils Antiderivative((~(e') + ~(f') * ~x) ^ ~(m') * (~(a') + ~(b') * atanh(~c + ~(d') * ~x)) ^ ~p, ~x) => Unintegrable((~e + ~f * ~x) ^ ~m * (~a + ~b * atanh(~c + ~d * ~x)) ^ ~p, ~x) <-- FreeQ([~a, ~b, ~c, ~d, ~e, ~f, ~m, ~p], ~x) && Not(IGtQ(~p, 0))
    @apply_utils Antiderivative((~(e') + ~(f') * ~x) ^ ~(m') * (~(a') + ~(b') * acoth(~c + ~(d') * ~x)) ^ ~p, ~x) => Unintegrable((~e + ~f * ~x) ^ ~m * (~a + ~b * acoth(~c + ~d * ~x)) ^ ~p, ~x) <-- FreeQ([~a, ~b, ~c, ~d, ~e, ~f, ~m, ~p], ~x) && Not(IGtQ(~p, 0))
    @apply_utils Antiderivative(atanh(~c + ~(d') * ~x) / (~e + ~(f') * (~x) ^ ~(n')), ~x) => (1 / 2) * Antiderivative(log(1 + ~c + ~d * ~x) / (~e + ~f * (~x) ^ ~n), ~x) - (1 / 2) * Antiderivative(log((1 - ~c) - ~d * ~x) / (~e + ~f * (~x) ^ ~n), ~x) <-- FreeQ([~c, ~d, ~e, ~f], ~x) && RationalQ(~n)
    @apply_utils Antiderivative(acoth(~c + ~(d') * ~x) / (~e + ~(f') * (~x) ^ ~(n')), ~x) => (1 / 2) * Antiderivative(log((1 + ~c + ~d * ~x) / (~c + ~d * ~x)) / (~e + ~f * (~x) ^ ~n), ~x) - (1 / 2) * Antiderivative(log((-1 + ~c + ~d * ~x) / (~c + ~d * ~x)) / (~e + ~f * (~x) ^ ~n), ~x) <-- FreeQ([~c, ~d, ~e, ~f], ~x) && RationalQ(~n)
    @apply_utils Antiderivative(atanh(~c + ~(d') * ~x) / (~e + ~(f') * (~x) ^ ~n), ~x) => Unintegrable(atanh(~c + ~d * ~x) / (~e + ~f * (~x) ^ ~n), ~x) <-- FreeQ([~c, ~d, ~e, ~f, ~n], ~x) && Not(RationalQ(~n))
    @apply_utils Antiderivative(acoth(~c + ~(d') * ~x) / (~e + ~(f') * (~x) ^ ~n), ~x) => Unintegrable(acoth(~c + ~d * ~x) / (~e + ~f * (~x) ^ ~n), ~x) <-- FreeQ([~c, ~d, ~e, ~f, ~n], ~x) && Not(RationalQ(~n))
    @apply_utils Antiderivative((~(A') + ~(B') * ~x + ~(C') * (~x) ^ 2) ^ ~(q') * (~(a') + ~(b') * atanh(~c + ~(d') * ~x)) ^ ~(p'), ~x) => (1 / ~d) * Subst(Antiderivative((-(~C) / (~d) ^ 2 + (~C / (~d) ^ 2) * (~x) ^ 2) ^ ~q * (~a + ~b * atanh(~x)) ^ ~p, ~x), ~x, ~c + ~d * ~x) <-- FreeQ([~a, ~b, ~c, ~d, ~A, ~B, ~C, ~p, ~q], ~x) && (EqQ(~B * (1 - (~c) ^ 2) + 2 * ~A * ~c * ~d, 0) && EqQ(2 * ~c * ~C - ~B * ~d, 0))
    @apply_utils Antiderivative((~(A') + ~(B') * ~x + ~(C') * (~x) ^ 2) ^ ~(q') * (~(a') + ~(b') * acoth(~c + ~(d') * ~x)) ^ ~(p'), ~x) => (1 / ~d) * Subst(Antiderivative((~C / (~d) ^ 2 + (~C / (~d) ^ 2) * (~x) ^ 2) ^ ~q * (~a + ~b * acoth(~x)) ^ ~p, ~x), ~x, ~c + ~d * ~x) <-- FreeQ([~a, ~b, ~c, ~d, ~A, ~B, ~C, ~p, ~q], ~x) && (EqQ(~B * (1 - (~c) ^ 2) + 2 * ~A * ~c * ~d, 0) && EqQ(2 * ~c * ~C - ~B * ~d, 0))
    @apply_utils Antiderivative((~(e') + ~(f') * ~x) ^ ~(m') * (~(A') + ~(B') * ~x + ~(C') * (~x) ^ 2) ^ ~(q') * (~(a') + ~(b') * atanh(~c + ~(d') * ~x)) ^ ~(p'), ~x) => (1 / ~d) * Subst(Antiderivative(((~d * ~e - ~c * ~f) / ~d + (~f * ~x) / ~d) ^ ~m * (-(~C) / (~d) ^ 2 + (~C / (~d) ^ 2) * (~x) ^ 2) ^ ~q * (~a + ~b * atanh(~x)) ^ ~p, ~x), ~x, ~c + ~d * ~x) <-- FreeQ([~a, ~b, ~c, ~d, ~e, ~f, ~A, ~B, ~C, ~m, ~p, ~q], ~x) && (EqQ(~B * (1 - (~c) ^ 2) + 2 * ~A * ~c * ~d, 0) && EqQ(2 * ~c * ~C - ~B * ~d, 0))
    @apply_utils Antiderivative((~(e') + ~(f') * ~x) ^ ~(m') * (~(A') + ~(B') * ~x + ~(C') * (~x) ^ 2) ^ ~(q') * (~(a') + ~(b') * acoth(~c + ~(d') * ~x)) ^ ~(p'), ~x) => (1 / ~d) * Subst(Antiderivative(((~d * ~e - ~c * ~f) / ~d + (~f * ~x) / ~d) ^ ~m * (-(~C) / (~d) ^ 2 + (~C / (~d) ^ 2) * (~x) ^ 2) ^ ~q * (~a + ~b * acoth(~x)) ^ ~p, ~x), ~x, ~c + ~d * ~x) <-- FreeQ([~a, ~b, ~c, ~d, ~e, ~f, ~A, ~B, ~C, ~m, ~p, ~q], ~x) && (EqQ(~B * (1 - (~c) ^ 2) + 2 * ~A * ~c * ~d, 0) && EqQ(2 * ~c * ~C - ~B * ~d, 0))
end

max_stars_repo_path: src/rules/7 Inverse hyperbolic functions/7.3 Inverse hyperbolic tangent/7.3.5 u (a+b arctanh(c+d x))^p.jl


text:
"""
    simulateGS(p::Int,q::Int,file::String)

Optimizes the model at parameters (p,q) looking for a global minimum using the :cplex optimizer, and appends the result to ```file```.

"""
function simulateGS(p::Int,q::Int,file::String)
    open(file, "a") do io
        writedlm(io, [modelFindMinimum(p,q,optimizer=:cplex)], '|')
    end
end

"""
    simulateMinima(p::Int,q::Int,iter::Int,file::String)

Optimizes the model at parameters (p,q) looking for a local minimum using the :gradient optimizer, and appends the result to ```file```.
It iterates the search ```iter``` times starting from random initial conditions.
"""
function simulateMinima(p::Int,q::Int,iter::Int,file::String)
    result = []
    for i in 1:iter
        push!(result, modelFindMinimum(p,q))
    end
    open(file, "a") do io
        writedlm(io, result, '|')
    end
end

"""
    heatAndCool(p::Int,q::Int,ts::Vector{T};) where T <: Real
    heatAndCool(p::Int,q::Int,ts::Vector{T};eta::Real) where T <: Real

It performs an heating and cooling experiment of the model at parameters (p,q) by iterating through the given temperatures ```ts```, and minimizing locally the free energy using the :gradient optimizer.

Before being used as initial condition for a new minimization, the final condition of the previous minimization is perturbed by adding a noise of strength ```eta``` to avoid memory effects. ```eta``` defaults to ```1e-2```.

The initial condition is given by ```startingCondition```, which defaults to a random initial condition if nothing is specified.

"""
function heatAndCool(p::Int,q::Int,ts::Vector{T};
                     eta::Real=1e-2,
                     startingCondition::Vector{T} = rand(Float64,modelDimension(q)) .|> u -> 2*u-1
                    ) where T <: Real

    # record energies
    result = []
    current = startingCondition

    # change the temperature of the system and minimize it locally
    for t in ts
        current = modelFindMinimum(p,q; optimizer=:gradient,
                            temperature=t,
                            startingCondition=addBoxNoise(current;eta=eta)
                           )[5]
        push!(result,[t,current])
    end

    return result
end


text:
function icd10_add_dots(icd10::AbstractString)
    ### Input is too short for a dot
    r1 = r"^([A-Z][A-Z\d][A-Z\d])$"
    m1 = match(r1, icd10)
    if m1 !== nothing
        return "$(m1[1])"::String
    end

    ### Input contains a dot
    r2 = r"^([A-Z][A-Z\d][A-Z\d])\.([A-Z\d][A-Z\d]?[A-Z\d]?[A-Z\d]?)$"
    m2 = match(r2, icd10)
    if m2 !== nothing
        return "$(m2[1]).$(m2[2])"::String
    end

    ### Input does not contain a dot
    r3 = r"^([A-Z][A-Z\d][A-Z\d])([A-Z\d][A-Z\d]?[A-Z\d]?[A-Z\d]?)$"
    m3 = match(r3, icd10)
    if m3 !== nothing
        return "$(m3[1]).$(m3[2])"::String
    end

    msg = "Invalid ICD-10-CM code: $(icd10)"
    throw(ErrorException(msg))
end

function icd10_remove_dots(icd10::AbstractString)
    ### Input is too short for a dot
    r1 = r"^([A-Z][A-Z\d][A-Z\d])$"
    m1 = match(r1, icd10)
    if m1 !== nothing
        return "$(m1[1])"::String
    end

    ### Input contains a dot
    r2 = r"^([A-Z][A-Z\d][A-Z\d])\.([A-Z\d][A-Z\d]?[A-Z\d]?[A-Z\d]?)$"
    m2 = match(r2, icd10)
    if m2 !== nothing
        return "$(m2[1])$(m2[2])"::String
    end

    ### Input does not contain a dot
    r3 = r"^([A-Z][A-Z\d][A-Z\d])([A-Z\d][A-Z\d]?[A-Z\d]?[A-Z\d]?)$"
    m3 = match(r3, icd10)
    if m3 !== nothing
        return "$(m3[1])$(m3[2])"::String
    end

    msg = "Invalid ICD-10-CM code: $(icd10)"
    throw(ErrorException(msg))
end


text:
module RaftRPC
  include("raft_rpc_pb.jl")
end


text:
# checks whether the type is categorical value
iscatvalue(::Type) = false
iscatvalue(::Type{Union{}}) = false # prevent incorrect dispatch to Type{<:CatValue} method
iscatvalue(::Type{<:CatValue}) = true
iscatvalue(x::Any) = iscatvalue(typeof(x))

leveltype(::Type{<:CategoricalValue{T}}) where {T} = T
leveltype(::Type{<:CategoricalString}) = String
leveltype(::Type) = throw(ArgumentError("Not a categorical value type"))
leveltype(x::Any) = leveltype(typeof(x))

# integer type of category reference codes used by categorical value
reftype(::Type{<:CatValue{R}}) where {R} = R
reftype(x::Any) = reftype(typeof(x))

pool(x::CatValue) = x.pool
level(x::CatValue) = x.level

# extract the type of the original value from array eltype `T`
unwrap_catvaluetype(::Type{T}) where {T} = T
unwrap_catvaluetype(::Type{T}) where {T >: Missing} =
    Union{unwrap_catvaluetype(nonmissingtype(T)), Missing}
unwrap_catvaluetype(::Type{Union{}}) = Union{} # prevent incorrect dispatch to T<:CatValue method
unwrap_catvaluetype(::Type{Any}) = Any # prevent recursion in T>:Missing method
unwrap_catvaluetype(::Type{T}) where {T <: CatValue} = leveltype(T)

# get the categorical value type given value type `T` and reference type `R`
catvaluetype(::Type{T}, ::Type{R}) where {T >: Missing, R} =
    catvaluetype(nonmissingtype(T), R)
catvaluetype(::Type{T}, ::Type{R}) where {T <: CatValue, R} =
    catvaluetype(leveltype(T), R)
catvaluetype(::Type{Any}, ::Type{R}) where {R} =
    CategoricalValue{Any, R}  # prevent recursion in T>:Missing method
catvaluetype(::Type{T}, ::Type{R}) where {T, R} =
    CategoricalValue{T, R}
catvaluetype(::Type{<:AbstractString}, ::Type{R}) where {R} =
    CategoricalString{R}
# to prevent incorrect dispatch to T<:CatValue method
catvaluetype(::Type{Union{}}, ::Type{R}) where {R} = CategoricalValue{Union{}, R}

# get the categorical value type given value type `T`
catvaluetype(::Type{T}) where {T >: Missing} = catvaluetype(nonmissingtype(T))
catvaluetype(::Type{T}) where {T <: CatValue} = catvaluetype(leveltype(T))
catvaluetype(::Type{Any}) = CategoricalValue{Any}  # prevent recursion in T>:Missing method
catvaluetype(::Type{T}) where {T} = CategoricalValue{T}
catvaluetype(::Type{<:AbstractString}) = CategoricalString
# to prevent incorrect dispatch to T<:CatValue method
catvaluetype(::Type{Union{}}) where {R} = CategoricalValue{Union{}}

Base.get(x::CatValue) = index(pool(x))[level(x)]
order(x::CatValue) = order(pool(x))[level(x)]

# creates categorical value for `level` from the `pool`
# The result is of type `C` that has "categorical value" trait
catvalue(level::Integer, pool::CategoricalPool{T, R, C}) where {T, R, C} =
    C(convert(R, level), pool)

Base.promote_rule(::Type{C}, ::Type{T}) where {C <: CatValue, T} = promote_type(leveltype(C), T)
Base.promote_rule(::Type{C1}, ::Type{Union{C2, Missing}}) where {C1 <: CatValue, C2 <: CatValue} =
    Union{promote_type(C1, C2), Missing}
# To fix ambiguities with definitions from Base
Base.promote_rule(::Type{C}, ::Type{Missing}) where {C <: CatValue} = Union{C, Missing}
Base.promote_rule(::Type{C}, ::Type{Any}) where {C <: CatValue} = Any

Base.promote_rule(::Type{CategoricalValue{S, R1}},
                  ::Type{CategoricalValue{T, R2}}) where {S, T, R1<:Integer, R2<:Integer} =
    CategoricalValue{promote_type(S, T), promote_type(R1, R2)}
Base.promote_rule(::Type{CategoricalString{R1}},
                  ::Type{CategoricalString{R2}}) where {R1<:Integer, R2<:Integer} =
    CategoricalString{promote_type(R1, R2)}
Base.promote_rule(::Type{C1}, ::Type{C2}) where
    {R1<:Integer, R2<:Integer, C1<:CatValue{R1}, C2<:CatValue{R2}} =
    catvaluetype(promote_type(leveltype(C1), leveltype(C2)), promote_type(R1, R2))
Base.promote_rule(::Type{C1}, ::Type{C2}) where {C1<:CatValue, C2<:CatValue} =
    catvaluetype(promote_type(leveltype(C1), leveltype(C2)))

Base.convert(::Type{Ref}, x::CatValue) = RefValue{leveltype(x)}(x)
Base.convert(::Type{String}, x::CatValue) = convert(String, get(x))
Base.convert(::Type{Any}, x::CatValue) = x

# Defined separately to avoid ambiguities
Base.convert(::Type{AbstractString}, x::CategoricalString) = x
Base.convert(::Type{T}, x::T) where {T <: CatValue} = x
Base.convert(::Type{Union{T, Missing}}, x::T) where {T <: CatValue} = x
Base.convert(::Type{Union{T, Nothing}}, x::T) where {T <: CatValue} = x
# General fallbacks
Base.convert(::Type{S}, x::T) where {S, T <: CatValue} =
    T <: S ? x : convert(S, get(x))
Base.convert(::Type{Union{S, Missing}}, x::T) where {S, T <: CatValue} =
    T <: Union{S, Missing} ? x : convert(Union{S, Missing}, get(x))
Base.convert(::Type{Union{S, Nothing}}, x::T) where {S, T <: CatValue} =
    T <: Union{S, Nothing} ? x : convert(Union{S, Nothing}, get(x))

(::Type{T})(x::T) where {T <: CatValue} = x

Base.Broadcast.broadcastable(x::CatValue) = Ref(x)

if VERSION >= v"0.7.0-DEV.2797"
    function Base.show(io::IO, x::CatValue)
        if nonmissingtype(get(io, :typeinfo, Any)) === nonmissingtype(typeof(x))
            print(io, repr(x))
        elseif isordered(pool(x))
            @printf(io, "%s %s (%i/%i)",
                    typeof(x), repr(x),
                    order(x), length(pool(x)))
        else
            @printf(io, "%s %s", typeof(x), repr(x))
        end
    end
else
    function Base.show(io::IO, x::CatValue)
        if get(io, :compact, false)
            print(io, repr(x))
        elseif isordered(pool(x))
            @printf(io, "%s %s (%i/%i)",
                    typeof(x), repr(x),
                    order(x), length(pool(x)))
        else
            @printf(io, "%s %s", typeof(x), repr(x))
        end
    end
end

Base.print(io::IO, x::CatValue) = print(io, get(x))
Base.repr(x::CatValue) = repr(get(x))

@inline function Base.:(==)(x::CatValue, y::CatValue)
    if pool(x) === pool(y)
        return level(x) == level(y)
    else
        return get(x) == get(y)
    end
end

Base.:(==)(::CatValue, ::Missing) = missing
Base.:(==)(::Missing, ::CatValue) = missing

# To fix ambiguities with Base
Base.:(==)(x::CatValue, y::WeakRef) = get(x) == y
Base.:(==)(x::WeakRef, y::CatValue) = y == x

Base.:(==)(x::CatValue, y::AbstractString) = get(x) == y
Base.:(==)(x::AbstractString, y::CatValue) = y == x

Base.:(==)(x::CatValue, y::Any) = get(x) == y
Base.:(==)(x::Any, y::CatValue) = y == x

@inline function Base.isequal(x::CatValue, y::CatValue)
    if pool(x) === pool(y)
        return level(x) == level(y)
    else
        return isequal(get(x), get(y))
    end
end

Base.isequal(x::CatValue, y::Any) = isequal(get(x), y)
Base.isequal(x::Any, y::CatValue) = isequal(y, x)

Base.isequal(::CatValue, ::Missing) = false
Base.isequal(::Missing, ::CatValue) = false

Base.in(x::CatValue, y::AbstractRange{T}) where {T<:Integer} = get(x) in y

Base.hash(x::CatValue, h::UInt) = hash(get(x), h)

# Method defined even on unordered values so that sort() works
function Base.isless(x::CatValue, y::CatValue)
    if pool(x) !== pool(y)
        throw(ArgumentError("CategoricalValue objects with different pools cannot be tested for order"))
    else
        return order(x) < order(y)
    end
end

Base.isless(x::CatValue, y) = order(x) < order(x.pool[get(x.pool, y)])
Base.isless(x::CatValue, y::AbstractString) = order(x) < order(x.pool[get(x.pool, y)])
Base.isless(::CatValue, ::Missing) = true
Base.isless(y, x::CatValue) = order(x.pool[get(x.pool, y)]) < order(x)
Base.isless(y::AbstractString, x::CatValue) = order(x.pool[get(x.pool, y)]) < order(x)
Base.isless(::Missing, ::CatValue) = false

function Base.:<(x::CatValue, y::CatValue)
    if pool(x) !== pool(y)
        throw(ArgumentError("CategoricalValue objects with different pools cannot be tested for order"))
    elseif !isordered(pool(x)) # !isordered(pool(y)) is implied by pool(x) === pool(y)
        throw(ArgumentError("Unordered CategoricalValue objects cannot be tested for order using <. Use isless instead, or call the ordered! function on the parent array to change this"))
    else
        return order(x) < order(y)
    end
end

function Base.:<(x::CatValue, y)
    if !isordered(pool(x))
        throw(ArgumentError("Unordered CategoricalValue objects cannot be tested for order using <. Use isless instead, or call the ordered! function on the parent array to change this"))
    else
        return order(x) < order(x.pool[get(x.pool, y)])
    end
end

Base.:<(x::CatValue, y::AbstractString) = invoke(<, Tuple{CatValue, Any}, x, y)
Base.:<(::CatValue, ::Missing) = missing

function Base.:<(y, x::CatValue)
    if !isordered(pool(x))
        throw(ArgumentError("Unordered CategoricalValue objects cannot be tested for order using <. Use isless instead, or call the ordered! function on the parent array to change this"))
    else
        return order(x.pool[get(x.pool, y)]) < order(x)
    end
end

Base.:<(y::AbstractString, x::CatValue) = invoke(<, Tuple{Any, CatValue}, y, x)
Base.:<(::Missing, ::CatValue) = missing

# JSON of CatValue is JSON of the value it refers to
JSON.lower(x::CatValue) = JSON.lower(get(x))
DataAPI.defaultarray(::Type{CategoricalString{R}}, N) where {R} =
  CategoricalArray{String, N, R}
DataAPI.defaultarray(::Type{Union{CategoricalString{R}, Missing}}, N) where {R} =
  CategoricalArray{Union{String, Missing}, N, R}
DataAPI.defaultarray(::Type{CategoricalValue{T, R}}, N) where {T, R} =
  CategoricalArray{T, N, R}
DataAPI.defaultarray(::Type{Union{CategoricalValue{T, R}, Missing}}, N) where {T, R} =
  CategoricalArray{Union{T, Missing}, N, R}

# Efficient method equivalent to the fallback
Base.string(x::CategoricalString) = get(x)
Base.String(x::CategoricalString) = get(x)


text:
@testset "SGHMC" begin
    #-------------------
    # test inplace addition vs reference
    #-------------------
    function sample_reference!(s::SGHMCState, grad)
        #increase iteration counter
        s.t += 1

        # improve readability
        x = s.x
        p = s.p
        D = s.D
        t = s.t
        stepsize = s.stepsize(t)
        m = s.mass

        p[:] += - stepsize*(D .* p) ./ m + stepsize*grad(x) + sqrt(2*stepsize.*D).*randn(length(x))
        x[:] += stepsize*p ./ m
        s
    end

    srand(1)
    s = SGHMCState(ones(2),0.01)
    srand(1)
    ref_s = SGHMCState(ones(2),0.01)
    sghmcgrad(x) = x.*x

    srand(2)
    sample!(s,sghmcgrad)

    srand(2)
    sample_reference!(ref_s,sghmcgrad)

    @test s.x == ref_s.x

    #------------------
    # Gaussian sanity check.
    #------------------
    dm = GaussianMixtureModel([0.],[1.],[1.])

    gradient = DataModel.getgrad(dm)
    s = SGHMCState(zeros(1),0.1)

    x_avg = [0.]
    x2_avg = [0.]
    for i in 1:50000
        sample!(s,gradient)
        x_avg =(i-1)/i*x_avg + s.x/i
        x2_avg = (i-1)/i*x2_avg + s.x.^2/i
    end

    @test abs(x_avg-0)[1] < 0.1
    @test abs(x2_avg-1)[1] < 0.2
end


text:
export LinearVApproximator, LinearQApproximator

using LinearAlgebra: dot

"""
    LinearVApproximator(weights::Array{Float64, N}) -> LinearVApproximator{N}
Use the weighted sum to represent the estimation of a state.
The state is expected to have the same length with `weights`.
See also [`LinearQApproximator`](@ref)
"""
struct LinearVApproximator{N} <: AbstractApproximator
    weights::Array{Float64,N}
end

RLCore.ApproximatorStyle(::LinearVApproximator) = V_APPROXIMATOR

# TODO: support Vector
(V::LinearVApproximator)(s) = dot(s, V.weights)

function RLBase.update!(V::LinearVApproximator, correction::Pair)
    s, e = correction
    V.weights .+= s .* e
end

"""
    LinearQApproximator(weights::Vector{Float64}, feature_func::F, actions::Vector{Int}) -> LinearQApproximator{F}
Use weighted sum to represent the estimation given a state and an action.
# Fields
- `weights::Vector{Float64}`: the weight of each feature.
- `feature_func::Function`: decide how to generate a feature vector of `length(weights)` given a state and an action as parameters.
- `actions::Vector{Int}`: all possible actions.
See also [`LinearVApproximator`](@ref).
"""
Base.@kwdef struct LinearQApproximator{F} <: AbstractApproximator
    weights::Vector{Float64}
    feature_func::F
    actions::Vector{Int}
end

RLCore.ApproximatorStyle(::LinearQApproximator) = Q_APPROXIMATOR

(Q::LinearQApproximator)(s, a::Int) = dot(Q.weights, Q.feature_func(s, a))

(Q::LinearQApproximator)(s) = [Q(s, a) for a in Q.actions]

function RLBase.update!(Q::LinearQApproximator, correction::Pair)
    (s, a), e = correction
    xs = Q.feature_func(s, a)
    Q.weights .+= xs .* e
end


text:
module TLWE
    using ..Key
    using ..Utils
    using LinearAlgebra
    function tlweSymEncrypt(p::Float64,α::Float64,key::Array{UInt32,1})
        a::Array{UInt32,1} = rand(0:1,length(key))
        b::UInt32 = Utils.gaussian32(p,α)[1] + a ⋅ key
        return append!(a,b)
    end
    function tlweSymDecrypt(c::Array{UInt32,1},key::Array{UInt32,1})
        return 1 - ((c[end] - c[1:end-1] ⋅ key) > UInt32(2)^31)
    end

    function bootsSymEncrypt(p::Array{Int32,1},sk::Key.SecretKey)
        return [tlweSymEncrypt(x, sk.params.α, sk.key.tlwe) for x in (2*p .- 1)*Key.μ]
    end
    function bootsSymDecrypt(c::Array{Array{UInt32,1},1},sk::Key.SecretKey)
        return [tlweSymDecrypt(x,sk.key.tlwe) for x in c]
    end
end


text:
using Test

include("emoji-times.jl")

@testset "Addition" begin
    @test 🕐 + 🕒 == 🕓
    @test 🕒 + 🕧 == 🕞
end

@testset "Addition overflow" begin
    @test 🕚 + 🕚 == 🕙
end

@testset "Subtraction" begin
    @test 🕗 - 🕔 == 🕒
    @test 🕤 - 🕞 == 🕕
end

@testset "Subtraction underflow" begin
    @test 🕓 - 🕘 == 🕖
end

@testset "0-element" begin
    @test 🕚 + 🕛 == 🕚
    @test 🕚 - 🕛 == 🕚
end