text:
# Calculates percent identity of two aligned sequences
# No account of the positions with gaps in both sequences in the length
"""
seq1 and seq2 should have the same len
"""
function _percentidentity(seq1, seq2, len)
    count = zero(Int)
    colgap = zero(Int)
    @inbounds for i in 1:len
        if seq1[i] == seq2[i]
            count += one(Int)
            colgap += Int(seq1[i] == GAP)
        end
    end
    100.0 * (count-colgap)/(len-colgap)
end

"""
`percentidentity(seq1, seq2)`

Calculates the fraction of identities between two aligned sequences.
The identity value is calculated as the number of identical characters in the i-th position of both
sequences divided by the length of both sequences.
Positions with gaps in both sequences are not counted in the length of the sequence.
"""
function percentidentity(seq1, seq2)
    len = length(seq1)
    if len != length(seq2)
        throw(ErrorException("Sequences of different length, they aren't aligned or don't come from the same alignment"))
    end
    _percentidentity(seq1, seq2, len)
end

"""
`percentidentity(seq1, seq2, threshold)`

Computes quickly if two aligned sequences have a identity value greater than a given `threshold` value.
Returns a boolean value.
"""
function percentidentity(seq1, seq2, threshold)
    fraction = threshold / 100.0
    len = length(seq1)
    if len != length(seq2)
        throw(ErrorException("Sequences of different length, they aren't aligned or don't come from the same alignment"))
    end
    n = len
    limit_count = n * fraction
    diff = 0
    count = 0
    @inbounds for i in 1:len
        if seq1[i] == seq2[i]
            if seq1[i] != GAP
                count += 1
                if count >= limit_count
                    return(true)
                end
            else
                n -= 1
                limit_count = n * fraction
            end
        else
            diff += 1
            if diff > n - limit_count
                return(false)
            end
        end
    end
    (count/n) >= fraction
end

# percentidentity for a MSA
# =========================

"""
aln should be transpose(msa)
"""
function _percentidentity_kernel!(scores, aln, nseq, len)
    k = 0
    list = scores.list
    @inbounds for i in 1:(nseq-1)
        a = aln[i]
        for j in (i+1):nseq
            list[k += 1] = _percentidentity(a, aln[j], len)
        end
    end
    scores
end

"""
`percentidentity(msa[, out::Type=Float64])`

Calculates the identity between all the sequences on a MSA.
You can indicate the output element type with the last optional parameter (`Float64` by default).
For a MSA with a lot of sequences, you can use `Float32` or `Flot16` in order to avoid the `OutOfMemoryError()`.
"""
function percentidentity{T}(msa::AbstractMatrix{Residue}, out::Type{T}=Float64)
    aln = getresiduesequences(msa)
    nseq = length(aln)
    len = length(aln[1])
    scores = sequencepairsmatrix(msa, T, false, T(100.0))
    _percentidentity_kernel!(scores, aln, nseq, len)
    scores
end

# Mean Percent Identity of an MSA
# ===============================

"""
Returns the mean of the percent identity between the sequences of a MSA.
If the MSA has 300 sequences or less, the mean is exact.
If the MSA has more sequences and the `exact` keyword is `false` (defualt),
44850 random pairs of sequences are used for the estimation.
The number of samples can be changed using the second argument.
Use `exact=true` to perform all the pairwise comparison (the calculation could be slow).
"""
function meanpercentidentity(msa, nsamples::Int=44850; exact::Bool=false) # lengthlist(300, false) == 44850
    nseq, ncol = size(msa)
    nvalues = lengthlist(nseq, Val{false})
    sum = 0.0
    if !exact && nvalues >= nsamples
        samples = Set{Tuple{Int,Int}}()
        sizehint!(samples, nsamples)
        while length(samples) < nsamples
            i = rand(1:(nseq-1))
            j = rand((i+1):nseq)
            if !in((i,j), samples)
                push!(samples, (i,j))
                sum += percentidentity(msa[i,:], msa[j,:])
            end
        end
        return( (sum/nsamples) )
    else # exact and/or few sequences
        for i in 1:(nseq-1)
            @inbounds @simd for j in (i+1):nseq
                sum += percentidentity(msa[i,:], msa[j,:])
            end
        end
        return( (sum/nvalues) )
    end
end

# Similarity percent
# ==================

"""
Calculates the similarity percent between two aligned sequences.
The 100% is the length of the aligned sequences minus the number of columns with gaps in both sequences.
Two residues are considered similar if they below to the same group.
The `groups` (third positional argument) can be indicated with a vector of length 20,
having the group labels of each `Residue` in the following order (mandatory):
A, R, N, D, C, Q, E, G, H, I, L, K, M, F, P, S, T, W, Y, V

**MIToS default groups** are:

```julia
[
:nonpolar, # A
:positive, # R
:polar,    # N
:negative, # D
:cysteine, # C
:polar,    # Q
:negative, # E
:glycine,  # G
:positive, # H
:nonpolar, # I
:nonpolar, # L
:positive, # K
:nonpolar, # M
:aromatic, # F
:proline,  # P
:polar,    # S
:polar,    # T
:aromatic, # W
:aromatic, # Y
:nonpolar  # V
]

```

Other residue groups:

**SMS (Sequence Manipulation Suite)** Ident and Sim: GAVLI, FYW, ST, KRH, DENQ, P, CM

*Stothard P (2000)
The Sequence Manipulation Suite: JavaScript programs for analyzing and formatting protein and DNA sequences.
Biotechniques 28:1102-1104.*

```julia

[1, 4, 5, 5, 7, 5, 5, 1, 4, 1, 1, 4, 7, 2, 6, 3, 3, 2, 2, 1]

```

**Bio3D 2.2** seqidentity: GA, MVLI, FYW, ST, KRH, DE, NQ, P, C

*Grant, B.J. et al. (2006) Bioinformatics 22, 2695--2696.*

```julia

[1, 5, 7, 6, 9, 7, 6, 1, 5, 2, 2, 5, 2, 3, 8, 4, 4, 3, 3, 2]

```
"""
function percentsimilarity{T}(seq1::Vector{Residue}, seq2::Vector{Residue},
    groups::Vector{T} =
    [:nonpolar,    # A
        :positive, # R
        :polar,    # N
        :negative, # D
        :cysteine, # C
        :polar,    # Q
        :negative, # E
        :glycine,  # G
        :positive, # H
        :nonpolar, # I
        :nonpolar, # L
        :positive, # K
        :nonpolar, # M
        :aromatic, # F
        :proline,  # P
        :polar,    # S
        :polar,    # T
        :aromatic, # W
        :aromatic, # Y
        :nonpolar  # V
        ])

    if length(groups) != 20
        throw(ErrorException("Groups should have 20 labels for the residues: A, R, N, D, C, Q, E, G, H, I, L, K, M, F, P, S, T, W, Y, V"))
    end

    len = length(seq1)
    if len != length(seq2)
        throw(ErrorException("Sequences of different length, they aren't aligned or don't come from the same alignment"))
    end

    count = zero(Int)
    colgap = zero(Int)
    @inbounds for i in 1:len
        res1 = Int(seq1[i])
        res2 = Int(seq2[i])
        isgap1 = res1 == Int(GAP)
        isgap2 = res2 == Int(GAP)
        if isgap1 && isgap2
            colgap += 1
        end
        if !isgap1 && !isgap2
            count += Int(groups[res1] == groups[res2])
        end
    end

    (100.0 * count)/(len-colgap)
end

"""
Calculates the similarity percent between all the sequences on a MSA.
You can indicate the output element type with the `out` keyword argument (`Float64` by default).
For a MSA with a lot of sequences, you can use `out=Float32` or `out=Flot16` in order to avoid the `OutOfMemoryError()`.
"""
function percentsimilarity(msa::AbstractMatrix{Residue}, args...; out::Type=Float64)
    aln = getresiduesequences(msa)
    scores = sequencepairsmatrix(msa, out, false, out(100.0))
    @inbounds @iterateupper scores false list[k] = :($percentsimilarity)(:($aln)[i], :($aln)[j], :($args)...)
    scores
end


text:
# run via Julia REPL
using Clustering, Makie, DataFrames, RDatasets

const iris = dataset("datasets", "iris")
const colors = [:red, :green, :blue]
const plt = Vector{Any}(undef,2)

scene1 = Scene()
scene2 = Scene()

for (i, sp) in enumerate(unique(iris[:Species]))
    idx = iris[:Species] .== sp
    sel = iris[idx, [:SepalWidth, :SepalLength]]
    plt[1] = scatter!(scene1, sel[1], sel[2], color = colors[i],
        limits = FRect(1.5, 4.0, 3.0, 4.0))
end

features = permutedims(convert(Array, iris[1:4]), [2, 1])

# K Means ++
result = kmeans(features, 3, init = :kmpp)  # set to 3 clusters with kmeans++ :kmpp

for center in unique(result.assignments)
    idx = result.assignments .== center
    sel = iris[idx, [:SepalWidth, :SepalLength]]
    plt[2] = scatter!(scene2, sel[1], sel[2], color = colors[center],
        limits = FRect(1.5, 4.0, 3.0, 4.0))
end

scene2[Axis][:names][:axisnames] = scene1[Axis][:names][:axisnames] =
    ("Sepal Width", "Sepal Length")
t1 = text(Theme(), "Species Classification", camera=campixel!)
t2 = text(Theme(), "Kmeans Classification", camera=campixel!)
vbox(hbox(plt[1], t1), hbox(plt[2], t2))

ext: jl

max_stars_repo_name: ethansaxenian/RosettaDecode

max_issues_repo_name: ethansaxenian/RosettaDecode

max_forks_repo_name: ethansaxenian/RosettaDecode


text:
const N_A = 6.02214086e23 # mol-1
const k_B = 1.38064852e-23 # m2 kg s-2 K-1
const R̄   = N_A*k_B # m2 kg s-2 K-1 mol-1


text:
module BarotropicQG

export
  Problem,
  set_zeta!,
  updatevars!,

  energy,
  enstrophy,
  meanenergy,
  meanenstrophy,
  dissipation,
  work,
  drag,
  drag_ens,
  work_ens,
  dissipation_ens

using
  CUDA,
  Reexport

@reexport using FourierFlows

using FFTW: rfft
using LinearAlgebra: mul!, ldiv!
using FourierFlows: getfieldspecs, parsevalsum, parsevalsum2

nothingfunction(args...) = nothing

"""
    Problem(; parameters...)

Construct a BarotropicQG turbulence problem.
"""

function Problem(dev::Device=CPU();
  # Numerical parameters
          nx = 256,
          Lx = 2π,
          ny = nx,
          Ly = Lx,
          dt = 0.01,
  # Physical parameters
           β = 0.0,
         eta = nothing,
  # Drag and/or hyper-/hypo-viscosity
           ν = 0.0,
          nν = 1,
           μ = 0.0,
  # Timestepper and equation options
     stepper = "RK4",
      calcFU = nothingfunction,
      calcFq = nothingfunction,
  stochastic = false,
           T = Float64)

  # the grid
  grid = TwoDGrid(dev, nx, Lx, ny, Ly; T=T)
  x, y = gridpoints(grid)

  # topographic PV
  eta === nothing && ( eta = zeros(dev, T, (nx, ny)) )

  params = !(typeof(eta)<:ArrayType(dev)) ? Params(grid, β, eta, μ, ν, nν, calcFU, calcFq) : Params(β, eta, rfft(eta), μ, ν, nν, calcFU, calcFq)

  vars = (calcFq == nothingfunction && calcFU == nothingfunction) ? Vars(dev, grid) : (stochastic ? StochasticForcedVars(dev, grid) : ForcedVars(dev, grid))

  equation = Equation(params, grid)

  return FourierFlows.Problem(equation, stepper, dt, grid, vars, params, dev)
end

# ----------
# Parameters
# ----------

"""
    Params(g::TwoDGrid, β, FU, eta, μ, ν, nν, calcFU, calcFq)

Returns the params for an unforced two-dimensional barotropic QG problem.
"""
struct Params{T, Aphys, Atrans} <: AbstractParams
        β :: T            # Planetary vorticity y-gradient
      eta :: Aphys        # Topographic PV
     etah :: Atrans       # FFT of Topographic PV
        μ :: T            # Linear drag
        ν :: T            # Viscosity coefficient
       nν :: Int          # Hyperviscous order (nν=1 is plain old viscosity)
   calcFU :: Function     # Function that calculates the forcing F(t) on
                          # domain-averaged zonal flow U(t)
  calcFq! :: Function     # Function that calculates the forcing on QGPV q
end

"""
    Params(g::TwoDGrid, β, eta::Function, μ, ν, nν, calcFU, calcFq)

Constructor for Params that accepts a generating function for the topographic PV.
"""
function Params(grid::AbstractGrid{T, A}, β, eta::Function, μ, ν, nν::Int, calcFU, calcFq) where {T, A}
  etagrid = A([eta(grid.x[i], grid.y[j]) for i=1:grid.nx, j=1:grid.ny])
     etah = rfft(etagrid)
  return Params(β, etagrid, etah, μ, ν, nν, calcFU, calcFq)
end

# ---------
# Equations
# ---------

"""
    Equation(params, grid)

Returns the equation for two-dimensional barotropic QG problem with `params` and `grid`.
"""
function Equation(params::Params, grid::AbstractGrid)
  L = @. - params.μ - params.ν * grid.Krsq^params.nν + im * params.β * grid.kr * grid.invKrsq
  CUDA.@allowscalar L[1, 1] = 0
  return FourierFlows.Equation(L, calcN!, grid)
end

# ----
# Vars
# ----

abstract type BarotropicQGVars <: AbstractVars end

struct Vars{Ascalar, Aphys, Atrans, F, P} <: BarotropicQGVars
        U :: Ascalar
        q :: Aphys
     zeta :: Aphys
      psi :: Aphys
        u :: Aphys
        v :: Aphys
       qh :: Atrans
    zetah :: Atrans
     psih :: Atrans
       uh :: Atrans
       vh :: Atrans
      Fqh :: F
  prevsol :: P
end

const ForcedVars = Vars{<:AbstractArray, <:AbstractArray, <:AbstractArray, <:AbstractArray, Nothing}
const StochasticForcedVars = Vars{<:AbstractArray, <:AbstractArray, <:AbstractArray, <:AbstractArray, <:AbstractArray}

"""
    Vars(dev, grid)

Returns the vars for unforced two-dimensional barotropic QG problem on device `dev` and with `grid`
"""
function Vars(dev::Dev, grid::AbstractGrid) where Dev
  T = eltype(grid)
  U = ArrayType(dev, T, 0)(undef, ); CUDA.@allowscalar U[] = 0
  @devzeros Dev T (grid.nx, grid.ny) q u v psi zeta
  @devzeros Dev Complex{T} (grid.nkr, grid.nl) qh uh vh psih zetah
  Vars(U, q, zeta, psi, u, v, qh, zetah, psih, uh, vh, nothing, nothing)
end

