text:
# Copyright (c) 2021 California Institute of Technology (“Caltech”).
# U.S. Government sponsorship acknowledged.
#
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# - Redistributions of source code must retain the above copyright notice, this
#   list of conditions and the following disclaimer.
#
# - Redistributions in binary form must reproduce the above copyright notice,
#   this list of conditions and the following disclaimer in the documentation
#   and/or other materials provided with the distribution.
#
# - Neither the name of Caltech nor its operating division, the Jet Propulsion
#   Laboratory, nor the names of its contributors may be used to endorse or
#   promote products derived from this software without specific prior written
#   permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS
# IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO,
# THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
# OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
# WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR
# OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
# ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

@doc raw"""
This module implements functions to create your own fabulous basis functions,
    primarily for the analysis of galaxy redshift surveys.
"""
module CryoFaBs

export CryoFaB, AngularCryoFaB, AngRadCryoFaB
export isinsurvey
export estimate_density_contrast
export angular_cryo_window
export cryobin

using LinearAlgebra
using BlockDiagonals
using Healpix
using Statistics
using SparseArrays
using SpecialFunctions
using ThreadsX

include("Splines.jl")
using .Splines

using TimerOutputs

#include("SciPy.jl")
#using .SciPy
using FastTransforms
#using SphericalHarmonics

using Scratch
cachedir_greensmatrix = ""
function __init__()
    global cachedir_greensmatrix = @get_scratch!("greensmatrix")
end

################### CryoFaB basis ######################

@doc raw"""
The abstract type `CryoFaB` defines a common interface for the bases defined in
        this module.

            δnlm = cfb \ δθϕr[:]
"""
abstract type CryoFaB end  # Create Your Own Fabulous Basis
# TODO: CryoFaB should probably be an AbstractMatrix.

struct AngularCryoFaB <: CryoFaB
    # define angular mask
    mask::HealpixMap{Float64}
    nside::Int
    npixsurvey::Int
    hpix_full2survey::Array{Int,1}  # surveypixel = hpix_full2survey[fullskypixel]

    # transform
    ell::Vector{Float64}
    Trans::Matrix{Float64}
    TransInv::Matrix{Float64}
end

struct AngRadCryoFaB <: CryoFaB  # CryoFaB that separates angular and radial basis
    # angular transform
    angfab::AngularCryoFaB
    Angl::Vector{Float64}
    AngTrans::BlockDiagonal{Float64}
    AngTransInv::BlockDiagonal{Float64}

    sortbyl::Array{Int,1}  # to transpose the input vector between angular and radial transforms

    # define radial boundary
    rmin::Float64
    rmax::Float64
    rbounds::AbstractArray{Float64,1}  # bin boundaries
    reff::AbstractArray{Float64,1}  # effective center of bin

    # radial transform
    RadK::Array{Float64,1}
    RadTrans::BlockDiagonal{Float64}
    RadTransInv::BlockDiagonal{Float64}
end

include("io.jl")
include("cryofunks_angular.jl")
include("cryofunks_radial.jl")
include("cryowindow_angular.jl")
include("cryowindow_radial.jl")
include("cryowindow_sfb.jl")

function AngularCryoFaB(mask; nside_hires=npix2nside(length(mask)), method=:cryofunk)
    if !(typeof(mask) <: HealpixMap)
        mask = HealpixMap{Float64,Healpix.RingOrder}(mask)
    end
    nside = npix2nside(length(mask))
    npix = nside2npix(nside)
    npixsurvey = sum(mask .> 0)
    hpix_full2survey = fill(0, npix)
    @show typeof(hpix_full2survey) size(hpix_full2survey) npix
    p = 1
    for i=1:npix
        if mask[i] > 0
            hpix_full2survey[i] = p
            p += 1
        end
    end
    if method == :cryofunk
        ℓ, Y, Yinv = get_angular_cryofunks(mask; nside_hires=nside_hires)
    elseif method == :integerLM
        ℓ, Y, Yinv = get_angular_continuous_cryofunks(mask; nside_hires=nside_hires)
    elseif method == :continuousfunks
        ℓ, = get_angular_cryofunks(mask; nside_hires=nside_hires)
        Y, Yinv = get_angular_continuous_cryofunks(ℓ, mask; nside_hires=nside_hires)
    end
    AngularCryoFaB(mask, nside, npixsurvey, hpix_full2survey, ℓ, Y, Yinv)
end

function block_vector(v, n)
    bv = eltype(v)[]
    for i=1:n
        append!(bv, v)
    end
    return bv
end

function AngRadCryoFaB(mask, rmin, rmax, nbins)
    # angular part
    angfab = AngularCryoFaB(mask)
    Angℓ = block_vector(angfab.ell, nbins)
    AngTrans = BlockDiagonal(fill(angfab.Trans, nbins))
    AngTransInv = BlockDiagonal(fill(angfab.TransInv, nbins))

    # radial part
    Δr = (rmax - rmin) / nbins
    reff = range(rmin + Δr/2, rmax - Δr/2, length=nbins)
    kn = []
    gn = Matrix{Float64}[]
    gninv = Matrix{Float64}[]
    @timeit "radial cryofunk" for l=1:length(angfab.ell)
        k, g, ginv = get_radial_cryofunks(angfab.ell[l], reff, Δr)
        append!(kn, k)
        push!(gn, g)
        push!(gninv, ginv)
    end
    RadK = collect(kn)
    RadTrans = BlockDiagonal(gn)
    RadTransInv = BlockDiagonal(gninv)

    #sortbyl = sortperm(Angℓ)  # No! Need to keep the order of shells for equal ℓ the same.
    sortbyl = Int[]
    for iLM=1:length(angfab.ell)
        for ik=1:nbins
            j = (ik - 1) * angfab.npixsurvey + iLM
            push!(sortbyl, j)
        end
    end

    rbounds = range(rmin, rmax, length=nbins+1)

    AngRadCryoFaB(angfab, Angℓ, AngTrans, AngTransInv, sortbyl,
        rmin, rmax, rbounds, reff, RadK, RadTrans, RadTransInv)
end

Base.size(cfb::AngularCryoFaB) = size(cfb.Trans)
Base.size(cfb::AngularCryoFaB, i) = size(cfb.Trans,i)

Base.size(cfb::AngRadCryoFaB) = size(cfb.RadTrans,1), size(cfb.AngTrans,2)
Base.size(cfb::AngRadCryoFaB, i) = (i==1) ? size(cfb.RadTrans,1) : size(cfb.AngTrans,2)

function isinsurvey(cfb::AngularCryoFaB, θ, ϕ)
    p = ang2pix(cfb.mask, θ, ϕ)
    p = cfb.hpix_full2survey[p]
    if p == 0
        return false
    end
    return true
end

function coord2cell(cfb::AngularCryoFaB, θ, ϕ)
    p = ang2pix(cfb.mask, θ, ϕ)
    p = cfb.hpix_full2survey[p]
    if p == 0
        error("angle outside of mask: (θ,ϕ)=($θ,$ϕ)")
    end
    return p
end

function cell2coord(cfb::AngularCryoFaB, p)
    error("unimplemented")
end

function mode2idx(cfb::AngularCryoFaB, l)
    error("unimplemented")
end

function idx2mode(cfb::AngularCryoFaB, i)
    return cfb.ell[i]
end

function isinsurvey(cfb::AngRadCryoFaB, r, θ, ϕ)
    if !isinsurvey(cfb.angfab, θ, ϕ)
        return false
    end

    rbounds = get_rbounds(cfb)
    i = (r == cfb.rmax) ? length(cfb.reff) : searchsortedlast(rbounds, r)
    if !(1 <= i <= length(cfb.reff))
        return false
    end
    return true
end

function coord2cell(cfb::AngRadCryoFaB, r, θ, ϕ)
    p = coord2cell(cfb.angfab, θ, ϕ)

    rbounds = get_rbounds(cfb)
    i = (r == cfb.rmax) ? length(cfb.reff) : searchsortedlast(rbounds, r)
    if !(1 <= i <= length(cfb.reff))
        error("radial dist outside of survey: (r,θ,ϕ)=($r,$θ,$ϕ), $i, $(length(cfb.reff))")
    end

    v = p + (i - 1) * cfb.angfab.npixsurvey
    if v > size(cfb,2)
        @error r,θ,ϕ p i v size(cfb) cfb.angfab.npixsurvey
    end
    return v
end

function cell2coord(cfb::AngRadCryoFaB, v)
    error("unimplemented")
end

function idx2mode(cfb::AngRadCryoFaB, i)
    return cfb.RadK[i], cfb.Angl[cfb.sortbyl][i]
end

function mode2idx(cfb::AngRadCryoFaB, k, l)
    error("unimplemented")
end

Base.:\(cfb::AngularCryoFaB, v) = cfb.TransInv * v
Base.:*(cfb::AngularCryoFaB, v) = cfb.Trans * v
Base.adjoint(cfb::AngularCryoFaB) = cfb.TransInv'
Base.inv(cfb::AngularCryoFaB) = cfb.TransInv

Base.:\(cfb::AngRadCryoFaB, v::AbstractVector) = begin
    cfb.RadTransInv * (cfb.AngTransInv * v)[cfb.sortbyl]
end

Base.:*(cfb::AngRadCryoFaB, v::AbstractVector) = begin
    error("unimplemented")
end

function get_pixs_at_r(cfb::AngRadCryoFaB, i)
    return (1:cfb.angfab.npixsurvey) .+ (i-1) * cfb.angfab.npixsurvey
end

function get_rbounds(cfb::AngRadCryoFaB)
    range(cfb.rmin, cfb.rmax, length=length(cfb.reff)+1)
end

################### estimate density contrast ##########

@doc raw"""
calc_numgals_per_cell(rθϕ, cfb::CryoFaB)

Calculate the number of galaxies in `rθϕ` in each pixel or voxel in the survey
that `cfb` can act on. The radial coordinates are in `rθϕ[1,:]`, etc.
"""
function calc_numgals_per_cell_v1(rθϕ, cfb::CryoFaB)
    T = Float64  # need to allow floats for easy conversion to a density contrast
    N = fill(T(0), size(cfb,1))
    for i=1:size(rθϕ,2)
        p = coord2cell(cfb, rθϕ[:,i]...)
        #p = coord2cell(cfb, rθϕ[1,i], rθϕ[2,i], rθϕ[3,i])
        N[p] += 1
    end
    return N
end
function calc_numgals_per_cell(rθϕ, cfb::AngRadCryoFaB)
    T = Float64  # need to allow floats for easy conversion to a density contrast
    N = fill(T(0), size(cfb,1))
    #@time pix1 = coord2cell.(Ref(cfb), rθϕ[1,:], rθϕ[2,:], rθϕ[3,:])
    #@time pix2 = map((x,y,z)->coord2cell(cfb, x, y, z), rθϕ[1,:], rθϕ[2,:], rθϕ[3,:])
    @time pix3 = ThreadsX.map((r,θ,ϕ)->coord2cell(cfb, r, θ, ϕ), rθϕ[1,:], rθϕ[2,:], rθϕ[3,:])
    #@time pix4 = @strided coord2cell.(Ref(cfb), rθϕ[1,:], rθϕ[2,:], rθϕ[3,:])
    #@assert pix1 == pix2
    #@assert pix1 == pix3
    #@assert pix1 == pix4
    @time for p in pix3
        N[p] += 1
    end
    return N
end

function numgals2densitycontrast!(δ, cfb::AngRadCryoFaB)
    rbounds = get_rbounds(cfb)
    for i=2:length(rbounds)
        p = get_pixs_at_r(cfb, i-1)
        ntotz = sum(δ[p])
        nofz = ntotz / length(p)  # average number of galaxies per pixel
        @. δ[p] = δ[p] / nofz - 1
    end
    return δ
end

function numgals2densitycontrast(δ, cfb::AngRadCryoFaB)
    return numgals2densitycontrast!(deepcopy(δ), cfb)
end

@doc raw"""
estimate_density_contrast(rθϕ, cfb::CryoFaB)

Estimate the density contrast from the galaxy positions in `rθϕ`.
"""
function estimate_density_contrast(rθϕ, cfb::AngRadCryoFaB)
    δ = calc_numgals_per_cell(rθϕ, cfb)
    numgals2densitycontrast!(δ, cfb)
    return δ
end

function calc_angular_average(nr, cfb::CryoFaB)
    rbounds = get_rbounds(cfb)
    nbar = fill(float(eltype(nr)(NaN)), length(rbounds) - 1)
    for i=1:length(rbounds)-1
        p = get_pixs_at_r(cfb, i)
        nbar[i] = mean(nr[p])
    end
    return nbar
end

function cryobin(P, cfb::AngRadCryoFaB; ltol=0.2, ktol=minimum(cfb.RadK), ignore_l=[])
    kk, ll = idx2mode(cfb, 1:length(P))

    # ignore ℓ=0 and ℓ>2*nside
    s = fill(true, length(ll))
    sl = @. cfb.angfab.ell > 1.5 * cfb.angfab.nside
    ignore_l = [ignore_l; cfb.angfab.ell[sl]]
    for l in ignore_l
        @. s &= (ll != l)
    end

    # ignore large k
    Δr = (cfb.rmax - cfb.rmin) / length(cfb.reff)
    @. s &= (kk < 3 / Δr)

    return cryobin(P[s], kk[s], ll[s]; ltol=ltol, ktol=ktol)
end

function cryobin(P, kk, ll; ltol=0.5, ktol=minimum(abs,kk))
    kbin = float(eltype(kk))[]
    lbin = float(eltype(ll))[]
    Pbin = float(eltype(P))[]
    Nbin = Int[]
    while length(kk) >= 1
        k = kk[1]
        l = ll[1]
        sel = @. (abs(kk - k) <= ktol) & (abs(ll - l) <= ltol)
        #@show length(sel),sum(sel)
        #@show k,l kk kk[sel] ll[sel] P[sel]
        sel[1] = true  # include at least one in case 'sel' is false
        push!(kbin, mean(kk[sel]))
        push!(lbin, mean(ll[sel]))
        push!(Pbin, mean(P[sel]))
        push!(Nbin, sum(sel))
        kk = kk[.!sel]
        ll = ll[.!sel]
        P = P[.!sel]
    end
    return kbin, lbin, Pbin, Nbin
end

end

# vim: set sw=4 et sts=4 :


text:
# Write a program that takes the duration of a year (in fractional days) for an imaginary planet as
# an input and produces a leap-year rule that minimizes the difference to the planet’s solar year.

planets = [["Mercury", "Venus", "Earth", "Mars", "Jupiter", "Saturn", "Uranus", "Neptune", "Pluto"],
	  [87.97, 224.70, 365.2422, 1.88, 11.86, 29.46, 84.01, 164.79, 248.59 ]]

function leapYear(num)
    y = 0
    p = 0
    while p < 1
	p = planets[2][num] % 1 * y
        y += 1
    end
  println("$(planets[1][num]) Leap Year Every $(y-2) Years")
end

for i in 1:9
  leapYear(i)
end

ext: jl

