text:
abstract type AbstractCoorSystem{dim,T<:Number} <: Any end

"""
    CoorSystemSym(OM::AbstractTens{1,dim,Sym},coords::NTuple{dim,Sym},bnorm::AbstractBasis{dim,Sym},χᵢ::NTuple{dim},
                  tmp_coords::NTuple = (),params::NTuple = ();rules::Dict = Dict(),tmp_var::Dict = Dict(),to_coords::Dict = Dict()) where {dim}
    CoorSystemSym(OM::AbstractTens{1,dim,Sym},coords::NTuple{dim,Sym},
                  tmp_coords::NTuple = (),params::NTuple = ();rules::Dict = Dict(),tmp_var::Dict = Dict(),to_coords::Dict = Dict()) where {dim}

Defines a new coordinate system either from
1. the position vector `OM`, the coordinates `coords`, the basis of unit vectors (`𝐞ᵢ`) `bnorm` and the Lamé coefficients `χᵢ`

    In this case the natural basis is formed by the vectors `𝐚ᵢ = χᵢ 𝐞ᵢ` directly calculated from the input data.

1. or the position vector `OM` and the coordinates `coords`

    In this case the natural basis is formed by the vectors `𝐚ᵢ = ∂ᵢOM` i.e. by the derivative of the position vector with respect to the `iᵗʰ` coordinate

Optional parameters can be provided:
- `tmp_coords` contains temporary variables depending on coordinates (in order to allow symbolic simplifications)
- `params` contains possible parameters involved in `OM`
- `rules` contains a `Dict` with substitution rules to facilitate the simplification of formulas
- `tmp_var` contains a `Dict` with substitution of coordinates by temporary variables
- `to_coords` indicates how to eliminate the temporary variables to come back to the actual coordinates before derivation for Examples

# Examples
```julia
julia> ϕ, p = symbols("ϕ p", real = true) ;

julia> p̄, q, q̄, c = symbols("p̄ q q̄ c", positive = true) ;

julia> coords = (ϕ, p, q) ; tmp_coords = (p̄, q̄) ; params = (c,) ;

julia> OM = Tens(c * [p̄ * q̄ * cos(ϕ), p̄ * q̄ * sin(ϕ), p * q]) ;

julia> Spheroidal = CoorSystemSym(OM, coords, tmp_coords, params; tmp_var = Dict(1-p^2 => p̄^2, q^2-1 => q̄^2), to_coords = Dict(p̄ => √(1-p^2), q̄ => √(q^2-1))) ;
```
"""
struct CoorSystemSym{dim,VEC,BNORM,BNAT} <: AbstractCoorSystem{dim,Sym}
    OM::VEC
    coords::NTuple{dim,Sym}
    normalized_basis::BNORM
    natural_basis::BNAT
    aᵢ::NTuple{dim}
    χᵢ::NTuple{dim}
    aⁱ::NTuple{dim}
    eᵢ::NTuple{dim}
    Γ::Array{Sym,3}
    tmp_coords::NTuple
    params::NTuple
    rules::Dict
    tmp_var::Dict
    to_coords::Dict
    function CoorSystemSym(
        OM::VEC,
        coords::NTuple{dim,Sym},
        normalized_basis::AbstractBasis{dim,Sym},
        χᵢ::NTuple{dim},
        tmp_coords::NTuple = (),
        params::NTuple = ();
        rules::Dict = Dict(),
        tmp_var::Dict = Dict(),
        to_coords::Dict = Dict(),
    ) where {VEC,dim}
        simp(t) = length(rules) > 0 ? simplify(subs(simplify(t), rules...)) : simplify(t)
        eᵢ = ntuple(
            i -> Tens(
                Vec{dim}(j -> j == i ? one(Sym) : zero(Sym)),
                normalized_basis,
                (:cov,),
            ),
            dim,
        )
        aᵢ = ntuple(
            i -> Tens(Vec{dim}(j -> j == i ? χᵢ[i] : zero(Sym)), normalized_basis, (:cov,)),
            dim,
        )
        aⁱ = ntuple(
            i -> Tens(
                Vec{dim}(j -> j == i ? inv(χᵢ[i]) : zero(Sym)),
                normalized_basis,
                (:cont,),
            ),
            dim,
        )
        Γ = compute_Christoffel(
            coords,
            χᵢ,
            metric(normalized_basis, :cov),
            metric(normalized_basis, :cont),
        )
        natural_basis = Basis(normalized_basis, χᵢ)
        new{dim,typeof(OM),typeof(normalized_basis),typeof(natural_basis)}(
            OM,
            coords,
            normalized_basis,
            natural_basis,
            aᵢ,
            χᵢ,
            aⁱ,
            eᵢ,
            Γ,
            tmp_coords,
            params,
            rules,
            tmp_var,
            to_coords,
        )
    end
    function CoorSystemSym(
        OM::VEC,
        coords::NTuple{dim,Sym},
        tmp_coords::NTuple = (),
        params::NTuple = ();
        rules::Dict = Dict(),
        tmp_var::Dict = Dict(),
        to_coords::Dict = Dict(),
    ) where {VEC,dim}
        simp(t) = length(rules) > 0 ? simplify(subs(simplify(t), rules...)) : simplify(t)
        chvar(t, d) = length(d) > 0 ? subs(t, d...) : t
        OMc = chvar(OM, to_coords)
        aᵢ = ntuple(i -> simp(chvar(∂(OMc, coords[i]), tmp_var)), dim)
        χᵢ = ntuple(i -> simp(norm(aᵢ[i])), dim)
        eᵢ = ntuple(i -> simp(aᵢ[i] / χᵢ[i]), dim)
        χᵢ = ntuple(i -> simp(chvar(χᵢ[i], to_coords)), dim)
        eᵢ = ntuple(i -> simp(chvar(eᵢ[i], to_coords)), dim)
        normalized_basis = Basis(simplify(hcat(components_canon.(eᵢ)...)))
        eᵢ = ntuple(
            i -> Tens(
                Vec{dim}(j -> j == i ? one(Sym) : zero(Sym)),
                normalized_basis,
                (:cov,),
            ),
            dim,
        )
        aᵢ = ntuple(
            i -> Tens(Vec{dim}(j -> j == i ? χᵢ[i] : zero(Sym)), normalized_basis, (:cov,)),
            dim,
        )
        aⁱ = ntuple(
            i -> Tens(
                Vec{dim}(j -> j == i ? inv(χᵢ[i]) : zero(Sym)),
                normalized_basis,
                (:cont,),
            ),
            dim,
        )
        Γ = compute_Christoffel(
            coords,
            χᵢ,
            metric(normalized_basis, :cov),
            metric(normalized_basis, :cont),
        )
        natural_basis = Basis(normalized_basis, χᵢ)
        new{dim,typeof(OM),typeof(normalized_basis),typeof(natural_basis)}(
            OMc,
            coords,
            normalized_basis,
            natural_basis,
            aᵢ,
            χᵢ,
            aⁱ,
            eᵢ,
            Γ,
            tmp_coords,
            params,
            rules,
            tmp_var,
            to_coords,
        )
    end
end

with_tmp_var(CS::CoorSystemSym, t) = length(CS.tmp_var) > 0 ? subs(t, CS.tmp_var...) : t
only_coords(CS::CoorSystemSym, t) = length(CS.to_coords) > 0 ? subs(t, CS.to_coords...) : t

getcoords(CS::CoorSystemSym) = CS.coords
getcoords(CS::CoorSystemSym, i::Integer) = getcoords(CS)[i]

@pure getdim(::CoorSystemSym{dim}) where {dim} = dim

getOM(CS::CoorSystemSym) = CS.OM

get_normalized_basis(CS::CoorSystemSym) = CS.normalized_basis
get_natural_basis(CS::CoorSystemSym) = CS.natural_basis

getLame(CS::CoorSystemSym) = CS.χᵢ
getChristoffel(CS::CoorSystemSym) = CS.Γ

natvec(CS::CoorSystemSym, ::Val{:cov}) = CS.aᵢ
natvec(CS::CoorSystemSym, ::Val{:cont}) = CS.aⁱ
natvec(CS::CoorSystemSym, var = :cov) = natvec(CS, Val(var))
natvec(CS::CoorSystemSym, i::Integer, var = :cov) = natvec(CS, var)[i]

unitvec(CS::CoorSystemSym) = CS.eᵢ
unitvec(CS::CoorSystemSym, i::Integer) = unitvec(CS)[i]

function compute_Christoffel(coords, χ, γ, invγ)
    dim = length(χ)
    gᵢⱼ = [γ[i, j] * χ[i] * χ[j] for i ∈ 1:dim, j ∈ 1:dim]
    gⁱʲ = [invγ[i, j] / (χ[i] * χ[j]) for i ∈ 1:dim, j ∈ 1:dim]
    ∂g = [SymPy.diff(gᵢⱼ[i, j], coords[k]) for i ∈ 1:dim, j ∈ 1:dim, k ∈ 1:dim]
    Γᵢⱼₖ =
        [(∂g[i, k, j] + ∂g[j, k, i] - ∂g[i, j, k]) / 2 for i ∈ 1:dim, j ∈ 1:dim, k ∈ 1:dim]
    return ein"ijl,lk->ijk"(Γᵢⱼₖ, gⁱʲ)
end

"""
    ∂(t::AbstractTens{order,dim,Sym,A},xᵢ::Sym)

Returns the derivative of the tensor `t` with respect to the variable `x_i`

# Examples
```julia

julia> θ, ϕ, ℬˢ, 𝐞ᶿ, 𝐞ᵠ, 𝐞ʳ = init_spherical(symbols("θ ϕ", real = true)...) ;

julia> ∂(𝐞ʳ, ϕ) == sin(θ) * 𝐞ᵠ
true

julia> ∂(𝐞ʳ ⊗ 𝐞ʳ,θ)
Tens.TensRotated{2, 3, Sym, SymmetricTensor{2, 3, Sym, 6}}
# data: 3×3 SymmetricTensor{2, 3, Sym, 6}:
 0  0  1
 0  0  0
 1  0  0
# basis: 3×3 Tensor{2, 3, Sym, 9}:
 cos(θ)⋅cos(ϕ)  -sin(ϕ)  sin(θ)⋅cos(ϕ)
 sin(ϕ)⋅cos(θ)   cos(ϕ)  sin(θ)⋅sin(ϕ)
       -sin(θ)        0         cos(θ)
# var: (:cont, :cont)
```
"""
∂(t::AbstractTens{order,dim,Sym}, xᵢ...) where {order,dim} =
    change_tens(Tens(SymPy.diff(components_canon(t), xᵢ...)), getbasis(t), getvar(t))

∂(t::Sym, xᵢ...) = SymPy.diff(t, xᵢ...)

function ∂(
    t::AbstractTens{order,dim,Sym},
    i::Integer,
    CS::CoorSystemSym{dim},
) where {order,dim}
    t = only_coords(CS, t)
    ℬ = get_natural_basis(CS)
    var = ntuple(_ -> :cont, order)
    t = Array(components(t, ℬ, var))
    Γ = getChristoffel(CS)
    data = diff.(t, getcoords(CS, i))
    for o ∈ 1:order
        ec1 = ntuple(j -> j == o ? order + 1 : j, order)
        ec2 = (order + 1, o)
        ec3 = ntuple(j -> j, order)
        data += einsum(EinCode((ec1, ec2), ec3), (t, view(Γ, i, :, :)))
    end
    return change_tens(Tens(data, ℬ, var), get_normalized_basis(CS), var)
end

∂(t::Sym, i::Integer, CS::CoorSystemSym{dim}) where {dim} =
    SymPy.diff(only_coords(CS, t), getcoords(CS, i))

function ∂(
    t::AbstractTens{order,dim,Sym},
    x::Sym,
    CS::CoorSystemSym{dim},
) where {order,dim}
    ind = findfirst(i -> i == x, getcoords(CS))
    return ind === nothing ? zero(t) : ∂(t, ind, CS)
end

function ∂(
    t::Sym,
    x::Sym,
    CS::CoorSystemSym{dim},
) where {order,dim}
    ind = findfirst(i -> i == x, getcoords(CS))
    return ind === nothing ? zero(t) : ∂(t, ind, CS)
end

"""
    GRAD(T::Union{Sym,AbstractTens{order,dim,Sym}},CS::CoorSystemSym{dim}) where {order,dim}

Calculates the gradient of `T` with respect to the coordinate system `CS`
"""
GRAD(T::Union{Sym,AbstractTens{order,dim,Sym}}, CS::CoorSystemSym{dim}) where {order,dim} =
    sum([∂(T, i, CS) ⊗ natvec(CS, i, :cont) for i = 1:dim])

"""
    SYMGRAD(T::Union{Sym,AbstractTens{order,dim,Sym}},CS::CoorSystemSym{dim}) where {order,dim}

Calculates the symmetrized gradient of `T` with respect to the coordinate system `CS`
"""
SYMGRAD(
    T::Union{Sym,AbstractTens{order,dim,Sym}},
    CS::CoorSystemSym{dim},
) where {order,dim} = sum([∂(T, i, CS) ⊗ˢ natvec(CS, i, :cont) for i = 1:dim])

"""
    DIV(T::AbstractTens{order,dim,Sym},CS::CoorSystemSym{dim}) where {order,dim}

Calculates the divergence  of `T` with respect to the coordinate system `CS`
"""
DIV(T::AbstractTens{order,dim,Sym}, CS::CoorSystemSym{dim}) where {order,dim} =
    sum([∂(T, i, CS) ⋅ natvec(CS, i, :cont) for i = 1:dim])

"""
    LAPLACE(T::Union{Sym,AbstractTens{order,dim,Sym}},CS::CoorSystemSym{dim}) where {order,dim}

Calculates the Laplace operator of `T` with respect to the coordinate system `CS`
"""
LAPLACE(
    T::Union{Sym,AbstractTens{order,dim,Sym}},
    CS::CoorSystemSym{dim},
) where {order,dim} = DIV(GRAD(T, CS), CS)

"""
    HESS(T::Union{Sym,AbstractTens{order,dim,Sym}},CS::CoorSystemSym{dim}) where {order,dim}

Calculates the Hessian of `T` with respect to the coordinate system `CS`
"""
HESS(T::Union{Sym,AbstractTens{order,dim,Sym}}, CS::CoorSystemSym{dim}) where {order,dim} =
    GRAD(GRAD(T, CS), CS)

