text:
import ..SamplingConstraints:
    SequentialSamplingConstraint,
    OrderedSamplingAlgorithm

const AUD = AbstractUncertainDataset

const SC = Union{SamplingConstraint, Vector{S}} where S <: SamplingConstraint
const SEQ = SequentialSamplingConstraint{O} where O <: OrderedSamplingAlgorithm

const XD = Union{AbstractUncertainDataset, Vector{<:AbstractUncertainValue}}

"""
    resample(x, ssc::SequentialSamplingConstraint)
    resample(x, ssc::SequentialSamplingConstraint, c)
    resample(x::UncertainIndexValueDataset, ssc::SequentialSamplingConstraint, ic, vc)

Sample `x` element-wise such that the samples obey the sequential constraints given by `ssc`.
Alteratively, apply constrain(s) `c` to `x` *before* sequential sampling is performed.

A check is performed before sampling to ensure that such a sequence exists.
Before the check is performed, the distributions in `x` are truncated element-wise
to the quantiles provided by `c` to ensure they have finite supports.

If `x` is an uncertain index-value dataset, then the sequential constraint is only applied to
the indices. If one set of additional constraints are added to an uncertain index-value dataset,
then they are applied to both the indices and the values. It is also possible to give separate
index constraints `ic` and value constraints `vc`.

If constraints `c`, or `ic` and `vc`, are given, then `c`/`ic`/`vc` must be either a single constraint,
or a vector of constraints where constraints are applied element-wise to the distributions in `x`.

    resample!(s, x, ssc::SequentialSamplingConstraint, lqs, uqs)

The same as above, but store the sampled values a pre-allocated vector `s`, where
`length(x) == length(s)`. This avoids excessive memory allocations during repeated
resampling. This requires pre-computing the element-wise lower and upper quantiles
`lqs` and `uqs` for the initial truncation step.

This method *does not* check for the existence of a strictly increasing sequence in `x`.
To check that, use [`sequence_exists`](@ref).

See also: [`sequence_exists`](@ref), [`StrictlyIncreasing`](@ref),
[`StrictlyDecreasing`](@ref), [`StartToEnd`](@ref).

## Examples

```julia
N = 100
t = [UncertainValue(Normal, i, 2) for i in 1:N];
resample(t, StrictlyIncreasing(StartToEnd()))
```

```julia
N = 100
t = [UncertainValue(Normal, i, 2) for i in 1:N];

# Verify that an increasing sequence through `t` exists
c = StrictlyIncreasing(StartToEnd())
exists, lqs, uqs = sequence_exists(t, c)

# Pre-allocate sample vector
s = zeros(Float64, N)

if exists
    for i = 1:100
        resample!(s, t, c)
        # Do something with s
        # ...
    end
end
```
"""
function resample(x::XD, ssc::SEQ)
    exists, lqs, uqs = sequence_exists(x, ssc)
    exists || error("Sequence does not exist")
    _draw(x, ssc, lqs, uqs)
end

function resample!(s, x::XD, ssc::SEQ, lqs, uqs)
    _draw!(s, x, ssc, lqs, uqs)
end

resample(udata::XD, ssc::SEQ, c::Union{SC, AbstractVector{<:SC}}) =
    resample(constrain(udata, c), ssc)

resample(udata::UncertainIndexValueDataset, ssc::SEQ) =
    resample(udata.indices, ssc), resample(udata.values)

resample(udata::UncertainIndexValueDataset, ssc::SEQ, c::SC) =
    resample(constrain(udata.indices, c), ssc), resample(constrain(udata.values, c))

resample(udata::UncertainIndexValueDataset, ssc::SEQ, ic::SC, vc::SC) =
    resample(constrain(udata.indices, ic), ssc), resample(constrain(udata.values, vc))

ext: jl


text:
### A Pluto.jl notebook ###
# v0.12.7

using Markdown
using InteractiveUtils

# ╔═╡ 6fd616f6-1ea1-11eb-3814-8bfb4a096c49
using Pkg, DrWatson

# ╔═╡ 6ffe2628-1ea1-11eb-24ea-57f985146a72
begin
	@quickactivate "StatisticsWithJuliaPlutoNotebooks"
	using Distributions
end;

# ╔═╡ ed174bc4-1ea0-11eb-1e2f-a32874cec549
md"## Listing 3.10"

# ╔═╡ b7267278-2133-11eb-0f50-89cc660d3fc3
dist = TriangularDist(0,2,1)

# ╔═╡ b726b8f0-2133-11eb-090e-13eb003b6cf7
params(dist)

# ╔═╡ b7276444-2133-11eb-1818-95aafff7ec2b
Text("Central descriptors: $(mean(dist)), $(median(dist))")

# ╔═╡ ca5f5b34-2133-11eb-20c2-eda06044b2e9
Text("Dispersion descriptors: $(var(dist)), $(std(dist))")

# ╔═╡ b732099e-2133-11eb-01e3-bd42597ccab9
Text("Higher moment shape descriptors:  $(skewness(dist)), $(kurtosis(dist))")

# ╔═╡ decac540-2133-11eb-07af-3d70c3daf2fb
Text("Range:  $(minimum(dist)), $(maximum(dist))")

# ╔═╡ b73c5110-2133-11eb-3fae-65132a905cb1
Text("Mode: $(mode(dist)), Modes: $(modes(dist))")

# ╔═╡ 70182cd8-1ea1-11eb-094d-8d0d49cf15f3
md"## End of listing 3.10"

# ╔═╡ Cell order:
# ╟─ed174bc4-1ea0-11eb-1e2f-a32874cec549
# ╠═6fd616f6-1ea1-11eb-3814-8bfb4a096c49
# ╠═6ffe2628-1ea1-11eb-24ea-57f985146a72
# ╠═b7267278-2133-11eb-0f50-89cc660d3fc3
# ╠═b726b8f0-2133-11eb-090e-13eb003b6cf7
# ╠═b7276444-2133-11eb-1818-95aafff7ec2b
# ╠═ca5f5b34-2133-11eb-20c2-eda06044b2e9
# ╠═b732099e-2133-11eb-01e3-bd42597ccab9
# ╠═decac540-2133-11eb-07af-3d70c3daf2fb
# ╠═b73c5110-2133-11eb-3fae-65132a905cb1
# ╟─70182cd8-1ea1-11eb-094d-8d0d49cf15f3


text:
using FileUtils

template_package_name = "PACKAGE"

package_names = [
  "ITensors",
  "ITensorGaussianMPS",
  "ITensorGLMakie",
  "ITensorGPU",
  "ITensorMakie",
  "ITensorUnicodePlots",
  "ITensorVisualizationBase",
  "NDTensors",
]

bug_report_file(package_name::String) = "$(package_name)_bug_report.md"
feature_request_file(package_name::String) = "$(package_name)_feature_request.md"

for package_name in package_names
  @show package_name

  template_bug_report = bug_report_file(template_package_name)
  new_bug_report = bug_report_file(package_name)

  if isfile(new_bug_report)
    println("File $new_bug_report already exists, skipping")
  else
    println("Copying $template_bug_report to $new_bug_report")
    cp(template_bug_report, new_bug_report)

    println("Replace $template_package_name with $package_name in $new_bug_report")
    replace_in_file(new_bug_report, template_package_name => package_name)
  end

  template_feature_request = feature_request_file(template_package_name)
  new_feature_request = feature_request_file(package_name)

  if isfile(new_feature_request)
    println("File $new_feature_request already exists, skipping")
  else
    println("Copying $template_feature_request to $new_feature_request")
    cp(template_feature_request, new_feature_request)

    println("Replace $template_package_name with $package_name in $new_feature_request")
    replace_in_file(new_feature_request, template_package_name => package_name)
  end
end


text:
lines = readlines("input.txt")
field = zeros(Int, (1024,1024))
re = r"(\d+)"

for line = lines
    x1,y1,x2,y2 = 1 .+ parse.(Int, m.match for m = eachmatch(re, line))
    dx,dy = x1 < x2 ? 1 : -1, y1 < y2 ? 1 : -1
    if x1 != x2 && y1 != y2
        for ij = CartesianIndex.(zip(y1:dy:y2, x1:dx:x2))
            field[ij] += 1
        end
    else
        field[y1:dy:y2, x1:dx:x2] .+= 1
    end
end

println(sum(@. field > 1))


text:
@testset "mindist2cntrs" begin
    npts_per_cluster = 4
    r = 0.25
    c = convert(Matrix{Float},[1 1 -1 -1; 1 -1 1 -1])
    x = VectorDataUtils.generate_4_circle_clusters(Float, npts_per_cluster; r)
    md = VectorDataUtils.mindist2cntrs(c,x)
    for i ∈ eachindex(md)
        @test md[i] == r^2
    end
end

@testset "closest_centre" begin
    npts_per_cluster = 4
    r = 0.25
    c = convert(Matrix{Float},[1 1 -1 -1; 1 -1 1 -1])
    x = VectorDataUtils.generate_4_circle_clusters(Float, npts_per_cluster; r)
    cc = VectorDataUtils.closest_centre(c,x)
    for i ∈ axes(c,2)
        sindx = (i-1)*npts_per_cluster+1
        eindx = i*npts_per_cluster
        @test cc[sindx:eindx]==ones(npts_per_cluster)*i
    end
end

@testset "max_centre_shift" begin
    c = convert(Matrix{Float},[1 1 -1 -1; 1 -1 1 -1])
    ncntrs = size(c,2)
    r = rand(ncntrs)
    θ = rand(ncntrs)*2π
    shift = Matrix{Float}(undef,2,ncntrs)
    for i ∈ axes(c,2)
        shift[1,i] = r[i]*cos(θ[i])
        shift[2,i] = r[i]*sin(θ[i])
    end
    ĉ = c .+ shift
    mshift = VectorDataUtils.max_centre_shift(ĉ,c)
    @test mshift ≈ maximum(r)
end

@testset "kmeans_init" begin
    npts_per_cluster = 4
    r = 0.25
    x = VectorDataUtils.generate_4_circle_clusters(npts_per_cluster; r)

    cntrs = VectorDataUtils.kmeans_init(VectorDataUtils.kmpp(), x, 4)
    @test size(cntrs,1) == size(x,1)
    @test size(cntrs,2) == 4
    @test eltype(cntrs) == eltype(x)
    md = VectorDataUtils.mindist2cntrs(x,cntrs)
    @test md == zeros(4)

    cntrs = VectorDataUtils.kmeans_init(VectorDataUtils.kmrand(), x, 4)
    @test size(cntrs,1) == size(x,1)
    @test size(cntrs,2) == 4
    @test eltype(cntrs) == eltype(x)
    md = VectorDataUtils.mindist2cntrs(x,cntrs)
    @test md == zeros(4)
end

@testset "k-means" begin
    npts_per_cluster = 8
    r = 0.25
    x = VectorDataUtils.generate_4_circle_clusters(npts_per_cluster; r)
    cn = VectorDataUtils.kmeans_init(VectorDataUtils.kmpp(), x, 4)
    icn = copy(cn)
    cc = Vector{Int}(undef,size(x,2))
    VectorDataUtils.kmeans_update!(cc,cn,x)
    for i ∈ axes(cn,2)
        @test VectorDataUtils.mindist2cntrs(icn,cn[:,[i]]) > [0.0]
    end

    kcn = VectorDataUtils.kmeans(x,4)
    @test size(kcn,1) == size(x,1)
    @test size(kcn,2) == 4
    @test eltype(kcn) == eltype(x)
    VectorDataUtils.closest_centre!(cc,kcn,x)
    for i ∈ axes(kcn,2)
        @test kcn[:,i] == VectorDataUtils.running_mean(x[:,cc.==i])
    end
end


