text:
#=

Augmented Lagrangian solver.

It is not abstracted as to handle multiple constraints

=#

using LinearAlgebra

include("../auglag/auglag-core.jl")
include("../auglag/solverParams.jl")

# include("..\\constraints\\constraintManager.jl")

include("backtrackLineSearch.jl")
include("trustRegion.jl")

@doc raw"""
    newtonStepALP(y0::SOCP_primals, al::augLagQP_2Cone, delta = 1,
                      gamma = 1.5, epsilon = 0.5, condNumMax = 1e5, a = 1e-6)

Evaluates the Newton Step for an Augmented Lagrangian of an SOCP.

```math
y ← y - [H(φ(y)) + λI]^{-1} ∇φ(y)
```

This function internally calculates the step `dk` using Cholesky Decomposition.
```math
dk = [H(φ(y)) + λI] \ ∇φ(y)
```

## Core Arguments
- `y0::SOCP_primals`: The primal vectors of the SOCP.
- `al::augLagQP_2Cone`: The SOCP Augmented Lagrangian.
See: [`SOCP_primals`](@ref), [`augLagQP_2Cone`](@ref)

## Trust Region and Regularization Arguements
- `delta`: Current trust region size
- `gamma`: Parameter (> 1) for how aggressive to increase the trust region size
- `epsilon`: Parameter (> 0) weighes the gradient versus the hessian in
determining the damping ratio size
- `condNumMax`: The max condition number on the damped Hessian
- `a`: Parameter (> 0) How quickly to damp the Hessian

See: [`findDamping`](@ref) for more information

## Return Values
returns `dk`, `phiD`, `rCho`, `damping` which are
- `dk`: the Newton step
- `phiD`: the residual (``∇φ(y)``)
- `rCho`: the Cholesky decomposition of `H(φ(y)) + λI`
- `damping`: the damping parameter (`λ` above)

See: [`findDamping`](@ref) for more information

"""
function newtonStepALP(y0, al::augLag, delta = 1,
                    gamma = 1.5, epsilon = 0.5, condNumMax = 1e5, a = 1e-6)
    hess = evalHessAl(al, y0)
    phiD = evalGradAL(al, y0)

    if typeof(al.cM) == constraintManager_Dynamics
        dualSize = size(al.cM.affineLambdaList, 1)
    else
        dualSize = 0
    end

    damping, dk, rCho = findDamping(hess, phiD, dualSize, delta, gamma, epsilon,
                                            condNumMax, a)

    if norm(dk) > delta
        dk = (delta / norm(dk)) * dk
    end

    return dk, phiD, rCho, damping

end

"""
    newtonTRLS_ALP(y0::SOCP_primals, al::augLagQP_2Cone, sp::solverParams)

Performs a maximum on `N` newton steps as specified by [`solverParams`](@ref).
May return early if the residual is lower than the tolerance specified by
[`solverParams`](@ref).

The function operates on an augmented lagrangian (AL) with primals (P) only
(as opposed to a Primal-Dual set-up) for second-order cone programs (SOCP).
The function uses a combined trust region (TR) and line search (LS) approach
as specified in Nocedal et Yuan (1998).

returns (yNewtStates, residNewt) which are:
- `yNewtStates`: An array with one entry of the primal vector per Newton step
- `residNewt`: An array with one entry of the residual vector per Newton step

See also: [`augLagQP_2Cone`](@ref), [`SOCP_primals`](@ref)
"""
function newtonTRLS_ALP(y0, al::augLag, sp::solverParams)
    yNewtStates = []
    residNewt = []
    # push!(yNewtStates, y0)
    yCurr = y0

    trustDelta = sp.trSizeStart

    # For printing
    early = false

    lineSearchObj(v) = evalAL(al, v)
    lineSearchdfdx(v) = evalGradAL(al, v)

    # Now, we run through the iterations
    for i in 1:(sp.maxNewtonSteps)

        if sp.verbose ≥ DEBUG
            println("Currently at $yCurr")
        end

        currentObjVal = lineSearchObj(yCurr)

        if sp.verbose ≥ ALL
            println("ϕ(y) = $currentObjVal")
            println("∇ϕ(y) = $(lineSearchdfdx(yCurr))")
            cCurr = evalConstraints(al.cM, yCurr, al.rho)
            println("Constraints: $cCurr")
        end

        # Take a Newton Step
        # Negative sign addressed above
        # (dk, residual, rCho, damp) = newtonStepALP(yCurr, al, trustDelta)
        (dk, residual, BDamp, damp) = newtonStepALP(yCurr, al, trustDelta)
        push!(residNewt, residual)
        # LCho = sparse(rCho.L)

        if sp.verbose ≥ ALL
            println("\nNewton Direction: $dk")
            println("AL: $al")
        end

        if sp.verbose ≥ INTERIM
            println("Damping of ($damp) and trust size of ($trustDelta)")
        end

        # Determine if the trust region is sufficient
        y0New = yCurr + dk
        baseObjVal = lineSearchObj(y0New)

        if sp.verbose ≥ SNAPSHOTS
            println("Objective from $currentObjVal → $baseObjVal")
        end

        if baseObjVal ≤ currentObjVal
            # Trust region was a success
            # Step 4 of algorithm 3.1 in the Nocedal et Yuan paper

            if sp.verbose ≥ INTERIM
                println("Trust Region Success")
            end

            # fObjApprox(d) = residual'd + (1/2) * d'*(LCho * LCho')*d
            fObjApprox(d) = residual'd + (1/2) * d'*(BDamp)*d

            if sp.verbose ≥ SNAPSHOTS
                println("fObjApprox(d) = $(fObjApprox(dk))")
            end

            rho = (currentObjVal - baseObjVal) / (-fObjApprox(dk)[1])
            roundingErrorTol = UInt8(round(-log10(sp.rTol)) / 2)

            if sp.verbose ≥ SNAPSHOTS
                println("Rho = $rho")
            end

            roundedDkNorm = round(norm(dk), digits = roundingErrorTol)
            roundedTrustDelta = round(trustDelta, digits = roundingErrorTol)

            if sp.verbose ≥ SNAPSHOTS
                print("||dk|| = $(norm(dk)) → $roundedDkNorm")
                print(" vs Δ = $trustDelta → $roundedTrustDelta")
                println(" at $roundingErrorTol digits")
            end

            if rho ≥ sp.trc2 && roundedDkNorm < roundedTrustDelta
                # ρ ≥ c2 AND ||dk|| < Δ
                # Condition 1: No change
                # trustDelta = trustDelta # Probably not needed
                if sp.verbose ≥ INTERIM
                    println("Trust Region Size unchanged. Still at $trustDelta")
                end
            elseif rho < sp.trc2
                # ρ < c2, but ||dk|| ≤ Δ
                trustDelta = (sp.trc3 * norm(dk) +
                                            sp.trc4 * trustDelta) / 2
                if sp.verbose ≥ INTERIM
                    println("Trust Region Size decreased. Now $trustDelta")
                end
            else
                # Implies that ρ ≥ c2 AND ||dk|| = Δ
                trustDelta = sp.trc1 * trustDelta

                if sp.verbose ≥ INTERIM
                    println("Trust Region Size increased. Now $trustDelta")
                end
            end

        else
            # Trust region failed.

            if sp.verbose ≥ INTERIM
                print("Trust Region Failed - ")
                println("Trying Line Search.")
            end

            # println("Attempting Line Search: ")
            lc = lineSearchObj(yCurr)
            # println("Base Fun lineSearchObj -> $(size(lc)) -> $lc")
            lgc = lineSearchdfdx(yCurr)
            # println("Base Grad lineSearchdfdx -> $(size(lgc)) ->")
            # display(lgc)
            rlgc = dk'lgc
            # println("Results include $(size(rlgc)) ->")
            # display(rlgc)

            # Attempt a linesearch
            # Get the line search recommendation
            y0New, stepLS = backtrackLineSearch(yCurr, dk,
                            lineSearchObj, lineSearchdfdx,
                            sp.paramA, sp.paramB,
                            (sp.verbose ≥ DEBUG))

            if sp.verbose ≥ SNAPSHOTS
                println("    Recommended Line Search Step: $stepLS")
                if stepLS < (sp.paramB ^ 3)
                    println("    Very low step.")
                end
            end

            if sp.verbose ≥ DEBUG
                print("Expected x = ")
                println("$y0New ?= $(primalVec(yCurr) + stepLS * dk)\n")
            end

            # Now we update the trust region
            yChange = norm(stepLS * dk)
            trustDelta = (yChange + sp.trc4 * trustDelta) / 2

        end

        # Save the new state to the output list
        push!(yNewtStates, y0New)

        if sp.verbose ≥ INTERIM
            println("Added State")#"$y0New\n")
        end

        # Break by tolerance and trust region size
        if (norm(residual, 2) < sp.rTol)
            if sp.verbose ≥ INTERIM
                println("Ended from tolerance at $i Newton steps\n")
            end

            early = true
            break
        end

        # Update the current state with the new state
        yCurr = y0New

        if typeof(al.cM) == constraintManager_Dynamics
            aSize = size(al.cM.affineLambdaList, 1)
            ySize = size(y0New, 1)
            al.cM.affineLambdaList = y0New[ySize - aSize + 1:end]
        end

    end

    if !early && sp.verbose ≥ INTERIM
        println("Ended from max steps in $(sp.maxNewtonSteps) Newton Steps\n")
    end

    if sp.verbose ≥ DEBUG
        println("Ended Newton Method at $yCurr")
        println("ϕ(y) = $(lineSearchObj(primalVec(yCurr)))")
        println("∇ϕ(y) = $(lineSearchdfdx(primalVec(yCurr)))")
        cCurr = evalConstraints(al.cM, yCurr, al.rho)
        println("Constraints: $cCurr\n")
    end

    return yNewtStates, residNewt

end

"""
    ALPrimalNewtonMain(y0, al::augLag, sp::solverParams)

Performs a maximum on `N` outer steps as specified by [`solverParams`](@ref).
May return early if the residual is lower than the tolerance specified by
[`solverParams`](@ref) and the constraint penalty is high enough.

The function operates on an augmented lagrangian (AL) with primals (P) only
(as opposed to a Primal-Dual set-up) for second-order cone programs (SOCP).
See [`newtonTRLS_ALP`](@ref) for more information on the Newton Steps

returns (yStates, residuals) which are:
- `yStates`: An array with one entry of the primal vector per Newton step
- `residuals`: An array with one entry of the residual vector per Newton step

See also: [`augLag`](@ref)
"""
function ALPrimalNewtonMain(y0, al::augLag, sp::solverParams)

    yStates = []
    residuals = []
    push!(yStates, y0)

    for i in 1:(sp.maxOuterIters)

        # Update x at each iteration
        if sp.verbose ≥ SNAPSHOTS
            println("\n--------------------------------------")
            println("Next Full Update starting at $y0")
        end

        (yNewStates, resAtStates) = newtonTRLS_ALP(y0, al, sp)

        # Take each step in the arrays above and save it to the respective
        # overall arrays
        yStates = [yStates; yNewStates]
        residuals = [residuals; resAtStates]

        # Determine the new lambda and rho
        # λ ← λ + ρ c(x_k*)       - Which is to say update with prior x*
        # ρ ← min(ρ * 10, 10^6)   - Which is to say we bound ρ's growth by 10^6
        yNewest = yNewStates[end]

        yNewPrimals = getPrimals(al, yNewest)

        cCurr = evalConstraints(al.cM, yNewPrimals, al.rho)

        # Only turn on if the solver is not returning all of the points.
        # if false
        #     println()
        #     println("yStates = $yStates")
        #     println("yNewest = $yNewest")
        #     println("All residuals = $residuals")
        # end
        if sp.verbose ≥ ALL
            println("\n################")
            println("Former Lambda: $(al.lambda)")
        end

        updateDual!(al.cM, yNewPrimals, al.rho)
        al.rho = clamp(al.rho * sp.penaltyStep, 0, sp.penaltyMax)

        if sp.verbose ≥ ALL
            println("New state added: (Newest) $yNewest")
            println("cCurr: $(cCurr)")
            println("Lambda Updated: $(al.lambda) vs. $lambdaNew")
            println("rho updated: $(al.rho)")
            println("################\n")
        end

        if (norm(residuals[end], 2) < sp.rTol) && (al.rho == sp.penaltyMax)
            if sp.verbose ≥ INTERIM
                println("Ended early at $i outer steps")
            end

            break
        else
            y0 = yNewest
        end
    end

    println("\nSOLVE COMPLETE\n")

    return yStates, residuals
end

max_stars_repo_name: danineamati/TrajOptSOCPs


text:
using JuMP, Base.Test, AmplNLWriter

# solver = AmplNLSolver(Ipopt.amplexe, ["print_level=0"])

# Note min and max not implemented in Couenne

## Solve test problem with simple min functions
 #
 #  max   min( x^2, x )
 #  s.t.  -0.5 <= x <= 0.5
 #
 #  The optimal objective value is 0.25.
 #      x = 0.5
##
@testset "example: maxmin" begin
    m = Model(solver=solver)
    @variable(m, -0.5 <= x <= 0.5, start = 0.25)
    @NLobjective(m, Max, min(x^2, 0.3, x))
    @test solve(m) == :Optimal
    @test isapprox(getobjectivevalue(m), 0.25, atol=1e-2)
    @test isapprox(getvalue(x), 0.5, atol=1e-2)
end

## Solve test problem with simple max functions
 #
 #  min   max( x^2, x )
 #  s.t.  -0.5 <= x <= 0.5
 #
 #  The optimal objective value is 0.
 #      x = 0.
##
@testset "example: minmax" begin
    m = Model(solver=solver)
    @variable(m, -1 <= x <= 1, start=-1)
    @NLobjective(m, Min, max(x^2, x, -1))
    @test solve(m) == :Optimal
    @test isapprox(getobjectivevalue(m), 0, atol=1e-2)
    @test isapprox(getvalue(x), 0, atol=1e-2)
end

ext: jl


text:
export GF256

"""
    struct GF256

Element of a finite field of size 256 with arithmetic as defined in rfc6330.

"""
struct GF256 <: Integer
    x::UInt8
end

GF256(x::Integer) = GF256(UInt8(x))
GF256(a::GF256) = GF256(a.x)

function show(io::IO, a::GF256)
    print(io, "GF256($(a.x))")
end

# TODO: remove
# Base.convert(::Type{Bool}, a::GF256) = Bool(a.x)
# Base.convert(::Type{Int}, a::GF256) = Int(a.x)
# Base.convert(::Type{GF256}, x::Int) = GF256(x)
# Base.convert(::Type{GF256}, x::UInt8) = GF256(x)
# Base.convert(::Type{GF256}, x::Bool) = x ? GF256(1) : GF256(0)
# Base.promote_rule(::Type{GF256}, ::Type{Bool}) = GF256
# Base.promote_rule(::Type{GF256}, ::Type{Int}) = Int