"""
    ForcedVars(dev, grid)

Returns the vars for forced two-dimensional barotropic QG problem on device dev and with `grid`.
"""
function ForcedVars(dev::Dev, grid::AbstractGrid) where Dev
  T = eltype(grid)
  U = ArrayType(dev, T, 0)(undef, ); CUDA.@allowscalar U[] = 0
  @devzeros Dev T (grid.nx, grid.ny) q u v psi zeta
  @devzeros Dev Complex{T} (grid.nkr, grid.nl) qh uh vh psih zetah Fqh
  return Vars(U, q, zeta, psi, u, v, qh, zetah, psih, uh, vh, Fqh, nothing)
end

"""
    StochasticForcedVars(dev, grid)

Returns the vars for stochastically forced two-dimensional barotropic QG problem on device dev and with `grid`.
"""
function StochasticForcedVars(dev::Dev, grid::AbstractGrid) where Dev
  T = eltype(grid)
  U = ArrayType(dev, T, 0)(undef, ); CUDA.@allowscalar U[] = 0
  @devzeros Dev T (grid.nx, grid.ny) q u v psi zeta
  @devzeros Dev Complex{T} (grid.nkr, grid.nl) qh uh vh psih zetah Fqh prevsol
  return Vars(U, q, zeta, psi, u, v, qh, zetah, psih, uh, vh, Fqh, prevsol)
end

# -------
# Solvers
# -------

function calcN_advection!(N, sol, t, clock, vars, params, grid)
  # Note that U = sol[1, 1]. For all other elements ζ = sol
  CUDA.@allowscalar vars.U[] = sol[1, 1].re
  @. vars.zetah = sol
  CUDA.@allowscalar vars.zetah[1, 1] = 0

  @. vars.uh =  im * grid.l  * grid.invKrsq * vars.zetah
  @. vars.vh = -im * grid.kr * grid.invKrsq * vars.zetah

  ldiv!(vars.zeta, grid.rfftplan, vars.zetah)
  ldiv!(vars.u, grid.rfftplan, vars.uh)
  vars.psih .= vars.vh # FFTW's irfft destroys its input; vars.vh is needed for N[1, 1]
  ldiv!(vars.v, grid.rfftplan, vars.psih)

  @. vars.q = vars.zeta + params.eta
  uq = vars.u                                            # use vars.u as scratch variable
  CUDA.@allowscalar @. uq = (vars.U[] + vars.u) * vars.q # (U+u)*q
  vq = vars.v                                            # use vars.v as scratch variable
  @. vq *= vars.q                                        # v*q

  uqh = vars.uh                                          # use vars.uh as scratch variable
  mul!(uqh, grid.rfftplan, uq)                           # \hat{(u+U)*q}

  # Nonlinear advection term for q (part 1)
  @. N = -im * grid.kr * uqh                             # -∂[(U+u)q]/∂x

  vqh = vars.uh                                          # use vars.uh as scratch variable
  mul!(vqh, grid.rfftplan, vq)                           # \hat{v*q}

  # Nonlinear advection term for q (part 2)
  @. N += - im * grid.l * vqh                            # -∂[vq]/∂y
  return nothing
end

function calcN!(N, sol, t, clock, vars, params, grid)
  calcN_advection!(N, sol, t, clock, vars, params, grid)
  addforcing!(N, sol, t, clock, vars, params, grid)
  if params.calcFU != nothingfunction
    # 'Nonlinear' term for U with topographic correlation.
    # Note: ⟨v*η⟩ =     sum(conj(v̂)*η̂) / (nx²ny²) if  fft is used,
    # while ⟨v*η⟩ = 2 * sum(conj(v̂)*η̂) / (nx²ny²) if rfft is used
    CUDA.@allowscalar N[1, 1] = params.calcFU(t) + 2 * sum(conj(vars.vh) .* params.etah).re / (grid.nx^2 * grid.ny^2)
  end
  return nothing
end

addforcing!(N, sol, t, clock, vars::Vars, params, grid) = nothing

function addforcing!(N, sol, t, clock, vars::ForcedVars, params, grid)
  params.calcFq!(vars.Fqh, sol, t, clock, vars, params, grid)
  @. N += vars.Fqh
  return nothing
end

function addforcing!(N, sol, t, clock, vars::StochasticForcedVars, params, grid)
  if t == clock.t # not a substep
    @. vars.prevsol = sol # sol at previous time-step is needed to compute budgets for stochastic forcing
    params.calcFq!(vars.Fqh, sol, t, clock, vars, params, grid)
  end
  @. N += vars.Fqh
  return nothing
end

# ----------------
# Helper functions
# ----------------

"""
    updatevars!(sol, vars, params, grid)

Update the variables in `vars` with the solution in `sol`.
"""
function updatevars!(sol, vars, params, grid)
  CUDA.@allowscalar vars.U[] = sol[1, 1].re
  @. vars.zetah = sol
  CUDA.@allowscalar vars.zetah[1, 1] = 0.0

  @. vars.psih = - vars.zetah * grid.invKrsq
  @. vars.uh = - im * grid.l  * vars.psih
  @. vars.vh =   im * grid.kr * vars.psih

  ldiv!(vars.zeta, grid.rfftplan, deepcopy(vars.zetah))
  ldiv!(vars.psi, grid.rfftplan, deepcopy(vars.psih))
  ldiv!(vars.u, grid.rfftplan, deepcopy(vars.uh))
  ldiv!(vars.v, grid.rfftplan, deepcopy(vars.vh))

  @. vars.q = vars.zeta + params.eta
  return nothing
end

updatevars!(prob) = updatevars!(prob.sol, prob.vars, prob.params, prob.grid)

"""
    set_zeta!(prob, zeta)
    set_zeta!(sol, vars, params, grid)

Set the solution `sol` as the transform of zeta and update variables `vars`
on the `grid`.
"""
function set_zeta!(sol, vars::Vars, params, grid, zeta)
  mul!(vars.zetah, grid.rfftplan, zeta)
  CUDA.@allowscalar vars.zetah[1, 1] = 0.0
  @. sol = vars.zetah

  updatevars!(sol, vars, params, grid)
  return nothing
end

function set_zeta!(sol, vars::Union{ForcedVars, StochasticForcedVars}, params, grid, zeta)
  CUDA.@allowscalar vars.U[] = deepcopy(sol[1, 1])
  mul!(vars.zetah, grid.rfftplan, zeta)
  CUDA.@allowscalar vars.zetah[1, 1] = 0.0
  @. sol = vars.zetah
  CUDA.@allowscalar sol[1, 1] = vars.U[]

  updatevars!(sol, vars, params, grid)
  return nothing
end

set_zeta!(prob, zeta) = set_zeta!(prob.sol, prob.vars, prob.params, prob.grid, zeta)

"""
    set_U!(prob, U)
    set_U!(sol, v, g, U)

Set the (kx, ky)=(0, 0) part of solution sol as the domain-average zonal flow U.
"""
function set_U!(sol, vars, params, grid, U::Float64)
  CUDA.@allowscalar sol[1, 1] = U
  updatevars!(sol, vars, params, grid)
  return nothing
end

set_U!(prob, U::Float64) = set_U!(prob.sol, prob.vars, prob.params, prob.grid, U)

"""
    energy(prob)
    energy(sol, grid)

Returns the domain-averaged kinetic energy of solution `sol`.
"""
energy(sol, grid::AbstractGrid) = 0.5 * ( parsevalsum2(grid.kr .* grid.invKrsq .* sol, grid) + parsevalsum2(grid.l .* grid.invKrsq .* sol, grid) ) / (grid.Lx * grid.Ly)
energy(prob) = energy(prob.sol, prob.grid)

"""
    enstrophy(prob)
    enstrophy(sol, grid, vars)

Returns the domain-averaged enstrophy of solution `sol`.
"""
function enstrophy(sol, grid, vars)
  @. vars.uh = sol
  CUDA.@allowscalar vars.uh[1, 1] = 0
  return 0.5*parsevalsum2(vars.uh, grid) / (grid.Lx * grid.Ly)
end
enstrophy(prob) = enstrophy(prob.sol, prob.grid, prob.vars)

"""
    meanenergy(prob)

Returns the energy of the domain-averaged U.
"""
meanenergy(prob) = CUDA.@allowscalar real(0.5 * prob.sol[1, 1]^2)

"""
    meanenstrophy(prob)

Returns the enstrophy of the domain-averaged U.
"""
meanenstrophy(prob) = CUDA.@allowscalar real(prob.params.β * prob.sol[1, 1])

"""
    energy_dissipation(prob)
    energy_dissipation(sol, vars, params, grid)

Returns the domain-averaged energy dissipation rate. nν must be >= 1.
"""
@inline function energy_dissipation(sol, vars, params, grid)
  energy_dissipationh = vars.uh # use vars.uh as scratch variable

  @. energy_dissipationh = params.ν * grid.Krsq^(params.nν-1) * abs2(sol)
  CUDA.@allowscalar energy_dissipationh[1, 1] = 0
  return 1 / (grid.Lx * grid.Ly) * parsevalsum(energy_dissipationh, grid)
end

@inline energy_dissipation(prob) = energy_dissipation(prob.sol, prob.vars, prob.params, prob.grid)

"""
    enstrophy_dissipation(prob)
    enstrophy_dissipation(sol, vars, params, grid)

Returns the domain-averaged enstrophy dissipation rate. nν must be >= 1.
"""
@inline function enstrophy_dissipation(sol, vars, params, grid)
  enstrophy_dissipationh = vars.uh # use vars.uh as scratch variable

  @. enstrophy_dissipationh = params.ν * grid.Krsq^params.nν * abs2(sol)
  CUDA.@allowscalar enstrophy_dissipationh[1, 1] = 0
  return 1 / (grid.Lx * grid.Ly) * parsevalsum(enstrophy_dissipationh, grid)
end

@inline enstrophy_dissipation(prob) = enstrophy_dissipation(prob.sol, prob.vars, prob.params, prob.grid)

"""
    energy_work(prob)
    energy_work(sol, vars, grid)

Returns the domain-averaged rate of work of energy by the forcing `Fqh`.
"""
@inline function energy_work(sol, vars::ForcedVars, grid)
  energy_workh = vars.uh # use vars.uh as scratch variable

  @. energy_workh = grid.invKrsq * sol * conj(vars.Fqh)
  return 1 / (grid.Lx * grid.Ly) * parsevalsum(energy_workh, grid)
end

@inline function energy_work(sol, vars::StochasticForcedVars, grid)
  energy_workh = vars.uh # use vars.uh as scratch variable

  @. energy_workh = grid.invKrsq * (vars.prevsol + sol)/2 * conj(vars.Fqh) # Stratonovich
  # @. energy_workh = grid.invKrsq * vars.prevsol * conj(vars.Fqh)             # Ito
  return 1 / (grid.Lx * grid.Ly) * parsevalsum(vars.uh, grid)
end

@inline energy_work(prob) = energy_work(prob.sol, prob.vars, prob.grid)

"""
    enstrophy_work(prob)
    enstrophy_work(sol, vars, grid)

Returns the domain-averaged rate of work of enstrophy by the forcing `Fqh`.
"""
@inline function enstrophy_work(sol, vars::ForcedVars, grid)
  enstrophy_workh = vars.uh # use vars.uh as scratch variable

  @. enstrophy_workh = sol * conj(vars.Fqh)
  return 1 / (grid.Lx * grid.Ly) * parsevalsum(enstrophy_workh, grid)
end

@inline function enstrophy_work(sol, vars::StochasticForcedVars, grid)
  enstrophy_workh = vars.uh # use vars.uh as scratch variable

  @. enstrophy_workh = (vars.prevsol + sol) / 2 * conj(vars.Fqh) # Stratonovich
  # @. enstrophy_workh = grid.invKrsq * vars.prevsol * conj(vars.Fh)           # Ito
  return 1 / (grid.Lx * grid.Ly) * parsevalsum(enstrophy_workh, grid)
end

@inline enstrophy_work(prob) = enstrophy_work(prob.sol, prob.vars, prob.grid)

"""
    energy_drag(prob)

Returns the extraction of domain-averaged energy by drag μ.
"""
@inline function energy_drag(prob)
  sol, vars, params, grid = prob.sol, prob.vars, prob.params, prob.grid
  energy_dragh = vars.uh # use vars.uh as scratch variable

  @. energy_dragh = params.μ * grid.invKrsq * abs2(sol)
  CUDA.@allowscalar energy_dragh[1, 1] = 0
  return  1 / (grid.Lx * grid.Ly) * parsevalsum(energy_dragh, grid)
end

"""
    enstrophy_drag(prob)

Returns the extraction of domain-averaged enstrophy by drag/hypodrag μ.
"""
@inline function enstrophy_drag(prob)
  sol, vars, params, grid = prob.sol, prob.vars, prob.params, prob.grid
  enstrophy_dragh = vars.uh # use vars.uh as scratch variable

  @. enstrophy_dragh = params.μ * abs2(sol)
  CUDA.@allowscalar enstrophy_dragh[1, 1] = 0
  return 1 / (grid.Lx * grid.Ly) * parsevalsum(enstrophy_dragh, grid)
end

end # module


text:
module GeneralizedDynamicsFromData

using Reexport

using LinearAlgebra
using Optim
using Statistics
using OrderedCollections
using FileIO

using JLD2
using HDF5

@reexport using OrdinaryDiffEq
@reexport using DiffEqFlux
@reexport using Flux

export regression_model, reshape_all_parameters, reshape_parameters, restructure_parameters
export CallbackLog, callbacks_to_hdf5
export repeat_experiment, grid_experiment
export InitializationLoader
include("utils.jl")

export fritzhugh_nagumo, roessler, genetic_toggle_switch, truscott_brindley, NPZ, selkov
include("eqlib.jl")

export mse_loss, normed_ld_loss, cosine_distance_loss, combined_loss, polar_loss, cosine_distance, normed_ld, mse_loss_norm
include("losslib.jl")

# export plot_weights
# include("plotlib.jl")

end # module

max_stars_repo_name: pnieters/GeneralizedDynamicsFromData

max_issues_repo_name: pnieters/GeneralizedDynamicsFromData

max_forks_repo_name: pnieters/GeneralizedDynamicsFromData


text:
export AbstractApproximationMethod
export approximation_name, approximation_short_name

abstract type AbstractApproximationMethod end

function approximation_name end
function approximation_short_name end

function approximate_meancov(method::AbstractApproximationMethod, g::Function, distribution)
    return approximate_meancov(method, g, mean(distribution), cov(distribution))
end

function approximate_meancov(method::AbstractApproximationMethod, g::Function, m::T, v::T) where { T <: Real }
    weights = getweights(method, m, v)
    points  = getpoints(method, m, v)

    cs   = Vector{eltype(m)}(undef, length(weights))
    norm = 0.0
    mean = 0.0

    for (index, (weight, point)) in enumerate(zip(weights, points))
        gv = g(point)
        cv = weight * gv

        mean += point * cv
        norm += cv

        @inbounds cs[index] = cv
    end

    mean /= norm

    var = 0.0
    for (index, (point, c)) in enumerate(zip(points, cs))
        point -= mean
        var += c * point ^ 2
    end

    var /= norm

    return mean, var
