text:
export GAN,
jsgan,
klgan,
wgan,
rklgan,
lsgan,
sample,
predict,
dcgan_generator,
dcgan_discriminator,
wgan_stable

mutable struct GAN
    latent_dim::Int64
    batch_size::Int64
    dim::Int64
    dat::PyObject
    generator::Union{Missing, Function}
    discriminator::Union{Missing, Function}
    loss::Function
    d_vars::Union{Missing,Array{PyObject}}
    g_vars::Union{Missing,Array{PyObject}}
    d_loss::Union{Missing,PyObject}
    g_loss::Union{Missing,PyObject}
    is_training::Union{Missing,PyObject, Bool}
    update::Union{Missing,Array{PyObject}}
    fake_data::Union{Missing,PyObject}
    true_data::Union{Missing,PyObject}
    ganid::String
    noise::Union{Missing,PyObject}
    ids::Union{Missing,PyObject}
    STORAGE::Dict{String, Any}
end

"""
    build!(gan::GAN)

Builds the GAN instances. This function returns `gan` for convenience.
"""
function build!(gan::GAN)
    gan.noise = placeholder(get_dtype(gan.dat), shape=(gan.batch_size, gan.latent_dim))
    gan.ids = placeholder(Int32, shape=(gan.batch_size,))
    variable_scope("generator_$(gan.ganid)") do
        gan.fake_data = gan.generator(gan.noise, gan)
    end
    gan.true_data = tf.gather(gan.dat,gan.ids-1)
    variable_scope("discriminator_$(gan.ganid)") do
        gan.d_loss, gan.g_loss = gan.loss(gan)
    end
    gan.d_vars = Array{PyObject}(get_collection("discriminator_$(gan.ganid)"))
    gan.d_vars = length(gan.d_vars)>0 ? gan.d_vars : missing
    gan.g_vars = Array{PyObject}(get_collection("generator_$(gan.ganid)"))
    gan.g_vars = length(gan.g_vars)>0 ? gan.g_vars : missing
    gan.update = Array{PyCall.PyObject}(get_collection(UPDATE_OPS))
    gan.update = length(gan.update)>0 ? gan.update : missing
    # gan.STORAGE["d_grad_magnitude"] = gradient_magnitude(gan.d_loss, gan.d_vars)
    # gan.STORAGE["g_grad_magnitude"] = gradient_magnitude(gan.g_loss, gan.g_vars)
    gan
end

@doc raw"""
    GAN(dat::Union{Array,PyObject}, generator::Function, discriminator::Function,
    loss::Union{Missing, Function}=missing; latent_dim::Union{Missing, Int64}=missing,
        batch_size::Int64=32)

Creates a GAN instance.

- `dat` ``\in \mathbb{R}^{n\times d}`` is the training data for the GAN, where ``n`` is the number of training data, and ``d`` is the dimension per training data.
- `generator```:\mathbb{R}^{d'} \rightarrow \mathbb{R}^d`` is the generator function, ``d'`` is the hidden dimension.
- `discriminator```:\mathbb{R}^{d} \rightarrow \mathbb{R}`` is the discriminator function.
- `loss` is the loss function. See [`klgan`](@ref), [`rklgan`](@ref), [`wgan`](@ref), [`lsgan`](@ref) for examples.
- `latent_dim` (default=``d``, the same as output dimension) is the latent dimension.
- `batch_size` (default=32) is the batch size in training.

# Example: Constructing a GAN
```julia
dat = rand(10000,10)
generator = (z, gan)->10*z
discriminator = (x, gan)->sum(x)
gan = GAN(dat, generator, discriminator, "wgan_stable")
```

# Example: Learning a Gaussian random variable
```julia
using ADCME
using PyPlot
using Distributions
dat = randn(10000, 1) * 0.5 .+ 3.0
function gen(z, gan)
    ae(z, [20,20,20,1], "generator_$(gan.ganid)", activation = "relu")
end
function disc(x, gan)
    squeeze(ae(x, [20,20,20,1], "discriminator_$(gan.ganid)", activation = "relu"))
end
gan = GAN(dat, gen, disc, g->wgan_stable(g, 0.001); latent_dim = 10)

dopt = AdamOptimizer(0.0002, beta1=0.5, beta2=0.9).minimize(gan.d_loss, var_list=gan.d_vars)
gopt = AdamOptimizer(0.0002, beta1=0.5, beta2=0.9).minimize(gan.g_loss, var_list=gan.g_vars)
sess = Session(); init(sess)
for i = 1:5000
    batch_x = rand(1:10000, 32)
    batch_z = randn(32, 10)
    for n_critic = 1:1
        global _, dl = run(sess, [dopt, gan.d_loss],
                feed_dict=Dict(gan.ids=>batch_x, gan.noise=>batch_z))
    end
    _, gl, gm, dm, gp = run(sess, [gopt, gan.g_loss,
        gan.STORAGE["g_grad_magnitude"], gan.STORAGE["d_grad_magnitude"],
        gan.STORAGE["gradient_penalty"]],
        feed_dict=Dict(gan.ids=>batch_x, gan.noise=>batch_z))
    mod(i, 100)==0 && (@info i, dl, gl, gm, dm, gp)
end

hist(run(sess, squeeze(rand(gan,10000))), bins=50, density = true)
nm = Normal(3.0,0.5)
x0 = 1.0:0.01:5.0
y0 = pdf.(nm, x0)
plot(x0, y0, "g")
```
"""
function GAN(dat::Union{Array,PyObject}, generator::Function, discriminator::Function,
    loss::Union{Missing, Function, String}=missing; latent_dim::Union{Missing, Int64}=missing,
        batch_size::Int64=32)
    dim = size(dat, 2)
    dat = convert_to_tensor(dat)
    if ismissing(latent_dim)
        latent_dim=dim
    end
    if ismissing(loss)
        loss = "jsgan"
    end

    if isa(loss, String)
        if loss=="jsgan"
            loss = jsgan
        elseif loss=="klgan"
            loss = klgan
        elseif loss=="wgan"
            loss = wgan
        elseif loss=="wgan_stable"
            loss = wgan_stable
        elseif loss=="rklgan"
            loss = rklgan
        elseif loss=="lsgan"
            loss = lsgan
        else
            error("loss function $loss not found!")
        end
    end
    gan = GAN(latent_dim, batch_size, dim, dat, generator, discriminator, loss, missing, missing, missing,
        missing, ADCME.options.training.training, missing, missing, missing, randstring(), missing, missing, Dict())
    build!(gan)
    gan
end
GAN(dat::Array{T}) where T<:Real = GAN(constant(dat))

function Base.:show(io::IO, gan::GAN)
    yes_or_no = x->ismissing(x) ? "✘" : "✔"
    print(
"""
( Generative Adversarial Network -- $(gan.ganid) )
==================================================
    $(gan.latent_dim) D          $(gan.dim) D
(Latent Space)---->(Data Space)
     batch_size = $(gan.batch_size)

loss function: $(gan.loss)
generator:     $(gan.generator)
discriminator: $(gan.discriminator)
d_vars     ... $(yes_or_no(gan.d_vars))
g_vars     ... $(yes_or_no(gan.g_vars))
d_loss     ... $(yes_or_no(gan.d_loss))
g_loss     ... $(yes_or_no(gan.g_loss))
update     ... $(yes_or_no(gan.update))
true_data  ... $(size(gan.true_data))
fake_data  ... $(size(gan.fake_data))
is_training... Placeholder (Bool), $(gan.is_training)
noise      ... Placeholder (Float64) of size $(size(gan.noise))
ids        ... Placeholder (Int32) of size $(size(gan.ids))
"""
    )
end
##################### GAN Library #####################
"""
    klgan(gan::GAN)

Computes the KL-divergence GAN loss function.
"""
function klgan(gan::GAN)
    P, Q = gan.true_data, gan.fake_data
    D_real = gan.discriminator(P, gan)
    D_fake = gan.discriminator(Q, gan)
    D_loss = -mean(log(D_real) + log(1-D_fake))
    G_loss = mean(log((1-D_fake)/D_fake))
    D_loss, G_loss
end

"""
    jsgan(gan::GAN)

Computes the vanilla GAN loss function.
"""
function jsgan(gan::GAN)
    P, Q = gan.true_data, gan.fake_data
    D_real = gan.discriminator(P, gan)
    D_fake = gan.discriminator(Q, gan)
    D_loss = -mean(log(D_real) + log(1-D_fake))
    G_loss = -mean(log(D_fake))
    D_loss, G_loss
end

"""
    wgan(gan::GAN)

Computes the Wasserstein GAN loss function.
"""
function wgan(gan::GAN)
    P, Q = gan.true_data, gan.fake_data
    D_real = gan.discriminator(P, gan)
    D_fake = gan.discriminator(Q, gan)
    D_loss = mean(D_fake)-mean(D_real)
    G_loss = -mean(D_fake)
    D_loss, G_loss
end

@doc raw"""
    wgan_stable(gan::GAN, λ::Float64)

Returns the discriminator and generator loss for the Wasserstein GAN loss with penalty parameter $\lambda$

The objective function is
```math
L = E_{\tilde x\sim P_g} [D(\tilde x)] - E_{x\sim P_r} [D(x)] + \lambda E_{\hat x\sim P_{\hat x}}[(||\nabla_{\hat x}D(\hat x)||^2-1)^2]
```
"""
function wgan_stable(gan::GAN, λ::Float64=1.0)
    real_data, fake_data = gan.true_data, gan.fake_data
    @assert length(size(real_data))==2
    @assert length(size(fake_data))==2
    D_logits = gan.discriminator(real_data, gan)
    D_logits_ = gan.discriminator(fake_data, gan)

    g_loss = -mean(D_logits_)
    d_loss_real = -mean(D_logits)
    d_loss_fake = tf.reduce_mean(D_logits_)
    α = tf.random_uniform(
        shape=(gan.batch_size, 1),
        minval=0.,
        maxval=1.,
        dtype=tf.float64
    )
    differences = fake_data - real_data
    interpolates = real_data + α .* differences # nb x dim
    d_inter = gan.discriminator(interpolates, gan)
    @assert length(size(d_inter))==1
    gradients = tf.gradients(d_inter, interpolates)[1] # ∇D(xt), nb x dim
    slopes = sqrt(sum(gradients^2, dims=2)) # ||∇D(xt)||, (nb,)
    gradient_penalty = mean((slopes-1.)^2)
    d_loss = d_loss_fake + d_loss_real + λ * gradient_penalty
    gan.STORAGE["gradient_penalty"] = mean(slopes)
    return d_loss, g_loss
end

"""
    rklgan(gan::GAN)

Computes the reverse KL-divergence GAN loss function.
"""
function rklgan(gan::GAN)
    P, Q = gan.true_data, gan.fake_data
    D_real = gan.discriminator(P, gan)
    D_fake = gan.discriminator(Q, gan)
    G_loss = mean(log((1-D_fake)/D_fake))
    D_loss = -mean(log(D_fake)+log(1-D_real))
    D_loss, G_loss
end

"""
    lsgan(gan::GAN)

Computes the least square GAN loss function.
"""
function lsgan(gan::GAN)
    P, Q = gan.true_data, gan.fake_data
    D_real = gan.discriminator(P, gan)
    D_fake = gan.discriminator(Q, gan)
    D_loss = mean((D_real-1)^2+D_fake^2)
    G_loss = mean((D_fake-1)^2)
    D_loss, G_loss
end
#######################################################

"""
    sample(gan::GAN, n::Int64)
    rand(gan::GAN, n::Int64)

Samples `n` instances from `gan`.
"""
function sample(gan::GAN, n::Int64)
    local out
    @info "Using a normal latent vector"
    noise = constant(randn(n, gan.latent_dim))
    variable_scope("generator_$(gan.ganid)") do
        out = gan.generator(noise, gan)
    end
    out
end

Base.:rand(gan::GAN) = squeeze(sample(gan, 1), dims=1)
Base.:rand(gan::GAN, n::Int64) = sample(gan, n)

"""
    predict(gan::GAN, input::Union{PyObject, Array})

Predicts the GAN `gan` output given input `input`.
"""
function predict(gan::GAN, input::Union{PyObject, Array})
    local out
    flag = false
    if length(size(input))==1
        flag = true
        input = reshape(input, 1, length(input))
    end
    input = convert_to_tensor(input)
    variable_scope("generator_$(gan.ganid)", initializer=random_uniform_initializer(0.0,1e-3)) do
        out = gan.generator(input, gan)
    end
    if flag; out = squeeze(out); end
    out
end

# adapted from https://github.com/aymericdamien/TensorFlow-Examples/blob/master/examples/3_NeuralNetworks/dcgan.py
function dcgan_generator(x::PyObject, n::Int64=1)
    if length(size(x))!=2
        error("ADCME: input must have rank 2, rank $(length(size(x))) received")
    end
    variable_scope("generator", reuse=AUTO_REUSE) do
        # TensorFlow Layers automatically create variables and calculate their
        # shape, based on the input.
        x = tf.layers.dense(x, units=6n * 6n * 128)
        x = tf.nn.tanh(x)
        # Reshape to a 4-D array of images: (batch, height, width, channels)
        # New shape: (batch, 6, 6, 128)
        x = tf.reshape(x, shape=[-1, 6n, 6n, 128])
        # Deconvolution, image shape: (batch, 14, 14, 64)
        x = tf.layers.conv2d_transpose(x, 64, 4, strides=2)
        # Deconvolution, image shape: (batch, 28, 28, 1)
        x = tf.layers.conv2d_transpose(x, 1, 2, strides=2)
        # Apply sigmoid to clip values between 0 and 1
        x = tf.nn.sigmoid(x)
    end
    return squeeze(x)
end

function dcgan_discriminator(x)
    variable_scope("Discriminator", reuse=AUTO_REUSE) do
        # Typical convolutional neural network to classify images.
        x = tf.layers.conv2d(x, 64, 5)
        x = tf.nn.tanh(x)
        x = tf.layers.average_pooling2d(x, 2, 2)
        x = tf.layers.conv2d(x, 128, 5)
        x = tf.nn.tanh(x)
        x = tf.layers.average_pooling2d(x, 2, 2)
        x = tf.contrib.layers.flatten(x)
        x = tf.layers.dense(x, 1024)
        x = tf.nn.tanh(x)
        # Output 2 classes: Real and Fake images
        x = tf.layers.dense(x, 2)
    end
    return x
