text:
@inline function sor_kernel(
    potential::AbstractArray{T, 4},
    imp_scale::AbstractArray{T, 4},
    point_types::AbstractArray{PointType, 4},
    volume_weights::AbstractArray{T, 4},
    q_eff_imp::AbstractArray{T, 4},
    q_eff_fix::AbstractArray{T, 4},
    ϵ_r::AbstractArray{T, 3},
    geom_weights::NTuple{3, AbstractArray{T, 2}},
    sor_const::AbstractArray{T, 1},
    update_even_points::Bool,
    depletion_handling_enabled::Bool,
    is_weighting_potential::Bool,
    only2d::Bool,
    ::Type{S},
    gpu_inds::NTuple{3, Int}
) where {T, S}
    # Comparison to CPU indices: (Cyl / Car)
    # i3 <-> idx3 / ir / iz
    # i2 <-> idx2 / iφ / iy
    # i1 <-> idx1 / iz / ix
    i1, in2, in3 = gpu_inds
    i1 += 1
    i2 = in2 + 1
    i3 = in3 + 1
    in1 = nidx(i1, update_even_points, iseven(i2 + i3))
    i1r = get_rbidx_right_neighbour(i1, update_even_points, iseven(i2 + i3))

    rb_tar_idx, rb_src_idx = update_even_points ? (rb_even::Int, rb_odd::Int) : (rb_odd::Int, rb_even::Int)

    geom_weights_3 = get_geom_weights_outerloop(geom_weights, in3, S)

    geom_weights_2 = prepare_weights_in_middleloop(
        geom_weights, S, i2, in2,
        geom_weights_3...,
        in3 == 1
    )

    weights = calculate_sor_weights(
        in1,
        S,
        ϵ_r,
        geom_weights[3],
        i2, in2, i3, in3,
        geom_weights_2...
    )

    old_potential = potential[i1, i2, i3, rb_tar_idx]

    neighbor_potentials = get_neighbor_potentials(
        potential, old_potential, i1, i2, i3, i1r, in2, in3, rb_src_idx, only2d
    )

    new_potential = calc_new_potential_by_neighbors_3D(
        volume_weights[i1, i2, i3, rb_tar_idx],
        weights,
        neighbor_potentials
    )

    update_point = point_types[i1, i2, i3, rb_tar_idx] & update_bit > 0

    if depletion_handling_enabled
        new_potential, new_imp_scale = handle_depletion(
            new_potential,
            imp_scale[i1, i2, i3, rb_tar_idx],
            r0_handling_depletion_handling(neighbor_potentials, S, in3),
            q_eff_imp[i1, i2, i3, rb_tar_idx],
            volume_weights[i1, i2, i3, rb_tar_idx]
        )
        imp_scale[i1, i2, i3, rb_tar_idx] = ifelse(update_point, new_imp_scale, imp_scale[i1, i2, i3, rb_tar_idx])
    end

    if !is_weighting_potential
        q_eff = if depletion_handling_enabled
            q_eff_fix[i1, i2, i3, rb_tar_idx]
        else
            q_eff_imp[i1, i2, i3, rb_tar_idx] + q_eff_fix[i1, i2, i3, rb_tar_idx]
        end
        new_potential += q_eff * volume_weights[i1, i2, i3, rb_tar_idx]
    end

    if is_weighting_potential || !depletion_handling_enabled
        new_potential = apply_over_relaxation(new_potential, old_potential, get_sor_constant(sor_const, S, in3))
    end

    potential[i1, i2, i3, rb_tar_idx] = ifelse(update_point, new_potential, old_potential)
    nothing
end

@inline function get_neighbor_potentials(
    potential::AbstractArray{T, 4},
    old_potential, i1, i2, i3, i1r, in2, in3, rb_src_idx, only2d::Bool
)::NTuple{6, T} where {T}
    @inbounds return ( # p: potential; 1: RB-dimension; l/r: left/right
    potential[i1r - 1,     i2,     i3, rb_src_idx], # p1l
    potential[    i1r,     i2,     i3, rb_src_idx], # p1r
    only2d ? old_potential : potential[ i1,    in2, i3, rb_src_idx], # p2l
    only2d ? old_potential : potential[ i1, i2 + 1, i3, rb_src_idx], # p2r
    potential[     i1,     i2,    in3, rb_src_idx], # p3l
    potential[     i1,     i2, i3 + 1, rb_src_idx], # p3r
    )
end

@inline function prepare_weights_in_middleloop(
    geom_weights::NTuple{3, <:AbstractArray{T, 2}}, ::Type{Cylindrical},
    i2, in2,
    pwwrr, pwwrl, r_inv_pwΔmpr, Δr_ext_inv_r_pwmprr, Δr_ext_inv_l_pwmprl, Δmpr_squared,
    is_r0::Bool
) where {T}
    pwwφr        = geom_weights[2][1, in2]
    pwwφl        = geom_weights[2][2, in2]
    pwΔmpφ       = geom_weights[2][3, in2]
    Δφ_ext_inv_r = geom_weights[2][4,  i2]
    Δφ_ext_inv_l = geom_weights[2][4, in2]

    if is_r0
        pwwφr = T(0.5)
        pwwφl = T(0.5)
        pwΔmpφ = T(2π)
        Δφ_ext_inv_r = inv(pwΔmpφ)
        Δφ_ext_inv_l = Δφ_ext_inv_r
    end
    pwwrr_pwwφr = pwwrr * pwwφr
    pwwrl_pwwφr = pwwrl * pwwφr
    pwwrr_pwwφl = pwwrr * pwwφl
    pwwrl_pwwφl = pwwrl * pwwφl

    pwΔmpφ_Δmpr_squared = pwΔmpφ * Δmpr_squared
    Δr_ext_inv_r_pwmprr_pwΔmpφ = Δr_ext_inv_r_pwmprr * pwΔmpφ
    Δr_ext_inv_l_pwmprl_pwΔmpφ = Δr_ext_inv_l_pwmprl * pwΔmpφ
    r_inv_pwΔmpr_Δφ_ext_inv_r = r_inv_pwΔmpr * Δφ_ext_inv_r
    r_inv_pwΔmpr_Δφ_ext_inv_l = r_inv_pwΔmpr * Δφ_ext_inv_l
    return (
        pwwrr, pwwrl, pwwφr, pwwφl,
        pwwrr_pwwφr, pwwrl_pwwφr, pwwrr_pwwφl, pwwrl_pwwφl,
        pwΔmpφ_Δmpr_squared,
        Δr_ext_inv_r_pwmprr_pwΔmpφ, Δr_ext_inv_l_pwmprl_pwΔmpφ,
        r_inv_pwΔmpr_Δφ_ext_inv_r, r_inv_pwΔmpr_Δφ_ext_inv_l
    )
end

"""
    function get_sor_kernel(::Type{S}, args...)

where S is either Cartesian or Cylindrical.

Developer notes:
Currently (February 2022), there are some limitations to the `@kernel` macro
of the package KernelAbstractions.jl. Especially, regarding usage of dispatch.

Thus, we have to write two kernel functions right now for the Cartesian & Cylindrical case:
`sor_cyl_gpu!` and `sor_car_gpu!`.

Inside kernel functions, everything is (and has to be) inlined and we can make use of multiple dispatch.
So in the end we only have to write one function for the kernel, `sor_kernel`,
which is then inlined inside the two kernel functions.

We can also use most of the CPU functions with the restriction that all
types have to be independent on the GPU-indices of the kernel.
E.g., making use of `i23_is_even_t = Val(iseven(i2 + i3))` and other similar statements is not possible,
which are used in the CPU implementation for optimization.
On the GPU (currently) those statements have to be calculated and booleans have to be passed.
Maybe this will change in the future.
"""
get_sor_kernel(::Type{Cylindrical}, dev, CPU_via_KernelAbstractions, args...) = sor_cyl_gpu!(dev, args...)
get_sor_kernel(::Type{Cartesian},   dev, CPU_via_KernelAbstractions, args...) = sor_car_gpu!(dev, args...)
get_sor_kernel(::Type{Cylindrical}, ::CPU, CPU_via_KernelAbstractions::Val{false}) = nothing
get_sor_kernel(::Type{Cartesian},   ::CPU, CPU_via_KernelAbstractions::Val{false}) = nothing
get_sor_kernel(::Type{Cylindrical}, dev::CPU, CPU_via_KernelAbstractions::Val{true}, args...) = sor_cyl_gpu!(dev, args...)
get_sor_kernel(::Type{Cartesian},   dev::CPU, CPU_via_KernelAbstractions::Val{true}, args...) = sor_car_gpu!(dev, args...)


text:
width = 10
height = 10
mines_count = 10
mines = falses(height, width)
opened = falses(height, width)

game_end = false

println("Minesweeper")

adj = [[-1,0],[-1,1],[0,1],[1,1],[1,0],[1,-1],[0,-1],[-1,-1]]
function mine_count(row, col)
    count = 0
    for a in adj
        r = row+a[1]
        c = col+a[2]
        if r>0 && c>0 && r<=height && c<=width && mines[CartesianIndex(r, c)] ==  true
            count += 1
        end
    end
    return count
end

function new_game()
    global mines = falses(height, width)
    global opened = falses(height, width)

    mc = mines_count
    while mc>0
        loc = rand(0:width*height-1)
        idx = CartesianIndex(loc÷height+1, loc%width+1)
        if !mines[idx]
            mines[idx] = true
            mc -= 1
        end
    end

    global game_end = false
end

function display()
    print("   ")
    for i = 1:width
        print(i, " ")
    end
    println()
    for j = 1:height
        print(lpad(j, 2, " ") * " ")
        for i = 1:width
            idx = CartesianIndex(j, i)
            if game_end && mines[idx]
                print("✽ ")
            elseif !opened[idx]
                print("◼ ")
            else
                if mines[idx]
                    print("✽ ")
                else
                    count = mine_count(j, i)
                    if count==0
                        print("◻ ")
                    else
                        print(count, " ")
                    end
                end
            end
        end
        println()
    end
    if !game_end
        println("To open: ", count(i->(!i), opened)-mines_count)
    end
end

function open(row, col)
    if !opened[CartesianIndex(row, col)]
        opened[CartesianIndex(row, col)] = true
        if count(i->(!i), opened)-mines_count == 0
            global game_end = true
            display()
            print("You won!")
        end
        if mine_count(row, col) > 0
            return
        end
        for a in adj
            r = row+a[1]
            c = col+a[2]
            if r>0 && c>0 && r<=height && c<=width
                open(r, c)
            end
        end
    end
end

function process(row, col)
    if mines[CartesianIndex(row, col)] == true
        opened[CartesianIndex(row, col)] = true
        global game_end = true
        display()
        println("You lost!")
    else
        open(row, col)
    end
end

new_game()

while !game_end
    display()
    print(">> ")
    user = strip(readline(stdin))
    if user == "end"
        global game_end = true
    else
        rc = split(user)
        if length(rc) != 2
            println("Incorrect input!")
            continue
        end
        row = parse(Int32, rc[1])
        col = parse(Int32, rc[2])
        process(row, col)
    end
end

ext: jl

max_stars_repo_name: aole/Games-in-Julia

max_issues_repo_name: aole/Games-in-Julia

max_forks_repo_name: aole/Games-in-Julia


text:
@testset "Testing yaml_import" begin
    path = joinpath(Dolo.pkg_path, "examples", "models")
    files = [f for f in readdir(path) if occursin("yaml", f)]
    for fname in files
        println("Importing ", fname)
        model = yaml_import(joinpath(path, fname))
        res = residuals(model)
        if !occursin("sudden_stop", fname)
            @test (maximum(res[:transition]) + maximum(res[:arbitrage]) )<1e-5
        else
            @test (maximum(res[:transition]) + maximum(res[:arbitrage])-0.215)<1e-5
        end
    end
end


text:
using Revise # reduce need for recompile
using Plots
using LinearAlgebra

push!(LOAD_PATH, "./src") # user defined modules
using CommonUtils
using Basis1D
using Basis2DTri
using UniformTriMesh

"Approximation parameters"
N   = 3 # The order of approximation
K1D = 16
CFL = .5
T   = 2.0 # endtime

"Mesh related variables"
VX,VY,EToV = uniform_tri_mesh(K1D,K1D)
FToF = connect_mesh(EToV,tri_face_vertices())
Nfaces,K = size(FToF)