"""
    coorsys_cartesian(coords = symbols("x y z", real = true))

Returns the cartesian coordinate system

# Examples
```julia
julia> Cartesian = coorsys_cartesian() ; 𝐗 = getcoords(Cartesian) ; 𝐄 = unitvec(Cartesian) ; ℬ = getbasis(Cartesian)

julia> 𝛔 = Tens(SymmetricTensor{2,3}((i, j) -> SymFunction("σ\$i\$j", real = true)(𝐗...))) ;

julia> DIV(𝛔, CScar)
Tens.TensCanonical{1, 3, Sym, Vec{3, Sym}}
# data: 3-element Vec{3, Sym}:
 Derivative(σ11(x, y, z), x) + Derivative(σ21(x, y, z), y) + Derivative(σ31(x, y, z), z)
 Derivative(σ21(x, y, z), x) + Derivative(σ22(x, y, z), y) + Derivative(σ32(x, y, z), z)
 Derivative(σ31(x, y, z), x) + Derivative(σ32(x, y, z), y) + Derivative(σ33(x, y, z), z)
# basis: 3×3 Tens.LazyIdentity{3, Sym}:
 1  0  0
 0  1  0
 0  0  1
# var: (:cont,)
```
"""
function coorsys_cartesian(coords = symbols("x y z", real = true))
    dim = length(coords)
    𝐗, 𝐄, ℬ = init_cartesian(coords)
    OM = sum([𝐗[i] * 𝐄[i] for i = 1:dim])
    χᵢ = ntuple(_ -> one(Sym), dim)
    return CoorSystemSym(OM, coords, ℬ, χᵢ)
end

"""
    coorsys_polar(coords = (symbols("r", positive = true), symbols("θ", real = true)); canonical = false)

Returns the polar coordinate system

# Examples
```julia
julia> Polar = coorsys_polar() ; r, θ = getcoords(Polar) ; 𝐞ʳ, 𝐞ᶿ = unitvec(Polar) ; ℬᵖ = getbasis(Polar)

julia> f = SymFunction("f", real = true)(r, θ) ;

julia> LAPLACE(f, Polar)
                               2
                              ∂
                             ───(f(r, θ))
                               2
               ∂             ∂θ
  2            ──(f(r, θ)) + ────────────
 ∂             ∂r                 r
───(f(r, θ)) + ──────────────────────────
  2                        r
∂r
```
"""
function coorsys_polar(
    coords = (symbols("r", positive = true), symbols("θ", real = true));
    canonical = false,
)
    (r, θ), (𝐞ʳ, 𝐞ᶿ), ℬᵖ = init_polar(coords, canonical = canonical)
    OM = r * 𝐞ʳ
    return CoorSystemSym(OM, coords, ℬᵖ, (one(Sym), r))
end

"""
    coorsys_cylindrical(coords = (symbols("r", positive = true), symbols("θ", real = true), symbols("z", real = true)); canonical = false)

Returns the cylindrical coordinate system

# Examples
```julia
julia> Cylindrical = coorsys_cylindrical() ; rθz = getcoords(Cylindrical) ; 𝐞ʳ, 𝐞ᶿ, 𝐞ᶻ = unitvec(Cylindrical) ; ℬᶜ = getbasis(Cylindrical)

julia> 𝐯 = Tens(Vec{3}(i -> SymFunction("v\$(rθz[i])", real = true)(rθz...)), ℬᶜ) ;

julia> DIV(𝐯, Cylindrical)
                                                  ∂
                                    vr(r, θ, z) + ──(vθ(r, θ, z))
∂                 ∂                               ∂θ
──(vr(r, θ, z)) + ──(vz(r, θ, z)) + ─────────────────────────────
∂r                ∂z                              r
```
"""
function coorsys_cylindrical(
    coords = (
        symbols("r", positive = true),
        symbols("θ", real = true),
        symbols("z", real = true),
    );
    canonical = false,
)
    (r, θ, z), (𝐞ʳ, 𝐞ᶿ, 𝐞ᶻ), ℬᶜ = init_cylindrical(coords, canonical = canonical)
    OM = r * 𝐞ʳ + z * 𝐞ᶻ
    return CoorSystemSym(OM, coords, ℬᶜ, (one(Sym), r, one(Sym)))
end

"""
    coorsys_spherical(coords = (symbols("θ", real = true), symbols("ϕ", real = true), symbols("r", positive = true)); canonical = false)

Returns the spherical coordinate system

# Examples
```julia
julia> Spherical = coorsys_spherical() ; θ, ϕ, r = getcoords(Spherical) ; 𝐞ᶿ, 𝐞ᵠ, 𝐞ʳ = unitvec(Spherical) ; ℬˢ = getbasis(Spherical)

julia> for σⁱʲ ∈ ("σʳʳ", "σᶿᶿ", "σᵠᵠ") @eval \$(Symbol(σⁱʲ)) = SymFunction(\$σⁱʲ, real = true)(\$r) end ;

julia> 𝛔 = σʳʳ * 𝐞ʳ ⊗ 𝐞ʳ + σᶿᶿ * 𝐞ᶿ ⊗ 𝐞ᶿ + σᵠᵠ * 𝐞ᵠ ⊗ 𝐞ᵠ ;

julia> div𝛔 = DIV(𝛔, Spherical)
Tens.TensRotated{1, 3, Sym, Vec{3, Sym}}
# data: 3-element Vec{3, Sym}:
                              (-σᵠᵠ(r) + σᶿᶿ(r))*cos(θ)/(r*sin(θ))
                                                                 0
 Derivative(σʳʳ(r), r) + (σʳʳ(r) - σᵠᵠ(r))/r + (σʳʳ(r) - σᶿᶿ(r))/r
# basis: 3×3 Tensor{2, 3, Sym, 9}:
 cos(θ)⋅cos(ϕ)  -sin(ϕ)  sin(θ)⋅cos(ϕ)
 sin(ϕ)⋅cos(θ)   cos(ϕ)  sin(θ)⋅sin(ϕ)
       -sin(θ)        0         cos(θ)
# var: (:cont,)

julia> div𝛔 ⋅ 𝐞ʳ
d            σʳʳ(r) - σᵠᵠ(r)   σʳʳ(r) - σᶿᶿ(r)
──(σʳʳ(r)) + ─────────────── + ───────────────
dr                  r                 r
```
"""
function coorsys_spherical(
    coords = (
        symbols("θ", real = true),
        symbols("ϕ", real = true),
        symbols("r", positive = true),
    );
    canonical = false,
)
    (θ, ϕ, r), (𝐞ᶿ, 𝐞ᵠ, 𝐞ʳ), ℬˢ = init_spherical(coords, canonical = canonical)
    OM = r * 𝐞ʳ
    rules = Dict(abs(sin(θ)) => sin(θ))
    return CoorSystemSym(OM, coords, ℬˢ, (r, r * sin(θ), one(Sym)); rules = rules)
end

"""
    coorsys_spheroidal(coords = (symbols("ϕ", real = true),symbols("p", real = true),symbols("q", positive = true),),
                            c = symbols("c", positive = true),tmp_coords = (symbols("p̄ q̄", positive = true)...,),)

Returns the spheroidal coordinate system

# Examples
```julia
julia> Spheroidal = coorsys_spheroidal() ; OM = getOM(Spheroidal)
Tens.TensCanonical{1, 3, Sym, Vec{3, Sym}}
# data: 3-element Vec{3, Sym}:
 c⋅p̄⋅q̄⋅cos(ϕ)
 c⋅p̄⋅q̄⋅sin(ϕ)
          c⋅p⋅q
# basis: 3×3 Tens.LazyIdentity{3, Sym}:
 1  0  0
 0  1  0
 0  0  1
# var: (:cont,)

julia> LAPLACE(OM[1]^2, Spheroidal)
2
```
"""
function coorsys_spheroidal(
    coords = (
        symbols("ϕ", real = true),
        symbols("p", real = true),
        symbols("q", positive = true),
    ),
    c = symbols("c", positive = true),
    tmp_coords = (symbols("p̄ q̄", positive = true)...,),
)
    ϕ, p, q = coords
    params = (c,)
    p̄, q̄ = tmp_coords
    OM = Tens(c * [p̄ * q̄ * cos(ϕ), p̄ * q̄ * sin(ϕ), p * q])
    # OM = Tens(c * [√(1 - p^2) * √(q^2 - 1) * cos(ϕ), √(1 - p^2) * √(q^2 - 1) * sin(ϕ), p * q])
    ℬ = RotatedBasis(
        Sym[
            -sin(ϕ) -p*sqrt(q^2 - 1)*cos(ϕ)/sqrt(q^2 - p^2) q*sqrt(1 - p^2)*cos(ϕ)/sqrt(q^2 - p^2)
            cos(ϕ) -p*sqrt(q^2 - 1)*sin(ϕ)/sqrt(q^2 - p^2) q*sqrt(1 - p^2)*sin(ϕ)/sqrt(q^2 - p^2)
            0 q*sqrt(1 - p^2)/sqrt(q^2 - p^2) p*sqrt(q^2 - 1)/sqrt(q^2 - p^2)
        ],
    )
    χᵢ = (
        c * sqrt(1 - p^2) * sqrt(q^2 - 1),
        c * sqrt(q^2 - p^2) / sqrt(1 - p^2),
        c * sqrt(q^2 - p^2) / sqrt(q^2 - 1),
    )
    tmp_var = Dict(1 - p^2 => p̄^2, q^2 - 1 => q̄^2)
    to_coords = Dict(p̄ => √(1 - p^2), q̄ => √(q^2 - 1))
    return CoorSystemSym(
        OM,
        coords,
        ℬ,
        χᵢ,
        tmp_coords,
        params;
        tmp_var = tmp_var,
        to_coords = to_coords,
    )
end

"""
    @set_coorsys CS
    @set_coorsys(CS)

Sets a coordinate system in order to avoid precising it in differential operators

# Examples
```julia
julia> Spherical = coorsys_spherical() ; θ, ϕ, r = getcoords(Spherical) ; 𝐞ᶿ, 𝐞ᵠ, 𝐞ʳ = unitvec(Spherical) ; vec = ("𝐞ᶿ", "𝐞ᵠ", "𝐞ʳ") ;

julia> @set_coorsys Spherical

julia> printvec(GRAD(𝐞ʳ),vec)
(1/r)𝐞ᶿ⊗𝐞ᶿ + (1/r)𝐞ᵠ⊗𝐞ᵠ

julia> printvec(DIV(𝐞ʳ ⊗ 𝐞ʳ),vec)
(2/r)𝐞ʳ

julia> LAPLACE(1/r)
0
```
"""
macro set_coorsys(CS = coorsys_cartesian(), vec = '𝐞', coords = nothing)
    m = @__MODULE__
    return quote
            $m.∂(t::AbstractTens{order,dim,Sym}, i::Integer) where {order,dim} = $m.∂(t, i, $(esc(CS)))
            $m.∂(t::AbstractTens{order,dim,Sym}, x::Sym) where {order,dim}  = $m.∂(t, x, $(esc(CS)))
            $m.∂(t::Sym, i::Integer) = $m.∂(t, i, $(esc(CS)))
            $m.∂(t::Sym, x::Sym) = $m.∂(t, x, $(esc(CS)))
            $m.GRAD(t::Union{Sym,AbstractTens}) = $m.GRAD(t, $(esc(CS)))
            $m.SYMGRAD(t::Union{Sym,AbstractTens}) = $m.SYMGRAD(t, $(esc(CS)))
            $m.DIV(t::AbstractTens) = $m.DIV(t, $(esc(CS)))
            $m.LAPLACE(t::Union{Sym,AbstractTens}) = $m.LAPLACE(t, $(esc(CS)))
            $m.HESS(t::Union{Sym,AbstractTens}) = $m.HESS(t, $(esc(CS)))

            coords = $(esc(coords)) === nothing ? string.(getcoords($(esc(CS)))) : ntuple(i -> i, getdim($(esc(CS))))
            ℬ = get_normalized_basis($(esc(CS)))

            Base.show(t::AbstractTens{order,dim,T}) where {order,dim,T} = printvec(change_tens(t, ℬ); vec = $(esc(vec)), coords = coords)
            Base.print(t::AbstractTens{order,dim,T}) where {order,dim,T} = printvec(change_tens(t, ℬ); vec = $(esc(vec)), coords = coords)
            Base.display(t::AbstractTens{order,dim,T}) where {order,dim,T} = printvec(change_tens(t, ℬ); vec = $(esc(vec)), coords = coords)

        end
end

export ∂, CoorSystemSym, getChristoffel
export GRAD, SYMGRAD, DIV, LAPLACE, HESS
export get_normalized_basis, get_natural_basis, natvec, unitvec, getcoords, getOM
export coorsys_cartesian, coorsys_polar, coorsys_cylindrical, coorsys_spherical, coorsys_spheroidal
export @set_coorsys


text:
using Plots, PairPlots
using Colors

##
N = 400_000
a = randn(N)
b = randn(N)
c = randn(N)
data = (;a,b,c)

##
for cname in (:green,:red,:purple)
    color = getproperty(Colors.JULIA_LOGO_COLORS, cname)
    corner(
        data,
        plotpercentiles=[],
        hist2d_kwargs=(;color, nbins=45, xguide="", yguide="", xticks=[], yticks=[]),
        contour_kwargs=(; xguide="", yguide="", xticks=[], yticks=[]),
        scatter_kwargs=(;color, alpha=0.1, xguide="", yguide="", xticks=[], yticks=[]),
        hist_kwargs=(;title="",xticks=[])
    )
    savefig("images/logo-$cname.png")
end

ext: jl

max_stars_repo_name: sefffal/CornerPlots

max_issues_repo_name: sefffal/CornerPlots


text:
using Test
using JuLIP

data = joinpath(dirname(pathof(JuLIP)), "..", "data") * "/"

@testset "EAM with ASE" begin
    using ASE
    import Random

    test_pots = joinpath(datadep"JuLIP_testpots", "JuLIP_data") * "/"
    alloy = test_pots * "PdAgH_HybridPd3Ag.eam.alloy"
    eam_fs = test_pots * "Fe-P.eam.fs"
    Ni = test_pots * "Ni.eam.fs"

    @test EAM(alloy) isa EAM{T, 1} where {T}
    @test EAM(eam_fs) isa EAM{T, 2} where {T}

    eam = pyimport("ase.calculators.eam")
    ase_calc = ASECalculator(eam.EAM(potential=Ni))
    # make sure we get a perfect fit ...
    julip_calc = EAM(Ni)

    atoms = bulk(:Ni) * 3
    Random.seed!(0)
    rattle!(atoms, 0.1)
    E_jl = energy(julip_calc, atoms)
    E_py = energy(ase_calc, atoms)
    # ... but even then we the evaluation codes aren't the
    #     same so we only get ca 1e-6 to 1e-7 match.
    @test abs(E_jl - E_py) < 1e-5

    # Test again with unordered species in parameter file
    ase_calc = ASECalculator(eam.EAM(potential=alloy))
    julip_calc = EAM(alloy)

    atoms = bulk(:Pd) * 3
    atoms.Z[1:3:end] .= AtomicNumber(:Ag)
    atoms.Z[1:5:end] .= AtomicNumber(:H)
    rattle!(atoms, 0.1)
    E_jl = energy(julip_calc, atoms)
    E_py = energy(ase_calc, atoms)
    @test E_jl ≈ E_py rtol=1e-6
