text:
@doc doc"""
    Metric

Abstract type for the pseudo-Riemannian metric tensor $g$, a family of smoothly
varying inner products on the tangent space. See [`inner`](@ref).
"""
abstract type Metric end

@doc doc"""
    RiemannianMetric <: Metric

Abstract type for Riemannian metrics, a family of positive definite inner
products. The positive definite property means that for $v \in T_x M$, the
inner product $g(v, v) > 0$ whenever $v$ is not the zero vector.
"""
abstract type RiemannianMetric <: Metric end

@doc doc"""
    LorentzMetric <: Metric

Abstract type for Lorentz metrics, which have a single time dimension. These
metrics assume the spacelike convention with the time dimension being last,
giving the signature $(++...+-)$.
"""
abstract type LorentzMetric <: Metric end

"""
    MetricManifold{M<:Manifold,G<:Metric} <: Manifold

Equip a [`Manifold`](@ref) explicitly with a [`Metric`](@ref) `G`.

For a Metric Manifold, by default, assumes, that you implement the linear form
from [`local_metric`](@ref) in order to evaluate the exponential map.

If the corresponding [`Metric`](@ref) `G` yields closed form formulae for e.g.
the exponential map and this is implemented directly (without solving the ode),
you can of course still implement that directly.

# Constructor
    MetricManifold(manifold, metric)

    """
struct MetricManifold{M<:Manifold,G<:Metric} <: Manifold
    manifold::M
    metric::G
end

is_decorator_manifold(M::MMT) where {MMT <: MetricManifold}= Val(true)

"""
    is_default_metric(M,G)

indicate, whether the [`Metric`](@ref) `G` is the default metric for
the [`Manifold`](@ref) `M`. This means that any occurence of
[`MetricManifold`](@ref)(M,G) where `typeof(is_default_metric(M,G)) = Val{true}`
falls back to just be called with `M` such that the [`Manifold`](@ref) `M`
implicitly has this metric, for example if this was the first one implemented
or is the one most commonly assumed to be used.
"""
is_default_metric(M::Manifold,G::Metric) = Val(false)
"""
    is_default_metric(MM)

indicate, whether the [`Metric`](@ref) `MM.G` is the default metric for
the [`Manifold`](@ref) `MM.M` within the [`MetricManifold`](@ref) `MM`.
This means that any occurence of
[`MetricManifold`](@ref)`(MM.M,MM.G)` where `typeof(is_default_metric(MM.M,MM.G)) = Val{true}`
falls back to just be called with `MM.MM` such that the [`Manifold`](@ref) `MM.M`
implicitly has the metric `MM.G`, for example if this was the first one
implemented or is the one most commonly assumed to be used.
"""
is_default_metric(M::MMT) where {MMT <: MetricManifold} = is_default_metric(base_manifold(M),metric(M))
# this automatically undecorates
# this should automatically decorate at least for simple cases
convert(T::Type{MetricManifold{MT,GT}},M::MT) where {MT,GT} = _convert_with_default(M,GT,is_default_metric(M,GT()))
_convert_with_default(M::MT,T::Type{<:Metric},::Val{true}) where {MT <: Manifold} = MetricManifold(M,T())
_convert_with_default(M::MT,T::Type{<:Metric},::Val{false}) where {MT <: Manifold} = error("Can not convert $(M) to a MetricManifold{$(MT),$(T)}, since $(T) is not the default metric.")

@doc doc"""
    metric(M::MetricManifold)

Get the metric $g$ of the manifold `M`.
"""
metric(M::MetricManifold) = M.metric

@doc doc"""
    local_metric(M::MetricManifold, x)

Local matrix representation at the point `x` of the metric tensor $g$ on the
manifold `M`, usually written $g_{ij}$. The matrix has the property that
$g(v, w)=v^T [g_{ij}] w = g_{ij} v^i w^j$, where the latter expression uses
Einstein summation convention.
"""
function local_metric(M::MetricManifold, x)
    error("Local metric not implemented on $(typeof(M)) for point $(typeof(x))")
end

@doc doc"""
    inverse_local_metric(M::MetricManifold, x)

Local matrix representation of the inverse metric (cometric) tensor, usually
written $g^{ij}$
"""
inverse_local_metric(M::MetricManifold, x) = inv(local_metric(M, x))

@doc doc"""
    det_local_metric(M::MetricManifold, x)

Determinant of local matrix representation of the metric tensor $g$
"""
det_local_metric(M::MetricManifold, x) = det(local_metric(M, x))

@doc doc"""
    log_local_metric_density(M::MetricManifold, x)

Return the natural logarithm of the metric density $\rho$ of `M` at `x`, which
is given by $\rho=\log \sqrt{|\det [g_{ij}]|}$.
"""
log_local_metric_density(M::MetricManifold, x) = log(abs(det_local_metric(M, x))) / 2

@doc doc"""
    local_metric_jacobian(M::MetricManifold, x)

Get partial derivatives of the local metric of `M` at `x` with respect to the
coordinates of `x`, $\frac{\partial}{\partial x^k} g_{ij} = g_{ij,k}$. The
dimensions of the resulting multi-dimensional array are ordered $(i,j,k)$.
"""
function local_metric_jacobian(M, x)
    error("local_metric_jacobian not implemented on $(typeof(M)) for point $(typeof(x)). For a suitable default, enter `using ForwardDiff`.")
end

@doc doc"""
    christoffel_symbols_first(M::MetricManifold, x)

Compute the Christoffel symbols of the first kind in local coordinates.
The Christoffel symbols are (in Einstein summation convention)

$\Gamma_{ijk} = \frac{1}{2} \left[g_{kj,i} + g_{ik,j} - g_{ij,k}\right],$

where $g_{ij,k}=\frac{\partial}{\partial x^k} g_{ij}$ is the coordinate
derivative of the local representation of the metric tensor. The dimensions of
the resulting multi-dimensional array are ordered $(i,j,k)$.
"""
function christoffel_symbols_first(M::MetricManifold, x)
    ∂g = local_metric_jacobian(M, x)
    n = size(∂g, 1)
    Γ = similar(∂g, Size(n, n, n))
    @einsum Γ[i,j,k] = 1 / 2 * (∂g[k,j,i] + ∂g[i,k,j] - ∂g[i,j,k])
    return Γ
end

@doc doc"""
    christoffel_symbols_second(M::MetricManifold, x)

Compute the Christoffel symbols of the second kind in local coordinates.
The Christoffel symbols are (in Einstein summation convention)

$\Gamma^{l}_{ij} = g^{kl} \Gamma_{ijk},$

where $\Gamma_{ijk}$ are the Christoffel symbols of the first kind, and
$g^{kl}$ is the inverse of the local representation of the metric tensor.
The dimensions of the resulting multi-dimensional array are ordered $(l,i,j)$.
"""
function christoffel_symbols_second(M::MetricManifold, x)
    ginv = inverse_local_metric(M, x)
    Γ₁ = christoffel_symbols_first(M, x)
    Γ₂ = similar(Γ₁)
    @einsum Γ₂[l,i,j] = ginv[k,l] * Γ₁[i,j,k]
    return Γ₂
end

@doc doc"""
    christoffel_symbols_second_jacobian(M::MetricManifold, x)

Get partial derivatives of the Christoffel symbols of the second kind
for manifold `M` at `x` with respect to the coordinates of `x`,
$\frac{\partial}{\partial x^l} \Gamma^{k}_{ij} = \Gamma^{k}_{ij,l}.$
The dimensions of the resulting multi-dimensional array are ordered $(i,j,k,l)$.
"""
function christoffel_symbols_second_jacobian(M, x)
    error("christoffel_symbols_second_jacobian not implemented on $(typeof(M)) for point $(typeof(x)). For a suitable default, enter `using ForwardDiff`.")
end

@doc doc"""
    riemann_tensor(M::MetricManifold, x)

Compute the Riemann tensor $R^l_{ijk}$, also known as the Riemann curvature
tensor, at the point `x`. The dimensions of the resulting multi-dimensional
array are ordered $(l,i,j,k)$.
"""
function riemann_tensor(M::MetricManifold, x)
    n = size(x, 1)
    Γ = christoffel_symbols_second(M, x)
    ∂Γ = christoffel_symbols_second_jacobian(M, x) ./ n
    R = similar(∂Γ, Size(n, n, n, n))
    @einsum R[l,i,j,k] = ∂Γ[l,i,k,j] - ∂Γ[l,i,j,k] + Γ[s,i,k] * Γ[l,s,j] - Γ[s,i,j] * Γ[l,s,k]
    return R
end

"""
    ricci_tensor(M::MetricManifold, x)

Compute the Ricci tensor, also known as the Ricci curvature tensor,
of the manifold `M` at the point `x`.
"""
function ricci_tensor(M::MetricManifold, x)
    R = riemann_tensor(M, x)
    n = size(R, 1)
    Ric = similar(R, Size(n, n))
    @einsum Ric[i,j] = R[l,i,l,j]
    return Ric
end

"""
    ricci_curvature(M::MetricManifold, x)

Compute the Ricci scalar curvature of the manifold `M` at the point `x`.
"""
function ricci_curvature(M::MetricManifold, x)
    ginv = inverse_local_metric(M, x)
    Ric = ricci_tensor(M, x)
    S = sum(ginv .* Ric)
    return S
end

"""
    gaussian_curvature(M::MetricManifold, x)

Compute the Gaussian curvature of the manifold `M` at the point `x`.
"""
gaussian_curvature(M::MetricManifold, x) = ricci_curvature(M, x) / 2

"""
    einstein_tensor(M::MetricManifold, x)

Compute the Einstein tensor of the manifold `M` at the point `x`.
"""
function einstein_tensor(M::MetricManifold, x)
    Ric = ricci_tensor(M, x)
    g = local_metric(M, x)
    ginv = inverse_local_metric(M, x)
    S = sum(ginv .* Ric)
    G = Ric - g .* S / 2
    return G
end

@doc doc"""
    solve_exp_ode(M::MetricManifold, x, v, tspan; solver=AutoVern9(Rodas5()), kwargs...)

Approximate the exponential map on the manifold over the provided timespan
assuming the Levi-Civita connection by solving the ordinary differential
equation

$\frac{d^2}{dt^2} x^k + \Gamma^k_{ij} \frac{d}{dt} x_i \frac{d}{dt} x_j = 0,$

where $\Gamma^k_{ij}$ are the Christoffel symbols of the second kind, and
the Einstein summation convention is assumed. The arguments `tspan` and
`solver` follow the `OrdinaryDiffEq` conventions. `kwargs...` specify keyword
arguments that will be passed to `OrdinaryDiffEq.solve`.

Currently, the numerical integration is only accurate when using a single
coordinate chart that covers the entire manifold. This excludes coordinates
in an embedded space.
"""
function solve_exp_ode(M, x, v, tspan; kwargs...)
    error("solve_exp_ode not implemented on $(typeof(M)) for point $(typeof(x)), vector $(typeof(v)), and timespan $(typeof(tspan)). For a suitable default, enter `using OrdinaryDiffEq`.")
end
exp(M::MMT, x, v, T::AbstractVector{T} where T) where {MMT <: MetricManifold} = exp(M, is_default_metric(M), x, v, T)
exp(M::MMT, ::Val{true}, x, v, T::AbstractVector{T} where T) where {MMT<:MetricManifold} = exp(base_manifold(M), x, v, T)
function exp(M::MMT, ::Val{false}, x, v, T::AbstractVector{T} where T) where {MMT<:MetricManifold}
    sol = solve_exp_ode(M, x, v, extrema(T); dense=false, saveat=T)
    n = length(x)
    return map(i -> sol.u[i][n+1:end], 1:length(T))
end

exp!(M::MMT, y, x, v) where {MMT <: MetricManifold}= exp!(M, is_default_metric(M), y, x, v)
exp!(M::MMT, ::Val{true}, y, x, v) where {MMT<:MetricManifold} = exp!(base_manifold(M),y,x,v)
"""
    exp(MM::MetricManifold, x, v, args...)

Numerically integrate the exponential map assuming the Levi-Civita connection.
See [`solve_exp_ode`](@ref) if the [`Metric`](@ref) `G` of the [`MetricManifold`](@ref){M,G}`
is not marked as the default one.

Currently, the numerical integration is only accurate when using a single
coordinate chart that covers the entire manifold. This excludes coordinates
in an embedded space.
"""
function exp!(M::MMT, ::Val{false}, y, x, v) where {MMT<:MetricManifold}
    tspan = (0.0, 1.0)
    sol = solve_exp_ode(M, x, v, tspan; dense=false, saveat=[1.0])
    n = length(x)
    y .= sol.u[1][n+1:end]
    return y
end

inner(M::MMT, x, v, w) where {MMT <: MetricManifold}= inner(M, is_default_metric(M), x, v, w)
inner(M::MMT, ::Val{false}, x, v, w) where {MMT <: MetricManifold}= dot(v, local_metric(M, x) * w)
inner(M::MMT, ::Val{true}, x, v, w) where {MMT <: MetricManifold} = inner(base_manifold(M), x,v,w)

log!(M::MMT, w, x, y) where {MMT <: MetricManifold} = log!(M, is_default_metric(M), w, x, y)
log!(M::MMT, ::Val{true}, w, x, y) where {MMT <: MetricManifold} = log!(base_manifold(M), w, x, y)
log!(M::MMT, ::Val{false}, w, x, y) where {MMT <: MetricManifold} = error("Logarithmic map not implemented on $(typeof(M)) for points $(typeof(x)) and $(typeof(y)).")