Base.Int128(a::GF256) = Int128(a.x)
Base.UInt128(a::GF256) = UInt128(a.x)
Base.Int64(a::GF256) = Int64(a.x)
Base.UInt64(a::GF256) = UInt64(a.x)
Base.Int32(a::GF256) = Int32(a.x)
Base.UInt32(a::GF256) = UInt32(a.x)
Base.Int16(a::GF256) = Int16(a.x)
Base.UInt16(a::GF256) = UInt16(a.x)
Base.UInt8(a::GF256) = UInt8(a.x)

# zeros/ones
Base.zero(::GF256) = GF256(0)
Base.zero(::Type{GF256}) = GF256(0)
Base.one(::GF256) = GF256(1)
Base.one(::Type{GF256}) = GF256(1)
Base.iszero(a::GF256) = iszero(a.x)
Base.isone(a::GF256) = isone(a.x)

function Base.zeros(::Type{GF256}, dims::Vararg{Union{Int, AbstractUnitRange},N} where N)
    rv = Array{GF256}(undef, dims...)
    rv .= zero(GF256)
end

function Base.ones(::Type{GF256}, dims::Vararg{Union{Int, AbstractUnitRange},N} where N)
    rv = Array{GF256}(undef, dims...)
    rv .= one(GF256)
end

# iteration
Base.length(a::GF256) = 1
Base.iterate(a::GF256) = (a, nothing)
Base.iterate(a::GF256, ::Nothing) = nothing

# random numbers
Base.rand(rng::AbstractRNG, ::Random.SamplerType{GF256}) = GF256(rand(rng, UInt8))

"""logarithm as defined by rfc6330. used for multiplication and division."""
function Base.log(a::GF256)::GF256
    !iszero(a) || throw(DomainError(GF256(0), "log(GF256(0)) is undefined"))
    GF256(RQ_OCT_LOG[a.x])
end

"""exponentiation as defined by rfc6330. used for multiplication and division."""
function Base.exp(a::GF256)::GF256
    a.x != 255 || throw(DomainError(GF256(255)), "exp(GF256(255)) is undefined")
    RQ_OCT_EXP[a.x+1]
end

# TODO: remove
# Base.:+(a::Bool, b::GF256) = +(promote(a, b)...)
# Base.:-(a::Bool, b::GF256) = -(promote(a, b)...)
# Base.:+(a::GF256, b::Bool) = +(promote(a, b)...)
# Base.:-(a::GF256, b::Bool) = -(promote(a, b)...)
# Base.:/(a::GF256, b::Bool) = /(promote(a, b)...)
# Base.:*(a::Bool, b::GF256) = a ? b : zero(GF256)
# Base.:*(a::GF256, b::Bool) = b*a

# TODO: faulty implementation
# Base.:^(a::GF256, b::Integer) = b < 0 ? throw(DomainError(b, "^$b undefined")) : exprq(b * logrq(a) % 510)

# arithmetic
Base.:+(a::GF256, b::GF256) = GF256(xor(a.x, b.x))
Base.:-(a::GF256, b::GF256) = a + b
function Base.:/(a::GF256, b::GF256)
    !iszero(b) || throw(DivideError())
    iszero(a) ? a : RQ_OCT_EXP[RQ_OCT_LOG[a.x] + 255 - RQ_OCT_LOG[b.x] + 1]
end
Base.:*(a::GF256, b::GF256) = iszero(a) || iszero(b) ? zero(GF256) : RQ_OCT_EXP[RQ_OCT_LOG[a.x] + RQ_OCT_LOG[b.x] + 1]
Base.sub_with_overflow(a::GF256, b::GF256) = (a-b, false)
Base.add_with_overflow(a::GF256, b::GF256) = (a+b, false)

# comparison
Base.:<(a::GF256, b::GF256) = a.x < b.x
Base.:<=(a::GF256, b::GF256) = a.x <= b.x
Base.:>(a::GF256, b::GF256) = a.x > b.x
Base.:>=(a::GF256, b::GF256) = a.x >= b.x
Base.:(==)(a::GF256, b::GF256) = a.x == b.x

"""

GF256 exponentiation lookup table (table 5.7.3 of rfc6330). The table is zero-indexed, i.e., one
must add `1` when indexing into it.
"""
const RQ_OCT_EXP = Vector{GF256}([
    1, 2, 4, 8, 16, 32, 64, 128, 29, 58, 116, 232, 205, 135, 19, 38, 76,
    152, 45, 90, 180, 117, 234, 201, 143, 3, 6, 12, 24, 48, 96, 192, 157,
    39, 78, 156, 37, 74, 148, 53, 106, 212, 181, 119, 238, 193, 159, 35,
    70, 140, 5, 10, 20, 40, 80, 160, 93, 186, 105, 210, 185, 111, 222,
    161, 95, 190, 97, 194, 153, 47, 94, 188, 101, 202, 137, 15, 30, 60,
    120, 240, 253, 231, 211, 187, 107, 214, 177, 127, 254, 225, 223, 163,
    91, 182, 113, 226, 217, 175, 67, 134, 17, 34, 68, 136, 13, 26, 52,
    104, 208, 189, 103, 206, 129, 31, 62, 124, 248, 237, 199, 147, 59,
    118, 236, 197, 151, 51, 102, 204, 133, 23, 46, 92, 184, 109, 218,
    169, 79, 158, 33, 66, 132, 21, 42, 84, 168, 77, 154, 41, 82, 164, 85,
    170, 73, 146, 57, 114, 228, 213, 183, 115, 230, 209, 191, 99, 198,
    145, 63, 126, 252, 229, 215, 179, 123, 246, 241, 255, 227, 219, 171,
    75, 150, 49, 98, 196, 149, 55, 110, 220, 165, 87, 174, 65, 130, 25,
    50, 100, 200, 141, 7, 14, 28, 56, 112, 224, 221, 167, 83, 166, 81,
    162, 89, 178, 121, 242, 249, 239, 195, 155, 43, 86, 172, 69, 138, 9,
    18, 36, 72, 144, 61, 122, 244, 245, 247, 243, 251, 235, 203, 139, 11,
    22, 44, 88, 176, 125, 250, 233, 207, 131, 27, 54, 108, 216, 173, 71,
    142, 1, 2, 4, 8, 16, 32, 64, 128, 29, 58, 116, 232, 205, 135, 19, 38,
    76, 152, 45, 90, 180, 117, 234, 201, 143, 3, 6, 12, 24, 48, 96, 192,
    157, 39, 78, 156, 37, 74, 148, 53, 106, 212, 181, 119, 238, 193, 159,
    35, 70, 140, 5, 10, 20, 40, 80, 160, 93, 186, 105, 210, 185, 111,
    222, 161, 95, 190, 97, 194, 153, 47, 94, 188, 101, 202, 137, 15, 30,
    60, 120, 240, 253, 231, 211, 187, 107, 214, 177, 127, 254, 225, 223,
    163, 91, 182, 113, 226, 217, 175, 67, 134, 17, 34, 68, 136, 13, 26,
    52, 104, 208, 189, 103, 206, 129, 31, 62, 124, 248, 237, 199, 147,
    59, 118, 236, 197, 151, 51, 102, 204, 133, 23, 46, 92, 184, 109, 218,
    169, 79, 158, 33, 66, 132, 21, 42, 84, 168, 77, 154, 41, 82, 164, 85,
    170, 73, 146, 57, 114, 228, 213, 183, 115, 230, 209, 191, 99, 198,
    145, 63, 126, 252, 229, 215, 179, 123, 246, 241, 255, 227, 219, 171,
    75, 150, 49, 98, 196, 149, 55, 110, 220, 165, 87, 174, 65, 130, 25,
    50, 100, 200, 141, 7, 14, 28, 56, 112, 224, 221, 167, 83, 166, 81,
    162, 89, 178, 121, 242, 249, 239, 195, 155, 43, 86, 172, 69, 138, 9,
    18, 36, 72, 144, 61, 122, 244, 245, 247, 243, 251, 235, 203, 139, 11,
    22, 44, 88, 176, 125, 250, 233, 207, 131, 27, 54, 108, 216, 173, 71,
    142
])

"""

GF256 logarithm lookup table (table 5.7.4 of rfc6330). The table is one-indexed.
"""
const RQ_OCT_LOG = Vector{Int16}([
    0, 1, 25, 2, 50, 26, 198, 3, 223, 51, 238, 27, 104, 199, 75, 4, 100,
    224, 14, 52, 141, 239, 129, 28, 193, 105, 248, 200, 8, 76, 113, 5,
    138, 101, 47, 225, 36, 15, 33, 53, 147, 142, 218, 240, 18, 130, 69,
    29, 181, 194, 125, 106, 39, 249, 185, 201, 154, 9, 120, 77, 228, 114,
    166, 6, 191, 139, 98, 102, 221, 48, 253, 226, 152, 37, 179, 16, 145,
    34, 136, 54, 208, 148, 206, 143, 150, 219, 189, 241, 210, 19, 92,
    131, 56, 70, 64, 30, 66, 182, 163, 195, 72, 126, 110, 107, 58, 40,
    84, 250, 133, 186, 61, 202, 94, 155, 159, 10, 21, 121, 43, 78, 212,
    229, 172, 115, 243, 167, 87, 7, 112, 192, 247, 140, 128, 99, 13, 103,
    74, 222, 237, 49, 197, 254, 24, 227, 165, 153, 119, 38, 184, 180,
    124, 17, 68, 146, 217, 35, 32, 137, 46, 55, 63, 209, 91, 149, 188,
    207, 205, 144, 135, 151, 178, 220, 252, 190, 97, 242, 86, 211, 171,
    20, 42, 93, 158, 132, 60, 57, 83, 71, 109, 65, 162, 31, 45, 67, 216,
    183, 123, 164, 118, 196, 23, 73, 236, 127, 12, 111, 246, 108, 161,
    59, 82, 41, 157, 85, 170, 251, 96, 134, 177, 187, 204, 62, 90, 203,
    89, 95, 176, 156, 169, 160, 81, 11, 245, 22, 235, 122, 117, 44, 215,
    79, 174, 213, 233, 230, 231, 173, 232, 116, 214, 244, 234, 168, 80,
    88, 175
])

max_stars_repo_name: severinson/RaptorCodes

max_issues_repo_name: severinson/RaptorCodes

max_forks_repo_name: severinson/RaptorCodes


text:
#=
Numerical_Cont.jl is able to do numerical continuation of system with Hopf bifurcation starting from the stable LCO solutions.
Equation of motion should be expressed in nonmutating form as flutter_eq_CBC_nonmu(u, p, t)
and the dimension of the equation should be defined as: dimension(::typeof(flutter_eq_CBC_nonmu)) = 6
=#

function fourier_diff(N::Integer)
    # For a MATLAB equivalent see http://appliedmaths.sun.ac.za/~weideman/research/differ.html
    h = 2π/N
    D = zeros(N, N)
    # First column
    n = ceil(Int, (N-1)/2)
    if (N % 2) == 0
        for i in 1:n
            D[i+1, 1] = -((i % 2) - 0.5)*cot(i*h/2)
        end
    else
        for i in 1:n
            D[i+1, 1] = -((i % 2) - 0.5)*csc(i*h/2)
        end
    end
    for i in n+2:N
        D[i, 1] = -D[N-i+2, 1]
    end
    # Other columns (circulant matrix)
    for j in 2:N
        D[1, j] = D[N, j-1]
        for i in 2:N
            D[i, j] = D[i-1, j-1]
        end
    end
    return D
end

function periodic_zero_problem(rhs, u, p, ee)
    # Assume that u has dimensions nN + 1 where n is the dimension of the ODE
    n = dimension(rhs)
    N = (length(u)-1)÷n
    D = fourier_diff(N)
    T = u[end]/(2π)  # the differentiation matrix is defined on [0, 2π]
    # Evaluate the right-hand side of the ODE
    res=Vector{Float64}(undef, n*N+1)
    for i in 1:n:n*N
        res[i:i+n-1] = T*rhs(@view(u[i:i+n-1]), p, 0)  # use a view for speed
    end
    # Evaluate the derivative; for-loop for speed equivalent to u*Dᵀ
    for (i, ii) in pairs(1:n:n*N)
        # Evaluate the derivative at the i-th grid point
        for (j, jj) in pairs(1:n:n*N)
            for k = 1:n
                res[ii+k-1] -= D[i, j]*u[jj+k-1]
            end
        end
    end
    res[end] = (u[1] - ee)  # phase condition - assumes that the limit cycle passes through the Poincare section u₁=1e-4; using a non-zero value prevents convergence to the trivial solution
    return res
end

function periodic_zero_problem2(rhs, u, p, ee, dv, pu, ds) # Periodic zero problem including Pseudo arc-length equation
    # Assume that u has dimensions nN + 1 where n is the dimension of the ODE
    n = dimension(rhs)
    N = (length(u)-1)÷n
    res=Vector{Float64}(undef, n*N+2)
    res1=periodic_zero_problem(rhs, u, p, ee)
    p1=p[1];
    uu=[u;p1]
    du=uu-pu

    res[1:end-1]=res1
    arclength=norm(transpose(dv)*du)
    res[end] = arclength-ds  #Pseudo-arclength equation
    return res
end

function periodic_zero_J(jeq, em, u, p, ee) # Jacobian of periodic zero problem
    n = dimension(em)
    N = (length(u)-1)÷n
    T = u[end]/(2π)
    J = zeros(n*N+1,n*N+1)
    D = fourier_diff(N)
    for i in 1:n:n*N
        J[i:i+n-1,i:i+n-1] = T*jeq(@view(u[i:i+n-1]), p, 0)
        J[i:i+n-1,n*N+1] = em(@view(u[i:i+n-1]), p, 0)
    end
    for (i, ii) in pairs(1:n:n*N)
        for (j, jj) in pairs(1:n:n*N)
            for k = 1:n
                J[ii+k-1,jj+k-1] -= D[i, j]
            end
        end
    end
    J[n*N+1,1:n*N+1] = zeros(1,n*N+1)
    J[n*N+1,1] = 1
    return J
end

function periodic_zero_J2(jeq, jeq2, em, u, p, ee, dv)  # Jacobian of periodic zero problem 2
    n=dimension(em)
    N = (length(u)-1)÷n
    J = zeros(n*N+2,n*N+2)
    T = u[end]/(2π)
    J[1:n*N+1,1:n*N+1]=periodic_zero_J(jeq, em, u, p, ee)
    for i in 1:n:n*N
        J[i:i+n-1,n*N+2] = T*jeq2(@view(u[i:i+n-1]), p, 0)
    end
    J[n*N+2,:]=dv
    return J
end