text:
#--------------------------------------------------------------------------------------------------
@kernel function compute_vec_grad_kernel!(
    sv::AbstractArray{FT},
    D::AbstractArray{FT, 2},
    vgeo::AbstractArray{FT},
    g::AbstractArray{FT, 4},
    vgrad_data::AbstractArray{FT, 3},
    _ρ::Int,
    _ρu::Int,
    _ρv::Int,
    _ρw::Int,
    ::Val{qm1},
) where {qm1, FT <: AbstractFloat}

    e = @index(Group, Linear)
    i, j = @index(Local, NTuple)

    s_D = @localmem FT (qm1, qm1)
    s_U = @localmem FT (qm1, qm1)
    s_V = @localmem FT (qm1, qm1)
    s_W = @localmem FT (qm1, qm1)

    s_D[i, j] = D[i, j]
    @synchronize
    # computing derivatives with respect to ξ1
    for t in 1:qm1, s in 1:qm1
        ijk = j + ((s - 1) + (t - 1) * qm1) * qm1
        s_U[i, j] = s_D[i, j] * (sv[ijk, _ρu, e] / sv[ijk, _ρ, e])
        s_V[i, j] = s_D[i, j] * (sv[ijk, _ρv, e] / sv[ijk, _ρ, e])
        s_W[i, j] = s_D[i, j] * (sv[ijk, _ρw, e] / sv[ijk, _ρ, e])
        @synchronize
        if j == 1
            for r in 2:qm1
                s_U[i, 1] += s_U[i, r]
                s_V[i, 1] += s_V[i, r]
                s_W[i, 1] += s_W[i, r]
            end
        end
        @synchronize
        if j == 1
            g[i + ((s - 1) + (t - 1) * qm1) * qm1, e, 1, 1] = s_U[i, 1] # ∂u₁∂ξ₁
            g[i + ((s - 1) + (t - 1) * qm1) * qm1, e, 2, 1] = s_V[i, 1] # ∂u₂∂ξ₁
            g[i + ((s - 1) + (t - 1) * qm1) * qm1, e, 3, 1] = s_W[i, 1] # ∂u₃∂ξ₁
        end
    end
    @synchronize
    # computing derivatives with respect to ξ2
    for t in 1:qm1, r in 1:qm1
        ijk = r + ((j - 1) + (t - 1) * qm1) * qm1
        s_U[i, j] = s_D[i, j] * (sv[ijk, _ρu, e] / sv[ijk, _ρ, e])
        s_V[i, j] = s_D[i, j] * (sv[ijk, _ρv, e] / sv[ijk, _ρ, e])
        s_W[i, j] = s_D[i, j] * (sv[ijk, _ρw, e] / sv[ijk, _ρ, e])
        @synchronize
        if j == 1
            for s in 2:qm1
                s_U[i, 1] += s_U[i, s]
                s_V[i, 1] += s_V[i, s]
                s_W[i, 1] += s_W[i, s]
            end
        end
        @synchronize
        if j == 1
            g[r + ((i - 1) + (t - 1) * qm1) * qm1, e, 1, 2] = s_U[i, 1] # ∂u₁∂ξ₂
            g[r + ((i - 1) + (t - 1) * qm1) * qm1, e, 2, 2] = s_V[i, 1] # ∂u₂∂ξ₂
            g[r + ((i - 1) + (t - 1) * qm1) * qm1, e, 3, 2] = s_W[i, 1] # ∂u₃∂ξ₂
        end
    end
    @synchronize
    # computing derivatives with respect to ξ3
    for s in 1:qm1, r in 1:qm1
        ijk = r + ((s - 1) + (j - 1) * qm1) * qm1
        s_U[i, j] = s_D[i, j] * (sv[ijk, _ρu, e] / sv[ijk, _ρ, e])
        s_V[i, j] = s_D[i, j] * (sv[ijk, _ρv, e] / sv[ijk, _ρ, e])
        s_W[i, j] = s_D[i, j] * (sv[ijk, _ρw, e] / sv[ijk, _ρ, e])
        @synchronize
        if j == 1
            for t in 2:qm1
                s_U[i, 1] += s_U[i, t]
                s_V[i, 1] += s_V[i, t]
                s_W[i, 1] += s_W[i, t]
            end
        end
        @synchronize
        if j == 1
            g[r + ((s - 1) + (i - 1) * qm1) * qm1, e, 1, 3] = s_U[i, 1] # ∂u₁∂ξ₃
            g[r + ((s - 1) + (i - 1) * qm1) * qm1, e, 2, 3] = s_V[i, 1] # ∂u₂∂ξ₃
            g[r + ((s - 1) + (i - 1) * qm1) * qm1, e, 3, 3] = s_W[i, 1] # ∂u₃∂ξ₃
        end
    end
    @synchronize

    ∂₁u₁, ∂₂u₁, ∂₃u₁ = 1, 2, 3
    ∂₁u₂, ∂₂u₂, ∂₃u₂ = 4, 5, 6
    ∂₁u₃, ∂₂u₃, ∂₃u₃ = 7, 8, 9

    for k in 1:qm1
        ijk = i + ((j - 1) + (k - 1) * qm1) * qm1

        ξ1x1 = vgeo[ijk, _ξ1x1, e]
        ξ1x2 = vgeo[ijk, _ξ1x2, e]
        ξ1x3 = vgeo[ijk, _ξ1x3, e]
        ξ2x1 = vgeo[ijk, _ξ2x1, e]
        ξ2x2 = vgeo[ijk, _ξ2x2, e]
        ξ2x3 = vgeo[ijk, _ξ2x3, e]
        ξ3x1 = vgeo[ijk, _ξ3x1, e]
        ξ3x2 = vgeo[ijk, _ξ3x2, e]
        ξ3x3 = vgeo[ijk, _ξ3x3, e]

        vgrad_data[ijk, ∂₁u₁, e] =
            g[ijk, e, 1, 1] * ξ1x1 +
            g[ijk, e, 1, 2] * ξ2x1 +
            g[ijk, e, 1, 3] * ξ3x1
        vgrad_data[ijk, ∂₂u₁, e] =
            g[ijk, e, 1, 1] * ξ1x2 +
            g[ijk, e, 1, 2] * ξ2x2 +
            g[ijk, e, 1, 3] * ξ3x2
        vgrad_data[ijk, ∂₃u₁, e] =
            g[ijk, e, 1, 1] * ξ1x3 +
            g[ijk, e, 1, 2] * ξ2x3 +
            g[ijk, e, 1, 3] * ξ3x3

        vgrad_data[ijk, ∂₁u₂, e] =
            g[ijk, e, 2, 1] * ξ1x1 +
            g[ijk, e, 2, 2] * ξ2x1 +
            g[ijk, e, 2, 3] * ξ3x1
        vgrad_data[ijk, ∂₂u₂, e] =
            g[ijk, e, 2, 1] * ξ1x2 +
            g[ijk, e, 2, 2] * ξ2x2 +
            g[ijk, e, 2, 3] * ξ3x2
        vgrad_data[ijk, ∂₃u₂, e] =
            g[ijk, e, 2, 1] * ξ1x3 +
            g[ijk, e, 2, 2] * ξ2x3 +
            g[ijk, e, 2, 3] * ξ3x3

        vgrad_data[ijk, ∂₁u₃, e] =
            g[ijk, e, 3, 1] * ξ1x1 +
            g[ijk, e, 3, 2] * ξ2x1 +
            g[ijk, e, 3, 3] * ξ3x1
        vgrad_data[ijk, ∂₂u₃, e] =
            g[ijk, e, 3, 1] * ξ1x2 +
            g[ijk, e, 3, 2] * ξ2x2 +
            g[ijk, e, 3, 3] * ξ3x2
        vgrad_data[ijk, ∂₃u₃, e] =
            g[ijk, e, 3, 1] * ξ1x3 +
            g[ijk, e, 3, 2] * ξ2x3 +
            g[ijk, e, 3, 3] * ξ3x3
    end
end
#--------------------------------------------------------------------------------------------------
@kernel function compute_vorticity_kernel!(
    vgrad_data::AbstractArray{FT, 3},
    vort_data::AbstractArray{FT, 3},
    ::Val{qm1},
) where {qm1, FT <: AbstractFloat}

    e = @index(Group, Linear)
    i, j = @index(Local, NTuple)

    ∂₁u₁, ∂₂u₁, ∂₃u₁ = 1, 2, 3
    ∂₁u₂, ∂₂u₂, ∂₃u₂ = 4, 5, 6
    ∂₁u₃, ∂₂u₃, ∂₃u₃ = 7, 8, 9
    Ω₁, Ω₂, Ω₃ = 1, 2, 3

    for k in 1:qm1
        ijk = i + ((j - 1) + (k - 1) * qm1) * qm1

        vort_data[ijk, Ω₁, e] =
            vgrad_data[ijk, ∂₃u₂, e] - vgrad_data[ijk, ∂₂u₃, e]
        vort_data[ijk, Ω₂, e] =
            vgrad_data[ijk, ∂₁u₃, e] - vgrad_data[ijk, ∂₃u₁, e]
        vort_data[ijk, Ω₃, e] =
            vgrad_data[ijk, ∂₂u₁, e] - vgrad_data[ijk, ∂₁u₂, e]
    end
end
#--------------------------------------------------------------------------------------------------

max_stars_repo_name: yairchn/CLIMA

max_issues_repo_name: yairchn/CLIMA

max_forks_repo_name: yairchn/CLIMA


text:
__precompile__(false)
module CUDAnativelib
  export CURANDkernel

  module Version
    using CUDAdrv
    using Cxx

    export __triple__, __sm__, CUDA_HOME, CUDA_VERSION, LLVM_VERSION

    # TODO: make this configurable
    const __cap__ = capability(CuDevice(0))
    const __sm__ = "sm_$(__cap__.major)$(__cap__.minor)"
    const __triple__ = Int === Int64 ? "nvptx64-nvidia-cuda" : "nvptx-nvidia-cuda"

    const candidates = ["/usr/local/cuda", "/opt/cuda", "/usr"]
    CUDA_HOME = ""
    if haskey(ENV, "CUDA_HOME")
        CUDA_HOME = ENV["CUDA_HOME"]
    else
      for dir in candidates
        if isdir(dir) && isfile(joinpath(dir, "include", "cuda.h"))
          CUDA_HOME = dir
          break
        end
      end
    end

    isempty(CUDA_HOME) && error("Please set CUDA_HOME to the location of your CUDA installation.")
    addHeaderDir(joinpath(CUDA_HOME, "include"), kind = C_System)

    cxxinclude("cuda.h")
    const CUDA_VERSION = icxx"CUDA_VERSION;"
    const LLVM_VERSION = VersionNumber(Base.libllvm_version)
  end

  include("codegen.jl")
  include("curand.jl")
end # module


text:
#########################################################################################
# module StateSelection
#
# Functions to transform the ODAE generated by Pantelides to a special index-1 DAE
# using the static dummy derivative method with tearing to select dummy states and a
# generalization of the Gear-Gupta-Leimkuhler stabilization of multibody systems to handle
# the remaining potential states (that would otherwise be treated with the dynamic
# dummy derivative method). Details are described in the paper:
#
#    Otter, Elmqvist (2017): Transformation of Differential Algebraic Array Equations to
#                            Index One Form. Modelica'2017 Conference.
#
# Author: Martin Otter, DLR-SR (first version: Jan. 11, 2017)
#########################################################################################
module StateSelection

export SortedEquationGraph
export getSortedEquationGraph!
export getSortedEquationGraph
export printSortedEquationGraph
export newRaggedIntMatrix

include("Tearing.jl")

"""
    v = newRaggedIntMatrix(n)

Generate a new ragged int matrix with n rows and every row is initialized with a zero-sized Int vector.
"""
newRaggedIntMatrix(n) = [ fill(0, 0) for i = 1:n ]

"""
    Arev = revertAssociation(A,nArev)

Reverts the association Vector A[i] = j, such that Arev[j] = i (A[i]=0 is allowed and is ignored)
nArev is the dimension of Arev (nArev >= largest value of j appearing in A[i] = j).
"""
function revertAssociation(A::Vector{Int}, nArev::Int)::Vector{Int}
    Arev = fill(0, nArev)
    for i in eachindex(A)
        if A[i] != 0
            Arev[ A[i] ] = i
        end
    end
    return Arev
end

"""
    (eConstraints,vConstraints) = getConstraintSets(eBLT,Eassign,Arev,Brev)

Determines the set of constraint equations and their unknowns for BLT component eBLT.

Input arguments:

- `eBLT`: Vector of equation indices of BLT component c.
- `Eassign`: vc = Eassign(ec[i]) is variable vc assigned to equation ec[i].
- `Arev`: Reverted variable association: `Arev[i] = der(V[k]) == V[i] ? k : 0`.
- `Brev`: Reverted equation association: `Erev[i] = der(E[k]) == E[i] ? k : 0`.

Output arguments:
- `eConstraints`: eConstraints[1] are the lowest-order, eConstraints[end-1] are the highest-order constraint equations
   and eConstraints[end] = eBLT
- `vConstraints`: vConstraints[i] are the unknowns of eConstraints[i].
"""
function getConstraintSets(eBLT::Vector{Int}, Eassign::Vector{Int}, Arev::Vector{Int}, Brev::Vector{Int})
    # Determine unknowns of BLT component c
    vBLT = fill(0, length(eBLT))
    for i in eachindex(eBLT)
        vBLT[i] = Eassign[eBLT[i]]
    end

    # Determine constraint sets
    eConstraints = [eBLT]
    vConstraints = [vBLT]
    while true
        # Determine constraints at one differentiation order less
        ceq = fill(0, 0)
        for eq in eConstraints[1]
            if Brev[eq] > 0
                push!(ceq, Brev[eq])
            end
        end
        if length(ceq) == 0; break; end

        @static if VERSION < v"0.7.0-DEV.2005"
            unshift!(eConstraints, ceq)   # move ceq to the beginning of the constraints vector
        else
            pushfirst!(eConstraints, ceq)   # move ceq to the beginning of the constraints vector
        end
        # Determine unknowns of constraints at one differentiation order less
        veq = fill(0, 0)

        for vc in vConstraints[1]
            if Arev[vc] > 0
                push!(veq, Arev[vc])
            end

            if length(veq) == 0;
                error("Error should not occur: eBLT equations and vBLT variables have different differentiation orders");
            end
        end

        @static if VERSION < v"0.7.0-DEV.2005"
            unshift!(vConstraints, veq)   # move veq to the beginning of the constraints vector
        else
            pushfirst!(vConstraints, veq)   # move veq to the beginning of the constraints vector
        end
    end

    @assert(length(eConstraints) == length(vConstraints))
    return (eConstraints, vConstraints)
end