end


text:
abstract type Jewel end

mutable struct RoseQuartz <: Jewel
    carats::Float64
    quality::String
end

mutable struct Sapphire <: Jewel
    color::String
    carats::Float64
    quality::String
end

color(j::RoseQuartz) = "rosepink"
color(j::Jewel) = "Use the loupe."
color(j::Sapphire) = j.color

function testtypecopy()
    a = Sapphire("blue", 5.0, "good")
    b = RoseQuartz(3.5, "excellent")

    j::Jewel = deepcopy(b)

    println("a is a Jewel? ", a isa Jewel)
    println("b is a Jewel? ", a isa Jewel)
    println("j is a Jewel? ", a isa Jewel)
    println("a is a Sapphire? ", a isa Sapphire)
    println("a is a RoseQuartz? ", a isa RoseQuartz)
    println("b is a Sapphire? ", b isa Sapphire)
    println("b is a RoseQuartz? ", b isa RoseQuartz)
    println("j is a Sapphire? ", j isa Sapphire)
    println("j is a RoseQuartz? ", j isa RoseQuartz)

    println("The color of j is ", color(j), ".")
    println("j is the same as b? ", j == b)
end

testtypecopy()

max_stars_repo_name: ethansaxenian/RosettaDecode

max_issues_repo_name: ethansaxenian/RosettaDecode

max_forks_repo_name: ethansaxenian/RosettaDecode


text:
include("types.jl")
include("homo.jl")
include("getprop.jl")

SeisForwExpt(attrib_mod::Union{AcousticBorn,ElasticBorn}, args1...; args2...) =
    PBorn(attrib_mod, args1...; args2...)

function PBorn(
    attrib_mod;
    medium::Medium = nothing,
    Q = nothing,
    medium_pert::Medium = medium,
    scatter_flag::Bool = false,
    tgrid::StepRangeLen = nothing,
    ageom::AGeom = nothing,
    srcwav::SrcWav = nothing,
    rfields::Vector{Symbol} = [:vz],
    sflag::Int64 = 2,
)
    issimilar(ageom, srcwav) || error("ageom and srcwav mismatch")

    nt = (length(tgrid) == length(srcwav[1].grid)) ? length(tgrid) : error("srcwav tgrid")
    np2 = nextpow(2, 2 * nt)
    fnpow2grid = FFTW.fftfreq(np2, inv(step(tgrid)))

    if (:Q ∈ names(medium.m)[1])
        vp0 = medium.ref[:vp]
        rho0 = medium.ref[:rho]
        Ku = rho0 * vp0 * vp0
        tau_epsilon = medium[:tau_epsilon]
        tau_sigma = medium[:tau_sigma]

        Kc = complexK(Ku, abs.(2.0 * pi .* fnpow2grid), tau_sigma, tau_epsilon)

        rho0 = fill(rho0, length(fnpow2grid))
        vp0 = sqrt.(Kc .* inv.(rho0))
    else
        m = [
            NamedArray(complex.([medium.ref[:vp], medium.ref[:rho]]), [:vp, :rho]) for
            i in fnpow2grid
        ]
    end

    cc = NamedArray([complex(0.0, 0.0)], [:dummy])

    if (scatter_flag)
        mesh_x = medium_pert.mgrid[2]
        mesh_z = medium_pert.mgrid[1]
        nz = length(mesh_z)
        nx = length(mesh_x)
        δx = step(mesh_x)
        δz = step(mesh_z)

        δmodKI = medium_pert[:KI] - (vp0 * vp0 * rho0)^(-1)
        δmodrr = medium_pert[:rhoI] - (rho0)^(-1)
    end

    ic = NamedArray(
        vcat(length.(medium.mgrid), [length(tgrid), np2]),
        vcat(dim_names(ndims(medium), "n"), [:nt, :np2]),
    )

    fc = NamedArray([step(tgrid)], [:dt])

    data = Records(tgrid, ageom, rfields)
    w = NamedArray(vcat([complex.(zeros(np2), zeros(np2)) for i = 1:3]), [:G, :S, :D])

    return PBorn(attrib_mod, medium, ageom, srcwav, fc, ic, cc, fnpow2grid, w, m, data)

end

function update!(pa::PBorn)
    ageom = pa.ageom
    nss = length(ageom)
    data = pa.data
    m = pa.m
    D = pa.w[:D]
    S = pa.w[:S]
    G = pa.w[:G]

    for iss = 1:nss
        for rfield in names(data[iss].d)[1]
            d = data[iss].d[rfield]
            for ir = 1:ageom[iss].nr
                fill!(D, complex(0.0))
                for is = 1:ageom[iss].ns
					off=offset(ageom[iss], is, ir)
					(sum(abs2.(off))==0) && error("distance cannot be zero")
                    sfield = names(pa.srcwav[iss].d)[1][1] # only one field for the source?
                    s = pa.srcwav[iss].d[sfield]
                    fill!(S, complex(0.0))
                    fill!(G, complex(0.0))
                    # zero pad wavelet
                    for it = 1:pa.ic[:nt]
                        S[it] = complex(s[it, is])
                    end
                    FFTW.fft!(S) # source wavelet in the frequency domain
                    # analytical expression for every frequency, except zero
                    G[1] = complex.(0.0, 0.0)
                    for iω = 2:pa.ic[:np2]
                        ω = 2.0 * pi * abs(pa.fnpow2grid[iω])
                        k = ω * inv(m[iω][:vp])
    					(k == 0.0) && error("wavenumber cannot be zero")
                        # if (scatter_flag)
                        #     term = complex(0.0, 0.0)
                        #     for ix = 1:nx
                        #         @simd for iz = 1:nz
                        #             if (δmodKI[iz, ix] ≠ 0.0)
                        #                 term +=
                        #                     (
                        #                         G0_homo_acou(
                        #                             sx[is] - mesh_x[ix],
                        #                             sz[is] - mesh_z[iz],
                        #                             k,
                        #                             rho0,
                        #                         )[1] * G0_homo_acou(
                        #                             rx[ir] - mesh_x[ix],
                        #                             rz[ir] - mesh_z[iz],
                        #                             k,
                        #                             rho0,
                        #                         )[1] .* ω .* ω .* δmodKI[iz, ix]
                        #                     ) *
                        #                     δx *
                        #                     δz
                        #                 # factor due to integration
                        #             end
                        #         end
                        #     end

                        # else
                        # if (src_flag == 2)
                        term = G0(
                            pa.attrib_mod,
                            eval(sfield)(),
                            eval(rfield)(),
                            k,
                            m[iω],
                            off...,
                        )
                        # elseif (src_flag == 1)
                        # G[iω] =
                        # G0_homo_acou(x, z, k, rho0[iω]) * im * abs(pa.fnpow2grid[iω])
                        # else
                        # error("invalid src_flag")
                        # end
                        # end
                        if (pa.fnpow2grid[iω] > 0)
                            G[iω] = term
                        elseif (pa.fnpow2grid[iω] < 0)
                            G[iω] = conj(term)
                        end

                    end
                    # convolution and stack over simultaneous sources
                    D += G .* S
                end

                # back to time domain
                FFTW.ifft!(D)

                # truncate
                for it = 1:pa.ic[:nt]
                    d[it, ir] = real(D[it])
                end
            end
        end
    end
    return data
end


text:
export save, load, @save, @load
import FileIO # save, load
using JLD2: JLD2, JLDWriteSession, jldopen, isgroup, lookup_offset
#include("serialize.jl") ## serialize  ## TODO: extend Base.Serialization, cover CuArrays

"""
    Knet.save(filename, args...; kwargs...)

Call `FileIO.save` after serializing Knet specific args.

File format is determined by the filename extension. JLD and JLD2 are supported. Other formats
may work if supported by FileIO, please refer to the documentation of FileIO and the specific
format.  Example:

    Knet.save("foo.jld2", "name1", value1, "name2", value2)
"""
function save(fname,args...;kwargs...)
     FileIO.save(fname,serialize.(args)...;kwargs...)
end

"""
    Knet.load(filename, args...; kwargs...)

Call `FileIO.load` then deserialize Knet specific values.

File format is determined by FileIO. JLD and JLD2 are supported. Other formats may work if
supported by FileIO, please refer to the documentation of FileIO and the specific format.
Example:

    Knet.load("foo.jld2")           # returns a ("name"=>value) dictionary
    Knet.load("foo.jld2", "name1")  # returns the value of "name1" in "foo.jld2"
    Knet.load("foo.jld2", "name1", "name2")   # returns tuple (value1, value2)
"""
function load(fname,args...;kwargs...)
     serialize(FileIO.load(fname,args...;kwargs...))
end

"""
    Knet.@save "filename" variable1 variable2...

Save the values of the specified variables to filename in JLD2 format.

When called with no variable arguments, write all variables in the global scope of the current
module to filename.  See [JLD2](https://github.com/JuliaIO/JLD2.jl).
"""
macro save(filename, vars...)
    if isempty(vars)
        # Save all variables in the current module
        quote
            let
                m = $(__module__)
                f = JLD2.jldopen($(esc(filename)), "w")
                wsession = JLD2.JLDWriteSession()
                try
                    for vname in names(m; all=true)
                        s = string(vname)
                        if !occursin(r"^_+[0-9]*$", s) # skip IJulia history vars
                            v = getfield(m, vname)
                            if !isa(v, Module)
                                try
                                    write(f, s, serialize(v), wsession)
                                catch e
                                    if isa(e, PointerException)
                                        @warn("skipping $vname because it contains a pointer")
                                    else
                                        rethrow(e)
                                    end
                                end
                            end
                        end
                    end
                finally
                    close(f)
                end
            end
        end
    else
        writeexprs = Vector{Expr}(undef, length(vars))
        for i = 1:length(vars)
            writeexprs[i] = :(write(f, $(string(vars[i])), serialize($(esc(vars[i]))), wsession))
        end

        quote
            JLD2.jldopen($(esc(filename)), "w") do f
                wsession = JLD2.JLDWriteSession()
                $(Expr(:block, writeexprs...))
            end
        end
    end
end

"""
    Knet.@load "filename" variable1 variable2...

Load the values of the specified variables from filename in JLD2 format.

When called with no variable arguments, load all variables in filename.  See
[JLD2](https://github.com/JuliaIO/JLD2.jl).
"""
macro load(filename, vars...)
    if isempty(vars)
        if isa(filename, Expr)
            throw(ArgumentError("filename argument must be a string literal unless variable names are specified"))
        end
        # Load all variables in the top level of the file
        readexprs = Expr[]
        vars = Symbol[]
        f = JLD2.jldopen(filename)
        try
            for n in keys(f)
                if !JLD2.isgroup(f, JLD2.lookup_offset(f.root_group, n))
                    push!(vars, Symbol(n))
                end
            end
        finally
            close(f)
        end
    end
    return quote
        ($([esc(x) for x in vars]...),) = JLD2.jldopen($(esc(filename))) do f
            ($([:(serialize(read(f, $(string(x))))) for x in vars]...),)
        end
        $(Symbol[v for v in vars]) # convert to Array
    end
end


text:
using StaticArrays

########################################
# SArrays and SVectors
########################################
function cembed(ts::Vector{SArray{Size,T,N,L}},
        in_which_pos,
        at_what_lags) where {Size, T, N, L}
    dim = length(in_which_pos)
    minlag, maxlag = minimum(at_what_lags), maximum(at_what_lags)
    npts = length(ts[1]) - (maxlag + abs(minlag))
    E = zeros(T, dim, npts)

    for i in 1:length(in_which_pos)
        ts_ind = in_which_pos[i]
        TS = ts[ts_ind]
        lag = at_what_lags[i]

        if lag > 0
            E[i, :] = TS[((1 + abs(minlag)) + lag):(end - maxlag) + lag]
        elseif lag < 0
            E[i, :] = TS[((1 + abs(minlag)) - abs(lag)):(end - maxlag - abs(lag))]
        elseif lag == 0
            E[i, :] = TS[(1 + abs(minlag)):(end - maxlag)]
        end
    end
    embeddingdata = EmbeddingData{dim, T}(
        ts,  # the dataseries
        ["" for i = 1:length(ts)], # empty labels by default
        ["" for i = 1:length(ts)], # empty descriptions by default
        EmbeddingPositions(in_which_pos), # the positions in which the dataseries appear
        EmbeddingLags(at_what_lags) # the embedding lag for each column
        )
    Embedding(Array(E), embeddingdata)
end

"""
    cembed(data::SArray) where T

Construct an embedding from data series gathered in a `SArray`.
"""
function cembed(data::SArray)
	if size(data, 1) > size(data, 2)
		#info("Treating each row as a point")
		dim = size(data, 2)
        which_pos = [i for i = 1:dim]
        which_lags = [0 for i in 1:dim]
		cembed([data[:, i] for i = 1:dim], which_pos, which_lags)
	else
		#info("Treating each column of data as a point")
		dim = size(data, 1)
        which_pos = [i for i = 1:dim]
        which_lags = [0 for i in 1:dim]
		cembed([data[i, :] for i = 1:dim], which_pos, which_lags)
	end
end

"""
    cembed(data::SArray) where T

Construct an embedding from data series gathered in a `SArray`, specifying the
dimensionality, which variables in the embedding are represented by which
variables in the `data` array, and what the embedding lags should be.
"""
function cembed(data::SArray, which_pos, which_lags)
    if size(data, 1) > size(data, 2)
		#info("Treating each row as a point")
		dim = size(data, 2)
		cembed([data[:, i] for i = 1:dim], which_pos, which_lags)
	else
		#info("Treating each column of data as a point")
		dim = size(data, 1)
		cembed([data[i, :] for i = 1:dim], which_pos, which_lags)
	end
end

export cembed