end


text:
using BinaryBuilder, Pkg

name = "MSTM"
version = v"4.0.0"

# Collection of sources required to complete build
sources = [
    GitSource("https://github.com/dmckwski/MSTM.git", "b15798c1ef216cd5e6b6022463dd972b058d372d")
]

# Bash recipe for building across all platforms
script = raw"""
if [[ "$target" == *-mingw* ]]; then
    cd ${includedir}
    cp ${prefix}/src/mpi.f90 .
    gfortran -c -DWIN${nbits} -DINT_PTR_KIND=8 -fno-range-check mpi.f90
    cd ${WORKSPACE}/srcdir/MSTM/code
    if [[ ${target} == x86_64-* ]]; then
        cfg_stub="void __guard_check_icall_fptr(unsigned long ptr) { }"
        msmpifec=msmpifec64
        msmpi=msmpi64
    else
        cfg_stub="void __guard_check_icall_fptr(unsigned long ptr) { } void __security_check_cookie(void) { }"
        msmpifec=msmpifec
        msmpi=msmpi
    fi
    echo "${cfg_stub}" | gcc -x c -c -o cfg_stub.o -
    gfortran -O2 -fno-range-check mpidefs-parallel.f90 mstm-intrinsics.f90 mstm-v4.0.f90 cfg_stub.o -L${prefix}/lib -I${includedir} -l${msmpifec} -l${msmpi} -o "${bindir}/mstm${exeext}"
    rm ${includedir}/mpi.f90 ${includedir}/*.mod ${includedir}/*.o
else
    cd ${WORKSPACE}/srcdir/MSTM/code
    mpifort -O2 -fno-range-check mpidefs-parallel.f90 mstm-intrinsics.f90 mstm-v4.0.f90 -o "${bindir}/mstm${exeext}"
fi
install_license ${WORKSPACE}/srcdir/MSTM/LICENSE
"""

# These are the platforms we will build for by default, unless further
# platforms are passed in on the command line
platforms = expand_gfortran_versions(supported_platforms())

# The products that we will ensure are always built
products = [
    ExecutableProduct("mstm", :mstm)
]

# Dependencies that must be installed before this package can be built
dependencies = [
    Dependency(PackageSpec(name = "MPICH_jll", uuid = "7cb0a576-ebde-5e09-9194-50597f1243b4"); platforms = filter(!Sys.iswindows, platforms)),
    Dependency(PackageSpec(name = "MicrosoftMPI_jll", uuid = "9237b28f-5490-5468-be7b-bb81f5f5e6cf"); platforms = filter(Sys.iswindows, platforms)),
    Dependency(PackageSpec(name = "CompilerSupportLibraries_jll", uuid = "e66e0078-7015-5450-92f7-15fbd957f2ae"))
]

# Build the tarballs, and possibly a `build.jl` as well.
build_tarballs(ARGS, name, version, sources, script, platforms, products, dependencies; julia_compat = "1.6")

ext: jl

max_stars_repo_name: jakebolewski/Yggdrasil

max_issues_repo_name: jakebolewski/Yggdrasil

max_forks_repo_name: jakebolewski/Yggdrasil


text:
module Appveyor

if VERSION < v"0.7.0-DEV.914"
    iswindows(k::Symbol=Base.Sys.KERNEL) = k in (:Windows, :NT)
else
    using Base.Sys: iswindows
end

if iswindows()
    # Allows us to verify Codecov upload via Appveyor is working
    const IS_WINDOWS = true
end

end # module


text:
__precompile__()

module TypeStability

    using Compat

    include("StabilityAnalysis.jl")
    include("InlineChecker.jl")
    include("Utils.jl")

end # module


text:
const TupleN{T,N} = NTuple{N,T}

# convenience functions to get generic Inf and NaN
inf(::Union{Type{T}, T}) where {T<:Number} = convert(T, Inf)
two(::Union{Type{T}, T}) where {T<:Number} = convert(T, 2)

@doc raw"""
    wmean(A; wts = nothing, dims = :)

Weighted mean of `A` along dimensions `dims`, weighted by `wts`.

```math
\frac{\sum_i A_i w_i}{\sum_i w_i}
```
"""
function wmean(A::AbstractArray; wts::Union{AbstractArray,Nothing} = nothing, dims = :)
    if isnothing(wts)
        # if no weights are given, fallback to unweighted mean
        return mean(A; dims)
    end

    if dims === (:)
        @assert size(wts) == size(A)
        w = wts
    else
        @assert size(wts) == ntuple(d -> size(A, dims[d]), length(dims))
        # insert singleton dimensions in weights, corresponding to reduced dimensions of `A`
        wsz = ntuple(ndims(A)) do i
            i ∈ dims ? size(A, i) : 1
        end
        w = reshape(wts, wsz)
    end

    return mean(A .* w; dims) ./ mean(wts)
end

@doc raw"""
    wsum(A; wts = nothing, dims = :)

Weighted sum of `A` along dimensions `dims`, weighted by `wts`.

```math
\frac{\sum_i A_i w_i}
```
"""
function wsum(A::AbstractArray; wts::Union{AbstractArray,Nothing} = nothing, dims = :)
    if isnothing(wts)
        return sum(A; dims)
    else
        return wmean(A; wts, dims) * sum(wts)
    end
end

"""
    generate_sequences(n, A = 0:1)

Retruns an iterator over all sequences of length `n` out of the alphabet `A`.
"""
function generate_sequences(n::Int, A = 0:1)
    return (collect(seq) for seq in Iterators.product(ntuple(_ -> A, n)...))
end

"""
    broadlike(A, B...)

Broadcasts `A` into the size of `A .+ B .+ ...` (without actually doing a sum).
"""
broadlike(A, B...) = first_argument.(A, B...)
first_argument(x, y...) = x

# convert to common eltype before matrix multiply, to make sure we hit BLAS
activations_convert_maybe(::AbstractArray{T}, x::AbstractArray{T}) where {T<:AbstractFloat} = x
activations_convert_maybe(::AbstractArray{T}, x::AbstractArray) where {T<:AbstractFloat} = map(T, x)
activations_convert_maybe(::AbstractArray, x::AbstractArray) = x

"""
    reshape_maybe(x, shape)

Like `reshape(x, shape)`, except that zero-dimensional outputs are returned as scalars.
"""
reshape_maybe(x::Number, ::Tuple{}) = x
reshape_maybe(x::AbstractArray, ::Tuple{}) = only(x)
reshape_maybe(x::AbstractArray, sz::TupleN{Int}) = reshape(x, sz)
reshape_maybe(x::Union{Number,AbstractArray}, sz::Int...) = reshape(x, sz)

"""
    repeat_size(sz, r...)

Returns `size(repeat(A, r...))`, provided `size(A) == sz`.
"""
repeat_size(sz::NTuple{N,Int}, r::Int...) where {N} = repeat_size(sz, r)
repeat_size(n::NTuple{N,Int}, r::NTuple{R,Int}) where {N,R} = ntuple(max(N, R)) do d
    if d ≤ N && d ≤ R
        n[d] * r[d]
    elseif R < d ≤ N
        n[d]
    elseif N < d ≤ R
        r[d]
    end
end

sizedims(A::AbstractArray, dims::Int...) = sizedims(A, dims)
sizedims(A::AbstractArray, dims::Tuple{Vararg{Int}}) = map(d -> size(A, d), dims)
sizedims(A::AbstractArray, ::Colon) = size(A)

"""
    moving_average(A, m)

Moving average of `A` with window size `m`.
"""
function moving_average(A::AbstractArray, m::Int)
    out = similar(A)
    R = CartesianIndices(A)
    Ifirst, Ilast = first(R), last(R)
    I1 = m÷2 * oneunit(Ifirst)
    for I in R
        n, s = 0, zero(eltype(out))
        for J in max(Ifirst, I-I1):min(Ilast, I+I1)
            s += A[J]
            n += 1
        end
        out[I] = s/n
    end
    return out
end