tangent_orthonormal_basis(M::MMT, x, v) where {MMT <: MetricManifold} = tangent_orthonormal_basis(M, is_default_metric(M), x, v)
tangent_orthonormal_basis(M::MMT, ::Val{true}, x, v) where {MMT <: MetricManifold} = tangent_orthonormal_basis(base_manifold(M), x, v)
tangent_orthonormal_basis(M::MMT, ::Val{false}, x, v) where {MMT <: MetricManifold} = error("tangent_orthogonal_basis not implemented on $(typeof(M)) for point $(typeof(x)) and tangent vector $(typeof(v)).")

project_point!(M::MMT, y, x) where {MMT <: MetricManifold} = project_point!(M, is_default_metric(M), y, x)
project_point!(M::MMT, ::Val{true}, y, x) where {MMT <: MetricManifold} = project_point!(base_manifold(M), y, x)
project_point!(M::MMT, ::Val{false}, y, x) where {MMT <: MetricManifold} = error("project_point! not implemented on $(typeof(M)) for point $(typeof(x))")

project_tangent!(M::MMT, w, x, v) where {MMT <: MetricManifold} = project_tangent!(M, is_default_metric(M), w, x, v)
project_tangent!(M::MMT, ::Val{true}, w, x, v) where {MMT <: MetricManifold} = project_tangent!(base_manifold(M), w, x, v)
project_tangent!(M::MMT, ::Val{false}, w, x, v) where {MMT <: MetricManifold} = error("project_tangent! not implemented for a $(typeof(M)) and tangent $(typeof(v)) at point $(typeof(x)).")

vector_transport_to!(M::MMT, vto, x, v, y, m::AbstractVectorTransportMethod) where {MMT <: MetricManifold} = vector_transport_to!(M, is_default_metric(M), vto, x, v, y, m)
vector_transport_to!(M::MMT, ::Val{true}, vto, x, v, y, m)  where {MMT <: MetricManifold}= vector_transport_to!(base_manifold(M), vto, x, v, y, m)
vector_transport_to!(M::MMT, ::Val{false}, vto, x, v, y, m)  where {MMT <: MetricManifold}= error("vector transport from a point of type $(typeof(x)) to a type $(typeof(y)) on a $(typeof(M)) for a vector of type $(v) and the $(typeof(m)) not yet implemented.")

projected_distribution(M::MMT, d, x) where {MMT <: MetricManifold} = projected_distribution(M, is_default_metric(M), d,x)
projected_distribution(M::MMT, ::Val{true}, d, x) where {MMT <: MetricManifold} = projected_distribution(base_manifold(M), d,x)
projected_distribution(M::MMT, ::Val{false}, d, x) where {MMT <: MetricManifold} =error("projected_distribution not implemented for a $(typeof(M)) and with $(typeof(d)) at point $(typeof(x)).")

projected_distribution(M::MMT, d) where {MMT <: MetricManifold} = projected_distribution(M, is_default_metric(M), d)
projected_distribution(M::MMT, ::Val{true}, d) where {MMT <: MetricManifold} = projected_distribution(base_manifold(M), d)
projected_distribution(M::MMT, ::Val{false}, d) where {MMT <: MetricManifold} =error("projected_distribution not implemented for a $(typeof(M)) with $(typeof(d)).")

normal_tvector_distribution(M::MMT, x, σ) where {MMT <: MetricManifold} = normal_tvector_distribution(M, is_default_metric(M), x, σ)
normal_tvector_distribution(M::MMT, ::Val{true}, x, σ) where {MMT <: MetricManifold} = normal_tvector_distribution(base_manifold(M), x, σ)
normal_tvector_distribution(M::MMT, ::Val{false}, x, σ) where {MMT <: MetricManifold} =error("normal_tvector_distribution not implemented for a $(typeof(M)) at point $(typeof(x)) with standard deviation $(typeof(σ)).")

function inner(B::VectorBundleFibers{<:CotangentSpaceType, MMT}, x, v, w) where {MMT<:MetricManifold}
    ginv = inverse_local_metric(B.M, x)
    return dot(v, ginv * w)
end

function flat!(M::MMT, v::FVector{CotangentSpaceType}, x, w::FVector{TangentSpaceType}) where {MMT<:MetricManifold}
    g = local_metric(M, x)
    copyto!(v.data, g*w.data)
    return v
end

function sharp!(M::MMT, v::FVector{TangentSpaceType}, x, w::FVector{CotangentSpaceType}) where {MMT<:MetricManifold}
    ginv = inverse_local_metric(M, x)
    copyto!(v.data, ginv*w.data)
    return v
end

# These are independent of the metric and hence can always fall back to M
injectivity_radius(M::MMT, args...) where {MMT <: MetricManifold} = injectivity_radius(base_manifold(M), args...)
zero_tangent_vector!(M::MMT, v, x) where {MMT <: MetricManifold} = zero_tangent_vector!(base_manifold(M), v, x)
check_manifold_point(M::MMT, x; kwargs...) where {MMT <: MetricManifold} = check_manifold_point(base_manifold(M), x; kwargs...)
check_tangent_vector(M::MMT, x, v; kwargs...) where {MMT <: MetricManifold} = check_tangent_vector(base_manifold(M), x, v; kwargs...)


text:
# ---
# title: Wide data
# cover: assets/wide_data.png
# description: Working with data in the wide format
# author: "[Pietro Vertechi](https://github.com/piever)"
# ---

using AlgebraOfGraphics, CairoMakie
using AlgebraOfGraphics: density
set_aog_theme!() #src

df = (a=randn(100), b=randn(100), c=randn(100))
labels = ["Trace 1", "Trace 2", "Trace 3"]
plt = data(df) *
    density() *
    mapping([:a, :b, :c] .=> "some label") *
    mapping(color=dims(1) => renamer(labels))
draw(plt)

#

df = (a=rand(100), b=rand(100), c=rand(100), d=rand(100))
labels = ["Trace 1", "Trace 2", "Trace 3"]
layers = linear() + visual(Scatter)
plt = data(df) * layers * mapping(1, 2:4 .=> "value", color=dims(1) => renamer(labels))
draw(plt)

# The wide format is combined with broadcast semantics.

df = (sepal_length=rand(100), sepal_width=rand(100), petal_length=rand(100), petal_width=rand(100))
xvars = ["sepal_length", "sepal_width"]
yvars = ["petal_length" "petal_width"]
layers = linear() + visual(Scatter)
plt = data(df) * layers * mapping(xvars, yvars, col=dims(1), row=dims(2))
fg = draw(plt)

# save cover image #src
mkpath("assets") #src
save("assets/wide_data.png", fg) #src

ext: jl


text:
using Plots, JLD2, Langevin, FluxReconstruction

begin
    set = Setup(
        "gas",
        "cylinder",
        "2d0f",
        "hll",
        "nothing",
        1, # species
        3, # order of accuracy
        "positivity", # limiter
        "euler",
        0.1, # cfl
        1.0, # time
    )
    ps = FRPSpace2D(0.0, 2.0, 100, 0.0, 1.0, 50, set.interpOrder-1, 1, 1)
    vs = nothing
    gas = Gas(
        1e-6,
        1.12, # Mach
        1.0,
        3.0, # K
        7/5,
        0.81,
        1.0,
        0.5,
    )
    ib = nothing
    ks = SolverSet(set, ps, vs, gas, ib)

    uqMethod = "galerkin"
    nr = 5
    nRec = 10
    opType = "uniform"
    parameter1 = 0.95
    parameter2 = 1.05

    uq = UQ1D(nr, nRec, parameter1, parameter2, opType, uqMethod)
end

cd(@__DIR__)

begin
    @load "fvm-fine/fvm1.jld2" x
    xc = deepcopy(x)

    fvm = []
    for i = 1:9
        filename = "fvm-fine/fvm" * string(i) * ".jld2"
        @load filename x sol
        push!(fvm, sol)
    end

    solc = zeros(size(fvm[1])..., 2)
    for i in axes(solc, 1), j in axes(solc, 2), k in axes(solc, 3)
        uRan = [fvm[idx][i, j, k] for idx = 1:9]
        uChaos = ran_chaos(uRan, uq)

        solc[i, j, k, 1] = mean(uChaos, uq.op)
        solc[i, j, k, 2] = std(uChaos, uq.op)
    end
end

begin
    @load "fvm/fvm1.jld2" x
    xr = deepcopy(x)

    fvm = []
    for i = 1:9
        filename = "fvm/fvm" * string(i) * ".jld2"
        @load filename x sol
        push!(fvm, sol)
    end

    solr = zeros(size(fvm[1])..., 2)
    for i in axes(solr, 1), j in axes(solr, 2), k in axes(solr, 3)
        uRan = [fvm[idx][i, j, k] for idx = 1:9]
        uChaos = ran_chaos(uRan, uq)

        solr[i, j, k, 1] = mean(uChaos, uq.op)
        solr[i, j, k, 2] = std(uChaos, uq.op)
    end
end

@load "galerkin/sol.jld2" x u#sol
begin
    sol = zeros(size(x)..., 4, 2)
    for i = 1:ps.nx, j = 1:ps.ny
        idx0 = (i - 1) * (ps.deg+1)
        idy0 = (j - 1) * (ps.deg+1)

        for k = 1:ps.deg+1, l = 1:ps.deg+1
            idx = idx0 + k
            idy = idy0 + l

            uChaos = uq_conserve_prim(u[i, j, k, l, :, :], ks.gas.γ, uq)
            uChaos[4, :] .= lambda_tchaos(Array(uChaos[4, :]), 1.0, uq)

            for s = 1:4
                sol[idx, idy, s, 1] = mean(uChaos[s, :], uq.op)
                sol[idx, idy, s, 2] = std(uChaos[s, :], uq.op)
            end
        end
    end
end

# density
plot(
    x[:, 1],
    0.5 .* (sol[:, end÷2, 1, 1] .+ sol[:, end÷2+1, 1, 1]),
    lw = 1.5,
    label = "Current",
    xlabel = "x",
    ylabel = "density",
)
scatter!(
    xc[1:2:end, 1],
    0.5 .* (solc[1:2:end, end÷2, 1, 1] .+ solc[1:2:end, end÷2+1, 1, 1]),
    label = "Collocation",
    alpha = 0.8,
)
plot!(
    xr,
    0.5 .* (solr[:, end÷2, 1, 1] .+ solr[:, end÷2+1, 1, 1]),
    lw = 1.5,
    label = "FVM",
    line = :dash,
    color = :gray27,
)
savefig("sv_t1_n_mean.pdf")

plot(
    x[:, 1],
    0.5 .* (sol[:, end÷2, 1, 2] .+ sol[:, end÷2+1, 1, 2]),
    lw = 1.5,
    label = "Current",
    xlabel = "x",
    ylabel = "density",
)
scatter!(
    xc[1:2:end, 1],
    0.5 .* (solc[1:2:end, end÷2, 1, 2] .+ solc[1:2:end, end÷2+1, 1, 2]),
    label = "Collocation",
    alpha = 0.8,
)
plot!(
    xr,
    0.5 .* (solr[:, end÷2, 1, 2] .+ solr[:, end÷2+1, 1, 2]),
    lw = 1.5,
    label = "FVM",
    line = :dash,
    color = :gray27,
)
savefig("sv_t1_n_std.pdf")

# temperature
plot(
    x[:, 1],
    0.5 .* (sol[:, end÷2, 4, 1] .+ sol[:, end÷2+1, 4, 1]),
    lw = 1.5,
    label = "Current",
    xlabel = "x",
    ylabel = "temperature",
)
scatter!(
    xc[1:2:end, 1],
    0.5 .* (solc[1:2:end, end÷2, 4, 1] .+ solc[1:2:end, end÷2+1, 4, 1]),
    label = "Collocation",
    alpha = 0.8,
)
plot!(
    xr,
    0.5 .* (solr[:, end÷2, 4, 1] .+ solr[:, end÷2+1, 4, 1]),
    lw = 1.5,
    label = "FVM",
    line = :dash,
    color = :gray27,
)
savefig("sv_t1_t_mean.pdf")

plot(
    x[:, 1],
    0.5 .* (sol[:, end÷2, 4, 2] .+ sol[:, end÷2+1, 4, 2]),
    lw = 1.5,
    label = "Current",
    xlabel = "x",
    ylabel = "temperature",
)
scatter!(
    xc[1:2:end, 1],
    0.5 .* (solc[1:2:end, end÷2, 4, 2] .+ solc[1:2:end, end÷2+1, 4, 2]),
    label = "Collocation",
    alpha = 0.8,
)
plot!(
    xr,
    0.5 .* (solr[:, end÷2, 4, 2] .+ solr[:, end÷2+1, 4, 2]),
    lw = 1.5,
    label = "FVM",
    line = :dash,
    color = :gray27,
)
savefig("sv_t1_t_std.pdf")

# ---
# contour
# ---
@load "iter_0.3.jld2" u#sol
begin
    sol = zeros(size(x)..., 4, 2)
    for i = 1:ps.nx, j = 1:ps.ny
        idx0 = (i - 1) * (ps.deg+1)
        idy0 = (j - 1) * (ps.deg+1)

        for k = 1:ps.deg+1, l = 1:ps.deg+1
            idx = idx0 + k
            idy = idy0 + l

            uChaos = uq_conserve_prim(u[i, j, k, l, :, :], ks.gas.γ, uq)
            uChaos[4, :] .= lambda_tchaos(Array(uChaos[4, :]), 1.0, uq)

            for s = 1:4
                sol[idx, idy, s, 1] = mean(uChaos[s, :], uq.op)
                sol[idx, idy, s, 2] = std(uChaos[s, :], uq.op)
            end
        end
    end
