text:
VcovFormula(::Type{Val{:cluster}}, x) = VcovClusterFormula(Terms(@eval(@formula($nothing ~ $x))).terms)

struct VcovClusterFormula <: AbstractVcovFormula
    _::Vector{Any}
end
allvars(x::VcovClusterFormula) =  vcat([allvars(a) for a in x._]...)

struct VcovClusterMethod <: AbstractVcovMethod
    clusters::DataFrame
end

function VcovMethod(df::AbstractDataFrame, vcovcluster::VcovClusterFormula)
    clusters = vcovcluster._
    vclusters = DataFrame(Matrix{Vector}(undef, size(df, 1), 0))
    for c in clusters
        if isa(c, Symbol)
            isa(df[c], CategoricalVector) || error("Cluster variable $(c) is of type $(typeof(df[c])), but should be a CategoricalVector.")
            vclusters[c] = group(df[c])
        elseif isa(c, Expr)
            factorvars, interactionvars = _split(df, allvars(c))
            vclusters[_name(factorvars)] = group((df[v] for v in factorvars)...)
        end
    end
    return VcovClusterMethod(vclusters)
end

function df_FStat(v::VcovClusterMethod, ::VcovData, ::Bool)
    minimum((length(v.clusters[c].pool) for c in names(v.clusters))) - 1
end

function vcov!(v::VcovClusterMethod, x::VcovData)
    S = shat!(v, x)
    return pinvertible(sandwich(x.crossmatrix, S))
end

function shat!(v::VcovClusterMethod, x::VcovData{T, N}) where {T, N}
    # Cameron, Gelbach, & Miller (2011): section 2.3
    dim = size(x.regressors, 2) * size(x.residuals, 2)
    S = zeros(dim, dim)
    G = typemax(Int)
    for c in combinations(names(v.clusters))
        # no need for group in case of one fixed effect, since was already done in VcovMethod
        f = (length(c) == 1) ? v.clusters[c[1]] : group((v.clusters[var] for var in c)...)
        # capture length of smallest dimension of multiway clustering in G
        G = min(G, length(f.pool))
        S += (-1)^(length(c) - 1) * helper_cluster(x.regressors, x.residuals, f)
    end
    # scale total vcov estimate by ((N-1)/(N-K)) * (G/(G-1))
    # another option would be to adjust each matrix given by helper_cluster by number of its categories
    # both methods are presented in Cameron, Gelbach and Miller (2011), section 2.3
    # I choose the first option following reghdfe
    rmul!(S, (size(x.regressors, 1) - 1) / x.dof_residual * G / (G - 1))
end

# Matrix version is used for IV
function helper_cluster(X::Matrix{Float64}, res::Union{Vector{Float64}, Matrix{Float64}}, f::CategoricalVector)
    X2 = fill(zero(Float64), length(f.pool), size(X, 2) * size(res, 2))
    index = 0
    for k in 1:size(res, 2)
        for j in 1:size(X, 2)
            index += 1
            @inbounds @simd for i in 1:size(X, 1)
                X2[f.refs[i], index] += X[i, j] * res[i, k]
            end
        end
    end
    return X2' * X2
end

function pinvertible(A::Matrix, tol = eps(real(float(one(eltype(A))))))
    eigval, eigvect = eigen(Symmetric(A))
    small = eigval .<= tol
    if any(small)
        @warn "estimated covariance matrix of moment conditions not of full rank.
                 model tests should be interpreted with caution."
        eigval[small] .= 0
        return eigvect' * Diagonal(eigval) * eigvect
    else
        return A
    end
end


text:
import Base

"""
Struct for WilsonFermion
"""
struct WilsonFermion_4D_wing_Adjoint{NC,NDW,NumofBasis} <: WilsonFermion_4D{NC} #AbstractFermionfields_4D{NC}

    f::Array{ComplexF64,6}
    NC::Int64
    NX::Int64
    NY::Int64
    NZ::Int64
    NT::Int64
    NG::Int64
    NDW::Int64
    Dirac_operator::String
    NumofBasis::Int64
    #BoundaryCondition::Vector{Int8}

    function WilsonFermion_4D_wing_Adjoint(NC::T,NX::T,NY::T,NZ::T,NT::T) where T<: Integer
        NumofBasis = ifelse(NC == 1,1,NC^2-1)
        NG = 4
        NDW = 1
        #@assert NDW == 1 "only NDW = 1 is supported. Now NDW = $NDW"
        f = zeros(ComplexF64,NumofBasis,NX+2NDW,NY+2NDW,NZ+2NDW,NT+2NDW,NG)
        Dirac_operator = "Wilson"
        return new{NC,NDW,NumofBasis}(f,NC,NX,NY,NZ,NT,NG,NDW,Dirac_operator,NumofBasis)
    end

    function WilsonFermion_4D_wing_Adjoint{NC}(NX::T,NY::T,NZ::T,NT::T) where {T<: Integer,NC}
        NumofBasis = ifelse(NC == 1,1,NC^2-1)
        NG = 4
        NDW = 1
        #@assert NDW == 1 "only NDW = 1 is supported. Now NDW = $NDW"
        f = zeros(ComplexF64,NumofBasis,NX+2NDW,NY+2NDW,NZ+2NDW,NT+2NDW,NG)
        Dirac_operator = "Wilson"
        return new{NC,NDW,NumofBasis}(f,NC,NX,NY,NZ,NT,NG,NDW,Dirac_operator,NumofBasis)
    end

end

function initialize_Adjoint_fermion(x::WilsonFermion_4D_wing{NC,NDW}) where {NC,NDW}
    xadj = WilsonFermion_4D_wing_Adjoint(x.NC,x.NX,x.NY,x.NZ,x.NT)
end

function Base.length(x::WilsonFermion_4D_wing_Adjoint{NC,NDW,NumofBasis}) where {NC,NDW,NumofBasis}
    return NC*x.NX*x.NY*x.NZ*x.NT*x.NG
end

function Base.setindex!(x::WilsonFermion_4D_wing_Adjoint{NC,NDW,NumofBasis},v,i1,i2,i3,i4,i5,i6)  where {NC,NDW,NumofBasis}
    @inbounds x.f[i1,i2 + NDW,i3 + NDW,i4 + NDW,i5 + NDW,i6] = v
end

function Base.getindex(x::WilsonFermion_4D_wing_Adjoint{NC,NDW,NumofBasis},i1,i2,i3,i4,i5,i6) where {NC,NDW,NumofBasis}
    @inbounds return x.f[i1,i2 .+ NDW,i3 .+ NDW,i4 .+ NDW,i5 .+ NDW,i6]
end

function Base.getindex(x::WilsonFermion_4D_wing_Adjoint{NC,NDW,NumofBasis},i1::N,i2::N,i3::N,i4::N,i5::N,i6::N) where {NC,NDW,NumofBasis,N<: Integer}
    @inbounds return  x.f[i1,i2 + NDW,i3 + NDW,i4 + NDW,i5 + NDW,i6]
end

function Base.similar(x::T) where T <: WilsonFermion_4D_wing_Adjoint
    return WilsonFermion_4D_wing_Adjoint(x.NC,x.NX,x.NY,x.NZ,x.NT)
end

function set_wing_fermion!(a::WilsonFermion_4D_wing_Adjoint{NC,NDW,NumofBasis},boundarycondition) where {NC,NDW,NumofBasis}
    NT = a.NT
    NZ = a.NZ
    NY = a.NY
    NX = a.NX

    #!  X-direction
    for ialpha=1:4
        for it=1:NT
            for iz = 1:NZ
                for iy=1:NY
                    @simd for k=1:NumofBasis
                        a[k,0,iy,iz,it,ialpha] = boundarycondition[1]*a[k,NX,iy,iz,it,ialpha]
                    end
                end
            end
        end
    end

    for ialpha=1:4
        for it=1:NT
            for iz=1:NZ
                for iy=1:NY
                    @simd for k=1:NumofBasis
                        a[k,NX+1,iy,iz,it,ialpha] = boundarycondition[1]*a[k,1,iy,iz,it,ialpha]
                    end
                end
            end
        end
    end

    #Y-direction
    for ialpha = 1:4
        for it=1:NT
            for iz=1:NZ
                for ix=1:NX
                    @simd for k=1:NumofBasis
                        a[k,ix,0,iz,it,ialpha] = boundarycondition[2]*a[k,ix,NY,iz,it,ialpha]
                    end
                end
            end
        end
    end

    for ialpha=1:4
        for it=1:NT
            for iz=1:NZ
                for ix=1:NX
                    @simd for k=1:NumofBasis
                        a[k,ix,NY+1,iz,it,ialpha] = boundarycondition[2]*a[k,ix,1,iz,it,ialpha]
                    end
                end
            end
        end
    end

    for ialpha=1:4
        # Z-direction
        for it=1:NT
            for iy=1:NY
                for ix=1:NX
                    @simd for k=1:NumofBasis
                        a[k,ix,iy,0,it,ialpha] = boundarycondition[3]*a[k,ix,iy,NZ,it,ialpha]
                        a[k,ix,iy,NZ+1,it,ialpha] = boundarycondition[3]*a[k,ix,iy,1,it,ialpha]

                    end
                end
            end
        end

        #T-direction
        for iz=1:NZ
            for iy=1:NY
                for ix=1:NX
                    @simd for k=1:NumofBasis
                        a[k,ix,iy,iz,0,ialpha] = boundarycondition[4]*a[k,ix,iy,iz,NT,ialpha]
                        a[k,ix,iy,iz,NT+1,ialpha] =boundarycondition[4]*a[k,ix,iy,iz,1,ialpha]
                    end
                end
            end
        end

    end

end

function set_wing_fermion!(a::WilsonFermion_4D_wing_Adjoint{NC,NDW,NumofBasis},boundarycondition,iseven::Bool) where {NC,NDW,NumofBasis}
    NT = a.NT
    NZ = a.NZ
    NY = a.NY
    NX = a.NX

    #!  X-direction
    @inbounds for ialpha=1:4
        for it=1:NT
            for iz = 1:NZ
                for iy=1:NY
                    ix = NX
                    evenodd = ifelse((ix+iy+iz+it) % 2 == 0,true,false)
                    if evenodd == iseven
                        @simd for k=1:NumofBasis
                            a[k,0,iy,iz,it,ialpha] = boundarycondition[1]*a[k,NX,iy,iz,it,ialpha]
                        end
                    end
                end
            end
        end
    end

    @inbounds for ialpha=1:4
        for it=1:NT
            for iz=1:NZ
                for iy=1:NY
                    ix = 1
                    evenodd = ifelse((ix+iy+iz+it) % 2 == 0,true,false)
                    if evenodd == iseven
                        @simd for k=1:NumofBasis
                            a[k,NX+1,iy,iz,it,ialpha] = boundarycondition[1]*a[k,1,iy,iz,it,ialpha]
                        end
                    end
                end
            end
        end
    end

    #Y-direction
    @inbounds for ialpha = 1:4
        for it=1:NT
            for iz=1:NZ
                for ix=1:NX
                    iy = NY
                    evenodd = ifelse((ix+iy+iz+it) % 2 == 0,true,false)
                    if evenodd == iseven

                        @simd for k=1:NumofBasis
                            a[k,ix,0,iz,it,ialpha] = boundarycondition[2]*a[k,ix,NY,iz,it,ialpha]
                        end
                    end
                end

            end
        end
    end

    @inbounds for ialpha=1:4
        for it=1:NT
            for iz=1:NZ
                for ix=1:NX
                    iy = 1
                    evenodd = ifelse((ix+iy+iz+it) % 2 == 0,true,false)
                    if evenodd == iseven

                        @simd for k=1:NumofBasis
                            a[k,ix,NY+1,iz,it,ialpha] = boundarycondition[2]*a[k,ix,1,iz,it,ialpha]
                        end
                    end
                end
            end
        end
    end

    @inbounds for ialpha=1:4
        # Z-direction
        for it=1:NT
            for iy=1:NY
                for ix=1:NX
                    iz = NX
                    evenodd = ifelse((ix+iy+iz+it) % 2 == 0,true,false)
                    if evenodd == iseven
                        @simd for k=1:NumofBasis
                            a[k,ix,iy,0,it,ialpha] = boundarycondition[3]*a[k,ix,iy,NZ,it,ialpha]
                        end
                    end

                    iz = 1
                    evenodd = ifelse((ix+iy+iz+it) % 2 == 0,true,false)
                    if evenodd == iseven
                        @simd for k=1:NumofBasis
                            a[k,ix,iy,NZ+1,it,ialpha] = boundarycondition[3]*a[k,ix,iy,1,it,ialpha]
                        end
                    end
                end
            end
        end

        #T-direction
        for iz=1:NZ
            for iy=1:NY
                for ix=1:NX
                    it = NT
                    evenodd = ifelse((ix+iy+iz+it) % 2 == 0,true,false)
                    if evenodd == iseven
                        @simd for k=1:NumofBasis
                            a[k,ix,iy,iz,0,ialpha] = boundarycondition[4]*a[k,ix,iy,iz,NT,ialpha]
                        end
                    end

                    it = 1
                    evenodd = ifelse((ix+iy+iz+it) % 2 == 0,true,false)
                    if evenodd == iseven
                        @simd for k=1:NumofBasis
                            a[k,ix,iy,iz,NT+1,ialpha] =boundarycondition[4]*a[k,ix,iy,iz,1,ialpha]
                        end
                    end

                end
            end
        end

    end

end