"""
    eqInit = SortedEquationGraph(G,BLT,assign,A,B,VNames)

Initialize data structure to determine the sorted equation graph

Input arguments:

- `Gorigin`: Gorigin[i] is the vector of variable indices of equation i.
- `BLT`: BLT[i] is the vector of equations belonging to BLT-block i
- `assign`: ei = assign[vi] is equation ei assigned to variable vi
- `A`: A-Vector of Pantelides
- `B`: B-Vector of Pantelides
- `VNames`: VNames[i] is the name of variable i (only used for potential debug and print output)

Output arguments:

- `eqInit`: Initialized data structure to determine the sorted equations.
            eqInit.Gunknowns is a subset of Gorigin that contains the variables
            of Gorigin that are treated as unknowns (e.g. used in tearing)

"""
mutable struct SortedEquationGraph
    first::Bool
    Gorigin::Vector{Any}            # Gorigin[i] is the vector of variable indices of equation i
    Gunknowns::Vector{Vector{Int}}  # Gunknowns[i] is the vector of variables of equation i that are treated as unknowns
    Gsolvable::Vector{Any}          # Gsolvable[i] contains the variables that can be
                                    # solved for equation i (used for tearing)
    eConstraintsVec::Vector{Vector{Vector{Int}}}   # ec = eConstraintsVec[i]: constraint sets of BLT block i;
                                                   #      ec[1]  : lowest derivative equations
                                                   #      ec[end]: highest derivative equations (= BLT block i)
    vConstraintsVec::Vector{Vector{Vector{Int}}}   # vs = vConstraintsVec[i]: unknowns of ec

    assign::Vector{Int}         # ei = assign[vi] is equation ei assigned to variable vi
    A::Vector{Int}              # A-Vector of Pantelides
    B::Vector{Int}              # B-Vector of Pantelides
    VNames::Vector{String}      # VNames[i] is the name of variable i

    Arev::Vector{Int}           # Reverted vector A
    Brev::Vector{Int}           # Reverted vector B
    Eassign::Vector{Int}        # Reverted vector assign; vi = Eassign[ei] is equation ei assigned to variable vi

    isorted::Int                # The last element of ESorted that is defined
    irbeg::Int                  # The next element of the residue vector
    ESorted::Vector{Int}        # Sorted equations (equations must be generated in the order ESorted[1],[2],..)
    ESolved::Vector{Int}        # If ESolved[i] > 0 then equation ESorted[i] is explicitly solved for variable ESolved[i]
                               # If ESolved[i] < 0 then ESorted[i] is a residue equation to compute _r[ abs(ESolved[i]) ]
    rcat::Vector{Vector{Int}}

    Vx::Vector{Int}             # Variable j = Vx[i] is part of vector x; if j=0, this element of x is a dummy variable (not used)
    VxRev::Vector{Int}          # If Vxrev[j] = i, then Vx[i] = j
                               # If Vxref[j] = 0, then variable j is not part of x
    Vderx::Vector{Int}          # Variable j = Vderx[i] is part of vector der_x; if j=0, this element of der_x is a dummy variable (not used)
    VderxRev::Vector{Int}       # If VderxRev[j] = i, then Vderx[i] = j
                               # If VderxRef[j] = 0, then variable j is not part of der_x
    Vmue::Vector{Int}           # Equation j = Vmue[i] is the equation that is associated to mue[i] (mue[i] has the same sizes has equation j)

    td::TraverseDAG             # Holds information about tearing the equations

    nc::Int                     # Number of residue constraints, so part that depends on (x,t) but not on der(x); 0 <= nc <= nx
    Er::Vector{Int}             # Equation j = Er[i] is part of residue vector r; if j=0, this element is part of ider0n2/ider1n1
    EAlgebraic::Vector{Bool}    # EAlgebraic[j] = true: equation j is an algebraic equation (cat = RA)
    VAlgebraic::Vector{Bool}    # VAlgebraic[j] = true: variable j is an algebraic variable (not a lambda variable)

   # The residue equation
   #   residue[1:nider] = der_x[ider0n2] - x[ider1n1]
   # must be added to the generated code
    ider0n2::Vector{Int}        # index vector der(0:n-2)
    ider1n1::Vector{Int}        # index vector der(1:n-1)

    function SortedEquationGraph(Gorigin, BLT, assign, A, B, VNames)
        @assert(length(assign) == length(VNames))
        @assert(length(A)      == length(VNames))
        @assert(length(B)      == length(Gorigin))

      # Revert association vectors
        Arev    = revertAssociation(A, length(A))
        Brev    = revertAssociation(B, length(B))
        Eassign = revertAssociation(assign, length(B))

      # Generate vectors with default elements
        ESorted = fill(0, 0)
        ESolved = fill(0, 0)

      # Residue vector structure
        rcat = newRaggedIntMatrix(6)

      # Structure of x and der_x vectors
        Vx       = fill(0, 0)
        VxRev    = fill(0, length(VNames))
        Vderx    = fill(0, 0)
        VderxRev = fill(0, length(VNames))
        Vmue     = fill(0, 0)

      # Algebraic equations/variables
        EAlgebraic = fill(false, length(Gorigin))
        VAlgebraic = fill(false, length(VNames))

      # Determine equation/variable constraint sets from BLT
        Gunknowns = newRaggedIntMatrix(length(Gorigin))
        eConstraintsVec = Vector{Vector{Int}}[]
        vConstraintsVec = Vector{Vector{Int}}[]
        lowerDerivativeEquationsInBLT = false
        c_ignore = false

        for c in BLT
         # Ignore block c if lower derivative equations
            c_ignore = false
            for ceq in c
                if B[ceq] != 0
                    lowerDerivativeEquationsInBLT = true
                    c_ignore = true
                end
            end

            if c_ignore
                continue
            end

         # Get all equation sets eConstraints and their corresponding unknowns vConstraints
         # from lowest to highest differentiation order (eConstraints[end] is c)
            (eConstraints, vConstraints) = getConstraintSets(c, Eassign, Arev, Brev)
         #println("... eConstraints = ", eConstraints)
         #println("... vConstraints = ", vConstraints)
            push!(eConstraintsVec, eConstraints)
            push!(vConstraintsVec, vConstraints)

         # Construct Gunknowns
            for i in eachindex(eConstraints)
                for eq in eConstraints[i]
                    Gunknowns[eq] = intersect(Gorigin[eq], vConstraints[i])
                end
            end
        end

        if lowerDerivativeEquationsInBLT
            println("\n... Warning from SortedEquationGraph(..) (in StateSelection.jl):",
                "\n... BLT blocks with lower derivative equations have been ignored.\n")
        end

      # Tearing information
        td = TraverseDAG(Gunknowns, length(VNames))

        new(true, Gorigin, Gunknowns, Any[], eConstraintsVec, vConstraintsVec, assign, A, B, VNames, Arev, Brev, Eassign, 0, 1, ESorted, ESolved,
          rcat, Vx, VxRev, Vderx, VderxRev, Vmue, td, 0, fill(0, 0), EAlgebraic, VAlgebraic)
    end
end

"""
rCategorie:

- `RD`: Original, non-differentiated equation with derivative variables as unknowns and no constraints
- `RA`: Original, non-differentiated algebraic equation (only algebraic variables in the equation) and no constraints
- `RDER0`: Constraint equation that is not differentiated
- `RDER1`: Constraint equation is differentiated at least once, but not the highest derivative level
- `RDERN`: Constraint equation on highest derivative level
"""
@enum(rCategorie, RD = 1, RA, RDER0, RDER1, RDERN)

"""
    appendToSortedEquations!(eq,eSolved,vSolved,eResidue,vTear,highestDerivative)

Append eSolved and eResidue equations to sorted equations eq.ESorted,
store variables vSolved appropriatey, and store variables vTear in x or der_x vectors.
"""
function appendToSortedEquations!(eq::SortedEquationGraph, eSolved::Vector{Int}, vSolved::Vector{Int},
                                  eResidue::Vector{Int}, vTear::Vector{Int}, highestDerivative::Bool)

   # println("... highestDerivative = ", highestDerivative, ", eSolved = ", eSolved, ", vSolved = ", vSolved, ",
   #          eResidue = ", eResidue, ", vTear = ", vTear)
    cat::rCategorie = RD

   # Store solved equations in the right order in ESorted
    for i in eachindex(eSolved)
        push!(eq.ESorted, eSolved[i])
        vs = vSolved[i]

        if highestDerivative && eq.Brev[eSolved[i]] == 0 && eq.Arev[vs] > 0
         # Highest derivative equation, that is not differentiated and vs is a derivative
         # -> change solved equation to residue
            push!(eq.ESolved, 0)
            push!(eq.rcat[Int(RD)], length(eq.ESorted))
            vsInt = eq.Arev[vs]
            if eq.VxRev[vsInt] == 0
            # vsInt is not yet in x; store vsInt in x and vs in der_x
                push!(eq.Vx, vsInt)
                eq.VxRev[vsInt] = length(eq.Vx)
                push!(eq.Vderx, vs)
                eq.VderxRev[vs] = length(eq.Vderx)
            end
        else
         # Algebraic or constraint equation -> explicitly solved local equation
            push!(eq.ESolved, vs)
        end
    end

   # Store residue equations in ESorted
    for er in eResidue
        push!(eq.ESorted, er)
        push!(eq.ESolved, 0)

        if highestDerivative
            if eq.EAlgebraic[er]
                cat = RA
            else
                cat = RD
            end
        else
            cat = eq.Brev[er] == 0 ? RDER0 : RDER1
        end
        push!(eq.rcat[Int((cat))], length(eq.ESorted))

        if cat == RDER1
         # Define mue variable associated with equation er
            push!(eq.Vmue, er)
        end
    end

   # Use tearing variables in x or in der_x-vector
    for vt in vTear
        if highestDerivative
            # Highest derivative equation
            if eq.VAlgebraic[vt]
            # vt is an algebraic unknown -> vt is part of x
                push!(eq.Vx, vt)
                push!(eq.Vderx, 0)
                eq.VxRev[vt] = length(eq.Vx)
            elseif eq.Arev[vt] == 0
            # vt is not a differentiated variable -> vt is part of lambda (so of der_x)
                push!(eq.Vx, 0)
                push!(eq.Vderx, vt)
                eq.VderxRev[vt] = length(eq.Vderx)
            else
            # vt is a differentiated variable -> Arev(vt) is part of x
                vtInt = eq.Arev[vt]
                if eq.VxRev[vtInt] == 0  # vtInt is not yet in x
                    push!(eq.Vx, vtInt)
                    eq.VxRev[vtInt] = length(eq.Vx)
                    push!(eq.Vderx, vt)
                    eq.VderxRev[vt] = length(eq.Vderx)
                end
            end
        else
         # Lower derivative constraint equation
            push!(eq.Vx, vt)
            eq.VxRev[vt] = length(eq.Vx)
            der_vt = eq.A[vt]
            if der_vt != 0 && eq.A[der_vt] == 0
            # der(vt) is highest derivative variable
                push!(eq.Vderx, der_vt)
                eq.VderxRev[vt] = length(eq.Vderx)
            else
            # der(vt) is not highest derivative variable
                push!(eq.Vderx, 0)
            end
        end
    end

end

"""
    determineAlgebraicProperty!(eq, ec)

Determine whether all equations ec contain only algebraic variables.
If yes, mark this in eq.EAlgebraic and eq.VAlgebraic
"""
function determineAlgebraicProperty!(eq::SortedEquationGraph, ec::Vector{Int})
    if length(ec) == 0
        error("... Internal error in checkAlgebraicProperty!(..) in StateSelection.jl: length(ec) = 0")
    end

    for eci in ec                  # for all equations in ec
        for vci in eq.Gorigin[eci]  # for all variables in equation eci = ec[i]
            if eq.A[vci] > 0 || eq.Arev[vci] > 0
               # vci is not an algebraic variable
                return nothing
            end
        end
    end

    # ec is a set of algebraic variables with algebraic unknowns
    for eci in ec
        eq.EAlgebraic[eci] = true
        for vci in eq.Gorigin[eci]
            eq.VAlgebraic[vci] = true
        end
    end
    return nothing
end

isAlgebraic(v::Int, A::Vector{Int}, Arev::Vector{Int}) = A[v] == 0 && Arev[v] == 0

function deduceHigher(vLower::Vector{Int}, A)
    v = fill(0, length(vLower))
    for i in eachindex(vLower)
        derv = A[ vLower[i] ]
        if derv == 0
            error("... Internal error in StateSelection.jl: vLower/eLower = ", vLower[i], " has no derivative.")
        end
        v[i] = derv
    end
    return v
end

function getSortedEquationGraph(G, Gsolvable, BLT, assign, A, B, VNames)
    eqInit = StateSelection.SortedEquationGraph(G, BLT, assign, A, B, VNames)
    eq = getSortedEquationGraph!(eqInit, Gsolvable)
    return eq
end

