text:
"""
    eⱽ(M, r, a, θ)
Modified from Cunningham et al. (1975) eq. (A2a):
```math
e^\\nu = \\sqrt{\\frac{\\Delta \\Sigma}{A}}.
```
"""
eⱽ(M, r, a, θ) = √(Σ(r, a, θ) * Δ(M, r, a) / A(M, r, a, θ))

"""
    eᶲ(M, r, a, θ)
Modified from Cunningham et al. (1975) eq. (A2b):
```math
e^\\Phi = \\sin \\theta \\sqrt{\\frac{A}{\\Sigma}}.
```
"""
eᶲ(M, r, a, θ) = sin(θ) * √(A(M, r, a, θ) / Σ(r, a, θ))

"""
    ω(M, r, a, θ)
From Cunningham et al. (1975) eq. (A2c):
```math
\\omega = \\frac{2 a M r}{A}.
```
"""
ω(M, r, a, θ) = 2 * a * M * r / A(M, r, a, θ)

"""
    Ωₑ(M, r, a)
Coordinate angular velocity of an accreting gas.
Taken from Cunningham et al. (1975) eq. (A7b):
```math
\\Omega_e = \\frac{\\sqrt{M}}{a \\sqrt{M} + r_e^{3/2}}.
```
# Notes
Fanton et al. (1997) use
```math
\\Omega_e = \\frac{\\sqrt{M}}{a \\sqrt{M} \\pm r_e^{3/2}},
```
where the sign is dependent on co- or contra-rotation. This function may be extended in the future to support this definition.
"""
Ωₑ(M, r, a) = √M / (r^1.5 + a * √M)

"""
    Vₑ(M, r, a, θ)
Velocity of an accreting gas in a locally non-rotating reference frame (see Bardeen et al. 1973).
Taken from Cunningham et al. (1975) eq. (A7b):
```math
V_e = (\\Omega_e - \\omega) e^{\\Phi - \\nu}.
```
"""
Vₑ(M, r, a, θ) = (Ωₑ(M, r, a) - ω(M, r, a, θ)) * eᶲ(M, r, a, θ) / eⱽ(M, r, a, θ)

"""
    Lₑ(M, rms, a)
Angular momentum of an accreting gas within ``r_ms``.
Taken from Cunningham et al. (1975) eq. (A11b):
```math
L_e = \\sqrt{M} \\frac{
        r_{\\text{ms}}^2 - 2 a \\sqrt{M r_{\\text{ms}}} + a^2
    }{
        r_{\\text{ms}}^{3/2} - 2 M \\sqrt{r_{\\text{ms}}} + a \\sqrt{M}
    }.
```
"""
Lₑ(M, rms, a) = √M * (rms^2 - 2 * a * √(M * rms) + a^2) / (rms^1.5 - 2 * M * √rms + a * √M)

"""
    H(M, rms, r, a)
Taken from Cunningham et al. (1975) eq. (A12e):
```math
H = \\frac{2 M r_e - a \\lambda_e}{\\Delta},
```
where we distinguing ``r_e`` as the position of the accreting gas.
"""
H(M, rms, r, a) = (2 * M * r - a * Lₑ(M, rms, a)) / Δ(M, r, a)

"""
    γₑ(M, rms)
Taken from Cunningham et al. (1975) eq. (A11c):
```math
\\gamma_e = \\sqrt{1 - \\frac{
        2M
    }{
        3 r_{\\text{ms}}
    }}.
```
"""
γₑ(M, rms) = √(1 - (2 * M) / (3 * rms))

"""
    uʳ(M, rms, r)
Taken from Cunningham et al. (1975) eq. (A12b):
```math
u^r = - \\sqrt{\\frac{
        2M
    }{
        3 r_{\\text{ms}}
    }} \\left(
        \\frac{ r_{\\text{ms}} }{r_e} - 1
    \\right)^{3/2}.
```
"""
uʳ(M, rms, r) = -√((2 * M) / (3 * rms)) * (rms / r - 1)^1.5

"""
    uᶲ(M, rms, r, a)
Taken from Cunningham et al. (1975) eq. (A12c):
```math
u^\\phi = \\frac{\\gamma_e}{r_e^2} \\left(
        L_e + aH
    \\right).
```
"""
uᶲ(M, rms, r, a) = γₑ(M, rms) / r^2 * (Lₑ(M, rms, a) + a * H(M, rms, r, a))

"""
    uᵗ(M, rms, r, a)
Taken from Cunningham et al. (1975) eq. (A12b):
```math
u^t = \\gamma_e \\left(
        1 + \\frac{2 M (1 + H)}{r_e}
    \\right).
```
"""
uᵗ(M, rms, r, a) = γₑ(M, rms) * (1 + 2 * M * (1 + H(M, rms, r, a)) / r)

"""
Experimental API:
"""
function regular_pdotu_inv(L, M, r, a, θ)
    (eⱽ(M, r, a, θ) * √(1 - Vₑ(M, r, a, θ)^2)) / (1 - L * Ωₑ(M, r, a))
end

@inline function regular_pdotu_inv(m::AbstractMetricParams{T}, u, v) where {T}
    metric = GeodesicBase.metric(m, u)
    # reverse signs of the velocity vector
    # since we're integrating backwards
    p = @inbounds @SVector [-v[1], 0, 0, -v[4]]

    # TODO: this only works for Kerr
    disc_norm = (
        AccretionFormulae.eⱽ(m.M, u[2], m.a, u[3]) *
        √(1 - AccretionFormulae.Vₑ(m.M, u[2], m.a, u[3])^2)
    )

    u_disc =
        @SVector [1 / disc_norm, 0, 0, AccretionFormulae.Ωₑ(m.M, u[2], m.a) / disc_norm]

    # use Tullio to do the einsum
    @tullio g := metric[i, j] * (u_disc[i]) * p[j]
    1 / g
end

function plunging_p_dot_u(E, a, M, L, Q, rms, r, sign_r)
    inv(
        uᵗ(M, rms, r, a) - uᶲ(M, rms, r, a) * L -
        sign_r * uʳ(M, rms, r) * Σδr_δλ(E, L, M, Q, r, a) / Δ(M, r, a),
    )
end

function plunging_p_dot_u(m::AbstractMetricParams{T}, u, v, rms) where {T}
    metric = GeodesicBase.metric(m, u)

    # reverse signs of the velocity vector
    # since we're integrating backwards
    p = @inbounds @SVector [-v[1], v[2], 0, -v[4]]

    u_disc = @inbounds @SVector [
        uᵗ(m.M, rms, u[2], m.a),
        uʳ(m.M, rms, u[2]),
        0,
        uᶲ(m.M, rms, u[2], m.a),
    ]

    @tullio g := metric[i, j] * u_disc[i] * p[j]
    1 / g
end

@inline function redshift_function(m::CarterMethodBL{T}, u, p, λ) where {T}
    isco = rms(m.M, m.a)
    if u[2] > isco
        @inbounds regular_pdotu_inv(p.L, m.M, u[2], m.a, u[3])
    else
        # change sign if we're after the sign flip
        # TODO: this isn't robust to multiple sign changes
        # TODO: i feel like there should be a better way than checking this with two conditions
        #       used to have λ > p.changes[1] to make sure we're ahead of the time flip (but when wouldn't we be???)
        #       now p.changes[1] > 0.0 to make sure there was a time flip at all
        sign_r = (p.changes[1] > 0.0 ? 1 : -1) * p.r
        @inbounds plunging_p_dot_u(m.E, m.a, m.M, p.L, p.Q, isco, u[2], sign_r)
    end
end

@inline function redshift_function(m::AbstractMetricParams{T}, u, v) where {T}
    isco = rms(m.M, m.a)
    if u[2] > isco
        regular_pdotu_inv(m, u, v)
    else
        plunging_p_dot_u(m, u, v, isco)
    end
end

# point functions exports

function _redshift_guard(
    m::CarterMethodBL{T},
    gp::CarterGeodesicPoint{T},
    max_time,
) where {T}
    redshift_function(m, gp.u, gp.p, gp.t)
end
function _redshift_guard(m::AbstractMetricParams{T}, gp, max_time) where {T}
    redshift_function(m, gp.u, gp.v)
end

const redshift = PointFunction(_redshift_guard)


text:
tests = ["lp_01",
#        "lp_01a",
#        "lp_01b",
         "lp_02",
         "lp_03",
         "mip_01",
         "qp_01",
         "qp_02",
         "qcqp_01",
         "mathprog"]

for t in tests
    fp = "$(t).jl"
    println("running $(fp) ...")
    evalfile(fp)
end

ext: jl


text:
module ResourceContexts

export @context, @!, @defer, enter_do, ResourceContext

abstract type AbstractContext end

using Base.Meta: isexpr

using Logging: Logging, @logmsg

mutable struct ResourceContext <: AbstractContext
    resources::Vector{Any}
    is_detached::Bool
end

function ResourceContext(needs_finalizer::Bool=true)
    c = ResourceContext(Vector{Any}(), false)
    if needs_finalizer
        finalizer(cleanup!, c)
    end
    c
end

cleanup!(x) = close(x)
cleanup!(f::Function) = f()

# This partially recursive arrangement does two things:
# 1. Avoids paying the try/catch setup cost more than once when no exceptions
#    occur.
# 2. Arranges all exceptions on the exception stack, reflecting the cleanup
#    ordering and avoiding the need for CompositeException
function _cleanup!(resources, i)
    try
        while i > 0
            cleanup!(resources[i])
            i -= 1
        end
        i == 0 && empty!(resources)
    catch exc
        _cleanup!(resources, i-1)
        rethrow()
    end
end

function cleanup!(context::ResourceContext, unscope_cleanup_point=true)
    if !context.is_detached || unscope_cleanup_point
        # Clean up resources, last to first.
        _cleanup!(context.resources, length(context.resources))
    end
    nothing
end

"""
    defer(f::Function, `ctx`)

Defer the call to `f` until `ctx` is cleaned up.
"""
function defer(f::Function, ctx::ResourceContext)
    push!(ctx.resources, f)
    nothing
end

# Name of the context variable
const _context_name = :var"#context"

function _current_context_expr(__module__, __source__)
    quote
        if $(Expr(:islocal, esc(_context_name)))
            # Or this Expr(:isdefined)? Really, we'd like
            #   * Strict `:islocal` if inside a function
            #   * `:isdefined` for top level exprs
            $(esc(_context_name))
        else
            # Beautiful hack from Simeon Schaub here:
            # https://github.com/JuliaLang/julia/issues/6733#issuecomment-827441915
            in_function = $(Expr(:isdefined, esc(:var"#self#")))
            global_context($__module__, $(QuoteNode(__source__.file)),
                           $(__source__.line), in_function)
        end
    end
end

"""
    @defer expression

Defers execution of the cleanup `expression` until the exit of the current
`@context` block.
"""
macro defer(ex)
    quote
        ctx = $(_current_context_expr(__module__, __source__))
        defer(()->$(esc(ex)), ctx)
    end
end

"""
    @context begin ... end
    @context function f() ... end

`@context` creates a local context and runs the provided code within that
context. When the code exits, any resources registered with the context will be
cleaned up with `ResourceContexts.cleanup!()`.

When the code is a `function` definition, the context block is inserted around
the function body.
"""
macro context(ex)
    if ex.head == :function
        ex.args[2] = quote
            $(_context_name) = $ResourceContexts.ResourceContext(false)
            try
                $(ex.args[2])
            finally
                $ResourceContexts.cleanup!($(_context_name), false)
            end
        end
        esc(ex)
    else
        quote
            let $(esc(_context_name)) = ResourceContext(false)
                try
                    $(esc(ex))
                finally
                    cleanup!($(esc(_context_name)), false)
                end
            end
        end
    end
end

"""
    @! func(args)

The `@!` macro calls `func(ctx, args)`, where `ctx` is the "current context" as
created by the `@context` macro.  If `@!` is used outside a `@context` block, a
warning is emitted and the global context is used.

---

    @! function my_func(...)
        ...
    end

This form adds a `context::AbstractContext` argument as the first argument to
`my_func`. This allows the responsibility for resource cleanup to be passed
back to the caller by using `@defer` within `my_func` or `@!` to further chain
the resource handling.
"""
macro !(ex)
    if isexpr(ex, :call)
        i = length(ex.args) > 1 && isexpr(ex.args[2], :parameters) ? 3 : 2
        map!(a->esc(a), ex.args, ex.args)
        insert!(ex.args, i, :ctx)
        quote
            ctx = $(_current_context_expr(__module__, __source__))
            $ex
        end
    elseif isexpr(ex, :function)
        # Insert context argument
        callargs = isexpr(ex.args[1], :where) ?
                   ex.args[1].args[1].args : ex.args[1].args
        i = length(callargs) > 1 && isexpr(callargs[2], :parameters) ? 3 : 2
            # handle keywords
        insert!(callargs, i, :($_context_name::$ResourceContexts.AbstractContext))
        esc(ex)
    else
        error("Expected call or function definition as arguments to `@!`")
    end