text:
@testset "Cholesky" begin

    import Nabla: level2partition, level3partition
    let rng = MersenneTwister(123456), N = 5
        A = randn(rng, N, N)
        r, d, B2, c = level2partition(A, 4, false)
        R, D, B3, C = level3partition(A, 4, 4, false)
        @test all(r .== R')
        @test all(d .== D)
        @test B2[1] == B3[1]
        @test all(c .== C)

        # Check that level2partition with 'U' is consistent with 'L'.
        rᵀ, dᵀ, B2ᵀ, cᵀ = level2partition(transpose(A), 4, true)
        @test r == rᵀ
        @test d == dᵀ
        @test B2' == B2ᵀ
        @test c == cᵀ

        # Check that level3partition with 'U' is consistent with 'L'.
        R, D, B3, C = level3partition(A, 2, 4, false)
        Rᵀ, Dᵀ, B3ᵀ, Cᵀ = level3partition(transpose(A), 2, 4, true)
        @test transpose(R) == Rᵀ
        @test transpose(D) == Dᵀ
        @test transpose(B3) == B3ᵀ
        @test transpose(C) == Cᵀ
    end

    import Nabla: chol_unblocked_rev, chol_blocked_rev
    let rng = MersenneTwister(123456), N = 10
        A, Ā = Matrix.(LowerTriangular.(randn.(Ref(rng), [N, N], [N, N])))
        # NOTE: BLAS gets angry if we don't materialize the Transpose objects first
        B, B̄ = Matrix.(transpose.([A, Ā]))
        @test chol_unblocked_rev(Ā, A, false) ≈ chol_blocked_rev(Ā, A, 1, false)
        @test chol_unblocked_rev(Ā, A, false) ≈ chol_blocked_rev(Ā, A, 3, false)
        @test chol_unblocked_rev(Ā, A, false) ≈ chol_blocked_rev(Ā, A, 5, false)
        @test chol_unblocked_rev(Ā, A, false) ≈ chol_blocked_rev(Ā, A, 10, false)
        @test chol_unblocked_rev(Ā, A, false) ≈ transpose(chol_unblocked_rev(B̄, B, true))

        @test chol_unblocked_rev(B̄, B, true) ≈ chol_blocked_rev(B̄, B, 1, true)
        @test chol_unblocked_rev(B̄, B, true) ≈ chol_blocked_rev(B̄, B, 5, true)
        @test chol_unblocked_rev(B̄, B, true) ≈ chol_blocked_rev(B̄, B, 10, true)
    end

    # Check sensitivities for lower-triangular version.
    let rng = MersenneTwister(123456), N = 10
        for _ in 1:10
            B, VB = randn.(Ref(rng), [N, N], [N, N])
            Ū = UpperTriangular(randn(rng, N, N))
            @test check_errs(B->cholesky(B'B + I).U, Ū, B, 1e-2 .* VB)
        end
    end

    let
        X_ = Matrix{Float64}(I, 5, 5)
        X = Leaf(Tape(), X_)
        C = cholesky(X)
        @test C isa Branch{<:Cholesky}
        @test getfield(C, :f) == LinearAlgebra.cholesky
        U = C.U
        @test U isa Branch{<:UpperTriangular}
        @test getfield(U, :f) == Base.getproperty
        @test unbox(U) ≈ cholesky(X_).U

        @test_throws ArgumentError ∇(X->cholesky(X).info)(X_)
    end

    let
        X_ = Matrix{Float64}(I, 5, 5)
        X = Leaf(Tape(), X_)
        U = cholesky(X).U
        C = Cholesky(U, 'U', 0)
        @test C isa Branch{<:Cholesky}
        @test getfield(C, :f) == LinearAlgebra.Cholesky
        @test unbox(C) == Cholesky(UpperTriangular(X_), 'U', 0)
        # Ensure we can still directly extract the .U field
        UU = C.U
        @test UU isa Branch{<:UpperTriangular}
        # And access .L as well
        LL = C.L
        @test LL isa Branch{<:LowerTriangular}
        # Make sure that computing the Cholesky and already having the Cholesky
        # produce the same results
        expected = Matrix(0.5I, 5, 5)
        @test ∇(X->det(cholesky(X).U))(X_)[1] ≈ expected
        @test ∇(X->det(cholesky(X).L))(X_)[1] ≈ expected
        @test ∇(X->det(Cholesky(cholesky(X).U, :U, 0).U))(X_)[1] ≈ expected
        @test ∇(X->det(Cholesky(cholesky(X).L, 'L', 0).U))(X_)[1] ≈ expected
    end
end


text:
function run_simulation(
    controlODE,
    params;
    noise::Union{Nothing, Real}=nothing,
    vars::Union{Nothing, AbstractArray{<:Integer}}=nothing,
    callback::Union{Nothing, DECallback}=nothing,
)
    if !isnothing(noise) && !isnothing(vars)
        @argcheck noise > zero(noise)
        @argcheck all(var in eachindex(controlODE.u0) for var in vars)

        function noiser(u, t, integrator)
            for var in vars
                u[var] += noise
            end
        end
        callback = FunctionCallingCallback(
            noiser;
            # funcat=tsteps,
            func_everystep=true,
        )
    end

    # integrate with given parameters
    solution = solve(controlODE, params; callback)

    # construct arrays with the same type used by the integrator
    elements_type = eltype(solution.t)
    states = Array(solution)
    total_steps = size(states, 2)
    # state_dimension = size(states, 1)
    control_dimension = length(controlODE.controller(solution.u[1], params))

    # regenerate controls from controlODE.controller
    controls = zeros(elements_type, control_dimension, total_steps)
    for (step, state) in enumerate(solution.u)
        controls[:, step] = controlODE.controller(state, params)
    end
    return solution.t, states, controls
end

# TODO: use DrWatson jl
function store_simulation(
    filename::Union{Nothing,String},
    controlODE::ControlODE,
    params::AbstractVector{<:Real};
    metadata=nothing::Union{Nothing,Dict},
    datadir=nothing::Union{Nothing,String},
    store_policy=true::Bool,
)
    if isnothing(datadir) || isnothing(filename)
        @info "Results not stored due to missing filename/datadir." maxlog = 1
        return nothing
    end

    if store_policy
        policy_path = joinpath(datadir, filename * ".jls")
        open(io -> serialize(io, controlODE.controller), policy_path, "w")

        # weights_path = joinpath(datadir, filename * "_nnweights.csv")
        # CSV.write(weights_path, Tables.table(initial_params(controlODE.controller)), writeheader=false)
    end

    times, states, controls = run_simulation(controlODE, params)

    state_headers = ["x$i" for i in 1:size(states, 1)]
    control_headers = ["c$i" for i in 1:size(controls, 1)]

    full_data = table(
        hcat(times, states', controls'); header=vcat(["t"], state_headers, control_headers)
    )

    CSV.write(joinpath(datadir, filename * ".csv"), full_data)

    if !isnothing(metadata)
        open(joinpath(datadir, filename * "_meta.json"), "w") do f
            JSON3.pretty(f, JSON3.write(metadata))
            println(f)
        end
    end
end


text:
module FromInteractionOp

import ElectronicStructure
using ..Utils: pow_of_minus_one
import ..AbstractFermiOp, ..OpSum, ..OpTerm, ..FermiOps
import ..FermiOps: NumberOp, Raise, Lower
import ..add!

"""
    FermiSumA(iop::ElectronicStructure.InteractionOperator)

Convert `iop` to a `FermiSum`.
"""
function OpSum{FermiOpT}(iop::ElectronicStructure.InteractionOperator) where {FermiOpT<:AbstractFermiOp}
    fsum = OpSum{FermiOpT}(iop.one_body_tensor, iop.two_body_tensor)
    add!(fsum, iop.nuclear_repulsion * one(fsum[1]))
    return fsum
end

function OpSum{FermiOpT}(tensor::AbstractArray{<:Number}, tensors::AbstractArray{<:Number}...) where {FermiOpT<:AbstractFermiOp}
    fsum = OpSum(Vector{FermiOpT}[], eltype(tensor)[])
    for tensor in (tensor, tensors...)
        tensor_to_fermi_sum!(fsum, tensor)
    end
    return fsum
end

function index_to_ops(i, j, k, l)
    ## We filter this elsewhere, but it's dangerous to remove this
    # if i == j || k == l
    #     return ()  # unstable, use a barrier or move this out
    # end
    if i == k || i == l
        op1 = (op=NumberOp, ind=i)
    else
        op1 = (op=Raise, ind=i)
    end
    if j == k || j == l
        op2 = (op=NumberOp, ind=j)
    else
        op2 = (op=Raise, ind=j)
    end
    if k != i && k != j
        op3 = (op=Lower, ind=k)
    else
        op3 = (op=FermiOps.NoOp, ind=k)
    end
    if l != i && l != j
        op4 = (op=Lower, ind=l)
    else
        op4 = (op=FermiOps.NoOp, ind=k)
    end
    return (op1, op2, op3, op4)
end

function index_to_ops(i, j)
    if i == j
        return ((op=NumberOp, ind=i), )
    end
    if i < j
        return ((op=Raise, ind=i), (op=Lower, ind=j))
    else
        return ((op=Lower, ind=i), (op=Raise, ind=j))
    end
end

index_to_ops(inds) = index_to_ops(inds...)

function count_permutations(inds::NTuple{2})
    if inds[2] < inds[1]
        return 1
    else
        return 0
    end
end

"""
    count_permutations(inds::NTuple{4})

Return the permuation signature of `inds`, that is the number of
pair permutations required to order `inds`.
"""
function count_permutations(inds::NTuple{4})
    count = 0
    @inbounds i = inds[1]
    @inbounds for j in (2, 3, 4)
        if i > inds[j]
            count += 1
        end
    end
    @inbounds i = inds[2]
    @inbounds for j in (3, 4)
        if i > inds[j]
            count += 1
        end
    end
    @inbounds i = inds[3]
    @inbounds if i > inds[4]
        count += 1
    end
    return count
end

_index_phase(inds::NTuple) = pow_of_minus_one(count_permutations(inds))

index_to_ops_phase(inds) = (ops=index_to_ops(inds), phase=_index_phase(inds))

_skip_inds(i, j, k, l) = (i == j || k == l) # skip if two raising or two lowering ops on one index
_skip_inds(i, j) = false # always one raising and one lowering, so never skip

## Embed `ops` in a string of width `n_modes`

"""
    _fermi_term(ops::NTuple, phase::Integer, coeff, n_modes::Integer)

Convert the sparse representation of a string of ops, `ops` to a dense representation.
Each element of `ops` is a two-element `Tuple` giving an operator and index. The returned
`Vector` contains each operator at the corresponding index, with unspecified indices
set to the identity.
"""
function _fermi_term(ops::NTuple, phase::Integer, coeff, n_modes::Integer)
    factors = fill(FermiOps.I, n_modes)
    for (op, ind) in ops
        if op === FermiOps.NoOp
            continue
        end
        factors[ind] = op
    end
    return (factors, phase * coeff)
end

function tensor_term_to_fermi_term(::Type{T}, inds::NTuple{N, Int}, coeff, n_modes::Integer) where {T<:AbstractFermiOp, N}
    (ops, phase) = index_to_ops_phase(inds)
    (factors, new_coeff) = _fermi_term(ops, phase, coeff, n_modes)
    return OpTerm{T}(factors, new_coeff)
end

function tensor_to_fermi_sum!(fsum::OpSum{T}, tensor) where T
    n_modes = first(size(tensor))
    @inbounds for ind in CartesianIndices(tensor)
        if _skip_inds((ind.I)...) || iszero(tensor[ind])
            continue
        end
        add!(fsum, tensor_term_to_fermi_term(T, ind.I, tensor[ind], n_modes))
    end
    return fsum
end

end # module FromInteractionOp


text:
using SparseArrays

import MathOptInterface
const MOI = MathOptInterface

# TODO: use a trait instead
const NONLINEAR_CONE = Union{
    MOI.SecondOrderCone, MOI.RotatedSecondOrderCone
}
const POLYHEDRAL_CONE = Union{
    MOI.Zeros, MOI.Nonnegatives, MOI.Nonpositives, MOI.Reals
}
const MOI_CONES = Union{POLYHEDRAL_CONE, NONLINEAR_CONE}
const SCALAR_SETS = Union{
    MOI.EqualTo{Float64}, MOI.GreaterThan{Float64}, MOI.LessThan{Float64}
}
const SUPPORTED_SETS = Union{
    SCALAR_SETS, MOI_CONES, MOI.Integer
}

mutable struct StandardProblem
    model::MOI.ModelLike

    A::SparseMatrixCSC{Float64, Int}  # constraint matrix
    b::Vector{Float64}  # Right-hand sides
    c::Vector{Float64}  # Objective
    c0::Float64  # Objective constant

    var_indices::Vector{MOI.VariableIndex}  # List of MOI variable index

    cones::Vector{Tuple{Vector{Int}, MOI.AbstractSet}}  # cones
    vartypes::Vector{Bool}  # Integer flags
    vartypes_implied::Vector{Bool}
end

"""
    convert_to_standard_form(m::MOI.ModelLike)

Convert instance to standard form.
"""
function build_standard_form(
    optimizer::MOI.OptimizerWithAttributes,
    src::MOI.ModelLike;
    bridge_type::Union{Nothing, Type}=nothing
)
    # Some solvers do not support Rotated Second Order cones directly,
    # so we enable bridges as a (hopefully short-term) work-around
    dst = MOI.instantiate(optimizer, with_bridge_type=bridge_type)

    b = Float64[]
    arows = Int[]
    acols = Int[]
    avals = Float64[]

    cones = Tuple{Vector{Int}, MOI.AbstractSet}[]

    # First, copy all variables
    x_old = MOI.get(src, MOI.ListOfVariableIndices())
    nvar = length(x_old)  # Original number of variables
    x = MOI.add_variables(dst, nvar)  # variables in new model
    vartypes = zeros(Bool, nvar)  # integer flags
    @info "Original problem has $nvar variables"

    var_indices = copy(x)

    old_variable_idx_map = Dict(x_old .=> 1:nvar)  # Map from old indices to new ones

    # Objective coefficients
    c = zeros(Float64, nvar)
    fobj = MOI.get(src, MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}())
    c0 = fobj.constant
    for t in fobj.terms
        j = old_variable_idx_map[t.variable_index]
        c[j] = t.coefficient
    end

    MOI.set(dst,
        MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}(),
        MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.(c, x), c0)
    )

    # Objective sense
    MOI.set(dst, MOI.ObjectiveSense(), MOI.get(src, MOI.ObjectiveSense()))

    # Now, constraints
    ncon = 0
    contypes = MOI.get(src, MOI.ListOfConstraints())
    b = Float64[]  # This will be the right-hand side

    for (F, S) in contypes
        S <: SUPPORTED_SETS || error("Set $S is not supported.")

        # Get list of corresponding constraint indices
        conindices = MOI.get(src, MOI.ListOfConstraintIndices{F, S}())

        for cidx in conindices
            fun = MOI.get(src, MOI.ConstraintFunction(), cidx)
            set = MOI.get(src, MOI.ConstraintSet(), cidx)

            # Update standard form
            add_to_standard_form!(dst,
                fun, set,
                old_variable_idx_map, var_indices,
                c, b, arows, acols, avals, cones, vartypes
            )
        end
    end

    nvar = length(c)
    ncon = length(b)

    # Add Free cones for free variables
    cflags = ones(Bool, nvar)
    for (kidx, k) in cones
        all(cflags[kidx]) || error("Some variables are in two cones")
        cflags[kidx] .= false
    end
    for (j, flag) in enumerate(cflags)
        flag && push!(cones, ([j], MOI.Reals(1)))
    end

    @info "Problem stats:\n\tVariables  : $nvar\n\tConstraints: $ncon\n\tCones      : $(length(cones))"

    return StandardProblem(dst,
        sparse(arows, acols, avals, ncon, nvar), b,
        c, c0,
        var_indices,
        cones,
        vartypes, copy(vartypes)
    )
end

"""

Convert integer restriction
"""
function add_to_standard_form!(
    dst::MOI.ModelLike,
    fun::MOI.SingleVariable,
    ::MOI.Integer,
    old_variable_idx_map, var_indices,
    c, b, arows, acols, avals, cones, vartypes
)

    j = old_variable_idx_map[fun.variable]
    vartypes[j] = true

    MOI.add_constraint(dst, MOI.SingleVariable(MOI.VariableIndex(j)), MOI.Integer())

    return nothing
end

"""

Convert vector polyhedral cone into series of smaller cones
"""
function add_to_standard_form!(
    dst::MOI.ModelLike,
    fun::MOI.VectorOfVariables,
    set::POLYHEDRAL_CONE,
    old_variable_idx_map, var_indices,
    c, b, arows, acols, avals, cones, vartypes
)
    if isa(set, MOI.Zeros)
        set_ = MOI.EqualTo(0.0)
        k = MOI.Zeros(1)
    elseif isa(set, MOI.Nonnegatives)
        set_ = MOI.GreaterThan(0.0)
        k = MOI.Nonnegatives(1)
    elseif isa(set, MOI.Nonpositives)
        set_ = MOI.LessThan(0.0)
        k = MOI.Nonpositives(1)
    else
        error("")
    end

    for j_old in fun.variables
        j = old_variable_idx_map[j_old]

        MOI.add_constraint(dst,
            MOI.SingleVariable(MOI.VariableIndex(j)),
            set_
        )

        push!(cones, ([j], k))
    end

    return nothing
end

"""

Convert vector polyhedral constraint into series of smaller cones.

This adds `m` variables, where `m` is the dimension of the conic constraint
"""
function add_to_standard_form!(
    dst::MOI.ModelLike,
    fun::MOI.VectorAffineFunction,
    set::POLYHEDRAL_CONE,
    old_variable_idx_map, var_indices,
    c, b, arows, acols, avals, cones, vartypes
)
    nvar = length(c)
    ncon = length(b)

    m = MOI.dimension(set)

    if isa(set, MOI.Nonnegatives)
        set_ = MOI.GreaterThan(0.0)
        k = MOI.Nonnegatives(1)
    elseif isa(set, MOI.Nonpositives)
        set_ = MOI.LessThan(0.0)
        k = MOI.Nonpositives(1)
    else
        error("")
    end

    # Add slack variables
    s = MOI.add_variables(dst, m)
    append!(var_indices, s)

    # Set bounds on slack variables
    for i in 1:m
        MOI.add_constraint(dst,
            MOI.SingleVariable(s[i]),
            set_
        )
        push!(cones, ([nvar + i], k))
    end
    append!(c, zeros(m))
    append!(vartypes, zeros(Bool, m))

    # Update matrix coefficients
    append!(arows, ncon .+ collect(1:m))
    append!(acols, nvar .+ collect(1:m))
    append!(avals, -ones(m))

    # Extract terms of VectorAffineFunction
    funs_ = [
        MOI.ScalarAffineFunction([MOI.ScalarAffineTerm(-1.0, s[i])], 0.0)
        for i in 1:m
    ]
    for t in fun.terms
        i = t.output_index
        j = old_variable_idx_map[t.scalar_term.variable_index]
        v = t.scalar_term.coefficient
        push!(funs_[i].terms, MOI.ScalarAffineTerm(v, MOI.VariableIndex(j)))

        # Update coefficients of A
        push!(arows, ncon + i)
        push!(acols, j)
        push!(avals, v)
    end

    # Add each constraint separately
    for i in 1:m
        # We add each equality constraint separately
        rhs = -fun.constants[i]
        push!(b, rhs)

        # Add equality constraint
        MOI.add_constraint(dst, funs_[i], MOI.EqualTo(rhs))
    end

    return nothing
end