"""
    sortedEquationGraph = getSortedEquationGraph(eqInit::SortedEquationGraph,Gsolvable)

Constructs all information for the sorted equation graph.

Input arguments:

- `eqInit`: Initialized data structure instantiated with SortedEquationGraph()
- `Gsolvable`: Subset of eqInit.Gunknowns that contains the unknowns that can be solved for (used for tearing)

Output arguments in sortedEquationGraph:

- `Vx      `: Variable j = Vx[i] is part of vector x; if j=0, this element of x is a dummy variable (do not use in model code).
- `VxRev   `: If Vxrev[j] = i, then Vx[i] = j; if Vxref[j] = 0, then variable j is not part of x.
- `Vderx   `: Variable j = Vderx[i] is part of vector der_x; if j=0, this element of der_x is a dummy variable (not used).
- `VderxRev`: If VderxRev[j] = i, then Vderx[i] = j. If VderxRef[j] = 0, then variable j is not part of der_x.
- `Er      `: Equation j = Er[i] is part of residue vector r; if j=0, this element is part of ider0n2/ider1n1
- `ider0n2 `: Index vector, such that _residue[1:length(ider0n2)] = _der_x[ider0n2] - _x[ider1n1].
- `ider1n1 `: Index vector, such that _residue[1:length(ider0n2)] = _der_x[ider0n2] - _x[ider1n1].
- `ESorted `: Sorted equations (equations must be generated in the order ESorted[1],[2],..).
- `ESolved `: If ESolved[i] > 0 then equation ESorted[i] is explicitly solved for variable ESolved[i].
              If ESolved[i] == 0, ESorted[i] is a residue equation.
- `nc      `: Number of residue constraints, so part that depends on (x,t) but not on der(x); 0 <= nc <= nx
- `nmue    `: Number of mue variables
"""
function getSortedEquationGraph!(eq::SortedEquationGraph, Gsolvable)
    if !eq.first
        error("... getSortedEquationGraph(..) of StateSelection.jl is called twice with the same sortedEquationGraph object.\n",
            "    This is not possible. You need to instanciate SortedEquationGraph once for every getSortedEquationGraph(..) call.")
    else
        eq.first = false
    end

    eq.Gsolvable = Gsolvable
    td = eq.td      # tearing data structure
    eSolved  = Int[]
    vSolved  = Int[]
    eResidue = Int[]
    vTear    = Int[]
    eSolvedFixedHighest  = Int[]
    vSolvedFixedHighest  = Int[]
    eResidueFixedHighest = Int[]
    vTearFixedHighest    = Int[]
    eConstraintsHighest  = Int[]
    vConstraintsHighest  = Int[]

    # Inspect every BLT component c in sequence
    for j in eachindex(eq.eConstraintsVec)
        eConstraints = eq.eConstraintsVec[j]
        vConstraints = eq.vConstraintsVec[j]

       # Analyze all equation sets eConstraints from lowest-order to highest-order derivatives
        for i in eachindex(eConstraints)
            # println("... i = ", i, ", eConstraints[",i,"] = ", eConstraints[i], ", vConstraints = ", vConstraints[i])
            if i > 1
            # eConstraints[i] is the derivative of eConstraints[i-1] + potentially additional equations
                eSolvedFixed = deduceHigher(eSolved, eq.B)
                vSolvedFixed = deduceHigher(vSolved, eq.A)
                vTearFixed   = deduceHigher(vTear, eq.A)
            end

            if i == length(eConstraints)
            # Highest derivative equations
                if i > 1
                    # Part of the equations are differentiated
                    if length(eResidue) > 0
                        # Remove differentiated residues equations of eConstraints[i-1]) from the highest derivative equations
                        eResidueFixed = deduceHigher(eResidue, eq.B)
                        eConstraints[i] = setdiff(eConstraints[i], eResidueFixed)
                    end
                    append!(eConstraintsHighest, eConstraints[i])
                    append!(vConstraintsHighest, vConstraints[i])
                    append!(eSolvedFixedHighest, eSolvedFixed)
                    append!(vSolvedFixedHighest, vSolvedFixed)
                    append!(vTearFixedHighest, vTearFixed)
                else
               # Original, undifferentiated equations
                    determineAlgebraicProperty!(eq, eConstraints[i])
                    append!(eConstraintsHighest, eConstraints[i])
                    append!(vConstraintsHighest, vConstraints[i])
                end

            else
                # Constraint equations, but not on highest level
                if i == 1
                    # Lowest derivative constraints
                    (eSolved, vSolved, eResidue, vTear) = tearEquations!(td, Gsolvable, eConstraints[i], vConstraints[i])
                else
                    # Higher derivative constraint (but not highest level)
                    (eSolved, vSolved, eResidue, vTear) = tearEquations!(td, Gsolvable, eConstraints[i], vConstraints[i];
                                                                    eSolvedFixed=eSolvedFixed,
                                                                    vSolvedFixed=vSolvedFixed,
                                                                    vTearFixed=vTearFixed)
                end
                appendToSortedEquations!(eq, eSolved, vSolved, eResidue, vTear, false)
            end
        end
    end

   # Tear equations on highest derivative level
    (eSolved, vSolved, eResidue, vTear) = tearEquations!(td, Gsolvable, eConstraintsHighest, vConstraintsHighest;
                                                        eSolvedFixed=eSolvedFixedHighest,
                                                        vSolvedFixed=vSolvedFixedHighest,
                                                        vTearFixed=vTearFixedHighest)

   # Append equations to sorted equations
    appendToSortedEquations!(eq, eSolved, vSolved, eResidue, vTear, true)

   # Determine index vectors ider0n2, ider1n1
    eq.ider0n2 = fill(0, 0)
    eq.ider1n1 = fill(0, 0)
    for vx in eq.Vx
        if vx > 0
            der_vx = eq.A[vx]

            if der_vx > 0 && eq.A[der_vx] > 0
            # der(der(vx)) exists -> der(vx) is within x
                push!(eq.ider0n2, vx)
                push!(eq.ider1n1, der_vx)
            end
        end
    end

    eq.Er = fill(0, length(eq.ider0n2))
    append!(eq.Er, eq.rcat[Int(RD)])
    append!(eq.Er, eq.rcat[Int(RA)])
    append!(eq.Er, eq.rcat[Int(RDER0)])
    append!(eq.Er, eq.rcat[Int(RDER1)])

    for i = length(eq.ider0n2) + 1:length(eq.Er)
        eq.ESolved[ eq.Er[i] ] = -i
    end

    eq.nc = length(eq.rcat[Int(RA)]) +
           length(eq.rcat[Int(RDER0)]) +
           length(eq.rcat[Int(RDER1)])

    return eq
end

#============================ Functions for debugging/testing ================================#
"""
    Ader = derAssociation(A,N)

Determine derivatives of A[i] with respect to N base variables/equation
"""
function derAssociation(A, N)
   # Ader[i,1]: Base variable/equation
   # Ader[i,2]: Derivative of Base variable/equation
    Ader = [1:length(A) fill(0, length(A))]
    for i = 1:N
        der_v  = A[i]
        der_nr = 1
        while der_v > 0
            Ader[der_v,1] = i
            Ader[der_v,2] = der_nr
            der_nr += 1
            der_v  = A[der_v]
        end
    end
    return Ader
end

function printEquation(ec, Bder)
    der_nr = Bder[ec,2]
    if der_nr == 0
        print("eq.", ec)
    else
        if der_nr == 1
            print("eq.", ec, " = der(eq.", Bder[ec,1], ")")
        else
            print("eq.", ec, " = der", der_nr, "(eq.", Bder[ec,1], ")")
        end
    end
    return nothing
end

"""
    printSortedEquationGraph(eq)

Print information about the sorted equation graph
"""
function printSortedEquationGraph(eq::SortedEquationGraph)
   # Determine original number of equations (that are not differentiated)
    neqOrig = 0
    i = 0
    while i < length(eq.Brev)
        i += 1
        if eq.Brev[i] == 0
            neqOrig += 1
        else
            break
        end
    end

    # Build information about original equation derivatives
    Bder = derAssociation(eq.B, neqOrig)

    # Print variables stored in x vector
    println("\n  Variables of _x vector (length=", length(eq.Vx) + length(eq.Vmue), "):")
    for i in eachindex(eq.Vx)
        vx = eq.Vx[i]
        print("     _x[", i, "]: ")
        if vx != 0
            print(eq.VNames[vx])
            if isAlgebraic(vx, eq.A, eq.Arev)
                print("      # algebraic variable")
            elseif eq.Arev[vx] > 0
                vxInt_index = eq.VxRev[ eq.Arev[vx] ]
                if vxInt_index != 0
                    print("      # = der(_x[", vxInt_index, "])")
                else
                    print("      # error (integral ", eq.VNames[eq.Arev[vx]], " not stored in _x)")
                end
            end
        else
            print("---      # integral of lambda variable")
        end
        print("\n")
    end
    nvx = length(eq.Vx)
    for i in eachindex(eq.Vmue)
        println("     _x[", nvx + i, "]: ---      # integral of mue variable")
    end

    # Print variables stored in der_x vector
    println("\n  Variables of _der_x vector (length=", length(eq.Vderx) + length(eq.Vmue), "):")
    for i in eachindex(eq.Vderx)
        vderx = eq.Vderx[i]
        print("     _der_x[", i, "]: ")
        if vderx != 0
            print(eq.VNames[vderx])
            if isAlgebraic(vderx, eq.A, eq.Arev)
                print("     # lambda variable")
            end
        else
            if isAlgebraic(eq.Vx[i], eq.A, eq.Arev)
                print("---      # derivative of algebraic variable")
            else
                vx = eq.Vx[i]
                dervx = eq.A[vx]
                if dervx == 0
                    print("---      # error: vx = ", eq.VNames[vx], ", but der(vx) is not defined")
                else
                    der_vx_index = eq.VxRev[ dervx ]
                    if der_vx_index != 0
                        print("---      # = _x[", der_vx_index, "] = ", eq.VNames[ eq.Vx[der_vx_index] ])
                    else
                        print("---      # error")
                    end
                end
            end
        end
        print("\n")
    end
    nvderx = length(eq.Vderx)
    for i in eachindex(eq.Vmue)
        emue = eq.Vmue[i]
        print("     _der_x[", nvderx + i, "]: ---      # mue variable associated with equation ")
        printEquation(emue, Bder)
        print("\n")
    end

    # Print sorted equations
    println("\n  Sorted equations (length(_r) = ", length(eq.Er), ", nc = ", eq.nc, "):")

    for i in eachindex(eq.ider0n2)
        iderx = eq.VxRev[ eq.ider0n2[i] ]
        ix    = eq.VxRev[ eq.ider1n1[i] ]
        println("     _r[", i, "]   = _der_x[", iderx, "] - _x[", ix, "]")
    end

    for i in eachindex(eq.ESorted)
        es = eq.ESorted[i]
        ev = eq.ESolved[i]
        if ev > 0  # local equation
            print("     ", eq.VNames[ev], "   = < solved from ")
        else # residue equation
            print("     _r[", -ev, "]   = < residue of ")
        end
        printEquation(es, Bder)
        print(" >\n")
    end

   # Raise error, if dimensions do not agree
    if length(eq.Vx) + length(eq.Vmue) != length(eq.Er)
        error("... length(_x) != length(_r)")
    end
    if eq.nc > length(eq.Er)
        error("... nc > length(Er)")
    end

end

end


text:
#########################
# assignment interface #
#########################

"""
    get_address_schema(T)

Return a shallow, compile-time address schema.

    Base.isempty(assignment)

Are there any primitive random choice anywhere in the hierarchy?

    has_internal_node(assignment, addr)

    get_internal_node(assignment, addr)

    get_internal_nodes(assignment)

Return an iterator over top-level internal nodes

    has_leaf_node(assignment, addr)

    get_leaf_node(assignment, addr)

    get_leaf_nodes(assignment)

Return an iterator over top-level leaf nodes

    Base.haskey(assignment, addr)

Alias for has_leaf_node

    Base.getindex(assignment, addr)

Alias for get_leaf_node

    merge(assignment1, assignment2)

Merge two assignments.
"""
abstract type Assignment end

has_internal_node(assignment::Assignment, addr) = false
get_internal_node(assignment::Assignment, addr) = throw(KeyError(addr))
has_leaf_node(assignment::Assignment, addr) = false
get_leaf_node(assignment::Assignment, addr) = throw(KeyError(addr))
Base.haskey(assignment::Assignment, addr) = has_leaf_node(assignment, addr)
Base.getindex(assignment::Assignment, addr) = get_leaf_node(assignment, addr)

function _has_leaf_node(assignment::T, addr::Pair) where {T <: Assignment}
    (first, rest) = addr
    if has_internal_node(assignment, first)
        node = get_internal_node(assignment, first)
        has_leaf_node(node, rest)
    else
        false
    end
end

function _get_leaf_node(assignment::T, addr::Pair) where {T <: Assignment}
    (first, rest) = addr
    node = get_internal_node(assignment, first)
    get_leaf_node(node, rest)
end

function _has_internal_node(assignment::T, addr::Pair) where {T <: Assignment}
    (first, rest) = addr
    if has_internal_node(assignment, first)
        node = get_internal_node(assignment, first)
        has_internal_node(node, rest)
    else
        false
    end
end

function _get_internal_node(assignment::T, addr::Pair) where {T <: Assignment}
    (first, rest) = addr
    node = get_internal_node(assignment, first)
    get_internal_node(node, rest)
end

function _print(io::IO, assignment::Assignment, pre, vert_bars::Tuple)
    VERT = '\u2502'
    PLUS = '\u251C'
    HORZ = '\u2500'
    LAST = '\u2514'
    indent_vert = vcat(Char[' ' for _ in 1:pre], Char[VERT, '\n'])
    indent_vert_last = vcat(Char[' ' for _ in 1:pre], Char[VERT, '\n'])
    indent = vcat(Char[' ' for _ in 1:pre], Char[PLUS, HORZ, HORZ, ' '])
    indent_last = vcat(Char[' ' for _ in 1:pre], Char[LAST, HORZ, HORZ, ' '])
    for i in vert_bars
        indent_vert[i] = VERT
        indent[i] = VERT
        indent_last[i] = VERT
    end
    indent_vert_str = join(indent_vert)
    indent_vert_last_str = join(indent_vert_last)
    indent_str = join(indent)
    indent_last_str = join(indent_last)
    leaf_nodes = collect(get_leaf_nodes(assignment))
    internal_nodes = collect(get_internal_nodes(assignment))
    n = length(leaf_nodes) + length(internal_nodes)
    cur = 1
    for (key, value) in leaf_nodes
        print(io, indent_vert_str)
        print(io, (cur == n ? indent_last_str : indent_str) * "$(repr(key)) : $value\n")
        cur += 1
    end
    for (key, node) in internal_nodes
        print(io, indent_vert_str)
        print(io, (cur == n ? indent_last_str : indent_str) * "$(repr(key))\n")
        _print(io, node, pre + 4, cur == n ? (vert_bars...,) : (vert_bars..., pre+1))
        cur += 1
    end
end

function Base.print(io::IO, assignment::Assignment)
    _print(io, assignment, 0, ())
