text:
"""
NumDiffTools.jacobian(f, x::AbstractArray, [abstol, reltol, h, t, maxiters])

Returns jacobian `J(f)` of `f` evaluted at `x`, assuming `f` is called as `f(x)`.

Optional Arguments:\n
`abstol`: Target absolute error in derivative estimate (default `1e-12`)\n
`reltol`: Target relative error in derivative estimate (default `1e-12`)\n
`h`: Initial step size (default `0.1`)\n
`t`: Ratio of step sizes for each Richardson iteration (default `2.0001`.  Should not be an exact integer for periodic functions)\n
`maxiters`: Maximum number of Richardson iterations to perform (default `6`)\n
"""
function jacobian(f, x::AbstractArray; abstol::Real = 1e-12, reltol::Real = 1e-12, h::Real = 0.1, t::Real = 2.0001, maxiters::Integer = 6)
    fx = f(x)

    # Initialize Work Arrays
    A1 = Vector{typeof(fx)}(undef, maxiters)
    A2 = Vector{typeof(fx)}(undef, maxiters)

    # Vector of step sizes
    hh = [h/t^(i - 1) for i = 1:maxiters]
    # Vector of orders
    k = 1:maxiters

    result = Matrix{eltype(fx)}(undef, length(fx), length(x))
    dx = zero(x)
    for i = 1:length(x)
        dx[i] = 1

        result[:, i] .= extrapolate(h->(f(x + h*dx) - fx)/h,          # Function to extrapolate
                                    A1, A2,                           # Work Arrays
                                    hh, k,                            # Sequence of step sizes (h)
                                    abstol, reltol                    # tolerances
                                    )
        dx[i] = 0
    end
    result
end

"""
NumDiffTools.jacobian(f!, y::AbstractArray, x::AbstractArray, [abstol, reltol, h, t, maxiters])

Returns jacobian `J(f!)` of `f!` evaluted at `x`, assuming `f` is called as `f!(y, x)` where the result is stored in y.
"""
function jacobian(f!, y::AbstractArray, x::AbstractArray; abstol::Real = 1e-12, reltol::Real = 1e-12, h::Real = 0.1, t::Real = 2.0001, maxiters::Integer = 6)

    # Initialize Work Arrays
    A1 = Vector{typeof(y)}(undef, maxiters)
    A2 = Vector{typeof(y)}(undef, maxiters)

    # Vector of step sizes
    hh = [h/t^(i - 1) for i = 1:maxiters]
    # Vector of orders
    k = 1:maxiters

    fx = similar(y)
    f!(fx, x)
    function A(h, x, dx)
        f!(y, x + h*dx)
        return (y - fx)/h
    end

    result = Matrix{eltype(y)}(undef, length(y), length(x))
    dx = zero(x)
    for i = 1:length(x)
        dx[i] = 1
        result[:,i] .= extrapolate(h->A(h, x, dx),                   # Function to extrapolate
                                   A1, A2,                           # Work Arrays
                                   hh, k,                            # Sequence of step sizes (h)
                                   abstol, reltol                    # tolerances
                                   )
        dx[i] = 0
    end
    result
end

"""
NumDiffTools.jacobian!(result::AbstractArray, f, x::AbstractArray, [abstol, reltol, h, t, maxiters])

Stores jacobian `J(f)` of `f` evaluted at `x` in `result`, assuming `f` is called as `f(x)`.

Optional Arguments:\n
`abstol`: Target absolute error in derivative estimate (default `1e-12`)\n
`reltol`: Target relative error in derivative estimate (default `1e-12`)\n
`h`: Initial step size (default `0.1`)\n
`t`: Ratio of step sizes for each Richardson iteration (default `2.0001`.  Should not be an exact integer for periodic functions)\n
`maxiters`: Maximum number of Richardson iterations to perform (default `6`)\n
"""
function jacobian!(result::AbstractArray, f, x::AbstractArray; abstol::Real = 1e-12, reltol::Real = 1e-12, h::Real = 0.1, t::Real = 2.0001, maxiters::Integer = 6)
    fx = f(x)

    # Initialize Work Arrays
    A1 = Vector{typeof(fx)}(undef, maxiters)
    A2 = Vector{typeof(fx)}(undef, maxiters)

    # Vector of step sizes
    hh = [h/t^(i - 1) for i = 1:maxiters]
    # Vector of orders
    k = 1:maxiters

    dx = zero(x)
    for i = 1:length(x)
        dx[i] = 1

        result[:, i] .= extrapolate(h->(f(x + h*dx) - fx)/h,          # Function to extrapolate
                                    A1, A2,                           # Work Arrays
                                    hh, k,                            # Sequence of step sizes (h)
                                    abstol, reltol                    # tolerances
                                    )
        dx[i] = 0
    end
    result
end

"""
NumDiffTools.jacobian!(result::AbstractArray, f!, y::AbstractArray, x::AbstractArray, [abstol, reltol, h, t, maxiters])

Stores jacobian `J(f!)` of `f!` evaluted at `x` in `result`, assuming `f` is called as `f!(y, x)` where the result is stored in y.
"""
function jacobian!(result::AbstractArray, f!, y::AbstractArray, x::AbstractArray; abstol::Real = 1e-12, reltol::Real = 1e-12, h::Real = 0.1, t::Real = 2.0001, maxiters::Integer = 6)

    # Initialize Work Arrays
    A1 = Vector{typeof(y)}(undef, maxiters)
    A2 = Vector{typeof(y)}(undef, maxiters)

    # Vector of step sizes
    hh = [h/t^(i - 1) for i = 1:maxiters]
    # Vector of orders
    k = 1:maxiters

    fx = similar(y)
    f!(fx, x)
    function A(h, x, dx)
        f!(y, x + h*dx)
        return (y - fx)/h
    end

    dx = zero(x)
    for i = 1:length(x)
        dx[i] = 1
        result[:,i] .= extrapolate(h->A(h, x, dx),                   # Function to extrapolate
                                   A1, A2,                           # Work Arrays
                                   hh, k,                            # Sequence of step sizes (h)
                                   abstol, reltol                    # tolerances
                                   )
        dx[i] = 0
    end
    result
end


text:
struct AA
    name::String
    three_letter_code::String
    one_letter_code::String
    type::String
    polar::Bool
    hydrophobic::Bool
    charge::Int
    mono_isotopic_mass::Float64
    mass::Float64
end

const ALA = AA("Alanine", "ALA", "A", "Aliphatic", false, false, 0, 71.037114, 71.0779)
const ARG = AA("Arginine", "ARG", "R", "Basic", true, false, 1, 156.101111, 156.1857)
const ASN = AA("Asparagine", "ASN", "N", "Amide", true, false, 0, 114.042927, 114.1026)
const ASP = AA("Aspartic acid", "ASP", "D", "Acidic", true, false, -1, 115.026943, 115.0874)
const CYS = AA("Cysteine", "CYS", "C", "Sulfuric", false, false, 0, 103.009185, 103.1429)
const GLN = AA("Glutamine", "GLN", "Q", "Amide", true, false, 0, 128.058578, 128.1292)
const GLU = AA("Glutamic acid", "GLU", "E", "Acidic", true, false, -1, 129.042593, 129.1140)
const GLY = AA("Glycine", "GLY", "G", "Aliphatic", false, false, 0, 57.021464, 57.0513)
const HIS = AA("Histidine", "HIS", "H", "Aromatic", true, false, 0, 137.058912, 137.1393)
const HSD = AA("Histidine", "HSD", "H", "Aromatic", true, false, 1, 138.067000, 138.1470)
const HSE = AA("Histidine", "HSE", "H", "Aromatic", true, false, 1, 138.067000, 138.1470)
const ILE = AA("Isoleucine", "ILE", "I", "Aliphatic", false, true, 0, 113.084064, 113.1576)
const LEU = AA("Leucine", "LEU", "L", "Aliphatic", false, true, 0, 113.084064, 113.1576)
const LYS = AA("Lysine", "LYS", "K", "Basic", true, false, 1, 128.094963, 128.1723)
const MET = AA("Methionine", "MET", "M", "Sulfuric", false, false, 0, 131.040485, 131.1961)
const PHE =
    AA("Phenylalanine", "PHE", "F", "Aromatic", false, true, 0, 147.068414, 147.1739)
const PRO = AA("Proline", "PRO", "P", "Cyclic", false, false, 0, 97.052764, 97.1152)
const SER = AA("Serine", "SER", "S", "Hydroxylic", true, false, 0, 87.032028, 87.07730)
const THR = AA("Threonine", "THR", "T", "Hydroxylic", true, false, 0, 101.047679, 101.1039)
const TRP = AA("Tryptophan", "TRP", "W", "Aromatic", false, true, 0, 186.079313, 186.2099)
const TYR = AA("Tyrosine", "TYR", "Y", "Aromatic", true, false, 0, 163.063320, 163.1733)
const VAL = AA("Valine", "VAL", "V", "Aliphatic", false, true, 0, 99.068414, 99.1311)
const natural_aminoacids = [
    ALA,
    ARG,
    ASN,
    ASP,
    CYS,
    GLN,
    GLU,
    GLY,
    HIS,
    HSD,
    HSE,
    ILE,
    LEU,
    LYS,
    MET,
    PHE,
    PRO,
    SER,
    THR,
    TRP,
    TYR,
    VAL,
]

"""

```
Sequence
```

Wrapper for strings, or vectors of chars, strings, or residue names, to dispatch on
functions that operate on amino acid sequences.

# Example

```julia-repl
julia> seq = ["Alanine", "Glutamic acid", "Glycine"];

julia> mass(Sequence(seq))
257.2432

julia> seq = "AEG";

julia> mass(Sequence(seq))
257.2432
```

"""
struct Sequence{T}
    s::T
end

"""

```
mass(s::Sequence)
```

Returns the mass of a sequence of amino acids, given a `Sequence` struct type.

# Examples

```julia-repl
julia> seq = ["Alanine", "Glutamic acid", "Glycine"];

julia> mass(Sequence(seq))
257.2432

julia> seq = "AEG";

julia> mass(Sequence(seq))
257.2432

julia> seq = ["ALA", "GLU", "GLY"];

julia> mass(Sequence(seq))
257.2432
```

"""
function mass(s::Sequence)
    m = 0.0
    for aa in s.s
        if length(aa) == 1
            iaa = findfirst(x -> x.one_letter_code == string(aa), natural_aminoacids)
        elseif length(aa) == 3
            iaa = findfirst(x -> x.three_letter_code == string(aa), natural_aminoacids)
        else
            iaa = findfirst(x -> x.name == string(aa), natural_aminoacids)
        end
        if isnothing(iaa)
            error("Could not residue $aa in natural amino acid list.")
        end
        m += natural_aminoacids[iaa].mass
    end
    return m
end

ext: jl

max_stars_repo_name: mcubeg/PDBTools

max_issues_repo_name: mcubeg/PDBTools

max_forks_repo_name: mcubeg/PDBTools


text: Dict{String, Any}("valid key" => Dict{String, Any}())


text:
using ANOVADDPTest
using ANOVADDPTest: dp_mass, n_clusters,  cluster_sizes
using ANOVADDPTest: active_clusters, passive_clusters, cluster_capacity, update!
using Distributions
using SpecialFunctions
using StatsBase: mode
using Statistics: mean, var
using Random
using Test

include("utils.jl")
include("normal.jl")
include("poisson.jl")
include("bernoulli.jl")


text:
# Test SPN learn of mixtures of Chow-Liu trees
@testset "Chow-Liu" begin
    import SPNetworks.BayesianNetworks: BayesianTree, BTRoot, BTNode, BTLeaf, compile
    # Example from Strudel paper (PGM 2020)
    tree = BTRoot(
        4,
        [0.6, 0.4],
        [
            BTNode(
                3,
                [0.2 0.7;
                 0.8 0.3],
                [
                    BTLeaf(
                        1,
                        [0.3 0.6;
                         0.7 0.4]
                    ),
                    BTLeaf(
                        2,
                        [0.5 0.1;
                         0.5 0.9]
                    )
                ]
            )
        ]
    )
    spn = compile(tree)
    @debug spn
    @test length(spn) == 25
    @test length(scope(spn)) == 4
    @test spn([NaN,NaN,NaN,NaN]) ≈ 1.0
    # joint probability values
    values =  Dict( (1,1,1,1) => 0.6*0.2*0.3*0.5, # x1=1, x2=1, x3=1, x4=1
                    (1,1,1,2) => 0.4*0.7*0.3*0.5, # x1=1, x2=1, x3=1, x4=2
                    (1,1,2,1) => 0.6*0.8*0.6*0.1, # x1=1, x2=1, x3=2, x4=1
                    (1,1,2,2) => 0.4*0.3*0.6*0.1, # x1=1, x2=1, x3=2, x4=2
                    (1,2,1,1) => 0.6*0.2*0.3*0.5,
                    (1,2,1,2) => 0.4*0.7*0.3*0.5,
                    (1,2,2,1) => 0.6*0.8*0.6*0.9,
                    (1,2,2,2) => 0.4*0.3*0.6*0.9,
                    (2,1,1,1) => 0.6*0.2*0.7*0.5, # x1=2, x2=1, x3=1, x4=1
                    (2,1,1,2) => 0.4*0.7*0.7*0.5, # x1=2, x2=1, x3=1, x4=2
                    (2,1,2,1) => 0.6*0.8*0.4*0.1, # x1=2, x2=1, x3=2, x4=1
                    (2,1,2,2) => 0.4*0.3*0.4*0.1, # x1=2, x2=1, x3=2, x4=2
                    (2,2,1,1) => 0.6*0.2*0.7*0.5,
                    (2,2,1,2) => 0.4*0.7*0.7*0.5,
                    (2,2,2,1) => 0.6*0.8*0.4*0.9,
                    (2,2,2,2) => 0.4*0.3*0.4*0.9,
                )
    for x1=1:2, x2=1:2, x3=1:2, x4=1:2
        @test spn(Float64[x1,x2,x3,x4]) ≈ values[x1,x2,x3,x4]
    end