end

contourf(
    x[:, 1],
    y[1, :],
    sol[:, :, 1, 1]',
    ratio=1,
    ylims=(0,1),
    xlabel="x",
    ylabel="y",
)

savefig("sv_t03_mean.pdf")

contourf(
    x[:, 1],
    y[1, :],
    sol[:, :, 1, 2]',
    ratio=1,
    ylims=(0,1),
    xlabel="x",
    ylabel="y",
)

savefig("sv_t03_std.pdf")

max_stars_repo_name: CSMMLab/FRSG

max_issues_repo_name: CSMMLab/FRSG

max_forks_repo_name: CSMMLab/FRSG


text:
module MaxHelpingHandSecretBerry

using ..Ahorn, Maple

@mapdef Entity "MaxHelpingHand/SecretBerry" SecretBerry(x::Integer, y::Integer,
    checkpointID::Integer=-1, order::Integer=-1, strawberrySprite::String="strawberry", ghostberrySprite::String="ghostberry",
    strawberryPulseSound::String="event:/game/general/strawberry_pulse", strawberryBlueTouchSound::String="event:/game/general/strawberry_blue_touch",
    strawberryTouchSound::String="event:/game/general/strawberry_touch", strawberryGetSound::String="event:/game/general/strawberry_get", countTowardsTotal::Bool=false,
    particleColor1::String="FF8563", particleColor2::String="FFF4A8", ghostParticleColor1::String="6385FF", ghostParticleColor2::String="72F0FF")

const placements = Ahorn.PlacementDict(
    "Secret Berry (max480's Helping Hand)" => Ahorn.EntityPlacement(
        SecretBerry,
        "point"
    )
)

Ahorn.editingOrder(entity::SecretBerry) = String["x", "y", "strawberrySprite", "ghostberrySprite", "strawberryPulseSound", "strawberryTouchSound", "strawberryBlueTouchSound", "strawberryGetSound",
    "particleColor1", "particleColor2", "ghostParticleColor1", "ghostParticleColor2"]

sprite = "collectables/moonBerry/normal00"

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

    return Ahorn.getSpriteRectangle(sprite, x, y)
end

Ahorn.render(ctx::Ahorn.Cairo.CairoContext, entity::SecretBerry, room::Maple.Room) = Ahorn.drawSprite(ctx, sprite, 0, 0)

end

max_stars_repo_name: Kitty-Cats/MaxHelpingHand

max_issues_repo_name: Kitty-Cats/MaxHelpingHand

max_forks_repo_name: Kitty-Cats/MaxHelpingHand


text:
export Segment

convert(::Type{ChebyshevInterval}, ::AnyDomain) = ChebyshevInterval()
convert(::Type{ChebyshevInterval{T}}, ::AnyDomain) where T = ChebyshevInterval{T}()

## Standard interval
# T Must be a Vector space
"""
	Segment(a,b)

represents a line segment from `a` to `b`.  In the case where `a` and `b`
are real and `a < b`, then this is is equivalent to an `Interval(a,b)`.
"""
struct Segment{T} <: AbstractSegment{T}
	a::T
	b::T
	Segment{T}(a,b) where {T} = new{T}(a,b)
end

Segment(a::Complex{IT1}, b::Complex{IT2}) where {IT1<:Integer,IT2<:Integer} =
	Segment(ComplexF64(a), ComplexF64(b)) #convenience method
Segment(a::Integer, b::Integer) = Segment(Float64(a),Float64(b)) #convenience method
Segment(a::Complex{IT}, b) where {IT<:Integer} = Segment(ComplexF64(a),b) #convenience method
Segment(a, b::Complex{IT}) where {IT<:Integer} = Segment(a,ComplexF64(b)) #convenience method
Segment(a, b) = Segment{promote_type(typeof(a),typeof(b))}(a,b)
Segment(a::Tuple, b::Tuple) = Segment(Vec(a...),Vec(b...))

convert(::Type{Domain{T}}, d::Segment) where {T<:Number} = Segment{T}(leftendpoint(d),rightendpoint(d))
convert(::Type{Domain{T}}, d::Segment) where {T<:SVector} = Segment{T}(leftendpoint(d),rightendpoint(d))
convert(::Type{Segment{T}}, d::Segment) where {T<:Number} = Segment{T}(leftendpoint(d),rightendpoint(d))
convert(::Type{Segment}, d::AbstractInterval) = Segment(leftendpoint(d), rightendpoint(d))
convert(::Type{Segment{T}}, d::AbstractInterval) where T =convert(Segment{T}, convert(Segment, d))

Segment(d::AbstractInterval) = convert(Segment, d)

AnySegment(::Type{T}) where {T} = Segment{T}(NaN,NaN)
AnySegment() = AnySegment(Float64)
isambiguous(d::IntervalOrSegmentDomain) = all(isnan(leftendpoint(d))) && all(isnan(rightendpoint(d)))
convert(::Type{Segment{T}},::AnyDomain) where {T<:Number} = AnySegment(T)
convert(::Type{Segment},::AnyDomain) = AnySegment()
convert(::Type{Interval}, d::Segment{<:Real}) = d.a < d.b ? d.a .. d.b : d.b .. d.a
convert(::Type{ClosedInterval}, ::AnyDomain) = NaN..NaN
convert(::Type{ClosedInterval{T}}, ::AnyDomain) where T = T(NaN)..T(NaN)
Interval(d::Segment) = convert(Interval, d)

## Information
@inline leftendpoint(d::Segment) = d.a
@inline rightendpoint(d::Segment) = d.b
@inline endpoints(d::Segment) = d.a, d.b

@inline minimum(d::Segment) = min(leftendpoint(d),rightendpoint(d))
@inline maximum(d::Segment) = max(leftendpoint(d),rightendpoint(d))

isempty(d::Segment) = isapprox(leftendpoint(d), rightendpoint(d); atol=200eps(eltype(d)))

issubset(a::Segment,b::Segment) = leftendpoint(a)∈b && rightendpoint(a)∈b

arclength(d::AbstractInterval) = width(d)
arclength(d::Segment) = norm(complexlength(d))
complexlength(d::IntervalOrSegment) = rightendpoint(d)-leftendpoint(d)
mean(d::IntervalOrSegment) = (rightendpoint(d)+leftendpoint(d))/2
angle(d::IntervalOrSegment) = angle(complexlength(d))
sign(d::IntervalOrSegment) = sign(complexlength(d))

## Map interval
# The first definition  is the more general

mobius(S::Space,x...) = mobius(domain(S),x...)

mobius(d::ChebyshevInterval{T},x) where {T<:Number} = x
fromcanonical(d::ChebyshevInterval{T},x) where {T<:Number} = x
fromcanonicalD(d::ChebyshevInterval{T},x) where {T<:Number} = one(x)
tocanonical(d::ChebyshevInterval{T},x) where {T<:Number} = x
tocanonicalD(d::ChebyshevInterval{T},x) where {T<:Number} = one(x)

tocanonical(d::IntervalOrSegment{T},x) where {T} = 2norm(x-leftendpoint(d))/arclength(d)-1
tocanonical(d::IntervalOrSegment{T},x::Number) where {T<:Complex} = 2norm(x-leftendpoint(d))/arclength(d)-1
mobius(d::IntervalOrSegment,x) = (2x - leftendpoint(d) - rightendpoint(d))/complexlength(d)
tocanonical(d::IntervalOrSegment{T},x) where {T<:Real} = mobius(d,x)
tocanonicalD(d::IntervalOrSegment{T},x) where {T<:Real} = 2/complexlength(d)
fromcanonical(d::IntervalOrSegment{T},x) where {T<:Number} = mean(d) + complexlength(d)x/2
fromcanonical(d::IntervalOrSegment{T},x) where {T<:Vec} = mean(d) + complexlength(d)x/2
fromcanonicalD(d::IntervalOrSegment,x) = complexlength(d) / 2

==(d::Segment, m::Segment) = (isambiguous(d) && isambiguous(m)) || (leftendpoint(d) == leftendpoint(m) && rightendpoint(d) == rightendpoint(m))
function isapprox(d::Segment, m::Segment)
    tol=10E-12
    norm(leftendpoint(d)-leftendpoint(m))<tol && norm(rightendpoint(d)-rightendpoint(m))<tol
end

for op in (:(==), :isapprox)
    @eval begin
        $op(d::Segment, m::AbstractInterval) = $op(d, Segment(m))
        $op(m::AbstractInterval, d::Segment) = $op(Segment(m), d)
    end
end

## algebra

for op in (:*,:+,:-)
    @eval begin
        $op(c::Number,d::Segment) = broadcast($op,c,d)
        $op(d::Segment,c::Number) = broadcast($op,d,c)
        broadcasted(::typeof($op), c::Number, d::Segment) = Segment($op(c,leftendpoint(d)),$op(c,rightendpoint(d)))
        broadcasted(::typeof($op), d::Segment, c::Number) = Segment($op(leftendpoint(d),c),$op(rightendpoint(d),c))
    end
end

broadcasted(::typeof(^),c::Number,d::Segment) = Segment(c^leftendpoint(d),c^rightendpoint(d))
function broadcasted(::typeof(^),d::Segment,c::Number)
    a,b = endpoints(d)
    if a < 0 < b
        Segment(0, b^c)
    elseif b < 0 < a
        Segment(a^c, 0)
    else
        Segment(a^c,b^c)
    end
end

broadcasted(::typeof(Base.literal_pow), ::typeof(^), d::Segment, ::Val{K}) where K =
    broadcasted(^, d, K)

/(d::Segment,c::Number) = broadcast(/,d,c)
broadcasted(::typeof(/), d::Segment,c::Number) = Segment(leftendpoint(d)/c,rightendpoint(d)/c)

sqrt(d::Segment) = broadcast(sqrt, d)
broadcasted(::typeof(sqrt), d::Segment)=Segment(sqrt(leftendpoint(d)),sqrt(rightendpoint(d)))

+(d1::Segment,d2::Segment)=Segment(d1.a+d2.a,d1.b+d2.b)
broadcasted(::typeof(+),d1::Segment,d2::Segment) = Segment(d1.a+d2.a,d1.b+d2.b)

DomainSets.map_domain(map::DomainSets.AbstractAffineMap, domain::AbstractSegment) =
    Segment(map(leftendpoint(domain)),map(rightendpoint(domain)))

## intersect/union

reverseorientation(d::IntervalOrSegment) = Segment(rightendpoint(d),leftendpoint(d))

intersect(a::Segment{<:Real}, b::Segment{<:Real}) = intersect(Interval(a), Interval(b))
intersect(a::AbstractInterval, b::Segment{<:Real}) = intersect(a, Interval(b))
intersect(a::Segment{<:Real}, b::AbstractInterval) = intersect(Interval(a), b)
setdiff(a::Segment{<:Real}, b::Segment{<:Real})  = setdiff(Interval(a), Interval(b))
setdiff(a::AbstractInterval, b::Segment{<:Real})  = setdiff(a, Interval(b))
setdiff(a::Segment{<:Real}, b::AbstractInterval)  = setdiff(Interval(a), b)

## sort
isless(d1::Segment{T1},d2::Segment{T2}) where {T1<:Real,T2<:Real} =
    d1 ≤ leftendpoint(d2) && d1 ≤ rightendpoint(d2)
isless(d1::Segment{T},x::Real) where {T<:Real}=leftendpoint(d1) ≤ x && rightendpoint(d1) ≤ x
isless(x::Real,d1::Segment{T}) where {T<:Real}=x≤leftendpoint(d1) && x≤rightendpoint(d1)


text:
using DynamicPPL
using Turing

using Random
using Test

setprogress!(false)

Random.seed!(100)

# load test utilities
include(joinpath(pathof(Turing), "..", "..", "test", "test_utils", "numerical_tests.jl"))

@testset "Turing" begin
    include("compiler.jl")
    include("loglikelihoods.jl")
    include("model.jl")
    include("prob_macro.jl")
    include("varinfo.jl")
end


text:
mutable struct CodonGraph{C<:Codon}
    ref::CodonGraphReference{C}
    vertices::Vector{C}
    edges::BitVector
    nv::Int
end

function CodonGraph{C}(ref::CodonGraphReference{C}) where C <: Codon
    return CodonGraph{C}(ref, Vector{C}(64), falses(2016), 0)
end

function CodonGraph{C}(code::GeneticCode) where C <: Codon
    reference = CodonGraphReference{C}(code)
    return CodonGraph{C}(reference)
end

function reset!(cg::CodonGraph, cs::CodonSet)
    # Fill the vertices buffer with codons from the codon set.
    cg.nv = length(cs)
    copy!(cg.vertices, cs)
    # Reset and turn on the bits indicating which edges we have.
    # Order of bits is not the same as cg.ref.edges, but is dictated by
    # cg.ref.edge_order. This removes the need to sort a lot, removing a key
    # bottleneck in Kruskal's algorithm.
    fill!(cg.edges, false)
    @inbounds for i in 1:cg.nv
        x = cg.vertices[i]
        for j in (i + 1):(cg.nv)
            y = cg.vertices[j]
            e = lookup_order(cg.ref, x, y)
            cg.edges[e] = true
        end
    end
