text:
using Revise # reduce recompilation time
using Plots
# using Documenter
using LinearAlgebra
using SparseArrays
using BenchmarkTools
using UnPack
using ToeplitzMatrices

push!(LOAD_PATH, "./src")
using CommonUtils
using Basis1D
using Basis2DQuad
using UniformQuadMesh

using SetupDG

# TODO: refactor the code
function fS(uL, uR)
    return 1/6*(uL^2+uR^2+uL*uR)
end

function burgers_exact_sol_2D(u0,x,y,T,dt)
    Nsteps = ceil(Int,T/dt)
    dt = T/Nsteps
    u = u0(x,y) # computed at input points
    for i = 1:Nsteps
        t = i*dt
        u .= @. u0(x-u*t,y) # evolve solution at quadrature points using characteristics
    end
    return u
end

S_N(x) = @. sin(pi*x/h)/(2*pi/h)/tan(x/2)
"""
Vandermonde matrix of sinc basis functions determined by h,
evaluated at r
"""
function vandermonde_Sinc(h,r)
    N = convert(Int, 2*pi/h)
    V = zeros(length(r),N)
    for n = 1:N
        V[:,n] = S_N(r.-n*h)
    end
    V[1,1] = 1
    V[end,end] = 1 # TODO: fix other vandermonde?
    return V
end

"Constants"
const sp_tol = 1e-12
u0(x,y) = @. -sin(pi*x)*sin(pi*y)
# u0(x,y) = @. exp(-1*(y^2+x^2))

"Program parameters"
compute_L2_err = false

"Approximation Parameters"
N_P = 2;    # The order of approximation in polynomial dimension
Np_P = N_P+1;
Np_F = 16;    # The order of approximation in Fourier dimension
K    = 30;   # Number of elements in polynomial dimension
CFL  = 0.1;
T    = 0.5;  # End time

"Time integration Parameters"
rk4a,rk4b,rk4c = rk45_coeffs()
CN = (N_P+1)*(N_P+2)/2  # estimated trace constant for CFL
dt = CFL * 2 / CN / K
Nsteps = convert(Int,ceil(T/dt))
dt = T/Nsteps

"Initialize Reference Element in Fourier dimension"
h = 2*pi/Np_F
column = [0; .5*(-1).^(1:Np_F-1).*cot.((1:Np_F-1)*h/2)]
column2 = [-pi^2/3/h^2-1/6; -((-1).^(1:Np_F-1)./(2*(sin.((1:Np_F-1)*h/2)).^2))]
Ds = Toeplitz(column,column[[1;Np_F:-1:2]])
D2s = Toeplitz(column2,column2[[1;Np_F:-1:2]])
s = LinRange(h,2*pi,Np_F)

"Mesh related variables"
VX = repeat(LinRange(-1,1,K+1),2)
VY = [(-1+2/Np_F)*ones(K+1,1);ones(K+1,1)]
fv = quad_face_vertices()
Nfaces = length(fv)
EToV = hcat(1:K, K+2:2*K+1, 2:K+1, K+3:2*K+2)

"Initialize Reference Element in polynomial dimension"
rd = init_reference_interval(N_P);
@unpack r,rq,wq,rf,Dr,VDM,Vq,Vf,Vp,M,Pq,LIFT,nrJ = rd
wf = [1;1]
Nq_P = length(rq)
Nfp_P = length(rf)
Nh_P = Nq_P+Nfp_P # Number of hybridized points
rp = LinRange(-1,1,100)
Vp = vandermonde_1D(N_P,rp)/VDM # Redefine Vp
Lq = LIFT

"Initialize Mesh"
FToF = connect_mesh(EToV,fv)
r1 = [-1; -1; 1; 1]
s1 = [h; 2*pi; h; 2*pi]
@unpack r = rd
s,r = meshgrid(LinRange(h,2*pi,Np_F),r)
V1 = vandermonde_2D(1,r,s)/vandermonde_2D(1,r1,s1)
x = reshape(V1*VX[transpose(EToV)],Np_P,Np_F,K)
y = reshape(V1*VY[transpose(EToV)],Np_P,Np_F,K)
mapM = reshape(collect(1:2*K*Np_F),2,K*Np_F)
mapP = [[mapM[2,end-Np_F+1:end]; mapM[2,1:end-Np_F]]'; [mapM[1,Np_F+1:end]; mapM[1,1:Np_F]]']

# Geometric Factors
xr = mapslices((x->Dr*x),x,dims=(1,2))
rxJ = ones(size(xr))
nxJ = repeat(nrJ,1,Np_F,K)
ys = 1/pi*ones(size(xr)) # TODO: hard coded
JP = xr # TODO: redundant, for clarity
JF = ys
J = @. JP*JF
#TODO: only assume uniform affine mesh, need to test for nonuniform meshes
#TODO: a lot of redundancy
JPq,JFq,Jq = (x->mapslices((x->Vq*x),x,dims=(1,2))).((JP,JF,J))
JPf,JFf,Jf = (x->mapslices((x->Vf*x),x,dims=(1,2))).((JP,JF,J))

# scale by Fourier dimension
M = h*M
wq = h*wq
wf = h*wf

# Hybridized operators
Ef = Vf*Pq
Br = diagm(wf.*nrJ)
Qr = Pq'*M*Dr*Pq
Qs = Ds
Drh = [Vq*Dr*Pq-1/2*Vq*Lq*diagm(nrJ)*Vf*Pq  1/2*Vq*Lq*diagm(nrJ);
       -1/2*diagm(nrJ)*Vf*Pq                1/2*diagm(nrJ)]
Qrh = 1/2*[Qr-Qr' Ef'*Br;
           -Br*Ef Br]