function add_fermion!(c::WilsonFermion_4D_wing_Adjoint{NC,NDW,NumofBasis},α::Number,a::T1,β::Number,b::T2,iseven) where {NC,NDW,NumofBasis,T1 <: Abstractfermion,T2 <: Abstractfermion}#c += alpha*a + beta*b
    n1,n2,n3,n4,n5,n6 = size(c.f)

    @inbounds for i6=1:n6
        for i5=1:n5
            it = i5 -NDW
            for i4=1:n4
                iz = i4 -NDW
                for i3=1:n3
                    iy = i3 - NDW
                    for i2=1:n2
                        ix = i2 - NDW
                        evenodd = ifelse((ix + iy + iz + it) % 2 == 0,true,false)
                        if evenodd == iseven
                            @simd for i1=1:NumofBasis
                                #println(a.f[i1,i2,i3,i4,i5,i6],"\t",b.f[i1,i2,i3,i4,i5,i6] )
                                c.f[i1,i2,i3,i4,i5,i6] += α*a.f[i1,i2,i3,i4,i5,i6] + β*b.f[i1,i2,i3,i4,i5,i6]
                            end
                        end
                    end
                end
            end
        end
    end
    return
end

function add_fermion!(c::WilsonFermion_4D_wing_Adjoint{NC,NDW,NumofBasis},α::Number,a::T1,iseven::Bool) where {NC,NDW,NumofBasis,T1 <: Abstractfermion,T2 <: Abstractfermion}#c += alpha*a + beta*b
    n1,n2,n3,n4,n5,n6 = size(c.f)

    @inbounds for i6=1:n6
        for i5=1:n5
            it = i5 -NDW
            for i4=1:n4
                iz = i4 -NDW
                for i3=1:n3
                    iy = i3 - NDW
                    for i2=1:n2
                        ix = i2 - NDW
                        evenodd = ifelse((ix + iy + iz + it) % 2 == 0,true,false)
                        if evenodd == iseven
                            @simd for i1=1:NumofBasis
                                #println(a.f[i1,i2,i3,i4,i5,i6],"\t",b.f[i1,i2,i3,i4,i5,i6] )
                                c.f[i1,i2,i3,i4,i5,i6] += α*a.f[i1,i2,i3,i4,i5,i6]
                            end
                        end
                    end
                end
            end
        end
    end
    return
end

function clear_fermion!(a::WilsonFermion_4D_wing_Adjoint{NC,NDW,NumofBasis} ,iseven) where {NC,NDW,NumofBasis}
    n1,n2,n3,n4,n5,n6 = size(a.f)
    @inbounds for i6=1:n6
        for i5=1:n5
            it = i5-NDW
            for i4=1:n4
                iz = i4-NDW
                for i3=1:n3
                    iy = i3 - NDW
                    for i2=1:n2
                        ix = i2 - NDW
                        evenodd = ifelse((ix+iy+iz+it) % 2 == 0,true,false)
                        if evenodd == iseven
                            @simd for i1=1:NumofBasis
                                a.f[i1,i2,i3,i4,i5,i6]= 0
                            end
                        end
                    end
                end
            end
        end
    end
end

function LinearAlgebra.mul!(x::WilsonFermion_4D_wing_Adjoint{NC,NDW,NumofBasis},A::TA) where {TA <: AbstractMatrix, NC,NDW,NumofBasis}
    NX = x.NX
    NY = x.NY
    NZ = x.NZ
    NT = x.NT

    #n6 = size(x.f)[6]
    #f = zeros(ComplexF64,4)
    #e = zeros(ComplexF64,4)

    for ic=1:NumofBasis
        for it=1:NT
            for iz=1:NZ
                for iy=1:NY
                    @simd for ix=1:NX
                            e1 = x[ic,ix,iy,iz,it,1]
                            e2 = x[ic,ix,iy,iz,it,2]
                            e3 = x[ic,ix,iy,iz,it,3]
                            e4 = x[ic,ix,iy,iz,it,4]

                            x[ic,ix,iy,iz,it,1] = A[1,1]*e1+A[1,2]*e2+A[1,3]*e3+A[1,4]*e4
                            x[ic,ix,iy,iz,it,2] = A[2,1]*e1+A[2,2]*e2+A[2,3]*e3+A[2,4]*e4
                            x[ic,ix,iy,iz,it,3] = A[3,1]*e1+A[3,2]*e2+A[3,3]*e3+A[3,4]*e4
                            x[ic,ix,iy,iz,it,4] = A[4,1]*e1+A[4,2]*e2+A[4,3]*e3+A[4,4]*e4

                    end
                end
            end
        end
    end

end

function LinearAlgebra.mul!(x::WilsonFermion_4D_wing_Adjoint{NC,NDW,NumofBasis},A::TA,iseven) where {TA <: AbstractMatrix, NC,NDW,NumofBasis}
    NX = x.NX
    NY = x.NY
    NZ = x.NZ
    NT = x.NT

    #n6 = size(x.f)[6]
    #f = zeros(ComplexF64,4)
    #e = zeros(ComplexF64,4)

    for ic=1:NumofBasis
        for it=1:NT
            for iz=1:NZ
                for iy=1:NY
                    @simd for ix=1:NX
                        evenodd = ifelse((ix+iy+iz+it) % 2 == 0,true,false)
                        if evenodd == iseven
                            e1 = x[ic,ix,iy,iz,it,1]
                            e2 = x[ic,ix,iy,iz,it,2]
                            e3 = x[ic,ix,iy,iz,it,3]
                            e4 = x[ic,ix,iy,iz,it,4]

                            x[ic,ix,iy,iz,it,1] = A[1,1]*e1+A[1,2]*e2+A[1,3]*e3+A[1,4]*e4
                            x[ic,ix,iy,iz,it,2] = A[2,1]*e1+A[2,2]*e2+A[2,3]*e3+A[2,4]*e4
                            x[ic,ix,iy,iz,it,3] = A[3,1]*e1+A[3,2]*e2+A[3,3]*e3+A[3,4]*e4
                            x[ic,ix,iy,iz,it,4] = A[4,1]*e1+A[4,2]*e2+A[4,3]*e3+A[4,4]*e4
                        end
                    end
                end
            end
        end
    end

end

function LinearAlgebra.mul!(xout::WilsonFermion_4D_wing_Adjoint{NC,NDW,NumofBasis},A::TA,x::WilsonFermion_4D_wing_Adjoint{NC,NDW,NumofBasis}) where {TA <: AbstractMatrix, NC,NDW,NumofBasis}
    NX = x.NX
    NY = x.NY
    NZ = x.NZ
    NT = x.NT

    #n6 = size(x.f)[6]
    #f = zeros(ComplexF64,4)
    #e = zeros(ComplexF64,4)

    for ic=1:NumofBasis
        for it=1:NT
            for iz=1:NZ
                for iy=1:NY
                    @simd for ix=1:NX
                            e1 = x[ic,ix,iy,iz,it,1]
                            e2 = x[ic,ix,iy,iz,it,2]
                            e3 = x[ic,ix,iy,iz,it,3]
                            e4 = x[ic,ix,iy,iz,it,4]

                            xout[ic,ix,iy,iz,it,1] = A[1,1]*e1+A[1,2]*e2+A[1,3]*e3+A[1,4]*e4
                            xout[ic,ix,iy,iz,it,2] = A[2,1]*e1+A[2,2]*e2+A[2,3]*e3+A[2,4]*e4
                            xout[ic,ix,iy,iz,it,3] = A[3,1]*e1+A[3,2]*e2+A[3,3]*e3+A[3,4]*e4
                            xout[ic,ix,iy,iz,it,4] = A[4,1]*e1+A[4,2]*e2+A[4,3]*e3+A[4,4]*e4

                    end
                end
            end
        end
    end

end

function LinearAlgebra.mul!(xout::WilsonFermion_4D_wing_Adjoint{NC,NDW,NumofBasis},A::TA,x::WilsonFermion_4D_wing_Adjoint{NC,NDW,NumofBasis},iseven) where {TA <: AbstractMatrix, NC,NDW,NumofBasis}
    NX = x.NX
    NY = x.NY
    NZ = x.NZ
    NT = x.NT

    #n6 = size(x.f)[6]
    #f = zeros(ComplexF64,4)
    #e = zeros(ComplexF64,4)

    for ic=1:NumofBasis
        for it=1:NT
            for iz=1:NZ
                for iy=1:NY
                    @simd for ix=1:NX
                        evenodd = ifelse((ix+iy+iz+it) % 2 == 0,true,false)
                        if evenodd == iseven
                            e1 = x[ic,ix,iy,iz,it,1]
                            e2 = x[ic,ix,iy,iz,it,2]
                            e3 = x[ic,ix,iy,iz,it,3]
                            e4 = x[ic,ix,iy,iz,it,4]

                            xout[ic,ix,iy,iz,it,1] = A[1,1]*e1+A[1,2]*e2+A[1,3]*e3+A[1,4]*e4
                            xout[ic,ix,iy,iz,it,2] = A[2,1]*e1+A[2,2]*e2+A[2,3]*e3+A[2,4]*e4
                            xout[ic,ix,iy,iz,it,3] = A[3,1]*e1+A[3,2]*e2+A[3,3]*e3+A[3,4]*e4
                            xout[ic,ix,iy,iz,it,4] = A[4,1]*e1+A[4,2]*e2+A[4,3]*e3+A[4,4]*e4
                        end

                    end
                end
            end
        end
    end

end

function LinearAlgebra.mul!(xout::WilsonFermion_4D_wing_Adjoint{NC,NDW,NumofBasis},x::WilsonFermion_4D_wing_Adjoint{NC,NDW,NumofBasis},A::TA) where {TA <: AbstractMatrix, NC,NDW,NumofBasis}
    NX = x.NX
    NY = x.NY
    NZ = x.NZ
    NT = x.NT

    #n6 = size(x.f)[6]
    #f = zeros(ComplexF64,4)
    #e = zeros(ComplexF64,4)

    for ic=1:NumofBasis
        for it=1:NT
            for iz=1:NZ
                for iy=1:NY
                    @simd for ix=1:NX
                            e1 = x[ic,ix,iy,iz,it,1]
                            e2 = x[ic,ix,iy,iz,it,2]
                            e3 = x[ic,ix,iy,iz,it,3]
                            e4 = x[ic,ix,iy,iz,it,4]

                            xout[ic,ix,iy,iz,it,1] = A[1,1]*e1+A[2,1]*e2+A[3,1]*e3+A[4,1]*e4
                            xout[ic,ix,iy,iz,it,2] = A[1,2]*e1+A[2,2]*e2+A[3,2]*e3+A[4,2]*e4
                            xout[ic,ix,iy,iz,it,3] = A[1,3]*e1+A[2,3]*e2+A[3,3]*e3+A[4,3]*e4
                            xout[ic,ix,iy,iz,it,4] = A[1,4]*e1+A[2,4]*e2+A[3,4]*e3+A[4,4]*e4

                    end
                end
            end
        end
    end

end

function LinearAlgebra.mul!(xout::WilsonFermion_4D_wing_Adjoint{NC,NDW,NumofBasis},x::WilsonFermion_4D_wing_Adjoint{NC,NDW,NumofBasis},A::TA,iseven) where {TA <: AbstractMatrix, NC,NDW,NumofBasis}
    NX = x.NX
    NY = x.NY
    NZ = x.NZ
    NT = x.NT

    #n6 = size(x.f)[6]
    #f = zeros(ComplexF64,4)
    #e = zeros(ComplexF64,4)

    for ic=1:NumofBasis
        for it=1:NT
            for iz=1:NZ
                for iy=1:NY
                    @simd for ix=1:NX
                        evenodd = ifelse((ix+iy+iz+it) % 2 == 0,true,false)
                        if evenodd == iseven
                            e1 = x[ic,ix,iy,iz,it,1]
                            e2 = x[ic,ix,iy,iz,it,2]
                            e3 = x[ic,ix,iy,iz,it,3]
                            e4 = x[ic,ix,iy,iz,it,4]

                            xout[ic,ix,iy,iz,it,1] = A[1,1]*e1+A[2,1]*e2+A[3,1]*e3+A[4,1]*e4
                            xout[ic,ix,iy,iz,it,2] = A[1,2]*e1+A[2,2]*e2+A[3,2]*e3+A[4,2]*e4
                            xout[ic,ix,iy,iz,it,3] = A[1,3]*e1+A[2,3]*e2+A[3,3]*e3+A[4,3]*e4
                            xout[ic,ix,iy,iz,it,4] = A[1,4]*e1+A[2,4]*e2+A[3,4]*e3+A[4,4]*e4
                        end
                    end
                end
            end
        end
    end

end

function LinearAlgebra.dot(a::WilsonFermion_4D_wing_Adjoint{NC,NDW,NumofBasis},b::WilsonFermion_4D_wing_Adjoint{NC,NDW,NumofBasis},iseven::Bool) where {NC,NDW,NumofBasis}
    NT = a.NT
    NZ = a.NZ
    NY = a.NY
    NX = a.NX
    NG = a.NG

    c = 0.0im
    @inbounds for α=1:NG
        for it=1:NT
            for iz=1:NZ
                for iy=1:NY
                    for ix=1:NX
                        evenodd = ifelse((ix+iy+iz+it) % 2 == 0,true,false)
                        if evenodd == iseven
                            @simd for ic=1:NumofBasis
                                c+= conj(a[ic,ix,iy,iz,it,α])*b[ic,ix,iy,iz,it,α]
                            end
                        end
                    end
                end
            end
        end
    end
    return c
end

#Overwrite Y with X*a + Y*b, where a and b are scalars. Return Y.
function LinearAlgebra.axpby!(a::Number, X::T, b::Number, Y::WilsonFermion_4D_wing_Adjoint{NC,NDW,NumofBasis},iseven::Bool) where {NC,NDW,NumofBasis,T <: AbstractFermionfields_4D}
    n1,n2,n3,n4,n5,n6 = size(Y.f)
    #println("axpby")

    @inbounds for i6=1:n6
        for i5=1:n5
            it = i5+NDW
            for i4=1:n4
                iz = i4+NDW
                for i3=1:n3
                    iy = i3+NDW
                    for i2=1:n2
                        ix = i2+NDW
                        evenodd = ifelse((ix+iy+iz+it) % 2 == 0,true,false)
                        if evenodd == iseven
                            @simd for i1=1:NumofBasis
                                Y.f[i1,i2,i3,i4,i5,i6] = a*X.f[i1,i2,i3,i4,i5,i6] + b*Y.f[i1,i2,i3,i4,i5,i6]
                            end
                        end
                    end
                end
            end
        end
    end