text:
using Statistics, SparseArrays
using MAT, DelimitedFiles
"""
```
eqcond(m::TwoAssetHANK)
```
Expresses the equilibrium conditions in canonical form using Γ0, Γ1, C, Ψ, and Π matrices.

# Outputs
* `Γ0` (`n_states` x `n_states`) holds coefficients of current time states.
* `Γ1` (`n_states` x `n_states`) holds coefficients of lagged states.
* `C`  (`n_states` x `1`) is a vector of constants
* `Ψ`  (`n_states` x `n_shocks_exogenous`) holds coefficients of iid shocks.
* `Π`  (`n_states` x `n_states_expectational`) holds coefficients of expectational states.
"""
function eqcond(m::TwoAssetHANK)

    # Read in parameters
    α  = m[:aalpha].value::Float64
    δ  = m[:ddelta].value::Float64
    ρ  = m[:rrho].value::Float64
    γ  = m[:ggamma].value::Float64
    κ  = m[:kappa].value::Float64
    ξ  = m[:xxi].value::Float64
    ϕ  = m[:pphi].value::Float64
    χ0 = m[:chi0].value::Float64
    χ1 = m[:chi1].value::Float64
    χ2 = m[:chi2].value::Float64

    ddeath = m[:ddeath].value::Float64
    a_lb   = m[:a_lb].value::Float64
    pam    = m[:pam].value::Float64
    τ_I    = m[:tau_I].value::Float64
    trans  = m[:trans].value::Float64
    n_SS   = m[:n_SS].value::Float64
    ν_aggZ = m[:nnu_aggZ].value::Float64
    σ_aggZ = m[:ssigma_aggZ].value::Float64

    # Set liquid rates
    r_b_SS       = m[:r_b_SS].value::Float64
    r_b_borr_SS  = m[:r_b_borr_SS].value::Float64
    borrwedge_SS = m[:borrwedge_SS].value::Float64

    K_liquid = get_setting(m, :K_liquid)::Bool
    r_b_fix  = get_setting(m, :r_b_fix)::Bool
    r_b_phi  = get_setting(m, :r_b_phi)::Bool
    B_fix    = get_setting(m, :B_fix)::Bool

    aggregate_variables        = get_setting(m, :aggregate_variables)::Bool
    distributional_variables   = get_setting(m, :distributional_variables)::Bool
    distributional_variables_1 = get_setting(m, :distributional_variables_1)::Bool
    permanent                  = get_setting(m, :permanent)::Bool

    I   = get_setting(m, :I)::Int64
    J   = get_setting(m, :J)::Int64
    I_g = get_setting(m, :I_g)::Int64
    J_g = get_setting(m, :J_g)::Int64
    N   = get_setting(m, :y_size)::Int64

    a   = get_setting(m, :a)::Vector{Float64}
    b   = get_setting(m, :b)::Vector{Float64}
    a_g = get_setting(m, :a_g)::Vector{Float64}
    b_g = get_setting(m, :b_g)::Vector{Float64}

    agrid_new = get_setting(m, :agrid_new)::Int64
    bgrid_new = get_setting(m, :bgrid_new)::Int64
    ygrid_new = get_setting(m, :ygrid_new)::Int64

    amin = get_setting(m, :amin)
    bmin = get_setting(m, :bmin)
    amax = get_setting(m, :amax)
    bmax = get_setting(m, :bmax)

    KL      = get_setting(m, :KL_0)::Float64

    n_v = get_setting(m, :n_v)::Int64
    n_g = get_setting(m, :n_g)::Int64
    n_p = get_setting(m, :n_p)::Int64
    n_Z = get_setting(m, :n_Z)::Int64
    nVars    = get_setting(m, :nVars)::Int64
    nEErrors = get_setting(m, :nEErrors)::Int64

    #--- Taken from what used to be inside
    a, a_g, a_g_0pos_arr          = create_a_grid(agrid_new, J, J_g, amin, amax)
    _, _, _, b, b_g, b_g_0pos_arr = create_b_grid(bgrid_new, I, I_g)
    lambda, y, y_mean, y_dist, _  = create_y_grid(N, ygrid_new)
    a_g_0pos = first(a_g_0pos_arr)
    b_g_0pos = first(b_g_0pos_arr)

    # Construct problem functions
    util, deposit, cost = construct_problem_functions(γ, χ0, χ1, χ2, a_lb)

    vars_SS     = vec(m[:vars_SS].value)::Vector{Float64}
    #vars_SS = vec(load("data/vars_SS.jld2","vars_SS"))
    V_SS   = vars_SS[1:n_v]
    g_SS   = vars_SS[n_v + 1 : n_v + n_g]
    K_SS   = vars_SS[n_v + n_g + 1]
    r_b_SS = vars_SS[n_v + n_g + 2]

    # Aggregate output and aggregate consumption
    aggY_SS = aggregate_variables ? vars_SS[n_v+n_g+3] : 0.0
    aggC_SS = aggregate_variables ? vars_SS[n_v+n_g+4] : 0.0

    # Consumption and earnings inequality
    C_Var_SS    = distributional_variables ? vars_SS[n_v+n_g+3] : 0.0
    earn_Var_SS = distributional_variables ? vars_SS[n_v+n_g+4] : 0.0

    # Consumption of wealthy and poor hand-to-mouth
    C_WHTM_SS = distributional_variables_1 ? vars_SS[n_v+n_g+3] : 0.0
    C_PHTM_SS = distributional_variables_1 ? vars_SS[n_v+n_g+4] : 0.0

    aggZ_SS = vars_SS[n_v+n_g+n_p+1] # Aggregate Z

    @inline function get_residuals(vars::Vector{T}) where {T<:Real}

        # ------- Unpack variables -------
        V      = reshape(vars[1:n_v] .+ V_SS, I, J, N) # value function
        g      = vars[n_v + 1 : n_v + n_g] .+ g_SS     # distribution
        K::T   = vars[n_v + n_g + 1]        + K_SS     # aggregate capital
        r_b::T = vars[n_v + n_g + 2]        + r_b_SS

        # Aggregate output, aggregate consumption
        aggY::T = aggregate_variables ? vars[n_v+n_g+3] + aggY_SS : zero(T)
        aggC::T = aggregate_variables ? vars[n_v+n_g+4] + aggC_SS : zero(T)

        # Consumption and earnings inequality
        C_Var::T    = distributional_variables ? vars[n_v+n_g+3] + C_Var_SS    : zero(T)
        earn_Var::T = distributional_variables ? vars[n_v+n_g+4] + earn_Var_SS : zero(T)

        # Consumption of wealthy and poor hand-to-mouth
        C_WHTM::T = distributional_variables_1 ? vars[n_v+n_g+3] + C_WHTM_SS : zero(T)
        C_PHTM::T = distributional_variables_1 ? vars[n_v+n_g+4] + C_PHTM_SS : zero(T)

        aggZ::T       = vars[n_v + n_g + n_p + 1] + aggZ_SS
        V_Dot         = vars[nVars +       1 : nVars + n_v]
        g_Dot         = vars[nVars + n_v + 1 : nVars + n_v + n_g]
        aggZ_Dot::T   = vars[nVars + n_v + n_g + n_p + 1]
        VEErrors      = vars[2*nVars + 1 : 2 * nVars + n_v]
        aggZ_Shock::T = vars[2*nVars + nEErrors + 1]

        # Prices
        w   = (1-α) * (K ^ α) * n_SS ^ (-α) * (!permanent ? exp(aggZ) ^ (1-α) : 1.0)
        r_a = α * (K ^ (α-1)) * ((!permanent ? exp(aggZ) : 1.0) * n_SS) ^ (1-α) - δ

        # Set liquid rates
        r_b_borr = r_b .+ borrwedge_SS

        # Initialize grids
        dab_g = reshape(repeat(backward_difference(a_g, b_g), N, 1), I_g, J_g, N)
        a_gg  = repeat(repeat(a_g, inner=I_g), outer=N)
        b_gg  = repeat(repeat(b_g, outer=J_g), outer=N)

        g_end = (1 - sum(g .* vec(dab_g)[1:end-1])) / dab_g[I_g, J_g, N]
        gg    = vcat(g, g_end)

        loc = findall(b .== 0)
        dab_g_small          = vec(dab_g[:,:,1])
        dab_g_small          = dab_g_small ./ dab_g_small[loc] * ddeath
        dab_g_small[loc]    .= 0.0
        death_process        = -ddeath * my_speye(I_g * J_g)
        death_process[loc,:] = vec(dab_g_small)
        death_process        = kron(my_speye(N), death_process)

        r_b_vec = r_b .* (b .>= 0) + r_b_borr .* (b .< 0)

        # Other necessary objects
        y_shock      = y .* exp.(κ * aggZ * (y .- y_mean) ./ std(y))
        y_shock_mean = dot(y_shock, y_dist)
        y_shock      = real(y_shock ./ y_shock_mean .* y_mean)

        println("Timing: solve_hjb()")
        @time c, s, d  = solve_hjb(V, a_lb, γ, ddeath, pam, trans, ξ, τ_I, aggZ, w,
                                   r_b_vec, y_shock, a, b, cost, util, deposit;
                                   permanent=permanent)

        interp_decision = kron(my_speye(N), interp(b_g, a_g, b, a))
        d_g = reshape(interp_decision * vec(d), I_g, J_g, N)
        s_g = reshape(interp_decision * vec(s), I_g, J_g, N)
        c_g = reshape(interp_decision * vec(c), I_g, J_g, N)

        # Derive transition matrices
        println("Timing: transition()")
        @time A, AT = transition(ddeath, pam, ξ, w, a_lb, aggZ, d, d_g, s, s_g, r_a,
                                 a, a_g, b, b_g, y_shock, cost; permanent=permanent)

        cc  = kron(lambda, my_speye(I*J))
        ccu = kron(lambda, my_speye(I_g*J_g))

        # Full transition matrix
        A  = A + cc
        AT = (AT + ccu)'

        #----------------------------------------------------------------
        # KFE
        #----------------------------------------------------------------
        ggdab_g = gg .* vec(dab_g)
        gIntermediate = (1.0 ./ vec(dab_g)) .* (AT * ggdab_g) + death_process * gg

        #----------------------------------------------------------------
        # Compute equilibrium conditions
        #----------------------------------------------------------------
        # HJB equation
        perm_mult   = !permanent ? ρ + ddeath : ρ + ddeath - (1 - γ) * aggZ
        hjbResidual = vec(util.(c)) + A * vec(V) + V_Dot + VEErrors - perm_mult * vec(V)

        # KFE
        gResidual  = g_Dot - gIntermediate[1:n_g]
        K_Residual = (K_liquid ? sum((a_gg .+ b_gg) .* ggdab_g) : sum(a_gg .* ggdab_g)) - K
        r_b_Residual = if r_b_fix
                           r_b_SS - r_b
                       elseif r_b_phi
                           sum(b_gg .* ggdab_g) - B_SS * exp(1/ϕ * (r_b - r_b_SS))
                       elseif B_fix
                           -dot(gIntermediate, vec(dab_g) .* b_gg)
                       elseif K_liquid
                           r_a_out - illiquid_wedge - r_b
                       end

        # Law of motion for aggregate TFP shock
        aggZ_Residual = aggZ_Dot - (-ν_aggZ * aggZ + σ_aggZ * aggZ_Shock)

        # Return equilibrium conditions
        if aggregate_variables

            aggY_out = (K ^ α) * (!permanent ? exp(aggZ) * n_SS : n_SS ) ^ (1 - α)
            aggC_out = sum(vec(c_g) .* ggdab_g)
            Y_Residual = aggY_out - aggY
            C_Residual = aggC_out - aggC

            return [hjbResidual; gResidual; K_Residual; r_b_Residual; Y_Residual;
                    C_Residual; aggZ_Residual]

        elseif distributional_variables

            r_b_g = repeat(repeat(r_b .* (b_g .>= 0) + r_b_borr .* (b_g .< 0), outer=J_g), outer=N)
            C_Var_out = sum(log.(vec(c_g)).^2 .* ggdab_g) - sum(log.(vec(c_g)) .* ggdab_g) ^ 2

            earn = vec(log.((1-τ_I) * w * repeat(repeat(vec(y), inner=I_g), inner=J_g) .+ b_gg .*
                            r_b_g .+ ddeath*pam) .+ trans .+ a_gg .* (r_a + ddeath*pam))
            earn_Var_out = sum(earn.^2 .* ggdab_g) - sum(earn .* ggdab_g) ^ 2

            C_Var_Residual    = C_Var_out - C_Var
            earn_Var_Residual = earn_Var_out - earn_Var

            return [hjbResidual; gResidual; K_Residual; r_b_Residual; C_Var_Residual;
                    earn_Var_Residual; aggZ_Residual]

        elseif distributional_variables_1

            WHTM_indicator  = zeros(I_g,J_g,N)
            WHTM_indicator[b_g_0pos:b_g_0pos+1, a_g_0pos+2:end, :] .= 1.
            C_WHTM_out      = sum(vec(WHTM_indicator) .* vec(c_g) .* ggdab_g)

            PHTM_indicator  = zeros(I_g,J_g,N)
            PHTM_indicator[b_g_0pos:b_g_0pos+1, a_g_0pos:a_g_0pos+2:end, :] .= 1.
            C_PHTM_out      = sum(vec(PHTM_indicator) .* vec(c_g) .* ggdab_g)

            C_WHTM_Residual = C_WHTM_out - C_WHTM
            C_PHTM_Residual = C_PHTM_out - C_PHTM

            return [hjbResidual; gResidual; K_Residual; r_b_Residual; C_WHTM_Residual;
                    C_PHTM_Residual; aggZ_Residual]
        end
        return [hjbResidual; gResidual; K_Residual; r_b_Residual; aggZ_Residual]
    end

    out = get_residuals(zeros(Float64, 2 * nVars + nEErrors + 1))

    #JLD2.jldopen("reference/eqcond_after_.jld2", true, true, true, IOStream) do file
    #    file["residuals"] = out
    #end

    my_out = vec(DelimitedFiles.readdlm("reference/my_residuals.csv", ','))

    @show maximum(abs.(my_out - vec(out)))#, length(findall(x->x>1e-5, abs.(my_out - out)))
    #@assert isapprox(my_out, vec(out), rtol=1e-4)

    nstates = nVars    # n_states(m)
    n_s_exp = nEErrors # n_shocks_expectational(m)
    n_s_exo = n_Z      # n_shocks_exogenous(m)

    @time get_residuals(zeros(Float64, 2 * nstates + n_s_exp + 1))

    x = zeros(Float64, 2 * nstates + n_s_exp + 1)
    @time derivs = ForwardDiff.sparse_jacobian(get_residuals, x)

    # vars = zeros(Float64, 2 * nstates + n_s_exp + n_s_exo)

    Γ1 = -derivs[:, 1:nstates]
    Γ0 =  derivs[:,   nstates +           1:2 * nstates]
    Π  = -derivs[:, 2*nstates +           1:2 * nstates + n_s_exp]
    Ψ  = -derivs[:, 2*nstates + n_s_exp + 1:2 * nstates + n_s_exp + n_s_exo]
    C  =  spzeros(Float64, nstates)

    if typeof(Ψ) == Vector{Float64}
        Ψ = reshape(Ψ, length(Ψ), n_s_exo)
    end

    test_out = load("data/eqcond_output_matlab.jld2")
    @assert test_out["g0"] == Γ0
    @show maximum(abs.(test_out["g1"] - Γ1))
    @assert test_out["g1"] ≈ Γ1
    @assert test_out["psi"] == Ψ
    @assert test_out["pi"] == Π

    return Γ0, Γ1, Ψ, Π, C
