text:
"""
    setEvaluator(fitnessFunctions::Array{FitnessFunction}; genj::GeneticJuliaStructure = GenJ,
                globalFitnessFunction::FitnessFunction = FitnessFunction(noFunc),
                compareFunction::Function=noFunc, compareFunctionArgs::String = "rawW")

Sets the evaluator for `genj`, receiving an array with structures for all the fitness functions
(See [FitnessFunction](@ref)). This function constructs a structure of type
`Evaluator` and adds it to the main structure.

# Arguments
- `fitnessFunctions::Array{FitnessFunction}`: each function must evaluate an individual,
    receiving the individual's genotype and returning a fitness.
    (For information about genotypes look [`Genotype`](@ref))

# Keyword Arguments
- `genj::GeneticJuliaStructure = GenJ`: the main structure. For `code user`: don't
    modify the default value unless multiple experiments are going to be run.
- `globalFitnessFunction::FitnessFunction = FitnessFunction(noFunc)`: second-phase
    fitness function that is going to be evaluated once the other 'raw' fitness
    functions have been. For example, for [SPEA](@ref) algorithm, this fitness
    function serves to assign each individual its fitness level in the pareto front.
- `compareFunction::Function=noFunc`: function used to compare between individuals.
- `compareFunctionArgs::String = "rawW"`: tells the evaluator what is going to be
    compared. "raw" is for fitness comparison, "rawW" is for weighted fitness comparison,
    "global" is for second-phase fitness comparison, "globalW" is for weighted
    second-phase fitness comparison, "both" is for both fitness and second-phase
    fitness comparison and "bothW" is for both weighted fitness and weighted
    second-phase fitness comparison.

See also: [`Evaluator`](@ref), [FitnessFunction](@ref)
"""
function setEvaluator(fitnessFunctions::Array{FitnessFunction}; genj::GeneticJuliaStructure = GenJ,
                      globalFitnessFunction::FitnessFunction = FitnessFunction(noFunc),
                      compareFunction::Function=noFunc, compareFunctionArgs::String = "rawW")

    if findfirst(x-> x==compareFunctionArgs, ["raw","rawW", "global", "globalW", "both", "bothW"]) == nothing
        error("Invalid compareFunctionArgs ($compareFunctionArgs), it should be either raw, rawW, global, globalW, both or bothW")
    end
    compareFitness = getParametrizedCompareFunction(compareFunction, compareFunctionArgs, getWeights(fitnessFunctions))

    genj._evaluator = Evaluator(fitnessFunctions, globalFitnessFunction, compareFitness, compareFunctionArgs)
end # function

"""
    setEvaluator(args)

Sets the evaluator for `genj`, receiving a single structure for the fitness function
(See [FitnessFunction](@ref)). This function constructs a structure of type
`Evaluator` and adds it to the main structure.

# Arguments
- `fitnessFunction::FitnessFunction`: the fitness function, which must receive
    the individual's genotype and return a fitness. (For information about
    genotypes look [`Genotype`](@ref))

# Keyword Arguments
- `genj::GeneticJuliaStructure = GenJ`: the main structure. For `code user`: don't
    modify the default value unless multiple experiments are going to be run.
- `globalFitnessFunction::FitnessFunction = FitnessFunction(noFunc)`: second-phase
    fitness function that is going to be evaluated once the other 'raw' fitness
    functions have been. For example, for [SPEA](@ref) algorithm, this fitness
    function serves to assign each individual its fitness level in the pareto front.
- `compareFunction::Function=noFunc`: function used to compare between individuals.
- `compareFunctionArgs::String = "rawW"`: tells the evaluator what is going to be
    compared. "raw" is for fitness comparison, "rawW" is for weighted fitness comparison,
    "global" is for second-phase fitness comparison, "globalW" is for weighted
    second-phase fitness comparison, "both" is for both fitness and second-phase
    fitness comparison and "bothW" is for both weighted fitness and weighted
    second-phase fitness comparison.

See also: [`Evaluator`](@ref), [FitnessFunction](@ref)
"""
function setEvaluator(fitnessFunction::FitnessFunction; genj::GeneticJuliaStructure = GenJ,
                      globalFitnessFunction::FitnessFunction = FitnessFunction(noFunc),
                      compareFunction::Function=noFunc, compareFunctionArgs::String = "rawW")
    setEvaluator([fitnessFunction], genj=genj, globalFitnessFunction=globalFitnessFunction,
                 compareFunction=compareFunction, compareFunctionArgs=compareFunctionArgs)
end # function

"""
    setEvaluator(functions::Array{Function}, weights::Array{Real}; genj::GeneticJuliaStructure = GenJ,
                 globalFitnessFunction::FitnessFunction = FitnessFunction(nullGlobalFunction),
                 compareFunction::Function=noFunc, compareFunctionArgs::String = "rawW")

Sets the evaluator for `genj`, receiving an array with a function per fitness
and an array with a weight per fitness. This function constructs a structure of
type `Evaluator` and adds it to the main structure.

# Arguments
- `functions::Array{Function}`: each function must evaluate an individual,
    receiving the individual's genotype and returning a fitness.
    (For information about genotypes look [`Genotype`](@ref))
- `weights::Array{Real}`: each weight, in absolute value, represents the
    importance of its fitness in the system, whereas the sign is representative
    of maximization (positive) or minimization (negative).

# Keyword Arguments
- `genj::GeneticJuliaStructure = GenJ`: the main structure. For `code user`: don't
    modify the default value unless multiple experiments are going to be run.
- `globalFitnessFunction::FitnessFunction = FitnessFunction(noFunc)`: second-phase
    fitness function that is going to be evaluated once the other 'raw' fitness
    functions have been. For example, for [SPEA](@ref) algorithm, this fitness
    function serves to assign each individual its fitness level in the pareto front.
- `compareFunction::Function=noFunc`: function used to compare between individuals.
- `compareFunctionArgs::String = "rawW"`: tells the evaluator what is going to be
    compared. "raw" is for fitness comparison, "rawW" is for weighted fitness comparison,
    "global" is for second-phase fitness comparison, "globalW" is for weighted
    second-phase fitness comparison, "both" is for both fitness and second-phase
    fitness comparison and "bothW" is for both weighted fitness and weighted
    second-phase fitness comparison.

See also: [`Evaluator`](@ref), [FitnessFunction](@ref)
"""
function setEvaluator(functions::Array{Function}, weights::Array{Real}; genj::GeneticJuliaStructure = GenJ,
                      globalFitnessFunction::FitnessFunction = FitnessFunction(nullGlobalFunction),
                      compareFunction::Function=noFunc, compareFunctionArgs::String = "rawW")
    if length(functions) != length(weights)
        error("The number of fitness functions must be the same as the number of weights")
    end
    fitnessFunctions = Array{FitnessFunction}(undef, length(functions))

    for i=1:length(functions)
        fitnessFunctions[i] = FitnessFunction(functions[i], weight = weights[i])
    end

    setEvaluator(fitnessFunctions, genj=genj, globalFitnessFunction=globalFitnessFunction, compareFunction=compareFunction, compareFunctionArgs=compareFunctionArgs)
end # function

"""
    setEvaluator(functions::Array{Function}; genj::GeneticJuliaStructure = GenJ,
                 globalFitnessFunction::FitnessFunction = FitnessFunction(nullGlobalFunction),
                 compareFunction::Function=noFunc, compareFunctionArgs::String = "rawW")

Sets the evaluator for `genj`, receiving an array with a function per fitness,
stablishing all the weights to one, meaning that all the fitnesses need to be
maximized and have the same importance. This function constructs a structure of
type `Evaluator` and adds it to the main structure.

# Arguments
- `functions::Array{Function}`: each function must evaluate an individual,
    receiving the individual's genotype and returning a fitness.
    (For information about genotypes look [`Genotype`](@ref))

# Keyword Arguments
- `genj::GeneticJuliaStructure = GenJ`: the main structure. For `code user`: don't modify
    the default value unless multiple experiments are going to be run.
- `globalFitnessFunction::FitnessFunction = FitnessFunction(noFunc)`: second-phase
    fitness function that is going to be evaluated once the other 'raw' fitness
    functions have been. For example, for [SPEA](@ref) algorithm, this fitness
    function serves to assign each individual its fitness level in the pareto front.
- `compareFunction::Function=noFunc`: function used to compare between individuals.
- `compareFunctionArgs::String = "rawW"`: tells the evaluator what is going to be
    compared. "raw" is for fitness comparison, "rawW" is for weighted fitness comparison,
    "global" is for second-phase fitness comparison, "globalW" is for weighted
    second-phase fitness comparison, "both" is for both fitness and second-phase
    fitness comparison and "bothW" is for both weighted fitness and weighted
    second-phase fitness comparison.

See also: [`Evaluator`](@ref), [FitnessFunction](@ref)
"""
function setEvaluator(functions::Array{Function}; genj::GeneticJuliaStructure = GenJ,
                      globalFitnessFunction::FitnessFunction = FitnessFunction(nullGlobalFunction),
                      compareFunction::Function=noFunc, compareFunctionArgs::String = "rawW")
    # The weights are stablished to one, which mean, the fitnesses are going to be maximized
    weights = ones(Real, length(functions))
    setEvaluator(functions, weights, genj=genj, globalFitnessFunction=globalFitnessFunction,
                 compareFunction=compareFunction, compareFunctionArgs=compareFunctionArgs)
end # function

"""
    setEvaluator(fitnessFunction1::Function, weights1::Real, varArgs...; genj::GeneticJuliaStructure = GenJ,
                 globalFitnessFunction::FitnessFunction = FitnessFunction(noFunc),
                 compareFunction::Function=noFunc, compareFunctionArgs::String = "rawW")

Sets the evaluator for `genj`, receiving pair of functions and a related weights per fitness.
This function constructs a structure of type `Evaluator` and adds it to the main structure.

# Arguments
- `fitnessFunction1::Function`: function for the first fitness that must
    evaluate an individual, receiving the individual's genotype and
    returning a fitness. (For information about genotypes look [`genotype`](@ref))
- `weights1::Real`: weight for the first fitness, in absolute value, it
    represents the importance of its fitness in the system, whereas the sign is
    representative of maximization (positive) or minimization (negative).
- `varArgs`: pairs of fitness function and asociate weights. See `"Examples"`. It can be empty.

# Keyword Arguments
- `genj::GeneticJuliaStructure = GenJ`: the main structure. For `code user`: don't modify
    the default value unless multiple experiments are going to be run.
- `globalFitnessFunction::FitnessFunction = FitnessFunction(noFunc)`: second-phase
    fitness function that is going to be evaluated once the other 'raw' fitness
    functions have been. For example, for [SPEA](@ref) algorithm, this fitness
    function serves to assign each individual its fitness level in the pareto front.
- `compareFunction::Function=noFunc`: function used to compare between individuals.
- `compareFunctionArgs::String = "rawW"`: tells the evaluator what is going to be
    compared. "raw" is for fitness comparison, "rawW" is for weighted fitness comparison,
    "global" is for second-phase fitness comparison, "globalW" is for weighted
    second-phase fitness comparison, "both" is for both fitness and second-phase
    fitness comparison and "bothW" is for both weighted fitness and weighted
    second-phase fitness comparison.

# Examples
```jldoctest
julia> setEvaluator(myComputeFitnessFunction, 1.0)

julia> setEvaluator(myFitnesFunction1, 1.0, myFitnessFunction2, 0.7, myFitnessFunction3, -1.0)
```

See also: [`Evaluator`](@ref), [FitnessFunction](@ref)
"""
function setEvaluator(fitnessFunction1::Function, weights1::Real, varArgs...; genj::GeneticJuliaStructure = GenJ,
                      globalFitnessFunction::FitnessFunction = FitnessFunction(noFunc),
                      compareFunction::Function=noFunc, compareFunctionArgs::String = "rawW")
    len = length(varArgs)
    if isodd(len)
        error("The number of fitness functions must be the same as the number of weights")
    else
        functions = Array{Function}(undef,1+div(len, 2))
        weights = Array{Real}(undef,1+div(len, 2))

        functions[1] = fitnessFunction1
        weights[1] = weights1

        for i=1:len
            if isodd(i)
                if typeof(varArgs[i]) <: Function
                    functions[div(i,2)+2] = varArgs[i]
                else
                    error("The odd positions in the arguments of the evaluator must be fitness functions")
                end
            else
                if typeof(varArgs[i]) <: Real
                    weights[div(i,2)+1] = varArgs[i]
                else
                    error("The even positions in the arguments of the evaluator must be weights")
                end
            end
        end

        setEvaluator(functions, weights, genj=genj, globalFitnessFunction=globalFitnessFunction,
                     compareFunction=compareFunction, compareFunctionArgs=compareFunctionArgs)
    end
end # function

"""
    setEvaluator(fitnessFunction1::Function, varArgs...; genj::GeneticJuliaStructure = GenJ,
                 globalFitnessFunction::FitnessFunction = FitnessFunction(noFunc),
                 compareFunction::Function=noFunc, compareFunctionArgs::String = "rawW")

Sets the evaluator for `genj`, receiving a set of function or a individual one,
stablishing the weights to one, meaning that all the fitnesses need to be maximized.
This function constructs a structure of type `Evaluator` and adds it to the main structure.

# Arguments
- `fitnessFunction1::Function`: function for the first fitness that must
    evaluate an individual, receiving the individual's genotype and returning
    a `fitness`. (For information about genotypes look [`Genotype`](@ref))
- `varArgs`: list of fitness function. See `"Examples"`. It can be empty.

# Keyword Arguments
- `genj::GeneticJuliaStructure = GenJ`: the main structure. For `code user`: don't modify
    the default value unless multiple experiments are going to be run.
- `globalFitnessFunction::FitnessFunction = FitnessFunction(noFunc)`: second-phase
    fitness function that is going to be evaluated once the other 'raw' fitness
    functions have been. For example, for [SPEA](@ref) algorithm, this fitness
    function serves to assign each individual its fitness level in the pareto front.
- `compareFunction::Function=noFunc`: function used to compare between individuals.
- `compareFunctionArgs::String = "rawW"`: tells the evaluator what is going to be
    compared. "raw" is for fitness comparison, "rawW" is for weighted fitness comparison,
    "global" is for second-phase fitness comparison, "globalW" is for weighted
    second-phase fitness comparison, "both" is for both fitness and second-phase
    fitness comparison and "bothW" is for both weighted fitness and weighted
    second-phase fitness comparison.

# Examples
```jldoctest
julia> setEvaluator(myComputeFitnessFunction)

julia> setEvaluator(myFitnesFunction1, myFitnessFunction2, myFitnessFunction3)
```

See also: [`Evaluator`](@ref), [FitnessFunction](@ref)
"""
function setEvaluator(fitnessFunction1::Function, varArgs...; genj::GeneticJuliaStructure = GenJ,
                      globalFitnessFunction::FitnessFunction = FitnessFunction(noFunc),
                      compareFunction::Function=noFunc, compareFunctionArgs::String = "rawW")
    len = length(varArgs)
    functions = Array{Function}(undef,1+len)
    weights = ones(Real, 1+len)
    functions[1] = fitnessFunction1

    for i=1:len
        if typeof(varArgs[i]) <: Function
            functions[i+1] = varArgs[i]
        else
            error("The arguments of the evaluator must be fitness functions")
        end
    end

    setEvaluator(functions, weights, genj=genj, globalFitnessFunction=globalFitnessFunction,
                 compareFunction=compareFunction, compareFunctionArgs=compareFunctionArgs)