end

"""
c--------------------------------------------------------------------------c
c     y = gamma_5 * x
c     here
c                  ( -1       )
c        GAMMA5 =  (   -1     )
c                  (     +1   )
c                  (       +1 )
c--------------------------------------------------------------------------c
    """
    function mul_γ5x!(y::WilsonFermion_4D_wing_Adjoint{NC,NDW,NumofBasis},x::WilsonFermion_4D_wing_Adjoint{NC,NDW,NumofBasis}) where {NC,NDW,NumofBasis}
        n1,n2,n3,n4,n5,n6 = size(x.f)
        @inbounds for i6=1:n6
            for i5=1:n5
                #it = i5+NDW
                for i4=1:n4
                    #iz = i4+NDW
                    for i3=1:n3
                        #iy = i3+NDW
                        for i2=1:n2
                            #ix = i2+NDW
                            @simd for ic=1:NumofBasis
                                y.f[ic,i2,i3,i4,i5,i6]=x.f[ic,i2,i3,i4,i5,i6]*ifelse(i6 <= 2,-1,1)
                            end
                        end
                    end
                end
            end
        end

    end

    function apply_γ5!(x::WilsonFermion_4D_wing_Adjoint{NC,NDW,NumofBasis}) where {NC,NDW,NumofBasis}
        n1,n2,n3,n4,n5,n6 = size(x.f)
        #println("axpby")

        @inbounds for i6=1:n6
            for i5=1:n5
                #it = i5+NDW
                for i4=1:n4
                    #iz = i4+NDW
                    for i3=1:n3
                        #iy = i3+NDW
                        for i2=1:n2
                            #ix = i2+NDW
                            @simd for ic=1:NumofBasis
                                x.f[ic,i2,i3,i4,i5,i6]=x.f[ic,i2,i3,i4,i5,i6]*ifelse(i6 <= 2,-1,1)
                            end
                        end
                    end
                end
            end
        end

    end

    function mul_1plusγ5x!(y::WilsonFermion_4D_wing_Adjoint{NC,NDW,NumofBasis},x::WilsonFermion_4D_wing_Adjoint{NC,NDW,NumofBasis})  where {NC,NDW,NumofBasis}#(1+gamma_5)/2
        NX = x.NX
        NY = x.NY
        NZ = x.NZ
        NT = x.NT
        #NC = x.NC
        for ic=1:NumofBasis
            for it=1:NT
                for iz=1:NZ
                    for iy=1:NY
                        @simd for ix=1:NX
                            y[ic,ix,iy,iz,it,1] = 0#-1*x[ic,ix,iy,iz,it,1]
                            y[ic,ix,iy,iz,it,2] = 0#-1*x[ic,ix,iy,iz,it,2]
                            y[ic,ix,iy,iz,it,3] = x[ic,ix,iy,iz,it,3]
                            y[ic,ix,iy,iz,it,4] = x[ic,ix,iy,iz,it,4]
                        end
                    end
                end
            end
        end
    end

    function mul_1plusγ5x_add!(y::WilsonFermion_4D_wing_Adjoint{NC,NDW,NumofBasis},x::WilsonFermion_4D_wing_Adjoint{NC,NDW,NumofBasis},factor) where {NDW,NumofBasis,NC}#x = x +(1+gamma_5)/2
        NX = x.NX
        NY = x.NY
        NZ = x.NZ
        NT = x.NT
        #NC = x.NC
        for ic=1:NumofBasis
            for it=1:NT
                for iz=1:NZ
                    for iy=1:NY
                        @simd for ix=1:NX
                            y[ic,ix,iy,iz,it,3] += factor*x[ic,ix,iy,iz,it,3]
                            y[ic,ix,iy,iz,it,4] += factor*x[ic,ix,iy,iz,it,4]
                        end
                    end
                end
            end
        end
    end

    function mul_1minusγ5x!(y::WilsonFermion_4D_wing_Adjoint{NC,NDW,NumofBasis},x::WilsonFermion_4D_wing_Adjoint{NC,NDW,NumofBasis}) where {NC,NDW,NumofBasis}#(1-gamma_5)/2
        NX = x.NX
        NY = x.NY
        NZ = x.NZ
        NT = x.NT
        #NC = x.NC
        for ic=1:NumofBasis
            for it=1:NT
                for iz=1:NZ
                    for iy=1:NY
                        @simd for ix=1:NX
                            y[ic,ix,iy,iz,it,1] = x[ic,ix,iy,iz,it,1]
                            y[ic,ix,iy,iz,it,2] = x[ic,ix,iy,iz,it,2]
                            y[ic,ix,iy,iz,it,3] = 0#x[ic,ix,iy,iz,it,3]
                            y[ic,ix,iy,iz,it,4] = 0#x[ic,ix,iy,iz,it,4]
                        end
                    end
                end
            end
        end
    end

    function mul_1minusγ5x_add!(y::WilsonFermion_4D_wing_Adjoint{NC,NDW,NumofBasis},x::WilsonFermion_4D_wing_Adjoint{NC,NDW,NumofBasis},factor) where {NC,NDW,NumofBasis}#+(1-gamma_5)/2
        NX = x.NX
        NY = x.NY
        NZ = x.NZ
        NT = x.NT
        #NC = x.NC
        for ic=1:NumofBasis
            for it=1:NT
                for iz=1:NZ
                    for iy=1:NY
                        @simd for ix=1:NX
                            y[ic,ix,iy,iz,it,1] += factor*x[ic,ix,iy,iz,it,1]
                            y[ic,ix,iy,iz,it,2] += factor*x[ic,ix,iy,iz,it,2]
                        end
                    end
                end
            end
        end
    end

    function LinearAlgebra.mul!(y::WilsonFermion_4D_wing_Adjoint{NC,NDW,NumofBasis},A::T,x::T3) where {NC,NDW,NumofBasis,T<:Abstractfields,T3 <:Abstractfermion}
        #@assert NC == x.NC "dimension mismatch! NC in y is $NC but NC in x is $(x.NC)"
        #@assert NC != 3 "NC should not be 3"
        NX = y.NX
        NY = y.NY
        NZ = y.NZ
        NT = y.NT
        NG = y.NG

        @inbounds for ialpha=1:NG
            for it=1:NT
                for iz=1:NZ
                    for iy=1:NY
                        for ix=1:NX
                            for k1=1:NumofBasis
                                y[k1,ix,iy,iz,it,ialpha] = 0
                                @simd for k2=1:NumofBasis
                                    y[k1,ix,iy,iz,it,ialpha] += A[k1,k2,ix,iy,iz,it]*x[k2,ix,iy,iz,it,ialpha]
                                end
                            end
                        end
                    end
                end
            end
        end
    end

#=
function LinearAlgebra.mul!(y::WilsonFermion_4D_wing_Adjoint{3,NDW,NumofBasis},A::T,x::T3) where {NDW,NumofBasis,T<:Abstractfields,T3 <:Abstractfermion}
    mul!(y,A,x)
end

function LinearAlgebra.mul!(y::WilsonFermion_4D_wing_Adjoint{3,NDW,NumofBasis},A::T,x::T3,iseven::Bool) where {NDW,NumofBasis,T<:Abstractfields,T3 <:Abstractfermion}
    mul!(y,A,x,iseven)
end

function LinearAlgebra.mul!(y::WilsonFermion_4D_wing_Adjoint{2,NDW,NumofBasis},A::T,x::T3) where {NDW,NumofBasis,T<:Abstractfields,T3 <:Abstractfermion}
    mul!(y,A,x)
end

function LinearAlgebra.mul!(y::WilsonFermion_4D_wing_Adjoint{2,NDW,NumofBasis},A::T,x::T3,iseven::Bool) where {NDW,NumofBasis,T<:Abstractfields,T3 <:Abstractfermion}
    mul!(y,A,x,iseven)
end
=#

function LinearAlgebra.mul!(y::WilsonFermion_4D_wing_Adjoint{NC,NDW,NumofBasis},x::T3,A::T) where {NC,NDW,NumofBasis,T<:Abstractfields,T3 <:Abstractfermion}
    #@assert NC == x.NC "dimension mismatch! NC in y is $NC but NC in x is $(x.NC)"
    NX = y.NX
    NY = y.NY
    NZ = y.NZ
    NT = y.NT
    NG = y.NG

    @inbounds for ialpha=1:NG
        for it=1:NT
            for iz=1:NZ
                for iy=1:NY
                    for ix=1:NX
                        for k1=1:NumofBasis
                            y[k1,ix,iy,iz,it,ialpha] = 0
                            @simd for k2=1:NumofBasis
                                y[k1,ix,iy,iz,it,ialpha] += x[k1,ix,iy,iz,it,ialpha]*A[k1,k2,ix,iy,iz,it]
                            end
                        end
                    end
                end
            end
        end
    end
end

function LinearAlgebra.mul!(y::WilsonFermion_4D_wing_Adjoint{NC,NDW,NumofBasis},x::T3,A::T,iseven::Bool) where {NC,NDW,NumofBasis,T<:Abstractfields,T3 <:Abstractfermion}
    #@assert NC == x.NC "dimension mismatch! NC in y is $NC but NC in x is $(x.NC)"
    NX = y.NX
    NY = y.NY
    NZ = y.NZ
    NT = y.NT
    NG = y.NG

    @inbounds for ialpha=1:NG
        for it=1:NT
            for iz=1:NZ
                for iy=1:NY
                    for ix=1:NX
                        evenodd = ifelse((ix + iy + iz + it) % 2 == 0,true,false)
                        if evenodd == iseven

                            for k1=1:NumofBasis
                                y[k1,ix,iy,iz,it,ialpha] = 0
                                @simd for k2=1:NumofBasis
                                    y[k1,ix,iy,iz,it,ialpha] += x[k1,ix,iy,iz,it,ialpha]*A[k1,k2,ix,iy,iz,it]
                                end
                            end
                        end
                    end
                end
            end
        end
    end
end

function LinearAlgebra.mul!(y::WilsonFermion_4D_wing_Adjoint{3,NDW,NumofBasis},x::T3,A::T) where {NDW,NumofBasis,T<:Abstractfields,T3 <:Abstractfermion}
    mul!(y,x,A)
end

function LinearAlgebra.mul!(y::WilsonFermion_4D_wing_Adjoint{2,NDW,NumofBasis},x::T3,A::T) where {NDW,NumofBasis,T<:Abstractfields,T3 <:Abstractfermion}
    mul!(y,x,A)
end

function LinearAlgebra.mul!(y::WilsonFermion_4D_wing_Adjoint{NC,NDW,NumofBasis},A::T,x::T3) where {NC,NDW,NumofBasis,T<:Number,T3 <:Abstractfermion}
    @assert NC == x.NC "dimension mismatch! NC in y is $NC but NC in x is $(x.NC)"
    NX = y.NX
    NY = y.NY
    NZ = y.NZ
    NT = y.NT
    NG = y.NG

    @inbounds for ialpha=1:NG
        for it=1:NT
            for iz=1:NZ
                for iy=1:NY
                    for ix=1:NX
                        for k1=1:NumofBasis
                            y[k1,ix,iy,iz,it,ialpha] = A*x[k1,ix,iy,iz,it,ialpha]
                        end
                    end
                end
            end
        end
    end
end

ext: jl


text:
module TestParticle

using LinearAlgebra: norm, ×
using Meshes
using Interpolations
using StaticArrays

export prepare, trace_numeric!, trace_analytic!, trace_analytic_relativistic!
export trace_numeric, trace_analytic, trace_analytic_relativistic
export Proton, Electron, Ion, User

include("utility/utility.jl")

"""
Type for the particles, `Proton`, `Electron`, `Ion`, or `User`.
"""
@enum Species Proton Electron Ion User

function getchargemass(species::Species, q, m)

   if species == Proton
      q = qᵢ
      m = mᵢ
   elseif species == Electron
      q = qₑ
      m = mₑ
   end
   q, m
end

function makegrid(grid)

   gridmin = coordinates(minimum(grid))
   gridmax = coordinates(maximum(grid))
   Δx = spacing(grid)

   gridx = range(gridmin[1], gridmax[1], step=Δx[1])
   gridy = range(gridmin[2], gridmax[2], step=Δx[2])
   gridz = range(gridmin[3], gridmax[3], step=Δx[3])

   gridx, gridy, gridz
end

function getinterp(A, gridx, gridy, gridz)

   @assert size(A,1) == 3 && ndims(A) == 4 "Only support 3D force field!"

   Ax = @view A[1,:,:,:]
   Ay = @view A[2,:,:,:]
   Az = @view A[3,:,:,:]

   itp = extrapolate(interpolate(Ax, BSpline(Cubic(Interpolations.Line(OnGrid())))), NaN)
   interpx = scale(itp, gridx, gridy, gridz)

   itp = extrapolate(interpolate(Ay, BSpline(Cubic(Interpolations.Line(OnGrid())))), NaN)
   interpy = scale(itp, gridx, gridy, gridz)

   itp = extrapolate(interpolate(Az, BSpline(Cubic(Interpolations.Line(OnGrid())))), NaN)
   interpz = scale(itp, gridx, gridy, gridz)

   interpx, interpy, interpz
end

"""
    prepare(grid, E, B; species=Proton) -> (q, m, E, B)

Return a tuple consists of particle charge, mass for a prescribed `species` and interpolated
EM field functions.
"""
function prepare(grid::CartesianGrid, E, B; species::Species=Proton, q=1.0, m=1.0)

   q, m = getchargemass(species, q, m)

   gridx, gridy, gridz = makegrid(grid)

   interpEx, interpEy, interpEz = getinterp(E, gridx, gridy, gridz)
   interpBx, interpBy, interpBz = getinterp(B, gridx, gridy, gridz)

   q, m, (interpEx, interpEy, interpEz), (interpBx, interpBy, interpBz)
end