end


text:
#**************************************************************************************
# make.jl
# =============== part of the GeoEfficiency.jl package.
#
# script for building documentation of the GeoEfficiency.jl package.
#
#  _args = ["clean", "doctest"]
#  include(raw"C:\Users\Mohamed\.julia\dev\GeoEfficiency\docs\make.jl")
#**************************************************************************************

using Documenter, DocumenterLaTeX
using GeoEfficiency

_args = @isdefined(_args) ? _args : ARGS
const PAGES = Any[
    "Home" => "index.md",
    "Manual" => [
		"manual/GeoEfficiency.md",
        #"manual/Error.md",
        #"manual/Input_Console.md",
        "manual/Physics_Model.md",
        "manual/Calculations.md",
        "manual/Output_Interface.md",
        "manual/Input_Batch.md",
     ],

    "Development" => "manual/Development.md",
    "Index" => "list.md",
    "LICENSE" => "LICENSE.md",
]

const formats = Any[
    Documenter.HTML(
        prettyurls = false,  #get(ENV, "CI", nothing) == "true",
        canonical = "https://github.com/DrKrar/GeoEfficiency.jl/v0.9.4-dev/",

    ),

]

if "pdf" in _args
    Sys.iswindows() ?   push!(formats, LaTeX(platform = "native")) :
                        push!(formats, LaTeX(platform = "docker"))
end

makedocs(
    format  = formats,
    modules = [GeoEfficiency],
    clean   = "clean" in _args,
    doctest = "doctest" in _args,
    sitename= "GeoEfficiency.jl",
    authors = "Mohamed E. Krar",
    pages   = PAGES,
    assets  = ["assets/custom.css"],
)

const REPO = "github.com/DrKrar/GeoEfficiency.jl.git" # "github.com/GeoEfficiency/GeoEfficiency.github.io.git"
const VERSIONS =["stable" => "v^", "v#.#", "dev" => "dev"]  # order of versions in drop down menu.
const BRANCH = "gh-pages" # "master"

mktempdir() do tmp
    # Hide the PDF from html-deploydocs
    build = joinpath(@__DIR__, "build")
    files = readdir(build)
    idx = findfirst(f -> startswith(f, "GeoEfficiency.jl") && endswith(f, ".pdf"), files)
    pdf = idx === nothing ? nothing : joinpath(build, files[idx])
    if pdf !== nothing
        pdf = mv(pdf, joinpath(tmp, basename(pdf)))
    end
    # Deploy HTML pages
    @info "Deploying HTML pages"
    deploydocs(
        repo = REPO,
        branch = BRANCH, #"gh-pages"
        versions = VERSIONS,
    )

    # Put back PDF into docs/build/pdf
    mkpath(joinpath(build, "pdf"))
    if pdf !== nothing
        pdf = mv(pdf, joinpath(build, "pdf", basename(pdf)))
    end
    # Deploy PDF
    @info "Deploying PDF"
    deploydocs(
        repo = REPO,
        target = "build/pdf",
        branch = BRANCH * "-pdf", #"gh-pages-pdf",
        forcepush = true,
    )
end


text:
"""
    ParticleSwarm(n_particles = 3,
                  w = 1.0,
                  c1 = 2.0,
                  c2 = 2.0,
                  prob_shift = 0.25,
                  rng = Random.GLOBAL_RNG)

Instantiate a particle swarm optimization tuning strategy. A swarm is initiated
by sampling hyperparameters with their customizable priors, and new models are
generated by referencing each member's and the swarm's best models so far.

### Supported Ranges

A single one-dimensional range or vector of one-dimensional ranges can be
specified. `ParamRange` objects are constructed using the `range` method. If not
paired with a prior, then one is fitted and truncated if bounded, as follows:

| Range Types             | Default Distribution                        |
|:----------------------- |:------------------------------------------- |
| `NominalRange`          | `Dirichlet([1, 1, ..., 1])`                 |
| Bounded `NumericRange`  | `Uniform(lower, upper)`                     |
| Positive `NumericRange` | `Gamma(α=(origin/unit)^2, θ=unit^2/origin`) |
| Other `NumericRange`    | `Normal(origin, unit)`                      |

Specifically, in `ParticleSwarm`, the `range` field of a `TunedModel` instance
can be:

- a single one-dimensional range (`ParamRange` object) `r`

- a pair of the form `(r, d)`, with `r` as above and where `d` is:

    - a Dirichlet distribution with the same number of categories as `r.values`
      (for `NominalRange` `r`)

    - any `Distributions.UnivariateDistribution` *instance* (for `NumericRange`
      `r`)

    - one of the distribution *types* in the table below, for automatic fitting
      using `Distributions.fit(d, r)` to a distribution whose support always
      lies between `r.lower` and `r.upper` (for `NumericRange` `r`)

- any vector of objects of the above form

| Range Types             | Distribution Types                                                                           |
|:----------------------- |:-------------------------------------------------------------------------------------------- |
| Bounded `NumericRange`  | `Arcsine`, `Uniform`, `Biweight`, `Cosine`, `Epanechnikov`, `SymTriangularDist`, `Triweight` |
| Positive `NumericRange` | `Gamma`, `InverseGaussian`, `Poisson`                                                        |
| Any `NumericRange`      | `Normal`, `Logistic`, `LogNormal`, `Cauchy`, `Gumbel`, `Laplace`                             |

### Examples

    using Distributions

    range1 = range(model, :hyper1, lower=0, upper=1)

    range2 = [(range(model, :hyper1, lower=1, upper=10), Arcsine),
              range(model, :hyper2, lower=2, upper=Inf, unit=1, origin=3),
              (range(model, :hyper2, lower=2, upper=4), Normal(0, 3)),
              range(model, :hyper3, values=[:ball, :tree])]

### Algorithm

Hyperparameter ranges are sampled and concatenated into position vectors for
each swarm particle. Velocity is initiated to be zeros, and in each iteration,
every particle's position is updated to approach its personal best and the
swarm's best models so far with the equations:

\$vₖ₊₁ = w⋅vₖ + c₁⋅rand()⋅(pbest - xₖ) + c₂⋅rand()⋅(gbest - xₖ)\$

\$xₖ₊₁ = xₖ + vₖ₊₁\$

New models are then generated for evaluation by mutating the fields of a deep
copy of `model`. If the corresponding range has a specified `scale` function,
then the transformation is applied before the hyperparameter is returned. If
`scale` is a symbol (eg, `:log`), it is ignored.

### Discrete Hyperparameter Handling

Since particle swarm is an optimization method for continuous problems, integer
and nominal hyperparameters require special handling: they are converted to
continuous values, and transformed back to their original domains at each step
for evaluation.

For integer `NumericRange`s, a continuous distribution is fitted to generate
initial values for the swarm. They are then rounded when each particle is mapped
to the corresponding candidate model.

`NominalRange`s on the other hand are represented as categorical distributions
over their values. Hence, we use Dirichlet prior distributions to initialize a
probability vector for each particle, defaulting to the uniform distribution
Dirichlet([1, 1, ..., 1]). The same velocity and position updates apply, but
probability values are further clamped in the range [0, 1] and normalized to sum
up to 1. When a better model is found, we replicate both its probability vector
and sampled value by shifting unchosen categories' weights towards the selected
one for pbest and gbest models:

\$pᵢ = (1 - prob_shift) * pᵢ\$

\$pₛ = pₛ + prob_shift\$

where pₛ is the probability of the sampled hyperparameter value. For more
information, refer to "A New Discrete Particle Swarm Optimization Algorithm" by
Strasser, Goodman, Sheppard, and Butcher.
"""
mutable struct ParticleSwarm{R<:AbstractRNG} <: AbstractParticleSwarm
    n_particles::Int
    w::Float64
    c1::Float64
    c2::Float64
    prob_shift::Float64
    rng::R
    # TODO: topology
end

function ParticleSwarm(;
    n_particles=3,
    w=1.0,
    c1=2.0,
    c2=2.0,
    prob_shift=0.25,
    rng::R=Random.GLOBAL_RNG
) where {R}
    swarm = ParticleSwarm{R}(n_particles, w, c1, c2, prob_shift, rng)
    message = MLJTuning.clean!(swarm)
    isempty(message) || @warn message
    return swarm
end

function MLJTuning.clean!(tuning::ParticleSwarm)
    warning = ""
    if tuning.n_particles < 3
        warning *= "ParticleSwarm requires at least 3 particles. Resetting n_particles=3. "
        tuning.n_particles = 3
    end
    if tuning.w < 0
        warning *= "ParticleSwarm requires w ≥ 0. Resetting w=1. "
        tuning.w = 1
    end
    if tuning.c1 < 0
        warning *= "ParticleSwarm requires c1 ≥ 0. Resetting c1=2. "
        tuning.c1 = 2
    end
    if tuning.c2 < 0
        warning *= "ParticleSwarm requires c2 ≥ 0. Resetting c2=2. "
        tuning.c2 = 2
    end
    if !(0 ≤ tuning.prob_shift < 1)
        warning *= "ParticleSwarm requires 0 ≤ prob_shift < 1. Resetting prob_shift=0.25. "
        tuning.prob_shift = 0.25
    end
    return warning
end

function MLJTuning.setup(tuning::ParticleSwarm, model, ranges, n, verbosity)
    return initialize(ranges, tuning)
end

function MLJTuning.models(
    tuning::ParticleSwarm,
    model,
    history,
    state::ParticleSwarmState{T},
    n_remaining,
    verbosity
) where {T}
    n_particles = tuning.n_particles
    if !isnothing(history)
        sign = MLJTuning.signature(history[1].measure[1])
        measurements = Vector{T}(undef, n_particles)
        map(history[end-n_particles+1:end]) do h
            measurements[h.metadata] = sign * h.measurement[1]
        end
        pbest!(state, measurements, tuning)
        gbest!(state)
        move!(tuning.rng, state, T(tuning.w), T(tuning.c1), T(tuning.c2))
    end
    retrieve!(state, tuning)
    fields = getproperty.(state.ranges, :field)
    new_models = map(1:n_particles) do i
        clone = deepcopy(model)
        for (field, param) in zip(fields, getindex.(state.parameters, i))
            recursive_setproperty!(clone, field, param)
        end
        (clone, i)
    end
    return new_models, state
end

function MLJTuning.tuning_report(tuning::ParticleSwarm, history, state)
    fields = getproperty.(state.ranges, :field)
    scales = MLJBase.scale.(state.ranges)
    return (; plotting = MLJTuning.plotting_report(fields, scales, history))
end


text:
# Note that this script can accept some limited command-line arguments, run
# `julia build_tarballs.jl --help` to see a usage message.
using BinaryBuilder, Pkg

name = "MAFFT"
version = v"7.490"

# Notes
# - build fails on windows, include file sys/resource.h is missing

# Collection of sources required to complete build
sources = [
    ArchiveSource("https://mafft.cbrc.jp/alignment/software/mafft-$(version.major).$(version.minor)-with-extensions-src.tgz", "d6eef33d8b9e282e20f9b25b6b6fb2757b9b6900e397ca621d56da86d9976541")
]

# Bash recipe for building across all platforms
script = raw"""
cd $WORKSPACE/srcdir
cd mafft-*/

# The mafft shell script is generated from mafft.tmpl. Change the
# absolute path hardcoded at build-time to a relative path generated
# at run-time.
sed -i -e 's|prefix=_LIBDIR|prefix="$(dirname "$(realpath "$0")")/../libexec/mafft"|' core/mafft.tmpl

cd core
make CC=${CC} CXX=${CXX} PREFIX=${prefix}
cd ..

cd extensions
make CC=${CC} CXX=${CXX} PREFIX=${prefix}
cd ..

cd core
make CC=${CC} CXX=${CXX} PREFIX=${prefix} install
cd ..

cd extensions
make CC=${CC} CXX=${CXX} PREFIX=${prefix} install
cd ..

install_license license
mv extensions/mxscarna_src/README extensions/mxscarna_src/README-mxscarna
mv extensions/mxscarna_src/probconsRNA/README extensions/mxscarna_src/probconsRNA/README-probcons
mv extensions/mxscarna_src/vienna/COPYING extensions/mxscarna_src/vienna/COPYING-ViennaRNA
install_license extensions/mxscarna_src/README-mxscarna
install_license extensions/mxscarna_src/probconsRNA/README-probcons
install_license extensions/mxscarna_src/vienna/COPYING-ViennaRNA
"""

# These are the platforms we will build for by default, unless further
# platforms are passed in on the command line
platforms = supported_platforms(; experimental=true, exclude=Sys.iswindows)
platforms = expand_cxxstring_abis(platforms)

# The products that we will ensure are always built
products = [
    ExecutableProduct("einsi", :einsi),
    ExecutableProduct("fftns", :fftns),
    ExecutableProduct("fftnsi", :fftnsi),
    ExecutableProduct("ginsi", :ginsi),
    ExecutableProduct("linsi", :linsi),
    ExecutableProduct("mafft", :mafft),
    ExecutableProduct("mafft-distance", :mafft_distance),
    ExecutableProduct("mafft-einsi", :mafft_einsi),
    ExecutableProduct("mafft-fftns", :mafft_fftns),
    ExecutableProduct("mafft-fftnsi", :mafft_fftnsi),
    ExecutableProduct("mafft-ginsi", :mafft_ginsi),
    ExecutableProduct("mafft-homologs.rb", :mafft_homologs_rb),
    ExecutableProduct("mafft-linsi", :mafft_linsi),
    ExecutableProduct("mafft-nwns", :mafft_nwns),
    ExecutableProduct("mafft-nwnsi", :mafft_nwnsi),
    ExecutableProduct("mafft-profile", :mafft_profile),
    ExecutableProduct("mafft-qinsi", :mafft_qinsi),
    ExecutableProduct("mafft-sparsecore.rb", :mafft_sparsecore_rb),
    ExecutableProduct("mafft-xinsi", :mafft_xinsi),
    ExecutableProduct("nwns", :nwns),
    ExecutableProduct("nwnsi", :nwnsi),
]

# Dependencies that must be installed before this package can be built
dependencies = Dependency[
]