"Construct matrices on reference elements"
r, s = nodes_2D(N)
V = vandermonde_2D(N, r, s)
Vr, Vs = grad_vandermonde_2D(N, r, s)
M = inv(V*V')
Dr = Vr/V
Ds = Vs/V

"Nodes on faces, and face node coordinate"
r1D, w1D = gauss_quad(0,0,N)
Nfp = length(r1D)
e = ones(Nfp,1)
z = zeros(Nfp,1)
rf = [r1D; -r1D; -e];
sf = [-e; r1D; -r1D];
wf = vec(repeat(w1D,3,1));
nrJ = [z; e; -e]
nsJ = [-e; e; z]
Vf = vandermonde_2D(N,rf,sf)/V
LIFT = M\(transpose(Vf)*diagm(wf)) # lift matrix

"Construct global coordinates"
r1,s1 = nodes_2D(1)
V1 = vandermonde_2D(1,r,s)/vandermonde_2D(1,r1,s1)
x = V1*VX[transpose(EToV)]
y = V1*VY[transpose(EToV)]

"Connectivity maps"
xf,yf = (x->Vf*x).((x,y))
mapM,mapP,mapB = build_node_maps((xf,yf),FToF)
mapM = reshape(mapM,Nfp*Nfaces,K)
mapP = reshape(mapP,Nfp*Nfaces,K)

"Make periodic"
LX,LY = (x->maximum(x)-minimum(x)).((VX,VY)) # find lengths of domain
mapPB = build_periodic_boundary_maps(xf,yf,LX,LY,Nfaces*K,mapM,mapP,mapB)
mapP[mapB] = mapPB

"Geometric factors and surface normals"
rxJ, sxJ, ryJ, syJ, J = geometric_factors(x, y, Dr, Ds)
nxJ = (Vf*rxJ).*nrJ + (Vf*sxJ).*nsJ;
nyJ = (Vf*ryJ).*nrJ + (Vf*syJ).*nsJ;
sJ = @. sqrt(nxJ^2 + nyJ^2)

"initial conditions"
u = @. exp(-25*(x^2+y^2))

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

"pack arguments into tuples"
ops = (Dr,Ds,LIFT,Vf)
vgeo = (rxJ,sxJ,ryJ,syJ,J)
fgeo = (nxJ,nyJ,sJ)
nodemaps = (mapP,mapB)

function rhs(u,ops,vgeo,fgeo,nodemaps)
    # unpack args
    Dr,Ds,LIFT,Vf = ops
    rxJ,sxJ,ryJ,syJ,J = vgeo
    nxJ,nyJ,sJ = fgeo
    (mapP,mapB) = nodemaps

    uM = Vf*u # can replace with nodal extraction
    uP = uM[mapP]
    du = uP-uM

    ur = Dr*u
    us = Ds*u
    ux = @. rxJ*ur + sxJ*us

    tau = .5
    rhsu = ux + LIFT*(@. .5*du*nxJ - tau*du*abs(nxJ))

    return -rhsu./J
end

resu = zeros(size(x))
for i = 1:Nsteps
    global u, resu # for scoping - these variables are updated

    for INTRK = 1:5
        rhsu = rhs(u,ops,vgeo,fgeo,nodemaps)
        @. resu = rk4a[INTRK]*resu + dt*rhsu
        @. u += rk4b[INTRK]*resu
    end

    if i%10==0 || i==Nsteps
        println("Number of time steps $i out of $Nsteps")
    end
end

"plotting nodes"

gr(aspect_ratio=1,legend=false,
   markerstrokewidth=0,markersize=2)

rp, sp = equi_nodes_2D(10)
Vp = vandermonde_2D(N,rp,sp)/V
vv = Vp*u
scatter(Vp*x,Vp*y,vv,zcolor=vv,camera=(0,90))

max_stars_repo_name: jlchan/ESDG-CNS

max_issues_repo_name: jlchan/ESDG-CNS

max_forks_repo_name: jlchan/ESDG-CNS


text:
# Notations from Walker & Ni, "Anderson acceleration for fixed-point iterations", SINUM 2011
# Attempts to accelerate the iteration xₙ₊₁ = xₙ + beta*f(xₙ)

struct AndersonCache{Tx,To,Tdg,Tg,TQ,TR} <: AbstractSolverCache
    x::Tx
    g::Tx
    fxold::To
    gold::To
    Δgs::Tdg
    γs::Tg
    Q::TQ
    R::TR
end

function AndersonCache(df, m)
    x = copy(df.x_f)
    g = copy(x)

    if m > 0
        fxold = copy(x)
        gold = copy(x)

        # maximum size of history
        mmax = min(length(x), m)

        # buffer storing the differences between g of the iterates, from oldest to newest
        Δgs = [copy(x) for _ in 1:mmax]

        T = eltype(x)
        γs = Vector{T}(undef, mmax) # coefficients obtained from the least-squares problem

        # matrices for QR decomposition
        Q = Matrix{T}(undef, length(x), mmax)
        R = Matrix{T}(undef, mmax, mmax)
    else
        fxold = nothing
        gold = nothing
        Δgs = nothing
        γs = nothing
        Q = nothing
        R = nothing
    end

    AndersonCache(x, g, fxold, gold, Δgs, γs, Q, R)
end

@views function anderson_(df::Union{NonDifferentiable, OnceDifferentiable},
                             initial_x::AbstractArray,
                             xtol::Real,
                             ftol::Real,
                             iterations::Integer,
                             store_trace::Bool,
                             show_trace::Bool,
                             extended_trace::Bool,
                             beta::Real,
                             aa_start::Integer,
                             droptol::Real,
                             cache::AndersonCache)
    copyto!(cache.x, initial_x)
    tr = SolverTrace()
    tracing = store_trace || show_trace || extended_trace

    aa_iteration = cache.γs !== nothing
    x_converged, f_converged, converged = false, false, false
    m = aa_iteration ? length(cache.γs) : 0
    aa_iteration && (m_eff = 0)

    iter = 0
    while iter < iterations
        iter += 1

        # evaluate function
        value!!(df, cache.x)
        fx = value(df)

        # check that all values are finite
        check_isfinite(fx)

        # compute next iterate of fixed-point iteration
        @. cache.g = cache.x + beta * fx

        # save trace
        if tracing
            dt = Dict()
            if extended_trace
                dt["x"] = copy(cache.x)
                dt["f(x)"] = copy(fx)
            end
            update!(tr,
                    iter,
                    maximum(abs, fx),
                    iter > 1 ? sqeuclidean(cache.g, cache.x) : convert(real(eltype(initial_x)), NaN),
                    dt,
                    store_trace,
                    show_trace)
        end

        # check convergence
        x_converged, f_converged = assess_convergence(cache.g, cache.x, fx, xtol, ftol)
        if any(isnan, cache.x) || any(isnan, value(df))
          break
        end
        converged = x_converged || f_converged
        converged && break

        # update current iterate
        copyto!(cache.x, cache.g)

        # perform Anderson acceleration
        if aa_iteration
            if iter == aa_start
                # initialize cache of residuals and g
                copyto!(cache.fxold, fx)
                copyto!(cache.gold, cache.g)
            elseif iter > aa_start
                # increase size of history
                m_eff += 1

                # remove oldest history if maximum size is exceeded
                if m_eff > m
                    # circularly shift differences of g
                    ptr = cache.Δgs[1]
                    for i in 1:(m-1)
                        cache.Δgs[i] = cache.Δgs[i + 1]
                    end
                    cache.Δgs[m] = ptr

                    # delete left-most column of QR decomposition
                    qrdelete!(cache.Q, cache.R, m)

                    # update size of history
                    m_eff = m
                end

                # update history of differences of g
                @. cache.Δgs[m_eff] = cache.g - cache.gold

                # replace/add difference of residuals as right-most column to QR decomposition
                @. cache.fxold = fx - cache.fxold
                qradd!(cache.Q, cache.R, vec(cache.fxold), m_eff)

                # update cached values
                copyto!(cache.fxold, fx)
                copyto!(cache.gold, cache.g)

                # define current Q and R matrices
                Q, R = view(cache.Q, :, 1:m_eff), UpperTriangular(view(cache.R, 1:m_eff, 1:m_eff))

                # check condition (TODO: incremental estimation)
                if droptol > 0
                    while cond(R) > droptol && m_eff > 1
                        qrdelete!(cache.Q, cache.R, m_eff)
                        m_eff -= 1
                        Q, R = view(cache.Q, :, 1:m_eff), UpperTriangular(view(cache.R, 1:m_eff, 1:m_eff))
                    end
                end

                # solve least squares problem
                γs = view(cache.γs, 1:m_eff)
                ldiv!(R, mul!(γs, Q', vec(fx)))

                # update next iterate
                for i in 1:m_eff
                    @. cache.x -= cache.γs[i] * cache.Δgs[i]
                end
            end
        end
    end

    return SolverResults("Anderson m=$m beta=$beta aa_start=$aa_start droptol=$droptol",
                         initial_x, copy(cache.x), norm(value(df), Inf),
                         iter, x_converged, xtol, f_converged, ftol, tr,
                         first(df.f_calls), 0)
end

function anderson(df::Union{NonDifferentiable, OnceDifferentiable},
                  initial_x::AbstractArray,
                  xtol::Real,
                  ftol::Real,
                  iterations::Integer,
                  store_trace::Bool,
                  show_trace::Bool,
                  extended_trace::Bool,
                  m::Integer,
                  beta::Real,
                  aa_start::Integer,
                  droptol::Real)
    anderson(df, initial_x, xtol, ftol, iterations, store_trace, show_trace, extended_trace, beta, aa_start, droptol, AndersonCache(df, m))
end

function anderson(df::Union{NonDifferentiable, OnceDifferentiable},
                  initial_x::AbstractArray,
                  xtol::Real,
                  ftol::Real,
                  iterations::Integer,
                  store_trace::Bool,
                  show_trace::Bool,
                  extended_trace::Bool,
                  beta::Real,
                  aa_start::Integer,
                  droptol::Real,
                  cache::AndersonCache)
    anderson_(df, initial_x, convert(real(eltype(initial_x)), xtol), convert(real(eltype(initial_x)), ftol), iterations, store_trace, show_trace, extended_trace, beta, aa_start, droptol, cache)
end


text:
@inline function ExaTronKernels.dtron(n::Int, x::ROCDeviceArray{Float64,1}, xl::ROCDeviceArray{Float64,1},
               xu::ROCDeviceArray{Float64,1}, f::Float64, g::ROCDeviceArray{Float64,1},
               A::ROCDeviceArray{Float64,2}, frtol::Float64, fatol::Float64,
               fmin::Float64, cgtol::Float64, itermax::Int, delta::Float64, task::Int,
               B::ROCDeviceArray{Float64,2}, L::ROCDeviceArray{Float64,2},
               xc::ROCDeviceArray{Float64,1}, s::ROCDeviceArray{Float64,1},
               indfree::ROCDeviceArray{Int,1}, gfree::ROCDeviceArray{Float64,1},
               isave::ROCDeviceArray{Int,1}, dsave::ROCDeviceArray{Float64,1},
               wa::ROCDeviceArray{Float64,1}, iwa::ROCDeviceArray{Int,1},
               wa1::ROCDeviceArray{Float64,1}, wa2::ROCDeviceArray{Float64,1},
               wa3::ROCDeviceArray{Float64,1}, wa4::ROCDeviceArray{Float64,1},
               wa5::ROCDeviceArray{Float64,1})
    zero = 0.0
    p5 = 0.5
    one = 1.0

    # Parameters for updating the iterates.

    eta0 = 1.0e-4
    eta1 = 0.25
    eta2 = 0.75

    # Parameters for updating the trust region size delta.

    sigma1 = 0.25
    sigma2 = 0.5
    sigma3 = 4.0

    work = 0

    # Initialization section.

    if task == 0  # "START"

        # Initialize local variables.

        iter = 1
        iterscg = 0
        alphac = one
        work = 1  # "COMPUTE"

    else

        @inbounds begin
            # Restore local variables.

            work = isave[1]
            iter = isave[2]
            iterscg = isave[3]
            fc = dsave[1]
            alphac = dsave[2]
            prered = dsave[3]
        end
    end

    AMDGPU.sync_workgroup()

    # Search for a lower function value.

    search = true
    while search

        # Compute a step and evaluate the function at the trial point.

        if work == 1 # "COMPUTE"

            # Save the best function value, iterate, and gradient.

            fc = f
            dcopy(n,x,1,xc,1)

            # Compute the Cauchy step and store in s.

            alphac = dcauchy(n,x,xl,xu,A,g,delta,
                             alphac,s,wa)

            # Compute the projected Newton step.

            info,iters = dspcg(n, delta, cgtol, itermax,
                               x, xl, xu, A, g, s,
                               B, L,
                               indfree, gfree, wa, iwa,
                               wa1, wa2, wa3, wa4, wa5)

            # Compute the predicted reduction.

            dssyax(n, A, s, wa)
            prered = -(ddot(n,s,1,g,1) + p5*ddot(n,s,1,wa,1))
            iterscg = iterscg + iters

            # Set task to compute the function.

            task = 1 # 'F'
        end

        # Evaluate the step and determine if the step is successful.

        if work == 2 # "EVALUATE"

            # Compute the actual reduction.

            actred = fc - f

            # On the first iteration, adjust the initial step bound.

            snorm = dnrm2(n,s,1)
            if iter == 1
                delta = min(delta,snorm)
            end

            # Update the trust region bound.

            g0 = ddot(n,g,1,s,1)
            if f-fc-g0 <= zero
                alpha = sigma3
            else
                alpha = max(sigma1,-p5*(g0/(f-fc-g0)))
            end

            # Update the trust region bound according to the ratio
            # of actual to predicted reduction.

            if actred < eta0*prered
                delta = min(max(alpha,sigma1)*snorm,sigma2*delta)
            elseif actred < eta1*prered
                delta = max(sigma1*delta,min(alpha*snorm,sigma2*delta))
            elseif actred < eta2*prered
                delta = max(sigma1*delta,min(alpha*snorm,sigma3*delta))
            else
                delta = max(delta,min(alpha*snorm,sigma3*delta))
            end

            # Update the iterate.

            if actred > eta0*prered

                # Successful iterate.

                task = 2 # 'G' or 'H'
                iter = iter + 1

            else

                # Unsuccessful iterate.

                task = 1 # 'F'
                dcopy(n,xc,1,x,1)
                f = fc

            end

            # Test for convergence.

            if f < fmin
                task = 10 # "WARNING: F .LT. FMIN"
            end
            if abs(actred) <= fatol && prered <= fatol
                task = 4 # "CONVERGENCE: FATOL TEST SATISFIED"
            end
            if abs(actred) <= frtol*abs(f) && prered <= frtol*abs(f)
                task = 4 # "CONVERGENCE: FRTOL TEST SATISFIED"
            end
        end

        # Test for continuation of search

        if task == 1 && work == 2 # Char(task[1]) == 'F' && work == "EVALUATE"
            search = true
            work = 1 # "COMPUTE"
        else
            search = false
        end
    end

    if work == 3 # "NEWX"
        task = 3 # "NEWX"
    end

    # Decide on what work to perform on the next iteration.

    if task == 1 && work == 1 # Char(task[1]) == 'F' && work == "COMPUTE"
        work = 2 # "EVALUATE"
    elseif task == 1 && work == 2 # Char(task[1]) == 'F' && work == "EVALUATE"
        work = 1 # "COMPUTE"
    elseif task == 2 # unsafe_string(pointer(task),2) == "GH"
        work = 3 # "NEWX"
    elseif task == 3 # unsafe_string(pointer(task),4) == "NEWX"
        work = 1 # "COMPUTE"
    end

    @inbounds begin
        # Save local variables.

        isave[1] = work
        isave[2] = iter
        isave[3] = iterscg

        dsave[1] = fc
        dsave[2] = alphac
        dsave[3] = prered
    end

    AMDGPU.sync_workgroup()

    return delta, task
end


text:
###
### This file contains methods defined on `AbstractDVec`
### The type definition and relevant methods are found in the file "src/Interfaces/dictvectors.jl"
###
function Base.show(io::IO, dvec::AbstractDVec)
    summary(io, dvec)
    limit, _ = displaysize()
    for (i, p) in enumerate(pairs(dvec))
        if length(dvec) > i > limit - 4
            print(io, "\n  ⋮   => ⋮")
            break
        else
            print(io, "\n  ", p)
        end
    end
end

###
### Types
###
Base.keytype(::Type{<:AbstractDVec{K}}) where {K} = K
Base.keytype(dv::AbstractDVec) = keytype(typeof(dv))
Base.valtype(::Type{<:AbstractDVec{<:Any,V}}) where {V} = V
Base.valtype(dv::AbstractDVec) = valtype(typeof(dv))
Base.eltype(::Type{<:AbstractDVec{K,V}}) where {K,V} = Pair{K,V}
Base.eltype(dv::AbstractDVec) = eltype(typeof(dv))

Base.isreal(v::AbstractDVec) = valtype(v)<:Real
Base.ndims(::AbstractDVec) = 1

###
### copy*, zero*
###
"""
    zero!(v)

Replace `v` by a zero vector as an inplace operation. For `AbstractDVec` types it means
removing all non-zero elements. For `AbstractArrays`, it sets all of the values to zero.
"""
zero!(v::AbstractVector{T}) where {T} = v .= zero(T)
zero!(v::AbstractDVec) = empty!(v)

Base.zero(dv::AbstractDVec) = empty(dv)

function Base.similar(dvec::AbstractDVec, args...; kwargs...)
    return sizehint!(empty(dvec, args...; kwargs...), length(dvec))
end

@inline function Base.copyto!(w::AbstractDVec, v)
    sizehint!(w, length(v))
    for (key, val) in pairs(v)
        w[key] = val
    end
    return w
end
@inline function Base.copy!(w::AbstractDVec, v)
    empty!(w)
    return copyto!(w, v)
end
Base.copy(v::AbstractDVec) = copyto!(empty(v), v)

###
### Linear algebra
###
function Base.sum(f, x::AbstractDVec)
    return sum(f, values(x))
end

function LinearAlgebra.norm(x::AbstractDVec, p::Real=2)
    if p === 1
        return float(sum(abs, values(x)))
    elseif p === 2
        return sqrt(sum(abs2, values(x)))
    elseif p === Inf
        return float(mapreduce(abs, max, values(x), init=real(zero(valtype(x)))))
    else
        error("$p-norm of $(typeof(x)) is not implemented.")
    end
end

@inline function LinearAlgebra.mul!(w::AbstractDVec, v::AbstractDVec, α)
    empty!(w)
    sizehint!(w, length(v))
    for (key, val) in pairs(v)
        w[key] = val * α
    end
    return w
end

# copying multiplication with scalar
function Base.:*(α::T, x::AbstractDVec{<:Any,V}) where {T,V}
    return mul!(similar(x, promote_type(T, V)), x, α)
end
Base.:*(x::AbstractDVec, α) = α * x

"""
    add!(x::AbstractDVec,y::AbstactDVec)

Inplace add `x+y` and store result in `x`.
"""
@inline function add!(x::AbstractDVec{K}, y::AbstractDVec{K}) where {K}
    for (k, v) in pairs(y)
        x[k] += v
    end
    return x
end
add!(x::AbstractVector, y) = x .+= values(y)

@inline function LinearAlgebra.axpy!(α, x::AbstractDVec, y::AbstractDVec)
    for (k, v) in pairs(x)
        y[k] += α * v
    end
    return y
end

function LinearAlgebra.rmul!(x::AbstractDVec, α)
    for (k, v) in pairs(x)
        x[k] = v * α
    end
    return x
end

# BLAS-like function: y = α*x + β*y
function LinearAlgebra.axpby!(α, x::AbstractDVec, β, y::AbstractDVec)
    rmul!(y, β) # multiply every non-zero element
    axpy!(α, x, y)
    return y
end

function LinearAlgebra.dot(x::AbstractDVec, y::AbstractDVec)
    # For mixed value types
    result = zero(promote_type(valtype(x), valtype(y)))
    if length(x) < length(y) # try to save time by looking for the smaller vec
        for (key, val) in pairs(x)
            result += conj(val) * y[key]
        end
    else
        for (key, val) in pairs(y)
            result += conj(x[key]) * val
        end
    end
    return result # the type is promote_type(T1,T2) - could be complex!
end
# For MPI version see mpi_helpers.jl

# threaded dot()
function LinearAlgebra.dot(x::AbstractDVec{K,T1}, ys::NTuple{N, AbstractDVec{K,T2}}) where {N, K, T1, T2}
    results = zeros(promote_type(T1,T2), N)
    Threads.@threads for i in 1:N
        results[i] = x⋅ys[i]
    end
    return sum(results)
end

Base.isequal(x::AbstractDVec{K1}, y::AbstractDVec{K2}) where {K1,K2} = false
function Base.isequal(x::AbstractDVec{K}, y::AbstractDVec{K}) where {K}
    x === y && return true
    length(x) != length(y) && return false
    for (k, v) in pairs(x)
        !isequal(y[k], v) && return false
    end
    return true
end

Base.:(==)(x::AbstractDVec, y::AbstractDVec) = isequal(x, y)

# Define this type union for local (non-MPI) data
const DVecOrVec = Union{AbstractDVec,AbstractVector}

"""
Abstract supertype for projectors to be used in in lieu of DVecs or Vectors.
"""
abstract type AbstractProjector end

"""
    UniformProjector()
Represents a vector with all elements 1. To be used with [`dot()`](@ref).
Minimizes memory allocations.

```julia
UniformProjector()⋅v == sum(v)
dot(UniformProjector(), LO, v) == sum(LO*v)
```

See also [`ReportingStrategy`](@ref) for use
of projectors in FCIQMC.
"""
struct UniformProjector <: AbstractProjector end

LinearAlgebra.dot(::UniformProjector, y::DVecOrVec) = sum(values(y))
# a specialised fast and non-allocating method for
# `dot(::UniformProjector, A::AbstractHamiltonian, y)` is defined in `Hamiltonians.jl`
Base.getindex(::UniformProjector, add) = 1

"""
    NormProjector()
Results in computing the one-norm when used in `dot()`. E.g.
```julia
dot(NormProjector(),x)
-> norm(x,1)
```
`NormProjector()` thus represents the vector `sign.(x)`.

See also [`ReportingStrategy`](@ref) for use
of projectors in FCIQMC.
"""
struct NormProjector <: AbstractProjector end

LinearAlgebra.dot(::NormProjector, y::DVecOrVec) = norm(y, 1)

"""
    Norm2Projector()
Results in computing the two-norm when used in `dot()`. E.g.
```julia
dot(NormProjector(),x)
-> norm(x,2) # with type Float64
```

See also [`ReportingStrategy`](@ref) for use
of projectors in FCIQMC.
"""
struct Norm2Projector <: AbstractProjector end

LinearAlgebra.dot(::Norm2Projector, y::DVecOrVec) = norm(y, 2)
# NOTE that this returns a `Float64` opposite to the convention for
# dot to return the promote_type of the arguments.

"""
    Norm1ProjectorPPop()
Results in computing the one-norm per population when used in `dot()`. E.g.
```julia
dot(Norm1ProjectorPPop(),x)
-> norm(real.(x),1) + im*norm(imag.(x),1)
```

See also [`ReportingStrategy`](@ref) for use
of projectors in FCIQMC.
"""
struct Norm1ProjectorPPop <: AbstractProjector end

function LinearAlgebra.dot(::Norm1ProjectorPPop, y::DVecOrVec)
    T = float(valtype(y))
    if T <: Complex
        return T(sum(values(y)) do p
            abs(real(p)) + im*abs(imag(p))
        end)
    else
        return dot(NormProjector(), y)
    end
end

# NOTE that this returns a `Float64` opposite to the convention for
# dot to return the promote_type of the arguments.
# NOTE: This operation should work for `MPIData` and is MPI synchronizing

"""
    PopsProjector()
Results in computing the projection of one population on the other
when used in `dot()`. E.g.
```julia
dot(PopsProjector(),x)
-> real(x) ⋅ imag(x)
```

See also [`ReportingStrategy`](@ref) for use
of projectors in FCIQMC.
"""
struct PopsProjector <: AbstractProjector end

function LinearAlgebra.dot(::PopsProjector, y::DVecOrVec)
    T = float(real(valtype(y)))
    return T(sum(values(y)) do p
        real(p) * imag(p)
    end)
end

"""
    walkernumber(w)

Compute the number of walkers in `w`. It is used for updating the shift. Overload this
function for modifying population control.

In most cases `walkernumber(w)` is identical to `norm(w,1)`. For `AbstractDVec`s with
complex coefficients it reports the one norm separately for the real and the imaginary part
as a `ComplexF64`. See [`Norm1ProjectorPPop`](@ref).
"""
walkernumber(w) = walkernumber(StochasticStyle(w), w)
# use StochasticStyle trait for dispatch
walkernumber(::StochasticStyle, w) = dot(Norm1ProjectorPPop(), w)
# complex walkers as two populations
# the following default is fast and generic enough to be good for real walkers and

"""
    FrozenDVec

See: [`freeze`](@ref).
"""
struct FrozenDVec{K,V}
    pairs::Vector{Pair{K,V}}
end
Base.keytype(::FrozenDVec{K}) where {K} = K
Base.valtype(::FrozenDVec{<:Any,V}) where {V} = V
Base.eltype(::FrozenDVec{K,V}) where {K,V} = Pair{K,V}
Base.pairs(fd::FrozenDVec) = fd.pairs

freeze(dv) = FrozenDVec(collect(pairs(dv)))

freeze(p::AbstractProjector) = p

function LinearAlgebra.dot(fd::FrozenDVec, dv::AbstractDVec)
    result = zero(promote_type(valtype(fd), valtype(dv)))
    for (k, v) in pairs(fd)
        result += dv[k] ⋅ v
    end
    return result
end


text:
function getrk!(tempk,press,rk)

#=**************************************************************************

 This is a function in an ozone NOx formaldehyde mechanism, adapted from Fortran
 by Obin Sturm (UC Davis) in October 2019.
 The original comments and file description are below.

 created by: Mike Kleeman
             ECI 298 Air Quality Modeling
             April 1999

 The purpose of this subroutine is to calculate the reaction rate constants
 for the ozone NOx formaldehyde system.  The order of equations is assumed
 to be:
         1) NO2 + HV = NO + O
         2) O + O2 = O3
         3) O3 + NO = NO2 + O2
         4) HCHO + HV = 2 HO2. + CO
         5) HCHO + HV = H2 + CO
         6) HCHO + HO. = HO2. + CO + H2O
         7) HO2. + NO = HO. + NO2
         8) HO. + NO2 = HNO3
         9) HO2H + HV = 2 HO.
        10) HO2H + HO. = HO2. + H2O

 It is assumed that the photolytic constants are set at a constant value
 appropriate for mid-day in Southern California.

 Inputs:
  tempk - atmospheric temperature in degrees Kelvin
  press - atmospheric pressure in atmospheres

 Outputs:
  rk    - reaction rate constants in ppm min units