function LCO_solve(u1,ini_T,eq,jeq,p,ee,z_tol) # Solve the periodic zero problem with Newton methods
    u0=vcat(u1,ini_T)
    err=[]
    for i in 1:40
        res=periodic_zero_problem(eq, u0, p, ee)
        u1=u0-periodic_zero_J(jeq, eq, u0, p, ee)\res
        u0=u1
        er=transpose(res)*res
        err=vcat(err,er)
        if er<z_tol
            break
        end
    end
    return (u=u0, err=err)
end

function LCO_solve2(u1,ini_T,eq,jeq,jeq2,p,ee,z_tol,dv,pu,ds) # Solve the periodic zero problem 2 with Newton methods
    p0=p[1]
    u0=[u1;ini_T]
    uu0=[u0;p0]
    err=[]
    np=p
    for i in 1:40
        res=periodic_zero_problem2(eq, u0, np, ee, dv, pu, ds)
        uu1=uu0-periodic_zero_J2(jeq, jeq2, eq, u0, np, ee, dv)\res
        uu0=uu1
        u0=uu0[1:end-1]
        p0=uu0[end]
        np=vcat(p0,p[2:end])
        er=norm(transpose(res)*res)
        err=vcat(err,er)
        if er<z_tol
            break
        end
    end
    return (u=uu0, err=err)
end

function get_sol(u0,N,ind1,ind2) # convert solution of collocation -> Array form
    # returns phase angle, amplitude
    dim=Int((length(u0)-1)/N)
    u=u0[1:end-1]
    T=u0[end]
    uu=Array{Float64}(undef,dim,N)
    theta=Array{Float64}(undef,N)
    r=Array{Float64}(undef,N)
    for i in 1:dim
        uu[i,:]=u[i:dim:end]
    end
    for i in 1:N
        theta[i]=atan(uu[ind2,i],uu[ind1,i])
        r[i]=sqrt(uu[ind1,i]^2+uu[ind2,i]^2)
    end
    return (u=uu,T=T,t=theta,r=r)
end

function get_sol(u0,dim,N,ind1,ind2)  # Convert solution of periodic zero-problem -> array solution and amplitude, phase angle of measured coordinate (z_1,z_2)
    u=u0[1:end-1]
    T=u0[end]
    uu=Array{Float64}(undef,N,dim)
    theta=Array{Float64}(undef,N)
    r=Array{Float64}(undef,N)
    for i in 1:dim
        uu[:,i]=u[i:dim:end]
    end
    for i in 1:N
        theta[i]=atan(uu[i,ind2],uu[i,ind1])
        r[i]=sqrt(uu[i,ind1]^2+uu[i,ind2]^2)
    end
    return (u=uu,T=T,t=theta,r=r)
end

function get_sol2(sol,ind) # Convert solution of ODE solver to array solution
    lu=length(sol.u)
    u=Vector{Float64}(undef, lu)
    for i in 1:lu
        uv=sol.u[i]
        u[i]=uv[ind]
    end
    return u
end

function get_stable_LCO(p,u0,tl,tol,eq,stol,rp) # Get a stable LCO from the numerical integration
    u=u0
    dim=length(u0)
    prob = ODEProblem(eq,u,(0,tl*rp),p)
    sol = solve(prob,Tsit5(),reltol=stol,abstol=stol)
    vP=1
    P=0
    T=0
    while vP>tol
        u=sol.u[end]
        prob = ODEProblem(eq,u,(0,tl),p)
        sol = solve(prob,Tsit5(),reltol=stol,abstol=stol)
        z=zero_measure(sol.u,1,sol.t)
        vP=var(z.hp)
        P=z.P[1]
        T=z.T
    end
    tl=length(sol)
    uu=Array{Float64}(undef,tl,dim)
    for i in 1:dim
        uu[:,i]=get_sol2(sol,i)
    end
    t=Array{Float64}(undef,length(sol))
    r=Array{Float64}(undef,length(sol))
    ind1=1;ind2=3;u=uu[:,ind1];v=uu[:,ind2]
    for i in 1:length(u)
        t[i]=atan(v[i],u[i])
        r[i]=sqrt(u[i]^2+v[i]^2)
    end
    return (u=uu,t=t,r=r,P=P,T=T)
end

function get_stable_LCO(p,u0,tl,tol,eq,stol,rp,ind1,ind2,u₀,v₀,st) # Get a stable LCO from numerical integration
    u=u0
    dim=length(u0)
    prob = ODEProblem(eq,u,(0,tl*rp),p)
    sol = DifferentialEquations.solve(prob,Tsit5(),reltol=stol,abstol=stol,saveat=st)
    vP=1
    P=0
    T=0
    while vP>tol
        u=sol.u[end]
        prob = ODEProblem(eq,u,(0,tl),p)
        sol = DifferentialEquations.solve(prob,Tsit5(),reltol=stol,abstol=stol,saveat=st)
        z=zero_measure(sol.u,1,sol.t)
        vP=Statistics.var(z.hp)
        P=z.P[1]
        T=z.T
    end
    tl=length(sol)
    uu=Array{Float64}(undef,tl,dim)
    for i in 1:dim
        uu[:,i]=get_sol2(sol,i)
    end
    t=Array{Float64}(undef,length(sol))
    r=Array{Float64}(undef,length(sol))
    u=uu[:,ind1];v=uu[:,ind2]
    for i in 1:length(u)
        t[i]=atan(v[i]-v₀,u[i]-u₀)
        r[i]=sqrt((u[i]-u₀)^2+(v[i]-v₀)^2)
    end
    return (u=uu,t=t,r=r,P=P,T=T)
end

function Newton_CBC(u1,ini_T,eq,jeq,p,ee,z_tol,c_tol,ind1,ind2) # Newton method to solve zero-problem of CBC
    p0=p[1:3]
    c0=p[4:end]
    nh=Int((length(c0)-1)/2)
    dm=dimension(eq)
    N=Int(length(u1)/dm)
    s=LCO_solve(u1,ini_T,eq,jeq,p,ee,z_tol)
    g=get_sol(s.u,6,N,ind1,ind2)
    r=LS_harmonics(g.r,g.t,1,nh)
    c=r.coeff
    del_c0=c-c0
    err=transpose(del_c0)*del_c0
    er=err
    for i in 1:10
        if err<c_tol
            break
        end
        J=Jac_CBC(p,u1,ini_T,eq,jeq,ee,ind1,ind2,z_tol)
        del_c=inv(J)*del_c0
        c1=c0-del_c
        p=vcat(p0,c1)
        s=LCO_solve(u1,ini_T,eq,jeq,p,ee,z_tol)
        g=get_sol(s.u,6,N,ind1,ind2)
        r=LS_harmonics(g.r,g.t,1,nh)
        del_c1=r.coeff-c1
        del_c0=del_c1
        c0=c1
        err=transpose(del_c1)*del_c1
        er=vcat(er,err)
    end
    return (s=s,p=p,e=er,u1=s.u[1:end-1],T1=s.u[end])
end

function Newton_CBC(u1,ini_T,eq,jeq,p,ee,z_tol,c_tol,ind1,ind2,dv,pu,ds)  # Newton method to solve zero-problem of CBC including Pseudo arclength equation
    p0=p[1:3]
    c0=p[4:end]
    nh=Int((length(c0)-1)/2)
    dm=dimension(eq)
    N=Int(length(u1)/dm)
    rr=res_CBC(p,u1,ini_T,eq,jeq,ee,ind1,ind2,z_tol,dv,pu,ds)
    res=rr.res
    u1=rr.u1;ini_T=rr.T1;
    err=transpose(res)*res
    er=err
    for i in 1:30
        if err<c_tol
            break
        end
        J=Jac_CBC(p,u1,ini_T,eq,jeq,ee,ind1,ind2,z_tol,dv)
        del_u=inv(J)*res
        c1=c0-del_u[1:end-1]
        p0[1]=p0[1]-del_u[end]
        p=vcat(p0,c1)
        rr=res_CBC(p,u1,ini_T,eq,jeq,ee,ind1,ind2,z_tol,dv,pu,ds)
        res=rr.res
        u1=rr.u1;ini_T=rr.T1;
        c0=c1
        err=transpose(res)*res
        er=vcat(er,err)
    end
    return (u1=u1,T1=ini_T,p=p,e=er)
end

function Jac_CBC(p,u0,ini_T,eq,jeq,ee,ind1,ind2,z_tol) # Jacobian of CBC zero-problem
    pp=p[1:3]
    c0=p[4:end]
    dm=dimension(eq)
    nh=Int((length(c0)-1)/2)
    N=Int(length(u0)/dm)
    s=LCO_solve(u0,ini_T,eq,jeq,p,ee,z_tol)
    g=get_sol(s.u,dm,N,ind1,ind2)
    c=LS_harmonics(g.r,g.t,1,nh).coeff
    del_c0=c-c0
    J=Array{Float64}(undef,2*nh+1,2*nh+1)
    eye=SMatrix{2*nh+1,2*nh+1}(1I)
    delta=1e-11
    for i in 1:2*nh+1
        c=c0+eye[i,:]*delta
        np=vcat(pp,c)
        s=LCO_solve(u0,ini_T,eq,jeq,np,ee,z_tol)
        g=get_sol(s.u,dm,N,ind1,ind2)
        dc=LS_harmonics(g.r,g.t,1,nh).coeff
        del_c1=dc-c
        dp=(del_c1-del_c0)/delta
        J[:,i]=dp
    end
    return J
end

function Jac_CBC(p,u0,ini_T,eq,jeq,ee,ind1,ind2,z_tol,dv) #Jacobian of CBC zero problem with Pseudo arclength
    pp=p[1:3]
    c0=p[4:end]
    dm=dimension(eq)
    nh=Int((length(c0)-1)/2)
    N=Int(length(u0)/dm)
    s=LCO_solve(u0,ini_T,eq,jeq,p,ee,z_tol)
    g=get_sol(s.u,dm,N,ind1,ind2)
    c=LS_harmonics(g.r,g.t,1,nh).coeff
    del_c0=c-c0
    J=Array{Float64}(undef,2*nh+2,2*nh+2)
    eye=SMatrix{2*nh+1,2*nh+1}(1I)
    delta=1e-11
    for i in 1:2*nh+1
        c=c0+eye[i,:]*delta
        np=vcat(pp,c)
        s=LCO_solve(u0,ini_T,eq,jeq,np,ee,z_tol)
        g=get_sol(s.u,dm,N,ind1,ind2)
        dc=LS_harmonics(g.r,g.t,1,nh).coeff
        del_c1=dc-c
        dp=(del_c1-del_c0)/delta
        J[1:end-1,i]=dp
    end
    npp=pp
    npp[1]=npp[1]+delta
    np=vcat(npp,c0)
    s=LCO_solve(u0,ini_T,eq,jeq,np,ee,z_tol)
    g=get_sol(s.u,dm,N,ind1,ind2)
    dc=LS_harmonics(g.r,g.t,1,nh).coeff
    del_c1=dc-c0
    dp=(del_c1-del_c0)/delta
    J[1:end-1,end]=dp
    J[end,:]=dv
    return J
end

function res_CBC(p,u0,ini_T,eq,jeq,ee,ind1,ind2,z_tol,dv,pu,ds) # Residue of zero problem of CBC
    pp=p[1:3]
    c0=p[4:end]
    dm=dimension(eq)
    nh=Int((length(c0)-1)/2)
    N=Int(length(u0)/dm)
    #res=Array{Float64}(undef,2*nh+2,1)
    res=Vector{Float64}(undef, 2*nh+2)
    s=LCO_solve(u0,ini_T,eq,jeq,p,ee,z_tol)
    g=get_sol(s.u,dm,N,ind1,ind2)
    c=LS_harmonics(g.r,g.t,1,nh).coeff
    del_c0=c-c0
    uu=[c0;pp[1]]
    du=uu-pu
    res[1:end-1]=del_c0
    arclength=transpose(du)*dv
    res[end] = arclength-ds
    return (res=res,u1=s.u[1:end-1],T1=s.u[end])
end

function zero_measure(u,ind,t) # Measuring the time of zero crossing points from numerical continuation
    l=length(u)
    zero=Array{Float64}(undef, 0)
    T=Array{Float64}(undef, 0)
    low_p=Array{Float64}(undef, 0)
    high_p=Array{Float64}(undef, 0)
    Ti=Array{Int64}(undef, 0)
    for i in 2:l-1
        sign_con2=u[i][ind+1]*u[i-1][ind+1]
        if sign_con2 < 0
            if (u[i][ind]+u[i-1][ind])/2 < 0
                low_p=vcat(low_p,(u[i][ind]+u[i-1][ind])/2)
            else
                high_p=vcat(high_p,(u[i][ind]+u[i-1][ind])/2)
            end
        end
    end
    h₀=mean(high_p)+mean(low_p)
    h₀=h₀/2
    for i in 2:l-1
        sign_con=(u[i][ind]-h₀)*(u[i+1][ind]-h₀)
        if sign_con < 0
            if (u[i][ind+1]+u[i+1][ind+1])/2 > 0
                zero=vcat(zero,(u[i][ind]+u[i-1][ind])/2)
                Ti=vcat(Ti,i)
                T=vcat(T,t[i])
            end
        end
    end
    t_l=length(T)
    P=Array{Float64}(undef, t_l-1)
    for j in 2:t_l
        P[j-1]=T[j]-T[j-1]
    end
    return (T=Ti, P=P, hp=high_p, lp=low_p, h₀=h₀)
end

function MonoD(eq,dim,u,p,t) # Variational equation for monodromy matrix computation (ForwardDiff is used for jacobian computation)
    jeq=(u,p,t) -> ForwardDiff.jacobian(u -> eq(u,p,t), u)
    n=dim
    v=u[1:Int(n*n)]
    uu=u[Int(n*n+1):end]
    M=reshape(v, n, n)
    J=jeq(uu,p,t)
    dM=J*M
    dv=vec(dM)
    duu=eq(uu,p,t)
    duu=vec(duu)
    du=[dv;duu]
    return du
end

function Monodromy_compute(eq,u,p,N,ind) # Monodromy matrix computation using numerical integration
    n=dimension(eq)
    eye=1.0*Matrix(I, n, n)
    v0=vec(eye)
    g=get_sol(u,n,N,ind[1],ind[2])
    T=g.T
    uu=g.u
    M=eye
    tl2=T/N
    for i in 1:N
        uu0=uu[i,:]
        u0=[v0;uu0]
        prob = ODEProblem((u,p,t) -> MonoD(eq,n,u,p,t),u0,(0,tl2),p)
        sol = solve(prob,Tsit5(),reltol=1e-14,abstol=1e-14)
        w=sol.u[end]
        m1=w[1:Int(n*n)]
        M1=reshape(m1, n, n)
        M=M1*M
    end
    Eig1=eigen(M)
    μ=Eig1.values
    return μ
end