function add_to_standard_form!(
    dst::MOI.ModelLike,
    fun::MOI.VectorAffineFunction,
    set::MOI.Zeros,
    old_variable_idx_map, var_indices,
    c, b, arows, acols, avals, cones, vartypes
)
    nvar = length(c)
    ncon = length(b)

    m = MOI.dimension(set)

    set_ = MOI.EqualTo(0.0)
    k = MOI.Zeros(1)

    # No need to add slack variables

    # Extract terms of VectorAffineFunction
    funs_ = [
        MOI.ScalarAffineFunction(MOI.ScalarAffineTerm{Float64}[], 0.0)
        for i in 1:m
    ]
    for t in fun.terms
        i = t.output_index
        j = old_variable_idx_map[t.scalar_term.variable_index]
        v = t.scalar_term.coefficient
        push!(funs_[i].terms, MOI.ScalarAffineTerm(v, MOI.VariableIndex(j)))

        # Update coefficients of A
        push!(arows, ncon + i)
        push!(acols, j)
        push!(avals, v)
    end

    # Add each constraint separately
    for i in 1:m
        # We add each equality constraint separately
        rhs = -fun.constants[i]
        push!(b, rhs)

        # Add equality constraint
        MOI.add_constraint(dst, funs_[i], MOI.EqualTo(rhs))
    end

    return nothing
end

# TODO: Add option to use extended formulation for second order cone
"""

Convert vector non-linear conic constraint into standard form

This adds `m` variables, where `m` is the dimension of the conic constraint.
"""
function add_to_standard_form!(
    dst::MOI.ModelLike,
    fun::MOI.VectorAffineFunction,
    set::NONLINEAR_CONE,
    old_variable_idx_map, var_indices,
    c, b, arows, acols, avals, cones, vartypes
)
    nvar = length(c)
    ncon = length(b)

    m = MOI.dimension(set)

    # Add slack variables
    s = MOI.add_variables(dst, m)
    append!(var_indices, s)
    append!(c, zeros(m))
    append!(vartypes, zeros(Bool, m))

    # Update matrix coefficients
    append!(arows, ncon .+ collect(1:m))
    append!(acols, nvar .+ collect(1:m))
    append!(avals, -ones(m))

    # Set conic domain on slacks
    MOI.add_constraint(dst, MOI.VectorOfVariables(s), set)
    push!(cones, (nvar .+ collect(1:m), set))

    # Extract terms of VectorAffineFunction
    funs_ = [
        MOI.ScalarAffineFunction([MOI.ScalarAffineTerm(-1.0, s[i])], 0.0)
        for i in 1:m
    ]
    for t in fun.terms
        i = t.output_index
        j = old_variable_idx_map[t.scalar_term.variable_index]
        v = t.scalar_term.coefficient
        push!(funs_[i].terms, MOI.ScalarAffineTerm(v, MOI.VariableIndex(j)))

        # Update coefficients of A
        push!(arows, ncon + i)
        push!(acols, j)
        push!(avals, v)
    end

    # Add each constraint separately
    for i in 1:m
        # We add each equality constraint separately
        rhs = -fun.constants[i]
        push!(b, rhs)

        # Add equality constraint
        MOI.add_constraint(dst, funs_[i], MOI.EqualTo(rhs))
    end

    return nothing
end

function add_to_standard_form!(
    dst::MOI.ModelLike,
    fun::MOI.VectorOfVariables,
    set::NONLINEAR_CONE,
    old_variable_idx_map, var_indices,
    c, b, arows, acols, avals, cones, vartypes
)
    m = MOI.dimension(set)

    x_ = [
        MOI.VariableIndex(old_variable_idx_map[v])
        for v in fun.variables
    ]

    # Add conic constraint
    MOI.add_constraint(dst, MOI.VectorOfVariables(x_), set)

    # Record cone
    push!(cones, ([old_variable_idx_map[v] for v in fun.variables], set))

    return nothing
end

max_stars_repo_name: mtanneau/CLaP

max_issues_repo_name: mtanneau/CLaP

max_forks_repo_name: mtanneau/CLaP


text:
using Vlasiator, SHA, LazyArtifacts
using Test

group = get(ENV, "TEST_GROUP", :all) |> Symbol

function nanmaximum(x::AbstractArray{T}) where T<:AbstractFloat
   result = convert(eltype(x), NaN)
   for i in x
      if !isnan(i)
         if isnan(result) || i > result
            result = i
         end
      end
   end
   result
end

@testset "Vlasiator.jl" begin
   rootpath = artifact"testdata"

   files = joinpath.(rootpath, ("bulk.1d.vlsv", "bulk.2d.vlsv", "bulk.amr.vlsv"))
   meta1 = load(files[1])
   meta2 = load(files[2])
   meta3 = load(files[3])

   if group in (:read, :all)
      @testset "Reading files" begin
         @test_throws ArgumentError load("data")
         meta = meta1
         @test ndims(meta) == 1
         @test startswith(repr(meta), "File: bulk.1d.vlsv")
         @test size(meta) == 5168730
         # Variable strings reading
         @test meta.variable[end] == "vg_rhom"
         # Variable info reading
         varinfo = readvariablemeta(meta, "proton/vg_rho")
         @test startswith(repr(varinfo), "Variable in LaTeX")
         @test varinfo.unit == "1/m^3"
         # Parameter checking
         @test hasparameter(meta, "dt") == true
         # Parameter reading
         t = readparameter(meta, "time")
         @test t == 8.0
         @test_throws ArgumentError meta["nonsense"]
         # unsorted ID
         @test readvariable(meta, "CellID", false) == UInt64[10, 9, 8, 7, 2, 1, 3, 4, 5, 6]
         indexRef = [6, 5, 7, 8, 9, 10, 4, 3, 2, 1]
         @test meta.cellindex == indexRef
         # sorted var by default
         @test meta["vg_boundarytype"] == [4, 4, 1, 1, 1, 1, 1, 1, 3, 3]
         # ID finding
         loc = [2.0, 0.0, 0.0]
         id = getcell(meta, loc)
         coords = getcellcoordinates(meta, id)
         @test coords == [2.5, 0.0, 0.0]
         @test readvariable(meta, "proton/vg_rho", id)[1,1] ≈ 1.77599 atol=1e-5
         # ID in a line
         point1 = [-2.0, 0.0, 0.0]
         point2 = [2.0, 0.0, 0.0]
         cellids, _, _ = getcellinline(meta, point1, point2)
         @test cellids == collect(4:7)
         point1 = [-5.1, 0.0, 0.0]
         @test_throws DomainError getcellinline(meta, point1, point2)
         point2 = [5.1, 0.0, 0.0]
         @test_throws DomainError getcellinline(meta, point1, point2)

         # Nearest ID with VDF stored
         @test getnearestcellwithvdf(meta, id) == 8

         # velocity space reading
         vcellids, vcellf = readvcells(meta, 2; species="proton")
         V = getvcellcoordinates(meta, vcellids; species="proton")
         @test V[end] == Float32[2.45, 1.95, 1.95]
         @test_throws ArgumentError readvcells(meta, 20)
         f = Vlasiator.flatten(meta.meshes["proton"], vcellf)
         @test f[CartesianIndex(24, 18, 7)] == 1.0f-45

         # AMR data reading, DCCRG grid
         metaAMR = meta3
         sliceoffset = abs(metaAMR.coordmin[2])
         idlist, indexlist = getslicecell(metaAMR, sliceoffset, 2,
            metaAMR.coordmin[2], metaAMR.coordmax[2])

         data = readvariable(metaAMR, "proton/vg_rho")
         dataslice = refineslice(metaAMR, idlist, data[indexlist], :y)
         @test sum(dataslice) ≈ 7.690352275026747e8
         let err = nothing
            try
               getslicecell(metaAMR, sliceoffset, 1, -2., -1.)
            catch err
            end
            @test err isa Exception
         end
         # AMR ID finding
         loc = [0.0, 0.0, 0.0]
         @test getcell(metaAMR, loc) == 0x0000000000001591

         # AMR level
         @test metaAMR.maxamr == 2
         @test getlevel(metaAMR, idlist[1]) == 1

         # DCCRG utilities
         @test isparent(metaAMR, 1)
         @test !isparent(metaAMR, 1080)
         @test getchildren(metaAMR, 1) == getsiblings(metaAMR, 129)
         @test getparent(metaAMR, 129) == 1
         @test_throws ArgumentError getparent(metaAMR, 5)
         @test_throws ArgumentError getsiblings(metaAMR, 1)

         # FS grid
         data = readvariable(metaAMR, "fg_e")
         ncells, namr = metaAMR.ncells, metaAMR.maxamr
         @test size(data) == (3, ncells[1]*namr^2, ncells[2]*namr^2, ncells[3]*namr^2)
         @test data[:,16,8,8] == [-1.3758785f-7, 3.2213068f-4, -3.1518404f-4]

         # Compare two VLSV files
         @test issame(files[1], files[1])
      end
   end

   if group in (:derive, :all)
      @testset "Derived variables" begin
         meta = meta1
         @test meta["Vmag"] |> sortperm == [7, 6, 5, 4, 3, 1, 2, 8, 9, 10]

         meta = meta2
         @test meta["Bmag"][4] == 3.0052159f-9

         @test meta["Emag"][1,10,99] == 2.6120072f-6

         @test meta["VS"] |> nanmaximum == 1.3726345956957596e6

         @test meta["VA"] |> nanmaximum == 2.3202628822256166e8

         @test meta["MA"][end] == 10.700530839822328

         @test meta["MS"][end] == 16.9375888861409

         @test meta["Vpar"][1] == 698735.3f0

         @test meta["Vperp"][1] == 40982.48f0

         @test meta["T"][1] == 347619.9817319378

         @test meta["Pdynamic"][1] == 8.204415428337215e-10

         @test meta["Pb"][1] == 3.59502523919876e-12

         @test meta["Beta"][1] == 1.3359065984817116

         @test meta["Poynting"][:,10,10] == [-3.677613f-11, 8.859047f-9, 2.4681486f-9]

         @test meta["IonInertial"][1] == 8.578087716535188e7

         @test readvariable(meta, "Larmor", UInt64[1])[1] == 322324.70603759587

         @test meta["Gyroperiod"][1] == 21.834297799454554

         @test meta["Gyrofrequency"][1] == 0.04579950356933307

         @test meta["J"][1,1000] == 2.314360722590665e-14

         @test meta["Omegap"][1] == 209.5467447842415

         @test meta["Plasmaperiod"][1] == 0.0047722048893178645

         #@test meta["Panisotropy"]

         #@test meta["Tanisotropy"]

         #@test meta["Agyrotropy"]

      end
   end

   if group in (:utility, :all)
      @testset "Rotation" begin
         using LinearAlgebra
         T = Diagonal([1.0, 2.0, 3.0])
         B = [0.0, 1.0, 0.0]
         Vlasiator.rotateWithB!(T, B)
         @test T == Diagonal([1.0, 3.0, 2.0])
         @test Vlasiator.rotateWithB(T, B) == Diagonal([1.0, 2.0, 3.0])

         v = fill(1/√3, 3)
         θ = π / 4
         R = Vlasiator.getRotationMatrix(v, θ)
         @test inv(R) ≈ R'

         Rᵀ = Vlasiator.rotateTensorToVectorZ(R, v)
         @test Rᵀ ≈ [1/√2 -1/√2 0.0; 1/√2 1/√2 0.0; 0.0 0.0 1.0]
      end
      @testset "Curvature" begin
         let dx = ones(Float32, 3)
            A = ones(Float32, 3,3,3,3)
            @test sum(Vlasiator.curl(dx, A)) == 0.0
            A = ones(Float32, 3,3,1,3)
            @test sum(Vlasiator.curl(dx, A)) == 0.0
         end
      end
   end

   if group in (:vtk, :all)
      @testset "VTK" begin
         meta = meta2 # no amr
         data, ghostType = Vlasiator.fillmesh(meta, "proton/vg_rho")
         @test size(data[1][1]) == (1, 63, 100, 1)

         meta = meta3 # amr
         write_vtk(meta)
         sha_str = bytes2hex(open(sha1, "bulk.amr_1.vti"))
         @test sha_str == "bafb52747908cd05c78218078054aa4268f453d9"

         filesaved = ["bulk.amr.vthb", "bulk.amr_1.vti", "bulk.amr_2.vti", "bulk.amr_3.vti"]
         rm.(filesaved, force=true)
      end
   end

   if group in (:log, :all)
      @testset "Log" begin
         file = joinpath(rootpath, "logfile.txt")
         timestamps, speed = readlog(file)
         @test length(speed) == 50 && speed[end] == 631.2511f0
      end
   end

   if group in (:plot, :all)
      @testset "PyPlot" begin
         using PyPlot
         ENV["MPLBACKEND"]="agg" # no GUI
         # 1D
         meta = meta1
         plot(meta, "proton/vg_rho")
         line = gca().lines[1]
         @test line.get_ydata() == meta["proton/vg_rho"]
         centers = plotmesh(meta, projection="y")
         points = centers.get_offsets()
         @test size(points) == (10, 2)
         fig = plt.figure()
         ax = fig.add_subplot(projection="3d")
         centers = plotmesh(meta, projection="3d")
         points = centers.get_offsets() # only 2D from 3D coords, might be improved
         @test size(points) == (10, 2)
         close(fig)

         @test_throws ArgumentError pcolormesh(meta, "proton/vg_rho")

         loc = [2.0, 0.0, 0.0]
         var = vdfslice(meta, loc).get_array()
         @test var[786] == 229.89486651619364
         @test_throws ArgumentError vdfslice(meta, loc, species="helium")

         # 2D
         meta = meta2
         var = pcolormesh(meta, "proton/vg_rho").get_array()
         @test var[end-2] == 999535.8f0 && length(var) == 6300
         var = pcolormesh(meta, "fg_b").get_array()
         @test var[1] == 3.0058909f-9
         var = pcolormesh(meta, "proton/vg_v", op=:x, colorscale=SymLog).get_array()
         @test var[2] == -699935.2f0
         p = streamplot(meta, "proton/vg_v", comp="xy")
         @test typeof(p) == PyPlot.PyObject
         var = quiver(meta, "proton/vg_v", axisunit=SI, stride=1).get_offsets()
         @test size(var) == (6300, 2)

         # 3D AMR
         meta = meta3
         var = pcolormesh(meta, "proton/vg_rho").get_array()
         @test var[255] == 1.0483886f6 && length(var) == 512
         var = pcolormesh(meta, "proton/vg_v").get_array()
         @test var[255] == 99992.586f0
         @test_throws ArgumentError pcolormesh(meta, "fg_b")
      end

      @testset "Plots" begin
         isdefined(Main, :PyPlot) && import Vlasiator.vdfslice
         include("../src/plot/plots.jl")
         RecipesBase.is_key_supported(k::Symbol) = true
         # 1D
         meta = meta1
         rec = RecipesBase.apply_recipe(Dict{Symbol, Any}(), meta, "proton/vg_rho")
         @test getfield(rec[1], 1)[:seriestype] == :line &&
            rec[1].args[1] isa LinRange

         rec = RecipesBase.apply_recipe(Dict{Symbol, Any}(), VDFSlice((meta,[0.0,0.0,0.0])))
         @test getfield(rec[1], 1)[:seriestype] == :histogram2d

         # 2D
         meta = meta2
         rec = RecipesBase.apply_recipe(Dict{Symbol, Any}(), meta, "proton/vg_rho")
         @test getfield(rec[1], 1)[:seriestype] == :heatmap &&
            rec[1].args[1] isa LinRange
      end
   end
   for meta in (meta1, meta2, meta3)
      close(meta.fid) # required for Windows?
   end