****************************************************************************=#

# function definition to convert from cm3 molec-1 s-1 to ppm min-1
    rkppmmin(rk1,press,tempk) = rk1 * 60 *  7.34e+15 * press / tempk

# begin by setting the rate constants for the photolytic reactions
    rk[1] = 0.5
    rk[4] = 0.015
    rk[5] = 0.022
    rk[9] = 0.0003

# reaction 2
    rk[2] = 6.0e-34 * ( tempk/300. )^(-2.3)
    rm = 7.34e+21 * press / tempk
    rk[2] = rk[2] * rm
    rk[2] = rkppmmin(rk[2],press,tempk)

# reaction 3
    rk[3] = 2.e-12 * exp(-1400/tempk)
    rk[3] = rkppmmin(rk[3],press,tempk)

# reaction 6
    rk[6] = 1.2e-14 * tempk^1 * exp(+287/tempk)
    rk[6] = rkppmmin(rk[6],press,tempk)

# reaction 7
    rk[7] = 3.7e-12 * exp(+250/tempk)
    rk[7] = rkppmmin(rk[7],press,tempk)

# reaction 8
# note: the following 2 lines are adapted from comments from the o.g. Fortran module
#      rk0 = 2.6e-30 * ( tempk/300 )^-3.2
#      rki = 2.4e-11 * ( tempk/300 )^-1.3
    rk0 = 2.0e-30 * ( tempk/300 )^(-3.0)
    rki = 2.5e-11 * ( tempk/300 )^0
    rm = 7.34e+21 * press / tempk
    rb = 0.6^(1.0 / ( 1.0 + (log10(rk0*rm/rki))^2.0 ))
    rk[8] = rk0 * rm * rb / (1.0 + rk0*rm/rki)
    rk[8] = rkppmmin(rk[8],press,tempk)