function LS_harmonics(r,t,ω,N) # Computing Fourier coefficients of the amplitude in the measued state-variable coordinates
    # Fourier coefficients are computed in least square sence
    c=Array{Float64}(undef,2*N+1)
    M=Array{Float64}(undef,1,2*N+1)
    tM=Array{Float64}(undef,0,2*N+1)
    tl=length(t)
    rr=Array{Float64}(undef,tl)
    M[1]=1
    for j in 1:tl
        for i in 1:N
            M[1+i]=cos(ω*t[j]*i)
            M[1+N+i]=sin(ω*t[j]*i)
        end
        tM=vcat(tM,M)
    end
    MM=transpose(tM)*tM
    rN=transpose(tM)*r
    c=inv(MM)*rN
    for j in 1:tl
        rr[j]=c[1]
        for i in 1:N
            rr[j]+=c[i+1]*cos(ω*t[j]*i)
            rr[j]+=c[i+1+N]*sin(ω*t[j]*i)
        end
    end
    return (coeff=c,rr=rr)
end

function LS_h_recon(tl,c,N) # Computing Fourier coefficients of the amplitude in the measued state-variable coordinates
    # Fourier coefficients are computed in least square sence
    t=range(0.0,2π,length=tl+1)
    t=t[1:end-1]
    rr=zeros(length(t))
    for j in 1:tl
        rr[j]=c[1]
        for i in 1:N
            rr[j]+=c[i+1]*cos(1*t[j]*i)
            rr[j]+=c[i+1+N]*sin(1*t[j]*i)
        end
    end
    return (r=rr,t=t)
end

function ini_val(p,tl,tol,N,eq,u0,stol,rp) #Get stable LCO from the ODE numerical solutions- initial collocation points
    dim=length(u0)
    s0=get_stable_LCO(p,u0,tl,tol,eq,stol,rp)
    pp=s0.P[1]
    ini_T=pp
    mu=s0.u[:,1]
    u=mu-1e-3*ones(length(mu))
    u=broadcast(abs, u)
    ps=findall(isequal(minimum(u)), u)
    ps=ps[1]
    u0= s0.u[ps,:]
    #Get a time series of periodic solution
    tl2=pp/N
    #ini_u=Array{Float64}(undef,Int(dim*N))
    ini_u=Vector{Float64}(undef, Int(dim*N))
    ini_u[1:dim]=u0
    eq=flutter_eq_CBC;
    #uu=pu+ds*dv;
    for i in 1:N-1
        prob = ODEProblem(eq,u0,(0,tl2),p)
        sol = solve(prob,Tsit5(),reltol=1e-11,abstol=1e-11)
        u0=sol.u[end]
        ini_u[Int(dim*i+1):Int(dim*i+dim)]=u0
    end
    return (u=ini_u,T=ini_T)
end

function get_CBC_u(p)
    c0=p[4:end]
    nu=p[1]
    u=vcat(c0,nu)
    return u
end

function get_CBC_p(nu,Kp,Kd,N)
    c=nu[1:end-1]
    p=vcat([nu[end],Kp,Kd],c)
    return p
end

function get_dv_u(u,T1,pu,ds)
    uu=vcat(u,T1)
    dv=uu-pu;dv=dv/norm(dv)
    nu=uu+dv*ds
    u1=nu[1:end-1]
    T1=nu[end]
    return (u=u1,T=T1)
end

function amp_LCO(U,ind,N) #Extract amplitude from collocation points
    l1=length(U)
    amp=Vector{Float64}(undef, l1)
    for i in 1:l1
        g=get_sol(U[i],6,N,1,3)
        uu=g.u[:,ind]
        amp[i]=maximum(uu)-minimum(uu)
    end
    return amp
end

function Hopf_point(eq,p) # Compute the Hopf point
t=0.0;u=zeros(dimension(eq))
jeq=(u,p,t) -> ForwardDiff.jacobian(u -> eq(u,p,t), u)
s=nlsolve(p0->maximum(real(eigen(jeq(u,p0,t)).values))-1e-8,p)
    return s.zero[1]
end
## Numerical continuation of Flutter model and CBC zero problem (Examples)

function flutter_eq_CBC(u, p, t) #  Equation of motion of Flutter model with CBC
    μ=p[1]
    ind1=1;ind2=3;
    theta=atan(u[ind2],u[ind1])
    if length(p)>=2
        Kp=p[2]
        Kd=p[3]
        c=p[4:end]
        r=c[1]
        h=c[2:end]
        nh=Int(length(h)/2)
        for i=1:nh
            r+=h[i]*cos(i*theta)+h[i+nh]*sin(theta*i)
        end
    else
        Kp=0
        Kd=0
        r=0
        theta=0
    end
#theta is computed from two state variables

    b=0.15; a=-0.5; rho=1.204; x__alpha=0.2340
    c__0=1; c__1=0.1650; c__2=0.0455; c__3=0.335; c__4=0.3; c__alpha=0.562766779889303; c__h=15.4430
    I__alpha=0.1726; k__alpha=54.116182926744390; k__h=3.5294e+03; m=5.3; m__T=16.9
    U=μ

    MM = [b ^ 2 * pi * rho + m__T -a * b ^ 3 * pi * rho + b * m * x__alpha 0; -a * b ^ 3 * pi * rho + b * m * x__alpha I__alpha + pi * (0.1e1 / 0.8e1 + a ^ 2) * rho * b ^ 4 0; 0 0 1;]
    DD = [c__h + 2 * pi * rho * b * U * (c__0 - c__1 - c__3) (1 + (c__0 - c__1 - c__3) * (1 - 2 * a)) * pi * rho * b ^ 2 * U 2 * pi * rho * U ^ 2 * b * (c__1 * c__2 + c__3 * c__4); -0.2e1 * pi * (a + 0.1e1 / 0.2e1) * rho * (b ^ 2) * (c__0 - c__1 - c__3) * U c__alpha + (0.1e1 / 0.2e1 - a) * (1 - (c__0 - c__1 - c__3) * (1 + 2 * a)) * pi * rho * (b ^ 3) * U -0.2e1 * pi * rho * (U ^ 2) * (b ^ 2) * (a + 0.1e1 / 0.2e1) * (c__1 * c__2 + c__3 * c__4); -1 / b a - 0.1e1 / 0.2e1 (c__2 + c__4) * U / b;]
    KK = [k__h 2 * pi * rho * b * U ^ 2 * (c__0 - c__1 - c__3) 2 * pi * rho * U ^ 3 * c__2 * c__4 * (c__1 + c__3); 0 k__alpha - 0.2e1 * pi * (a + 0.1e1 / 0.2e1) * rho * (c__0 - c__1 - c__3) * (b ^ 2) * (U ^ 2) -0.2e1 * pi * rho * b * (U ^ 3) * (a + 0.1e1 / 0.2e1) * c__2 * c__4 * (c__1 + c__3); 0 -U / b c__2 * c__4 * U ^ 2 / b ^ 2;]

    K1=-inv(MM)*KK;
    D1=-inv(MM)*DD;

    J1=[0 1 0 0 0 0]
    J2=[K1[1,1] D1[1,1] K1[1,2] D1[1,2] K1[1,3] D1[1,3]]
    J3=[0 0 0 1 0 0]
    J4=[K1[2,1] D1[2,1] K1[2,2] D1[2,2] K1[2,3] D1[2,3]]
    J5=[0 0 0 0 0 1]
    J6=[K1[3,1] D1[3,1] K1[3,2] D1[3,2] K1[3,3] D1[3,3]]

    J=[J1;J2;J3;J4;J5;J6]
    du=J*u
    #Control added on heave
    M2=inv(MM)
    control=Kp*(r*cos(theta)-u[ind1])+Kd*(r*sin(theta)-u[ind2])
    ka2=751.6; ka3=5006;
    nonlinear_h=-M2[1,2]*(ka2*u[3]^2+ka3*u[3]^3)
    nonlinear_theta=-M2[2,2]*(ka2*u[3]^2+ka3*u[3]^3)

    du[2]+=nonlinear_h+M2[1,1]*control
    du[4]+=nonlinear_theta+M2[2,1]*control
    return du
end

dimension(::typeof(flutter_eq_CBC)) = 6 # dimension should be defined

function continuation_flutter(N,sp,sU,ds) #Get stable LCO from the ODE numerical solutions- initial collocation points
    μ=sU+0.02;p=[μ,0,0,0,0,0];ee=1e-3;
    eq=flutter_eq_CBC;
    tol=1e-7;
    jeq=(u,p,t) -> ForwardDiff.jacobian(u -> eq(u,p,t), u)
    jeq2=(u,p,t) -> ForwardDiff.jacobian(p -> eq(u,p,t), p)[:,1]
    stol=1e-10;rp=5;u0=[0.1;0.1;0.1;0;0;0];tl=10.0
    s1=ini_val(p,tl,tol,N,eq,u0,stol,rp)
    s=LCO_solve(s1.u,s1.T,eq,jeq,p,s1.u[1],tol)
    μ=sU;p=[μ,0,0,0,0,0];
    s2=LCO_solve(s1.u,s1.T,eq,jeq,p,s1.u[1],tol)

    du=s2.u-s.u
    du=[du;-0.02]
    dv=du/norm(du)
    dim=6;ll=dim*N+1;
    V = [zeros(ll) for _ in 1:sp]
    P = zeros(sp)
    pu=s2.u;pu=[pu;sU]
    V[1]=vec(s2.u);P[1]=sU

    uu=pu+ds*dv
    μ=uu[end];p=[μ,0,0,0,0,0];u=uu[1:end-1];
    u1=u[1:end-1];ini_T=u[end];z_tol=tol*0.1;
    for i in 2:sp
        s2=LCO_solve2(u1,ini_T,eq,jeq,jeq2,p,ee,z_tol,dv,pu,ds)
        s2.err
        du=s2.u-pu;dv=du/norm(du)
        V[i]=s2.u[1:end-1];P[i]=s2.u[end]
        pu=s2.u;
        uu=pu+ds*dv;
        p=[uu[end],0,0,0,0,0];u=uu[1:end-1];
        u1=u[1:end-1];ini_T=u[end];z_tol=tol*0.1;
    end
    return (V=V,P=P)
end

function hhcat(amp) # vcat vector of vectors
    aamp=amp[1]
    for ii=2:length(amp)
        aamp=hcat(aamp,amp[ii])
    end
    return aamp
end

function generate_data(vel_l,Vel,nh,st,s_ind,U,N, θ_l) # Training data
    #Generate training data
    u0=Float32[2e-1,0,2e-1,0,0,0];
    tol=1e-7;stol=1e-8
    eq=flutter_eq_CBC;
    rp=5;ind1=1;ind2=3;
    pp=[zeros(10) for i in 1:vel_l]
    AA=zeros(vel_l,Int(nh*2+1))

    p_=zeros(9)
    p_=vcat(Vel[1],p_)
    tl=10.0
    g=get_stable_LCO(p_,u0,tl,tol,eq,stol,rp,ind1,ind2,0.0,0.0,st)

    u₀=mean(g.u[:,1]);v₀=mean(g.u[:,3])
    for i in 1:vel_l
        pp[i][1]=Vel[i]
        g=get_stable_LCO(pp[i],u0,tl,tol,eq,stol,rp,ind1,ind2,u₀,v₀,st)
        r=g.r;t=g.t;
        c=LS_harmonics(r,t,1,nh).coeff
        AA[i,:]=c
    end
    cc=[LS_harmonics(get_sol(U[s_ind[i]],N,1,3).r,get_sol(U[s_ind[i]],N,1,3).t,1,nh).coeff for i in 1:length(s_ind)]
    cc=hhcat(cc)

    AA=transpose(AA)
    Al=AA
    AA=hcat(AA,cc)
    t_series=[Transpose(get_stable_LCO([Vel[i] transpose(zeros(9))],u0,tl,tol,eq,stol,rp,ind1,ind2,u₀,v₀,st).u[:,[1,3]]) for i in 1:vel_l]
    θ_series=[get_stable_LCO([Vel[i] transpose(zeros(9))],u0,tl,tol,eq,stol,rp,ind1,ind2,u₀,v₀,st).t for i in 1:vel_l]
    θ=range(0, stop = 2π, length = θ_l)
    coθ=cos.(θ)
    siθ=sin.(θ)
    return (data=AA,ts=t_series,d0=[u₀,v₀],data2=Al,coθ=coθ,siθ=siθ,theta_s=θ_series)
end

export continuation_flutter


text:
module ExponentialAction

using LinearAlgebra
using ChainRulesCore: ChainRulesCore
using SparseArrays: sparse

include("util.jl")
include("coefficients.jl")
include("parameters.jl")
include("expv.jl")

export expv

end


text:
import Literate
Literate.notebook(joinpath("src", "stuff-i-like-about-julia.jl"))

max_stars_repo_name: Hasnep/stuff-i-like-about-julia

max_issues_repo_name: Hasnep/stuff-i-like-about-julia

max_forks_repo_name: Hasnep/stuff-i-like-about-julia


text:
export DatasetStore, studydir, path, readonly, remove

include("Utils.jl")

abstract type DatasetStore end

# mandatory functions
@mustimplement studydir(::DatasetStore)
@mustimplement readonly(::DatasetStore)

# functions for writable store
@mustimplement Base.empty!(d::DatasetStore)
@mustimplement changeParam(exp, paramName::AbstractString, paramValue)

include("Study.jl")
include("Experiment.jl")
include("MDFDatasetStore.jl")
include("Reconstruction.jl")
include("Visualization.jl")
include("BrukerDatasetStore.jl")
include("Export.jl")
include("SFDatabase.jl")


text:
# conda install -c mjirik io3d
# python -m io3d -l 3Dircadb1.1

include("../src/LarSurf.jl")
using Plasm
using LinearAlgebraicRepresentation
Lar = LinearAlgebraicRepresentation
# using LarSurf

# pth = Pio3d.datasets_join_path("medical/orig/3Dircadb1.1/MASKS_DICOM/liver")
pth = Pio3d.datasets_join_path("medical/orig/3Dircadb1.1/MASKS_DICOM/leftkidney")
datap = Pio3d.read3d(pth)

data3d = datap["data3d"]
voxelsize_mm = datap["voxelsize_mm"]
print("input voxelsize [mm]: ", voxelsize_mm)

block_size = [50, 50, 50]
margin_size = 0
step = [1, 1, 1]

## Artifical data
# data3d = ones(Int16, 3, 3, 3)
# data3d[2, 2, 2] = 10
# data3d

# voxelsize_mm = [1.0, 1.0, 1.0]
threshold=0

println("min max ", Int16(minimum(data3d)), " ", Int16(maximum(data3d)))

blocks_number, blocks_number_axis = LarSurf.number_of_blocks_per_axis(
size(data3d), block_size)
#
#
block1 = LarSurf.get_block(
    1, data3d, block_size, margin_size, blocks_number_axis, false
)