end # function

max_stars_repo_name: SergioGmezM/GeneticJulia

max_issues_repo_name: SergioGmezM/GeneticJulia

max_forks_repo_name: SergioGmezM/GeneticJulia


text:
#=
function exp_series(A::Array, N::Integer)
    ret = A^0
    M = deepcopy(A)
    ret = zero(A)
    fac = 1
    for k in 1:N
        # M = M * A / k
        fac *= k
        ret += A^k / fac
    end
    return ret
end

function exp_series(A::T, N::Integer) where {B<:Basis,T<:Operator{B,B}}
    ret = exp_series(A.data, N)
    return DenseOperator(A.basis_l, A.basis_r, ret)
end
=#

"""
    thermal_state(T, mu_array, Ham, aggIndices;
    \tboltzmann_const = 0.69503476, diagonalize = false, diagonal = false)

Get initial state as thermal state excited with ultra-fast laser pulse. In this version
we suppose that after the thermal state is excited with laser pulse, the whole population
of ground state is distributed over electric states in `mu_array`. We assume
Condon approximation.

``\\rho_\\text{thermal} = \\exp( -\\frac{i}{\\hbar} H ), \\quad \\hbar = 1```.

# Arguments
* `T`: Temperature of the initial thermal state.
* `mu_array`: Vector of electric states in local basis, see [`electronicIndices`](@ref). The first
        index is for ground state of the aggregate the rest are first excited states in local basis.
* `Ham`: Arbitrary operator specifying the Hamiltonian.
* `aggIndices`: Aggregate indices, see [`getIndices`](@ref).
* `boltzmann_const`: Boltzmann const in ``\\mathrm{cm^{-1}}``.
* `diagonalize`: Decompose Hamiltonian into ``\\lambda_i, S, S^{-1}``
        and calculate the exponential using eigenvalue decomposition.
* `diagonal`: Return only the diagonal part of the excited density matrix.
"""
function thermal_state(
    T,
    mu_array,
    aggCore::AggregateCore,
    aggTools::AggregateTools,
    aggOperators::AggregateOperators;
    boltzmann_const::Float64 = 0.69503476,
    diagonalize::Bool = false,
    diagonal = false
)
    vibindices = aggTools.indicesMap
    Ham = aggOperators.Ham
    a1 = vibindices[1][1]
    a2 = vibindices[1][end]
    Ham_g = Ham.data[a1:a2, a1:a2]
    data = -Ham_g / (T * boltzmann_const)
    if diagonal
        data = Diagonal(data)
    end
    if !diagonalize
        data = exp(data)
    else
        H_lambda, H_S = eigen(data)
        H_lambda = diagm(H_lambda)
        data = H_S * exp(H_lambda / (T * boltzmann_const)) * inv(H_S)
    end
    W0 = DenseOperator(Ham.basis_r, Ham.basis_l, zero(Ham.data))
    excitedElInd = 0
    for elInd in 1:length(vibindices)
        excitedElInd += elInd * (mu_array[1][elInd] - 1)
    end
    a1 = vibindices[excitedElInd][1]
    a2 = vibindices[excitedElInd][end]
    W0.data[a1:a2, a1:a2] = data

    normalize!(W0)
    return W0
end

thermal_state(T, mu_array, agg::Aggregate; boltzmann_const::Float64 = 0.69503476, diagonalize::Bool = false, diagonal = false
) = thermal_state(
    T,
    mu_array,
    agg.core,
    agg.tools,
    agg.operators;
    boltzmann_const=boltzmann_const,
    diagonalize=diagonalize,
    diagonal=diagonal
)

"""
    thermal_state_composite(T, mu_weighted, Ham, aggIndices;
    \tboltzmann_const::Float64 = 0.69503476, diagonalize::Bool=false, diagonal=false)

Functionality of this method is similar to [`thermal_state`](@ref), but the final
state is constructed from partial `thermal_states` with weight specified in `mu_weighted`.
For example

```julia
thermal_state_composite(T, [0.0, 0.8, 0.2], ...) = 0.8 * thermal_state(T, [1, 2, 1], ...) + 0.2 * thermal_state(T, [1, 1, 2], ...)
````

``\\rho_\\text{thermal} = \\exp( -\\frac{i}{\\hbar} H ), \\quad \\hbar = 1``.

# Arguments
* `T`: Temperature of the initial thermal state.
* `mu_weighted`: Vector of weights of electric states in local basis, see [`electronicIndices`](@ref).
* `Ham`: Arbitrary operator specifying the Hamiltonian.
* `aggIndices`: Aggregate indices, see [`getIndices`](@ref).
* `boltzmann_const`: Boltzmann const in ``\\mathrm{cm^{-1}}``.
* `diagonalize`: Decompose Hamiltonian into ``\\lambda_i, S, S^{-1}``
        and calculate the exponential using eigenvalue decomposition.
* `diagonal`: Return only the diagonal part of the excited density matrix.
"""
function thermal_state_composite(
    T,
    mu_weighted,
    aggCore::AggregateCore,
    aggTools::AggregateTools,
    aggOperators::AggregateOperators;
    boltzmann_const::Float64 = 0.69503476,
    diagonalize::Bool = false,
    diagonal = false,
)
    vibindices = aggTools.indicesMap
    aggIndices = aggTools.indices
    Ham = aggOperators.Ham
    a1 = vibindices[1][1]
    a2 = vibindices[1][end]
    Ham_g = Ham.data[a1:a2, a1:a2]
    data = -Ham_g / (T * boltzmann_const)
    if diagonal
        data = Diagonal(data)
    end
    if !diagonalize
        data = exp(data)
    else
        H_lambda, H_S = eigen(data)
        H_lambda = diagm(H_lambda)
        data = H_S * exp(H_lambda / (T * boltzmann_const)) * inv(H_S)
    end
    W0 = DenseOperator(Ham.basis_r, Ham.basis_l, zero(Ham.data))

    for excitedElInd = 1:length(mu_weighted)
        a1 = vibindices[excitedElInd][1]
        a2 = vibindices[excitedElInd][end]
        W0.data[a1:a2, a1:a2] = data * mu_weighted[excitedElInd]
    end

    normalize!(W0)
    return W0
end

thermal_state_composite(T, mu_weighted, agg::Aggregate; boltzmann_const::Float64 = 0.69503476, diagonalize::Bool = false, diagonal = false
) = thermal_state_composite(
    T,
    mu_weighted,
    agg.core,
    agg.tools,
    agg.operators;
    boltzmann_const=boltzmann_const,
    diagonalize=diagonalize,
    diagonal=diagonal
)

function ultrafast_laser_excitation(T::AbstractFloat, weights::Array, agg::Aggregate; diagonalize = true)
    molCount = agg.core.molCount
    mu_array = [ones(Int64, molCount+1)]

    mu_array_tmp = deepcopy(mu_array)
    mu_array_tmp[1][1] = 2
    W0 = weights[1] * thermal_state(T, mu_array_tmp, agg.core, agg.tools, agg.operators; diagonalize = diagonalize)

    for i in 2:molCount+1
        mu_array_tmp = deepcopy(mu_array)
        mu_array_tmp[1][i] = 2
        W0.data[:, :] += weights[i] * thermal_state(T, mu_array_tmp, agg.core, agg.tools, agg.operators; diagonalize = diagonalize).data
    end
    normalize!(W0)
    W0 = DenseOperator(W0.basis_l, W0.basis_r, complex(W0.data))

    rho0 = trace_bath(W0, agg.core, agg.tools; vib_basis=agg.operators.vib_basis)
    rho0 = DenseOperator(rho0.basis_l, rho0.basis_r, complex(rho0.data))

    W0_bath = get_rho_bath(W0, agg.core, agg.tools; vib_basis=agg.operators.vib_basis)
    W0_bath = DenseOperator(W0_bath.basis_l, W0_bath.basis_r, complex(W0_bath.data))
    return W0, rho0, W0_bath
end

ext: jl

max_stars_repo_name: detrin/OpenQuantumSystems

max_issues_repo_name: detrin/OpenQuantumSystems

max_forks_repo_name: detrin/OpenQuantumSystems


text:
using LinearAlgebra, ForwardDiff, Distributions, StaticArrays
using Convex, SCS, ECOS
using IterativeSolvers
using Plots

# noc
function κ_no(z)
    max.(0.0, z)
end

function Jκ_no(z)
    p = zero(z)
    for (i, pp) in enumerate(z)
        println(pp)
        println(i)
        if pp >= 0.0
            p[i] = 1.0
        end
    end
    return Diagonal(p)
end

# soc
function κ_soc(z)
    z1 = z[1:end-1]
    z2 = z[end]

    z_proj = zero(z)

    if norm(z1) <= z2
        z_proj = copy(z)
    elseif norm(z1) <= -z2
        z_proj = zero(z)
    else
        a = 0.5 * (1.0 + z2 / norm(z1))
        z_proj[1:end-1] = a * z1
        z_proj[end] = a * norm(z1)
    end
    return z_proj
end

function Jκ_soc(z)
    z1 = z[1:end-1]
    z2 = z[end]
    m = length(z)

    if norm(z1) <= z2
        return Diagonal(ones(m))
    elseif norm(z1) <= -z2
        return Diagonal(zeros(m))
    else
        D = zeros(m, m)
        for i = 1:m
            if i < m
                D[i, i] = 0.5 + 0.5 * z2 / norm(z1) - 0.5 * z2 * ((z1[i])^2.0) / norm(z1)^3.0
            else
                D[i, i] = 0.5
            end
            for j = 1:m
                if j > i
                    if j < m
                        D[i, j] = -0.5 * z2 * z1[i] * z1[j] / norm(z1)^3.0
                        D[j, i] = -0.5 * z2 * z1[i] * z1[j] / norm(z1)^3.0
                    elseif j == m
                        D[i, j] = 0.5 * z1[i] / norm(z1)
                        D[j, i] = 0.5 * z1[i] / norm(z1)
                    end
                end
            end
        end
        return D
    end
end

# friction problem
v = 1.0 * ones(2)
y = 1.0 #1.0e-3

n = 3
m = 1
c = [v; 0.0]
A = [0.0; 0.0; 1.0]'
b = copy(y)

θ_friction = [vec(A); b; c]

function r_friction(z)
    x = z[1:3]
    μ = z[3 + 1]
    λ = z[3 + 1 .+ (1:3)]

    [[v; 0.0] - [0.0; 0.0; 1.0] * μ - λ;
     y - [0.0; 0.0; 1.0]' * x;
     x - κ_soc(x - λ)]
end

function r_friction(z, θ)
    x = z[1:3]
    μ = z[3 + 1]
    λ = z[3 + 1 .+ (1:3)]

    A = reshape(θ[1:1 * 3], 1, 3)
    b = θ[1 * 3 .+ (1:1)]
    c = θ[1 * 3 + 1 .+ (1:3)]

    [c - A' * μ - λ;
     b - A * x;
     x - κ_soc(x - λ)]
end

function R_friction(z)
    x = z[1:3]
    μ = z[3 + 1]
    λ = z[3 + 1 .+ (1:3)]
    J = Jκ_soc(x - λ)

    [zeros(3, 3) -[0.0; 0.0; 1.0]'' -Diagonal(ones(3));
     -[0.0; 0.0; 1.0]' zeros(1, 1) zeros(1, 3);
     (Diagonal(ones(3)) - J) zeros(3, 1) J]
end

function solve()
    z = 0.001 * rand(3 + 1 + 3)

    extra_iters = 0

    for i = 1:10
        _F = r_friction(z)
        _J = R_friction(z) #R_friction(z)
        Δ = gmres(_J, 1.0 * _F, abstol = 1.0e-12, maxiter = i + extra_iters)
        # Δ = (_J' * _J + 1.0e-5 * I) \ (_J' * _F)
        iter = 0
        α = 1.0
        while norm(r_friction(z - α * Δ))^2.0 >= (1.0 - 0.001 * α) * norm(_F)^2.0 && α > 1.0e-4
            α = 0.5 * α
            # println("   α = $α")
            iter += 1
            if iter > 100
                @error "line search fail"

                return z
            end
        end

        if α <= 1.0e-4
            extra_iters += 1
        end

        println("iter ($i) - norm: $(norm(r_friction(z)))")

        z .-= α * Δ
    end

    return z
end

z_sol = solve()

@show b_sol = z_sol[1:2]

θ = [vec(A); b; c]
r_friction(z_sol, θ)

_rz(w) = r_friction(w, θ)
_rθ(w) = r_friction(z_sol, w)

Rz = ForwardDiff.jacobian(_rz, z_sol)
Rθ = ForwardDiff.jacobian(_rθ, θ)

dθ = ((Rz' * Rz + 1.0e-5 * I) \ (Rz' * Rθ))[1:2, m * n .+ (1:m+(n-1))]

# impact problem
# rod
struct Particle
      m # mass
      g # gravity

      nq # configuration dimension
end

# jacobian
function jacobian(model, q)
    Diagonal(@SVector ones(model.nq))
end

# mass matrix
function mass_matrix(model)
      m = model.m

      Diagonal(@SVector [m, m, m])
end

# gravity
function gravity(model, q)
      m = model.m
      g = model.g

      @SVector [0.0, 0.0, m * g]
end