"""
    prepare(grid, E, B, F; species=Proton, q=1.0, m=1.0) -> (q, m, E, B, F)

Return a tuple consists of particle charge, mass for a prescribed `species` of charge `q`
and mass `m`, analytic EM field functions, and external force `F`.
"""
function prepare(grid::CartesianGrid, E, B, F; species::Species=Proton, q=1.0, m=1.0)

   q, m = getchargemass(species, q, m)

   gridx, gridy, gridz = makegrid(grid)

   interpEx, interpEy, interpEz = getinterp(E, gridx, gridy, gridz)
   interpBx, interpBy, interpBz = getinterp(B, gridx, gridy, gridz)
   interpFx, interpFy, interpFz = getinterp(F, gridx, gridy, gridz)

   q, m, (interpEx, interpEy, interpEz), (interpBx, interpBy, interpBz),
   (interpFx, interpFy, interpFz)
end

"""
    prepare(E, B; species=Proton, q=1.0, m=1.0) -> (q, m, E, B)

Return a tuple consists of particle charge, mass for a prescribed `species` of charge `q`
and mass `m` and analytic EM field functions. Prescribed `species` are `Electron` and
`Proton`; other species can be manually specified with `species=Ion/User`, `q` and `m`.
"""
function prepare(E, B; species::Species=Proton, q=1.0, m=1.0)

   q, m = getchargemass(species, q, m)

   q, m, E, B
end

"""
    prepare(E, B, F; species=Proton, q=1.0, m=1.0) -> (q, m, E, B, F)

Return a tuple consists of particle charge, mass for a prescribed `species` of charge `q`
and mass `m`, analytic EM field functions, and external force `F`.
"""
function prepare(E, B, F; species::Species=Proton, q=1.0, m=1.0)
   t = prepare(E, B; species, q, m)
   push!(t, F)
   t
end

"ODE equations for charged particle moving in static numerical EM field."
function trace_numeric!(dy, y, p, t)
   q, m, interpE, interpB = p
   dy[1:3] = y[4:6]
   dy[4:6] = q/m*(getE(y, interpE) + y[4:6] × getB(y, interpB))
end

function trace_numeric(y, p, t)
   q, m, interpE, interpB = p
   dx, dy, dz = y[4:6]
   dux, duy, duz = q/m*(getE(y, interpE) + y[4:6] × getB(y, interpB))
   SVector{6}(dx, dy, dz, dux, duy, duz)
end

"ODE equations for charged particle moving in static numerical EM field and
external force field."
function trace_numeric_full!(dy, y, p, t)
   q, m, interpE, interpB, interpF = p
   dy[1:3] = y[4:6]
   dy[4:6] = (q*(getE(y, interpE) + y[4:6] × getB(y, interpB)) + getF(t, interpF)) / m
end

function trace_numeric_full(y, p, t)
   q, m, interpE, interpB, interpF = p
   dx, dy, dz = y[4:6]
   dux, duy, duz = (q*(getE(y, interpE) + y[4:6] × getB(y, interpB)) + getF(t, interpF)) / m
   SVector{6}(dx, dy, dz, dux, duy, duz)
end

"ODE equations for relativistic charged particle moving in static numerical EM field."
function trace_numeric_relativistic!(dy, y, p, t)
   q, m, E, B = p
   if y[4]*y[4] + y[5]*y[5] + y[6]*y[6] ≥ c^2
      throw(ArgumentError("Particle faster than the speed of light!"))
   end
   γInv = √(1.0 - (y[4]*y[4] + y[5]*y[5] + y[6]*y[6])/c^2)
   dy[1:3] = y[4:6]
   dy[4:6] = q/m*γInv*(getE(y, interpE) + y[4:6] × getB(y, interpB))
end

function trace_numeric_relativistic(y, p, t)
   q, m, E, B = p
   if y[4]*y[4] + y[5]*y[5] + y[6]*y[6] ≥ c^2
      throw(ArgumentError("Particle faster than the speed of light!"))
   end
   γInv = √(1.0 - (y[4]*y[4] + y[5]*y[5] + y[6]*y[6])/c^2)
   dx, dy, dz = y[4:6]
   dux, duy, duz = q/m*γInv*(getE(y, interpE) + y[4:6] × getB(y, interpB))
   SVector{6}(dx, dy, dz, dux, duy, duz)
end

"ODE equations for charged particle moving in static analytical EM field."
function trace_analytic!(dy, y, p, t)
   q, m, E, B = p
   dy[1:3] = y[4:6]
   dy[4:6] = q/m*(E(y) + y[4:6] × (B(y[1:3])))
end

function trace_analytic(y, p, t)
   q, m, E, B = p
   dx, dy, dz = y[4:6]
   dux, duy, duz = q/m*(E(y) + y[4:6] × (B(y[1:3])))
   SVector{6}(dx, dy, dz, dux, duy, duz)
end

"ODE equations for charged particle moving in static analytical EM field and external force
field."
function trace_analytic_full!(dy, y, p, t)
   q, m, E, B, F = p
   dy[1:3] = y[4:6]
   dy[4:6] = (q*(E(y) + y[4:6] × (B(y[1:3]))) + F) / m
end

function trace_analytic_full(y, p, t)
   q, m, E, B, F = p
   dx, dy, dz = y[4:6]
   dux, duy, duz = (q*(E(y) + y[4:6] × (B(y[1:3]))) + F) / m
   SVector{6}(dx, dy, dz, dux, duy, duz)
end

"ODE equations for relativistic charged particle moving in static analytical EM field."
function trace_analytic_relativistic!(dy, y, p, t)
   q, m, E, B = p

   if y[4]*y[4] + y[5]*y[5] + y[6]*y[6] ≥ c^2
      throw(ArgumentError("Particle faster than the speed of light!"))
   end
   γInv = √(1.0 - (y[4]*y[4] + y[5]*y[5] + y[6]*y[6])/c^2)
   dy[1:3] = y[4:6]
   dy[4:6] = q/m*γInv*(E(y) + y[4:6] × (B(y[1:3])))
end

function trace_analytic_relativistic(y, p, t)
   q, m, E, B = p

   if y[4]*y[4] + y[5]*y[5] + y[6]*y[6] ≥ c^2
      throw(ArgumentError("Particle faster than the speed of light!"))
   end
   γInv = √(1.0 - (y[4]*y[4] + y[5]*y[5] + y[6]*y[6])/c^2)
   dx, dy, dz = y[4:6]
   dux, duy, duz = q/m*γInv*(E(y) + y[4:6] × (B(y[1:3])))
   SVector{6}(dx, dy, dz, dux, duy, duz)
end

# Return eletric field at a given location.
function getE(xu, interpE)
   x = @view xu[1:3]
   u = @view xu[4:6]

   SA[interpE[1](x...), interpE[2](x...), interpE[3](x...)]
end

# Return magnetic field at a given location.
function getB(xu, interpB)
   x = @view xu[1:3]
   u = @view xu[4:6]

   SA[interpB[1](x...), interpB[2](x...), interpB[3](x...)]
end

# Return force at a given location.
function getF(xu, interpF)
   x = @view xu[1:3]
   u = @view xu[4:6]

   SA[interpF[1](x...), interpF[2](x...), interpF[3](x...)]
end

end


text:
module Util

export Option, apply_temperature

import Random
using Distributions: Categorical

const Option{T} = Union{T, Nothing}

infinity(::Type{R}) where R <: Real = one(R) / zero(R)

"""
    concat_columns(cols) == hcat(cols...) # but faster
"""
function concat_columns(cols)
  @assert !isempty(cols)
  nsamples = length(cols)
  excol = first(cols)
  sdim = length(excol)
  arr = similar(excol, (sdim, nsamples))
  for (i, col) in enumerate(cols)
    arr[:,i] = col
  end
  return arr
end

"""
    superpose(xs) == cat(xs..., dims=ndims(first(xs))+1) # but faster
"""
function superpose(arrays)
  n = length(arrays)
  @assert n > 0
  ex = first(arrays)
  dest = similar(ex, size(ex)..., n)
  i = 1
  for src in arrays
    for j in eachindex(src)
      dest[i] = src[j]
      i += 1
    end
  end
  return dest
end

"""
    @printing_errors expr

Evaluate expression `expr` while printing any uncaught exception on `stderr`.

This is useful to avoid silent falure of concurrent tasks, as explained in
[this issue](https://github.com/JuliaLang/julia/issues/10405).
"""
macro printing_errors(expr)
  return quote
    try
      $(esc(expr))
    catch e
      showerror(stderr, e, catch_backtrace())
    end
  end
end

"""
    generate_update_constructor(T)

Generate a new constructor for immutable type `T` that enables copying
an existing structure while only updating a subset of its fields.

For example, given the following struct:

    struct Point
      x :: Int
      y :: Int
    end

The generated code is equivalent to:

    Point(pt; x=pt.x, y=pt.y) = Point(x, y)

** This function may be deprecated in favor of Setfield.jl in the future.**
"""
function generate_update_constructor(T)
  fields = fieldnames(T)
  Tname = Symbol(split(string(T), ".")[end])
  base = :_old_
  @assert base ∉ fields
  fields_withdef = [Expr(:kw, f, :($base.$f)) for f in fields]
  quote
    #$Tname(;$(fields...)) = $Tname($(fields...))
    $Tname($base::$Tname; $(fields_withdef...)) = $Tname($(fields...))
  end
end

"""
    fix_probvec(π)

Convert probability vector `π` to type `Vector{Float32}` and renormalize it
if necessary.

This is useful as `Distributions.isprobvec` can be picky about its
input when it does not sum to one due to numerical errors.
"""
function fix_probvec(π)
  π = convert(Vector{Float32}, π)
  s = sum(π)
  if !(s ≈ 1)
    if iszero(s)
      n = length(π)
      π = ones(Float32, n) ./ n
    else
      π ./= s
    end
  end
  return π
end

"""
Draw a sample from a categorical distribution represented as a probability vector.
See [`fix_probvec`](@ref).
"""
function rand_categorical(π)
  π = fix_probvec(π)
  return rand(Categorical(π))
end

"""
    apply_temperature(π, τ)

Apply temperature `τ` to probability distribution `π`.
Handle the limit case where `τ=0`.
"""
function apply_temperature(π, τ)
  if isone(τ)
    return π
  elseif iszero(τ)
    res = zeros(eltype(π), length(π))
    res[argmax(π)] = 1
    return res
  else
    res = π .^ inv(τ)
    res ./= sum(res)
    return res
  end
end

"""
Same smoothing function that is used by Temsorboard to smooth time series.
"""
function momentum_smoothing(x, μ)
  sx = similar(x)
  isempty(x) && return x
  v = x[1]
  for i in eachindex(x)
    v = μ * x[i] + (1-μ) * v
    sx[i] = v
  end
  return sx
end

"""
    batches(X, batchsize; partial=false)

Take a data tensor `X` and split it into batches of fixed size along the
last dimension of `X`.

If `partial=true` and the number of samples in `X` is
not a multiple of `batchsize`, then an additional smaller batch is added
at the end (otherwise, it is discarded).
"""
function batches(X, batchsize; partial=false)
  n = size(X)[end]
  b = batchsize
  nbatches = n ÷ b
  # The call to `copy` after selectdim is important because Flux does not
  # deal well with views.
  select(a, b) = copy(selectdim(X, ndims(X), a:b))
  batches = [select(1+b*(i-1), b*i) for i in 1:nbatches]
  if partial && n % b > 0
    # If the number of samples is not a multiple of the batch size
    push!(batches, select(b*nbatches+1, n))
  end
  return batches
end

function batches_tests()
  @assert batches(collect(1:5), 2, partial=true) == [[1, 2], [3, 4], [5]]
end

"""
    random_batches(convert, data::Tuple, batchsize; partial=false)

Take a tuple of data tensors, shuffle its samples according to a random
permutation and split them into a sequence of minibatches.

The result is a lazy iterator that calls `convert` on the tensors of each
new batch right before returning it. The `convert` function is typically
used to transfer data onto the GPU.

!!! note
    In the future, it may be good to deprecate this function along with
    [`random_batches_stream`](@ref) and use a standard solution instead,
    such as `Flux.DataLoader`.
"""
function random_batches(
  convert, data::Tuple, batchsize; partial=false)
  n = size(data[1])[end]
  perm = Random.randperm(n)
  batchs = map(data) do x
    batches(selectdim(x, ndims(x), perm), batchsize, partial=partial)
  end
  batchs = collect(zip(batchs...))
  return (convert.(b) for b in batchs)
end

"""
    random_batches_stream(convert, data::Tuple, batchsize)

Generate an infinite stateful iterator of random batches by calling
[`random_batches`](@ref) repeatedly. Every sample is guaranteed to be drawn
exactly once per epoch.
"""
function random_batches_stream(convert, data::Tuple, batchsize)
  partial = size(data[1])[end] < batchsize
  return Iterators.Stateful(Iterators.flatten((
    random_batches(convert, data, batchsize, partial=partial)
    for _ in Iterators.repeated(nothing))))
end

#####
##### Multithreading utilities
#####

function threads_pmap(f, xs)
  return fetch.([Threads.@spawn @printing_errors f(x) for x in xs])
end

# TODO: the `mapreduce` function should ultimately be removed and replaced by a
# more sandard solution.

"""
    mapreduce(make_worker, args, num_workers, combine, init)

In spirit, this computes `reduce(combine, map(f, args); init)` on `num_workers`
where `f` is defined below.

The `make_worker` function must create a worker `w` with two fields:
  - a `process` function such that `w.process(x)` evaluates to `f(x)`
  - a `terminate` function to be called with no argument when the worker terminates.

!!! note

    This function makes one call to `combine` per computed element
    and so it should only be used when the associated synchronization cost
    is small compared to the cost of computing individual elements.
"""
function mapreduce(make_worker, args, num_workers, combine, init)
  next = 1
  ret = init
  lock = ReentrantLock()
  tasks = []
  for w in 1:num_workers
    task = Threads.@spawn Util.@printing_errors begin
      local k = 0
      worker = make_worker()
      while true
        Base.lock(lock)
        if next > length(args)
          Base.unlock(lock)
          break
        end
        k = next
        next += 1
        Base.unlock(lock)
        y = worker.process(args[k])
        Base.lock(lock)
        ret = combine(ret, y)
        Base.unlock(lock)
      end
      worker.terminate()
    end
    push!(tasks, task)
  end
  wait.(tasks)
  return ret