working_voxelsize_mm = voxelsize_mm .* step

verts, trifaces = LarSurf.import_data3d(data3d[1:step[1]:end, 1:step[2]:end, 1:step[3]:end], working_voxelsize_mm, 100)

print("working voxelsize [mm]: ", working_voxelsize_mm)
# println(verts)

# println(trifaces)

println("======= To LAR ========")

V, EV, FE = LarSurf.to_lar(verts, trifaces)

# println(V)
# println(EV)
# println(FE)

println("V: ", typeof(V), size(V))
println("EV: ", typeof(EV))
print("FV: ", typeof(FE))

@show EV;

# convert to list of lists
# trifaces_list = [trifaces[k, :] for k=1:size(trifaces,1)]
trifaces_list = [trifaces[k, 3:-1:1] for k=1:size(trifaces,1)]
Vt = permutedims(V, [2,1])

# Plasm.view(Lar.cuboid([1,1,1]))

Plasm.view(Vt, trifaces_list)
# Plasm.viewexploded(Vt, trifaces_list)(2,2,2)


text:
abstract GeometryPrimitive

#Geometry type.
immutable Geometry{P <: GeometryPrimitive} <: Compose3DNode
	primitives::Vector{P}
end

#Cube

immutable BoxPrimitive <: GeometryPrimitive
	center::Vec3
	width::Length
	height::Length
	depth::Length
end

typealias Box Geometry{BoxPrimitive}

BoxPrimitive(x::Length,y::Length,z::Length,w::Length,h::Length,d::Length) =
    BoxPrimitive((x,y,z),w,h,d)
BoxPrimitive(x::Length,y::Length,z::Length,side::Length) =
    BoxPrimitive((x,y,z),side,side,side)

cube(x::Length,y::Length,z::Length,side::Length) = Box([BoxPrimitive(x,y,z,side)])

cube(side::Length) = Box([BoxPrimitive(0mm,0mm,0mm,side)])

box(x::Length,y::Length,z::Length,w::Length,h::Length,d::Length) =
    Box([BoxPrimitive(x,y,z,w,h,d)])

box(w::Length,h::Length,d::Length) = Box([BoxPrimitive(0mm,0mm,0mm,w,h,d)])

function resolve(box::Absolute3DBox, cube::BoxPrimitive)
	rescenter = resolve(box,cube.center)
	width = resolve(box, cube.width)
	height = resolve(box, cube.height)
	depth = resolve(box, cube.depth)
	BoxPrimitive(rescenter,width,height,depth)
end

#Sphere

immutable SpherePrimitive <: GeometryPrimitive
	center::Vec3
	radius::Length
end

typealias Sphere Geometry{SpherePrimitive}

SpherePrimitive(x::Length,y::Length,z::Length,radius::Length) =
    SpherePrimitive((x,y,z),radius)

sphere(x::Length,y::Length,z::Length,radius::Length) =
    Sphere([SpherePrimitive(x,y,z,radius)])

sphere(radius::Length) = Sphere([SpherePrimitive(0mm,0mm,0mm,radius)])

function resolve(box::Absolute3DBox, sphere::SpherePrimitive)
	abscenter = resolve(box,sphere.center)
	absradius = resolve(box, sphere.radius)
	SpherePrimitive(abscenter,absradius)
end

#Pyramid

immutable PyramidPrimitive <: GeometryPrimitive
	corner::Vec3
	base::Length
	height::Length
end

typealias Pyramid Geometry{PyramidPrimitive}

PyramidPrimitive(x::Length,y::Length,z::Length,base::Length,height::Length) =
    PyramidPrimitive((x,y,z),base,height)

pyramid(x::Length,y::Length,z::Length,base::Length,height::Length) =
    Pyramid([PyramidPrimitive((x,y,z),base,height)])
pyramid(base::Length, height::Length) =
    Pyramid([PyramidPrimitive(0mm,0mm,0mm,base,height)])

function resolve(box::Absolute3DBox, pyramid::PyramidPrimitive)
	abscorner = resolve(box,pyramid.corner)
	absbase = resolve(box, pyramid.base)
	absheight = resolve(box, pyramid.height)
	PyramidPrimitive(abscorner,absbase,absheight)
end

#Cylinder

immutable CylinderPrimitive <: GeometryPrimitive
    center::Vec3
    topradius::Length
    bottomradius::Length
    height::Length
end

typealias Cylinder Geometry{CylinderPrimitive}

CylinderPrimitive(x::Length, y::Length, z::Length, top::Length, bottom::Length, height::Length) =
    CylinderPrimitive((x,y,z),top,bottom,height)

cylinder(x::Length,y::Length,z::Length,top::Length,bottom::Length,height::Length) =
    Cylinder([CylinderPrimitive(x,y,z,top,bottom,height)])
cylinder(top::Length,bottom::Length,height::Length) =
    Cylinder([CylinderPrimitive(0mm,0mm,0mm,top,bottom,height)])

function resolve(box::Absolute3DBox, cylinder::CylinderPrimitive)
    abscenter = resolve(box, cylinder.center)
    abstop = resolve(box, cylinder.topradius)
    absbottom = resolve(box, cylinder.bottomradius)
    absheight = resolve(box, cylinder.height)
    CylinderPrimitive(abscenter, abstop, absbottom, absheight)
end

immutable TorusPrimitive <: GeometryPrimitive
    center::Vec3
    radius::Length
    tubediameter::Length
end

typealias Torus Geometry{TorusPrimitive}

TorusPrimitive(x::Length,y::Length,z::Length,r::Length,tube::Length) =
    TorusPrimitive((x,y,z),r,tube)

torus(x::Length,y::Length,z::Length,r::Length,tube::Length) =
    Torus([TorusPrimitive(x,y,z,r,tube)])
torus(r::Length,tube::Length) = Torus([TorusPrimitive(0mm,0mm,0mm,r,tube)])

function resolve(box::Absolute3DBox, torus::TorusPrimitive)
    abscenter = resolve(box, torus.center)
    absradius = resolve(box, torus.radius)
    abstube = resolve(box, torus.tubediameter)
    TorusPrimitive(abscenter, absradius, abstube)
end

immutable ParametricPrimitive <: GeometryPrimitive
    origin::Vec3
    f::Function
    slices::Integer
    stacks::Integer
    xrange::Range
    yrange::Range
    mesh::Bool
end

typealias Parametric Geometry{ParametricPrimitive}

ParametricPrimitive(
    x::Length,y::Length,z::Length,f::Function,slices::Integer,stacks::Integer,
    xrange::Range, yrange::Range, mesh::Bool)=
    ParametricPrimitive((x,y,z),f,slices,stacks, xrange, yrange, mesh)

parametric(
    x::Length,y::Length,z::Length,f::Function,slices::Integer,stacks::Integer,
    xrange::Range, yrange::Range; mesh::Bool=false)=
        Parametric([ParametricPrimitive(x,y,z,f,slices,stacks,mesh)])

parametric(f::Function,slices::Integer,stacks::Integer,x::Range,y::Range; mesh::Bool=false)=
        Parametric([ParametricPrimitive(0mm,0mm,0mm,f,slices,stacks,x,y,mesh)])

function resolve(box::Absolute3DBox, parametricsurf::ParametricPrimitive)
    absorigin = resolve(box, parametricsurf.origin)
    ParametricPrimitive(
        absorigin,
        parametricsurf.f,
        parametricsurf.slices,
        parametricsurf.stacks,
        parametricsurf.xrange,
        parametricsurf.yrange,
        parametricsurf.mesh
    )
end


text:
LazyBox(T2::Type{T}) where {T<:Real} = LazyBox{T2}(Dict())

"""
    getindex(o::LazyBox{D}, key::Int) where D

Gets nth dim of box, creates it as unit-interval if it doesn't exist
"""
function getindex(o::LazyBox{D}, key::Int) where D
  if haskey(o.intervals,key)
    o.intervals[key]
  else
    o.intervals[key] = unit(Interval{D})
  end
end

setindex!(b::LazyBox{T}, val::Interval{T}, key::Int) where T = b.intervals[key] = val
ndims(b::LazyBox) = length(keys(b.intervals))
dims(b::LazyBox) = keys(b.intervals)
hasdim(b::LazyBox, i::Int) = haskey(b.intervals,i)

## Conversion
## ==========
convert(::Type{Vector{Interval{T}}}, b::LazyBox{T}) where T = collect(values(b.intervals))
convert(::Type{Vector{Interval}}, b::LazyBox{T}) where T = collect(values(b.intervals))
convert(::Type{Vector{Interval{T}}}, b::LazyBox{T}, dims::Vector) where T = Interval{T}[b[d] for d in dims]
convert(::Type{Vector{Interval}}, b::LazyBox{T}, dims::Vector) where T = Interval{T}[b[d] for d in dims]
convert(::Type{HyperBox}, l::LazyBox) = HyperBox(convert(Vector{Interval},l))

# ## Splitting
# ## =========
mid(b::LazyBox{T}) where T =
  Dict(dim => mid(interval) for (dim,interval) in b.intervals)

function split_box(b::LazyBox{T}, split_point::Dict{Int,T}) where T
  ks = collect(keys(b.intervals))
  intervals_set = product_boxes(b,split_point)
  [LazyBox(Dict{Int,Interval{T}}(zip(ks,intervals))) for intervals in intervals_set]
end

partial_split_box(b::LazyBox{T}, split_point::Dict{Int,T}) where T =
  split_box(b,split_point)

## Rand
"""
    struct LazyRandomVector{T<:Real}

A Vector of whose values are sampled uniformly from [0,1], but are not created until accessed (hence Lazy).
"""
struct LazyRandomVector{T<:Real}
  samples::Dict{Int64,T}
end
LazyRandomVector(T1::Type{T}) where {T<:Real} = LazyRandomVector(Dict{Int64,T1}())

function getindex(o::LazyRandomVector{T}, key::Int) where T
  if haskey(o.samples,key)
    o.samples[key]
  else
    i = rand(T)
    o.samples[key] = i
    i
  end
end

function setindex!(o::LazyRandomVector{T}, val::T, key::Int) where T
  o.samples[key] = val
end

function rand(b::LazyBox{T}) where T <: Real
  l = LazyRandomVector(T)
  for (dim,interval) in b.intervals
    l[dim] = rand(interval)
  end
  l
end

## Print
## =====
string(b::LazyBox) = b.intervals
print(io::IO, b::LazyBox) = print(io, string(b))
show(io::IO, b::LazyBox) = print(io, string(b))
showcompact(io::IO, b::LazyBox) = print(io, string(b))


text:
#
# This file is a part of MolecularGraph.jl
# Licensed under the MIT License http://opensource.org/licenses/MIT
#

export
    SmartsAtom, SmartsBond, SMARTS,
    SmilesParser, SmartsParser,
    smilestomol, smartstomol,
    associate_operations

struct SmartsAtom <: QueryAtom
    query::QueryFormula
end

struct SmartsBond <: QueryBond
    query::QueryFormula
end

SmartsBond() = SmartsBond(QueryFormula(:any, true))

SMARTS = QueryMol{SmartsAtom,SmartsBond}

mutable struct SmartsParserState{N<:AbstractNode,E<:UndirectedEdge}
    input::String
    allow_disconnected::Bool

    pos::Int
    done::Bool
    node::Int # No. of current node
    branch::Int # No. of node at the current branch root
    root::Int # No. of node at the current tree root
    ringlabel::Dict{Int,Tuple{Int,Union{E,Symbol,Nothing}}} # TODO: strange union

    edges::Vector{Tuple{Int,Int}}
    nodeattrs::Vector{N}
    edgeattrs::Vector{E}
    connectivity::Vector{Vector{Int}}

    function SmartsParserState{N,E}(str, disconn
            ) where {N<:AbstractNode,E<:UndirectedEdge}
        new(str, disconn, 1, false, 0, 1, 1, Dict(), [], [], [], [])
    end
end

SmilesParser = SmartsParserState{SmilesAtom,SmilesBond}
SmartsParser = SmartsParserState{SmartsAtom,SmartsBond}

function Base.parse(::Type{SMILES}, str::AbstractString)
    state = SmilesParser(str, true)
    fragment!(state)
    return graphmol(state.edges, state.nodeattrs, state.edgeattrs)
end

function Base.parse(::Type{SMARTS}, str::AbstractString)
    if occursin('.', str)
        state = SmartsParser(str, true)
        componentquery!(state)
    else
        state = SmartsParser(str, false)
        fragment!(state)
    end
    qmol = querymol(state.edges, state.nodeattrs, state.edgeattrs, state.connectivity)
    resolvedefaultbond!(qmol)
    return qmol
end

"""
    smilestomol(smiles::AbstractString) -> GraphMol{SmilesAtom,SmilesBond}

Parse SMILES string into `GraphMol` object.

The syntax of SMILES in this library follows both Daylight SMILES and OpenSMILES.

# References

1. OpenSMILES Specification http://opensmiles.org/spec/open-smiles.html
1. Daylight Tutorials https://www.daylight.com/dayhtml_tutorials/index.html
"""
function smilestomol(smiles::AbstractString)
    mol = parse(SMILES, smiles)
    kekulize!(mol)
    setdiastereo!(mol)
    return mol
end

"""
    smartstomol(smarts::AbstractString) -> QueryMol{SmartsAtom,SmartsBond}

Parse SMARTS string into `QueryMol` object.
"""
smartstomol(smarts::AbstractString) = parse(SMARTS, smarts)

function lookahead(state::SmartsParserState, pos::Int)
    # Negative pos can be used
    newpos = state.pos + pos
    if newpos > length(state.input) || newpos < 1
        return Char(0)
    else
        c = state.input[newpos]
        if isascii(c)
            return state.input[newpos]
        else
            throw(ErrorException("invalid charactor $(c)"))
        end
    end
end

Base.read(state) = lookahead(state, 0)  # TODO: rename

function forward!(state::SmartsParserState, num::Int)
    # Negative num can be used
    state.pos += num
    if state.pos > length(state.input)
        state.done = true
    elseif state.done
        throw(ErrorException("charactors in the buffer were used up"))
    elseif state.pos < 1
        throw(ErrorException("no more backtracking!"))
    end
end

forward!(state) = forward!(state, 1)
backtrack!(state, num) = forward!(state, -num)
backtrack!(state) = backtrack!(state, 1)

"""
    resolvedefaultbond(qmol::QueryMol) -> Nothing

Resolve default SMARTS bonds.
"""
function resolvedefaultbond!(qmol::QueryMol)
    for e in 1:edgecount(qmol)
        q = edgeattr(qmol, e).query
        q == QueryFormula(:defaultbond, true) || continue
        (u, v) = getedge(qmol, e)
        uq = nodeattr(qmol, u).query
        uarom = findformula(uq, :isaromatic)
        vq = nodeattr(qmol, v).query
        varom = findformula(vq, :isaromatic)
        arombond = QueryFormula(:isaromaticbond, true)
        single = QueryFormula(:bondorder, 1)
        if uarom === true && varom === true
            setedgeattr!(qmol, e, SmartsBond(arombond))
        elseif uarom === false || varom === false
            setedgeattr!(qmol, e, SmartsBond(single))
        else
            setedgeattr!(qmol, e, SmartsBond(QueryFormula(:or, Set([single, arombond]))))
        end
    end