end

function approximate_meancov(method::AbstractApproximationMethod, g::Function, m::AbstractVector{T}, P::AbstractMatrix{T}) where { T <: Real }
    ndims = length(m)

    weights = getweights(method, m, P)
    points  = getpoints(method, m, P)

    cs = similar(m, eltype(m), length(weights))
    norm = zero(T)
    mean = zeros(T, ndims)

    for (index, (weight, point)) in enumerate(zip(weights, points))
        gv = g(point)
        cv = weight * gv

        # mean = mean + point * weight * g(point)
        broadcast!(*, point, point, cv)  # point *= cv
        broadcast!(+, mean, mean, point) # mean += point
        norm += cv

        @inbounds cs[index] = cv
    end

    broadcast!(/, mean, mean, norm)

    cov = zeros(T, ndims, ndims)
    foreach(enumerate(zip(points, cs))) do (index, (point, c))
        broadcast!(-, point, point, mean)                # point -= mean
        mul!(cov, point, reshape(point, (1, ndims)), c, 1.0) # cov = cov + c * (point)⋅(point)' where c = weight * g(point)
    end

    broadcast!(/, cov, cov, norm)

    return mean, cov
end

function approximate_kernel_expectation(method::AbstractApproximationMethod, g::Function, distribution)
    return approximate_kernel_expectation(method, g, mean(distribution), cov(distribution))
end

function approximate_kernel_expectation(method::AbstractApproximationMethod, g::Function, m::AbstractVector{T}, P::AbstractMatrix{T}) where { T <: Real }
    ndims = length(m)

    weights = getweights(method, m, P)
    points  = getpoints(method, m, P)

    gbar = zeros(ndims, ndims)
    foreach(zip(weights, points)) do (weight, point)
        axpy!(weight, g(point), gbar) # gbar = gbar + weight * g(point)
    end

    return gbar
end


text:
function FieldPolynomialSplineFunction(arg0::Vector{RealFieldElement}, arg1::Vector{FieldPolynomialFunction})
    return FieldPolynomialSplineFunction((Vector{RealFieldElement}, Vector{FieldPolynomialFunction}), arg0, arg1)
end

function get_field(obj::FieldPolynomialSplineFunction)
    return jcall(obj, "getField", Field, ())
end

function get_knots(obj::FieldPolynomialSplineFunction)
    return jcall(obj, "getKnots", Vector{RealFieldElement}, ())
end

function get_n(obj::FieldPolynomialSplineFunction)
    return jcall(obj, "getN", jint, ())
end

function get_polynomials(obj::FieldPolynomialSplineFunction)
    return jcall(obj, "getPolynomials", Vector{FieldPolynomialFunction}, ())
end

function is_valid_point(obj::FieldPolynomialSplineFunction, arg0::RealFieldElement)
    return jcall(obj, "isValidPoint", jboolean, (RealFieldElement,), arg0)
end

function polynomial_spline_derivative(obj::FieldPolynomialSplineFunction)
    return jcall(obj, "polynomialSplineDerivative", FieldPolynomialSplineFunction, ())
end

function value(obj::FieldPolynomialSplineFunction, arg0::RealFieldElement)
    return jcall(obj, "value", RealFieldElement, (RealFieldElement,), arg0)
end

function value(obj::FieldPolynomialSplineFunction, arg0::jdouble)
    return jcall(obj, "value", RealFieldElement, (jdouble,), arg0)
end


text:
"""
    LGMModel(X, Y, groupindex, ncomponent)

Constructing an modelling object.
Optional keywords arguments:

 - `X` and `Y` are the covariates and response.
 - `groupindex` is the random effect groups, should be an interger vector
 - `ncomponent` is the number of components to try
 - `ngh` means the number of points used in Gaussian-Hermite quadrature approximation in calculating the marginal log likelihood.
 - `taufixed` specifies if fixing the ratio of `p[whichtosplit]` and `p[whichtosplit+1])` to be `tau/(1-tau)`
 - `sn` is the standard deviation to used in penalty function
 - `an` is the penalty factor

The following methods are available:

  - `nobs` returns the number of random effect levels
  - `model_response` returns the response `Y`
  - `coef` returns the fixed effects `β`
  - `ranef!` return the predict random effects
  - `stderror` gives the standard error of fixed effects
  - `confint` calculates the confidence interval
  - `coeftable` prints the fixed effects and their p values
  - `loglikelihood` calculates the log marginal likelihood
  - `vcov` returns the covariance matrix of fixed effects
  -  `asymptoticdistribution` returns the simulated asymptotic distribution of the restricted likelihood ratio test
  - `predict` computes the probability of `Y` being 1 at given new data
  - `FDR` detect the "outstanding" random effects while controlling the False  Discovery Rate.

"""
mutable struct LGMModel <: RegressionModel
    X::Matrix{Float64}
    Y::Vector{Bool}
    groupindex::IntegerVector
    ncomponent::Int
    p::Vector{Float64}
    μ::Vector{Float64}
    σ::Vector{Float64}
    β::Vector{Float64}
    n::Int
    ngh::Int
    ghx::Vector
    ghw::Vector
    μ_lb::Vector
    μ_ub::Vector
    σ_lb::Vector
    σ_ub::Vector
    Wim::Matrix{Float64}
    Wm::Matrix{Float64}
    gammaprediction::Vector{Float64}
    lln::Vector{Float64}
    llN::Vector{Float64}
    llN2::Vector{Float64}
    llN3::Vector{Float64}
    Xscratch::Matrix{Float64}
    xb::Vector{Float64}
    gammaM::Vector{Float64}
    XWX::Matrix{Float64}
    XWY::Vector{Float64}
    ll::Float64
    sn::Vector{Float64}
    an::Float64
    taufixed::Bool
    whichtosplit::Int
    tau::Real
    fit::Bool

    function LGMModel(X::Matrix{Float64}, Y::Vector{Bool}, groupindex::IntegerVector, ncomponent::Int;
        ngh::Int=100,
        taufixed::Bool=false, whichtosplit=1, tau=.5,
        sn=ones(ncomponent), an=1.0/maximum(groupindex))

        n = maximum(groupindex)
        N, J = size(X)
        Wim=zeros(n, ncomponent*ngh)
        Wm=zeros(1, ncomponent*ngh)
        gammaprediction = zeros(n)
        lln=zeros(n)
        llN=zeros(N)
        llN2=zeros(N)
        llN3=zeros(N)
        Xscratch=copy(X)
        xb::Vector{Float64}=zeros(N)
        gammaM=zeros(ncomponent*ngh)
        XWX = zeros(J, J)
        XWY = zeros(J)
        ghx, ghw = gausshermite(ngh)
        p = ones(ncomponent) ./ ncomponent
        μ = sort!(randn(ncomponent))
        σ = ones(ncomponent)
        β = randn(J)
        ll = -Inf
        μ_lb = μ .- 10
        μ_ub = μ .+ 10
        σ_lb = .25 .* σ
        σ_ub = 2.0 .* σ

        new(X, Y, groupindex, ncomponent, p, μ, σ, β, n, ngh, ghx, ghw, μ_lb, μ_ub, σ_lb, σ_ub, Wim, Wm, gammaprediction, lln, llN, llN2, llN3, Xscratch, xb, gammaM, XWX, XWY, ll, sn, an, taufixed, whichtosplit, tau, false)
    end
end
function show(io::IO, obj::LGMModel)
    if !obj.fit
        @warn("Model has not been fit")
        return nothing
    end

    println(io, "$(typeof(obj)):\n\nFixed Coefficients:\n", coeftable(obj))
    println(io, "Random Effects:")
    latexprint(io, ranefmixture(obj), surrounding=false)
    println(io, "\nNumber of random effect levels: ", obj.n)
    println(io, "Log likelihood is: ", obj.ll)
end
function initialize!(m::LGMModel)

    m0 = LGMModel(m.X, m.Y, m.groupindex, 1)
    fit!(m0)
    ranef!(m0);
    m.p, m.μ, m.σ, ml_tmp = gmm(m0.gammaprediction, m.ncomponent)
    m.β = m0.β
    or = sortperm(m.μ)
    m.p = m.p[or]
    m.μ = m.μ[or]
    m.σ = m.σ[or]

    mingamma = minimum(m0.gammaprediction) - 3 * std(m0.gammaprediction)
    maxgamma = maximum(m0.gammaprediction) + 3 * std(m0.gammaprediction)
    fill!(m.μ_lb, mingamma)
    fill!(m.μ_ub, maxgamma)
    copy!(m.σ_lb, .25 .* m.σ)
    copy!(m.σ_ub, 2.0 .* m.σ)
    m.sn = ones(m.ncomponent) .* std(m0.gammaprediction)
    m
end
const valid_opts = [:maxiteration, :tol, :debuginfo, :Qmaxiteration, :dotest, :epsilon, :updatebeta, :bn, :pl, :ptau]
const valid_opt_types = [Int, Real, Bool, Int, Bool, Real, Bool, Real, Bool, Bool]
const deprecated_opts = Dict()
function val_opts(opts)
    d = Dict{Symbol,Union{Bool,NTuple{2,Integer},Char,Integer}}()
    for (opt_name, opt_val) in opts
        !in(opt_name, valid_opts) && throw(ArgumentError("unknown option $opt_name"))
        opt_typ = valid_opt_types[findfirst(valid_opts, opt_name)]
        !isa(opt_val, opt_typ) && throw(ArgumentError("$opt_name should be of type $opt_typ, got $(typeof(opt_val))"))
        d[opt_name] = opt_val
        haskey(deprecated_opts, opt_name) && @warn("$opt_name is deprecated, use $(deprecated_opts[opt_name]) instead")
    end
    d
end