end

export Assignment
export get_address_schema
export has_internal_node
export get_internal_node
export get_internal_nodes
export has_leaf_node
export get_leaf_node
export get_leaf_nodes

# assignments that have static address schemas should also support faster
# accessors, which make the address explicit in the type (Val(:foo) instaed of
# :foo)
function static_get_leaf_node end
function static_get_internal_node end
export static_get_leaf_node
export static_get_internal_node

function _fill_array! end
function _from_array end

function to_array(assignment::Assignment, ::Type{T}) where {T}
    arr = Vector{T}(undef, 32)
    n = _fill_array!(assignment, arr, 0)
    @assert n <= length(arr)
    resize!(arr, n)
    arr
end

function from_array(proto_assignment::Assignment, arr)
    (n, assignment) = _from_array(proto_assignment, arr, 0)
    if n != length(arr)
        error("Dimension mismatch: $n, $(length(arr))")
    end
    assignment
end

export to_array, from_array

function Base.merge(assignment1::Assignment, assignment2::Assignment)
    assignment = DynamicAssignment()
    for (key, value) in get_leaf_nodes(assignment1)
        assignment.leaf_nodes[key] = value
    end
    for (key, node1) in get_internal_nodes(assignment1)
        if has_internal_node(assignment2, key)
            node2 = get_internal_node(assignment2, key)
            node = merge(node1, node2)
        else
            node = node1
        end
        assignment.internal_nodes[key] = node
    end
    for (key, value) in get_leaf_nodes(assignment2)
        if haskey(assignment.leaf_nodes, key)
            error("assignment1 has leaf node at $key and assignment2 has leaf node at $key")
        end
        if haskey(assignment.internal_nodes, key)
            error("assignment1 has internal node at $key and assignment2 has leaf node at $key")
        end
        assignment.leaf_nodes[key] = value
    end
    for (key, node) in get_internal_nodes(assignment2)
        if haskey(assignment.leaf_nodes, key)
            error("assignment1 has leaf node at $key and assignment2 has internal node at $key")
        end
        if !haskey(assignment.internal_nodes, key)
            # otherwise it should already be included
            assignment.internal_nodes[key] = node
        end
    end
    return assignment
end

######################
# static assignment #
######################

struct StaticAssignment{R,S,T,U} <: Assignment
    leaf_nodes::NamedTuple{R,S}
    internal_nodes::NamedTuple{T,U}
end

# TODO invariant: all internal_nodes are nonempty, but this is not verified at construction time

function get_address_schema(::Type{StaticAssignment{R,S,T,U}}) where {R,S,T,U}
    leaf_keys = Set{Symbol}()
    internal_keys = Set{Symbol}()
    for (key, _) in zip(R, S.parameters)
        push!(leaf_keys, key)
    end
    for (key, _) in zip(T, U.parameters)
        push!(internal_keys, key)
    end
    StaticAddressSchema(leaf_keys, internal_keys)
end

function Base.isempty(assignment::StaticAssignment)
    length(assignment.leaf_nodes) == 0 && length(assignment.internal_nodes) == 0
end

get_leaf_nodes(assignment::StaticAssignment) = pairs(assignment.leaf_nodes)
get_internal_nodes(assignment::StaticAssignment) = pairs(assignment.internal_nodes)
has_leaf_node(assignment::StaticAssignment, addr::Pair) = _has_leaf_node(assignment, addr)
get_leaf_node(assignment::StaticAssignment, addr::Pair) = _get_leaf_node(assignment, addr)
has_internal_node(assignment::StaticAssignment, addr::Pair) = _has_internal_node(assignment, addr)
get_internal_node(assignment::StaticAssignment, addr::Pair) = _get_internal_node(assignment, addr)

# NOTE: there is no static_has_internal_node or static_has_leaf_node because
# this is known from the static address schema

## has_internal_node ##

function has_internal_node(assignment::StaticAssignment, key::Symbol)
    haskey(assignment.internal_nodes, key)
end

## has_leaf_node ##

function has_leaf_node(assignment::StaticAssignment, key::Symbol)
    haskey(assignment.leaf_nodes, key)
end

## get_internal_node ##

function get_internal_node(assignment::StaticAssignment, key::Symbol)
    assignment.internal_nodes[key]
end

function static_get_internal_node(assignment::StaticAssignment, ::Val{A}) where {A}
    assignment.internal_nodes[A]
end

## get_leaf_node ##

function get_leaf_node(assignment::StaticAssignment, key::Symbol)
    assignment.leaf_nodes[key]
end

function static_get_leaf_node(assignment::StaticAssignment, ::Val{A}) where {A}
    assignment.leaf_nodes[A]
end

# convert from any other schema that has only Val{:foo} addresses
function StaticAssignment(other::Assignment)
    leaf_keys_and_nodes = collect(get_leaf_nodes(other))
    internal_keys_and_nodes = collect(get_internal_nodes(other))
    if length(leaf_keys_and_nodes) > 0
        (leaf_keys, leaf_nodes) = collect(zip(leaf_keys_and_nodes...))
    else
        (leaf_keys, leaf_nodes) = ((), ())
    end
    if length(internal_keys_and_nodes) > 0
        (internal_keys, internal_nodes) = collect(zip(internal_keys_and_nodes...))
    else
        (internal_keys, internal_nodes) = ((), ())
    end
    StaticAssignment(
        NamedTuple{leaf_keys}(leaf_nodes),
        NamedTuple{internal_keys}(internal_nodes))
end

function pair(a, b, key1::Symbol, key2::Symbol)
    StaticAssignment(NamedTuple(), NamedTuple{(key1,key2)}((a, b)))
end

function unpair(assignment, key1::Symbol, key2::Symbol)
    if length(get_leaf_nodes(assignment)) != 0 || length(get_internal_nodes(assignment)) > 2
        error("Not a pair")
    end
    a = has_internal_node(assignment, key1) ? get_internal_node(assignment, key1) : EmptyAssignment()
    b = has_internal_node(assignment, key2) ? get_internal_node(assignment, key2) : EmptyAssignment()
    (a, b)
end

# TODO make it a generated function?
function _fill_array!(assignment::StaticAssignment, arr::Vector{T}, start_idx::Int) where {T}
    if length(arr) < start_idx + length(assignment.leaf_nodes)
        resize!(arr, 2 * length(arr))
    end
    for (i, value) in enumerate(assignment.leaf_nodes)
        arr[start_idx + i] = value
    end
    idx = start_idx + length(assignment.leaf_nodes)
    for (i, node) in enumerate(assignment.internal_nodes)
        n_written = _fill_array!(node, arr, idx)
        idx += n_written
    end
    idx - start_idx
end

@generated function _from_array(proto_assignment::StaticAssignment{R,S,T,U}, arr::Vector{V}, start_idx::Int) where {R,S,T,U,V}
    leaf_node_keys = proto_assignment.parameters[1]
    leaf_node_types = proto_assignment.parameters[2].parameters
    internal_node_keys = proto_assignment.parameters[3]
    internal_node_types = proto_assignment.parameters[4].parameters

    # leaf nodes
    leaf_node_refs = Expr[]
    for (i, key) in enumerate(leaf_node_keys)
        expr = Expr(:ref, :arr, Expr(:call, :(+), :start_idx, QuoteNode(i)))
        push!(leaf_node_refs, expr)
    end

    # internal nodes
    internal_node_blocks = Expr[]
    internal_node_names = Symbol[]
    for (key, typ) in zip(internal_node_keys, internal_node_types)
        node = gensym()
        push!(internal_node_names, node)
        push!(internal_node_blocks, quote
            (n_read, $node::$typ) = _from_array(proto_assignment.internal_nodes.$key, arr, idx)
            idx += n_read
        end)
    end

    quote
        n_read = $(QuoteNode(length(leaf_node_keys)))
        leaf_nodes_field = NamedTuple{R,S}(($(leaf_node_refs...),))
        idx::Int = start_idx + n_read
        $(internal_node_blocks...)
        internal_nodes_field = NamedTuple{T,U}(($(internal_node_names...),))
        assignment = StaticAssignment{R,S,T,U}(leaf_nodes_field, internal_nodes_field)
        (idx - start_idx, assignment)
    end
end

@generated function Base.merge(assignment1::StaticAssignment{R,S,T,U},
                               assignment2::StaticAssignment{W,X,Y,Z}) where {R,S,T,U,W,X,Y,Z}

    # unpack first assignment type parameters
    leaf_node_keys1 = assignment1.parameters[1]
    leaf_node_types1 = assignment1.parameters[2].parameters
    internal_node_keys1 = assignment1.parameters[3]
    internal_node_types1 = assignment1.parameters[4].parameters
    keys1 = (leaf_node_keys1..., internal_node_keys1...,)

    # unpack second assignment type parameters
    leaf_node_keys2 = assignment2.parameters[1]
    leaf_node_types2 = assignment2.parameters[2].parameters
    internal_node_keys2 = assignment2.parameters[3]
    internal_node_types2 = assignment2.parameters[4].parameters
    keys2 = (leaf_node_keys2..., internal_node_keys2...,)

    # leaf vs leaf collision is an error
    colliding_leaf_leaf_keys = intersect(leaf_node_keys1, leaf_node_keys2)
    if !isempty(colliding_leaf_leaf_keys)
        error("assignment1 and assignment2 both have leaf nodes at key(s): $colliding_leaf_leaf_keys")
    end

    # leaf vs internal collision is an error
    colliding_leaf_internal_keys = intersect(leaf_node_keys1, internal_node_keys2)
    if !isempty(colliding_leaf_internal_keys)
        error("assignment1 has leaf node and assignment2 has internal node at key(s): $colliding_leaf_internal_keys")
    end

    # internal vs leaf collision is an error
    colliding_internal_leaf_keys = intersect(internal_node_keys1, leaf_node_keys2)
    if !isempty(colliding_internal_leaf_keys)
        error("assignment1 has internal node and assignment2 has leaf node at key(s): $colliding_internal_leaf_keys")
    end

    # internal vs internal collision is not an error, recursively call merge
    colliding_internal_internal_keys = (intersect(internal_node_keys1, internal_node_keys2)...,)
    internal_node_keys1_exclusive = (setdiff(internal_node_keys1, internal_node_keys2)...,)
    internal_node_keys2_exclusive = (setdiff(internal_node_keys2, internal_node_keys1)...,)

    # leaf nodes named tuple
    leaf_node_keys = (leaf_node_keys1..., leaf_node_keys2...,)
    leaf_node_types = map(QuoteNode, (leaf_node_types1..., leaf_node_types2...,))
    leaf_node_values = Expr(:tuple,
        [Expr(:(.), :(assignment1.leaf_nodes), QuoteNode(key))
            for key in leaf_node_keys1]...,
        [Expr(:(.), :(assignment2.leaf_nodes), QuoteNode(key))
            for key in leaf_node_keys2]...)
    leaf_nodes = Expr(:call,
        Expr(:curly, :NamedTuple,
            QuoteNode(leaf_node_keys),
            Expr(:curly, :Tuple, leaf_node_types...)),
        leaf_node_values)

    # internal nodes named tuple
    internal_node_keys = (internal_node_keys1_exclusive...,
                          internal_node_keys2_exclusive...,
                          colliding_internal_internal_keys...)
    internal_node_values = Expr(:tuple,
        [Expr(:(.), :(assignment1.internal_nodes), QuoteNode(key))
            for key in internal_node_keys1_exclusive]...,
        [Expr(:(.), :(assignment2.internal_nodes), QuoteNode(key))
            for key in internal_node_keys2_exclusive]...,
        [Expr(:call, :merge,
                Expr(:(.), :(assignment1.internal_nodes), QuoteNode(key)),
                Expr(:(.), :(assignment2.internal_nodes), QuoteNode(key)))
            for key in colliding_internal_internal_keys]...)
    internal_nodes = Expr(:call,
        Expr(:curly, :NamedTuple, QuoteNode(internal_node_keys)),
        internal_node_values)

    # construct assignment from named tuples
    Expr(:call, :StaticAssignment, leaf_nodes, internal_nodes)
end

export StaticAssignment
export pair, unpair

#######################
# dynamic assignment #
#######################

struct DynamicAssignment <: Assignment
    leaf_nodes::Dict{Any,Any}
    internal_nodes::Dict{Any,Any}
end

DynamicAssignment() = DynamicAssignment(Dict(), Dict())
get_address_schema(::Type{DynamicAssignment}) = DynamicAddressSchema()

# invariant: all internal nodes are nonempty
Base.isempty(assignment::DynamicAssignment) = isempty(assignment.leaf_nodes) && isempty(assignment.internal_nodes)
get_leaf_nodes(assignment::DynamicAssignment) = assignment.leaf_nodes
get_internal_nodes(assignment::DynamicAssignment) = assignment.internal_nodes
has_leaf_node(assignment::DynamicAssignment, addr::Pair) = _has_leaf_node(assignment, addr)
get_leaf_node(assignment::DynamicAssignment, addr::Pair) = _get_leaf_node(assignment, addr)
has_internal_node(assignment::DynamicAssignment, addr::Pair) = _has_internal_node(assignment, addr)
get_internal_node(assignment::DynamicAssignment, addr::Pair) = _get_internal_node(assignment, addr)

function Base.values(assignment::DynamicAssignment)
    iterators::Vector = collect(map(values, assignment.internal_nodes))
    push!(iterators, values(assignment.leaf_nodes))
    Iterators.flatten(iterators)
end

function has_internal_node(assignment::DynamicAssignment, addr)
    haskey(assignment.internal_nodes, addr)
end

function get_internal_node(assignment::DynamicAssignment, addr)
    assignment.internal_nodes[addr]