# reaction 10
    rk[10] = 3.3e-12 * exp(-200/tempk)
    rk[10] = rkppmmin(rk[10],press,tempk)

# return to the calling subroutine
    return rk
end


text:
using MPIMeasurements
using PyPlot

function plotData(u, daq, fignum)
    figure(fignum)
    clf()
    subplot(4,1,1)
    plot(vec(u[:,1,:,:]),"r")
    subplot(4,1,2)
    semilogy(freq[idx]/1000, abs.(rfft(vec(u[:,1,:,:])))[idx],"ro-")
    xlabel("freq / kHz")
    subplot(4,1,3)
    plot(vec(u[:,2,:,:]),"b")
    subplot(4,1,4)
    semilogy(freq[idx]/1000, abs.(rfft(vec(u[:,2,:,:])))[idx],"bo-")
    xlabel("freq / kHz")
    savefig("images/control$(fignum).png")
end

include("config.jl")

scanner = MPIScanner("TestSingleRP2D.toml")
daq = getDAQ(scanner)

numFreq = div(daq.params.rxNumSamplingPoints,2)+1
freq = collect(0:(numFreq-1))./(numFreq-1).*daq.params.rxBandwidth
idx = 70:80

params = toDict(daq.params)

params["studyName"] = "Test"
params["studyDescription"] = "A very cool measurement"
params["scannerOperator"] = "Tester"
params["dfStrength"] = [0.05, 0.05]
params["dfPhase"] = [pi/2, pi/2]
params["dfWaveform"] = "SINE"
params["acqNumFrames"] = 1
params["acqNumAverages"] = 1
params["controlLoopAmplitudeAccuracy"] = 0.0001

params["correctCrossCoupling"] = false
params["controlPhase"] = false
uNoControl = measurement(daq, params)
plotData(uNoControl, daq, 1)

scanner = MPIScanner("TestSingleRP2D.toml")
daq = getDAQ(scanner)
params["correctCrossCoupling"] = false
params["controlPhase"] = true
uNoCorr = measurement(daq, params)
plotData(uNoCorr, daq, 2)

#=
scanner = MPIScanner("TestSingleRP2D.toml")
daq = getDAQ(scanner)
params["correctCrossCoupling"] = true
params["controlPhase"] = true
uCorr = measurement(daq, params)
plotData(uCorr, daq, 3)

figure(4)
clf()
subplot(2,1,1)
semilogy(freq[idx]/1000, abs.(rfft(vec(uCorr[:,1,:,:])))[idx],"ro-")
semilogy(freq[idx]/1000, abs.(rfft(vec(uNoCorr[:,1,:,:])))[idx],"rx--")
legend(["Coupling Corr", "No Coupling Corr"])
ylabel("DF Signal / a.u.")
xlabel("freq / kHz")
subplot(2,1,2)
semilogy(freq[idx]/1000, abs.(rfft(vec(uCorr[:,2,:,:])))[idx],"bo-")
semilogy(freq[idx]/1000, abs.(rfft(vec(uNoCorr[:,2,:,:])))[idx],"bx--")
legend(("Coupling Corr", "No Coupling Corr"))
ylabel("DF Signal / a.u.")
xlabel("freq / kHz")
subplots_adjust(left=0.15,right=0.95, bottom=0.1, top=0.95, hspace=0.4,wspace=0.4)
savefig("images/control4.png")
=#


text:
"""
Filter function for polynomial chaos

- @args p: λ, Δ, ℓ, op, ..., mode

"""
function Base.filter!(u::T, p...) where {T<:AbstractArray{<:AbstractFloat,1}}

    q0 = eachindex(u) |> first
    q1 = eachindex(u) |> last

    λ = p[1]
    mode = p[end]

    if mode == :l2
        if length(p) > 2
            Δ = p[2]
            op = p[3]
            filter_l2!(u, op, λ, Δ)
        else
            filter_l2!(u, λ)
        end
    elseif mode == :l1
        for i = q0+1:q1
            sc = 1.0 - 5.0 * λ * i * (i - 1) * ℓ[i] / abs(u[i])
            if sc < 0.0
                sc = 0.0
            end
            u[i] *= sc
        end
    elseif mode == :lasso
        nr = length(u)
        _λ = abs(u[end]) / (nr * (nr - 1) * ℓ[end])
        for i = q0+1:q1
            sc = 1.0 - _λ * i * (i - 1) * ℓ[i] / abs(u[i] + 1e-8)
            if sc < 0.0
                sc = 0.0
            end
            u[i] *= sc
        end
    elseif mode == :dev
        nr = length(u)
        for i = q0+1:q1
            _λ = λ * abs(u[i] / (u[1] + 1e-10)) / (nr * (nr - 1) * ℓ[i])
            u[i] /= (1.0 + _λ * i^2 * (i - 1)^2)
        end
    else
        throw("unavailable filter mode")
    end

end

function Base.filter!(u::AbstractArray{<:AbstractFloat,2}, dim::Integer, p...)

    if dim == 1
        for j in axes(u, 2)
            _u = @view u[:, j]
            filter!(_u, p...)
        end
    elseif dim == 2
        for i in axes(u, 1)
            _u = @view u[i, :]
            filter!(_u, p...)
        end
    end

end

function Base.filter!(u::AbstractArray{<:AbstractFloat,3}, dim::Integer, p...)

    if dim == 1
        for k in axes(u, 3), j in axes(u, 2)
            _u = @view u[:, j, k]
            filter!(_u, p...)
        end
    elseif dim == 2
        for k in axes(u, 3), i in axes(u, 1)
            _u = @view u[i, :, k]
            filter!(_u, p...)
        end
    elseif dim == 3
        for j in axes(u, 2), i in axes(u, 1)
            _u = @view u[i, j, :]
            filter!(_u, p...)
        end
    end

end

function filter_l2!(u::T, λ) where {T<:AbstractArray{<:AbstractFloat,1}}

    q0 = eachindex(u) |> first
    q1 = eachindex(u) |> last

    for i = q0+1:q1
        u[i] /= (1.0 + λ * i^2 * (i - 1)^2)
    end

    return nothing

end

function filter_l2!(
    u::T1,
    op::T2,
    λ,
    Δ,
) where {T1<:AbstractArray{<:AbstractFloat,1},T2<:AbstractOrthoPoly}

    q0 = eachindex(u) |> first
    q1 = eachindex(u) |> last

    uRan = evaluatePCE(u, op.quad.nodes, op)
    δ = 0.5 * (maximum(uRan) - minimum(uRan))
    _λ = λ * (exp(δ / Δ) - 1.0)
    #_λ = λ * δ / Δ

    for i = q0+1:q1
        u[i] /= (1.0 + _λ * i^2 * (i - 1)^2)
    end

    return nothing

end


text:
using Test, MLToolkit
using Distributions: Beta, Dirichlet
using Statistics: mean

@testset "Beta" begin
    d = 10
    n = 500_000

    @testset "Kumaraswamy" begin
        for (a, b) = [(1.0, 1.0), (1.0, 2.0), (2.0, 1.0)]
            beta_mean = mean(rand(Beta(a, b), n))

            a_mat = Matrix{FT}(undef, 1, 1)
            b_mat = Matrix{FT}(undef, 1, 1)
            a_mat[1,1] = a
            b_mat[1,1] = b
            kuma = BatchKumaraswamy{AT}(a_mat, b_mat)

            x = Array(hcat([rand(kuma) for _ = 1:n]...))
            @test mean(x) ≈ beta_mean atol=ATOL_RAND

            x = Array(hcat([exp.(logrand(kuma)) for _ = 1:n]...))
            @test mean(x) ≈ beta_mean atol=ATOL_RAND

            kuma2 = BatchKumaraswamy(a, b)

            x = Array(rand(kuma2, n))
            @test mean(x) ≈ beta_mean atol=ATOL_RAND

            x = Array(exp.(logrand(kuma2, n)))
            @test mean(x) ≈ beta_mean atol=ATOL_RAND
        end
        @warn "`logpdf(kuma::BatchKumaraswamy, x)` is not tested."
    end

    @testset "Dirichlet" begin
        for _  = 1:NUM_RANDTESTS
            dir1 = Dirichlet(rand(FT, d).^2 .+ ones(FT, d))
            dir2 = Dirichlet(rand(FT, d).^2 .+ ones(FT, d))

            x = rand(dir1, n)

            kl_12 = mean(logpdf(dir1, x) - logpdf(dir2, x))

            @test kldiv(dir1, dir2) ≈ kl_12 atol=(d * ATOL_RAND)
        end
    end

    @testset "Beta" begin
        # Test cases are taken from https://en.wikipedia.org/wiki/Beta_distribution
        for (a1, b1, a2, b2, ans) = [(1.0, 1.0, 3.0, 3.0, 0.598803),
                                     (3.0, 3.0, 1.0, 1.0, 0.267864),
                                     (3.0, 0.5, 0.5, 3.0, 7.21574),
                                     (0.5, 3.0, 3.0, 0.5, 7.21574)]
            @test kldiv(BatchBeta(a1, b1), BatchBeta(a2, b2)) ≈ ans atol=100ATOL
            @test kldiv(Dirichlet([a1, b1]), Dirichlet([a2, b2])) ≈ ans atol=100ATOL
        end

        param_list = []
        for _ = 1:NUM_RANDTESTS
            (a1, b1) = 1.0 .+ rand(2).^2
            (a2, b2) = [a1, b1] .+ rand(2).^2
            push!(param_list, (a1, b1, a2, b2))
        end

        @testset "KL[Dirichlet/Beta($a1,$b1) || Dirichlet/Beta($a2,$b2)]" for (a1, b1, a2, b2) in param_list
            # Generate a1, b1, a2, b2 that are not too far away from each other
            (a1, b1) = 1.0 .+ rand(2).^2
            (a2, b2) = [a1, b1] .+ rand(2).^2

            @test kldiv(Dirichlet([a1, b1]), Dirichlet([a2, b2])) ≈
                  kldiv(BatchBeta(a1, b1), BatchBeta(a2, b2)) atol=ATOL_RAND
        end

        @testset "KL[Kuma($a1,$b1) || Beta($a2,$b2)]" for (a1, b1, a2, b2) in param_list
            # Generate a1, b1, a2, b2 that are not too far away from each other
            (a1, b1) = 1.0 .+ rand(2).^2
            (a2, b2) = [a1, b1] .+ rand(2).^2

            kuma = BatchKumaraswamy(a1, b1)
            bb = BatchBeta(a2, b2)

            x = vec([rand(kuma) for _ = 1:n])
            kl_mc = mean(logpdf(kuma, x) - logpdf.(Beta(a2, b2), x))
            @test kldiv(kuma, bb) ≈ kl_mc atol=(2 * ATOL_RAND)
        end

        @warn "`logpdf(bb::BatchBeta, x)` is not tested."
    end
end


text:
# Lazy structure for an array whose real and fourier part can be accessed by A.real and A.fourier
# Used to represent potentials and densities
# For normalization conventions, see PlaneWaveBasis.jl

import Base.getproperty, Base.setproperty!, Base.propertynames
import Base.*, Base.+, Base./
import Base.eltype, Base.size, Base.length

## TODO these functions are minimally optimized for the moment
## TODO this always allocates both real and fourier parts, which might be a bit wasteful
## TODO implement broadcasting. This seems to be non-trivial, so I just implement a few simple methods
## TODO implement compressed representation in Fourier space of real arrays

@enum RFA_state RFA_only_real RFA_only_fourier RFA_both

"""
A structure to facilitate manipulations of an array of real-space type T, in both real
and fourier space. Create with `from_real` or `from_fourier`, and access
with `A.real` and `A.fourier`.
"""
mutable struct RealFourierArray{T <: Real,
                                TRealArray <: AbstractArray{T, 3},
                                TFourierArray <: AbstractArray{Complex{T}, 3}}
    basis::PlaneWaveBasis{T}
    _real::TRealArray
    _fourier::TFourierArray
    _state::RFA_state
end
# Type of the real part
Base.eltype(::RealFourierArray{T}) where {T} = T
Base.size(array::RealFourierArray) = size(array._real)
Base.size(array::RealFourierArray, i) = size(array._real, i)
Base.length(array::RealFourierArray) = length(array._real)

# Constructors
RealFourierArray(basis, real, fourier) = RealFourierArray(basis, real, fourier, RFA_both)
function RealFourierArray(basis::PlaneWaveBasis{T}) where {T}
    from_real(basis, zeros(T, basis.fft_size...))
end

function from_real(basis, real_part::AbstractArray{T}) where {T <: Real}
    RealFourierArray(basis, real_part, similar(real_part, complex(T)), RFA_only_real)
end
function from_fourier(basis, fourier_part::AbstractArray{T}; check_real=true) where {T <: Complex}
    if check_real
        # Go through G vectors and check c_{-G} = (c_G)' (if both G and -G are in the grid)
        # TODO check it's reasonably fast so we can make it the default
        # arr[1] is G=0, arr[1] is G=1, arr[N] is G=-1.
        # So 1 -> 1, 2 -> N, ..., N -> 2
        reflect(i, N) = i == 1 ? 1 : N-i+2
        # Eg if N = 3, we get [0 1 -1], if N = 4, we get [0 1 2 -1] => we check to div(N+1,2)
        for i = 1:div(basis.fft_size[1]+1, 2)
            for j = 1:div(basis.fft_size[2]+1, 2)
                for k = 1:div(basis.fft_size[3]+1, 2)
                    err = abs(fourier_part[i, j, k] -
                              conj(fourier_part[reflect(i, end), reflect(j, end), reflect(k, end)]))
                    err > sqrt(eps(real(T))) && error("Input array not real")
                end
            end
        end
    end
    RealFourierArray(basis, similar(fourier_part, real(T)), fourier_part, RFA_only_fourier)