end


text:
using Documenter, RangeEnclosures

DocMeta.setdocmeta!(RangeEnclosures, :DocTestSetup, :(using RangeEnclosures); recursive=true)

makedocs(
    sitename = "RangeEnclosures.jl",
    modules = [RangeEnclosures],
    format = Documenter.HTML(
        prettyurls = get(ENV, "CI", nothing) == "true",
        assets = ["assets/aligned.css"]),
    doctest = true,
    strict = true,
    pages = [
        "Home" => "index.md",
        "Library" => Any[
            "Types" => "lib/types.md",
            "Methods" => "lib/methods.md"],
        "About" => "about.md"
    ]
)

deploydocs(
    repo = "github.com/JuliaReach/RangeEnclosures.jl.git",
    push_preview=true
)


text:
# ==========
# Arithmetic
# ==========

SUITE["Arithmetic-Vector"] = BenchmarkGroup()

SUITE["Arithmetic-Scalar"] = BenchmarkGroup()

# The examples are taken from Althoff, M., Grebenyuk, D., & Kochdumper, N. (2018).
# Implementation of Taylor models in CORA 2018. In Proc. of the 5th International Workshop on
# Applied Verification for Continuous and Hybrid Systems.

t = Taylor1(6)

# taylor models of order 6
b1 = TaylorModel1(0.5 + 1.5t, 0..0, 0.5..0.5, -1..1)
b2 = TaylorModel1(2.5 + 0.5t, 0..0, 0.5..0.5, -1..1)
b3 = TaylorModel1(-5.0 + 1.0t, 0..0, 0.5..0.5, -1..1)
b4 = TaylorModel1(5.0 + 1.0t, 0..0, 0.5..0.5, -1..1)

B1 = [b1, b2]
B2 = [b3, b4]

SUITE["Arithmetic-Vector"]["addition"] = @benchmarkable $B1 + $B2
SUITE["Arithmetic-Vector"]["matrix multiplication"] = @benchmarkable transpose($B1) * $B2
SUITE["Arithmetic-Vector"]["pointwise multiplication"] = @benchmarkable B1 .* B2
SUITE["Arithmetic-Vector"]["division by scalar"] = @benchmarkable $B1 / 2
SUITE["Arithmetic-Vector"]["pointwise division"] = @benchmarkable $B1 ./ $B2
SUITE["Arithmetic-Vector"]["power function"] = @benchmarkable $B1.^3
SUITE["Arithmetic-Vector"]["sine function"] = @benchmarkable sin.($B1)
SUITE["Arithmetic-Vector"]["combination of functions"] = @benchmarkable sin($B1[1, 1]) + $B1[2, 1].^2  - transpose($B1) * $B2

SUITE["Arithmetic-Scalar"]["addition"] = @benchmarkable $b1 + $b2
SUITE["Arithmetic-Scalar"]["multiplication"] = @benchmarkable $b1 * $b2
SUITE["Arithmetic-Scalar"]["division"] = @benchmarkable $b1 / 2
SUITE["Arithmetic-Scalar"]["pointwise division"] = @benchmarkable $b1 / $b2
SUITE["Arithmetic-Scalar"]["power function"] = @benchmarkable $b1^3
SUITE["Arithmetic-Scalar"]["sine function"] = @benchmarkable sin($b1)
SUITE["Arithmetic-Scalar"]["combination of functions"] = @benchmarkable sin($b1[1]) + $b1[2]^2  - $b1 * $b2


text:
__precompile__()

module GenomeGraphs

export
    ###
    ### Re-exports of GenomeGraphs framework sub-components.
    ###

    # Re-exports of ReadDatastores
    ReadDatastore,
    PairedReads,
    LongReads,
    LinkedReads,
    buffer,
    FwRv,

    # Re-exports of FASTX
    FASTA,
    FASTQ,

    ###
    ### Re-exports of BioSequences.jl
    ###
    DNAAlphabet,
    BioSequence,
    LongSequence,
    AbstractMer,
    Mer,
    BigMer,
    DNAMer,
    DNAKmer,
    BigDNAMer,
    BigDNAKmer,

    ###
    ### Re-exports of KmerAnalysis.jl
    ###
    mer,
    freq,
    Canonical,
    NonCanonical,
    CANONICAL,
    NONCANONICAL,
    # Kmer counters
    serial_mem,
    dist_mem,
    spectra,

    empty_graph,
    #GRAPH_TYPE,

    ### WorkSpace
    #WorkSpace,
    #add_paired_reads!,
    #paired_reads,
    #add_mer_counts!,
    #mer_counts,
    read_datastore,
    ###
    ### Processes
    ###
    dbg,
    dbg!,
    remove_tips!

using BioSequences, FASTX, ReadDatastores, KmerAnalysis
import BioSequences.EveryMerIterator

include("Graphs.jl")       # Submodule defining the key Graph type and basic methods.
include("GraphIndexes.jl") # Submodule defining types that allow indexing of a graph.

#include("workspace/WorkSpace.jl")
#include("views/NodeView.jl")

# Utility function for more quickly making a read datastore.
function read_datastore(R1file::String, R2file::String, name::String, minlen::Int, maxlen::Int, insertlen::Int, mode::PairedReadOrientation)
    fwq = open(FASTQ.Reader, R1file)
    rvq = open(FASTQ.Reader, R2file)
    return PairedReads{DNAAlphabet{4}}(fwq, rvq, name, name, minlen, maxlen, insertlen, mode)
end

empty_graph(::Type{T}) where {T<:LongSequence} = Graphs.SequenceDistanceGraph{T}()

include("dbg.jl")
include("remove_tips.jl")

end # module GenomeGraphs

max_issues_repo_name: BioJulia/BioSequenceGraphs

max_forks_repo_name: BioJulia/BioSequenceGraphs


text:
@userplot PhasePlot
# args:  z_upperleft::Complex, z_lowerright::Complex, f::Function
@recipe function plot_portrait(P::PhasePlot;
                               no_pixels=(600,600), point_color=cs_j(),
                               no_ticks=7, ticks_sigdigits=2)
  length(no_ticks) == 1 && (no_ticks = (no_ticks, no_ticks))
  z_upperleft, z_lowerright, f = P.args[1:3]
  img = portrait(z_upperleft, z_lowerright, f;
                 no_pixels=no_pixels, point_color=point_color);

  seriestype := :heatmap
  xticks := (LinRange(0, no_pixels[1], no_ticks[1]),
             round.(LinRange(imag(z_lowerright), imag(z_upperleft), no_ticks[1]), sigdigits=ticks_sigdigits))
  yticks := (LinRange(0, no_pixels[2], no_ticks[2]),
             reverse!(round.(LinRange(real(z_upperleft), real(z_lowerright), no_ticks[2]), sigdigits=ticks_sigdigits)))

  img
end

"""
```
phaseplot(z_upperleft, z_lowerright, f)
```
Convinience function, automatically adds ticks and creates a `Plot` object.

Can be used in layouts. Examples:

```
using Plots, ComplexPortraits
@ComplexPortraits.import_huge

phaseplot(-2.0 + 2.0im, 2.0 - 2.0im, z -> z^2)

phaseplot(-2.0 + 2.0im, 2.0 - 2.0im, z -> z^2;
          no_pixels=(600, 600),
          point_color=cs_j(),
          no_ticks=(7, 7),
          ticks_sigdigits=2)
```

Todo: hack colorbar plotattribute to show colorwheel.
"""
phaseplot;

"""
tests:
function crop_to_circle!(A:Array; radius=0, background_color=RGB{Float64}(1.0,1.0,1.0))
    m, n = size(A)
    max_radius = min(m, n) / 2
    center = [m/2, n/2]
    (radius < 0 || radius > max_radius) && cv_error("plot image size too small. ",
                                                    "Maximum radius = ", max_radius)
    radius == 0 && (radius = max_radius)

    for i in Base.OneTo(m), j in Base.OneTo(n)
        sum(([i, j] .- center).^2) * π > radius && (A[i, j] = background_color)
    end

    return A
end
"""


text:
export ∞, asℝ, asℝ₊, asℝ₋, as𝕀, as_real, as_positive_real, as_negative_real,
    as_unit_interval

"""
$(TYPEDEF)

Transform a scalar (real number) to another scalar.

Subtypes mustdefine `transform`, `transform_and_logjac`, and `inverse`; other
methods of of the interface should have the right defaults.
"""
abstract type ScalarTransform <: AbstractTransform end

dimension(::ScalarTransform) = 1

function transform_with(flag::NoLogJac, t::ScalarTransform, x::AbstractVector, index)
    transform(t, @inbounds x[index]), flag, index + 1
end

function transform_with(::LogJac, t::ScalarTransform, x::AbstractVector, index)
    transform_and_logjac(t, @inbounds x[index])..., index + 1
end

function inverse_at!(x::AbstractVector, index, t::ScalarTransform, y::Real)
    x[index] = inverse(t, y)
    index + 1
end

inverse_eltype(t::ScalarTransform, y::T) where {T <: Real} = float(T)

random_arg(t::ScalarTransform; kwargs...) = random_real(; kwargs...)

####
#### identity
####

"""
$(TYPEDEF)

Identity ``x ↦ x``.
"""
struct Identity <: ScalarTransform end

transform(::Identity, x::Real) = x

transform_and_logjac(::Identity, x::Real) = x, zero(x)

inverse(::Identity, x::Real) = x

####
#### shifted exponential
####

"""
$(TYPEDEF)

Shifted exponential. When `D::Bool == true`, maps to `(shift, ∞)` using `x ↦
shift + eˣ`, otherwise to `(-∞, shift)` using `x ↦ shift - eˣ`.
"""
struct ShiftedExp{D, T <: Real} <: ScalarTransform
    shift::T
    function ShiftedExp{D,T}(shift::T) where {D, T <: Real}
        @argcheck D isa Bool
        new(shift)
    end
end

ShiftedExp(D::Bool, shift::T) where {T <: Real} = ShiftedExp{D,T}(shift)

transform(t::ShiftedExp{D}, x::Real) where D =
    D ? t.shift + exp(x) : t.shift - exp(x)

transform_and_logjac(t::ShiftedExp, x::Real) = transform(t, x), x

function inverse(t::ShiftedExp{D}, x::Real) where D
    @unpack shift = t
    if D
        @argcheck x > shift DomainError
        log(x - shift)
    else
        @argcheck x < shift DomainError
        log(shift - x)
    end
end

####
#### scaled and shifted logistic
####

"""
$(TYPEDEF)

Maps to `(scale, shift + scale)` using `x ↦ logistic(x)*scale + shift`.
"""
struct ScaledShiftedLogistic{T <: Real} <: ScalarTransform
    scale::T
    shift::T
    function ScaledShiftedLogistic{T}(scale::T, shift::T) where {T <: Real}
        @argcheck scale > 0
        new(scale, shift)
    end
end

ScaledShiftedLogistic(scale::T, shift::T) where {T <: Real} =
    ScaledShiftedLogistic{T}(scale, shift)

ScaledShiftedLogistic(scale::Real, shift::Real) =
    ScaledShiftedLogistic(promote(scale, shift)...)

transform(t::ScaledShiftedLogistic, x::Real) = logistic(x) * t.scale + t.shift

# NOTE: would prefer fma(logistic(x), t.scale, t.shift) for all types, but cf
# https://github.com/JuliaDiff/DiffRules.jl/issues/28
transform(t::ScaledShiftedLogistic, x::AbstractFloat) = fma(logistic(x), t.scale, t.shift)

transform_and_logjac(t::ScaledShiftedLogistic, x) =
    transform(t, x), log(t.scale) + logistic_logjac(x)

function inverse(t::ScaledShiftedLogistic, y)
    @argcheck y > t.shift           DomainError
    @argcheck y < t.scale + t.shift DomainError
    logit((y - t.shift)/t.scale)
end

# NOTE: inverse_and_logjac interface experimental and sporadically implemented for now
function inverse_and_logjac(t::ScaledShiftedLogistic, y)
    @argcheck y > t.shift           DomainError
    @argcheck y < t.scale + t.shift DomainError
    z = (y - t.shift) / t.scale
    logit(z), logit_logjac(z) - log(t.scale)
end

####
#### to_interval interface
####

struct Infinity{ispositive}
    Infinity{T}() where T = (@argcheck T isa Bool; new{T}())
end

"""
Placeholder representing of infinity for specifing interval boundaries. Supports
the `-` operator, ie `-∞`.
"""
const ∞ = Infinity{true}()

Base.show(::Infinity{T}) where T = print(io, T ? "∞" : "-∞")

Base.:(-)(::Infinity{T}) where T = Infinity{!T}()

"""
    as(Real, left, right)

Return a transformation that transforms a single real number to the given (open)
interval.

`left < right` is required, but may be `-∞` or `∞`, respectively, in which case
the appropriate transformation is selected. See [`∞`](@ref).

Some common transformations are predefined as constants, see [`ℝ`](@ref),
[`ℝ₋`](@ref), [`ℝ₊`](@ref), [`𝕀`](@ref).
"""
as(::Type{Real}, left, right) =
    throw(ArgumentError("($(left), $(right)) must be an interval"))

as(::Type{Real}, ::Infinity{false}, ::Infinity{true}) = Identity()

as(::Type{Real}, left::Real, ::Infinity{true}) = ShiftedExp(true, left)

as(::Type{Real}, ::Infinity{false}, right::Real) = ShiftedExp(false, right)