end


text:
"""

Constructor for SimMesh object.
"""
struct SimMesh <: Mesh
    Nodes::Array{Float64,1}
    NBases::Int
    Fil::Dict{String,BitArray{1}}
end
function SimMesh(
    model::Model,
    Nodes::Array{Float64,1},
    NBases::Int;
    Fil::Dict{String,BitArray{1}}=Dict{String,BitArray{1}}(),
)
    NIntervals = length(Nodes) - 1 # the number of intervals
    Δ = (Nodes[2:end] - Nodes[1:end-1]) # interval width
    CellNodes = zeros(Float64, NBases, NIntervals)
    for i = 1:NIntervals
        CellNodes[:, i] .= (Nodes[i+1] + Nodes[i]) / 2
    end
    TotalNBases = NBases * NIntervals

    ## Construct the sets Fᵐ = ⋃ᵢ Fᵢᵐ, global index for sets of type m
    if isempty(Fil)
        idxPlus = model.r.r(Nodes[1:end-1].+Δ[:]/2).>0
        idxZero = model.r.r(Nodes[1:end-1].+Δ[:]/2).==0
        idxMinus = model.r.r(Nodes[1:end-1].+Δ[:]/2).<0
        for i in 1:NPhases(model)
            Fil[string(i)*"+"] = idxPlus[:,i]
            Fil[string(i)*"0"] = idxZero[:,i]
            Fil[string(i)*"-"] = idxMinus[:,i]
            if model.C[i] .<= 0
                Fil["p"*string(i)*"+"] = [model.r.r(model.Bounds[1,1])[i]].>0
                Fil["p"*string(i)*"0"] = [model.r.r(model.Bounds[1,1])[i]].==0
                Fil["p"*string(i)*"-"] = [model.r.r(model.Bounds[1,1])[i]].<0
            end
            if model.C[i] .>= 0
                Fil["q"*string(i)*"+"] = [model.r.r(model.Bounds[1,end])[i]].>0
                Fil["q"*string(i)*"0"] = [model.r.r(model.Bounds[1,end])[i]].==0
                Fil["q"*string(i)*"-"] = [model.r.r(model.Bounds[1,end])[i]].<0
            end
        end
    end
    CurrKeys = keys(Fil)
    for ℓ in ["+", "-", "0"], i = 1:NPhases(model)
        if !in(string(i) * ℓ, CurrKeys)
            Fil[string(i)*ℓ] = falses(NIntervals)
        end
        if !in("p" * string(i) * ℓ, CurrKeys) && model.C[i] <= 0
            Fil["p"*string(i)*ℓ] = falses(1)
        end
        if !in("p" * string(i) * ℓ, CurrKeys) && model.C[i] > 0
            Fil["p"*string(i)*ℓ] = falses(0)
        end
        if !in("q" * string(i) * ℓ, CurrKeys) && model.C[i] >= 0
            Fil["q"*string(i)*ℓ] = falses(1)
        end
        if !in("q" * string(i) * ℓ, CurrKeys) && model.C[i] < 0
            Fil["q"*string(i)*ℓ] = falses(0)
        end
    end
    for ℓ in ["+", "-", "0"]
        Fil[ℓ] = falses(NIntervals * NPhases(model))
        Fil["p"*ℓ] = trues(0)
        Fil["q"*ℓ] = trues(0)
        for i = 1:NPhases(model)
            idx = findall(Fil[string(i)*ℓ]) .+ (i - 1) * NIntervals
            Fil[string(ℓ)][idx] .= true
            Fil["p"*ℓ] = [Fil["p"*ℓ]; Fil["p"*string(i)*ℓ]]
            Fil["q"*ℓ] = [Fil["q"*ℓ]; Fil["q"*string(i)*ℓ]]
        end
    end
    return SimMesh(NBases, CellNodes, Fil, Δ, NIntervals, Nodes, TotalNBases, "lagrange")
end
function SimMesh()
    return SimMesh(
        0,
        Array{Real,2}(undef,0,0),
        Dict{String,BitArray{1}}(),
        Array{Float64,1}(undef,0),
        0,
        Array{Float64,1}(undef,0),
        0,
        "lagrange",
    )
end

"""
Simulates a SFM defined by `Model` until the `StoppingTime` has occured,
given the `InitialCondition` on (φ(0),X(0)).

    SimSFM(
        model::SFFM.Model,
        StoppingTime::Function,
        InitCondition::NamedTuple{(:φ, :X)},
    )

# Arguments
- `model`: A Model object
- `StoppingTime`: A function which takes the value of the process at the current
    time and at the time of the last jump of the phase process, as well as the
    `model` object.
    i.e. `StoppingTime(;model,SFM,SFM0)` where `SFM` and `SFM0` are tuples with
    keys `(:t::Float64, :φ::Int, :X::Float64, :n::Int)` which are the value of
    the SFM at the current time, and time of the previous jump of the phase
    process, repsectively. The `StoppingTime` must return a
    `NamedTuple{(:Ind, :SFM)}` type where `:Ind` is a `:Bool` value stating
    whether the stopping time has occured or not and `:SFM` is a tuple in the
    same form as the input `SFM` but which contains the value of the `SFM` at
    the stopping time.
- `InitCondition`: `NamedTuple` with keys `(:φ, :X)`, `InitCondition.φ` is a
    vector of length `M` of initial states for the phase, and `InitCondition.X`
    is a vector of length `M` of initial states for the level. `M` is the number
    of simulations to be done.

# Output
- a tuple with keys
    - `t::Array{Float64,1}` a vector of length `M` containing the values of
        `t` at the `StoppingTime`.
    - `φ::Array{Float64,1}` a vector of length `M` containing the values of
        `φ` at the `StoppingTime`.
    - `X::Array{Float64,1}` a vector of length `M` containing the values of
        `X` at the `StoppingTime`.
    - `n::Array{Float64,1}` a vector of length `M` containing the number of
        transitions of `φ` at the `StoppingTime`
"""
function SimSFM(
    model::SFFM.Model,
    StoppingTime::Function,
    InitCondition::NamedTuple{(:φ, :X)},
)
    # the transition matrix of the jump chain
    d = LinearAlgebra.diag(model.T)
    P = (model.T - LinearAlgebra.diagm(0 => d)) ./ -d
    CumP = cumsum(P, dims = 2)
    Λ = LinearAlgebra.diag(model.T)

    M = length(InitCondition.φ)
    tSims = zeros(Float64, M)
    φSims = zeros(Float64, M)
    XSims = zeros(Float64, M)
    nSims = zeros(Float64, M)

    for m = 1:M
        SFM0 = (t = 0.0, φ = InitCondition.φ[m], X = InitCondition.X[m], n = 0)
        while 1 == 1
            S = log(rand()) / Λ[SFM0.φ] # generate exp random variable
            t = SFM0.t + S
            X = UpdateXt(model, SFM0, S)
            φ = findfirst(rand() .< CumP[SFM0.φ, :])
            n = SFM0.n + 1
            SFM = (t = t, φ = φ, X = X, n = n)
            τ = StoppingTime(model, SFM, SFM0)
            if τ.Ind # if the stopping time occurs
                (tSims[m], φSims[m], XSims[m], nSims[m]) = τ.SFM
                break
            end
            SFM0 = SFM
        end
    end
    return (t = tSims, φ = φSims, X = XSims, n = nSims)
end

"""
Simulates a SFFM defined by `model` until the `StoppingTime` has occured,
given the `InitialCondition` on (φ(0),X(0),Y(0)).

    SimSFFM(
        model::SFFM.Model,
        StoppingTime::Function,
        InitCondition::NamedTuple{(:φ, :X, :Y)},
    )

# Arguments
- `model`: A Model object
- `StoppingTime`: A function which takes the value of the process at the current
    time and at the time of the last jump of the phase process, as well as the
    `model` object.
    i.e. `StoppingTime(;model,SFFM,SFFM0)` where `SFFM` and `SFFM0` are tuples
    with keys `(:t::Float64, :φ::Int, :X::Float64, :Y::Float64, :n::Int)` which
    are the value of the SFFM at the current time, and time of the previous jump
    of the phase process, repsectively. The `StoppingTime` must return a
    `NamedTuple{(:Ind, :SFFM)}` type where `:Ind` is a `:Bool` value stating
    whether the stopping time has occured or not and `:SFFM` is a tuple in the
    same form as the input `SFFM` but which contains the value of the SFFM at
    the stopping time.
- `InitCondition`: `NamedTuple` with keys `(:φ, :X, :Y)`, `InitCondition.φ` is a
    vector of length `M` of initial states for the phase, `InitCondition.X` is a
    vector of length `M` of initial states for the X-level, `InitCondition.Y` is
    a vector of length `M` of initial states for the Y-level. `M` is the number
    of simulations to be done.

# Output
- a tuple with keys
    - `t::Array{Float64,1}` a vector of length `M` containing the values of
        `t` at the `StoppingTime`.
    - `φ::Array{Float64,1}` a vector of length `M` containing the values of
        `φ` at the `StoppingTime`.
    - `X::Array{Float64,1}` a vector of length `M` containing the values of
        `X` at the `StoppingTime`.
    - `Y::Array{Float64,1}` a vector of length `M` containing the values of
        `Y` at the `StoppingTime`.
    - `n::Array{Float64,1}` a vector of length `M` containing the number of
        transitions of `φ` at the `StoppingTime`
"""
function SimSFFM(
    model::SFFM.Model,
    StoppingTime::Function,
    InitCondition::NamedTuple{(:φ, :X, :Y)},
)
    d = LinearAlgebra.diag(model.T)
    P = (model.T - LinearAlgebra.diagm(0 => d)) ./ -d
    CumP = cumsum(P, dims = 2)
    Λ = LinearAlgebra.diag(model.T)

    M = length(InitCondition.φ)
    tSims = Array{Float64,1}(undef, M)
    φSims = Array{Float64,1}(undef, M)
    XSims = Array{Float64,1}(undef, M)
    YSims = Array{Float64,1}(undef, M)
    nSims = Array{Float64,1}(undef, M)

    for m = 1:M
        SFFM0 = (
            t = 0.0,
            φ = InitCondition.φ[m],
            X = InitCondition.X[m],
            Y = InitCondition.Y[m],
            n = 0.0,
        )
        if !(model.Bounds[1, 1] <= SFFM0.X <= model.Bounds[1, 2]) ||
           !(model.Bounds[2, 1] <= SFFM0.Y <= model.Bounds[2, 2]) ||
           !in(SFFM0.φ, 1:NPhases(model))
            (tSims[m], φSims[m], XSims[m], YSims[m], nSims[m]) =
                (t = NaN, φ = NaN, X = NaN, Y = NaN, n = NaN)
        else
            while 1 == 1
                S = log(rand()) / Λ[SFFM0.φ]
                t = SFFM0.t + S
                X = UpdateXt(model, SFFM0, S)
                Y = UpdateYt(model, SFFM0, S)
                φ = findfirst(rand() .< CumP[SFFM0.φ, :])
                n = SFFM0.n + 1.0
                SFFM = (t = t, φ = φ, X = X, Y = Y, n = n)
                τ = StoppingTime(model, SFFM, SFFM0)
                if τ.Ind
                    (tSims[m], φSims[m], XSims[m], YSims[m], nSims[m]) = τ.SFFM
                    break
                end
                SFFM0 = SFFM
            end
        end
    end
    return (t = tSims, φ = φSims, X = XSims, Y = YSims, n = nSims)
end

"""
Returns ``X(t+S) = min(max(X(t) + cᵢS,0),U)`` where ``U`` is some upper bound
on the process.

    UpdateXt(
        model::SFFM.Model,
        SFM0::NamedTuple,
        S::Real,
    )

# Arguments
- `model`: a Model object
- `SFM0::NamedTuple` containing at least the keys `:X` giving the value of
    ``X(t)`` at the current time, and `:φ` giving the value of
    ``φ(t)`` at the current time.
- `S::Real`: an elapsed amount of time to evaluate ``X`` at, i.e. ``X(t+S)``.
"""
function UpdateXt(
    model::SFFM.Model,
    SFM0::NamedTuple,
    S::Real,
)
    # given the last position of a SFM, SFM0, a time step of size s, find the
    # position of X at time t
    X = min(
        max(SFM0.X + model.C[SFM0.φ] * S, model.Bounds[1, 1]),
        model.Bounds[1, 2],
    )
    return X
end

"""
Returns ``Y(t+S)`` given ``Y(t)``.

    UpdateYt(
        model::SFFM.Model,
        SFFM0::NamedTuple{(:t, :φ, :X, :Y, :n)},
        S::Real,
    )

# Arguments
- `model`: a Model object
- `SFFM0::NamedTuple` containing at least the keys `:X` giving the value of
    ``X(t)`` at the current time, and `:Y` giving the value of ``Y(t)`` at the
    current time, and `:φ` giving the value of `φ(t)`` at the current time.
- `S::Real`: an elapsed amount of time to evaluate ``X`` at, i.e. ``X(t+S)``.
"""
function UpdateYt(
    model::SFFM.Model,
    SFFM0::NamedTuple{(:t, :φ, :X, :Y, :n)},
    S::Real,
)
    # given the last position of a SFFM, SFFM0, a time step of size s, find the
    # position of Y at time t
    if model.C[SFFM0.φ] == 0
        Y = SFFM0.Y + S * model.r.r(SFFM0.X)[SFFM0.φ]
    else
        X = UpdateXt(model, SFFM0, S)
        ind = (X.==model.Bounds[1, :])[:]
        if any(ind)
            # time at which Xt hits u or v
            t0 = (model.Bounds[1, ind][1] - SFFM0.X) / model.C[SFFM0.φ]
            Y =
                SFFM0.Y +
                (model.r.R(X)[SFFM0.φ] - model.r.R(SFFM0.X)[SFFM0.φ]) /
                model.C[SFFM0.φ] +
                model.r.r(model.Bounds[1, ind][1])[SFFM0.φ] * (S - t0)
        else
            Y =
                SFFM0.Y +
                (model.r.R(X)[SFFM0.φ] - model.r.R(SFFM0.X)[SFFM0.φ]) /
                model.C[SFFM0.φ]
        end
    end
    return Y
end