end

function Base.propertynames(array::RealFourierArray, private=false)
    ret = [:basis, :real, :fourier]
    private ? append!(ret, fieldnames(RealFourierArray)) : ret
end

function Base.getproperty(A::RealFourierArray{T, TReal, TFourier}, x::Symbol) where {T, TReal, TFourier}
    if x == :real
        if A._state == RFA_only_fourier
            r = G_to_r(A.basis, A._fourier)
            if eltype(A) <: Real
                A._real .= real(r)
            else
                A._real .= r
            end
            setfield!(A, :_state, RFA_both)
        end
        return A._real::TReal  # the compiler does not figure it out otherwise
    elseif x == :fourier
        if A._state == RFA_only_real
            r_to_G!(A._fourier, A.basis, complex.(A._real))
            setfield!(A, :_state, RFA_both)
        end
        return A._fourier::TFourier
    else
        getfield(A, x)
    end
end
function Base.setproperty!(::RealFourierArray, ::Any)
    error("RealFourierArray is intended to be read-only." *
          "(This can be bypassed by `setfield!` if you really want to).")
end

# Algebraic operations
function Base.:+(A::RealFourierArray, B::RealFourierArray)
    from_real(A.basis, A.real + B.real)
end
function Base.:-(A::RealFourierArray, B::RealFourierArray)
    from_real(A.basis, A.real - B.real)
end
function Base.:*(α::Number, A::RealFourierArray)
    from_real(A.basis, α * A.real)
end
Base.:*(A::RealFourierArray, α::Number) = α * A
Base.:/(A::RealFourierArray, α::Number) = A * inv(α)


text:
#  ____                                    ____  ____  _
# |  _ \ _   _ _ __   __ _ _ __ ___   ___ |  _ \/ ___|| |
# | | | | | | | '_ \ / _` | '_ ` _ \ / _ \| | | \___ \| |
# | |_| | |_| | | | | (_| | | | | | | (_) | |_| |___) | |___
# |____/ \__, |_| |_|\__,_|_| |_| |_|\___/|____/|____/|_____|
#        |___/

# This file defines some types and serializers which allow for the
# easy writing and manipulation of dynamodb conditional write
# expressions ("only write if X is true"), key condition expressions
# ("only read objects from the disk which match X"), and filter
# expressions ("ignore values read from disk which don't match X").

# The reason this is preferable than just writing queries by hand as
# strings is that dynamo requires a strange serialization format for
# types, and has surprising keyword requirements for unescaped
# fieldnames. Doing it this way makes those bits non-issues.

# a reference to a dynamodb record field
abstract type DynamoValue end
abstract type DynamoReference <: DynamoValue end

struct DynamoAttribute <: DynamoReference
    name :: AbstractString
end
attribute(name :: AbstractString) = DynamoAttribute(name)
attribute(name :: Symbol) = attribute(string(name))
attr(name :: AbstractString) = attribute(name)
attr(name :: Symbol) = attribute(name)

# a reference to dynamo sub-documents... eg "foo.bar" in {foo => {bar => 3}}
struct NestedDynamoAttribute <: DynamoReference
    attrs :: Array{DynamoAttribute}
end
attribute(names...) = NestedDynamoAttribute([attr(e) for e=names])
attr(names...) = attribute(names...)

struct DynamoListElement <: DynamoReference
    attr :: DynamoReference
    idx
end
import Base.getindex
getindex(attr :: DynamoReference, idx :: Int64) = DynamoListElement(attr, idx)
getindex(attr :: DynamoReference, idx :: Int32) = DynamoListElement(attr, idx)

struct DynamoLiteralValue <: DynamoValue
    val
end
value_or_literal(a :: DynamoValue) = a
value_or_literal(other) = DynamoLiteralValue(other)

#   ____                _ _ _   _                   _
#  / ___|___  _ __   __| (_) |_(_) ___  _ __   __ _| |
# | |   / _ \| '_ \ / _` | | __| |/ _ \| '_ \ / _` | |
# | |__| (_) | | | | (_| | | |_| | (_) | | | | (_| | |
#  \____\___/|_| |_|\__,_|_|\__|_|\___/|_| |_|\__,_|_|
#  _____                              _
# | ____|_  ___ __  _ __ ___  ___ ___(_) ___  _ __  ___
# |  _| \ \/ / '_ \| '__/ _ \/ __/ __| |/ _ \| '_ \/ __|
# | |___ >  <| |_) | | |  __/\__ \__ \ | (_) | | | \__ \
# |_____/_/\_\ .__/|_|  \___||___/___/_|\___/|_| |_|___/
#            |_|

# https://docs.aws.amazon.com/amazondynamodb/latest/developerguide/Expressions.SpecifyingConditions.html
# https://docs.aws.amazon.com/amazondynamodb/latest/developerguide/ExpressionPlaceholders.html

abstract type CEBoolean end
abstract type CEFnVal end
const  CEVal  =  Union{DynamoValue, CEFnVal}

struct CETrue <: CEBoolean ; end
no_conditions() = CETrue()

struct CESize <: CEFnVal
    attr :: DynamoReference
end
import Base.size
size(attr :: DynamoReference) = CESize(attr)

struct CEBeginsWith <: CEBoolean
    attr :: DynamoReference
    val
end
begins_with(attr :: DynamoReference, val) = CEBeginsWith(attr, value_or_literal(val))

struct CEContains <: CEBoolean
    attr :: DynamoReference
    val
end
import Base.contains
contains(attr :: DynamoReference, val) = CEContains(attr, value_or_literal(val))

struct CEAttributeType <: CEBoolean
    attr :: DynamoReference
    ty :: AbstractString
end
is_string(attr :: DynamoReference) = CEAttributeType(attr, "S")
is_string_set(attr :: DynamoReference) = CEAttributeType(attr, "SS")
is_real(attr :: DynamoReference) = CEAttributeType(attr, "N")
is_real_set(attr :: DynamoReference) = CEAttributeType(attr, "NS")
is_binary(attr :: DynamoReference) = CEAttributeType(attr, "B")
is_binary_set(attr :: DynamoReference) = CEAttributeType(attr, "BS")
is_bool(attr :: DynamoReference) = CEAttributeType(attr, "BOOL")
is_null(attr :: DynamoReference) = CEAttributeType(attr, "NULL")
is_list(attr :: DynamoReference) = CEAttributeType(attr, "L")
is_map(attr :: DynamoReference) = CEAttributeType(attr, "M")
is_document(attr :: DynamoReference) = CEAttributeType(attr, "M")

struct CEExists <: CEBoolean
    attr :: DynamoReference
end
exists(attr :: DynamoReference) = CEExists(attr)

struct CENotExists <: CEBoolean
    attr :: DynamoReference
end
not_exists(attr :: DynamoReference) = CENotExists(attr)

struct CEBinaryOp <: CEBoolean
    op :: AbstractString
    lhs
    rhs
end
import Base.<
<(lhs :: CEVal, rhs :: CEVal) = CEBinaryOp("<", lhs, rhs)
<(lhs :: CEVal, rhs) = CEBinaryOp("<", lhs, value_or_literal(rhs))
<(lhs, rhs :: CEVal) = CEBinaryOp("<", value_or_literal(lhs), rhs)

import Base.>
>(lhs :: CEVal, rhs :: CEVal) = CEBinaryOp(">", lhs, rhs)
>(lhs :: CEVal, rhs) = CEBinaryOp(">", lhs, value_or_literal(rhs))
>(lhs, rhs :: CEVal) = CEBinaryOp(">", value_or_literal(lhs), rhs)

import Base.<=
<=(lhs :: CEVal, rhs :: CEVal) = CEBinaryOp("<=", lhs, rhs)
<=(lhs :: CEVal, rhs) = CEBinaryOp("<=", lhs, value_or_literal(rhs))
<=(lhs, rhs :: CEVal) = CEBinaryOp("<=", value_or_literal(lhs), rhs)

import Base.>=
>=(lhs :: CEVal, rhs :: CEVal) = CEBinaryOp(">=", lhs, rhs)
>=(lhs :: CEVal, rhs) = CEBinaryOp(">=", lhs, value_or_literal(rhs))
>=(lhs, rhs :: CEVal) = CEBinaryOp(">=", value_or_literal(lhs), rhs)

import Base.==
==(lhs :: CEVal, rhs :: CEVal) = CEBinaryOp("=", lhs, rhs)
==(lhs :: CEVal, rhs :: WeakRef) = CEBinaryOp("=", lhs, value_or_literal(rhs.value))
==(lhs :: CEVal, rhs) = CEBinaryOp("=", lhs, value_or_literal(rhs))
==(lhs :: WeakRef, rhs :: CEVal) = CEBinaryOp("=", value_or_literal(lhs.value), rhs)
==(lhs, rhs :: CEVal) = CEBinaryOp("=", value_or_literal(lhs), rhs)

import Base.!=
!=(lhs :: CEVal, rhs :: CEVal) = CEBinaryOp("<>", lhs, rhs)
!=(lhs :: CEVal, rhs) = CEBinaryOp("<>", lhs, value_or_literal(rhs))
!=(lhs, rhs :: CEVal) = CEBinaryOp("<>", value_or_literal(lhs), rhs)

and(lhs :: CETrue, rhs :: CETrue) = lhs
and(lhs :: CEBoolean, rhs :: CETrue) = lhs
and(lhs :: CETrue, rhs :: CEBoolean) = rhs
and(lhs :: CEBoolean, rhs :: CEBoolean) = CEBinaryOp("AND", lhs, rhs)
and(lhs :: DynamoReference, rhs :: CEBoolean) = CEBinaryOp("AND", lhs, rhs)
and(lhs :: CEBoolean, rhs :: DynamoReference) = CEBinaryOp("AND", lhs, rhs)
and(lhs :: DynamoReference, rhs :: DynamoReference) = CEBinaryOp("AND", lhs, rhs)

or(lhs :: CETrue, rhs :: CETrue) = lhs
or(lhs :: CETrue, rhs :: CEBoolean) = lhs
or(lhs :: CEBoolean, rhs :: CETrue) = rhs
or(lhs :: CETrue, rhs :: DynamoReference) = lhs
or(lhs :: DynamoReference, rhs :: CETrue) = rhs
or(lhs :: CEBoolean, rhs :: CEBoolean) = CEBinaryOp("OR", lhs, rhs)
or(lhs :: DynamoReference, rhs :: CEBoolean) = CEBinaryOp("OR", lhs, rhs)
or(lhs :: CEBoolean, rhs :: DynamoReference) = CEBinaryOp("OR", lhs, rhs)
or(lhs :: DynamoReference, rhs :: DynamoReference) = CEBinaryOp("OR", lhs, rhs)

struct CENot <: CEBoolean
    exp
end

import Base.(!)
!(exp :: CEBoolean) = CENot(exp)
!(exp :: DynamoAttribute) = CENot(exp)
not(exp :: CEBoolean) = CENot(exp)
not(exp :: DynamoAttribute) = CENot(exp)

struct CEBetween <: CEBoolean
    attr :: DynamoAttribute
    min
    max
end
between(attr :: DynamoAttribute, min, max) =
    CEBetween(attr, value_or_literal(min), value_or_literal(max))

#  _   _           _       _
# | | | |_ __   __| | __ _| |_ ___
# | | | | '_ \ / _` |/ _` | __/ _ \
# | |_| | |_) | (_| | (_| | ||  __/
#  \___/| .__/ \__,_|\__,_|\__\___|
#       |_|
#  _____                              _
# | ____|_  ___ __  _ __ ___  ___ ___(_) ___  _ __  ___
# |  _| \ \/ / '_ \| '__/ _ \/ __/ __| |/ _ \| '_ \/ __|
# | |___ >  <| |_) | | |  __/\__ \__ \ | (_) | | | \__ \
# |_____/_/\_\ .__/|_|  \___||___/___/_|\___/|_| |_|___/
#            |_|

# https://docs.aws.amazon.com/amazondynamodb/latest/developerguide/Expressions.Modifying.html

# SET
# REMOVE
# ADD
# DELETE

abstract type DynamoUpdateExpression end
abstract type DUFnVal end
const  DUVal  =  Union{DynamoValue, DUFnVal}

struct DefaultValue <: DUFnVal
    attr :: DynamoReference
    val
end
get_or_else(attr :: DynamoReference, val) = DefaultValue(attr, value_or_literal(val))

struct ListAppend <: DynamoUpdateExpression
    attr :: DynamoReference
    val
end
append_to_list(attr :: DynamoReference, val) = ListAppend(attr, value_or_literal(val))

struct AssignExpression <: DynamoUpdateExpression
    attr :: DynamoReference
    val
end
assign(attr :: DynamoReference, val) = AssignExpression(attr, value_or_literal(val))

struct SetAddExpression <: DynamoUpdateExpression
    attr :: DynamoReference
    val
end
add_to_set(attr :: DynamoAttribute, val) = SetAddExpression(attr, value_or_literal(val))

struct SetRemoveExpression <: DynamoUpdateExpression
    attr :: DynamoReference
    val
end
remove_from_set(attr :: DynamoReference, val) = SetRemoveExpression(attr, value_or_literal(val))

struct DeleteExpression <: DynamoUpdateExpression
    attr :: DynamoReference
end
delete(attr :: DynamoReference) = DeleteExpression(attr)

# Dictionary of references, used to keep the serialization orderly and
# compact.
mutable struct DynamoAttrAndValReferences
    gensym_n :: Int64

    attrs
    attrs_reversed # used to avoid double-sending duplicate values

    vals
    vals_reversed # used to avoid double-sending duplicate values