end

function has_leaf_node(assignment::DynamicAssignment, addr)
    haskey(assignment.leaf_nodes, addr)
end

function get_leaf_node(assignment::DynamicAssignment, addr)
    assignment.leaf_nodes[addr]
end

# mutation (not part of the assignment interface)

function set_leaf_node!(assignment::DynamicAssignment, addr, value)
    delete!(assignment.internal_nodes, addr)
    assignment.leaf_nodes[addr] = value
end

function set_leaf_node!(assignment::DynamicAssignment, addr::Pair, value)
    (first, rest) = addr
    if haskey(assignment.leaf_nodes, first)
        # we are not writing to the address directly, so we error instead of
        # delete the existing node.
        error("Tried to create assignment at $first but there was already a value there.")
    end
    if haskey(assignment.internal_nodes, first)
        node = assignment.internal_nodes[first]
    else
        node = DynamicAssignment()
        assignment.internal_nodes[first] = node
    end
    node = assignment.internal_nodes[first]
    set_leaf_node!(node, rest, value)
end

function set_internal_node!(assignment::DynamicAssignment, addr, new_node)
    delete!(assignment.leaf_nodes, addr)
    delete!(assignment.internal_nodes, addr)
    if !isempty(new_node)
        assignment.internal_nodes[addr] = new_node
    end
end

function set_internal_node!(assignment::DynamicAssignment, addr::Pair, new_node)
    (first, rest) = addr
    if haskey(assignment.leaf_nodes, first)
        # we are not writing to the address directly, so we error instead of
        # delete the existing node.
        error("Tried to create assignment at $first but there was already a value there.")
    end
    if haskey(assignment.internal_nodes, first)
        node = assignment.internal_nodes[first]
    else
        node = DynamicAssignment()
        assignment.internal_nodes[first] = node
    end
    set_internal_node!(node, rest, new_node)
end

Base.setindex!(assignment::DynamicAssignment, value, addr) = set_leaf_node!(assignment, addr, value)

# b should implement the choice-assignment interface
function Base.merge!(a::DynamicAssignment, b)
    for (key, value) in get_leaf_nodes(b)
        a.leaf_nodes[key] = value
    end
    for (key, b_node) in get_internal_nodes(b)
        @assert !isempty(b_node)
        if haskey(a.internal_nodes, key)
            # NOTE: error if a.internal_nodes[key] does not implement merge!
            @assert !isempty(a.internal_nodes[key])
            merge!(a.internal_nodes[key], b_node)
        else
            a.internal_nodes[key] = b_node
        end
    end
    a
end

function _fill_array!(assignment::DynamicAssignment, arr::Vector{T}, start_idx::Int) where {T}
    if length(arr) < start_idx + length(assignment.leaf_nodes)
        resize!(arr, 2 * length(arr))
    end
    leaf_keys_sorted = sort(collect(keys(assignment.leaf_nodes)))
    internal_node_keys_sorted = sort(collect(keys(assignment.internal_nodes)))
    for (i, key) in enumerate(leaf_keys_sorted)
        arr[start_idx + i] = assignment.leaf_nodes[key]
    end
    idx = start_idx + length(assignment.leaf_nodes)
    for key in internal_node_keys_sorted
        n_written = _fill_array!(get_internal_node(assignment, key), arr, idx)
        idx += n_written
    end
    idx - start_idx
end

function _from_array(proto_assignment::DynamicAssignment, arr::Vector{T}, start_idx::Int) where {T}
    @assert length(arr) >= start_idx
    assignment = DynamicAssignment()
    leaf_keys_sorted = sort(collect(keys(proto_assignment.leaf_nodes)))
    internal_node_keys_sorted = sort(collect(keys(proto_assignment.internal_nodes)))
    for (i, key) in enumerate(leaf_keys_sorted)
        assignment.leaf_nodes[key] = arr[start_idx + i]
    end
    idx = start_idx + length(assignment.leaf_nodes)
    for key in internal_node_keys_sorted
        (n_read, node) = _from_array(get_internal_node(proto_assignment, key), arr, idx)
        idx += n_read
        assignment.internal_nodes[key] = node
    end
    (idx - start_idx, assignment)
end

export DynamicAssignment

#######################################
## vector combinator for assignments #
#######################################

# TODO implement LeafVectorAssignment, which stores a vector of leaf nodes

struct InternalVectorAssignment{T} <: Assignment
    internal_nodes::Vector{T}
    is_empty::Bool
end

# note some internal nodes may be empty

get_address_schema(::Type{InternalVectorAssignment}) = VectorAddressSchema()

function vectorize_internal(nodes::Vector{T}) where {T}
    is_empty = all(map(isempty, nodes))
    InternalVectorAssignment(nodes, is_empty)
end

Base.isempty(assignment::InternalVectorAssignment) = assignment.is_empty

has_leaf_node(assignment::InternalVectorAssignment, addr::Pair) = _has_leaf_node(assignment, addr)
get_leaf_node(assignment::InternalVectorAssignment, addr::Pair) = _get_leaf_node(assignment, addr)
has_internal_node(assignment::InternalVectorAssignment, addr::Pair) = _has_internal_node(assignment, addr)
get_internal_node(assignment::InternalVectorAssignment, addr::Pair) = _get_internal_node(assignment, addr)

function has_internal_node(assignment::InternalVectorAssignment, addr::Int)
    n = length(assignment.internal_nodes)
    addr >= 1 && addr <= n
end

function get_internal_node(assignment::InternalVectorAssignment, addr::Int)
    assignment.internal_nodes[addr]
end

function get_internal_nodes(assignment::InternalVectorAssignment)
    # TODO need to check they are not empty
    ((i, assignment.internal_nodes[i]) for i=1:length(assignment.internal_nodes))
end

get_leaf_nodes(assignment::InternalVectorAssignment) = ()

function _fill_array!(assignment::InternalVectorAssignment, arr::Vector{T}, start_idx::Int) where {T}
    idx = start_idx
    for key=1:length(assignment.internal_nodes)
        n = _fill_array!(assignment.internal_nodes[key], arr, idx)
        idx += n
    end
    idx - start_idx
end

function _from_array(proto_assignment::InternalVectorAssignment{U}, arr::Vector{T}, start_idx::Int) where {T,U}
    @assert length(arr) >= start_idx
    nodes = Vector{U}(undef, length(proto_assignment.internal_nodes))
    idx = start_idx
    for key=1:length(proto_assignment.internal_nodes)
        (n_read, nodes[key]) = _from_array(proto_assignment.internal_nodes[key], arr, idx)
        idx += n_read
    end
    assignment = InternalVectorAssignment(nodes, proto_assignment.is_empty)
    (idx - start_idx, assignment)
end

export InternalVectorAssignment
export vectorize_internal

####################
# empty assignment #
####################

struct EmptyAssignment <: Assignment end

Base.isempty(::EmptyAssignment) = true
get_address_schema(::Type{EmptyAssignment}) = EmptyAddressSchema()
get_internal_nodes(::EmptyAssignment) = ()
get_leaf_nodes(::EmptyAssignment) = ()

_fill_array!(assignment::EmptyAssignment, arr::Vector, start_idx::Int) = 0
_from_array(proto_assignment::EmptyAssignment, arr::Vector, start_idx::Int) = (0, EmptyAssignment())

export EmptyAssignment

max_stars_repo_name: BoltzmannBrain/Gen

max_issues_repo_name: BoltzmannBrain/Gen

max_forks_repo_name: BoltzmannBrain/Gen


text:
using FileIO

function predict_q(w, x; nh=1)
    inp = x
    for i=1:nh
        inp = relu.(w["w_$i"] * inp .+ w["b_$i"])
    end
    q = w["w_out"] * inp .+ w["b_out"]
    return q
end

function init_weights(input, hiddens, nout, atype)
    w = Dict()
    inp = input
    for i=1:length(hiddens)
        w["w_$i"] = xavier(hiddens[i], inp)
        w["b_$i"] = zeros(hiddens[i])
        inp = hiddens[i]
    end

    w["w_out"] = xavier(nout, hiddens[end])
    w["b_out"] = zeros(nout, 1)

    for k in keys(w)
        w[k] = convert(atype, w[k])
    end
    return w
end

function save_model(w, fname)
    tmp = Dict()
    for k in keys(w)
        tmp[k] = convert(Array{Float32}, w[k])
    end
    save(fname, tmp)
end

function load_model(fname, atype)
    w = load(fname)
    for k in keys(w)
        w[k] = convert(atype, w[k])
    end
    return w
end


text:
# Verify discretization technique by simulating income generating process
# and markov chain simulation for T_ret periods and then comparing distributions

using Plots, LinearAlgebra, Statistics, Random, Roots, QuantEcon, Interpolations
using BenchmarkTools

# Simulate income generating process
# Parameters
## Demographics
N = 1000000
T_ret = 10
κ = zeros(T_ret)

# Directly from KV2010
σ_η = 0.01
σ_z0 = 0.15
σ_ε = 0.05

function net_labor_income(N, T_ret, σ_ε, σ_η, σ_z0, κ)
    Y = zeros(Float64, N, T_ret-1)
    z_t1 = sqrt(σ_z0) .* randn(N)

    for t in 1:(T_ret-1)
        z_t = z_t1 + sqrt(σ_η) .* randn(N)
        Y[:,t] = exp.( κ[t] .+ z_t + sqrt(σ_ε) .* randn(N) )
        z_t1 = copy(z_t)
    end
    return Y
end

Y_igp = net_labor_income(N, T_ret, σ_ε, σ_η, σ_z0, κ)
histogram(Y_igp[:,1], alpha = 0.3, bins = 100)
histogram!(Y_igp[:,end], alpha = 0.3, bins = 100)
means = mean(Y_igp,dims=1)'
vars = var(Y_igp,dims=1)'
plot(1:(T_ret-1), means, label="means")
plot!(1:(T_ret-1), vars, label="vars", xlabel="t")

# Tauchen discretization
tauc_ε = tauchen(19,0.,sqrt(σ_ε))
tauc_η = tauchen(39,0.99999,sqrt(σ_η))

tauc_ε.p

max_stars_repo_path: test/persistent discretization test.jl

max_issues_repo_path: test/persistent discretization test.jl

max_forks_repo_path: test/persistent discretization test.jl


text:
using Test
include("HelmholtzTest.jl");
include("ShiftedLaplacianTest.jl");
include("AttenuationTest.jl");
include("testFictitiousSource2D.jl")
include("testHybridKacmarzHelmholtz.jl")


text:
using Test
using ITensorsGateEvolution
using ITensors

@testset "apply" begin

  @testset "Simple on-site state evolution" begin
    N = 3

    pos = ProductOps()
    pos *= "Z", 3
    pos *= "Y", 2
    pos *= "X", 1

    s = siteinds("qubit", N)
    gates = ops(s, pos)
    ψ0 = productMPS(s, "0")

    # Apply the gates
    ψ = apply(gates, ψ0)

    # Move site 1 to position 3
    ψ′ = movesite(ψ, 1, 3)
    @test siteind(ψ′, 1) == s[2]
    @test siteind(ψ′, 2) == s[3]
    @test siteind(ψ′, 3) == s[1]
    @test prod(ψ) ≈ prod(ψ′)

    # Move the site back
    ψ′′ = movesite(ψ′, 3, 1)
    @test siteind(ψ′′, 1) == s[1]
    @test siteind(ψ′′, 2) == s[2]
    @test siteind(ψ′′, 3) == s[3]
    @test prod(ψ) ≈ prod(ψ′′)
  end

  @testset "More complex evolution" begin
    N = 7

    osX = ProductOps()
    for n in 1:N
      osX *= "X", n
    end

    osZ = ProductOps()
    for n in 1:N
      osZ *= "Z", n
    end

    osRand = ProductOps()
    for n in 1:N
      osRand *= "rand", n
    end

    osSw = ProductOps()
    for n in 1:N-2
      osSw *= "Sw", n, n+2
    end

    osCx = ProductOps()
    for n in 1:N-3
      osCx *= "Cx", n, n+3
    end

    osT = ProductOps()
    for n in 1:N-3
      osT *= "T", n, n+1, n+3
    end

    osRx = ProductOps()
    for n in 1:N
      osRx *= "Rx", n, (θ = π,)
    end

    os_noise = ProductOps()
    for n in 1:N-4
      os_noise *= "noise", n, n+2, n+4
    end

    os = osRand * osX * osSw * osRx * osZ * osCx * osT
    #os = osX * osSw * osRx * osZ * osCx * osT
    s = siteinds("qubit", N)
    gates = ops(s, os)

    @testset "Pure state evolution" begin
      ψ0 = productMPS(s, "0")
      maxdim = prod(dim(siteinds(ψ0, j)) for j in 1:N÷2)
      ψ = apply(gates, ψ0; cutoff = 1e-15, maxdim = maxdim)
      @test maxlinkdim(ψ) == maxdim
      prodψ = apply(gates, prod(ψ0))
      @test prod(ψ) ≈ prodψ
    end

    M0 = MPO(s, "Id")
    maxdim = prod(dim(siteinds(M0, j)) for j in 1:N÷2)

    @testset "Mixed state evolution" begin
      M = apply(gates, M0; cutoff = 1e-15, maxdim = maxdim)
      @test maxlinkdim(M) == 30
      sM0 = siteinds(M0)
      sM = siteinds(M)
      for n in 1:N
        @test hassameinds(sM[n], sM0[n])
      end
      set_warn_itensor_order!(15)
      prodM = apply(gates, prod(M0))
      @test prod(M) ≈ prodM
      reset_warn_itensor_order!()
    end

    @testset "Mixed state noisy evolution" begin
      os *= os_noise
      gates = ops(s, os)
      M = apply(gates, M0; apply_dag = true,
                           cutoff = 1e-15, maxdim = maxdim)
      @test maxlinkdim(M) == maxdim
      sM0 = siteinds(M0)
      sM = siteinds(M)
      for n in 1:N
        @test hassameinds(sM[n], sM0[n])
      end
      set_warn_itensor_order!(16)
      prodM = apply(gates, prod(M0); apply_dag = true)
      @test prod(M) ≈ prodM
      reset_warn_itensor_order!()
    end

    @testset "Mixed state noisy evolution" begin
      os *= os_noise
      gates = ops(s, os)
      M = apply(gates, M0; apply_dag = true,
                           cutoff = 1e-15, maxdim = maxdim-1)
      @test maxlinkdim(M) == maxdim-1
      sM0 = siteinds(M0)
      sM = siteinds(M)
      for n in 1:N
        @test hassameinds(sM[n], sM0[n])
      end
      set_warn_itensor_order!(16)
      prodM = apply(gates, prod(M0); apply_dag = true)
      @test prod(M) ≈ prodM rtol = 1e-1
      reset_warn_itensor_order!()
    end

  end

  @testset "movesites" begin
    N = 6
    s = siteinds("qubit", N)
    ψ0 = randomMPS(s)
    ψ = ITensorsGateEvolution.movesites(ψ0, 1:N, reverse(1:N))
    for n in 1:N
      @test siteind(ψ, n) == s[N-n+1]
    end
  end