end


text:
using TuringModels
#using StatsPlots

# Define a simple Normal model with unknown mean and variance.

@model gdemo(x, y) = begin
  s ~ InverseGamma(2, 3)
  m ~ Normal(0, sqrt(s))
  x ~ Normal(m, sqrt(s))
  y ~ Normal(m, sqrt(s))
end

#  Run sampler, collect results

chn = sample(gdemo(1.5, 2), NUTS(0.65), 1000)

# Summarise results (currently requires the master branch from MCMCChains)

describe(chn) |> display

# Plot and save results

p = plot(chn)
#savefig("basic-example-plot.png")


text:
using Turing
using LinearAlgebra

using BenchmarkHelper

include("lr_helper.jl")

if !haskey(BenchmarkSuite, "nuts")
    BenchmarkSuite["nuts"] = BenchmarkGroup(["nuts"])
end

x, y = readlrdata()

@model function hlr_nuts(x, y, θ)

    N,D = size(x)

    σ² ~ Exponential(θ)
    α ~ Normal(0, sqrt(σ²))
    β ~ MvNormal(zeros(D), σ² * I)

    for n = 1:N
        y[n] ~ BinomialLogit(1, dot(x[n,:], β) + α)
    end
end

# Sampling parameter settings
n_samples = 10_000

# Sampling
BenchmarkSuite["nuts"]["hrl"] = @benchmarkable sample(hlr_nuts(x, y, 1/0.1), NUTS(0.65), n_samples)