Qrh_skew = 1/2*(Qrh-Qrh')
Qsh = Qs # Not the SBP operator, weighted when flux differencing

# TODO: assume mesh uniform affine, so Jacobian are constants
# TODO: fix other Jacobian parts
JP = 1/K
JF = 1/pi
J = JF*JP
wq = J*wq
wf = JF*wf
Lq = 1/JP*Lq
Qrh = JF*Qrh
Qsh = JP*Qsh

# TODO: refactor
ops = (Vq,Vf,wq,wf,Pq,Lq,Qrh,Qsh)
mesh = (nrJ,JP,JF,J,h,mapM,mapP)
param = (K,Np_P,Np_F,Nh_P)
function rhs(u,ops,mesh,param,compute_rhstest)
    Vq,Vf,wq,wf,Pq,Lq,Qrh,Qsh = ops
    nrJ,JP,JF,J,h,mapM,mapP = mesh
    K,Np_P,Np_F = param
    uq,uf = (A->A*u).((Vq,Vf))
    uh = [uq;uf]
    Nh = size(uh,1)
    w = [wq;wf]
    ∇fh = zeros(size(uh))

    # TODO: fix Jacobian
    # Flux differencing in x direction
    for k = 1:K
        for nf = 1:Np_F
            j = nf+Np_F*(k-1)
            ∇fh[:,j] += 2*Qrh.*[fS(uL,uR) for uL in uh[:,j], uR in uh[:,j]]*ones(size(uh,1),1)
        end
    end

    # Flux differencing in y direction
    for k = 1:K
        for nh = 1:Nh_P
            ∇fh[nh,(k-1)*Np_F+1:k*Np_F] += 2*[wq/J;0;0][nh]*Qsh.*[fS(uL,uR) for uL in uh[nh,(k-1)*Np_F+1:k*Np_F], uR in uh[nh,(k-1)*Np_F+1:k*Np_F]]*ones(Np_F,1)
        end
    end

    # Spatial term
    ∇f = [Pq Lq]*diagm(1 ./ w)*∇fh
    # Flux term
    uf = Vf*u
    uM = uf[mapM]
    uP = uM[mapP]
    LF = @. max(abs(uP),abs(uM))*(uP-uM)
    uflux = diagm(nrJ)*(@. fS(uP,uM)-uM*uM/2)-LF
    rhsu = -(∇f+Lq*uflux)

    # Artificial Viscosity
    # ϵ = 0.05
    ϵ = 0.1
    Δf = zeros(size(u))
    for k = 1:K
        Δf[:,(k-1)*Np_F+1:k*Np_F] = ϵ*u[:,(k-1)*Np_F+1:k*Np_F]*D2s'
    end
    rhsu = -(∇f+Lq*uflux-Δf)

    rhstest = 0
    if compute_rhstest
        rhstest += sum(diagm(wq)*uq.*(Vq*rhsu))
    end

    return rhsu,rhstest
end

# Time stepping
u = u0(x,y)
u = reshape(u,size(u,1),size(u,2)*size(u,3)) # reshape for faster matrix multiplication
resu = zeros(size(u))
for i = 1:Nsteps
    rhstest = 0
    for INTRK = 1:5
        compute_rhstest = INTRK==5
        rhsu,rhstest = rhs(u,ops,mesh,param,compute_rhstest)
        @. resu = rk4a[INTRK]*resu + dt*rhsu
        @. u += rk4b[INTRK]*resu
    end

    if i%10==0 || i==Nsteps
        println("Time step: $i out of $Nsteps with rhstest = $rhstest")
    end
end

# Plotting
gr(aspect_ratio=1, legend=false,
   markerstrokewidth=0, markersize=3)
plot()

u = reshape(u,Np_P,Np_F,K)
yp = LinRange(-1+2/Np_F,1,200)
sp = LinRange(h,2*pi,200)
VDM_F = vandermonde_Sinc(h,sp)
V2 = vandermonde_1D(1,sp)/vandermonde_1D(1,LinRange(h,2*pi,Np_F))
Vp = vandermonde_1D(N_P,LinRange(-1,1,10))/VDM
x_plot = mapslices((x->Vp*x*V2'),x,dims=(1,2))
y_plot = mapslices((x->Vp*x*V2'),y,dims=(1,2))
u_plot = mapslices((x->Vp*x*VDM_F'),u,dims=(1,2))
plt = scatter(vec(x_plot),vec(y_plot),vec(u_plot),zcolor=vec(u_plot),camera=(0,90))
display(plt)

max_stars_repo_name: yiminllin/ESDG-Fourier

max_issues_repo_name: yiminllin/ESDG-Fourier

max_forks_repo_name: yiminllin/ESDG-Fourier


text:
include("economic_dispatch.jl")
include("unit_commitment.jl")
include("opf.jl")

ext: jl


text:
include("area.jl")
include("count.jl")
include("mass.jl")
include("ratio.jl")
include("carbon.jl")
#include("nitrogen.jl")
include("water.jl")
include("weight.jl")
include("density.jl")

include("radiation.jl")
include("gasexchange/gasexchange.jl")
include("photosynthesis.jl")

include("plant.jl")


text:
##### Numeric simplification

"""
    simplify(x, [rules=SIMPLIFY_RULES]; fixpoint=true, applyall=true, recurse=true)

Apply a `RuleSet` of rules provided in `rules`. By default
these rules are `SymbolicUtils.SIMPLIFY_RULES`. If `fixpoint=true`
repeatedly applies the set of rules until there are no changes.
Applies them once if `fixpoint=false`.

The `applyall` and `recurse` keywords are forwarded to the enclosed
`RuleSet`.
"""
function simplify(x, rules=SIMPLIFY_RULES; fixpoint=true, applyall=true, recurse=true)
    if fixpoint
        SymbolicUtils.fixpoint(rules; recurse=recurse, applyall=recurse)(x)
    else
        rules(x; recurse=recurse, applyall=recurse)
    end
end

### Predicates

sym_isa(::Type{T}) where {T} = @nospecialize(x) -> x isa T || symtype(x) <: T
is_operation(f) = @nospecialize(x) -> (x isa Term) && (operation(x) == f)

isnumber(x) = x isa Number

_iszero(t) = false
_iszero(x::Number) = iszero(x)
_isone(t) = false
_isone(x::Number) = isone(x)

# A total ordering
<ₑ(a::Real,    b::Real) = abs(a) < abs(b)
<ₑ(a::Complex, b::Complex) = (abs(real(a)), abs(imag(a))) < (abs(real(b)), abs(imag(b)))
<ₑ(a::Real,    b::Complex) = true
<ₑ(a::Complex, b::Real) = false

<ₑ(a::Symbolic, b::Number) = false
<ₑ(a::Number,   b::Symbolic) = true

arglength(a) = length(arguments(a))
function <ₑ(a::Sym, b::Term)
    args = arguments(b)
    if length(args) === 2
        n1, n2 = !isnumber(args[1]) , !isnumber(args[2])
        if n1 && n2
            # both subterms are terms, so it's definitely firster
            return true
        elseif n1
            return a <ₑ args[1]
        elseif n2
            return a <ₑ args[2]
        else
            # both arguments are not numbers
            # This case when a <ₑ Term(^, [1,-1])
            # so this term should go to the left.
            return false
        end
    elseif length(args) === 1
        # make sure a < sin(a) < b^2 < b
        if isequal(a, args[1])
            return true # e.g sin(a)*a should become a*sin(a)
        else
            return a<ₑargs[1]
        end
    else
        # variables to the right
        return false
    end
end

<ₑ(a::Symbolic, b::Sym) = !(b <ₑ a)

function <ₑ(a::Symbol, b::Symbol)
    # Enforce the order [+,-,\,/,^,*]
    if b === :*
        a in (:^, :/, :\, :-, :+)
    elseif b === :^
        a in (:/, :\, :-, :+) && return true
    elseif b === :/
        a in (:\, :-, :+) && return true
    elseif b === :\
        a in (:-, :+) && return true
    elseif b === :-
        a === :+ && return true
    elseif a in (:*, :^, :/, :-, :+)
        false
    else
        a < b
    end
end

<ₑ(a::Sym, b::Sym) = a.name < b.name
<ₑ(a::T, b::S) where {T, S} = T===S ? isless(a, b) : nameof(T) < nameof(S)

function <ₑ(a::Term, b::Term)
    if arglength(a) === 0
        return operation(a) <ₑ b
    elseif arglength(b) === 0
        return a <ₑ operation(b)
    end

    if 0 < arglength(a) <= 2 && 0 < arglength(b) <= 2
        # e.g. a < sin(a) < b ^ 2 < b
        @goto compare_args
    end
    na = nameof(operation(a))
    nb = nameof(operation(b))
    if na !== nb
        return na <ₑ nb
    elseif arglength(a) != arglength(b)
        return arglength(a) < arglength(b)
    else
        @label compare_args
        aa, ab = arguments(a), arguments(b)
        if length(aa) !== length(ab)
            return length(aa) < length(ab)
        else
            terms = zip(Iterators.filter(!isnumber, aa), Iterators.filter(!isnumber, ab))

            for (x,y) in terms
                if x <ₑ y
                    return true
                elseif y <ₑ x
                    return false
                end
            end

            # compare the numbers
            nums = zip(Iterators.filter(isnumber, aa),
                       Iterators.filter(isnumber, ab))
            return any(a <ₑ b for (a, b) in nums)
        end
        return na <ₑ nb
    end
end

issortedₑ(args) = issorted(args, lt=<ₑ)

# are there nested ⋆ terms?
function isnotflat(⋆)
    function (args)
        for t in args
            if t isa Term && operation(t) === (⋆)
                return true
            end
        end
        return false
    end
end

function hasrepeats(x)
    length(x) <= 1 && return false
    for i=1:length(x)-1
        if isequal(x[i], x[i+1])
            return true
        end
    end
    return false
end

function merge_repeats(merge, xs)
    length(xs) <= 1 && return false
    merged = Any[]
    i=1

    while i<=length(xs)
        l = 1
        for j=i+1:length(xs)
            if isequal(xs[i], xs[j])
                l += 1
            else
                break
            end
        end
        if l > 1
            push!(merged, merge(xs[i], l))
        else
            push!(merged, xs[i])
        end
        i+=l
    end
    return merged
end

### Simplification rules

pow(x,y) = y==0 ? 1 : y<0 ? inv(x)^(-y) : x^y
pow(x::Symbolic,y) = y==0 ? 1 : Base.:^(x,y)

# Numbers to the back
function flatten_term(⋆, args)
    # flatten nested ⋆
    flattened_args = []
    for t in args
        if t isa Term && operation(t) === (⋆)
            append!(flattened_args, arguments(t))
        else
            push!(flattened_args, t)
        end
    end
    Term(⋆, flattened_args)
end

function sort_args(f, args)
    if length(args) < 2
        return Term(f, args)
    elseif length(args) == 2
        x, y = args
        return Term(f, x <ₑ y ? [x,y] : [y,x])
    end
    args = args isa Tuple ? [args...] : args
    Term(f, sort(args, lt=<ₑ))
end


text:
module AlphaStableDistributions

using LinearAlgebra, Statistics, Random
using StatsBase, Distributions, StaticArrays
using MAT, SpecialFunctions, ToeplitzMatrices
using Interpolations

export AlphaStable, SymmetricAlphaStable, AlphaSubGaussian, fit

Base.@kwdef struct AlphaStable{T} <: Distributions.ContinuousUnivariateDistribution
    α::T = 1.5
    β::T = zero(α)
    scale::T = one(α)
    location::T = zero(α)
end

# sampler(d::AlphaStable) = error("Not implemented")
# pdf(d::AlphaStable, x::Real) = error("Not implemented")
# logpdf(d::AlphaStable, x::Real) = error("Not implemented")
# cdf(d::AlphaStable, x::Real) = error("Not implemented")
# quantile(d::AlphaStable, q::Real) = error("Not implemented")
# minimum(d::AlphaStable) = error("Not implemented")
# maximum(d::AlphaStable) = error("Not implemented")
# insupport(d::AlphaStable, x::Real) = error("Not implemented")
Statistics.mean(d::AlphaStable) = d.α > 1 ? d.location : error("Not defined")
Statistics.var(d::AlphaStable) = d.α == 2 ? 2d.scale^2 : Inf
# modes(d::AlphaStable) = error("Not implemented")
# mode(d::AlphaStable) = error("Not implemented")
# skewness(d::AlphaStable) = error("Not implemented")
# kurtosis(d::Distribution, ::Bool) = error("Not implemented")
# entropy(d::AlphaStable, ::Real) = error("Not implemented")
# mgf(d::AlphaStable, ::Any) = error("Not implemented")
# cf(d::AlphaStable, ::Any) = error("Not implemented")

# lookup tables from McCulloch (1986)
const _ena = [
    2.4388
    2.5120
    2.6080
    2.7369
    2.9115
    3.1480
    3.4635
    3.8824
    4.4468
    5.2172
    6.3140
    7.9098
    10.4480
    14.8378
    23.4831
    44.2813
]
const _να = [
    2.439
    2.500
    2.600
    2.700
    2.800
    3.000
    3.200
    3.500
    4.000
    5.000
    6.000
    8.000
    10.000
    15.000
    25.000
]
const _νβ = [
    0.0
    0.1
    0.2
    0.3
    0.5
    0.7
    1.0
]
const ψ₁ = [
    2.000 2.000 2.000 2.000 2.000 2.000 2.000
    1.916 1.924 1.924 1.924 1.924 1.924 1.924
    1.808 1.813 1.829 1.829 1.829 1.829 1.829
    1.729 1.730 1.737 1.745 1.745 1.745 1.745
    1.664 1.663 1.663 1.668 1.676 1.676 1.676
    1.563 1.560 1.553 1.548 1.547 1.547 1.547
    1.484 1.480 1.471 1.460 1.448 1.438 1.438
    1.391 1.386 1.378 1.364 1.337 1.318 1.318
    1.279 1.273 1.266 1.250 1.210 1.184 1.150
    1.128 1.121 1.114 1.101 1.067 1.027 0.973
    1.029 1.021 1.014 1.004 0.974 0.935 0.874
    0.896 0.892 0.887 0.883 0.855 0.823 0.769
    0.818 0.812 0.806 0.801 0.780 0.756 0.691
    0.698 0.695 0.692 0.689 0.676 0.656 0.595
    0.593 0.590 0.588 0.586 0.579 0.563 0.513
]
const ψ₂ = [
    0.0 2.160 1.000 1.000 1.000 1.000 1.000
    0.0 1.592 3.390 1.000 1.000 1.000 1.000
    0.0 0.759 1.800 1.000 1.000 1.000 1.000
    0.0 0.482 1.048 1.694 2.229 1.000 1.000
    0.0 0.360 0.760 1.232 2.229 1.000 1.000
    0.0 0.253 0.518 0.823 1.575 1.000 1.000
    0.0 0.203 0.410 0.632 1.244 1.906 1.000
    0.0 0.165 0.332 0.499 0.943 1.560 1.000
    0.0 0.136 0.271 0.404 0.689 1.230 2.195
    0.0 0.109 0.216 0.323 0.539 0.827 1.917
    0.0 0.096 0.190 0.284 0.472 0.693 1.759
    0.0 0.082 0.163 0.243 0.412 0.601 1.596
    0.0 0.074 0.147 0.220 0.377 0.546 1.482
    0.0 0.064 0.128 0.191 0.330 0.478 1.362
    0.0 0.056 0.112 0.167 0.285 0.428 1.274
]
const _α =[
    0.5
    0.6
    0.7
    0.8
    0.9
    1.0
    1.1
    1.2
    1.3
    1.4
    1.5
    1.6
    1.7
    1.8
    1.9
    2.0
]
const _β = [
    0.0
    0.25
    0.50
    0.75
    1.00
]
const ϕ₃ = [
    2.588  3.073  4.534  6.636  9.144
    2.337  2.635  3.542  4.808  6.247
    2.189  2.392  3.004  3.844  4.775
    2.098  2.244  2.676  3.265  3.912
    2.040  2.149  2.461  2.886  3.356
    2.000  2.085  2.311  2.624  2.973
    1.980  2.040  2.205  2.435  2.696
    1.965  2.007  2.125  2.294  2.491
    1.955  1.984  2.067  2.188  2.333
    1.946  1.967  2.022  2.106  2.211
    1.939  1.952  1.988  2.045  2.116
    1.933  1.940  1.962  1.997  2.043
    1.927  1.930  1.943  1.961  1.987
    1.921  1.922  1.927  1.936  1.947
    1.914  1.915  1.916  1.918  1.921
    1.908  1.908  1.908  1.908  1.908
]
const ϕ₅ = [
    0.0 0.0    0.0    0.0    0.0
    0.0 -0.017 -0.032 -0.049 -0.064
    0.0 -0.030 -0.061 -0.092 -0.123
    0.0 -0.043 -0.088 -0.132 -0.179
    0.0 -0.056 -0.111 -0.170 -0.232
    0.0 -0.066 -0.134 -0.206 -0.283
    0.0 -0.075 -0.154 -0.241 -0.335
    0.0 -0.084 -0.173 -0.276 -0.390
    0.0 -0.090 -0.192 -0.310 -0.447
    0.0 -0.095 -0.208 -0.346 -0.508
    0.0 -0.098 -0.223 -0.383 -0.576
    0.0 -0.099 -0.237 -0.424 -0.652
    0.0 -0.096 -0.250 -0.469 -0.742
    0.0 -0.089 -0.262 -0.520 -0.853
    0.0 -0.078 -0.272 -0.581 -0.997
    0.0 -0.061 -0.279 -0.659 -1.198
]

"""
    fit(d::Type{<:AlphaStable}, x; alg=QuickSort)

Fit an α stable distribution to data.

returns `AlphaStable`

α∈[0.6,2.0], β∈[-1,1] , c∈[0,∞] and δ∈[-∞,∞] are the characteristic exponent,
skewness parameter, scale parameter (dispersion^1/α) and location parameter respectively.

α, β, c and δ are computed based on McCulloch (1986) fractile.
"""
function Distributions.fit(::Type{<:AlphaStable}, x, alg=QuickSort)
    sx = sort(x, alg=alg)
    p = quantile.(Ref(sx), (0.05, 0.25, 0.28, 0.5, 0.72, 0.75, 0.95), sorted=true)
    να = (p[7]-p[1]) / (p[6]-p[2])
    νβ = (p[7]+p[1]-2p[4]) / (p[7]-p[1])
    (να < _να[1]) && (να = _να[1])
    (να > _να[end]) && (να = _να[end])

    itp₁ = interpolate((_να, _νβ), ψ₁, Gridded(Linear()))
    α = itp₁(να, abs(νβ))
    itp₂ = interpolate((_να, _νβ), ψ₂, Gridded(Linear()))
    β = sign(νβ) * itp₂(να, abs(νβ))

    (β > 1.0) && (β = 1.0)
    (β < -1.0) && (β = -1.0)
    itp₃ = interpolate((_α, _β), ϕ₃, Gridded(Linear()))
    c = (p[6]-p[2]) / itp₃(α, abs(β))
    itp₄ = interpolate((_α, _β), ϕ₅, Gridded(Linear()))
    ζ = p[4] + c * sign(β) * itp₄(α, abs(β))
    if abs(α - 1.0) < 0.1
        δ = ζ
    else
        δ = ζ - β * c * tan(π*α/2)
    end
    return AlphaStable(α=α, β=β, scale=c, location=oftype(α, δ))
end

Base.@kwdef struct SymmetricAlphaStable{T} <: Distributions.ContinuousUnivariateDistribution
    α::T = 1.5
    scale::T = one(α)
    location::T = zero(α)
end

"""
    fit(d::Type{<:SymmetricAlphaStable}, x; alg=QuickSort)

Fit a symmetric α stable distribution to data.

returns `SymmetricAlphaStable`

α∈[1,2], c∈[0,∞] and δ∈[-∞,∞] are the characteristic exponent, scale parameter
(dispersion^1/α) and location parameter respectively.

α is computed based on McCulloch (1986) fractile.
scale is computed based on Fama & Roll (1971) fractile.
location is the 50% trimmed mean of the sample.
"""
function Distributions.fit(::Type{<:SymmetricAlphaStable}, x, alg=QuickSort)
    sx = sort(x, alg=alg)
    δ = mean(@view(sx[end÷4:(3*end)÷4]))
    p = quantile.(Ref(sx), (0.05, 0.25, 0.28, 0.72, 0.75, 0.95), sorted=true)
    c = (p[4]-p[3])/1.654
    an = (p[6]-p[1])/(p[5]-p[2])
    if an < 2.4388
        α = 2.
    else
        α = 0.
        j = findfirst(>=(an), _ena) # _np.where(an <= _ena[:,0])[0]
        (j === nothing || j == length(_ena)) && (j = length(_ena))
        t = (an-_ena[j-1])/(_ena[j]-_ena[j-1])
        α = (22-j-t)/10
    end
    if α < 0.5
        α = 0.5
    end
    return SymmetricAlphaStable(α=α, scale=c, location=oftype(α, δ))
end

"""
Generate independent stable random numbers.

:param α: characteristic exponent (0.1 to 2.0)
:param β: skew (-1 to +1)
:param scale: scale parameter
:param loc: location parameter (mean for α > 1, median/mode when β=0)

This implementation is based on the method in J.M. Chambers, C.L. Mallows
and B.W. Stuck, "A Method for Simulating Stable Random Variables," JASA 71 (1976): 340-4.
McCulloch's MATLAB implementation (1996) served as a reference in developing this code.
"""
function Base.rand(rng::AbstractRNG, d::AlphaStable)
    α=d.α; β=d.β; scale=d.scale; loc=d.location
    (α < 0.1 || α > 2) && throw(DomainError(α, "α must be in the range 0.1 to 2"))
    abs(β) > 1 && throw(DomainError(β, "β must be in the range -1 to 1"))
    ϕ = (rand(rng) - 0.5) * π
    if α == 1 && β == 0
        return loc + scale*tan(ϕ)
    end
    w = -log(rand(rng))
    α == 2 && (return loc + 2*scale*sqrt(w)*sin(ϕ))
    β == 0 && (return loc + scale * ((cos((1-α)*ϕ) / w)^(1.0/α - 1) * sin(α * ϕ) / cos(ϕ)^(1.0/α)))
    cosϕ = cos(ϕ)
    if abs(α-1) > 1e-8
        ζ = β * tan(π*α/2)
        aϕ = α * ϕ
        a1ϕ = (1-α) * ϕ
        return loc + scale * (( (sin(aϕ)+ζ*cos(aϕ))/cosϕ * ((cos(a1ϕ)+ζ*sin(a1ϕ))) / ((w*cosϕ)^((1-α)/α)) ))
    end
    bϕ = π/2 + β*ϕ
    x = 2/π * (bϕ*tan(ϕ) - β*log(π/2*w*cosϕ/bϕ))
    α == 1 || (x += β * tan(π*α/2))

    return loc + scale*x
end

"""

Generate alpha-sub-Gaussian (aSG) random numbers.

The implementation is based on https://github.com/ahmd-mahm/alpha-SGNm/blob/master/asgn.m

Reference:
A. Mahmood and M. Chitre, "Generating random variates for stable sub-Gaussian processes
with memory", Signal Processing, Volume 131, Pages 271-279, 2017.
(https://doi.org/10.1016/j.sigpro.2016.08.016.)

# Arguments
- `α`: characteristic exponent associated with the aSGN(m) process. This is
a scalar input and should lie within `collect(1.1:0.01:1.98)`.
- `R`: covariance matrix of any adjacent `m+1` samples in an aSGN(m) process.
The dimension of `R` is equal to `m+1`. It should be a symmetric toeplitz matrix.
The maximum acceptable size of `R` is `10x10`
- `n`: number of samples required

# Examples
```jldoctest
julia> x = rand(AlphaSubGaussian(n=1000))
```
"""
Base.@kwdef struct AlphaSubGaussian{T,M<:AbstractMatrix} <: Distributions.ContinuousUnivariateDistribution
    α::T = 1.50
    R::M = SMatrix{5,5}(collect(SymmetricToeplitz([1.0000, 0.5804, 0.2140, 0.1444, -0.0135])))
    n::Int
end

"""
Generates the conditional probability f(X2|X1) if [X1, X2] is a sub-Gaussian
stable random vector such that X1(i)~X2~S(alpha,delta) and rho is the correlation
coefficient of the underlying Gaussian vector. We assume the joint-probabiluty is given by f(X1,X2).
"""
function subgausscondprobtabulate(α, x1, x2_ind, invRx1, invR, vjoint, nmin, nmax, step, rind, kappa, k1, k2, kmarg)::Float64
    m = length(x1)
    r1 = sqrt(x1'*invRx1*x1)
    x = SVector{length(x1)+1, Float64}(x1..., x2_ind)
    r = sqrt(x'*invR*x)

    if r1<nmin
        grad = (vjoint[m, 1]-k2[1])/nmin
        cons = k2[1]
        vjointR1 = grad*r1+cons
    elseif r1>nmax
        vjointR1 = α*k1[1]*(r1^(-α-m))
    else
        ti = (log10(r1)-log10(nmin))/step+1
        tempind = (floor(Int, ti), ceil(Int, ti))
        grad = (vjoint[m, tempind[1]]-vjoint[m, tempind[2]])/(rind[tempind[1]]-rind[tempind[2]])
        cons = vjoint[m, tempind[1]]-grad*rind[tempind[1]]
        vjointR1 = grad*r1+cons
    end

    if r<nmin
        grad = (vjoint[m+1, 1]-k2[2])/nmin
        cons = k2[2]
        vjointR = grad*r+cons
    elseif r>nmax
        vjointR = α*k1[2]*(r^(-α-m-1))
    else
        ti = (log10(r)-log10(nmin))/step+1
        tempind = (floor(Int, ti), ceil(Int, ti))
        grad = (vjoint[m+1, tempind[1]]-vjoint[m+1, tempind[2]])/(rind[tempind[1]]-rind[tempind[2]])
        cons = vjoint[m+1, tempind[1]]-grad*rind[tempind[1]]
        vjointR = grad*r+cons
    end
    (1/sqrt(kappa))*kmarg*vjointR/vjointR1
end

function Random.rand!(rng::AbstractRNG, d::AlphaSubGaussian, x::AbstractArray)
    α=d.α; R=d.R; n=d.n
    length(x) >= n || throw(ArgumentError("length of x must be at least n"))
    α ∈ 1.10:0.01:1.98 || throw(DomainError(α, "α must lie within `1.10:0.01:1.98`"))
    m = size(R, 1)-1
    funk1 = x -> (2^α)*sin(π*α/2)*gamma((α+2)/2)*gamma((α+x)/2)/(gamma(x/2)*π*α/2)
    funk2 = x -> 4*gamma(x/α)/((α*2^2)*gamma(x/2)^2)
    funkmarg = x -> gamma(x/2)/(gamma((x-1)/2)*sqrt(π))
    c = 1.2
    k1 = (funk1(m), funk1(m+1))
    k2 = (funk2(m), funk2(m+1))
    kmarg = funkmarg(m+1)
    onetom = StaticArrays.SOneTo(m)
    kappa = det(R)/det(R[onetom, onetom])
    invR = inv(R)
    invRx1 = inv(R[onetom, onetom])
    sigrootx1 = cholesky(R[onetom, onetom]).L
    modefactor = R[end, onetom]'/R[onetom, onetom]
    matdict = matread(joinpath(@__DIR__(),"vr_repo/vr_alpha=$(α).mat"))
    nmax, nmin, res, rind, vjoint = matdict["Nmax"]::Float64, matdict["Nmin"]::Float64, matdict["res"]::Float64, vec(matdict["rind"])::Vector{Float64}, matdict["vJoint"]::Matrix{Float64}
    step = (log10(nmax)-log10(nmin))/res
    m>size(vjoint, 1)-1 && throw(DomainError(R, "The dimensions of `R` exceed the maximum possible 10x10"))
    A = rand(AlphaStable(α/2, 1.0, 2*cos(π*α/4)^(2.0/α), 0.0))
    T = rand(Chisq(m))
    S = randn(m)
    S = S/sqrt(sum(abs2,S))
    xtmp = ((sigrootx1*sqrt(A*T))*S)'
    if n<=m
        copyto!(x, @view(xtmp[1:n]))
    else
        # x = zeros(n)
        x[onetom] = xtmp
        vstud = α+m
        norms = pdf(TDist(vstud), 0.0)
        @inbounds for i = m+1:n
            x1 = SVector{m,Float64}(view(x,i-m:i-1))
            mode = modefactor*x1
            norm1 = subgausscondprobtabulate(α, x1, mode, invRx1, invR, vjoint, nmin, nmax, step, rind, kappa, k1, k2, kmarg)
            notaccept = true
            while notaccept
                u = rand()
                v = (norms/norm1)*rand(TDist(vstud)) + mode
                gv = (norm1/norms)*pdf(TDist(vstud), (v-mode)*(norm1/norms))
                fv = subgausscondprobtabulate(α, x1, v, invRx1, invR, vjoint, nmin, nmax, step, rind, kappa, k1, k2, kmarg)
                if c*u <= fv/gv
                    x[i] = v
                    notaccept = false
                end
            end
        end
    end
    x
end

Base.rand(rng::AbstractRNG, d::AlphaSubGaussian) = rand!(rng, d, zeros(d.n))

"""
    fit(d::Type{<:AlphaSubGaussian}, x, m; p=1.0)

Fit an aSGN(m) model to data via the covariation method.

The covariation method requires an additional parameter `p`. Ideally, 1 < p < α. In most practical impulsive scenarios p=1.0 is sufficient.
`m` is the number of lags in the covariance matrix.

The implementation is based on https://github.com/ahmd-mahm/alpha-SGNm/blob/master/param_est/asgnfit.m

Reference:
A. Mahmood and M. Chitre, "Generating random variates for stable sub-Gaussian processes
with memory", Signal Processing, Volume 131, Pages 271-279, 2017.
(https://doi.org/10.1016/j.sigpro.2016.08.016.)
"""
function Distributions.fit(d::Type{<:AlphaSubGaussian}, x::AbstractVector{T}, m::Integer; p=one(T)) where T
    d1   = fit(AlphaStable, x)
    α    = d1.α; scale=d1.scale
    cov  = zeros(T, m+1, m+1)
    xlen = length(x)
    c    = ((sum(x->abs(x)^p, x)/xlen)^(1/p))/scale
    for i in 1:m
        tempxlen = xlen-mod(xlen, i)
        xtemp = reshape(x[1:end-mod(xlen, i)], i, tempxlen÷i)
        if mod(tempxlen÷i, 2) != 0
            xtemp = xtemp[:, 1:end-1]
            tempxlen = size(xtemp, 1)*size(xtemp, 2)
        end
        xtemp = reshape(xtemp', 2, tempxlen÷2)
        @views r = (2/(c^p))*(scale^(2-p))*(xtemp[1, :]'*((sign.(xtemp[2, :]).*(abs.(xtemp[2, :]).^(p-1)))))/(tempxlen/2)
        cov[diagind(cov, i)] .+= r
    end
    cov = (cov+cov')+2*(scale^2)*I(m+1)
    cov ./= 2*scale^2
    AlphaSubGaussian(α=α, R=cov, n=length(x))
end

end # module


text:
module TestScikitLearn

using MLJBase
using Test

import MLJModels
import ScikitLearn
using MLJModels.ScikitLearn_
using CategoricalArrays

task = load_iris();
X, y = X_and_y(task)
train, test = partition(eachindex(y), 0.6) # levels of y are split across split

plain_classifier = SVMClassifier()
nu_classifier = SVMNuClassifier()
linear_classifier = SVMLClassifier(max_iter=10000)

plain_regressor = SVMRegressor()
nu_regressor = SVMNuRegressor()
linear_regressor = SVMLRegressor(max_iter=10000)

# Xr = [[0, 0], [1, 1]]
# yr = [0, 1]
# targetr = [[1, 1]]

task = load_boston()
Xr, yr = X_and_y(task)

fitresultC, cacheC, reportC = MLJBase.fit(plain_classifier, 1,
                                          selectrows(X, train), y[train]);
fitresultCnu, cacheCnu, reportCnu = MLJBase.fit(nu_classifier, 1,
                                          selectrows(X, train), y[train]);
fitresultCL, cacheCL, reportCL = MLJBase.fit(linear_classifier, 1,
                                          selectrows(X, train), y[train]);

fitresultR, cacheR, reportR = MLJBase.fit(plain_regressor, 1, Xr, yr);
fitresultRnu, cacheRnu, reportRnu = MLJBase.fit(nu_regressor, 1, Xr, yr);
fitresultRL, cacheRL, reportRL = MLJBase.fit(linear_regressor, 1, Xr, yr);

pcpred = predict(plain_classifier, fitresultC, selectrows(X, test));
nucpred = predict(nu_classifier, fitresultCnu, selectrows(X, test));
lcpred = predict(linear_classifier, fitresultCL, selectrows(X, test));

@test levels(pcpred) == levels(y[train])
@test levels(nucpred) == levels(y[train])
@test levels(lcpred) == levels(y[train])

rpred = predict(plain_regressor, fitresultR, targetr);
nurpred = predict(nu_regressor, fitresultRnu, targetr);
Lrpred = predict(linear_regressor, fitresultRL, targetr);

end
true

max_stars_repo_name: ablaom/MLJModels

max_issues_repo_name: ablaom/MLJModels

max_forks_repo_name: ablaom/MLJModels


text:
include("runtests_preamble.jl")

using Oceananigans.Units
using LESbrary.IdealizedExperiments: three_layer_constant_fluxes_simulation

@testset "Three layer constant flux simulation" begin
    simulation = three_layer_constant_fluxes_simulation(size = (1, 1, 32),
                                                        stop_time = 2.0,
                                                        snapshot_time_interval = 1.0,
                                                        averages_time_interval = 1.0,
                                                        averages_time_window = 1.0)
    run!(simulation)

    @test simulation.stop_time == 2.0
end

max_stars_repo_name: ali-ramadhan/LESbrary

max_issues_repo_name: ali-ramadhan/LESbrary

max_forks_repo_name: ali-ramadhan/LESbrary


text:
using GalerkinSparseGrids
using Base.Test
using Cubature

#--------------------------------------
# Elementary Tests
#--------------------------------------

ε = eps(Float64)

print("Testing cell_index function... ")
import GalerkinSparseGrids.cell_index
for l in 1:5
	@test cell_index(1.3,l)==(1<<(l-1))
end
for l in 1:5
	@test cell_index(.01,l)==1
end
@test cell_index(.3,1)==1
@test cell_index(.3,2)==1
@test cell_index(.3,3)==2
@test cell_index(.3,4)==3
@test cell_index(.3,5)==5
println("Tests Passed.")

import GalerkinSparseGrids.inner_product

print("Testing 1D inner product... ")
@test inner_product(x->x[1]^2, x->x[1]^3, (0,),CartesianIndex((1,))) ≈ (1/6) atol=2*ε
@test inner_product(x->sin(pi*x[1]), x->cos(pi*x[1]), (0,),CartesianIndex((1,))) ≈ 0 atol=2*ε
println("Tests Passed.")

print("Testing 2D inner product... ")
@test inner_product(x->(x[1]^2+x[2]^2), x->x[1]^3, (0,0),CartesianIndex((1,1))) ≈ .25 atol=2*ε
println("Tests Passed.")

print("Testing 3D inner product... ")
@test inner_product(x->(x[1]^2+x[2]^2-x[3]^2), x->x[1]^3+x[2]+x[3], (0,0,0),CartesianIndex((1,1,1))) ≈ .5 atol=2*ε
println("Tests Passed.")


text:
used_bits(x::Union{Signed, Unsigned}) = sizeof(x)*8 - leading_zeros(x)

"""
    kronecker_product(a, b)

Returns `a ⊗ b`, see: https://en.wikipedia.org/wiki/Kronecker_product.
"""
function kronecker_product(a::T, b::T) where T <: AbstractArray
    out = similar(a, size(a) .* size(b))
    for i in eachindex(IndexCartesian(), a)
        for j in eachindex(IndexCartesian(), b)
            out[CartesianIndex((Tuple(i) .- 1) .* size(b) .+ Tuple(j))] = a[i]*b[j]
        end
    end
    out
end

"""
    kronecker_power(a, power)

returns `a` kroneker producted with itself `power` times. Does not copy when power == 1.

See also: [`kronecker_product`](@ref)
"""
function kronecker_power(a::AbstractArray, power::Integer)
    if power < 0
        throw(ArgumentError("Power must be non-negative, got $power < 0."))
    elseif power == 0
        out = similar(a, ntuple(_->1, ndims(a)))
        out[firstindex(out)] = one(eltype(a))
        return out
    else power == 1
        out = a
        for i in 2:power
            out = kronecker_product(out, a)
        end
        out
    end
end

max_stars_repo_name: LilithHafner/FBD

max_issues_repo_name: LilithHafner/FBD

max_forks_repo_name: LilithHafner/FBD


text:
reset_default_graph(); sess = Session()

@testset "test_jacobian" begin
    function f(x)
        x.^3, 3*diagm(0=>x.^2)
    end
    err1, err2 = test_jacobian(f, rand(10))
    @test err2[end] < err1[end]
end

@testset "lineview" begin
    @test begin
        a = placeholder(rand(10))
        loss = sum((a+1)^2)
        close("all")
        lineview(sess, a, loss, -ones(10))
        savefig("lineview.png")
        true
    end
end

@testset "meshview" begin
    @test begin
        a = placeholder(rand(10))
        loss = sum((a+1)^2)
        close("all")
        meshview(sess, a, loss, -ones(10))
        savefig("meshview.png")
        true
    end
    @test begin
        a = placeholder(rand(10))
        loss = sum((a+1)^2)
        close("all")
        pcolormeshview(sess, a, loss, -ones(10))
        savefig("pcolormeshview.png")
        true
    end
end

@testset "gradview" begin
    @test begin
        a = placeholder(rand(10))
        loss = sum((a+1.0)^2)
        close("all")
        gradview(sess, a, loss, rand(10))
        true
    end
end

@testset "jacview" begin
    @test begin
        u0 = rand(10)
        function verify_jacobian_f(θ, u)
            r = u^3+u - u0
            r, spdiag(3u^2+1.0)
        end
        close("all")
        jacview(sess, verify_jacobian_f, missing, u0)
        savefig("jacview.png")
        true
    end
end

@testset "PCLview" begin
end

@testset "animate" begin
    @test begin
        θ = LinRange(0, 2π, 100)
        x = cos.(θ)
        y = sin.(θ)
        pl, = plot([], [], "o-")
        t = title("0")
        xlim(-1.2,1.2)
        ylim(-1.2,1.2)
        function update(i)
            t.set_text("$i")
            pl.set_data([x[1:i] y[1:i]]'|>Array)
        end
        p = animate(update, 1:100)
        saveanim(p, "anim.gif")
        true
    end
end


text:
"""
    VFE(fz::FiniteGP)

The "Variational Free Energy" sparse approximation [1], used to construct an
approximate posterior with inducing inputs `fz.x`. See [`posterior(v::VFE,
fx::FiniteGP, y::AbstractVector{<:Real})`](@ref) for a usage example.

[1] - M. K. Titsias. "Variational learning of inducing variables in sparse Gaussian
processes". In: Proceedings of the Twelfth International Conference on Artificial
Intelligence and Statistics. 2009.
"""
struct VFE{Tfz<:FiniteGP}
    fz::Tfz
end

const DTC = VFE

struct ApproxPosteriorGP{Tapprox,Tprior,Tdata} <: AbstractGP
    approx::Tapprox
    prior::Tprior
    data::Tdata
end

"""
    posterior(vfe::VFE, fx::FiniteGP, y::AbstractVector{<:Real})

Compute the optimal approximate posterior [1] over the process `f = fx.f`, given observations `y`
of `f` at `x`, and inducing points `vfe.fz.x`.

```jldoctest
julia> f = GP(Matern52Kernel());

julia> x = randn(1000);

julia> z = range(-5.0, 5.0; length=13);

julia> vfe = VFE(f(z));

julia> y = rand(f(x, 0.1));

julia> post = posterior(vfe, f(x, 0.1), y);

julia> post(z) isa AbstractGPs.FiniteGP
true
```

[1] - M. K. Titsias. "Variational learning of inducing variables in sparse Gaussian
processes". In: Proceedings of the Twelfth International Conference on Artificial
Intelligence and Statistics. 2009.
"""
function posterior(vfe::VFE, fx::FiniteGP, y::AbstractVector{<:Real})
    @assert vfe.fz.f === fx.f

    U_y = _cholesky(_symmetric(fx.Σy)).U
    U = cholesky(_symmetric(cov(vfe.fz))).U

    B_εf = U' \ (U_y' \ cov(fx, vfe.fz))'

    b_y = U_y' \ (y - mean(fx))

    D = B_εf * B_εf' + I
    Λ_ε = cholesky(_symmetric(D))

    m_ε = Λ_ε \ (B_εf * b_y)

    cache = (m_ε=m_ε, Λ_ε=Λ_ε, U=U, α=U \ m_ε, b_y=b_y, B_εf=B_εf, x=fx.x, Σy=fx.Σy)
    return ApproxPosteriorGP(vfe, fx.f, cache)
end

"""
    function update_posterior(
        f_post_approx::ApproxPosteriorGP{<:VFE},
        fx::FiniteGP,
        y::AbstractVector{<:Real}
    )

Update the `ApproxPosteriorGP` given a new set of observations. Here, we retain the same
set of pseudo-points.
"""
function update_posterior(
    f_post_approx::ApproxPosteriorGP{<:VFE}, fx::FiniteGP, y::AbstractVector{<:Real}
)
    @assert f_post_approx.prior === fx.f

    U = f_post_approx.data.U
    z = inducing_points(f_post_approx)

    U_y₂ = _cholesky(_symmetric(fx.Σy)).U

    temp = zeros(size(f_post_approx.data.Σy, 1), size(fx.Σy, 2))
    Σy = [f_post_approx.data.Σy temp; temp' fx.Σy]

    b_y = vcat(f_post_approx.data.b_y, U_y₂ \ (y - mean(fx)))

    B_εf₂ = U' \ (U_y₂' \ cov(fx.f, fx.x, z))'
    B_εf = hcat(f_post_approx.data.B_εf, B_εf₂)

    Λ_ε = f_post_approx.data.Λ_ε

    for col in eachcol(B_εf₂)
        lowrankupdate!(Λ_ε, col)
    end

    m_ε = Λ_ε \ (B_εf * b_y)
    α = U \ m_ε
    x = vcat(f_post_approx.data.x, fx.x)

    cache = (m_ε=m_ε, Λ_ε=Λ_ε, U=U, α=α, z=z, b_y=b_y, B_εf=B_εf, x=x, Σy=Σy)
    return ApproxPosteriorGP(f_post_approx.approx, fx.f, cache)
end

"""
    function update_posterior(
        f_post_approx::ApproxPosteriorGP{<:VFE},
        z::FiniteGP,
    )

Update the `ApproxPosteriorGP` given a new set of pseudo-points to append to the existing
set of pseudo-points.
"""
function update_posterior(f_post_approx::ApproxPosteriorGP{<:VFE}, fz::FiniteGP)
    @assert f_post_approx.prior === fz.f

    z_old = inducing_points(f_post_approx)
    z = fz.x

    U11 = f_post_approx.data.U
    C12 = cov(f_post_approx.prior, z_old, z)
    C22 = _symmetric(cov(f_post_approx.prior, z))
    U = update_chol(Cholesky(U11, 'U', 0), C12, C22).U
    U22 = U[(end - length(z) + 1):end, (end - length(z) + 1):end]
    U12 = U[1:length(z_old), (end - length(z) + 1):end]

    B_εf₁ = f_post_approx.data.B_εf

    Cu1f = cov(f_post_approx.prior, z_old, f_post_approx.data.x)
    Cu2f = cov(f_post_approx.prior, z, f_post_approx.data.x)

    U_y = _cholesky(_symmetric(f_post_approx.data.Σy)).U

    B_εf₂ = U22' \ (Cu2f * inv(U_y) - U12' * B_εf₁)
    B_εf = vcat(B_εf₁, B_εf₂)

    Λ_ε = update_chol(f_post_approx.data.Λ_ε, B_εf₁ * B_εf₂', B_εf₂ * B_εf₂' + I)

    m_ε = Λ_ε \ (B_εf * f_post_approx.data.b_y)

    α = U \ m_ε

    z_new = vcat(z_old, z)
    vfe = f_post_approx.approx
    fz_new = vfe.fz.f(z_new, vfe.fz.Σy)

    cache = (
        m_ε=m_ε,
        Λ_ε=Λ_ε,
        U=U,
        α=α,
        b_y=f_post_approx.data.b_y,
        B_εf=B_εf,
        x=f_post_approx.data.x,
        Σy=f_post_approx.data.Σy,
    )
    return ApproxPosteriorGP(VFE(fz_new), f_post_approx.prior, cache)
end

# AbstractGP interface implementation.

function Statistics.mean(f::ApproxPosteriorGP{<:VFE}, x::AbstractVector)
    return mean(f.prior, x) + cov(f.prior, x, inducing_points(f)) * f.data.α
end

function Statistics.cov(f::ApproxPosteriorGP{<:VFE}, x::AbstractVector)
    A = f.data.U' \ cov(f.prior, inducing_points(f), x)
    return cov(f.prior, x) - At_A(A) + Xt_invA_X(f.data.Λ_ε, A)
end

function Statistics.var(f::ApproxPosteriorGP{<:VFE}, x::AbstractVector)
    A = f.data.U' \ cov(f.prior, inducing_points(f), x)
    return var(f.prior, x) - diag_At_A(A) + diag_Xt_invA_X(f.data.Λ_ε, A)
end

function Statistics.cov(f::ApproxPosteriorGP{<:VFE}, x::AbstractVector, y::AbstractVector)
    A_zx = f.data.U' \ cov(f.prior, inducing_points(f), x)
    A_zy = f.data.U' \ cov(f.prior, inducing_points(f), y)
    return cov(f.prior, x, y) - A_zx'A_zy + Xt_invA_Y(A_zx, f.data.Λ_ε, A_zy)
end

function StatsBase.mean_and_cov(f::ApproxPosteriorGP{<:VFE}, x::AbstractVector)
    A = f.data.U' \ cov(f.prior, inducing_points(f), x)
    m_post = mean(f.prior, x) + A' * f.data.m_ε
    C_post = cov(f.prior, x) - At_A(A) + Xt_invA_X(f.data.Λ_ε, A)
    return m_post, C_post
end

function StatsBase.mean_and_var(f::ApproxPosteriorGP{<:VFE}, x::AbstractVector)
    A = f.data.U' \ cov(f.prior, inducing_points(f), x)
    m_post = mean(f.prior, x) + A' * f.data.m_ε
    c_post = var(f.prior, x) - diag_At_A(A) + diag_Xt_invA_X(f.data.Λ_ε, A)
    return m_post, c_post
end

inducing_points(f::ApproxPosteriorGP{<:VFE}) = f.approx.fz.x

"""
    elbo(vfe::VFE, fx::FiniteGP, y::AbstractVector{<:Real})

The Titsias Evidence Lower BOund (ELBO) [1]. `y` are observations of `fx`, and `v.z`
are inducing points.

```jldoctest
julia> f = GP(Matern52Kernel());

julia> x = randn(1000);

julia> z = range(-5.0, 5.0; length=13);

julia> v = VFE(f(z));

julia> y = rand(f(x, 0.1));

julia> elbo(v, f(x, 0.1), y) < logpdf(f(x, 0.1), y)
true
```

[1] - M. K. Titsias. "Variational learning of inducing variables in sparse Gaussian
processes". In: Proceedings of the Twelfth International Conference on Artificial
Intelligence and Statistics. 2009.
"""
function elbo(vfe::VFE, fx::FiniteGP, y::AbstractVector{<:Real})
    @assert vfe.fz.f === fx.f
    _dtc, A = _compute_intermediates(fx, y, vfe.fz)
    return _dtc - (tr_Cf_invΣy(fx, fx.Σy) - sum(abs2, A)) / 2
end

"""
    dtc(v::VFE, fx::FiniteGP, y::AbstractVector{<:Real})

The Deterministic Training Conditional (DTC) [1]. `y` are observations of `fx`, and `v.z`
are inducing points.

```jldoctest
julia> f = GP(Matern52Kernel());

julia> x = randn(1000);

julia> z = range(-5.0, 5.0; length=256);

julia> v = VFE(f(z));

julia> y = rand(f(x, 0.1));

julia> isapprox(dtc(v, f(x, 0.1), y), logpdf(f(x, 0.1), y); atol=1e-6, rtol=1e-6)
true
```

[1] - M. Seeger, C. K. I. Williams and N. D. Lawrence. "Fast Forward Selection to Speed Up
Sparse Gaussian Process Regression". In: Proceedings of the Ninth International Workshop on
Artificial Intelligence and Statistics. 2003
"""
function dtc(vfe::VFE, fx::FiniteGP, y::AbstractVector{<:Real})
    @assert vfe.fz.f === fx.f
    _dtc, _ = _compute_intermediates(fx, y, vfe.fz)
    return _dtc
end

# Factor out computations common to the `elbo` and `dtc`.
function _compute_intermediates(fx::FiniteGP, y::AbstractVector{<:Real}, fz::FiniteGP)
    consistency_check(fx, y)
    chol_Σy = _cholesky(fx.Σy)

    A = cholesky(_symmetric(cov(fz))).U' \ (chol_Σy.U' \ cov(fx, fz))'
    Λ_ε = cholesky(Symmetric(A * A' + I))
    δ = chol_Σy.U' \ (y - mean(fx))

    tmp = logdet(chol_Σy) + logdet(Λ_ε) + sum(abs2, δ) - sum(abs2, Λ_ε.U' \ (A * δ))
    _dtc = -(length(y) * typeof(tmp)(log2π) + tmp) / 2
    return _dtc, A
end

function consistency_check(fx, y)
    @assert length(fx) == length(y)
end

function tr_Cf_invΣy(f::FiniteGP, Σy::Diagonal)
    return sum(var(f.f, f.x) ./ diag(Σy))
end


text:
# Unit tests for JUDI Geometry structure
# Philipp Witte (pwitte.slim@gmail.com)
# May 2018
#

using JUDI.TimeModeling, SeisIO, Test, LinearAlgebra

@testset "Geometry Unit Test" begin

    # Number of sources
    nsrc = 2

    # Constructor if nt is not passed
    xsrc = convertToCell(range(100f0, stop=1100f0, length=nsrc))
    ysrc = convertToCell(range(0f0, stop=0f0, length=nsrc))
    zsrc = convertToCell(range(20f0, stop=20f0, length=nsrc))

    geometry =  Geometry(xsrc, ysrc, zsrc; dt=2f0, t=1000f0)

    @test isequal(typeof(geometry.xloc), Array{Any, 1})
    @test isequal(typeof(geometry.yloc), Array{Any, 1})
    @test isequal(typeof(geometry.zloc), Array{Any, 1})
    @test isequal(typeof(geometry.nt), Array{Any, 1})
    @test isequal(typeof(geometry.dt), Array{Any, 1})
    @test isequal(typeof(geometry.t), Array{Any, 1})

    # Constructor if coordinates are not passed as a cell arrays
    xsrc = range(100f0, stop=1100f0, length=nsrc)
    ysrc = range(0f0, stop=0f0, length=nsrc)
    zsrc = range(20f0, stop=20f0, length=nsrc)

    geometry = Geometry(xsrc, ysrc, zsrc; dt=4f0, t=1000f0, nsrc=nsrc)

    @test isequal(typeof(geometry.xloc), Array{Any, 1})
    @test isequal(typeof(geometry.yloc), Array{Any, 1})
    @test isequal(typeof(geometry.zloc), Array{Any, 1})
    @test isequal(typeof(geometry.nt), Array{Any, 1})
    @test isequal(typeof(geometry.dt), Array{Any, 1})
    @test isequal(typeof(geometry.t), Array{Any, 1})

    # Set up source geometry object from in-core data container
    block = segy_read("../data/unit_test_shot_records.segy")
    src_geometry = Geometry(block; key="source", segy_depth_key="SourceSurfaceElevation")
    rec_geometry = Geometry(block; key="receiver", segy_depth_key="RecGroupElevation")

    @test isequal(typeof(src_geometry), GeometryIC)
    @test isequal(typeof(rec_geometry), GeometryIC)
    @test isequal(get_header(block, "SourceSurfaceElevation")[1], src_geometry.zloc[1])
    @test isequal(get_header(block, "RecGroupElevation")[1], rec_geometry.zloc[1][1])
    @test isequal(get_header(block, "SourceX")[1], src_geometry.xloc[1])
    @test isequal(get_header(block, "GroupX")[1], rec_geometry.xloc[1][1])

    # Set up geometry summary from out-of-core data container
    container = segy_scan("../data/", "unit_test_shot_records", ["GroupX", "GroupY", "RecGroupElevation", "SourceSurfaceElevation", "dt"])
    src_geometry = Geometry(container; key="source", segy_depth_key="SourceSurfaceElevation")
    rec_geometry = Geometry(container; key="receiver", segy_depth_key="RecGroupElevation")

    @test isequal(typeof(src_geometry), GeometryOOC)
    @test isequal(typeof(rec_geometry), GeometryOOC)
    @test isequal(src_geometry.key, "source")
    @test isequal(rec_geometry.key, "receiver")
    @test isequal(src_geometry.segy_depth_key, "SourceSurfaceElevation")
    @test isequal(rec_geometry.segy_depth_key, "RecGroupElevation")
    @test isequal(prod(size(block.data)), sum(rec_geometry.nsamples))

    # Set up geometry summary from out-of-core data container passed as cell array
    container_cell = Array{SeisIO.SeisCon}(undef, nsrc)
    for j=1:nsrc
        container_cell[j] = split(container, j)
    end

    src_geometry = Geometry(container_cell; key="source", segy_depth_key="SourceSurfaceElevation")
    rec_geometry = Geometry(container_cell; key="receiver", segy_depth_key="RecGroupElevation")

    @test isequal(typeof(src_geometry), GeometryOOC)
    @test isequal(typeof(rec_geometry), GeometryOOC)
    @test isequal(src_geometry.key, "source")
    @test isequal(rec_geometry.key, "receiver")
    @test isequal(src_geometry.segy_depth_key, "SourceSurfaceElevation")
    @test isequal(rec_geometry.segy_depth_key, "RecGroupElevation")
    @test isequal(prod(size(block.data)), sum(rec_geometry.nsamples))

    # Load geometry from out-of-core Geometry container
    src_geometry_ic = Geometry(src_geometry)
    rec_geometry_ic = Geometry(rec_geometry)

    @test isequal(typeof(src_geometry_ic), GeometryIC)
    @test isequal(typeof(rec_geometry_ic), GeometryIC)
    @test isequal(get_header(block, "SourceSurfaceElevation")[1], src_geometry_ic.zloc[1])
    @test isequal(get_header(block, "RecGroupElevation")[1], rec_geometry_ic.zloc[1][1])
    @test isequal(get_header(block, "SourceX")[1], src_geometry_ic.xloc[1])
    @test isequal(get_header(block, "GroupX")[1], rec_geometry_ic.xloc[1][1])

    # Subsample in-core geometry structure
    src_geometry_sub = subsample(src_geometry_ic, 1)
    @test isequal(typeof(src_geometry_sub), GeometryIC)
    @test isequal(length(src_geometry_sub.xloc), 1)

    src_geometry_sub = subsample(src_geometry_ic, 1:2)
    @test isequal(typeof(src_geometry_sub), GeometryIC)
    @test isequal(length(src_geometry_sub.xloc), 2)

    # Subsample out-of-core geometry structure
    src_geometry_sub = subsample(src_geometry, 1)
    @test isequal(typeof(src_geometry_sub), GeometryOOC)
    @test isequal(length(src_geometry_sub.dt), 1)
    @test isequal(src_geometry_sub.segy_depth_key, "SourceSurfaceElevation")

    src_geometry_sub = subsample(src_geometry, 1:2)
    @test isequal(typeof(src_geometry_sub), GeometryOOC)
    @test isequal(length(src_geometry_sub.dt), 2)
    @test isequal(src_geometry_sub.segy_depth_key, "SourceSurfaceElevation")

    # Compare if geometries match
    @test compareGeometry(src_geometry_ic, src_geometry_ic)
    @test compareGeometry(rec_geometry_ic, rec_geometry_ic)

    @test compareGeometry(src_geometry, src_geometry)
    @test compareGeometry(rec_geometry, rec_geometry)

end


text:
"""
	get_population(rho_vec::Array{ComplexF64,1})

# Argument:
- rho_vec, density matrix in the vector form
# Output:
- the population in an array
"""
function get_population(rho_vec::Array{ComplexF64,1})
	N = size(rho_vec)[1]
	N = Int64(sqrt(N))
    P = zeros(Float64,N)
	# compute the trace for R
    for j = 1:N
        P[j] = P[j] + real(rho_vec[1+(j-1)*(N+1)])
    end
    return P
end

function get_population(rho_vec::Array{Float64,1})
	N = size(rho_vec)[1]
	N = Int64(sqrt(N))
    P = zeros(Float64,N)
	# compute the trace for R
    for j = 1:N
        P[j] = P[j] + rho_vec[1+(j-1)*(N+1)]
    end
    return P
end

function get_population_UQ_vector(rho_vec)
	N = size(rho_vec)[1]
	N = Int64(sqrt(N))
    #P = zeros(Float64,N)
	# compute the trace for R
	P = rho_vec[1]
    for j = 2:N
        P = [P;rho_vec[1+(j-1)*(N+1)]]
    end
    return P
end
"""
	get_population(RhoRealHistory::Array{Float64,2})

# Argument:
- RhoRealHistory: time history for the real part of the density matrix (saved as a flattened vector), j-th column corresponding to the j-th time point

rho_history = rhou_history - i rhov_history.

Since the density matrix should be Hermitian, it is good enough to just have the
real part.

# Output:
- The population history P which is a 2D array, P[i,j] is corresponding
to time point i and state j
"""
function get_population(rho_u_history::Array{Float64,2})
	# rho_u_history is RhoRealHistory
	N,nt = size(rho_u_history)
	N = Int64(sqrt(N))
    return transpose(rho_u_history[1:N+1:end,:])
end

function get_population(rho_u_history::Array{Array{Float64}})
	N,nt = size(rho_u_history)
	N = Int64(sqrt(N))
    return transpose(rho_u_history[1:N+1:end,:])
end

function get_population_UQ_matrix(rho_u_history)
	N,nt = size(rho_u_history)
	N = Int64(sqrt(N))
    return transpose(rho_u_history[1:N+1:end,:])
end

"""
	get_population(RhoVecHistory::Array{ComplexF64,2})

# Argument:
- RhoVecHistory, time history for the vectorized density matrix,
i-th column corresponding to the i-th time point
# Output:
- The population history P which is a 2D array, P[i,j] is corresponding
to time point i and state j
"""
function get_population(rho_vec_history::Array{ComplexF64,2})
	N,nt = size(rho_vec_history)
	N = Int64(sqrt(N))
	P = zeros(Float64,nt,N)
	# compute the trace
    return transpose(real(rho_vec_history[1:N+1:end,:]))
end

function get_population(rho_vec_history::Array{Array{ComplexF64}})
	N,nt = size(rho_vec_history)
	N = Int64(sqrt(N))
	P = zeros(Float64,nt,N)
	# compute the trace for R
	return transpose(real(rho_vec_history[1:N+1:end,:]))
end
"""
	convert_state_to_density(u::Array{Float64,1},v::Array{Float64,1})

# Argument:
- state vector: u-iv
# Output:
- vectorized density matrix: rho_u-i rho_v
"""
function convert_state_to_density(u::Array{Float64,1},v::Array{Float64,1})
	return (u*transpose(u)+v*transpose(v))[:],(v*transpose(u)-u*transpose(v))[:]
end
"""
	hello_world():
A hello world function
"""
function hello_world()
	println("Hello Quantum World!")
end


text:
mutable struct gpuCTM <: GPUTopicModel
	K::Int
	M::Int
	V::Int
	N::Vector{Int}
	C::Vector{Int}
	B::Int
	corp::Corpus
	batches::Vector{UnitRange{Int}}
	topics::VectorList{Int}
	mu::Vector{Float32}
	sigma::Matrix{Float32}
	beta::Matrix{Float32}
	lambda::VectorList{Float32}
	vsq::VectorList{Float32}
	lzeta::Vector{Float32}
	phi::MatrixList{Float32}
	invsigma::Matrix{Float32}
	newbeta::Void
	Npsums::VectorList{Int}
	Jpsums::VectorList{Int}
	terms::VectorList{Int}
	counts::VectorList{Int}
	words::VectorList{Int}
	device::cl.Device
	context::cl.Context
	queue::cl.CmdQueue
	mukern::cl.Kernel
	betakern::cl.Kernel
	betanormkern::cl.Kernel
	newbetakern::cl.Kernel
	lambdakern::cl.Kernel
	vsqkern::cl.Kernel
	lzetakern::cl.Kernel
	phikern::cl.Kernel
	phinormkern::cl.Kernel
	Cbuf::cl.Buffer{Int}
	Npsumsbuf::cl.Buffer{Int}
	Jpsumsbuf::cl.Buffer{Int}
	termsbuf::cl.Buffer{Int}
	countsbuf::cl.Buffer{Int}
	wordsbuf::cl.Buffer{Int}
	newtontempbuf::cl.Buffer{Float32}
	newtongradbuf::cl.Buffer{Float32}
	newtoninvhessbuf::cl.Buffer{Float32}
	mubuf::cl.Buffer{Float32}
	sigmabuf::cl.Buffer{Float32}
	invsigmabuf::cl.Buffer{Float32}
	betabuf::cl.Buffer{Float32}
	newbetabuf::cl.Buffer{Float32}
	lambdabuf::cl.Buffer{Float32}
	vsqbuf::cl.Buffer{Float32}
	lzetabuf::cl.Buffer{Float32}
	phibuf::cl.Buffer{Float32}
	elbo::Float32
	newelbo::Float32

	function gpuCTM(corp::Corpus, K::Integer, batchsize::Integer=length(corp))
		@assert !isempty(corp)
		@assert all(ispositive.([K, batchsize]))
		checkcorp(corp)

		M, V, U = size(corp)
		N = [length(doc) for doc in corp]
		C = [size(doc) for doc in corp]

		batches = partition(1:M, batchsize)
		B = length(batches)

		topics = [collect(1:V) for _ in 1:K]

		mu = zeros(K)
		sigma = eye(K)
		beta = rand(Dirichlet(V, 1.0), K)'
		lambda = [zeros(K) for _ in 1:M]
		vsq = [ones(K) for _ in 1:M]
		lzeta = zeros(M)
		phi = [ones(K, N[d]) / K for d in batches[1]]

		model = new(K, M, V, N, C, B, copy(corp), batches, topics, mu, sigma, beta, lambda, vsq, lzeta, phi)
		fixmodel!(model, check=false)

		for (b, batch) in enumerate(batches)
			model.phi = [ones(K, N[d]) / K for d in batch]
			updateNewELBO!(model, b)
		end
		model.phi = [ones(K, N[d]) / K for d in batches[1]]
		updateELBO!(model)
		return model
	end
end

function Elogpeta(model::gpuCTM, d::Int)
	x = 0.5 * (logdet(model.invsigma) - model.K * log(2pi) - dot(diag(model.invsigma), model.vsq[d]) - dot(model.lambda[d] - model.mu, model.invsigma * (model.lambda[d] - model.mu)))
	return x
end

function Elogpz(model::gpuCTM, d::Int, m::Int)
	counts = model.corp[d].counts
	x = dot(model.phi[m]' * model.lambda[d], counts) + model.C[d] * model.lzeta[d]
	return x
end

function Elogpw(model::gpuCTM, d::Int, m::Int)
	terms, counts = model.corp[d].terms, model.corp[d].counts
	x = sum(model.phi[m] .* log.(@boink model.beta[:,terms]) * counts)
	return x
end

function Elogqeta(model::gpuCTM, d::Int)
	x = -entropy(MvNormal(model.lambda[d], diagm(model.vsq[d])))
	return x
end

function Elogqz(model::gpuCTM, d::Int, m::Int)
	counts = model.corp[d].counts
	x = -sum([c * entropy(Categorical(model.phi[m][:,n])) for (n, c) in enumerate(counts)])
	return x
end

function updateELBO!(model::gpuCTM)
	model.elbo = model.newelbo
	model.newelbo = 0
	return model.elbo
end

function updateNewELBO!(model::gpuCTM, b::Int)
	batch = model.batches[b]
	for (m, d) in enumerate(batch)
		model.newelbo += (Elogpeta(model, d)
						+ Elogpz(model, d, m)
						+ Elogpw(model, d, m)
						- Elogqeta(model, d)
						- Elogqz(model, d, m))
	end
end

const CTM_MU_cpp =
"""
kernel void
updateMu(long K,
			long M,
			const global float *lambda,
			global float *mu)

			{
			long i = get_global_id(0);

			float acc = 0.0f;

			for (long d=0; d<M; d++)
				acc += lambda[K * d + i];

			mu[i] = acc / M;
			}
			"""

function updateMu!(model::gpuCTM)
	model.queue(model.mukern, model.K, nothing, model.K, model.M, model.lambdabuf, model.mubuf)
end

function updateSigma!(model::gpuCTM)
	@host model.mubuf
	@host model.lambdabuf
	@host model.vsqbuf

	model.sigma = diagm(sum(model.vsq)) / model.M + Base.covm(hcat(model.lambda...), model.mu, 2, false)
	model.invsigma = inv(model.sigma)

	@buf model.sigma
	@buf model.invsigma
end

const CTM_BETA_cpp =
"""
kernel void
updateBeta(long K,
			global float *newbeta,
			global float *beta)

			{
			long i = get_global_id(0);
			long j = get_global_id(1);

			beta[K * j + i] = newbeta[K * j + i];
			newbeta[K * j + i] = 0.0f;
			}
			"""

const CTM_BETA_NORM_cpp =
"""
kernel void
normalizeBeta(long K,
				long V,
				global float *beta)

				{
				long i = get_global_id(0);

				float normalizer = 0.0f;

				for (long j=0; j<V; j++)
					normalizer += beta[K * j + i];

				for (long j=0; j<V; j++)
					beta[K * j + i] /= normalizer;
				}
				"""

function updateBeta!(model::gpuCTM)
	model.queue(model.betakern, (model.K, model.V), nothing, model.K, model.newbetabuf, model.betabuf)
	model.queue(model.betanormkern, model.K, nothing, model.K, model.V, model.betabuf)
end

const CTM_NEWBETA_cpp =
"""
kernel void
updateNewbeta(long K,
			const global long *Jpsums,
			const global long *counts,
			const global long *words,
			const global float *phi,
			global float *newbeta)

			{
			long i = get_global_id(0);
			long j = get_global_id(1);

			float acc = 0.0f;

			for (long w=Jpsums[j]; w<Jpsums[j+1]; w++)
				acc += counts[words[w]] * phi[K * words[w] + i];

			newbeta[K * j + i] += acc;
			}
			"""

function updateNewbeta!(model::gpuCTM)
	model.queue(model.newbetakern, (model.K, model.V), nothing, model.K, model.Jpsumsbuf, model.countsbuf, model.wordsbuf, model.phibuf, model.newbetabuf)
end

const CTM_LAMBDA_cpp =
"""
$(RREF_cpp)
$(NORM2_cpp)

kernel void
updateLambda(long niter,
				float ntol,
				long K,
				long F,
				global float *A,
				global float *lambdaGrad,
				global float *lambdaInvHess,
				const global long *C,
				const global long *Npsums,
				const global long *counts,
				const global float *mu,
				const global float *sigma,
				const global float *invsigma,
				const global float *vsq,
				const global float *lzeta,
				const global float *phi,
				global float *lambda)

				{
				long d = get_global_id(0);

				long D = K * K * d;

				for (long _=0; _<niter; _++)
				{
					for (long i=0; i<K; i++)
					{
						float acc = 0.0f;

						for (long l=0; l<K; l++)
						{
							acc += invsigma[K * l + i] * (mu[l] - lambda[K * (F + d) + l]);
							A[D + K * l + i] = -C[d] * sigma[K * l + i] * exp(lambda[K * (F + d) + l] + 0.5f * vsq[K * (F + d) + l] - lzeta[F + d]);
						}

						for (long n=Npsums[d]; n<Npsums[d+1]; n++)
							acc += phi[K * n + i] * counts[n];

						lambdaGrad[K * d + i] = acc - C[d] * exp(lambda[K * (F + d) + i] + 0.5f * vsq[K * (F + d) + i] - lzeta[F + d]);
						A[D + K * i + i] -= 1.0f;
					}

					for (long l=0; l<K; l++)
						for (long i=0; i<K; i++)
							lambdaInvHess[D + K * l + i] = sigma[K * l + i];

					rref(K, D, A, lambdaInvHess);

					for (long l=0; l<K; l++)
						for (long i=0; i<K; i++)
							lambda[K * (F + d) + i] -= lambdaInvHess[D + K * l + i] * lambdaGrad[K * d + l];

					float lgnorm = norm2(K, d, lambdaGrad);
					if (lgnorm < ntol)
						break;
				}
				}
				"""

function updateLambda!(model::gpuCTM, b::Int, niter::Int, ntol::Float32)
	batch = model.batches[b]
	model.queue(model.lambdakern, length(batch), nothing, niter, ntol, model.K, batch[1] - 1, model.newtontempbuf, model.newtongradbuf, model.newtoninvhessbuf, model.Cbuf, model.Npsumsbuf, model.countsbuf, model.mubuf, model.sigmabuf, model.invsigmabuf, model.vsqbuf, model.lzetabuf, model.phibuf, model.lambdabuf)
end

const CTM_VSQ_cpp =
"""
$(NORM2_cpp)

kernel void
updateVsq(long niter,
			float ntol,
			long K,
			long F,
			global float *p,
			global float *vsqGrad,
			const global long *C,
			const global float *invsigma,
			const global float *lambda,
			const global float *lzeta,
			global float *vsq)

			{
			long d = get_global_id(0);

			float vsqInvHess;

			for (long _=0; _<niter; _++)
			{
				float rho = 1.0f;

				for (long i=0; i<K; i++)
				{
					vsqGrad[K * d + i] = -0.5f * (invsigma[K * i + i] + C[d] * exp(lambda[K * (F + d) + i] + 0.5f * vsq[K * (F + d) + i] - lzeta[F + d]) - 1 / vsq[K * (F + d) + i]);
					vsqInvHess = -1 / (0.25f * C[d] * exp(lambda[K * (F + d) + i] + 0.5f * vsq[K * (F + d) + i] - lzeta[F + d]) + 0.5f / (vsq[K * (F + d) + i] * vsq[K * (F + d) + i]));

					p[K * d + i] = vsqGrad[K * d + i] * vsqInvHess;
					while (vsq[K * (F + d) + i] - rho * p[K * d + i] <= 0)
						rho *= 0.5f;
				}

				for (long i=0; i<K; i++)
					vsq[K * (F + d) + i] -= rho * p[K * d + i];

				float vgnorm = norm2(K, d, vsqGrad);
				if (vgnorm < ntol)
					break;
			}

			for (long i=0; i<K; i++)
				vsq[K * (F + d) + i] += $(EPSILON32);
			}
			"""

function updateVsq!(model::gpuCTM, b::Int, niter::Int, ntol::Float32)
	batch = model.batches[b]
	model.queue(model.vsqkern, length(batch), nothing, niter, ntol, model.K, batch[1] - 1, model.newtontempbuf, model.newtongradbuf, model.Cbuf, model.invsigmabuf, model.lambdabuf, model.lzetabuf, model.vsqbuf)
end

const CTM_LZETA_cpp =
"""
kernel void
updateLzeta(long K,
			long F,
			const global float *lambda,
			const global float *vsq,
			global float *lzeta)

			{
			long d = get_global_id(0);

			float maxval = 0.0f;

			for (long i=0; i<K; i++)
			{
				float x = lambda[K * (F + d) + i] + 0.5f * vsq[K * (F + d) + i];
				if (x > maxval)
					maxval = x;
			}

			float acc = 0.0f;

			for (long i=0; i<K; i++)
				acc += exp(lambda[K * (F + d) + i] + 0.5f * vsq[K * (F + d) + i] - maxval);

			lzeta[F + d] = maxval + log(acc);
			}
			"""

function updateLzeta!(model::gpuCTM, b::Int)
	batch = model.batches[b]
	model.queue(model.lzetakern, length(batch), nothing, model.K, batch[1] - 1, model.lambdabuf, model.vsqbuf, model.lzetabuf)
end

const CTM_PHI_cpp =
"""
kernel void
updatePhi(long K,
			long F,
			const global long *Npsums,
			const global long *terms,
			const global float *beta,
			const global float *lambda,
			global float *phi)

			{
			long i = get_global_id(0);
			long d = get_global_id(1);

			for (long n=Npsums[d]; n<Npsums[d+1]; n++)
				phi[K * n + i] = log(beta[K * terms[n] + i]) + lambda[K * (F + d) + i];
			}
			"""

const CTM_PHI_NORM_cpp =
"""
kernel void
normalizePhi(long K,
				global float *phi)

				{
				long dn = get_global_id(0);

				float maxval = 0.0f;

				for (long i=0; i<K; i++)
					if (phi[K * dn + i] > maxval)
						maxval = phi[K * dn + i];

				float normalizer = 0.0f;

				for (long i=0; i<K; i++)
				{
					phi[K * dn + i] -= maxval;
					normalizer += exp(phi[K * dn + i]);
				}

				for (long i=0; i<K; i++)
					phi[K * dn + i] = exp(phi[K * dn + i]) / normalizer;
				}
				"""

function updatePhi!(model::gpuCTM, b::Int)
	batch = model.batches[b]
	model.queue(model.phikern, (model.K, length(batch)), nothing, model.K, batch[1] - 1, model.Npsumsbuf, model.termsbuf, model.betabuf, model.lambdabuf, model.phibuf)
	model.queue(model.phinormkern, sum(model.N[batch]), nothing, model.K, model.phibuf)
end

function train!(model::gpuCTM; iter::Integer=150, tol::Real=1.0, niter::Integer=1000, ntol::Real=1/model.K^2, viter::Integer=10, vtol::Real=1/model.K^2, chkelbo::Integer=1)
	@assert all(.!isnegative.([tol, ntol, vtol]))
	@assert all(ispositive.([iter, niter, viter, chkelbo]))
	niter, ntol = Int(niter), Float32(ntol)
	lowVRAM = model.B > 1

	for k in 1:iter
		chk = (k % chkelbo == 0)
		for b in 1:model.B
			for _ in 1:viter
				oldlambda = @host model.lambdabuf
				updateLzeta!(model, b)
				updatePhi!(model, b)
				updateLambda!(model, b, niter, ntol)
				updateVsq!(model, b, niter, ntol)
				lambda = @host model.lambdabuf
				if sum([norm(diff) for diff in oldlambda - lambda]) < length(model.batches[b]) * vtol
					break
				end
			end
			updateNewbeta!(model)
			if chk
				updateHost!(model, b)
				updateNewELBO!(model, b)
			end
			lowVRAM && updateBuf!(model, b)
		end
		updateMu!(model)
		updateSigma!(model)
		updateBeta!(model)
		if checkELBO!(model, k, chk, tol)
			break
		end
	end
	updateHost!(model, 1)
	model.topics = [reverse(sortperm(vec(model.beta[i,:]))) for i in 1:model.K]
	nothing
end


text:
"""
    couple_models(aero::LiftingLine, dyn::RigidBody, flap::LiftingLineFlaps, ctrl::Trim)

Create a coupled model using a lifting line aerodynamic model, a rigid
body dynamics model, a lifting line control surface model, and a controller which
maintains trimmed operating conditions.  This model introduces additional
parameters corresponding to the length, position, and orientation of each lifting
line element ``L, p_e, e_1, e_2, e_3``, followed by the freestream air density
``rho_\\infty``, rigid body inertial properties ``m, Ixx, Iyy, Izz, Ixz, Ixy,
Iyz``, additional forces/moments applied on the body ``F_x, F_y, F_z, M_x, M_y,
M_z``, and the control inputs which are not defined by the controller.

**NOTE: When using this model, the local frame for each lifting line element should be
oriented with the x-axis in the chordwise direction, the y-axis in the spanwise
direction (out the right wing), and the z-axis in the airfoil normal direction**
"""
function couple_models(aero::LiftingLine, dyn::RigidBody, flap::LiftingLineFlaps,
    ctrl::Trim)

    return (aero, dyn, flap, ctrl)
end

# --- traits --- #

function number_of_additional_parameters(::LiftingLine{NA,TA}, ::RigidBody,
    ::LiftingLineFlaps{NF,NG,TF,TG}, ::Trim{NC}) where {NA,NF,NG,NC,TA,TF,TG}

    return 13*NA + 14 + NG - NC
end

function coupling_inplaceness(::Type{<:LiftingLine}, ::Type{<:RigidBody},
    ::Type{<:LiftingLineFlaps}, ::Type{<:Trim})

    return InPlace()
end

function coupling_rate_jacobian_type(::Type{LiftingLine{NA,TA}}, ::Type{<:RigidBody},
    ::Type{LiftingLineFlaps{NF,NG,TF,TG}}, ::Type{<:Trim}) where {NA,NF,NG,TA,TF,TG}

    aero_model_types = TA.parameters
    flap_model_types = TF.parameters
    if all(islinear.(coupling_rate_jacobian_type.(aero_model_types, Ref(LiftingLineSection),
        flap_model_types, Ref(LiftingLineSectionControl))))
        return Linear()
    else
        return Nonlinear()
    end
end

function coupling_state_jacobian_type(::Type{LiftingLine{NA,TA}}, ::Type{<:RigidBody},
    ::Type{LiftingLineFlaps{NF,NG,TF,TG}}, ::Type{<:Trim}) where {NA,NF,NG,TA,TF,TG}

    aero_model_types = TA.parameters
    flap_model_types = TF.parameters
    if all(islinear.(coupling_state_jacobian_type.(aero_model_types, Ref(LiftingLineSection),
        flap_model_types, Ref(LiftingLineSectionControl))))
        return Linear()
    else
        return Nonlinear()
    end
end

function coupling_parameter_jacobian_type(::Type{<:LiftingLine}, ::Type{<:RigidBody},
    ::Type{<:LiftingLineFlaps}, ::Type{<:Trim})

    return Nonlinear()
end

function coupling_time_gradient_type(::Type{LiftingLine{NA,TA}}, ::Type{<:RigidBody},
    ::Type{LiftingLineFlaps{NF,NG,TF,TG}}, ::Type{<:Trim}) where {NA,NF,NG,TA,TF,TG}

    aero_model_types = TA.parameters
    flap_model_types = TF.parameters
    if all(isempty.(coupling_time_gradient_type.(aero_model_types, Ref(LiftingLineSection),
        flap_model_types, Ref(LiftingLineSectionControl))))
        return Empty()
    elseif all(iszero.(coupling_time_gradient_type.(aero_model_types, Ref(LiftingLineSection),
        flap_model_types, Ref(LiftingLineSectionControl))))
        return Zeros()
    elseif all(isinvariant.(coupling_time_gradient_type.(aero_model_types, Ref(LiftingLineSection),
        flap_model_types, Ref(LiftingLineSectionControl))))
        return Invariant()
    elseif all(isconstant.(coupling_time_gradient_type.(aero_model_types, Ref(LiftingLineSection),
        flap_model_types, Ref(LiftingLineSectionControl))))
        return Constant()
    elseif all(islinear.(coupling_time_gradient_type.(aero_model_types, Ref(LiftingLineSection),
        flap_model_types, Ref(LiftingLineSectionControl))))
        return Linear()
    else
        return Nonlinear()
    end
end

# --- methods --- #

function get_coupling_inputs!(y, aero::LiftingLine{NA,TA}, dyn::RigidBody,
    flap::LiftingLineFlaps{NF,NG,TF,TG}, ctrl::Trim{NC}, dx, x, p, t) where {NA,NF,NG,NC,TA,TF,TG}

    # extract number of state variables, inputs, and parameters
    nua = number_of_states(aero) # number of aerodynamic states
    nud = number_of_states(dyn) # number of rigid body states
    nuf = number_of_states(flap) # number of control surface states
    nuc = number_of_states(ctrl) # number of control surface states
    nya = number_of_inputs(aero) # number of aerodynamic inputs
    nyd = number_of_inputs(dyn) # number of rigid body inputs
    nyf = number_of_inputs(flap) # number of control surface inputs
    nyc = number_of_inputs(ctrl) # number of control surface inputs
    npa = number_of_parameters(aero) # number of aerodynamic parameters
    npd = number_of_parameters(dyn) # number of rigid body parameters
    npf = number_of_parameters(flap) # number of control surface parameters
    npc = number_of_parameters(ctrl) # number of control surface parameters
    npadd = number_of_additional_parameters(aero, dyn, flap, ctrl) # number of additional parameters

    # get indices for state variables, inputs, and parameters
    iua = 1:nua # indices of aerodynamic states
    iya = 1:nya # indices of aerodynamic inputs
    ipa = 1:npa # indices of aerodynamic parameters
    iud = nua + 1 : nua + nud # indices of rigid body states
    iyd = nya + 1 : nya + nyd # indices of rigid body inputs
    ipd = npa + 1 : npa + npd # indices of rigid body parameters
    iuf = nua + nud + 1 : nua + nud + nuf # indices of control surface states
    iyf = nya + nyd + 1 : nya + nyd + nyf # indices of control surface inputs
    ipf = npa + npd + 1 : npa + npd + npf # indices of control surface parameters
    iuc = nua + nud + nuf + 1 : nua + nud + nuf + nuc # indices of controller states
    iyc = nya + nyd + nyf + 1 : nya + nyd + nyf + nyc # indices of controller inputs
    ipc = npa + npd + npf + 1 : npa + npd + npf + npc # indices of controller parameters
    ipadd = npa + npd + npf + npc + 1 : npa + npd + npf + npc + npadd # indices of additional parameters
    iuas = state_indices(aero.models) # indices of aerodynamic states for each section
    iyas = input_indices(aero.models) # indices of aerodynamic inputs for each section
    ipas = parameter_indices(aero.models) # indices of aerodynamic parameters for each section
    iufs = state_indices(flap.models) # indices of control surface states for each section
    iyfs = input_indices(flap.models) # indices of control surface inputs for each section
    ipfs = parameter_indices(flap.models) # indices of control surface parameters for each section

    # separate rates, states, inputs, and parameters
    dua = view(dx, iua) # aerodynamic rates
    ua = view(x, iua) # aerodynamic state variables
    ya = view(y, iya) # aerodynamic inputs
    pa = view(p, ipa) # aerodynamic parameters
    dud = view(dx, iud) # rigid body rates
    ud = view(x, iud) # rigid body state variables
    yd = view(y, iyd) # rigid body inputs
    pd = view(p, ipd) # rigid body parameters
    duf = view(dx, iuf) # control surface rates
    uf = view(x, iuf) # control surface states
    yf = view(y, iyf) # control surface inputs
    pf = view(p, ipf) # control surface parameters
    duc = view(dx, iuc) # controller rates
    uc = view(x, iuc) # controller states
    yc = view(y, iyc) # controller inputs
    pc = view(p, ipc) # controller parameters
    padd = view(p, ipadd) # additional parameters
    duas = view.(Ref(dua), iuas) # aerodynamic rates for each section
    uas = view.(Ref(ua), iuas) # aerodynamic states for each section
    yas = view.(Ref(ya), iyas) # aerodynamic inputs for each section
    pas = view.(Ref(pa), ipas) # aerodynamic parameters for each section
    dufs = view.(Ref(duf), iufs) # control surface rates for each section
    ufs = view.(Ref(uf), iufs) # control surface state variables for each section
    yfs = view.(Ref(yf), iyfs) # control surface inputs for each section
    pfs = view.(Ref(pf), ipfs) # control surface parameters for each section

    # rigid body rates and states
    dxr, dyr, dzr, dϕr, dθr, dψr, dur, dvr, dwr, dpr, dqr, drr = dud
    xr, yr, zr, ϕr, θr, ψr, ur, vr, wr, pr, qr, rr = ud

    # controller state variables
    δc = uc

    # global parameters
    ρ = padd[13*NA+1]
    m = padd[13*NA+2]
    Ixx = padd[13*NA+3]
    Iyy = padd[13*NA+4]
    Izz = padd[13*NA+5]
    Ixz = padd[13*NA+6]
    Ixy = padd[13*NA+7]
    Iyz = padd[13*NA+8]
    Fbx = padd[13*NA+9]
    Fby = padd[13*NA+10]
    Fbz = padd[13*NA+11]
    Mbx = padd[13*NA+12]
    Mby = padd[13*NA+13]
    Mbz = padd[13*NA+14]
    δp = padd[SVector{NG-NC}(13*NA + 14 + 1 : 13*NA + 14 + NG - NC)]

    # commanded control deflections
    T = promote_type(eltype(x), eltype(p))
    δ = zeros(SVector{NG,T})
    iδu = iδp = 1
    for i = 1:NG
        if i in ctrl.state_indices
            δ = setindex(δ, δc[iδu], i)
            iδu += 1
        else
            δ = setindex(δ, δp[iδp], i)
            iδp += 1
        end
    end
    dδ = zero(δ)

    # body linear and angular velocity
    V = SVector(ur, vr, wr)
    Ω = SVector(pr, qr, rr)

    # body linear and angular acceleration
    dV = SVector(dur, dvr, dwr)
    dΩ = SVector(dpr, dqr, drr)

    # section forces and moments
    Fa = @SVector zeros(3)
    Ma = @SVector zeros(3)
    for i = 1:NA
        # models for this section
        section_aero = aero.models[i]
        section_stru = LiftingLineSection()
        section_flap = flap.models[i]
        section_ctrl = LiftingLineSectionControl()
        section_models = (section_aero, section_stru, section_flap, section_ctrl)

        # model dimensions for this section
        Nuai = number_of_states(section_aero)
        Nyai = number_of_inputs(section_aero)
        Npai = number_of_parameters(section_aero)
        Nusi = number_of_states(section_stru)
        Nysi = number_of_inputs(section_stru)
        Npsi = number_of_parameters(section_stru)
        Nufi = number_of_states(section_flap)
        Nyfi = number_of_inputs(section_flap)
        Npfi = number_of_parameters(section_flap)
        Nuci = number_of_states(section_ctrl)
        Nyci = number_of_inputs(section_ctrl)
        Npci = number_of_parameters(section_ctrl)

        # local section properties
        offset = 13*(i-1)
        ΔL = padd[offset+1] # section spanwise length
        pe = SVector(padd[offset+2], padd[offset+3], padd[offset+4]) # local frame position
        e1 = SVector(padd[offset+5], padd[offset+6], padd[offset+7]) # local frame x-axis
        e2 = SVector(padd[offset+8], padd[offset+9], padd[offset+10]) # local frame y-axis
        e3 = SVector(padd[offset+11], padd[offset+12], padd[offset+13]) # local frame z-axis
        Cab = [e1 e2 e3] # local to body frame transformation matrix
        Cba = Cab' # body to local frame transformation matrix
        vi = -Cba*(V + cross(Ω, pe)) # local linear freestream velocity
        ωi = Cba*Ω # local angular freestream velocity
        dvi =-Cba*(dV + cross(dΩ, pe)) # local linear freestream acceleration
        dωi = Cba*dΩ # local angular freestream acceleration

        # section rates
        dxai = SVector{Nuai}(duas[i]) # aerodynamic state variables
        dxsi = vcat(dvi, dωi) # rigid body state variables
        dxfi = SVector{Nufi}(dufs[i]) # control surface state variables
        dxci = SMatrix{1,NG}(getindex.(flap.gains, i)) * dδ # controller state variables
        dxi = vcat(dxai, dxsi, dxfi, dxci)

        # section states
        xai = SVector{Nuai}(uas[i]) # aerodynamic state variables
        xsi = vcat(vi, ωi) # rigid body state variables
        xfi = SVector{Nufi}(ufs[i]) # control surface state variables
        xci = SMatrix{1,NG}(getindex.(flap.gains, i)) * δ # controller state variables
        xi = vcat(xai, xsi, xfi, xci)

        # section parameters
        pai = SVector{Npai}(pas[i]) # aerodynamic parameters
        psi = SVector(ρ) # rigid body parameters
        pfi = SVector{Npfi}(pfs[i]) # control surface parameters
        pci = SVector{0,Float64}() # controller parameters
        pi = vcat(pai, psi, pfi, pci)

        # section inputs
        yi = get_coupling_inputs(section_models, dxi, xi, pi, t)

        # separate inputs
        yai = view(yi, 1:Nyai)
        ysi = view(yi, Nyai + 1 : Nyai + Nysi)
        yfi = view(yi, Nyai + Nysi + 1 : Nyai + Nysi + Nyfi)
        yci = view(yi, Nyai + Nysi + Nyfi + 1 : Nyai + Nysi + Nyfi + Nyci)

        # save local inputs
        y[iya[iyas[i]]] .= yai # aerodynamic inputs
        y[iyf[iyfs[i]]] .= yfi # control surface inputs

        # add to global inputs

        # section aerodynamic loads (in body frame)
        Fi = Cab*ΔL*SVector(ysi[1], ysi[2], ysi[3])
        Mi = Cab*ΔL*SVector(ysi[4], ysi[5], ysi[6])

        # add loads to total aerodynamic loads
        Fa += Fi
        Ma += cross(pe, Fi) + Mi
    end

    # save rigid body inputs
    y[iyd[1]] = m
    y[iyd[2]] = Ixx
    y[iyd[3]] = Iyy
    y[iyd[4]] = Izz
    y[iyd[5]] = Ixz
    y[iyd[6]] = Ixy
    y[iyd[7]] = Iyz
    y[iyd[8]] = Fa[1] + Fbx
    y[iyd[9]] = Fa[2] + Fby
    y[iyd[10]] = Fa[3] + Fbz
    y[iyd[11]] = Ma[1] + Mbx
    y[iyd[12]] = Ma[2] + Mby
    y[iyd[13]] = Ma[3] + Mbz

    # save controller inputs
    loads = SVector(
        Fa[1] + Fbx,
        Fa[2] + Fby,
        Fa[3] + Fbz,
        Ma[1] + Mbx,
        Ma[2] + Mby,
        Ma[3] + Mbz)

    for i = 1:NC
        idx = ctrl.force_indices[i]
        y[iyc[i]] = loads[idx]
    end

    return y
end

# --- Performance Overloads --- #

# TODO: Rate, state, input, and parameter jacobians

# --- Convenience Methods --- #

function set_additional_parameters!(padd, model::LiftingLine{NA,TA},
    dyn::RigidBody, flap::LiftingLineFlaps{NF,NG,TF,TG}, ctrl::Trim{NC};
    section_parameters, rho, m, Ixx, Iyy, Izz, Ixz, Ixy, Iyz, Fbx, Fby, Fbz,
    Mbx, Mby, Mbz, delta) where {NA,NF,NG,NC,TA,TF,TG}

    for i = 1:NA
        padd[13*(i-1)+1] = section_parameters[i].L
        padd[13*(i-1)+2 : 13*(i-1)+4] = section_parameters[i].pe
        padd[13*(i-1)+5 : 13*(i-1)+7] = section_parameters[i].e1
        padd[13*(i-1)+8 : 13*(i-1)+10] = section_parameters[i].e2
        padd[13*(i-1)+11 : 13*(i-1)+13] = section_parameters[i].e3
    end

    padd[13*NA+1] = rho
    padd[13*NA+2] = m
    padd[13*NA+3] = Ixx
    padd[13*NA+4] = Iyy
    padd[13*NA+5] = Izz
    padd[13*NA+6] = Ixz
    padd[13*NA+7] = Ixy
    padd[13*NA+8] = Iyz
    padd[13*NA+9] = Fbx
    padd[13*NA+10] = Fby
    padd[13*NA+11] = Fbz
    padd[13*NA+12] = Mbx
    padd[13*NA+13] = Mby
    padd[13*NA+14] = Mbz
    padd[13*NA+14+1 : 13*NA+14+NG-NC] = delta

    return padd
end

function separate_additional_parameters(model::LiftingLine{NA,TA},
    dyn::RigidBody, flap::LiftingLineFlaps{NF,NG,TF,TG}, trim::Trim{NC},
    padd) where {NA,NF,NG,NC,TA,TF,TG,TI}

    # section parameters
    return (
        section_parameters = ntuple(i -> (
            L = padd[13*(i-1)+1], # element length
            pe = view(padd, 13*(i-1)+2 : 13*(i-1)+4), # local frame position
            e1 = view(padd, 13*(i-1)+5 : 13*(i-1)+7), # local frame x-axis
            e2 = view(padd, 13*(i-1)+8 : 13*(i-1)+10), # local frame y-axis
            e3 = view(padd, 13*(i-1)+11 : 13*(i-1)+13), # local frame z-axis
            ), NA),
        rho = padd[13*NA+1],
        m = padd[13*NA+2],
        Ixx = padd[13*NA+3],
        Iyy = padd[13*NA+4],
        Izz = padd[13*NA+5],
        Ixz = padd[13*NA+6],
        Ixy = padd[13*NA+7],
        Iyz = padd[13*NA+8],
        Fbx = padd[13*NA+9],
        Fby = padd[13*NA+10],
        Fbz = padd[13*NA+11],
        Mbx = padd[13*NA+12],
        Mby = padd[13*NA+13],
        Mbz = padd[13*NA+14],
        delta = view(padd, 13*NA+14+1 : 13*NA+14+NG-NC),
        )
end


text:
using DistributionsAD

using ChainRulesCore
using ChainRulesTestUtils
using Combinatorics
using Distributions
using FiniteDifferences
using PDMats

using Random, LinearAlgebra, Test

using Distributions: meanlogdet
using DistributionsAD: TuringUniform, TuringMvNormal, TuringMvLogNormal,
                       TuringPoissonBinomial, TuringDirichlet
using StatsBase: entropy
using StatsFuns: StatsFuns, logsumexp, logistic

import NNlib

Random.seed!(1) # Set seed that all testsets should reset to.

const FDM = FiniteDifferences
const GROUP = get(ENV, "GROUP", "All")

# Figure out which AD backend to test
const AD = get(ENV, "AD", "All")
if AD == "All" || AD == "ForwardDiff"
    @eval using ForwardDiff
end
if AD == "All" || AD == "Zygote"
    @eval using Zygote
end
if AD == "All" || AD == "ReverseDiff"
    @eval using ReverseDiff
end
if AD == "All" || AD == "Tracker"
    @eval using Tracker
end

if GROUP == "All" || GROUP == "Others"
    include("others.jl")
end

if GROUP == "All" || GROUP == "AD"
    # Create positive definite matrix
    to_posdef(A::AbstractMatrix) = A * A' + I
    to_posdef_diagonal(a::AbstractVector) = Diagonal(a.^2 .+ 1)

    # Create vectors in probability simplex.
    to_simplex(x::AbstractArray) = NNlib.softmax(x; dims=1)
    to_simplex(x::AbstractArray{<:AbstractArray}) = to_simplex.(x)

    if AD == "All" || AD == "ReverseDiff"
        @eval begin
            # Define adjoint for ReverseDiff
            function to_simplex(x::AbstractArray{<:ReverseDiff.TrackedReal})
                return ReverseDiff.track(to_simplex, x)
            end
            ReverseDiff.@grad function to_simplex(x)
                _x = ReverseDiff.value(x)
                y = to_simplex(_x)
                function pullback(∇)
                    return (NNlib.∇softmax(∇, _x, y; dims=1),)
                end
                return y, pullback
            end
        end
    end

    if AD == "All" || AD == "Tracker"
        @eval begin
            # Define adjoints for Tracker
            to_posdef(A::Tracker.TrackedMatrix) = Tracker.track(to_posdef, A)
            Tracker.@grad function to_posdef(A::Tracker.TrackedMatrix)
                data_A = Tracker.data(A)
                S = data_A * data_A' + I
                function pullback(∇)
                    return ((∇ + ∇') * data_A,)
                end
                return S, pullback
            end
        end
    end

    include("ad/utils.jl")
    include("ad/chainrules.jl")
    include("ad/distributions.jl")
end


text:
using JSON

export MSA, ColorScheme, Conf, Columns, Vis, Zoomer, plot

struct MSA
    seq::AbstractString
    options::Dict
    div::String
    div_seqs::String

    function MSA(seq::AbstractString, options...)
        if isempty(options)
            new(seq, Dict(), Base.Random.randstring(), Base.Random.randstring())
        else
            if isavalidtype(options)
                new(seq, merge_options!(options), Base.Random.randstring(), Base.Random.randstring())
            end
        end
    end
end

struct Columns <: Options
    fields::Dict{AbstractString, Any}
    function Columns(;kwargs...)
        opts = merge_options!("columns", kwargs)
        new(opts)
    end
end

struct ColorScheme <: Options
    fields::Dict{AbstractString, Any}
    function ColorScheme(;kwargs...)
        opts = merge_options!("colorscheme", kwargs)
        new(opts)
    end
end

struct Conf <: Options
    fields::Dict{AbstractString, Any}
    function Conf(;kwargs...)
        opts = merge_options!("conf", kwargs)
        new(opts)
    end
end

struct Vis <: Options
    fields::Dict{AbstractString, Any}
    function Vis(;kwargs...)
        opts = merge_options!("vis", kwargs)
        new(opts)
    end
end

struct Zoomer <: Options
    fields::Dict{AbstractString, Any}
    function Zoomer(;kwargs...)
        opts = merge_options!("zoomer", kwargs)
        new(opts)
    end
end

function plot(msa::MSA)
    plot_msa(msa)
end

function plot_msa(msa::MSA)
    options = """
    {
       el: $(msa.div),
       seqs: seqs,
       bootstrapMenu: true,
       $(JSON.json(msa.options)[2:end-1])
    }
    """

    txt_html = """

$(msa.seq)

    """
    display("text/html", txt_html)
end


text:
"""
    add_reactions(
        m::CoreModel,
        s::V1,
        b::V2,
        c::AbstractFloat,
        xl::AbstractFloat,
        xu::AbstractFloat;
        check_consistency = false,
    ) where {V1<:VecType,V2<:VecType}

Add reaction(s) to a `CoreModel` model `m`.

"""
function add_reactions(
    m::CoreModel,
    s::V1,
    b::V2,
    c::AbstractFloat,
    xl::AbstractFloat,
    xu::AbstractFloat;
    check_consistency = false,
) where {V1<:VecType,V2<:VecType}
    return add_reactions(
        m,
        sparse(reshape(s, (length(s), 1))),
        sparse(b),
        sparse([c]),
        sparse([xl]),
        sparse([xu]),
        check_consistency = check_consistency,
    )
end

"""
    add_reactions(
        m::CoreModel,
        s::V1,
        b::V2,
        c::AbstractFloat,
        xl::AbstractFloat,
        xu::AbstractFloat,
        rxn::String,
        mets::K;
        check_consistency = false,
    ) where {V1<:VecType,V2<:VecType,K<:StringVecType}

"""
function add_reactions(
    m::CoreModel,
    s::V1,
    b::V2,
    c::AbstractFloat,
    xl::AbstractFloat,
    xu::AbstractFloat,
    rxn::String,
    mets::K;
    check_consistency = false,
) where {V1<:VecType,V2<:VecType,K<:StringVecType}
    return add_reactions(
        m,
        sparse(reshape(s, (length(s), 1))),
        sparse(b),
        sparse([c]),
        sparse([xl]),
        sparse([xu]),
        [rxn],
        mets,
        check_consistency = check_consistency,
    )
end

"""
    add_reactions(
        m::CoreModel,
        Sp::M,
        b::V,
        c::V,
        xl::V,
        xu::V;
        check_consistency = false,
    ) where {M<:MatType,V<:VecType}

"""
function add_reactions(
    m::CoreModel,
    Sp::M,
    b::V,
    c::V,
    xl::V,
    xu::V;
    check_consistency = false,
) where {M<:MatType,V<:VecType}
    rxns = ["r$x" for x = length(m.rxns)+1:length(m.rxns)+length(xu)]
    mets = ["m$x" for x = length(m.mets)+1:length(m.mets)+size(Sp)[1]]
    return add_reactions(
        m,
        Sp,
        b,
        c,
        xl,
        xu,
        rxns,
        mets,
        check_consistency = check_consistency,
    )
end

"""
    add_reactions(m1::CoreModel, m2::CoreModel; check_consistency = false)

Add all reactions from `m2` to `m1`.

"""
function add_reactions(m1::CoreModel, m2::CoreModel; check_consistency = false)
    return add_reactions(
        m1,
        m2.S,
        m2.b,
        m2.c,
        m2.xl,
        m2.xu,
        m2.rxns,
        m2.mets,
        check_consistency = check_consistency,
    )
end

"""
    add_reactions(
        m::CoreModel,
        Sp::M,
        b::V,
        c::V,
        xl::V,
        xu::V,
        rxns::K,
        mets::K;
        check_consistency = false,
    ) where {M<:MatType,V<:VecType,K<:StringVecType}

"""
function add_reactions(
    m::CoreModel,
    Sp::M,
    b::V,
    c::V,
    xl::V,
    xu::V,
    rxns::K,
    mets::K;
    check_consistency = false,
) where {M<:MatType,V<:VecType,K<:StringVecType}

    Sp = sparse(Sp)
    b = sparse(b)
    c = sparse(c)
    xl = sparse(xl)
    xu = sparse(xu)

    all([length(b), length(mets)] .== size(Sp, 1)) ||
        throw(DimensionMismatch("inconsistent number of metabolites"))
    all(length.([c, xl, xu, rxns]) .== size(Sp, 2)) ||
        throw(DimensionMismatch("inconsistent number of reactions"))

    new_reactions = findall(Bool[!(rxn in m.rxns) for rxn in rxns])
    new_metabolites = findall(Bool[!(met in m.mets) for met in mets])

    if check_consistency
        (newReactions1, newMetabolites1) = verify_consistency(
            m,
            Sp,
            b,
            c,
            xl,
            xu,
            rxns,
            mets,
            new_reactions,
            new_metabolites,
        )
    end

    new_mets = vcat(m.mets, mets[new_metabolites])

    zero_block = spzeros(length(new_metabolites), n_reactions(m))
    ext_s = vcat(sparse(m.S), zero_block)

    mapping = [findfirst(isequal(met), new_mets) for met in mets]
    (I, J, elements) = findnz(sparse(Sp[:, new_reactions]))
    ext_sp = spzeros(length(new_mets), length(new_reactions))
    for (k, i) in enumerate(I)
        new_i = mapping[i]
        ext_sp[new_i, J[k]] = elements[k]
    end

    new_s = hcat(ext_s, ext_sp)
    newb = vcat(m.b, b[new_metabolites])
    #new_c = hcat(m.C, spzeros(size(m.C, 1), length(new_reactions)))
    newc = vcat(m.c, c[new_reactions])
    newxl = vcat(m.xl, xl[new_reactions])
    newxu = vcat(m.xu, xu[new_reactions])
    new_rxns = vcat(m.rxns, rxns[new_reactions])
    #new_lp = CoreModel(new_s, newb, new_c, m.cl, m.cu, newc, newxl, newxu, new_rxns, new_mets)
    new_lp = CoreModel(new_s, newb, newc, newxl, newxu, new_rxns, new_mets)

    if check_consistency
        return (new_lp, new_reactions, new_metabolites)
    else
        return new_lp
    end
end

"""
    verify_consistency(
        m::CoreModel,
        Sp::M,
        b::V,
        c::V,
        xl::V,
        xu::V,
        names::K,
        mets::K,
        new_reactions,
        new_metabolites,
    ) where {M<:MatType,V<:VecType,K<:StringVecType}

Check the consistency of given reactions with existing reactions in `m`.

TODO: work in progress, doesn't return consistency status.

"""
function verify_consistency(
    m::CoreModel,
    Sp::M,
    b::V,
    c::V,
    xl::V,
    xu::V,
    names::K,
    mets::K,
    new_reactions,
    new_metabolites,
) where {M<:MatType,V<:VecType,K<:StringVecType}

    if !isempty(new_reactions)
        statuses = Vector{ReactionStatus}(undef, length(names))
        for (i, name) in enumerate(names)
            rxn_index = findfirst(isequal(name), m.rxns)
            reaction = Sp[:, i]
            stoich_index = findfirst(Bool[reaction == m.S[:, j] for j = 1:size(m.S, 2)])
            if isnothing(rxn_index) & isnothing(stoich_index)
                statuses[i] = ReactionStatus(false, 0, "new")
            end

            if !isnothing(rxn_index) & isnothing(stoich_index)
                statuses[i] = ReactionStatus(true, 0, "same name")
            end

            if isnothing(rxn_index) & !isnothing(stoich_index)
                statuses[i] = ReactionStatus(true, 0, "same stoichiometry")
            end

            if !isnothing(rxn_index) & !isnothing(stoich_index)
                statuses[i] = ReactionStatus(true, 0, "same name, same stoichiometry")
            end
        end
    end

    return (new_reactions, new_metabolites)
end

"""
    remove_metabolites(model::CoreModel, metabolites)

Removes a set of `metabolites` from the `model` of type `CoreModel` and returns
a new `CoreModel` without those metabolites. Here, `metabolites` can be either a
string, a vector of strings, an index or a vector of indices. Also removes any
reactions that have no associated metabolites after the metabolites have been
removed.

# Example
```
model = load_model(CoreModel, "e_coli_core.json")

m1 = remove_metabolites(model, ["glc__D_e", "for_c"])
m2 = remove_metabolites(model, "glc__D_e")
m3 = remove_metabolites(model, indexin(["glc__D_e", "for_c"], metabolites(model)))
m4 = remove_metabolites(model, first(indexin(["glc__D_e"], metabolites(model))))
```
"""
function remove_metabolites(model::CoreModel, mets)
    mets_to_keep = filter(x -> x ∉ mets, 1:n_metabolites(model))
    temp_S = model.S[mets_to_keep, :]

    (I, rxns_to_keep, val) = findnz(temp_S)
    sort!(rxns_to_keep)
    unique!(rxns_to_keep)
    new_S = model.S[mets_to_keep, rxns_to_keep]
    new_b = model.b[mets_to_keep]
    new_c = model.c[rxns_to_keep]
    new_lbs = model.xl[rxns_to_keep]
    new_ubs = model.xu[rxns_to_keep]
    new_rxns = model.rxns[rxns_to_keep]
    new_mets = model.mets[mets_to_keep]

    return CoreModel(new_S, new_b, new_c, new_lbs, new_ubs, new_rxns, new_mets)
end

function remove_metabolites(model::CoreModel, met::Integer)
    return remove_metabolites(model, [met])
end

function remove_metabolites(model::CoreModel, met::String)
    return remove_metabolites(model, [met])
end

function remove_metabolites(model::CoreModel, mets::Vector{String})
    met_indices = filter(!isnothing, indexin(mets, metabolites(model)))
    if isempty(met_indices)
        return model
    else
        return remove_metabolites(model, met_indices)
    end
end

"""
Removes a set of reactions from a CoreModel.
Also removes the metabolites not involved in any reaction.
"""
function remove_reactions(m::CoreModel, rxns::Vector{Int})
    rxns_to_keep = filter(e -> e ∉ rxns, 1:n_reactions(m))
    temp_s = m.S[:, rxns_to_keep]

    (mets_to_keep, J, val) = findnz(temp_s)
    sort!(mets_to_keep)
    unique!(mets_to_keep)
    new_s = m.S[mets_to_keep, rxns_to_keep]
    newb = m.b[mets_to_keep]
    newc = m.c[rxns_to_keep]
    newxl = m.xl[rxns_to_keep]
    newxu = m.xu[rxns_to_keep]
    new_rxns = m.rxns[rxns_to_keep]
    new_mets = m.mets[mets_to_keep]
    new_model = CoreModel(new_s, newb, newc, newxl, newxu, new_rxns, new_mets)
    return new_model
end

"""
    remove_reactions(m::CoreModel, rxn::Integer)

"""
function remove_reactions(m::CoreModel, rxn::Integer)
    return remove_reactions(m, [rxn])
end

"""
    remove_reactions(m::CoreModel, rxn::String)

"""
function remove_reactions(m::CoreModel, rxn::String)
    return remove_reactions(m, [rxn])
end

"""
    remove_reactions(m::CoreModel, rxns::Vector{String})

"""
function remove_reactions(m::CoreModel, rxns::Vector{String})
    rxn_indices = [findfirst(isequal(name), m.rxns) for name in intersect(rxns, m.rxns)]
    if isempty(rxn_indices)
        return m
    else
        return remove_reactions(m, rxn_indices)
    end
end

"""
    change_bounds!(
        model::CoreModel,
        rxns::Vector{Int};
        xl::V = Float64[],
        xu::V = Float64[],
    ) where {V<:VecType}

Change the lower and/or upper bounds ('xl' and 'xu') for given reactions.

"""
function change_bounds!(
    model::CoreModel,
    rxns::Vector{Int};
    xl::V = Float64[],
    xu::V = Float64[],
) where {V<:VecType}
    found = [index ∈ 1:n_reactions(model) for index in rxns]
    length(rxns[found]) == length(unique(rxns[found])) ||
        error("`rxns` appears to contain duplicates")

    if !isempty(xl)
        length(rxns) == length(xl) ||
            throw(DimensionMismatch("`rxns` size doesn't match with `xl`"))
        model.xl[rxns[found]] = xl[found]
    end
    if !isempty(xu)
        length(rxns) == length(xu) ||
            throw(DimensionMismatch("`rxns` size doesn't match with `xu`"))
        model.xu[rxns[found]] = xu[found]
    end
end

"""
    change_bounds!(
        model::CoreModel,
        rxns::Vector{String};
        xl::V = Float64[],
        xu::V = Float64[],
    ) where {V<:VecType}

"""
function change_bounds!(
    model::CoreModel,
    rxns::Vector{String};
    xl::V = Float64[],
    xu::V = Float64[],
) where {V<:VecType}
    found = [name ∈ model.rxns for name in rxns]
    rxn_indices = zeros(Int, length(rxns))
    rxn_indices[found] = [findfirst(isequal(name), model.rxns) for name in rxns[found]]
    change_bounds!(model, rxn_indices, xl = xl, xu = xu)
end


text:
# Use baremodule to shave off a few KB from the serialized `.ji` file
baremodule ZITSOL_1_jll
using Base
using Base: UUID
import JLLWrappers

JLLWrappers.@generate_main_file_header("ZITSOL_1")
JLLWrappers.@generate_main_file("ZITSOL_1", UUID("dd94201c-7435-5bf3-a1a1-33dc1dbf1631"))
end  # module ZITSOL_1_jll


text:
function smacof(Δ::AbstractMatrix{T}; p = 2, wt = nothing, Xinit = nothing, anchors = nothing, ε = 1e-6, itmax = 100, verbose = false, return_hist::Bool = false) where T
    n = size(Δ, 1)

    # Check if Δ is a dissimilarity matrix.
    !issymmetric(Δ) && error("First argument is not a dissimilarity matrix.")

    # Check if anchors is correct dims
    if !isnothing(anchors)
        if isa(anchors, NamedTuple)
            if size(anchors.pos) != (n, p)
                error("`anchors` has incorrect dimensions. $(size(anchors.pos)) != ($n, $p)")
            end
        else
            if size(anchors) != (n, p)
                error("`anchors` has incorrect dimensions. $(size(anchors)) != ($n, $p)")
            end
        end
    end

    # Normalize weights
    if isnothing(wt)
        wt = ones(n, n)
        Vinv = I(n) / n
    else
        V = - copy(wt)
        for i in 1:size(V, 1)
            V[i, i] = -sum(view(V, :, i))
        end
        Vinv = pinv(V)
    end

    # Use classical scaling to get initial mapping if Xinit is not provided
    X = isnothing(Xinit) ? classical_scaling(Δ, p) : Xinit
    return_hist && (Xhist = [X])
    D = dist(X)
    B = zeros(T, size(Δ))
    update!(B, Δ, D, wt)
    σ = zeros(T, itmax + 1)
    σ[1] = stress(Δ, D, wt)

    Xcache = zeros(size(X))
    for i in 1:itmax
        log_output(i, σ, verbose)
        # Gutman transform
        mul!(Xcache, B, X)
        mul!(X, Vinv, Xcache)
        return_hist && push!(Xhist, copy(X))
        dist!(D, X)
        σ[i + 1] = stress(Δ, D, wt)
        abs(σ[i] - σ[i + 1]) < ε * σ[i] && break
        update!(B, Δ, D, wt)
    end

    # Use anchors if provided
    align!(X, anchors)
    return_hist ? (hist = Xhist, σ = σ) : X
end

function update!(B, Δ, D, W)
    for ij in eachindex(B)
        B[ij] = W[ij] * Δ[ij] / (D[ij] + eps())
    end

    for i in 1:size(B, 1)
        B[i, i] = 0.0
        B[i, i] = -sum(view(B, :, i))
    end
end