# Build the tarballs, and possibly a `build.jl` as well.
build_tarballs(ARGS, name, version, sources, script, platforms, products, dependencies; julia_compat="1.6")

max_stars_repo_name: sharanry/Yggdrasil

max_issues_repo_name: sharanry/Yggdrasil

max_forks_repo_name: sharanry/Yggdrasil


text:
function predict(net::Vector{Layer},input)
  T=sum(layer->typeof(layer)==JuLearn.dataLayer,net)
  L=div(length(net),T)
  prediction=Vector{Array}(T)
  Ttrain=size(input,2)
  for t=0:Ttrain-1 net[1+t*L].rₒ[:,1]=input[:,1+t] end
  forward!(net)
  for t=0:Ttrain-1 prediction[t+1]=net[t*L+L-1].rₒ[:,1] end

  for t=Ttrain:T-1
  net[1+t*L].rₒ[:,1]=prediction[t]
  forward!(net)
  prediction[t+1]=net[t*L+L-1].rₒ[:,1]
  end

  return prediction
end


text:
__precompile__()

module NearestNeighbors

using Distances
import Distances: Metric, result_type, eval_reduce, eval_end, eval_op, eval_start, evaluate

using StaticArrays

import Base.show
import Compat.view

export BruteTree, KDTree, BallTree, DataFreeTree
export knn, inrange # TODOs? , allpairs, distmat, npairs
export injectdata

export Euclidean,
       Cityblock,
       Minkowski,
       Chebyshev,
       Hamming

# Change this to enable debugging
const DEBUG = false

abstract NNTree{V <: AbstractVector, P <: Metric}

typealias MinkowskiMetric Union{Euclidean, Chebyshev, Cityblock, Minkowski}

function check_input{V1, V2 <: AbstractVector}(::NNTree{V1}, ::Vector{V2})
    if length(V1) != length(V2)
        throw(ArgumentError(
            "dimension of input points:$(length(V2)) and tree data:$(length(V1)) must agree"))
    end
end

function check_input{V1, V2 <: Number}(::NNTree{V1}, point::Vector{V2})
    if length(V1) != length(point)
        throw(ArgumentError(
            "dimension of input points:$(length(point)) and tree data:$(length(V1)) must agree"))
    end
end

get_T{T <: AbstractFloat}(::Type{T}) = T
get_T{T}(::T) = Float64

include("debugging.jl")
include("evaluation.jl")
include("tree_data.jl")
include("datafreetree.jl")
include("knn.jl")
include("inrange.jl")
include("hyperspheres.jl")
include("hyperrectangles.jl")
include("utilities.jl")
include("brute_tree.jl")
include("kd_tree.jl")
include("ball_tree.jl")
include("tree_ops.jl")

end # module


text:
using pwr, Plots

## Exercise 8.1 P. 357 from Cohen (1988)
#pwr.anova.test(f=0.28,k=4,n=20,sig.level=0.05)
t = pwr.AnovaTest(f=0.28,k=4,n=20,power=0.0)
plot(t)

## Exercise 6.1 p. 198 from Cohen (1988)
#pwr.2p.test(h=0.3,n=80,sig.level=0.05,alternative="greater")
t2 = pwr.TwoPTest(h=0.3,n=80,power=0.0,alternative="greater")
plot(t2)

## Exercise 7.3 p. 251
#pwr.chisq.test(w=0.346,df=(2-1)*(3-1),N=140,sig.level=0.01)
t3 = pwr.ChisqTest(w=0.346,df=(2-1)*(3-1),N=140,alpha=0.01,power=0.0)
plot(t3)

## Exercise 6.5 p. 203 from Cohen (1988)
#pwr.p.test(h=0.2,n=60,sig.level=0.05,alternative="two.sided")
t4 = pwr.PTest(h=0.2,n=60,alpha=0.05,alternative="two",power=0.0)
plot(t4,backend=plotly)

## medium effect size for the correlation test
#cohen.ES(test="r", size="medium")
cohenES(test="r", size="medium")

## sample size for a medium size effect in the two-sided correlation test
## using the conventional power of 0.80
#pwr.r.test(r=cohen.ES(test="r",size="medium")$effect.size,power=0.80, sig.level=0.05, alternative="two.sided")
t5 = pwr.RTest(n=0,r=cohenES(test="r",size="medium"),power=0.80, alpha=0.05, alternative="two.sided")
plot(t5)

## Exercise 6.5 p. 203 from Cohen
#h<-ES.h(0.5,0.4)
#h
h = ESh(0.5,0.4)

#pwr.p.test(h=h,n=60,sig.level=0.05,alternative="two.sided")
t6 = pwr.PTest(h=h,n=60,alpha=0.05,power=0.0,alternative="two")
plot(t6)

## Exercise 7.1 p. 249 from Cohen
#P0<-rep(1/4,4)
P0 = fill(1/4,4)

#P1<-c(0.375,rep((1-0.375)/3,3))
P1 = vcat(0.375,fill((1-0.375)/3,3))

#ES.w1(P0,P1)
w = ESw1(P0,P1)

#pwr.chisq.test(w=ES.w1(P0,P1),N=100,df=(4-1))
t7 = pwr.ChisqTest(w=ESw1(P0,P1),N=100,df=(4-1),power=0.0)
plot(t7)

# ES.w2
#prob<-matrix(c(0.225,0.125,0.125,0.125,0.16,0.16,0.04,0.04),nrow=2,byrow=TRUE)
prob = reshape([0.225,0.125,0.125,0.125,0.16,0.16,0.04,0.04],4,2)'

#ES.w2(prob)
ESw2(prob)

#pwr.chisq.test(w=ES.w2(prob),df=(2-1)*(4-1),N=200)
t8 = pwr.ChisqTest(w=ESw2(prob),df=(2-1)*(4-1),N=200,power=0.0)
plot(t8)

## Exercise 6.1 p. 198 from Cohen (1988)
#pwr.2p.test(h=0.3,n=80,sig.level=0.05,alternative="greater")
t9 = pwr.TwoPTest(h=0.3,n=80,alpha=0.05,alternative="greater",power=0.0)
plot(t9)

## Exercise 6.3 P. 200 from Cohen (1988)
#pwr.2p2n.test(h=0.30,n1=80,n2=245,sig.level=0.05,alternative="greater")
t10 = pwr.TwoP2nTest(h=0.30,n1=80,n2=245,alpha=0.05,power=0.0,alternative="greater")
plot(t10)

## Exercise 6.7 p. 207 from Cohen (1988)
#pwr.2p2n.test(h=0.20,n1=1600,power=0.9,sig.level=0.01,alternative="two.sided")
t11 =pwr.TwoP2nTest(h=0.20,n1=1600,power=0.9,alpha=0.01,alternative="two.sided")
plot(t11)

## Exercise 8.1 P. 357 from Cohen (1988)
#pwr.anova.test(f=0.28,k=4,n=20,sig.level=0.05)
t12 = pwr.AnovaTest(f=0.28,k=4,n=20,alpha=0.05)
plot(t12)

## Exercise 8.10 p. 391
#pwr.anova.test(f=0.28,k=4,power=0.80,sig.level=0.05)
t13 = pwr.AnovaTest(f=0.28,k=4,power=0.80,alpha=0.05)
plot(t13)

## Exercise 7.1 P. 249 from Cohen (1988)
#pwr.chisq.test(w=0.289,df=(4-1),N=100,sig.level=0.05)
t14 = pwr.ChisqTest(w=0.289,df=(4-1),N=100,alpha=0.05)
plot(t14)

## Exercise 7.3 p. 251
#pwr.chisq.test(w=0.346,df=(2-1)*(3-1),N=140,sig.level=0.01)
t15 = pwr.ChisqTest(w=0.346,df=(2-1)*(3-1),N=140,alpha=0.01)
plot(t15)

## Exercise 7.8 p. 270
#pwr.chisq.test(w=0.1,df=(5-1)*(6-1),power=0.80,sig.level=0.05)
t16 = pwr.ChisqTest(w=0.1,df=(5-1)*(6-1),power=0.80,alpha=0.05)
plot(t16)

## Exercise 9.1 P. 424 from Cohen (1988)
#pwr.f2.test(u=5,v=89,f2=0.1/(1-0.1),sig.level=0.05)
t17 = pwr.F2Test(u=5,v=89,f2=0.1/(1-0.1),alpha=0.05)

## Power at mu=105 for H0:mu=100 vs. H1:mu>100 (sigma=15) 20 obs. (alpha=0.05)
# sigma<-15
σ = 15
# c<-100
c = 100
# mu<-105
μ = 105
# d<-(mu-c)/sigma
d = (μ - c)/σ
# pwr.norm.test(d=d,n=20,sig.level=0.05,alternative="greater")
pwr.NormTest(d=d,n=20,alpha=0.05,alternative="greater")

## Sample size of the test for power=0.80
#pwr.norm.test(d=d,power=0.8,sig.level=0.05,alternative="greater")
pwr.NormTest(d=d,power=0.8,alpha=0.05,alternative="greater")

## Power function of the same test
#mu<-seq(95,125,l=100)
μ = collect(linspace(95,125,100))

#d<-(mu-c)/sigma
d = (μ - c)/σ

# plot(d,pwr.norm.test(d=d,n=20,sig.level=0.05,alternative="greater")$power,
# type="l",ylim=c(0,1))
# abline(h=0.05)
# abline(h=0.80)
power = [powerNormTest(d=x,n=20,alternative="greater") for x in d]
plot(d,power,ylim=[0,1],legend=false)
hline!([0.05,0.8])

## Power function for the two-sided alternative
# plot(d,pwr.norm.test(d=d,n=20,sig.level=0.05,alternative="two.sided")$power,
#       type="l",ylim=c(0,1))
# abline(h=0.05)
# abline(h=0.80)
plot(d,[powerNormTest(d=x,n=20,alpha=0.05,alternative="two.sided") for x in d],ylim=[0,1],legend=false)
hline!([0.05,0.8])

## Exercise 6.5 p. 203 from Cohen
# h<-ES.h(0.5,0.4)
h = ESh(0.5,0.4)
# h
# pwr.p.test(h=h,n=60,sig.level=0.05,alternative="two.sided")
pwr.PTest(h=h,n=60,alpha=0.05,alternative="two.sided")

## Exercise 6.8 p. 208
#pwr.p.test(h=0.2,power=0.95,sig.level=0.05,alternative="two.sided")
pwr.PTest(h=0.2,power=0.95,alpha=0.05,alternative="two.sided")

## Exercise 3.1 p. 96 from Cohen (1988)
#pwr.r.test(r=0.3,n=50,sig.level=0.05,alternative="two.sided")
pwr.RTest(r=0.3,n=50,alpha=0.05,alternative="two.sided")

#pwr.r.test(r=0.3,n=50,sig.level=0.05,alternative="greater")
pwr.RTest(r=0.3,n=50,alpha=0.05,alternative="greater")

## Exercise 3.4 p. 208
#pwr.r.test(r=0.3,power=0.80,sig.level=0.05,alternative="two.sided")
pwr.RTest(r=0.3,power=0.80,alpha=0.05,alternative="two.sided")

#pwr.r.test(r=0.5,power=0.80,sig.level=0.05,alternative="two.sided")
pwr.RTest(r=0.5,power=0.80,alpha=0.05,alternative="two.sided")

#pwr.r.test(r=0.1,power=0.80,sig.level=0.05,alternative="two.sided")
pwr.RTest(r=0.1,power=0.80,alpha=0.05,alternative="two.sided")

## One sample (power)
## Exercise 2.5 p. 47 from Cohen (1988)
#pwr.t.test(d=0.2,n=60,sig.level=0.10,type="one.sample",alternative="two.sided")
pwr.TTest(d=0.2,n=60,alpha=0.10,power=0.,sampletype="one.sample",alternative="two.sided")

## Paired samples (power)
## Exercise p. 50 from Cohen (1988)
#d<-8/(16*sqrt(2*(1-0.6)))
d = 8/(16*sqrt(2*(1-0.6)))
#pwr.t.test(d=d,n=40,sig.level=0.05,type="paired",alternative="two.sided")
pwr.TTest(d=d,n=40,alpha=0.05,sampletype="paired",power=0.,alternative="two.sided")

## Two independent samples (power)
## Exercise 2.1 p. 40 from Cohen (1988)
#d<-2/2.8
d = 2/2.8
#pwr.t.test(d=d,n=30,sig.level=0.05,type="two.sample",alternative="two.sided")
pwr.TTest(d=d,n=30,alpha=0.05,sampletype="two.sample",power=0.,alternative="two.sided")

## Two independent samples (sample size)
## Exercise 2.10 p. 59
#pwr.t.test(d=0.3,power=0.75,sig.level=0.05,type="two.sample",alternative="greater")
pwr.TTest(d=0.3,power=0.75,alpha=0.05,sampletype="two.sample",alternative="greater")

## Exercise 2.3 p. 437 from Cohen (1988)
#pwr.t2n.test(d=0.6,n1=90,n2=60,alternative="greater")
pwr.T2nTest(d=0.6,n1=90,n2=60,alternative="greater")

max_stars_repo_name: mwsohn/pwr

max_issues_repo_name: mwsohn/pwr

max_forks_repo_name: mwsohn/pwr