text:
@testset "ModelKit - SLP" begin

    @testset "CompiledHomotopy/InterpretedHomotopy" begin
        @var x y a b c
        f = [(2 * x^2 + b^2 * y^3 + 2 * a * x * y)^3, (a + c)^4 * x + y^2]

        @var s sp[1:3] sq[1:3]
        g = subs(f, [a, b, c] => s .* sp .+ (1 .- s) .* sq)

        h = Homotopy(g, [x, y], s, [sp; sq])

        p = [5.2, -1.3, 9.3]
        q = [2.6, 3.3, 2.3]

        for H in [InterpretedHomotopy(h), CompiledHomotopy(h)]
            @test size(H) == (2, 2)

            v = [0.192, 2.21]
            t = 0.232

            u = zeros(ComplexF64, 2)
            U = zeros(ComplexF64, 2, 2)

            evaluate!(u, H, v, t, [p; q])
            @test u ≈ f([x, y] => v, [a, b, c] => t * p + (1 - t) * q)

            taylor!(u, Val(1), H, TaylorVector{1}(Matrix(v')), t, [p; q])
            @test u ≈ let
                @var s sp[1:3] sq[1:3]
                pt = s .* sp .+ (1 .- s) .* sq
                differentiate(subs(f, [a, b, c] => pt), s)(
                    [x, y] => v,
                    sp => p,
                    sq => q,
                    s => t,
                )
            end

            evaluate_and_jacobian!(u, U, H, v, t, [p; q])
            @test U ≈
                  differentiate(f, [x, y])([x, y] => v, [a, b, c] => t * p + (1 - t) * q)
        end
    end

    @testset "Codegen helpers" begin
        @test ModelKit.sqr(3 + 2im) == (3 + 2im)^2
        @test ModelKit.sqr(3) == 3^2
    end

    @testset "Homotopy codegen (Katsura(3))" begin
        n = 3
        @var x[0:n] ẋ[0:n] ẍ[0:n] x3[0:n] t γ
        K = [
            (
                sum(x[abs(l)+1] * x[abs(m - l)+1] for l = -n:n if abs(m - l) <= n) - x[m+1] for m = 0:n-1
            )...,
            x[1] + 2 * sum(x[i+1] for i = 1:n) - 1,
        ]

        h = γ .* t .* [x[1:n] .^ 2 .- 1; x[n+1] - 1] + (1 - t) .* K
        H = ModelKit.Homotopy(h, x, t, [γ])
        TH = CompiledHomotopy(H)

        tx3 = TaylorVector{4}(zeros(Expression, 4, 4))
        y, y1, y2, y3 = vectors(tx3)
        y .= x
        y1 .= ẋ
        y2 .= ẍ
        y3 .= x3

        dt1 =
            ModelKit.taylor!(zeros(Expression, 4), Val(1), TH, TaylorVector{1}(tx3), t, [γ])
        dt2 =
            ModelKit.taylor!(zeros(Expression, 4), Val(2), TH, TaylorVector{2}(tx3), t, [γ])
        dt3 =
            ModelKit.taylor!(zeros(Expression, 4), Val(3), TH, TaylorVector{3}(tx3), t, [γ])
        dt4 =
            ModelKit.taylor!(zeros(Expression, 4), Val(4), TH, TaylorVector{4}(tx3), t, [γ])

        @var λ
        Hd1 = subs(H.expressions, t => t + λ)
        true_dt1 = subs(differentiate(Hd1, λ, 1), λ => 0)

        Hd2 = subs(H.expressions, x => x .+ λ .* ẋ, t => t + λ)
        true_dt2 = subs(differentiate(Hd2, λ, 2), λ => 0) / 2

        Hd3 = subs(H.expressions, x => x .+ λ .* ẋ .+ λ^2 .* ẍ, t => t + λ)
        true_dt3 = subs(differentiate(Hd3, λ, 3), λ => 0) / 6

        Hd4 = subs(H.expressions, x => x .+ λ .* ẋ .+ λ^2 .* ẍ .+ λ^3 .* x3, t => t + λ)
        true_dt4 = subs(differentiate(Hd4, λ, 4), λ => 0) / 24

        @test expand.(TH(x, t, [γ])) == expand.(h)
        @test expand.(ModelKit.jacobian(TH, x, t, [γ])) == expand.(differentiate(h, x))
        @test expand.(dt1) == expand.(true_dt1)
        @test expand.(dt2) == expand.(true_dt2)
        @test expand.(dt3) == expand.(true_dt3)
        @test expand.(dt4) == expand.(true_dt4)
    end

    @testset "taylor! - system" begin
        @var x[1:2] ẋ[1:2] ẍ[1:2] x3[1:2] p[1:2] ṗ[1:2]
        f = [(x[1] + x[2])^3 + x[1]^2 + x[1] + 5 * x[2] + 3 * p[1], 2 * x[1]^2 + p[2]]
        F = System(f, x, p)
        TF = CompiledSystem(F)

        tx = TaylorVector{4}([x'; ẋ'; ẍ'; x3'])
        tv = TaylorVector{5}(Expression, 2)
        v, v1, v2, v3, v4 = vectors(tv)
        @var λ

        ModelKit.taylor!(TaylorVector{3}(tv), Val(2), TF, TaylorVector{2}(tx), p)
        Fd2 = subs(F.expressions, x => x .+ λ .* ẋ)
        true_v1 = subs(differentiate(Fd2, λ, 1), λ => 0)
        true_v2 = subs(differentiate(Fd2, λ, 2), λ => 0) / 2
        @test expand.(v) == expand.(f)
        @test expand.(v1) == expand.(true_v1)
        @test expand.(v2) == expand.(true_v2)

        ModelKit.taylor!(TaylorVector{4}(tv), Val(3), TF, TaylorVector{3}(tx), p)
        Fd3 = subs(F.expressions, x => x .+ λ .* ẋ .+ λ^2 .* ẍ)
        true_v1 = subs(differentiate(Fd3, λ, 1), λ => 0)
        true_v2 = subs(differentiate(Fd3, λ, 2), λ => 0) / 2
        true_v3 = subs(differentiate(Fd3, λ, 3), λ => 0) / 6
        @test expand.(v) == expand.(f)
        @test expand.(v1) == expand.(true_v1)
        sub = variables(v2) => rand(1:10_000, 6)
        @test v2(sub) == true_v2(sub)
        @test v3(sub) == true_v3(sub)
    end

    @testset "Monomials" begin
        @var x y
        F = System([x^2 - 1, y])
        u = zeros(ComplexF64, 2)
        U = zeros(ComplexF64, 2, 2)
        w = randn(ComplexF64, 2)
        evaluate_and_jacobian!(u, U, CompiledSystem(F), w)
        @test sum(abs, U) > 0
        U .= 0
        evaluate_and_jacobian!(u, U, InterpretedSystem(F), w)
        @test sum(abs, U) > 0
    end

    @testset "Evaluation of Arb" begin
        @var x y
        # define the polynomials
        f₁ = (x^4 + y^4 - 1) * (x^2 + y^2 - 2) + x^5 * y
        f₂ = x^2 + 2x * y^2 - 2 * y^2 - 1 / 2
        F = System([f₁, f₂])
        for I in [
            Arblib.AcbMatrix(randn(ComplexF64, 2, 1)),
            Arblib.AcbRefVector(randn(ComplexF64, 2)),
        ]
            @test all(!iszero, F(I))
            @test F(real.(I)) isa Arblib.ArbVector
            @test all(!iszero, System(F)(real.(I)))

            u = Arblib.AcbVector(2)
            evaluate!(u, InterpretedSystem(System(F)), I)
            @test all(!iszero, u)
            v = Arblib.ArbVector(2)
            evaluate!(v, InterpretedSystem(System(F)), real.(I))
            @test all(!iszero, v)
            U = Arblib.AcbMatrix(2, 2)
            jacobian!(U, InterpretedSystem(System(F)), I)
            @test all(!iszero, U)
            V = Arblib.ArbMatrix(2, 2)
            jacobian!(V, InterpretedSystem(System(F)), real.(I))
            @test all(!iszero, V)
        end
    end
end


text:
"""
    upload(upfun::Function, gl::GitLink; verbose = false, tout = 60.0)

Allow to modify the stage directory of the GitLink.
The function `upfun(stage_dir)` will be executed and it should copy/create/modify
the files into `stage_dir`.
It is recommended that `upfun` not to be an expensive function.
Then an attempt is made to synchronize the changes with the upstream repo.
This method will update all events tokens `pull/stage/push`.
This method will sleep till (or timeout `tout`) the GitLink lock is free (which must by must of the time, but...).
Returns `true` if the action was successful.
"""
function upload(upfun::Function, gl::GitLink; verbose = false, tout = 60.0)
    stage(upfun, gl; tout) || return false
    return sync_link(gl; verbose, tout, force = true)
end


text:
export DiffEqLinearOperator, ConstrainedODEFunction

DEFAULT_PARAM_UPDATE_FUNC(q,u,p,t) = q
DEFAULT_PARAM_UPDATE_FUNC(u,p,t) = p

### Abstract algorithms ###
abstract type ConstrainedOrdinaryDiffEqAlgorithm <: OrdinaryDiffEq.OrdinaryDiffEqAlgorithm end

### Abstract caches ###
# The sc parameter specifies whether it contains static constraint operators or not
# If false, then it expects that the state vector contains a component for updating the opertors
abstract type ConstrainedODEMutableCache{sc,solverType} <: OrdinaryDiffEqMutableCache end
abstract type ConstrainedODEConstantCache{sc,solverType} <: OrdinaryDiffEqConstantCache end

#### Operator and function types ####

mutable struct DiffEqLinearOperator{T,aType} <: AbstractDiffEqLinearOperator{T}
    L :: aType
    DiffEqLinearOperator(L::aType; update_func=DEFAULT_PARAM_UPDATE_FUNC,
                          dtype=Float64) where {aType} = new{dtype,aType}(L)
end

(f::DiffEqLinearOperator)(du,u,p,t) = (dy = state(du); mul!(dy,f.L,state(u)))
(f::DiffEqLinearOperator)(u,p,t) = (du = deepcopy(u); zero_vec!(du); dy = state(du); mul!(dy,f.L,state(u)); return du)

import Base: exp
exp(f::DiffEqLinearOperator,args...) = exp(f.L,args...)
exp(f::ArrayPartition,t,x::ArrayPartition) =
                ArrayPartition((exp(Li,t,xi) for (Li,xi) in zip(f.L.x,x.x))...)

exp(f::ODEFunction,args...) = exp(f.f,args...)

has_exp(::DiffEqLinearOperator) = true

exp(L::AbstractMatrix,t,x) = exp(factorize(L)*t)
exp(L::UniformScaling,t,x) = exp(Diagonal(L,length(x))*t)

#=
ConstrainedODEFunction

A function of this type should be able to take arguments (du,u,p,t) (for in-place)
or (u,p,t) for out-of-place, and distribute the parts of u and du as
needed to the component functions. u and du will both be of type ArrayPartition,
with two parts: state(u) and constraint(u).

In some cases we wish to solve a separate (unconstrained) system, e.g., to
update the constraint operators B1 and B2. This update is done via the
parameters, using param_update_func. In this case, u and du are ArrayPartition
with the first part corresponding to the constrained system (and of ArrayPartition type)
and the second part is for the unconstrained system. We supply `r1` as
an ArrayPartition of the component functions (in the same order). The arguments of
each component function of `r1` should only take in and return their own parts
of this state.

In the future we can expand this to take in multiple (coupled) systems of this form,
as in FSI problems. In these cases, the component functions should be supplied as
ArrayPartition's of the functions for each system. Each of the component functions of each system
should be able to take in the *full* state and/or constraint and return the
*full* state/constraint, as appropriate, in order to enable coupling of the
systems.
=#

"""
        ConstrainedODEFunction(r1,r2,B1,B2[,L])

This specifies the functions and operators that comprise an ODE problem with the form

``
\\dfrac{dy}{dt} = Ly - B_1 z + r_1(y,t)
``

``
B_2 y = r_2(t)
``

The optional linear operator `L` defaults to zeros. The `B1` and `B2` functions must be of the respective
in-place forms `B1(dy,z,p)` (to compute the action of `B1` on `z`) and `B2(dz,y,p)` (to compute that action
of `B2` on `y`). The function `r1` must of the in-place form `r1(dy,y,p,t)`, and `r2` must be in the in-place form
`r2(dz,p,t)`. Alternatively, one can supply out-of-place forms, respectively, as `B1(z,p)`, `B2(y,p)`,
`r1(y,p,t)` and `r2(p,t)`.

An optional keyword argument `param_update_func` can be used to set a function that updates problem parameters with
the current solution. This function must take the in-place form `f(q,u,p,t)` or out of place form
`f(u,p,t)` to create some `q` based on `u`. (Note that `q` might enter the function simply as `p`.) This function can
be used to update `B1` and `B2` with state information, for example.

We can also include another (unconstrained) set of equations to the set above:

``
\\dfrac{dx}{dt} = r_{1x}(x,t)
``

In this case, we would pass the pair of `r1` functions as an `ArrayPartition`.
"""
struct ConstrainedODEFunction{iip,static,F1,F2,TMM,C,Ta,Tt,TJ,JVP,VJP,JP,SP,TW,TWt,TPJ,S,TCV,PF} <: AbstractODEFunction{iip}
    odef :: F1
    conf :: F2
    mass_matrix::TMM
    cache::C
    analytic::Ta
    tgrad::Tt
    jac::TJ
    jvp::JVP
    vjp::VJP
    jac_prototype::JP
    sparsity::SP
    Wfact::TW
    Wfact_t::TWt
    paramjac::TPJ
    syms::S
    colorvec::TCV
    param_update_func :: PF
end

function ConstrainedODEFunction(r1,r2,B1,B2,L=DiffEqLinearOperator(0*I);
                                                 param_update_func = DEFAULT_PARAM_UPDATE_FUNC,
                                                 mass_matrix=I,_func_cache=nothing,
                                                 analytic=nothing,
                                                 tgrad = nothing,
                                                 jac = nothing,
                                                 jvp=nothing,
                                                 vjp=nothing,
                                                 jac_prototype = nothing,
                                                 sparsity=jac_prototype,
                                                 Wfact = nothing,
                                                 Wfact_t = nothing,
                                                 paramjac = nothing,
                                                 syms = nothing,
                                                 colorvec = nothing)

    allempty(r2,B1,B2) || noneempty(r2,B1,B2) || error("Inconsistent null operators")
    unconstrained = allempty(r2,B1,B2)

    allinplace(r1,r2,B1,B2) || alloutofplace(r1,r2,B1,B2) || error("Inconsistent function signatures")
    iip = allinplace(r1,r2,B1,B2)

    static = param_update_func == DEFAULT_PARAM_UPDATE_FUNC ? true : false

    L_local = (L isa DiffEqLinearOperator) ? L : DiffEqLinearOperator(L)

    local_cache = deepcopy(_func_cache)
    zero_vec!(local_cache)

    odef_nl = SplitFunction(_complete_r1(r1,Val(iip),_func_cache=local_cache),
                            _complete_B1(B1,Val(iip));_func_cache=deepcopy(local_cache))
    odef = SplitFunction(L_local, odef_nl ;_func_cache=deepcopy(local_cache))
    conf = SplitFunction(_complete_r2(r2,Val(iip)),_complete_B2(B2,Val(iip));_func_cache=deepcopy(local_cache))

    ConstrainedODEFunction{iip,static,typeof(odef),
                           typeof(conf),typeof(mass_matrix),typeof(local_cache),
                           typeof(analytic),typeof(tgrad),typeof(jac),typeof(jvp),typeof(vjp),
                           typeof(jac_prototype),typeof(sparsity),
                           typeof(Wfact),typeof(Wfact_t),typeof(paramjac),typeof(syms),
                           typeof(colorvec),typeof(param_update_func)}(odef,conf,mass_matrix,local_cache,
                           analytic,tgrad,jac,jvp,vjp,jac_prototype,
                           sparsity,Wfact,Wfact_t,paramjac,syms,colorvec,param_update_func)
end

# For unconstrained systems
ConstrainedODEFunction(r1,L=DiffEqLinearOperator(0*I);kwargs...) =
      ConstrainedODEFunction(r1,nothing,nothing,nothing,L;kwargs...)

function Base.show(io::IO, m::MIME"text/plain",f::ConstrainedODEFunction{iip,static}) where {iip,static}
    iips = iip ? "in-place" : "out-of-place"
    statics = static ? "static" : "variable"
    println(io,"Constrained ODE function of $iips type and $statics constraints")
end

# Here is where we define the structure of the function
@inline _fetch_ode_L(f::ConstrainedODEFunction) = f.odef.f1
@inline _fetch_ode_r1(f::ConstrainedODEFunction) = f.odef.f2.f.f1
@inline _fetch_ode_neg_B1(f::ConstrainedODEFunction) = f.odef.f2.f.f2
@inline _fetch_constraint_neg_B2(f::ConstrainedODEFunction) = f.conf.f2
@inline _fetch_constraint_r2(f::ConstrainedODEFunction) = f.conf.f1

for fcn in (:_ode_L,:_ode_r1,:_ode_neg_B1,:_constraint_neg_B2,:_constraint_r2)
  fetchfcn = Symbol("_fetch",string(fcn))
  iipfcn = Symbol(string(fcn),"!")
  @eval $iipfcn(du,f::ConstrainedODEFunction,u,p,t) = $fetchfcn(f)(du,u,p,t)
  @eval $fcn(f::ConstrainedODEFunction,u,p,t) = $fetchfcn(f)(u,p,t)
end

allempty(::Nothing,::Nothing,::Nothing) = true
allempty(r2,B1,B2) = false
noneempty(r2,B1,B2) = !(isnothing(r2) || isnothing(B1) || isnothing(B2))

allinplace(r1,r2,B1,B2) = _isinplace_r1(r1) && _isinplace_r2(r2) && _isinplace_B1(B1) && _isinplace_B2(B2)
alloutofplace(r1,r2,B1,B2) = _isoop_r1(r1) && _isoop_r2(r2) && _isoop_B1(B1) && _isoop_B2(B2)

allinplace(r1,::Nothing,::Nothing,::Nothing) = _isinplace_r1(r1)
alloutofplace(r1,::Nothing,::Nothing,::Nothing) = _isoop_r1(r1)

for (f,nv) in ((:r1,4),(:r2,3),(:B1,3),(:B2,3))
  iipfcn = Symbol("_isinplace_",string(f))
  oopfcn = Symbol("_isoop_",string(f))
  completefcn = Symbol("_complete_",string(f))
  @eval $iipfcn(fcn) = isinplace(fcn,$nv)
  @eval $iipfcn(fcn::ArrayPartition) = all(($iipfcn(x) for x in fcn.x))
  @eval $oopfcn(fcn) = numargs(fcn) == $(nv-1)
  @eval $oopfcn(fcn::ArrayPartition) = all(($oopfcn(x) for x in fcn.x))
  @eval $completefcn(fcn,::Val{iip};_func_cache=nothing) where {iip} = $completefcn(fcn,Val(iip),_func_cache)
  @eval $completefcn(::Nothing,::Val{false},_func_cache) = (u,p,t) -> zero(u)
end

_complete_r1(r1,::Val{true},_func_cache) = (du,u,p,t) -> (dy = state(du); r1(dy,state(u),p,t))
_complete_r1(r1,::Val{false},_func_cache) = (u,p,t) -> (du = deepcopy(u); zero_vec!(du); state(du) .= r1(state(u),p,t); return du)
_complete_r1(r1::ArrayPartition,::Val{true},_func_cache) =
            SplitFunction((du,u,p,t) ->(dy = state(du); dx = aux_state(du); zero_vec!(dx); state_r1(r1)(dy,state(u),p,t)),
                          (du,u,p,t) ->(dy = state(du); dx = aux_state(du); zero_vec!(dy); aux_r1(r1)(dx,aux_state(u),p,t));
                          _func_cache=deepcopy(_func_cache))
_complete_r1(r1::ArrayPartition,::Val{false},_func_cache) =
            SplitFunction((u,p,t) -> (du = deepcopy(u); zero_vec!(du); state(du) .= state_r1(r1)(state(u),p,t); return du),
                          (u,p,t) -> (du = deepcopy(u); zero_vec!(du); aux_state(du) .= aux_r1(r1)(aux_state(u),p,t); return du))

_complete_r2(r2,::Val{true},_func_cache) = (du,u,p,t) -> (dz = constraint(du); r2(dz,p,t))
_complete_r2(::Nothing,::Val{true},_func_cache) = (du,u,p,t) -> zero_vec!(constraint(du))
_complete_r2(r2,::Val{false},_func_cache) = (u,p,t) -> (du = deepcopy(u); zero_vec!(du); constraint(du) .= r2(p,t); return du)

_complete_B1(B1,::Val{true},_func_cache) = (du,u,p,t) -> (dy = state(du); dx = aux_state(du); zero_vec!(dx); B1(dy,constraint(u),p); dy .*= -1.0)
_complete_B1(::Nothing,::Val{true},_func_cache) = (du,u,p,t) -> (zero_vec!(aux_state(du)); zero_vec!(state(du)))
_complete_B1(B1,::Val{false},_func_cache) = (u,p,t) -> (du = deepcopy(u); zero_vec!(du); state(du) .= -B1(constraint(u),p); return du)

_complete_B2(B2,::Val{true},_func_cache) = (du,u,p,t) -> (dz = constraint(du); B2(dz,state(u),p); dz .*= -1.0)
_complete_B2(::Nothing,::Val{true},_func_cache) = (du,u,p,t) -> zero_vec!(constraint(du))
_complete_B2(B2,::Val{false},_func_cache) = (u,p,t) -> (du = deepcopy(u); zero_vec!(du); constraint(du) .= -B2(state(u),p); return du)

function (f::ConstrainedODEFunction)(du,u,p,t)
    zero_vec!(f.cache)
    zero_vec!(du)
    f.odef(f.cache,u,p,t)
    f.conf(du,u,p,t)
    du .+= f.cache
end

(f::ConstrainedODEFunction)(u,p,t) = f.odef(u,p,t) + f.conf(u,p,t)

@inline isstatic(f::ConstrainedODEFunction) = f.param_update_func == DEFAULT_PARAM_UPDATE_FUNC ? true : false


text:
using AstroTime: Epoch
using ..Bodies:
    CelestialBody,
    NAIFId,
    SolarSystemBarycenter,
    naifid,
    path_ids,
    ssb
import ..Interfaces:
    AbstractEphemeris,
    position,
    position!,
    state,
    state!,
    velocity,
    velocity!

for f in (:position, :velocity)
    fmut = Symbol(f, "!")
    @eval begin
        function $fmut(arr,
                       eph::AbstractEphemeris,
                       ep::Epoch,
                       from::T,
                       to::T;
                       kwargs...
                      ) where T<:CelestialBody
            arr .+= 0.0
        end

        function $fmut(arr,
                       eph::AbstractEphemeris,
                       ep::Epoch,
                       from::CelestialBody,
                       to::CelestialBody;
                       kwargs...
                      )
            path = path_ids(from, to)
            for (origin, target) in zip(path[1:end-1], path[2:end])
                $fmut(arr, eph, ep, origin, target; kwargs...)
            end
            arr
        end

        function $f(eph::AbstractEphemeris,
                    ep::Epoch,
                    from::CelestialBody,
                    to::CelestialBody;
                    kwargs...
                   )
            $fmut(zeros(3), eph, ep, from, to; kwargs...)
        end

        function $fmut(arr, eph::AbstractEphemeris, ep::Epoch, to::CelestialBody; kwargs...)
            $fmut(arr, eph, ep, ssb, to; kwargs...)
        end

        function $f(eph::AbstractEphemeris, ep::Epoch, to::CelestialBody; kwargs...)
            $fmut(zeros(3), eph, ep, ssb, to; kwargs...)
        end
    end
end

function state!(pos, vel, eph::AbstractEphemeris, ep::Epoch, from::NAIFId, to::NAIFId; kwargs...)
    position!(pos, eph, ep, from, to; kwargs...), velocity!(vel, eph, ep, from, to; kwargs...)
end

function state!(pos,
                vel,
                eph::AbstractEphemeris,
                ep::Epoch,
                from::CelestialBody,
                to::CelestialBody;
                kwargs...)
    path = path_ids(from, to)
    for (origin, target) in zip(path[1:end-1], path[2:end])
        state!(pos, vel, eph, ep, origin, target; kwargs...)
    end
    pos, vel
end

state!(pos,
       vel,
       ::AbstractEphemeris,
       ::Epoch,
       ::T, ::T; kwargs...) where {T<:CelestialBody} = pos, vel

function state(eph::AbstractEphemeris, ep::Epoch, from::CelestialBody, to::CelestialBody; kwargs...)
    state!(zeros(3), zeros(3), eph, ep, from, to; kwargs...)
end

function state!(pos, vel, eph::AbstractEphemeris, ep::Epoch, to::CelestialBody; kwargs...)
    state!(pos, vel, eph, ep, ssb, to; kwargs...)
end

function state(eph::AbstractEphemeris, ep::Epoch, to::CelestialBody; kwargs...)
    state!(zeros(3), zeros(3), eph, ep, ssb, to; kwargs...)
end


text:
const DOC_IMAGES = OrderedDict(

    "lorenz_attractor" => [quote
        using Plots
        gr()

        # define the Lorenz attractor
        mutable struct Lorenz
            dt; σ; ρ; β; x; y; z
        end

        function step!(l::Lorenz)
            dx = l.σ * (l.y - l.x)       ; l.x += l.dt * dx
            dy = l.x * (l.ρ - l.z) - l.y ; l.y += l.dt * dy
            dz = l.x * l.y - l.β * l.z   ; l.z += l.dt * dz
        end

        attractor = Lorenz((dt = 0.02, σ = 10., ρ = 28., β = 8//3, x = 1., y = 1., z = 1.)...)

        # initialize a 3D plot with 1 empty series
        plt = plot3d(1, xlim = (-25, 25), ylim = (-25, 25), zlim=(0, 50), title = "Lorenz Attractor", marker = 2)

        # build an animated gif by pushing new points to the plot, saving every 10th frame
        anim = @gif for i in 1:1500
            step!(attractor)
            push!(plt, attractor.x, attractor.y, attractor.z)
        end every 10
    end],

    "waves" => [quote
        using Plots, ProgressMeter
        no_initialize_backend(:pyplot)
        default(leg = false, ticks = nothing) #change to the pyplot backend and define some defaults
        x = y = range(-5, stop = 5, length = 40)
        zs = zeros(0,40)
        n = 100

        # create a progress bar for tracking the animation generation
        prog = Progress(n,1)

        anim = @gif for i in range(0, stop = 2π, length = n)
            f(x,y) = sin(x + 10 * sin(i)) + cos(y)

            # create a plot with 3 subplots and a custom layout
            l = @layout [a{0.7w} b; c{0.2h}]
            p = plot(x, y, f, st = [:surface, :contourf], layout=l)

            # induce a slight oscillating camera angle sweep, in degrees (azimuth, altitude)
            plot!(p[1], camera=(15 * cos(i), 40))

            # add a tracking line
            fixed_x = zeros(40)
            z = map(f, fixed_x, y)
            plot!(p[1], fixed_x, y, z, line = (:black, 5, 0.2))
            vline!(p[2], [0], line = (:black, 5))

            # add to and show the tracked values over time
            global zs = vcat(zs, z')
            plot!(p[3], zs, alpha = 0.2, palette = cgrad(:blues).colors)

            # increment the progress bar
            next!(prog)
        end
    end],

    "decision" => [quote
        gr()
        Random.seed!(7) # for a consistent result

        P = 40;  R = 50;  N = P*R;  r = 0:0.004:1
        points = rand(ComplexF64, P, R)

        anim = @animate(for t = 0:0.03:13
            # create a simple classifier to return the region for any point (x, y)
            midpoints = vec(sum(points; dims=1)) / P
            classify(x, y) = argmin(abs.(x + y*im .- midpoints))

            # draw decision boundary and points
            contour(r, r, classify, c=:cyclic2, fill=true, nlev=R, leg=:none)
            scatter!(reim(points)..., c=cvec(:cyclic2, R)', lims=(0,1))

            # update position of points
            target(d) = 0.65*cis(4*sin(t/2+d)+d) + 0.5 + 0.5im
            points[:] .+= 0.01*(target.(0:2π/(N-1):2π) .- points[:])
        end)
    end],

    "decision-poster" => [quote
        gr()
        Random.seed!(7) # for a consistent result

        P = 40;  R = 50;  N = P*R;  r = 0:0.004:1
        points = rand(ComplexF64, P, R)

        for t = 0:0.03:13
            # create a simple classifier to return the region for any point (x, y)
            midpoints = vec(sum(points; dims=1)) / P
            classify(x, y) = argmin(abs.(x + y*im .- midpoints))

            # draw decision boundary and points
            contour(r, r, classify, c=:cyclic2, fill=true, nlev=R, leg=:none)
            scatter!(reim(points)..., c=cvec(:cyclic2, R)', lims=(0,1))

            # poster image, shown before video starts
            if t ≈ 7.5
                contour(r, r, classify, c=:cyclic2, fill=true, nlev=R, leg=:none)
                return scatter!(reim(points)..., c=cvec(:cyclic2, R)', lims=(0,1))
            end

            # update position of points
            target(d) = 0.65*cis(4*sin(t/2+d)+d) + 0.5 + 0.5im
            points[:] .+= 0.01*(target.(0:2π/(N-1):2π) .- points[:])
        end
        png(local_path("PlotDocs", "index", "decision-poster.png"))
    end],

    "iris" => [quote
        using StatsPlots, RDatasets
        gr()

        # load a dataset
        iris = dataset("datasets", "iris")

        # Scatter plot with some custom settings
        @df iris scatter(:SepalLength, :SepalWidth,
            group = :Species,
            title = "My awesome plot",
            xlabel = "Length",
            ylabel = "Width",
            m = (0.5, [:cross :hex :star7], 12),
            bg = RGB(0.2, 0.2, 0.2))
    end],

    "lines_1" => [quote
        gr()
        x = 1:10; y = rand(10) # These are the plotting data
        plot(x, y)
    end],

    "lines_2" => [quote
        gr()
    end,
    quote
        x = 1:10; y = rand(10, 2) # 2 columns means two lines
        plot(x, y)
    end],

    "lines_3" => [quote
        gr()
        x = 1:10; y = rand(10, 2)
    end,
    quote
        z = rand(10)
        plot!(x, z)
    end],

    "lines_4" => [quote
        gr()
    end,
    quote
        x = 1:10; y = rand(10,2) # 2 columns means two lines
        p = plot(x, y)
        z = rand(10)
        plot!(p, x, z)
    end],

    "attr_1" => [quote
        gr()
    end,
    quote
        x = 1:10; y = rand(10,2) # 2 columns means two lines
        plot(x, y, title = "Two Lines", label = ["Line 1" "Line 2"], lw = 3)
    end],

    "attr_2" => [quote
        gr()
    end,
    quote
        xlabel!("My x label")
    end],

    "backends_1" => [quote
        using Plots
        no_initialize_backend(:plotlyjs)
        x = 1:10; y = rand(10,2) # 2 columns means two lines
        plot(x, y, title = "This is Plotted using Plotly")
    end],

    "backends_2" => [quote
        gr() # Set the backend to GR
        plot(x, y, title = "This is Plotted using GR") # This plots using GR
    end],

    "scatter_1" => [quote
        gr() # We will continue onward using the GR backend
        plot(x, y, seriestype = :scatter, title = "My Scatter Plot")
    end],

    "scatter_2" => [quote
        gr()
    end,
    quote
        scatter(x, y, title = "My Scatter Plot")
    end],

    "subplots_1" => [quote
        gr()
    end,
    quote
        y = rand(10, 4)
        plot(x, y, layout = (4, 1))
    end],

    "subplots_2" => [quote
        gr()
    end,
    quote
        p1 = plot(x, y) # Make a line plot
        p2 = scatter(x, y) # Make a scatter plot
        p3 = plot(x, y, xlabel = "This one is labelled", lw = 3, title = "Subtitle")
        p4 = histogram(x, y) # Four histograms each with 10 points? Why not!
        plot(p1, p2, p3, p4, layout = (2, 2), legend = false)
    end],

    "user_recipes_1" => [quote
        gr()
    end,
    quote
        # Pkg.add("StatsPlots")
        using StatsPlots # Required for the DataFrame user recipe
        # Now let's create the DataFrame
        using DataFrames
        df = DataFrame(a = 1:10, b = 10 * rand(10), c = 10 * rand(10))
        # Plot the DataFrame by declaring the points by the column names
        @df df plot(:a, [:b :c]) # x = :a, y = [:b :c]. Notice this is two columns!
    end],

    "user_recipes_2" => [quote
        gr()
        df = DataFrame(a = 1:10, b = 10 * rand(10), c = 10 * rand(10))
    end,
    quote
        @df df scatter(:a, :b, title = "My DataFrame Scatter Plot!") # x = :a, y = :b
    end],

    "type_recipes" => [quote
        gr()
    end,
    quote
        using Distributions
        plot(Normal(3, 5), lw = 3)
    end],

    "plot_recipes" => [quote
        gr()
    end,
    quote
        #Pkg.add("RDatasets")
        using RDatasets
        iris = dataset("datasets", "iris")
        @df iris marginalhist(:PetalLength, :PetalWidth)
    end],

    "series_recipes_1" => [quote
        using StatsPlots
        gr()
    end,
    quote
        y = rand(100, 4) # Four series of 100 points each
        StatsPlots.violin(["Series 1" "Series 2" "Series 3" "Series 4"], y, leg = false)
    end],

    "series_recipes_2" => [quote
        using StatsPlots
        gr()
    end,
    quote
        StatsPlots.boxplot!(["Series 1" "Series 2" "Series 3" "Series 4"], y, leg = false)
    end],

    "columns_are_series" => [quote
        gr()
    end,
    quote
        # 10 data points in 4 series
        xs = 0:(2π / 10):2π
        data = [sin.(xs) cos.(xs) 2sin.(xs) 2cos.(xs)]

        # We put labels in a row vector: applies to each series
        labels = ["Apples" "Oranges" "Hats" "Shoes"]

        # Marker shapes in a column vector: applies to data points
        markershapes = [:circle, :star5]

        # Marker colors in a matrix: applies to series and data points
        markercolors = [:green :orange :black :purple :red :yellow :brown :white]

        plot(xs, data, label = labels, shape = markershapes, color = markercolors, markersize = 10)
    end],

    "groups" => [quote
        no_initialize_backend(:plotlyjs)
    end,
    quote
        function rectangle_from_coords(xb,yb,xt,yt)
            [
                xb yb
                xt yb
                xt yt
                xb yt
                xb yb
                NaN NaN
            ]
        end

        some_rects=[
            rectangle_from_coords(1, 1, 5, 5)
            rectangle_from_coords(10, 10, 15, 15)
        ]
        other_rects=[
            rectangle_from_coords(1, 10, 5, 15)
            rectangle_from_coords(10, 1, 15, 5)
        ]

        plot(some_rects[:,1], some_rects[:,2], label = "some group")
        plot!(other_rects[:,1], other_rects[:,2], label = "other group")
    end],

    "dataframes" => [quote
        gr() # maybe remove again
    end,
    quote
        using StatsPlots, RDatasets
        iris = dataset("datasets", "iris")
        @df iris scatter(:SepalLength, :SepalWidth, group = :Species,
            m = (0.5, [:+ :h :star7], 12), bg = RGB(0.2, 0.2, 0.2))
    end],

    "batman" => [quote
        no_initialize_backend(:pyplot)
        default(leg = false, framestyle = :none, size=(500,500))
    end,
    quote
        P2 = Plots.P2
        function make_batman()
            p = P2[(0,0), (0.5, 0.2), (1, 0), (1,2),  (0.3,1.2), (0.2,2), (0,1.7)]
            m = P2[(p[i]+p[i+1])/2 for i=1:length(p)-1]
            m += P2[(0.2, 1), (0.4, 1), (2, 0), (0.5, -0.6), (0,0), (0,-0.15)]

            pts = P2[]
            for (i,mi) in enumerate(m)
                append!(pts, map(BezierCurve(P2[p[i], m[i], p[i+1]]), range(0, 1, length = 30)))
            end
            x, y = Plots.unzip(pts)
            Shape(vcat(x, -reverse(x)), vcat(y, reverse(y)))
        end
    end,
    quote
        # background and limits
        plt = plot(bg = :black, xlim = (0.1,0.9), ylim = (0.2,1.5))

        # create an ellipse in the sky
        pts = Plots.partialcircle(0,2π,100,0.1)
        x, y = Plots.unzip(pts)
        x = 1.5x .+ 0.7
        y .+= 1.3
        pts = collect(zip(x, y))

        # beam
        beam = Shape([(0.3, 0.0), pts[95], pts[50], (0.3, 0.0)])
        plot!(beam, fillcolor = plot_color(:yellow, 0.3))

        # spotlight
        plot!(Shape(x,y), c=:yellow)
    end,
    quote
        # buildings
        rect(w, h, x, y) = Shape(x .+ [0,w,w,0,0], y .+ [0,0,h,h,0])
        gray(pct) = RGB(pct,pct,pct)
        function windowrange(dim, denom)
            range(0, 1, length = max(3, round(Int, dim/denom)))[2:end-1]
        end

        for k=1:50
            w, h, x, y = 0.1rand()+0.05, 0.8rand()+0.3, rand(), 0.0
            shape = rect(w, h, x, y)
            graypct = 0.3rand()+0.3
            plot!(shape, c = gray(graypct))

            # windows
            I = windowrange(w,0.015)
            J = windowrange(h,0.04)
            pts = vec([(Float64(x+w*i), Float64(y+h*j)) for i=I,j=J])
            windowcolors = Symbol[rand()<0.2 ? :yellow : :black for i=1:length(pts)]
            scatter!(pts, marker = (stroke(0), :rect, windowcolors))
        end
        plt
    end,
    quote
        # Holy plotting, Batman!
        batman = Plots.scale(make_batman(), 0.07, 0.07, (0,0))
        batman = translate(batman, 0.7, 1.23)
        plot!(batman, fillcolor = :black,  size=(500,500))
    end],

    "my_count_1" => [quote
        using Plots, StatsBase
        gr()
        default(size = (300, 300), leg = false)

        @userplot MyCount
        @recipe function f(mc::MyCount)
            # get the array from the args field
            arr = mc.args[1]

            T = typeof(arr)
            if T.parameters[1] == Float64
                seriestype := :histogram
                arr
            else
                seriestype := :bar
                cm = countmap(arr)
                x = sort!(collect(keys(cm)))
                y = [cm[xi] for xi=x]
                x, y
            end
        end
        mycount(rand(500))
    end],

    "my_count_2" => [quote
        gr()
        default(size = (300, 300), leg = false)
    end,
    quote
        mycount(rand(["A", "B", "C"], 100))
    end],

    "my_count_3" => [quote
        gr()
        default(size = (300, 300), leg = false)
    end,
    quote
        mycount(rand(1:500, 500))
    end],

    "pipeline_1" => [quote
        n = 100
        x, y = range(0, stop=1, length=n), randn(n, 3)
        using Plots
        no_initialize_backend(:pyplot)
        plot(x, y, line = (0.5, [4 1 0], [:path :scatter :density]),
            marker = (10, 0.5, [:none :+ :none]), fill = 0.5,
            orientation = :h, title = "My title")
    end],

    "pipeline_2" => [quote
        gr()
        mutable struct MyVecWrapper
            v::Vector{Float64}
        end
        myvec = MyVecWrapper(rand(100))

        @recipe function f(myvec::MyVecWrapper)
            markershape --> :circle
            markersize  --> 30
            myvec.v
        end

        plot(
            plot(myvec.v),
            plot(myvec)
        )
    end],

    "pipeline_3" => [quote
        gr()
    end,
    quote
        scatter(rand(100), group = rand(1:3, 100), marker = (10, 0.3, [:s :o :x]))
    end],

    "layout_1" => [quote
        gr()
    end,
    quote
        # create a 2x2 grid, and map each of the 4 series to one of the subplots
        plot(rand(100, 4), layout = 4)
    end],

    "layout_2" => [quote
        gr()
    end,
    quote
        # create a 4x1 grid, and map each of the 4 series to one of the subplots
        plot(rand(100,4), layout = (4, 1))
    end],

    "layout_3" => [quote
        gr()
    end,
    quote
        plot(rand(100, 4), layout = grid(4, 1, heights = [0.1, 0.4, 0.4, 0.1]))
    end],

    "layout_4" => [quote
        gr()
    end,
    quote
        l = @layout [a{0.3w} [grid(3, 3)
                                b{0.2h}]]
        plot(
            rand(10, 11),
            layout = l, legend = false, seriestype = [:bar :scatter :path],
            title = ["($i)" for j = 1:1, i = 1:11], titleloc = :right, titlefont = font(8)
        )
    end],

    "layout_5" => [quote
        # boxplot is defined in StatsPlots
        using StatsPlots
        gr()
        default(leg = false, bg = :lightgrey)

        # Create a filled contour and boxplot side by side.
        plot(contourf(randn(10, 20)), boxplot(rand(1:4, 1000), randn(1000)))

        # Add a histogram inset on the heatmap.
        # We set the (optional) position relative to bottom-right of the 1st subplot.
        # The call is `bbox(x, y, width, height, origin...)`, where numbers are treated as "percent of parent"
        histogram!(randn(1000), inset = (1, bbox(0.05, 0.05, 0.5, 0.25, :bottom, :right)), ticks = nothing, subplot = 3, bg_inside = nothing)

        # Add sticks floating in the window (inset relative to the window, as opposed to being relative to a subplot)
        sticks!(randn(100), inset = bbox(0,-0.2, 200Plots.px, 100Plots.px, :center), ticks = nothing, subplot = 4)
    end],
)

const DOC_ANIMATION_IMAGES = ["lorenz_attractor", "waves", "rand_anim"]
const DOC_MP4_IMAGES = ["decision"]

const DOC_IMAGE_FILES = Dict(
    "lorenz_attractor" => "index",
    "waves" => "index",
    "decision" => "index",
    "decision-poster" => "index",
    "iris" => "index",
    "lines_1" => "tutorial",
    "lines_2" => "tutorial",
    "lines_3" => "tutorial",
    "lines_4" => "tutorial",
    "attr_1" => "tutorial",
    "attr_2" => "tutorial",
    "backends_1" => "tutorial",
    "backends_2" => "tutorial",
    "scatter_1" => "tutorial",
    "scatter_2" => "tutorial",
    "subplots_1" => "tutorial",
    "subplots_2" => "tutorial",
    "user_recipes_1" => "tutorial",
    "user_recipes_2" => "tutorial",
    "type_recipes" => "tutorial",
    "plot_recipes" => "tutorial",
    "series_recipes_1" => "tutorial",
    "series_recipes_2" => "tutorial",
    "columns_are_series" => "input",
    "groups" => "input",
    "dataframes" => "input",
    "batman" => "input",
    "rand_anim" => "animations",
    "my_count_1" => "contributing",
    "my_count_2" => "contributing",
    "my_count_3" => "contributing",
    "pipeline_1" => "pipeline",
    "pipeline_2" => "pipeline",
    "pipeline_3" => "pipeline",
    "layout_1" => "layout",
    "layout_2" => "layout",
    "layout_3" => "layout",
    "layout_4" => "layout",
    "layout_5" => "layout",
)


text:
"""
    newton(d::Union{ChiSquare, ModifiedChiSquare}, m::DRO; maxiter::Integer = 1000, atol::Real = 1e-6)

A simple bisection method for finding the root of the function `h` for the DRO problem with
χ²-distance or modified χ²-distance.
"""
function newton(d::Union{ChiSquare, ModifiedChiSquare}, m::DRO; maxiter::Integer = 1000, atol::Real = 1e-6)
    stats.optimizer  = "newton"
    f(λ)  = h(d, m, λ)
    ∇f(λ) = ∇h(d, m, λ)

    λ     = initial(d,m)
    f_val = f(λ)
    add_eval()

    f_val == 0 && return λ

    for k in 1:maxiter
        abs(f_val) > atol || return λ

        λ    -= f_val/∇f(λ)
        f_val = f(λ)

        add_eval()
    end
    @warn "Newton method for the DRO model with $(name(d)) reached maximal
           number of iteration. The solution may not be accurate: h(μ) = $(f(λ))"
    return λ
end

"""
    newton(d::Ltwo, m::DRO; maxiter::Integer = 1000, atol::Real = 1e-6)

A simple bisection method for finding the root of the function `h` for the DRO problem with
l-2 norm.
"""
function newton(d::Ltwo, m::DRO; maxiter::Integer = 1000, atol::Real = 1e-6)
    stats.optimizer  = "newton"
    f(μ, λ)  = h(d, m, μ; λ = λ)
    ∇f(μ, λ) = ∇h(d, m, μ; λ = λ)

    μ     = initial(d,m)
    λ     = g(d, m, μ)
    f_val = f(μ, λ)
    add_eval()

    f_val == 0 && return μ

    for k in 1:maxiter
        abs(f_val) > atol || return μ

        μ    -= f_val/∇f(μ, λ)
        λ     = g(d, m, μ)
        f_val = f(μ, λ)
        add_eval()
    end
    @warn "Newton method for the DRO model with $(name(d)) reached maximal
           number of iteration. The solution may not be accurate: h(μ) = $(f(μ))"
    return μ
end

"""
    bisection(d::Union{KullbackLeibler, Burg, Hellinger}, m::DRO; kwargs...)

A simple bisection method for finding the root of the function `h` for the DRO problem with
Kullback-Leibler divergence, Burg entropy or Hellinger distance.
"""
function bisection(d::Union{KullbackLeibler, Burg, Hellinger}, m::DRO; kwargs...)
    f(μ) = h(d, m, μ)
    a, b = bounds(d,m)

    return bisection(f, a, b; kwargs...)
end

"""
    bisection(m::Simplex; kwargs...)

A simple bisection method for finding the root of the function `h` for the Simplex problems.
"""
function bisection(m::Simplex; kwargs...)
    f(μ) = h(m, μ)
    a, b = bounds(m)

    return bisection(f, a, b; kwargs...)
end

"""
    bisection(f::Function, a::Real, b::Real; maxiter::Integer = 1000, atol::Real = 1e-8)

A simple bisection method for finding the root of the function `f`.
"""
function bisection(f::Function, a::Real, b::Real; maxiter::Integer = 1000, atol::Real = 1e-8)
    stats.optimizer  = "bisection"
    local c

    f_a = f(a)
    f_b = f(b)
    add_eval(2)
    @assert f_a*f_b <= 0 "The interval [a,b] is not a bracketing interval."

    f_a == 0 && return a
    f_b == 0 && return b

    for k in 1:maxiter
        abs(b - a) > atol || return c

        c   = Statistics.middle(a, b)
        f_c = f(c)
        add_eval()

        f_c == 0 && return c

        if f_a*f_c > 0
            a, f_a = c, f_c
        else
            b = c
        end
    end
   @warn "Bisection method reached maximal number of iteration.
          The solution may not be accurate: h(μ) = $(f(c))"
    return c
end

max_stars_repo_name: VaclavMacha/Projections

max_issues_repo_name: VaclavMacha/Projections

max_forks_repo_name: VaclavMacha/Projections


text:
using DFTK
using JLD2
using WriteVTK
using MPI
import DFTK: ScfDefaultCallback, ScfSaveCheckpoints
using Test
include("testcases.jl")

function test_scfres_agreement(tested, ref)
    @test tested.basis.model.lattice           == ref.basis.model.lattice
    @test tested.basis.model.smearing          == ref.basis.model.smearing
    @test tested.basis.model.symmetries        == ref.basis.model.symmetries
    @test tested.basis.model.spin_polarization == ref.basis.model.spin_polarization

    @test tested.basis.model.positions == ref.basis.model.positions
    @test atomic_symbol.(tested.basis.model.atoms) == atomic_symbol.(ref.basis.model.atoms)

    @test tested.basis.Ecut      == ref.basis.Ecut
    @test tested.basis.kweights  == ref.basis.kweights
    @test tested.basis.fft_size  == ref.basis.fft_size

    kcoords_test = getproperty.(tested.basis.kpoints, :coordinate)
    kcoords_ref  = getproperty.(ref.basis.kpoints, :coordinate)
    @test kcoords_test == kcoords_ref

    # Note: Apart from the energy (which is recomputed on loading) the other quantities
    #       should be exactly as stored. For MPI runs it is possible that the density differs
    #       slightly for non-master processes as the version from the master process is stored.
    @test tested.n_iter         == ref.n_iter
    @test tested.energies.total ≈  ref.energies.total atol=1e-13
    @test tested.eigenvalues    == ref.eigenvalues
    @test tested.occupation     == ref.occupation
    @test tested.ψ              == ref.ψ
    @test tested.ρ              ≈  ref.ρ     rtol=1e-14
end

@testset "Test checkpointing" begin
    O = ElementPsp(o2molecule.atnum, psp=load_psp("hgh/pbe/O-q6.hgh"))
    model = model_PBE(o2molecule.lattice, [O, O], o2molecule.positions;
                      temperature=0.02, smearing=Smearing.Gaussian(),
                      magnetic_moments=[1., 1.], symmetries=false)

    kgrid = [1, mpi_nprocs(), 1]   # Ensure at least 1 kpt per process
    basis  = PlaneWaveBasis(model; Ecut=4, kgrid)

    # Run SCF and do checkpointing along the way
    mktempdir() do tmpdir
        checkpointfile = joinpath(tmpdir, "scfres.jld2")
        checkpointfile = MPI.bcast(checkpointfile, 0, MPI.COMM_WORLD)  # master -> everyone

        callback = ScfDefaultCallback() ∘ ScfSaveCheckpoints(checkpointfile; keep=true)
        scfres = self_consistent_field(basis, tol=5e-2, callback=callback)
        test_scfres_agreement(scfres, load_scfres(checkpointfile))
    end
end

@testset "Test serialisation" begin
    model = model_LDA(silicon.lattice, silicon.atoms, silicon.positions;
                      spin_polarization=:collinear, temperature=0.01)
    kgrid = [2, 3, 4]
    basis = PlaneWaveBasis(model; Ecut=5, kgrid)
    scfres = self_consistent_field(basis)

    @test_throws ErrorException save_scfres("MyVTKfile.random", scfres)
    @test_throws ErrorException save_scfres("MyVTKfile", scfres)

    @testset "JLD2" begin
        mktempdir() do tmpdir
            dumpfile = joinpath(tmpdir, "scfres.jld2")
            dumpfile = MPI.bcast(dumpfile, 0, MPI.COMM_WORLD)  # master -> everyone

            save_scfres(dumpfile, scfres)
            @test isfile(dumpfile)
            test_scfres_agreement(scfres, load_scfres(dumpfile))
        end
    end

    @testset "VTK" begin
        mktempdir() do tmpdir
            dumpfile = joinpath(tmpdir, "scfres.vts")
            dumpfile = MPI.bcast(dumpfile, 0, MPI.COMM_WORLD)  # master -> everyone

            save_scfres(dumpfile, scfres; save_ψ=true)
            @test isfile(dumpfile)
        end
    end
end


text:
"""
Dynamic Programming Tree Solver
===============================

GOAL: avoid reallocation of memory as much as possible.
"""
module DPTree

include("queues.jl")
include("event.jl")
include("clip.jl")

import GraphIdx.Tree: RootedTree
import GraphIdx.Tree: ChildrenIndex, reset!, dfs_walk_rev
import .QueueUnion: Queues, merge, reset!, Range

const Tree = RootedTree

function _alloc_queues(n::Int)
    pq = Vector{Range}(undef, n)
    proc_order = Vector{Int}()
    stack = Vector{Int}()
    sizehint!(stack, n)
    sizehint!(proc_order, n)
    childs = ChildrenIndex(n)
    return pq, proc_order, stack, childs
end

function init_queues(parent::Vector{Int}, root::Int = 1)
    n = length(parent)
    pq, proc_order, stack, childs = _alloc_queues(n)
    return _init_queues(parent, root, pq, proc_order, stack, childs)
end

function _init_queues(parent, root, pq, proc_order, stack, childs)
    reset!(childs, parent, root)

    n = length(parent)
    @assert length(childs) == n
    @assert length(childs.idx) == n+1

    empty!(proc_order)
    sizehint!(proc_order, n)
    @assert isempty(proc_order)

    local t::Ref{Int} = Ref{Int}(1)
    dfs_walk_rev(childs, stack) do v::Int
        t[] += 1
        if v >= 0
            push!(proc_order, v)
            pq[v] = Range(t[], t[]-1)
        end
    end

    return pq, proc_order
end

@inline function merge2(buf::Vector{Event}, elements::Vector{Event},
                        parent::Range, child::Range)::Range
    if parent.start <= parent.stop
        gap = child.start - parent.stop - 1
        old_stop = parent.stop
        res = Range(parent.start, child.stop - gap)

        i = 1
        if length(buf) < length(parent)
            resize!(buf, length(parent))
        end
        for j in parent.start:parent.stop
            buf[i] = elements[j]
            i += 1
        end
        buf_end = i
        l = 1
        r = child.start
        el = buf[l]::Event
        er = elements[r]::Event
        for k in res.start:res.stop
            if l < buf_end && (r > child.stop || el.x < er.x)
                elements[k] = el
                l += 1
                el = buf[l]::Event
            else
                elements[k] = er
                r += 1
                er = elements[r]::Event
            end
        end

        return res
    end
    return child
end

# Actual Implementation -----------------------------------------------
dp_tree(y::Vector{Float64}, λ::Float64, t::Tree) =
    dp_tree(y, i->λ, i->1.0, t)

dp_tree(y::Matrix{Float64}, λ::Float64, t::Tree) =
    reshape(dp_tree(vec(y), λ, t), size(y))

function dp_tree(y, λ, µ, t::Tree)
    n = length(y)
    mem = DPMem(n)
    pq, proc_order, stack, childs = _alloc_queues(n)
    return _dp_tree(y, λ, μ, t.root, t.parent,
                    mem.lb, mem.ub, mem.elements,
                    pq, proc_order, stack, childs)
end

struct DPMem
    lb::Vector{Float64}
    ub::Vector{Float64}
    elements::Vector{Event}
    DPMem(n) =
        new(Vector{Float64}(undef, n),
            Vector{Float64}(undef, n),
            Vector{Event}(undef, 2n))
end

_dp_tree(
    y::Vector{Float64},
    λ::Function,
    μ::Function,
    root::Int,
    parent::Vector{Int},
    mem::DPMem,
    pq::Vector{Range},
    proc_order::Vector{Int},
    stack::Vector{Int},
    childs::ChildrenIndex) =
        _dp_tree(y, λ, μ, root, parent,
                 mem.lb, mem.ub, mem.elements, pq,
                 proc_order, stack, childs)

function _dp_tree(y::Vector{Float64},
                  λ::Function,
                  μ::Function,
                  root::Int,
                  parent::Vector{Int},
                  lb::Vector{Float64},
                  ub::Vector{Float64},
                  elements::Vector{Event},
                  pq::Vector{Range},
                  proc_order::Vector{Int},
                  stack::Vector{Int},
                  childs::ChildrenIndex)
    sig = lb
    _init_queues(parent, root, pq, proc_order, stack, childs)
    sig .= 0.0
    for i in @view proc_order[1:end-1]
        sig[parent[i]] += λ(i)
        sig_i = sig[i]
        lb[i] = clip_front(elements, pq, i, µ(i), -µ(i)*y[i] -sig_i, -λ(i))
        ub[i] = clip_back( elements, pq, i, µ(i), -µ(i)*y[i] +sig_i, +λ(i))
        pq[parent[i]] = merge(elements, pq[parent[i]], pq[i])
    end

    x = ub
    x[root] = clip_front(elements, pq, root,
                         µ(root), -µ(root)*y[root] -sig[root], 0.0)
    for v in @view proc_order[end-1:-1:1]
        x[v] = clamp(x[parent[v]], lb[v], ub[v])
    end
    return x
end

# Debug ---------------------------------------------------------------
#=
function _print_elments(elements)
    println("[", join([@sprintf("Event(% 4.2f, % 4.2f % 4.2f)\n",
                                e.x, e.slope, intercept(e))
                       for e in elements], ", "), "]")

end
=#

end

max_stars_repo_name: EQt/treelas

max_issues_repo_name: EQt/treelas

max_forks_repo_name: EQt/treelas


text:
# run.jl
# Julia script for producing "watermarked" pdf solutions for students

using Pkg
Pkg.activate(".") # be sure to use local Manifest

using FileIO
using ColorTypes, ColorVectorSpace
using ColorVectorSpace: ⊙ # hadamard

# convert black&white watermark to rainbow colors
function torainbow(black::AbstractArray)
    dims = size(black)
    # random RGB raised to the 4th power (cf gamma "correction") so fainter:
    r = rand(dims...) .^4
    b = rand(dims...) .^4
    g = rand(dims...) .^4
    mask = RGB.(r, b, g) # random RGB colors
    white = RGB(1,1,1)
    return white .- ((white .- black) .⊙ mask)
end

# extract "name" from "name@place.edu"
function email2name(email)
    at = findfirst('@', email)
    return email[1:(at-1)]
end

# read list of emails from file
emailfile = "emails.txt"
emails = read(emailfile, String)
emails = split(emails)

for email in emails
    name = email2name(email)
    @show email, name

#   file = email
    file = name

    # run pdflatex to create initial "watermark.pdf" with email in it
    cm = Cmd(`pdflatex "\\def\\email{$email} \\input{watermark}"`)
#   @show cm
    run(cm)

    # convert pdf to png for processing
    # uses 'convert' from ImageMagick
    # todo: size is ignored!?
    cm = Cmd(`convert -size 2001x2590 watermark.pdf tmp/watermark-black.png`)
    run(cm)

    black = load("tmp/watermark-black.png")
    rainbow = torainbow(black) # convert to rainbow
    save("tmp/watermark-rainbow.pdf", rainbow)

    cm = Cmd(`merger.py`) # merge solution.pdf with watermark
    run(cm)

    mv("tmp/tmp.pdf", "output/$file.pdf", force=true)
end

# At this point in the script, director output/ has pdf files for each solution.
# The rest of this script is for emailing those solutions to the students.
# This part is not yet completed...
for email in emails
    name = email2name(email)
# todo
end

max_stars_repo_name: JeffFessler/WatermarkSolutions

max_issues_repo_name: JeffFessler/WatermarkSolutions

max_forks_repo_name: JeffFessler/WatermarkSolutions


text:
@testset "AtomsData internals" begin

    @test length(PDB._3_letter_aa) == 20

    s = Set{Tuple{String,String}}()
    PDB._add_CTER_O!(s)
    @test length(s) == 20 * 3
    for aa in PDB._3_letter_aa
        @test (aa,"OXT") in s
        @test (aa,"OT2") in s
        @test (aa,"OT1") in s
    end

    d = Dict{Tuple{String,String},Float64}()
    PDB._add_CTER_O!(d, 1.0)
    @test length(d) == 20 * 3
    for aa in PDB._3_letter_aa
        @test d[(aa,"OXT")] == 1.0
        @test d[(aa,"OT2")] == 1.0
        @test d[(aa,"OT1")] == 1.0
    end

    value = Set(String["OXT", "OT2", "OT1"])

    gd = Dict{String,Set{String}}()
    PDB._generate_dict!(gd, d)
    @test length(gd) == 20
    for aa in PDB._3_letter_aa
        @test gd[aa] == value
    end

    gs = Dict{String,Set{String}}()
    PDB._generate_dict!(gs, d)
    @test length(gs) == 20
    for aa in PDB._3_letter_aa
        @test gs[aa] == value
    end

    as = PDB._generate_interaction_keys(d, s, s, s, s)
    @test length(as) == 20
    for aa in PDB._3_letter_aa
        @test as[aa] == value
    end

    @test isa(PDB._interaction_keys, Dict{String,Set{String}})
end


text:
"""
   compute the partial derivative in vertical direction
"""
function vertical_partial_derivative!(r::Vector{Tv}, p::Vector{Tv}, dpdz::SparseMatrixCSC{Tv,Ti},
         nz::Ti, nx::Ti, tmp1::Vector{Tv}, tmp2::Vector{Tv}; iflag=1) where {Ti<:Int64, Tv<:AbstractFloat}

    istart = 1
    for ix = 1 : nx

        # process column by column
        copyto!(tmp1, 1, p, istart, nz)

        # main part
        if iflag == 1
           A_mul_b!(tmp2, dpdz, tmp1)

        elseif iflag == 2
           At_mul_b!(tmp2, dpdz, tmp1)
        end

        # save the result
        copyto!(r, istart, tmp2, 1, nz)

        # prepare for next column
        istart = istart + nz
    end

    return nothing
end

"""
   compute the partial derivative in horizontal direction
"""
function horizontal_partial_derivative!(r::Vector{Tv}, p::Vector{Tv}, dpdx::SparseMatrixCSC{Tv,Ti},
         nz::Ti, nx::Ti, tmp1::Vector{Tv}, tmp2::Vector{Tv}; iflag=1) where {Ti<:Int64, Tv<:AbstractFloat}

    # the address of vector p and r
    sp = pointer(p)
    sr = pointer(r)

    # electronic size
    esize = sizeof(Tv)
    shift = 0

    for iz = 1 : nz

        # extract one row
        BLAS.blascopy!(nx, sp+shift, nz, tmp1, 1)

        # main part
        if iflag == 1
           A_mul_b!(tmp2, dpdx, tmp1)

        elseif iflag == 2
           At_mul_b!(tmp2, dpdx, tmp1)
        end

        # save the result
        BLAS.blascopy!(nx, tmp2, 1, sr+shift, nz)

        # prepare for next row
        shift = shift + esize
    end

    return nothing
end

"""
   one step forward modelling
"""
function one_step_forward!(spt2::Snapshot{Tv}, spt1::Snapshot{Tv}, params::TdParams{Ti,Tv},
         tmp_z1::Vector{Tv}, tmp_z2::Vector{Tv},
         tmp_x1::Vector{Tv}, tmp_x2::Vector{Tv}) where {Ti<:Int64, Tv<:AbstractFloat}

    for ix = 1 : params.Nx
        ilower = (ix-1) * params.Nz + 1
        iupper = ilower + params.Nz - 1

        # copy one column
        idx    = 0
        for iz = ilower : iupper
            idx= idx + 1
            tmp_z1[idx] = spt1.pz[iz] + spt1.px[iz]  # p1 = pz1 + px1
        end

        # partial derivative in Z direction
        A_mul_b!(tmp_z2, params.dpdz, tmp_z1)        # dpdz
        # mul!(tmp_z2, params.dpdz, tmp_z1)

        # update one column
        idx    = 0
        for iz = ilower : iupper
            idx = idx + 1
            spt2.vz[iz] = params.MvzBvz[idx]*spt1.vz[iz] + params.MvzBp[iz]*tmp_z2[idx]
        end
    end

    # update vx row-by-row
    for iz = 1 : params.Nz

        idx = iz
        for ix = 1 : params.Nx
            tmp_x1[ix] = spt1.pz[idx] + spt1.px[idx] # p1 = pz1 + px1
            idx= idx + params.Nz
        end

        A_mul_b!(tmp_x2, params.dpdx, tmp_x1)        # dpdx
        # mul!(tmp_x2, params.dpdx, tmp_x1)        # dpdx

        idx = iz
        for ix = 1 : params.Nx
            spt2.vx[idx] = params.MvxBvx[ix]*spt1.vx[idx] + params.MvxBp[idx]*tmp_x2[ix]
            idx= idx + params.Nz
        end
    end

    # update pz column-by-column
    for ix = 1 : params.Nx
        ilower = (ix-1) * params.Nz + 1
        iupper = ilower + params.Nz - 1

        idx    = 0
        for iz = ilower : iupper
            idx= idx + 1
            tmp_z1[idx] = spt2.vz[iz]
        end

        A_mul_b!(tmp_z2, params.dvdz, tmp_z1)        # dpdz
        # mul!(tmp_z2, params.dvdz, tmp_z1)        # dpdz

        idx    = 0
        for iz = ilower : iupper
            idx = idx + 1
            spt2.pz[iz] = params.MpzBpz[idx]*spt1.pz[iz] + params.MpzBvz[iz]*tmp_z2[idx]
        end
    end

    # update px row-by-row
    for iz = 1 : params.Nz

        idx = iz
        for ix = 1 : params.Nx
            tmp_x1[ix] = spt2.vx[idx]
            idx= idx + params.Nz
        end

        A_mul_b!(tmp_x2, params.dvdx, tmp_x1)         # dvdx
        # mul!(tmp_x2, params.dvdx, tmp_x1)         # dvdx

        idx = iz
        for ix = 1 : params.Nx
            spt2.px[idx] = params.MpxBpx[ix]*spt1.px[idx] + params.MpxBvx[idx]*tmp_x2[ix]
            idx= idx + params.Nz
        end
    end

    return nothing
end

"""

        Wave_field= multi_step_forward(sources, params; <keyword arguments>)

Forward modelling for acoustic wavefield using finite difference method. This function returns a Recording structure, to acess the wavefield values do:

    Wf_values = Wave_field.p

Please see the extended help with `?? Source

#Fields

- `src::Ts<:Union{Source, Vector{Source}}`: Sources for generating the wavefield.
- `params::TdParams`: Structure that contain time domain model parameters.

#Keyword Arguments

-`rz`: Position for the receiver (sample of the wavefield) in the direction z.
-`rx`: Position for the receiver (sample of the wavefield) in the direction x.
-`path_shot!= "NULL"` : save the records.
-`path_spt != "NULL"`` : vz, vx, pz, px include PML part, save every save_interval step.
-`path_wfd != "NULL"` : vz, vx, pz+px  without PML part, save every save_interval step.
-`path_pre != "NULL"` : p=pz+px without PML bounary part, save every save_interval step.
-`path_bnd != "NULL"` : Save boundary of wavefield at each step, it will be used for source-side wavefield reconstruction.
-`path_lwfd != "NULL"` : Save the last wavefield, it will be used for source-side wavefield reconstruction.
-`path_sws != "NULL"` : Save the source-side wavefield strength (preconditioner for imaging).

# Extended help

        Wave_field= multi_step_forward(sources, params; rx, rz)

Given a source or a vector of sources ( of type Vector{Source{Int64, Float64}}) and the parameter for time-domain modeling, this function will
generate a wavefield sampled for all rx and rz locations. It is possible to use the function inside a control flow structure such as "for loop"
to generate a set of sampled wavefields.

# Example

Given a vector of sources `sx` and the number of receivers in x direction `n_rx` and `n_rz` direction:

#shot gather volume
sgv=zeros(Float64,4001,length(n_rx),length(sources));

#loop

for i=1:length(sx)
    Wave_field[i]= multi_step_forward(sources[i], params; rz=n_rz, rx=n_rx) ;
    sgv[:,:,i].= shot_gather[i].p; # access
end

The cube `sgv` contains all the smapled wavefields for the given sources and receiver geometry.

"""
function multi_step_forward(src::Ts, params::TdParams;
                            rz="NULL", rx="NULL", location_flag="index", path_shot="NULL",
                            path_spt="NULL", path_wfd="NULL" , path_pre="NULL", interval=1,
                            path_bnd="NULL", path_lwfd="NULL", path_sws="NULL") where {Ts<:Union{Source, Vector{Source}}}

    # allocate space for recordings
    if rz != "NULL" && rx != "NULL" && length(rz) == length(rx)
       rec = Recordings(rz, rx, params; location_flag=location_flag)
    end

    # initialize some variables
    spt1   = Snapshot(params)
    spt2   = Snapshot(params)
    tmp_z1 = zeros(params.data_format, params.Nz)
    tmp_z2 = zeros(params.data_format, params.Nz)
    tmp_x1 = zeros(params.data_format, params.Nx)
    tmp_x2 = zeros(params.data_format, params.Nx)

    # add source to the first snapshot
    add_source!(spt1, src, 1)

    # sample wavefield to get recordings
    rz != "NULL" && sample_spt2rec!(rec, spt1, 1)

    # save snapshot
    if path_spt != "NULL"
       hdr_spt   = snapshot_header(params, interval)
       fid_spt   = write_RSheader(path_spt, hdr_spt)
       append_one_snapshot(fid_spt, spt1)
    end

    # save wavefield
    if path_wfd != "NULL"
       hdr_wfd   = wavefield_header(params, interval)
       fid_wfd   = write_RSheader(path_wfd, hdr_wfd)
       append_one_wavefield(fid_wfd, spt1, params)
    end

    # save pressure
    if path_pre != "NULL"
       hdr_pre   = pressure_header(params, interval)
       fid_pre   = write_RSheader(path_pre, hdr_pre)
       append_one_pressure(fid_pre, spt1, params)
    end

    # boundary of wavefield
    if path_bnd != "NULL"
       hdr_bnd = boundary_header(params)
       fid_bnd = write_RSheader(path_bnd, hdr_bnd)
       append_one_boundary(fid_bnd, spt1, params)
    end

    # strength of source-side wavefield
    if path_sws != "NULL"
       N = params.nz * params.nx
       strength = zeros(params.data_format, N)
    end

    # applying time stepping nt-1 times
    for it = 2 : params.nt

        # one step forward
        one_step_forward!(spt2, spt1, params, tmp_z1, tmp_z2, tmp_x1, tmp_x2)

        # compute the strength of source-side wavefield
        path_sws != "NULL" && source_strength!(strength, spt1, spt2, params)

        # inject source term
        add_source!(spt2, src, it)

        # sample wavefield to get recordings
        rz != "NULL" && sample_spt2rec!(rec, spt2, it)

        # write current wavefield to disk
        if mod(it-1, interval) == 0
           path_spt != "NULL" && append_one_snapshot(fid_spt, spt2)
           path_wfd != "NULL" && append_one_wavefield(fid_wfd, spt2, params)
           path_pre != "NULL" && append_one_pressure(fid_pre, spt2, params)
        end

        # save boundary of wavefield
        path_bnd != "NULL" && append_one_boundary(fid_bnd, spt2, params)

        # prepare for next time stepping
        copy_snapshot!(spt1, spt2)
    end

    # save the source-side wavefield strength to disk
    if path_sws != "NULL"
       strength  = reshape(strength, params.nz, params.nx)
       hdr_sws   = RegularSampleHeader(strength; title="source strength")
       write_RSdata(path_sws, hdr_sws, strength)
    end

    # save the last wavefield
    if path_lwfd != "NULL"
       wfd = sample_spt2wfd(spt1, params)
       write_wavefield(path_lwfd, wfd, params)
    end

    # close file
    path_spt != "NULL" && close(fid_spt)
    path_wfd != "NULL" && close(fid_wfd)
    path_pre != "NULL" && close(fid_pre)
    path_bnd != "NULL" && close(fid_bnd)

    if rz != "NULL"
       if path_shot == "NULL"
          return rec
       else
          write_recordings(path_shot, rec)
          return path_shot
       end

    else
       return nothing
    end
end

# define a auxillary function to improve efficiency
function source_strength!(strength, spt1, spt2, params)

    # total number of elements
    N = params.nz * params.nx

    for i = 1 : N
        j = params.spt2wfd[i]

        p2= spt2.pz[j]+spt2.px[j]
        p1= spt1.pz[j]+spt1.px[j]

        tmp = 2.0 * (p2 - p1) / params.vel[i]
        strength[i] += tmp * tmp
    end
    return nothing
end

# ==============================================================================
#          one step forward modelling with saved boundary value
# ==============================================================================
"""
   one step-forward reconstruction of wavefield via recording the boundary value
"""
function one_step_forward!(wfd2::Wavefield, wfd1::Wavefield,
         bnd::WavefieldBound, params::TdParams,
         tmp_z1::Vector{Tv}, tmp_z2::Vector{Tv},
         tmp_x1::Vector{Tv}, tmp_x2::Vector{Tv}) where {Ti<:Int64, Tv<:AbstractFloat}

    # update vz column-by-column
    for ix = 1 : params.nx
        ilower = (ix-1) * params.nz + 1
        iupper = ilower + params.nz - 1

        idx    = 0
        for iz = ilower : iupper
            idx= idx + 1
            tmp_z1[idx] = wfd1.p[iz]
        end

        A_mul_b!(tmp_z2, params.rpdz, tmp_z1)        # dpdz

        idx    = 0
        for iz = ilower : iupper
            idx = idx + 1
            wfd2.vz[iz] = wfd1.vz[iz] + params.RvzBp[iz]*tmp_z2[idx]
        end
    end

    # update vx row-by-row
    for iz = 1 : params.nz

        idx = iz
        for ix = 1 : params.nx
            tmp_x1[ix] = wfd1.p[idx]
            idx= idx + params.nz
        end

        A_mul_b!(tmp_x2, params.rpdx, tmp_x1)         # dpdx

        idx = iz
        for ix = 1 : params.nx
            wfd2.vx[idx] = wfd1.vx[idx] + params.RvxBp[idx]*tmp_x2[ix]
            idx= idx + params.nz
        end
    end

    # correct for boundary part
    for i = 1 : length(params.wfd2bnd)
        j = params.wfd2bnd[i]
        wfd2.vz[j] = bnd.vz[i]
        wfd2.vx[j] = bnd.vx[i]
    end

    # update pz column-by-column
    for ix = 1 : params.nx
        ilower = (ix-1) * params.nz + 1
        iupper = ilower + params.nz - 1

        idx    = 0
        for iz = ilower : iupper
            idx= idx + 1
            tmp_z1[idx] = wfd2.vz[iz]
        end

        A_mul_b!(tmp_z2, params.rvdz, tmp_z1)        # dpdz

        idx    = 0
        for iz = ilower : iupper
            idx = idx + 1
            wfd2.p[iz] = wfd1.p[iz] + params.RpBv[iz]*tmp_z2[idx]
        end
    end

    # update px row-by-row
    for iz = 1 : params.nz

        idx = iz
        for ix = 1 : params.nx
            tmp_x1[ix] = wfd2.vx[idx]
            idx= idx + params.nz
        end

        A_mul_b!(tmp_x2, params.rvdx, tmp_x1)         # dvdx

        idx = iz
        for ix = 1 : params.nx
            wfd2.p[idx] = wfd2.p[idx] + params.RpBv[idx]*tmp_x2[ix]
            idx= idx + params.nz
        end
    end

    # correct for boundary part
    for i = 1 : length(params.wfd2bnd)
        j = params.wfd2bnd[i]
        wfd2.p[j] = bnd.p[i]
    end

    return nothing
end

"""
   Forward reconstruct acoustic pressure field using the saved boundary wavefield value and
the source (used for concept proofing)
"""
function pressure_reconstruct_forward(path_bnd::Tp, src::Ts,
         params::TdParams) where {Tp<:String, Ts<:Union{Source,Vector{Source}}}

    # make the source as a vector has one element
    if Ts <: Source
       srcs    = Vector{Source}(undef,1)
       srcs[1] = src
    else
       srcs    = src
    end
    ns = length(srcs)

    # length of one-step pressure field
    N = params.nz * params.nx

    # initialize intermediate variables
    wfd1   = Wavefield(params)
    wfd2   = Wavefield(params)
    tmp_z1 = zeros(params.data_format, params.nz)
    tmp_z2 = zeros(params.data_format, params.nz)
    tmp_x1 = zeros(params.data_format, params.nx)
    tmp_x2 = zeros(params.data_format, params.nx)

    # initialize boundary value
    fid_bnd = open(path_bnd, "r")
    bnd = WavefieldBound(params)

    # allocate memory for saving pressure field
    pre = zeros(params.data_format, N * params.nt)

    # add source to the first wavefield
    add_source!(wfd1, srcs, 1)

    # save pressure field
    copyto!(pre, 1, wfd1.p, 1, N)

    # the start index for saving the next pressure field
    idx_o = N+1

    # loop over time stepping
    for it = 2 : params.nt

        # read the boundary value
        read_one_boundary!(bnd, fid_bnd, it, params)

        # forward time steping and correcting boundaries
        one_step_forward!(wfd2, wfd1, bnd, params,
                          tmp_z1, tmp_z2, tmp_x1, tmp_x2)

        # loop over sources
        for i = 1 : ns
            if (params.order+1 <= srcs[i].isz <= params.nz-params.order &&
                params.order+1 <= srcs[i].isx <= params.nx-params.order)

               add_source!(wfd2, srcs[i], it)
            end
        end

        # save the pressure field
        copyto!(pre, idx_o, wfd2.p, 1, N)
        idx_o = idx_o + N

        # prepare for next step
        copy_wavefield!(wfd1, wfd2)
    end

    close(fid_bnd)
    return reshape(pre, params.nz, params.nx, params.nt)
end

# this version of one step forward is little bit slower than the used one
# since the overhead of calling an function
# function one_step_forward_new!(spt2::Snapshot{Tv}, spt1::Snapshot{Tv}, params::TdParams{Ti,Tv},
#                                tmp_z1::Vector{Tv}, tmp_z2::Vector{Tv},
#                                tmp_x1::Vector{Tv}, tmp_x2::Vector{Tv}) where {Ti<:Int64, Tv<:AbstractFloat}
#
#     # compute pressure field
#     spt2.pz .= spt1.pz .+ spt1.px
#
#     # update particle velocity
#     vertical_partial_derivative!(spt2.vz, spt2.pz, params.dpdz, params.Nz, params.Nx, tmp_z1, tmp_z2; iflag=1)
#     horizontal_partial_derivative!(spt2.vx, spt2.pz, params.dpdx, params.Nz, params.Nx, tmp_x1, tmp_x2; iflag=1)
#     idx = 0
#     for ix = 1 : params.Nx
#         for iz = 1 : params.Nz
#             idx= idx + 1
#             spt2.vz[idx] = params.MvzBvz[iz]*spt1.vz[idx] + params.MvzBp[idx]*spt2.vz[idx]
#             spt2.vx[idx] = params.MvxBvx[ix]*spt1.vx[idx] + params.MvxBp[idx]*spt2.vx[idx]
#         end
#     end
#
#     # update pressure field
#     vertical_partial_derivative!(spt2.pz, spt2.vz, params.dvdz, params.Nz, params.Nx, tmp_z1, tmp_z2; iflag=1)
#     horizontal_partial_derivative!(spt2.px, spt2.vx, params.dvdx, params.Nz, params.Nx, tmp_x1, tmp_x2; iflag=1)
#     idx = 0
#     for ix = 1 : params.Nx
#         for iz = 1 : params.Nz
#             idx= idx + 1
#             spt2.pz[idx] = params.MpzBpz[iz]*spt1.pz[idx] + params.MpzBvz[idx]*spt2.pz[idx]
#             spt2.px[idx] = params.MpxBpx[ix]*spt1.px[idx] + params.MpxBvx[idx]*spt2.px[idx]
#         end
#     end
#
#     return nothing
# end

# function one_step_forward!(wfd2::Wavefield, wfd1::Wavefield,
#          bnd::WavefieldBound, params::TdParams,
#          tmp_a1::Vector{Tv}, tmp_a2::Vector{Tv},
#          tmp_z1::Vector{Tv}, tmp_z2::Vector{Tv},
#          tmp_x1::Vector{Tv}, tmp_x2::Vector{Tv}) where {Ti<:Int64, Tv<:AbstractFloat}
#
#     # total number of field elements
#     N = params.nz * params.nx
#
#     # partial derivative in vertical direction
#     vertical_partial_derivative!(tmp_a1, wfd1.p, params.rpdz,
#     params.nz, params.nx, tmp_z1, tmp_z2)
#
#     # partial derivative in horizontal direction
#     horizontal_partial_derivative!(tmp_a2, wfd1.p, params.rpdx,
#     params.nz, params.nx, tmp_x1, tmp_x2)
#
#     # update vz
#     for i = 1 : N
#         wfd2.vz[i] = wfd1.vz[i] + params.RvzBp[i]*tmp_a1[i]
#         wfd2.vx[i] = wfd1.vx[i] + params.RvxBp[i]*tmp_a2[i]
#     end
#
#     # correct for boundary part
#     for i = 1 : length(params.wfd2bnd)
#         j = params.wfd2bnd[i]
#         wfd2.vz[j] = bnd.vz[i]
#         wfd2.vx[j] = bnd.vx[i]
#     end
#
#     # partial derivative in vertical direction
#     vertical_partial_derivative!(tmp_a1, wfd2.vz, params.rvdz,
#     params.nz, params.nx, tmp_z1, tmp_z2)
#
#     # partial derivative in horizontal direction
#     horizontal_partial_derivative!(tmp_a2, wfd2.vx, params.rvdx,
#     params.nz, params.nx, tmp_x1, tmp_x2)
#
#     # update for pressure
#     for i = 1 : N
#         wfd2.p[i] = wfd1.p[i] + params.RpBv[i] * (tmp_a1[i] + tmp_a2[i])
#     end
#
#     # correct for boundary part
#     for i = 1 : length(params.wfd2bnd)
#         j = params.wfd2bnd[i]
#         wfd2.p[j] = bnd.p[i]
#     end
#
#     return nothing
# end


text:
using Catch22
using StatsBase
using Colors
using LinearAlgebra
using ColorVectorSpace
using Clustering
using Tullio

# function horizontal_legend(N, ax=PyPlot.gca())
#     PyPlot.svg(true)
#     l = PyPlot.gca().legend_
#     # https://stackoverflow.com/questions/23689728/how-to-modify-matplotlib-legend-after-it-has-been-created
#     defaults = Dict(:loc => l._loc,
#         :numpoints => l.numpoints,
#         :markerscale => l.markerscale,
#         :scatterpoints => l.scatterpoints,
#         :scatteryoffsets => l._scatteryoffsets,
#         :prop => l.prop,
#         # fontsize = None,
#         :borderpad => l.borderpad,
#         :labelspacing => l.labelspacing,
#         :handlelength => l.handlelength,
#         :handleheight => l.handleheight,
#         :handletextpad => l.handletextpad,
#         :borderaxespad => l.borderaxespad,
#         :columnspacing => l.columnspacing,
#         :ncol => N,
#         :mode => l._mode,
#         :shadow => l.shadow,
#         :title => l._legend_title_box.get_visible() ? l.get_title().get_text() : nothing,
#         :framealpha => l.get_frame().get_alpha(),
#         :bbox_to_anchor => l.get_bbox_to_anchor()._bbox,
#         :bbox_transform => l.get_bbox_to_anchor()._transform)
#     PyPlot.legend(("1", "2", "3"), defaults)
#     f = PyPlot.gcf()
# end

# ------------------------------------------------------------------------------------------------ #
#                                       Plot a feature matrix                                      #
# ------------------------------------------------------------------------------------------------ #
@userplot FeatureImage
@recipe function f(FI::FeatureImage)
    if length(FI.args) == 3
        x = FI.args[1]
        y = FI.args[2]
        F = FI.args[3]
    elseif length(FI.args) == 2
        F = FI.args[2]
        x = FI.args[1]
        if F isa AbstractFeatureArray
            y = string.(Catch22.featureDims(F))
        else
            y = string.(1:size(F, 1))
        end
    elseif length(FI.args) == 1
        F = FI.args[1]
        if F isa AbstractFeatureArray
            y = string.(Catch22.featureDims(F))
        else
            y = string.(1:size(F, 1))
        end
        x = 1:size(F, 2)#string.(dims(F, :timeseries).val)
    end
    @series begin
        seriestype := :heatmap
        framestyle --> :box
        yflip --> true
        xaxis --> nothing
        yticks --> :all
        clim = max(abs.(extrema(F))...)
        seriescolor --> palette(:Blues_9, 7)
        if min(F...) > 0.0
            clims := (0.0, clim)
        elseif max(F...) < 0.0
            clims := (-clim, 0.0)
        else
            clims := (-clim, clim)
            seriescolor --> palette(:RdYlBu_11, 7)
        end
        (x, y, X) = (pysafelabel.(x), pysafelabel.(y), Array(F))
    end
end

@userplot DistanceImage
@recipe function f(DI::DistanceImage; metric=StatsBase.corspearman, annotatedist=false)
    if length(DI.args) == 2
        F = DI.args[2]
        f = DI.args[1]
    elseif length(DI.args) == 1
        F = DI.args[1]
        if F isa AbstractFeatureArray
            f = string.(Catch22.featureDims(F))
        else
            f = string.(1:size(F, 1))
        end
    end

    Df = 1.0.-abs.(metric(F'))
    idxs = Clustering.hclust(Df; linkage=:average, branchorder=:optimal).order
    Df′ = Df[idxs, idxs]

    #plot!(yticks=(1:size(Df, 1), replace.(string.(Catch22.featurenames[idxs]), '_'=>"\\_")), size=(800, 400), xlims=[0.5, size(Df, 1)+0.5], ylims=[0.5, size(Df, 1)+0.5], box=:on, )
    ann = [(x-0.5, y-0.5, text("$(round(Df′[x, y], digits=2))", :white)) for x ∈ 1:length(idxs) for y ∈ 1:length(idxs)]

    @series begin
        seriestype := :heatmap
        framestyle --> :box
        xaxis --> nothing
        yflip --> true
        lims --> (0, length(idxs))
        aspect_ratio --> :equal
        size --> (800, 400)
        if length(f) ≤ 60
            yticks --> :all
        end
        if annotatedist == true
            annotations --> ann
            colorbar --> nothing
        end
        seriescolor --> cgrad(:Greys_9, rev=true)
        colorbar_title --> "1-|ρ|"
        clims --> (0.0, 1.0)
        (x, y, X) = (pysafelabel.(f[idxs]), pysafelabel.(f[idxs]), Df′)
    end
end


text:
using EconPDEs, Distributions

struct AchdouHanLasryLionsMoll_TwoAssetsModel
    # income process parameters
    κy::Float64
    ybar::Float64
    σy::Float64

    r::Float64
    μR::Float64
    σR::Float64

    # utility parameters
    ρ::Float64
    γ::Float64

    amin::Float64
    amax::Float64
end

function AchdouHanLasryLionsMoll_TwoAssetsModel(;κy = 0.1, ybar = 1.0, σy = 0.07, r = 0.03, μR = 0.04, σR = 0.1, ρ = 0.05, γ = 2.0, amin = 0.0, amax = 1000.0)
    AchdouHanLasryLionsMoll_TwoAssetsModel(κy, ybar, σy, r, μR, σR, ρ, γ, amin, amax)
end

function initialize_stategrid(m::AchdouHanLasryLionsMoll_TwoAssetsModel; yn = 5, an = 100)
    κy = m.κy ; ybar = m.ybar ; σy = m.σy  ; ρ = m.ρ ; γ = m.γ ; amin = m.amin ; amax = m.amax

    distribution = Gamma(2 * κy * ybar / σy^2, σy^2 / (2 * κy))
    ymin = quantile(distribution, 0.001)
    ymax = quantile(distribution, 0.999)
    ys = range(ymin, stop = ymax, length = yn)
    as = range(amin, stop = amax, length = an)
    OrderedDict(:y => ys, :a => as)
end

function initialize_y(m::AchdouHanLasryLionsMoll_TwoAssetsModel, stategrid::OrderedDict)
    OrderedDict(:v => [log(y + a) for y in stategrid[:y], a in stategrid[:a]])
end

function (m::AchdouHanLasryLionsMoll_TwoAssetsModel)(state::NamedTuple, value::NamedTuple)
    κy = m.κy ; σy = m.σy ; ybar = m.ybar ; r = m.r ; μR = m.μR ; σR = m.σR ; ρ = m.ρ ; γ = m.γ ; amin = m.amin ; amax = m.amax
    y, a = state.y, state.a
    v, vy, va, vyy, vya, vaa = value.v, value.vy, value.va, value.vyy, value.vya, value.vaa
    μy = κy * (ybar - y)
    va = max(va, eps())

    c = va^(-1 / γ)
    k = (μR - r) / σR^2 * (- va / vaa)
    k = clamp(k, 0.0, a - amin)
    μa = y + r * a + (μR - r) * k - c
    σa = k * σR
    # There is no second derivative at 0 so just specify first order derivative
    if (a ≈ amin) && (μa <= 0.0)
        va = (y + r * amin)^(-γ)
        k = 0.0
        c = y + r * amin
        μa = 0.0
    end
    # this branch is unnecessary when individuals dissave at the top (default)
    if (a ≈ amax) && (μa >= 0.0)
        va = ((ρ - r) / γ + r - (1-γ) / (2 * γ) * (μR - r)^2 / (γ * σR^2))^(-γ) * a^(-γ)
    end
    # Since first derivative is upwinded, I can directly impose value of second derivative
    if (a ≈ amax)
        vaa = - m.γ * va / a
    end
    c = va^(-1 / γ)
    k = (μR - r) / σR^2 * (- va / vaa)
    k = clamp(k, 0.0, a - amin)
    μa = y + r * a + (μR - r) * k - c
    σa = k * σR
    vt = c^(1 - γ) / (1 - γ) + va * μa + 0.5 * vaa * σa^2 + vy * μy + 0.5 * vyy * σy^2 - ρ * v
    return (vt,), (μy, μa), (v = v, c = c, k = k, va = va, vaa = vaa, vy = vy, y = y, a = a, μa = μa)
end

m = AchdouHanLasryLionsMoll_TwoAssetsModel()
stategrid = initialize_stategrid(m)
y0 = initialize_y(m, stategrid)
y, result, distance = pdesolve(m, stategrid, y0)
#
# # Important: check marginal value of wealth converges to 1.0
# # This happens ONLY if a >= 1000.0. Otherwise with 300 it does not work. This is interesting. Maybe it means there should be a better way to have bordering condition at # top
# b = ((m.r + (m.ρ - m.r)/m.γ - (1-m.γ) / (2 * m.γ) * (m.μR - m.r)^2 / (m.γ * m.σR^2)))^(1/(1 - 1/m.γ))
# pw = (result[:v] * (1-m.γ)).^(1/(1-m.γ)-1) .* result[:va] ./ b


text:
module POMDPs

import Base.rand

export
    # Abstract type
    POMDP,
    DiscretePOMDP,

    # Discrete Functions
    n_states,
    n_actions,
    n_observations,

    # Model functions
    discount,
    states,
    actions,
    observations,
    transition,
    observation,
    reward,
    isterminal,

    # Need below?;
    create_state,
    create_observation,
    create_action,

    # Spaces, Distributions and accessor functions
    AbstractDistribution,
    DiscreteDistribution,
    AbstractSpace,

    # Discrete Functions
    length,
    index,
    weight,

    # Common Functions
    rand,
    pdf,
    dimensions,
    lowerbound,
    upperbound,
    getindex,
    iterator,
    create_transition_distribution,
    create_observation_distribution,
    create_belief,
    initial_belief,

    # Solver types
    Solver,
    solve,

    # States
    State,

    # Actions
    Action,
    create_action,

    # Observations
    Observation,

    # Rewards
    Reward,

    # Beliefs
    Belief,
    BeliefUpdater,
    initial_belief,
    update,
    convert_belief,

    # Policy
    Policy,
    action,
    updater,
    value,

    # Simulation
    Simulator,
    simulate

    # Utilities
    #add not explicitly imported

include("errors.jl")
include("constants.jl")
include("utils.jl")
include("pomdp.jl")
include("distribution.jl")
include("belief.jl")
include("space.jl")
include("solver.jl")
include("policy.jl")
include("simulator.jl")
include("docs.jl")

end

max_stars_repo_name: sisl/ExprSearch

max_issues_repo_name: sisl/ExprSearch