function as(::Type{Real}, left::Real, right::Real)
    @argcheck left < right "the interval ($(left), $(right)) is empty"
    ScaledShiftedLogistic(right - left, left)
end

"""
Transform to a positive real number.

`asℝ₊` and `as_positive_real` are equivalent alternatives.
"""
const asℝ₊ = as(Real, 0.0, ∞)

const as_positive_real = asℝ₊

"""
Transform to a negative real number.

`asℝ₋` and `as_negative_real` are equivalent alternatives.
"""
const asℝ₋ = as(Real, -∞, 0.0)

const as_negative_real = asℝ₋

"""
Transform to the unit interval `(0, 1)`.

`as𝕀` and `as_unit_interval` are equivalent alternatives.
"""
const as𝕀 = as(Real, 0.0, 1.0)

const as_unit_interval = as𝕀

"""
Transform to the real line (identity).

`asℝ` and `as_real` are equivalent alternatives.
"""
const asℝ = as(Real, -∞, ∞)

const as_real = asℝ

Base.show(io::IO, t::ShiftedExp) =
    if t === asℝ₊
        print(io, "asℝ₊")
    elseif t === asℝ₋
        print(io, "asℝ₋")
    elseif t isa ShiftedExp{true}
        print(io, "as(Real, ", t.shift, ", ∞)")
    else
        print(io, "as(Real, -∞, ", t.shift, ")")
    end

Base.show(io::IO, t::ScaledShiftedLogistic) =
    if t === as𝕀
        print(io, "as𝕀")
    else
        print(io, "as(Real, ", t.shift, ", ", t.shift + t.scale, ")")
    end

Base.show(io::IO, t::Identity) = print(io, "asℝ")


text:
"""
    tern2cart(a, b, c)

Takes a point in the ternary coordinate system and returns the cartesian co-ordinate of this point.
The ternary axes are all between 0 and 1 and the cartesian co-ordinates form an equilateral triangle at (0, 0), (1, 0), (1/2, √3/2)
"""
tern2cart(a, b, c) = (1 / 2 * (2b + c) / (a + b + c), √3 / 2 * (c / (a + b + c)))
tern2cart(vec) = tern2cart(vec[1], vec[2], vec[3])

"""
    cart2tern(x, y)

Takes a point in the cartesian coordinate system and returns the ternary co-ordinate of this point.
"""
function cart2tern(x, y)
    c = (2 * y) / √3
    b = x - c / 2
    a = 1 - b - c
    return (a, b, c)
end
cart2tern(array) = cart2tern(array[1], array[2])


text:
@testset "mcpc3ds" begin
# Data
fn_phase = "data/small/Phase.nii"
fn_mag = "data/small/Mag.nii"
phase_nii = readphase(fn_phase)
mag_nii = readmag(fn_mag)
complex = mag_nii .* exp.(1im .* phase_nii)
TEs = 4:4:12

mcpc3ds(complex; TEs=TEs)
mcpc3ds(phase_nii, mag_nii; TEs=TEs)
mcpc3ds(phase_nii; TEs=TEs)

end


text:
#!/usr/bin/env julia

include("MPC.jl")

# Black magic to make Julia and Python happy
using Distributed

using Pkg
Distributed.@everywhere using Pkg

Distributed.@everywhere begin
  ENV["PYTHON"] = "/usr/bin/python3"
  Pkg.build("PyCall")
end

using PyCall
Distributed.@everywhere using PyCall

#Import Statements for RobotOS
using RobotOS

@rosimport trajectory_msgs.msg: JointTrajectory, JointTrajectoryPoint
@rosimport arthur_planning.msg: arthur_traj
@rosimport sensor_msgs.msg: JointState
rostypegen()
using .trajectory_msgs.msg
using .arthur_planning.msg
using .sensor_msgs.msg

function callback(msg::JointTrajectory, X, U, initialized, N, idx, lastTime)
    if msg.header.stamp != lastTime[1]
        lastTime[1] = msg.header.stamp
        N .= length(msg.points)
        for k = 1:N[1]
            X[k][1:7] .= msg.points[k].positions
            X[k][8:14] .= msg.points[k].velocities
            U[k][1:end] .= msg.points[k].effort
        end
        if !initialized[1] && N[1] > 0
            initialized .= true
        end
        idx .= 1
    end
end

function traj_callback(msg::arthur_traj, Xref)
    Xref[1] = typeof(zeros(14))[]
    for k=1:length(msg.traj.points)
        q = msg.traj.points[k].positions
        q̇ = msg.traj.points[k].velocities
        push!(Xref[1], [q; q̇])
    end
    # RobotOS.loginfo("In Callback")
    # println(Xref[1][1])
end

function callback_x0(msg::JointState, x0)
    x0_new = zeros(14)
    x0_new[1:7] .= msg.position
    x0_new[8:14] .= msg.velocity
    copy!(x0, x0_new)
end

function loop(sub_obj1, sub_obj2, sub_obj3, pub_obj, x0, X, U, params, initialized, N, Xref, mpc_controller, idx, starting)
    loop_rate = Rate(20)
    while ! is_shutdown()
        RobotOS._run_callbacks(sub_obj1)
        RobotOS._run_callbacks(sub_obj2)
        RobotOS._run_callbacks(sub_obj3)
        if initialized[1] && mpc_controller[1]
            i = argmin(norm.([(X[k] - x0) for k=1:N[1]]))
            if norm(x0[1:7] - Xref[1][end][1:7]) < 0.05
                mpc_controller .= false
            end
            # if !isnan(x0[1])
            #     RobotOS.loginfo("MPC Control")
            #     println(U[i])
            #     println(x0)
            #     println(norm(x0 - Xref[1][end]))
            # end
            JointTrajectoryOutput = JointTrajectoryPoint()
            JointTrajectoryOutput.effort = U[i]
            RobotOS.loginfo("MPC Control")
            println(i)
            println(U[i])
            println(x0)
            # JointTrajectoryOutput.positions = x0[1:7]
            # JointTrajectoryOutput.velocities = x0[8:14]
            publish(pub_obj, JointTrajectoryOutput)
            # idx[1] = min(N[1], idx[1] + 1)
        # elseif norm(x0[8:14]) < 1e-6
        #     set_configuration!(params.state, x0[1:7])
        #     set_velocity!(params.state, zeros(7))
        #     params.v̇ .= zeros(7)
        #     U_grav = inverse_dynamics(params.state, params.v̇)
        #     RobotOS.loginfo("Gravity Compensation Control")
        #     println(U_grav)
        #     println(x0)
        #     # println(RobotDynamics.dynamics(params.model, x0, U_grav))
        #     JointTrajectoryOutput = JointTrajectoryPoint()
        #     JointTrajectoryOutput.effort = U_grav
        #     publish(pub_obj, JointTrajectoryOutput)
        elseif starting[1]
            if norm(x0 - Xref[1][1]) < 1e-2
                starting .= false
            end
            set_configuration!(params.state, x0[1:7])
            set_velocity!(params.state, x0[8:14])
            i = 1
            ∆x = Xref[1][i][1:7] - x0[1:7]
            # println(∆x)
            Kp = [30, 20, 20, 20, 20, 50, 20]
            Kd = [20, 20, 20, 20, 20, 20, 20]
            params.v̇ .= (Kp .* ∆x) .- (Kd .* x0[8:14])
            U_PD = inverse_dynamics(params.state, params.v̇)
            RobotOS.loginfo("PD Control")
            println(U_PD)
            println(x0)
            println(norm(x0 - Xref[1][1]))
            JointTrajectoryOutput = JointTrajectoryPoint()
            JointTrajectoryOutput.effort = U_PD
            publish(pub_obj, JointTrajectoryOutput)
        else
            set_configuration!(params.state, x0[1:7])
            set_velocity!(params.state, x0[8:14])

            # if initialized[1]
            #     i = argmin(norm.([(Xref[1][k][1:7] - x0[1:7]) for k=1:length(Xref[1])]))
            # else
            #     i = 1
            #     # println(Xref[1][i][1:7])
            #     # println(x0[1:7])
            # end

            i = min(argmin(norm.([(Xref[1][k][1:7] - x0[1:7]) for k=1:length(Xref[1])])) + 1, length(Xref[1]))

            ∆x = Xref[1][i][1:7] - x0[1:7]
            # println(∆x)
            Kp = [100, 100, 100, 100, 100, 100, 100]
            Kd = [20, 20, 20, 20, 20, 20, 20]
            params.v̇ .= (Kp .* ∆x) .- (Kd .* x0[8:14])
            U_PD = inverse_dynamics(params.state, params.v̇)
            # x0 .= rk4(params.model, x0, U_PD, params.dt)
            # if !isnan(x0[1])
            #     RobotOS.loginfo("PD Control")
            #     println(params.v̇[1:7])
            #     println(U_PD[1:7])
            #     println(x0[1:7])
            #     println(norm(x0 - Xref[end]))
            # end
            RobotOS.loginfo("PD Tracking Control")
            println(U_PD)
            println(x0)
            println(i)
            JointTrajectoryOutput = JointTrajectoryPoint()
            JointTrajectoryOutput.effort = U_PD
            publish(pub_obj, JointTrajectoryOutput)
        end
        rossleep(loop_rate)
    end
end

function main()
    init_node("Controller_Node")
    params = MPC_Params()

    # Xref = get_trajectory(params)
    # x0 = copy(Xref[1])
    Xref = [typeof(zeros(14))[]]
    x0 = zeros(14)
    # N = [params.H-1]
    N = [params.LQRH]
    X = [zeros(14) for k = 1:params.H-1]
    U = [zeros(7) for k = 1:params.H-1]
    initialized = [false]
    mpc_controller = [false]
    starting = [true]
    idx = [1]
    lastTime = [RobotOS.now()]

    sub_x0 = Subscriber{JointState}("/my_gen3/joint_states", callback_x0, (x0,), queue_size=1)
    sub = Subscriber{JointTrajectory}("joint_torques",callback,(X, U, initialized, N, idx, lastTime),queue_size=1)
    sub_traj = Subscriber{arthur_traj}("/my_gen3/arthur_traj",traj_callback,(Xref,),queue_size=1)

    pub = Publisher{JointTrajectoryPoint}("mpc_torques",queue_size=1)

    while length(Xref[1]) == 0
        RobotOS._run_callbacks(sub_traj)
    end
    while norm(x0 - zeros(14)) == 0
        RobotOS._run_callbacks(sub_x0)
    end

    loop(sub, sub_traj, sub_x0, pub, x0, X, U, params, initialized, N, Xref, mpc_controller, idx, starting)
end

if !isinteractive()
    main()
end


text:
function maskedSsdDistance(referenceImage::ImageMeta,templateImage::ImageMeta,
                           transformedGrid::Array{Float64,1},
                           mask::Array{Float64,1},
                           o::regOptions)

    # interpolation of the template image at transformed grid points
    transformedImage, dY_transformedImage, dX_transformedImage =
      linearImageInterpolationAtGridWithDerivative(templateImage,transformedGrid)

    dY_transformedImage = mask[:] .* dY_transformedImage[:]
    dX_transformedImage = mask[:] .* dX_transformedImage[:]

    # measure the ssd distance
    N = prod(size(referenceImage))
    h = prod(pixelspacing(referenceImage))
    residual = Array{Float64,N}
    residual = mask[:] .* (transformedImage[:] .- referenceImage.data[:])
	  functionValue = 0.5 .* h .* residual' * residual

    # check if only the distance is wanted
    if(~o.doDerivative)
	      return functionValue
    end

    # rearrange the image derivative
    dTransformedImage = spdiagm((dX_transformedImage[:],dY_transformedImage[:]),[0,N])

    # first derivative
    dFunctionValue    = h .* dTransformedImage' * residual

    # approximate second derivative
    d2FunctionValue   = h .* dTransformedImage' * dTransformedImage

    return functionValue,dFunctionValue,d2FunctionValue
end


text:
using Gridap

using GridapGmsh

model = GmshDiscreteModel("light-mesh/light-mesh.msh")

model.grids[2]


text:
using IntervalArithmetic, IntervalRootFinding
using Base.Test

@testset "`bisect` function" begin
    X = 0..1
    @test bisect(X) == (0..0.5, 0.5..1)
    @test bisect(X, 0.25) == (0..0.25, 0.25..1)

    X = -∞..∞
    @test bisect(X) == (-∞..0, 0..∞)
    @test bisect(X, 0.75) == (-∞..0, 0..∞)

    X = 1..∞
    @test bisect(X) == (Interval(1, prevfloat(∞)), Interval(prevfloat(∞), ∞))

    X = (0..1) × (0..2)
    @test bisect(X) == ( (0..1) × (0..1), (0..1) × (1..2) )
    @test bisect(X, 0.25) == ( (0..1) × (0..0.5), (0..1) × (0.5..2) )
    @test bisect(X, 1, 0.5) == ( (0..0.5) × (0..2), (0.5..1) × (0..2) )
    @test bisect(X, 1, 0.25) == ( (0..0.25) × (0..2), (0.25..1) × (0..2) )

    X = (-∞..∞) × (-∞..∞)
    @test bisect(X) == ( (-∞..0) × (-∞..∞), (0..∞) × (-∞..∞))
end


text:
#=
# Cachesim example
=#

### If the Finch package has already been added, use this line #########
using Finch # Note: to add the package, first do: ]add "https://github.com/paralab/Finch.git"

### If not, use these four lines (working from the examples directory) ###
# if !@isdefined(Finch)
#     include("../Finch.jl");
#     using .Finch
# end
##########################################################################

init_finch("cachesim");
useLog("cachesimlog")

#################################################
# For cache sim
cachesim(true); # Select cachesim. Normal solve won't happen
l1 = build_cache_level(1, 32, 8, 64, "LRU");    #
l2 = build_cache_level(2, 256, 4, 64, "LRU");   # Set up the cache specs (this is my laptop)
l3 = build_cache_level(3, 12288, 16, 64, "LRU");# (level, sets, ways, line_size, policy)
build_cache([l1, l2, l3]);
#################################################

n = 10;
ord = 2;

# Set up the configuration
domain(3)
functionSpace(order=ord)

# Specify the problem
mesh(HEXMESH, elsperdim=n)