end

mutable struct UpdateWriter
    refs :: DynamoAttrAndValReferences

    sets :: Array{Any}
    removes :: Array{Any}
    adds :: Array{Any}
    deletes :: Array{Any}
end

function incr_gensym(w :: DynamoAttrAndValReferences)
    n = w.gensym_n
    w.gensym_n += 1
    n
end

refs_tracker() = DynamoAttrAndValReferences(1, Dict{AbstractString, Any}(), Dict{Any, AbstractString}(),
                                            Dict{AbstractString, Any}(), Dict{Any, AbstractString}())

can_write_expression(::Void) = false
can_write_expression(::CETrue) = false
can_write_expression(x) = true

# Okay, with that out of the way, let's get serializin'
write_expression(w :: DynamoAttrAndValReferences, e :: CESize) =
    "size($(write_expression(w, e.attr)))"
write_expression(w :: DynamoAttrAndValReferences, e :: CEBeginsWith) =
    "begins_with($(write_expression(w, e.attr)), $(write_expression(w, e.val)))"
write_expression(w :: DynamoAttrAndValReferences, e :: CEContains) =
    "contains($(write_expression(w, e.attr)), $(write_expression(w, e.val)))"
write_expression(w :: DynamoAttrAndValReferences, e :: CEAttributeType) =
    "attribute_type($(write_expression(w, e.attr)), \"$(e.ty)\")"
write_expression(w :: DynamoAttrAndValReferences, e :: CENotExists) =
    "attribute_not_exists($(write_expression(w, e.attr)))"
write_expression(w :: DynamoAttrAndValReferences, e :: CEExists) =
    "attribute_exists($(write_expression(w, e.attr)))"
write_expression(w :: DynamoAttrAndValReferences, e :: CEBinaryOp) =
    "($(write_expression(w, e.lhs))) $(e.op) ($(write_expression(w, e.rhs)))"
write_expression(w :: DynamoAttrAndValReferences, e :: CEBetween) =
    "$(write_expression(w, e.attr)) BETWEEN $(write_expression(w, e.min)) AND $(write_expression(w, e.max))"
write_expression(w :: DynamoAttrAndValReferences, e :: CENot) =
    "NOT ($(write_expression(w, e.exp)))"

serialize_expression(expr :: CEBoolean, refs = refs_tracker()) =
    write_expression(refs, expr)

write_expression(w :: DynamoAttrAndValReferences, v :: DefaultValue) =
    "if_not_exists($(write_expression(w, v.attr)), $(write_expression(w, v.val)))"
write_update(w :: UpdateWriter, v :: ListAppend) =
    push!(w.sets, "$(write_expression(w.refs, v.attr)) = list_append($(write_expression(w.refs, v.attr)), $(write_expression(w.refs, v.val)))")
write_update(w :: UpdateWriter, v :: AssignExpression) =
    push!(w.sets, "$(write_expression(w.refs, v.attr)) = $(write_expression(w.refs, v.val))")
write_update(w :: UpdateWriter, v :: SetAddExpression) =
    push!(w.adds, "$(write_expression(w.refs, v.attr)) $(write_expression(w.refs, v.val))")
write_update(w :: UpdateWriter, v :: SetRemoveExpression) =
    push!(w.deletes, "$(write_expression(w.refs, v.attr)) $(write_expression(w.refs, v.val))")
write_update(w :: UpdateWriter, v :: DeleteExpression) =
    push!(w.removes, "$(write_expression(w.refs, v.attr))")

function serialize_updates(arr :: Array, refs = refs_tracker())
    w = UpdateWriter(refs, [], [], [], [])
    for e=arr
        write_update(w, e)
    end

    parts = []
    if w.sets != []
        push!(parts, string("SET ", join(w.sets, ", ")))
    end
    if w.removes != []
        push!(parts, string("REMOVE ", join(w.removes, ", ")))
    end
    if w.adds != []
        push!(parts, string("ADD ", join(w.adds, ", ")))
    end
    if w.deletes != []
        push!(parts, string("DELETE ", join(w.deletes, ", ")))
    end

    join(parts, " ")
end

write_expression(w :: DynamoAttrAndValReferences, nested :: NestedDynamoAttribute) =
    join([write_expression(w, e) for e=nested.attrs], ".")

write_expression(w :: DynamoAttrAndValReferences, v :: DynamoListElement) =
    "$(write_expression(w, v.attr))[$(v.idx)]"

function write_expression(w :: DynamoAttrAndValReferences, val :: DynamoLiteralValue)
    if haskey(w.vals_reversed, val.val)
        return w.vals_reversed[val.val]
    end

    name = ":$(incr_gensym(w))"
    w.vals[name] = null_or_val(val.val)
    w.vals_reversed[val.val] = name
    name
end

function write_expression(w :: DynamoAttrAndValReferences, attr :: DynamoAttribute)
    if haskey(w.attrs_reversed, attr.name)
        return w.attrs_reversed[attr.name]
    end

    name = "#$(incr_gensym(w))"
    w.attrs[name] = attr.name
    w.attrs_reversed[attr.name] = name
    name
end


text:
# This file is a part of Julia. License is MIT: https://julialang.org/license

#const Profile1 = Base.root_module(Base.PkgId(Base.UUID("9abbd945-dff8-562f-b5e8-e1ebf5ef1b79"), "Profile"))
using Profile
using Profile: StackFrameTree, LineInfoDict, LineInfoFlatDict, StackFrame

function tree_format(io::IO, dom::Vector{Node}, frame::StackFrameTree)
    li = frame.frame
    let id = startNode(dom, io, li, li.from_c ? "frame-c" : "frame")
        print(io, frame.count)
        print(io, ' ')
        show(io, li; full_path = true)
        #if li == Profile.UNKNOWN
        #    print(io, " unknown stackframe")
        #else
        #    if li.line == li.pointer
        #        print(io, " unknown function (pointer: 0x",
        #            string(li.pointer, base = 16, pad = 2*sizeof(Ptr{Cvoid})),
        #            ")")
        #    else
        #        fname = string(li.func)
        #        if !li.from_c && li.linfo !== nothing
        #            fname = sprint(Profile.show_spec_linfo, li)
        #        end
        #        print(io, " ", string(li.file), ":",
        #            li.line == -1 ? "?" : string(li.line),
        #            "; ", fname)
        #    end
        #end
        endNode(dom, io, id)
    end
end

# Print a "branch" starting at a particular level. This gets called recursively.
function tree(io::IO, dom::Vector{Node}, bt::StackFrameTree)
    isempty(bt.down) && return
    let id = startNode(dom, io, bt, "level")
        # Order the line information
        nexts = collect(values(bt.down))
        lilist = collect(frame.frame for frame in nexts)
        # Generate the string for each line
        # Recurse to the next level
        for i in Profile1.liperm(lilist)
            nextChild(dom, io, id)
            down = nexts[i]
            tree_format(io, dom, down)
            tree(io, dom, down)
        end
        endNode(dom, io, id)
    end
end

function drop_shadow!(node::StackFrameTree)
    empty!(node.builder_key)
    empty!(node.builder_value)
    foreach(drop_shadow!, values(node.down))
    nothing
end

function tree(io::IO, dom::Vector{Node}, data::Vector{UInt64}, lidict::LineInfoFlatDict)
    combine = true
    if combine
        root = Profile1.tree!(StackFrameTree{StackFrame}(), data, lidict, true)
    else
        root = Profile1.tree!(StackFrameTree{UInt64}(), data, lidict, true)
    end
    drop_shadow!(root)
    let id = startNode(dom, io, root, "ProfileTree")
        tree(io, dom, root)
        endNode(dom, io, id)
    end
end

function tree(io::IO, dom::Vector{Node}, data::Vector{UInt64}, lidict::LineInfoDict)
    newdata, newdict = Profile.flatten(data, lidict)
    tree(io, dom, newdata, newdict)
    nothing
end

function richprofile(data::Vector{<:Unsigned} = Profile.fetch(), lidict::LineInfoDict = Profile.getdict(data))
    dom = Vector{Node}()
    io = IOBuffer()
    tree(convert(IOContext, io), dom, data, lidict)
    return annotate!(take!(io), dom)
end

max_stars_repo_name: vtjnash/Firelight


text:
# Support for ZeroMQ, a network and interprocess communication library

module ZMQ

using Base
import Base.convert, Base.ref, Base.get, Base.bytestring

export
    #Types
    ZMQStateError,ZMQContext,ZMQSocket,ZMQMessage,
    #functions
    version,close, get, set, bind, connect,send,recv,convert, ref,
    #Constants
    ZMQ_IO_THREADS,ZMQ_MAX_SOCKETS,ZMQ_PAIR,ZMQ_PUB,ZMQ_SUB,ZMQ_REQ,ZMQ_REP,ZMQ_DEALER,ZMQ_DEALER,ZMQ_PULL,ZMQ_PUSH,ZMQ_XPUB,ZMQ_XPUB,ZMQ_XREQ,ZMQ_XREP,ZMQ_UPSTREAM,ZMQ_DOWNSTREAM,ZMQ_MORE,ZMQ_MORE,ZMQ_SNDMORE,ZMQ_POLLIN,ZMQ_POLLOUT,ZMQ_POLLERR,ZMQ_STREAMER,ZMQ_FORWARDER,ZMQ_QUEUE

const ZMQ_MAX_VSM_SIZE = 30
# This next should be replaced by a ccall, when packages can have C code
const _jl_zmq_msg_t_size = ZMQ_MAX_VSM_SIZE + sizeof(Uint) + 2

# A server will report most errors to the client over a ZMQSocket, but
# errors in ZMQ state can't be reported because the socket may be
# corrupted. Therefore, we need an exception type for errors that
# should be reported locally.
type ZMQStateError <: Exception
    msg::String
end
show(io, thiserr::ZMQStateError) = print(io, "ZMQ: ", thiserr.msg)

# Basic functions
function jl_zmq_error_str()
    errno = ccall((:zmq_errno, :libzmq), Int32, ())
    c_strerror = ccall ((:zmq_strerror, :libzmq), Ptr{Uint8}, (Int32,), errno)
    if c_strerror != C_NULL
        strerror = bytestring(c_strerror)
        return strerror
    else
        return "Unknown error"
    end
end

let major = zeros(Int32, 1), minor = zeros(Int32, 1), patch = zeros(Int32, 1)
global version
function version()
    ccall((:zmq_version, :libzmq), Void, (Ptr{Int32}, Ptr{Int32}, Ptr{Int32}), major, minor, patch)
    return (major[1], minor[1], patch[1])
end
end

_zmq_major, _zmq_minor, _zmq_path = version()

## Contexts ##
# Provide the same constructor API for version 2 and version 3, even
# though the underlying functions are changing
if _zmq_major < 3
    global ZMQContext
    type ZMQContext
        data::Ptr{Void}

        function ZMQContext(n::Integer)
            p = ccall((:zmq_init, :libzmq), Ptr{Void},  (Int32,), n)
            if p == C_NULL
    	    throw(ZMQStateError(jl_zmq_error_str()))
            end
            zctx = new(p)
            finalizer(zctx, close)
            return zctx
        end
    end
else  # Versions 3 and higher
    global ZMQContext
    type ZMQContext
        data::Ptr{Void}

        function ZMQContext(n::Integer)
            p = ccall((:zmq_ctx_new, :libzmq), Ptr{Void},  ())
            if p == C_NULL
    	    throw(ZMQStateError(jl_zmq_error_str()))
            end
            zctx = new(p)
            finalizer(zctx, close)
            set(zctx, ZMQ_IO_THREADS, n)
            return zctx
        end
    end
end
ZMQContext() = ZMQContext(1)
if _zmq_major < 3
    global close
    function close(ctx::ZMQContext)
        rc = ccall((:zmq_term, :libzmq), Int32,  (Ptr{Void},), ctx.data)
        if rc != 0
            throw(ZMQStateError(jl_zmq_error_str()))
        end
    end
else  # Versions 3 and higher
    global close
    function close(ctx::ZMQContext)
        rc = ccall((:zmq_ctx_destroy, :libzmq), Int32,  (Ptr{Void},), ctx.data)
        if rc != 0
            throw(ZMQStateError(jl_zmq_error_str()))
        end
    end
    global get
    function get(ctx::ZMQContext, option::Integer)
        val = ccall((:zmq_ctx_get, :libzmq), Int32, (Ptr{Void}, Int32), ctx.data, option)
        if val < 0
            throw(ZMQStateError(jl_zmq_error_str()))
        end
        return val
    end
    global set
    function set(ctx::ZMQContext, option::Integer, value::Integer)
        rc = ccall((:zmq_ctx_set, :libzmq), Int32, (Ptr{Void}, Int32, Int32), ctx.data, option, value)
        if rc != 0
            throw(ZMQStateError(jl_zmq_error_str()))
        end
    end
end
term(ctx::ZMQContext) = close(ctx)

## Sockets ##
type ZMQSocket
    data::Ptr{Void}

    function ZMQSocket(ctx::ZMQContext, typ::Integer)
        p = ccall((:zmq_socket, :libzmq), Ptr{Void},  (Ptr{Void}, Int32), ctx.data, typ)
        if p == C_NULL
	    throw(ZMQStateError(jl_zmq_error_str()))
        end
        socket = new(p)
        finalizer(socket, close)
        return socket
    end
end

function close(socket::ZMQSocket)
    rc = ccall((:zmq_close, :libzmq), Int32,  (Ptr{Void},), socket.data)
    if rc != 0
	throw(ZMQStateError(jl_zmq_error_str()))
    end
end