# dynamics
function dynamics(model, q1, q2, q3, λ, h)
      nq = model.nq
      SVector{nq}(mass_matrix(model) * (2.0 * q2 - q1 - q3) / h
            - h * gravity(model, q2)
            + h * jacobian(model, q3)' * λ)
end

# create model
model = Particle(1.0, 9.81, 3)

q1 = [0.0; 0.0; 1.0]
q2 = [1.0; 1.0; 1.0]

function r_action(z)
    q3 = z[1:3]
    q3z = z[3]
    λ = z[4]

    [dynamics(model, q1, q2, q3, [0.0; 0.0; λ], 0.1);
     q3z - κ_no(q3z - λ)]
end

function solve()
    z = [q2; 0.001]

    extra_iters = 0

    for i = 1:10
        _F = r_action(z)
        _J = ForwardDiff.jacobian(r_action, z)
        Δ = gmres(_J, 1.0 * _F, abstol = 1.0e-12, maxiter = i + extra_iters)
        # Δ = (_J' * _J + 1.0e-5 * I) \ (_J' * _F)
        iter = 0
        α = 1.0
        while norm(r_action(z - α * Δ))^2.0 >= (1.0 - 0.001 * α) * norm(_F)^2.0 && α > 1.0e-4
            α = 0.5 * α
            # println("   α = $α")
            iter += 1
            if iter > 100
                @error "line search fail"

                return z
            end
        end

        if α <= 1.0e-4
            extra_iters += 1
        end

        println("iter ($i) - norm: $(norm(r_action(z)))")

        z .-= α * Δ
    end

    return z
end

z_sol = solve()

function r_friction(z)
    x = z[1:3]
    μ = z[3 + 1]
    λ = z[3 + 1 .+ (1:3)]

    [[v; 0.0] - [0.0; 0.0; 1.0] * μ - λ;
     y - [0.0; 0.0; 1.0]' * x;
     x - κ_soc(x - λ)]
end

# impact and friction

# function r_impact_friction(z)
#     q3 = z[1:3]
#     q3z = z[3]
#     v = (q3[1:2] - q2[1:2]) ./ 0.1
#     λ = z[4]
#     b = z[5:6]
#     b̄ = z[5:7]
#     μ = z[8]
#     η = z[9:11]
#
#     [dynamics(model, q1, q2, q3, [b; λ], 0.1);
#      q3z - κ_no(q3z - λ);
#
#      [v; 0.0] - [0.0; 0.0; 1.0] * μ - η;
#       λ - [0.0; 0.0; 1.0]' * b̄;
#       b̄ - κ_soc(b̄ - η)]
# end
#
# function R_impact_friction(z)
#     R = ForwardDiff.jacobian(r_impact_friction, z)
#
#     # fix projection
#     b̄ = z[5:7]
#     μ = z[8]
#     η = z[9:11]
#     J = Jκ_soc(b̄ - η)
#
#     R[(end-2):end, 5:11] = [(Diagonal(ones(3)) - J) zeros(3, 1) J]
#
#     return R
# end

function solve(q1, q2, h)
    function r(z)
        q3 = z[1:3]
        ϕ = z[3]
        λ = z[4]
        b = z[5:6]
        b̄ = z[5:7]
        μ = z[8]
        η = z[9:11]

        [dynamics(model, q1, q2, q3, [b; λ], h);
         ϕ - κ_no(ϕ - λ);

         [(q3[1:2] - q2[1:2]) ./ h; 0.0] - [0.0; 0.0; 1.0] * μ - η;
          λ - [0.0; 0.0; 1.0]' * b̄;
          b̄ - κ_soc(b̄ - η)]
    end

    function R(z)
        _R = ForwardDiff.jacobian(r, z)

        # fix projection
        b̄ = z[5:7]
        μ = z[8]
        η = z[9:11]
        J = Jκ_soc(b̄ - η)

        _R[(end-2):end, 5:11] = [(Diagonal(ones(3)) - J) zeros(3, 1) J]

        return _R
    end

    z = 1.0e-6 * rand(11)
    # z = zeros(11)
    z[1:3] = copy(q2)

    extra_iters = 0

    for i = 1:100
        _F = r(z)
        _J = R(z)
        Δ = gmres(_J, 1.0 * _F, abstol = 1.0e-12, maxiter = i + extra_iters)
        # Δ = (_J' * _J + 1.0e-6 * I) \ (_J' * _F)
        iter = 0
        α = 1.0
        while norm(r(z - α * Δ))^2.0 >= (1.0 - 0.001 * α) * norm(_F)^2.0 && α > 1.0e-4
            α = 0.5 * α
            # println("   α = $α")
            iter += 1
            if iter > 100
                @error "line search fail"

                return z
            end
        end

        if α <= 1.0e-4
            extra_iters += 1
        end

        println("iter ($i) - norm: $(norm(r(z)))")

        z .-= α * Δ
    end

    if norm(r(z)) < 1.0e-5
        status = true
    else
        status = false
    end

    return z, status
end

h = 0.01
v1 = [1.0; 1.0; 1.0]
q2 = [1.0; -2.0; 1.0]
q1 = q2 - h * v1

function simulate(q1, q2, T, h)
    println("simulation")
    q = [q1, q2]
    y = [0.0]
    b = [zeros(2)]
    for t = 1:T
        println("   t = $t")
        z_sol, status = solve(q[end-1], q[end], h)

        if !status
            @error "failed step (t = $t)"
            return q, y, b
        else
            push!(q, z_sol[1:3])
            push!(y, z_sol[4])
            push!(b, z_sol[5:6])
        end
    end

    return q, y, b
end

q_sol, y_sol, b_sol = simulate(q1, q2, 100, h)
# q_sol[end]
# q_sol[end-1]
plot(hcat(q_sol...)[3:3, :]', xlabel = "")
plot!(h .* hcat(y_sol...)', xlabel = "", linetype = :steppost)

plot(hcat(q_sol...)[1:2, :]', xlabel = "")
plot!(h * hcat(b_sol...)', linetype = :steppost)


text:
using ExoplanetsSysSim

start_dir = pwd()
data_dir_home = joinpath(dirname(pathof(ExoplanetsSysSim)),"..")
pkg_dev_home = joinpath(dirname(pathof(ExoplanetsSysSim)),"..")

cd(pkg_dev_home)
#= Disabled to prevent merge conflicts
println("# Pulling most recent version, just to be sure.")
flush(stdout)
run(`git pull`)
=#

# Disabled, so can create modules manually in another directory and make data a symlink
# include("init_modules.jl")

cd(start_dir)


text:
using LinearAlgebra: Transpose

@testset "dimension reducing methods" begin
    a = [1 2; 3 4]
    dimz = (X((143, 145)), Y((-38, -36)))
    da = DimensionalArray(a, dimz)

    @test sum(da; dims=X()) == sum(da; dims=1)
    @test sum(da; dims=Y()) == sum(da; dims=2)
    @test dims(sum(da; dims=Y())) == (X(LinRange(143.0, 145.0, 2)), Y(-38.0))
    @test prod(da; dims=X) == [3 8]
    @test prod(da; dims=Y()) == [2 12]'
    @test dims(prod(da; dims=X())) == (X(143.0), Y(LinRange(-38.0, -36.0, 2)))
    @test maximum(da; dims=X) == [3 4]
    @test maximum(da; dims=Y()) == [2 4]'
    @test minimum(da; dims=X) == [1 2]
    @test minimum(da; dims=Y()) == [1 3]'
    @test dims(minimum(da; dims=X())) == (X(143.0), Y(LinRange(-38.0, -36.0, 2)))
    @test mean(da; dims=X) == [2.0 3.0]
    @test mean(da; dims=Y()) == [1.5 3.5]'
    @test dims(mean(da; dims=Y())) == (X(LinRange(143.0, 145.0, 2)), Y(-38.0))
    @test mapreduce(x -> x > 3, +, da; dims=X) == [0 1]
    @test mapreduce(x -> x > 3, +, da; dims=Y()) == [0 1]'
    @test dims(mapreduce(x-> x > 3, +, da; dims=Y())) == (X(LinRange(143.0, 145.0, 2)), Y(-38.0))
    @test reduce(+, da) == reduce(+, a)
    @test reduce(+, da; dims=X) == [4 6]
    @test reduce(+, da; dims=Y()) == [3 7]'
    @test dims(reduce(+, da; dims=Y())) == (X(LinRange(143.0, 145.0, 2)), Y(-38.0))
    @test std(da; dims=X()) == [1.4142135623730951 1.4142135623730951]
    @test std(da; dims=Y()) == [0.7071067811865476 0.7071067811865476]'
    @test var(da; dims=X()) == [2.0 2.0]
    @test var(da; dims=Y()) == [0.5 0.5]'
    @test dims(var(da; dims=Y())) == (X(LinRange(143.0, 145.0, 2)), Y(-38.0))
    a = [1 2 3; 4 5 6]
    da = DimensionalArray(a, dimz)
    @test median(da; dims=Y()) == [2.0 5.0]'
    @test median(da; dims=X()) == [2.5 3.5 4.5]
end

@testset "dimension dropping methods" begin
    a = [1 2 3; 4 5 6]
    dimz = (X((143, 145)), Y((-38, -36)))
    da = DimensionalArray(a, dimz)
    # Dimensions must have length 1 to be dropped
    @test dropdims(da[X(1:1)]; dims=X) == [1, 2, 3]
    @test dropdims(da[2:2, 1:1]; dims=(X(), Y()))[] == 4
    @test typeof(dropdims(da[2:2, 1:1]; dims=(X(), Y()))) <: DimensionalArray{Int,0,Tuple{}}
    @test refdims(dropdims(da[X(1:1)]; dims=X)) == (X(143.0),)
end

if VERSION > v"1.1-"
    @testset "iteration methods" begin
        a = [1 2 3 4
             3 4 5 6
             5 6 7 8]
        # eachslice
        da = DimensionalArray(a, (Y(10:30), Time(1:4)))
        @test [mean(s) for s in eachslice(da; dims=Time)] == [3.0, 4.0, 5.0, 6.0]
        slices = [s .* 2 for s in eachslice(da; dims=Y)]
        @test map(sin, da) == map(sin, a)
        @test slices[1] == [2, 4, 6, 8]
        @test slices[2] == [6, 8, 10, 12]
        @test slices[3] == [10, 12, 14, 16]
        dims(slices[1]) == (Time(1.0:1.0:4.0),)
        slices = [s .* 2 for s in eachslice(da; dims=Time)]
        @test slices[1] == [2, 6, 10]
        dims(slices[1]) == (Y(10.0:10.0:30.0),)
    end
end

@testset "simple dimension reordering methods" begin
    da = DimensionalArray(zeros(5, 4), (Y(10:20), X(1:4)))
    tda = transpose(da)
    @test tda == transpose(parent(da))
    @test dims(tda) == (X(LinRange(1.0, 4.0, 4)), Y(LinRange(10.0, 20.0, 5)))
    @test size(tda) == (4, 5)

    tda = Transpose(da)
    @test tda == Transpose(parent(da))
    @test dims(tda) == (X(LinRange(1.0, 4.0, 4)), Y(LinRange(10.0, 20.0, 5)))
    @test size(tda) == (4, 5)
    @test typeof(tda) <: DimensionalArray

    ada = adjoint(da)
    @test ada == adjoint(parent(da))
    @test dims(ada) == (X(LinRange(1.0, 4.0, 4)), Y(LinRange(10.0, 20.0, 5)))
    @test size(ada) == (4, 5)

    dsp = permutedims(da)
    @test permutedims(parent(da)) == parent(dsp)
    @test dims(dsp) == reverse(dims(da))
end

@testset "dimension reordering methods with specified permutation" begin
    da = DimensionalArray(ones(5, 2, 4), (Y(10:20), Time(10:11), X(1:4)))
    dsp = permutedims(da, [3, 1, 2])

    @test permutedims(da, [X, Y, Time]) == permutedims(da, (X, Y, Time))
    @test permutedims(da, [X(), Y(), Time()]) == permutedims(da, (X(), Y(), Time()))
    dsp = permutedims(da, (X(), Y(), Time()))
    @test dsp == permutedims(parent(da), (3, 1, 2))
    @test dims(dsp) == (X(LinRange(1.0, 4.0, 4)), Y(LinRange(10.0, 20.0, 5)), Time(LinRange(10.0, 11.0, 2)))

    dsp = PermutedDimsArray(da, (3, 1, 2))
    @test dsp == PermutedDimsArray(parent(da), (3, 1, 2))
    @test typeof(dsp) <: DimensionalArray
end

@testset "dimension mirroring methods" begin
    # Need to think about dims for these, currently (Y, Y) etc.
    # But you can't index (Y, Y) with dims as you get the
    # first Y both times. It will plot correctly at least.
    a = rand(5, 4)
    da = DimensionalArray(a, (Y(10:20), X(1:4)))

    cvda = cov(da; dims=X)
    @test cvda == cov(a; dims=2)
    @test dims(cvda) == (X(LinRange(1.0, 4.0, 4)), X(LinRange(1.0, 4.0, 4)))
    crda = cor(da; dims=Y)
    @test crda == cor(a; dims=1)
    @test dims(crda) == (Y(LinRange(10.0, 20.0, 5)), Y(LinRange(10.0, 20.0, 5)))
end

@testset "mapslices" begin
    a = [1 2 3 4
         3 4 5 6
         5 6 7 8]
    da = DimensionalArray(a, (Y(10:30), Time(1:4)))
    ms = mapslices(sum, da; dims=Y)
    @test ms == [9 12 15 18]
    @test dims(ms) == (Y(10.0),Time(LinRange(1.0, 4.0, 4)))
    @test refdims(ms) == ()
    ms = mapslices(sum, da; dims=Time)
    @test parent(ms) == [10 18 26]'
end

# These need fixes in base. kwargs are ::Integer so we can't add methods
# or dispatch on AbstractDimension in underscore _methods
accumulate
cumsum
cumprod


text:
function add_box_con!(sim_data::SimData, name::Symbol, var_name::Symbol, min::AbstractArray{T}, max::AbstractArray{T}, range::UnitRange) where T
    for n = range
        add_box_con!(sim_data, Symbol(name, n), Symbol(var_name, n), min, max)
    end

    sim_data.con_fns
end

function add_box_con_snopt!(x_min::AbstractArray{T}, x_max::AbstractArray{T}, sim_data::SimData, var_name::Symbol, min::AbstractArray{T}, max::AbstractArray{T}, range::UnitRange) where T
    for n = range
        x_min[sim_data.vs(Symbol(var_name, n))] .= min
        x_max[sim_data.vs(Symbol(var_name, n))] .= max
    end
end

function trajopt(sim_data::SimData;
                 x0=nothing,quaternion_state=false,x_min=nothing,x_max=nothing,
                 opt_tol=1e-6,major_feas=1e-6,minor_feas=1e-6,max_iter=10000,verbose=0,callback_fn=nothing)

    solver_fn = eval(sim_data.generate_solver_fn)(sim_data)

    if isa(x0, Nothing)
        # should start with some linear interpolation here

        # need a better way to do this...
        x0 = zeros(sim_data.vs.num_vars)

        # TODO use the normalize_configuration!
        if quaternion_state
            for n = 1:sim_data.N
                x0[sim_data.vs(Symbol("q", n))] .= [1., 0., 0., 0., 0., 0., 0.]
            end
        end

        for n = 1:sim_data.N-1
            x0[sim_data.vs(Symbol("h", n))] .= sim_data.Δt
        end
    end

    options = Dict{String, Any}()
    options["Derivative option"] = 1
    options["Verify level"] = -1 # -1 => 0ff, 0 => cheap
    options["Major optimality tolerance"] = opt_tol
    options["Major feasibility tolerance"] = major_feas
    options["Minor feasibility tolerance"] = minor_feas
    options["Iterations limit"] = max_iter

    xopt, info = snopt(solver_fn, sim_data.cs.num_eqs, sim_data.cs.num_ineqs, x0, options, x_min=x_min,x_max=x_max,callback_fn=callback_fn)

    if verbose >= 1
        println(info)
    end

    sol = eval(sim_data.extract_sol)(sim_data, xopt)

    sol
end


text:
using Test

using ClimateMachine.SurfaceFluxes
using ClimateMachine.SurfaceFluxes.Nishizawa2018
using ClimateMachine.SurfaceFluxes.Byun1990
using ClimateMachine.Thermodynamics
using RootSolvers
using CLIMAParameters
struct EarthParameterSet <: AbstractEarthParameterSet end
const param_set = EarthParameterSet()

# FIXME: Use realistic values / test for correctness
# These tests have been run to ensure they do not fail,
# but they need further testing for correctness.

FT = Float32
rtol = 10 * eps(FT)

@testset "SurfaceFluxes" begin
    shf, lhf, T_b, q_pt, α_0 =
        FT(60), FT(50), FT(350), PhasePartition{FT}(0.01, 0.002, 0.0001), FT(1)
    buoyancy_flux =
        SurfaceFluxes.compute_buoyancy_flux(param_set, shf, lhf, T_b, q_pt, α_0)
    @test buoyancy_flux ≈ 0.0017808608107074118
    @test buoyancy_flux isa FT
end

@testset "SurfaceFluxes.Byun1990" begin
    u, flux = FT(0.1), FT(0.2)
    MO_len = Byun1990.monin_obukhov_len(param_set, u, flux)
    @test MO_len ≈ -0.0125

    u_ave, buoyancy_flux, z_0, z_1 = FT(0.1), FT(0.2), FT(2), FT(5)
    γ_m, β_m = FT(15), FT(4.8)
    tol_abs, iter_max = FT(1e-3), 10
    u_star = Byun1990.compute_friction_velocity(
        param_set,
        u_ave,
        buoyancy_flux,
        z_0,
        z_1,
        β_m,
        γ_m,
        tol_abs,
        iter_max,
    )
    @test u_star ≈ 0.201347256193615 atol = tol_abs
    @test u_star isa FT

    Ri, z_b, z_0, Pr_0 = FT(10), FT(2), FT(5.0), FT(0.74)
    γ_m, γ_h, β_m, β_h = FT(15), FT(9), FT(4.8), FT(7.8)
    cm, ch, L_mo = Byun1990.compute_exchange_coefficients(
        param_set,
        Ri,
        z_b,
        z_0,
        γ_m,
        γ_h,
        β_m,
        β_h,
        Pr_0,
    )
    @test cm ≈ 19.700348427787368
    @test ch ≈ 3.3362564728997803
    @test L_mo ≈ -14.308268023583906
    @test cm isa FT
    @test ch isa FT
    @test L_mo isa FT

    Ri, z_b, z_0, Pr_0 = FT(-10), FT(10.0), FT(1.0), FT(0.74)
    γ_m, γ_h, β_m, β_h = FT(15.0), FT(9.0), FT(4.8), FT(7.8)
    cm, ch, L_mo = Byun1990.compute_exchange_coefficients(
        param_set,
        Ri,
        z_b,
        z_0,
        γ_m,
        γ_h,
        β_m,
        β_h,
        Pr_0,
    )
    @test cm ≈ 0.33300280321092746 rtol = rtol
    @test ch ≈ 1.131830939627489 rtol = rtol
    @test L_mo ≈ -0.3726237964444814 rtol = rtol
    @test cm isa FT
    @test ch isa FT
    @test L_mo isa FT

end

@testset "SurfaceFluxes.Nishizawa2018" begin
    FT = Float32
    rtol = 10 * eps(FT)
    u, θ, flux = FT(2), FT(350), FT(20)
    MO_len = Nishizawa2018.monin_obukhov_len(param_set, u, θ, flux)
    @test MO_len ≈ -35.67787971457696

    u_ave, θ, flux, Δz, z_0, a =
        FT(110.0), FT(350.0), FT(20.0), FT(100.0), FT(0.01), FT(5.0)
    Ψ_m_tol, tol_abs, iter_max = FT(1e-3), FT(1e-3), 10
    u_star = Nishizawa2018.compute_friction_velocity(
        param_set,
        u_ave,
        θ,
        flux,
        Δz,
        z_0,
        a,
        Ψ_m_tol,
        tol_abs,
        iter_max,
    )
    @test u_star ≈ 5.526644550864822 atol = tol_abs
    @test u_star isa FT

    FT = Float64
    rtol = 10 * eps(FT)
    z, F_m, F_h, a, u_star, θ, flux, Pr =
        FT(1), FT(2), FT(3), FT(5), FT(110), FT(350), FT(20), FT(0.74)

    K_m, K_h, L_mo = Nishizawa2018.compute_exchange_coefficients(
        param_set,
        z,
        F_m,
        F_h,
        a,
        u_star,
        θ,
        flux,
        Pr,
    )
    @test K_m ≈ -11512.071612359368 rtol = rtol
    @test K_h ≈ -6111.6196776263805 rtol = rtol
    @test L_mo ≈ -5.935907237512742e6 rtol = rtol

    # Test type-stability:
    FT = Float32
    rtol = 10 * eps(FT)
    z, F_m, F_h, a, u_star, θ, flux, Pr =
        FT(1), FT(2), FT(3), FT(5), FT(110), FT(350), FT(20), FT(0.74)

    K_m, K_h, L_mo = Nishizawa2018.compute_exchange_coefficients(
        param_set,
        z,
        F_m,
        F_h,
        a,
        u_star,
        θ,
        flux,
        Pr,
    )
    @test K_m isa FT
    @test K_h isa FT
    @test L_mo isa FT

end

max_stars_repo_name: leios/CLIMA

max_issues_repo_name: leios/CLIMA

max_forks_repo_name: leios/CLIMA


text:
# But, we have 16 processes, shouldn't it be 16x faster?

# Let's try a heavier example.

using Statistics: std, mean, median
using StatsBase: skewness, kurtosis

function stats_by_security(valuation, funcs)
    (nstates, nattr, n) = size(valuation)
    result = zeros(n, nattr, length(funcs))
    for i in 1:n
        for j in 1:nattr
            for (k, f) in enumerate(funcs)
                result[i, j, k] = f(valuation[:, j, i])
            end
        end
    end
    return result
end

funcs = (std, skewness, kurtosis);
@time result = stats_by_security(valuation, funcs);
#=
julia> funcs = (std, skewness, kurtosis);

julia> @time result = stats_by_security(valuation, funcs);
 20.271371 seconds (3.72 M allocations: 67.162 GiB, 4.84% gc time)
=#

# Distributed version.

@everywhere using Statistics: std, mean, median
@everywhere using StatsBase: skewness, kurtosis

function stats_by_security2(valuation, funcs)
    (nstates, nattr, n) = size(valuation)
    result = SharedArray{Float64}((n, nattr, length(funcs)))
    @sync @distributed for i in 1:n
        for j in 1:nattr
            for (k, f) in enumerate(funcs)
                result[i, j, k] = f(valuation[:, j, i])
            end
        end
    end
    return result
end

@time result = stats_by_security2(valuation, funcs);
#=
julia> @time result = stats_by_security2(valuation, funcs);
  2.351749 seconds (4.89 k allocations: 234.469 KiB)
=#

# How much faster?
#=
julia> 20.271371 / 2.351749
8.619700061528675
=#


text:
# -*- coding: utf-8 -*-
# %%
# Copyright 2020-2021 John T. Foster
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# %%
include("bspline.jl")
using LinearAlgebra, IterativeSolvers
export construct_spline_matrix, reconstruct_trajectory, construct_helix, L2error

function construct_helix(n::Integer = 100)
    """
    Helper function which outputs an nx3 array of form [cos.(x)' sin.(x)'  linspace(0,1,n+1)]
    """
    t = LinRange(0, 4*π, n+1)
    arr = zeros(Float64, (length(t), 3))
    arr[:, 1] = cos.(t)
    arr[:, 2] = sin.(t)
    arr[:, 3] = LinRange(0, 1, n+1)
    arr
end

function create_knot_vector(Qk::Vector{<:Float64}, p::Integer = 3)
    ū = create_ūk(Qk)
    n = length(ū)
    m = length(ū) + p + 1
    kv = zeros(m)
    for j = 2:(n - p)
        kv[j+p] = sum(ū[j:(j+p-1)]) / float(p)
    end
    kv[(end - p):end] .= 1
    kv
end

#Knots = m
#From definition, m = n + p + 1
#set m = length(tangents) so the linear algebra works out
#then set n (num samples) = m - p - 1
function reconstruct_trajectory(tangents::Matrix{<:Float64}, p::Integer=3)
    control_points = zeros(size(tangents))
    control_points[:,1] = reconstruct_trajectory_1d(tangents[:,1], p)
    control_points[:,2] = reconstruct_trajectory_1d(tangents[:,2], p)
    control_points[:,3] = tangents[:,3] #Don't reconstruct the domain
    kv = create_knot_vector(tangents[:,1], p)
    basis = BSplineBasis(kv, p, k=2)
    curve = BSplineCurve(basis, control_points)
    return curve
end

function reconstruct_trajectory_1d(tangents::Vector{<:Float64}, p::Integer=3)
    kv = create_knot_vector(tangents, p)
    ū = create_ūk(tangents) # n = m - p - 1
    basis = BSplineBasis(kv, p, k=2)
    N, Nprime = construct_spline_matrix(basis, ū, length(kv), p)
    control_points = lsmr(N, tangents)
    return control_points
end

function construct_spline_matrix(basis::BSplineBasis, samples::Vector{<:Float64}, num_knots::Integer, p::Integer=3)
    rows, cols = length(samples), num_knots - p - 1
    N, Nprime = zeros(Float64, (rows, cols)), zeros(Float64, (rows, cols))
    N[1, 1] = 1
    N[end, end] = 1
    Nprime[1, 1] = 1
    Nprime[end, end] = 1
    for i in p:rows
        evals = basis(samples[i])
        column = find_knot_span(basis, samples[i])
        N[i,column-p:column] = evals[1,:]#evals[1, :][evals[1, :] .> 0]
        Nprime[i,column-p:column] = evals[2,:]#evals[2, :][evals[2, :] .> 0]
    end
    N, Nprime
end

function L2error(C::BSplineCurve, tangents::Matrix{<:Float64})
    """
    Returns a vector with the cumulative sum of the L2 error for each dimension.
    """
    evals = evaluate(C, length(tangents[:,1]))
    error = zeros(length(tangents[:,1]))
    for (i, u) in enumerate(evals[:,3]')
        index = findall(x-> x>=u, tangents[:,3])[1] #Find index of tangents vec closest to z axis of our control point
        error[i] = abs(tangents[index,2] - evals[i,2]) + abs(tangents[index,3] - evals[i,3])
    end
    return error
end

function main(n::Integer=100, p::Integer=3)
    Q = construct_helix(n)
    Curve = reconstruct_trajectory(Q, p)
    Q, Curve
end


text:
# Read a,b,K,U,lambda from dat file
a = readdlm("/Users/niladridas/Documents/GITHUB/JOT/data/a.dat")
b = readdlm("/Users/niladridas/Documents/GITHUB/JOT/data/b.dat")
K = readdlm("/Users/niladridas/Documents/GITHUB/JOT/data/K.dat")
U = readdlm("/Users/niladridas/Documents/GITHUB/JOT/data/U.dat")
Lam = readdlm("/Users/niladridas/Documents/GITHUB/JOT/data/lambda.dat")
##
include("/Users/niladridas/Documents/GITHUB/JOT/ot/sinkhornTransport.jl")

D, L, u, v = @time sinkhornTransport(a,b,K,U,Lam,"distanceRelativeDecrease",Inf,.5e-2,500,1)
#
# # Profile.clear()
# # @profile sinkhornTransport(a,b,K,U,Lam,"marginalDifference",Inf,.5e-2,500,0)
#
# @code_warntype sinkhornTransport(a,b,K,U,Lam,"marginalDifference",Inf,.5e-2,500,0)
p_matOTSink = diagm(vec(u))*K*diagm(vec(v))
M = size(a)[1]
p_matOTSink = p_matOTSink.*size(a)[1]
dist_OTSink = D
x_aOTSink = zeros(M)
for i= 1:M
    x_aOTSink[i] = sum(p_matOTSink[i,j]*samples[j] for j=1:M)
end


text:
# energy calculation using linked lists method
function utotal_parallel(box, DATA, first_atom, next_atom, nc)

  nthreads = Threads.nthreads()     # number of threads available
  Ut = zeros(nthreads)              # vector with nthreads positions

 Threads.@threads  for iat in 1:DATA.N                  # paralelization through the particles
    icell = trunc(Int64,box[iat][1]/DATA.cutoff) + 1
    jcell = trunc(Int64,box[iat][2]/DATA.cutoff) + 1
    for i in icell-1:icell+1
      id = Threads.threadid()
      for j in jcell-1:jcell+1
        iw, jw = wrapcell(nc, i, j)
        jat = first_atom[iw,jw]
        while jat > 0
          if jat > iat
            rij = pbcseparation(box[jat],box[iat],DATA.side)
            if rij <= 2.
              Ut[id] += upair(rij,DATA)
            end
          end
          jat = next_atom[jat]
        end
      end
    end
  end
  return sum(Ut)
end


text:
using LinearAlgebra

export flattened_tup_chain, flattened_named_tuple, flattened_tuple
export FlattenType, FlattenArr, RetainArr

abstract type FlattenType end

"""
    FlattenArr

Flatten arrays in `flattened_tup_chain`
and `flattened_named_tuple`.
"""
struct FlattenArr <: FlattenType end

"""
    RetainArr

Do _not_ flatten arrays in `flattened_tup_chain`
and `flattened_named_tuple`.
"""
struct RetainArr <: FlattenType end

# The Vars instance has many empty entries.
# Keeping all of the keys results in many
# duplicated values. So, it's best we
# "prune" the tree by removing the keys:
flattened_tup_chain(
    ::Type{NamedTuple{(), Tuple{}}},
    ::FlattenType = FlattenArr();
    prefix = (Symbol(),),
) = ()

flattened_tup_chain(
    ::Type{T},
    ::FlattenType;
    prefix = (Symbol(),),
) where {T <: Real} = (prefix,)

flattened_tup_chain(
    ::Type{T},
    ::RetainArr;
    prefix = (Symbol(),),
) where {T <: SArray} = (prefix,)
flattened_tup_chain(
    ::Type{T},
    ::FlattenArr;
    prefix = (Symbol(),),
) where {T <: SArray} = ntuple(i -> (prefix..., i), length(T))

flattened_tup_chain(
    ::Type{T},
    ::RetainArr;
    prefix = (Symbol(),),
) where {T <: SHermitianCompact} = (prefix,)
flattened_tup_chain(
    ::Type{T},
    ::FlattenType;
    prefix = (Symbol(),),
) where {T <: SHermitianCompact} =
    ntuple(i -> (prefix..., i), length(StaticArrays.lowertriangletype(T)))

flattened_tup_chain(
    ::Type{T},
    ::RetainArr;
    prefix = (Symbol(),),
) where {N, TA, T <: Diagonal{N, TA}} = (prefix,)
flattened_tup_chain(
    ::Type{T},
    ::FlattenArr;
    prefix = (Symbol(),),
) where {N, TA, T <: Diagonal{N, TA}} = ntuple(i -> (prefix..., i), length(TA))

flattened_tup_chain(::Type{T}, ::FlattenType; prefix = (Symbol(),)) where {T} =
    (prefix,)

"""
    flattened_tup_chain(::Type{T}) where {T <: Union{NamedTuple,NTuple}}

An array of tuples, containing symbols
and integers for every combination of
each field in the `Vars` array.
"""
function flattened_tup_chain(
    ::Type{T},
    ft::FlattenType = FlattenArr();
    prefix = (Symbol(),),
) where {T <: Union{NamedTuple, NTuple}}
    map(1:fieldcount(T)) do i
        Ti = fieldtype(T, i)
        name = fieldname(T, i)
        sname = name isa Int ? name : Symbol(name)
        flattened_tup_chain(
            Ti,
            ft;
            prefix = prefix == (Symbol(),) ? (sname,) : (prefix..., sname),
        )
    end |>
    Iterators.flatten |>
    collect
end
flattened_tup_chain(
    ::AbstractVars{S},
    ft::FlattenType = FlattenArr(),
) where {S} = flattened_tup_chain(S, ft)

"""
    flattened_named_tuple

A flattened NamedTuple, given a
`Vars` or nested `NamedTuple` instance.

# Example:

```julia
using Test
using ClimateMachine.VariableTemplates
nt = (x = 1, a = (y = 2, z = 3, b = ((a = 1,), (a = 2,), (a = 3,))));
fnt = flattened_named_tuple(nt);
@test keys(fnt) == (:x, :a_y, :a_z, :a_b_1_a, :a_b_2_a, :a_b_3_a)
@test length(fnt) == 6
@test fnt.x == 1
@test fnt.a_y == 2
@test fnt.a_z == 3
@test fnt.a_b_1_a == 1
@test fnt.a_b_2_a == 2
@test fnt.a_b_3_a == 3
```
"""
function flattened_named_tuple end

function flattened_named_tuple(v::AbstractVars, ft::FlattenType = FlattenArr())
    ftc = flattened_tup_chain(v, ft)
    keys_ = Symbol.(join.(ftc, :_))
    vals = map(x -> getproperty(v, wrap_val.(x)), ftc)
    length(keys_) == length(vals) || error("key-value mismatch")
    return (; zip(keys_, vals)...)
end

function flattened_named_tuple(nt::NamedTuple, ft::FlattenType = FlattenArr())
    ftc = flattened_tup_chain(typeof(nt), ft)
    keys_ = Symbol.(join.(ftc, :_))
    vals = flattened_tuple(ft, nt)
    length(keys_) == length(vals) || error("key-value mismatch")
    return (; zip(keys_, vals)...)
end

flattened_tuple(::FlattenArr, a::AbstractArray) = tuple(a...)
flattened_tuple(::RetainArr, a::AbstractArray) = tuple(a)

flattened_tuple(::FlattenArr, a::Diagonal) = tuple(a.diag...)
flattened_tuple(::RetainArr, a::Diagonal) = tuple(a.diag)

flattened_tuple(::FlattenArr, a::SHermitianCompact) = tuple(a.lowertriangle...)
flattened_tuple(::RetainArr, a::SHermitianCompact) = tuple(a.lowertriangle)

# when we splat an empty tuple `b` into `flattened_tuple(ft, b...)`
flattened_tuple(::FlattenType) = ()

# for structs
flattened_tuple(::FlattenType, a) = (a,)

# Divide and conquer:
flattened_tuple(ft::FlattenType, a, b...) =
    tuple(flattened_tuple(ft, a)..., flattened_tuple(ft, b...)...)

flattened_tuple(ft::FlattenType, a::Tuple) = flattened_tuple(ft, a...)

flattened_tuple(ft::FlattenType, a::NamedTuple) = flattened_tuple(ft, Tuple(a))


text:
immutable PoissonRmathSampler <: Sampleable{Univariate,Discrete}
    mu::Float64
end

rand(s::PoissonRmathSampler) =
    int(ccall((:rpois, "libRmath-julia"), Float64, (Float64,), s.mu))

function poissonpvec(μ::Float64, n::Int)
    # Poisson probabilities, from 0 to n
    pv = Array(Float64, n+1)
    @inbounds pv[1] = p = exp(-μ)
    for i = 1:n
        @inbounds pv[i+1] = (p *= (μ / i))
    end
    return pv
end

# Naive sampler by counting exp variables
#
# Suitable for small μ
#
immutable PoissonCountSampler <: Sampleable{Univariate,Discrete}
    μ::Float64
end

function rand(s::PoissonCountSampler)
    μ = s.μ
    n = 0
    c = randexp()
    while c < μ
        n += 1
        c += randexp()
    end
    return n
end

# Algorithm from:
#
#   J.H. Ahrens, U. Dieter (1982)
#   "Computer Generation of Poisson Deviates from Modified Normal Distributions"
#   ACM Transactions on Mathematical Software, 8(2):163-179
#
#   For μ sufficiently large, (i.e. >= 10.0)
#
immutable PoissonADSampler <: Sampleable{Univariate,Discrete}
    μ::Float64
    s::Float64
    d::Float64
    L::Int
end

PoissonADSampler(μ::Float64) =
    PoissonADSampler(μ,sqrt(μ),6.0*μ^2,ifloor(μ-1.1484))

function rand(s::PoissonADSampler)
    # Step N
    G = s.μ + s.s*randn()

    if G >= 0.0
        K = ifloor(G)
        # Step I
        if K >= s.L
            return K
        end

        # Step S
        U = rand()
        if s.d*U >= (s.μ-K)^3
            return K
        end

        # Step P
        px,py,fx,fy = procf(s.μ,K,s.s)

        # Step Q
        if fy*(1-U) <= py*exp(px-fx)
            return K
        end
    end

    while true
        # Step E
        E = randexp()
        U = 2.0*rand()-1.0
        T = 1.8+copysign(E,U)
        if T <= -0.6744
            continue
        end

        K = ifloor(s.μ + s.s*T)
        px,py,fx,fy = procf(s.μ,K,s.s)
        c = 0.1069/s.μ

        # Step H
        if c*abs(U) <= py*exp(px+E)-fy*exp(fx+E)
            return K
        end
    end
end

# Procedure F
function procf(μ::Float64, K::Int, s::Float64)
    # can be pre-computed, but does not seem to affect performance
    ω = 0.3989422804014327/s
    b1 = 0.041666666666666664/μ
    b2 = 0.3*b1*b1
    c3 = 0.14285714285714285*b1*b2
    c2 = b2 - 15.0*c3
    c1 = b1 - 6.0*b2 + 45.0*c3
    c0 = 1.0 - b1 + 3.0*b2 - 15.0*c3

    if K < 10
        px = -μ
        py = μ^K/factorial(K)
    else
        δ = 0.08333333333333333/K
        δ -= 4.8*δ^3
        V = (μ-K)/K
        px = K*log1pmx(V) - δ # avoids need for table
        py = 0.3989422804014327/sqrt(K)

    end
    X = (K-μ+0.5)/s
    X2 = X^2
    fx = -0.5*X2 # missing negation in pseudo-algorithm, but appears in fortran code.
    fy = ω*(((c3*X2+c2)*X2+c1)*X2+c0)
    return px,py,fx,fy
end


text:
module TestShow

using DataFrames, Random, Test

function capture_stdout(f::Function)
    oldstdout = stdout
    rd, wr = redirect_stdout()
    f()
    redirect_stdout(oldstdout)
    size = displaysize(rd)
    close(wr)
    str = read(rd, String)
    close(rd)
    str, size
end

# In the future newline character \n should be added to this test case
df = DataFrame(A = Int64[1:4;], B = ["x\"", "∀ε>0: x+ε>x", "z\$", "A\nC"],
               C = Float32[1.0, 2.0, 3.0, 4.0])

refstr = """
4×3 DataFrame
│ Row │ A     │ B           │ C       │
│     │ Int64 │ String      │ Float32 │
├─────┼───────┼─────────────┼─────────┤
│ 1   │ 1     │ x"          │ 1.0     │
│ 2   │ 2     │ ∀ε>0: x+ε>x │ 2.0     │
│ 3   │ 3     │ z\$          │ 3.0     │
│ 4   │ 4     │ A\\nC        │ 4.0     │"""

for allrows in [true, false], allcols in [true, false]
    io = IOBuffer()
    show(io, df, allcols=allcols)
    str = String(take!(io))
    @test str == refstr
end

Random.seed!(1)
df_big = DataFrame(rand(25,5))

io = IOContext(IOBuffer(), :displaysize=>(11,40), :limit=>true)
show(io, df_big)
str = String(take!(io.io))
@test str == """
25×5 DataFrame. Omitted printing of 2 columns
│ Row │ x1       │ x2       │ x3       │
│     │ Float64  │ Float64  │ Float64  │
├─────┼──────────┼──────────┼──────────┤
│ 1   │ 0.236033 │ 0.644883 │ 0.440897 │
⋮
│ 24  │ 0.278582 │ 0.241591 │ 0.990741 │
│ 25  │ 0.751313 │ 0.884837 │ 0.550334 │"""

io = IOContext(IOBuffer(), :displaysize=>(11,40), :limit=>true)
show(io, df_big, allcols=true)
str = String(take!(io.io))
@test str == """
25×5 DataFrame
│ Row │ x1       │ x2       │ x3       │
│     │ Float64  │ Float64  │ Float64  │
├─────┼──────────┼──────────┼──────────┤
│ 1   │ 0.236033 │ 0.644883 │ 0.440897 │
⋮
│ 24  │ 0.278582 │ 0.241591 │ 0.990741 │
│ 25  │ 0.751313 │ 0.884837 │ 0.550334 │

│ Row │ x4       │ x5       │
│     │ Float64  │ Float64  │
├─────┼──────────┼──────────┤
│ 1   │ 0.580782 │ 0.138763 │
⋮
│ 24  │ 0.762276 │ 0.755415 │
│ 25  │ 0.339081 │ 0.649056 │"""

io = IOContext(IOBuffer(), :displaysize=>(11,40), :limit=>true)
show(io, df_big, allrows=true, allcols=true)
str = String(take!(io.io))
@test str == """
25×5 DataFrame
│ Row │ x1         │ x2        │
│     │ Float64    │ Float64   │
├─────┼────────────┼───────────┤
│ 1   │ 0.236033   │ 0.644883  │
│ 2   │ 0.346517   │ 0.0778264 │
│ 3   │ 0.312707   │ 0.848185  │
│ 4   │ 0.00790928 │ 0.0856352 │
│ 5   │ 0.488613   │ 0.553206  │
│ 6   │ 0.210968   │ 0.46335   │
│ 7   │ 0.951916   │ 0.185821  │
│ 8   │ 0.999905   │ 0.111981  │
│ 9   │ 0.251662   │ 0.976312  │
│ 10  │ 0.986666   │ 0.0516146 │
│ 11  │ 0.555751   │ 0.53803   │
│ 12  │ 0.437108   │ 0.455692  │
│ 13  │ 0.424718   │ 0.279395  │
│ 14  │ 0.773223   │ 0.178246  │
│ 15  │ 0.28119    │ 0.548983  │
│ 16  │ 0.209472   │ 0.370971  │
│ 17  │ 0.251379   │ 0.894166  │
│ 18  │ 0.0203749  │ 0.648054  │
│ 19  │ 0.287702   │ 0.417039  │
│ 20  │ 0.859512   │ 0.144566  │
│ 21  │ 0.0769509  │ 0.622403  │
│ 22  │ 0.640396   │ 0.872334  │
│ 23  │ 0.873544   │ 0.524975  │
│ 24  │ 0.278582   │ 0.241591  │
│ 25  │ 0.751313   │ 0.884837  │

│ Row │ x3        │ x4        │
│     │ Float64   │ Float64   │
├─────┼───────────┼───────────┤
│ 1   │ 0.440897  │ 0.580782  │
│ 2   │ 0.404673  │ 0.768359  │
│ 3   │ 0.736787  │ 0.519525  │
│ 4   │ 0.953803  │ 0.514863  │
│ 5   │ 0.0951856 │ 0.998136  │
│ 6   │ 0.519675  │ 0.603682  │
│ 7   │ 0.0135403 │ 0.758775  │
│ 8   │ 0.303399  │ 0.590953  │
│ 9   │ 0.702557  │ 0.722086  │
│ 10  │ 0.596537  │ 0.953207  │
│ 11  │ 0.638935  │ 0.384411  │
│ 12  │ 0.872347  │ 0.320011  │
│ 13  │ 0.548635  │ 0.865625  │
│ 14  │ 0.262992  │ 0.45457   │
│ 15  │ 0.526443  │ 0.420287  │
│ 16  │ 0.465019  │ 0.225151  │
│ 17  │ 0.275519  │ 0.286169  │
│ 18  │ 0.461823  │ 0.309144  │
│ 19  │ 0.951861  │ 0.170391  │
│ 20  │ 0.288737  │ 0.147162  │
│ 21  │ 0.661232  │ 0.230063  │
│ 22  │ 0.194568  │ 0.0929292 │
│ 23  │ 0.393193  │ 0.681415  │
│ 24  │ 0.990741  │ 0.762276  │
│ 25  │ 0.550334  │ 0.339081  │

│ Row │ x5        │
│     │ Float64   │
├─────┼───────────┤
│ 1   │ 0.138763  │
│ 2   │ 0.456446  │
│ 3   │ 0.739918  │
│ 4   │ 0.816004  │
│ 5   │ 0.114529  │
│ 6   │ 0.748928  │
│ 7   │ 0.878108  │
│ 8   │ 0.930481  │
│ 9   │ 0.896291  │
│ 10  │ 0.663145  │
│ 11  │ 0.472799  │
│ 12  │ 0.880525  │
│ 13  │ 0.0141033 │
│ 14  │ 0.502774  │
│ 15  │ 0.224851  │
│ 16  │ 0.287858  │
│ 17  │ 0.104033  │
│ 18  │ 0.475749  │
│ 19  │ 0.416681  │
│ 20  │ 0.521387  │
│ 21  │ 0.908499  │
│ 22  │ 0.102832  │
│ 23  │ 0.670421  │
│ 24  │ 0.755415  │
│ 25  │ 0.649056  │"""

io = IOContext(IOBuffer(), :displaysize=>(11,40), :limit=>true)
show(io, df_big, allrows=true, allcols=false)
str = String(take!(io.io))
@test str == """
25×5 DataFrame. Omitted printing of 3 columns
│ Row │ x1         │ x2        │
│     │ Float64    │ Float64   │
├─────┼────────────┼───────────┤
│ 1   │ 0.236033   │ 0.644883  │
│ 2   │ 0.346517   │ 0.0778264 │
│ 3   │ 0.312707   │ 0.848185  │
│ 4   │ 0.00790928 │ 0.0856352 │
│ 5   │ 0.488613   │ 0.553206  │
│ 6   │ 0.210968   │ 0.46335   │
│ 7   │ 0.951916   │ 0.185821  │
│ 8   │ 0.999905   │ 0.111981  │
│ 9   │ 0.251662   │ 0.976312  │
│ 10  │ 0.986666   │ 0.0516146 │
│ 11  │ 0.555751   │ 0.53803   │
│ 12  │ 0.437108   │ 0.455692  │
│ 13  │ 0.424718   │ 0.279395  │
│ 14  │ 0.773223   │ 0.178246  │
│ 15  │ 0.28119    │ 0.548983  │
│ 16  │ 0.209472   │ 0.370971  │
│ 17  │ 0.251379   │ 0.894166  │
│ 18  │ 0.0203749  │ 0.648054  │
│ 19  │ 0.287702   │ 0.417039  │
│ 20  │ 0.859512   │ 0.144566  │
│ 21  │ 0.0769509  │ 0.622403  │
│ 22  │ 0.640396   │ 0.872334  │
│ 23  │ 0.873544   │ 0.524975  │
│ 24  │ 0.278582   │ 0.241591  │
│ 25  │ 0.751313   │ 0.884837  │"""

# Test two-argument show
str1, size = capture_stdout() do
    show(df)
end
io = IOContext(IOBuffer(), :limit=>true, :displaysize=>size)
show(io, df)
str2 = String(take!(io.io))
@test str1 == str2

str1, size = capture_stdout() do
    show(df_big)
end
io = IOContext(IOBuffer(), :limit=>true, :displaysize=>size)
show(io, df_big)
str2 = String(take!(io.io))
@test str1 == str2

subdf = view(df, [2, 3], :)
io = IOBuffer()
show(io, subdf, allrows=true, allcols=false)
str = String(take!(io))
@test str == """
2×3 SubDataFrame
│ Row │ A     │ B           │ C       │
│     │ Int64 │ String      │ Float32 │
├─────┼───────┼─────────────┼─────────┤
│ 1   │ 2     │ ∀ε>0: x+ε>x │ 2.0     │
│ 2   │ 3     │ z\$          │ 3.0     │"""
show(io, subdf, allrows=true)
show(io, subdf, allcols=true)
show(io, subdf, allcols=true, allrows=true)

df = DataFrame(A = Vector{String}(undef, 3))

A = DataFrames.StackedVector(Any[[1, 2, 3], [4, 5, 6], [7, 8, 9]])
show(io, A)
A = DataFrames.RepeatedVector([1, 2, 3], 5, 1)
show(io, A)
A = DataFrames.RepeatedVector([1, 2, 3], 1, 5)
show(io, A)

#Test colored show output (for REPL and similar)
df = DataFrame(Fish = ["Suzy", "Amir"], Mass = [1.5, missing])
@test sprint(show, df, context=:color=>true) == """
    2×2 DataFrame
    │ Row │ Fish   │ Mass     │
    │     │ \e[90mString\e[39m │ \e[90mFloat64⍰\e[39m │
    ├─────┼────────┼──────────┤
    │ 1   │ Suzy   │ 1.5      │
    │ 2   │ Amir   │ \e[90mmissing\e[39m  │"""

# Test showing missing
df = DataFrame(A = [:Symbol, missing, :missing],
               B = [missing, "String", "missing"],
               C = [:missing, "missing", missing])
@test sprint(show, df, context=:color=>true) == """
    3×3 DataFrame
    │ Row │ A       │ B       │ C       │
    │     │ \e[90mSymbol⍰\e[39m │ \e[90mString⍰\e[39m │ \e[90mAny\e[39m     │
    ├─────┼─────────┼─────────┼─────────┤
    │ 1   │ Symbol  │ \e[90mmissing\e[39m │ missing │
    │ 2   │ \e[90mmissing\e[39m │ String  │ missing │
    │ 3   │ missing │ missing │ \e[90mmissing\e[39m │"""

# Test showing nothing
df_nothing = DataFrame(A = [1.0, 2.0, 3.0], B = ["g", "g", nothing])
io = IOBuffer()
show(io, df_nothing)
str = String(take!(io))
@test str == """
3×2 DataFrame
│ Row │ A       │ B      │
│     │ Float64 │ Union… │
├─────┼─────────┼────────┤
│ 1   │ 1.0     │ g      │
│ 2   │ 2.0     │ g      │
│ 3   │ 3.0     │        │"""

# Test computing width for Array{String} columns
df = DataFrame([["a"]], [:x])
io = IOBuffer()
show(io, df)
str = String(take!(io))
@test str == """
1×1 DataFrame
│ Row │ x      │
│     │ String │
├─────┼────────┤
│ 1   │ a      │"""

# Test escape characters
df = DataFrame(a = ["1\n1", "2\t2", "3\r3", "4\$4", "5\"5", "6\\6"])
io = IOBuffer()
show(io, df)
str = String(take!(io))
@test str == """
6×1 DataFrame
│ Row │ a      │
│     │ String │
├─────┼────────┤
│ 1   │ 1\\n1   │
│ 2   │ 2\\t2   │
│ 3   │ 3\\r3   │
│ 4   │ 4\$4    │
│ 5   │ 5"5    │
│ 6   │ 6\\\\6   │"""

# Test categorical values
df = DataFrame(a = categorical([1,2,3]), b = categorical(["a", "b", missing]))
io = IOBuffer()
show(io, df)
str = String(take!(io))
@test str == """
3×2 DataFrame
│ Row │ a            │ b             │
│     │ Categorical… │ Categorical…⍰ │
├─────┼──────────────┼───────────────┤
│ 1   │ 1            │ a             │
│ 2   │ 2            │ b             │
│ 3   │ 3            │ missing       │"""

# Test BigFloat
df = DataFrame(a = [big(1.0), missing])
io = IOBuffer()
show(io, df)
str = String(take!(io))
@test str == """
2×1 DataFrame
│ Row │ a         │
│     │ BigFloat⍰ │
├─────┼───────────┤
│ 1   │ 1.0       │
│ 2   │ missing   │"""

end # module


text:
# This file was generated by the Julia Swagger Code Generator
# Do not modify this file directly. Modify the swagger specification instead.

mutable struct ApplicationGatewayAuthenticationCertificatePropertiesFormat <: SwaggerModel
    data::Any # spec type: Union{ Nothing, String } # spec name: data
    provisioningState::Any # spec type: Union{ Nothing, ProvisioningState } # spec name: provisioningState

    function ApplicationGatewayAuthenticationCertificatePropertiesFormat(;data=nothing, provisioningState=nothing)
        o = new()
        validate_property(ApplicationGatewayAuthenticationCertificatePropertiesFormat, Symbol("data"), data)
        setfield!(o, Symbol("data"), data)
        validate_property(ApplicationGatewayAuthenticationCertificatePropertiesFormat, Symbol("provisioningState"), provisioningState)
        setfield!(o, Symbol("provisioningState"), provisioningState)
        o
    end
end # type ApplicationGatewayAuthenticationCertificatePropertiesFormat

const _property_map_ApplicationGatewayAuthenticationCertificatePropertiesFormat = Dict{Symbol,Symbol}(Symbol("data")=>Symbol("data"), Symbol("provisioningState")=>Symbol("provisioningState"))
const _property_types_ApplicationGatewayAuthenticationCertificatePropertiesFormat = Dict{Symbol,String}(Symbol("data")=>"String", Symbol("provisioningState")=>"ProvisioningState")
Base.propertynames(::Type{ ApplicationGatewayAuthenticationCertificatePropertiesFormat }) = collect(keys(_property_map_ApplicationGatewayAuthenticationCertificatePropertiesFormat))
Swagger.property_type(::Type{ ApplicationGatewayAuthenticationCertificatePropertiesFormat }, name::Symbol) = Union{Nothing,eval(Base.Meta.parse(_property_types_ApplicationGatewayAuthenticationCertificatePropertiesFormat[name]))}
Swagger.field_name(::Type{ ApplicationGatewayAuthenticationCertificatePropertiesFormat }, property_name::Symbol) =  _property_map_ApplicationGatewayAuthenticationCertificatePropertiesFormat[property_name]

function check_required(o::ApplicationGatewayAuthenticationCertificatePropertiesFormat)
    true
end

function validate_property(::Type{ ApplicationGatewayAuthenticationCertificatePropertiesFormat }, name::Symbol, val)
end


text:
function generate(::Type{Event},
                  rates::EventRates{T},
                  time::Float64) where T <: EpidemicModel
  totals = Weights([sum(rates[state]) for state in _state_progressions[T][2:end]])
  if sum(totals) == Inf
    new_state = sample(_state_progressions[T][2:end], totals)
    id = sample(1:rates.individuals, Weights(rates[new_state]))
    return Event{T}(time, id, new_state)
  elseif sum(totals) == 0.0
    return NoEvent()
  else
    # Generate new state
    new_state = sample(_state_progressions[T][2:end], totals)
    # Generate event individual
    id = sample(1:rates.individuals, Weights(rates[new_state]))
    return Event{T}(time + rand(Exponential(1.0 / sum(totals))), id, new_state)
  end
end

function generate(::Type{Transmission},
                  tr::TransmissionRates,
                  event::Event{T}) where T <: EpidemicModel
  id = event.individual
  if _new_transmission(event)
    external_or_internal = Weights([tr.external[id]; sum(tr.internal[:,id])])
    if sum(external_or_internal) == 0.0
      @error "All transmission rates = 0.0, No transmission can be generated"
      return NoTransmission()
    elseif sample([true; false], external_or_internal)
      @debug "Exogenous tranmission generated"
      return ExogenousTransmission(id)
    else
      source = sample(1:individuals(tr), Weights(tr.internal[:, id]))
      @debug "Endogenous transmission generated (source id = $source)"
      return EndogenousTransmission(id, source)
    end
  else
    @debug "No transmission generated"
    return NoTransmission()
  end
end


text:
### A Pluto.jl notebook ###
# v0.12.17

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : missing
        el
    end
end

# ╔═╡ b88b727a-3f3e-11eb-3e3c-bb62b380a257
begin
	using HDF5
	using Plots
	using PlutoUI
	using LaTeXStrings
	using Printf
	import Pda
end

# ╔═╡ 92ae91ea-3f52-11eb-12f9-497e9140820d
@bind frame PlutoUI.Slider(0:99, show_value=true)

# ╔═╡ a6760a18-3f5d-11eb-3cac-152259f7d4c8
file = "U1.000tauR1.000disk.h5"

# ╔═╡ c2d093dc-3f3e-11eb-187e-195e52e1fd31
begin
	x = Pda.get_h5data(file, "config/$(frame)/x")
	y = Pda.get_h5data(file, "config/$(frame)/y")
end

# ╔═╡ d2bd59f4-3f5c-11eb-3201-9fdd56f3174b
function average_radial_density(file)
	bins =100
	loc = zeros(bins)
	n = zeros(bins)
	frames = 100
	for frame in 0:frames-1
		x = Pda.get_h5data(file, "config/$(frame)/x")
		y = Pda.get_h5data(file, "config/$(frame)/y")
		r = sqrt.(x.^2 + y.^2)
		loc, n_ = Pda.density1d(r, 0, maximum(r), bins=100, normalize=false)
		n += n_
	end
	return loc, n/frames
end

# ╔═╡ 13220fe6-3f5d-11eb-3b9a-3d925df3f57c
loc,n = average_radial_density(file)

# ╔═╡ 18248b9a-3f3f-11eb-018f-cdf6768c2ac4
scatter(loc[2:end-1], (n./loc)[2:end-1], framestyle=:box, label=nothing)

# ╔═╡ Cell order:
# ╠═b88b727a-3f3e-11eb-3e3c-bb62b380a257
# ╠═92ae91ea-3f52-11eb-12f9-497e9140820d
# ╠═a6760a18-3f5d-11eb-3cac-152259f7d4c8
# ╠═c2d093dc-3f3e-11eb-187e-195e52e1fd31
# ╠═d2bd59f4-3f5c-11eb-3201-9fdd56f3174b
# ╠═13220fe6-3f5d-11eb-3b9a-3d925df3f57c
# ╠═18248b9a-3f3f-11eb-018f-cdf6768c2ac4


text:
module Trees

import Base: parent, show, start, next, done, eltype, length, map
export Tree, data, children, isleaf, isroot, parent, add_child!, tree, leaf_data
export @tree_str

mutable struct Tree{T}
    data    :: T
    parent  :: Nullable{Tree{T}}
    children:: Vector{Tree{T}}
end

Tree{T}(data, parent::Tree{T}) = Tree(data, Nullable(parent), Vector{Tree{T}}())
Tree(data, T=typeof(data)) = Tree(data, Nullable{Tree{T}}(), Vector{Tree{T}}())

data(tree::Tree) = tree.data
children(tree::Tree) = tree.children
isleaf(tree::Tree) = isempty(tree.children)
isroot(tree::Tree) = isnull(tree.parent)
parent(tree::Tree) = get(tree.parent)

function show(io::IO, tree::Tree)
    print(io, '[', tree.data)
    for child in tree.children
      print(io, child)
    end
    print(io, ']')
end

add_child!{T}(tree::Tree{T}, data::T) = push!(tree.children, Tree(data, tree))
function add_child!{T}(tree::Tree{T}, child::Tree{T})
    child.parent = Nullable(tree)
    push!(tree.children, child)
end

function tree(str::AbstractString)
    node = Tree("")
    str  = replace(str, " ", "")[2:end-1]

    for c in str
        if c == '['
            add_child!(node, "")
            node = children(node)[end]
        elseif c == ']'
            node = parent(node)
        else
            node.data *= c
        end
    end
    node
end

macro tree_str(str) tree(str) end

start(tree::Tree) = [tree]
next{T}(::Tree{T}, list::Vector{Tree{T}}) =
    list[1], prepend!(list[2:end], children(list[1]))
done{T}(::Tree{T}, list::Vector{Tree{T}}) = isempty(list)
eltype{T}(::Type{Tree{T}}) = Tree{T}
length(tree::Tree) = isleaf(tree) ? 1 : 1 + sum(length(c) for c in children(tree))

leaf_data(tree::Tree) = [node.data for node in tree if isleaf(node)]

function Base.map(f, t::Tree; uniform_type=true)
    n = uniform_type ? Tree(f(t.data)) : Tree(f(t.data), Any)
    for c in t.children
        add_child!(n, map(f, c, uniform_type=uniform_type))
    end
    n
end

end # module

# tests
# using Trees
# t = Tree("head")
# add_child!(t, "left")
# add_child!(t, "right")
# string(t) == "[head[left][right]]"
# string(tree"[C[G][C]]") == "[C[G][C]]"
# data.(collect(tree"[C[G][C]]")) == ["C", "G", "C"]
# leaf_data(tree"[C[G][C]]") == ["G", "C"]
# string(map(str->str*"m", tree"[C[G][C]]")) == "[Cm[Gm][Cm]]"


text:
module SNOPT7

using SparseArrays

export initialize, snopt, readOptions, setOption
export snoptWorkspace

function __init__()
    # Set up library path
end

const libsnopt7 = "libsnopt7"

mutable struct snoptWorkspace
    status::Int

    # Workspace
    leniw::Int
    lenrw::Int
    iw::Vector{Int32}
    rw::Vector{Float64}

    leniu::Int
    lenru::Int
    iu::Vector{Int32}
    ru::Vector{Float64}

    x::Vector{Float64}
    lambda::Vector{Float64}
    obj_val::Float64

    function snoptWorkspace(leniw::Int,lenrw::Int)
        prob = new(0,leniw, lenrw,
                   zeros(Int,leniw), zeros(Float64,lenrw), 0, 0)
        finalizer(freeWorkspace,prob)
        prob
    end
end

function freeWorkspace(prob::snoptWorkspace)
    ccall((:f_snend, libsnopt7),
          Cvoid, (Ptr{Cint}, Cint, Ptr{Float64}, Cint),
          prob.iw, prob.leniw, prob.rw, prob.lenrw)
end

# Exit codes
SNOPT_status = Dict(
    1=>:Solve_Succeeded,
    2=>:Feasible_Point_Found,
    3=>:Solved_To_Acceptable_Level,
    4=>:Solved_To_Acceptable_Level,
    5=>:Solved_To_Acceptable_Level,
    6=>:Solved_To_Acceptable_Level,
    11=>:Infeasible_Problem_Detected,
    12=>:Infeasible_Problem_Detected,
    13=>:Infeasible_Problem_Detected,
    14=>:Infeasible_Problem_Detected,
    15=>:Infeasible_Problem_Detected,
    16=>:Infeasible_Problem_Detected,
    21=>:Unbounded_Problem_Detected,
    22=>:Unbounded_Problem_Detected,
    31=>:Maximum_Iterations_Exceeded,
    32=>:Maximum_Iterations_Exceeded,
    34=>:Maximum_CpuTime_Exceeded,
    41=>:Numerical_Difficulties,
    42=>:Numerical_Difficulties,
    43=>:Numerical_Difficulties,
    44=>:Numerical_Difficulties,
    45=>:Numerical_Difficulties,
    71=>:User_Requested_Stop,
    72=>:User_Requested_Stop,
    73=>:User_Requested_Stop,
    74=>:User_Requested_Stop,
    81=>:Insufficient_Memory,
    82=>:Insufficient_Memory,
    83=>:Insufficient_Memory,
    91=>:Invalid_Problem_Definition,
    92=>:Invalid_Problem_Definition,
    999=>:Internal_Error)

# Callbacks
function obj_wrapper(mode_::Ptr{Cint}, nnobj::Cint, x_::Ptr{Float64},
                     f_::Ptr{Float64}, g_::Ptr{Float64},
                     status::Cint)
    x    = unsafe_wrap(Array, x_, Int(nnobj))
    mode = unsafe_load(mode_)

    if mode == 0 || mode == 2
        obj = convert(Float64, eval_f(x)) :: Float64
        unsafe_store!(f_,obj)
    end

    if mode == 1 || mode == 2
        g = unsafe_wrap(Array,g_,Int(nnobj))
        eval_grad_f(x,g)
    end
    return
end

function con_wrapper(mode_::Ptr{Cint}, nncon::Cint, nnjac::Cint, negcon::Cint,
                     x_::Ptr{Float64}, c_::Ptr{Float64}, J_::Ptr{Float64},
                     status::Cint)

    x    = unsafe_wrap(Array, x_, Int(nnjac))
    mode = unsafe_load(mode_)

    if mode == 0 || mode == 2
        c = unsafe_wrap(Array, c_, Int(nncon))
        eval_g(x, c)
    end

    if mode == 1 || mode == 2
        J = unsafe_wrap(Array,J_,Int(negcon))
        eval_jac_g(x, J)
    end
    return
end

# SNOPT7 routines
function initialize(printfile::String, summfile::String)
    prob = snoptWorkspace(30500,3000)

    ccall((:f_sninitx, libsnopt7), Cvoid,
          (Ptr{UInt8}, Cint, Ptr{UInt8}, Cint,
           Ptr{Cint}, Cint, Ptr{Cdouble}, Cint),
          printfile, length(printfile), summfile, length(summfile),
          prob.iw, prob.leniw, prob.rw, prob.lenrw)
    return prob
end

function initialize(printfile::String, summfile::String,
                    leniw::Int, lenrw::Int)
    prob = snoptWorkspace(leniw, lenrw)

    ccall((:f_sninitx, libsnopt7), Cvoid,
          (Ptr{UInt8}, Cint, Cint,
           Ptr{UInt8}, Cint, Cint,
           Ptr{Cint}, Cint, Ptr{Cdouble}, Cint),
          printfile, plen, prob.iprint, summfile, slen, prob.isumm,
          prob.iw, prob.leniw, prob.rw, prob.lenrw)
    return prob
end

function readOptions(prob::snoptWorkspace, specsfile::String)
    status = [0]
    ccall((:f_snspecf, libsnopt7), Cvoid,
          (Ptr{UInt8}, Cint, Ptr{Cint},
           Ptr{Cint}, Cint, Ptr{Cdouble}, Cint),
          specsfile, length(specsfile), status,
          prob.iw, prob.leniw, prob.rw, prob.lenrw)
    prob.status = status[1]
    return Int(prob.status)
end

function snopt(prob::snoptWorkspace, start::String, name::String,
               m::Int, n::Int, nnCon::Int, nnObj::Int, nnJac::Int,
               fObj::Float64, iObj::Int,
               confun::Function, objfun::Function,
               #eval_f::Function, eval_grad_f::Function,
               #eval_g::Function, eval_jac_g::Function,
               J::SparseMatrixCSC, bl::Vector{Float64}, bu::Vector{Float64},
               hs::Vector{Int}, x::Vector{Float64})

    @assert n+m == length(x) == length(bl) == length(bu)
    @assert n+m == length(hs)

    prob.iu = [0]
    prob.ru = [0.]

    prob.x      = copy(x)
    prob.lambda = zeros(Float64,n+m)
    pi          = zeros(Float64,m)

    obj_callback = @cfunction($objfun, Cvoid,
                             (Ptr{Cint}, Ptr{Cint}, Ptr{Cdouble},
                              Ptr{Cdouble}, Ptr{Cdouble}, Ptr{Cint}))
    con_callback = @cfunction($confun, Cvoid,
                             (Ptr{Cint}, Ptr{Cint}, Ptr{Cint}, Ptr{Cint}, Ptr{Cdouble},
                              Ptr{Cdouble}, Ptr{Cdouble}, Ptr{Cint}))

    valJ = J.nzval
    indJ = convert(Array{Cint}, J.rowval)
    locJ = convert(Array{Cint}, J.colptr)
    neJ  = length(valJ)

    status  = [0]
    nS      = [0]
    nInf    = [0]
    sInf    = [0.0]
    obj_val = [0.0]
    miniw   = [0]
    minrw   = [0]

    ccall((:f_snoptb, libsnopt7), Cvoid,
          (Ptr{UInt8}, Ptr{UInt8},
           Cint, Cint, Cint, Cint, Cint, Cint,
           Cint, Cdouble,
           Ptr{Cvoid}, Ptr{Cvoid},
           Ptr{Cdouble}, Ptr{Cint}, Ptr{Cint},
           Ptr{Float64}, Ptr{Float64}, Ptr{Cint},
           Ptr{Float64}, Ptr{Float64}, Ptr{Float64},
           Ptr{Cint}, Ptr{Cint}, Ptr{Cint}, Ptr{Cdouble}, Ptr{Cdouble},
           Ptr{Cint}, Ptr{Cint},
           Ptr{Cint}, Cint, Ptr{Cdouble}, Cint,
           Ptr{Cint}, Cint, Ptr{Cdouble}, Cint),
          start, name, m, n, neJ, nnCon, nnObj, nnJac,
          iObj, fObj,
          con_callback, obj_callback,
          valJ, indJ, locJ,
          bl, bu, hs, prob.x, pi, prob.lambda,
          status, nS, nInf, sInf, obj_val,
          miniw, minrw,
          prob.iu, prob.leniu, prob.ru, prob.lenru,
          prob.iw, prob.leniw, prob.rw, prob.lenrw)

    prob.status  = status[1]
    prob.obj_val = obj_val[1]

    return Int(prob.status)
end

function setOption(prob::snoptWorkspace, optstring::String)
    # Set SNOPT7 option via string
    if !isascii(optstring)
        error("SNOPT7: Non-ASCII parameters not supported")
    end

    errors = [0]
    ccall((:f_snset, libsnopt7), Cvoid,
          (Ptr{UInt8}, Cint, Ptr{Cint},
           Ptr{Cint}, Cint, Ptr{Cdouble}, Cint),
          optstring, length(optstring), errors,
          prob.iw, prob.leniw, prob.rw, prob.lenrw)
    return errors[1]
end

function setOption(prob::snoptWorkspace, keyword::String, value::Int)
    # Set SNOPT7 integer option
    if !isascii(keyword)
        error("SNOPT7: Non-ASCII parameters not supported")
    end

    errors = [0]
    ccall((:f_snseti, libsnopt7), Cvoid,
          (Ptr{UInt8}, Cint, Cint, Ptr{Cint},
           Ptr{Cint}, Cint, Ptr{Cdouble}, Cint),
          optstring, length(optstring), value, errors,
          prob.iw, prob.leniw, prob.rw, prob.lenrw)
    return errors[1]
end

function setOption(prob::snoptWorkspace, keyword::String, value::Float64)
    # Set SNOPT7 real option
    if !isascii(keyword)
        error("SNOPT7: Non-ASCII parameters not supported")
    end

    errors = [0]
    ccall((:f_snseti, libsnopt7), Cvoid,
          (Ptr{UInt8}, Cint, Cdouble, Ptr{Cint},
           Ptr{Cint}, Cint, Ptr{Cdouble}, Cint),
          optstring, length(optstring), value, errors,
          prob.iw, prob.leniw, prob.rw, prob.lenrw)
    return errors[1]
end

include("MOIWrapper.jl")

end # module


text:
@testset "Regular matrix algebra" begin
  G = small_group(8, 4)
  A = AlgGrp(FlintQQ, G)
  B, BtoA = Hecke.regular_matrix_algebra(A)

  @test dim(B) == dim(A)

  for i = 1:dim(A)
    for j = 1:dim(A)
      @test BtoA(B[i])*BtoA(B[j]) == BtoA(B[i]*B[j])
      @test (BtoA\A[i])*(BtoA\A[j]) == BtoA\(A[i]*A[j])
    end
  end
end


text:
export RunSummary, DataFrame

using TimeZones, DataFrames

struct RunSummary
    time::Vector{Float64} # seconds
    dist::Vector{Float64}
    alt::Vector{Float64}
    hr::Vector{Float64}
    unit_time::Bool
    start_time::ZonedDateTime

    function RunSummary(n::Int, time::ZonedDateTime; unit_time::Bool = false)
        return new(nan_vec(n), nan_vec(n), nan_vec(n), nan_vec(n), unit_time, time)
    end

end

function nan_vec(n::Int)
    v = Vector{Float64}(undef, n)
    fill!(v, NaN)
    return v
end

import Base.length
function length(rs::RunSummary)
    return length(rs.time)
end

# function copyrow!(dest::RunSummary, src::RunSummary, ind::Int)
#     dest.time[i] = src.time[i]
#     dest.dist[i] = src.dist[i]
#     dest.hr[i] = src.hr[i]
#     dest.alt[i] = src.alt[i]
# end

function unit_run_sum(rs::RunSummary; unit_time = 1.0)
    n = Int(round(sum(rs.time)))

    unit_rs = RunSummary(n, rs.start_time, unit_time = true)

    m = length(rs)

    leftover = 0.0
    ind = 1
    res_time = 0.0
    res_dist = 0.0
    old_unit_dist = 0.0
    old_hr = rs.hr[1]
    old_alt = rs.alt[1]
    old_unit_Δhr = 0.0

    for i = 1:m
        t = rs.time[i] + res_time

        if t < unit_time
            res_time = t
            res_dist += rs.dist[i]
            continue
        end

        unit_dist = rs.dist[i] / rs.time[i]
        unit_alt = rs.alt[i] / rs.time[i]
        unit_Δhr = (rs.hr[i] - old_hr) / rs.time[i]

        α = res_time / unit_time
        β = unit_time - α

        unit_rs.time[ind] = unit_time
        unit_rs.dist[ind] = α * old_unit_dist + β * unit_dist
        unit_rs.hr[ind] = old_hr + α * old_unit_Δhr + β * unit_Δhr
        unit_rs.alt[ind] = α * old_alt + β * unit_alt

        ind += 1

        t -= unit_time

        while t >= 1.0
            unit_rs.time[ind] = unit_time
            unit_rs.dist[ind] = unit_dist
            unit_rs.hr[ind] = unit_rs.hr[ind - 1] + unit_Δhr
            unit_rs.alt[ind] = unit_alt

            ind += 1
            t -= unit_time
        end

        res_time = t
        old_unit_dist = unit_dist
        old_hr = unit_rs.hr[ind - 1]
        unit_rs.hr[ind - 1]
        old_alt = unit_alt

    end

    return unit_rs
end

function RunMetrics.DataFrame(rs::RunSummary)
    df = DataFrames.DataFrame(time = rs.time,
                   dist = rs.dist,
                   alt = rs.alt,
                   hr = rs.hr,
                   start_time = fill(rs.start_time, length(rs)))
    return df
end


text:
using Revise, Distributions, InventoryModels, Test

@testset "InventoryModels.jl" begin
    include("test_inventory.jl")
    include("test_leadtime.jl")
    include("test_supplier.jl")
    include("test_market.jl")
    include("test_assembly.jl")
    include("test_item.jl")
    include("test_endproduct.jl")
    include("test_inventory_system.jl")
    include("test_constraints.jl")
end


text:
get_dim_lyap(integrator::Union{DEIntegrator, DiscreteIterator}) =
    size(init_tangent_state(integrator))[2]

function get_dim_lyap(prob::Union{LEProblem, CLVProblem})
    Q0 = @view prob.tangent_prob.u0[:, 2:end]
    return size(Q0, 2)
end


text:
@testset "EvalTSort" begin

    @testset "tsort" begin
        # For test, reverse node order to expect output in same order
        _tsort(a) = tsort(length(a), i->a[i])

        @test _tsort([
            [3],
            [3],
            [],
        ]) == [2, 1, 3]

        @test _tsort([
            [4],
            [4, 5],
            [5, 8],
            [6, 7, 8],
            [7],
            [],
            [],
            [],
        ]) == [3, 2, 5, 1, 4, 6, 7, 8]

        @test _tsort([
            [],
            [],
            [4],
            [2],
            [1, 2],
            [1, 3],
        ]) == [6, 5, 3, 4, 2, 1]

        @test _tsort([
            [3, 4],
            [4, 5, 8],
            [6, 9],
            [6, 7],
            [7],
            [9],
            [8],
            [10],
            [10],
            [],
        ]) == [2, 5, 1, 3, 4, 6, 9, 7, 8, 10]
    end

    @testset "fire" begin
        # TODO can share one example/test between DFS and TSORT to show diff?
        d = Dag()

        rn1 = node!(d, Int64)
        rn2 = node!(d, Int64)

        an = add!(d, [
            ("left", rn1, latest(0), true),
            ("right", rn2, latest(0), true),
        ])

        res = []
        onfire!(d, rn1.nid, v -> push!(res, "rn1:$v"))
        onfire!(d, rn2.nid, v -> push!(res, "rn2:$v"))
        onfire!(d, an.nid, v -> push!(res, "an:$v"))

        eve = EvalTSort()
        fire(eve, d, rn1.nid, 42)
        @test res == ["rn1:42", "an:42"]
    end

    @testset "eliding" begin
        d = Dag()

        rn1 = node!(d, Tuple{Int64, Int64})
        rn11 = first!(d, rn1)
        rn12 = last!(d, rn1)

        an = add!(d, [
            ("left", rn11, latest(0), true),
            ("right", rn12, latest(0), true),
        ])

        res = []
        onfire!(d, an.nid, v -> push!(res, v))

        eve = EvalTSort()
        fire(eve, d, rn1.nid, (1, 2))
        @test res == [3]
    end

end


text:
# ---
# title: 501. Find Mode in Binary Search Tree
# id: problem501
# author: Tian Jun
# date: 2020-10-31
# difficulty: Easy
# categories: Tree
# link: <https://leetcode.com/problems/find-mode-in-binary-search-tree/description/>
# hidden: true
# ---
#
# Given a binary search tree (BST) with duplicates, find all the
# [mode(s)](https://en.wikipedia.org/wiki/Mode_\(statistics\)) (the most
# frequently occurred element) in the given BST.
#
# Assume a BST is defined as follows:
#
#   * The left subtree of a node contains only nodes with keys **less than or equal to** the node's key.
#   * The right subtree of a node contains only nodes with keys **greater than or equal to** the node's key.
#   * Both the left and right subtrees must also be binary search trees.
#
#
#
# For example:
# Given BST `[1,null,2,2]`,
#
#
#
#        1
#         \
#          2
#         /
#        2
#
#
#
#
# return `[2]`.
#
# **Note:** If a tree has more than one mode, you can return them in any order.
#
# **Follow up:** Could you do that without using any extra space? (Assume that
# the implicit stack space incurred due to recursion does not count).
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


text:
pathwalk(f, ex, path = []) = f(path, ex)

pathwalk(f, ex::AbstractVector, path = []) =
  map((i, x) -> pathwalk(f, x, [path..., i]), 1:length(ex), ex)

pathwalk(f, ex::Expr, path = []) =
  f(path, Expr(ex.head, pathwalk(f, ex.args, path)...))

function apply!(f, path, patch)
  sourcemap(f) do ex
    pathwalk(ex) do p, x
      p == path ? patch(x) : x
    end
  end
end

function index(ex, fs)
  idx = Dict()
  pathwalk(ex) do p, x
    for f in fs
      if matches(f, x)
        push!(get!(idx, p, []), f)
      end
    end
    return x
  end
  return idx
end

parsefile(f) = Expr(CSTParser.parse(String(read(f)), true))

function indices(dir, fs)
  idx = []
  for (root, dirs, files) in walkdir(dir), f in files
    endswith(f, ".jl") || continue
    p = joinpath(root, f)
    fidx = index(parsefile(p), fs)
    isempty(fidx) || push!(idx, [p, filesize(p), fidx])
  end
  return idx
end

function pick(idx)
  file = wsample([size for (_, size, _) in idx])
  (file, _, idx) = idx[file]
  (path, patches) = rand(collect(idx))
  patch = rand(patches)
  return file, path, patch
end

pick!(idx) = apply!(pick(idx)...)


text:
a = r"TEXT .*?(!|;)|.(param)[ ]+([A-Z0-9]*)[= ]*([0-9.e+-]*)(.*?)(?:\\n|$)"im
b = r".(measure|meas)[ ]+(?:ac|dc|op|tran|tf|noise)[ ]+(\S+)[ ]+"im
c = r".(step)[ ]+(oct|param|[])[ ]+([a-z0-9]*)[ ]+(list|[])[ ]+([0-9.e+-]*)[ ]+([0-9.e+-]*)[ ]+([0-9.e+-]*)"im

match_tags = r"""(
                ^TEXT .*?(!|;)|
                .(param)[ ]+([A-Za-z0-9]*)[= ]*([0-9.eE+-]*)(\w)|
                .(measure|meas)[ ]+(?:ac|dc|op|tran|tf|noise)[ ]+(\w)[ ]+|
                .(step)[ ]+(oct |param ){0,1}[ ]*(\w)[ ]+(list ){0,1}[ ]*([0-9.e+-]*(\w)[ ]*)*
                )"""imx

d = r".(step)[ ]+(oct |param ){0,1}[ ]*([a-z0-9]*)[ ]+(list ){0,1}[ ]*([0-9.e+-]*[ ]*)*"im


text:
# This file was generated by the Julia Swagger Code Generator
# Do not modify this file directly. Modify the swagger specification instead.

mutable struct IoK8sApiextensionsApiserverPkgApisApiextensionsV1CustomResourceConversion <: SwaggerModel
    strategy::Any # spec type: Union{ Nothing, String } # spec name: strategy
    webhook::Any # spec type: Union{ Nothing, IoK8sApiextensionsApiserverPkgApisApiextensionsV1WebhookConversion } # spec name: webhook

    function IoK8sApiextensionsApiserverPkgApisApiextensionsV1CustomResourceConversion(;strategy=nothing, webhook=nothing)
        o = new()
        validate_property(IoK8sApiextensionsApiserverPkgApisApiextensionsV1CustomResourceConversion, Symbol("strategy"), strategy)
        setfield!(o, Symbol("strategy"), strategy)
        validate_property(IoK8sApiextensionsApiserverPkgApisApiextensionsV1CustomResourceConversion, Symbol("webhook"), webhook)
        setfield!(o, Symbol("webhook"), webhook)
        o
    end
end # type IoK8sApiextensionsApiserverPkgApisApiextensionsV1CustomResourceConversion

const _property_map_IoK8sApiextensionsApiserverPkgApisApiextensionsV1CustomResourceConversion = Dict{Symbol,Symbol}(Symbol("strategy")=>Symbol("strategy"), Symbol("webhook")=>Symbol("webhook"))
const _property_types_IoK8sApiextensionsApiserverPkgApisApiextensionsV1CustomResourceConversion = Dict{Symbol,String}(Symbol("strategy")=>"String", Symbol("webhook")=>"IoK8sApiextensionsApiserverPkgApisApiextensionsV1WebhookConversion")
Base.propertynames(::Type{ IoK8sApiextensionsApiserverPkgApisApiextensionsV1CustomResourceConversion }) = collect(keys(_property_map_IoK8sApiextensionsApiserverPkgApisApiextensionsV1CustomResourceConversion))
Swagger.property_type(::Type{ IoK8sApiextensionsApiserverPkgApisApiextensionsV1CustomResourceConversion }, name::Symbol) = Union{Nothing,eval(Base.Meta.parse(_property_types_IoK8sApiextensionsApiserverPkgApisApiextensionsV1CustomResourceConversion[name]))}
Swagger.field_name(::Type{ IoK8sApiextensionsApiserverPkgApisApiextensionsV1CustomResourceConversion }, property_name::Symbol) =  _property_map_IoK8sApiextensionsApiserverPkgApisApiextensionsV1CustomResourceConversion[property_name]

function check_required(o::IoK8sApiextensionsApiserverPkgApisApiextensionsV1CustomResourceConversion)
    (getproperty(o, Symbol("strategy")) === nothing) && (return false)
    true
end

function validate_property(::Type{ IoK8sApiextensionsApiserverPkgApisApiextensionsV1CustomResourceConversion }, name::Symbol, val)
end


text:
function freqanalysisvdv(signal,Ts,iweight)
include("weight_factors_mod.jl")
include("Butterworth_filter_function_alt.jl")
(fwl,fw,fwu,wk,wd,wf,wc,we,wj,wb,www)=weight_factors_mod(iweight);
aw=zeros(length(signal),1);
iband=3;  # bandpass filtering
iphase=1; # refiltering for phase correction
for i=1:44
    fh=fwl[i];  # highpass filter frequency
    fl=fwu[i];  # lowpass filter frequency
    sr=1/Ts;
    if(fl<(sr/2.1))
        (y,mu,sd,rms)=Butterworth_filter_function_alt(signal,Ts,iband,fl,fh,iphase);
        aw=aw+y*www[i];
    end
end
n=length(aw);
AW=std(aw);
VDV=0;
for i=1:n
    VDV=VDV+aw[i]^4;
end
VDV=(VDV*Ts)^0.25;

# max, rms_a, rms_aw, vdv, crest
maxv=maximum(abs.(signal));
rms_a=std(signal);
rms_aw=AW;
# VDV;
crest=maximum(abs.(aw))/std(signal);

return (maxv=maxv,rms_a=rms_a,rms_aw=rms_aw,VDV=VDV,crest=crest)

# aa=[signal,aw];

# test=1;
end


text:
using MUMPS
if VERSION >= v"0.5.0-dev+7720"
    using Base.Test
else
    using BaseTestNext
    const Test = BaseTestNext
end
include("getDivGrad.jl");

A = getDivGrad(32,32,16);
n = size(A,1);

# REAL: test for single rhs
println("Test for real SPD matrix: one sparse random rhs");
rhs = sprandn(n,1,0.05);

x = solveMUMPS(A,rhs,1);
err  =  norm(A*x-full(rhs)) / norm(full(rhs));
@test err < 1e-14

# REAL: test for multiple rhs
println("Test for real SPD matrix: multiple sparse random rhs");
nrhs = 10;
rhs = sprandn(n,nrhs,0.03);

x = solveMUMPS(A,rhs,1);

err = zeros(nrhs)
for i=1:nrhs
        err[i] =  norm(A*x[:,i]-full(rhs[:,i])) / norm(full(rhs[:,i]));
end
@test eltype(x) == Float64
@test maximum(err) < 1e-14

println("Test for real SPD matrix: multiple sparse delta function rhs");
rhs = speye(n)
rhs = rhs[:,1:nrhs]

x = solveMUMPS(A,rhs,1);

err = zeros(nrhs)
for i=1:nrhs
        err[i] =  norm(A*x[:,i]-full(rhs[:,i])) / norm(full(rhs[:,i]));
end
@test eltype(x) == Float64
@test maximum(err) < 1e-14

# COMPLEX: test for one rhs
println("Test for complex SPD matrix: one rhs");
r = rand(n);
A = A + im*spdiagm(r,0);

rhs = sprandn(n,1,0.03) + im*sprandn(n,1,0.03);

x = solveMUMPS(A,rhs,2);

err=  norm(A*x-full(rhs)) / norm(full(rhs))
@test eltype(x) == Complex128
@test err < 1e-14

# COMPLEX : test for multiple rhs
println("Test for complex SPD matrix: multiple sparse random rhs");
nrhs = 10;
rhs = sprandn(n,nrhs,0.03) + im*sprandn(n,nrhs,0.03);

x = solveMUMPS(A,rhs,2);
err = zeros(nrhs)
for i=1:nrhs
        err[i] =  norm(A*x[:,i]-full(rhs[:,i])) / norm(full(rhs[:,i]));
end
@test eltype(x) == Complex128
@test maximum(err) < 1e-14

println("DONE!")


text:
## HIDDEN SECTION
#

using LinearAlgebra
using Random
using Makie, GLMakie, AbstractPlotting
#import Plots
AbstractPlotting.inline!(false)

## VISIBLE SECTION
#

mutable struct Boid{T <: Float64}
  id::Int; pos::Array{T,1}; vel::Array{T,1}
  dv::T; av::T; ds::T
  w0::T; ws::T; wa::T; wc::T
  vmax::T; vmin::T
  bmin::Array{T,1}; bmax::Array{T,1}
  ww::T; wb::T; wp::T; we::T
  nc::Int; dp::T; de::T;
end

dis(bi, bj) = norm(bi.pos - bj.pos)
ang(bi, bj) = acos(dot(bi.vel, bj.pos)/(norm(bi.vel)*norm(bj.pos)))
dim(b) = length(b.pos)

## HIDDEN SECTION
#

macro reload()
    :( include("boids.jl"); println("Reloaded.") )
end

mutable struct World
  xy; sc; ms
  World() = new()
end

include("models/basic.jl")
include("models/modded.jl")
include("models/pureva.jl")


text:
module ReactionNetworkImporters

using DataStructures, Catalyst, SparseArrays
using Symbolics: operation, unwrap

# creates a ModelingToolkit function-like Symbol
# can then do stuff like
# @parameters t
# S₁ = funcsym(S,1)
# u = S₁(t)
function funcsym(S::Symbol, t, args...)
    S = Symbol(S,args...)
    (@variables $(S)(t))[1]
end

abstract type NetworkFileFormat end

# exported data types
#struct RSSANetwork <: NetworkFileFormat end
struct BNGNetwork <: NetworkFileFormat end

struct ParsedReactionNetwork
    "Catalyst Network"
    rn::ReactionSystem

    "Dict mapping initial condition symbolic variables to values."
    u₀

    "Dict mapping parameter symbolic variables to values."
    p

    "Dict mapping symbolic variable for species names to full string for species name"
    varstonames

    "Dict from group name (as string) to corresponding symbolic variable"
    groupstosyms

end
ParsedReactionNetwork(rn::ReactionSystem; u₀=nothing, p=nothing, varstonames=nothing, groupstosyms=nothing) =
                        ParsedReactionNetwork(rn, u₀, p, varstonames, groupstosyms)

export BNGNetwork, MatrixNetwork, ParsedReactionNetwork, ComplexMatrixNetwork

# parsers
include("parsing_routines_bngnetworkfiles.jl")
include("parsing_routines_matrixnetworks.jl")

export loadrxnetwork

end # module


text:
@testset "@β macro" begin
    #=
    @test MacroTools.striplines(@macroexpand @β(["commonwealth"], ["massachusetts", "virginia"])) ==
        MacroTools.striplines(:(βcommonwealth_massachusetts * massachusetts + βcommonwealth_virginia * virginia))

    @test MacroTools.striplines(@macroexpand @β(["massachusetts", "virginia", "kentucky", "pennsylvania"], "commonwealth")) ==
        MacroTools.striplines(:(
            βcommonwealth_massachusetts * massachusetts +
            βcommonwealth_virginia * virginia +
            βcommonwealth_kentucky * kentucky +
            βcommonwealth_pennsylvania * pennsylvania
        ))

    @test MacroTools.striplines(@macroexpand @β(["carolina"], "north")) ==
        MacroTools.striplines(:(βnorth_carolina * carolina))

    @test MacroTools.striplines(@macroexpand @β([], "south")) == MacroTools.striplines(quote end)
    =#
end