u = variable("u")
testSymbol("v")

boundary(u, 1, DIRICHLET, 0)

# Write the weak form
coefficient("f", "-3*pi*pi*sin(pi*x)*sin(pi*y)*sin(pi*z)")
weakForm(u, "-dot(grad(u), grad(v)) - f*v")

cachesimSolve(u);

finalize_finch()

max_stars_repo_name: paralab/Finch

max_issues_repo_name: paralab/Finch

max_forks_repo_name: paralab/Finch


text:
# Tests for Profit Additive DEA Models
@testset "ProfitAdditiveDEAModel" begin

    # Test using Book data
    X = [2; 4; 8; 12; 6; 14; 14; 9.412];
    Y = [1; 5; 8; 9; 3; 7; 9; 2.353];
    W = [1; 1; 1; 1; 1; 1; 1; 1];
    P = [2; 2; 2; 2; 2; 2; 2; 2];

    # With :Ones
    deaprofitaddones = deaprofitadd(X, Y, W, P, :Ones)

    @test typeof(deaprofitaddones) == ProfitAdditiveDEAModel

    @test nobs(deaprofitaddones) == 8
    @test ninputs(deaprofitaddones) == 1
    @test noutputs(deaprofitaddones) == 1
    @test ismonetary(deaprofitaddones) == false

    @test efficiency(deaprofitaddones) == efficiency(deaprofitaddones, :Economic)
    @test efficiency(deaprofitaddones, :Economic)   ≈ [8.0; 2.0; 0.0; 2.0; 8.0; 8.0; 4.0; 12.706]
    @test efficiency(deaprofitaddones, :Technical)  ≈ [0.0; 0.0; 0.0; 0.0; 4.0; 22/3; 2.0; 8059/1000]
    @test efficiency(deaprofitaddones, :Allocative) ≈ [8.0; 2.0; 0.0; 2.0; 4.0; 2/3; 2.0; 4647/1000]
    @test normfactor(deaprofitaddones) ≈ [1.0; 1.0; 1.0; 1.0; 1.0; 1.0; 1.0; 1.0]

    # Default is :Ones
    @test efficiency(deaprofitadd(X, Y, W, P)) == efficiency(deaprofitaddones)

    # With :MIP
    deaprofitaddmip = deaprofitadd(X, Y, W, P, :MIP)

    @test efficiency(deaprofitaddmip, :Economic)   ≈ [4.0; 0.5; 0.0; 1/6; 4/3; 4/7; 2/7; 2.699958] atol = 1e-5
    @test efficiency(deaprofitaddmip, :Technical)  ≈ [0.0; 0.0; 0.0; 0.0; 7/6; 4/7; 1/7; 2.549936] atol = 1e-5
    @test efficiency(deaprofitaddmip, :Allocative) ≈ [4.0; 0.5; 0.0; 1/6; 1/6; 0.0; 1/7; 0.150021] atol = 1e-5
    @test normfactor(deaprofitaddmip) ≈ [2.0; 4.0; 8.0; 12.0; 6.0; 14.0; 14.0; 4.706]

    # With :MIP and monetary = true
    deaprofitaddmipmon = deaprofitadd(X, Y, W, P, :MIP, monetary = true)

    @test ismonetary(deaprofitaddmipmon) == true

    @test efficiency(deaprofitaddmipmon) ≈ efficiency(deaprofitaddmip) .* normfactor(deaprofitaddmip)
    @test efficiency(deaprofitaddmipmon, :Technical) ≈ efficiency(deaprofitaddmip, :Technical) .* normfactor(deaprofitaddmip)
    @test efficiency(deaprofitaddmipmon, :Allocative) ≈ efficiency(deaprofitaddmip, :Allocative) .* normfactor(deaprofitaddmip)

    # With custom weights
    deaprofitaddcustom = deaprofitadd(X, Y, W, P, rhoX = 1 ./ X, rhoY = 1 ./ Y)

    @test efficiency(deaprofitaddcustom) ≈ efficiency(deaprofitaddmip)

    # Print
    show(IOBuffer(), deaprofitaddones)

    # Test errors
    @test_throws DimensionMismatch deaprofitadd([1; 2 ; 3], [4; 5], [1; 1; 1], [1; 1], :Ones) #  Different number of observations
    @test_throws DimensionMismatch deaprofitadd([1; 2; 3], [4; 5; 6], [1; 2; 3; 4], [1; 1; 1], :Ones) # Different number of observation in input prices
    @test_throws DimensionMismatch deaprofitadd([1; 2; 3], [4; 5; 6], [1; 1; 1], [1; 2; 3; 4], :Ones) # Different number of observation in output prices
    @test_throws DimensionMismatch deaprofitadd([1 1; 2 2; 3 3 ], [4; 5; 6], [1 1 1; 2 2 2; 3 3 3], [1; 1; 1], :Ones) # Different number of input prices and inputs
    @test_throws DimensionMismatch deaprofitadd([1; 2; 3 ], [4; 5; 6], [1; 1; 1], [1 1 1; 2 2 2; 3 3 3], :Ones) # Different number of output prices and outputs
    @test_throws DimensionMismatch deaprofitadd([1; 2; 3] , [4; 5; 6], [1; 1; 1], [1; 1; 1], rhoX = [1 1; 2 2; 3 3], rhoY = [1; 1; 1]) # Different size of inputs weights
    @test_throws DimensionMismatch deaprofitadd([1; 2; 3] , [4; 5; 6], [1; 1; 1], [1; 1; 1], rhoX = [1; 1; 1], rhoY = [1 1; 2 2; 3 3]) # Different size of output weights
    @test_throws ArgumentError deaprofitadd([1; 2; 3], [1; 2; 3], [1; 1; 1], [1; 1; 1], :Error) # Invalid inuts weights
    @test_throws ArgumentError deaprofitadd([1; 2; 3], [1; 2; 3], [1; 1; 1], [1; 1; 1], :Ones, rhoX= [1; 1; 1], rhoY = [1; 1; 1]) # Weights not allowed if model != :Custom

end


text:
abstract type AbstractStorageForm <: AbstractDeviceFormulation end

abstract type BookKeepingModel <: AbstractStorageForm end

# storage variables

function activepower_variables(ps_m::CanonicalModel, devices::Array{T,1}, time_range::UnitRange{Int64}) where {T <: PSY.Storage}

    add_variable(ps_m, devices, time_range, :Psin, false,:var_active, -1)
    add_variable(ps_m, devices, time_range, :Psout, false, :var_active)

    return

end

function reactivepower_variables(ps_m::CanonicalModel, devices::Array{T,1}, time_range::UnitRange{Int64}) where {T <: PSY.Storage}

    add_variable(ps_m, devices, time_range, :Qst, false, :var_reactive)

    return

end

function energystorage_variables(ps_m::CanonicalModel, devices::Array{T,1}, time_range::UnitRange{Int64}) where T <: PSY.Storage

    add_variable(ps_m, devices, time_range,:Est, false)

    return

end

function storagereservation_variables(ps_m::CanonicalModel, devices::Array{T,1}, time_range::UnitRange{Int64}) where T <: PSY.Storage

    add_variable(ps_m, devices, time_range, :Sst, true)

    return

end

# output constraints

function activepower_constraints(ps_m::CanonicalModel, devices::Array{St,1}, device_formulation::Type{D}, system_formulation::Type{S}, time_range::UnitRange{Int64}) where {St <: PSY.Storage, D <: AbstractStorageForm, S <: PM.AbstractPowerFormulation}

    range_data_in = [(s.name, s.inputactivepowerlimits) for s in devices]

    range_data_out = [(s.name, s.outputactivepowerlimits) for s in devices]

    device_semicontinuousrange(ps_m, range_data_in, time_range, :storage_inputpower_range, :Psin, :Est)

    reserve_device_semicontinuousrange(ps_m, range_data_in, time_range, :storage_outputpower_range, :Psout, :Sst)

    return

end

"""
This function adds the reactive  power limits of generators when there are CommitmentVariables
"""
function reactivepower_constraints(ps_m::CanonicalModel, devices::Array{St,1}, device_formulation::Type{D}, system_formulation::Type{S}, time_range::UnitRange{Int64}) where {St <: PSY.Storage, D <: AbstractStorageForm, S <: PM.AbstractPowerFormulation}

    range_data = [(s.name, s.reactivepowerlimits) for s in devices]

    device_range(ps_m, range_data , time_range, :storage_reactive_range, :Qst)

    return

end

# book keeping constraints

function energy_balance_constraint(ps_m::CanonicalModel, devices::Array{T,1}, device_formulation::Type{D}, system_formulation::Type{S}, time_range::UnitRange{Int64}, initial_conditions::Array{Float64,1}) where {T <: PSY.Storage, D <: AbstractStorageForm, S <: PM.AbstractPowerFormulation}

    named_initial_conditions = [(d.name, initial_conditions[ix]) for (ix, d) in enumerate(devices)]

    p_eff_data = [ (d.name,d.energy) for d in devices if !isa(d.energy, Nothing)]

    if !isempty(p_eff_data)

        energy_balance(ps_m,time_range,named_initial_conditions,p_eff_data, :energy_balance,(:Psout,:Psin,:Est))

    else
        @warn "Data doesn't contain Storage efficiency , consider adjusting your formulation"
    end

    return

end


text:
using ArgParse
using Distributed

@everywhere using DataFrames
@everywhere using Statistics
@everywhere using Printf
@everywhere using GZip
@everywhere using CSV

@everywhere const re_MolName = r"^#+ +Name: +"
@everywhere const re_Number = r"^#+ +Number: +"
@everywhere const re_OXR = r"^#+ +OXR +"
@everywhere const re_Total = r"^#+ +Total Energy: +"

@everywhere function build_gzfn(sub_idx, cls_idx)
    # zero padded cls_idx
    zp_idx = @sprintf "%04d" parse(Int64, cls_idx)

    mol2_fn = "$sub_idx/test.$zp_idx.mol2.gz"

    return mol2_fn
end

@everywhere function AUC(x, y)
    """
    Calculates AUC using the trapezoidal rule,
    x = FPR
    y = TPR
    """

    s = length(x)

    a = zeros(s)
    for i in 1:s-1
        x1, x2 = x[i:i+1]
        y1, y2 = y[i:i+1]

        a[i] = (x2-x1) * ((y1 + y2) / 2)
    end

    return sum(a)
end

@everywhere function LogAUC(x, y, lambda = 0.001)
    """
    Calculated Adjusted LogAUC as decribed in Mysinger/Shoichet 2010
    doi://10.1021/ci100214a
    """

    logauc_max = 1.0
    random_logauc = (logauc_max-lambda)/log.(10)/log10.(logauc_max/lambda)

    # identifies starting point for arrays (first nonzero FPR)
    start = 0
    for i in 1:length(x)
        if x[i] >= lambda
            start = i
            break
        end
    end

    # slice arrays to correct size
    new_x = x[start:end]
    new_y = y[start:end]

    # take log of x
    log_x = log10.(new_x)

    # center log10 of x
    log_x = log_x .+ abs(min(log_x...))

    # bound log_x to 0,1
    log_x = log_x ./ max(log_x...)

    # calculate AUC and adjust for random_logauc
    log_auc = AUC(log_x, new_y) - random_logauc

    return log_auc
end

@everywhere function ReadName(name_fn)
    names = []
    open(name_fn) do io
        while !eof(io)
            line = readline(io)
            push!(names, line)
        end
    end
    return Set(names)
end

@everywhere function label!(frame, l_set, d_set)
    frame[!, :Type] = map(
        x -> if in(x, l_set) "ligand" elseif in(x, d_set) "decoy" else missing end,
        frame[!, :mol_name]
    )
end

@everywhere function ROC(frame)
    """
    Requires frame with :Type annotated as ligand/decoy

    FPR : False Positive Rate (FP / FP + TN)
    TPR : True Positive Rate (TP / TP + FN)
    """

    roc = Dict(:FPR => Float64[0.0], :TPR => Float64[0.0])

    FP = 0
    TP = 0

    P = sum(frame[:, :Type] .== "ligand")
    N = sum(frame[:, :Type] .== "decoy")

    for t in frame[:, :Type]
        if t == "ligand"
            TP += 1
        else
            FP += 1
        end

        fpr = FP / N
        tpr = TP / P

        push!(roc[:FPR], fpr)
        push!(roc[:TPR], tpr)

    end

    return DataFrame(roc, copycols=false)
end

@everywhere function CalculateAUC(score_frame, ligand_set, decoy_set)

    label!(score_frame, ligand_set, decoy_set)
    sort!(score_frame, :Total)

    roc = ROC(score_frame)
    auc = AUC(roc[!, :FPR], roc[!, :TPR])
    log_auc = LogAUC(roc[!, :FPR], roc[!, :TPR])

    return [@sprintf "%.4f" i*100 for i in [auc, log_auc]]
end

@everywhere function AddMolecule(line_name, io, mol_set, sub_idx, cls_idx, coord_frame)
    values = [sub_idx, cls_idx]

    name = replace(strip(line_name), re_MolName => "")

    # skip 4
    [readline(io) for i in 1:4]
    number = replace(strip(readline(io)), re_Number => "")

    # exit if number is not found
    if !in(number, mol_set)
        return false

    else
        # skip 4
        [readline(io) for i in 1:4]

        # prepare OXR
        oxr = []
        for i in 1:4
            line = split(replace(readline(io), re_OXR => ""))
            append!(oxr, line)
        end

        # skip 14
        [readline(io) for i in 1:14]

        # parse Total
        total = replace(strip(readline(io)), re_Total=>"")

        append!(values, [name, number, total])
        append!(values, oxr)

        push!(coord_frame, values)

        return true
    end
end

@everywhere function ParseMol2(sub_idx, cls_idx, mol_frame, coord_frame)

    gz_fn = build_gzfn(sub_idx, cls_idx)
    mol_set = Set(mol_frame.mol_idx)

    GZip.open(gz_fn) do io
        while !eof(io)
            line = readline(io)
            if occursin(re_MolName, line)
                AddMolecule(line, io, mol_set, sub_idx, cls_idx, coord_frame)
            end
        end
    end
end