end


text:
include("Tests.jl")
include("PTests.jl")

using Tests: tests
using PTests: ptests

tests()
ptests()


text:
module TalkSetup
using Revise
using Plots
gr()
using Distributions
using AverageShiftedHistograms
using BenchmarkTools
using RCall
using GLM
using DataFrames
using Query
end


text:
# Gloriously inefficient ways of computing running & rolled quantities.

function _naive_inception_reduce(T, f::Function, block::Block, min_window::Int)
    @assert min_window > 0
    times = DateTime[]
    values = T[]
    buffer = value_type(block)[]
    for (i, (time, value)) in enumerate(block)
        # Push a new value onto the buffer.
        push!(buffer, value)
        if i < min_window
            # We need to wait longer before emitting.
            continue
        end
        push!(times, time)
        push!(values, f(buffer))
    end

    return Block(times, values)
end

function _naive_inception_reduce(
    T, f::Function, block_a::Block, block_b::Block, min_window::Int
)
    @assert min_window > 0
    @assert block_a.times == block_b.times  # Don't attempt to perform alignment
    times = DateTime[]
    values = T[]
    buffer_a = value_type(block_a)[]
    buffer_b = value_type(block_b)[]
    for i in 1:length(block_a)
        time, a = block_a[i]
        _, b = block_b[i]
        # Push new values onto the buffer.
        push!(buffer_a, a)
        push!(buffer_b, b)
        if i < min_window
            # We need to wait longer before emitting.
            continue
        end
        push!(times, time)
        push!(values, f(buffer_a, buffer_b))
    end

    return Block(times, values)
end

function _naive_window_reduce(
    T, f::Function, block::Block, window::Int, emit_early::Bool, min_window::Int
)
    @assert window > 0
    @assert min_window > 0
    times = DateTime[]
    values = T[]
    buffer = CircularBuffer{value_type(block)}(window)
    for (i, (time, value)) in enumerate(block)
        # Push a new value onto the buffer. (If the buffer is full, it will overwrite the
        # value that is falling out of the buffer, because it is circular.)
        push!(buffer, value)
        if !emit_early && !isfull(buffer)
            # The buffer is not full, and we require that it be full for us to emit a knot.
            continue
        end
        if i < min_window
            # We need to wait longer before emitting.
            continue
        end
        push!(times, time)
        push!(values, f(buffer))
    end

    return Block(times, values)
end

function _naive_window_reduce(
    T,
    f::Function,
    block_a::Block,
    block_b::Block,
    window::Int,
    emit_early::Bool,
    min_window::Int,
)
    @assert window > 0
    @assert min_window > 0
    @assert block_a.times == block_b.times  # Don't attempt to perform alignment

    times = DateTime[]
    values = T[]
    buffer_a = CircularBuffer{value_type(block_a)}(window)
    buffer_b = CircularBuffer{value_type(block_b)}(window)
    for i in 1:length(block_a)
        time, value_a = block_a[i]
        _, value_b = block_b[i]
        # Push a new value onto the buffer. (If the buffer is full, it will overwrite the
        # value that is falling out of the buffer, because it is circular.)
        push!(buffer_a, value_a)
        push!(buffer_b, value_b)
        if !emit_early && !isfull(buffer_a)
            @assert !isfull(buffer_b)
            # The buffer is not full, and we require that it be full for us to emit a knot.
            continue
        end
        if i < min_window
            # We need to wait longer before emitting.
            continue
        end
        push!(times, time)
        push!(values, f(buffer_a, buffer_b))
    end

    return Block(times, values)
end

function _test_inception_op(
    T, f_normal::Function, f_timedag::Function=f_normal; min_window=1, block::Block=b4
)
    n_in = block_node(block)
    @test _eval(f_timedag(n_in)) ≈ _naive_inception_reduce(T, f_normal, block, min_window)
end

function _test_binary_inception_op(
    T, f_normal::Function, f_timedag::Function=f_normal; min_window=1
)
    # TODO test other alignments
    # pre-align outputs so the reduction operation is simpler
    na, nb = coalign(n1, n4)
    block_a, block_b = _eval_fast([na, nb])
    @test isapprox(
        _eval(f_timedag(n1, n4)),
        _naive_inception_reduce(T, f_normal, block_a, block_b, min_window);
        nans=true,
    )
end

function _test_window_op(
    T, f_normal::Function, f_timedag::Function=f_normal; min_window=1, block::Block=b4
)
    n_in = block_node(block)
    for window in min_window:(length(block) + 2)
        n = f_timedag(n_in, window)

        block_default = _eval(n)
        block_ee_false = _eval(f_timedag(n_in, window; emit_early=false))
        block_ee_true = _eval(f_timedag(n_in, window; emit_early=true))

        @test isequal(block_default, block_ee_false)
        @test isapprox(
            block_ee_false,
            _naive_window_reduce(T, f_normal, block, window, false, min_window);
            nans=true,
        )
        @test isapprox(
            block_ee_true,
            _naive_window_reduce(T, f_normal, block, window, true, min_window);
            nans=true,
        )
    end
    return nothing
end

function _test_binary_window_op(
    T, f_normal::Function, f_timedag::Function=f_normal; min_window=1
)
    na = n4
    nb = TimeDag.align(n1, n4)
    block_a, block_b = _eval_fast([na, nb])
    for window in min_window:(length(b4) + 2)
        n = f_timedag(n1, n4, window)

        block = _eval(n)
        block_ee_false = _eval(f_timedag(n1, n4, window; emit_early=false))
        block_ee_true = _eval(f_timedag(n1, n4, window; emit_early=true))

        @test isequal(block, block_ee_false)
        @test isapprox(
            block_ee_false,
            _naive_window_reduce(T, f_normal, block_a, block_b, window, false, min_window);
            nans=true,
        )
        @test isapprox(
            block_ee_true,
            _naive_window_reduce(T, f_normal, block_a, block_b, window, true, min_window);
            nans=true,
        )
    end
    return nothing
end

@testset "sum" begin
    @testset "inception" begin
        @test sum(constant(42)) == constant(42)
        _test_inception_op(Int64, sum)
    end

    @testset "window" begin
        _test_window_op(Int64, sum)
    end
end

@testset "prod" begin
    @testset "inception" begin
        @test prod(constant(42)) == constant(42)
        _test_inception_op(Int64, prod)
    end

    @testset "window" begin
        _test_window_op(Int64, prod)
    end
end

@testset "mean" begin
    @testset "inception" begin
        @test mean(constant(42)) == constant(42)
        _test_inception_op(Float64, mean)
    end

    @testset "window" begin
        _test_window_op(Float64, mean)
    end
end

@testset "var" begin
    @testset "inception" begin
        @test_throws ArgumentError var(constant(42.0))
        _test_inception_op(Float64, var; min_window=2)
        _test_inception_op(Float64, partial(var; corrected=false); min_window=2)
        _test_inception_op(Float64, partial(var; corrected=true); min_window=2)
    end

    @testset "window" begin
        # Windows that are too small should trigger an exception when attempting to create
        # the node.
        @test_throws ArgumentError var(n4, 0)
        @test_throws ArgumentError var(n4, 1)
        @test_throws ArgumentError var(constant(42.0), 2)

        _test_window_op(Float64, var; min_window=2)
        _test_window_op(Float64, partial(var; corrected=false); min_window=2)
        _test_window_op(Float64, partial(var; corrected=true); min_window=2)
    end
end

@testset "std" begin
    @testset "inception" begin
        @test_throws ArgumentError std(constant(42.0))
        _test_inception_op(Float64, std; min_window=2)
        _test_inception_op(Float64, partial(std; corrected=false); min_window=2)
        _test_inception_op(Float64, partial(std; corrected=true); min_window=2)
    end

    @testset "window" begin
        # Windows that are too small should trigger an exception when attempting to create
        # the node.
        @test_throws ArgumentError std(n4, 0)
        @test_throws ArgumentError std(n4, 1)
        @test_throws ArgumentError std(constant(42.0), 2)

        _test_window_op(Float64, std; min_window=2)
        _test_window_op(Float64, partial(std; corrected=false); min_window=2)
        _test_window_op(Float64, partial(std; corrected=true); min_window=2)
    end
end

@testset "cov" begin
    @testset "inception" begin
        @test_throws ArgumentError cov(constant(42.0), constant(24.0))
        _test_binary_inception_op(Float64, cov; min_window=2)
        _test_binary_inception_op(Float64, partial(cov; corrected=false); min_window=2)
        _test_binary_inception_op(Float64, partial(cov; corrected=true); min_window=2)
    end

    @testset "window" begin
        # Windows that are too small should trigger an exception when attempting to create
        # the node.
        @test_throws ArgumentError cov(n1, n4, 0)
        @test_throws ArgumentError cov(n1, n4, 1)
        @test_throws ArgumentError cov(constant(42.0), constant(24.0), 2)

        _test_binary_window_op(Float64, cov; min_window=2)
        _test_binary_window_op(Float64, partial(cov; corrected=false); min_window=2)
        _test_binary_window_op(Float64, partial(cov; corrected=true); min_window=2)
    end
end

@testset "cov matrix static" begin
    dim = 3
    n_obs = 20
    block = _get_rand_svec_block(MersenneTwister(42), dim, n_obs)
    T = SMatrix{dim,dim,eltype(value_type(block)),dim * dim}

    @testset "inception" begin
        @test_throws ArgumentError cov(constant(SVector((1.0, 2.0, 3.0))))
        _test_inception_op(T, cov; min_window=2, block)
        _test_inception_op(T, partial(cov; corrected=false); min_window=2, block)
        _test_inception_op(T, partial(cov; corrected=true); min_window=2, block)
    end

    @testset "window" begin
        # Windows that are too small should trigger an exception when attempting to create
        # the node.
        @test_throws ArgumentError cov(block_node(block), 0)
        @test_throws ArgumentError cov(block_node(block), 1)
        @test_throws ArgumentError cov(constant(SVector((1, 2, 3))), 2)

        _test_window_op(T, cov; min_window=2, block)
        _test_window_op(T, partial(cov; corrected=false); min_window=2, block)
        _test_window_op(T, partial(cov; corrected=true); min_window=2, block)
    end
end

@testset "cov matrix" begin
    dim = 3
    n_obs = 20
    block = _get_rand_vec_block(MersenneTwister(42), dim, n_obs)
    T = Matrix{eltype(value_type(block))}

    @testset "inception" begin
        @test_throws ArgumentError cov(constant([1.0, 2.0, 3.0]))
        _test_inception_op(T, cov; min_window=2, block)
        _test_inception_op(T, partial(cov; corrected=false); min_window=2, block)
        _test_inception_op(T, partial(cov; corrected=true); min_window=2, block)
    end

    @testset "window" begin
        # Windows that are too small should trigger an exception when attempting to create
        # the node.
        @test_throws ArgumentError cov(block_node(block), 0)
        @test_throws ArgumentError cov(block_node(block), 1)
        @test_throws ArgumentError cov(constant([1, 2, 3]), 2)

        _test_window_op(T, cov; min_window=2, block)
        _test_window_op(T, partial(cov; corrected=false); min_window=2, block)
        _test_window_op(T, partial(cov; corrected=true); min_window=2, block)
    end
end

@testset "cor" begin
    @testset "inception" begin
        @test_throws ArgumentError cor(constant(42.0), constant(24.0))
        _test_binary_inception_op(Float64, cor; min_window=2)
    end

    @testset "window" begin
        # Windows that are too small should trigger an exception when attempting to create
        # the node.
        @test_throws ArgumentError cor(n1, n4, 0)
        @test_throws ArgumentError cor(n1, n4, 1)
        @test_throws ArgumentError cor(constant(42.0), constant(24.0), 2)

        _test_binary_window_op(Float64, cor; min_window=2)
    end
end