max_stars_repo_name: atott/Simple-Programming-Problems-Julia

max_issues_repo_name: atott/Simple-Programming-Problems-Julia

max_forks_repo_name: atott/Simple-Programming-Problems-Julia


text:
include("../src/iterative_lanczos.jl")
using LinearAlgebra
using SparseArrays
using Test

@testset begin
    function Test1(n=5)
        @info "Testing Iterative Lanczos on a ($n,$n) matrix. "
        A = rand(n, n)
        A = A' + A
        b = rand(n)
        il = IterativeLanczos(A, b)
        betas = il(n - 1)
        T = GetTMatrix(il)
        Q = GetQMatrix(il)
        println("The T matrix is: ")
        display(T)
        println("The Q matrix is: ")
        display(Q)
        println("The Q'AQ is: ")
        QAQ = Q'*A*Q
        display(QAQ)
        error = norm(QAQ - T, Inf)
        if error >= 1e-10
            @warn "The inf error of Q'AQ - T is $error, which is lager than expected, please check the code. "
        end
        return error < 1e-10
    end

    function Test2(n=5)
        @info "Testing it with a ($n, $n) Complex Hermitian Matrix"
        A = rand(n, n) + rand(n, n)*im
        A = A' + A
        b = rand(n)
        il = IterativeLanczos(A, b)
        betas = il(n - 1)
        T = GetTMatrix(il)
        Q = GetQMatrix(il)
        println("The T matrix is: ")
        display(T)
        println("The Q matrix is: ")
        display(Q)
        println("The Q'AQ is: ")
        QAQ = Q'*A*Q
        display(QAQ)
        error = norm(QAQ - T, Inf)
        if error >= 1e-10
            @warn "The inf error of Q'AQ - T is $error, which is lager than expected, please check the code. "
        end
        return error < 1e-10
        return true
    end

    function Test3(n=5)
        @info "Testing it with a ($n, $n) Exact Hermitian Matrix"
        A = rand(n, n)
        A = A' + A
        A = convert(Matrix{Rational{BigInt}}, A)
        b = convert(Vector{Rational{BigInt}}, rand(n))
        il = IterativeLanczos(A, b)
        il(n - 1)
        T = GetTMatrix(il)
        Q = GetQMatrix(il)
        println("The T matrix is: ")
        display(T)
        println("The Q matrix is: ")
        display(Q)
        println("The Q'AQ is: ")
        QAQ = Q'*A*Q
        display(QAQ)
        error = norm(QAQ - T, Inf)
        if error >= 1e-10
            @warn "The inf error of Q'AQ - T is $error, which is lager than expected, please check the code. "
        end
        return error < 1e-10
    end

    """
    Testing the fully re-orthgonalization functionality of the algorithm.
    """
    function test4(n=128)
        @info "testin fully re-orthgonalization"
        A = rand(n, n)
        A = A' + A
        A = convert(Matrix{Rational{BigInt}}, A)
        b = convert(Vector{Rational{BigInt}}, rand(n))
        il = IterativeLanczos(A, b)
        il.reorthogonalize = true
        il(n - 1)
        T = GetTMatrix(il)
        Q = GetQMatrix(il)
        println("The T matrix is: ")
        display(T)
        println("The Q matrix is: ")
        display(Q)
        println("The Q'AQ is: ")
        QAQ = Q'*A*Q
        display(QAQ)
        error = norm(QAQ - T, Inf)
        if error >= 1e-10
            @warn "The inf error of Q'AQ - T is $error, which is lager than expected, please check the code. "
        end
        return error < 1e-10
    end

    @test Test1()
    @test Test2()
    @test Test3()
    @test test4()

end


text:
"""
	CBFQP <: Policy

Struct representing a Control Barrier Function based quadratic program (CBFQP).

# Fields
- `control`: function u = control(x) describing the control law.
"""
struct CBFQP <: Policy
    control
end

"""
	(κ::CBFQP)(x)

Solve CBF-QP at state x.
"""
function (κ::CBFQP)(x)
    return κ.control(x)
end