end

@inline Base.start(x::CodonGraph) = findnext(x.edges, 1)
@inline function Base.next(x::CodonGraph, state::Int)
    e = lookup_edge(x.ref, state)
    s = findnext(x.edges, state + 1)
    return e, s
end
@inline Base.done(x::CodonGraph, state::Int) = state == 0
Base.eltype(x::CodonGraph{C}) where C <: Codon = Tuple{C, C, Int, Int}
Base.iteratorsize(::Type{CodonGraph{C}}) where C <: Codon = HasLength()
Base.iteratoreltype(::Type{CodonGraph{C}}) where C <: Codon = HasEltype()
Base.length(x::CodonGraph) = count(x.edges)


text:
module ADI

using Parameters
using Reexport
using Statistics
using UnPack

@reexport using HCIToolbox

export reconstruct,
       decompose,
       Median,
       PCA,
       GreeDS,
       TPCA,
       NMF,
       SingleSDI,
       DoubleSDI,
       SliceSDI

"""
    ADI.ADIAlgorithm

This abstract type is used for defining ADI algorithms. See the extended help (`??ADIAlgorithm`) for interface details.

# Extended help
## Interface
To extend `ADIAlgorithm` you may implement the following

| function | default | description |
|----------|---------|-------------|
| [`reconstruct`](@ref) | | Subroutine for creating the full reconstructed cube with the PSF |
| `(::ADIAlgorithm)` | subtracts output of `reconstruct`, then derotates and collapses | Subroutine for returning the reduced residual cube |
"""
abstract type ADIAlgorithm end

"""
    reconstruct(::ADIAlgorithm, cube, angles, [cube_ref]; kwargs...)

Reconstruct the PSF approximation for the given algorithm, using `cube_ref` as the reference cube if given.

# Examples

```julia
julia> cube, angles = # load data

julia> S = reconstruct(PCA(10), cube, angles);

julia> size(S) == size(cube)
true

julia> flat_res = collapse(cube .- S, angles); # form resid, derotate, and combine
```
"""
function reconstruct end

"""
    (::ADIAlgorithm)(cube, angles, [cube_ref]; kwargs...)

Fully process an ADI data cube using [`reconstruct`](@ref) and collapsing the residuals. Keyword arguments will be passed to `HCIToolbox.collapse!`.
"""
function (alg::ADIAlgorithm)(cube::AbstractArray{T,3}, angles, args...; kwargs...) where T
    S = reconstruct(alg, cube, angles, args...; kwargs...)
    residual_cube = cube .- S
    return collapse!(residual_cube, angles; kwargs...)
end

include("median.jl")

"""
    ADI.LinearAlgorithm <: ADI.ADIAlgorithm

This abstract type is used for defining linear ADI algorithms. See the extended help (`??LinearAlgorithm`) for interface details.

# Extended help
## Interface
To extend `LinearAlgorithm` you may implement the following

| function | default | description |
|----------|---------|-------------|
| [`decompose`](@ref) | | Subroutine for fitting the linear basis and coefficients as unrolled matrices |
| [`reconstruct`](@ref) | Computes the inner product of the design matrix and weights from `decompose` | Subroutine for creating the full reconstructed cube with the PSF |
| `(::LinearAlgorithm)` | subtracts output of `reconstruct`, then derotates and collapses | Subroutine for returning the reduced residual cube |
"""
abstract type LinearAlgorithm <: ADIAlgorithm end

"""
    ADI.decompose(alg::LinearAlgorithm, cube, angles, [cube_ref]; kwargs...)

Return the basis and coefficients fit by `alg`. This process generally requires fitting a basis from `cube` (or `cube_ref`) and then fitting the coefficients that project `cube` onto that basis. The inner-product of the two outputs should for an approximate reconstruction of the reference used.

# Examples

```julia
julia> cube, angles = # load data

julia> A, w = decompose(PCA(5), cube, angles);

julia> S = reconstruct(PCA(5), cube, angles);

julia> S ≈ expand(w * A) # inner product reshaped into cube
true
```
"""
function decompose end

function reconstruct(alg::LinearAlgorithm, cube::AbstractArray{T,3}, angles, args...; kwargs...) where T
    # assumed sizes are (n, Npx) (M, n)
    basis, weights = decompose(alg, cube, angles, args...; kwargs...)
    return reconstruct(alg, basis, weights; kwargs...)
end

function reconstruct(alg::LinearAlgorithm, basis::AbstractMatrix, weights::AbstractMatrix; kwargs...)
    # assumed sizes are (n, Npx) (M, n)
    return weights * basis |> expand
end
# The core decomposition routines
include("pca.jl")
include("greeds.jl")
include("nmf.jl")
include("sdi.jl")

include("metrics/Metrics.jl")
@reexport using .Metrics

end


text:
InputFolder = './Images/ManyCells/';
OutputFolder = './Results/Images/ManyCells/';
Fill = -1;

@iA = '*.tif';		% Image filter

@fxg_sBlobDetRay [iA] > [S];
params.Scale = 0.4;
params.Thr = 30;
params.NAngles = 18;
params.L = 9;
params.Step = 1;
params.Fraction = 0.25;
params.NoiseTol = 0.1;
/endf

/show iA > S;
/keep S > tif;

max_stars_repo_name: SebastienTs/LOBSTER

max_issues_repo_name: SebastienTs/LOBSTER

max_forks_repo_name: SebastienTs/LOBSTER


text:
@doc doc"""
point_estimate_avg(psm::Array{Float64, 2})

	Find optimal partition which minimizes the lower bound to the Variation of Information
	obtain from Jensen's inequality where the expectation and log are reversed.

	Code based on R implementation by Sara Wade <sara.wade@eng.cam.ac.uk>

""" ->
function point_estimate_hclust(psm::Array{Float64, 2}; maxk = -1, method = :average)

	havg = hclust(1-psm, method)
	cls = reduce(hcat, map(k -> cutree(havg; k = k), 1:maxk))

	# compute Variation of Information lower bound
	vi = variation_of_information_lb(cls, psm)

	idx = findfirst(minimum(vi) .== vi)

	if ndims(cls) == 2
		cl = cls[:,idx]
	else
		cl = cls
	end

	return (cl, minimum(vi))

end

@doc doc"""
point_estimate(psm::Array{Float64, 2})

	Find optimal partition which minimizes the lower bound to the Variation of Information
	obtain from Jensen's inequality where the expectation and log are reversed.

	Code based on R implementation by Sara Wade <sara.wade@eng.cam.ac.uk>

""" ->
function point_estimate(psm::Array{Float64, 2}; method = :avg, maxk = -1)

	methods = [:avg, :comp, :greedy]

	@assert method in methods "Method must be :avg , :comp or :greedy."

	if maxk == -1
		maxk = convert(Int, ceil(size(psm)[1] / 4.0))
	end

	if method == :avg
		return point_estimate_hclust(psm, maxk = maxk, method = :average)
	elseif method == :comp
		return point_estimate_hclust(psm, maxk = maxk, method = :complete)
	else
		println("not implemented")
	end
end