text:
"""
    _add_enemies!(ata_model::AbstractModel)

# Description

Add enemy sets to the `ATA.AbstractModel` as specified in the `settings_file` loaded by [`start_ata`](#ATA.start_ata) function.
It requires the [`start_ata`](#ATA.start_ata) step.

# Arguments

- **`ata_model::AbstractModel`** : Required. The model built with `start_ata()` and with settings loaded by [`start_ata`](#ATA.start_ata) function.
"""
function _add_enemies!(ata_model::AbstractModel)
    bank = ata_model.settings.bank
    n_items = ata_model.settings.n_items
    if size(ata_model.constraints[1].constr_b, 1) > 0
        A = Vector{Matrix{Float64}}(undef, ata_model.settings.T)
        b = Vector{Vector{Float64}}(undef, ata_model.settings.T)
        for t = 1:ata_model.settings.T
            A[t] = ata_model.constraints[t].constr_A
            b[t] = ata_model.constraints[t].constr_b
        end

    else
        A = Vector{Matrix{Float64}}(undef, ata_model.settings.T)
        b = Vector{Vector{Float64}}(undef, ata_model.settings.T)
        for t = 1:ata_model.settings.T
            A[t] = Matrix{Float64}(undef, 0, ata_model.settings.n_items)
            b[t] = Vector{Float64}(undef, 0)
        end
    end

    if size(ata_model.settings.es.var, 1) > 0
        enemy_sets_var = ata_model.settings.es.var
        EnemySets = unique(
            vcat(
                [
                    unique(skipmissing(bank[!, (enemy_sets_var[isv])])) for
                    isv = 1:(size(enemy_sets_var, 1))
                ]...,
            ),
        )
        DelimitedFiles.writedlm("EnemySets.csv", EnemySets)
        nes = size(EnemySets, 1)
        sets = Vector{Vector{Int64}}(undef, nes)
        for es = 1:nes
            sets[es] = Int64[]
            for t = 1:ata_model.settings.T
                A[t] = vcat(A[t], zeros(ata_model.settings.n_items)')
            end
            for i = 1:n_items
                if any(
                    skipmissing(
                        vcat(
                            [
                                (
                                    (bank[!, (enemy_sets_var[isv])][i] == EnemySets[es]) ==
                                    true
                                ) for isv = 1:(size(enemy_sets_var, 1))
                            ]...,
                        ),
                    ),
                )
                    #esMaata_model.settings.Tg[n_items, es] = 1
                    for t = 1:ata_model.settings.T
                        A[t][end, i] = 1
                    end
                    sets[es] = vcat(sets[es], i)
                end
            end
            for t = 1:ata_model.settings.T
                push!(b[t], 1)
            end
        end
        for t = 1:ata_model.settings.T
            DelimitedFiles.writedlm("opt/A_$t.csv", A[t])
            DelimitedFiles.writedlm("opt/b_$t.csv", b[t])
        end
        #update model
        ata_model.settings.es.names = EnemySets
        ata_model.settings.es.sets = sets
        for t = 1:ata_model.settings.T
            ata_model.constraints[t].constr_b = b[t]
            ata_model.constraints[t].constr_A = A[t]
        end
        JLD2.@save "opt/ata_model.jld2" AbstractModel
        success!(ata_model, string("", nes, " enemy sets added and constrained."))

    end
    return nothing
end


text:
module H5SectionsArrays
using HDF5

const WAIVER_ID = 1
const H5_DATASET_NAME = "img"
const H5_DATASET_ELEMENT_TYPE = UInt8
const DATASET_NDIMS = 3
# const GLOBAL_OFFSET = [16384,16384,16384]

export H5SectionsArray, boundingbox, get_section_filename

# register item of one section / hdf5 file
const Tsecreg = Dict{Symbol, Union{AbstractString, Int}}
# the whole register records filename, xstart, ystart, xdim, ydim
const Tregister = Dict{Int, Tsecreg}

mutable struct H5SectionsArray{T, N} <: AbstractArray{T,N}
    register::Tregister
    function (::Type{H5SectionsArray})( register::Tregister )
        new{UInt8, 3}(register)
    end
end

"""
construct from a register file,
which was the final output of registration pipeline in seunglab
"""
function H5SectionsArray(fregister::AbstractString)
    if isdir( fregister )
        fregister = joinpath(fregister, "registry.txt")
    end
    @assert isfile(fregister)

    f = open(fregister)
    lines = readlines(f)
    close(f)
    register = Tregister()
    # sizehint!(Tregister, length(lines))
    z = 0
    for i in eachindex(lines)
        z += 1
        # initialize the registration of a section image
        d = Tsecreg()
        line = lines[i]
        if length(split(line)) == 7
            registerFile, tmpZero, xoff, yoff, xdim, ydim, tf = split(line)
        elseif length(split(line))==6
            registerFile, xoff, yoff, xdim, ydim, tf = split(line)
        else
            error("unsupported format of register file: $(line)")
        end
        d[:secIDinWaver] = parse(split(split(registerFile,'_')[1], ',')[2])
        d[:waiverID] = parse(split(split(registerFile,'_')[1], ',')[1])
        d[:registerFile] = joinpath(dirname(fregister), registerFile * ".h5")
        d[:xoff] = parse(xoff)
        d[:yoff] = parse(yoff)
        d[:zoff] = z-1
        d[:xdim] = parse(xdim)
        d[:ydim] = parse(ydim)
        d[:zdim] = 1
        register[z] = d
    end
    H5SectionsArray(register)
end

"""
get the section file name of the z coordinate
"""
function get_section_filename( ba::H5SectionsArray, z::Int )
    return ba.register[z][:registerFile]
end

"""
specialized for UInt8 raw image data type
"""
function Base.eltype(A::H5SectionsArray{T,N}) where {T,N}
    return T
    # return H5_DATASET_ELEMENT_TYPE
    # for key in keys(A.register)
    #     registerFile = A.register[key][:registerFile]
    #     if isfile(registerFile) && ishdf5(registerFile)
    #         f = h5open(registerFile)
    #         ret = eltype(f[])
    #         close(f)
    #         return ret
    #     end
    # end
end

"""
number of dimension
"""
function Base.ndims(A::H5SectionsArray)
    return DATASET_NDIMS
    # for key in keys(A.register)
    #     registerFile = A.register[key][:registerFile]
    #     if isfile(registerFile) && ishdf5(registerFile)
    #         f = h5open(registerFile)
    #         ret = ndims(f[H5_DATASET_NAME])
    #         close(f)
    #         # don't forget the z dimension!
    #         return ret+1
    #     end
    # end
end

"""
bounding box of the whole volume
"""
function boundingbox(A::H5SectionsArray; zmin=-Inf, zmax =Inf)
    x1 = Inf;   x2 = -Inf;
    y1 = Inf;   y2 = -Inf;
    z1 = Inf;   z2 = -Inf;
    for d in values( A.register )
        if d[:zoff] < zmin-1 || d[:zoff] > zmax-1
            continue
        end
        @show d[:zoff]
        x1 = min(x1, d[:xoff] + 1)
        y1 = min(y1, d[:yoff] + 1)
        z1 = min(z1, d[:zoff] + 1)
        x2 = max(x2, d[:xoff] + d[:xdim])
        y2 = max(y2, d[:xoff] + d[:ydim])
        z2 = max(z2, d[:zoff] + d[:zdim])
    end
    (Int64(x1):Int64(x2), Int64(y1):Int64(y2), Int64(z1):Int64(z2))
end

"""
compute size from bounding box
"""
function Base.size(A::H5SectionsArray)
    bb = boundingbox(A)
    map(length, bb)
end

function Base.size(A::H5SectionsArray, i::Int)
    size(A)[i]
end

function Base.show(A::H5SectionsArray)
    println("type: $(typeof(A))")
    println("bounding box: $(boundingbox(A))")
    println("the data is in disk, can not ")
end

"""
read out part of the image from HDF5 file
"""
function read_subimage!(buf,
                        registerFile::AbstractString,
                        sizeX::Integer,
                        sizeY::Integer,
                        xidxSection::Union{Int, UnitRange},
                        yidxSection::Union{Int, UnitRange},
                        zidxSection::Integer)
    @assert ishdf5(registerFile)

    # the explicit coordinate range
    x1 = max(1, first(xidxSection));   x2 = min(sizeX, last(xidxSection));
    y1 = max(1, first(yidxSection));   y2 = min(sizeY, last(yidxSection));
    if x1>x2 || y1>y2
        warn("no overlaping region in this section: $(registerFile)")
    else
        # index in buffer
        bufx1 = x1 - first(xidxSection) + 1;
        bufx2 = x2 - first(xidxSection) + 1;
        bufy1 = y1 - first(yidxSection) + 1;
        bufy2 = y2 - first(yidxSection) + 1;
        buf[bufx1:bufx2, bufy1:bufy2, zidxSection] = h5read(registerFile, H5_DATASET_NAME, (x1:x2, y1:y2))
    end
end

function Base.getindex(A::H5SectionsArray, r::CartesianRange)
    getindex(A, cartesian_range2unitrange(r)...)
end

"""
extract chunk from a bigarray
only works for 3D now.
"""
function Base.getindex(A::H5SectionsArray, idxes::Union{UnitRange, Int}...)
    @assert length(idxes) == 3
    # allocate memory
    sx = length(idxes[1])
    sy = length(idxes[2])
    sz = length(idxes[3])
    # create buffer
    buf = zeros(H5_DATASET_ELEMENT_TYPE, (sx,sy,sz))
    @show idxes
    @sync begin
        for z in idxes[3]
            @async begin
                if haskey(A.register, z)
                    registerFile = A.register[z][:registerFile]
                    xidxSection = idxes[1] - A.register[z][:xoff]
                    yidxSection = idxes[2] - A.register[z][:yoff]
                    zidxSection = z  - first(idxes[3]) + 1
                    read_subimage!( buf,
                                    registerFile,
                                    A.register[z][:xdim],
                                    A.register[z][:ydim],
                                    xidxSection, yidxSection, zidxSection)
                else
                    warn("section file not exist: $z")
                end
            end
        end
    end
    buf
end

function Base.CartesianRange( ba::H5SectionsArray, z::Int )
    d = ba.register[z]
    start = CartesianIndex( d[:xoff]+1, d[:yoff]+1, z )
    stop  = CartesianIndex( d[:xoff]+d[:xdim], d[:yoff]+d[:ydim], z)
    return CartesianRange( start, stop )
end

end # end of module


text:
using KernelAbstractions
using CUDAKernels

# only expect performance on the gpu

@kernel function random_phase_kernel!(field, @Const(A), @Const(𝓀ˣ), @Const(𝓀ʸ), @Const(x), @Const(y), @Const(φ), Nx, Ny)
    i, j = @index(Global, NTuple)

    tmp_sum = zero(eltype(field)) # create temporary array
    xx = x[i]
    yy = y[j]
    for ii in 1:Nx, jj in 1:Ny
        tmp_sum += A[ii, jj] * cos(𝓀ˣ[ii] * xx + 𝓀ʸ[jj] * yy + φ[ii, jj])
    end

    field[i, j] = tmp_sum
end

function stream_function!(field, A, 𝓀ˣ, 𝓀ʸ, x, y, φ; comp_stream = Event(CUDADevice()))
    kernel! = random_phase_kernel!(CUDADevice(), 256)
    Nx = length(𝓀ˣ)
    Ny = length(𝓀ʸ)
    event = kernel!(field, A, 𝓀ˣ, 𝓀ʸ, x, y, φ, Nx, Ny, ndrange = size(field), dependencies = (comp_stream,))
    # event = kernel!(field, A, 𝓀ˣ, 𝓀ʸ, x, y, φ, Nx, Ny, ndrange = size(field))
    # wait(event) here on event = kernel! causes the gpu to hang, need to wait outside
    return event
end

# Expensive on GPU but better on CPU
function random_phase!(field, A, 𝓀ˣ, 𝓀ʸ, x, y, φ)
    field .= 0.0
    for i in eachindex(𝓀ˣ), j in eachindex(𝓀ʸ)
        @. field += A[i, j] * cos(𝓀ˣ[i] * x + 𝓀ʸ[j] * y + φ[i, j])
    end
end

function θ_rhs!(θ̇, θ, params)
    #(; ψ, A, 𝓀ˣ, 𝓀ʸ, x, y, φ, u, v, ∂ˣθ, ∂ʸθ, s, P, P⁻¹, filter) = params
    ψ, A, 𝓀ˣ, 𝓀ʸ, x, y, φ, u, v, ∂ˣθ, ∂ʸθ, s, P, P⁻¹, filter = params
    P * ψ # in place fft
    P * θ # in place fft
    # ∇ᵖψ
    @. u = filter * -1.0 * (∂y * ψ)
    @. v = filter * (∂x * ψ)
    # ∇θ
    @. ∂ˣθ = filter * ∂x * θ
    @. ∂ʸθ = filter * ∂y * θ
    @. κΔθ = κ * Δ * θ
    # go back to real space
    P⁻¹ * ψ
    P⁻¹ * θ
    P⁻¹ * u
    P⁻¹ * v
    P⁻¹ * ∂ˣθ
    P⁻¹ * ∂ʸθ
    P⁻¹ * κΔθ
    # Assemble RHS
    @. θ̇ = -u * ∂ˣθ - v * ∂ʸθ + κΔθ + s
    return nothing
end

function θ_rhs_new!(θ̇, θ, params)
    #(; ψ, A, 𝓀ˣ, 𝓀ʸ, x, y, φ, u, v, ∂ˣθ, ∂ʸθ, s, P, P⁻¹, filter) = params
    ψ, A, 𝓀ˣ, 𝓀ʸ, x, y, φ, u, v, ∂ˣθ, ∂ʸθ, s, P, P⁻¹, filter = params
    event = stream_function!(ψ, A, 𝓀ˣ, 𝓀ʸ, x, y, φ)
    wait(event)
    P * ψ # in place fft
    P * θ # in place fft
    # ∇ᵖψ
    @. u = filter * -1.0 * (∂y * ψ)
    @. v = filter * (∂x * ψ)
    # ∇θ
    @. ∂ˣθ = filter * ∂x * θ
    @. ∂ʸθ = filter * ∂y * θ
    @. κΔθ = κ * Δ * θ
    # go back to real space
    P⁻¹ * ψ
    P⁻¹ * θ
    P⁻¹ * u
    P⁻¹ * v
    P⁻¹ * ∂ˣθ
    P⁻¹ * ∂ʸθ
    P⁻¹ * κΔθ
    # Assemble RHS
    @. θ̇ = -u * ∂ˣθ - v * ∂ʸθ + κΔθ + s
    return nothing
end

function θ_rhs_zeroth!(θ̇, θ, params)
    #(; ψ, A, 𝓀ˣ, 𝓀ʸ, x, y, φ, u, v, ∂ˣθ, ∂ʸθ, s, P, P⁻¹, filter) = params
    ψ, A, 𝓀ˣ, 𝓀ʸ, x, y, φ, u, v, ∂ˣθ, ∂ʸθ, s, P, P⁻¹, filter = params
    P * ψ # in place fft
    P * θ # in place fft
    # ∇ᵖψ
    @. u = filter * -1.0 * (∂y * ψ)
    @. v = filter * (∂x * ψ)
    # ∇θ
    @. ∂ˣθ = filter * ∂x * θ
    @. ∂ʸθ = filter * ∂y * θ
    @. κΔθ = κ * Δ * θ
    # go back to real space
    P⁻¹ * ψ
    P⁻¹ * θ
    P⁻¹ * u
    P⁻¹ * v
    P⁻¹ * ∂ˣθ
    P⁻¹ * ∂ʸθ
    P⁻¹ * κΔθ
    # Assemble RHS
    @. θ̇ = -u * ∂ˣθ - v * ∂ʸθ + κΔθ - u
    return nothing
end

function θ_rhs_convergent!(θ̇, θ, params)
    ψ, A, 𝓀ˣ, 𝓀ʸ, x, y, φ, u, v, ∂ˣuθ, ∂ʸvθ, s, P, P⁻¹, filter, uᶜ, vᶜ = params
    P * ψ # in place fft
    # ∇ᵖψ
    @. u = filter * -1.0 * (∂y * ψ)
    @. v = filter * (∂x * ψ)
    P⁻¹ * u
    P⁻¹ * v
    # ∇⋅(u⃗ θ) + κΔ
    @. ∂ˣuθ = (u + uᶜ) * θ
    @. ∂ʸvθ = (v + vᶜ) * θ
    P * ∂ˣuθ # in place fft
    P * ∂ʸvθ # in place fft
    P * θ  # in place fft
    @. ∂ˣuθ = ∂x * ∂ˣuθ
    @. ∂ʸvθ = ∂y * ∂ʸvθ
    @. κΔθ = κ * Δ * θ
    # go back to real space
    P⁻¹ * ∂ˣuθ
    P⁻¹ * ∂ʸvθ
    P⁻¹ * θ
    P⁻¹ * κΔθ
    P⁻¹ * ψ # technically not necessary
    # Assemble RHS
    @. θ̇ = -∂ˣuθ - ∂ʸvθ + κΔθ
    return nothing
end

function φ_rhs!(φ̇, φ, rng)
    rand!(rng, φ̇) # can use randn(rng, φ̇); @. φ̇ *= sqrt(1/12)
    φ̇ .-= 0.5
    return nothing
end

function φ_rhs_normal!(φ̇, φ, rng)
    randn!(rng, φ̇)
    φ̇ .*= sqrt(1 / 12)
    return nothing
end

max_stars_repo_name: sandreza/FourierCore

max_issues_repo_name: sandreza/FourierCore

max_forks_repo_name: sandreza/FourierCore


text:
# 3-D variant of vtk_grid with 4-D array xyz.
function vtk_grid{T<:AbstractFloat}(
        filename_noext::AbstractString, xyz::AbstractArray{T,4};
        compress::Bool=true, append::Bool=true, extent=nothing)

    Ncomp, Ni, Nj, Nk = size(xyz)
    @assert Ncomp == 3  # three components (x, y, z)

    xvtk = XMLDocument()

    Npts = Ni*Nj*Nk
    Ncls = num_cells_structured(Ni, Nj, Nk)

    ext = extent_attribute(Ni, Nj, Nk, extent)

    vtk = DatasetFile(xvtk, filename_noext*".vts", "StructuredGrid",
                      Npts, Ncls, compress, append)

    # VTKFile node
    xroot = vtk_xml_write_header(vtk)

    # StructuredGrid node
    xGrid = new_child(xroot, vtk.gridType_str)
    set_attribute(xGrid, "WholeExtent", ext)

    # Piece node
    xPiece = new_child(xGrid, "Piece")
    set_attribute(xPiece, "Extent", ext)

    # Points node
    xPoints = new_child(xPiece, "Points")

    # DataArray node
    data_to_xml(vtk, xPoints, xyz, "Points", 3)

    return vtk::DatasetFile
end

# 2-D variant of vtk_grid with 3-D array xy
function vtk_grid{T<:AbstractFloat}(
        filename_noext::AbstractString, xy::AbstractArray{T,3};
        compress::Bool=true, append::Bool=true, extent=nothing)

    dim, Ni, Nj = size(xy); Nk = 1
    dim == 2 || throw(ArgumentError("The size of xy needs to be (2,Ni,Nj)"))
    xyz = zeros(T,3,Ni,Nj,Nk)
    xyz[1:2,:,:,1] = xy

    return vtk_grid(filename_noext, xyz,
                    compress=compress, append=append, extent=extent)
end

# 3-D variant of vtk_grid with 3-D arrays x, y, z.
function vtk_grid{T<:AbstractFloat}(
        filename_noext::AbstractString,
        x::AbstractArray{T,3}, y::AbstractArray{T,3}, z::AbstractArray{T,3};
        compress::Bool=true, append::Bool=true, extent=nothing)
    @assert size(x) == size(y) == size(z)
    Ni, Nj, Nk = size(x)
    xyz = Array(T, 3, Ni, Nj, Nk)
    for k = 1:Nk, j = 1:Nj, i = 1:Ni
        xyz[1, i, j, k] = x[i, j, k]
        xyz[2, i, j, k] = y[i, j, k]
        xyz[3, i, j, k] = z[i, j, k]
    end
    return vtk_grid(filename_noext, xyz, compress=compress,
                    append=append, extent=extent)::DatasetFile
end

# 2-D variant of vtk_grid with 2-D arrays x, y.
function vtk_grid{T<:AbstractFloat}(
        filename_noext::AbstractString,
        x::AbstractArray{T,2}, y::AbstractArray{T,2};
        compress::Bool=true, append::Bool=true, extent=nothing)
    @assert size(x) == size(y)
    Ni, Nj = size(x)
    Nk = 1
    xyz = zeros(T, 3, Ni, Nj, Nk)
    for j = 1:Nj, i = 1:Ni
        xyz[1, i, j, 1] = x[i, j]
        xyz[2, i, j, 1] = y[i, j]
    end
    return vtk_grid(filename_noext, xyz,
                    compress=compress, append=append, extent=extent)::DatasetFile
end


text:
function unix2second(ts)
    dt = unix2datetime(ts)
end

function maketimetable(path::String)
    allfields = [:mtime, :user, :iter, :success, :error, :time]
    groupby = GroupBy([:mtime, :user, :success, :time],
                      groupers=[:mtime => unix2second,
                                :user => identity,
                                :success => identity])
    open(path) do f
        for line in eachline(f)
            (mtime, user, iter, success, error, time) = split(line, ',')
            if !bool(success)
                # TODO: parse :error and save error reasons with count
            end
            update!(groupby, [mtime, user, success, time])
        end
    end
    df = todataframe(groupby) # TODO
    # TODO: return one more dataframe for errors?
end


text:
using DataFrames, Gadfly

set_default_plot_size(9inch, 3inch)

shapes = [Shape.square, Shape.diamond, Shape.cross, Shape.xcross, Shape.utriangle, Shape.dtriangle,
    Shape.star1, Shape.star2, Shape.hexagon, Shape.octagon, Shape.hline, Shape.vline]

 D = DataFrame(x=[1:5; 1:4; 1:3], y=[5:-1:1; 4:-1:1; 3:-1:1], g=string.(1:12))

function plotf(coord::Coord.cartesian)
    plot(D, coord,
        layer(x=:x, y=:y, style(default_color=colorant"black", point_size=2px), order=2),
        layer(x=:x, y=:y, shape=:g, color=:g, style(point_size=4mm, point_shapes=shapes)),
    Theme(key_position=:none, plot_padding=[0mm], background_color="white"))
end

coords = [Coord.cartesian(ymin=0, ymax=6), Coord.cartesian(xflip=true, ymin=0, ymax=6), Coord.cartesian(yflip=true, ymin=0, ymax=6)]
plots = plotf.(coords)
hstack(plots...)


text:
@doc raw"""
    LinearScalarAdvectionEquation3D

The linear scalar advection equation
```math
\partial_t u + a_1 \partial_1 u + a_2 \partial_2 u + a_3 \partial_3 u = 0
```
in two space dimensions with constant velocity `a`.
"""
struct LinearScalarAdvectionEquation3D <: AbstractLinearScalarAdvectionEquation{3, 1}
  sources::String
  advectionvelocity::SVector{3, Float64}
end

function LinearScalarAdvectionEquation3D()
  sources = parameter("sources", "none")
  a = convert(SVector{3,Float64}, parameter("advectionvelocity"))
  LinearScalarAdvectionEquation3D(sources, a)
end

get_name(::LinearScalarAdvectionEquation3D) = "LinearScalarAdvectionEquation3D"
varnames_cons(::LinearScalarAdvectionEquation3D) = SVector("scalar")
varnames_prim(::LinearScalarAdvectionEquation3D) = SVector("scalar")

# Set initial conditions at physical location `x` for time `t`
function initial_conditions_gauss(x, t, equation::LinearScalarAdvectionEquation3D)
  # Store translated coordinate for easy use of exact solution
  x_trans = x - equation.advectionvelocity * t

  return @SVector [exp(-(x_trans[1]^2 + x_trans[2]^2 + x_trans[3]^2))]
end

function initial_conditions_convergence_test(x, t, equation::LinearScalarAdvectionEquation3D)
  # Store translated coordinate for easy use of exact solution
  x_trans = x - equation.advectionvelocity * t

  c = 1.0
  A = 0.5
  L = 2
  f = 1/L
  omega = 2 * pi * f
  scalar = c + A * sin(omega * sum(x_trans))
  return @SVector [scalar]
end

function initial_conditions_sin_periodic(x, t, equation::LinearScalarAdvectionEquation3D)
  # Store translated coordinate for easy use of exact solution
  x_trans = x - equation.advectionvelocity * t

  scalar = sin(2 * pi * x_trans[1]) * sin(2 * pi * x_trans[2]) * sin(2 * pi * x_trans[3])
  return @SVector [scalar]
end

function initial_conditions_constant(x, t, equation::LinearScalarAdvectionEquation3D)
  # Store translated coordinate for easy use of exact solution
  x_trans = x - equation.advectionvelocity * t

  return @SVector [2.0]
end

function initial_conditions_linear_z(x, t, equation::LinearScalarAdvectionEquation3D)
  # Store translated coordinate for easy use of exact solution
  x_trans = x - equation.advectionvelocity * t

  return @SVector [x_trans[3]]
end

# Apply boundary conditions
function boundary_conditions_linear_z(u_inner, orientation, direction, x, t,
                                              surface_flux_function,
                                              equation::LinearScalarAdvectionEquation3D)
  u_boundary = initial_conditions_linear_z(x, t, equation)

  # Calculate boundary flux
  if direction in (2, 4, 6) # u_inner is "left" of boundary, u_boundary is "right" of boundary
    flux = surface_flux_function(u_inner, u_boundary, orientation, equation)
  else # u_boundary is "left" of boundary, u_inner is "right" of boundary
    flux = surface_flux_function(u_boundary, u_inner, orientation, equation)
  end

  return flux
end

# Pre-defined source terms should be implemented as
# function source_terms_WHATEVER(ut, u, x, element_id, t, n_nodes, equation::LinearScalarAdvectionEquation3D)

# Calculate 1D flux in for a single point
@inline function calcflux(u, orientation, equation::LinearScalarAdvectionEquation3D)
  a = equation.advectionvelocity[orientation]
  return a * u
end

function flux_lax_friedrichs(u_ll, u_rr, orientation, equation::LinearScalarAdvectionEquation3D)
  a = equation.advectionvelocity[orientation]
  return 0.5 * ( a * (u_ll + u_rr) - abs(a) * (u_rr - u_ll) )
end

# Determine maximum stable time step based on polynomial degree and CFL number
function calc_max_dt(u, element_id, invjacobian, cfl,
                     equation::LinearScalarAdvectionEquation3D, dg)
  λ_max = maximum(abs, equation.advectionvelocity)
  return cfl * 2 / (nnodes(dg) * invjacobian * λ_max)
end

# Convert conservative variables to primitive
@inline cons2prim(u, equation::LinearScalarAdvectionEquation3D) = u

# Convert conservative variables to entropy variables
@inline cons2entropy(u, equation::LinearScalarAdvectionEquation3D) = u

# Calculate entropy for a conservative state `cons`
@inline entropy(u::Real, ::LinearScalarAdvectionEquation3D) = 0.5 * u^2
@inline entropy(u, equation::LinearScalarAdvectionEquation3D) = entropy(u[1], equation)

# Calculate total energy for a conservative state `cons`
@inline energy_total(u::Real, ::LinearScalarAdvectionEquation3D) = 0.5 * u^2
@inline energy_total(u, equation::LinearScalarAdvectionEquation3D) = energy_total(u[1], equation)


text:
import StatsBase

function quad_knapsack(solver, seed)
    rng = Random.MersenneTwister(seed)
    if Base.libblas_name == "libmkl_rt"
        model = Model()
    else
        model = Model(solver=solver)
    end

    # Instance size
    n = 20
    # k-item capacity
    # k = Int(n / 10)
    k = 1
    # Frequency of nonzero weights
    delta = 0.5

    # Build weights and capacity
    a = zeros(n)
    for i in 1:n
        a[i] = Random.rand(rng, 1:50)
    end
    a = ones(n)
    b = Random.rand(rng, 100:sum(a)+100)

    # Profits
    C = zeros((n, n))
    for i in 1:n
        for j in 1:n
            if StatsBase.sample(rng, [1, 0], StatsBase.Weights([delta, 1.0 - delta])) == 1
                c_ = - Random.rand(rng, 1:100)
                C[i, j] = c_
                C[j, i] = c_
            end
        end
    end

    # Decision variable
    if Base.libblas_name == "libmkl_rt"
        @variable(model, X[1:n+1, 1:n+1], PSD)
    else
        @variable(model, X[1:n+1, 1:n+1], SDP)
    end
    @objective(model, Min, sum(C[i, j] * X[i+1, j+1] for i in 1:n, j in 1:n))
    # Capacity constraint
    # @constraint(model, cap, sum(a[i] * X[i+1, i+1] for i in 1:n) <= b)
    # k-item constraint
    @constraint(model, k_knap, sum(X[i+1, i+1] for i in 1:n) == k)

    # @constraint(model, bla, X[1, 1] == 1)

    # @constraint(model, lb[i in 1:n+1, j in 1:n+1], X[i, j] >= 0.0)
    # @constraint(model, ub[i in 1:n+1, j in 1:n+1], X[i, j] <= 1.0)

    if Base.libblas_name == "libmkl_rt"
        JuMP.attach(model, solver)
    end

    @time teste = JuMP.solve(model)

    if Base.libblas_name == "libmkl_rt"
        XX = getvalue2.(X)
    else
        XX = getvalue.(X)
    end
    rank = length([eig for eig in LinearAlgebra.eigen(XX).values if eig > 1e-10])
    @show rank
    @show diag(XX)
end

getvalue2(var::JuMP.Variable) = (m=var.m;m.solverinstance.primal[m.solverinstance.varmap[m.variabletosolvervariable[var.instanceindex]]])


text:
module ComplexDiff

include("unicomplex.jl")
include("bicomplex.jl")
include("tricomplex.jl")
include("quadcomplex.jl")
include("pentacomplex.jl")

include("jacobian.jl")

export image2
export bicomplex, tricomplex, quadcomplex, pentacomplex
export derivative, biderivative, triderivative, quadderivative, pentaderivative
export jacobian

end


text:
using Desktop
using Test

@testset "Desktop.jl" begin
    @test hasdesktop() in (true, false)
end