end

# Global context
function __init__()
    atexit() do
        global_cleanup!()
    end
end

_global_context = ResourceContext(false)

@noinline function global_context(_module, file, line, in_function)
    # Heuristically reduce the log level in interactive usage because sloppy
    # resource management is probably ok there.
    filestr = string(file)
    level = ((_module === Main || startswith(filestr, "REPL")) && !in_function) ?
            Logging.Debug : Logging.Warn
    @logmsg level """Using global `ResourceContext` — use a `@context` block to avoid this message.
             Use `ResourceContexts.global_cleanup!()` to clean up the resource.""" #=
        =# _module=_module _group="context" _file=filestr _line=line
    return _global_context
end

function global_cleanup!()
    cleanup!(_global_context)
    empty!(_global_context.resources)
    nothing
end

#-------------------------------------------------------------------------------
# Utilities for interoperating with context-unaware code.

"""
    @! enter_do(func, args...)

`enter_do` transforms do-block-based resource management into context-based
resource management. That is, if the user is expected to write

```
func(args...) do x,y
    # do stuff with x,y
end
```

they can instead use `func` with a context, as in

```
x,y = @! enter_do(func, args...)
```
"""
@! function enter_do(func::Function, args...; kws...)
    # We need a separate stack to run the do block, so we use a separate task
    # here for that purpose. We know exactly which order the parent and child
    # tasks should run in so we should be able to use explicit yields and not
    # involve the main scheduler.
    parent_task = current_task()
    child_task = @task begin
        # cleanup_caller is the task which triggered yieldto(child_task) in the
        # @defer'd cleanup code further down.
        cleanup_caller::Task = parent_task
        try
            func(args...; kws...) do resources...
                cleanup_caller = yieldto(parent_task, resources)
            end
            # Success path: yield back to whoever triggered the cleanup.
            # `child_task` will never complete, but we've already run whatever
            # cleanup code in `func` so that should be ok.
            yieldto(cleanup_caller, nothing)
        catch
            # Failure path
            yieldto(cleanup_caller, :failed)
            rethrow()
        end
    end
    res = yieldto(child_task)
    if res === :failed
        # There was an exception during setup of the resource. We don't have a
        # value to return so force a TaskFailedException, ensuring that the
        # child task is scheduled so that it can run to completion while we
        # wait.
        schedule(child_task)
        wait(child_task)
    end
    @defer begin
        # Allow child task to free any resources
        #
        # If the user has called detach_context_cleanup(), this code may run in
        # a task other than parent_task and we need to make sure we yield back
        # appropriately.
        res = yieldto(child_task, current_task())
        if res === :failed
            # There was an exception during resource cleanup - report this with
            # a TaskFailedException.
            schedule(child_task)
            wait(child_task)
        end
    end
    return res
end

"""
    @! detach_context_cleanup(x)

Defer the cleanup of `context` to the time at which `x` is finalized by the
garbage collector and return `x`. This transforms context-based resource
handling into finalizer-based resource handling.

For this to work, `x` must be mutable. If `x` is immutable you could consider
making a wrapper type with the same API.

!!! note
    This function is best avoided if possible because it makes any failures during
    context cleanup impossible to handle neatly — instead, they are caught and
    logged asynchronously. However, it's very useful for interacting with
    context-unaware code.

# Examples

Create a temporary directory with two files in it, and return the directory
name as a string.  Cleanup of the directory is associated with finalization of
the string `dir`.

```
dir = @context begin
    dir = @! mktempdir()
    write(joinpath(dir, "file1.txt"), "Some content")
    write(joinpath(dir, "file2.txt"), "Some other content")
    @! ResourceContexts.detach_context_cleanup(dir)
end
```
"""
function detach_context_cleanup(ctx::AbstractContext, x)
    ctx.is_detached = true
    finalizer(x) do _
        # Must be async, as the finalizer itself isn't allowed to task switch
        # and context cleanup may involve several task-switchy things
        @async try
            cleanup!(ctx)
        catch exc
            @error "Error cleaning up context" exception=(exc, catch_backtrace())
            rethrow()
        end
    end
    x
end

#-------------------------------------------------------------------------------
include("base_interop.jl")

end


text:
include("../src/sweep_operator.jl")

@testset "Sweep Operator Correctness" begin
    correct_result = [1.1666666666666667 -0.5 -0.0 1.5; -0.5 0.25 -0.0 0.25; 0.0 -0.0 0.16666666666666666 0.3333333333333333; -1.5 -0.25 -0.3333333333333333 3.0833333333333335]
    M  = [  1.   1.   1.   1.
            1.   2.   1.   3.
            1.   3.   1.   3.
            1.   1.  -1.   2.
            1.   2.  -1.   2.
            1.   3.  -1.   1. ]
    M0 = M' * M
    sweep_op!(M0, 1:3)

    @test isapprox(correct_result, M0)
end


text:
# tests rule choice points

@generator RCShortGen begin # prefix RC (for Rule Choice points) to avoid type name clashes
    start() = x()
    x() = 'a'
    x() = 'b'
    x() = 'c'
end

@generator RCLongGen begin
    start() = x()
    function x()
        'a'
    end
    function x()
        'b'
    end
    function x()
        'c'
    end
end

@generator RCMixedGen begin
    start() = x()
    x() = 'a'
    function x()
        'b'
    end
    x() = 'c'
end

@testset "rule choice points" begin

	@testset "rule choice point using short form function definitions" begin

	    gn = RCShortGen()

		@mtestset reps=Main.REPS alpha=Main.ALPHA begin
        	td = choose(gn)
			@mtest_values_are ['a','b','c'] td
		end

	end

	@testset "rule choice point using long form function definitions" begin

	    gn = RCLongGen()

		@mtestset reps=Main.REPS alpha=Main.ALPHA begin
        	td = choose(gn)
			@mtest_values_are ['a','b','c'] td
		end

	end

	@testset "rule choice point using mixture of short and long form function definitions" begin

	    gn = RCMixedGen()

		@mtestset reps=Main.REPS alpha=Main.ALPHA begin
        	td = choose(gn)
			@mtest_values_are ['a','b','c'] td
		end

	end

end


text:
"""
    smoothconv(z,nas)

Smoothen field `z(ns,ns)` with a spectral method at scale `ns/nas`
Takes into account missing values.
"""
function smoothspec(zi,nas)

        @compat iinan=findall(isnan.(zi))
        @compat iinotnan=findall(.~isnan.(zi))
        zi[iinan]=0.

        nss=size(zi);
        ns=nss[1];
        sdim=div(ns,nas); # the smoothing radius is one large scale pixel wide (setting the diameter to 1 pixel is wrong)

        kx=[collect(0:ns/2)' collect(-ns/2+1:-1)']';
        kx=kx*ones(1,ns);
        kx=kx.^2+permutedims(kx,[2 1]).^2;

        zif=fft(zi);
        zif[kx[:].>(nas/2).^2]=0.0;
        zif[kx[:].==(nas/2).^2]=real(zif[kx[:].==(nas/2).^2]*2.0);
        zif=real(ifft(zif));

        return zif
end


text:
# originally from Ashok Litwin-Kumar (Eichler et al. 2017)
# with slight modification (refer to larva/learn.jl)

using Ipopt,MathProgBase,LinearAlgebra
#return maximal margin solution
function quadprogmargin(labels,pats)
	N = size(pats,1)
	P = length(labels)