end

# Same semantics than mapreduce but sequential: to be used for
# debugging purposes only.
function mapreduce_sequential(make_worker, args, num_workers, combine, unit)
  ys = map(args) do x
    worker = f()
    y = worker.process(x)
    worker.terminate()
    return y
  end
  return reduce(combine, ys, init=unit)
end

end


text:
include("HWLC.jl")
using .HWLC
using Plots, Distributions, MAT

#Set plot parameters
gr(xlabel="F(X) (pN)", label="",ylabel="F'(X) (pN/nm)",
    grid=false,guidefont=("Arial",10),
    legend=:topleft,
    legendfont=("Arial",10),scale=:log10)

#Import experiment data for comparison
data=matread("Experimental_data/U2OS_ForceStiffnessStd.mat")
mean_experiment=vec(data["k_avg"])
std_exp=vec(data["k_err"])
fs_exp=vec(data["f_ax"])

#Set parameters
kT=4.22
Marko_Siggia=true

f_min,f_max,N_f=extrema(fs_exp)...,500
forces=exp.(range(log(f_min),log(f_max),length=N_f))

N=7        #Number of component chains
L=1200.     #Total length of assembly

#Power-law distribution for persistence lengths
slope=0.86        #Slope of stiffness-force graph
# Set parameters for f_c distribution
minfc, maxfc=9.,1000.
λ=160.

"""
Show individual and total curves for one assembly
"""
fcs=rand(truncated(Exponential(λ),minfc,maxfc),N)
persistence_lengths=fcs.^(-1)*kT
#ls=ones(N)*L/N  #All equal length
ls=rand(Normal(L/N,L/N*0.1),N) #Alternatively sample from normal distr

total_ext,extensions=force_extensions(forces,ls,persistence_lengths,kT,Marko_Siggia)
total_stiffness=stiffnesses(forces,total_ext)[1]
individual_stiffnesses=stiffnesses(forces,extensions)

#Plot individual and total response
plot(forces[2:end],total_stiffness, label="Assembly");
plot!([20,200],[10^-1.5,(200/20)^slope*10^-1.5], color="black",label="Slope $slope");
for ind in individual_stiffnesses
    plot!(forces[2:end],ind, color="lightgrey")
end
display(current())
png("Plots/exp_distr_individual_responses")

"""
Sample many assemblies, and compare mean
stiffening behaviour to data
"""
n_samples=15
sim_data=zeros(length(forces)-1,n_samples)
for i in 1:n_samples
    #sample assembly
    sample_ls=rand(Normal(L/N,L/N*0.1),N) #comment if using constant lengths
    sample_ps=rand(truncated(Exponential(λ),minfc,maxfc),N).^(-1)*kT
    tot,exts=force_extensions(forces,sample_ls,sample_ps,kT,Marko_Siggia)
    sim_data[:,i]=stiffnesses(forces,tot)[1]
end

mean_sim=mean(sim_data,dims=2)
std_sim=std(sim_data,dims=2)

plot(fs_exp,mean_experiment, label="Experimental mean",legend=:topleft, color="blue");
plot!(fs_exp,mean_experiment.+std_exp, color="lightblue");
plot!(fs_exp[findall(>(0.),mean_experiment.-std_exp)],(mean_experiment.-std_exp)[findall(>(0.),mean_experiment.-std_exp)], color="lightblue");
plot!(forces[2:end],mean_sim, label="Sim mean", color="orange", ribbon=std_sim)
png("Plots/exp_distr_compare_data_sim")


text:
#################### Posterior Statistics ####################
"""
    autocor(chn;
        lags=[1, 5, 10, 50],
        demean=true,
        relative=true
        showall=false,
        append_chains=true,
        digits::Int=4,
        sections=[:parameters])

Compute the autocorrelation of each parameter for the chain. Setting `append_chains=false` will return a vector of dataframes containing the autocorrelations for each chain.

The `digits` keyword may be a(n)
- `Integer`, which sets rounds all numerical columns to the value)
- `NamedTuple`, which only rounds the named column to the specified digits, as with `(mean=2, std=3)`. This would round the `mean` column to 2 digits and the `std` column to 3 digits.
- `Dict`, with a similar structure as `NamedTuple`. `Dict(mean => 2, std => 3)` would set `mean` to two digits and `std` to three digits.
"""
function autocor(chn::Chains;
        lags::AbstractVector{<:Integer}=[1, 5, 10, 50],
        demean::Bool=true,
        relative::Bool=true,
        showall=false,
        append_chains = false,
        sections::Union{Symbol, Vector{Symbol}}=Symbol[:parameters],
        digits::Int=4,
        kwargs...)
    funs = Function[]
    func_names = @. Symbol("lag ", lags)
    for i in lags
        push!(funs, x -> autocor(x, [i], demean=demean)[1])
    end
    return summarize(chn, funs...;
        func_names = func_names,
        showall = showall,
        sections = sections,
        append_chains = append_chains,
        name = "Autocorrelation",
        digits=digits)
end

"""
    cor(chn; showall=false, append_chains=true, sections=[:parameters], digits::Int=4)

Compute the Pearson correlation matrix for the chain. Setting `append_chains=false` will
return a vector of dataframes containing a correlation matrix for each chain.

The `digits` keyword may be a(n)
- `Integer`, which sets rounds all numerical columns to the value)
- `NamedTuple`, which only rounds the named column to the specified digits, as with `(mean=2, std=3)`. This would round the `mean` column to 2 digits and the `std` column to 3 digits.
- `Dict`, with a similar structure as `NamedTuple`. `Dict(mean => 2, std => 3)` would set `mean` to two digits and `std` to three digits.
"""
function cor(chain::Chains;
        showall=false,
        append_chains=true,
        sections::Union{Symbol, Vector{Symbol}}=:parameters,
        digits::Int=4,
        kwargs...
)
    # Obtain interesting subset of the chain.
    chn = showall ? chain : Chains(chain, sections)

    # Obstain names of parameters.
    names_of_params = Symbol.(names(chn))

    if append_chains
        df = chaindataframe_cor("Correlation", names_of_params, to_matrix(chn);
                                digits = digits)
        return df
    else
        vector_of_df = [
            chaindataframe_cor(
                "Correlation - Chain $i", names_of_params, data; digits = digits
            )
            for (i, data) in enumerate(to_vector_of_matrices(chn))
        ]
        return vector_of_df
    end
end

function chaindataframe_cor(name, names_of_params, chains::AbstractMatrix; kwargs...)
    # Compute the correlation matrix.
    cormat = cor(chains)

    # Summarize the results in a named tuple.
    nt = (; parameters = names_of_params,
          zip(names_of_params, (cormat[:, i] for i in axes(cormat, 2)))...)

    # Create a ChainDataFrame.
    return ChainDataFrame(name, nt; kwargs...)
end

"""
    changerate(chn;
        append_chains=true,
        showall=false,
        sections=[:parameters],
        digits::Int=4)

Computes the change rate for the chain. Setting `append_chains=false` will
return a vector of dataframes containing the change rates for each chain.

The `digits` keyword may be a(n)
- `Integer`, which sets rounds all numerical columns to the value)
- `NamedTuple`, which only rounds the named column to the specified digits, as with `(mean=2, std=3)`. This would round the `mean` column to 2 digits and the `std` column to 3 digits.
- `Dict`, with a similar structure as `NamedTuple`. `Dict(mean => 2, std => 3)` would set `mean` to two digits and `std` to three digits.
"""
function changerate(chains::Chains{<:Real};
    append_chains=true,
    showall=false,
    sections::Union{Symbol, Vector{Symbol}}=Symbol[:parameters],
    digits::Int=4,
    kwargs...
)
    # Obtain interesting subset of the chains.
    chn = showall ? chains : Chains(chains, sections)

    # Obstain names of parameters.
    names_of_params = Symbol.(names(chn))

    if append_chains
        df = chaindataframe_changerate("Change Rate", names_of_params, chn.value.data)
        return df
    else
        vector_of_df = [
            chaindataframe_changerate(
                "Change Rate - Chain $i", names_of_params, data; digits = digits
            )
            for (i, data) in enumerate(to_vector_of_matrices(chn))
        ]
        return vector_of_df
    end
end

function chaindataframe_changerate(name, names_of_params, chains; kwargs...)
    # Compute the change rates.
    changerates, mvchangerate = chains

    # Summarize the results in a named tuple.
    nt = (; zip(names_of_params, changerates)..., multivariate = mvchangerate)

    # Create a ChainDataFrame.
    return ChainDataFrame(name, nt; kwargs...)
end

changerate(chains::AbstractMatrix{<:Real}) = changerate(reshape(chains, Val(3)))
function changerate(chains::AbstractArray{<:Real,3})
    niters, nparams, nchains = size(chains)

    changerates = zeros(nparams)
    mvchangerate = 0.0

    for chain in 1:nchains, iter in 2:niters
        isanychanged = false

        for param in 1:nparams
            # update if the sample is different from the one in the previous iteration
            if chains[iter-1, param, chain] != chains[iter, param, chain]
                changerates[param] += 1
                isanychanged = true
            end
        end

        mvchangerate += isanychanged
    end

    factor = nchains * (niters - 1)
    changerates ./= factor
    mvchangerate /= factor

    changerates, mvchangerate
end

describe(c::Chains; args...) = describe(stdout, c; args...)

"""
    describe(io::IO,
        c::Chains;
        q = [0.025, 0.25, 0.5, 0.75, 0.975],
        etype=:bm,
        showall=false,
        sections::Union{Symbol, Vector{Symbol}}=Symbol[:parameters],
        digits::Int=4,
        args...)

Prints the summary statistics and quantiles for the chain.

The `digits` keyword may be a(n)
- `Integer`, which sets rounds all numerical columns to the value)
- `NamedTuple`, which only rounds the named column to the specified digits, as with `(mean=2, std=3)`. This would round the `mean` column to 2 digits and the `std` column to 3 digits.
- `Dict`, with a similar structure as `NamedTuple`. `Dict(mean => 2, std => 3)` would set `mean` to two digits and `std` to three digits.
"""
function describe(io::IO,
                  c::Chains;
                  q = [0.025, 0.25, 0.5, 0.75, 0.975],
                  etype=:bm,
                  showall::Bool=false,
                  sections::Union{Symbol, Vector{Symbol}}=Symbol[:parameters],
                  digits::Int=4,
                  sorted=false,
                  args...
                 )
    dfs = vcat(summarystats(c,
                showall=showall,
                sections=sections,
                etype=etype,
                digits=digits,
                sorted=sorted,
                args...),
           quantile(c,
                showall=showall,
                sections=sections,
                q=q,
                digits=digits,
                sorted=sorted))
    return dfs
end

function _hpd(x::AbstractVector{<:Real}; alpha::Real=0.05)
    n = length(x)
    m = max(1, ceil(Int, alpha * n))

    y = sort(x)
    a = y[1:m]
    b = y[(n - m + 1):n]
    _, i = findmin(b - a)

    return [a[i], b[i]]
end

function hpd(chn::Chains; alpha::Real=0.05, kwargs...)
    labels = [:upper, :lower]
    l(x) = _hpd(x, alpha=alpha)[1]
    u(x) = _hpd(x, alpha=alpha)[2]
    return summarize(chn, u, l; name = "HPD", func_names = labels, kwargs...)
end

"""
    quantile(chn;
        q=[0.025, 0.25, 0.5, 0.75, 0.975],
        append_chains=true,
        showall=false,
        sections=[:parameters],
        digits::Int=4)

Computes the quantiles for each parameter in the chain. Setting `append_chains=false` will
return a vector of dataframes containing the quantiles for each chain.

The `digits` keyword may be a(n)
- `Integer`, which sets rounds all numerical columns to the value)
- `NamedTuple`, which only rounds the named column to the specified digits, as with `(mean=2, std=3)`. This would round the `mean` column to 2 digits and the `std` column to 3 digits.
- `Dict`, with a similar structure as `NamedTuple`. `Dict(mean => 2, std => 3)` would set `mean` to two digits and `std` to three digits.
"""
function quantile(chn::Chains;
        q::Vector=[0.025, 0.25, 0.5, 0.75, 0.975],
        append_chains=true,
        showall=false,
        sections::Union{Symbol, Vector{Symbol}}=Symbol[:parameters],
        digits::Int=4,
        sorted=false)
    # compute quantiles
    funs = Function[]
    func_names = @. Symbol(100 * q, :%)
    for i in q
        push!(funs, x -> quantile(cskip(x), i))
    end

    return summarize(chn, funs...;
        func_names=func_names,
        showall=showall,
        sections=sections,
        name="Quantiles",
        digits=digits,
        append_chains=append_chains,
        sorted=sorted)
end