@testset "ema" begin
    n_obs = 20
    block = _get_rand_block(MersenneTwister(42), n_obs)
    n = block_node(block)

    @test_throws ArgumentError ema(n, 0.0)
    @test_throws ArgumentError ema(n, 0)
    @test_throws ArgumentError ema(n, 1.0)
    @test_throws ArgumentError ema(n, 1)

    for alpha in [0.1, 0.5, 0.9]
        # Naive computation.
        weighted_sum = 0.0
        weighted_count = 0.0
        expected_values = Float64[]
        for x in block.values
            weighted_sum = x + (1 - alpha) * weighted_sum
            weighted_count = 1 + (1 - alpha) * weighted_count
            push!(expected_values, weighted_sum / weighted_count)
        end

        @test _eval(ema(n, alpha)) == Block(block.times, expected_values)
    end

    @test ema(n, 0.5) === ema(n, 3)
end


text:
import Base

"""Base class for a single optimization objective.

All objectives must have a field `initial_state` and a field `generator`, at
minimum. Also, objectives must be able to be instantiated via keyword
arguments.
"""
abstract type AbstractControlObjective end

"""Standard optimization objective.

```julia
Objective(;
    initial_state=<initial_state>,
    generator=<generator>,
    target_state=<target_state>
)
```

describes an optimization objective where the time evaluation of the given
`initial_state` under the given `generator` aims towards `target_state`. The
`generator` here is e.g. a time-dependent Hamiltonian or Liouvillian.

The most common control problems in quantum control, e.g. state-to-state
transitions or quantum gate implementations can be expressed by simultaneously
fulfilling multiple objectives of this type.

Note that the objective can only be instantiated via keyword arguments.
"""
struct Objective{ST,GT} <: AbstractControlObjective
    initial_state::ST
    generator::GT
    target_state::ST
    function Objective(; initial_state::ST, generator::GT, target_state::ST) where {ST,GT}
        new{ST,GT}(initial_state, generator, target_state)
    end
end

"""Minmal optimization objective (initial state and dynamical generator only).

```julia
Objective(;
    initial_state=<initial_state>,
    generator=<generator>,
)
```

describes and optimization objective like the standard [`Objective`](@ref),
except for functionals that are not expressed with respect to some
`target_state`. Having only an `initial_state` and a `generator`, this is the
minimal data structure that is a valid instance of
[`AbstractControlObjective`](@ref).
"""
struct MinimalObjective{ST,GT} <: AbstractControlObjective
    initial_state::ST
    generator::GT
    function MinimalObjective(; initial_state::ST, generator::GT) where {ST,GT}
        new{ST,GT}(initial_state, generator)
    end
end

"""Standard optimization objective with a weight.

```julia
WeightedObjective(;
    initial_state=<initial_state>,
    generator=<genenerator>,
    target_state=<target_state>,
    weight=<weight>
)
```

initializes a control objective like [`Objective`](@ref), but with an
additional `weight` parameter (a float generally between 0 and 1) that weights
the objective relative to other objectives that are part of the same control
problem.
"""
struct WeightedObjective{ST,GT} <: AbstractControlObjective
    initial_state::ST
    generator::GT
    target_state::ST
    weight::Float64
    function WeightedObjective(;
        initial_state::ST,
        generator::GT,
        target_state::ST,
        weight::Float64
    ) where {ST,GT}
        new{ST,GT}(initial_state, generator, target_state, weight)
    end
end

"""A full control problem with multiple objectives.

```julia
ControlProblem(
   objectives=<list of objectives>,
   pulse_options=<dict of controls to pulse options>,
   tlist=<time grid>,
   kwargs...
)
```

Note that the control problem can only be instantiated via keyword arguments.

The `objectives` are a list of [`AbstractControlObjective`](@ref) instances,
each defining an initial state and a dynamical generator for the evolution of
that state. Usually, the objective will also include a target state (see
[`Objective`](@ref)) and possibly a weight (see [`WeightedObjective`](@ref)).

The `pulse_options` are a dictionary (`IdDict`) mapping controls that occur in
the `objectives` to properties specific to the control method.

The `tlist` is the time grid on which the time evolution of the initial states
of each objective should be propagated.

The remaining `kwargs` are keyword arguments that are passed directly to the
optimal control method. These typically include e.g. the optimization
functional.

The control problem is solved by finding a set of controls that simultaneously
fulfill all objectives.
"""
struct ControlProblem{OT<:AbstractControlObjective,OPT<:AbstractDict}
    objectives::Vector{OT}
    pulse_options::OPT # TODO: make this optional?
    tlist::Vector{Float64}
    kwargs::Dict{Symbol,Any}
    function ControlProblem(; objectives, pulse_options, tlist, kwargs...)
        kwargs_dict = Dict{Symbol,Any}(kwargs)  # make the kwargs mutable
        new{typeof(objectives[1]),typeof(pulse_options)}(
            objectives,
            pulse_options,
            tlist,
            kwargs_dict
        )
    end
end

function Base.copy(problem::ControlProblem)
    return ControlProblem(
        objectives=problem.objectives,
        pulse_options=problem.pulse_options,
        tlist=problem.tlist;
        problem.kwargs...
    )
end

# adjoint for the nested-tuple dynamical generator (e.g. `(H0, (H1, ϵ))`)
function dynamical_generator_adjoint(G::Tuple)
    result = []
    for part in G
        # `copy` materializes the `adjoint` view, so we don't end up with
        # unnecessary `Adjoint{Matrix}` instead of Matrix, for example
        if isa(part, Tuple)
            push!(result, (copy(Base.adjoint(part[1])), part[2]))
        else
            push!(result, copy(Base.adjoint(part)))
        end
    end
    return Tuple(result)
end

# fallback adjoint
dynamical_generator_adjoint(G) = Base.adjoint(G)

"""Construct the adjoint of an optimization objective.

```julia
adjoint(objective)
```

Adjoint of a control objective. The adjoint objective contains the adjoint of
the dynamical generator `obj.generator`. It also contains the adjoints of all
other fields (`initial_state`, etc.) if they are defined, or a copy of the
original field value otherwise.

The primary purpose of this adjoint is to facilitate the backward propagation
under the adjoint generator that is central to gradient-based optimization
methods such as GRAPE and Krotov's method.
"""
function Base.adjoint(obj::AbstractControlObjective)
    fields = propertynames(obj)
    adjoints = Dict{Symbol,Any}()  # field => adjoint value
    for field ∈ fields
        if field == :generator
            # For the generator, the adjoint *must* be defined. GRAPE and
            # Krotov critically depend on this for the backward prop
            adjoints[field] = dynamical_generator_adjoint(obj.generator)
        else
            # Any other field, it doesn't really matter too much whether
            # we take the adjoint or not (none of the normal optimization
            # methods depend on anything but the generator being the adjoint)
            adj_value = getproperty(obj, field)
            try
                adj_value = Base.adjoint(adj_value)
            catch
                # if we can't do an adjoint, we'll keep the original value
            end
            try
                adjoints[field] = copy(adj_value)
            catch
                # `copy` isn't available e.g. for Strings
                adjoints[field] = adj_value
            end
        end
    end
    return typeof(obj).name.wrapper(; adjoints...)
end


text:
using StaticArrays
using Distributions
using Setfield

abstract type Particle end
abstract type Process end

struct Photon <: Particle
    E::Float64
    position::SVector{3,Float64}
    direction::SVector{3,Float64}
end

struct Electron <: Particle
    E::Float64
    position::SVector{3,Float64}
    direction::SVector{3,Float64}
end

struct Positron <: Particle
    E::Float64
    position::SVector{3,Float64}
    direction::SVector{3,Float64}
end

struct PhotoEffect <: Process end
struct ComptonEffect <: Process end

struct State
    photons::Vector{Photon}
    electrons::Vector{Electron}
end

function stack(::Type{Photon}, s::State)
    s.photons
end
function stack(::Type{Electron}, s::State)
    s.electrons
end

function processes(::Type{Photon}, s::State)
    PhotoEffect(), ComptonEffect()
end

Base.push!(s::State, p) = push!(stack(typeof(p), s), p)
Base.pop!(s::State, ::Type{P}) where {P} = pop!(stack(P, s))

function simulate_event!(state::State, p_init)
    push!(state, p_init)
    while true
        if !isempty(stack(Photon, state))
            cutoff_or_step!(state, pop!(stack(Photon, state)))
        elseif !isempty(stack(Electron, state))
            cutoff_or_step!(state, pop!(stack(Electron, state)))
        else
            break
        end
    end
end

function mean_free_path(state, p::Photon, proc::PhotoEffect)
    1/(p.E)^3
end

function mean_free_path(state, p::Photon, proc::ComptonEffect)
    p.E < 0.5 ? 1000. : 1.
end

function roll_stepsize(state, p::Particle, proc::Process)
    lambda = mean_free_path(state,p,proc)
    rand(Exponential(lambda))
end

function roll_iproc_stepsize(state, p::Particle)
    procs = processes(typeof(p), state)
    stepsizes = map(procs) do proc
        roll_stepsize(state, p, proc)
    end
    iproc = argmax(stepsizes)
    stepsize = stepsizes[iproc]
    # runtime dispatch
    iproc, stepsize
end

function dispatch_process(f, procs, i)
    if i == 1
        f(procs[1])
    elseif i == 2
        f(procs[2])
    elseif i == 3
        f(procs[3])
    else
        error()
    end
    nothing
end

function run_process!(state,p::Photon,proc::ComptonEffect, stepsize)
    E2 = rand() * p.E
    E1 = p.E - E2
    pos = p.position + stepsize * p.direction
    dir1 = @SVector randn(3)
    dir2 = @SVector randn(3)
    push!(state, Electron(E1, pos, dir1))
    push!(state, Photon(E2, pos, dir2))
end
function run_process!(state, p::Photon,proc::PhotoEffect, stepsize)
    E2 = rand() * p.E
    E1 = p.E - E2
    pos = p.position + stepsize * p.direction
    dir1 = @SVector randn(3)
    dir2 = @SVector randn(3)
    push!(state, Electron(E1, pos, dir1))
    push!(state, Photon(E2, pos, dir2))
end

function needscutoff(state, p::Photon)
    p.E < 0.1
end
function needscutoff(state, p::Particle)
    true
end
function edep!(state, p, E)
    println("edep!($state, $p)")
end

function cutoff_or_step!(state, p)
    if needscutoff(state,p)
        edep!(state, p, p.E)
    else
        step!(state,p)
    end
end

function step!(state, p::Photon)
    i, stepsize = roll_iproc_stepsize(state, p)
    procs = processes(Photon, state)
    dispatch_process(procs, i) do proc
        run_process!(state, p, proc, stepsize)
    end
end


text:
using DiffEqFlux, Optim, OrdinaryDiffEq, DiffEqSensitivity, Zygote, Test

function f(du,u,p,t)
  du[1] = u[2]
  du[2] = -p[1]
end

function condition(u,t,integrator) # Event when event_f(u,t) == 0
  u[1]
end

function affect!(integrator)
  integrator.u[2] = -integrator.p[2]*integrator.u[2]
end

cb = ContinuousCallback(condition,affect!)
u0 = [50.0,0.0]
tspan = (0.0,15.0)
p = [9.8, 0.8]
prob = ODEProblem(f,u0,tspan,p)
sol = solve(prob,Tsit5(),callback=cb)

function loss1(θ)
  sol = solve(prob,Tsit5(),p=[9.8,θ[1]],saveat=0.1,callback=cb,sensealg=ForwardDiffSensitivity())
  target = 20.0
  abs2(sol[end][1] - target)
end

function loss2(θ)
  sol = solve(prob,Tsit5(),p=[9.8,θ[1]],saveat=0.1,callback=cb,sensealg=ReverseDiffAdjoint())
  target = 20.0
  abs2(sol[end][1] - target)
end

#Zygote.gradient(loss1,[0.8])

#Zygote.gradient(loss2,[0.8])

#res = DiffEqFlux.sciml_train(loss1,[0.8],BFGS())
#@test loss1(res.minimizer) < 1

#res = DiffEqFlux.sciml_train(loss2,[0.8],BFGS())
#@test loss2(res.minimizer) < 1


text:
using Revise
using AdFem

mesh = Mesh3(10,10,10,0.1)
nu = constant(rand(get_ngauss(mesh)))
u = constant(rand(mesh.ndof))

f = compute_fem_laplace_term1(u, nu, mesh)

sess = Session(); init(sess)
run(sess, f)


text:
using DFTK
using Test
include("testcases.jl")

@testset "Symmetrization and not symmetrization yields the same density and energy" begin
    args = ((kgrid=[2, 2, 2], kshift=[1/2, 0, 0]),
            (kgrid=[2, 2, 2], kshift=[1/2, 1/2, 0]),
            (kgrid=[2, 2, 2], kshift=[0, 0, 0]),
            (kgrid=[3, 2, 3], kshift=[0, 0, 0]),
            (kgrid=[3, 2, 3], kshift=[0, 1/2, 1/2]),
            )
    for a in args
        Si = ElementPsp(silicon.atnum, psp=load_psp(silicon.psp))
        atoms = [Si => [(ones(3)) / 8, -ones(3) / 8]]
        model = model_DFT(silicon.lattice, atoms, :lda_xc_teter93)

        basis = PlaneWaveBasis(model; Ecut=5, kgrid=(1, 1, 1), use_symmetry=false, a...)
        scfres = self_consistent_field(basis; is_converged=DFTK.ScfConvergenceDensity(1e-10))
        ρ1 = scfres.ρ
        E1 = scfres.energies.total

        basis = PlaneWaveBasis(model; Ecut=5, kgrid=(1, 1, 1), use_symmetry=true, a...)
        scfres = self_consistent_field(basis; is_converged=DFTK.ScfConvergenceDensity(1e-10))
        ρ2 = scfres.ρ
        E2 = scfres.energies.total

        @test abs(E1 - E2) < 1e-10
        @test norm(ρ1 - ρ2) .* sqrt(basis.dvol) < 1e-8
    end
end