#	Q = diagm(ones(N+1))
	Q = convert(Array, Diagonal(ones(N+1)))
	Q[N+1,N+1] = 0

	c = zeros(N+1)

	b = ones(P)

	A = zeros(N+1,P)

	for ni = 1:N
		for qi = 1:P
			A[ni,qi] = pats[ni,qi]*labels[qi]
		end
	end
	A[N+1,:] = labels

	sol = quadprog(c,Q,A','>',b,-Inf,Inf,IpoptSolver(print_level=0,max_iter=400))
	if length(sol.sol) > 0
		w = sol.sol[1:N]
		bias = sol.sol[N+1]
		return w/norm(w),bias/norm(w)
	else
		return zeros(N),0
	end
end


text:
module KoalaLow

import Koala: setup, fit, transform, inverse_transform

export setup, fit, transform, inverse_transform

end


text:
module SSLTestDataRead
using SimpleSDMLayers
using Test

file = joinpath(dirname(pathof(SimpleSDMLayers)), "..", "data", "connectivity.tiff")
struct MyConnectivityMap <: SimpleSDMLayers.SimpleSDMSource end
SimpleSDMLayers.latitudes(::Type{MyConnectivityMap}) = (-10.0, 10.0)
SimpleSDMLayers.longitudes(::Type{MyConnectivityMap}) = (-20.0, 20.0)
mp = SimpleSDMLayers.raster(SimpleSDMResponse, MyConnectivityMap(), file)

@test typeof(mp) <: SimpleSDMResponse
@test size(mp) == (1255, 1206)

end


text:
using ComradeBase
using Test

@testset "ComradeBase.jl" begin
    # Write your tests here.
end


text:
@testset "IO" begin

N = 256
I = shepp_logan(N)

# simulation parameters
params = Dict{Symbol, Any}()
params[:simulation] = "fast"
params[:trajName] = "Radial"
params[:numProfiles] = floor(Int64, pi/2*N)
params[:numSamplingPerProfile] = 2*N

# do simulation
acqData = simulation(I, params)

# reco parameters
params = Dict{Symbol, Any}()
params[:reco] = "direct"
params[:reconSize] = (N,N)
params[:alpha] = 1.75
params[:m] = 4.0
params[:K] = 28
Ireco = reconstruction(acqData, params)

saveImage( joinpath(tmpdir, "testReco.nii"), Ireco, false)
IrecoLoaded = loadImage( joinpath(tmpdir, "testReco.nii") )
@test (norm(vec(Ireco)-vec(IrecoLoaded))/norm(vec(Ireco))) < 1e-5

saveImage( joinpath(tmpdir, "testRecoAbs.nii"), Ireco, true)
IrecoLoadedAbs = loadImage( joinpath(tmpdir, "testRecoAbs.nii"))
@test (norm(vec(abs.(Ireco))-vec(IrecoLoadedAbs))/norm(vec(abs.(Ireco)))) < 1e-5

end


text:
using Flux, XLA, Test

x, y = randn(2), randn(1, 2)
@test collect(xla(*)(x, y)) ≈ x*y

f(x) = gradient(x -> sum(x.*x), x)

@test collect(xla(f)([1, 2, 3])[1]) == [2, 4, 6]

f(x) = gradient(x -> sum(x.+x), x)

@test collect(xla(f)([1, 2, 3])[1]) == [2, 2, 2]

W = rand(2, 3)

f(x) = gradient(x -> sum(W*x), x)

@test collect(xla(f)([1.0, 2, 3])[1]) == f([1.0, 2, 3])[1]

@test xla(sin'')(0.5) == -sin(0.5)

x = rand(10)
m = Chain(Dense(10, 5), Dense(5, 2)) |> f64

f(x) = gradient(m -> sum(m(x)), m)[1]

xf = xla(f)

@test collect(xf(x).layers[1].W) ≈ f(x).layers[1].W


text:
# https://github.com/joan2937/pigpio/blob/master/pigpio.py
module PiGPIO

export Pi

using Sockets

include("constants.jl")
include("pi.jl")
include("wave.jl")
include("i2c.jl")
include("spiSerial.jl")

end


text:
p_school  = 0.3
p_work 	  = 0.03
p_leisure = 0.003

function start_contact_generation(df::DataFrame, parameters::Contact_simulation_options)
	# Create the check-in writer
	event_writer = EventWriter("resources/generated_contact.csv")
	# Define simulation start time
	day = 1
	month = 1
	year = 2021
	simulation_day = DateTime(year,month,day)
	households = get_households(df,parameters)
	cont = 0
	simulation_num_days = parameters.days

	for day in 1:simulation_num_days
		############ CONTACT WITH HOUSEHOLD ###########
		workers = Set{Person}()
		students = Set{Person}()
		all_people = Set{Person}()
		for household in households
			workers = workers ∪ household.workers
			students = students ∪ household.students
			family = union(household.workers,household.students,household.unemployeds)
			all_people = all_people ∪ family
			for person in family
				add_event(event_writer,Event(person.id,household.id,PLACE_HOME,DateTime(year,month,day,DAY_START)))
				cont += 1
			end
		end
		write_events(event_writer)

		println("|People| = $(length(all_people)) - |students| = $(length(students)) - |workers| = $(length(workers))")

		# ########## WORK CONTACT #############
		if !_is_holiday(simulation_day)
			cont = 0
			for workplace in 1:parameters.work_place_num
				workers_same_workplace = filter(w -> w.work_school_place == workplace,workers)
				println("Workplace $workplace: |Size| = $(length(workers_same_workplace))")

				combinations_iterator = combinations(collect(workers_same_workplace),2)
				for couple in combinations_iterator
					if rand() > p_work
						continue
					end
					worker_1 = couple[1]
					worker_2 = couple[2]
					contact_hour = rand(WORK_START:WORK_END)
					add_event(event_writer,Event(worker_1.id,worker_1.work_school_place,PLACE_WORK, simulation_day + Dates.Hour(contact_hour)))
					add_event(event_writer,Event(worker_2.id,worker_2.work_school_place,PLACE_WORK, simulation_day + Dates.Hour(contact_hour)))
					cont += 2
				end
				write_events(event_writer)
			end
			println("Work contact: $(cont)")
		end

		# pre_school_students, primary_students, secondary_students, higher_students = divide_students_by_age(students)

		######### SCHOOL CONTACT #############
		if !_is_holiday(simulation_day)
			cont = 0
			for school in 1:parameters.school_place_num
				students_same_school = filter(s -> s.work_school_place == school, students)
				println("School $school: |Size| = $(length(students_same_school))")

				combinations_iterator = combinations(collect(students_same_school), 2)
				for couple in combinations_iterator
						if rand() > p_school #Skip if the probability check fails
							continue
						end
						student_1 = couple[1]
						student_2 = couple[2]
						contact_hour = rand(SCHOOL_START:SCHOOL_END)
						add_event(event_writer,Event(student_1.id,student_1.work_school_place,PLACE_SCHOOL,simulation_day + Dates.Hour(contact_hour)))
						add_event(event_writer,Event(student_2.id,student_2.work_school_place,PLACE_SCHOOL,simulation_day + Dates.Hour(contact_hour)))
						cont += 2
				end
				write_events(event_writer)
			end
			println("School contact: $(cont)")
		end

		######### LEISURE CONTACT #############
		println(" |Size| = $(length(all_people))")

		combinations_iterator = combinations(collect(all_people),2)
		cont = 0
		for couple in combinations_iterator
			if rand() > p_leisure
				continue
			end

			person_1 = couple[1]
			person_2 = couple[2]

			if !_is_holiday(simulation_day) && (person_1 ∈ workers || person_2 ∈ workers)
				contact_hour = rand(WORK_END:DAY_END)
			elseif !_is_holiday(simulation_day) && (person_1 ∈ students || person_2 ∈ students)
				contact_hour = rand(SCHOOL_END:DAY_END)
			else
				contact_hour = rand((DAY_START+1):DAY_END)
			end

			contact_place = rand(1:parameters.leisure_place_num)
			add_event(event_writer,Event(person_1.id,contact_place,PLACE_LEISURE,simulation_day + Dates.Hour(contact_hour)))
			add_event(event_writer,Event(person_2.id,contact_place,PLACE_LEISURE,simulation_day + Dates.Hour(contact_hour)))
			cont += 2
		end
		println("Leisure contact: $(cont)")

		simulation_day = simulation_day + Dates.Day(1)
	end

	close_writer(event_writer)
end

function get_random_subset(origin::Set, num::Int)
	tmp = origin
	subset = Set()
	for i in 1:num
		elem = rand(tmp,1)[1];
		push!(subset,elem)
		tmp = setdiff(tmp,[elem])
	end
	subset
end

function equalize_sets(set1::Set, set2::Set, bucket::Set)
	while length(set1) < length(set2)
		push!(set1,rand(setdiff(bucket,set1 ∪ set2)))
	end
end

max_stars_repo_name: Luigi-Crisci/Hypergraphs-Epidemics

max_issues_repo_name: Luigi-Crisci/Hypergraphs-Epidemics

max_forks_repo_name: Luigi-Crisci/Hypergraphs-Epidemics


text:
using Dates
using TimeFrames
using TimeFrames: TimePeriodFrame, DatePeriodFrame
using TimeFrames: period_step, _period_step
using TimeFrames: CustomTimeFrame, tonext

using Test

@testset "low level (period frame)" begin
    @testset "TimePeriodFrame" begin
        tf = TimePeriodFrame{Dates.Hour}()
        @test tf.period.value == 1
        @test tf.boundary == Begin

        tf = TimePeriodFrame{Dates.Hour}(5)
        @test tf.period.value == 5
        @test tf.boundary == Begin

        tf2 = TimePeriodFrame{Dates.Hour}(5)
        @test tf == tf2

        tf3 = TimePeriodFrame{Dates.Hour}(3)
        @test tf != tf3
    end

    @testset "DatePeriodFrame" begin
        tf = DatePeriodFrame{Dates.Month}()
        @test tf.period.value == 1
        @test tf.boundary == Begin

        tf = DatePeriodFrame{Dates.Month}(5)
        @test tf.period.value == 5
        @test tf.boundary == Begin

        tf2 = DatePeriodFrame{Dates.Month}(5)
        @test tf == tf2

        tf3 = DatePeriodFrame{Dates.Month}(3)
        @test tf != tf3
    end

    @testset "NoTimeFrame" begin
        tf = TimeFrame()
        @test typeof(tf) == NoTimeFrame

        tf2 = TimeFrame()
        @test tf == tf2
    end

    @testset "_period_step" begin
        @test _period_step(DateTime) == period_step
        @test _period_step(Date) == Dates.Day(1)
    end
end

@testset "high level" begin
    @testset "YearEnd, Minute, ..." begin
        tf = YearEnd()
        @test tf.period.value == 1

        @test YearEnd() == YearEnd()
        @test YearEnd() != YearBegin()
        @test YearEnd() != YearEnd(5)

        tf = TimeFrames.Minute()
        @test tf.period.value == 1

        tf = TimeFrames.Minute(15)
        @test tf.period.value == 15

        tf1 = TimeFrames.Minute(15)
        tf2 = TimeFrames.Minute(15)
        @test tf1 == tf2

        tf1 = TimeFrames.Minute(15)
        tf2 = TimeFrames.Minute(30)
        @test tf1 != tf2
    end

    @testset "to string" begin
        tf = TimeFrames.Minute()
        @test String(tf) == "T"

        tf = TimeFrames.Minute(15)
        @test String(tf) == "15T"
    end

    @testset "parse" begin

        @testset "simple parse" begin
            tf = TimeFrame("15T")
            @test tf.period.value == 15
            @test typeof(tf) == TimeFrames.Minute

            tf = TimeFrame("T")
            @test tf.period.value == 1
            @test typeof(tf) == TimeFrames.Minute

            tf = TimeFrame("15Min")
            @test tf.period.value == 15
            @test String(tf) == "15T"
            @test typeof(tf) == TimeFrames.Minute

            tf = TimeFrame("5H")
            @test tf.period.value == 5
            @test typeof(tf) == TimeFrames.Hour
        end

        @testset "NoTimeFrame parse" begin
            tf = TimeFrame("")
            @test typeof(tf) == NoTimeFrame
            @test NoTimeFrame() == NoTimeFrame(1,2,3)
        end

        @testset "boundary" begin
            tf = TimeFrame("3A")
            @test tf == YearEnd(3)
            @test tf.period == Dates.Year(3)
            @test tf.boundary == End

            tf = TimeFrame("3AS")
            @test tf == YearBegin(3)
            @test tf.period == Dates.Year(3)
            #@test tf.boundary == Begin  # ToFix

            tf = TimeFrame("3M")
            @test tf == MonthEnd(3)
            @test tf.period == Dates.Month(3)
            @test tf.boundary == End

            tf = TimeFrame("3MS")
            @test tf == MonthBegin(3)
            @test tf.period == Dates.Month(3)
            #@test tf.boundary == Begin  # ToFix
        end

        @testset "*" begin
            @testset "TimePeriodFrame * n / n * TimePeriodFrame" begin
                tf = TimeFrame("2H")
                @test tf * 3 == TimeFrame("6H")

                tf = TimeFrame("2H")
                @test 3 * tf == TimeFrame("6H")
            end
            @testset "DatePeriodFrame * n / n * DatePeriodFrame" begin
                tf = TimeFrame("2M")
                @test tf * 3 == TimeFrame("6M")

                tf = TimeFrame("2M")
                @test 3 * tf == TimeFrame("6M")
            end
        end

        @testset "grouper/apply" begin
            @test apply(MonthEnd(), Date(2010, 2, 20)) == Date(2010, 2, 28)
            @test apply(MonthEnd(), DateTime(2010, 2, 20)) == DateTime(2010, 3, 1) - period_step

            d = Date(2016, 7, 20)
            dt = DateTime(2016, 7, 20, 13, 24, 35, 245)

            tf = TimeFrame(dt -> floor(dt, Dates.Minute(15)))  # custom TimeFrame with lambda function as DateTime grouper
            @test apply(tf, dt) == DateTime(2016, 7, 20, 13, 15, 0, 0)
            @test typeof(tf) == CustomTimeFrame

            tf = TimeFrame(Dates.Minute(15))  # TimePeriodFrame using TimePeriod
            @test apply(tf, dt) == DateTime(2016, 7, 20, 13, 15, 0, 0)

            tf = TimeFrame(Dates.Day(1))  # TimePeriodFrame using DatePeriod
            @test apply(tf, dt) == DateTime(2016, 7, 20, 0, 0, 0, 0)

            tf = YearBegin()
            #@test apply(tf, dt) == DateTime(2016, 1, 1, 0, 0, 0, 0)
            @test apply(tf, dt) == Date(2016, 1, 1)
            #@test typeof(f_group(dt)) == Date

            #tf = YearEnd(boundary=End)
            #@test apply(tf, dt) == DateTime(2016, 1, 1, 0, 0, 0, 0)
            #@test apply(tf, dt) == Date(2016, 1, 1)

            tf = MonthBegin()
            #@test apply(tf, dt) == DateTime(2016, 7, 1, 0, 0, 0, 0)
            @test apply(tf, dt) == Date(2016, 7, 1)

            tf = TimeFrames.Week()
            #@test apply(tf, dt) == DateTime(2016, 7, 18, 0, 0, 0, 0)
            @test apply(tf, dt) == Date(2016, 7, 18)

            tf = TimeFrames.Day()
            #@test apply(tf, dt) == DateTime(2016, 7, 20, 0, 0, 0, 0)
            @test apply(tf, dt) == Date(2016, 7, 20)

            tf = TimeFrames.Hour()
            @test apply(tf, dt) == DateTime(2016, 7, 20, 13, 0, 0, 0)

            tf = TimeFrames.Minute()
            @test apply(tf, dt) == DateTime(2016, 7, 20, 13, 24, 0, 0)

            tf = TimeFrames.Second()
            @test apply(tf, dt) == DateTime(2016, 7, 20, 13, 24, 35, 0)

            tf = TimeFrames.Millisecond()
            @test apply(tf, dt) == DateTime(2016, 7, 20, 13, 24, 35, 245)

            tf = YearBegin(10)
            @test apply(tf, dt) == DateTime(2010, 1, 1, 0, 0, 0, 0)

            tf = YearEnd(10)
            @test apply(tf, d) == DateTime(2019, 12, 31)
            @test apply(tf, dt) == DateTime(2020, 1, 1) - period_step

            tf = TimeFrames.Minute(15)
            @test apply(tf, dt) == DateTime(2016, 7, 20, 13, 15, 0, 0)
        end

        @testset "tonext" begin
            dt = DateTime(2010, 1, 1, 10, 30)
            @test tonext(TimeFrame("2H"), dt) == DateTime(2010, 1, 1, 12, 0)

            dt = DateTime(2010, 1, 1, 0, 0)  # ; same=false
            @test tonext(TimeFrame("2H"), dt) == DateTime(2010, 1, 1, 2, 0)

            dt = DateTime(2010, 1, 1, 0, 0)
            @test tonext(TimeFrame("2H"), dt; same=true) == DateTime(2010, 1, 1, 0, 0)

            dt = DateTime(2010, 1, 1, 10, 30)
            @test tonext(TimeFrame("1D"), dt) == DateTime(2010, 1, 2, 0, 0)

            dt = DateTime(2010, 1, 1, 10, 30)
            @test tonext(TimeFrame("1MS"), dt) == DateTime(2010, 2, 1, 0, 0)

            dt = DateTime(2010, 1, 1, 10, 30)
            @test tonext(TimeFrame("1M"), dt) == DateTime(2010, 2, 1) - period_step

            dt = DateTime(2010, 6, 5, 10, 30)
            @test tonext(TimeFrame("1A"), dt) == DateTime(2011, 1, 1) - period_step

            dt = DateTime(2010, 6, 5, 10, 30)
            @test tonext(TimeFrame("1AS"), dt) == DateTime(2011, 1, 1)

          end

        @testset "range" begin
            dt1 = DateTime(2010, 1, 1, 20)
            dt2 = DateTime(2010, 1, 14, 16)
            tf = Dates.Day(1)
            rng = range(dt1, TimeFrame(tf), dt2)
            @test rng[1] == DateTime(2010, 1, 1)
            @test rng[end] == DateTime(2010, 1, 14)

            rng = range(dt1, TimeFrame(tf), dt2, apply_tf=false)
            @test rng[1] == DateTime(2010, 1, 1, 20)
            @test rng[end] == DateTime(2010, 1, 13, 20)

            tf = Dates.Day(1)
            N = 5
            rng = range(dt1, TimeFrame(tf), N)
            @test length(rng) == N
            @test rng[1] == DateTime(2010, 1, 1, 20)
            @test rng[end] == DateTime(2010, 1, 5, 20)

            tf = Dates.Day(1)
            N = 5
            rng = range(TimeFrame(tf), dt2, N)
            @test length(rng) == N
            @test rng[end] == DateTime(2010, 1, 13, 16)
            @test rng[1] == DateTime(2010, 1, 9, 16)

            dt1 = DateTime(2010, 1, 1, 20)
            dt2 = DateTime(2010, 1, 14, 16)
            tf = NoTimeFrame()
            rng = range(dt1, tf, dt2)
            @test length(rng) == 1
            @test rng[1] == dt1
        end

        @testset "@tf" begin
            @test tf"1T" == TimeFrame("1T")
            @test tf"1M" == TimeFrame("M")
        end  # @testset "@tf
    end

    @testset "+" begin
        d = Date(2017, 12, 1)
        let tf = tf"1A", ans = Date(2018, 12, 1)
            @test isdefined(tf, :period)
            @test d + tf == ans
            @test tf + d == ans
            @inferred d + tf
            @inferred tf + d
        end

        let tf = tf"1AS", ans = Date(2018, 12, 1)
            @test isdefined(tf, :period)
            @test d + tf == ans
            @test tf + d == ans
            @inferred d + tf
            @inferred tf + d
        end

        let tf = tf"1M", ans = Date(2018, 1, 1)
            @test isdefined(tf, :period)
            @test d + tf == ans
            @test tf + d == ans
            @inferred d + tf
            @inferred tf + d
        end

        let tf = tf"1MS", ans = Date(2018, 1, 1)
            @test isdefined(tf, :period)
            @test d + tf == ans
            @test tf + d == ans
            @inferred d + tf
            @inferred tf + d
        end

        let tf = tf"1W", ans = Date(2017, 12, 8)
            @test isdefined(tf, :period)
            @test d + tf == ans
            @test tf + d == ans
            @inferred d + tf
            @inferred tf + d
        end

        let tf = tf"1D", ans = Date(2017, 12, 2)
            @test isdefined(tf, :period)
            @test d + tf == ans
            @test tf + d == ans
            @inferred d + tf
            @inferred tf + d
        end

        let tf = tf"1H", ans = DateTime(2017, 12, 1, 1, 0, 0, 0)
            @test isdefined(tf, :period)
            @test d + tf == ans
            @test tf + d == ans
            @inferred d + tf
            @inferred tf + d
        end

        let tf = tf"1T", ans = DateTime(2017, 12, 1, 0, 1, 0, 0)
            @test isdefined(tf, :period)
            @test d + tf == ans
            @test tf + d == ans
            @inferred d + tf
            @inferred tf + d
        end

        let tf = tf"1S", ans = DateTime(2017, 12, 1, 0, 0, 1, 0)
            @test isdefined(tf, :period)
            @test d + tf == ans
            @test tf + d == ans
            @inferred d + tf
            @inferred tf + d
        end

        let tf = tf"1L", ans = DateTime(2017, 12, 1, 0, 0, 0, 1)
            @test isdefined(tf, :period)
            @test d + tf == ans
            @test tf + d == ans
            @inferred d + tf
            @inferred tf + d
        end

        d = Dates.Time(0, 1, 2, 3)
        let tf = tf"1H", ans = Dates.Time(1, 1, 2, 3)
            @test isdefined(tf, :period)
            @test d + tf == ans
            @test tf + d == ans
            @inferred d + tf
            @inferred tf + d
        end

        let tf = tf"1T", ans = Dates.Time(0, 2, 2, 3)
            @test isdefined(tf, :period)
            @test d + tf == ans
            @test tf + d == ans
            @inferred d + tf
            @inferred tf + d
        end

        let tf = tf"1S", ans = Dates.Time(0, 1, 3, 3)
            @test isdefined(tf, :period)
            @test d + tf == ans
            @test tf + d == ans
            @inferred d + tf
            @inferred tf + d
        end

        let tf = tf"1L", ans = Dates.Time(0, 1, 2, 4)
            @test isdefined(tf, :period)
            @test d + tf == ans
            @test tf + d == ans
            @inferred d + tf
            @inferred tf + d
        end

        d = Dates.Time(0, 1, 2, 3)
        @test_throws InexactError d + tf"1A"
        @test_throws InexactError d + tf"1AS"
        @test_throws InexactError d + tf"1M"
        @test_throws InexactError d + tf"1MS"
        @test_throws InexactError d + tf"1W"
        @test_throws InexactError d + tf"1D"

        @test_throws InexactError tf"1A"  + d
        @test_throws InexactError tf"1AS" + d
        @test_throws InexactError tf"1M"  + d
        @test_throws InexactError tf"1MS" + d
        @test_throws InexactError tf"1W"  + d
        @test_throws InexactError tf"1D"  + d
    end  # @testset "+"

    @testset "-" begin
        d = Date(2017, 12, 1)
        let tf = tf"1A"
            @test isdefined(tf, :period)
            @test d - tf == Date(2016, 12, 1)
            @inferred d - tf
        end

        let tf = tf"1AS"
            @test isdefined(tf, :period)
            @test d - tf == Date(2016, 12, 1)
            @inferred d - tf
        end

        let tf = tf"1M"
            @test isdefined(tf, :period)
            @test d - tf == Date(2017, 11, 1)
            @inferred d - tf
        end

        let tf = tf"1MS"
            @test isdefined(tf, :period)
            @test d - tf == Date(2017, 11, 1)
            @inferred d - tf
        end

        let tf = tf"1W"
            @test isdefined(tf, :period)
            @test d - tf == Date(2017, 11, 24)
            @inferred d - tf
        end

        let tf = tf"1D"
            @test isdefined(tf, :period)
            @test d - tf == Date(2017, 11, 30)
            @inferred d - tf
        end

        let tf = tf"1H"
            @test isdefined(tf, :period)
            @test d - tf == DateTime(2017, 11, 30, 23, 0, 0, 0)
            @inferred d - tf
        end

        let tf = tf"1T"
            @test isdefined(tf, :period)
            @test d - tf == DateTime(2017, 11, 30, 23, 59, 0, 0)
            @inferred d - tf
        end

        let tf = tf"1S"
            @test isdefined(tf, :period)
            @test d - tf == DateTime(2017, 11, 30, 23, 59, 59, 0)
            @inferred d - tf
        end

        let tf = tf"1L"
            @test isdefined(tf, :period)
            @test d - tf == DateTime(2017, 11, 30, 23, 59, 59, 999)
            @inferred d - tf
        end

        d = Dates.Time(0, 1, 2, 3)
        let tf = tf"1H"
            @test isdefined(tf, :period)
            @test d - tf == Dates.Time(23, 1, 2, 3)
            @inferred d - tf
        end

        let tf = tf"1T"
            @test isdefined(tf, :period)
            @test d - tf == Dates.Time(0, 0, 2, 3)
            @inferred d - tf
        end

        let tf = tf"1S"
            @test isdefined(tf, :period)
            @test d - tf == Dates.Time(0, 1, 1, 3)
            @inferred d - tf
        end

        let tf = tf"1L"
            @test isdefined(tf, :period)
            @test d - tf == Dates.Time(0, 1, 2, 2)
            @inferred d - tf
        end

        d = Dates.Time(0, 1, 2, 3)
        @test_throws InexactError d - tf"1A"
        @test_throws InexactError d - tf"1AS"
        @test_throws InexactError d - tf"1M"
        @test_throws InexactError d - tf"1MS"
        @test_throws InexactError d - tf"1W"
        @test_throws InexactError d - tf"1D"
    end  # @testset "-"
end  # @testset "high level"


text:
using PyPlot
using Redpitaya

# OLD FILE: DO NOT USE!!!

#rp = RedPitaya("10.167.6.99")
rp = RedPitaya("192.168.1.26")

baseFreq = 125e6
dec = 8

numSampPerPeriod = 100
numPeriods = 10
numSamp = numSampPerPeriod*numPeriods
freq = div(baseFreq,dec*numSampPerPeriod)#*dec

println(freq)

# start sending
send(rp,"GEN:RST")
sendAnalogSignal(rp,1,"SINE",freq,0.4,numPeriods*2)

#sendAnalogSignal(rp,2,"TRIANGLE",25000,0.4)

# prepare acquisition
# Set decimation vale (sampling rate) in respect to you
# acquired signal frequency
acqReset(rp)
decimation(rp,dec)
send(rp,"ACQ:TRIG:LEV 0.01")

# Set trigger delay to 0 samples
# 0 samples delay set trigger to center of the buffer
# Signal on your graph will have trigger in the center (symmetrical)
# Samples from left to the center are samples before trigger
# Samples from center to the right are samples after trigger
send(rp,"ACQ:TRIG:DLY 8190")
#send(rp,"ACQ:TRIG:DLY 0")

## Start & Trigg
# Trigger source setting must be after ACQ:START
# Set trigger to source 1 positive edge

acqStart(rp)
# After acquisition is started some time delay is needed in order to acquire fresh samples in to buffer
# Here we have used time delay of one second but you can calculate exact value taking in to account buffer
# length and smaling rate
sleep(0.1)

send(rp,"ACQ:TRIG AWG_PE")
send(rp,"SOUR1:TRIG:IMM")
#send(rp,"SOUR2:TRIG:IMM")

# Wait for trigger
# Until trigger is true wait with acquiring
# Be aware of while loop if trigger is not achieved
# Ctrl+C will stop code executing in Julia

while true
  send(rp,"ACQ:TRIG:STAT?")
  trig_rsp = receive(rp)
  println(trig_rsp)
  if trig_rsp[1:2] == "TD"
     break
   end
end
#sleep(0.1)
#acqStop(rp)

println( query(rp,"ACQ:TPOS?") )
println( query(rp,"ACQ:WPOS?") )

tpos = parse(Int64,query(rp,"ACQ:TPOS?"))

# Read data from buffer
u1=receiveAnalogSignalLowLevel(rp,1)
u2=receiveAnalogSignalLowLevel(rp,1,0,numSamp,"STA")
u3=receiveAnalogSignalLowLevel(rp,1,0,numSamp,"OLD")
u4=receiveAnalogSignalLowLevel(rp,1,0,numSamp,"LAT")
u5=receiveAnalogSignalLowLevel(rp,1,tpos,numSamp,"STA")

figure(1)
clf()
subplot(2,1,1)
plot(u1)
subplot(2,1,2)
semilogy(abs(rfft(u1))[1:numPeriods*4],"o-b",lw=2)

figure(2)
clf()
subplot(5,1,1)
plot(u1,"g")
subplot(5,1,2)
plot(u2,"b")
subplot(5,1,3)
plot(u3,"r")
subplot(5,1,4)
plot(u4,"y")
subplot(5,1,5)
plot(u5,"y")
#subplot(2,1,2)
#plot(u2)


text:
using FFTW, .ForwardDiff
import FFTW.fft!;

function fft!(x::AbstractArray{T1}) where {T1 <: Complex{cpx}} where {cpx <: ForwardDiff.Dual{Tg, T, N}} where {Tg, T <: Real, N}
    Xcomplex = [ForwardDiff.value(real(el)) + 1im * ForwardDiff.value(imag(el)) for el in x]
    Xder1 = reduce(hcat, [ForwardDiff.partials(real(el)) + 1im * ForwardDiff.partials(imag(el)) for el in x])'
    planned_fft = plan_fft!(similar(Xcomplex))
    Xcomplex .= planned_fft * Xcomplex
    for col in eachcol(Xder1)
        col .= planned_fft * col
    end
    npartials = size(Xder1, 2)
    epss = [ForwardDiff.Dual{Tg}(0, ntuple(i -> i == k, Val(npartials))) for k = 1:npartials]
    x .= Xcomplex .+ Xder1 * epss
    nothing
end


text:
@testset "Edit distance" begin

    gtri = random_regular_graph(3, 2)
    gquad = random_regular_graph(4, 2)
    gpent = random_regular_graph(5, 2)

    @testset "edit_distance $triangle, $quadrangle, $pentagon" for triangle in testgraphs(gtri), quadrangle in testgraphs(gquad), pentagon in testgraphs(gpent)
        d, λ = @inferred(edit_distance(triangle, quadrangle, subst_cost=MinkowskiCost(1:3, 1:4)))
        @test d == 1.0
        @test λ == Tuple[(1, 1), (2, 2), (3, 3), (0, 4)]

        d, λ = @inferred(edit_distance(quadrangle, triangle, subst_cost=MinkowskiCost(1:4, 1:3)))
        @test d == 1.0
        @test λ == Tuple[(1, 1), (2, 2), (3, 3), (4, 0)]

        d, λ = @inferred(edit_distance(triangle, pentagon, subst_cost=MinkowskiCost(1:3, 1:5)))
        @test d == 2.0
        @test λ == Tuple[(1, 1), (2, 2), (3, 3), (0, 4), (0, 5)]

        d, λ = @inferred(edit_distance(pentagon, triangle, subst_cost=MinkowskiCost(1:5, 1:3)))
        @test d == 2.0
        @test λ == Tuple[(1, 1), (2, 2), (3, 3), (4, 0), (5, 0)]
    end

    @testset "Minkowski cost / bounded Minkowski" begin
        cost = @inferred(MinkowskiCost(1:3, 1:3))
        bcost = @inferred(BoundedMinkowskiCost(1:3, 1:3))
        for i = 1:3
          @test cost(i, i) == 0.
          @test bcost(i, i) == 2 / 3
        end
    end

    g1c = complete_graph(4)
    g2c = complete_graph(4)
    rem_edge!(g2c, 1, 2)
    @testset "edit_distance $g1, $g2" for g1 in testgraphs(g1c), g2 in testgraphs(g2c)
        d, λ = @inferred(edit_distance(g1, g2))
        @test d == 2.0
        @test λ == Tuple[(1, 1), (2, 2), (3, 3), (4, 4)]
    end
end


text:
const PLOT_TYPES = (
    scatter=Scatter,
    lines=Lines,
    barplot=BarPlot,
    boxplot=BoxPlot,
    violin=Violin,
    heatmap=Heatmap,
    contour=Contour,
    contourf=Contourf,
    wireframe=Wireframe,
    surface=Surface,
)

const ANALYSES = (
    histogram=AlgebraOfGraphics.histogram,
    density=AlgebraOfGraphics.density,
    frequency=AlgebraOfGraphics.frequency,
    expectation=AlgebraOfGraphics.expectation,
    linear=AlgebraOfGraphics.linear,
    smooth=AlgebraOfGraphics.smooth,
)

const STYLES = (:color, :marker, :markersize, :linestyle, :col, :row, :layout, :side, :dodge, :stack)

struct PlotSpecs
    names::Observable{Vector{String}}
    attributes::Observable{String}
    layers::Observable{String}
end

function PlotSpecs(df::Observable)
    return PlotSpecs(lift(colnames, df), Observable(""), Observable(""))
end

is_set(p::PlotSpecs) = !isempty(p.attributes[]) || !isempty(p.layers[])

function to_visual(sym)
    plt = get(PLOT_TYPES, sym, nothing)
    isnothing(plt) || return visual(plt)
    an = get(ANALYSES, sym, nothing)
    isnothing(an) || return an()
    msg = "plot or analysis not available"
    throw(ArgumentError(msg))
end

to_symbolpair((a, b)) = Symbol(a) => Symbol(b)

function to_algebraic(str::AbstractString)
    calls = compute_calls(str)
    layers = map(calls) do call
        visual = mapreduce(to_visual∘Symbol, *, call.fs, init=mapping())
        args = map(Symbol, call.positional)
        attrs = map(to_symbolpair, call.named)
        return visual * mapping(args...; attrs...)
    end
    return Layers(layers)
end

to_algebraic(v::Union{Tuple, AbstractArray}) = mapreduce(to_algebraic, *, v, init=mapping())

to_algebraic(specs::PlotSpecs) = to_algebraic(map(getindex, (specs.attributes, specs.layers)))

function layers_options()
    options = vecmap(((keys(PLOT_TYPES)..., keys(ANALYSES)...))) do key
        string(key) => String[]
    end
    pushfirst!(options, "+" => String[])
    return options
end

function style_options(names)
    options = vecmap(style -> string(style) => names, STYLES)
    pushfirst!(options, "" => names)
    return options
end

function specs_options(session::Session, specs::PlotSpecs; name)
    return map(session, specs.names; result=Observable{AutocompleteOptions}()) do names
        options = name == :layers ? [layers_options(); style_options(names)] : style_options(names)
        return to_autocomplete_options(options)
    end
end

struct Chart <: AbstractVisualization
    table::Observable{SimpleTable}
    plotspecs::PlotSpecs
end

Chart(table::Observable{SimpleTable}) = Chart(table, PlotSpecs(table))
Chart(pipelines::AbstractVector) = Chart(output(last(pipelines)))

to_algebraic(chart::Chart) = data(chart.table[]) * to_algebraic(chart.plotspecs)

defaultplot() = Figure(; backgroundcolor=colorant"#F3F4F6")

function on_pixelratio(f, session; once=false)
    pixelratio = Observable(1.0)
    flag = false
    on(pixelratio) do pr
        once && flag && return
        f(pr)
        flag = true
    end
    evaljs(session, js"$(UtilitiesJS).trackPixelRatio($(pixelratio))")
    return
end

function jsrender(session::Session, chart::Chart)

    specs_widgets = map([:attributes, :layers]) do name
        label = DOM.p(
            class="text-blue-800 text-xl font-semibold p-4 w-full text-left",
            uppercasefirst(string(name))
        )
        options = specs_options(session, chart.plotspecs; name)
        textbox = jsrender(session, Autocomplete(getproperty(chart.plotspecs, name), options))
        return DOM.div(label, DOM.div(class="pl-4", textbox))
    end

    plot = Observable{Figure}(defaultplot())

    width, height = Observable(350), Observable(350)
    on_pixelratio(session) do pr
        width[] = round(Int, 350pr)
        height[] = round(Int, 350pr)
    end

    reset_plot!(_) = plot[] = defaultplot()
    function update_plot!(_)
        is_set(chart.plotspecs) || return
        plt = to_algebraic(chart)
        axis = (width=width[], height=height[])
        fg = draw(plt; axis)
        plot[] = fg.figure
    end

    plot_button = Button("Plot", class=buttonclass(true))
    clear_button = Button("Clear", class=buttonclass(false))
    ui = scrollable_component(
        DOM.div(class="grid grid-cols-2 gap-8", specs_widgets),
        DOM.div(class="mt-8 pl-4", plot_button, clear_button),
        DOM.div(class="mt-12 pl-4", plot)
    )

    tryon(update_plot!, session, plot_button.value)
    tryon(update_plot!, session, chart.plotspecs.names) # gets updated when table changes
    tryon(reset_plot!, session, clear_button.value)

    return jsrender(session, ui)
end


text:
# compiler driver and main interface

# NOTE: the keyword arguments to compile/codegen control those aspects of compilation that
#       might have to be changed (e.g. set libraries=false when recursing, or set
#       strip=true for reflection). What remains defines the compilation job itself,
#       and those values are contained in the CompilerJob struct.

# (::CompilerJob)
const compile_hook = Ref{Union{Nothing,Function}}(nothing)

"""
    compile(target::Symbol, job::CompilerJob;
            libraries=true, deferred_codegen=true,
            optimize=true, strip=false, ...)

Compile a function `f` invoked with types `tt` for device capability `cap` to one of the
following formats as specified by the `target` argument: `:julia` for Julia IR, `:llvm` for
LLVM IR and `:asm` for machine code.

The following keyword arguments are supported:
- `libraries`: link the GPU runtime and `libdevice` libraries (if required)
- `deferred_codegen`: resolve deferred compiler invocations (if required)
- `optimize`: optimize the code (default: true)
- `strip`: strip non-functional metadata and debug information (default: false)
- `validate`: validate the generated IR before emitting machine code (default: true)
- `only_entry`: only keep the entry function, remove all others (default: false).
  This option is only for internal use, to implement reflection's `dump_module`.

Other keyword arguments can be found in the documentation of [`cufunction`](@ref).
"""
function compile(target::Symbol, @nospecialize(job::CompilerJob);
                 libraries::Bool=true, deferred_codegen::Bool=true,
                 optimize::Bool=true, strip::Bool=false, validate::Bool=true,
                 only_entry::Bool=false)
    if compile_hook[] !== nothing
        compile_hook[](job)
    end

    return codegen(target, job;
                   libraries, deferred_codegen, optimize, strip, validate, only_entry)
end

function codegen(output::Symbol, @nospecialize(job::CompilerJob);
                 libraries::Bool=true, deferred_codegen::Bool=true, optimize::Bool=true,
                 strip::Bool=false, validate::Bool=true, only_entry::Bool=false)
    ## Julia IR

    method_instance, world = emit_julia(job)

    output == :julia && return method_instance

    ## LLVM IR

    ir, kernel = emit_llvm(job, method_instance, world;
                           libraries, deferred_codegen, optimize, only_entry)

    if output == :llvm
        if strip
            @timeit_debug to "strip debug info" strip_debuginfo!(ir)
        end

        return ir, kernel
    end

    ## machine code

    format = if output == :asm
        LLVM.API.LLVMAssemblyFile
    elseif output == :obj
        LLVM.API.LLVMObjectFile
    else
        error("Unknown assembly format $output")
    end
    code = emit_asm(job, ir, kernel; strip, validate, format)

    undefined_fns = LLVM.name.(decls(ir))
    undefined_gbls = map(x->(name=LLVM.name(x),type=llvmtype(x),external=isextinit(x)), LLVM.globals(ir))

    (output == :asm || output == :obj) && return code, LLVM.name(kernel), undefined_fns, undefined_gbls

    error("Unknown compilation output $output")
end

@locked function emit_julia(@nospecialize(job::CompilerJob))
    @timeit_debug to "validation" check_method(job)

    @timeit_debug to "Julia front-end" begin

        # get the method instance
        world = Base.get_world_counter()
        meth = which(job.source.f, job.source.tt)
        sig = Base.signature_type(job.source.f, job.source.tt)::Type
        (ti, env) = ccall(:jl_type_intersection_with_env, Any,
                          (Any, Any), sig, meth.sig)::Core.SimpleVector
        meth = Base.func_for_method_checked(meth, ti, env)
        method_instance = ccall(:jl_specializations_get_linfo, Ref{Core.MethodInstance},
                      (Any, Any, Any, UInt), meth, ti, env, world)

        for var in env
            if var isa TypeVar
                throw(KernelError(job, "method captures a typevar (you probably use an unbound type variable)"))
            end
        end
    end

    return method_instance, world
end

# primitive mechanism for deferred compilation, for implementing CUDA dynamic parallelism.
# this could both be generalized (e.g. supporting actual function calls, instead of
# returning a function pointer), and be integrated with the nonrecursive codegen.
const deferred_codegen_jobs = Vector{Tuple{Core.Function,Type}}()
@generated function deferred_codegen(::Val{f}, ::Val{tt}) where {f,tt}
    push!(deferred_codegen_jobs, (f,tt))
    id = length(deferred_codegen_jobs)

    quote
        # TODO: add an edge to this method instance to support method redefinitions
        ccall("extern deferred_codegen", llvmcall, Ptr{Cvoid}, (Int,), $id)
    end
end

@locked function emit_llvm(@nospecialize(job::CompilerJob), @nospecialize(method_instance), world;
                           libraries::Bool=true, deferred_codegen::Bool=true, optimize::Bool=true,
                           only_entry::Bool=false)
    @timeit_debug to "IR generation" begin
        ir, kernel = irgen(job, method_instance, world)
        ctx = context(ir)
        kernel_fn = LLVM.name(kernel)
    end

    # always preload the runtime, and do so early; it cannot be part of any timing block
    # because it recurses into the compiler
    if libraries
        runtime = load_runtime(job, ctx)
        if haskey(globals(runtime), "llvm.used")
            # the runtime shouldn't link-in stuff that gets preserved in the output. this is
            # a hack to get rid of the device function slots emitted by the PTX back-end,
            # but it also makes sense.
            gv = globals(runtime)["llvm.used"]
            LLVM.unsafe_delete!(runtime, gv)
        end
        runtime_fns = LLVM.name.(defs(runtime))
    end

    @timeit_debug to "LLVM middle-end" begin
        # target-specific libraries
        if libraries
            undefined_fns = LLVM.name.(decls(ir))
            @timeit_debug to "target libraries" link_libraries!(job, ir, undefined_fns)
        end

        if optimize
            @timeit_debug to "optimization" optimize!(job, ir)

            # optimization may have replaced functions, so look the entry point up again
            kernel = functions(ir)[kernel_fn]
        end

        if libraries
            undefined_fns = LLVM.name.(decls(ir))
            if any(fn -> fn in runtime_fns, undefined_fns)
                @timeit_debug to "runtime library" link_library!(ir, runtime)
            end
        end

        if ccall(:jl_is_debugbuild, Cint, ()) == 1
            @timeit_debug to "verification" verify(ir)
        end

        if only_entry
            # replace non-entry function definitions with a declaration
            for f in functions(ir)
                f == kernel && continue
                isdeclaration(f) && continue
                LLVM.isintrinsic(f) && continue
                # FIXME: expose llvm::Function::deleteBody with a C API
                fn = LLVM.name(f)
                LLVM.name!(f, "")
                f′ = LLVM.Function(ir, fn, eltype(llvmtype(f)))
                # copying attributes is broken due to maleadt/LLVM.jl#186,
                # but that doesn't matter because `only_entry` is only used for reflection,
                # and the emitted code has already been optimized at this point.
                replace_uses!(f, f′)
            end
        end

        # remove everything except for the kernel and any exported global variables
        @timeit_debug to "clean-up" begin
            exports = String[kernel_fn]
            for gvar in globals(ir)
                push!(exports, LLVM.name(gvar))
            end

            ModulePassManager() do pm
                internalize!(pm, exports)

                # eliminate all unused internal functions
                global_optimizer!(pm)
                global_dce!(pm)
                strip_dead_prototypes!(pm)

                # merge constants (such as exception messages) from the runtime
                constant_merge!(pm)

                run!(pm, ir)
            end
        end
    end

    # deferred code generation
    if !only_entry && deferred_codegen && haskey(functions(ir), "deferred_codegen")
        dyn_marker = functions(ir)["deferred_codegen"]

        cache = Dict{CompilerJob, String}(job => kernel_fn)

        # iterative compilation (non-recursive)
        changed = true
        while changed
            changed = false

            # find deferred compiler
            # TODO: recover this information earlier, from the Julia IR
            worklist = MultiDict{CompilerJob, LLVM.CallInst}()
            for use in uses(dyn_marker)
                # decode the call
                call = user(use)::LLVM.CallInst
                id = convert(Int, first(operands(call)))

                global deferred_codegen_jobs
                dyn_f, dyn_tt = deferred_codegen_jobs[id]
                dyn_job = similar(job, FunctionSpec(dyn_f, dyn_tt, #=kernel=# true))
                push!(worklist, dyn_job => call)
            end

            # compile and link
            for dyn_job in keys(worklist)
                # cached compilation
                dyn_kernel_fn = get!(cache, dyn_job) do
                    dyn_ir, dyn_kernel = codegen(:llvm, dyn_job; optimize,
                                                 deferred_codegen=false)
                    dyn_kernel_fn = LLVM.name(dyn_kernel)
                    @assert context(dyn_ir) == ctx
                    link!(ir, dyn_ir)
                    changed = true
                    dyn_kernel_fn
                end
                dyn_kernel = functions(ir)[dyn_kernel_fn]

                # insert a pointer to the function everywhere the kernel is used
                T_ptr = convert(LLVMType, Ptr{Cvoid}, ctx)
                for call in worklist[dyn_job]
                    Builder(ctx) do builder
                        position!(builder, call)
                        fptr = ptrtoint!(builder, dyn_kernel, T_ptr)
                        replace_uses!(call, fptr)
                    end
                    unsafe_delete!(LLVM.parent(call), call)
                end
            end
        end

        # merge constants (such as exception messages) from each kernel
        ModulePassManager() do pm
            constant_merge!(pm)

            run!(pm, ir)
        end

        # all deferred compilations should have been resolved
        @compiler_assert isempty(uses(dyn_marker)) job
        unsafe_delete!(ir, dyn_marker)
    end

    return ir, kernel
end

@locked function emit_asm(@nospecialize(job::CompilerJob), ir::LLVM.Module, kernel::LLVM.Function;
                          strip::Bool=false, validate::Bool=true, format::LLVM.API.LLVMCodeGenFileType)
    finish_module!(job, ir)

    if validate
        @timeit_debug to "validation" begin
            check_invocation(job, kernel)
            check_ir(job, ir)
        end
    end

    # NOTE: strip after validation to get better errors
    if strip
        @timeit_debug to "strip debug info" strip_debuginfo!(ir)
    end

    @timeit_debug to "LLVM back-end" begin
        @timeit_debug to "preparation" prepare_execution!(job, ir)

        code = @timeit_debug to "machine-code generation" mcgen(job, ir, kernel, format)
    end

    return code
end


text:
# This file is a part of ValueShapes.jl, licensed under the MIT License (MIT).

using ValueShapes
using Test
using ArraysOfArrays

@testset "const_value_shape" begin

    @inferred(size(ConstValueShape(42))) == ()
    @inferred(eltype(ConstValueShape(42))) == Int
    @inferred(totalndof(ConstValueShape(42))) == 0

    @inferred(size(ConstValueShape(rand(2,3)))) == (2,3)
    @inferred(ValueShapes.default_unshaped_eltype(ConstValueShape(rand(Float32,2,3)))) == Float32
    @inferred(totalndof(ConstValueShape(rand(2,3)))) == 0

    @test @inferred(ConstValueShape([1 4; 3 2])(undef)) == [1 4; 3 2]
    @test @inferred(ConstValueShape([1 4; 3 2])(Int[])) == [1 4; 3 2]

    data = [1 4; 3 2]
    shape = ConstValueShape([1 4; 3 2])

    @test typeof(@inferred(Vector{Int32}(undef, shape))) == Vector{Int32}
    @test size(@inferred(Vector{Int32}(undef, shape))) == (0,)

    @test @inferred(length(shape)) == 4

    @test @inferred(ValueShapes.shaped_type(shape, Real)) == typeof(data)

    @test @inferred (ConstValueShape(4.0) <= ConstValueShape(4.0)) == true
    @test @inferred (ConstValueShape(4.0) >= ConstValueShape{AbstractFloat}(4.0)) == false
    @test @inferred (ConstValueShape{AbstractFloat}(4.0) >= ConstValueShape(4.0)) == true
    @test @inferred (ConstValueShape(4) <= ConstValueShape(5)) == false

    @test @inferred(unshaped(4.2, ConstValueShape(4.2))) == Float32[]
    @test_throws ArgumentError unshaped(4.3, ConstValueShape(4.2))

    vecs_of_vecs = VectorOfSimilarVectors(reshape(collect(1:22), 11, 2))
    va = ValueAccessor(ArrayShape{Real}(11,1), 0)
    bcv = ValueShapes._bcasted_view(vecs_of_vecs, va)
    for (index,value) in enumerate(bcv[1])
        @test value == vecs_of_vecs[1][index]
    end
    aosv = ArrayOfSimilarVectors([ [1,2,3] [4,5,6] ])
    va = ValueAccessor(ConstValueShape{Real}(1), 0)
    inds = getindex.(aosv, Ref(va))
    for i in inds
        @test i == 1
    end
    views = view.(aosv, Ref(va))
    @test views === inds
end


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

using Base.Test
using TimerOutputs
const to = TimerOutput()

test_files = String[]
push!(test_files, "test_assembly.jl")
push!(test_files, "test_elements.jl")
push!(test_files, "test_fields.jl")
push!(test_files, "test_integrate.jl")
push!(test_files, "test_problems.jl")
push!(test_files, "test_sparse.jl")
push!(test_files, "test_solvers.jl")
push!(test_files, "test_analysis.jl")
push!(test_files, "test_test.jl")
push!(test_files, "test_types.jl")

@testset "FEMBase.jl" begin
    for fn in test_files
        timeit(to, fn) do
            include(fn)
        end
    end
end
println()
println("Test statistics:")
println(to)


text:
"""
    start_ata(;
        settings::InputSettings = InputSettings(),
        bank::DataFrames.DataFrame = DataFrames.DataFrame(),
        settings_file = "settings_ata.jl",
        bank_file = "bank.csv",
        bank_delim = ";",
    )

# Description

Initialize an empty ATA model and load the test assembly settings contained in the `settings_file`
using item data in file `bank_file` which values are separated by `bank_delim`.
Alternatively, settings object and bank dataframe can be passed using the arguments `settings` and `bank`.

# Arguments

- **`settings::InputSettings`** : Optional. Default:  `InputSettings()` object with custom ATA settings.
- **`bank::DataFrames.DataFrame`** : Optional. Default: `DataFrame()`. A dataframe containing data about the items.
- **`settings_file`** : Optional. Default: "settings_ata.jl". The path of the file containing the ATA settings in the form of an `InputSettings` struct.
- **`bank_file`** : Optional. Default: "bank.csv". The path of the file containing the item pool/bank in the form of custom-separated values.
- **`bank_delim`** : Optional. Default: ";". The custom-separator for the bank_file.

# Output

- An ATA model.

"""
function start_ata(;
    settings::InputSettings = InputSettings(),
    bank::DataFrames.DataFrame = DataFrames.DataFrame(),
    settings_file = "settings_ata.jl",
    bank_file = "bank.csv",
    bank_delim = ";",
)

    ata_model = NoObjModel()

    try
        #LOAD INPUT SETTINGS, if error return ata_model
        Inputs = InputSettings()
        if size(settings.T, 1) > 0
            Inputs = settings
            success!(ata_model, "InputSettings object loaded.")
        elseif !isfile(settings_file)
            error!(
                ata_model,
                string(
                    "Settings file with name ",
                    settings_file,
                    " does not exist. Provide a valid InputSettings object or a name of an existing file.",
                ),
            )
        else
            try
                Inputs = read_settings_file(settings_file)
            catch e
                error!(
                    ata_model,
                    string(
                        "Error in reading the input settings file:\n  ",
                        sprint(showerror, e),
                        ".",
                    ),
                )
            end
            success!(ata_model, "Settings file loaded.")
        end

        #SELECT MODEL TYPE
        infos = ata_model.output.infos
        if Inputs.obj_type != ""
            if Inputs.obj_type == "maximin"
                ata_model = MaximinModel()
            elseif Inputs.obj_type == "minimax"
                ata_model = MinimaxModel()
            elseif Inputs.obj_type == "cc_maximin"
                ata_model = CCMaximinModel()
            elseif Inputs.obj_type == "soyster_maximin"
                ata_model = SoysterMaximinModel()
            elseif Inputs.obj_type == "de_jong_maximin"
                ata_model = DeJongMaximinModel()
            elseif Inputs.obj_type == "robust_maximin"
                ata_model = RobustMaximinModel()
            elseif Inputs.obj_type == "custom"
                ata_model = CustomModel()
            else
                error!(ata_model,
                    "Only \"maximin\", \"minimax\", \"cc_maximin\", \"soyster_maximin\", \"de_jong_maximin\", \"custom\" and \"\" objective types are supported."
                )
            end
            success!(ata_model, string(ata_model.obj.name, " model initialized."))
        end
        ata_model.output.infos = infos

        #START ASSIGNING SETTINGS
        if Inputs.n_groups > 0
            ata_model.settings.n_groups = Inputs.n_groups
        else
            success!(ata_model, "n_groups has been set to 1.")
        end
        if Inputs.n_items > 0
            ata_model.settings.n_items = Inputs.n_items
        else
            error!(ata_model, "n_items must be grater than 0.")
        end
        if Int(sum(Inputs.T)) > 0
            ata_model.settings.T = Int(sum(Inputs.T))
            ata_model.settings.Tg = Inputs.T
        else
            error!(ata_model, "Sum of T must be grater than 0.")
        end

        #INITIALIZE CONSTRAINTS
        ata_model.constraints = [Constraint() for t = 1:ata_model.settings.T]
        ata_model.settings.n_groups = size(Inputs.T, 1)
        x_forced0 = Vector{Vector{Bool}}(undef, ata_model.settings.T)
        for t = 1:ata_model.settings.T
            x_forced0[t] = fill(true, ata_model.settings.n_items)
            ata_model.constraints[t].constr_A =
                zeros(Float64, 0, ata_model.settings.n_items)
        end

        if !isdir("opt")
            mkdir("opt")
        end

        #ADD BANK
        add_bank!(ata_model, bank = bank, bank_file = bank_file, bank_delim = bank_delim)
        #open("opt/Settings.jl", "w") do f
        #write(f, "#Settings \n")

        #IRT
        if (Inputs.irt_model != "")
            ata_model.settings.irt.model = Inputs.irt_model
            ata_model.settings.irt.parametrization = Inputs.irt_parametrization
            ata_model.settings.irt.D = Inputs.irt_D
            if (size(ata_model.settings.bank, 1) == ata_model.settings.n_items)
                _add_irt!(ata_model)
            end
        end

        #TEST LENGTH
        if Inputs.length_min != Int64[]
            lengthmin = zeros(Int64, ata_model.settings.T)
            lengthweight = ones(Int64, ata_model.settings.T)
            t1 = 1
            for g = 1:ata_model.settings.n_groups, t = 1:ata_model.settings.Tg[g]
                lengthmin[t1] = Int(Inputs.length_min[g])
                lengthweight[t1] = Inputs.length_weight[g]
                t1 += 1
            end
            for t = 1:ata_model.settings.T
                ata_model.constraints[t].constr_A = vcat(
                    ata_model.constraints[t].constr_A,
                    (-lengthweight[t]) .* ones(Float64, ata_model.settings.n_items)',
                )
                ata_model.constraints[t].constr_b =
                    vcat(ata_model.constraints[t].constr_b, -lengthmin[t] * lengthweight[t])
                ata_model.constraints[t].length_min = lengthmin[t]
            end
            #write(f, "length_min = $lengthmin")
            success!(ata_model, "Minimum length of tests constrained.")
        end

        if Inputs.length_max != Int64[]
            lengthmax = zeros(Int64, ata_model.settings.T)
            lengthweight = ones(Int64, ata_model.settings.T)
            t1 = 1
            for g = 1:ata_model.settings.n_groups, t = 1:ata_model.settings.Tg[g]
                lengthmax[t1] = Int(Inputs.length_max[g])
                lengthweight[t1] = Inputs.length_weight[g]
                t1 += 1
            end
            for t = 1:ata_model.settings.T
                ata_model.constraints[t].constr_A = vcat(
                    ata_model.constraints[t].constr_A,
                    (lengthweight[t]) .* ones(ata_model.settings.n_items)',
                )
                ata_model.constraints[t].constr_b =
                    vcat(ata_model.constraints[t].constr_b, lengthmax[t] * lengthweight[t])
                ata_model.constraints[t].length_max = lengthmax[t]
            end
            #write(f, "length_max = $lengthmax")
            success!(ata_model, "Maximum length of tests constrained.")
        end

        #FICTIOUS FRIEND SETS
        ata_model.settings.n_fs = ata_model.settings.n_items
        ata_model.settings.fs.counts = ones(ata_model.settings.n_items)
        ata_model.settings.fs.sets = string.(collect(1:ata_model.settings.n_items))
        ata_model.settings.fs.items = [[i] for i = 1:ata_model.settings.n_items]
        ata_model.settings.forced0 = x_forced0

        if Inputs.friend_sets_var != String[]
            ata_model.settings.fs.var = Symbol.(Inputs.friend_sets_var)
            if size(ata_model.settings.bank, 1) == ata_model.settings.n_items
                val = Symbol.(ata_model.settings.fs.var)
                _add_friends!(ata_model)
            else
                error!(
                    ata_model,
                    string(
                        "Variable for Friend sets NOT loaded because the provided item bank has length not equal to ",
                        ata_model.settings.n_items,
                        ".",
                    ),
                )
            end
        end
        #enemy sets
        if Inputs.enemy_sets_var != String[]
            ata_model.settings.es.var = Symbol.(Inputs.enemy_sets_var)
            if size(ata_model.settings.bank, 1) == ata_model.settings.n_items
                _add_enemies!(ata_model)
            else
                error!(
                    ata_model,
                    string(
                        "Variable for Enemy sets NOT loaded because the provided item bank has length not equal to ",
                        ata_model.settings.n_items,
                        ".",
                    ),
                )
            end
        end

        #EXPECTED SCORE
        if Inputs.expected_score_var != String[]
            t1 = 1
            for g = 1:ata_model.settings.n_groups, t = 1:ata_model.settings.Tg[g]
                ata_model.constraints[t1].expected_score.var =
                    Symbol.(Inputs.expected_score_var[g])
                t1 += 1
            end
            success!(ata_model, "Expected score variable loaded.")
        end
        if Inputs.expected_score_pts != Float64[]
            t1 = 1
            for g = 1:ata_model.settings.n_groups, t = 1:ata_model.settings.Tg[g]
                ata_model.constraints[t1].expected_score.pts = Inputs.expected_score_pts[g]
                t1 += 1
            end
        else
            t1 = 1
            for g = 1:ata_model.settings.n_groups, t = 1:ata_model.settings.Tg[g]
                ata_model.constraints[t1].expected_score.pts = [0.0]
                t1 += 1
            end
            success!(ata_model, "Expected score points loaded.")
        end
        min_exp_score = false
        max_exp_score = false
        if size(Inputs.expected_score_min, 1) > 0
            if any(vcat(Inputs.expected_score_min...) .> 0)
                min_exp_score = true
                t1 = 1
                for g = 1:ata_model.settings.n_groups, t = 1:ata_model.settings.Tg[g]
                    ata_model.constraints[t1].expected_score.min =
                        Inputs.expected_score_min[g]
                    t1 += 1
                end
                success!(ata_model, "Lower bounds for expected score loaded.")
            end
        end
        if size(Inputs.expected_score_max, 1) > 0
            if any(vcat(Inputs.expected_score_max...) .< 1.00)
                max_exp_score = true
                t1 = 1
                for g = 1:ata_model.settings.n_groups, t = 1:ata_model.settings.Tg[g]
                    ata_model.constraints[t1].expected_score.max =
                        Inputs.expected_score_max[g]
                    t1 += 1
                end
                success!(ata_model, "Upper bounds for expected score loaded.")
            end
        end
        if (max_exp_score || min_exp_score)
            if size(ata_model.settings.bank, 1) == ata_model.settings.n_items
                _add_exp_score!(ata_model)
            else
                error!(
                    ata_model,
                    string(
                        "Expected scores (ICFs) based on IRT parameters NOT computed because the provided item bank has length not equal to ",
                        ata_model.settings.n_items,
                        ".",
                    ),
                )
            end
        end

        #ITEM USE
        if size(Inputs.item_use_min, 1) > 0
            if any(Inputs.item_use_min .> 0)
                ata_model.settings.iu.min = Inputs.item_use_min
                ata_model.settings.to_apply[2] = true
                success!(ata_model, "Minimum item use constrained.")
            end
        end
        if size(Inputs.item_use_max, 1) > 0
            if any(Inputs.item_use_max .< ata_model.settings.T)
                ata_model.settings.iu.max = Inputs.item_use_max
                for v = 1:ata_model.settings.T
                    x_forced0[v][findall(ata_model.settings.iu.max .< 1)] .= false
                end
                ata_model.settings.to_apply[1] = true
                success!(ata_model, "Maximum item use constrained.")
            end
        end

        #OBJECTIVE FUNCTION
        if Inputs.obj_type in ["maximin"]
            ata_model.obj.cores =
                [MaximinObjectiveCore() for t = 1:sum(ata_model.settings.Tg)]
            if size(Inputs.obj_points, 1) > 0
                t1 = 1
                for g = 1:ata_model.settings.n_groups, t = 1:ata_model.settings.Tg[g]
                    ata_model.obj.cores[t1].points = Inputs.obj_points[g]
                    t1 += 1
                end
                success!(ata_model, "Optimization points loaded.")
            else
                error!(
                    ata_model,
                    "maximin, cc_maximin, soyster_maximin, de_jong_maximin, robust_maximin, and minimax objective functions require optimization points. Use obj_points field in input settings.",
                )
                return nothing
            end
        elseif Inputs.obj_type == "cc_maximin"
            ata_model.obj.cores =
                [CCMaximinObjectiveCore() for t = 1:sum(ata_model.settings.Tg)]
            if size(Inputs.obj_points, 1) > 0
                t1 = 1
                for g = 1:ata_model.settings.n_groups, t = 1:ata_model.settings.Tg[g]
                    ata_model.obj.cores[t1].points = Inputs.obj_points[g]
                    ata_model.obj.cores[t1].alpha = Inputs.obj_aux_float
                    t1 += 1
                end
                ata_model.obj.R = Inputs.obj_aux_int
                success!(
                    ata_model,
                    string(
                        "Optimization points, number of samples (",
                        ata_model.obj.R,
                        "), and alpha (",
                        Inputs.obj_aux_float,
                        ") loaded.",
                    ),
                )
            else
                error!(
                    ata_model,
                    "maximin, cc_maximin, soyster_maximin, de_jong_maximin, robust_maximin, and minimax objective functions require optimization points. Use obj_points field in input settings.",
                )
                return nothing
            end
        elseif Inputs.obj_type in ["soyster_maximin", "de_jong_maximin"]
            ata_model.obj.cores =
                [MaximinObjectiveCore() for t = 1:sum(ata_model.settings.Tg)]
            if size(Inputs.obj_points, 1) > 0
                t1 = 1
                for g = 1:ata_model.settings.n_groups, t = 1:ata_model.settings.Tg[g]
                    ata_model.obj.cores[t1].points = Inputs.obj_points[g]
                    t1 += 1
                end
                ata_model.obj.R = Inputs.obj_aux_int
                success!(
                    ata_model,
                    "Optimization points and number of replications loaded.",
                )
            else
                error!(
                    ata_model,
                    "maximin, cc_maximin, soyster_maximin, de_jong_maximin, robust_maximin, and minimax objective functions require optimization points. Use obj_points field in input settings.",
                )
                return nothing
            end
        elseif Inputs.obj_type in ["robust_maximin"]
            ata_model.obj.cores =
                [RobustMaximinObjectiveCore() for t = 1:sum(ata_model.settings.Tg)]
            if size(Inputs.obj_points, 1) > 0
                t1 = 1
                for g = 1:ata_model.settings.n_groups, t = 1:ata_model.settings.Tg[g]
                    ata_model.obj.cores[t1].points = Inputs.obj_points[g]
                    t1 += 1
                end
                ata_model.obj.R = Inputs.obj_aux_int
                ata_model.obj.Gamma = Inputs.obj_aux_float
                success!(
                    ata_model,
                    string(
                        "Optimization points, number of samples (",
                        ata_model.obj.R,
                        "), and Gamma (",
                        ata_model.obj.Gamma,
                        ") loaded.",
                    ),
                )
            else
                error!(
                    ata_model,
                    "maximin, cc_maximin, soyster_maximin, de_jong_maximin, robust_maximin, and minimax objective functions require optimization points. Use obj_points field in input settings.",
                )
                return nothing
            end
        elseif Inputs.obj_type == "minimax"
            ata_model.obj.cores =
                [MinimaxObjectiveCore() for t = 1:sum(ata_model.settings.Tg)]
            if size(Inputs.obj_points, 1) > 0
                t1 = 1
                for g = 1:ata_model.settings.n_groups, t = 1:ata_model.settings.Tg[g]
                    ata_model.obj.cores[t1].points = Inputs.obj_points[g]
                    t1 += 1
                end
                success!(ata_model, "Optimization points loaded.")
            else
                error!(
                    ata_model,
                    "maximin, cc_maximin, soyster_maximin, de_jong_maximin, robust_maximin, and minimax objective functions require optimization points. Use obj_points field in input settings.",
                )
                return nothing
            end
            if (size(Inputs.obj_targets, 1) > 0)
                t1 = 1
                for g = 1:ata_model.settings.n_groups
                    if size(Inputs.obj_targets[g], 1) != size(Inputs.obj_points[g], 1)
                        error!(
                            ata_model,
                            string(
                                "For group ",
                                g,
                                " size of obj_targets (",
                                size(Inputs.obj_targets[g], 1),
                                ") is not the same as size of obj_points (",
                                size(Inputs.obj_points[g], 1),
                                ").",
                            ),
                        )
                        return nothing
                    end
                    for t = 1:ata_model.settings.Tg[g]
                        ata_model.obj.cores[t1].targets = Inputs.obj_targets[g]
                        t1 += 1
                    end
                end
                success!(ata_model, "Targets loaded.")
            else
                error!(
                    ata_model,
                    "minimax objective function requires targets. Use obj_targets field in input settings.",
                )
                return nothing
            end
        end

        if Inputs.categories != String[]
            val = Symbol.(Inputs.categories)
            ata_model.output.categories = copy(val)
            #write(f, "categories = $val\n")
            success!(ata_model, "Categories for output loaded.")
        end

        #overlap matrix
        ata_model.settings.ol_max =
            zeros(Float64, ata_model.settings.T, ata_model.settings.T)
        #end
        success!(
            ata_model,
            string(
                "ASSEMBLE ",
                ata_model.settings.T,
                " FORMS DIVIDED INTO ",
                ata_model.settings.n_groups,
                " GROUPS.",
            ),
        )
        #update model
        JLD2.@save "opt/ata_model.jld2" ata_model
    catch e
        error!(ata_model, string(sprint(showerror, e)))
    end
    return ata_model
end


text:
"""
    ArcEager()

Arc-Eager transition system for dependency parsing.

# Transitions

| Transition  | Definition                                    |
|:----------- |:--------------------------------------------- |
| LeftArc(l)  | (σ\\|s, b\\|β, A) → (σ, b\\|β, A ∪ (b, l, s)) |
| RightArc(l) | (σ\\|s, b\\|β, A) → (σ, b\\|β, A ∪ (b, l, s)) |
| Reduce      | (σ\\|s, β,  A) → (σ, β,   A)                  |
| Shift	      | (σ,  b\\|β, A) → (σ\\|b, β, A)                |

# Preconditions

| Transition  | Condition                        |
|:----------- |:-------------------------------- |
| LeftArc(l)  | ¬[s = 0], ¬∃k∃l'[(k, l', i) ϵ A] |
| RightArc(l) | ¬∃k∃l'[(k, l', j) ϵ A]           |
| Reduce      | ∃k∃l[(k, l, i) ϵ A]              |

# References

[Nivre 2003](http://stp.lingfil.uu.se/~nivre/docs/iwpt03.pdf), [Nivre 2008](https://www.aclweb.org/anthology/J08-4003.pdf).
"""
struct ArcEager <: AbstractTransitionSystem end

initconfig(::ArcEager, graph::DependencyTree) =
    ArcEagerConfig(graph)

transition_space(::ArcEager, labels=[]) =
    isempty(labels) ? [LeftArc(), RightArc(), Reduce(), Shift()] :
    [LeftArc.(labels)..., RightArc.(labels)..., Reduce(), Shift()]

projective_only(::ArcEager) = true

struct ArcEagerConfig <: AbstractParserConfiguration
    stack::Vector{Int}
    buffer::Vector{Int}
    A::Vector{Token}
end

ArcEagerConfig(sentence) = stack_buffer_config(ArcEagerConfig, sentence)

buffer(cfg::ArcEagerConfig) = cfg.buffer
stack(cfg::ArcEagerConfig)  = cfg.stack
tokens(cfg::ArcEagerConfig) = cfg.A

function apply_transition(f, cfg::ArcEagerConfig, a...; k...)
    σ, β, A = f(cfg.stack, cfg.buffer, cfg.A, a...; k...)
    return ArcEagerConfig(σ, β, A)
end

leftarc(cfg::ArcEagerConfig, args...; kwargs...) =
    apply_transition(leftarc_reduce, cfg, args...; kwargs...)

rightarc(cfg::ArcEagerConfig, args...; kwargs...) =
    apply_transition(rightarc_shift, cfg, args...; kwargs...)

reduce(cfg::ArcEagerConfig) = apply_transition(reduce, cfg)

shift(cfg::ArcEagerConfig) = apply_transition(shift, cfg)

isfinal(cfg::ArcEagerConfig) = all(has_head, cfg.A)

has_head(cfg::ArcEagerConfig, k) = has_head(token(cfg, k))

"""
    static_oracle(cfg::ArcEagerConfig, gold, arc=untyped)

Default static oracle function for arc-eager dependency parsing.

See [Goldberg & Nivre 2012](https://www.aclweb.org/anthology/C12-1059.pdf).
(Also called Arc-Eager-Reduce in [Qi & Manning 2017](https://nlp.stanford.edu/pubs/qi2017arcswift.pdf)).
"""
function static_oracle(cfg::ArcEagerConfig, gold, arc=untyped)
    if stacklength(cfg) >= 1
        (σ, s) = popstack(cfg)
        if bufferlength(cfg) >= 1
            (b, β) = shiftbuffer(cfg)
            has_arc(gold, b, s) && return LeftArc(arc(gold[s])...)
            has_arc(gold, s, b) && return RightArc(arc(gold[b])...)
        end
        if all(k -> k > 0 && has_head(cfg, k), [s ; deps(gold, s)])
            return Reduce()
        end
    end
    return Shift()
end

"""
    static_oracle_prefer_shift(cfg::ArcEagerConfig, tree, arc=untyped)

Static oracle for arc-eager dependency parsing. Similar to the
"regular" static oracle, but always Shift when ambiguity is present.

See [Qi & Manning 2017](https://nlp.stanford.edu/pubs/qi2017arcswift.pdf).
"""
function static_oracle_prefer_shift(cfg::ArcEagerConfig, tree, arc=untyped)
    l = i -> arc(token(tree, i))
    gold_arc = (a, b) -> has_arc(tree, a, b)
    (σ, s), (b, β) = popstack(cfg), shiftbuffer(cfg)
    gold_arc(b, s) && return LeftArc(l(s)...)
    gold_arc(s, b) && return RightArc(l(b)...)
    must_reduce = false
    for k in stack(cfg)
        if gold_arc(k, b) || gold_arc(b, k)
            must_reduce = true
            break
        elseif has_head(token(cfg, k), -1)
            break
        end
    end
    has_right_children = any(k -> s in rightdeps(tree, k), buffer(cfg))
    if !must_reduce || s > 0 && !has_head(cfg, s) || has_right_children
        return Shift()
    else
        return Reduce()
    end
end

"""
    dynamic_oracle(cfg::ArgEagerConfig, tree, arc=untyped)

Dynamic oracle function for arc-eager parsing.

For details, see [Goldberg & Nivre 2012](https://aclweb.org/anthology/C12-1059).
"""
dynamic_oracle(cfg::ArcEagerConfig, tree, arc=untyped) =
    filter(t -> cost(t, cfg, tree) == 0, possible_transitions(cfg, tree, arc))

# see figure 2 in goldberg & nivre 2012 "a dynamic oracle..."
possible_transitions(cfg::ArcEagerConfig, graph::DependencyTree, arc=untyped) =
    possible_transitions(cfg, arc)

function possible_transitions(cfg::ArcEagerConfig, arc=untyped)
    ts = TransitionOperator[]
    if is_possible(LeftArc(), cfg)
        s = last(stack(cfg))
        push!(ts, LeftArc(arc(token(cfg, s))...))
    end
    if is_possible(RightArc(), cfg)
        b = first(buffer(cfg))
        push!(ts, RightArc(arc(token(cfg, b))...))
    end
    is_possible(Reduce(), cfg) && push!(ts, Reduce())
    is_possible(Shift(), cfg) && push!(ts, Shift())
    return ts
end

function cost(t::LeftArc, cfg::ArcEagerConfig, gold)
    # left arc cost: num of arcs (k,l',s), (s,l',k) s.t. k ϵ β
    σ, s = popstack(cfg)
    b, β = shiftbuffer(cfg)
    if has_arc(gold, b, s)
        0
    else
        count(k -> has_arc(gold, k, s) || has_arc(gold, s, k), β)
    end
end

function cost(t::RightArc, cfg::ArcEagerConfig, gold)
    # right arc cost: num of gold arcs (k,l',b), s.t. k ϵ σ or k ϵ β,
    #                 plus num of gold arcs (b,l',k) s.t. k ϵ σ
    σ, s = popstack(cfg)
    b, β = shiftbuffer(cfg)
    if has_arc(gold, s, b)
        0
    else
        count(k -> has_arc(gold, k, b), [σ ; β]) + count(k -> has_arc(gold, b, k), σ)
    end
end

function cost(t::Reduce, cfg::ArcEagerConfig, gold)
    # num of gold arcs (s,l',k) s.t. k ϵ b|β
    σ, s = popstack(cfg)
    count(k -> has_arc(gold, s, k), buffer(cfg))
end

function cost(t::Shift, cfg::ArcEagerConfig, gold)
    # num of gold arcs (k,l',b), (b,l',k) s.t. k ϵ s|σ
    b, β = shiftbuffer(cfg)
    count(k -> has_arc(gold, k, b) || has_arc(gold, b, k), stack(cfg))
end

function is_possible(::LeftArc, cfg::ArcEagerConfig)
    if length(cfg.stack) > 0 && length(cfg.buffer) > 0
        s = last(stack(cfg))
        return s != 0 && !has_head(token(cfg, s))
    end
    return false
end

function is_possible(::RightArc, cfg::ArcEagerConfig)
    return length(cfg.stack) > 0 && length(cfg.buffer) > 0 &&
        !has_head(token(cfg, first(buffer(cfg))))
end

is_possible(::Reduce, cfg::ArcEagerConfig) =
    stacklength(cfg) > 0 && has_head(token(cfg, last(stack(cfg))))

is_possible(::Shift, cfg::ArcEagerConfig) = bufferlength(cfg) > 0

==(cfg1::ArcEagerConfig, cfg2::ArcEagerConfig) =
    cfg1.stack == cfg2.stack && cfg1.buffer == cfg2.buffer && cfg1.A == cfg2.A

Base.getindex(cfg::ArcEagerConfig, i) = token(cfg, i)


text:
@testset "model" begin
    data_gen = DataGenerator(
        Sawtooth(),
        batch_size=4,
        x_context=Uniform(-2, 2),
        x_target=Uniform(-2, 2),
        num_context=DiscreteUniform(0, 3),
        num_target=DiscreteUniform(1, 3),
        σ²=1e-4
    )

    # Construct CNPs:
    cnps = [
        cnp_1d(
            dim_embedding=10,
            num_encoder_layers=2,
            num_decoder_layers=2
        ),
        acnp_1d(
            dim_embedding=10,
            num_encoder_layers=2,
            num_encoder_heads=3,
            num_decoder_layers=2
        ),
        convcnp_1d(
            receptive_field=1f0,
            num_layers=2,
            num_channels=2,
            points_per_unit=5f0
        )
    ]

    # Construct NPs:
    nps = Any[]
    for noise_type in ["fixed", "amortised", "het"], pooling_type in ["mean", "sum"]
        push!(nps, np_1d(
            dim_embedding=10,
            num_encoder_layers=2,
            num_decoder_layers=2,
            noise_type=noise_type,
            pooling_type=pooling_type
        ))
        push!(nps, anp_1d(
            dim_embedding=10,
            num_encoder_layers=2,
            num_encoder_heads=3,
            num_decoder_layers=2,
            noise_type=noise_type,
            pooling_type=pooling_type
        ))
        push!(nps, convnp_1d(
            receptive_field=1f0,
            num_encoder_layers=2,
            num_decoder_layers=3,
            num_encoder_channels=2,
            num_decoder_channels=1,
            num_latent_channels=3,
            num_global_channels=0,
            points_per_unit=5f0,
            noise_type=noise_type,
            pooling_type=pooling_type
        ))
        push!(nps, convnp_1d(
            receptive_field=1f0,
            num_encoder_layers=2,
            num_decoder_layers=3,
            num_encoder_channels=2,
            num_decoder_channels=1,
            num_latent_channels=3,
            num_global_channels=5,
            points_per_unit=5f0,
            noise_type=noise_type,
            pooling_type=pooling_type
        ))
    end

    # Construct all pairs of models and losses. CNPs:
    model_losses = Any[]
    for model in cnps
        push!(
            model_losses,
            (model, (xs...) -> NeuralProcesses.loglik(xs..., num_samples=1))
        )
    end

    # NPs:
    for model in nps
        push!(model_losses, (
            model,
            (xs...) -> NeuralProcesses.loglik(
                xs...,
                num_samples=2,
                batch_size=1,
                importance_weighted=false,
                fixed_σ_epochs=1
            )
        ))
        push!(model_losses, (
            model,
            (xs...) -> NeuralProcesses.loglik(
                xs...,
                num_samples=2,
                batch_size=1,
                importance_weighted=true,
                fixed_σ_epochs=1
            )
        ))
        push!(model_losses, (
            model,
            (xs...) -> NeuralProcesses.elbo(
                xs...,
                num_samples=2,
                fixed_σ_epochs=1
            )
        ))
    end

    for (model, loss) in model_losses
        # Test model training for a few epochs. The statements just have to not error.
        @test (report_num_params(model); true)
        @test (train!(
            model,
            loss,
            data_gen,
            ADAM(1e-4),
            bson=nothing,
            starting_epoch=1,
            epochs=2,
            tasks_per_epoch=200,
            path=nothing
        ); true)
    end
end


text:
using Yao, Zygote, YaoFlux, YaoArrayRegister, YaoBlocks
import Yao.YaoBlocks: RotationGate

# a, back = Zygote._forward(apply!, rand_state(1), Rx(0.1))

# back(rand_state(1))

Zygote.refresh()

r, back = Zygote._forward(0.1, 0.1) do x, y
    return apply!(zero_state(1), chain(Rx(x), Rz(y)))
end

back(YaoFlux.GradReg(r, rand_state(1)))

_, back = Zygote._forward(fidelity, rand_state(1), rand_state(1))
back([1])

t = rand_state(1)
r = zero_state(1)
gradient(x->fidelity(t, apply!(r, Rx(x)))[], 0.1)

Zygote._forward(Rx, 0.1)

_, back = Zygote._forward((x, y)->abs(statevec(x)' * statevec(y)), rand_state(1), rand_state(1))
back(1)

_, back = Zygote._forward(state, rand_state(1))
back(rand(2))