"""
	ess(chn::Chains;
		showall=false,
		sections::Union{Symbol, Vector{Symbol}}=Symbol[:parameters],
		maxlag = 250,
        digits::Int=4)

Compute a chain's number of effective samples. More information can be found in the Gelman et al. (2014) book "Bayesian Data Analysis", or in [this article](https://arxiv.org/abs/1903.08008).

The `digits` keyword may be a(n)
- `Integer`, which sets rounds all numerical columns to the value)
- `NamedTuple`, which only rounds the named column to the specified digits, as with `(mean=2, std=3)`. This would round the `mean` column to 2 digits and the `std` column to 3 digits.
- `Dict`, with a similar structure as `NamedTuple`. `Dict(mean => 2, std => 3)` would set `mean` to two digits and `std` to three digits.
"""
function ess(chn::Chains;
    showall=false,
    sections::Union{Symbol, Vector{Symbol}}=Symbol[:parameters],
    maxlag = 250,
    digits::Int=4,
    sorted=false
)
	param = showall ? names(chn) : names(chn, sections)
	n_chain_orig = size(chn, 3)

	# Split the chains.
	parameter_vec = Vector(undef, length(param))
    midpoint = Int32(floor(size(chn, 1) / 2))
    n = size(chn, 1)
	for i in 1:length(param)
		parameter_vec[i] = []
		for j in 1:n_chain_orig
            c1 = vec(cskip(chn[1:midpoint, param[i], j].value.data))
            c2 = vec(cskip(chn[midpoint+1:end, param[i], j].value.data))
            n = min(n, length(c1), length(c2))
            push!(parameter_vec[i], c1, c2)
		end
	end

    # Misc allocations.
    m = n_chain_orig * 2
    maxlag = min(maxlag, n-1)
    lags = 0:maxlag

    # Preallocate B, W, varhat, and Rhat vectors for each param.
    B = Vector(undef, length(param))
    W = Vector(undef, length(param))
    varhat = Vector(undef, length(param))
    Rhat = Vector{Float64}(undef, length(param))

    # calculate B, W, varhat, and Rhat for each param.
    for i in 1:length(param)
		draw = parameter_vec[i]
		p = param[i]
        allchain = mean(vcat([d for d in draw]...))
        eachchain = [mean(draw[j]) for j in 1:m]
        s = [sum((draw[j] .- eachchain[j]).^2) / (n-1) for j in 1:m]
        B[i] = (n / (m - 1)) * sum((eachchain .- allchain).^2)
        W[i] = sum(s) / m
        varhat[i] = (n-1)/n * W[i] + (B[i] / n)
        Rhat[i] = sqrt(varhat[i] / W[i])
    end

	V = Vector(undef, length(param))
    ρ = Vector(undef, length(param))
    for p in eachindex(V)
        V[p] = Vector(undef, length(lags))
		ρ[p] = Vector(undef, length(lags))
    end

    # Calculate ρ
    c_autocor = Vector(undef, length(param))
    for i in 1:length(param)
        c_autocor[i] = [0.0]
    end

    for t in eachindex(lags)
        lag = lags[t]
        range1 = lag+1:n
        range2 = 1:(n-lag)
        for i in 1:length(param)
            draw = parameter_vec[i]
			p = param[i]
            z = [draw[j][range1] .- draw[j][range2] for j in 1:m]
			z = sum([zi .^ 2 for zi in z])
            V[i][t] = 1 / (m * (n-lag)) * sum(z)
            autocors = [_autocorrelation(draw[j], lag) for j in 1:m]
			ρ[i][t] = 1 - V[i][t] / (2 * varhat[i])
        end
    end

	# Find first odd positive integer where ρ[p][T+1] + ρ[p][T+2] is negative
    P = Vector(undef, length(param))
    ess = Vector{Float64}(undef, length(param))
	for i in 1:length(param)
        big_P = 0.0
		ρ_val = Float64.(ρ[i])

        # Big P.
        P[i] = Float64[ρ_val[1]]
        k = tprime = 1
        for tprime in 1:Int(floor((length(lags)/2 - 1)))
            sumvals = ρ_val[2*tprime] + ρ_val[2*tprime+1]
            push!(P[i], sumvals)
            k = tprime

            if sumvals < 0
                break
            end
        end

        # Create monotone.
        P_monotone = [min(P[i][t], P[i][1:t]...) for t in 1:length(P[i])]

        ess[i] = (n*m) / (-1 + 2*sum(P_monotone))
	end

    df = (parameters = string.(param), ess = ess, r_hat = Rhat)
	return ChainDataFrame("ESS", df; digits=digits)
end

# this function is sourced from https://github.com/tpapp/MCMCDiagnostics.jl/blob/master/src/MCMCDiagnostics.jl
function _autocorrelation(x::AbstractVector, k::Integer, v = var(x), kwargs...)
    x1 = @view(x[1:(end-k)])
    x2 = @view(x[(1+k):end])
    V = sum((x1 .- x2).^2) / length(x1)
    1 - V / (2*v)
end

"""
    summarystats(chn;
        append_chains=true,
        showall=false,
        sections=[:parameters],
        digits::Int=4,
        args...)

Computes the mean, standard deviation, naive standard error, Monte Carlo standard error, and effective sample size for each parameter in the chain. Setting `append_chains=false` will return a vector of dataframes containing the summary statistics for each chain. `args...` is passed to the `msce` function.

The `digits` keyword may be a(n)
- `Integer`, which sets rounds all numerical columns to the value)
- `NamedTuple`, which only rounds the named column to the specified digits, as with `(mean=2, std=3)`. This would round the `mean` column to 2 digits and the `std` column to 3 digits.
- `Dict`, with a similar structure as `NamedTuple`. `Dict(mean => 2, std => 3)` would set `mean` to two digits and `std` to three digits.
"""
function summarystats(chn::Chains;
        append_chains::Bool=true,
        showall::Bool=false,
        sections::Union{Symbol, Vector{Symbol}}=Symbol[:parameters],
        etype=:bm,
        digits::Int=4,
        sorted=false,
        args...
    )

    # Make some functions.
    df_mcse(x) = length(x) < 200 ?
        missing :
        mcse(cskip(x), etype, args...)

    # Store everything.
    funs = [mean∘cskip, std∘cskip, sem∘cskip, df_mcse]
    func_names = [:mean, :std, :naive_se, :mcse]

    # Caluclate ESS separately.
    ess_df = ess(chn, sections=sections, showall=showall, sorted=sorted)

    # Summarize.
    summary_df = summarize(chn, funs...;
        func_names=func_names,
        showall=showall,
        sections=sections,
        name="Summary Statistics",
        additional_df = ess_df,
        digits=digits,
        append_chains=append_chains,
        sorted=sorted)

    return summary_df
end

"""
    mean(chn::Chains;
            append_chains::Bool=true,
            showall::Bool=false,
            sections::Union{Symbol, Vector{Symbol}}=Symbol[:parameters],
            digits::Int=4,
            args...)
    mean(chn::Chains, ss::Vector{Symbol})
    mean(chn::Chains, s::Symbol)
    mean(chn::Chains, s::String)
    mean(chn::Chains, ss::Vector{String})

Calculates the mean of a `Chains` object, or a specific parameter.
"""
function mean(chn::Chains;
        append_chains::Bool=true,
        showall::Bool=false,
        sections::Union{Symbol, Vector{Symbol}}=Symbol[:parameters],
        digits::Int=4,
        args...
    )
    # Store everything.
    funs = [mean∘cskip]
    func_names = [:mean]

    # Summarize.
    summary_df = summarize(chn, funs...;
        func_names=func_names,
        showall=showall,
        sections=sections,
        name="Mean",
        digits=digits)

    return summary_df
end

mean(chn::Chains, ss::Vector{Symbol}) = mean(chn[:, ss, :])
mean(chn::Chains, s::String) = mean(chn, Symbol(s))
mean(chn::Chains, ss::Vector{String}) = mean(chn, Symbol.(ss))
function mean(chn::Chains, s::Symbol)
    syms = _sym2index(chn, [s])

    if length(syms) == 0
        throw(ArgumentError("Symbol :$s not found in chain."))
    end

    if length(syms) > 1
        return mean(chn[:, syms, :])
    else
        return mean(chn[:,syms,:].value.data)
    end
end


text:
module ClusteringMetrics
using StatsBase
using MultivariateStats
using LinearAlgebra