"""
    fit!(m)

Fits the model.

 - `maxiteration` is the maximum MCEM iterations allowed. `tol` means the stopping criteria. `initial_iteration` is the number of initial iterations using smaller `Mmax`. Initial iterations are used to save time in the first few iterations when the parameters are still far from the truth
 - `tol` the stoppting criteria
 - `debuginfo` is the switch to print more information to help debug
 - `Qmaxiteration` the number of iterations used in Iteratively Reweighted Least Squares, for updating β
 - `bn` the prespecified minimum requirement on component prior
 - `pl` and `ptau` should the penalty on σ and `tau` be included in loglikelihood
"""
function StatsBase.fit!(m::LGMModel;
    maxiteration::Int=2000, tol::Real=.001,
    debuginfo::Bool=false, Qmaxiteration::Int=30,
    dotest::Bool=false, epsilon::Real=1e-6,
    updatebeta::Bool=true, bn::Real=0.01,
    pl::Bool=false, ptau::Bool=false)

    m.fit && @warn("The model is already fit") && return(m)
    N, J = size(m.X)
    n = m.n
    ncomponent = m.ncomponent
    ngh = m.ngh
    M = ncomponent*m.ngh

    p_old = copy(m.p)
    μ_old = copy(m.μ)
    σ_old = copy(m.σ)
    β_old = copy(m.β)

    ll0=-Inf
    alreadystable = false
    if maxiteration <= 3
        alreadystable = true
    end
    for iter_em in 1:maxiteration
        if any(isnan.(m.β)) || any(isnan.(m.p))
            error("NaN in paramters!")
        end
        for ix in 1:ngh, jcom in 1:ncomponent
            ixM = ix+ngh*(jcom-1)
            m.gammaM[ixM] = m.ghx[ix]*m.σ[jcom]*sqrt(2)+m.μ[jcom]
        end

        copy!(p_old, m.p)
        copy!(μ_old, m.μ)
        copy!(σ_old, m.σ)

        LinearAlgebra.mul!(m.xb, m.X, m.β)
        m.ll=integralweight!(m.Wim, m.X, m.Y, m.groupindex, m.gammaM, m.p, m.ghw, m.llN, m.llN2, m.xb, N, J, n, ncomponent, ngh)
        lldiff = m.ll - ll0
        ll0 = m.ll
        if dotest
            if (lldiff < epsilon) && (iter_em > 3)
                break
            end
        end
        if debuginfo
            println("At $(iter_em)th iteration:")
        end
        if updatebeta && (mod1(iter_em, 3) == 1 || alreadystable)
            copy!(β_old, m.β)
            updateβ!(m.β, m.X, m.Y, m.groupindex, .001, .001,
            m.XWX, m.XWY, m.Xscratch, m.gammaM, m.Wim, m.lln,
            m.llN, m.llN2, m.llN3,
            m.xb, N, J, n, ncomponent, ngh, Qmaxiteration)
            debuginfo && println("beta=", m.β)
        end
        updateθ!(m.p, m.μ, m.σ, m.X, m.Y, m.groupindex,
        m.gammaM, m.Wim, m.Wm, m.sn, m.an, N, J, n, ncomponent, ngh)

        for kcom in 1:ncomponent
            if m.p[kcom] < bn
                m.p[kcom] = bn
            end
        end
        tmpsum = sum(m.p)
        for kcom in 1:ncomponent
            m.p[kcom] = m.p[kcom] / tmpsum
        end

        if m.taufixed
            Yeppp.max!(m.μ, m.μ, m.μ_lb)
            Yeppp.min!(m.μ, m.μ, m.μ_ub)
            p_tmp = m.p[m.whichtosplit]+m.p[m.whichtosplit+1]
            m.p[m.whichtosplit] = p_tmp*m.tau
            m.p[m.whichtosplit+1] = p_tmp*(1-m.tau)
        end
        if debuginfo
            println("p=$(m.p)")
            println("μ=$(m.μ)")
            println("σ=$(m.σ)")
            println("loglikelihood=$(m.ll)")
        end

        if !dotest
            if stopRule(vcat(m.β, m.p, m.μ, m.σ), vcat(β_old, p_old, μ_old, σ_old), tol=tol)
                if debuginfo
                    println("Converged at $(iter_em)th iteration.")
                end
                break
            end
        end
        if (iter_em == maxiteration) && (maxiteration > 50)
            @warn("Fail to converge with $(iter_em) iterations. The taufixed is $(m.taufixed).")
            println("Current parameters are
            $(m.p), $(m.μ), $(m.σ), $(m.β).")
            println(" Current likelihood is $(m.ll). The likelihood increase from last iteration is $(lldiff).")
        end
    end
    if pl
        m.ll += sum(pn(m.σ, m.sn, an=m.an))
    end
    if ptau
        tau2 = m.p[m.whichtosplit] / (m.p[m.whichtosplit]+m.p[m.whichtosplit+1])
        m.ll += log(1 - abs(1 - 2*tau2))
    end
    m.fit=true
    return(m)
end

function multiplefit!(m::LGMModel, ntrials::Int=25; kwargs...)
    pm = repeat(m.p, 1, ntrials)
    μm = repeat(m.μ, 1, ntrials)
    σm = repeat(m.σ, 1, ntrials)
    βm = repeat(m.β, 1, ntrials)
    ml = -Inf .* ones(ntrials)

    for im in 1:ntrials
        m.p = rand(Dirichlet(m.ncomponent, 1.0))
        m.μ = rand(m.ncomponent) .* (m.μ_ub .- m.μ_lb) .+ m.μ_lb
        m.σ = rand(m.ncomponent) .* (m.σ_ub .- m.σ_lb) .+ m.σ_lb
        m.fit=false

        fit!(m; kwargs...)
        pm[:, im] = m.p
        μm[:, im] = m.μ
        σm[:, im] = m.σ
        βm[:, im] = m.β
        ml[im] = m.ll
    end
    mlmax, imax = findmax(ml)
    or = sortperm(μm[:, imax])
    m.p = pm[or, imax]
    m.μ = μm[or, imax]
    m.σ = σm[or, imax]
    m.β = βm[:, imax]
    m.ll = mlmax
    return m
end

function EMtest(m::LGMModel, ntrials::Int=25, vtau::Vector{Float64}=[0.5;]; kwargs...)
    C0 = m.ncomponent
    C1 = C0 + 1
    n = m.n
    debuginfo = get(val_opts(kwargs), :debuginfo, false)

    initialize!(m)
    multiplefit!(m, ntrials; kwargs...)
    C0 > 1 && (trand=asymptoticdistribution(m))
    debuginfo && println("loglikelihood of C0:", m.ll)

    ranef!(m)
    mingamma = minimum(m.gammaprediction) - 3 * std(m.gammaprediction)
    maxgamma = maximum(m.gammaprediction) + 3 * std(m.gammaprediction)
    an = decidepenalty(m.p, m.μ, m.σ, n)
    m1 = LGMModel(m.X, m.Y, m.groupindex, C1, taufixed=true, an=an)
    lrv = zeros(length(vtau), C0)
    for whichtosplit in 1:C0, i in 1:length(vtau)
        ind = [1:whichtosplit; whichtosplit:C0;]
        if C1==2
            fill!(m1.μ_lb, mingamma)
            fill!(m1.μ_ub, maxgamma)
        elseif C1>2
            mu_lb = [mingamma; (m.μ[1:(C0-1)] .+ m.μ[2:C0])./2;]
            mu_ub = [(m.μ[1:(C0-1)] .+ m.μ[2:C0])./2; maxgamma;]
            copy!(m1.μ_lb, mu_lb[ind])
            copy!(m1.μ_ub, mu_ub[ind])
        end
        m1.σ_lb = 0.25 .* m.σ[ind]
        m1.σ_ub = 2 .* m.σ[ind]
        m1.whichtosplit = whichtosplit
        m1.tau = vtau[i]
        m1.sn = m.σ[ind]
        m1.β = m.β

        multiplefit!(m1, ntrials; kwargs...)
        m1.fit=false
        fit!(m1, maxiteration=2, tol=0.0)
        debuginfo && show(m1)
        debuginfo && println(whichtosplit, " ", vtau[i], "->", m1.ll)
        lrv[i, whichtosplit] = m1.ll
    end
    lr = maximum(lrv)
    if debuginfo
        println("lr=", lr)
    end
    Tvalue = 2*(lr - m.ll)
    if C0 == 1
        pvalue = 1 - cdf(Chisq(2), Tvalue)
    else
        pvalue = mean(trand .> Tvalue)
    end
    println("The Test Statistic is ", Tvalue)
    println("The p vlaue of the test is ", pvalue)
    return(Tvalue, pvalue)
end

function ranef!(m::LGMModel)
    ncomponent = m.ncomponent
    n = m.n
    ngh = m.ngh
    M = ngh*ncomponent
    N, J = size(m.X)

    for ix in 1:ngh, jcom in 1:ncomponent
        ixM = ix+ngh*(jcom-1)
        m.gammaM[ixM] = m.ghx[ix]*m.σ[jcom]*sqrt(2)+m.μ[jcom]
    end
    LinearAlgebra.mul!(m.xb, m.X, m.β)
    integralweight!(m.Wim, m.X, m.Y, m.groupindex, m.gammaM, m.p, m.ghw, m.llN, m.llN2, m.xb, N, J, n, ncomponent, ngh)
    fill!(m.gammaprediction, 0.0)
    for i in 1:n
        for j in 1:M
            m.gammaprediction[i] += m.gammaM[j] * m.Wim[i,j]
        end
    end
    return m.gammaprediction
end

nobs(m::LGMModel) = m.n
model_response(m::LGMModel) = m.Y
coef(m::LGMModel) = m.β
ranefmixture(m::LGMModel) = MixtureModel(map((u, v) -> Normal(u, v), m.μ, m.σ), m.p)
#deviance(m::LGMModel) = -2*loglikelihood(m)
function stderror(m::LGMModel)
    vc = vcov(m;includelambda=false)
    J = size(m.X, 2)
    sqrt.(diag(vc))[1:J]
end
function confint(m::LGMModel, level::Real)
    hcat(coef(m),coef(m)) + stderror(m)*quantile(Normal(),(1. -level)/2.)*[1. -1.]
end
confint(m::LGMModel) = confint(m, 0.95)
function coeftable(m::LGMModel)
    cc = coef(m)
    se = stderror(m)
    zz = cc ./ se
    CoefTable(hcat(cc,se,zz,2.0 * ccdf.(Normal(), abs.(zz))),
              ["Estimate","Std.Error","z value", "Pr(>|z|)"],
              ["x$i" for i = 1:size(m.X, 2)], 4)
end

function loglikelihood(m::LGMModel)
    N, J = size(m.X)
    LinearAlgebra.mul!(m.xb, m.X, m.β)
    for ix in 1:m.ngh, jcom in 1:m.ncomponent
        ixM = ix+m.ngh*(jcom-1)
        m.gammaM[ixM] = m.ghx[ix]*m.σ[jcom]*sqrt(2)+m.μ[jcom]
    end

    m.ll=integralweight!(m.Wim, m.X, m.Y, m.groupindex, m.gammaM, m.p, m.ghw, m.llN, m.llN2, m.xb, N, J, m.n, m.ncomponent, m.ngh)
    m.ll
end

function infomatrix(m::LGMModel; debuginfo::Bool=false, includelambda::Bool=true)

    N,J = size(m.X)
    n = m.n
    ngh=m.ngh
    C = m.ncomponent
    fill!(m.Wim, 0.0)
    summat_beta = zeros(n, ngh*C, J)
    llnC = zeros(n, C)
    S_β = zeros(n, J)
    S_π = zeros(n, C-1)
    S_μσ = zeros(n, 2*C)
    if includelambda
        S_λ = zeros(n, 2*C)
    end
    ml = zeros(n)
    LinearAlgebra.mul!(m.xb, m.X, m.β)
    for jcom in 1:C
        for ix in 1:ngh
            ixM = ix+ngh*(jcom-1)
            m.gammaM[ixM] = m.ghx[ix]*m.σ[jcom]*sqrt(2)+m.μ[jcom]
            for i in 1:N
                @inbounds m.llN[i] = ifelse(m.Y[i], -m.gammaM[ixM] - m.xb[i], m.gammaM[ixM] + m.xb[i])
            end
            copy!(m.llN2, m.llN)
            logistic!(m.llN2)
            negateiffalse!(m.llN2, m.Y)
            for i in 1:N
                @inbounds ind = m.groupindex[i]
                for j in 1:J
                    @inbounds summat_beta[ind, ixM, j] += m.llN2[i] * m.X[i,j]
                end
            end
            log1pexp!(m.llN, m.llN, m.llN3, N)

            for i in 1:N
                @inbounds m.Wim[m.groupindex[i], ixM] -= m.llN[i]
            end
            for i in 1:n
                @inbounds m.Wim[i, ixM] +=  log(m.ghw[ix])
            end
        end
    end
    for i in 1:n
        u = maximum(m.Wim[i, :])
        for jcol in 1:C*ngh
            @inbounds m.Wim[i, jcol] = m.Wim[i, jcol] - u
        end
    end
    for i in 1:n
        for kcom in 1:C
            @inbounds llnC[i, kcom] = sumexp(m.Wim[i,(1+ngh*(kcom-1)):ngh*kcom])
        end
    end
    for i in 1:n
        for jcom in 1:C
            for ix in 1:ngh
                @inbounds m.Wim[i, ix+ngh*(jcom-1)] += log(m.p[jcom])
            end
        end
        ml[i]=sumexp(m.Wim[i, :])
    end
    for kcom in 1:(C-1)
        S_π[:, kcom] = (llnC[:, kcom] .- llnC[:, C]) ./ ml
    end
    for i in 1:n
        for kcom in 1:C
            ind = (1+ngh*(kcom-1)):ngh*kcom
            S_μσ[i, 2*kcom-1] = sumexp(m.Wim[i, ind], H1(m.gammaM[ind], m.μ[kcom], m.σ[kcom])) / ml[i]
            S_μσ[i, 2*kcom] = sumexp(m.Wim[i, ind], H2(m.gammaM[ind], m.μ[kcom], m.σ[kcom]))/ml[i]
            if includelambda
                S_λ[i, 2*kcom-1] = sumexp(m.Wim[i, ind], H3(m.gammaM[ind], m.μ[kcom], m.σ[kcom]))/ml[i]
                S_λ[i, 2*kcom] = sumexp(m.Wim[i, ind], H4(m.gammaM[ind], m.μ[kcom], m.σ[kcom]))/ml[i]
            end
        end
        for j in 1:J
            S_β[i, j] = sumexp(m.Wim[i,:], summat_beta[i, :, j])/ml[i]
        end
    end
    S_η = hcat(S_β, S_π, S_μσ)
    debuginfo && println(round.(S_η[1:5,:], 5))
    debuginfo && println(round.(sum(S_η, 1)./sqrt(n), 6))
    I_η = S_η'*S_η./n
    if includelambda
        I_η = S_η'*S_η./n
        I_λη = S_λ'*S_η./n
        I_λ = S_λ'*S_λ./n
        I_all = vcat(hcat(I_η, I_λη'), hcat(I_λη, I_λ))
    else
        I_all = I_η
    end
    if 1/cond(I_all) < eps(Float64)
        @warn("Information Matrix is singular!")
        D, V = eigen(I_all)
        debuginfo && println(D)
        tol2 = maximum(abs.(D)) * 1e-14
        D[D.<tol2] = tol2
        I_all = V*Diagonal(D)*V'
    end
    return I_all
    #return inv(I_all)[1:J,1:J]./n
end

function vcov(m::LGMModel; debuginfo::Bool=false,includelambda::Bool=false)
    J = size(m.X, 2)
    I_all = infomatrix(m, debuginfo=debuginfo,includelambda=includelambda)
    return inv(I_all)[1:J, 1:J] ./ m.n
end

function asymptoticdistribution(m::LGMModel; debuginfo::Bool=false, nrep::Int=10000)
    N,J = size(m.X)
    n = m.n
    ngh=m.ngh
    C = m.ncomponent
    if C == 1
        return rand(Chisq(2), nrep)
    end

    I_all = infomatrix(m, includelambda=true, debuginfo=debuginfo)
    I_λ_η = I_all[(J+3*C):(J+5*C-1), (J+3*C):(J+5*C-1)] - I_all[(J+3*C):(J+5*C-1), 1:(J+3*C-1)] * inv(I_all[1:(J+3*C-1), 1:(J+3*C-1)]) * I_all[1:(J+3*C-1),(J+3*C):(J+5*C-1) ]
    debuginfo && println(round.(I_λ_η, 6))
    D, V = eigen(I_λ_η)
    D[D.<0.] = 0.
    I_λ_η2 = V * Diagonal(sqrt.(D)) * V'
    u = randn(nrep, 2*C) * I_λ_η2
    EM = zeros(nrep, C)
    T = zeros(nrep)
    for kcom in 1:C
        EM[:, kcom] = sum(u[:, (2*kcom-1):(2*kcom)] * inv(I_λ_η[(2*kcom-1):(2*kcom), (2*kcom-1):(2*kcom)]) .* u[:, (2*kcom-1):(2*kcom)], dims=2)
    end
    for i in 1:nrep
        T[i] = maximum(EM[i, :])
    end
    T
end

function predict(m::LGMModel, newX::Matrix{Float64}, newgroup::IntegerVector)
    !m.fit && @warn("The model is not fitted!")
    N, J = size(m.X)
    N2, J2 = size(newX)
    J != J2 && error("Dimension Dismatch in new data")
    newxb = zeros(N2)
    newgamma = zeros(N2)
    LinearAlgebra.mul!(newxb, newX, m.β)
    ranef!(m)
    overallmean = sum(m.μ .* m.p)
    for i in eachindex(newgroup)
        if newgroup[i] > m.n
            newgamma[i] = overallmean
        else
            newgamma[i] = m.gammaprediction[newgroup[i]]
        end
    end
    for i in 1:N2
        newxb[i] += newgamma[i]
    end
    logistic!(newxb)
end
"""
    latentgmm(Y~x1+x2+(1|groupindex), fr, ncomponent)

 - `fr` is a `DataFrame` containing the field `Y`, `x1`, `x2` and `groupindex`
  - `ncomponent` is the number of components in random effects.
"""
function latentgmm(f::Formula, fr::AbstractDataFrame, ncomponent::Int; kwargs...)
    mf = ModelFrame(f,fr)
    X = ModelMatrix(mf).m
    Y = convert(Vector{Bool}, getindex(fr, f.lhs))
    retrms = filter(x->Meta.isexpr(x,:call) && x.args[1] == :|, mf.terms.terms)
    if length(retrms) ≤ 0
        throw(ArgumentError("$f has no random-effects terms"))
    end
    # groupindex =  sort!([remat(e,mf.df) for e in retrms]; by = nlevs, rev = true) #or remat(retrms[1], mf.df).f
    groupindex = convert(IntegerVector, getindex(mf.df, retrms[1].args[3]))
    m = LGMModel(X, Y, groupindex, ncomponent; kwargs...)
    return m
end

"""
    FDR(m, CNull)

 - `CNull` is a vector of numbers, indicating which components should be viewd as "null"; then the elements with smallest density ratio f_{null}/f_{alternative} will be rejected
 - `alphalevel` is the False Discovery Rate
"""
function FDR(m::LGMModel, C0::IntegerVector=[findmax(m.p)[2];])
    !m.fit && @warn("The model is not yet fitted!")
    n = m.n
    X, Y, groupindex, ncomponent = m.X, m.Y, m.groupindex, m.ncomponent
    ghw, ghx = m.ghw, m.ghx
    gammaM = m.gammaM
    Wim = m.Wim
    ngh = m.ngh
    M = m.ngh*ncomponent
    N, J = size(m.X)
    if maximum(C0) > ncomponent || minimum(C0) < 1
        error("C0 must be within 1 to number of components")
    end

    piposterior = zeros(n, ncomponent)
    clFDR = zeros(n)

    for ix in 1:ngh, jcom in 1:ncomponent
        ixM = ix+ngh*(jcom-1)
        m.gammaM[ixM] = ghx[ix]*m.σ[jcom]*sqrt(2)+m.μ[jcom]
    end

    LinearAlgebra.mul!(m.xb, m.X, m.β)
    integralweight!(Wim, X, Y, groupindex, gammaM, m.p, ghw, m.llN, m.llN2, m.xb, N, J, n, ncomponent, ngh)
    for i in 1:n
        for jcom in 1:ncomponent
            for ix in 1:ngh
                ixM = ix+ngh*(jcom-1)
                piposterior[i, jcom] += Wim[i,ixM]
            end
        end
    end
    for i in 1:n
        for jcom in 1:ncomponent
            clFDR[i] = sum(piposterior[i, C0])#/sum(m.p[C0])
        end
    end
    return clFDR
end

function detect(m::LGMModel, C0::IntegerVector=[findmax(m.p)[2];]; alphalevel::Real=.05)

    clFDR = FDR(m, C0)
    order = sortperm(clFDR)
    #n0 = sum(clFDR .< alphalevel)
    Ni=counts(m.groupindex)
    n0=0
    for i in 1:m.n
        if mean(clFDR[order[1:i]]) < alphalevel
            n0 = i
        else
            break
        end
    end

    ranef!(m)
    ids = order[1:n0]
    ni = ones(Int, n0)
    yr = zeros(n0)
    for i in eachindex(ids)
        ni[i] = sum(m.groupindex .== ids[i])
        yr[i] = 1 - mean(m.Y[m.groupindex .== ids[i]])
    end
    CoefTable(Any[round(clFDR[ids], 4), round(m.gammaprediction[ids],4), ni, round(yr, 4)],
              ["FDR", "γ Prediction","Sample Size", "Survival Rate"],
              ["\\#$i" for i = ids])
end


text:
"""
Authors: Various
Compiled by: Eric Ford (see each function for credits)
Created: August 2020
Contact: https://github.com/eford/
"""

"""
   `searchsortednearest(a<:AbstractVector, x::Real; assume_sorted = false )`
   `searchsortednearest(a<:AbstractVector, x<:AbstractVector; assume_sorted = false )`

   Find the index of vector a where the value of a is closest to x.
   All vectors are assumed to already be sorted.
   To turn off assertions, set assume_sorted to true.

Credit: traktofon @ https://discourse.julialang.org/t/findnearest-function/4143/4
Vector Vector version by Christian Gilbertson?
issorted assertion and optional assume_sorted added by Eric Ford
"""
function searchsortednearest end

function searchsortednearest(a::AbstractVector{T} where T<:Real, x::Real ; assume_sorted::Bool = false)
   if !assume_sorted
	   @assert issorted(a)
   end
   idx = searchsortedfirst(a,x)
   if (idx==1); return idx; end
   if (idx>length(a)); return length(a); end
   if (a[idx]==x); return idx; end
   #if (abs(a[idx]-x) < abs(a[idx-1]-x))
   if (abs2(a[idx]-x) < abs2(a[idx-1]-x))
      return idx
   else
      return idx-1
   end
end

function searchsortednearest(x::T, a::AbstractVector{T}; assume_sorted::Bool = false) where T
    @warn "Did you mean to reverse the order of x and a?"
    return searchsortednearest(a, x, assume_sorted=assume_sorted)
end

function searchsortednearest(a::AbstractVector{T1}, x::AbstractVector{T2}; assume_sorted::Bool = false) where { T1, T2 }
   if !assume_sorted
	   @assert issorted(a)
   	   @assert issorted(x)
   end
   len_x = length(x)
   len_a = length(a)
   idxs = zeros(Int64, len_x)
   idxs[1] = searchsortednearest(a, x[1])
   for i in 2:len_x
	   idxs[i] = idxs[i-1] + searchsortednearest(view(a, idxs[i-1]:len_a), x[i]) - 1
   end
   return idxs
end

"""A generalized version of the built in append!() function
By Christian Gilbertson?
# TODO:  Ask Christian what the purpose of this is relative to std append
"""
function multiple_append!(a::Vector{T}, b...) where {T<:Real}
    for i in 1:length(b)
        append!(a, b[i])
    end
    return a
end

""" Return true if all elements of array are equal to each other. """
@inline function allequal(x::AbstractArray{T,1}) where {T<:Real}
    length(x) < 2 && return true
    e1 = x[1]
    i = 2
    @inbounds for i=2:length(x)
        x[i] == e1 || return false
    end
    return true
end

""" `findargminmax(a)`
Return (argmin, min, argmax, max)
Adapapted from https://github.com/JuliaLang/julia/blob/697e782ab86bfcdd7fd15550241fe162c51d9f98/base/array.jl#L2191
"""
findargminmax(a) = _findminmax(a, :)
function _findminmax(a, ::Colon)
    p = pairs(a)
    y = iterate(p)
    if y === nothing
        throw(ArgumentError("collection must be non-empty"))
    end
    (mi, m), s = y
    i = mi
	argmin = mi
	valmin = m
	argmax = mi
	valmax = m
    while true
        y = iterate(p, s)
        y === nothing && break
        valmin != valmin && break
		valmax != valmax && break
        (i, ai), s = y
        if ai != ai || isless(ai, valmin)
            valmin = ai
            argmin = i
        end
		if ai != ai || isless(valmax, ai)
			valmax = ai
			argmax = i
		end

    end
    return (argmin=argmin, min=valmin, argmax=argmax, max=valmax)
end

""" `interp_linear(;x1,x2,y1,y2,xpred)
Return result of simple linear interpolant at xpred.
Does not test that xpred is between x1 and x2.
"""
function interp_linear(;x1::T1,x2::T1,y1::T2,y2::T2,xpred::T1) where { T1<:Real, T2<:Real }
  ypred = y1+(y2-y1)*((xpred-x1)/(x2-x1))
end

max_stars_repo_name: RvSpectML/RvSpectMLTemplate

max_issues_repo_name: RvSpectML/RvSpectMLTemplate

max_forks_repo_name: RvSpectML/RvSpectMLTemplate


text:
let
    # Initial Conditions
    epc = Epoch(2019, 1, 1)
    oe  = [R_EARTH + 500e3, 0.01, 90.0, 45.0, 75.0, 180.0]
    eci = sOSCtoCART(oe, use_degrees=true)

    # Initialize orbit
    orb = EarthInertialState(epc, eci, dt=1.0,
            mass=1.0, n_grav=0, m_grav=0,
            drag=false, srp=false,
            moon=false, sun=false,
            relativity=false
    )

    # Set propagation period
    epcf = epc + orbit_period(oe[1])

    # Run propagation
    stepto!(orb, epcf)

    # Final state should exactly equal the inittial state
    ecif = orb.x
    oef  = sCARTtoOSC(orb.x, use_degrees=true)
    for i in 1:6
        @test isapprox(ecif[i], eci[i], atol=1e-5)
        @test isapprox(oef[i], oe[i], atol=1e-5)
    end
end

# Test integration of variational equations
let
    # Initial Conditions
    epc = Epoch(2019, 1, 1)
    oe  = [R_EARTH + 500e3, 0.01, 90.0, 45.0, 75.0, 180.0]
    eci = sOSCtoCART(oe, use_degrees=true)

    # Initialize orbit
    orb = EarthInertialState(epc, eci, dt=1.0)
    reinit!(orb) # Initialize STM to enable propagation

    # Set propagation period
    epcf = epc + orbit_period(oe[1])

    # Run propagation
    stepto!(orb, epcf)

    # Test to confirm non-zero STM
    for i in 1:length(eci)
        for j in 1:length(eci)
            @test orb.phi[i,j] != 0.0
        end
    end
end

# Test agreegation of output information
let
    # Initial Conditions
    epc = Epoch(2019, 1, 1)
    oe  = [R_EARTH + 500e3, 0.01, 90.0, 45.0, 75.0, 180.0]
    eci = sOSCtoCART(oe, use_degrees=true)

    # Initialize orbit
    orb = EarthInertialState(epc, eci, dt=1.0)
    reinit!(orb) # Initialize STM to enable propagation

    # Set propagation period
    epcf = 5

    # Simulate output to agregate state information
    sim_t, sim_epc, sim_x, sim_phi = sim!(orb, epcf)

    for i in 1:(epcf+1)
        @test sim_t[i] == i-1
    end
end


text:
using Random, Test, ModelParams, DistributionsLH

mdl = DistributionsLH;

function distrib_test(switches :: AbstractDistributionSwitches{T1}) where T1
    um = init_distribution(ObjectId(:distrib), switches);
    println(um)
    T2 = eltype(um);
    @test T1 == T2

    if isbounded(um)
        @test isa(max_value(um), T1)
        @test min_value(um) < max_value(um)
    end

    rng = MersenneTwister(123);

    sizeV = (4, 3);
    uDrawM = draw(um, sizeV, rng);
    @test check_draws(um, uDrawM)
    @test size(uDrawM) == sizeV
    @test eltype(uDrawM) == T1

    uDraw = draw(um, rng);
    @test isa(uDraw, T1)

    if isbounded(um)
        qV = quantiles(um, 0.0 : 0.2 : 1.0);
        @test isa(qV, Vector{T1})
        @test isapprox(qV[1], min_value(um))
        @test isapprox(qV[end], max_value(um))
        @test all(diff(qV) .> 0.0)
    end

    inM = randn(rng, sizeV...);
    drawM = scale_normal_draws(um, inM);
    @test size(inM) == size(drawM)
    @test check_draws(um, drawM)
end

@testset "Uniform marginal" begin
    for switches in [mdl.make_test_uniform_switches(),
        mdl.make_test_uniform_centered_switches()]

        distrib_test(switches);
    end
end

@testset "Beta" begin
    for switches in [mdl.make_test_beta_switches()]
        distrib_test(switches);
    end
end

@testset "Normal" begin
    for switches in (mdl.make_test_normal_switches(),
        mdl.make_test_log_normal_switches())
        distrib_test(switches);
    end
end

# ----------


text:
const loglevels = [Logging.Debug, Logging.Info, Logging.Warn, Logging.Error]

@enum Solver glpk scip cplex

@enum SolvingMode mtz scf mcf cec dcc

max_stars_repo_name: Entze/jkMST

max_issues_repo_name: Entze/jkMST

max_forks_repo_name: Entze/jkMST


text:
"""
    LeadSolver <: MetaSolver
Solver managed remotely by a MainSolver. Can manage its own set of local sub solvers.
"""
mutable struct LeadSolver <: MetaSolver
    meta_local_id::Tuple{Int, Int}
    model::_Model
    options::Options
    pool::Pool
    rc_report::RemoteChannel
    rc_sol::RemoteChannel
    rc_stop::RemoteChannel
    state::State
    strategies::MetaStrategy
    subs::Vector{_SubSolver}
end

function solver(mlid, model, options, pool, rc_report, rc_sol, rc_stop, strats, ::Val{:lead})
    l_options = deepcopy(options)
    set_option!(options, "print_level", :silent)
    ss = Vector{_SubSolver}()
    return LeadSolver(mlid, model, l_options, pool, rc_report, rc_sol, rc_stop, state(), strats, ss)
end

function _init!(s::LeadSolver)
    _init!(s, :meta)
    _init!(s, :local)
end

stop_while_loop(s::LeadSolver, ::Atomic{Bool}, ::Int, ::Float64) = isready(s.rc_stop)

function remote_stop!(s::LeadSolver)
    isready(s.rc_stop) && take!(s.rc_stop)
    sat = is_sat(s)
    if !sat || !has_solution(s)
        while isready(s.rc_report)
            wait(s.rc_sol)
            t = take!(s.rc_sol)
            update_pool!(s, t)
            sat && has_solution(t) && break
            take!(s.rc_report)
        end
    end
end


text:
using Documenter, HighFrequencyCovariance

makedocs(
    format = Documenter.HTML(),
    sitename = "HighFrequencyCovariance.jl",
    modules = [HighFrequencyCovariance],
    pages = Any[
        "Introduction" => "index.md",
        "First Steps with HighFrequencyCovariance.jl" => Any[
            "Algorithms" => "1_algorithms.md",
            "Data Structures" => "2_data_structures.md",
            "Writing Code" => "3_WritingCode.md",
            "References" => "9_references.md"
        ],
        "API" => Any[
            "Types" => "types.md",
            "Estimation Functions" => "functions.md",
            "Helper Functions" => "helper_functions.md",
            "Internal Functions" => "internals.md"
        ]
    ],
)

deploydocs(
    repo   = "github.com/s-baumann/HighFrequencyCovariance.jl.git",
    target = "build",
    deps   = nothing,
    make   = nothing,
    devbranch = "main"
)


text:
###########################################################
### Manifold Truncated Newton
###########################################################

abstract type TruncationStrategy end
@with_kw struct Newton <: TruncationStrategy
    ε_CGres::Float64 = 1e-13
end
@with_kw struct TruncatedNewton <: TruncationStrategy
    ν_reductionfactor::Float64 = 0.1
end

"""
    ManTruncatedNewton

Optimizer performing inexact Newton steps on a prescribed manifold. The inexactness in
the approximation of the Newton step is controlled by `νₖ`: one expects a direction `dᴺ`
such that
        ||∇²(f+g)(x)[dᴺ] + ∇(f+g)|| ≤ νₖ ||∇(f+g)||
"""
@with_kw struct ManTruncatedNewton{T} <: ManifoldUpdate
    linesearch::ManifoldLinesearch = Armijo()
    CG_maxiter::Int64 = 400
    truncationstrat::T = TruncatedNewton()
end
Base.summary(::ManTruncatedNewton{Newton}) = "ManNewton"
Base.summary(::ManTruncatedNewton{TruncatedNewton}) = "ManTruncatedNewton"

ManNewton(;kwargs...) = ManTruncatedNewton(truncationstrat=Newton(); kwargs...)

@with_kw mutable struct ManTruncatedNewtonState{Tlsstate} <: AbstractUpdateState
    norm_∇fgₘ::Float64 = -1.0
    νₖ::Float64 = 1.0
    CG_niter::Int64 = -1
    CG_ε::Float64 = -1
    CG_residual::Float64 = -1
    d_type::Symbol = :Unsolved
    norm_dᴺ::Float64 = -1
    cosθ::Float64 = -1
    λ_x::Float64 = -1
    ls_niter::Int64 = -1
    linesearch_state::Tlsstate
end

initial_state(manTN::ManTruncatedNewton, x, reg) = ManTruncatedNewtonState(linesearch_state=initial_state(manTN.linesearch))

# str_updatelog(o::ManifoldGradient, t::ManifoldGradientState) = @sprintf "ls-nit %2i\t\t||gradₘ f+g|| %.3e" t.ls_niter t.norm_∇fgₘ

function str_updatelog(::ManTruncatedNewton, t::ManTruncatedNewtonState)
    resstyle = (t.CG_residual <= t.CG_ε) ? "0" : "4;1"
    dtype_style = "0"
    (t.d_type == :MaxIter) && (dtype_style = "33")
    (t.d_type == :QuasiNegCurvature) && (dtype_style = "34")
    return string(
        (@sprintf "CG: nit:%3i  " t.CG_niter),
        "\033[$(dtype_style)m",
        string(t.d_type),
        "\033[0m",
        (@sprintf " νₖ:%.1e ε:%.1e res:" t.νₖ t.CG_ε),
        "\033[$(resstyle)m",
        (@sprintf "%.1e" t.CG_residual),
        "\033[0m",
        (@sprintf "     |dᴺ|:%.3e   cos(θₖ):%.1e  λ(x):%.1e   ls-nit%2i" t.norm_dᴺ t.cosθ t.λ_x t.ls_niter)
    )
end

#
### Truncation strategies
#
function update_CG_ε!(state_TN::ManTruncatedNewtonState, n::Newton, norm_rgrad)
    state_TN.CG_ε = n.ε_CGres
    return
end
function update_CG_ν!(state_TN, n::Newton, norm_rgrad, nit_ls)
    state_TN.νₖ = 1e-15
    return
end

function update_CG_ε!(state_TN::ManTruncatedNewtonState, tn::TruncatedNewton, norm_rgrad::Float64)
    state_TN.CG_ε = min(0.5, sqrt(norm_rgrad)) * norm_rgrad    # Forcing sequence as sugested in NW, p. 168
    return
end
function update_CG_ν!(state_TN, n::TruncatedNewton, norm_rgrad, nit_ls)
    # ν_strat = :ls
    # if ν_strat == :ls
    if nit_ls == 0
        state_TN.νₖ *= n.ν_reductionfactor
    end
    # state_TN.νₖ = 1e-6

    # if ν_strat == :lsnormgrad
    # if nit_ls <= 0 && state_TN.νₖ ≥ 1e-2 * norm_rgrad^2
    #     state_TN.νₖ *= 0.1
    # end
    # else
    #     @error "unknown reduction strategy"
    # end
    return
end

#
### Point update
#
function update_iterate!(state::PartlySmoothOptimizerState{Tx}, pb, o::ManTruncatedNewton) where Tx
    ###
    grad_fgₖ = state.tempvec_man
    M = state.M
    x = state.x.man_repr

    @assert state.x.repr == manifold_repr
    @assert is_manifold_point(M, x)

    state_TN = state.update_to_updatestate[o]

    ## 1. Get first & second order information
    # TODO: remove intermediate alloc from .+= op.
    grad_fgₖ = egrad_to_rgrad(M, x, state.∇f_x) + ∇M_g(pb, M, x); state.ncalls_gradₘF += 1
    function hessfg_x_h(ξ)
        state.ncalls_HessₘF += 1
        return ehess_to_rhess(M, x, state.∇f_x, ∇²f_h(pb, x, ξ), ξ) + ∇²M_g_ξ(pb, M, x, ξ)
    end

    norm_rgrad = norm(M, x, grad_fgₖ);
    state_TN.norm_∇fgₘ = norm_rgrad
    # check_tangent_vector(M, x, grad_fgₖ)

    ## 2. Get Truncated Newton direction
    update_CG_ε!(state_TN, o.truncationstrat, norm_rgrad)

    dᴺ, state_TN.CG_niter, state_TN.d_type = solve_tCG(M, x, grad_fgₖ, hessfg_x_h, ϵ_residual = state_TN.CG_ε, ν = state_TN.νₖ, printlev=0, maxiter=o.CG_maxiter)

    # check_tangent_vector(M, x, dᴺ)

    ## 3. Execute linesearch
    # TODO: make linesearch inplace for x, return status.
    hist_ls = Dict()
    x_ls = linesearch(o.linesearch, state, state_TN.linesearch_state, pb, M, x, grad_fgₖ, dᴺ, hist=hist_ls)

    x = x_ls
    state_TN.ls_niter = hist_ls[:niter]

    ## 4. Update CG precision
    update_CG_ν!(state_TN, o.truncationstrat, norm_rgrad, hist_ls[:niter])
    state.x.man_repr = x_ls

    ### Logging data
    state_TN.CG_residual = norm(M, x, hessfg_x_h(dᴺ) + grad_fgₖ)
    state_TN.norm_dᴺ = norm(M, x, dᴺ)
    state_TN.cosθ = inner(M, x, -dᴺ, grad_fgₖ)/(norm(M, x, dᴺ)*norm_rgrad)
    state_TN.λ_x = inner(M, x, -dᴺ, grad_fgₖ)
    state_TN.ls_niter = hist_ls[:niter]

    return
end


text:
srand(0)

facts("qraddcol") do

    context("β = 0") do
        m = 100
        A = randn(m,0)
        R = Array{Float64}(0,0)

        for i in 1:m
            a = randn(m)
            R = qraddcol(A, R, a)
            A = [A a]
            @fact  vecnorm( R'*R - A'*A ) --> less_than(1e-5)
        end
    end

    context("β > 0") do
        m = 100
        A = randn(m,0)
        R = Array{Float64}(0,0)
        β = 0.1
        for i in 1:m
            a = randn(m)
            B = [A a]
            R = qraddcol(A, R, a, β)
            @fact vecnorm( R'*R - B'*B - β^2*I ) --> less_than(1e-5)
        end
    end

end


text:
include("GyreInABox.jl")

using Test

if !isempty(ARGS)
  stommel = Bool(parse(Int,ARGS[1]))
  linear  = Bool(parse(Int,ARGS[2]))
  test    =      parse(Int,ARGS[3])
else
  stommel = true
  linear  = true
  test    = 1
end

###################
# PARAM SELECTION #
###################
FT = Float64

const τₒ = 2e-4 # value includes τₒ, g, and ρ
const fₒ = 1e-4
const β  = 1e-11
const λ  = 1e-6
const ν  = 1e4

const Lˣ = 1e6
const Lʸ = 1e6
const timeend = 100 * 24 * 60 * 60
const H = 1000
const c = sqrt(grav * H)

if stommel
  gyre = "stommel"
else
  gyre = "munk"
end

if linear
  advec = "linear"
else
  advec = "nonlinear"
end

outname = "vtk_new_dt_" * gyre * "_" * advec

function setup_model(FT, stommel, linear, τₒ, fₒ, β, γ, ν, Lˣ, Lʸ, H)
  problem = GyreInABox{FT}(τₒ, fₒ, β, Lˣ, Lʸ, H)

  if stommel
    turbulence = LinearDrag{FT}(λ)
  else
    turbulence = ConstantViscosity{FT}(ν)
  end

  if linear
    advection = nothing
  else
    advection = NonLinearAdvection()
  end

  model = SWModel(problem, turbulence, advection, c)
end

function shallow_init_state!(p::GyreInABox, T::TurbulenceClosure, state,
                           aux, coords, t)
  if t == 0
    null_init_state!(p, T, state, aux, coords, 0)
  else
    gyre_init_state!(p, T, state, aux, coords, t)
  end
end

function shallow_init_aux!(p::GyreInABox, aux, geom)
  @inbounds y  = geom.coord[2]

  Lʸ = p.Lʸ
  τₒ = p.τₒ
  fₒ = p.fₒ
  β  = p.β

  aux.τ = @SVector [-τₒ * cos(π * y / Lʸ), 0, 0]
  aux.f = @SVector [0, 0, fₒ + β * (y - Lʸ/2)]

  return nothing
end

#########################
# Timestepping function #
#########################

function run(mpicomm, topl, ArrayType, N, dt, FT, model, test)
  grid = DiscontinuousSpectralElementGrid(topl,
                                          FloatType = FT,
                                          DeviceArray = ArrayType,
                                          polynomialorder = N,
                                         )

  dg = DGModel(model,
               grid,
               Rusanov(),
               CentralNumericalFluxDiffusive(),
               CentralGradPenalty())

  Q  = init_ode_state(dg, FT(0))
  Qe = init_ode_state(dg, FT(timeend))

  lsrk = LSRK144NiegemannDiehlBusch(dg, Q; dt = dt, t0 = 0)

  cb = ()

  if test > 2
    outprefix = @sprintf("ic_mpirank%04d_ic", MPI.Comm_rank(mpicomm))
    statenames = flattenednames(vars_state(model, eltype(Q)))
    auxnames = flattenednames(vars_aux(model, eltype(Q)))
    writevtk(outprefix, Q, dg, statenames, dg.auxstate, auxnames)

    outprefix = @sprintf("exact_mpirank%04d", MPI.Comm_rank(mpicomm))
    statenames = flattenednames(vars_state(model, eltype(Qe)))
    auxnames = flattenednames(vars_aux(model, eltype(Qe)))
    writevtk(outprefix, Qe, dg, statenames, dg.auxstate, auxnames)

    step = [0]
    vtkpath = outname
    mkpath(vtkpath)
    cbvtk = GenericCallbacks.EveryXSimulationSteps(1000)  do (init=false)
      outprefix = @sprintf("%s/mpirank%04d_step%04d",vtkpath,
                           MPI.Comm_rank(mpicomm), step[1])
      @debug "doing VTK output" outprefix
      statenames = flattenednames(vars_state(model, eltype(Q)))
      auxnames = flattenednames(vars_aux(model, eltype(Q)))
      writevtk(outprefix, Q, dg, statenames, dg.auxstate, auxnames)
      step[1] += 1
      nothing
    end
    cb = (cb..., cbvtk)
  end

  solve!(Q, lsrk; timeend=timeend, callbacks=cb)

  error = euclidean_distance(Q, Qe) / norm(Qe)
  @info @sprintf("""Finished
                 error = %.16e
                 """, error)

  return error
end

################
# Start Driver #
################

let
  MPI.Initialized() || MPI.Init()
  mpicomm = MPI.COMM_WORLD

  ll = uppercase(get(ENV, "JULIA_LOG_LEVEL", "INFO"))
  loglevel = ll == "DEBUG" ? Logging.Debug :
  ll == "WARN"  ? Logging.Warn  :
  ll == "ERROR" ? Logging.Error : Logging.Info
  logger_stream = MPI.Comm_rank(mpicomm) == 0 ? stderr : devnull
  global_logger(ConsoleLogger(logger_stream, loglevel))
  @static if haspkg("CUDAnative")
    device!(MPI.Comm_rank(mpicomm) % length(devices()))
  end

  @static if haspkg("CuArrays")
    ArrayType = CuArray
  else
    ArrayType = Array
  end

  model = setup_model(FT, stommel, linear, τₒ, fₒ, β, λ, ν, Lˣ, Lʸ, H)

  if test == 1
    cellsrange = 10:10
    orderrange = 4:4
    testval = 1.6068814534535144e-03
  elseif test == 2
    cellsrange = 5:5:10
    orderrange = 3:4
  elseif test > 2
    cellsrange = 5:5:25
    orderrange = 6:10
  end

  errors = zeros(FT, length(cellsrange), length(orderrange))
  for (i, Ne) in enumerate(cellsrange)
    brickrange = (range(FT(0); length=Ne+1, stop=Lˣ),
                  range(FT(0); length=Ne+1, stop=Lʸ))
    topl = BrickTopology(mpicomm, brickrange,
                         periodicity = (false, false))

    for (j, N) in enumerate(orderrange)
      @info "running Ne $Ne and N $N with"
      dt = (Lˣ / c) / Ne / N^2
      @info @sprintf("\n dt = %f", dt)
      errors[i, j] = run(mpicomm, topl, ArrayType, N, dt, FT, model, test)
    end
  end

  @test errors[end,end] ≈ testval

  #=
  msg = ""
  for i in length(cellsrange)-1
    rate = log2(errors[i, end] - log2(errors[i+1, end]))
    msg *= @sprintf("\n rate for Ne %d = %e", cellsrange[i], rate)
  end
  @info msg

  msg = ""
  for j in length(orderrange)-1
    rate = log2(errors[end, j] - log2(errors[end, j+1]))
    msg *= @sprintf("\n rate for N  %d = %e", orderrange[j], rate)
  end
  @info msg
  =#

end

max_stars_repo_name: akshaysridhar/CLIMA

max_issues_repo_name: akshaysridhar/CLIMA

max_forks_repo_name: akshaysridhar/CLIMA


text:
module Caesar

import Tokenize
const TM = Tokenize
const TS = Tokenize.Tokens

using ReplMaker: initrepl

export tokenize, parse, interp!, interp_sexp_string!, enable_repl

struct SchemeToken
    kind::Any
    startpos::Any
    endpos::Any
    startbyte::Any
    endbyte::Any
    val::Any
end
kind(t::SchemeToken) = t.kind
untokenize(t::SchemeToken) = t.val

@enum(
    Kind,
    ENDMARKER, # EOF
    LPAREN,
    RPAREN,
    SYMBOL,
    INTEGER,
    FLOAT,
    STRING,
    CHAR,
    WHITESPACE,
)

function scheme_kind(tokenize_kind::TS.Kind)
    if tokenize_kind === TS.INTEGER
        return INTEGER
    elseif tokenize_kind === TS.FLOAT
        return FLOAT
    elseif tokenize_kind === TS.LPAREN
        return LPAREN
    elseif tokenize_kind === TS.RPAREN
        return RPAREN
    elseif tokenize_kind === TS.STRING
        return STRING
    elseif tokenize_kind === TS.CHAR
        return CHAR
    elseif tokenize_kind === TS.WHITESPACE
        return WHITESPACE
    elseif tokenize_kind === TS.ENDMARKER
        return ENDMARKER
    else
        return SYMBOL
    end
end

function tokenize(string)
    tokens = TM.tokenize(string)
    raw_tokens = []
    last_token = nothing
    for t in tokens
        if is_identifier_part(TS.kind(t))
            if is_identifier_part(TS.kind(last_token))
                push!(raw_tokens[end], t)
            else
                push!(raw_tokens, [t])
            end
        else
            push!(raw_tokens, t)
        end
        last_token = t
    end

    scheme_tokens = Array{SchemeToken}(undef, length(raw_tokens))
    for i in eachindex(raw_tokens)
        t = raw_tokens[i]
        if t isa TS.Token
            scheme_tokens[i] = SchemeToken(
                scheme_kind(TS.kind(t)),
                TS.startpos(t),
                TS.endpos(t),
                TS.startbyte(t),
                TS.endbyte(t),
                TS.untokenize(t),
            )
        else
            # We have a list of tokens that should become a Symbol.
            val = join(map(TS.untokenize, t))
            scheme_tokens[i] = SchemeToken(
                SYMBOL,
                TS.startpos(t[1]),
                TS.endpos(t[end]),
                TS.startbyte(t[1]),
                TS.endbyte(t[end]),
                val,
            )
        end
    end
    return scheme_tokens
end

function is_identifier_part(tokenize_kind)
    return scheme_kind(tokenize_kind) === SYMBOL
end

struct SchemeData{T}
    val::T
    token::Union{SchemeToken,Nothing}
end
value(s::SchemeData) = s.val

Base.show(io::IO, s::SchemeData) = show(io, s.val)

struct SExpr
    terms::Array{Union{SchemeData,SExpr}}
end
car(s::SExpr) = s.terms[1]
cdr(s::SExpr) = SExpr(@view s.terms[2:end])
cadr(s::SExpr) = car(cdr(s))
cddr(s::SExpr) = cdr(cdr(s))
caddr(s::SExpr) = car(cdr(cdr(s)))
cadddr(s::SExpr) = car(cdr(cdr(cdr(s))))
null(s::SExpr) = length(s.terms) == 0
Base.length(s::SExpr) = length(s.terms)
Base.push!(s::SExpr, a) = push!(s.terms, a)
Base.iterate(s::SExpr) = iterate(s.terms)
Base.iterate(s::SExpr, i::Int) = iterate(s.terms, i)

struct SchemeEnvironment
    bindings::Dict{Symbol,Any}
    parent::Union{SchemeEnvironment,Nothing}
end
SchemeEnvironment() = SchemeEnvironment(Dict{Symbol,Any}(), nothing)
extend!(env::SchemeEnvironment, name, value) = env.bindings[name] = value
function Base.haskey(env::SchemeEnvironment, name)
    if haskey(env.bindings, name)
        return true
    elseif env.parent !== nothing
        return haskey(env.parent, name)
    else
        throw(ArgumentError("Unbound variable `$(name)`."))
    end
end

function Base.get(env::SchemeEnvironment, name)
    if haskey(env.bindings, name)
        return env.bindings[name]
    elseif env.parent !== nothing
        return get(env.parent, name)
    else
        throw(ArgumentError("Unbound variable `$(name)`."))
    end
end

function set!(env::SchemeEnvironment, name, value)
    if haskey(env.bindings, name)
        return env.bindings[name] = value
    elseif env.parent !== nothing
        return set!(env.parent, name, value)
    else
        throw(ArgumentError("Unbound variable `$(name)`."))
    end
end

SchemeEnvironment(parent::SchemeEnvironment) = SchemeEnvironment(Dict{Symbol,Any}(), parent)

struct SchemeProcedure
    env::SchemeEnvironment
    formals::Union{Array{Symbol},Symbol}
    body::SExpr
    SchemeProcedure(parent::SchemeEnvironment, formals, body) =
        new(SchemeEnvironment(parent), formals, body)
end
function (proc::SchemeProcedure)(args...)
    if proc.formals isa Symbol
        # varargs
        extend!(proc.env, proc.formals, args)
    else
        for (variable, value) in zip(proc.formals, args)
            extend!(proc.env, variable, value)
        end
    end
    return interp!(proc.env, proc.body)
end

function Base.show(io::IO, s::SExpr)
    write(io, "(")
    n = length(s.terms)
    for (i, t) in enumerate(s.terms)
        show(io, t)
        i < n && write(io, " ")
    end
    write(io, ")")
end

function Base.parse(tokens::Array{SchemeToken})
    toplevel = SExpr([SchemeData(:begin, nothing)])
    prev = []
    current = toplevel
    for t in tokens
        if kind(t) === WHITESPACE
            continue
        elseif kind(t) === LPAREN
            push!(prev, current)
            new = SExpr(SchemeData[])
            push!(current, new)
            current = new
        elseif kind(t) === RPAREN
            current = pop!(prev)
        elseif kind(t) === ENDMARKER
            break
        else
            push!(current, parse(t))
        end
    end
    return toplevel
end

function Base.parse(token::SchemeToken)
    if kind(token) === SYMBOL
        return SchemeData(Symbol(untokenize(token)), token)
    elseif kind(token) === INTEGER
        return SchemeData(parse(Int, untokenize(token)), token)
    elseif kind(token) === FLOAT
        return SchemeData(parse(Float64, untokenize(token)), token)
    elseif kind(token) === STRING
        return SchemeData(strip(untokenize(token), '\"'), token)
    elseif kind(token) === CHAR
        return SchemeData(first(strip(untokenize(token), '\'')), token)
    else
        throw(ArgumentError("Unknown token kind found: `$(token)`"))
    end
end

function interp!(env::SchemeEnvironment, s::SExpr)
    # @show env s
    result = nothing
    if null(s)
        return nothing
    end

    head = car(s)
    if typeof(head) !== SchemeData{Symbol}
        throw(ArgumentError("Invalid s-expression. Must begin with symbol. Found `$(head)`"))
    end

    head = value(head)
    if head === :define
        result = interp_define!(env, s)
    elseif head === :set!
        result = interp_set!(env, s)
    elseif head === :begin
        result = interp_begin!(env, s)
    elseif head === :lambda
        result = interp_lambda!(env, s)
    elseif head === :if
        result = interp_if!(env, s)
    else
        # must be a procedure call
        result = interp_call!(env, s)
    end
    return result
end

function interp!(env::SchemeEnvironment, s::SchemeData)
    # @show env s
    if typeof(s) === SchemeData{Symbol}
        return get(env, value(s))
    else
        # literal value
        return value(s)
    end
end

function interp_define!(env::SchemeEnvironment, s::SExpr)
    if length(s) != 3
        throw(ArgumentError("`define` expression must have exactly 3 terms `(define <var> <value>)`."))
    end

    variable = cadr(s)
    if typeof(variable) != SchemeData{Symbol}
        throw(ArgumentError("`<var>` in `(define <var> <value>)` must be an identifier. Got `$(variable)`"))
    end
    variable = value(variable)
    result = interp!(env, caddr(s))
    extend!(env, variable, result)
    return nothing
end

function interp_set!(env::SchemeEnvironment, s::SExpr)
    if length(s) != 3
        throw(ArgumentError("`set!` expression must have exactly 3 terms `(set! <var> <value>)`."))
    end

    variable = cadr(s)
    if typeof(variable) != SchemeData{Symbol}
        throw(ArgumentError("`<var>` in `(set! <var> <value>)` must be an identifier. Got `$(variable)`"))
    end
    variable = value(variable)
    if haskey(env, variable)
        result = interp!(env, caddr(s))
        set!(env, variable, result)
    else
        throw(ArgumentError("`$variable` is unbound."))
    end
    return nothing
end

function interp_begin!(env::SchemeEnvironment, s::SExpr)
    if length(s) < 1
        throw(ArgumentError("`begin` expression must have form `(begin <expr>*)`."))
    end

    result = nothing
    for expr in cdr(s)
        result = interp!(env, expr)
    end
    return result
end

function interp_lambda!(env::SchemeEnvironment, s::SExpr)
    if length(s) < 3
        throw(ArgumentError("`lambda` expression must have form `(lambda <formals> <body-expr>*)`."))
    end

    function _push_formal!(f, v, err)
        if v isa SchemeData{Symbol}
            push!(f, value(v))
        else
            throw(ArgumentError(err))
        end
    end

    formals = Symbol[]
    formals_expr = cadr(s)
    if formals_expr isa SchemeData{Symbol}
        # (lambda <var> <body>*)
        _push_formal!(formals, formals_expr, "`<formals>` must be a symbol.")
    else
        # (lambda (<var>*) <body>*)
        for var in formals_expr
            _push_formal!(formals, var, "`<formals>` must be a list of symbols.")
        end
    end
    body = SExpr([SchemeData(:begin, nothing), cddr(s)...])
    return SchemeProcedure(env, formals, body)
end

function interp_if!(env::SchemeEnvironment, s::SExpr)
    if length(s) < 3 || length(s) > 4
        throw(ArgumentError("`if` expression must have form `(if <pred> <true-case> [<false-case>])`."))
    end

    pred = cadr(s)
    pred_value = interp!(env, pred)
    if pred_value
        true_case = caddr(s)
        return interp!(env, true_case)
    else
        if length(s) == 4
            false_case = cadddr(s)
            return interp!(env, false_case)
        end
    end

    return nothing
end

function interp_call!(env::SchemeEnvironment, s::SExpr)
    procedure = interp!(env, car(s))
    args = [interp!(env, expr) for expr in cdr(s)]
    result = procedure(args...)
    return result
end

interp_sexp_string!(env, sexp_str) = interp!(env, parse(tokenize(sexp_str)))

function init_env!(env)
    Caesar.extend!(env, :+, +)
    Caesar.extend!(env, :-, -)
    Caesar.extend!(env, :*, *)
    Caesar.extend!(env, :display, print)
    Caesar.extend!(env, :newline, () -> println())
    Caesar.extend!(env, Symbol("eq?"), ==)
end

### REPL functionality
function enable_repl()
    env = Caesar.SchemeEnvironment()
    init_env!(env)

    initrepl(
        (x) -> interp_sexp_string!(env, x);
        prompt_text = "caesar> ",
        prompt_color = :blue,
        start_key = ')',
        mode_name = "Caesar_mode",
    )
end

end # module


text:
#################################################################
#
#    1st order derivation testing (types)
#
#################################################################

module Sandbox
    type Abcd
        a::Float64
        b::Vector{Float64}
    end
    foo(t::Abcd) = t.a + t.b[2]
end

m.@deriv_rule Sandbox.Abcd(a,b) a ds[1]
m.@deriv_rule Sandbox.Abcd(a,b) b ds[2]
m.@deriv_rule Sandbox.foo(t)    t Any[ ds, (a=zeros(length(t.b)) ; a[2]=ds ; a) ]

t = Sandbox.Abcd(1., [2., 3.])

@compare t.a  * x        v0ref
@compare sum(t.b .* x)   v0ref
@compare sum(t.b .+ [x,x])   v0ref

t = Sandbox.Abcd(1., ones(v1ref))

ex = quote
    z = Sandbox.Abcd(1., x)
    Sandbox.foo(z)
end
compare(ex, v1ref)

ex = quote
    z = Sandbox.Abcd(0., [x*x])
    z.a
end
compare(ex, v0ref)

ex = quote
    z = Sandbox.Abcd(0., [x*x])
    sum(z.b)
end
compare(ex, v0ref)

ex = quote
    z = Sandbox.Abcd(0., [0.])
    z.a = x*x
    z.a
end
compare(ex, v0ref)

ex = quote
    z = Sandbox.Abcd(1., [x, x])
    z.a = 2x
    z.b[2]
end
compare(ex, v0ref)

ex = quote
    z = Sandbox.Abcd(1., [x, x])
    z.a = 2x
    y = z.b
    y[2]
end
compare(ex, v0ref)

ex = quote
    z = Sandbox.Abcd(1., [x, x])
    z.a = x
    Sandbox.foo(z)
end
compare(ex, v0ref)

ex = quote
    z = Sandbox.Abcd(1., x)
    z.a = x[3]
    Sandbox.foo(z)
end
compare(ex, v1ref)


text:
include("common.jl")
@testset "Constrained time step" begin
    # Check that numerical solutions approximate analytical solutions,
    # independent of problem structure

    alg = MethodOfSteps(BS3(); constrained=true)

    # Single constant delay
    @testset "single constant delay" begin
        ## Not in-place function with scalar history function
        prob = prob_dde_1delay_scalar_notinplace
        dde_int = init(prob, alg; dt=0.1)
        sol = solve!(dde_int)

        @test sol.errors[:l∞] < 3e-5
        @test sol.errors[:final] < 2.1e-5
        @test sol.errors[:l2] < 1.3e-5

        ## Not in-place function with vectorized history function
        prob = prob_dde_1delay_notinplace
        dde_int = init(prob, alg; dt=0.1)
        sol2 = solve!(dde_int)

        @test sol.t ≈ sol2.t && sol.u ≈ sol2[1, :]

        ## In-place function
        prob = prob_dde_1delay
        dde_int = init(prob, alg; dt=0.1)
        sol2 = solve!(dde_int)

        @test sol.t == sol2.t && sol.u == sol2[1, :]
    end

    # Two constant delays
    @testset "two constant delays" begin
        ## Not in-place function with scalar history function
        prob = prob_dde_2delays_scalar_notinplace
        dde_int = init(prob, alg; dt=0.1)
        sol = solve!(dde_int)

        @test sol.errors[:l∞] < 4.1e-6
        @test sol.errors[:final] < 1.5e-6
        @test sol.errors[:l2] < 2.3e-6

        ## Not in-place function with vectorized history function
        prob = prob_dde_2delays_notinplace
        dde_int = init(prob, alg; dt=0.1)
        sol2 = solve!(dde_int)

        @test sol.t ≈ sol2.t && sol.u ≈ sol2[1, :]

        ## In-place function
        prob = prob_dde_2delays
        dde_int = init(prob, alg; dt=0.1)
        sol2 = solve!(dde_int)

        @test sol.t == sol2.t && sol.u == sol2[1, :]
    end
end


text:
@testset "SGLD" begin
    #-----------------------
    # testing inplace addition
    #-----------------------
    function sample_reference!(s::SGLDState,grad::Function)
        # increment iteration
        s.t += 1

        # improve readability
        t = s.t
        x = s.x
        stepsize = s.stepsize(t)

        x[:] += stepsize*grad(x) + sqrt(2.0*stepsize)*randn(length(x))
        s
    end

    s = SGLDState(ones(2),0.01)
    ref_s = SGLDState(ones(2),0.01)
    sgldgrad(x) = x.*x

    srand(1)
    sample!(s,sgldgrad)

    srand(1)
    sample_reference!(ref_s,sgldgrad)

    @test s.x == ref_s.x

    #----------------------
    # Gaussian sanity check.
    #----------------------
    dm = GaussianMixtureModel([0.],[1.],[1.])

    gradient = DataModel.getgrad(dm)
    s = SGLDState(zeros(1),0.1)

    x_avg = [0.]
    x2_avg = [0.]
    for i in 1:50000
        sample!(s,gradient)
        x_avg =(i-1)/i*x_avg + s.x/i
        x2_avg = (i-1)/i*x2_avg + s.x.^2/i
    end

    @test abs(x_avg-0)[1] < 0.1
    @test abs(x2_avg-1)[1] < 0.2
end


text:
module MoltenSaltThermodynamics
using LsqFit
using Statistics: mean
using Plots
using LinearAlgebra
using LaTeXStrings
using ForwardDiff
using Debugger
using NLsolve
using FiniteDifferences
using Infiltrator

const R_GAS = 1.380649*6.02214076 # Boltzmann*Avogadro, values from CODATA18
const FARADAY = 96485.3365

include("extrapolation.jl")
include("modifiedquasichemical.jl")
include("numericalmethods.jl")
include("test_mqcm.jl")

export infinitedilution_extrapolate, plot_model, compare_actcoeffs, mytest1

end # module


text:
# fft
a = rand(8) + im*rand(8)
@test norm(ifft(fft(a)) - a) < 1e-8

m4 = [16.    2     3    13;
    5    11    10     8;
    9     7     6    12;
    4    14    15     1]

b = rand(17,14)
b[3:6,9:12] = m4
sm4 = slice(b,3:6,9:12)

pm4 = plan_fft(m4,1)

fft_m4 = fft(m4,1)
fftd2_m4 = fft(m4,2)
ifft_fft_m4 = ifft(fft(m4,1),1)
fftn_m4 = fft(m4)
ifftn_fftn_m4 = ifft(fft(m4))

fft!_m4 = complex(m4); fft!(fft!_m4,1)
fft!d2_m4 = complex(m4); fft!(fft!d2_m4,2)
ifft!_fft_m4 = fft(m4,1); ifft!(ifft!_fft_m4,1)
fft!n_m4 = complex(m4); fft!(fft!n_m4)
ifft!n_fftn_m4 = fft(m4); ifft!(ifft!n_fftn_m4)

pfft_m4 = plan_fft(m4,1)(m4)
pfftd2_m4 = plan_fft(m4,2)(m4)
pifft_fft_m4 = plan_ifft(fft_m4,1)(fft_m4)
pfftn_m4 = plan_fft(m4)(m4)
pifftn_fftn_m4 = plan_ifft(fftn_m4)(fftn_m4)

pfft!_m4 = complex(m4); plan_fft!(pfft!_m4,1)(pfft!_m4)
pfft!d2_m4 = complex(m4); plan_fft!(pfft!d2_m4,2)(pfft!d2_m4)
pifft!_fft_m4 = fft(m4,1); plan_ifft!(pifft!_fft_m4,1)(pifft!_fft_m4)
pfft!n_m4 = complex(m4); plan_fft!(pfft!n_m4)(pfft!n_m4)
pifft!n_fftn_m4 = fft(m4); plan_ifft!(pifft!n_fftn_m4)(pifft!n_fftn_m4)

sfftn_m4 = fft(sm4)
psfftn_m4 = plan_fft(sm4)(sm4)
sfft!n_b = complex128(b)
sfft!n_m4 = slice(sfft!n_b,3:6,9:12); fft!(sfft!n_m4)
psfft!n_b = complex128(b)
psfft!n_m4 = slice(psfft!n_b,3:6,9:12); plan_fft!(psfft!n_m4)(psfft!n_m4)

true_fft_m4 = [
    34.            34.            34.            34.;
     7. - 1.im  -5. + 3.im  -3. + 5.im   1. - 7.im;
    16.           -16.           -16.            16.;
     7. + 1.im  -5. - 3.im  -3. - 5.im   1. + 7.im ]

true_fftn_m4 = [
 136.        0          0         0 ;
   0.       20          8 + 8im   0 - 12im ;
   0.       32 + 32im   0        32 - 32im ;
   0.        0 + 12im   8 - 8im  20 ]

true_fftd2_m4 = [
   34.   13 + 11im    4   13 - 11im ;
   34.   -5 -  3im   -4   -5 +  3im ;
   34.    3 +  5im   -4    3 -  5im ;
   34.  -11 - 13im    4  -11 + 13im ]

for i = 1:length(m4)
    @test_approx_eq fft_m4[i] true_fft_m4[i]
    @test_approx_eq fftd2_m4[i] true_fftd2_m4[i]
    @test_approx_eq ifft_fft_m4[i] m4[i]
    @test_approx_eq fftn_m4[i] true_fftn_m4[i]
    @test_approx_eq ifftn_fftn_m4[i] m4[i]

    @test_approx_eq fft!_m4[i] true_fft_m4[i]
    @test_approx_eq fft!d2_m4[i] true_fftd2_m4[i]
    @test_approx_eq ifft!_fft_m4[i] m4[i]
    @test_approx_eq fft!n_m4[i] true_fftn_m4[i]
    @test_approx_eq ifft!n_fftn_m4[i] m4[i]

    @test_approx_eq pfft_m4[i] true_fft_m4[i]
    @test_approx_eq pfftd2_m4[i] true_fftd2_m4[i]
    @test_approx_eq pifft_fft_m4[i] m4[i]
    @test_approx_eq pfftn_m4[i] true_fftn_m4[i]
    @test_approx_eq pifftn_fftn_m4[i] m4[i]

    @test_approx_eq pfft!_m4[i] true_fft_m4[i]
    @test_approx_eq pfft!d2_m4[i] true_fftd2_m4[i]
    @test_approx_eq pifft!_fft_m4[i] m4[i]
    @test_approx_eq pfft!n_m4[i] true_fftn_m4[i]
    @test_approx_eq pifft!n_fftn_m4[i] m4[i]

    @test_approx_eq sfftn_m4[i] true_fftn_m4[i]
    @test_approx_eq sfft!n_m4[i] true_fftn_m4[i]
    @test_approx_eq psfftn_m4[i] true_fftn_m4[i]
    @test_approx_eq psfft!n_m4[i] true_fftn_m4[i]
end

ifft!(sfft!n_m4)
plan_ifft!(psfft!n_m4)(psfft!n_m4)
@test norm(sfft!n_m4 - m4) < 1e-8
@test norm(psfft!n_m4 - m4) < 1e-8

m3d = float32(reshape(1:5*3*2, 5, 3, 2))
ifft3_fft3_m3d = ifft(fft(m3d))

fftd3_m3d = fft(m3d,3)
ifftd3_fftd3_m3d = ifft(fftd3_m3d,3)

fft!d3_m3d = complex(m3d); fft!(fft!d3_m3d,3)
ifft!d3_fftd3_m3d = copy(fft!d3_m3d); ifft!(ifft!d3_fftd3_m3d,3)

pfftd3_m3d = plan_fft(m3d,3)(m3d)
pifftd3_fftd3_m3d = plan_ifft(fftd3_m3d,3)(fftd3_m3d)

pfft!d3_m3d = complex(m3d); plan_fft!(pfft!d3_m3d,3)(pfft!d3_m3d)
pifft!d3_fftd3_m3d = copy(fft!d3_m3d); plan_ifft!(pifft!d3_fftd3_m3d,3)(pifft!d3_fftd3_m3d)

@test isa(fftd3_m3d, Array{Complex64,3})
@test isa(ifftd3_fftd3_m3d, Array{Complex64,3})
@test isa(fft!d3_m3d, Array{Complex64,3})
@test isa(ifft!d3_fftd3_m3d, Array{Complex64,3})
@test isa(pfftd3_m3d, Array{Complex64,3})
@test isa(pifftd3_fftd3_m3d, Array{Complex64,3})
@test isa(pfft!d3_m3d, Array{Complex64,3})
@test isa(pifft!d3_fftd3_m3d, Array{Complex64,3})

true_fftd3_m3d = Array(Float32, 5, 3, 2)
true_fftd3_m3d[:,:,1] = 17:2:45
true_fftd3_m3d[:,:,2] = -15

for i = 1:length(m3d)
    @test_approx_eq fftd3_m3d[i] true_fftd3_m3d[i]
    @test_approx_eq ifftd3_fftd3_m3d[i] m3d[i]
    @test_approx_eq ifft3_fft3_m3d[i] m3d[i]

    @test_approx_eq fft!d3_m3d[i] true_fftd3_m3d[i]
    @test_approx_eq ifft!d3_fftd3_m3d[i] m3d[i]

    @test_approx_eq pfftd3_m3d[i] true_fftd3_m3d[i]
    @test_approx_eq pifftd3_fftd3_m3d[i] m3d[i]
    @test_approx_eq pfft!d3_m3d[i] true_fftd3_m3d[i]
    @test_approx_eq pifft!d3_fftd3_m3d[i] m3d[i]
end

# rfft/rfftn

rfft_m4 = rfft(m4,1)
rfftd2_m4 = rfft(m4,2)
rfftn_m4 = rfft(m4)

prfft_m4 = plan_rfft(m4,1)(m4)
prfftd2_m4 = plan_rfft(m4,2)(m4)
prfftn_m4 = plan_rfft(m4)(m4)

srfftn_m4 = rfft(sm4)
psrfftn_m4 = plan_rfft(sm4)(sm4)

for i = 1:3, j = 1:4
    @test_approx_eq rfft_m4[i,j] true_fft_m4[i,j]
    @test_approx_eq rfftd2_m4[j,i] true_fftd2_m4[j,i]
    @test_approx_eq rfftn_m4[i,j] true_fftn_m4[i,j]

    @test_approx_eq prfft_m4[i,j] true_fft_m4[i,j]
    @test_approx_eq prfftd2_m4[j,i] true_fftd2_m4[j,i]
    @test_approx_eq prfftn_m4[i,j] true_fftn_m4[i,j]

    @test_approx_eq srfftn_m4[i,j] true_fftn_m4[i,j]
    @test_approx_eq psrfftn_m4[i,j] true_fftn_m4[i,j]
end

irfft_rfft_m4 = irfft(rfft_m4,size(m4,1),1)
irfft_rfftd2_m4 = irfft(rfftd2_m4,size(m4,2),2)
irfftn_rfftn_m4 = irfft(rfftn_m4,size(m4,1))

pirfft_rfft_m4 = plan_irfft(rfft_m4,size(m4,1),1)(rfft_m4)
pirfft_rfftd2_m4 = plan_irfft(rfftd2_m4,size(m4,2),2)(rfftd2_m4)
pirfftn_rfftn_m4 = plan_irfft(rfftn_m4,size(m4,1))(rfftn_m4)

for i = 1:length(m4)
    @test_approx_eq irfft_rfft_m4[i] m4[i]
    @test_approx_eq irfft_rfftd2_m4[i] m4[i]
    @test_approx_eq irfftn_rfftn_m4[i] m4[i]

    @test_approx_eq pirfft_rfft_m4[i] m4[i]
    @test_approx_eq pirfft_rfftd2_m4[i] m4[i]
    @test_approx_eq pirfftn_rfftn_m4[i] m4[i]
end

rfftn_m3d = rfft(m3d)
rfftd3_m3d = rfft(m3d,3)
@test size(rfftd3_m3d) == size(fftd3_m3d)
irfft_rfftd3_m3d = irfft(rfftd3_m3d,size(m3d,3),3)
irfftn_rfftn_m3d = irfft(rfftn_m3d,size(m3d,1))
for i = 1:length(m3d)
    @test_approx_eq rfftd3_m3d[i] true_fftd3_m3d[i]
    @test_approx_eq irfft_rfftd3_m3d[i] m3d[i]
    @test_approx_eq irfftn_rfftn_m3d[i] m3d[i]
end

fftn_m3d = fft(m3d)
@test size(fftn_m3d) == (5,3,2)
rfftn_m3d = rfft(m3d)
@test size(rfftn_m3d) == (3,3,2)
for i = 1:3, j = 1:3, k = 1:2
    @test_approx_eq rfftn_m3d[i,j,k] fftn_m3d[i,j,k]
end

max_stars_repo_name: jakebolewski/julia

max_issues_repo_name: jakebolewski/julia

max_forks_repo_name: jakebolewski/julia


text:
# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     formats: ipynb,jl:hydrogen
#     text_representation:
#       extension: .jl
#       format_name: hydrogen
#       format_version: '1.3'
#       jupytext_version: 1.11.2
#   kernelspec:
#     display_name: Julia 1.8.0-DEV
#     language: julia
#     name: julia-1.8
# ---

# %%
using MetaUtils

# %% [markdown]
# * https://github.com/mitmath/18S096/blob/master/lectures/other/Syntax%20trees%20in%20Julia.ipynb

# %%
# https://github.com/mitmath/18S096/blob/master/lectures/other/Syntax%20trees%20in%20Julia.ipynb

using LightGraphs
using TikzGraphs
using MacroTools

"Return the number of the new vertex"
add_numbered_vertex!(g) = (add_vertex!(g); top = nv(g))

"Convert the current node into a label"
function label(sym)
    sym == :(^) && return "\\textasciicircum"  # TikzGraphs chokes on ^
    return string("\\texttt{", sym, "}")
end

"""
Traverse the expression and build up the graph g and the set of labels for each node. (These both get modified in place.) `show_call` specifies whether to include `call` nodes in the graph in function calls. Including them represents the Julia AST more precisely, but adds a lot of visual noise to the display.

traverse! returns the number of the vertex at the top of the subexpression.
"""
function traverse!(g, labels, ex, show_call=true)
    top_vertex = add_numbered_vertex!(g)
    start_argument = 1  # which argument to start with

    if !(show_call) && ex.head == :call
        f = ex.args[1]   # the function name
        push!(labels, label(f))
        start_argument = 2   # drop "call" from tree
    else
        push!(labels, label(ex.head))
    end

    for i in start_argument:length(ex.args)
        if isa(ex.args[i], Expr)
            child = traverse!(g, labels, ex.args[i], show_call)
            add_edge!(g, top_vertex, child)
        else
            n = add_numbered_vertex!(g)
            add_edge!(g, top_vertex, n)
            push!(labels, label(ex.args[i]))
        end
    end

    return top_vertex
end

function make_graph(ex::Expr, show_call=false)
    # https://github.com/FluxML/MacroTools.jl/blob/master/src/utils.jl#L87
    ex = MacroTools.striplines(ex)
    g = Graph()
    labels = String[]
    traverse!(g, labels, ex, show_call)
    return g, labels
end

function draw_syntax_tree(ex::Expr, show_call=false)
    TikzGraphs.plot(make_graph(ex::Expr, show_call)...)
end

macro tree(ex::Expr)
    draw_syntax_tree(ex)
end

macro tree_with_call(ex::Expr)
    draw_syntax_tree(ex, true)
end

# %%
t = @tree z = 3x + (2y+1)^2

# %%
t.data

# %%
@tree_with_call z = 3x + (2y+1)^2

# %%
@show_tree z = 3x + (2y+1)^2

# %%
@show_texpr z = 3x + (2y+1)^2

# %%
@tree z = f(g(f(x, y)))

# %%
@tree begin
    a = b + c
    d = 2a^2 + a
end

# %%
@tree (f(x^2 + 2x), g(y - (x^2 + 2x)))

# %%
@tree @time sin(10)

# %%
@show_tree @time sin(10)

# %%
@show_texpr @time sin(10)

# %%
g, labels = make_graph(:(x^2 + y^2))

for i in 1:length(labels)
    labels[i] = string("\$v_", i, "\$: \$", labels[i], "\$")
end

TikzGraphs.plot(g, labels)

# %%
@tree x*y - sin(z)

# %%
@tree (c = 0; for k in 1:n c += rand()^2 + rand()^2 <= 1 end; 4c/n)

# %%
@show_tree (c = 0; for k in 1:n c += rand()^2 + rand()^2 ≤ 1 end; 4c/n)

# %%
@show_texpr (c = 0; for k in 1:n c += rand()^2 + rand()^2 ≤ 1 end; 4c/n)

# %%

max_stars_repo_path: 0008/@tree and @tree_with_call macros.jl

max_stars_repo_name: genkuroki/public

max_issues_repo_name: genkuroki/public

max_forks_repo_name: genkuroki/public


text:
println("Loading ADCP data")
d1 = load_data(deps[1])
d2 = load_data(deps[2])


text:
immutable Event
    activate_at::ActivationPoint
    action  # A DataType that implements Base.call(T, env, sched)
end

isless(a::Event, b::Event) = a.activate_at < b.activate_at


text:
function equivalentLatticeNetwork(
	network::RegularNetwork;
	seed::Integer = -1,
	numReshuffle::Integer = 10,
	maxTentatives::Integer = (numconnections(network) * numReshuffle * 300)
)
	return network
end