"""
	CBFQP(Σ::ControlAffineSystem, CBF::ControlBarrierFunction)
	CBFQP(Σ::ControlAffineSystem, CBF::ControlBarrierFunction, A, b)
    CBFQP(Σ::ControlAffineSystem, CBF::ControlBarrierFunction, κ::FeedbackPolicy)
	CBFQP(Σ::ControlAffineSystem, CBF::ControlBarrierFunction, κ::FeedbackPolicy, A, b)
	CBFQP(Σ::ControlAffineSystem, CBF::ControlBarrierFunction, κ::CLFQP)
	CBFQP(Σ::ControlAffineSystem, CBF::ControlBarrierFunction, κ::CLFQP, A, b)
	CBFQP(Σ::ControlAffineSystem, CBF::ControlBarrierFunction, CLF::ControlLyapunovFunction, p::Float64=10e2)
	CBFQP(Σ::ControlAffineSystem, CBF::ControlBarrierFunction, CLF::ControlLyapunovFunction, A, b, p::Float64=10e2)
	CBFQP(Σ::ControlAffineSystem, CBFs::Vector{ControlBarrierFunction})
	CBFQP(Σ::ControlAffineSystem, CBFs::Vector{ControlBarrierFunction}, A, b)
	CBFQP(Σ::ControlAffineSystem, CBFs::Vector{ControlBarrierFunction}, κ::FeedbackPolicy)
	CBFQP(Σ::ControlAffineSystem, CBFs::Vector{ControlBarrierFunction}, κ::FeedbackPolicy, A, b)
	CBFQP(Σ::ControlAffineSystem, CBFs::Vector{ControlBarrierFunction}, CLF::ControlLyapunovFunction, p::Float64=10e2)
	CBFQP(Σ::ControlAffineSystem, CBFs::Vector{ControlBarrierFunction}, CLF::ControlLyapunovFunction, A, b, p::Float64=10e2)

Construct a CBF-QP from a ControlAffineSystem and ControlBarrierFunction.
"""
function CBFQP(Σ::ControlAffineSystem, CBF::ControlBarrierFunction)
    function control(x)
        model = Model(OSQP.Optimizer)
        set_silent(model)
        Σ.m == 1 ? @variable(model, u) : @variable(model, u[1:(Σ.m)])
        @objective(model, Min, (1 / 2)u'u)
        @constraint(model, cbf, CBF.∇h(x) * (Σ.f(x) + Σ.g(x) * u) >= -CBF.α(CBF(x)))
        optimize!(model)

        return Σ.m == 1 ? value(u) : value.(u)
    end

    return CBFQP(control)
end

function CBFQP(Σ::ControlAffineSystem, CBF::ControlBarrierFunction, A, b)
    function control(x)
        model = Model(OSQP.Optimizer)
        set_silent(model)
        Σ.m == 1 ? @variable(model, u) : @variable(model, u[1:(Σ.m)])
        @objective(model, Min, (1 / 2)u'u)
        @constraint(model, cbf, CBF.∇h(x) * (Σ.f(x) + Σ.g(x) * u) >= -CBF.α(CBF(x)))
        @constraint(model, A * u .<= b)
        optimize!(model)

        return Σ.m == 1 ? value(u) : value.(u)
    end

    return CBFQP(control)
end

function CBFQP(Σ::ControlAffineSystem, CBF::ControlBarrierFunction, κ::FeedbackPolicy)
    function control(x)
        model = Model(OSQP.Optimizer)
        set_silent(model)
        Σ.m == 1 ? @variable(model, u) : @variable(model, u[1:(Σ.m)])
        @objective(model, Min, (1 / 2)u'u - κ(x)'u)
        @constraint(model, cbf, CBF.∇h(x) * (Σ.f(x) + Σ.g(x) * u) >= -CBF.α(CBF(x)))
        optimize!(model)

        return Σ.m == 1 ? value(u) : value.(u)
    end

    return CBFQP(control)
end

function CBFQP(Σ::ControlAffineSystem, CBF::ControlBarrierFunction, κ::FeedbackPolicy, A, b)
    function control(x)
        model = Model(OSQP.Optimizer)
        set_silent(model)
        Σ.m == 1 ? @variable(model, u) : @variable(model, u[1:(Σ.m)])
        @objective(model, Min, (1 / 2)u'u - κ(x)'u)
        @constraint(model, cbf, CBF.∇h(x) * (Σ.f(x) + Σ.g(x) * u) >= -CBF.α(CBF(x)))
        @constraint(model, A * u .<= b)
        optimize!(model)

        return Σ.m == 1 ? value(u) : value.(u)
    end

    return CBFQP(control)
end

function CBFQP(Σ::ControlAffineSystem, CBF::ControlBarrierFunction, κ::CLFQP)
    function control(x)
        model = Model(OSQP.Optimizer)
        set_silent(model)
        Σ.m == 1 ? @variable(model, u) : @variable(model, u[1:(Σ.m)])
        @objective(model, Min, (1 / 2)u'u - κ(x)'u)
        @constraint(model, cbf, CBF.∇h(x) * (Σ.f(x) + Σ.g(x) * u) >= -CBF.α(CBF(x)))
        optimize!(model)

        return Σ.m == 1 ? value(u) : value.(u)
    end

    return CBFQP(control)
end

function CBFQP(Σ::ControlAffineSystem, CBF::ControlBarrierFunction, κ::CLFQP, A, b)
    function control(x)
        model = Model(OSQP.Optimizer)
        set_silent(model)
        Σ.m == 1 ? @variable(model, u) : @variable(model, u[1:(Σ.m)])
        @objective(model, Min, (1 / 2)u'u - κ(x)'u)
        @constraint(model, cbf, CBF.∇h(x) * (Σ.f(x) + Σ.g(x) * u) >= -CBF.α(CBF(x)))
        @constraint(model, A * u .<= b)
        optimize!(model)

        return Σ.m == 1 ? value(u) : value.(u)
    end

    return CBFQP(control)
end

function CBFQP(
    Σ::ControlAffineSystem,
    CBF::ControlBarrierFunction,
    CLF::ControlLyapunovFunction,
    p::Float64=10e2,
)
    function control(x)
        model = Model(OSQP.Optimizer)
        set_silent(model)
        Σ.m == 1 ? @variable(model, u) : @variable(model, u[1:(Σ.m)])
        @variable(model, δ)
        @objective(model, Min, (1 / 2)u'u + p * δ^2)
        @constraint(model, cbf, CBF.∇h(x) * (Σ.f(x) + Σ.g(x) * u) >= -CBF.α(CBF(x)))
        @constraint(model, clf, CLF.∇V(x) * (Σ.f(x) + Σ.g(x) * u) <= -CLF.α(CLF(x)) + δ)
        optimize!(model)

        return Σ.m == 1 ? value(u) : value.(u)
    end

    return CBFQP(control)
end

function CBFQP(
    Σ::ControlAffineSystem,
    CBF::ControlBarrierFunction,
    CLF::ControlLyapunovFunction,
    A,
    b,
    p::Float64=10e2,
)
    function control(x)
        model = Model(OSQP.Optimizer)
        set_silent(model)
        Σ.m == 1 ? @variable(model, u) : @variable(model, u[1:(Σ.m)])
        @variable(model, δ)
        @objective(model, Min, (1 / 2)u'u + p * δ^2)
        @constraint(model, cbf, CBF.∇h(x) * (Σ.f(x) + Σ.g(x) * u) >= -CBF.α(CBF(x)))
        @constraint(model, clf, CLF.∇V(x) * (Σ.f(x) + Σ.g(x) * u) <= -CLF.α(CLF(x)) + δ)
        @constraint(model, A * u .<= b)
        optimize!(model)

        return Σ.m == 1 ? value(u) : value.(u)
    end

    return CBFQP(control)
end

function CBFQP(Σ::ControlAffineSystem, CBFs::Vector{ControlBarrierFunction})
    function control(x)
        model = Model(OSQP.Optimizer)
        set_silent(model)
        Σ.m == 1 ? @variable(model, u) : @variable(model, u[1:(Σ.m)])
        @objective(model, Min, (1 / 2)u'u)
        for CBF in CBFs
            @constraint(model, CBF.∇h(x) * (Σ.f(x) + Σ.g(x) * u) >= -CBF.α(CBF(x)))
        end
        optimize!(model)

        return Σ.m == 1 ? value(u) : value.(u)
    end

    return CBFQP(control)
end

function CBFQP(Σ::ControlAffineSystem, CBFs::Vector{ControlBarrierFunction}, A, b)
    function control(x)
        model = Model(OSQP.Optimizer)
        set_silent(model)
        Σ.m == 1 ? @variable(model, u) : @variable(model, u[1:(Σ.m)])
        @objective(model, Min, (1 / 2)u'u)
        for CBF in CBFs
            @constraint(model, CBF.∇h(x) * (Σ.f(x) + Σ.g(x) * u) >= -CBF.α(CBF(x)))
        end
        @constraint(model, A * u .<= b)
        optimize!(model)

        return Σ.m == 1 ? value(u) : value.(u)
    end

    return CBFQP(control)
end

function CBFQP(
    Σ::ControlAffineSystem, CBFs::Vector{ControlBarrierFunction}, κ::FeedbackPolicy
)
    function control(x)
        model = Model(OSQP.Optimizer)
        set_silent(model)
        Σ.m == 1 ? @variable(model, u) : @variable(model, u[1:(Σ.m)])
        @objective(model, Min, (1 / 2)u'u - κ(x)'u)
        for CBF in CBFs
            @constraint(model, CBF.∇h(x) * (Σ.f(x) + Σ.g(x) * u) >= -CBF.α(CBF(x)))
        end
        optimize!(model)

        return Σ.m == 1 ? value(u) : value.(u)
    end

    return CBFQP(control)
end

function CBFQP(
    Σ::ControlAffineSystem, CBFs::Vector{ControlBarrierFunction}, κ::FeedbackPolicy, A, b
)
    function control(x)
        model = Model(OSQP.Optimizer)
        set_silent(model)
        Σ.m == 1 ? @variable(model, u) : @variable(model, u[1:(Σ.m)])
        @objective(model, Min, (1 / 2)u'u - κ(x)'u)
        for CBF in CBFs
            @constraint(model, CBF.∇h(x) * (Σ.f(x) + Σ.g(x) * u) >= -CBF.α(CBF(x)))
        end
        @constraint(model, A * u .<= b)
        optimize!(model)

        return Σ.m == 1 ? value(u) : value.(u)
    end

    return CBFQP(control)
end

function CBFQP(
    Σ::ControlAffineSystem,
    CBFs::Vector{ControlBarrierFunction},
    CLF::ControlLyapunovFunction,
    p::Float64=10e2,
)
    function control(x)
        model = Model(OSQP.Optimizer)
        set_silent(model)
        Σ.m == 1 ? @variable(model, u) : @variable(model, u[1:(Σ.m)])
        @variable(model, δ)
        @objective(model, Min, (1 / 2)u'u + p * δ^2)
        for CBF in CBFs
            @constraint(model, CBF.∇h(x) * (Σ.f(x) + Σ.g(x) * u) >= -CBF.α(CBF(x)))
        end
        @constraint(model, clf, CLF.∇V(x) * (Σ.f(x) + Σ.g(x) * u) <= -CLF.α(CLF(x)) + δ)
        optimize!(model)

        return Σ.m == 1 ? value(u) : value.(u)
    end

    return CBFQP(control)
end

function CBFQP(
    Σ::ControlAffineSystem,
    CBFs::Vector{ControlBarrierFunction},
    CLF::ControlLyapunovFunction,
    A,
    b,
    p::Float64=10e2,
)
    function control(x)
        model = Model(OSQP.Optimizer)
        set_silent(model)
        Σ.m == 1 ? @variable(model, u) : @variable(model, u[1:(Σ.m)])
        @variable(model, δ)
        @objective(model, Min, (1 / 2)u'u + p * δ^2)
        for CBF in CBFs
            @constraint(model, CBF.∇h(x) * (Σ.f(x) + Σ.g(x) * u) >= -CBF.α(CBF(x)))
        end
        @constraint(model, clf, CLF.∇V(x) * (Σ.f(x) + Σ.g(x) * u) <= -CLF.α(CLF(x)) + δ)
        @constraint(model, A * u .<= b)
        optimize!(model)

        return Σ.m == 1 ? value(u) : value.(u)
    end

    return CBFQP(control)
end


text:
using Random, LinearAlgebra, Distributions
import Distributions: MvNormal
import LinearAlgebra: cholesky, Symmetric

using SEC

⊗(a, b) = a * b'
Random.seed!(1234)

# Here we go through an example where the theory "works perfectly"

# Define a linear forward map
# input Space
Mo = 2
# output space
M = 2
H = randn(M, Mo)
forward_map(x) = H * x

# Create artificial solution
x = randn(Mo) .- 30
y̅ = forward_map(x)

# Define an initial ensemble distribution
MoMo = randn(Mo, Mo)
Σ = MoMo' * MoMo + I
prior = MvNormal(Σ)
# this choice implies prior ∝ exp( - x' Σ⁻¹ x)

# Implicitly define the likelihood function via the covariance
MM = randn(M, M)
Γ = (MM' * MM + I)
# implies likelihood  exp(- T * dot(y - H * x, Γ⁻¹ * (y - H * x)) )
# where T is the final time of the simulation
# Observe that
# likelihood = exp(- dot(y - H * x, Γ⁻¹ * (y - H * x)) ) for T = 1

# number of steps
N = 100
# timestep size (h = 1/N implies T=1 at the final time)
Δt = 1 / N
# number of ensemble members
J = 4000 * Mo

# Construct Posterior
HΓHΣ = Symmetric((H' * (cholesky(Γ) \ H) + cholesky(Σ) \ I))
cHΓHΣ = cholesky(HΓHΣ) # LU fails
rhs = (H' * (cholesky(Γ) \ y̅[:]))
posterior_μ = cHΓHΣ \ rhs
posterior_Σ = cHΓHΣ \ I
posterior = MvNormal(posterior_μ, posterior_Σ)

# Construct empirical prior
u = [rand(prior) for i = 1:J]
u₀ = copy(u)

timeseries = []
push!(timeseries, copy(u))

# Viewed as a minimization problem, we are minimizing
# dot(y - H * x, Γ⁻¹ * (y - H * x)) + dot(x, Σ⁻¹ x)
# whose solution is
# (H' Γ⁻¹ H + Σ⁻¹)⁻¹ H' Γ⁻¹ y
# Alternatively we can view the posterior as
# posterior ∝ exp(-dot(y - H * x, Γ⁻¹ * (y - H * x)) - x' Σ⁻¹ x)
# which implies μ = (H' Γ⁻¹ H + Σ⁻¹)⁻¹ H' Γ⁻¹ y
# and           σ = (H' Γ⁻¹ H + Σ⁻¹)⁻¹

# EKI Algorithm
for _ = 1:N
    # Likelihood perturbation for algorithm, (constructed from before)
    ξ = MvNormal(1 / Δt * Γ)
    SEC.Calibrate.eki_step!(u, forward_map, y̅, J, ξ, Γ, Δt)
    push!(timeseries, copy(u))
end

# This
approx = mean(u)
if length(approx) < 4
    println("mean appro = ", approx)
end
# should converge to this
exact = posterior.μ
if length(exact) < 4
    println("mean exact = ", exact)
end
println("mean relative error = ", norm(approx - exact) / norm(exact))
println("loss relative error = ", norm(H * approx - H * exact) / norm(H * exact))

exact_cov = posterior.Σ
approx_cov = cov(u)
if length(exact) < 4
    @show approx_cov
    @show exact_cov
end
println("cov relative error = ", norm(approx_cov - exact_cov) / norm(exact_cov))

##
using GLMakie, Printf

fig = Figure()
ax = Axis(fig[2, 1])
sc_init = scatter!(ax, [(timeseries[1][i][1], timeseries[1][i][2]) for i in eachindex(u)], color = :red)
sc_final = scatter!(ax, [(timeseries[end][i][1], timeseries[end][i][2]) for i in eachindex(u)], color = :blue)

time_slider = Slider(fig, range = 1:length(timeseries), startvalue = 1)
ti = time_slider.value

ensemble = @lift [(timeseries[$ti][i][1], timeseries[$ti][i][2]) for i in eachindex(u)]
sc_transition = scatter!(ax, ensemble, color = :purple)

scatter!(ax, [(exact[1], exact[2])], marker = '⋆', color = :yellow, markersize = 30)

ax.xlabel = "c¹"
ax.ylabel = "c²"

time_string = @lift("Time = " * @sprintf("%0.2f", ($ti - 1) / (length(timeseries) - 1)))
fig[3, 1] = vgrid!(
    Label(fig, time_string, width = nothing),
    time_slider,
)

c¹ = @lift([timeseries[$ti][i][1] for i in eachindex(timeseries[end])])
c² = @lift([timeseries[$ti][i][2] for i in eachindex(timeseries[end])])

ax.xlabelsize = 25
ax.ylabelsize = 25

ax_above = Axis(fig[1, 1])
ax_side = Axis(fig[2, 2])

ax_above.ylabel = "probability density"
ax_side.xlabel = "probability density"

limits!(ax, -50, 10, -8, 4)
xlims!(ax_above, ax.limits[][1]...)
ylims!(ax_above, 0, 0.4)
ylims!(ax_side, ax.limits[][2]...)

hideydecorations!(ax_side, ticks = false, grid = false)
hidexdecorations!(ax_above, ticks = false, grid = false)

# colsize!(fig.layout, 1, Relative(2 / 3))
# rowsize!(fig.layout, 1, Relative(1 / 3))
colgap!(fig.layout, 10)
rowgap!(fig.layout, 10)

d1 = density!(ax_above, c¹, color = (:purple, 0.15), label = "density c¹", strokewidth = 1)
d1lims = ax_above.limits[][1][1]:ax_above.limits[][1][2]
tmpx = collect(range(d1lims[1], d1lims[end], length = 1000))
σ = prior.Σ[1, 1]
μ = prior.μ[1]
tmpy = @. 1 / sqrt(2 * π * σ) * exp(-0.5 * (tmpx - μ)^2 / σ)
ln1 = lines!(ax_above, tmpx, tmpy, linewidth = 4, color = :red)
σ = posterior.Σ[1, 1]
μ = posterior.μ[1]
tmpy = @. 1 / sqrt(2 * π * σ) * exp(-0.5 * (tmpx - μ)^2 / σ)
lines!(ax_above, tmpx, tmpy, linewidth = 4, color = :blue)

d2 = density!(ax_side, c², color = (:purple, 0.15), label = "density c²", strokewidth = 1, direction = :y)
dlims = ax_side.limits[][2][1]:ax_side.limits[][2][2]
tmpx = collect(range(dlims[1], dlims[end], length = 1000))
σ = prior.Σ[2, 2]
μ = prior.μ[2]
tmpy = @. 1 / sqrt(2 * π * σ) * exp(-0.5 * (tmpx - μ)^2 / σ)
lines!(ax_side, tmpy, tmpx, linewidth = 4, color = :red)
σ = posterior.Σ[2, 2]
μ = posterior.μ[2]
tmpy = @. 1 / sqrt(2 * π * σ) * exp(-0.5 * (tmpx - μ)^2 / σ)
ln2 = lines!(ax_side, tmpy, tmpx, linewidth = 4, color = :blue)

Legend(fig[1, 2],
    [sc_init, sc_final, sc_transition, ln1, ln2],
    ["prior ensemble", "posterior ensemble", "transition ensemble", "exact marginal prior", "exact marginal posterior"])

#update!(fig.scene)
display(fig)

seconds = 5
fps = 30
frames = round(Int, fps * seconds)
frames = 101
fps = 30
record(fig, pwd() * "/example.mp4"; framerate = fps) do io
    for i = 1:frames
        ti[] = i
        sleep(1 / fps)
        recordframe!(io)
    end
end

max_stars_repo_name: bischtob/SEC

max_issues_repo_name: bischtob/SEC

max_forks_repo_name: bischtob/SEC


text:
struct WeightedLeastSquares
    estimate::Array{Float64,1}
    rmse::Float64
    mae::Float64
    wrss::Float64
end

struct ErrorMetric
    rmse::Float64
    mae::Float64
    wrss::Float64
end

struct ErrorMetricWiden
    rmse::Float64
    mae::Float64
    wrss::Float64
    rmseGBPWLS::Float64
    maeGBPWLS::Float64
end

########## Compute the GBP error metrics ##########
function errorMetric(gbp::Union{ContinuousModel, ContinuousTreeModel})
    observationGBP = gbp.system.jacobian * gbp.inference.mean
    NactiveRows = activeRows(gbp)

    rmse = 0.0; mae = 0.0; wrss = 0.0
    @inbounds for i = 1:length(gbp.system.observation)
        rmse += (gbp.system.observation[i] - observationGBP[i])^2
        mae += abs(gbp.system.observation[i] - observationGBP[i])
        wrss += (gbp.system.observation[i] - observationGBP[i]) / gbp.system.variance[i]
    end
    rmse = (rmse / NactiveRows)^(1/2)
    mae = mae / NactiveRows

    return ErrorMetric(rmse, mae, wrss)
end

########## Compute the GBP error metrics and error according to the WLS ##########
@inline function errorMetric(gbp::Union{ContinuousModel, ContinuousTreeModel}, wls::WeightedLeastSquares)
    observationGBP = gbp.system.jacobian * gbp.inference.mean
    NactiveRows = activeRows(gbp)

    rmse = 0.0; mae = 0.0; wrss = 0.0
    @inbounds for i = 1:length(gbp.system.observation)
        rmse += (gbp.system.observation[i] - observationGBP[i])^2
        mae += abs(gbp.system.observation[i] - observationGBP[i])
        wrss += (gbp.system.observation[i] - observationGBP[i]) / gbp.system.variance[i]
    end
    rmse = (rmse / NactiveRows)^(1/2)
    mae = mae / NactiveRows

    rmseGBPWLS = 0.0; maeGBPWLS = 0.0
    @inbounds for i = 1:gbp.graph.Nvariable
        rmseGBPWLS += (wls.estimate[i] - gbp.inference.mean[i])^2
        maeGBPWLS += abs(wls.estimate[i] - gbp.inference.mean[i])
    end
    rmseGBPWLS = (rmseGBPWLS / gbp.graph.Nvariable)^(1/2)
    maeGBPWLS = maeGBPWLS / gbp.graph.Nvariable

    return ErrorMetricWiden(rmse, mae, wrss, rmseGBPWLS, maeGBPWLS)
end

########## Compute WLS solution and error metrics ##########
function wls(gbp::Union{ContinuousModel, ContinuousTreeModel})
    W = spdiagm(0 =>  @. 1.0 / sqrt(gbp.system.variance))
    A = W * gbp.system.jacobian
    G = A' * A
    b = A' * W * gbp.system.observation
    x = G \ b

    observationWLS = gbp.system.jacobian * x
    NactiveRows = activeRows(gbp)

    rmse = 0.0; mae = 0.0; wrss = 0.0
    @inbounds for i = 1:length(gbp.system.observation)
        rmse += (gbp.system.observation[i] - observationWLS[i])^2
        mae += abs(gbp.system.observation[i] - observationWLS[i])
        wrss += (gbp.system.observation[i] - observationWLS[i]) / gbp.system.variance[i]
    end
    rmse = (rmse / NactiveRows)^(1/2)
    mae = mae / NactiveRows

    return WeightedLeastSquares(x, rmse, mae, wrss)
end

########## Display: GBP and WLS ##########
function displayData(args...)
    gbpIdx = 0; wlsIdx = 0; errorIdx = 0; widenIdx = 0
    @inbounds for (k, i) in enumerate(args)
        T = split.(string(typeof(i)), ".")[end]
        if T == "ContinuousModel" || T == "ContinuousTreeModel"
            gbpIdx = k
        end
        if T == "WeightedLeastSquares"
            wlsIdx = k
        end
        if T == "ErrorMetric" || T == "ErrorMetricWiden"
            errorIdx = k
        end
    end
    if gbpIdx == 0
       error("FactorGraph model is missing.")
    end

    if errorIdx != 0 && wlsIdx == 0
        data = ["" "" ""
                "Error Metrics" "" ""
                "Root mean square error of the Gaussian belief propagation" "."^10 getfield(args[errorIdx], :rmse);
                "Mean absolute error of the Gaussian belief propagation" "."^10 getfield(args[errorIdx], :mae);
                "Weighted residual sum of squares error of the Gaussian belief propagation" "."^10 getfield(args[errorIdx], :wrss)]

        pretty_table(data, tf = tf_borderless, columns_width = [75, 11, 14],
        noheader = true, alignment = [:l, :l, :r], formatters = ft_printf("%1.4e", 3),
        highlighters = (hl_cell( [(2,1)], crayon"bold"), hl_col(2, crayon"dark_gray")),
        body_hlines = [2], body_hlines_format = Tuple('─' for _ = 1:4))
    end

    if errorIdx != 0 && wlsIdx != 0
        data = ["" "" ""
                "Error Metrics" "" ""
                "Root mean square error of the Gaussian belief propagation" "."^10 getfield(args[errorIdx], :rmse);
                "Root mean square error of the weighted least-squares method" "."^10  getfield(args[wlsIdx], :rmse);
                "Root mean square error ratio" "."^10 getfield(args[errorIdx], :rmse) / getfield(args[wlsIdx], :rmse);
                "" "" "";
                "Mean absolute error of the Gaussian belief propagation" "."^10 getfield(args[errorIdx], :mae);
                "Mean absolute error of the weighted least-squares method" "."^10 getfield(args[wlsIdx], :mae);
                "Mean absolute error ratio" "."^10 getfield(args[errorIdx], :mae) / getfield(args[wlsIdx], :mae);
                "" "" "";
                "Weighted residual sum of squares of the belief propagation" "."^10 getfield(args[errorIdx], :wrss);
                "Weighted residual sum of squares of the weighted least-squares method" "."^10 getfield(args[wlsIdx], :wrss);
                "Weighted residual sum of squares ratio" "."^10 getfield(args[errorIdx], :wrss) / getfield(args[wlsIdx], :wrss)]

        high = Highlighter((data,i,j) -> i == 5 && j == 3 && data[5,3] > 1.1, Crayon(bold = true, background = :red))
        pretty_table(data, tf = tf_borderless, columns_width = [75, 11, 14],
        noheader = true, alignment = [:l, :l, :r], formatters = ft_printf("%1.4e", 3),
        highlighters = (high, hl_cell([(2,1), (8,3), (5,3)], crayon"bold"), hl_col(2, crayon"dark_gray")),
        body_hlines = [2], body_hlines_format = Tuple('─' for _ = 1:4))
    end

    if wlsIdx == 0
        println("\n State Variables Display")
        pretty_table([collect(1:getfield(args[gbpIdx].graph, :Nvariable)) getfield(args[gbpIdx].inference, :mean)], header = ["State Variable", "Estimate"],
        alignment=[:r, :r], formatters = ft_printf(["%3.0f", "%3.6f"], [1, 2]))
    end

    if wlsIdx != 0
        A = [collect(1:getfield(args[gbpIdx].graph, :Nvariable)) getfield(args[gbpIdx].inference, :mean) getfield(args[wlsIdx], :estimate) getfield(args[gbpIdx].inference, :mean) ./ getfield(args[wlsIdx], :estimate)]
        A = A[reverse(sortperm(A[:, 4])),  :]

        println("\n State Variables Display")
        pretty_table(A, header = ["State Variable", "Belief Propagation Estimate", "Weighted Least-squares Estimate", "Estimates Ratio"],
        columns_width = [15, 28, 32, 20], alignment=[:r, :r, :r, :r], formatters = ft_printf(["%3.0f", "%3.6f", "%3.6f", "%3.6e"], [1, 2, 3, 4]))
    end
end

########## Find number of nonzeros rows ##########
@inline function activeRows(gbp)
    NactiveRows = length(gbp.system.observation)
    @inbounds for (k, i) in enumerate(gbp.system.observation)
        if i == 0.0
            if all(gbp.system.jacobianTranspose[:, k] .== 0.0)
                NactiveRows -= 1
            end
        end
    end

    return NactiveRows
end

########## Start row in xlsx-file ##########
function startxlsx(xf)
    start = 1
    @inbounds for r in XLSX.eachrow(xf)
        if !isa(r[1], String)
            start = XLSX.row_number(r)
            break
        end
    end
    return start
end

########## Check data for import ##########
function checkImportFile(args)
    #### Check the package is installed
    pathtoFactorGraph = Base.find_package("FactorGraph")
    if isnothing(pathtoFactorGraph)
        throw(ErrorException("FactorGraph not found in install packages"))
    end
    packagepath = abspath(joinpath(dirname(pathtoFactorGraph), ".."))

    extension = ".h5"; path = ""; dataname = ""; fullpath = ""

    @inbounds for i = 1:length(args)
        try
            extension = string(match(r"\.[A-Za-z0-9]+$", args[i]).match)
        catch
            extension = ""
        end
        if extension == ".h5" || extension == ".xlsx"
            fullpath = args[i]
            path = dirname(args[i])
            dataname = basename(args[i])
            break
        end
    end

    if isempty(extension)
        throw(ErrorException("the input DATA extension is not found"))
    elseif extension != ".h5" && extension != ".xlsx"
        throw(DomainError(extension, "the input DATA extension is not supported"))
    end

    if path == ""
        path = joinpath(packagepath, "src/example/")
        fullpath = joinpath(packagepath, "src/example/", dataname)
    end

    if !(dataname in cd(readdir, path))
        throw(DomainError(dataname, "the input DATA is not found"))
    end

    return fullpath, extension, dataname
end

########## Check keyword arguments ##########
function checkKeywords(prob, alpha, variance)
    #### Check convergence parameters
    if prob <= 0.0 || prob >= 1.0
        error("Invalid prob value.")
    end
    if alpha <= 0.0 || alpha >= 1.0
        error("Invalid alpha value.")
    end

    #### Check initial variance
    if variance < 0.0
        error("Invalid variance value.")
    end
end

########## Type of the system data ##########
function checkFileOrArguments(args)
    fromfile = false
    @inbounds for i in args
        if typeof(i) == String
            fromfile = true
            break
        end
    end

    return fromfile
end


text:
using MAT
using Random
include("../src/hypergraph-affinity-functions.jl")
D = matread("DBLP-Author-Gender_H.mat")
H = D["H"]
classes = D["class"]

n = size(H,2)
r = 4
alpha, B1, B2, H1, H2, R1, R2, N = Hypergraph_to_Scores(H,classes,r)
class1 = "female"
class2 = "male"

## Bootstrap
alpha, B1, B2, H1, H2, tR1, tR2, N = Hypergraph_to_Scores(H,classes,4)

for k =  2:4
B = 100
T1 = tR1[k,1:k]  # True affinity scores
T0 = tR2[k,1:k]

TA1 = H1[k,1:k]  # True affinity scores
TA0 = H2[k,1:k]

# Bootstraps
R0, R1, MR1, MR0, SR1, SR0, A0, A1, MA1, MA0, SA1, SA0 = Bootstrap_Affinities(H,k,B,classes)

# E1 = vec(maximum(R1,dims = 1))
# E1m = vec(minimum(R1,dims = 1))
# E0 = vec(maximum(R0,dims = 1))
# E0m = vec(minimum(R0,dims = 1))

## Plot affinity scores
lw = 2.5
ms = 8
gfs = 15
tfs = 15
titlesize = 14
s1 = 350
s2 = 300
sc = 2
color1 = :green
color2 = :blue
title = ""
leg = false
xlab = "Affinity type t"
ylab = "Affinity"
p = plot()
p = Plots.plot(fontfamily = "helvetica",linewidth = 2,yscale = :identity,legend = false,grid = false, size = (s1,s2),color = :gray)
title!(p,title,titlefontsize = titlesize)
xaxis!(p,xlab, xticks = 1:k,tickfontsize=tfs,guidefontsize = gfs)
yaxis!(p,ylab,tickfontsize=tfs,guidefontsize = gfs)

# [M0-2*SE0 M0+2*SE]
# [E0m E0]
# [E1m E1]
p = plot!(1:k,[TA0 TA0], fillrange=[MA0-sc*SA0 MA0+sc*SA0], fillalpha=0.3, color = color2, linewidth = lw, markersize = ms,
    markershape = :circle, markerstrokecolor = color2, label = class2)
p = plot!(1:k,[TA1 TA1], fillrange=[MA1-sc*SA1 MA1+sc*SA1],fillalpha=0.3,color = color1, linewidth = lw, markersize = ms,
    markershape = :circle,markerstrokecolor = color1, label = class1)

savefig("Figures/Affinities_$(k)_DBLP_Bootstrap.pdf")

## Ratio scores

xlab = "Affinity type t"
ylab = "Affinity / Baseline"
p = plot()
p = Plots.plot([1; k],[1; 1],fontfamily = "helvetica",linewidth = 2,yscale = :identity,legend = false,grid = false, size = (s1,s2),color = :gray)
title!(p,title,titlefontsize = titlesize)
xaxis!(p,xlab, xticks = 1:k,tickfontsize=tfs,guidefontsize = gfs)
yaxis!(p,ylab,tickfontsize=tfs,guidefontsize = gfs)

p = plot!(1:k,[T0 T0], fillrange=[MR0-sc*SR0 MR0+sc*SR0], fillalpha=0.3, color = color2, linewidth = lw, markersize = ms,
    markershape = :circle, markerstrokecolor = color2, label = class2)
p = plot!(1:k,[T1 T1], fillrange=[MR1-sc*SR1 MR1+sc*SR1],fillalpha=0.3,color = color1, linewidth = lw, markersize = ms,
    markershape = :circle,markerstrokecolor = color1, label = class1)
savefig("Figures/Ratios_$(k)_DBLP_Bootstrap.pdf")
end

max_stars_repo_name: nveldt/HypergraphHomophily

max_issues_repo_name: nveldt/HypergraphHomophily

max_forks_repo_name: nveldt/HypergraphHomophily


text:
test_date = Date(2012, 1 ,1)
test_datetime = DateTime(2012, 1, 1, 1, 1, 1, 1)
test_time = Time(1, 1, 1, 1)

@testset "InfExtendedTime" begin
    @testset "Base" begin
        @inferred InfExtendedTime{Date}(test_date)
        @inferred InfExtendedTime{Date}(∞)
        @inferred InfExtendedTime{Date}(-∞)
        @inferred InfExtendedTime{Time}(test_time)
        @inferred InfExtendedTime{Time}(∞)
        @inferred InfExtendedTime{Time}(-∞)
        @inferred InfExtendedTime{DateTime}(test_datetime)
        @inferred InfExtendedTime{DateTime}(∞)
        @inferred InfExtendedTime{DateTime}(-∞)
        @test_throws MethodError InfExtendedTime{Time}(test_date)
        @test_throws MethodError InfExtendedTime{Date}(test_time)
        @test_throws MethodError InfExtendedTime{DateTime}(test_time)

        d = InfExtendedTime{Date}(test_date)
        t = InfExtendedTime{Time}(test_time)
        dt = InfExtendedTime{DateTime}(test_datetime)
        @test InfExtendedTime{Date}(dt) == d
        @test InfExtendedTime{Date}(test_datetime) == d
        @test InfExtendedTime{Time}(dt) == t
        @test InfExtendedTime{Time}(test_datetime) == t
        @test InfExtendedTime{DateTime}(dt) == dt

        @test InfExtendedTime(Date) == InfExtendedTime{Date}
        @test InfExtendedTime(Time) == InfExtendedTime{Time}
        @test InfExtendedTime(DateTime) == InfExtendedTime(DateTime)

        @test InfExtendedTime(test_date) == InfExtendedTime{Date}(test_date)
        @test InfExtendedTime(test_time) == InfExtendedTime{Time}(test_time)
        @test InfExtendedTime(test_datetime) == InfExtendedTime{DateTime}(test_datetime)

        @test posinf(typeof(d)) == InfExtendedTime{Date}(∞)
        @test neginf(typeof(d)) == InfExtendedTime{Date}(-∞)
        @test isposinf(InfExtendedTime{Date}(∞))
        @test !isposinf(InfExtendedTime{Date}(-∞))
        @test !isposinf(d)
        @test !isneginf(InfExtendedTime{Date}(∞))
        @test isneginf(InfExtendedTime{Date}(-∞))
        @test !isneginf(d)

        @test posinf(typeof(t)) == InfExtendedTime{Time}(∞)
        @test neginf(typeof(t)) == InfExtendedTime{Time}(-∞)
        @test isposinf(InfExtendedTime{Time}(∞))
        @test !isposinf(InfExtendedTime{Time}(-∞))
        @test !isposinf(t)
        @test !isneginf(InfExtendedTime{Time}(∞))
        @test isneginf(InfExtendedTime{Time}(-∞))
        @test !isneginf(t)

        @test posinf(typeof(dt)) == InfExtendedTime{DateTime}(∞)
        @test neginf(typeof(dt)) == InfExtendedTime{DateTime}(-∞)
        @test isposinf(InfExtendedTime{DateTime}(∞))
        @test !isposinf(InfExtendedTime{DateTime}(-∞))
        @test !isposinf(dt)
        @test !isneginf(InfExtendedTime{DateTime}(∞))
        @test isneginf(InfExtendedTime{DateTime}(-∞))
        @test !isneginf(dt)

        inf = InfExtendedTime{Date}(∞)
        @test d.instant == test_date.instant
        @test t.instant == test_time.instant
        @test inf.instant == ∞
    end

    @testset "IO" begin
        d = InfExtendedTime{Date}(test_date)
        i = InfExtendedTime{Date}(∞)

        @test string(d) == "InfExtendedTime{Date}($(repr(test_date)))"
        @test sprint(show, d, context=:compact=>true) == "2012-01-01"
        @test sprint(show, d) == string(d)

        @test string(i) == "InfExtendedTime{Date}(∞)"
        @test sprint(show, i, context=:compact=>true) == "∞"
        @test sprint(show, i) == string(i)
    end

    @testset "Conversion" begin
        @test promote_rule(InfExtendedTime{DateTime}, InfExtendedTime{Date}) ==
            InfExtendedTime{DateTime}
        @test promote_rule(InfExtendedTime{DateTime}, Date) == InfExtendedTime{DateTime}
        @test promote_rule(InfExtendedTime{Date}, Infinite) == InfExtendedTime{Date}
        @test promote_rule(Infinite, Date) == InfExtendedTime{Date}

        @test convert(InfExtendedTime{Date}, test_datetime) ==
            InfExtendedTime{Date}(test_date)
        @test convert(InfExtendedTime{Date}, InfExtendedTime{DateTime}(test_datetime)) ==
            InfExtendedTime{Date}(InfExtendedTime{DateTime}(test_datetime))
        @test convert(InfExtendedTime{Date}, ∞) == InfExtendedTime{Date}(∞)

        @test promote_rule(InfExtendedTime{ZonedDateTime}, ZonedDateTime) ==
            InfExtendedTime{ZonedDateTime}
    end

    @testset "Parse" begin
        val = tryparse(InfExtendedTime{Date}, "∞")
        @test val == ∞
        @test val isa InfExtendedTime{Date}

        val = tryparse(InfExtendedTime{Date}, "2012-01-01")
        @test val == Date(2012, 1, 1)
        @test val isa InfExtendedTime{Date}

        val = tryparse(InfExtendedTime{Date}, "a")
        @test val == nothing
        @test val isa Nothing

        val = parse(InfExtendedTime{Date}, "∞")
        @test val == ∞
        @test val isa InfExtendedTime{Date}

        val = parse(InfExtendedTime{Date}, "2012-01-01")
        @test val == Date(2012, 1, 1)
        @test val isa InfExtendedTime{Date}

        @test_throws ArgumentError parse(InfExtendedTime{Date}, "a")
    end

    @testset "Comparison" begin
        @test isfinite(InfExtendedTime{Date}(test_date))
        @test !isfinite(InfExtendedTime{Date}(∞))

        @test isinf(InfExtendedTime{Date}(∞))
        @test isinf(InfExtendedTime{Date}(-∞))
        @test !isinf(InfExtendedTime{Date}(test_date))

        @test InfExtendedTime{Date}(test_date) == InfExtendedTime{Date}(test_date)
        @test InfExtendedTime{Date}(test_date) != InfExtendedTime{Date}(test_date + Day(1))
        @test InfExtendedTime{Date}(∞) == InfExtendedTime{Date}(∞)
        @test InfExtendedTime{Date}(-∞) == InfExtendedTime{Date}(-∞)
        @test InfExtendedTime{Date}(∞) != InfExtendedTime{Date}(-∞)
        @test InfExtendedTime{Date}(-∞) != InfExtendedTime{Date}(∞)
        @test InfExtendedTime{Date}(∞) != InfExtendedTime{Date}(test_date)
        @test InfExtendedTime{Date}(test_date) != InfExtendedTime{Date}(∞)

        @test hash(InfExtendedTime{Date}(∞)) == hash(InfExtendedTime{Date}(∞))
        @test hash(InfExtendedTime{Date}(-∞)) == hash(InfExtendedTime{Date}(-∞))
        @test hash(InfExtendedTime{Date}(∞)) != hash(InfExtendedTime{Date}(-∞))
        @test hash(InfExtendedTime{Date}(test_date)) ==
            hash(InfExtendedTime{Date}(test_date))
        @test hash(InfExtendedTime{Date}(test_date)) !=
            hash(InfExtendedTime{Date}(test_date + Day(1)))

        d1 = InfExtendedTime(test_date)
        d2 = InfExtendedTime(test_date - Day(1))
        d3 = InfExtendedTime(test_date + Day(1))
        @test d1 == d1
        @test d1 <= d1
        @test d1 != d2
        @test d2 < d1
        @test d1 < d3
        @test d2 < d3
        @test d3 > d2
        @test d3 > d1
        @test d1 >= d1

        inf = InfExtendedTime{Date}(∞)
        ninf = InfExtendedTime{Date}(-∞)
        @test ninf < inf
        @test ninf <= ninf
        @test ninf < inf
        @test inf > ninf
        @test inf >= inf

        @test ninf < d1
        @test d3 < inf
        @test inf > d1
        @test d3 > ninf

        @test ∞ == inf
        @test -∞ == ninf
        @test inf == ∞
        @test ninf == -∞

        @test -∞ < inf
        @test ninf < ∞
        @test inf <= ∞
        @test ninf <= -∞
        @test -∞ <= ninf
    end

    @testset "Arithmetic" begin
        @test typemin(InfExtendedTime{Date}) == InfExtendedTime{Date}(-∞)
        @test typemax(InfExtendedTime{Date}) == InfExtendedTime{Date}(∞)

        @test InfExtendedTime(test_date) + Day(1) == InfExtendedTime(test_date + Day(1))
        @test Day(1) + InfExtendedTime(test_date) == InfExtendedTime(test_date + Day(1))
        @test InfExtendedTime(test_date) - Day(1) == InfExtendedTime(test_date - Day(1))
        @test Day(1) - InfExtendedTime(test_date) == InfExtendedTime(test_date - Day(1))
        @test InfExtendedTime{Date}(∞) + Day(1) == InfExtendedTime{Date}(∞)
        @test InfExtendedTime{Date}(∞) - Day(1) == InfExtendedTime{Date}(∞)
        @test InfExtendedTime{Date}(-∞) + Day(1) == InfExtendedTime{Date}(-∞)
        @test InfExtendedTime{Date}(-∞) - Day(1) == InfExtendedTime{Date}(-∞)

        a = InfExtendedTime{Date}(∞)
        b = InfExtendedTime{Date}(-∞)
        c = InfExtendedTime(test_date)

        @test a + ∞ == ∞ + a == a
        @test_throws InfMinusInfError b + ∞
        @test_throws InfMinusInfError ∞ + b
        @test c + ∞ == ∞ + c == a

        @test_throws InfMinusInfError a - ∞
        @test_throws InfMinusInfError ∞ - a
        @test b - ∞ == b
        @test ∞ - b == a
        @test c - ∞ == b
        @test ∞ - c == a

        @test ∞ + test_date == ∞
        @test test_date + ∞ == ∞
        @test ∞ - test_date == ∞
        @test test_date - ∞ == -∞

        @test ∞ + Day(1) == ∞
        @test Day(1) + ∞ == ∞
        @test ∞ - Day(1) == ∞
        @test Day(1) - ∞ == -∞

        @test ∞ + test_date.instant == ∞
        @test test_date.instant + ∞ == ∞
        @test ∞ - test_date.instant == ∞
        @test test_date.instant - ∞ == -∞
    end
end


text:
# # Minimization of simple loss functions
#
# First we load the required packages.

using Distributions, LinearAlgebra, Random, Plots

using EnsembleKalmanProcesses.EnsembleKalmanProcessModule, EnsembleKalmanProcesses.ParameterDistributionStorage

# ## Loss function with single minimum
#
# Here, we minimize the loss function
# ```math
# G₁(u) = \|u - u_*\| ,
# ```
# where ``u`` is a 2-vector of parameters and ``u_*`` is given; here ``u_* = (-1, 1)``.
u★ = [1, -1]
G₁(u) = [sqrt((u[1] - u★[1])^2 + (u[2] - u★[2])^2)]
nothing # hide

# We set the seed for pseudo-random number generator for reproducibility.
rng_seed = 41
Random.seed!(rng_seed)
nothing # hide

# We set a stabilization level, which can aid the algorithm convergence
dim_output = 1
stabilization_level = 1e-3
Γ_stabilization = stabilization_level * Matrix(I, dim_output, dim_output)

# The functional is positive so to minimize it we may set the target to be 0,
G_target = [0]
nothing # hide

# ### Prior distributions
#
# As we work with a Bayesian method, we define a prior. This will behave like an "initial guess"
# for the likely region of parameter space we expect the solution to live in.
prior_distributions = [Parameterized(Normal(0, 1)), Parameterized(Normal(0, 1))]

constraints = [[no_constraint()], [no_constraint()]]

parameter_names = ["u1", "u2"]

prior = ParameterDistribution(prior_distributions, constraints, parameter_names)

# ### Calibration
#
# We choose the number of ensemble members and the number of iterations of the algorithm
N_ensemble = 20
N_iterations = 10
nothing # hide

# The initial ensemble is constructed by sampling the prior
initial_ensemble = EnsembleKalmanProcessModule.construct_initial_ensemble(prior, N_ensemble; rng_seed = rng_seed)

# We then initialize the Ensemble Kalman Process algorithm, with the initial ensemble, the
# target, the stabilization and the process type (for EKI this is `Inversion`, initialized
# with `Inversion()`).
ensemble_kalman_process =
    EnsembleKalmanProcessModule.EnsembleKalmanProcess(initial_ensemble, G_target, Γ_stabilization, Inversion())

# Then we calibrate by *(i)* obtaining the parameters, *(ii)* calculate the loss function on
# the parameters (and concatenate), and last *(iii)* generate a new set of parameters using
# the model outputs:
for i in 1:N_iterations
    params_i = get_u_final(ensemble_kalman_process)

    g_ens = hcat([G₁(params_i[:, i]) for i in 1:N_ensemble]...)

    EnsembleKalmanProcessModule.update_ensemble!(ensemble_kalman_process, g_ens)
end

# and visualize the results:
u_init = get_u_prior(ensemble_kalman_process)

anim_unique_minimum = @animate for i in 1:N_iterations
    u_i = get_u(ensemble_kalman_process, i)

    plot(
        [u★[1]],
        [u★[2]],
        seriestype = :scatter,
        markershape = :star5,
        markersize = 11,
        markercolor = :red,
        label = "optimum u⋆",
    )

    plot!(
        u_i[1, :],
        u_i[2, :],
        seriestype = :scatter,
        xlims = extrema(u_init[1, :]),
        ylims = extrema(u_init[2, :]),
        xlabel = "u₁",
        ylabel = "u₂",
        markersize = 5,
        markeralpha = 0.6,
        markercolor = :blue,
        label = "particles",
        title = "EKI iteration = " * string(i),
    )
end
nothing # hide

# The results show that the minimizer of ``G_1`` is ``u=u_*``.

gif(anim_unique_minimum, "unique_minimum.gif", fps = 1) # hide

# ## Loss function with two minima

# Now let's do an example in which the loss function has two minima. We minimize the loss
# function
# ```math
# G₂(u) = \|u - v_{*}\| \|u - w_{*}\| ,
# ```
# where again ``u`` is a 2-vector, and ``v_{*}`` and ``w_{*}`` are given 2-vectors. Here, we take ``v_{*} = (1, -1)`` and ``w_{*} = (-1, -1)``.

v★ = [1, -1]
w★ = [-1, -1]
G₂(u) = [sqrt(((u[1] - v★[1])^2 + (u[2] - v★[2])^2) * ((u[1] - w★[1])^2 + (u[2] - w★[2])^2))]
nothing # hide
#
# The procedure is same as the single-minimum example above.

# We set the seed for pseudo-random number generator for reproducibility,
rng_seed = 10
Random.seed!(rng_seed)
nothing # hide

# A positive function can be minimized with a target of 0,
G_target = [0]

# We choose the stabilization as in the single-minimum example

# ### Prior distributions
#
# We define the prior. We can place prior information on e.g., ``u₁``, demonstrating a belief
# that ``u₁`` is more likely to be negative. This can be implemented by setting a bias in the
# mean of its prior distribution to e.g., ``-0.5``:

prior_distributions = [Parameterized(Normal(-0.5, sqrt(2))), Parameterized(Normal(0, sqrt(2)))]

constraints = [[no_constraint()], [no_constraint()]]

parameter_names = ["u1", "u2"]

prior = ParameterDistribution(prior_distributions, constraints, parameter_names)

# ### Calibration
#
# We choose the number of ensemble members, the number of EKI iterations, construct our initial ensemble and the EKI with the `Inversion()` constructor (exactly as in the single-minimum example):
N_ensemble = 20
N_iterations = 20

initial_ensemble = EnsembleKalmanProcessModule.construct_initial_ensemble(prior, N_ensemble; rng_seed = rng_seed)

ensemble_kalman_process =
    EnsembleKalmanProcessModule.EnsembleKalmanProcess(initial_ensemble, G_target, Γ_stabilization, Inversion())

# We calibrate by *(i)* obtaining the parameters, *(ii)* calculating the
# loss function on the parameters (and concatenate), and last *(iii)* generate a new set of
# parameters using the model outputs:
for i in 1:N_iterations
    params_i = get_u_final(ensemble_kalman_process)

    g_ens = hcat([G₂(params_i[:, i]) for i in 1:N_ensemble]...)

    EnsembleKalmanProcessModule.update_ensemble!(ensemble_kalman_process, g_ens)
end

# and visualize the results:
u_init = get_u_prior(ensemble_kalman_process)

anim_two_minima = @animate for i in 1:N_iterations
    u_i = get_u(ensemble_kalman_process, i)

    plot(
        [v★[1]],
        [v★[2]],
        seriestype = :scatter,
        markershape = :star5,
        markersize = 11,
        markercolor = :red,
        label = "optimum v⋆",
    )

    plot!(
        [w★[1]],
        [w★[2]],
        seriestype = :scatter,
        markershape = :star5,
        markersize = 11,
        markercolor = :green,
        label = "optimum w⋆",
    )

    plot!(
        u_i[1, :],
        u_i[2, :],
        seriestype = :scatter,
        xlims = extrema(u_init[1, :]),
        ylims = extrema(u_init[2, :]),
        xlabel = "u₁",
        ylabel = "u₂",
        markersize = 5,
        markeralpha = 0.6,
        markercolor = :blue,
        label = "particles",
        title = "EKI iteration = " * string(i),
    )
end
nothing # hide

# Our bias in the prior shifts the initial ensemble into the negative ``u_1`` direction, and
# thus increases the likelihood (over different instances of the random number generator) of
# finding the minimizer ``u=w_*``.

gif(anim_two_minima, "two_minima.gif", fps = 1) # hide


text:
using CairoMakie, Random
x = 0:0.05:4π
y1 =  sin.(3x) ./ (cos.(x) .+ 2)./x
y2 =  -2.5cos.(x)./x
xl = LinRange(5,10,10)
yl = LinRange(1,0.5,10)

set_theme!(theme_dark())
fig = Figure(resolution = (600, 400))
ax = Axis(fig[1,1])
lines!(x, y1; color = :dodgerblue, label = "blue")
band!(x, y1 .- 0.1, y1 .+ 0.1; color = (:dodgerblue, 0.2), label = "blue")

lines!(x, y2; color = :red, label = "red")
band!(x, y2 .- 0.1, y2 .+ 0.1; color = (:red, 0.2), label = "red")

lines!(xl, yl; color = :slategray1, label = "slategray1")
band!(xl, fill(0.5,10), yl; color= (:slategray1, 0.05),label="slategray1")
ylims!(-0.55,1)
axislegend(ax, position = :rt, merge = true)
hidedecorations!(ax; grid = false)
save("band.png", fig, px_per_unit = 2)
display(fig)


text:
struct Mutation
    location::Int64
    ancestral_value::Float64
    new_value::Float64
end

struct Genotype
    grn::Array{Float64,2}
    fixed_mutations::Array{Mutation,1}
end


text:
"""
The data layout is:
diagonal elements
sub diagonal lower triangular part

The matrix is also padded.
Given avx2, a 7x7 matrix of Float64 would be stored as:
7 diagonal elements + 1 element padding

6x6 lower triangle...
column 1: 6 elements + 2 padding
column 2: 5 elements + 3 padding
column 3: 4 elements + 0 padding
column 4: 3 elements + 1 padding
column 4: 2 elements + 2 padding
column 4: 1 elements + 3 padding
"""
@inline function Base.getindex(S::AbstractSymmetricMatrixL{P,T,L}, i, j) where {P,T,L}
    j, i = minmax(j, i)
    @boundscheck i > P && ThrowBoundsError("max(i, j) = $j > $P.")
    @inbounds S.data[lt_sub2ind(P, i, j)]
end
@inline function Base.getindex(S::AbstractSymmetricMatrixU{P,T,L}, i, j) where {P,T,L}
    i, j = minmax(i, j)
    @boundscheck j > P && ThrowBoundsError("max(i, j) = $j > $P.")
    @inbounds S.data[ut_sub2ind(P, i, j)]
end

@inline function Base.setindex!(S::AbstractMutableSymmetricMatrixL{P,T,L}, v::T, i, j) where {P,T,L}
    j, i = minmax(j, i)
    @boundscheck i > P && ThrowBoundsError("max(i, j) = $j > $P.")
    VectorizationBase.store!(pointer(S) + sizeof(T) * (lt_sub2ind(P, i, j)-1), v)
    @inbounds S.data[lt_sub2ind(P, i, j)]
end
@inline function Base.setindex!(S::AbstractMutableSymmetricMatrixU{P,T,L}, v::T, i, j) where {P,T,L}
    i, j = minmax(i, j)
    @boundscheck j > P && ThrowBoundsError("max(i, j) = $j > $P.")
    VectorizationBase.store!(pointer(S) + sizeof(T) * (ut_sub2ind(P, i, j)-1), v)
end

#@inline function Base.setindex!(S::AbstractMutableSymmetricMatrixL{P,T,L}, v::T, i) where {P,T,L}
#    j, i = minmax(j, i)
#    @boundscheck i > P && ThrowBoundsError("i = $i > $L = L")
#    VectorizationBase.store!(pointer(S) + sizeof(T) * i, v)
#end
#@inline function Base.setindex!(S::AbstractMutableSymmetricMatrixU{P,T,L}, v::T, i) where {P,T,L}
#    @boundscheck i > L && ThrowBoundsError("i = $i > $L = L")
#    VectorizationBase.store!(pointer(S) + sizeof(T) * i, v)
#end

@generated function SymmetricMatrixL(S::PaddedMatrices.AbstractFixedSizeMatrix{M,M,T,P}) where {M,T,P}
    Lbase = binomial2(M+1)
    Wm1 = VectorizationBase.pick_vector_width(Lbase, T) - 1
    Lfull = (Lbase + Wm1) & ~Wm1

    outtup = Expr(:tuple)
    for m ∈ 1:M
        push!(outtup.args, :(S[ $(P*(m-1) + m) ]))
    end
    for mc ∈ 1:M-1
        for mr ∈ mc+1:M
            push!(outtup.args, :( S[$( P*(mc-1) + mr )] ))
        end
    end
    for l ∈ Lbase+1:Lfull
        push!(outtup.args, :($(zero(T))))
    end
    :(@inbounds SymmetricMatrixL{$M,$T,$Lfull}($outtup))
end
function MutableSymmetricMatrixL(S::PaddedMatrices.AbstractFixedSizeMatrix{M,M,T,P}) where {M,T,P}
    Lbase = binomial2(M+1)
#    Wm1 = VectorizationBase.pick_vector_width(Lbase, T) - 1
    #    Lfull = (Lbase + Wm1) & ~Wm1
    Sm = MutableSymmetricMatrixL{M,T,Lbase}(undef)
    @inbounds for m ∈ 1:M
        Sm[m] = S[m,m]
    end
    @inbounds for mc ∈ 1:M-1, mr ∈ mc+1:M
        Sm[mr,mc] = S[mr,mc]
    end
    Sm
end

# @generated function lower_cholesky(Σ::SymmetricMatrixL{P,T,L}) where {P,T,L}
#
# end

### Be more clever than this.
# @generated function quadform(x::AbstractFixedSizeVector{P,T,R}, Σ::SymmetricMatrixL{P,T,L}) where {P,T,R,L}
#     W = pick_vector_width(P, T)
#     q = quote
#         # vout = vbroadcast(Vec{$P,$T}, zero($T))
#         out_1 = zero($T)
#         @vectorize $T for i ∈ 1:$P
#             out_1 += x[i] * Σ[i]
#         end
#     end
#     for p ∈ 2:P
#         out_temp = Symbol(:out_, p)
#         push!(q.args, quote
#             @vectorize $T for i ∈ 1:$P
#                 $out_temp += x[i] * Σ[i]
#             end
#         end)
#     end
#     q
# end

# @generated function Base.:*(A::AbstractFixedSizeMatrix{M,P,T}, Σ::SymmetricMatrix{P,T,L}) where {M,P,T,L}
#     W, Wshift = VectorizationBase.pick_vector_width_shift(M, T)
#
#     quote
#
#     end
# end


text:
#MOBIL.jl
#A separate file for all the MOBIL stuff

###############
##MOBIL Model##
###############

struct MOBILParam <: FieldVector{3, Float64}
	p::Float64 #politeness factor
	b_safe::Float64 #safe braking value
	a_thr::Float64 #minimum accel
	#db::Float64 #lane bias #we follow the symmetric/USA lane change rule
end #MOBILParam
MOBILParam(;p::Float64=0.25,b_safe::Float64=4.,a_thr::Float64=0.2) = MOBILParam(p,b_safe,a_thr)

StaticArrays.similar_type(::Type{MOBILParam}, ::Type{Float64}, ::Size{(3,)}) = MOBILParam

nan(::Type{MOBILParam}) = MOBILParam(NaN, NaN, NaN)
#=
==(a::MOBILParam,b::MOBILParam) = (a.p==b.p) && (a.b_safe==b.b_safe) && (a.a_thr == b.a_thr)
Base.hash(a::MOBILParam,h::UInt64=zero(UInt64)) = hash(a.p,hash(a.b_safe,hash(a.a_thr,h)))

+(a::MOBILParam, b::MOBILParam) = MOBILParam(a.p+b.p, a.b_safe+b.b_safe, a.a_thr+b.a_thr)
-(a::MOBILParam, b::MOBILParam) = a+(-1.0*b)
*(a::Real, p::MOBILParam) = MOBILParam(a*p.p, a*p.b_safe, a*p.a_thr)
.*(p::MOBILParam, v::Vector{Float64}) = MOBILParam(v[1]*p.p, v[2]*p.b_safe, v[3]*p.a_thr)
.^(p::MOBILParam, n::Real) = MOBILParam(p.p^n, p.b_safe^n, p.a_thr^n)
sqrt(p::MOBILParam) = MOBILParam(sqrt(p.p), sqrt(p.b_safe), sqrt(p.a_thr))
zero(::Type{MOBILParam}) = MOBILParam(0.0, 0.0, 0.0)
nan(::Type{MOBILParam}) = MOBILParam(NaN, NaN, NaN)
abs(p::MOBILParam) = MOBILParam(abs(p.p), abs(p.b_safe), abs(p.a_thr))
max(a::MOBILParam, b::MOBILParam) = MOBILParam(max(a.p,b.p), max(a.b_safe, b.b_safe), max(a.a_thr, b.a_thr))
=#

function MOBILParam(s::AbstractString)
	#typical politeness range: [0.0,0.5]
	typedict = Dict{AbstractString,Float64}("cautious"=>0.5,"normal"=>0.25,"aggressive"=>0.0)
	p = get(typedict,s,-1.)
	@assert(p >= 0.)
	return MOBILParam(p=p)
end

function is_lanechange_dangerous(pp::PhysicalParam, s::MLState, nbhd::Array{Int,1}, idx::Int, dir::Real)

	#check if dir is oob
	dt = pp.dt
	lane_ = s.cars[idx].y + dir * dt # XXX this line is the issue!!!

	if (lane_ > pp.nb_lanes) || (lane_ < 1.)
		return true
	end
	#check if will hit car next to you?
	l_car = pp.l_car

	dvlf, slf = get_dv_ds(pp,s,nbhd,idx,round(Int, 5+sign(dir)))
	dvlb, slb = get_dv_ds(pp,s,nbhd,idx,round(Int, 2+sign(dir)))

	#distance at next time step
	#dv ref: ahead: me - him; behind: him - me
	#ds ref: ahead: him - me - l_car; behind: me - him - l_car
	dslf = slf-dvlf*dt
	dslb = slb-dvlb*dt
	diff = 0.0 #something >=0--the safety distance

	return (slf < diff*l_car) || (slb < diff*l_car) || dslb < diff*l_car ||
				dslf < diff*l_car

end

function get_neighborhood(pp::PhysicalParam,s::Union{MLState,MLObs},idx::Int)
	nbhd = zeros(Int,6)
	dists = Inf*ones(6)
	"""
	Each index corresponds to the index (in the car state) that corresponds to
	the position, 0 if there is no such car
	indices:
	6      |     3
	5    |car|   2 ->
	4      |     1
	where 1/2/3 is the front, 3/6 is left

    Note: The same car can occupy two spots if the vehicle is changing lanes.
	"""

	x = s.cars[idx]
	#rightmost lane: no one to the right
	if x.y <= 1.
		dists[[1;4]] = [-1.;-1.]
	#leftmost lane: no one to the left
	elseif x.y >= pp.nb_lanes
		dists[[3;6]] = [-1.;-1.]
	end

	for (i,car) in enumerate(s.cars)
		#i am not a neighbor
		if i == idx
			continue
		end
		pos = car.x
		lane = car.y

        @assert pos >= 0. # this should never happen

		dlane = lane - x.y #NOTE float: convert to int
		#too distant to be a neighbor
		# if abs(dlane) > 2.
		if abs(dlane) >= 2. # changed august 13 - don't remember why this was > 2
			continue
		# elseif abs(dlane) <= 0.5 # this is now handled below in the overlap section
		# 	dlane = 0
		else
			dlane = convert(Int,sign(dlane))
		end

		d = pos-x.x

		offset = d >= 0. ? 2 : 5

		if abs(d) < dists[offset+dlane]
			dists[offset+dlane] = abs(d)
			nbhd[offset+dlane] = i
		end

        if occupation_overlap(car.y, x.y) && abs(d) < dists[offset]
            dists[offset] = abs(d)
            nbhd[offset] = i
        end
	end

	return nbhd
end

function get_dv_ds(pp::PhysicalParam,s::MLState,nbhd::Array{Int,1},idx::Int,idy::Int)
	"""
	idx is the car from whom the perspective is
	idy is the other car
	"""
	car = s.cars[idx]
	nbr = nbhd[idy]
	#dv: if ahead: me - him; behind: him - me
	dv = nbr != 0 ? -1*sign((idy-3.5))*(car.vel - s.cars[nbr].vel) : 0.

	ds = nbr != 0 ? abs(s.cars[nbr].x - car.x) - pp.l_car : 1000.

	return dv::Float64, ds::Float64
end

"""
Return two numbers first if the lane change is made, second if the lane change is not made

(I think) - Zach (7/13/16)
"""
function get_rear_accel(pp::PhysicalParam,s::MLState,nbhd::Array{Int,1},idx::Int,dir::Int)

	#there is no car behind in that spot
	if nbhd[5+dir] == 0
		return 0.::Float64, 0.::Float64
	end

	v = s.cars[idx].vel

	#behind - me
	dv_behind, s_behind = get_dv_ds(pp, s, nbhd, idx, 5+dir)
	v_behind = v + dv_behind

	#me - front
	#what would the relative velocity, distance be if idx wasn't there
	dv_behind_, s_behind_ = get_dv_ds(pp, s, nbhd, idx, 2+dir)
	dv_behind_ += dv_behind
	s_behind_ += s_behind + pp.l_car

	dt = pp.dt
	# if !(typeof(s.cars[nbhd[5+dir]].behavior) <: IDMMOBILBehavior)
	# 	behind_idm = NORMAL_IDM
	# else
	# end
	behind_idm = s.cars[nbhd[5+dir]].behavior.p_idm::IDMParam
	a_follower = get_idm_dv(behind_idm,dt,v_behind,dv_behind,s_behind)/dt #distance behind is a negative number
	a_follower_ = get_idm_dv(behind_idm,dt,v_behind,dv_behind_,s_behind_)/dt

	return a_follower::Float64, a_follower_::Float64
end

# should return a float
function get_mobil_lane_change(behavior, pp::PhysicalParam,s::MLState,nbhd::Array{Int,1},idx::Int,rng::AbstractRNG=MersenneTwister(123))
	#TODO: catch if the parameters don't exist

	#need 6 distances: distance to person behind me, ahead of me
	#				potential distance to person behind me, ahead of me
	#				in other lane(s)
	#need sets of idm parameters
	dt = pp.dt
	state = s.cars[idx]
	p_idm_self = behavior.p_idm
	p_mobil = behavior.p_mobil
	#println(neighborhood)
	if sum(nbhd) == 0
		return 0. #no reason to change lanes if you're all alone
	end

	##if between lanes, return +1 if moving left, -1 if moving right
	if !isinteger(state.y)
		return state.lane_change #continue going in the direction you're going
	end

	v = state.vel
	#get predicted and potential accelerations

	#TODO generalize to get_dv()? # btw this takes a lot of processing time # should we cache?
    dv2, ds2 = get_dv_ds(pp,s,nbhd,idx,2)
    dv3, ds3 = get_dv_ds(pp,s,nbhd,idx,3)
    dv1, ds1 = get_dv_ds(pp,s,nbhd,idx,1)
	a_self = get_idm_dv(p_idm_self,dt,v,dv2,ds2)/dt
	a_self_left = get_idm_dv(p_idm_self,dt,v,dv3,ds3)/dt
	a_self_right = get_idm_dv(p_idm_self,dt,v,dv1,ds1)/dt

	a_follower, a_follower_ = get_rear_accel(pp,s,nbhd,idx,0)
	a_follower_left_, a_follower_left = get_rear_accel(pp,s,nbhd,idx,1)
	a_follower_right_, a_follower_right = get_rear_accel(pp,s,nbhd,idx,-1)

	#calculate incentives
	left_crit = a_self_left-a_self+p_mobil.p*(a_follower_left_-a_follower_left+a_follower_-a_follower)
	right_crit = a_self_right-a_self+p_mobil.p*(a_follower_right_-a_follower_right+a_follower_-a_follower)

	#check safety criterion, also check if there is physically space
	if (a_follower_right_ < -p_mobil.b_safe) && (a_follower_left_ < -p_mobil.b_safe)
		return 0. #neither safe
	end
	if is_lanechange_dangerous(pp,s,nbhd,idx,1) || (a_follower_left_ < -p_mobil.b_safe)
		left_crit = -Inf
	end

	if is_lanechange_dangerous(pp,s,nbhd,idx,-1) || (a_follower_right_ < -p_mobil.b_safe)
		right_crit = -Inf
	end

	#check if going left or right is preferable
    if left_crit >= right_crit
        dir_flag = 1
    else
        dir_flag = -1
    end

	#check incentive criterion
	if max(left_crit,right_crit) > p_mobil.a_thr
		return float(dir_flag)
	end
	return 0.
end


text:
# https://github.com/tbreloff/Qwt.jl

supported_attrs(::QwtBackend) = merge_with_base_supported([
    :annotations,
    :linecolor,
    :fillrange,
    :fillcolor,
    :label,
    :legend,
    :seriescolor, :seriesalpha,
    :linestyle,
    :linewidth,
    :markershape,
    :markercolor,
    :markersize,
    :bins,
    :pos,
    :title,
    :window_title,
    :guide, :lims, :ticks, :scale,
  ])
supported_types(::QwtBackend) = [:path, :scatter, :hexbin, :bar]
supported_markers(::QwtBackend) = [:none, :auto, :rect, :circle, :diamond, :utriangle, :dtriangle, :cross, :xcross, :star5, :star8, :hexagon]
supported_scales(::QwtBackend) = [:identity, :log10]
is_subplot_supported(::QwtBackend) = true

# --------------------------------------------------------------------------------------

function _initialize_backend(::QwtBackend; kw...)
  @eval begin
    warn("Qwt is no longer supported... many features will likely be broken.")
    import Qwt
    export Qwt
  end
end

# -------------------------------

const _qwtAliases = KW(
    :bins => :heatmap_n,
    :fillrange => :fillto,
    :linewidth => :width,
    :markershape => :marker,
    :hexbin => :heatmap,
    :path => :line,
    :steppost => :step,
    :steppre => :stepinverted,
    :star5 => :star1,
    :star8 => :star2,
  )

function fixcolors(d::KW)
  for (k,v) in d
    if typeof(v) <: ColorScheme
      d[k] = getColor(v)
    end
  end
end

function replaceQwtAliases(d, s)
  if haskey(_qwtAliases, d[s])
    d[s] = _qwtAliases[d[s]]
  end
end

function adjustQwtKeywords(plt::Plot{QwtBackend}, iscreating::Bool; kw...)
  d = KW(kw)
  st = d[:seriestype]
  if st == :scatter
    d[:seriestype] = :none
    if d[:markershape] == :none
      d[:markershape] = :circle
    end

  elseif st in (:hline, :vline)
    addLineMarker(plt, d)
    d[:seriestype] = :none
    d[:markershape] = :circle
    d[:markersize] = 1
    if st == :vline
      d[:x], d[:y] = d[:y], d[:x]
    end

  elseif !iscreating && st == :bar
    d = barHack(; kw...)
  elseif !iscreating && st == :histogram
    d = barHack(; histogramHack(; kw...)...)
  end

  replaceQwtAliases(d, :seriestype)
  replaceQwtAliases(d, :markershape)

  for k in keys(d)
    if haskey(_qwtAliases, k)
      d[_qwtAliases[k]] = d[k]
    end
  end

  d[:x] = collect(d[:x])
  d[:y] = collect(d[:y])

  d
end

# function _create_plot(pkg::QwtBackend, d::KW)
function _create_backend_figure(plt::Plot{QwtBackend})
  fixcolors(plt.attr)
  dumpdict(plt.attr,"\n\n!!! plot")
  o = Qwt.plot(zeros(0,0); plt.attr..., show=false)
  # plt = Plot(o, pkg, 0, d, KW[])
  # plt
end

# function _series_added(::QwtBackend, plt::Plot, d::KW)
function _series_added(plt::Plot{QwtBackend}, series::Series)
  d = adjustQwtKeywords(plt, false; series.d...)
  fixcolors(d)
  dumpdict(d,"\n\n!!! plot!")
  Qwt.oplot(plt.o; d...)
  # push!(plt.seriesargs, d)
  # plt
end

# ----------------------------------------------------------------

function updateLimsAndTicks(plt::Plot{QwtBackend}, d::KW, isx::Bool)
  lims = get(d, isx ? :xlims : :ylims, nothing)
  ticks = get(d, isx ? :xticks : :yticks, nothing)
  w = plt.o.widget
  axisid = Qwt.QWT.QwtPlot[isx ? :xBottom : :yLeft]

  if typeof(lims) <: Union{Tuple,AVec} && length(lims) == 2
    if isx
      plt.o.autoscale_x = false
    else
      plt.o.autoscale_y = false
    end
    w[:setAxisScale](axisid, lims...)
  end

  if typeof(ticks) <: Range
    if isx
      plt.o.autoscale_x = false
    else
      plt.o.autoscale_y = false
    end
    w[:setAxisScale](axisid, float(minimum(ticks)), float(maximum(ticks)), float(step(ticks)))
  elseif !(ticks in (nothing, :none, :auto))
    warn("Only Range types are supported for Qwt xticks/yticks. typeof(ticks)=$(typeof(ticks))")
  end

  # change the scale
  scalesym = isx ? :xscale : :yscale
  if haskey(d, scalesym)
    scaletype = d[scalesym]
    scaletype == :identity  && w[:setAxisScaleEngine](axisid, Qwt.QWT.QwtLinearScaleEngine())
    # scaletype == :log       && w[:setAxisScaleEngine](axisid, Qwt.QWT.QwtLogScaleEngine(e))
    # scaletype == :log2      && w[:setAxisScaleEngine](axisid, Qwt.QWT.QwtLogScaleEngine(2))
    scaletype == :log10     && w[:setAxisScaleEngine](axisid, Qwt.QWT.QwtLog10ScaleEngine())
    scaletype in supported_scales() || warn("Unsupported scale type: ", scaletype)
  end

end

function _update_plot_object(plt::Plot{QwtBackend}, d::KW)
  haskey(d, :title) && Qwt.title(plt.o, d[:title])
  haskey(d, :xguide) && Qwt.xlabel(plt.o, d[:xguide])
  haskey(d, :yguide) && Qwt.ylabel(plt.o, d[:yguide])
  updateLimsAndTicks(plt, d, true)
  updateLimsAndTicks(plt, d, false)
end

function _update_plot_pos_size(plt::AbstractPlot{QwtBackend}, d::KW)
  haskey(d, :size) && Qwt.resizewidget(plt.o, d[:size]...)
  haskey(d, :pos) && Qwt.movewidget(plt.o, d[:pos]...)
end

# ----------------------------------------------------------------

        # curve.setPen(Qt.QPen(Qt.QColor(color), linewidth, self.getLineStyle(linestyle)))
function addLineMarker(plt::Plot{QwtBackend}, d::KW)
  for yi in d[:y]
    marker = Qwt.QWT.QwtPlotMarker()
    ishorizontal = (d[:seriestype] == :hline)
    marker[:setLineStyle](ishorizontal ? 1 : 2)
    marker[ishorizontal ? :setYValue : :setXValue](yi)
    qcolor = Qwt.convertRGBToQColor(getColor(d[:linecolor]))
    linestyle = plt.o.widget[:getLineStyle](string(d[:linestyle]))
    marker[:setLinePen](Qwt.QT.QPen(qcolor, d[:linewidth], linestyle))
    marker[:attach](plt.o.widget)
  end

  # marker[:setValue](x, y)
  # marker[:setLabel](Qwt.QWT.QwtText(val))
  # marker[:attach](plt.o.widget)
end

function createQwtAnnotation(plt::Plot, x, y, val::PlotText)
  marker = Qwt.QWT.QwtPlotMarker()
  marker[:setValue](x, y)
  qwttext = Qwt.QWT.QwtText(val.str)
  qwttext[:setFont](Qwt.QT.QFont(val.font.family, val.font.pointsize))
  qwttext[:setColor](Qwt.convertRGBToQColor(getColor(val.font.color)))
  marker[:setLabel](qwttext)
  marker[:attach](plt.o.widget)
end

function createQwtAnnotation(plt::Plot, x, y, val::AbstractString)
  marker = Qwt.QWT.QwtPlotMarker()
  marker[:setValue](x, y)
  marker[:setLabel](Qwt.QWT.QwtText(val))
  marker[:attach](plt.o.widget)
end

function _add_annotations{X,Y,V}(plt::Plot{QwtBackend}, anns::AVec{Tuple{X,Y,V}})
  for ann in anns
    createQwtAnnotation(plt, ann...)
  end
end

# ----------------------------------------------------------------

# accessors for x/y data

function getxy(plt::Plot{QwtBackend}, i::Int)
  series = plt.o.lines[i]
  series.x, series.y
end

function setxy!{X,Y}(plt::Plot{QwtBackend}, xy::Tuple{X,Y}, i::Integer)
  series = plt.o.lines[i]
  series.x, series.y = xy
  plt
end

# -------------------------------

# -------------------------------

# # create the underlying object (each backend will do this differently)
# function _create_subplot(subplt::Subplot{QwtBackend}, isbefore::Bool)
#   isbefore && return false
#   i = 0
#   rows = Any[]
#   row = Any[]
#   for (i,(r,c)) in enumerate(subplt.layout)
#     push!(row, subplt.plts[i].o)
#     if c == ncols(subplt.layout, r)
#       push!(rows, Qwt.hsplitter(row...))
#       row = Any[]
#     end
#   end
#   # for rowcnt in subplt.layout.rowcounts
#   #   push!(rows, Qwt.hsplitter([plt.o for plt in subplt.plts[(1:rowcnt) + i]]...))
#   #   i += rowcnt
#   # end
#   subplt.o = Qwt.vsplitter(rows...)
#   # Qwt.resizewidget(subplt.o, getattr(subplt,1)[:size]...)
#   # Qwt.moveToLastScreen(subplt.o)  # hack so it goes to my center monitor... sorry
#   true
# end

function _expand_limits(lims, plt::Plot{QwtBackend}, isx::Bool)
  for series in plt.o.lines
    _expand_limits(lims, isx ? series.x : series.y)
  end
end

function _remove_axis(plt::Plot{QwtBackend}, isx::Bool)
end

# ----------------------------------------------------------------

function Base.show(io::IO, ::MIME"image/png", plt::Plot{QwtBackend})
  Qwt.refresh(plt.o)
  Qwt.savepng(plt.o, "/tmp/dfskjdhfkh.png")
  write(io, readall("/tmp/dfskjdhfkh.png"))
end

# function Base.show(io::IO, ::MIME"image/png", subplt::Subplot{QwtBackend})
#   for plt in subplt.plts
#     Qwt.refresh(plt.o)
#   end
#   Qwt.savepng(subplt.o, "/tmp/dfskjdhfkh.png")
#   write(io, readall("/tmp/dfskjdhfkh.png"))
# end

function Base.display(::PlotsDisplay, plt::Plot{QwtBackend})
  Qwt.refresh(plt.o)
  Qwt.showwidget(plt.o)
end

# function Base.display(::PlotsDisplay, subplt::Subplot{QwtBackend})
#   for plt in subplt.plts
#     Qwt.refresh(plt.o)
#   end
#   Qwt.showwidget(subplt.o)
# end


text:
@testset "gridstructure.jl" begin
    gs = GMRF.GridStructure((1,2), [[2], [1]], [[1],[2]], spzeros(1, 2), spzeros(1, 2))
    io = IOBuffer()

    @testset "showGridStructure(io, obj; prefix)" begin
        # does not throw
        @test_logs GMRF.showGridStructure(io, gs, prefix = "\t")

    end

    @testset "Base.show(io, obj)" begin
        # does not throw
        @test_logs Base.show(io, gs)

    end

end


text:
module hmodel

# Now for hydrostatic equilibrium, vertical temperature gradient.

export write_grid, write_model, write_lambda, write_dust, Parameters, Grid

# The double dot is because we are now inside the model module, and we want to import the
# constants module, which is part of the enclosing DiskJockey package.
using Dierckx
using ..constants

# Write the wavelength sampling file. Only run on setup
function write_lambda(lams::Array{Float64, 1}, basedir::AbstractString)
    fcam = open(basedir * "camera_wavelength_micron.inp", "w")
    nlam = length(lams)
    @printf(fcam, "%d\n", nlam)
    for lam in lams
        @printf(fcam, "%.9e\n", lam) # [microns]
    end

    close(fcam)
end

const eqmirror = true # mirror the grid about the z=0 midplane ?
# if we decide to mirror, then ncells = 1/2 of the true value

# Define a grid object which stores all of these variables
# This will not change for the duration of the run
immutable Grid
    nr::Int
    ntheta::Int
    nphi::Int
    ncells::Int
    # cell edges
    Rs::Vector{Float64}
    Thetas::Vector{Float64}
    Phis::Vector{Float64}
    # cell centers
    rs::Vector{Float64}
    thetas::Vector{Float64}
    phis::Vector{Float64}
end

function Grid(nr::Int, ntheta::Int, r_in::Real, r_out::Real, eqmirror::Bool)
    # Specify a 2D axisymmetric *separable* grid in spherical coordinates:
    # {r, theta, phi}, where theta is angle from zenith, phi is azimuth

    # Number of cells in each dimension
    nphi = 1 # axisymmetric disk
    ncells = nr * ntheta * nphi
    r_in = convert(Float64, r_in) * AU # Inner extent of disk
    r_out = convert(Float64, r_out) * AU # Outer extent of disk

    #Define the cell *walls*
    Rs = logspace(log10(r_in), log10(r_out), nr+1) # [cm] logarithmically spaced

    if eqmirror
        ped = 0.1
        #Thetas = linspace(0, pi/2., ntheta+1)
        # [rad] Angles are internally defined in radians, not degrees
        Thetas = pi/2. - (logspace(log10(ped), log10(pi/2. + ped), ntheta+1) - ped)[end:-1:1]
        #Spaced closer near the z=0
    else
        Thetas = linspace(0, pi, ntheta+1)
        # [rad] Angles are internally defined in radians, not degrees
    end

    Phis = Float64[0.0, 0.0] # [rad] cell walls for inactive coordinate

    #Define the cell centers as the average between walls
    rs = 0.5 * (Rs[1:end-1] + Rs[2:end])
    thetas = 0.5 * (Thetas[1:end-1] + Thetas[2:end])
    phis = Float64[0.0]

    return Grid(nr, ntheta, nphi, ncells, Rs, Thetas, Phis, rs, thetas, phis)

end

#This function only needs to be run once, upon setup.
function write_grid(basedir::AbstractString, grid::Grid)
    #amr_grid.inp
    f = open(basedir * "amr_grid.inp", "w")

    #Write the header
    @printf(f, "%d\n", 1) #iformat
    @printf(f, "%d\n", 0) #regular grid (no AMR or Oct-tree)
    @printf(f, "%d\n", 100) #spherical coordiantes
    @printf(f, "%d\n", 0) #gridinfo (none needed for now)
    #incl_r incl_phi incl_z #use this axis?
    @printf(f, "%d %d %d \n", 1, 1, 0) # 2D axisymmetric
    #n_r    n_phi   n_z #number of cells in this dimension
    @printf(f, "%d %d %d \n", grid.nr, grid.ntheta, grid.nphi)

    for R in grid.Rs
        @printf(f, "%.9e\n", R)
    end

    for Theta in grid.Thetas
        @printf(f, "%.9e\n", Theta)
    end

    for Phi in grid.Phis
        @printf(f, "%.9e\n", Phi)
    end

    close(f)
end

#Let's try defining a parameters type, an object of which gets passed around.

type Parameters
    M_star::Float64 # [M_sun] stellar mass
    r_c::Float64 # [AU] characteristic radius
    r_in::Float64 # [AU] inner radius of the grid
    r_out::Float64 # [AU] outer radius of the grid
    T_10m::Float64 # [K] temperature at 10 AU, midplane
    q_m::Float64 # midplane temperature gradient exponent
    T_10a::Float64 # [K] temperature at 10 AU, atmosphere
    q_a::Float64 # atmosphere temperature gradient exponent
    T_freeze::Float64 # [K] temperature below which to reduce CO abundance
    X_freeze::Float64 # [ratio] amount to reduce CO abundance
    sigma_s::Float64 # Photodissociation boundary in units of A_V.
    gamma::Float64 # surface temperature gradient exponent
    h::Float64 # Number of scale heights that z_q is at, currently fixed to 4
    delta::Float64 # Shape exponent, currently fixed to 2
    M_gas::Float64 # [M_Sun] disk mass of gas
    delta_gas::Float64 # Fraction by which gas is reduced inside the cavity
    r_cav::Float64 # [AU] the radius interior to which the gas density is reduced by delta_gas
    ksi::Float64 # [cm s^{-1}] micsroturbulence
    dpc::Float64 # [pc] distance to system
    incl::Float64 # [degrees] inclination 0 deg = face on, 90 = edge on.
    PA::Float64 # [degrees] position angle (East of North)
    vel::Float64 # [km/s] systemic velocity (positive is redshift/receeding)
    mu_RA::Float64 # [arcsec] central offset in RA
    mu_DEC::Float64 # [arcsec] central offset in DEC
end

# Assume all inputs to these functions are in CGS units and in *cylindrical* coordinates.
# Parametric type T allows passing individual Float64 or Vectors.
# Alternate functions accept pars passed around, where pars is in M_star, AU, etc...
function velocity(r::Float64, z::Float64, M_star::Float64)
    sqrt(G * M_star / (r^2 + z^2)^(3./2)) * r
end
velocity(r::Float64, z::Float64, pars::Parameters) = velocity(r, z, pars.M_star * M_sun)

# Midplane temperature
function T_mid{T}(r::T, T_10::Float64, q::Float64)
    T_10 * (r ./ (10. * AU)).^(-q)
end
T_mid{T}(r::T, pars::Parameters) = T_mid(r, pars.T_10m, pars.q_m)

# Atmosphere temperature
function T_atm{T}(r::T, T_10::Float64, q::Float64)
    T_10 * (r ./ (10. * AU)).^(-q)
end
T_atm{T}(r::T, pars::Parameters) = T_atm(r, pars.T_10a, pars.q_a)

# Scale height computed from midplane temperature
function Hp{T}(r::T, M_star::Float64, T_10::Float64, q::Float64)
    temp = T_mid(r, T_10, q)
    return sqrt((kB * temp .* r.^3.)/(mu_gas * m_H * G * M_star))
end
Hp{T}(r::T,  pars::Parameters) = Hp(r, pars.M_star * M_sun, pars.T_10m, pars.q_m)

# Calculate the gas surface density
function Sigma{T}(r::T, pars::Parameters)
    r_c = pars.r_c * AU
    Sigma_c = pars.M_gas * M_sun * (2 - pars.gamma) / (2 * pi * r_c^2)
    Sigma_c .* (r./r_c).^(-pars.gamma) .* exp(-(r./r_c).^(2 - pars.gamma))
end

# Atmosphere height computed as multiple of scale height (computed at midplane)
function z_q{T}(r::T, pars::Parameters)
    return pars.h * Hp(r, pars)
end

# No parametric type for temperature, because it is a 2D function.
# This is according to the Williams and Best 14 equations
function temperature(r::Float64, z::Float64, pars::Parameters)
    zq = z_q(r, pars)
    Ta = T_atm(r, pars)
    if z >= zq
        return Ta
    else
        Tm = T_mid(r, pars)
        # return Ta + (Tm - Ta) * cos(pi * z/ (2 * zq))^pars.delta
        return Tm + (Ta - Tm) * sin(pi * z/ (2 * zq))^(2 * pars.delta)
    end
end

# The temperature change
function dT(r::Float64, z::Float64, pars::Parameters)
    zq = z_q(r, pars)
    Ta = T_atm(r, pars)

    if z >= zq
        return 0.
    else
        Tm = T_mid(r, pars)
        # return - pi * pars.delta / (2 * zq) * (Tm - Ta) * (cos(pi * z/(2 * zq)))^(pars.delta - 1) * sin(pi * z / (2 * zq))
        return pars.delta * pi/zq * (Ta - Tm) * cos(pi * z/(2 * zq)) * (sin(pi * z/(2 * zq)))^(2 * pars.delta - 1)
    end
end

# The integrand
function dlnrho(r::Float64, z::Float64, pars::Parameters)
    T = temperature(r, z, pars)
    dTemp = dT(r, z, pars)
    return -(dTemp/T + (mu_gas * m_H)/(kB * T) * (G * pars.M_star * M_sun * z)/(r^2 + z^2)^1.5)
    # Vertically isothermal testcase
    # return -(mu_gas * m_H)/(kB * T) * (G * pars.M_star * M_sun * z)/r^3

end

# Delivers an unnormalized gas density. Needs to be multiplied by correction factor.
function un_lnrho(r::Float64, z::Float64, ztop::Float64, pars::Parameters)
    f(x) = dlnrho(r, x, pars)

    val, err = quadgk(f, ztop, z)
    return val
end

# Calculate a slice of density at a given radius
function density_slice(r::Float64, pars::Parameters)
    zq = z_q(r, pars) # The height of the disk "atmosphere"

    # 10 AU is an arbitrary number, but this is to make sure that we have a high enough bound
    # in the inner disk
    if (10 * zq) >= 10 * AU
        ztop = zq * 10
    else
        ztop = 10 * AU
    end

    # It is assumed that there is 0 gas above this upper bound, ztop.

    # Create an array of heights that goes from the midplane (0, 0.01, ..., ztop)
    nz = 64
    zs = cat(1, [0], logspace(log10(0.01 * AU), log10(ztop), nz-1))

    un_lnrhos = Array(Float64, nz)
    # For each height, integrate dlnrho/dz to yeild a yet-to-be normalized un_lnrho
    for i=1:nz
        un_lnrhos[i] = un_lnrho(r, zs[i], ztop, pars)
    end

    # Subtract the largest value from this array to avoid numberical overflow in subsequent steps
    bar_lnrho = un_lnrhos[1]

    # Remove the largest value so we don't get an overflow error
    un_lnrhos -= bar_lnrho

    # Convert from ln(rho) to rho, but this is still unnormalized.
    un_rhos = exp(un_lnrhos)

    # Now that we have the shape of the density structure but not the normalization,
    # integrate to get the normalization.
    # The density should be unitful, so we need to exponentiate (can't do this with un_lnrhos).
    spl = Spline1D(zs, un_rhos)

    # Integral is from -inf to inf
    tot = 2 * integrate(spl, 0., zs[end])

    # The actual value of the surface density
    S = Sigma(r, pars)

    # The normalized density slices
    rhos = S/tot * un_rhos # [g/cm^3]

    # Threshold column
    thresh = pars.sigma_s * Av_sigmaH # [g/cm^2]

    # Normalize the previous spline and integrate to get zboundary
    f(x) = S/tot * integrate(spl, x, zs[end])

    # Go through each of the z spacings, and use Riemann integration in a while loop to find at which z point we've finally
    column = 0.0 # total acumulated column
    icolumn::Int = 0 # index of the z point
    while column <= thresh

    end

    # fsolve, or some other root finding algorithm to find where f(x) = thresh.

    # Make all rhos above this equal to zero.

    return (zs, rhos)
end

function make_density_interpolator(pars::Parameters, grd::Grid)
    # Using the cell centers (in radius), go through and get a density slice.
    # Append the slice to the current list.
    # npoints = 64 * grd.nr
    rs = Array(Float64, (64, grid.nr))
    zs = Array(Float64, (64, grid.nr))
    rhos = Array(Float64, (64, grid.nr))
    for i=1:grid.nr
        r = grid.rs[i]
        rs[:,i] = r
        zs[:,i], rhos[:,i] = density_slice(r, pars)
    end

    # Now flatten these arrays and feed them to a 2D spline with Dierckx

end

# Calculate the multiplicative factor needed to normalize `un_rho` based upon the boundary
# condition that the integral of rho(r, z) with z from -inf to inf should be equal to Sigma(r)
function lncorrection_factor(r::Float64, pars::Parameters)
    S = Sigma(r, pars)
    zq = z_q(r, pars)
    # Most of the total density is near the midplane. This is necessary to create a logspaced array with 0 as the beginning.

    nz = 50

    # How best to choose the outer bound? Some multiple of zq is probably a good idea, but this
    # doesn't really go to high enough altitude in the inner disk. Instead, let's choose a minimum threshold.
    if (10 * zq) >= 10 * AU
        outer_reaches = zq * 10
    else
        outer_reaches = 10 * AU
    end

    zints = cat(1, [0], logspace(log10(0.01 * AU), log10(outer_reaches), nz-1))

    un_lnrhos = Array(Float64, nz)
    for i=1:nz
        un_lnrhos[i] = un_lnrho(r, zints[i], pars)
    end

    bar_lnrho = un_lnrhos[1]

    # Remove the largest value so we don't get an overflow error
    un_lnrhos -= bar_lnrho

    # Convert from ln(rho) to rho, but this is still unnormalized.
    un_rhos = exp(un_lnrhos)

    # Now that we have the shape of the density structure but not the normalization,
    # integrate to get the normalization
    # The density should be unitful, so we need to exponentiate.
    spl = Spline1D(zints, un_rhos)

    # Integral is from -inf to inf
    tot = 2 * integrate(spl, 0., zints[end])

    # println("Norm. tot: ", tot, " lntot:", log(tot))
    # println("bar_lnrho. exp: ", exp(bar_lnrho), " bar_lnrho: ", bar_lnrho)
    # println("With bar_lnrho :", tot * exp(bar_lnrho), " lntot :", log(tot) + bar_lnrho)

    # The correction factor is the ratio of the actual surface density to the value of the integral of our unnormalized density.
    lncor = log(S) - (log(tot) + bar_lnrho)

    # cor = S/tot
    # apply this as
    # rhos = cor * unnormed_rhos
    return lncor
end

# Now, create an interpolator for the the correction function. In spherical coordinates, that way, first evaluate it at every point (r, z=0) in the grid. Then, when we want to query the density of the disk at (r, z), which will correspond to a different r_cyl than the one used to compute the normalization, we can just use the spline interpolation and not have to redo the time consuming normalization integral.
function make_correction_interpolator(pars::Parameters, grid::Grid)

    nr = 64
    rs = logspace(log10(0.5 * AU), log10(grid.Rs[end]), nr)

    cors = Array(Float64, nr)

    for i=1:nr
        cors[i] = lncorrection_factor(rs[i], pars)
    end

    # An interpolating spline for this correction factor.
    spl = Spline1D(rs, cors)

    return spl
end

# spl is the spline object produced by a call to make_correction_interpolator()
function rho_gas(r::Float64, z::Float64, pars::Parameters, spl)

    # If we are inner to the radius used to compute the correction factors (in this case 0.5 AU),
    # just return 0.0 gas density
    if r < (0.5 * AU)
        return 0

    else
        # Evaluate the correction factor for this radius.
        lncor = evaluate(spl, r)

        # Calculate the unnormalized ln rho
        un = un_lnrho(r, z, pars)

        # rho = cor * exp(un)
        # apply the correction factor
        lnrho = lncor + un

        # Exponentiate it
        rho = exp(lnrho)

        # Check to make sure that this didn't return NaN, due to overflow or underflow errors somewhere in the integration pipeline. RADMC-3D will run with a NaN input to synthesize a blank image, and so the MCMC chain will gladly go on synthesizing blank images, even though the model is nonsense.
        if isnan(rho)
            println("Returned NaN for density.")
            throw(OverflowError)
        else
            return rho
        end
    end
end

n_12CO(r::Float64, z::Float64, pars::Parameters, spl) = number_densities["12CO"] * rho_gas(r, z, pars, spl)

n_13CO(r::Float64, z::Float64, pars::Parameters, spl) = number_densities["13CO"] * rho_gas(r, z, pars, spl)

n_C18O(r::Float64, z::Float64, pars::Parameters, spl) = number_densities["C18O"] * rho_gas(r, z, pars, spl)

# It is realistic to include freezout of CO onto dust grains.
# This is the amount by which the number density of the CO is reduced (X_freeze) relative to
# the nominal value.
function X_freeze(temp::Float64, pars::Parameters)
    # If it's cooler than the freezout temperature, reduce the number density by the given factor
    if temp <= pars.T_freeze
        return pars.X_freeze
    # Otherwise, just keep it as is.
    else
        return 1.0
    end
end

# Ksi is microturbulent broadining width in units of km/s. Output of this function
# is in cm/s for RADMC (RADMC manual, eqn 7.12)
function microturbulence(ksi::Float64)
    return ksi * 1.e5 # convert from km/s to cm/s
end

microturbulence(pars::Parameters) = microturbulence(pars.ksi)

function write_model(pars::Parameters, basedir::AbstractString, grid::Grid, species::AbstractString)

    funcs = Dict([("12CO", n_12CO), ("13CO", n_13CO), ("C18O", n_C18O)])
    n_CO = funcs[species]

    # numberdens_co.inp
    fdens = open(basedir * "numberdens_" * molnames[species] * ".inp", "w")
    @printf(fdens, "%d\n", 1) #iformat
    @printf(fdens, "%d\n", grid.ncells)

    # gas_velocity.inp
    fvel = open(basedir * "gas_velocity.inp", "w")
    @printf(fvel, "%d\n", 1) #iformat
    @printf(fvel, "%d\n", grid.ncells)

    # gas_temperature.inp
    ftemp = open(basedir * "gas_temperature.inp", "w")
    @printf(ftemp, "%d\n", 1) #iformat
    @printf(ftemp, "%d\n", grid.ncells)

    # microturbulence.inp
    fmicro = open(basedir * "microturbulence.inp", "w")
    @printf(fmicro, "%d\n", 1) #iformat
    @printf(fmicro, "%d\n", grid.ncells)

    # Calculate the correction function interpolator
    spl = make_correction_interpolator(pars, grid)

    # Now, we will need to write the three other files as a function of grid position.
    # Therefore we will do *one* loop over these indices, calculate the required value,
    # and write it to the appropriate file.

    #Looping over the cell centers
    for phi in grid.phis
        for theta in grid.thetas
            for r in grid.rs
                #Convert from spherical to cylindrical coordinates
                z = r * cos(theta)
                r_cyl = r * sin(theta)

                temp = temperature(r_cyl, z, pars)
                XF = X_freeze(temp, pars)
                @printf(fdens, "%.9e\n", XF * n_CO(r_cyl, z, pars, spl))
                @printf(fvel, "0 0 %.9e\n", velocity(r_cyl, pars))
                @printf(ftemp, "%.9e\n", temp)
                @printf(fmicro, "%.9e\n", microturbulence(pars))
            end
        end
    end

    close(fdens)
    close(fvel)
    close(ftemp)
    close(fmicro)

end

function write_dust(pars::Parameters, basedir::AbstractString, grid::Grid)
    fdens = open(basedir * "dust_density.inp", "w")
    @printf(fdens, "%d\n", 1) #iformat
    @printf(fdens, "%d\n", grid.ncells)
    @printf(fdens, "%d\n", 1) # number of dust species

    for phi in grid.phis
        for theta in grid.thetas
            for r in grid.rs
                #Convert from spherical to cylindrical coordinates
                z = r * cos(theta)
                r_cyl = r * sin(theta)

                @printf(fdens, "%.9e\n", rho_dust(r_cyl, z, pars))
            end
        end
    end

    close(fdens)
end

end

max_stars_repo_name: elnjensen/DiskJockey

max_issues_repo_name: elnjensen/DiskJockey

max_forks_repo_name: elnjensen/DiskJockey


text:
module MPIClusterManagers

export MPIManager, launch, manage, kill, procs, connect, mpiprocs, @mpi_do, TransportMode, MPI_ON_WORKERS, TCP_TRANSPORT_ALL, MPI_TRANSPORT_ALL

using Distributed, Serialization
import MPI

include("mpimanager.jl")

end # module


text:
@test latexify("α"; convert_unicode=false) == raw"$α$"

@test latexify(['α', :β, "γ/η"], transpose=true, convert_unicode=false) == replace(
raw"\begin{equation}
\left[
\begin{array}{ccc}
α & β & \frac{γ}{η} \\
\end{array}
\right]
\end{equation}
", "\r\n"=>"\n")

@test latexify("αaβ") == raw"${\alpha}a\beta$"