"""
Constructs the `StoppingTime` ``1(t>T)``

    FixedTime( T::Real)

# Arguments
- `T`: a time at which to stop the process

# Output
- `FixedTimeFun`: a function with two methods
    - `FixedTimeFun(
        model::SFFM.Model,
        SFM::NamedTuple{(:t, :φ, :X, :n)},
        SFM0::NamedTuple{(:t, :φ, :X, :n)},
    )`: a stopping time for a SFM.
    - `FixedTimeFun(
        model::SFFM.Model,
        SFFM::NamedTuple{(:t, :φ, :X, :Y, :n)},
        SFFM0::NamedTuple{(:t, :φ, :X, :Y, :n)},
    )`: a stopping time for a SFFM
"""
function FixedTime( T::Real)
    # Defines a simple stopping time, 1(t>T).
    # SFM method
    function FixedTimeFun(
        model::SFFM.Model,
        SFM::NamedTuple{(:t, :φ, :X, :n)},
        SFM0::NamedTuple{(:t, :φ, :X, :n)},
    )
        Ind = SFM.t > T
        if Ind
            s = T - SFM0.t
            X = UpdateXt(model, SFM0, s)
            SFM = (t = T, φ = SFM0.φ, X = X, n = SFM0.n)
        end
        return (Ind = Ind, SFM = SFM)
    end
    # SFFM METHOD
    function FixedTimeFun(
        model::SFFM.Model,
        SFFM::NamedTuple{(:t, :φ, :X, :Y, :n)},
        SFFM0::NamedTuple{(:t, :φ, :X, :Y, :n)},
    )
        Ind = SFFM.t > T
        if Ind
            s = T - SFFM0.t
            X = UpdateXt(model, SFFM0, s)
            Y = UpdateYt(model, SFFM0, s)
            SFFM = (T, SFFM0.φ, X, Y, SFFM0.n)
        end
        return (Ind = Ind, SFFM = SFFM)
    end
    return FixedTimeFun
end

"""
Constructs the `StoppingTime` ``1(N(t)>n)`` where ``N(t)`` is the number of
jumps of ``φ`` by time ``t``.

    NJumps( N::Int)

# Arguments
- `N`: a desired number of jumps

# Output
- `NJumpsFun`: a function with two methods
    - `NJumpsFun(
        model::SFFM.Model,
        SFM::NamedTuple{(:t, :φ, :X, :n)},
        SFM0::NamedTuple{(:t, :φ, :X, :n)},
    )`: a stopping time for a SFM.
    - `NJumpsFun(
        model::SFFM.Model,
        SFFM::NamedTuple{(:t, :φ, :X, :Y, :n)},
        SFFM0::NamedTuple{(:t, :φ, :X, :Y, :n)},
    )`: a stopping time for a SFFM
"""
function NJumps( N::Int)
    # Defines a simple stopping time, 1(n>N), where n is the number of jumps of φ.
    # SFM method
    function NJumpsFun(
        model::SFFM.Model,
        SFM::NamedTuple{(:t, :φ, :X, :n)},
        SFM0::NamedTuple{(:t, :φ, :X, :n)},
    )
        Ind = SFM.n >= N
        return (Ind = Ind, SFM = SFM)
    end
    # SFFM method
    function NJumpsFun(
        model::SFFM.Model,
        SFFM::NamedTuple{(:t, :φ, :X, :Y, :n)},
        SFFM0::NamedTuple{(:t, :φ, :X, :Y, :n)},
    )
        Ind = SFFM.n >= N
        return (Ind = Ind, SFFM = SFFM)
    end
    return NJumpsFun
end

"""
Constructs the `StoppingTime` which is the first exit of the process ``X(t)``
from the interval ``[u,v]``.

    FirstExitX( u::Real, v::Real)

# Arguments
- `u`: a lower boundary
- `v`: an upper boundary

# Output
- `FirstExitXFun`: a function with two methods
    - `FirstExitXFun(
        model::SFFM.Model,
        SFM::NamedTuple{(:t, :φ, :X, :n)},
        SFM0::NamedTuple{(:t, :φ, :X, :n)},
    )`: a stopping time for a SFM.
    - `FirstExitXFun(
        model::SFFM.Model,
        SFFM::NamedTuple{(:t, :φ, :X, :Y, :n)},
        SFFM0::NamedTuple{(:t, :φ, :X, :Y, :n)},
    )`: a stopping time for a SFFM
"""
function FirstExitX( u::Real, v::Real)
    # SFM Method
    function FirstExitXFun(
        model::SFFM.Model,
        SFM::NamedTuple{(:t, :φ, :X, :n)},
        SFM0::NamedTuple{(:t, :φ, :X, :n)},
    )
        Ind = ((SFM.X > v) || (SFM.X < u))
        if Ind
            if SFM.X > v
                X = v
            else
                X = u
            end
            s = (X - SFM0.X) / model.C[SFM0.φ] # can't exit with c = 0
            t = SFM0.t + s
            SFM = (t = t, φ = SFM0.φ, X = X, n = SFM0.n)
        end
        return (Ind = Ind, SFM = SFM)
    end
    # SFFM Method
    function FirstExitXFun(
        model::NamedTuple{(:T, :C, :r, :Bounds, :NPhases, :SDict, :TDict)},
        SFFM::NamedTuple{(:t, :φ, :X, :Y, :n)},
        SFFM0::NamedTuple{(:t, :φ, :X, :Y, :n)},
    )
        Ind = ((SFFM.X > v) || (SFFM.X < u))
        if Ind
            if (SFFM.X > v)
                X = v
            else
                X = u
            end
            s = (X - SFFM0.X) / model.C[SFFM0.φ] # can't exit with c = 0.
            t = SFFM0.t + s
            Y = UpdateYt(model, SFFM0, s)
            SFFM = (t, SFFM0.φ, X, Y, SFFM0.n)
        end
        return (Ind = Ind, SFFM = SFFM)
    end
    return FirstExitXFun
end

"""
Constructs the `StoppingTime` which is the first exit of the process ``Y(t)``
from the interval ``[u,v]``. ASSUMES ``Y(t)`` is monotonic between jumps.

    FirstExitY( u::Real, v::Real)

# Arguments
- `u`: a lower boundary
- `v`: an upper boundary

# Output
- `FirstExitYFun`: a function with one method
    - `FirstExitYFun(
        model::SFFM.Model,
        SFFM::NamedTuple{(:t, :φ, :X, :Y, :n)},
        SFFM0::NamedTuple{(:t, :φ, :X, :Y, :n)},
    )`: a stopping time for a SFFM
"""
function FirstExitY( u::Real, v::Real)
    # SFFM Method
    function FirstExitYFun(
        model::SFFM.Model,
        SFFM::NamedTuple{(:t, :φ, :X, :Y, :n)},
        SFFM0::NamedTuple{(:t, :φ, :X, :Y, :n)},
    )
        Ind = ((SFFM.Y < u) || (SFFM.Y > v))
        if Ind
            idx = [SFFM.Y < u; SFFM.Y > v]
            boundaryHit = [u;v][idx][1]
            YFun(t) = UpdateYt(model, SFFM0, t) - boundaryHit
            S = SFFM.t - SFFM0.t
            tstar = fzero(YFun, 0, S)
            X = UpdateXt(model, SFFM0, tstar)
            t = SFFM0.t + tstar
            Y = boundaryHit
            SFFM = (t, SFFM0.φ, X, Y, SFFM0.n)
        end
        return (Ind = Ind, SFFM = SFFM)
    end
    return FirstExitYFun
end

"""
Finds zero of `f` using the bisection method on the interval `[a,b]` with
error `err`.

    fzero( f::Function, a::Real, b::Real; err::Float64 = 1e-8)

"""
function fzero( f::Function, a::Real, b::Real; err::Float64 = 1e-6)
    # finds zeros of f using the bisection method
    c = a + (b - a) / 2
    while a < c < b
        fc = f(c)
        if abs(fc) < err
            break
        end
        p = f(a) * fc
        if p < 0
            a, b = a, c
        elseif p > 0
            a, b = c, b
        elseif fc == 0
            break
        else
            a = a+sqrt(eps())
        end
        c = a + (b - a) / 2
    end
    return c
end

"""
Convert from simulations of a SFM or SFFM to a distribution.

    Sims2Dist(
        model::SFFM.Model,
        mesh::SFFM.Mesh,
        sims::NamedTuple;
        type::String = "density",
    )

# Arguments
- `Model`: a Model object
- `mesh`: a Mesh object as output from `MakeMesh`
- `sims::Array`: a named tuple as output of `SFFMSim` or `SFMSim`
- `type::String`: an (optional) declaration of what type of distribution you
    want to convert to. Options are `"probability"` to return the probabilities
    ``P(X(t)∈ D_k, φ(t) = i)`` where ``D_k``is the kth cell, `"cumulative"` to
    return the CDF evaluated at cell edges, or `"density"` to return an
    approximation to the density ar at the CellNodes(mesh).

# Output
- a tuple with keys
(pm=pm, distribution=yvals, x=xvals, type=type)
    - `pm::Array{Float64}`: a vector containing the point masses, the first
        `sum(model.C.<=0)` entries are the left hand point masses and the last
        `sum(model.C.>=0)` are the right-hand point masses.
    - `distribution::Array{Float64,3}`:
        - if `type="cumulative"` returns a `2×NIntervals×NPhases` array
            containing the CDF evaluated at the cell edges as contained in
            `x` below. i.e. `distribution[1,:,i]` returns the cdf at the
            left-hand edges of the cells in phase `i` and `distribution[2,:,i]`
            at the right hand edges.
        - if `type="probability"` returns a `1×NIntervals×NPhases` array
            containing the probabilities ``P(X(t)∈ D_k, φ(t) = i)`` where ``D_k``
            is the kth cell.
        - if `type="density"` returns a `NBases×NIntervals×NPhases` array
            containing the kde estimate of the density function evaluated at the
            cell nodes as contained in `x` below.
    - `x::Array{Float64,2}`:
        - if `type="cumulative"` returns a `2×NIntervals×NPhases` array
            containing the cell edges as contained. i.e. `x[1,:]`
            returns the left-hand edges of the cells and `x[2,:]` at the
            right-hand edges.
        - if `type="probability"` returns a `1×NIntervals×NPhases` array
            containing the cell centers.
        - if `type="density"` returns a `NBases×NIntervals×NPhases` array
            containing the cell nodes.
    - `type`: as input in arguments.
"""
function Sims2Dist(
    model::SFFM.Model,
    mesh::SFFM.Mesh,
    sims::NamedTuple,
    type::Type{T} = SFFMProbability,
) where {T<:SFFMDistribution}

    if type == SFFMDensity
        distribution = zeros(Float64, NBases(mesh), NIntervals(mesh), NPhases(model))
    elseif type == SFFMProbability
        distribution = zeros(Float64, 1, NIntervals(mesh), NPhases(model))
    elseif type == SFFMCDF
        distribution = zeros(Float64, 2, NIntervals(mesh), NPhases(model))
    end
    pm = zeros(Float64, sum(model.C .<= 0) + sum(model.C .>= 0))
    pc = 0
    qc = 0
    xvals = CellNodes(mesh)
    for i = 1:NPhases(model)
        # find the simluated value of imterest for this iteration
        whichsims =
            (sims.φ .== i) .&
            (sims.X .!= model.Bounds[1, 1]) .&
            (sims.X .!= model.Bounds[1, end])
        data = sims.X[whichsims]
        totalprob = sum(whichsims) / length(sims.φ)
        if type == SFFMProbability
            if length(data)!=0
                h = StatsBase.fit(StatsBase.Histogram, data, mesh.Nodes)
                h = h.weights ./ sum(h.weights) * totalprob
                distribution[:, :, i] = h
            end
            if NBases(mesh) == 1
                xvals = CellNodes(mesh)
            else
                xvals = CellNodes(mesh)[1, :] + Δ(mesh) / 2
            end
        elseif type == SFFMDensity
            if length(data)!=0
                if model.Bounds[1, end] == Inf
                    U = KernelDensity.kde(data)
                else
                    U = KernelDensity.kde(
                        data,
                        boundary = (model.Bounds[1, 1], model.Bounds[1, end]),
                    )
                end
                distribution[:, :, i] =
                    reshape(
                        KernelDensity.pdf(U, CellNodes(mesh)[:])',
                        NBases(mesh),
                        NIntervals(mesh),
                    ) * totalprob
            end
        elseif type == SFFMCDF
            if length(data)!=0
                h = StatsBase.fit(StatsBase.Histogram, data, mesh.Nodes)
                tempDist = h.weights ./ sum(h.weights) * totalprob
                tempDist = cumsum(tempDist)
                distribution[1, 2:end, i] = tempDist[1:end-1]
                distribution[2, :, i] = tempDist
                if NBases(mesh) == 1
                    xvals = [mesh.Nodes[1:end-1]';mesh.Nodes[2:end]']
                else
                    xvals = CellNodes(mesh)[[1;end], :]
                end
            end
        end

        if model.C[i] <= 0
            pc = pc + 1
            whichsims = (sims.φ .== i) .& (sims.X .== model.Bounds[1, 1])
            p = sum(whichsims) / length(sims.φ)
            pm[pc] = p
            if type == SFFMCDF
                distribution[:,:,i] = distribution[:,:,i] .+ pm[pc]
            end
        end
        if model.C[i] >= 0
            qc = qc + 1
            whichsims = (sims.φ .== i) .& (sims.X .== model.Bounds[1, end])
            p = sum(whichsims) / length(sims.φ)
            if type == SFFMCDF
                pm[sum(model.C .<= 0)+qc] = p + distribution[end,end,i]
            else
                pm[sum(model.C .<= 0)+qc] = p
            end
        end
    end
    if type == SFFMDensity && NBases(mesh) == 1
        distribution = [1; 1] .* distribution
        xvals = [mesh.Nodes[1:end-1]';mesh.Nodes[2:end]']
    end
    return type(pm, distribution, xvals)
end

max_stars_repo_name: angus-lewis/SFFM

max_issues_repo_name: angus-lewis/SFFM

max_forks_repo_name: angus-lewis/SFFM


text:
using OrdinaryDiffEq, Test
using DiffEqProblemLibrary.ODEProblemLibrary: importodeproblems; importodeproblems()
import DiffEqProblemLibrary.ODEProblemLibrary: van