@doc doc"""
variation_of_information()
	Computes the lower bound to the posterior expected Variation of Information
""" ->
function variation_of_information_lb(cls::Array{Int, 2}, psm)

	N = size(psm)[1]
	F = zeros(size(cls)[2])

	for ci in 1:size(cls)[2]
		c = cls[:,ci]
		for i in 1:N
			ind = c .== c[i]
			F[ci] += (log2(sum(ind)) +log2(sum(psm[i,:])) -2 * log2(sum(ind' .* psm[i,:]))) / N
		end
	end

	return F

end

@doc doc"""
variation_of_information()
	Computes the lower bound to the posterior expected Variation of Information
""" ->
function variation_of_information_lb(cls::Vector{Int}, psm)

	N = size(psm)[1]
	F = zeros(1)

	for i in 1:N
		ind = cls .== cls[i]
		F[1] += (log2(sum(ind))+log2(sum(psm[i,:]))-2*log2(sum(ind' .* psm[i,:]))) / N
	end

	return F

end


text:
# method of moments estimators for a sample from Chi^2(theta)
# increase n to observe consistency
# note that the two estimators are different from one another
using Distributions
n = 30
theta = 3

# the distribution doesn't matter, what matters is
# that the moments are correctly specified
#y = rand(Chisq(theta), n)
y = rand(Normal(theta, sqrt(2*theta)), n)

# a MM estimator
thetahat = mean(y)
println("mm estimator based on mean: ", thetahat)

# a second MM estimator
thetahat = 0.5*var(y)
println("mm estimator based on variance: ", thetahat)

max_stars_repo_name: Hiroakiyusheng/Econometrics

max_issues_repo_name: Hiroakiyusheng/Econometrics

max_forks_repo_name: Hiroakiyusheng/Econometrics


text:
# Quality control criteria for results

"Return the ratio λ₂⁰/minimum(λ₂), where λ₂⁰ is the uncorrected λ₂.
This should be **maximised** when removing the best-fitting splitting operator,
turning elliptical particle motion linear."
lam2_ratio(s::Result) = first(s.lam2)/minimum(s.lam2)

"Return the ratio λ₁⁰/maximum(λ₁), where λ₁⁰ is the uncorrected λ₁.
This should be **minimised** when removing the best-fitting splitting operator,
turning elliptical particle motion linear."
lam1_ratio(s::Result) = first(s.lam1)/maximum(s.lam1)

"Return the ratio (λ₂⁰/λ₁⁰) / [minimum(λ₂)/maximum(λ₁)].
This should be **maximised** when removing the best-fitting splitting operator,
turning elliptical partice motion linear."
lam_ratio(s::Result) = lam2_ratio(s)/lam1_ratio(s)

"""
    quality(result) -> Q

Return the quality index `Q`, which is 1 for a perfect splitting measurement,
-1 for a perfect null, and 0 for a poor measurement.

See Wuestefeld et al. (Geophysical Prospecting, 2010) for details.
"""
function quality(s::Result)::Float64
    s.xcorr_dt === nothing &&
        throw(ArgumentError("rotation correlation map not computed for this result"))
    Δ = s.xcorr_dt_best/s.dt_best
    ϕdiff = abs(Seis.angle_difference(s.phi_best, s.xcorr_phi_best, true))
    ϕdiff > 45 && (ϕdiff = 90 - ϕdiff)
    Ω = ϕdiff/45
    d_null = sqrt(Δ^2 + (Ω - 1)^2)*√2
    d_null > 1 && (d_null = 1.0)
    d_good = sqrt((Δ - 1)^2 + Ω^2)*√2
    d_good > 1 && (d_good = 1.0)
    Q = d_null < d_good ? -(1 - d_null) : 1 - d_good
    dt_max = maximum(s.dt)
    Δ > 0.9 && s.dt_best/dt_max > 0.9 && (Q = -Q*(10 - 10*s.dt_best/dt_max))
    Q
end

"""
    snr_restivo_helffrich(result) -> SNR

Return the signal-to-noise ratio as defined by Restivo & Helffrich (GJI, 1999):
the signal is the maximum absolute amplitude of the corrected trace in the
source polarisation orientation, and the noise is the 2σ.
"""
function snr_restivo_helffrich(s::Result)
    N, E = deepcopy.((s.trace1, s.trace2))
    Seis.cut!.((N, E), s.window_start, s.window_end)
    rotate_traces!(N, E, -N.sta.azi) # Now in N-E frame
    apply_split!(N, E, s.phi_best, -s.dt_best)
    rotate_traces!(N, E, -s.spol) # Now in spol-(spol+90) frame
    signal = maximum(abs, Seis.trace(N))
    noise = 2*Statistics.std(Seis.trace(E))
    signal/noise
end

max_stars_repo_name: anowacki/SeisSplit

max_issues_repo_name: anowacki/SeisSplit

max_forks_repo_name: anowacki/SeisSplit


text:
"""
    HeteroscedasticLikelihood(λ::T=1.0)

## Arguments
- `λ::Real` : The maximum precision possible (this is optimized during training)

---

Gaussian with heteroscedastic noise given by another gp:
```math
    p(y|f,g) = N(y|f,(λ σ(g))⁻¹)
```
Where `σ` is the logistic function

The augmentation is not trivial and will be described in a future paper
"""
mutable struct HeteroscedasticLikelihood{T<:Real,A<:AbstractVector{T}} <:
               RegressionLikelihood{T}
    λ::T
    c::A
    ϕ::A
    γ::A
    θ::A
    σg::A
    function HeteroscedasticLikelihood{T}(λ::T) where {T<:Real}
        return new{T,Vector{T}}(λ)
    end
    function HeteroscedasticLikelihood{T}(
        λ::T, c::A, ϕ::A, γ::A, θ::A, σg::A
    ) where {T<:Real,A<:AbstractVector{T}}
        return new{T,A}(λ, c, ϕ, γ, θ, σg)
    end
end

function HeteroscedasticLikelihood(λ::T=1.0) where {T<:Real}
    return HeteroscedasticLikelihood{T}(λ)
end

implemented(::HeteroscedasticLikelihood, ::AnalyticVI) = true

function (l::HeteroscedasticLikelihood)(y::Real, f::AbstractVector)
    return pdf(Normal(y, inv(sqrt(l.λ * logistic(f[2])))), f[1])
end

function Distributions.loglikelihood(
    l::HeteroscedasticLikelihood, y::Real, f::AbstractVector
)
    return logpdf(Normal(y, inv(sqrt(l.λ * logistic(f[2])))), f[1])
end

function Base.show(io::IO, ::HeteroscedasticLikelihood{T}) where {T}
    return print(io, "Gaussian likelihood with heteroscedastic noise")
end

num_latent(::HeteroscedasticLikelihood) = 2

function treat_labels!(y::AbstractVector{<:Real}, likelihood::HeteroscedasticLikelihood)
    return y, 2, likelihood
end

function treat_labels!(y::AbstractVector, likelihood::HeteroscedasticLikelihood)
    return error("For regression target(s) should be real valued")
end

function init_likelihood(
    likelihood::HeteroscedasticLikelihood{T}, ::AbstractInference{T}, ::Int, nMinibatch::Int
) where {T<:Real}
    λ = likelihood.λ
    c = ones(T, nMinibatch)
    ϕ = ones(T, nMinibatch)
    γ = ones(T, nMinibatch)
    θ = ones(T, nMinibatch)
    σg = ones(T, nMinibatch)
    return HeteroscedasticLikelihood{T}(λ, c, ϕ, γ, θ, σg)
end

function compute_proba(
    l::HeteroscedasticLikelihood{T},
    μ::AbstractVector{<:AbstractVector},
    Σ::AbstractVector{<:AbstractVector},
) where {T}
    return μ[1], max.(Σ[1], zero(T)) .+ inv.(l.λ * logistic.(μ[2]))
end

function local_updates!(
    l::HeteroscedasticLikelihood,
    y::AbstractVector,
    μ::NTuple{2,<:AbstractVector},
    diagΣ::NTuple{2,<:AbstractVector},
)
    # gp[1] is f and gp[2] is g (for approximating the noise)
    @. l.ϕ = 0.5 * (abs2(μ[1] - y) + diagΣ[1])
    @. l.c = sqrt(abs2(μ[2]) + diagΣ[2])
    @. l.γ = 0.5 * l.λ * l.ϕ * safe_expcosh(-0.5 * μ[2], 0.5 * l.c)
    @. l.θ = 0.5 * (0.5 + l.γ) / l.c * tanh(0.5 * l.c)
    @. l.σg = expectation(logistic, μ[2], diagΣ[2])
    return l.λ = 0.5 * length(l.ϕ) / dot(l.ϕ, l.σg)
end

function variational_updates!(
    m::AbstractGP{T,<:HeteroscedasticLikelihood,<:AnalyticVI}
) where {T}
    local_updates!(likelihood(m), yview(m), mean_f(m), var_f(m))
    natural_gradient!(
        ∇E_μ(likelihood(m), opt_type(inference(m)), yview(m))[2],
        ∇E_Σ(likelihood(m), opt_type(inference(m)), yview(m))[2],
        getρ(inference(m)),
        opt_type(inference(m)),
        last(Zviews(m)),
        m.f[2],
    )
    global_update!(m.f[2], opt_type(inference(m)), inference(m))
    heteroscedastic_expectations!(likelihood(m), mean_f(m.f[2]), var_f(m.f[2]))
    natural_gradient!(
        ∇E_μ(likelihood(m), opt_type(inference(m)), yview(m))[1],
        ∇E_Σ(likelihood(m), opt_type(inference(m)), yview(m))[1],
        getρ(inference(m)),
        opt_type(inference(m)),
        first(Zviews(m)),
        m.f[1],
    )
    return global_update!(m.f[1], opt_type(inference(m)), inference(m))
end

function heteroscedastic_expectations!(
    l::HeteroscedasticLikelihood{T}, μ::AbstractVector, Σ::AbstractVector
) where {T}
    @. l.σg = expectation(logistic, μ, Σ)
    return l.λ = 0.5 * length(l.ϕ) / dot(l.ϕ, l.σg)
end

function expectation(f::Function, μ::Real, σ²::Real)
    x = pred_nodes * sqrt(max(σ², zero(σ²))) .+ μ
    return dot(pred_weights, f.(x))
end

@inline function ∇E_μ(l::HeteroscedasticLikelihood, ::AOptimizer, y::AbstractVector)
    return (0.5 * y .* l.λ .* l.σg, 0.5 * (0.5 .- l.γ))
end

@inline function ∇E_Σ(l::HeteroscedasticLikelihood, ::AOptimizer, ::AbstractVector)
    return (0.5 * l.λ .* l.σg, 0.5 * l.θ)
end

function proba_y(
    model::AbstractGP{T,HeteroscedasticLikelihood{T},AnalyticVI{T}},
    X_test::AbstractMatrix{T},
) where {T<:Real}
    (μf, σ²f), (μg, σ²g) = predict_f(model, X_test; cov=true)
    return μf, σ²f + expectation.(x -> inv(model.likelihood.λ * logistic(x)), μg, σ²g)
end

function expec_loglikelihood(
    l::HeteroscedasticLikelihood{T}, ::AnalyticVI, y::AbstractVector, μ, diag_cov
) where {T}
    tot = length(y) * (0.5 * log(l.λ) - log(2 * sqrt(twoπ)))
    tot += 0.5 * (dot(μ[2], (0.5 .- l.γ)) - dot(abs2.(μ[2]), l.θ) - dot(diag_cov[2], l.θ))
    tot -= PoissonKL(l, y, μ[1], diag_cov[1])
    return tot
end

AugmentedKL(l::HeteroscedasticLikelihood, ::AbstractVector) = PolyaGammaKL(l)

function PoissonKL(
    l::HeteroscedasticLikelihood{T}, y::AbstractVector, μ::AbstractVector, Σ::AbstractVector
) where {T}
    return PoissonKL(
        l.γ, 0.5 * l.λ * (abs2.(y - μ) + Σ), log.(0.5 * l.λ * (abs2.(μ - y) + Σ))
    )
end

function PolyaGammaKL(l::HeteroscedasticLikelihood{T}) where {T}
    return PolyaGammaKL(0.5 .+ l.γ, l.c, l.θ)
end


text:
@testset "Tokenizer" begin

    sample_text = "this is some sample text"

    tkns = TextAnalysis.tokenize(
        Languages.English(),
        sample_text
    )

    @test isequal(
        tkns,
        String["this", "is", "some", "sample", "text"]
    )
end


text:
"""
Return the rotation angle around the
reference axis for the first vector to be
aligned to the second vector.

Arguments
----------
:Type{ProjectedRotationAngle}
    must be type ``ProjectedRotationAngle``

input:AbstractArray
    the vector to be rotated

target:AbstractArray
    the target vector to be aligned to

ref_axis:AbstractArray
    the reference axis around which input vector will be rotated

tol_near_zero=1e-15:AbstractFloat
    (optional) tolerance for the length of the input vectors to be close to zero
"""
function gage(::Type{ProjectedRotationAngle}, input::AbstractArray, target::AbstractArray,
    ref_axis::AbstractArray, tol_near_zero::AbstractFloat=1e-15)

    @debug @assert length(input) == length(target)

    ref_axis_magnitude = norm(ref_axis, 2)
    input_magnitude = norm(input, 2)
    target_magnitude = norm(target, 2)

    if ref_axis_magnitude < tol_near_zero # this means the input vector is the origin
        return 0. # no need to rotate
    else
        unit_ref_axis = ref_axis / ref_axis_magnitude
    end

    if input_magnitude < tol_near_zero # this means the input vector is the origin
        return 0. # no need to rotate
    else
        unit_input = input / input_magnitude
    end

    if target_magnitude < tol_near_zero # this means the target vector is the origin
        return 0. # no need to rotate
    else
        unit_target = target / target_magnitude
    end

    # get the projection along a direction normal to the ref_axis
    function normal_projection(input::AbstractArray, unit_ref_axis::AbstractArray)
        return input - dot(input, unit_ref_axis) * unit_ref_axis
    end

    return gage(RotationAngle, normal_projection(unit_input, unit_ref_axis),
        normal_projection(unit_target, unit_ref_axis))
end


text:
# Test the extensions to LCRST
@testset "LCRST Extensions" begin
    # test additions to addchild
    @testset "addchild(parent, child)" begin
        # setup nodes
        p = Node(1)
        c1 = Node(2)
        c2 = Node(3)
        # test first addition of first child
        @test addchild(p, c1) == c1
        @test c1.parent == p
        @test lastsibling(c1) == c1
        @test c1.child == c1
        @test p.child == c1
        # test first addition of 2nd child
        @test addchild(p, c2) == c2
        @test c1.parent == p
        @test c2.parent == p
        @test c1.sibling == c2
        @test c2.sibling == c2
        @test c2.child == c2
        @test c1.child == c1
        @test p.child == c1
        # test subsequent addition of first child (invokes copy)
        @test addchild(p, c1) !== c1
        @test c1.parent == p
        @test lastsibling(c1) !== c2 && lastsibling(c1) !== c1
        @test c1.child == c1
        @test p.child == c1
        c3 = lastsibling(c1)
        @test c3.parent == p
        @test c2.sibling == c3
        @test c3.child == c3
        @test c3 !== c1
        @test c3.data == 2
        # test subsequent addition of child to a node with no children
        p2 = Node(4)
        @test addchild(p2, c1) !== c1
        c4 = p2.child
        @test c4.parent == p2
        @test c4.data == 2
        @test c4.sibling == c4
        @test c4.child == c4
    end
    @testset "addchild(parent, nothing, child)" begin
        # setup nodes
        p = Node(1)
        c = Node(2)
        # test first addition of first child
        @test addchild(p, nothing, c) == c
        @test c.parent == p
        @test lastsibling(c) == c
        @test c.child == c
        @test p.child == c
        # test subsequent addition of child to a node with no children
        p2 = Node(4)
        @test addchild(p2, nothing, c) !== c
        c2 = p2.child
        @test c2.parent == p2
        @test c2.data == 2
        @test c2.sibling == c2
        @test c2.child == c2
    end
    @testset "addchild(parent, prevchild, data)" begin
        # setup nodes
        p = Node(1)
        c1 = Node(2)
        p.child = c1
        c1.parent = p
        # test addition on first next sibling
        @test addchild(p, c1, 3) isa Node
        c2 = c1.sibling
        @test c1 !== c2
        @test c2.data == 3
        @test c2.parent == p
        @test c2.sibling == c2
        @test c2.child == c2
        @test c1.sibling == c2
        @test c1.parent == p
        @test c1.child == c1
        # test adding more in for loop
        prev = c2
        for i in 4:5
            prev = addchild(p, prev, i)
        end
        @test [n.data for n in p] == [2, 3, 4, 5]
        c3 = c2.sibling
        c4 = c3.sibling
        @test c2 !== c3 && c3 !== c4
        @test c3.parent == p
        @test c3.sibling == c4
        @test c3.child == c3
        @test c4.parent == p
        @test c4.sibling == c4
    end
    @testset "addchild(parent, prevchild, child)" begin
        # setup nodes
        p = Node(1)
        c1 = Node(2)
        p.child = c1
        c1.parent = p
        c2 = Node(3)
        # test invalid previous
        @test_throws AssertionError addchild(p, c2, c1)
        # test first addition of 2nd child
        @test addchild(p, c1, c2) == c2
        @test c1.parent == p
        @test c2.parent == p
        @test c1.sibling == c2
        @test c2.sibling == c2
        @test c2.child == c2
        @test c1.child == c1
        @test p.child == c1
        # test subsequent addition of first child (invokes copy)
        @test addchild(p, c2, c1) !== c1
        @test c1.parent == p
        @test lastsibling(c1) !== c2 && lastsibling(c1) !== c1
        @test c1.child == c1
        @test p.child == c1
        c3 = lastsibling(c1)
        @test c3.parent == p
        @test c2.sibling == c3
        @test c3.child == c3
        @test c3 !== c1
        @test c3.data == 2
        # test subsequent double addition of a child to a root
        p = Node(1)
        c1 = Node(2)
        p.child = c1
        c1.parent = p
        @test addchild(p, c1, c1) !== c1
        @test [n.data for n in p] == [2, 2]
        c2 = c1.sibling
        @test c1 !== c2
        @test c1.parent == p
        @test c2.parent == p
        @test c1.sibling == c2
        @test c2.sibling == c2
        @test c2.child == c2
        @test c1.child == c1
        @test p.child == c1
        # test for loop addition of children
        p = Node(1)
        cs = [Node(2), Node(3), Node(4)]
        prev = nothing
        for i in 1:3
            prev = addchild(p, prev, cs[i])
        end
        @test p.child == cs[1]
        for i in 1:3
            @test cs[i].data == i + 1
            @test cs[i].parent == p
            @test lastsibling(cs[i]) == cs[3]
        end
    end
    # test _map_tree
    @testset "_map_tree" begin
        # setup tree
        p = Node(1)
        c1 = addchild(p, 2)
        c2 = addchild(p, 3)
        c3 = addchild(c2, 4)
        # test simple map
        str_p = InfiniteOpt._map_tree(n -> Node(string(n.data)), p)
        @test isroot(str_p)
        @test str_p.data == "1"
        @test str_p.child.data == "2"
        @test isleaf(str_p.child)
        @test str_p.sibling == str_p
        @test str_p.child.sibling.data == "3"
        @test islastsibling(str_p.child.sibling)
        @test isleaf(str_p.child.sibling.child)
        @test str_p.child.sibling.child.data == "4"
        # test empty root
        p = Node(1)
        str_p = InfiniteOpt._map_tree(n -> Node(string(n.data)), p)
        @test isroot(str_p)
        @test isleaf(str_p)
        @test str_p.data == "1"
    end
    # test copy(Node)
    @testset "Base.copy" begin
        # setup tree
        p = Node(1)
        c1 = addchild(p, 2)
        c2 = addchild(p, 3)
        c3 = addchild(c2, 4)
        # test simple copy
        p2 = copy(p)
        @test p2 !== p
        @test isroot(p2)
        @test p2.data == 1
        @test p2.child.data == 2
        @test isleaf(p2.child)
        @test p2.child !== p.child
        @test p2.sibling == p2
        @test p2.child.sibling.data == 3
        @test islastsibling(p2.child.sibling)
        @test isleaf(p2.child.sibling.child)
        @test p2.child.sibling.child.data == 4
        # test empty root
        p = Node(1)
        p2 = copy(p)
        @test p !== p2
        @test isroot(p2)
        @test isleaf(p2)
        @test p2.data == 1
    end
    # test _merge_parent_and_child
    @testset "_merge_parent_and_child" begin
        # setup tree
        p = Node(1)
        c1 = addchild(p, 2)
        c2 = addchild(c1, 3)
        c3 = addchild(c2, 4)
        c4 = addchild(c2, 5)
        # test normal
        @test InfiniteOpt._merge_parent_and_child(c1) == c1
        @test c1.data == 3
        @test c1.child == c3
        @test c3.parent == c1
        @test c4.parent == c1
        @test c1.parent == p
        # test at root node
        @test InfiniteOpt._merge_parent_and_child(p) == p
        @test p.data == 3
        @test p.child == c3
        @test c3.parent == p
        @test c4.parent == p
        @test p.parent == p
        # test nothing happens
        @test InfiniteOpt._merge_parent_and_child(p) == p
        @test p.data == 3
        @test p.child == c3
        @test c3.parent == p
        @test c4.parent == p
        @test p.parent == p
    end
end

# Test the core NLP data structures and methods
@testset "Core Data Structures" begin
    # setup model data
    m = InfiniteModel()
    @infinite_parameter(m, t in [0, 1])
    @variable(m, y, Infinite(t))
    @variable(m, z)
    # test NodeData
    @testset "NodeData" begin
        @test NodeData(1).value == 1
    end
    # test _node_value
    @testset "_node_value" begin
        @test InfiniteOpt._node_value(NodeData(1)) == 1
    end
    # test _is_zero
    @testset "_is_zero" begin
        # test easy case
        @test InfiniteOpt._is_zero(Node(NodeData(0)))
        # test is leaf that isn't zero
        @test !InfiniteOpt._is_zero(Node(NodeData(1)))
        @test !InfiniteOpt._is_zero(Node(NodeData(y)))
        # test addition block
        p = Node(NodeData(:+))
        addchild(p, NodeData(y))
        addchild(p, NodeData(0))
        @test !InfiniteOpt._is_zero(p)
        p.child.data = NodeData(0)
        @test InfiniteOpt._is_zero(p)
        # test multiplication block
        p = Node(NodeData(:*))
        addchild(p, NodeData(y))
        addchild(p, NodeData(0))
        @test InfiniteOpt._is_zero(p)
        p.child.sibling.data = NodeData(t)
        @test !InfiniteOpt._is_zero(p)
        # test power/divide
        p = Node(NodeData(:^))
        addchild(p, NodeData(y))
        addchild(p, NodeData(2))
        @test !InfiniteOpt._is_zero(p)
        p.child.data = NodeData(0)
        @test InfiniteOpt._is_zero(p)
        # test function
        p = Node(NodeData(:max))
        addchild(p, NodeData(0))
        addchild(p, NodeData(0))
        @test InfiniteOpt._is_zero(p)
        p.child.data = NodeData(y)
        @test !InfiniteOpt._is_zero(p)
    end
    # test _drop_zeros!
    @testset "_drop_zeros!" begin
        # test simple case
        n = Node(NodeData(0))
        @test InfiniteOpt._drop_zeros!(n) == n
        @test n.data.value == 0
        # test addition with zeros
        p = Node(NodeData(:+))
        addchild(p, NodeData(0))
        addchild(p, NodeData(y))
        @test InfiniteOpt._drop_zeros!(p) == p
        @test isequal(p.data.value, y)
        # test case with subtraction and zero function
        p = Node(NodeData(:-))
        n1 = addchild(p, NodeData(:sin))
        n2 = addchild(p, NodeData(y))
        n3 = addchild(n1, NodeData(0))
        @test InfiniteOpt._drop_zeros!(p) == p
        @test isequal(p.child.data.value, y)
        # test case with substraction of zero
        p = Node(NodeData(:-))
        n1 = addchild(p, NodeData(:sin))
        n2 = addchild(p, NodeData(0))
        n3 = addchild(n1, NodeData(y))
        @test InfiniteOpt._drop_zeros!(p) == p
        @test p.data.value == :sin
        @test isequal(p.child.data.value, y)
        # test truncating leaf
        p = Node(NodeData(:^))
        n1 = addchild(p, NodeData(y))
        n2 = addchild(p, NodeData(:*))
        n3 = addchild(n2, NodeData(0))
        n4 = addchild(n2, NodeData(t))
        @test InfiniteOpt._drop_zeros!(p) == p
        @test p.child.sibling.data.value == 0
    end
    # test isequal for Nodes
    @testset "Base.isequal (Nodes)" begin
        # test simple case
        n1 = Node(NodeData(1))
        n2 = Node(NodeData(2))
        @test !isequal(n1, n2)
        n2.data = NodeData(1)
        @test isequal(n1, n2)
        # test unequal count case
        n1 = Node(NodeData(:+))
        addchild(n1, NodeData(1))
        n2 = Node(NodeData(:+))
        addchild(n2, NodeData(1))
        addchild(n2, NodeData(1))
        @test !isequal(n1, n2)
        # test complicated isequal case
        n1 = Node(NodeData(:+))
        addchild(n1, NodeData(1))
        addchild(n1, NodeData(y))
        n2 = Node(NodeData(:+))
        addchild(n2, NodeData(1))
        addchild(n2, NodeData(y))
        @test isequal(n1, n2)
        # test more complicated not equal case
        n2 = Node(NodeData(:+))
        addchild(n2, NodeData(1))
        addchild(n2, NodeData(t))
        @test !isequal(n1, n2)
    end
    # test NLPExpr
    @testset "NLPExpr" begin
        @test NLPExpr(Node(NodeData(0))).tree_root.data.value == 0
    end
    # test Base basics
    @testset "Base Basics (NLPExpr)" begin
        # setup expr
        p = Node(NodeData(:+))
        addchild(p, NodeData(y))
        addchild(p, NodeData(2))
        nlp = NLPExpr(p)
        # test broadcastable
        @test Base.broadcastable(nlp) isa Base.RefValue
        # test copy
        @test copy(nlp) !== nlp
        @test isequal(copy(nlp), nlp)
        # test zero and one
        @test zero(NLPExpr).tree_root.data.value == 0
        @test one(NLPExpr).tree_root.data.value == 1
        # test isequal
        p = Node(NodeData(:+))
        addchild(p, NodeData(y))
        addchild(p, NodeData(3))
        nlp2 = NLPExpr(p)
        @test !isequal(nlp, nlp2)
        p.child.sibling.data = NodeData(2)
        @test isequal(nlp, nlp2)
    end
    # test drop_zeros!
    @testset "drop_zeros!" begin
        # setup expr
        p = Node(NodeData(:^))
        n1 = addchild(p, NodeData(y))
        n2 = addchild(p, NodeData(:*))
        n3 = addchild(n2, NodeData(0))
        n4 = addchild(n2, NodeData(t))
        nlp = NLPExpr(p)
        @test drop_zeros!(nlp) === nlp
        @test nlp.tree_root.child.sibling.data.value == 0
    end
    # test isequal_canonical
    @testset "isequal_canonical" begin
        # test that they are equal
        p = Node(NodeData(:^))
        n1 = addchild(p, NodeData(y))
        n2 = addchild(p, NodeData(:*))
        n3 = addchild(n2, NodeData(0))
        n4 = addchild(n2, NodeData(t))
        nlp1 = NLPExpr(p)
        p = Node(NodeData(:^))
        n1 = addchild(p, NodeData(y))
        n2 = addchild(p, NodeData(0))
        nlp2 = NLPExpr(p)
        @test isequal_canonical(nlp1, nlp2)
        # test that they are not equal
        p = Node(NodeData(:^))
        n1 = addchild(p, NodeData(y))
        n2 = addchild(p, NodeData(t))
        nlp3 = NLPExpr(p)
        @test !isequal_canonical(nlp1, nlp3)
        @test !isequal_canonical(nlp2, nlp3)
    end
    # test tree printing
    @testset "Tree Printing" begin
        # setup
        p = Node(NodeData(:^))
        n1 = addchild(p, NodeData(y))
        n2 = addchild(p, NodeData(2))
        nlp = NLPExpr(p)
        # test with NLPExpre
        expected = "^\n├─ y(t)\n└─ 2\n"
        io_test(print_expression_tree, expected, nlp)
        # test with other
        io_test(print_expression_tree, "y(t)\n", y)
        # test again
        test_output = @capture_out print_expression_tree(nlp)
        @test test_output == expected
    end
    # test ast mapping
    @testset "ast mapping" begin
        @variable(Model(), x)
        # map tree of vars and constants
        p = Node(NodeData(:^))
        n1 = addchild(p, NodeData(y))
        n2 = addchild(p, NodeData(:*))
        addchild(n2, NodeData(t))
        addchild(n2, NodeData(0.2))
        nlp = NLPExpr(p)
        @test map_nlp_to_ast(v -> x, nlp) == :($x ^ ($x * 0.2))
        # map tree with affine expression
        p = Node(NodeData(:^))
        n1 = addchild(p, NodeData(2y + t - 42))
        n2 = addchild(p, NodeData(:*))
        addchild(n2, NodeData(t))
        addchild(n2, NodeData(0.2))
        nlp = NLPExpr(p)
        @test map_nlp_to_ast(v -> x, nlp) == :((2 * $x + $x + -42) ^ ($x * 0.2))
        # map tree with quadratic expression (no affine)
        p = Node(NodeData(:^))
        n1 = addchild(p, NodeData(2y^2 + t^2))
        n2 = addchild(p, NodeData(:*))
        addchild(n2, NodeData(t))
        addchild(n2, NodeData(0.2))
        nlp = NLPExpr(p)
        @test map_nlp_to_ast(v -> x, nlp) == :((2 * $x * $x + $x * $x) ^ ($x * 0.2))
        # map tree with quadratic expression (w/ affine)
        p = Node(NodeData(:^))
        n1 = addchild(p, NodeData(2y^2 + t^2 - 42))
        n2 = addchild(p, NodeData(:*))
        addchild(n2, NodeData(t))
        addchild(n2, NodeData(0.2))
        nlp = NLPExpr(p)
        @test map_nlp_to_ast(v -> x, nlp) == :((2 * $x * $x + $x * $x + -42) ^ ($x * 0.2))
    end
end

# Test the basic expression generation via operators
@testset "Operator Definition" begin
    # setup model data
    m = InfiniteModel()
    @infinite_parameter(m, t in [0, 1])
    @variable(m, y, Infinite(t))
    @variable(m, z)
    # test _process_child_input
    @testset "_process_child_input" begin
        # nlp
        p = Node(NodeData(:sin))
        addchild(p, NodeData(y))
        nlp = NLPExpr(p)
        @test InfiniteOpt._process_child_input(nlp) == p
        # expressions
        @test InfiniteOpt._process_child_input(y) == NodeData(y)
        @test isequal_canonical(InfiniteOpt._process_child_input(2y).value, 2y)
        @test isequal_canonical(InfiniteOpt._process_child_input(y^2).value, y^2)
        # function symbol
        @test InfiniteOpt._process_child_input(:^) == NodeData(:^)
        # constant
        @test InfiniteOpt._process_child_input(true) == NodeData(true)
        @test InfiniteOpt._process_child_input(3) == NodeData(3)
        # fallback
        @test_throws ErrorException InfiniteOpt._process_child_input("bad")
    end
    # test the basic graph builder
    @testset "_call_graph" begin
        # test simple
        p = Node(NodeData(:*))
        addchild(p, NodeData(y))
        addchild(p, NodeData(t))
        addchild(p, NodeData(1))
        @test isequal(InfiniteOpt._call_graph(:*, y, t, 1), p)
        # test with other expression
        p = Node(NodeData(:sin))
        addchild(p, NodeData(y))
        nlp = NLPExpr(p)
        p = Node(NodeData(:^))
        addchild(p, nlp.tree_root)
        addchild(p, NodeData(y^2))
        @test isequal(InfiniteOpt._call_graph(:^, nlp, y^2), p)
    end
    # test the sum
    @testset "sum" begin
        # test empty sums
        if Base.VERSION >= v"1.6"
            @test sum(i for i in Int[]) == 0
        end
        @test isequal(sum(NLPExpr[]), zero(NLPExpr))
        @test isequal(sum(GeneralVariableRef[]),
                      zero(GenericAffExpr{Float64, GeneralVariableRef}))
        # test NLP sums
        p = Node(NodeData(:sin))
        addchild(p, NodeData(y))
        nlp = NLPExpr(p)
        new = NLPExpr(InfiniteOpt._call_graph(:+, nlp, nlp))
        @test isequal(sum(i for i in [nlp, nlp]), new)
        @test isequal(sum([nlp, nlp]), new)
        # test other expressions
        @test isequal(sum(i for i in [y, t]), y + t)
        @test isequal(sum([y, t]), y + t)
        # normal sums
        @test sum(i for i in 1:3) == 6
    end
    # test the product
    @testset "prod" begin
        # test empty sums
        if Base.VERSION >= v"1.6"
            @test prod(i for i in Int[]) == 1
        end
        @test isequal(prod(NLPExpr[]), one(NLPExpr))
        # test NLP sums
        p = Node(NodeData(:sin))
        addchild(p, NodeData(y))
        nlp = NLPExpr(p)
        new = NLPExpr(InfiniteOpt._call_graph(:*, nlp, nlp))
        @test isequal(prod(i for i in [nlp, nlp]), new)
        @test isequal(prod([nlp, nlp]), new)
        # test other expressions
        new = NLPExpr(InfiniteOpt._call_graph(:*, y, t, z))
        @test isequal(prod(i for i in [y, t, z]), new)
        @test isequal(prod([y, t, z]), new)
        # test normal products
        @test prod(i for i in 1:3) == 6
    end
    # prepare the test grid
    aff = 2z - 2
    quad = y^2 + y
    p = Node(NodeData(:sin))
    addchild(p, NodeData(y))
    nlp = NLPExpr(p)
    # test the multiplication operator
    @testset "Multiplication Operator" begin
        for (i, iorder) in [(42, 0), (y, 1), (aff, 1), (quad, 2), (nlp, 3)]
            for (j, jorder) in [(42, 0), (y, 1), (aff, 1), (quad, 2), (nlp, 3)]
                if iorder + jorder >= 3
                    expected = NLPExpr(InfiniteOpt._call_graph(:*, i, j))
                    @test isequal(i * j, expected)
                end
            end
        end
        expected = NLPExpr(InfiniteOpt._call_graph(:*, y, t, z,nlp))
        @test isequal(y * t * z * nlp, expected)
        @test isequal(*(nlp), nlp)
    end
    # test division operator
    @testset "Division Operator" begin
        for i in [42, y, aff, quad, nlp]
            for j in [y, aff, quad, nlp]
                expected = NLPExpr(InfiniteOpt._call_graph(:/, i, j))
                @test isequal(i / j, expected)
            end
        end
        expected = NLPExpr(InfiniteOpt._call_graph(:/, nlp, 42))
        @test isequal(nlp / 42, expected)
        @test isequal(nlp / 1, nlp)
        @test_throws ErrorException nlp / 0
    end
    # test the power operator
    @testset "Power Operator" begin
        for i in [42, y, aff, quad, nlp]
            for j in [y, aff, quad, nlp]
                expected = NLPExpr(InfiniteOpt._call_graph(:^, i, j))
                @test isequal(i ^ j, expected)
            end
        end
        one_aff = one(JuMP.GenericAffExpr{Float64, GeneralVariableRef})
        for i in [y, aff, quad, nlp]
            for f in [Float64, Int]
                @test isequal(i^zero(f), one_aff)
                @test isequal(i^one(f), i)
                if i isa NLPExpr
                    expected = NLPExpr(InfiniteOpt._call_graph(:^, i, f(2)))
                    @test isequal(i ^ f(2), expected)
                else
                    @test isequal(i^f(2), i * i)
                end
                expected = NLPExpr(InfiniteOpt._call_graph(:^, i, f(3)))
                @test isequal(i ^ f(3), expected)
            end
        end
        # extra tests
        @test isequal(y^0, one_aff)
        @test isequal(y^1, y)
        @test isequal(y^2, y * y)
        @test isequal(y^0.0, one_aff)
        @test isequal(y^1.0, y)
        @test isequal(y^2.0, y * y)
    end
    # test the subtraction operator
    @testset "Subtraction Operator" begin
        for i in [42, y, aff, quad, nlp]
            expected = NLPExpr(InfiniteOpt._call_graph(:-, nlp, i))
            @test isequal(nlp - i, expected)
            expected = NLPExpr(InfiniteOpt._call_graph(:-, i, nlp))
            @test isequal(i - nlp, expected)
        end
        expected = NLPExpr(InfiniteOpt._call_graph(:-, nlp))
        @test isequal(-nlp, expected)
        @test isequal(y - y, zero(JuMP.GenericAffExpr{Float64, GeneralVariableRef}))
        @test (y - z).constant == 0
        @test (y - z).terms[y] == 1
        @test (y - z).terms[z] == -1
    end
    # test the addition operator
    @testset "Addition Operator" begin
        for i in [42, y, aff, quad, nlp]
            expected = NLPExpr(InfiniteOpt._call_graph(:+, nlp, i))
            @test isequal(nlp + i, expected)
            expected = NLPExpr(InfiniteOpt._call_graph(:+, i, nlp))
            @test isequal(i + nlp, expected)
        end
        @test isequal(+nlp, nlp)
    end
    # test the comparison operators
    @testset "Comparison Operators" begin
        for i in [42, y, aff, quad, nlp]
            for j in [42, y, aff, quad, nlp]
                for (n, f) in (:< => Base.:(<), :(==) => Base.:(==),
                               :> => Base.:(>), :<= => Base.:(<=),
                               :>= => Base.:(>=))
                    if !(isequal(i, 42) && isequal(j, 42)) && !(isequal(i, y) && isequal(j, y))
                        expected = NLPExpr(InfiniteOpt._call_graph(n, i, j))
                        @test isequal(f(i, j), expected)
                    end
                end
            end
        end
        for (n, f) in (:(==) => Base.:(==), :<= => Base.:(<=), :>= => Base.:(>=))
            expected = NLPExpr(InfiniteOpt._call_graph(n, y, z))
            @test isequal(f(y, z), expected)
            @test f(y, y)
        end
        for (n, f) in (:< => Base.:(<), :> => Base.:(>))
            expected = NLPExpr(InfiniteOpt._call_graph(n, y, z))
            @test isequal(f(y, z), expected)
            @test !f(y, y)
        end
        # extra tests
        @test (y == 0) isa NLPExpr
        @test (y <= 0) isa NLPExpr
        @test (y >= 0) isa NLPExpr
        @test (y > 0) isa NLPExpr
        @test (y < 0) isa NLPExpr
        @test (0 == y) isa NLPExpr
        @test (0 <= y) isa NLPExpr
        @test (0 >= y) isa NLPExpr
        @test (0 > y) isa NLPExpr
        @test (0 < y) isa NLPExpr
    end
    # test the logic operators
    @testset "Logic Operators" begin
        for i in [y, nlp]
            for j in [y, nlp]
                for (n, f) in (:&& => Base.:&, :|| => Base.:|)
                    expected = NLPExpr(InfiniteOpt._call_graph(n, i, j))
                    @test isequal(f(i, j), expected)
                end
            end
        end
        for i in [y, nlp]
            @test !(i & false)
            @test !(false & i)
            @test isequal(i & true, i)
            @test isequal(true & i, i)
            @test (i | true)
            @test (true | i)
            @test isequal(i | false, i)
            @test isequal(false | i, i)
        end
    end
    # test ifelse
    @testset "ifelse" begin
        for i in [42, y, aff, quad, nlp]
            for j in [42, y, aff, quad, nlp]
                expected = NLPExpr(InfiniteOpt._call_graph(:ifelse, nlp, i, j))
                @test isequal(InfiniteOpt.ifelse(nlp, i, j), expected)
            end
        end
        @test isequal(InfiniteOpt.ifelse(true, y, z), y)
        @test isequal(InfiniteOpt.ifelse(false, y, z), z)
    end
    # test the default functions
    @testset "Default Registered Functions" begin
        for i in [y, aff, quad, nlp]
            for (n, f) in InfiniteOpt._Base1ArgFuncList
            expected = NLPExpr(InfiniteOpt._call_graph(n, i))
            @test isequal(f(i), expected)
            end
        end
        for i in [y, aff, quad, nlp]
            for (n, f) in InfiniteOpt._Special1ArgFuncList
            expected = NLPExpr(InfiniteOpt._call_graph(n, i))
            @test isequal(f(i), expected)
            end
        end
        for i in [42, y, aff, quad, nlp]
            for j in [42, y, aff, quad, nlp]
                for (n, f) in [(:max, max), (:min, min)]
                    if !(isequal(i, 42) && isequal(j, 42))
                        expected = NLPExpr(InfiniteOpt._call_graph(n, i, j))
                        @test isequal(f(i, j), expected)
                    end
                end
            end
        end
    end
end

# Test the MutableArithmetics stuff
@testset "MutableArithmetics" begin
    # setup model data
    m = InfiniteModel()
    @infinite_parameter(m, t in [0, 1])
    @variable(m, y, Infinite(t))
    @variable(m, z)
    aff = 2z - 2
    quad = y^2 + y
    nlp = sin(y)
    # test MA.mutability
    @testset "MA.mutability" begin
        @test MA.mutability(NLPExpr) == MA.IsMutable()
    end
    # test MA.promote_operation
    @testset "MA.promote_operation" begin
        for i in (2.0, 2, y, aff, quad, nlp)
            for f in (+, -, *, /, ^)
                @test MA.promote_operation(f, typeof(i), NLPExpr) == NLPExpr
                @test MA.promote_operation(f, NLPExpr, typeof(i)) == NLPExpr
            end
        end
        for i in (y, aff, quad)
            for f in (*, /, ^)
                @test MA.promote_operation(f, typeof(i), typeof(quad)) == NLPExpr
                @test MA.promote_operation(f, typeof(quad), typeof(i)) == NLPExpr
            end
        end
        for i in (2, 2.0, y, aff, quad, nlp)
            for j in (y, aff, quad, nlp)
                for f in (/, ^)
                    @test MA.promote_operation(f, typeof(i), typeof(j)) == NLPExpr
                end
            end
        end
    end
    # test MA.scaling
    @testset "MA.scaling" begin
        @test_throws ErrorException MA.scaling(nlp)
        @test MA.scaling(one(NLPExpr)) == 1
    end
    # test mutable_copy
    @testset "MA.mutable_copy" begin
        @test MA.mutable_copy(nlp) === nlp
    end
    # test mutable_operate!
    @testset "MA.mutable_operate!" begin
        # test zero and one
        @test MA.mutable_operate!(zero, nlp) !== nlp
        @test isequal(MA.mutable_operate!(zero, nlp), zero(NLPExpr))
        @test MA.mutable_operate!(one, nlp) !== nlp
        @test isequal(MA.mutable_operate!(one, nlp), one(NLPExpr))
        # test operators
        for i in (2.0, 2, y, aff, quad, nlp)
            for f in (+, -, *, /, ^)
                @test isequal(MA.mutable_operate!(f, nlp, i), f(nlp, i))
                @test isequal(MA.mutable_operate!(f, i, nlp), f(i, nlp))
            end
        end
        # test AddSubMul
        @test isequal(MA.mutable_operate!(MA.add_mul, nlp, y, 2), nlp + y * 2)
        @test isequal(MA.mutable_operate!(MA.sub_mul, nlp, y, 2), nlp - y * 2)
    end
end

# Test LinearAlgebra stuff
@testset "Linear Algebra" begin
    # setup the model data
    m = InfiniteModel()
    @infinite_parameter(m, t in [0, 1])
    @variable(m, y, Infinite(t))
    @variable(m, z)
    aff = 2z - 2
    quad = y^2 + y
    nlp = sin(y)
    # test promotions
    @testset "Base.promote_rule" begin
        for i in (3, y, aff, quad)
            @test [nlp, i] isa Vector{NLPExpr}
        end
        # extra tests
        @test promote_rule(NLPExpr, typeof(quad)) == NLPExpr
    end
    # test dot
    @testset "LinearAlgebra.dot" begin
        for i in (2, y, aff, quad, nlp)
            for j in (2, y, aff, quad, nlp)
                @test isequal(dot(i, j), i * j)
            end
        end
    end
    # test linear algebra operations
    @testset "Operations" begin
        # make variables
        @variable(m, A[1:2, 1:3])
        @variable(m, x[1:2])
        @variable(m, w[1:3])
        # test expression
        @test x' * A * w isa NLPExpr # TODO fix with improved MA extension
    end
end

# Test registration utilities
@testset "Registration Methods" begin
    # setup model data
    m = InfiniteModel()
    @variable(m, y)
    # test name_to_function
    @testset "name_to_function" begin
        @test name_to_function(m, :tan, 1) == tan
        @test name_to_function(m, :+, 10) == +
        @test name_to_function(m, :*, 2) == *
    end
    # test all_registered_functions
    @testset "all_registered_functions" begin
        @test all_registered_functions(m) isa Vector{Function}
    end
    # test user_registered_functions
    @testset "user_registered_functions" begin
        @test user_registered_functions(m) == RegisteredFunction[]
    end
    # define functions for tests
    f(a) = a^3
    g(a::Int) = 42
    h(a, b) = 42
    f1(a) = 32
    f2(a) = 10
    h1(a, b) = 13
    function hg(v, a, b)
        v[1] = 1
        v[2] = 2
        return
    end
    # test creation helper errors
    @testset "Registration Helpers" begin
        @test_throws ErrorException RegisteredFunction(:a, 1, f, g)
        @test_throws ErrorException RegisteredFunction(:a, 2, f, g)
        @test_throws ErrorException RegisteredFunction(:a, 1, f, g, f)
        @test_throws ErrorException RegisteredFunction(:a, 1, f, f, g)
        @test_throws ErrorException InfiniteOpt._register(error, Main, m, :f, 1, 1)
        @test_throws ErrorException InfiniteOpt._register(error, Main, m, :sin, 1, sin)
        @test_throws ErrorException InfiniteOpt._register(error, Main, m, :g, 1, g)
        @test_throws ErrorException InfiniteOpt._register(error, Main, m, :eta, 1, eta)
    end
    # test @register
    @testset "@register" begin
        # test errors
        @test_macro_throws ErrorException @register(Model(), f(a))
        @test_macro_throws ErrorException @register(m, f(a), bad = 42)
        @test_macro_throws ErrorException @register(m, f(a), g, h, 3)
        @test_macro_throws ErrorException @register(m, f[i](a))
        @test_macro_throws ErrorException @register(m, f(a::Int))
        @test_macro_throws ErrorException @register(m, f(a), 2)
        @test_macro_throws ErrorException @register(m, g(a))
        @test_macro_throws ErrorException @register(m, eta(a))
        @test_macro_throws ErrorException @register(m, f(a), g)
        @test_macro_throws ErrorException @register(m, f(a), g, g)
        @test_macro_throws ErrorException @register(m, f(a), f, g)
        @test_macro_throws ErrorException @register(m, h(a, b), h)
        @test_macro_throws ErrorException @register(m, sin(a))
        # test univariate function with no gradient or hessian
        @test @register(m, f(a)) isa Function
        @test f(y) isa NLPExpr
        @test f(y).tree_root.data.value == :f
        @test isequal(f(y).tree_root.child.data.value, y)
        @test f(2) == 8
        @test length(user_registered_functions(m)) == 1
        @test user_registered_functions(m)[1].name == :f
        @test user_registered_functions(m)[1].func == f
        @test user_registered_functions(m)[1].num_args == 1
        @test user_registered_functions(m)[1].gradient isa Nothing
        @test user_registered_functions(m)[1].hessian isa Nothing
        @test name_to_function(m, :f, 1) == f
        # test univariate function with gradient
        @test @register(m, f1(a), f) isa Function
        @test f1(y) isa NLPExpr
        @test f1(y).tree_root.data.value == :f1
        @test isequal(f1(y).tree_root.child.data.value, y)
        @test f1(2) == 32
        @test length(user_registered_functions(m)) == 2
        @test user_registered_functions(m)[2].name == :f1
        @test user_registered_functions(m)[2].func == f1
        @test user_registered_functions(m)[2].num_args == 1
        @test user_registered_functions(m)[2].gradient == f
        @test user_registered_functions(m)[2].hessian isa Nothing
        @test name_to_function(m, :f1, 1) == f1
        # test univariate function with gradient and hessian
        @test @register(m, f2(a), f, f1) isa Function
        @test f2(y) isa NLPExpr
        @test f2(y).tree_root.data.value == :f2
        @test isequal(f2(y).tree_root.child.data.value, y)
        @test f2(2) == 10
        @test length(user_registered_functions(m)) == 3
        @test user_registered_functions(m)[3].name == :f2
        @test user_registered_functions(m)[3].func == f2
        @test user_registered_functions(m)[3].num_args == 1
        @test user_registered_functions(m)[3].gradient == f
        @test user_registered_functions(m)[3].hessian == f1
        @test name_to_function(m, :f2, 1) == f2
        # test multivariate function with no gradient
        @test @register(m, h(a, b)) isa Function
        @test h(2, y) isa NLPExpr
        @test h(2, y).tree_root.data.value == :h
        @test isequal(h(2, y).tree_root.child.data.value, 2)
        @test isequal(h(2, y).tree_root.child.sibling.data.value, y)
        @test h(y, y) isa NLPExpr
        @test h(y, 2) isa NLPExpr
        @test h(2, 2) == 42
        @test length(user_registered_functions(m)) == 4
        @test user_registered_functions(m)[4].name == :h
        @test user_registered_functions(m)[4].func == h
        @test user_registered_functions(m)[4].num_args == 2
        @test user_registered_functions(m)[4].gradient isa Nothing
        @test user_registered_functions(m)[4].hessian isa Nothing
        @test name_to_function(m, :h, 2) == h
        # test multivariate function with gradient
        @test @register(m, h1(a, b), hg) isa Function
        @test h1(2, y) isa NLPExpr
        @test h1(2, y).tree_root.data.value == :h1
        @test isequal(h1(2, y).tree_root.child.data.value, 2)
        @test isequal(h1(2, y).tree_root.child.sibling.data.value, y)
        @test h1(y, y) isa NLPExpr
        @test h1(y, 2) isa NLPExpr
        @test h1(2, 2) == 13
        @test length(user_registered_functions(m)) == 5
        @test user_registered_functions(m)[5].name == :h1
        @test user_registered_functions(m)[5].func == h1
        @test user_registered_functions(m)[5].num_args == 2
        @test user_registered_functions(m)[5].gradient == hg
        @test user_registered_functions(m)[5].hessian isa Nothing
        @test name_to_function(m, :h1, 2) == h1
        # test wrong model error
        @test_throws ErrorException f(@variable(InfiniteModel()))
        # test functional registration
        function registration_test()
            mt = InfiniteModel()
            @variable(mt, x)
            q(a) = 1
            @test @register(mt, q(a)) isa Function
            @test q(x) isa NLPExpr
            return
        end
        @test registration_test() isa Nothing
        @test registration_test() isa Nothing
    end
    # test add_registered_to_jump
    @testset "add_registered_to_jump" begin
        # test normal
        m1 = Model()
        @test add_registered_to_jump(m1, m) isa Nothing
        @test m1.nlp_data.largest_user_input_dimension == 2
        r1 = m1.nlp_data.user_operators
        @test length(keys(r1.univariate_operator_to_id)) == 3
        @test r1.univariate_operator_f == [f, f1, f2]
        @test r1.univariate_operator_fprime[2:end] == [f, f]
        @test r1.univariate_operator_fprimeprime[3] == f1
        @test length(r1.multivariate_operator_to_id) == 2
        # test error
        m2 = Model()
        h2(a, b) = 3
        @test @register(m, h2(a, b), hg, f) isa Function
        @test_throws ErrorException add_registered_to_jump(m2, m) isa Nothing
        r2 = m2.nlp_data.user_operators
        @test length(keys(r2.univariate_operator_to_id)) == 3
        @test r2.univariate_operator_f == [f, f1, f2]
        @test r2.univariate_operator_fprime[2:end] == [f, f]
        @test r2.univariate_operator_fprimeprime[3] == f1
        @test length(r2.multivariate_operator_to_id) == 2
    end
end

# Test string methods
@testset "String Methods" begin
    # setup the model data
    m = InfiniteModel()
    @infinite_parameter(m, t in [0, 1])
    @variable(m, y, Infinite(t))
    @variable(m, z)
    # test making strings
    @testset "String Creation" begin
        # test some simple ones
        @test string(sin(y) + 2) == "sin(y(t)) + 2"
        @test string((z*z) ^ 4) == "(z²)^4"
        @test string((cos(z) + sin(z)) / y) == "(cos(z) + sin(z)) / y(t)"
        @test string(-cos(z + y) * z^2.3) == "-cos(z + y(t)) * z^2.3"
        @test string((-max(0, z)) ^ 3) == "(-(max(0, z))^3"
        # test AffExpr cases
        aff0 = zero(GenericAffExpr{Float64, GeneralVariableRef})
        @test string(aff0^4 + (2z - 3y + 42) ^ (-1z)) == "0^4 + (2 z - 3 y(t) + 42)^(-z)"
        @test string((1z) / (2.3 * y)) == "z / (2.3 y(t))"
        # test QuadExpr cases
        quad0 = zero(GenericQuadExpr{Float64, GeneralVariableRef})
        @test string(quad0 / (z^2 + y * z)) == "0 / (z² + y(t)*z)"
        @test string((0z^2 + 42) * sin(y)) == "42 * sin(y(t))"
        @test string((z^2) / y) == "(z²) / y(t)"
    end
end


text:
"""

## linmul_sky!

This subroutine forms the product of symmetric matrix stored as
a skyline and a vector.

### Method
```julia
linmul_sky!(kv, disps, loads, kdiag)
```

### Arguments
```julia
* kv::Vector{Float64}       : Sparse stiffnes matrix (Skyline format)
* disps::Vector{Float64}    : Displacements
* loads::Vector{Float64}    : Loads vector (Updated)
* kdiag::Vector{Int}      : Bandwidth vector
```
"""
function linmul_sky!(kv::Vector{Float64}, disps::Vector{Float64},
  loads::Vector{Float64}, kdiag::Vector{Int})
  #
  # This subroutine forms the product of symmetric matrix stored as
  # a skyline and a vector.
  #
  # Updates loads[:]
  #
 n = size(kdiag, 1)
 for i in 1:n
   x = 0.0
   lup = kdiag[i]
   if i == 1
     low = lup
   end
   if i !== 1
     low = kdiag[i-1] + 1
   end
   for j in low:lup
     x += kv[j] * disps[i+j-lup]
   end
   loads[i] = x
   i == 1 && continue
   lup = lup - 1
   for j in low:lup
     k = i + j - lup - 1
     loads[k] += kv[j] * disps[i]
   end
 end
 loads
end


text:
function __add_node!(m::AbstractMultiScaleArray, node::AbstractMultiScaleArray)
    push!(m.end_idxs, m.end_idxs[end] + length(node) + length(m.values))
    push!(m.nodes, node)
    nothing
end

function __update_lengths(m::AbstractMultiScaleArray, beg, len)
    for j = beg:num_nodes(m)
        m.end_idxs[j] += len
    end
    isempty(m.values) || (m.end_idxs[end] += len)
    nothing
end

function __add_node!(m::AbstractMultiScaleArray, node::AbstractMultiScaleArray, i::Int)
    __add_node!(m.nodes[i], node)
    __update_lengths(m, i, length(node))
end

function __add_node!(m::AbstractMultiScaleArray, node::AbstractMultiScaleArray, i, I::Int...)
    __add_node!(m.nodes[i], node, I...)
    __update_lengths(m, i, length(node))
end

add_node!(m::AbstractMultiScaleArrayHead, node::AbstractMultiScaleArray) =
    __add_node!(m, node)

add_node!(m::AbstractMultiScaleArrayHead, node::AbstractMultiScaleArray, i::Int) =
    __add_node!(m, node, i)

add_node!(m::AbstractMultiScaleArrayHead, node::AbstractMultiScaleArray, i, I::Int...) =
    __add_node!(m, node, i, I...)

function __remove_node!(m::AbstractMultiScaleArray, i::Int)
    del_length = length(m.nodes[i])
    deleteat!(m.nodes, i)
    deleteat!(m.end_idxs, i)
    for j = i:num_nodes(m)
        m.end_idxs[j] -= del_length
    end
    isempty(m.values) || (m.end_idxs[end] -= del_length)
    del_length
end

function __remove_node!(m::AbstractMultiScaleArrayLeaf, i::Int)
    deleteat!(m.nodes, i)
    1
end

function remove_node!(m::AbstractMultiScaleArrayHead, i, I::Int...)
    del_length = __remove_node!(m.nodes[i], I...)
    for j = i:num_nodes(m)
        m.end_idxs[j] -= del_length
    end
    isempty(m.values) || (m.end_idxs[end] -= del_length)
    if size(m.nodes[i].nodes) == (0, )
        deleteat!(m.nodes, i)
        deleteat!(m.end_idxs, i)
    end
    nothing
end

function __remove_node!(m::AbstractMultiScaleArray, i, I::Int...)
    del_length = __remove_node!(m.nodes[i], I...)
    for j = i:num_nodes(m)
        m.end_idxs[j] -= del_length
    end
    isempty(m.values) || (m.end_idxs[end] -= del_length)
    isempty(m.nodes[i]) && deleteat!(m.nodes, i)
    del_length
end

remove_node!(m::AbstractMultiScaleArrayHead, i::Int) =
    (__remove_node!(m, i); nothing)


text:
push!(LOAD_PATH, "../src")

using CASIM

render_option = Not_Render()
sim_horizon = 5000
num_run = 3
init_num_ac = 100
const_num_ac = [150, 200, 250, 300, 350, 400, 450]

"""
	Corrected_VICAS
	num_ac
"""
exp = "ConstNumAC_CorrectedVICAS_10km_init_ac_" * string(init_num_ac) * "/" * correction_version * "_final"
if !ispath("../simulation_results/ConstNumAC_CorrectedVICAS_10km_init_ac_" * string(init_num_ac))
	mkdir("../simulation_results/ConstNumAC_CorrectedVICAS_10km_init_ac_" * string(init_num_ac))
end
cas = Corrected_VICAS()
cas_ctrl = Uniform_CAS(cas)
sensor = Multi_Threat_Tracker()
sensor_ctrl = Uniform_Sensor(sensor)

for run in 1:num_run
	for num_ac in const_num_ac
		scheme = "Corrected_num_ac_" * string(Int(num_ac)) * "_run_" * string(run)
		airspace_ctrl = Constant_Number_Controller(num_ac)
		simulate(exp, scheme, sim_horizon, airspace_ctrl, cas_ctrl, sensor_ctrl, render_option, correction=true)
	end
end

max_stars_repo_name: parachutel/DeepCorrectionCAS