"""
    get_cluster_distance(Z1::Matrix{T}, label1::Vector{T2}, Z2::Matrix{T}, label2::Vector{T2}) where T <: Real where T2

Compute the distance between the class clusters in `Z1` and `Z2` using class labels `label1` and `label2`, normalized by the the total intra-class variance.
"""
function get_cluster_distance(Z1::Matrix{T}, label1::Vector{T2}, Z2::Matrix{T}, label2::Vector{T2}) where T <: Real where T2
    # Compute class statistics on the projected trials
    nc1 = maximum(label1)
    nc2 = maximum(label2)
    ldastats_1 = MultivariateStats.multiclass_lda_stats(nc1, Z1, label1)
    ldastats_2 = MultivariateStats.multiclass_lda_stats(nc2, Z2, label2)

    # Compute a normalized distance between cluster centers
    Sw = Diagonal(diag(ldastats_1.Sw + ldastats_2.Sw))
    d = ldastats_1.cmeans - ldastats_2.cmeans
    D = diag(d'*inv(Sw)*d)
end

"""
    get_cluster_confusion(Z1::Matrix{T}, label1::Vector{T2}, Z2::Matrix{T}, label2::Vector{T2}) where T <: Real where T2

Compute the distance to the nearest cluster normalied to the distance to the true cluster.
"""
function get_cluster_confusion(Z1::Matrix{T}, label1::Vector{T2}, Z2::Matrix{T}, label2::Vector{T2}) where T <: Real where T2

    # Compute class statistics on the projected trials
    ldastats_1 = CodeMorphingCore.MultivariateStats.multiclass_lda_stats(4, Z1, label1)
    ldastats_2 = CodeMorphingCore.MultivariateStats.multiclass_lda_stats(4, Z2, label2)
    nc1= ldastats_1.nclasses
    nc2= ldastats_2.nclasses
    d = fill(0.0,nc1, nc2)
    for j in 1:nc2
        for i in 1:nc1
            d[i,j] = ldastats_1.cmeans[:,i] - ldastats_2.cmeans[:,j]
        end
    end

    d_true = diag(d)
    d_nearest = minimum(d, dims=2)
    d_nearest./d_true
end

"""
    get_isolation_distance(Z::Matrix{T}, label::Vector{T2}) where T <: Real where T2

For each point in `Z` compute the distance to its nearest neightbour in its own class and compare that to the its nearest neighbour outside its class.
"""
function get_isolation_distance(Z::Matrix{T}, label::Vector{T2}) where T <: Real where T2
    nc = maximum(label)
    d_in = fill(Inf, nc)
    d_out = fill(Inf, nc)
    np = size(Z,2)
    for i in 1:np
        li = label[i]
        dii = d_in[li]
        doi = d_out[li]
        for j in 1:np
            if i == j
                continue
            end
            lj = label[j]
            d = sum(x->x*x, Z[:,i] - Z[:,j])
            if li == lj
                dii = min(dii, d)
            else
                doi = min(doi, d)
            end
        end
        d_in[li] = dii
        d_out[li] = doi
    end
    d_out./d_in
end

"""
    et_isolation_distance(Z1::Matrix{T}, label1::Vector{T2}, Z2::Matrix{T}, label2::Vector{T2}) where T <: Real where T2

Compute the stimilarly of two clusterings by comparing the nearest point between identical clusters to nearest points between clusters.

A large value means that points for that class in one set are closer to the points of the same class in the other, compared to points in other classes. Intuitively, a large value means that the clustering generalizes across the two sets, while a small value means that the classes are mixed between the two sets.
"""
function get_isolation_distance(Z1::Matrix{T}, label1::Vector{T2}, Z2::Matrix{T}, label2::Vector{T2}) where T <: Real where T2
    nc = maximum(label1)
    nc == maximum(label2) || error("Cluster labels should match between the two inputs")

    np1 = size(Z1, 2)
    np2 = size(Z2, 2)

    n1 = StatsBase.countmap(label1)
    n2 = StatsBase.countmap(label2)
    d_in = fill(0.0, nc)
    d_out = fill(Inf,np1*np2,nc)
    k = 1
    for i in 1:np1
        li = label1[i]
        for j in 1:np2
            lj = label2[j]
            d = sum(x->x*x, Z1[:,i] - Z2[:,j])
            if li == lj
                d_in[li] += d # add to the mean for this class
            else
                d_out[k,li] = d
            end
            k += 1
        end
    end
    d_in ./= [n1[k]*n2[k] for k in 1:nc]
    sort!(d_out, dims=1)
    d_om = fill(0.0, nc)
    for i in 1:nc
        nc2 = 0
        for j in 1:nc
            if i!=j
                nc2 += n2[j]
            end
        end
        d_om[i] = mean(d_out[1:n1[i]*nc2 ,i])
    end
    d_om./d_in
end

"""
    get_cluster_isolation(Z1::Matrix{T}, label1::Vector{T2}, Z2::Matrix{T}, label2::Vector{T2}) where T <: Real where T2

For each cluster, return ratio of the number of points `Z1` that are closer to a point in the same cluster in `Z2` to the number of points that are closer to a different cluster in `Z2`.
"""
function get_cluster_isolation(Z1::Matrix{T}, label1::Vector{T2}, Z2::Matrix{T}, label2::Vector{T2}) where T <: Real where T2
    nc = maximum(label1)
    nc == maximum(label2) || error("Cluster labels should match between the two inputs")

    np1 = size(Z1, 2)
    np2 = size(Z2, 2)

    n1 = StatsBase.countmap(label1)
    n2 = StatsBase.countmap(label2)
    n_in = fill(0, nc)
    n_out = fill(0,nc)
    for i in 1:np1
        li = label1[i]
        d_in = Inf
        d_out = Inf
        for j in 1:np2
            lj = label2[j]
            d = sum(x->x*x, Z1[:,i] - Z2[:,j])
            if lj == li
                d_in = min(d_in, d)
            else
                d_out = min(d_out,d)
            end
        end
        if d_in < d_out
            n_in[li] += 1
        else
            n_out[li] += 1
        end
    end
    q = (n_in - n_out)./(n_in + n_out)
    #normalize from 0 to 1
    q = (q .+ 1)./2
end

end # module


text:
#NOTE: LogNormal_Models reuse Normal_Records

struct LogNormal_Model <: GMC_NS_Model
    trajectory::Integer #id integer assigned at construction
    i::Integer #id of parent model, if any

    θ::Vector{Float64} #model's parameter Vector
    log_Li::Float64 #model's likelihood, calculated by constructor

    pos::Vector{Float64}
    v::Vector{Float64} #model's velocity in parameter space
end

function construct_lognormal_model(trajectory::Integer, i::Integer, θ::Vector{Float64}, pos::Vector{Float64}, v::Vector{Float64}, obs::Vector{Float64}; v_init=false)
    μ,λ=θ
    lμ=log(μ^2/sqrt(μ^2 + inv(λ)))
    lσ=sqrt(log(1+(inv(λ)/μ^2)))
    mod_lnormal=LogNormal(lμ,lσ)
    log_lh=lps(logpdf.(mod_lnormal, obs))
    v_init && (v=rand(MvNormal(length(θ),1.)))

    LogNormal_Model(trajectory, i, θ, log_lh, pos, v)
end

function Base.show(io::IO, m::LogNormal_Model; xsteps=100)
    μ,λ=m.θ
    lμ=log(μ^2/sqrt(μ^2 + inv(λ)))
    lσ=sqrt(log(1+(inv(λ)/μ^2)))
    n=LogNormal(lμ,lσ)
    X=[quantile(n,.025):(quantile(n,.975)-quantile(n,.025))/xsteps:quantile(n,.975)...]
    y=pdf.(n,X)

    show(io, lineplot(X,y,xlabel="X",ylabel="p",title=title="LogNormal Model $(m.trajectory).$(m.i), log_Li $(m.log_Li)"))
    println()
    println("θ: $(m.θ)")
    println("v: $(m.v)")
end


text:
using Documenter, DataDrivenDiffEq

ENV["GKSwstype"] = "100"

makedocs(
    sitename="DataDrivenDiffEq.jl",
    authors="Julius Martensen, Chris Rackauckas",
    modules=[DataDrivenDiffEq],
    clean=true,doctest=false,
    format = Documenter.HTML(#analytics = "UA-90474609-3",
                             assets = ["assets/favicon.ico"],
                             canonical="https://datadriven.sciml.ai/stable/"),
    pages=[
        "Home" => "index.md",
        "Getting Started" => "quickstart.md",
        "Basis" => "basis.md",
        "Koopman Operators" => Any[
            "koopman/koopman.md",
            "koopman/dmd.md",
            "koopman/edmd.md",
            "koopman/dmdc.md",
            "koopman/algorithms.md"
        ],
        "Sparse Identification" => Any[
            "sparse_identification/sindy.md",
            "sparse_identification/isindy.md",
            "sparse_identification/optimizers.md"
        ],
        "Utilities" => "utils.md",
        "Contributing" => "contributions.md",
        "Extended Examples" => "extended_examples.md"
    ]
)

deploydocs(
   repo = "github.com/SciML/DataDrivenDiffEq.jl.git";
   push_preview = true
)


text:
include("backward_linear_steps.jl")

"""
compute the gradients (∇) of the parameters (master_cache) of the constructed model with respect to the cost of prediciton (Ŷ) in comparison with actual output (Y).
"""
function back_propagate_model_weights(Ŷ, Y, master_cache)
    # initiate the dictionary to store the gradients for all the components in each layer
    ∇ = Dict()

    L = length(master_cache)
    Y = reshape(Y, size(Ŷ))

    # Partial derivate of the output layer
    δŶ = - (Y / Ŷ) + ((1 .- Y) / (1 .- Ŷ))
    current_cache = master_cache[L]

    # backpropagate on the layer preceeding the output layer
    ∇[string("δW_", (L))], ∇[string("δb_", (L))], ∇[string("δA_", (L-1))] = linear_activation_backwards(δŶ, current_cache, activation_function = "softmax")

    # go backward in the layers and compute the partial derivates of each component
    for l = reverse(0:L-2)
        current_cache = master_cache[l+1]
        ∇[string("δW_", (l+1))], ∇[string("δb_", (l+1))], ∇[string("δA_", (l))] = linear_activation_backwards(∇[string("δA_", (l+1))], current_cache, activation_function="tanh")
    end # for

    return ∇
end # function

max_stars_repo_name: Squalm/NeuralNetworks

max_issues_repo_name: Squalm/NeuralNetworks

max_forks_repo_name: Squalm/NeuralNetworks


text:
using Base.Test
include("/opt/julia-0.6/share/julia/test/testenv.jl")
addprocs_with_testenv(3)
@test nworkers() == 3

using LShapedSolvers
using JuMP
using StochasticPrograms
using Gurobi

solver = GurobiSolver(OutputFlag=0)
dlsolvers = [(LShapedSolver(:dls,solver,κ=1.0,log=false),"L-Shaped"),
             (LShapedSolver(:drd,solver,crash=Crash.EVP(),autotune=true,κ=0.8,log=false),"RD L-Shaped"),
             (LShapedSolver(:dtr,solver,crash=Crash.EVP(),autotune=true,κ=0.8,log=false),"TR L-Shaped"),
             (LShapedSolver(:dlv,solver,crash=Crash.EVP(),λ=0.6,κ=0.8,log=false),"Leveled L-Shaped")]
lsolvers = [(LShapedSolver(:ls,solver,log=false),"L-Shaped"),
            (LShapedSolver(:rd,solver,crash=Crash.EVP(),autotune=true,log=false),"RD L-Shaped"),
            (LShapedSolver(:tr,solver,crash=Crash.EVP(),autotune=true,log=false),"TR L-Shaped"),
            (LShapedSolver(:lv,solver,crash=Crash.EVP(),λ=0.6,log=false),"Leveled L-Shaped")]

problems = Vector{Tuple{JuMP.Model,String}}()
info("Loading test problems...")
info("Loading simple...")
include("simple.jl")
info("Loading farmer...")
include("farmer.jl")

info("Test problems loaded. Starting test sequence.")
@testset "Distributed $lsname Solver with Distributed Data: $name" for (lsolver,lsname) in dlsolvers, (sp,name) in problems
    solve(sp,solver=solver)
    x̄ = copy(sp.colVal)
    Q̄ = copy(sp.objVal)
    solve(sp,solver=lsolver)
    @test norm(optimal_decision(sp) - x̄) <= 1e-2
    @test abs(optimal_value(sp) - Q̄) <= 1e-2
end

@testset "Distributed $lsname Solver: $name" for (lsolver,lsname) in dlsolvers, (sp,name) in problems
    sp_nondist = StochasticProgram(first_stage_data(sp),second_stage_data(sp),scenarios(sp),procs=[1])
    transfer_model!(stochastic(sp_nondist),stochastic(sp))
    generate!(sp_nondist)
    solve(sp_nondist,solver=solver)
    x̄ = copy(sp_nondist.colVal)
    Q̄ = copy(sp_nondist.objVal)
    solve(sp_nondist,solver=lsolver)
    @test norm(sp_nondist.colVal - x̄) <= 1e-2
    @test abs(sp_nondist.objVal - Q̄) <= 1e-2
end

@testset "$lsname Solver with Distributed Data: $name" for (lsolver,lsname) in lsolvers, (sp,name) in problems
    solve(sp,solver=solver)
    x̄ = copy(sp.colVal)
    Q̄ = copy(sp.objVal)
    solve(sp,solver=lsolver)
    @test norm(sp.colVal - x̄) <= 1e-2
    @test abs(sp.objVal - Q̄) <= 1e-2
end


text:
# scatter markers

#Jl Plots
using Plots

markers = filter((m->begin
                m in Plots.supported_markers()
            end), Plots._shape_keys)
markers = permutedims(markers)
n = length(markers)
x = (range(0, stop = 12, length = n + 2))[2:end - 1]
y = repeat(reshape(reverse(x), 1, :), n, 1)
Plots.scatter(x, y, m = markers, markersize = 6, lab = map(string, markers), bg = :honeydew, xlim = (0, 12), ylim = (0, 12))

# makie markers

using CairoMakie

markers_labels = [
    (:rect, ":rect"),
    (:star5, ":star5"),
    (:diamond, ":diamond"),
    (:hexagon, ":hexagon"),
    (:cross, ":cross"),
    (:xcross, ":xcross"),
    (:utriangle, ":utriangle"),
    (:dtriangle, ":dtriangle"),
    (:ltriangle, ":ltriangle"),
    (:rtriangle, ":rtriangle"),
    (:pentagon, ":pentagon"),
    (:star4, ":star4"),
    (:star8, ":star8"),
    (:vline, ":vline"),
    (:hline, ":hline"),
    (:x, ":x"),
    (:+, ":+"),
    (:circle, ":circle"),
    ('a', "'a'"),
    ('B', "'B'"),
    ('↑', "'\\uparrow'"),
    ('😄', "'\\:smile:'"),
    ('✈', "'\\:airplane:'"),
]

f = Figure()
ax = Axis(f[1, 1], yreversed = true,
    xautolimitmargin = (0.15, 0.15),
    yautolimitmargin = (0.15, 0.15)
)

for (i, (marker, label)) in enumerate(markers_labels)
    p = Point2f(fldmod1(i, 6)...)

    CairoMakie.scatter!(p, marker = marker, markersize = 20, color = :black)
    text!(label, position = p, color = :gray70, offset = (0, 20),
        align = (:center, :bottom))
end

max_stars_repo_name: Ellon-M/visualizations

max_issues_repo_name: Ellon-M/visualizations

max_forks_repo_name: Ellon-M/visualizations


text:
push!(LOAD_PATH, "./src")
using ParametricOptInterface, MathOptInterface, Ipopt

const POI = ParametricOptInterface
const MOI = MathOptInterface

MOI.Benchmarks.@add_benchmark function add_constrained_variable(new_model)
    model = new_model()
    for _ in 1:10_000
        MOI.add_constrained_variable(model, POI.Parameter(0))
    end
    return
end

MOI.Benchmarks.@add_benchmark function add_constraint_svf(new_model)
    model = new_model()
    x = MOI.add_variables(model, 10_000)
    for i in 1:10_000
        MOI.add_constraint(
            model,
            MOI.SingleVariable(x[i]),
            MOI.GreaterThan(1.0),
        )
    end
    return
end

MOI.Benchmarks.@add_benchmark function add_constraint_saf_1(new_model)
    model = new_model()
    x = MOI.add_variables(model, 10_000)
    y, _ = MOI.add_constrained_variable(model, POI.Parameter(0))
    for i in 1:10_000
        MOI.add_constraint(
            model,
            MOI.ScalarAffineFunction(
                MOI.ScalarAffineTerm.(ones(103), [x[i]; x[1:100]; y; y]),
                0.0,
            ),
            MOI.GreaterThan(1.0),
        )
    end
    return
end

MOI.Benchmarks.@add_benchmark function add_constraint_saf_2(new_model)
    model = new_model()
    x = MOI.add_variables(model, 10_000)
    y, _ = MOI.add_constrained_variable(model, POI.Parameter(0))
    for i in 1:10_000
        MOI.add_constraint(
            model,
            MOI.ScalarAffineFunction(
                MOI.ScalarAffineTerm.([1.0, 0.5], [x[i], y]),
                0.0,
            ),
            MOI.GreaterThan(1.0),
        )
    end
    return
end

MOI.Benchmarks.@add_benchmark function add_constraint_saf_3(new_model)
    model = new_model()
    x = MOI.add_variables(model, 10_000)
    for i in 1:10_000
        MOI.add_constraint(
            model,
            MOI.ScalarAffineFunction(
                MOI.ScalarAffineTerm.([1.0, 0.5], [x[i], x[1]]),
                0.0,
            ),
            MOI.GreaterThan(1.0),
        )
    end
    return
end

MOI.Benchmarks.@add_benchmark function add_constraint_sqf_1(new_model)
    model = new_model()
    x = MOI.add_variables(model, 10_000)
    for i in 1:10_000
        MOI.add_constraint(
            model,
            MOI.ScalarQuadraticFunction(
                MOI.ScalarAffineTerm.([1.0, 0.5], [x[i], x[1]]),
                [MOI.ScalarQuadraticTerm(1.0, x[1], x[i])],
                0.0,
            ),
            MOI.GreaterThan(1.0),
        )
    end
    return
end

MOI.Benchmarks.@add_benchmark function add_constraint_sqf_2(new_model)
    model = new_model()
    x = MOI.add_variables(model, 10_000)
    y, _ = MOI.add_constrained_variable(model, POI.Parameter(0))
    for i in 1:10_000
        MOI.add_constraint(
            model,
            MOI.ScalarQuadraticFunction(
                MOI.ScalarAffineTerm.([1.0, 0.5], [x[i], x[1]]),
                [MOI.ScalarQuadraticTerm(1.0, y, x[i])],
                0.0,
            ),
            MOI.GreaterThan(1.0),
        )
    end
    return
end

MOI.Benchmarks.@add_benchmark function add_constraint_sqf_3(new_model)
    model = new_model()
    x = MOI.add_variables(model, 10_000)
    y, _ = MOI.add_constrained_variable(model, POI.Parameter(0))
    for i in 1:10_000
        MOI.add_constraint(
            model,
            MOI.ScalarQuadraticFunction(
                MOI.ScalarAffineTerm.([1.0, 0.5], [x[i], x[1]]),
                [MOI.ScalarQuadraticTerm(1.0, x[i], y)],
                0.0,
            ),
            MOI.GreaterThan(1.0),
        )
    end
    return
end

MOI.Benchmarks.@add_benchmark function add_constraint_sqf_3(new_model)
    model = new_model()
    x = MOI.add_variables(model, 10_000)
    y, _ = MOI.add_constrained_variable(model, POI.Parameter(0))
    for i in 1:10_000
        MOI.add_constraint(
            model,
            MOI.ScalarQuadraticFunction(
                MOI.ScalarAffineTerm.([1.0, 0.5], [x[i], x[1]]),
                [MOI.ScalarQuadraticTerm(1.0, y, y)],
                0.0,
            ),
            MOI.GreaterThan(1.0),
        )
    end
    return
end

function print_help()
    return println("""
           Usage
               benchmark.jl [arg] [name]
           [arg]
               --new       Begin a new benchmark comparison
               --compare   Run another benchmark and compare to existing
           [name]          A name for the benchmark test
           Examples
               git checkout master
               julia benchmark.jl --new master
               git checkout approach_1
               julia benchmark.jl --new approach_1
               git checkout approach_2
               julia benchmark.jl --compare master
               julia benchmark.jl --compare approach_1
           """)
end

if length(ARGS) != 2
    print_help()
else
    const Benchmarks = MOI.Benchmarks
    const suite = Benchmarks.suite(
        () -> POI.Optimizer(Ipopt.Optimizer());
        exclude = [r"delete", r"copy"],
    )
    if ARGS[1] == "--new"
        Benchmarks.create_baseline(
            suite,
            ARGS[2];
            directory = @__DIR__,
            verbose = true,
        )
    elseif ARGS[1] == "--compare"
        Benchmarks.compare_against_baseline(
            suite,
            ARGS[2];
            directory = @__DIR__,
            verbose = true,
        )
    else
        print_help()
    end
end


text:
# This file is a part of JuliaFEM.
# License is MIT: see https://github.com/JuliaFEM/BoundingSphere.jl/blob/master/LICENSE

function welzl!(pts, bdry::BoundaryDevice, alg::WelzlMTF)
    bdry_len = length(bdry)

    support_count = 0

    ball = get_ball(bdry)
    if ismaxlength(bdry)
        support_count = 0
        return ball, support_count
    end

    for i in eachindex(pts)
        pt = pts[i]
        if !isinside(pt, ball)
            pts_i = prefix(pts, i-1)
            isstable = push_if_stable!(bdry, pt)
            if isstable
                ball, s = welzl!(pts_i, bdry, alg)
                @assert isinside(pt, ball, rtol=1e-2, atol=1e-10)
                pop!(bdry)
                move_to_front!(pts, i)
                support_count = s + 1
            end
        end
    end

    @assert bdry_len == length(bdry)
    ball, support_count
end

function find_max_excess(ball, pts, k1)
    T = eltype(first(pts))
    e_max = T(-Inf)
    k_max = k1 -1
    for k in k1:length(pts)
        pt = pts[k]
        e = sqdist(pt, center(ball)) - sqradius(ball)
        if  e > e_max
            e_max = e
            k_max = k
        end
    end
    e_max, k_max
end

function welzl!(pts, bdry, alg::WelzlPivot)
    t = 1
    alg_inner = WelzlMTF()
    @assert isempty(bdry)
    ball, s = welzl!(prefix(pts,t), bdry, alg_inner)
    for i in 1:alg.max_iterations
        @assert s <= t
        e, k = find_max_excess(ball, pts, t+1)

        P = eltype(pts)
        F = eltype(P)
        if e >  eps(F)  # TODO should this be a parameter of the algorithm?
            @assert t < k
            pt = pts[k]
            push_if_stable!(bdry, pt)
            ball_new, s_new = welzl!(prefix(pts,s), bdry, alg_inner)
            # @assert isinside(pt, ball_new, rtol=1e-6)
            pop!(bdry)
            @assert isempty(bdry)
            move_to_front!(pts,k)
            ball = ball_new
            t = s + 1
            s = s_new + 1
        else
            return ball
        end
    end
    return ball
end


text:
# if any of these data are modified make sure to run data/build_data.jl

# gibbs_test_name = "$(name)_samples.csv"
gibbs_test_samples = 1000000

# formulation_test_name = "$(name)_$(formulation)_learned.csv"
#formulations = [:RISE, :logRISE, :RPLE]
formulations = Dict(
    "RISE" => RISE(),
    "logRISE" => logRISE(),
    "RPLE" => RPLE()
)

gms = Dict(
    "a" => FactorGraph([
        0.0 0.1 0.2;
        0.1 0.0 0.3;
        0.2 0.3 0.0
    ]),
    "b" => FactorGraph([
        0.3 0.1 0.2;
        0.1 0.2 0.3;
        0.2 0.3 0.1
    ]),
    "c" => FactorGraph([
        0.0 0.1 0.2 0.3;
        0.1 0.0 0.2 0.3;
        0.2 0.2 0.0 0.3;
        0.3 0.3 0.3 0.0
    ])
)

mutable struct AccuracyTest
   formulation::GMLFormulation
   samples::Int
   threshold::Float64
end


text:
normal_kl_divergence(layer::Any) = 0
scale_mixture_kl_divergence(layer::Any) = 0

function normal_kl_divergence(layer::AbstractTrainableDist)
    posterior = DistributionsAD.TuringMvNormal(layer.mean, abs.(layer.stddev))
    posterior_loglikelihood = DistributionsAD.loglikelihood(posterior, layer.sample)

    # Calculate prior loglikelihood
    μ2 = layer.device(zeros(prod(layer.shape)))
    σ2 = layer.device(ones(prod(layer.shape)))
    prior = DistributionsAD.TuringMvNormal(μ2, σ2)
    prior_loglikelihood = DistributionsAD.loglikelihood(prior, layer.sample)

    return (posterior_loglikelihood - prior_loglikelihood)
end

function scale_mixture_kl_divergence(
    layer::AbstractTrainableDist;
    pi = 0.5,
    prior_sigma1 = 1.0,
    prior_sigma2 = 0.0025,
)
    posterior = DistributionsAD.TuringMvNormal(layer.mean, abs.(layer.stddev))
    posterior_loglikelihood = DistributionsAD.loglikelihood(posterior, layer.sample)

    μ = layer.device(zeros(prod(layer.shape)))
    σ1 = layer.device(ones(prod(layer.shape)) .* prior_sigma1)
    σ2 = layer.device(ones(prod(layer.shape)) .* prior_sigma2)
    prior = DistributionsAD.TuringMvNormal(μ, σ1)
    likelihood_n1 = exp.(DistributionsAD.loglikelihood(prior, layer.sample))
    prior = DistributionsAD.TuringMvNormal(μ, σ2)
    likelihood_n2 = exp.(DistributionsAD.loglikelihood(prior, layer.sample))
    p_scalemixture = pi * likelihood_n1 + (1 - pi) * likelihood_n2
    prior_loglikelihood = log(p_scalemixture)

    return (posterior_loglikelihood - prior_loglikelihood)
end


text:
# ============================= #
# ===                       === #
# === Predictive Likelihood === #
# ===                       === #
# ============================= #

"""
Computes:

    log p ( x | x_1, ..., x_m, θ, z_k)

where `x` is a new spike, {x_1, ..., x_m} are a set of spikes that are
currently assigned to a latent event `z_k`, and `θ` are global
parameters (neuron offsets, etc.).
"""
function log_posterior_predictive(event::SeqEvent, x::Spike, model::SeqModel)

    # Only valid for non-empty events.
    @assert event.spike_count > 0

    # Shorter names for convenience.
    n = x.neuron
    log_p_neuron = model.globals.neuron_response_log_proportions
    offsets = model.globals.neuron_response_offsets
    widths = model.globals.neuron_response_widths
    warps = model.priors.warp_values

    # Grab size (R,W) matrix (already pre-allocated)
    log_prob = model._RW_buffer

    # Compute p({x1...xm} | r)
    for r = 1:num_sequence_types(model)
        for w in 1:num_warp_values(model)

            m = event.summed_potentials[r, w]
            v = event.summed_precisions[r, w]
            μ = m / v
            σ2 = 1 / v

            log_prob[r, w] = (
                event.seq_type_posterior[r, w]   # p( r, w | x1 ... xm)
                + log_p_neuron[n, r]             # p( n | r)
                + normlogpdf(                    # ∫ p( t | s) p(s | x1...xm, r) dτ
                    μ + offsets[n, r] * warps[w],
                    sqrt(σ2 + widths[n, r] * (warps[w] ^ 2)),
                    x.timestamp)
            )

            # See pg. 41 of Gelman et al., Bayesian Data Analysis 3rd Ed. for
            # computing integral ∫ p( t | s) p(s | μ, σ2, r) dτ
        end
    end

    # Return log p(x | x1...xm, Θ)
    return logsumexp(log_prob)
end

"""
Computes:

    log p( {x} | θ )

Where `{x}` is a singleton cluster containing one spike `x`
and `θ` are global parameters (neuron offsets, etc.),
"""
function log_posterior_predictive(x::Spike, model::SeqModel)

    R = num_sequence_types(model)
    W = num_warp_values(model)
    n = x.neuron

    log_p_neuron = model.globals.neuron_response_log_proportions
    log_p_seqtype = model.globals.seq_type_log_proportions
    log_p_warps = model.priors.warp_log_proportions

    # Grab length-R array
    log_prob = model._R_buffer

    # Compute probability of singleton: p(x) = sum_r ( p(n | r) p(r) )
    # Note: no need to sum over warps as they do not affect the likelihood
    # of a spike at time t on neuron n.
    for r = 1:R
        log_prob[r] = log_p_seqtype[r] + log_p_neuron[n, r]
    end

    return logsumexp(log_prob) - log(model.max_time)

end

# =============================================== #
# ===                                         === #
# === Prior distribution on global parameters === #
# ===                                         === #
# =============================================== #

"""
Computes:

    log p({θ_r} | ξ)

where {θ_r} is the set of global parameters for each sequence
type r.
"""
function log_prior(model::SeqModel)

    priors = model.priors
    globals = model.globals

    R = priors.seq_type_proportions.dim  # num sequence types
    N = priors.bkgd_proportions.dim      # num neurons

    # Rate of background spikes.
    lp = logpdf(
        priors.bkgd_amplitude,           # gamma distribution.
        globals.bkgd_amplitude           # nonnegative rate parameter.
    )

    # Background spike proportions across neurons.
    lp += logpdf(
        priors.bkgd_proportions,            # dirichlet distribution.
        exp.(globals.bkgd_log_proportions)  # probability vector.
    )

    # Add contribution of each sequence type.
    for r = 1:R

        # Log prior on neuron amplitudes for each sequence type.
        lp += logpdf(
            priors.neuron_response_proportions,                  # dirichlet distribution.
            exp.(globals.neuron_response_log_proportions[:, r])  # probability vector.
        )

        # Neuron offsets and widths for each sequence type.
        for n = 1:N
            lp += logpdf(
                priors.neuron_response_profile,         # norm-inv-chisq distribution.
                globals.neuron_response_offsets[n, r],  # mean parameter.
                globals.neuron_response_widths[n, r]    # variance parameter.
            )
        end

    end

    return lp
end

# ==================================================== #
# ===                                              === #
# === Log probability density of the latent events === #
# ===                                              === #
# ==================================================== #

"""
Computes:

    log p({z_k} | {θ_r}, η)

where {z_k} is the current set of latent events, {θ_r}
is the set of global parameters for each sequence type r.
"""
function log_p_latents(model::SeqModel)

    priors = model.priors
    globals = model.globals
    sequence_events = model.sequence_events

    lp = 0.0
    α = priors.seq_event_amplitude.α
    β = priors.seq_event_amplitude.β

    # # === Log probability of the number and size of partitions === #
    #
    # # Add log V(t) term, in notation of Miller & Harrison (2018).
    # lp += log_p_num_events(model, num_sequence_events(model))
    #
    # # Add contribution of cluster sizes
    # lp -= lgamma(α) * num_sequence_events(model)
    #
    # for event in sequence_events
    #     lp += lgamma(event.spike_count + α)
    # end
    #
    # # === Log probability of sequence event types and amplitudes === #

    # For each latent event add...
    for event in sequence_events

        # the log prior on event amplitude.
        lp += logpdf(
            priors.seq_event_amplitude,  # gamma distribution.
            event.sampled_amplitude      # nonnegative rate parameter.
        )

        # the log probability of event type.
        lp += globals.seq_type_log_proportions[event.sampled_type]

        # the log probability of event warp.
        lp += priors.warp_log_proportions[event.sampled_warp]

    end

    return lp
end

# ====================== #
# ===                === #
# === Log Likelihood === #
# ===                === #
# ====================== #

"""
Computes the log likelihood of observed events, conditioned
on the current set of latent events and global parameters.

    log p({x_s} | {z_k}, {θ_r})
"""
function log_like(
        model::SeqModel,
        spikes::Vector{Spike}
    )

    max_time = model.max_time
    globals = model.globals
    sequence_events = model.sequence_events

    bkgd_amplitude = globals.bkgd_amplitude
    offsets = globals.neuron_response_offsets
    widths = globals.neuron_response_widths
    warps = model.priors.warp_values
    ll = 0.0

    # == FIRST TERM == #
    # -- Sum of Poisson Process intensity at all spikes -- #

    for x in spikes

        # Compute intensity.
        g = bkgd_amplitude * exp(globals.bkgd_log_proportions[x.neuron])
        for event in sequence_events
            w = warps[event.sampled_warp]
            μ = event.sampled_timestamp + offsets[x.neuron, event.sampled_type] * w
            σ2 = widths[x.neuron, event.sampled_type] * (w ^ 2)
            g += event.sampled_amplitude * normpdf(μ, sqrt(σ2), x.timestamp)
        end

        # Add term to log-likelihood.
        ll += log(g)
    end

    # == SECOND TERM == #
    # -- Penalty on integrated intensity function -- #

    ll -= bkgd_amplitude * max_time
    for event in sequence_events
        ll -= event.sampled_amplitude
    end

    return ll

end

# ================================ #
# ===                          === #
# === Per Event Log Likelihood === #
# ===                          === #
# ================================ #

function event_log_like(
        model::SeqModel,
        spikes::AbstractVector{Spike}
    )

    # Compute likelihood of initial spike.
    ll = log_posterior_predictive(spikes[1], model)

    # Return early if one spike.
    length(spikes) == 1 && return ll

    # Create singleton cluster (at index k).
    k = add_event!(model, spikes[1])
    event = model.sequence_events[k]

    # Account for remaining spikes.
    i = 2
    while true

        # Compute likelihood, conditioned on spikes <i.
        ll += log_posterior_predictive(event, spikes[i], model)

        # Return if we reached the last spike.
        if i == length(spikes)
            remove_event!(model.sequence_events, k)
            return ll
        end

        # Otherwise, add spike to event to update new posterior.
        add_datapoint!(model, spikes[i], k)

        # Increment to next spike.
        i += 1
    end

end