@testset "Stiffness Detection Tests" begin
  prob1 = ODEProblem(van,[0,2.],(0.0,6),inv(0.003))
  # out-of-place test
  function _van(u, p, t)
    μ = p[1]
    [μ*((1-u[2]^2)*u[1] - u[2]),
     1*u[1]]
  end
  prob2 = ODEProblem(_van,[0,2.],(0.0,6),inv(0.003))
  probArr = [prob1, prob2]
  # Test if switching back and forth
  is_switching_fb(sol) = maximum(diff(findall(x->x==2, sol.alg_choice))) > 5
  for prob in probArr
    alg = AutoTsit5(Rodas5(); maxstiffstep=5, maxnonstiffstep=5, stiffalgfirst=true)
    sol = solve(prob, alg)
    @test length(sol.t) < 280
    @test typeof(alg.algs[sol.alg_choice[1]]) <: Rodas5
    @test is_switching_fb(sol)
    sol = solve(prob, AutoDP5(Rodas5(); maxstiffstep=2, maxnonstiffstep=2,
                              stifftol=11//10, nonstifftol=9//10),
                              reltol=1e-5, abstol=1e-5)
    @test length(sol.t) < 625
    @test is_switching_fb(sol)

    sol = solve(prob,AutoVern6(Kvaerno3(); maxstiffstep=4, maxnonstiffstep=4))
    @test length(sol.t) < 690
    @test is_switching_fb(sol)
    sol = solve(prob,AutoVern7(Hairer42(); maxstiffstep=4, maxnonstiffstep=4))
    @test length(sol.t) < 540
    @test is_switching_fb(sol)
    sol = solve(prob,AutoVern8(Rosenbrock23(); maxstiffstep=4, maxnonstiffstep=4))
    @test length(sol.t) < 910
    @test is_switching_fb(sol)
    sol = solve(prob,AutoVern9(KenCarp3(); maxstiffstep=4, maxnonstiffstep=4))
    @test length(sol.t) < 470
    @test is_switching_fb(sol)
  end
end


text:
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#
# Description
# ==============================================================================
#
#   J2 orbit propagator algorithm.
#
#   This algorithm propagates the orbit considering the perturbed two-body
#   equations as presented in [1, p. 690-692]. It uses the first-order
#   approximation of Kepler's problem, considering the effects of secular
#   gravitational and drag perturbations.
#
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#
# References
# ==============================================================================
#
#   [1] Vallado, D. A (2013). Fundamentals of Astrodynamics and Applications.
#       Microcosm Press, Hawthorn, CA, USA.
#
#   [2] Wertz, J. R (1978). Spacecraft attitude determination and control.
#       Kluwer Academic Publishers, Dordrecht, The Netherlands.
#
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #

export j2_gc_egm08, j2_gc_egm96, j2_gc_jgm02, j2_gc_jgm03
export j2_gc_egm08_f32, j2_gc_egm96_f32, j2_gc_jgm02_f32, j2_gc_jgm03_f32
export j2_init, j2!

################################################################################
#                                     TODO
################################################################################
#
# 1. Analyze the reference frame representation of the inputs for this
#    algorithm.
#
#   The SGP4 algorithm expects that the input parameters are represented in the
#   TEME (true equator, mean equinox) reference frame. This J2 orbit propagator
#   model requires that the input parameters are consistent with the
#   gravitational perturbation theory in which the `J2` coefficient was
#   computed. Looking at [1, p. 642], it appears that the perturbations are
#   considering a frame in which the Z-axis is aligned with the CIP (Celestial
#   Intermediate Pole, or the Earth rotation axis). Hence, the J2 parameter is
#   defined based on the PEF. Since no rotations or adaptations are programmed,
#   then the input parameters for this propagator should be represented in any
#   reference frame with a true Equator, because of the symmetry.
#
#   This needs to be further analyzed and confirmed.
#
################################################################################

################################################################################
#                                  Constants
################################################################################

# These constants were obtained from the GFC files. Remember that:
#
#   J_n = -C_n,0 * sqrt(2n+1)
#

# EGM-08 Gravitational constants.
const j2_gc_egm08 = J2_GravCte(
    6378137.0,
    sqrt(3.986004415e14/6378137.0^3),
    0.0010826261738522227
)

const j2_gc_egm08_f32 = J2_GravCte{Float32}(
    6378137.0,
    sqrt(3.986004415e14/6378137.0^3),
    0.0010826261738522227
)

# EGM-96 Gravitational constants.
const j2_gc_egm96 = J2_GravCte(
    6378136.3,
    sqrt(3.986004415e14/6378136.3^3),
    0.0010826266835531513
)

const j2_gc_egm96_f32 = J2_GravCte{Float32}(
    6378136.3,
    sqrt(3.986004415e14/6378136.3^3),
    0.0010826266835531513
)

# JGM-02 Gravitational constants.
const j2_gc_jgm02 = J2_GravCte(
    6378136.3,
    sqrt(3.986004415e14/6378136.3^3),
    0.0010826269256388149
)

const j2_gc_jgm02_f32 = J2_GravCte{Float32}(
    6378136.3,
    sqrt(3.986004415e14/6378136.3^3),
    0.0010826269256388149
)

# JGM-03 Gravitational constants.
const j2_gc_jgm03 = J2_GravCte(
    6378136.3,
    sqrt(3.986004415e14/6378136.3^3),
    0.0010826360229829945
)

const j2_gc_jgm03_f32 = J2_GravCte{Float32}(
    6378136.3,
    sqrt(3.986004415e14/6378136.3^3),
    0.0010826360229829945
)

################################################################################
#                                  Functions
################################################################################

"""
    j2_init(epoch::Tepoch, a_0::Number, e_0::Number, i_0::Number, Ω_0::Number, ω_0::Number, f_0::Number, dn_o2::Number = 0, ddn_o6::Number = 0; j2_gc::J2_GravCte{T} = j2_gc_egm08) where {Tepoch, T}

Initialize the data structure of J2 orbit propagator algorithm.

!!! note
    The type used in the propagation will be the same as used to define the
    gravitational constants in the structure `j2_gc`.

# Args

- `epoch::Number`: Epoch of the initial mean orbital elements [Julian Day].
- `a_0::Number`: Initial mean semi-major axis [m].
- `e_0::Number`: Initial mean eccentricity.
- `i_0::Number`: Initial mean inclination [rad].
- `Ω_0::Number`: Initial mean right ascension of the ascending node [rad].
- `ω_0::Number`: Initial mean argument of perigee [rad].
- `f_0::Number`: Initial mean true anomaly [rad].
- `dn_o2::Number`: First time derivative of the mean motion divided by two
    [rad/s^2].
- `ddn_o6::Number`: Second time derivative of the mean motion divided by six
    [rad/s^3].

# Keywords

- `j2_gc::J2_GravCte`: J2 orbit propagator gravitational constants (see
    [`J2_GravCte`](@ref)). (**Default** = `j2_gc_egm08`)

# Returns

The structure [`J2_Structure`](@ref) with the initialized parameters.
"""
function j2_init(
    epoch::Tepoch,
    a_0::Number,
    e_0::Number,
    i_0::Number,
    Ω_0::Number,
    ω_0::Number,
    f_0::Number,
    dn_o2::Number = 0,
    ddn_o6::Number = 0;
    j2_gc::J2_GravCte{T} = j2_gc_egm08
) where {Tepoch, T}
    # Unpack the gravitational constants to improve code readability.
    @unpack_J2_GravCte j2_gc

    # Convert all inputs to the correct type.
    a_0    = T(a_0)
    e_0    = T(e_0)
    i_0    = T(i_0)
    Ω_0    = T(Ω_0)
    ω_0    = T(ω_0)
    f_0    = T(f_0)
    dn_o2  = T(dn_o2)
    ddn_o6 = T(ddn_o6)

    # Initial values.
    al_0 = a_0 / R0              # ............. Normalized semi-major axis [er]
    n_0  = μm / al_0^(T(3 / 2))  # ............. Unperturbed mean motion [rad/s]
    p_0  = al_0 * (1 - e_0^2)    # ...................... Semi-latus rectum [er]
    M_0  = f_to_M(e_0, f_0)      # .................. Initial mean anomaly [rad]

    # Auxiliary variables.
    dn   = 2dn_o2            # ..... Time-derivative of the mean motion [rad/s²]
    p_0² = p_0^2
    e_0² = e_0^2

    sin_i_0, cos_i_0 = sincos(i_0)
    sin_i_0² = sin_i_0^2

    # First-order time-derivative of the orbital elements.
    #
    # See [1, p 692].

    δa   = +T(2 / 3) * al_0 * dn / n_0
    δe   = +T(2 / 3) * (1 - e_0) * dn / n_0
    δΩ   = -T(3 / 2) * n_0 * J2 / p_0² * cos_i_0
    δω   = +T(3 / 4) * n_0 * J2 / p_0² * (4 - 5sin_i_0²)
    δM_0 = +T(3 / 4) * n_0 * J2 / p_0² * sqrt(1 - e_0²) * (2 - 3sin_i_0²)

    # Create the output structure with the data.
    return J2_Structure{Tepoch, T}(
        epoch  = epoch,
        al_0   = al_0,
        n_0    = n_0,
        e_0    = e_0,
        i_0    = i_0,
        Ω_0    = Ω_0,
        ω_0    = ω_0,
        f_0    = f_0,
        M_0    = M_0,
        dn_o2  = dn_o2,
        ddn_o6 = ddn_o6,
        j2_gc  = j2_gc,
        Δt     = 0,
        al_k   = al_0,
        e_k    = e_0,
        i_k    = i_0,
        Ω_k    = Ω_0,
        ω_k    = ω_0,
        f_k    = f_0,
        M_k    = M_0,
        δa     = δa,
        δe     = δe,
        δΩ     = δΩ,
        δω     = δω,
        δM_0   = δM_0,
    )
end

"""
    j2!(j2d::J2_Structure{Tepoch, T}, t::Number) where {Tepoch, T}

Propagate the orbit defined in `j2d` (see [`J2_Structure`](@ref)) until the time
`t` [s].

!!! note
    The internal values in `j2d` will be modified.

# Returns

- The position vector represented in the inertial frame at time `t` [m].
- The velocity vector represented in the inertial frame at time `t` [m/s]

# Remarks

The inertial frame in which the output is represented depends on which frame it
was used to generate the orbit parameters. Notice that the perturbation theory
requires an inertial frame with true equator.
"""
function j2!(j2d::J2_Structure{Tepoch, T}, t::Number) where {Tepoch, T}
    # Unpack the variables.
    @unpack al_0, n_0, e_0, i_0, Ω_0, ω_0, f_0, M_0, dn_o2, ddn_o6 = j2d
    @unpack δa, δe, δΩ, δω, δM_0, j2_gc = j2d
    @unpack R0, μm, J2 = j2_gc

    # Time elapsed since epoch.
    Δt = T(t)

    # Propagate the orbital elements.
    al_k = al_0 - δa * Δt
    e_k  = e_0 - δe * Δt
    i_k  = i_0
    Ω_k  = mod(Ω_0 + δΩ * Δt, T(2π))
    ω_k  = mod(ω_0 + δω * Δt, T(2π))

    # In [1, p. 692], it is mentioned that the mean anomaly must be updated
    # considering the equation:
    #                           .           ..
    #                          n_0          n_0
    #   M_k = M_0 + n_0 ⋅ Δt + ---- ⋅ Δt² + ---- ⋅ Δt³
    #                           2            6
    #
    # However, the time derivative of M_0 for perturbed orbits considering terms
    # up to J2 is not 0. This is clear in [1, p. 652].
    #
    # Hence, we will use the equation:
    #                                     .           ..
    #                .                    n_0          n_0
    #   M_k = M_0 + M_0 ⋅ Δt + n_0 ⋅ Δt + ---- ⋅ Δt² + ---- ⋅ Δt³
    #                                      2            6
    #

    M_k = mod(@evalpoly(Δt, M_0, (n_0 + δM_0), dn_o2, ddn_o6), T(2π))
    f_k = M_to_f(e_k, M_k)

    # Make sure that eccentricity is not lower than 0.
    e_k < 0 && (e_k = T(0))

    # Compute the position and velocity vectors given the orbital elements.
    r_i_k, v_i_k = kepler_to_rv(al_k * R0, e_k, i_k, Ω_k, ω_k, f_k)

    # Update the J2 orbit propagator structure.
    @pack! j2d = Δt, al_k, e_k, i_k, Ω_k, ω_k, M_k, f_k

    # Return the position and velocity vector represented in the inertial
    # reference frame.
    return r_i_k, v_i_k
end


text:
include("matrixcompletion.jl")

if !isdir("data")
  print("Downloading data...")
  run(`./get-data.sh`)
  run(`mkdir output`)
  println("done.")
end

if !isfile("c/libmc.so.1.0")
  print("Compiling c library...")
  run(`./make-lib.sh`)
  println("done.")
end

train = readdlm("data/jellyfish/datasets/netflix-dataset/netflix_train_set.dat",' ', Int64; dims=(100198805, 3))
test = readdlm("data/jellyfish/datasets/netflix-dataset/netflix_probe_set.dat",' ', Int64; dims=(281702, 3))

println(size(train))
println("training on $(size(train,1)) examples, testing on $(size(test,1)).")

m,k = maximum(train,1)[1],maximum(train,1)[2]

test_sim = MatrixCompletion(m,k,size(test,1),test[:,1],test[:,2])
sim = MatrixCompletion(m,k,size(train,1),train[:,1], train[:,2])

n_train = size(train,1)

target = copy(float(train[:,3]))
train_mean = 3.35
target.-= train_mean

n_test = size(test,1)
test = vec(test[:,3])

const tau = 100000.0

start_time = time()
history = zeros(4,0)
function evalTestRMSE(old_points, points, weights,output, ov)
  test_out = runForwardSimulator(test_sim, points)*weights .+ train_mean
  test_out_clip = copy(test_out)
  test_out_clip[test_out_clip .>= 5.0] = 5.0
  test_out_clip[test_out_clip .<= 1.0] = 1.0
  println(sum(weights)," : ",tau)
  t_rmse = sqrt(sumabs2(test_out_clip - test)/n_test)
  println("clipped test RMSE: ", t_rmse)
  println("objective value on: ", ov)
  println("rank: ", length(weights))
  println("time: $(time() - start_time)")
  global history = [history [ov, t_rmse, length(weights), time() - start_time]]
  return false
end

#run with improvement
(t,w) = ADCG(sim, LSLoss(), target , tau;  callback = evalTestRMSE,max_iters = 20, max_cd_iters = 200)
writecsv("output/fwami.csv",history')

#ablations
#fwa-m
history = zeros(4,0)
(t,w) = ADCG(sim, LSLoss(), target , tau;  callback = evalTestRMSE,max_iters = 20, max_cd_iters = 0)
writecsv("output/fwam.csv",history')

#gf
history = zeros(4,0)
(t,w) = ADCG(sim, LSLoss(), target , tau;  callback = evalTestRMSE,max_iters = 20, max_cd_iters = 1)
writecsv("output/gf.csv",history')


text:
@testset "1310.xor-queries-of-a-subarray.jl" begin
    @test xor_queries([1, 3, 4, 8], [[0, 1], [1, 2], [0, 3], [3, 3]]) == [2, 7, 14, 8]
    @test xor_queries([4, 8, 2, 10], [[2, 3], [1, 3], [0, 0], [0, 3]]) == [8, 0, 4, 4]
end