@everywhere function ParseOUTDOCK(subcluster, score_frame, time_frame)

    # define regex expressions
    re_header = r"mol#"
    re_elapsed = r"elapsed"
    re_values = r"^ +[0-9]"
    re_ignore = r"^ 9|colors|skip_size|poses|<|>|no_match|bump|clashes"
    re_timesub = r"(elapsed.+\(sec\): +| +\(hour\).+)"

    # define file name
    OUTDOCK = "$subcluster/OUTDOCK"

    # start file parse
    current_group = 0
    open(OUTDOCK) do io
        while !eof(io)
            line = readline(io)

            # new group found
            if occursin(re_header, line)
                current_group += 1

            elseif occursin(re_values, line)
                if occursin(re_ignore, line)
                    continue
                end
                values = split(line)
                prepend!(values, [subcluster, string(current_group)])
                push!(score_frame, values)

            # final line found
            elseif occursin(re_elapsed, line)
                time = replace(line, re_timesub => "")
                push!(time_frame, [subcluster, time])
            end

        end
    end
end

@everywhere function ProcessDir(dir_name, q)
    println(dir_name)

    # initialize dataframes
    score_frame = DataFrame(
        sub_idx = String[],
        cls_idx = String[],
        mol_idx = String[],
        mol_name = String[],
        flex_code = String[],
        matched = String[],
        nscored = String[],
        time = String[],
        hac = String[],
        setnum = String[],
        matnum = String[],
        rank = String[],
        cloud = String[],
        elect = String[],
        gist = String[],
        vdW = String[],
        psol = String[],
        asol = String[],
        inter = String[],
        rec_e = String[],
        rec_d = String[],
        r_hyd = String[],
        Total = String[]
    )
    time_frame = DataFrame(
        sub_idx = [],
        time = []
    )
    coord_frame = DataFrame(
        sub_idx = String[],
        cls_idx = String[],
        mol_name = String[],
        mol_idx = String[],
        mol_total = String[],
        OXR_X1 = String[],
        OXR_Y1 = String[],
        OXR_Z1 = String[],
        OXR_X2 = String[],
        OXR_Y2 = String[],
        OXR_Z2 = String[],
        OXR_X3 = String[],
        OXR_Y3 = String[],
        OXR_Z3 = String[],
        OXR_X4 = String[],
        OXR_Y4 = String[],
        OXR_Z4 = String[]
    )

    # generate sdi subcluster list
    subcluster_list = [
        joinpath(dir_name, i) for i in readdir(dir_name) if occursin("subcluster", i)
    ]

    # generate ligand/decoy filenames
    ligand_fn = joinpath(dir_name, "ligands.names")
    decoy_fn = joinpath(dir_name, "decoys.names")

    # generate ligand/decoy molecule name sets
    ligand_set = ReadName(ligand_fn)
    decoy_set = ReadName(decoy_fn)

    # parse each subcluster OUTDOCK
    for s in subcluster_list
        ParseOUTDOCK(s, score_frame, time_frame)
    end

    # convert types
    time_frame[!, :time] = map(x -> parse(Float64, x), time_frame[!, :time])
    score_frame[!, :Total] = map(x -> parse(Float64, x), score_frame[!, :Total])

    # sort on time and score/molecule
    sort!(time_frame, :time)
    sort!(score_frame, [:Total, :mol_name])

    # take best scoring unique molecules
    unique_score_frame = unique(score_frame, [:mol_name])

    # write dataframes
    CSV.write(
        "$dir_name/extract_all.sort.txt",
        select(score_frame, Not(:cls_idx)),
        delim="\t", header=false
        )
    CSV.write(
        "$dir_name/extract_all.sort.uniq.txt",
        select(unique_score_frame, Not(:cls_idx)),
        delim="\t", header=false
        )

    # Calculate AUC and LogAUC
    auc, log_auc = CalculateAUC(unique_score_frame, ligand_set, decoy_set)
    time_frame[!, :AUC] .= auc
    time_frame[!, :LogAUC] .= log_auc

    # calculate percentile
    pc = quantile!(unique_score_frame.Total, q)

    # filter scores for hit molecules
    top_score_frame = filter(x -> x.Total < pc, unique_score_frame)
    sort!(top_score_frame, [:sub_idx, :cls_idx, :mol_idx])

    # parse mol2 gzipped files for coordinates
    gdf = groupby(top_score_frame, [:sub_idx, :cls_idx])
    for (key, mol_frame) in pairs(gdf)
        ParseMol2(key.sub_idx, key.cls_idx, mol_frame, coord_frame)
    end

    # label type for ligand set
    label!(coord_frame, ligand_set, decoy_set)
    coord_frame[!, :mol_total] = map(
        x -> parse(Float64, x),
        coord_frame[!, :mol_total]
        )
    sort!(coord_frame, :mol_total)

    # write timing/enrichment dataframe
    CSV.write(
        "$dir_name/time_and_enrichment.tab",
        time_frame, delim="\t", header=false
        )

    # write coordinate dataframe
    CSV.write(
        "$dir_name/coords.tab",
        coord_frame, delim="\t"
    )
end

@everywhere function parallel_process(d, q)
    try
        ProcessDir(d, q)
    catch
        println("ERROR in {$d}")
        return
    end
end

function get_args()
    s = ArgParseSettings()

    @add_arg_table! s begin
        "--directories", "-i"
            help = "Directories to process"
            required = true
            nargs='+'
        "--quantile", "-q"
            help = "Quantile to filter scores (default = 0.1)"
            required = false
            default = 0.1
            arg_type = Float64
    end

    return parse_args(s)
end

function main()

    parsed_args = get_args()
    directories = parsed_args["directories"]

    arguments = [
        (d, parsed_args["quantile"]) for d in directories
    ]

    # Parallel Map Arguments
    pmap(
        (args)->parallel_process(args...),
        arguments
        )

end

main()

max_stars_repo_name: noamteyssier/ClusterDock

max_issues_repo_name: noamteyssier/ClusterDock

max_forks_repo_name: noamteyssier/ClusterDock


text: AbstractfLDA = Union{fLDA}


text:
using MPI
using NPZ
using Plots
using UltraDark

if ~MPI.Initialized()
    MPI.Init()
end
comm = MPI.COMM_WORLD
@info "Process ID $(MPI.Comm_rank(comm)+1) of $(MPI.Comm_size(comm)) with $(Threads.nthreads()) threads\n"

include(joinpath(@__DIR__, "init_soliton.jl"))

const resol = 64

const num_snapshots = 20
const output_times = LinRange(0, 5, num_snapshots)

const mass = 10
const position0 = [-4, 0, 0]
const velocity = [1, 0, 0]
const phase = 0
const t0 = 0

const output_dir = "output/vel"

function run_sim()

    output_config = OutputConfig(output_dir, output_times; box=true, slice=false)
    options = Config.SimulationConfig()

    grids = PencilGrids(10.0, resol)
    @info "Grids created" MPI.Comm_rank(comm) size(grids.ψk)

    mkpath(output_config.directory)
    npzwrite(joinpath(output_config.directory, "x.npy"), grids.x)
    npzwrite(joinpath(output_config.directory, "y.npy"), grids.y)
    npzwrite(joinpath(output_config.directory, "z.npy"), grids.z)

    add_soliton(grids, mass, position0, velocity, phase, t0)
    @info "Initialized" MPI.Comm_rank(comm)

    simulate(grids, options, output_config) == nothing
    @info "Simulation complete" MPI.Comm_rank(comm)

end

function plot_results()
    @info "Starting plots" MPI.Comm_rank(comm)
    pos_the = reshape(position0, (1, :)) .+ (reshape(velocity, (1, :)) .* reshape(output_times, (:, 1)))

    pos_sim = Array{Float64}(undef, (num_snapshots, 3))
    maxs = Array{Float64}(undef, num_snapshots)

    x = npzread(joinpath(output_dir, "x.npy"))
    y = npzread(joinpath(output_dir, "y.npy"))
    z = npzread(joinpath(output_dir, "z.npy"))

    for i in 1:num_snapshots
        rho = npzread(joinpath(output_dir, "rho_$i.npy"))

        rhomax, maxindex = findmax(rho)

        maxs[i] = rhomax
        pos_sim[i, :] .= [x[maxindex[1]], y[maxindex[2]], z[maxindex[3]]]

    end

    yerr = abs(x[2]-x[1])/2
    plot_pos = plot(legend=:bottomright)
    scatter!(output_times, pos_sim[:, 1], yerr=yerr, markershape=:auto, label="simulation")
    plot!(output_times, pos_the[:, 1], label="theory")
    plot!(xlabel=raw"$t$", ylabel=raw"$x$")

    savefig(plot_pos, "soliton_position.pdf")

    plot_max = plot(legend=:bottomright)
    plot!(output_times, maxs, label="theory")

    anim = @animate for i ∈ 1:num_snapshots
        rho = npzread(joinpath(output_dir, "rho_$i.npy"))
        psi = npzread(joinpath(output_dir, "psi_$i.npy"))

        density_plot = contourf(
                                rho[:, :, resol÷2],
                                title=raw"$\rho$",
                                clims=(0, maximum(maxs)),
                                showaxis=false,
                                ticks=false,
                                aspect_ratio=:equal,
                               )
        phase_plot = contourf(
                              angle.(psi[:, :, resol÷2]),
                              title="phase",
                              clims=(-π, +π),
                              showaxis=false,
                              ticks=false,
                              aspect_ratio=:equal,
                             )

        plot(density_plot, phase_plot, )
    end

    gif(anim, "soliton_velocity.gif", fps = 15)
end

run_sim()

if MPI.Comm_rank(comm) == 0
    plot_results()
end


text:
md"""
$$
  x[n] = A \cos(\omega n T + \phi) = A \cos(2\pi f n T + \phi)
$$
- $A$: amplitude
- $\omega$: angular frequency (`radians/sec`)
- $f = \omega / 2\pi$: frequency in Hertz (`cycles/sec`)
- $\phi$: initial phase (`radians`)
- $n$: time index
- $t = nT = n / f_s$, where $f_s$ is the sampling frequency

"""
begin
  import Pkg
  Pkg.activate(mktempdir())
  Pkg.add([
      "AudioDisplay",
      "PyPlot",
  ])
end

begin
  A = .8
  f = 1000    # i.e. 1000 cycles/sec
  ϕ = pi/2
  fs = 44100  # i.e. 44100 samples/sec
  t = range(-.002, .002, step=1/fs)
  x = A .* cos.(2πf.*t .+ ϕ)
  size(x)
end

A * cos.(2π*f.*t .+ ϕ)


text:
module aws_s3_control
include("AWSCorePrototypeServices.jl")
using .AWSCorePrototypeServices: s3_control

"""
Retrieves the PublicAccessBlock configuration for an Amazon Web Services account.
"""
GetPublicAccessBlock(AccountId) = s3_control("GET", "/v20180820/configuration/publicAccessBlock")

"""
Creates an access point and associates it with the specified bucket.
"""
CreateAccessPoint(AccountId, Name, Bucket) = s3_control("PUT", "/v20180820/accesspoint/$name")

"""
Returns the access point policy associated with the specified access point.
"""
GetAccessPointPolicy(AccountId, Name) = s3_control("GET", "/v20180820/accesspoint/$name/policy")

"""
Associates an access policy with the specified access point. Each access point can have only one policy, so a request made to this API replaces any existing policy associated with the specified access point.
"""
PutAccessPointPolicy(AccountId, Name, Policy) = s3_control("PUT", "/v20180820/accesspoint/$name/policy")

"""
Deletes the specified access point.
"""
DeleteAccessPoint(AccountId, Name) = s3_control("DELETE", "/v20180820/accesspoint/$name")

"""
Returns configuration information about the specified access point.
"""
GetAccessPoint(AccountId, Name) = s3_control("GET", "/v20180820/accesspoint/$name")

"""
Updates an existing job's priority.
"""
UpdateJobPriority(AccountId, JobId, Priority) = s3_control("POST", "/v20180820/jobs/$id/priority")

"""
Creates an Amazon S3 batch operations job.
"""
CreateJob(AccountId, Operation, Report, ClientRequestToken, Manifest, Priority, RoleArn) = s3_control("POST", "/v20180820/jobs")

"""
Returns a list of the access points currently associated with the specified bucket. You can retrieve up to 1000 access points per call. If the specified bucket has more than 1000 access points (or the number specified in maxResults, whichever is less), then the response will include a continuation token that you can use to list the additional access points.
"""
ListAccessPoints(AccountId) = s3_control("GET", "/v20180820/accesspoint")

"""
Deletes the access point policy for the specified access point.
"""
DeleteAccessPointPolicy(AccountId, Name) = s3_control("DELETE", "/v20180820/accesspoint/$name/policy")

"""
Indicates whether the specified access point currently has a policy that allows public access. For more information about public access through access points, see Managing Data Access with Amazon S3 Access Points in the Amazon Simple Storage Service Developer Guide.
"""
GetAccessPointPolicyStatus(AccountId, Name) = s3_control("GET", "/v20180820/accesspoint/$name/policyStatus")

"""
Retrieves the configuration parameters and status for a batch operations job.
"""
DescribeJob(AccountId, JobId) = s3_control("GET", "/v20180820/jobs/$id")

"""
Removes the PublicAccessBlock configuration for an Amazon Web Services account.
"""
DeletePublicAccessBlock(AccountId) = s3_control("DELETE", "/v20180820/configuration/publicAccessBlock")

"""
Lists current jobs and jobs that have ended within the last 30 days for the AWS account making the request.
"""
ListJobs(AccountId) = s3_control("GET", "/v20180820/jobs")

"""
Creates or modifies the PublicAccessBlock configuration for an Amazon Web Services account.
"""
PutPublicAccessBlock(PublicAccessBlockConfiguration, AccountId) = s3_control("PUT", "/v20180820/configuration/publicAccessBlock")

"""
Updates the status for the specified job. Use this operation to confirm that you want to run a job or to cancel an existing job.
"""
UpdateJobStatus(AccountId, JobId, RequestedJobStatus) = s3_control("POST", "/v20180820/jobs/$id/status")
end

max_stars_repo_name: nicoleepp/JuliaCloud-AWS-State

max_issues_repo_name: nicoleepp/JuliaCloud-AWS-State

max_forks_repo_name: nicoleepp/JuliaCloud-AWS-State