# Getting and setting socket options
# Socket options of integer type
let u64p = zeros(Uint64, 1), i64p = zeros(Int64, 1), ip = zeros(Int32, 1), u32p = zeros(Uint32, 1), sz = zeros(Uint, 1)
opslist = {
    (:set_affinity,                :get_affinity,                 4, u64p)
    (nothing,                      :get_fd,                      14,   ip)
    (:set_type,                    :get_type,                    16,   ip)
    (:set_linger,                  :get_linger,                  17,   ip)
    (:set_reconnect_ivl,           :get_reconnect_ivl,           18,   ip)
    (:set_backlog,                 :get_backlog,                 19,   ip)
    (:set_reconnect_ivl_max,       :get_reconnect_ivl_max,       21,   ip)
  }

major, minor, patch = version()
if major == 2
    opslist = vcat(opslist, {
    (:set_hwm,                     :get_hwm,                      1, u64p)
    (:set_swap,                    :get_swap,                     3, i64p)
    (:set_rate,                    :get_rate,                     8, i64p)
    (:set_recovery_ivl,            :get_recovery_ivl,             9, i64p)
    (:_zmq_setsockopt_mcast_loop,  :_zmq_getsockopt_mcast_loop,  10, i64p)
    (:set_sndbuf,                  :get_sndbuf,                  11, u64p)
    (:set_rcvbuf,                  :get_rcvbuf,                  12, u64p)
    (nothing,                      :_zmq_getsockopt_rcvmore,     13, i64p)
    (nothing,                      :get_events,                  15, u32p)
    (:set_recovery_ivl_msec,       :get_recovery_ivl_msec,       20, i64p)
    })
elseif major == 3
    opslist = vcat(opslist, {
    (:set_rate,                    :get_rate,                     8,   ip)
    (:set_recovery_ivl,            :get_recovery_ivl,             9,   ip)
    (:set_sndbuf,                  :get_sndbuf,                  11,   ip)
    (:set_rcvbuf,                  :get_rcvbuf,                  12,   ip)
    (nothing,                      :_zmq_getsockopt_rcvmore,     13,   ip)
    (nothing,                      :get_events,                  15,   ip)
    (:set_maxmsgsize,              :get_maxmsgsize,              22,   ip)
    (:set_sndhwm,                  :get_sndhwm,                  23,   ip)
    (:set_rcvhwm,                  :get_rcvhwm,                  24,   ip)
    (:set_multicast_hops,          :get_multicast_hops,          25,   ip)
    (:set_ipv4only,                :get_ipv4only,                31,   ip)
    (:set_tcp_keepalive,           :get_tcp_keepalive,           34,   ip)
    (:set_tcp_keepalive_idle,      :get_tcp_keepalive_idle,      35,   ip)
    (:set_tcp_keepalive_cnt,       :get_tcp_keepalive_cnt,       36,   ip)
    (:set_tcp_keepalive_intvl,     :get_tcp_keepalive_intvl,     37,   ip)
    })
end
if major > 2 || (major == 2 && minor > 1)
    opslist = vcat(opslist, {
    (:set_rcvtimeo,                :get_rcvtimeo,                27,   ip)
    (:set_sndtimeo,                :get_sndtimeo,                28,   ip)
    })
end

for (fset, fget, k, p) in opslist
    if fset != nothing
        @eval global ($fset)
        @eval function ($fset)(socket::ZMQSocket, option_val::Integer)
            ($p)[1] = option_val
            rc = ccall((:zmq_setsockopt, :libzmq), Int32,
                       (Ptr{Void}, Int32, Ptr{Void}, Uint),
                       socket.data, $k, $p, sizeof(eltype($p)))
            if rc != 0
                throw(ZMQStateError(jl_zmq_error_str()))
            end
        end
    end
    if fget != nothing
        @eval global($fget)
        @eval function ($fget)(socket::ZMQSocket)
            ($sz)[1] = sizeof(eltype($p))
            rc = ccall((:zmq_getsockopt, :libzmq), Int32,
                       (Ptr{Void}, Int32, Ptr{Void}, Ptr{Uint}),
                       socket.data, $k, $p, $sz)
            if rc != 0
                throw(ZMQStateError(jl_zmq_error_str()))
            end
            return int(($p)[1])
        end
    end
end
# For some functions, the publicly-visible versions should require &
# return boolean
if major == 2
    global set_mcast_loop
    set_mcast_loop(socket::ZMQSocket, val::Bool) = _zmq_setsockopt_mcast_loop(socket, val)
    global get_mcast_loop
    get_mcast_loop(socket::ZMQSocket) = bool(_zmq_getsockopt_mcast_loop(socket))
end
end  # let
# More functions with boolean prototypes
get_rcvmore(socket::ZMQSocket) = bool(_zmq_getsockopt_rcvmore(socket))
# And a convenience function
ismore(socket::ZMQSocket) = get_rcvmore(socket)

# Socket options of string type
let u8ap = zeros(Uint8, 255), sz = zeros(Uint, 1)
major, minor, patch = version()
opslist = {
    (:set_identity,                :get_identity,                5)
    (:set_subscribe,               nothing,                      6)
    (:set_unsubscribe,             nothing,                      7)
    }
if major == 3
    opslist = vcat(opslist, {
    (nothing,                      :get_last_endpoint,          32)
    (:set_tcp_accept_filter,       nothing,                     38)
    })
end
for (fset, fget, k) in opslist
    if fset != nothing
        @eval global ($fset)
        @eval function ($fset)(socket::ZMQSocket, option_val::ByteString)
            if length(option_val) > 255
                throw(ZMQStateError("option value too large"))
            end
            rc = ccall((:zmq_setsockopt, :libzmq), Int32,
                       (Ptr{Void}, Int32, Ptr{Uint8}, Uint),
                       socket.data, $k, option_val, length(option_val))
            if rc != 0
                throw(ZMQStateError(jl_zmq_error_str()))
            end
        end
    end
    if fget != nothing
        @eval global ($fget)
        @eval function ($fget)(socket::ZMQSocket)
            ($sz)[1] = length($u8ap)
            rc = ccall((:zmq_getsockopt, :libzmq), Int32,
                       (Ptr{Void}, Int32, Ptr{Uint8}, Ptr{Uint}),
                       socket.data, $k, $u8ap, $sz)
            if rc != 0
                throw(ZMQStateError(jl_zmq_error_str()))
            end
            return bytestring(convert(Ptr{Uint8}, $u8ap), int(($sz)[1]))
        end
    end
end
end  # let

function bind(socket::ZMQSocket, endpoint::String)
    rc = ccall((:zmq_bind, :libzmq), Int32, (Ptr{Void}, Ptr{Uint8}), socket.data, endpoint)
    if rc != 0
        throw(ZMQStateError(jl_zmq_error_str()))
    end
end

function connect(socket::ZMQSocket, endpoint::String)
    rc=ccall((:zmq_connect, :libzmq), Int32, (Ptr{Void}, Ptr{Uint8}), socket.data, endpoint)
    if rc != 0
        throw(ZMQStateError(jl_zmq_error_str()))
    end
end

## Messages ##
typealias ByteArray Union(Array{Uint8,1}, ByteString)
type ZMQMessage
    obj::Vector{Uint8}

    # Create an empty message (for receive)
    function ZMQMessage()
        obj = Array(Uint8, _jl_zmq_msg_t_size)
        rc = ccall((:zmq_msg_init, :libzmq), Int32, (Ptr{Void},), obj)
        if rc != 0
            throw(ZMQStateError(jl_zmq_error_str()))
        end
        zmsg = new(obj)
        finalizer(zmsg, close)
        return zmsg
    end
    # Create a message with a given buffer size (for send)
    function ZMQMessage(len::Integer)
        obj = Array(Uint8, _jl_zmq_msg_t_size)
        rc = ccall((:zmq_msg_init_size, :libzmq), Int32, (Ptr{Void}, Uint), obj, uint(len))
        if rc != 0
            throw(ZMQStateError(jl_zmq_error_str()))
        end
        zmsg = new(obj)
        finalizer(zmsg, close)
        return zmsg
    end
end
# Construct a message from a string (including copying the string)
# In many cases it's more efficient to allocate the zmsg first and
# then build the data in-place, but this is here for convenience
function ZMQMessage(data::ByteArray)
    len = length(data)
    zmsg = ZMQMessage(len)
    ccall(:memcpy, Ptr{Void}, (Ptr{Uint8}, Ptr{Uint8}, Uint),
          msg_data(zmsg), data, len)
    return zmsg
end

# Convert message to array of Uint8 with Uint8[zmsg]
# Copies the data
function ref(::Type{Uint8}, zmsg::ZMQMessage)
    len = msg_size(zmsg)
    data = Array(Uint8, len)
    ccall(:memcpy, Ptr{Void}, (Ptr{Void}, Ptr{Void}, Uint),
          data, msg_data(zmsg), len)
    return data
end
# Convert message to string with ASCIIString[zmsg]
# Copies the data
ref(::Type{ASCIIString}, zmsg::ZMQMessage) = bytestring(msg_data(zmsg), msg_size(zmsg))
bytestring(zmsg::ZMQMessage) = bytestring(msg_data(zmsg), msg_size(zmsg))
# Build an IOStream from a message
# Copies the data
function convert(::Type{IOStream}, zmsg::ZMQMessage)
    len = msg_size(zmsg)
    a = pointer_to_array(msg_data(zmsg), (len,))
    s = memio()
    write(s, a)
    return s
end
# Close a message. You should not need to call this manually (let the
# finalizer do it).
function close(zmsg::ZMQMessage)
    rc = ccall((:zmq_msg_close, :libzmq), Int32, (Ptr{Uint8},), zmsg.obj)
    if rc != 0
        throw(ZMQStateError(jl_zmq_error_str()))
    end
end
# Low-level functions
# Extract a pointer to the ByteArray data in a message
msg_data(zmsg::ZMQMessage) = ccall((:zmq_msg_data, :libzmq), Ptr{Uint8}, (Ptr{Uint8},), zmsg.obj)
# Determine the number of bytes in a message
msg_size(zmsg::ZMQMessage) = ccall((:zmq_msg_size, :libzmq), Int, (Ptr{Uint8},) , zmsg.obj)

if _zmq_major > 2

    global get
    function get(zmsg::ZMQMessage, property::Integer)
        val = ccall((:zmq_msg_get, :libzmq), Int32, (Ptr{Void}, Int32), zmsg.data, property)
        if val < 0
            throw(ZMQStateError(jl_zmq_error_str()))
        end
    end
    global set
    function set(zmsg::ZMQMessage, property::Integer, value::Integer)
        rc = ccall((:zmq_msg_set, :libzmq), Int32, (Ptr{Void}, Int32, Int32), zmsg.data, property, value)
        if rc < 0
            throw(ZMQStateError(jl_zmq_error_str()))
        end
    end
end
## Send/receive messages
#
# Julia defines two types of ZMQ messages: "raw" and "serialized". A "raw"
# message is just a plain ZeroMQ message, used for sending a sequence
# of bytes. You send these with the following:
#   send(socket, zmsg)
#   zmsg = recv(socket)
send(socket::ZMQSocket, zmsg::ZMQMessage) = send(socket, zmsg, int32(0))
function send(socket::ZMQSocket, zmsg::ZMQMessage, noblock::Bool, sndmore::Bool)

    flag::Int32 = 0;
    if (noblock) flag = flag | ZMQ_NOBLOCK ; end
    if (sndmore) flag = flag | ZMQ_SNDMORE ; end
    send(socket, zmsg, flag)
end

if _zmq_major == 2
    global send
    function send(socket::ZMQSocket, zmsg::ZMQMessage, flag::Integer)
        rc = ccall((:zmq_send, :libzmq), Int32, (Ptr{Void}, Ptr{Uint8}, Int32),
                   socket.data, zmsg.obj, flag)
        if rc != 0
            throw(ZMQStateError(jl_zmq_error_str()))
        end
    end
elseif _zmq_major == 3
    global send
    function send(socket::ZMQSocket, zmsg::ZMQMessage, flag::Integer)
        rc = ccall((:zmq_msg_send, :libzmq), Int32, (Ptr{Void}, Ptr{Uint8}, Int32),
                    zmsg.obj, socket.data, flag)
        if rc != 0
            throw(ZMQStateError(jl_zmq_error_str()))
        end
    end
end
recv(socket::ZMQSocket) = recv(socket, int32(0))
function recv(socket::ZMQSocket, noblock::Bool)
    flag::Int32 = 0;
    if (noblock) flag = flag | ZMQ_NOBLOCK ; end
    recv(socket, flag)
end

if _zmq_major == 2
    global recv
    function recv(socket::ZMQSocket, flag::Integer)
        zmsg = ZMQMessage()
        rc = ccall((:zmq_recv, :libzmq), Int32, (Ptr{Void}, Ptr{Void}, Int32),
                   socket.data, zmsg.obj, flag)
        if rc != 0
            throw(ZMQStateError(jl_zmq_error_str()))
        end
        return zmsg
    end
elseif _zmq_major == 3
    global recv
    function recv(socket::ZMQSocket, flag::Integer)
        zmsg = ZMQMessage()
        rc = ccall((:zmq_msg_recv, :libzmq), Int32, (Ptr{Void}, Ptr{Void}, Int32),
                    zmsg.obj, socket.data, flag)
        if rc != 0
            throw(ZMQStateError(jl_zmq_error_str()))
        end
        return zmsg
    end

end

# A "serialized" message includes information needed to interpret the
# data. For example, sending an array requires information about the
# element type and dimensions. See zmq_serialize.jl.

## Constants

# Context options
const ZMQ_IO_THREADS = 1
const ZMQ_MAX_SOCKETS = 2

#Socket Types
const ZMQ_PAIR = 0
const ZMQ_PUB = 1
const ZMQ_SUB = 2
const ZMQ_REQ = 3
const ZMQ_REP = 4
const ZMQ_DEALER = 5
const ZMQ_ROUTER = 6
const ZMQ_PULL = 7
const ZMQ_PUSH = 8
const ZMQ_XPUB = 9
const ZMQ_XSUB = 10
const ZMQ_XREQ = ZMQ_DEALER
const ZMQ_XREP = ZMQ_ROUTER
const ZMQ_UPSTREAM = ZMQ_PULL
const ZMQ_DOWNSTREAM = ZMQ_PUSH

#Message options
const ZMQ_MORE = 1

#Send/Recv Options
const ZMQ_NOBLOCK = 1
const ZMQ_DONTWAIT = 1
const ZMQ_SNDMORE = 2

#IO Multiplexing
const ZMQ_POLLIN = 1
const ZMQ_POLLOUT = 2
const ZMQ_POLLERR = 4

#Built in devices
const ZMQ_STREAMER = 1
const ZMQ_FORWARDER = 2
const ZMQ_QUEUE = 3

end


text:
using Distributions: Distribution
using ..Basis, ..Traits
export ctxapply, Vari

# # Dispatch
# In the contextual execution of a 15able, every intermediate variable application
# of the form `f(ω)` is __intercepted__.  This allows us to do all kinds of things
# such as do causal interventions, track loglikelihood information, etc
# Our implementation models Cassette.jl

# FIXME: Add Conditional to Vari

"Interceptable Variable"
const Vari = Union{Variable, Distribution, Mv, Member, PwVar}

(f::Vari)(ω::Ω) where {Ω <: AbstractΩ} = f(traits(Ω), ω)
(f::Vari)(traits::Trait, ω::Ω) where {Ω <: AbstractΩ} = ctxapply(traits, f, ω)

@inline function ctxapply(traits::Trait, f, ω::AbstractΩ)
  # FIXME: CAUSATION CAN prehook/recurse change traits?
  prehook(traits, f, ω)
  ret = recurse(f, ω)
  posthook(traits, ret, f, ω)
  ret
end

# by default, pre and posthooks do nothing
@inline prehook(traits, f, ω) = nothing
@inline posthook(traits, ret, f, ω) = nothing

# # ## Families

# (f::Vari)(id, ω::Ω) where {Ω <: AbstractΩ} = f(traits(Ω), id, ω)
# (f::Vari)(traits::Trait, id, ω::Ω) where {Ω <: AbstractΩ} = ctxapply(traits, f, id, ω)

# @inline function ctxapply(traits::Trait, f, id, ω::AbstractΩ)
#   # FIXME: CAUSATION CAN prehook/recurse change traits?
#   prehook(traits, f, id, ω)
#   ret = recurse(f, id, ω)
#   posthook(traits, ret, f, id, ω)
#   ret
# end

# # by default, pre and posthooks do nothing
# @inline prehook(traits::Trait, f, id, ω) = nothing
# @inline posthook(traits::Trait, ret, f, id, ω) = nothing


text:
export
    coordarray,
    rmsd,
    displacements,
    sqdistance,
    distance,
    bondangle,
    dihedralangle,
    omegaangle,
    phiangle,
    psiangle,
    ramachandranangles,
    contactmap

"""
Get the atomic coordinates of a `StructuralElementOrList` as a 2D `Array` with
each column corresponding to one atom.
Additional arguments are atom selector functions - only atoms that return
`true` from the functions are retained.
"""
function coordarray(el::StructuralElementOrList, atom_selectors::Function...)
    at_list = collectatoms(el, atom_selectors...)
    coords_out = zeros(3, length(at_list))
    for j in eachindex(at_list)
        coords_out[1,j] = x(at_list[j])
        coords_out[2,j] = y(at_list[j])
        coords_out[3,j] = z(at_list[j])
    end
    return coords_out
end

# Selector functions ignored
coordarray(coords_in::Array{Float64}, atom_selectors::Function...) = coords_in

"""
Get the root-mean-square deviation (RMSD) between two `StructuralElementOrList`s
or two coordinate `Array`s of the same size. Assumes they are already
superimposed.
Additional arguments are atom selector functions - only atoms that return
`true` from the functions are retained.
"""
function rmsd(coords_one::Array{Float64}, coords_two::Array{Float64})
    if size(coords_one) != size(coords_two)
        throw(ArgumentError("Sizes of coordinate arrays are different - cannot calculate RMSD"))
    end
    diff = coords_one - coords_two
    return sqrt.(dot(diff, diff) / size(coords_one, 2))
end

function rmsd(el_one::StructuralElementOrList,
            el_two::StructuralElementOrList,
            atom_selectors::Function...)
    return rmsd(
        coordarray(el_one, atom_selectors...),
        coordarray(el_two, atom_selectors...))
end

"""
Get the displacements between atomic coordinates from two
`StructuralElementOrList`s or two coordinate `Array`s of the same size. Assumes
they are already superimposed.
Additional arguments are atom selector functions - only atoms that return
`true` from the functions are retained.
"""
function displacements(coords_one::Array{Float64}, coords_two::Array{Float64})
    if size(coords_one) != size(coords_two)
        throw(ArgumentError("Sizes of coordinate arrays are different - cannot calculate displacements"))
    end
    diff = coords_one - coords_two
    return sqrt.(sum(diff .* diff, 1))[:]
end

function displacements(el_one::StructuralElementOrList,
                    el_two::StructuralElementOrList,
                    atom_selectors::Function...)
    return displacements(
        coordarray(el_one, atom_selectors...),
        coordarray(el_two, atom_selectors...))
end

"""
Get the minimum square distance between two `StructuralElementOrList`s.
Additional arguments are atom selector functions - only atoms that return
`true` from the functions are retained.
"""
function sqdistance(el_one::StructuralElementOrList,
                    el_two::StructuralElementOrList,
                    atom_selectors::Function...)
    coords_one = coordarray(el_one, atom_selectors...)
    coords_two = coordarray(el_two, atom_selectors...)
    min_sq_dist = Inf
    for i in 1:size(coords_one, 2)
        for j in 1:size(coords_two, 2)
            @inbounds sq_dist = (coords_one[1,i] - coords_two[1,j]) ^ 2 + (coords_one[2,i] - coords_two[2,j]) ^ 2 + (coords_one[3,i] - coords_two[3,j]) ^ 2
            if sq_dist < min_sq_dist
                min_sq_dist = sq_dist
            end
        end
    end
    return min_sq_dist
end

function sqdistance(at_one::AbstractAtom, at_two::AbstractAtom)
    return (x(at_one) - x(at_two)) ^ 2 +
           (y(at_one) - y(at_two)) ^ 2 +
           (z(at_one) - z(at_two)) ^ 2
end

"""
Get the minimum distance between two `StructuralElementOrList`s.
Additional arguments are atom selector functions - only atoms that return
`true` from the functions are retained.
"""
function Bio.distance(el_one::StructuralElementOrList,
                      el_two::StructuralElementOrList,
                      atom_selectors::Function...)
    return sqrt(sqdistance(el_one, el_two, atom_selectors...))
end

function Bio.distance(at_one::AbstractAtom, at_two::AbstractAtom)
    return sqrt(sqdistance(at_one, at_two))
end

"""
Calculate the bond angle in radians between three `AbstractAtom`s, where A-B and
B-C are assumed to be bonded.
"""
function bondangle(at_a::AbstractAtom,
                at_b::AbstractAtom,
                at_c::AbstractAtom)
    return bondangle(
        coords(at_a) - coords(at_b),
        coords(at_c) - coords(at_b)
    )
end

function bondangle(vec_a::Vector{Float64}, vec_b::Vector{Float64})
    return acos(dot(vec_a, vec_b) / (norm(vec_a) * norm(vec_b)))
end

"""
Calculate the dihedral angle in radians defined by four `AbstractAtom`s. This is
the angle between the planes defined by atoms (A,B,C) and (B,C,D).
"""
function dihedralangle(at_a::AbstractAtom,
            at_b::AbstractAtom,
            at_c::AbstractAtom,
            at_d::AbstractAtom)
    return dihedralangle(
        coords(at_b) - coords(at_a),
        coords(at_c) - coords(at_b),
        coords(at_d) - coords(at_c))
end

function dihedralangle(vec_a::Vector{Float64},
                    vec_b::Vector{Float64},
                    vec_c::Vector{Float64})
    return atan2(
        dot(cross(cross(vec_a, vec_b), cross(vec_b, vec_c)), vec_b / norm(vec_b)),
        dot(cross(vec_a, vec_b), cross(vec_b, vec_c)))
end

"""
Calculate the omega angle in radians for an `AbstractResidue`. Arguments are the
residue (atoms "N" and "CA" required) and the previous residue (atoms "CA" and
"C" required).
"""
function omegaangle(res::AbstractResidue, res_prev::AbstractResidue)
    if !("CA" in atomnames(res_prev))
        throw(ArgumentError("Atom with atom name \"CA\" not found in previous residue"))
    elseif !("C"  in atomnames(res_prev))
        throw(ArgumentError("Atom with atom name \"C\" not found in previous residue"))
    elseif !("N"  in atomnames(res))
        throw(ArgumentError("Atom with atom name \"N\" not found in residue"))
    elseif !("CA" in atomnames(res))
        throw(ArgumentError("Atom with atom name \"CA\" not found in residue"))
    end
    return dihedralangle(res_prev["CA"], res_prev["C"], res["N"], res["CA"])
end

"""
Calculate the phi angle in radians for an `AbstractResidue`. Arguments are the
residue (atoms "N", "CA" and "C" required) and the previous residue (atom "C"
required).
"""
function phiangle(res::AbstractResidue, res_prev::AbstractResidue)
    if !("C"  in atomnames(res_prev))
        throw(ArgumentError("Atom with atom name \"C\" not found in previous residue"))
    elseif !("N"  in atomnames(res))
        throw(ArgumentError("Atom with atom name \"N\" not found in residue"))
    elseif !("CA" in atomnames(res))
        throw(ArgumentError("Atom with atom name \"CA\" not found in residue"))
    elseif !("C"  in atomnames(res))
        throw(ArgumentError("Atom with atom name \"C\" not found in residue"))
    end
    return dihedralangle(res_prev["C"], res["N"], res["CA"], res["C"])
end

"""
Calculate the psi angle in radians for an `AbstractResidue`. Arguments are the
residue (atoms "N", "CA" and "C" required) and the next residue (atom "N"
required).
"""
function psiangle(res::AbstractResidue, res_next::AbstractResidue)
    if !("N"  in atomnames(res))
        throw(ArgumentError("Atom with atom name \"N\" not found in residue"))
    elseif !("CA" in atomnames(res))
        throw(ArgumentError("Atom with atom name \"CA\" not found in residue"))
    elseif !("C"  in atomnames(res))
        throw(ArgumentError("Atom with atom name \"C\" not found in residue"))
    elseif !("N"  in atomnames(res_next))
        throw(ArgumentError("Atom with atom name \"N\" not found in next residue"))
    end
    return dihedralangle(res["N"], res["CA"], res["C"], res_next["N"])
end

"""
Calculate the `Vector`s of phi and psi angles of a `StructuralElementOrList`.
The vectors have `NaN` for residues where an angle cannot be calculated,
e.g. due to missing atoms or lack of an adjacent residue.
Additional arguments are residue selector functions - only residues that return
`true` from the functions are retained.
"""
function ramachandranangles(el::StructuralElementOrList,
                    residue_selectors::Function...)
    res_list = collectresidues(el, residue_selectors...)
    if length(res_list) < 2
        throw(ArgumentError("Multiple residues required to calculate Ramachandran angles"))
    end
    phi_angles = Float64[NaN] # First res has no previous res
    psi_angles = Float64[]
    # Phi angles
    for i in 2:length(res_list)
        res = res_list[i]
        res_prev = res_list[i-1]
        if sequentialresidues(res_prev, res)
            try
                phi_angle = phiangle(res, res_prev)
                push!(phi_angles, phi_angle)
            catch ex
                isa(ex, ArgumentError) || rethrow()
                push!(phi_angles, NaN)
            end
        else
            push!(phi_angles, NaN)
        end
    end
    # Psi angles
    for i in 1:length(res_list)-1
        res = res_list[i]
        res_next = res_list[i+1]
        if sequentialresidues(res, res_next)
            try
                psi_angle = psiangle(res, res_next)
                push!(psi_angles, psi_angle)
            catch ex
                isa(ex, ArgumentError) || rethrow()
                push!(psi_angles, NaN)
            end
        else
            push!(psi_angles, NaN)
        end
    end
    push!(psi_angles, NaN) # Last res has no next res
    return phi_angles, psi_angles
end

"""
Calculate the contact map for a `StructuralElementOrList`, or between two
`StructuralElementOrList`s. This is an `BitArray{2}` with `true` where the
sub-elements are no further than the contact distance and `false` otherwise.
"""
function contactmap(el_one::StructuralElementOrList,
                el_two::StructuralElementOrList,
                contact_dist::Real)
    sq_contact_dist = contact_dist ^ 2
    contacts = falses(length(el_one), length(el_two))
    el_one_list = collect(el_one)
    el_two_list = collect(el_two)
    for i in 1:length(el_one)
        for j in 1:length(el_two)
            if sqdistance(el_one_list[i], el_two_list[j]) <= sq_contact_dist
                contacts[i,j] = true
            end
        end
    end
    return contacts
end

function contactmap(el::StructuralElementOrList, contact_dist::Real)
    sq_contact_dist = contact_dist ^ 2
    contacts = falses(length(el), length(el))
    el_list = collect(el)
    for i in 1:length(el)
        contacts[i,i] = true
        for j in 1:i-1
            if sqdistance(el_list[i], el_list[j]) <= sq_contact_dist
                contacts[i,j] = true
                contacts[j,i] = true
            end
        end
    end
    return contacts
end