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

using Markdown
using InteractiveUtils

# ╔═╡ 1619652c-2b63-11eb-2ec4-9bdbd10c1539
begin
	using SpecialFunctions
	using PyPlot
	using LaTeXStrings
	rc("text", usetex = true)
	ion()
end

# ╔═╡ 3f123684-2b63-11eb-23e2-9749d9d79fea
using HypergeometricFunctions

# ╔═╡ 5dc79c5c-2b88-11eb-240c-8bc0041210f6
using SpecialPolynomials

# ╔═╡ db77ccb8-2b5f-11eb-367b-2332d89b4e97
md"### Testing Julia packages and methods for producing Jacobi functions used to represent the eigenfunctions of the discrete and finite harmonic oscillator

##### Theory
In their 2016 article, Steffen Weimann _et. al_ [1], used the eigenfunctions of the angular momentum coupling $\hat{J}_{x}$ to represent the modes of the finite and discrete harmonic oscillator as

$$u_{n}^{q} = 2^{n}\sqrt{\frac{(j-n)!(j+n)!}{(j-q)!(j+q)!}}P_{n}^{q-n,-q-n}(0),$$

where $P_{n}^{\alpha,\beta}(0)$ is the Jacobi function at $z=0$; $n$ and $q$ are the _energy levels_ and _position index_, respectively such that $n, q\in\{-j,...,j\}$, for $j\in Z^{+}.$

##### The Task
For the purpose, we want to build a $d_{j}\times d_{j}$ matrix for $d_{j}=N+1=2j+1$ where the entries are labeled by the pair $(n,q)$. So, we need the numerical evaluation of the Jacobi function around $z=0$.
__________

[1] Weimann, S., Perez-Leija, A., Lebugle, M. et al. Implementation of quantum and classical discrete fractional Fourier transforms. Nat Commun 7, 11027 (2016). [https://doi.org/10.1038/ncomms11027](https://doi.org/10.1038/ncomms11027)
"

# ╔═╡ e5549ccc-2b62-11eb-24af-bd2ef6a166d7
begin
	# Since most of the representation of the discrete and finite harmonic oscillator are based on the analysis of the angular momentum algebras, we adopt the parameters
	j = 8 # Angular momentum number
	N = 2*j
	dim = N+1
end;

# ╔═╡ 64ad621c-2b61-11eb-339d-95be29513595
md"##### (I) Using `HypergeometricFunctions.jl`

Since this package hasn't implemented the function $P_{n}^{\alpha,\beta}(z)$, we use the definition in terms of the hypergeometric function $_{2}F_{1}(n,a,b;c)$ as

$$P_{n}^{\alpha,\beta}(z)\equiv\binom{n+\alpha}{n}_{2}F_{1}(-n,-n+\alpha+\beta+1,\alpha+1;\tfrac{1}{2}(1-z)).$$

Where we expect a nicely behaviour around $z=0$. The issue as you can see is that $\alpha,\beta>-1$, who is something not fullfilled when $\alpha=-q-n$.
"

# ╔═╡ 12149bc0-2b65-11eb-1152-7159f82080b3
function JHyp(n::Integer,α::Integer,β::Integer,z)
	binomial(n+α,n)*HypergeometricFunctions._₂F₁(-n,n+α+β+1,α+1,(1-z)/2)
end;

# ╔═╡ 63b06a18-2b65-11eb-02b4-0147f09abdfd
function JHypF(n::Integer,q::Integer)
	(2^Float64(n))*sqrt(gamma(big(j-n+1))*gamma(big(j+n+1))/(gamma(big(j-q+1))*gamma(big(j+q+1))))*JHyp(j+n,q-n,-q-n,0)
end;

# ╔═╡ 9946f82c-2b65-11eb-1472-2355273a3f18
@timed begin
	JHypM = zeros(Float64,(dim,dim))
	for l in -j:j
		for m in -j:j
			JHypM[l+j+1,m+j+1] = JHypF(l,m)
		end
	end
end

# ╔═╡ d8159ab8-2b65-11eb-3fb0-674f6f8d069b
begin
	imshow(JHypM,cmap="gray")
	savefig("JHypM.png",dpi=150)
end

# ╔═╡ 0808796e-2b9a-11eb-28bb-898219baca42
md"![](https://raw.githubusercontent.com/rurz/WG_FractionalEvolution/main/tests/JHypM.png)"

# ╔═╡ 47fcb3dc-2b66-11eb-0322-7160add87fd2
md"When we plot the matrix array, we see that only the upper left half of the values are obtained. Meaning that the method provided by `HypergeometricFunctions.jl` lacks on rules of parity, simmetry, and reflection. Examining the matrix we encounter NaN results, so we faced a problem of ill-defined numerical method when we don't fullfill the condition $\alpha,\beta>-1.$"

# ╔═╡ 4c6ceec6-2b88-11eb-36f8-f3992615474f
md"#### (II) Using `SpecialPolynomials.jl`
This package in particular takes two parameters, (α,β), and a vector of $n$ dimension to produce the related polynomial on $z$ who after needs to be evaluated.\
$$J_{n}^{\alpha,\beta}(z)\equiv$$ `Jacobi{α,β}(V)`, where $dim(V)=n\in\{0,...,N\}.$
"

# ╔═╡ 08fd30ea-2b8a-11eb-1474-ef09e8030734
begin
	function vn(n)
		vd = zeros(Int64,n+1)
		vd[end]=1
		return vd
	end
	function JSP(n,α,β,z)
		(SpecialPolynomials.Jacobi{α,β}(vn(n)))(z)
	end
end;

# ╔═╡ 07566168-2b8b-11eb-1067-b3327e8c3c06
function JSPF(n,q)
	(2^Float64(n))*sqrt(gamma(big(j-n+1))*gamma(big(j+n+1))/(gamma(big(j-q+1))*gamma(big(j+q+1))))*JSP(j+n,q-n,-q-n,0)
end;

# ╔═╡ 1f9a2ca8-2b8b-11eb-14a5-7f83166acde4
# Uncomment for evaluate. Carefull, it requires a lot of time to perform the calculations for j above 6
begin
	JSPFM = zeros(Float64,(dim,dim))
	for l in -j:j
		for m in -j:j
			JSPFM[l+j+1,m+j+1] = JSPF(l,m)
		end
	end
end

# ╔═╡ 80b72b96-2b8e-11eb-3150-6f7795a2d027
begin
	imshow(JSPFM,cmap="gray")
	savefig("JSPFM.png",dpi=150)
end

# ╔═╡ f17428a6-2b99-11eb-2e6f-73737235bc48
md"![](https://raw.githubusercontent.com/rurz/WG_FractionalEvolution/main/tests/JSPFM.png)"

# ╔═╡ 3813f4e0-2b94-11eb-3387-cf14cad21276
md"
In this particular implementation we see that a problem arise in the lower middle part of the matrix $(n,q)$, implying that the definition of the `Jacobi{α,β}` suffer an error for some combination of values, as we can see in the error message above. Moreover, the evaluation it's far from be fast because we need to evaluate every polynomial at the time.
"

# ╔═╡ d9e0a670-2b66-11eb-3181-e3a029a778cb
md"##### (III) Using series representation expansion around $z=0$

Provided that we can obtain a closed form of $P_{n}^{\alpha,\beta}(z)$ for $z\rightarrow 0$, we use this approach. Given in [Wolfram Function Library](https://functions.wolfram.com/Polynomials/JacobiP/06/01/04/01/0004/), the definition used is

$$P_{n}^{\alpha,\beta}(z)=2^{-n}\sum\limits_{k=0}^{n}\binom{\alpha+n}{k}\binom{\beta+n}{n-k}(z+1)^{k}(z-1)^{n-k},$$

who isn't ill-defined at $z\rightarrow 0$, and we expect to behave nicely for values $\alpha,\beta<-1$.
"

# ╔═╡ 980b12e2-2960-11eb-220a-7786e56f2261
function JSer(n::Integer,α::Integer,β::Integer,z)
	JSerM = zeros(Float64,n+1)
	for k in 0:n
		JSerM[k+1] = binomial(α+n,k)*binomial(β+n,n-k)*(z+1)^(k)*(z-1)^(n-k)
	end
	return (1/2^(n))*sum(JSerM)
end;

# ╔═╡ f2509478-2566-11eb-2aa0-0be2f2b03b3f
function JSerF(n::Integer,q::Integer)
	(2^Float64(n))*sqrt(gamma(big(j-n+1))*gamma(big(j+n+1))/(gamma(big(j-q+1))*gamma(big(j+q+1))))*JSer(j+n,q-n,-q-n,0)
end;

# ╔═╡ 49243b96-2955-11eb-1d0e-a3f83590722e
@timed begin
	JSerFM = zeros(Float64,(dim,dim))
	for l in -j:j
		for m in -j:j
			JSerFM[l+j+1,m+j+1] = JSerF(l,m)
		end
	end
end

# ╔═╡ 9b37f096-2962-11eb-0512-b702db069ddf
begin
	imshow(JSerFM,cmap="gray")
	savefig("JSerFM.png",dpi=150)
end

# ╔═╡ 824608be-2b99-11eb-3b4e-1752fd403812
md"![K](https://raw.githubusercontent.com/rurz/WG_FractionalEvolution/main/tests/JSerFM.png)"

# ╔═╡ 4fc55ac4-2b68-11eb-124e-675dece08b7c
md"This time, we obtain the matrix in complete form. Since this approach use a series representation, we have not problems dealing with symmetries and reflection changes. We then have a nicely numerical way to obtain the matrix. The cons are that we need to construct the series and stored in a function `JSer`, who need to be evaluated every time whe change the value of $j\in Z^{+}$."

# ╔═╡ 0045571e-2b69-11eb-18fa-9bb5003e3968
md"
**Aftermath**: A revision on the methods of `Hypergeometricfunctions.jl` and `SpecialPolynomials.jl` would be nice to overcome this problem and have a fully implementable function, who for sure, responds on this particular problem at the time.
"

# ╔═╡ de05e088-2b68-11eb-274c-b9be94d7f0a7
md"
Alejandro R. Urzúa (2020)

[Github](https://github.com/rurz)
"

# ╔═╡ Cell order:
# ╟─db77ccb8-2b5f-11eb-367b-2332d89b4e97
# ╠═1619652c-2b63-11eb-2ec4-9bdbd10c1539
# ╠═e5549ccc-2b62-11eb-24af-bd2ef6a166d7
# ╟─64ad621c-2b61-11eb-339d-95be29513595
# ╠═3f123684-2b63-11eb-23e2-9749d9d79fea
# ╠═12149bc0-2b65-11eb-1152-7159f82080b3
# ╠═63b06a18-2b65-11eb-02b4-0147f09abdfd
# ╠═9946f82c-2b65-11eb-1472-2355273a3f18
# ╠═d8159ab8-2b65-11eb-3fb0-674f6f8d069b
# ╟─0808796e-2b9a-11eb-28bb-898219baca42
# ╟─47fcb3dc-2b66-11eb-0322-7160add87fd2
# ╟─4c6ceec6-2b88-11eb-36f8-f3992615474f
# ╠═5dc79c5c-2b88-11eb-240c-8bc0041210f6
# ╠═08fd30ea-2b8a-11eb-1474-ef09e8030734
# ╠═07566168-2b8b-11eb-1067-b3327e8c3c06
# ╠═1f9a2ca8-2b8b-11eb-14a5-7f83166acde4
# ╠═80b72b96-2b8e-11eb-3150-6f7795a2d027
# ╟─f17428a6-2b99-11eb-2e6f-73737235bc48
# ╟─3813f4e0-2b94-11eb-3387-cf14cad21276
# ╟─d9e0a670-2b66-11eb-3181-e3a029a778cb
# ╠═980b12e2-2960-11eb-220a-7786e56f2261
# ╠═f2509478-2566-11eb-2aa0-0be2f2b03b3f
# ╠═49243b96-2955-11eb-1d0e-a3f83590722e
# ╠═9b37f096-2962-11eb-0512-b702db069ddf
# ╟─824608be-2b99-11eb-3b4e-1752fd403812
# ╟─4fc55ac4-2b68-11eb-124e-675dece08b7c
# ╟─0045571e-2b69-11eb-18fa-9bb5003e3968
# ╟─de05e088-2b68-11eb-274c-b9be94d7f0a7


text:
julia> println(parse(Int, readuntil(stdin, ' ')) + parse(Int, readuntil(stdin, '\n')))
1 2
3

ext: jl

max_stars_repo_name: ethansaxenian/RosettaDecode

max_issues_repo_name: ethansaxenian/RosettaDecode

max_forks_repo_name: ethansaxenian/RosettaDecode


text:
# function poiseuille(y::Real, canal::CanalConfig, μ, dpdx)
#     h = canal.H
#     vmax = -1 / μ * dpdx * h^2 / 2
#     vmax * (1 - (y / h)^2)
# end

# poiseuille(y::AbstractVector, c::CanalConfig, μ, dpdx) = poiseuille.(y, Ref(c), Ref(μ), Ref(dpdx))

function poiseuille(y::Real, gb::GridBox, Q)
    uc = 1.5 * Q / gb.D
    uc * (1 - (y / (0.5*gb.D))^2)
end

poiseuille(y::AbstractVector, c::GridBox, Q) = poiseuille.(y, Ref(c), Ref(Q))

function poiseuille!(grid::AbstractGrid, c::GridBox, Q::Real)
    ys = getys(grid)
    for i in eachindex(grid)
        grid[i] = poiseuille(ys[i[2]], c, Q)
    end
end

# poiseuille(y::Real, canal::CanalConfig, params::SimParam) = poiseuille(y, canal, params.Q)
poiseuille!(y::AbstractGrid, gb::GridBox, params::SimParam) = poiseuille!(y, gb, params.Q)

poiseuille!(sim::Simulation) = poiseuille!(sim.u, sim.u.grid, sim.params)
poiseuille(sim::Simulation, y) = poiseuille(y, sim.u.grid, sim.params.Q)


text:
# ---
# title: 503. Next Greater Element II
# id: problem503
# author: Indigo
# date: 2021-03-06
# difficulty: Medium
# categories: Stack
# link: <https://leetcode.com/problems/next-greater-element-ii/description/>
# hidden: true
# ---
#
# Given a circular array (the next element of the last element is the first
# element of the array), print the Next Greater Number for every element. The
# Next Greater Number of a number x is the first greater number to its
# traversing-order next in the array, which means you could search circularly to
# find its next greater number. If it doesn't exist, output -1 for this number.
#
# **Example 1:**
#
#
#
#     Input: [1,2,1]
#     Output: [2,-1,2]
#     Explanation: The first 1's next greater number is 2; The number 2 can't find next greater number; The second 1's next greater number needs to search circularly, which is also 2.
#
#
# **Note:** The length of given array won't exceed 10000.
#
#
## @lc code=start
using LeetCode

function next_greater_elements(nums::Vector{Int})
    stk = Int[]
    res = fill(-1, size(nums))
    for rd in 1:2, (idx, num) in enumerate(nums)
        while !isempty(stk) && nums[stk[end]] < num
            res[pop!(stk)] = num
        end
        (rd == 1) && push!(stk, idx)
    end
    return res
end
## @lc code=end


text:
@testset "construction" begin
    P = UnivariatePolynomial{ZZ{Int},:x}
    p = P([1, -2])
    q = P([1, -1])
    pq = Frac(p, q)
    F = Frac{P}
    @test basetype(F) == P
    @test depth(F) == 3
    @test !issimpler(pq, pq)
    @test lcunit(pq) == 1
    @test copy(pq) == pq
    @test F(pq) === pq
    @test Frac(p) == p
    P2 = UnivariatePolynomial{ZZ{Int8},:x}
    pq2 = P2([1, -2]) // P2([1, -1])
    @test F(pq2) == pq
    @test F(13) == F(13, 1)
    @test F(ZZ(11)) == F(ZZ(22), ZZ(2))
    @test basetype(Frac(Int8(11))) == ZZ{Int8}

    @test convert(F, pq) === pq
    @test convert(F, pq2) == pq2
    @test convert(F, p) == p
    @test convert(F, 42) == F(84, 2)
    @test convert(F, 42 // 8) == 21 // 4

    @test Frac(p, q) != nothing
    @test p // q == Frac(p, q)
    @test p // p == 1
    @test pq * pq == F(p * p, P([1, -2, 1]))
    @test pq / pq == 1
    @test isunit(F(-1))
    @test isunit(pq)
    @test isone(one(F))
    @test !isone(pq)
    @test one(F) == F(10, 10)
    @test iszero(zero(F))
    @test !iszero(F(1))
    @test F(1) + P(1) == 2
    @test P(1) + F(1) == 2

    @test sprint(show, F(p)) == "-2*x + 1"
    @test sprint(show, pq) == "(2*x - 1)/(x - 1)"

    Q = Frac(ZZ{Int})
    @test sprint(show, Q(12, 8)) == "(3)/(2)"

end

@testset "polynomial over GF(7)" begin
    G = GF(7)
    GX = G[:x]
    x = monom(GX)
    p = (x + 1) ^ 3 + 1
    q =  x + 2
    @test p // q == p / q
    @test p // q^2 == (x^2 + x + 1) // q
end

@testset "polynomial over GF(4)" begin
    G = GF(4)
    GX = G[:x]
    x = monom(GX)
    p = (x + 1) ^ 3 + 1
    q =  x + G[2]
    @test p // q == p / q
    @test p // q^2 == (x^2 + G[3]*x) // q
end


text:
using Kinetic, Plots
using KitBase.JLD2
using Base.Threads: @threads
using KitBase.ProgressMeter: @showprogress
cd(@__DIR__)

X = Float32.([[1e-4, 0.0, 0.0, 1e-4]; zeros(4); 1.0])
Y = Int32(0)

begin
    set = Setup(space = "1d1f3v", boundary = ["fix", "fix"], collision = "fsm")
    ps = PSpace1D(-0.5, 0.5, 250, 1)
    vs = VSpace3D(-6.0, 6.0, 28, -6.0, 6.0, 28, -6.0, 6.0, 28)
    gas = begin
        Kn = 5e-3
        Gas(Kn = Kn, K = 0.0, fsm = fsm_kernel(vs, ref_vhs_vis(Kn, 1.0, 0.5)))
    end
    fw = function(x)
        prim = zeros(5)
        if x <= 0
            prim .= [1.0, 0.0, 1.0, 0.0, 1.0]
        else
            prim .= [1.0, 0.0, -1.0, 0.0, 2.0]
        end

        return prim_conserve(prim, ks.gas.γ)
    end
    ib = IB1F(fw, vs, gas)
    ks = SolverSet(set, ps, vs, gas, ib)
    ctr, face = init_fvm(ks, ks.ps)
end

τ0 = vhs_collision_time(ctr[1].prim, ks.gas.μᵣ, ks.gas.ω)
tmax = 50 * τ0
t = 0.0
dt = timestep(ks, ctr, t)
nt = Int(tmax ÷ dt)
res = zero(ctr[1].w)

@showprogress for iter = 1:5
    evolve!(ks, ctr, face, dt)
    update!(ks, ctr, face, dt, res)

    if iter%5 == 0
        for i = 1:ks.ps.nx
            if abs(ctr[i].prim[2]) > 1e-6
                @show i
                sw = (ctr[i+1].w .- ctr[i-1].w) / ks.ps.dx[i] / 2
                x, y = regime_data(ks, ctr[i].w, sw, zero(sw), zero(sw), ctr[i].f)
                rx = [x[1:3]; x[5:8]; x[10:11]]
                global X = hcat(X, rx)
                global Y = hcat(Y, y)
            end
        end
    end

    global t += dt
end

"""
i: [121, 130]
"""

@load "model/nn_layer.jld2" nn
accuracy(nn, X, Y)

# field
sol = zeros(ks.ps.nx, 5)
for i in axes(sol, 1)
    sol[i, :] .= ctr[i].prim
    sol[i, end] = 1 / sol[i, end]
end

plot(ks.ps.x[1:ks.ps.nx], sol)

begin
    rg_ref = zeros(Int, ks.ps.nx)
    for i = 1:ks.ps.nx
        sw = (ctr[i+1].w .- ctr[i-1].w) / ks.ps.dx[i] / 2
        x, y = regime_data(ks, ctr[i].w, sw, zero(sw), zero(sw), ctr[i].f)
        rg_ref[i] = y
    end

    rg_nn = zeros(Int, ks.ps.nx)
    @inbounds Threads.@threads for i = 1:ks.ps.nx
        sw = abs.((ctr[i+1].w - ctr[i-1].w) / ks.ps.dx[i] / 2)
        tau = vhs_collision_time(ctr[i].prim, ks.gas.μᵣ, ks.gas.ω)
        rg_nn[i] = nn([ctr[i].w[1:3]; ctr[i].w[end]; sw[1:3]; sw[end]; tau])[1] |> round |> Int
    end

    plot(ks.ps.x[1:ks.ps.nx], rg_ref)
    plot!(ks.ps.x[1:ks.ps.nx], rg_nn)
end

#@load "data/dataset.jld2" X Y

nn(X[:, 9])[1] |> round

sw = (ctr[end÷2+1].w - ctr[end÷2-1].w) / (1e-6 + ks.ps.x[end÷2+1] - ks.ps.x[end÷2-1])
tau = vhs_collision_time(ctr[end÷2].prim, ks.gas.μᵣ, ks.gas.ω)
nn([ctr[end÷2].w[1:3]; ctr[end÷2].w[end]; sw[1:3]; sw[end]; tau])[1]
regime_data(ks, ctr[end÷2].w, sw, zero(sw), zero(sw), ctr[end÷2].f)[2]

i = 125
sw = (ctr[i+1].w - ctr[i-1].w) / (ks.ps.x[i+1] - ks.ps.x[i-1])
@show tau = vhs_collision_time(ctr[i].prim, ks.gas.μᵣ, ks.gas.ω)

max_stars_repo_name: CSMMLab/Flowmachine

max_issues_repo_name: CSMMLab/Flowmachine

max_forks_repo_name: CSMMLab/Flowmachine


text:
"""
    QDLDL inplace functionality
"""
mutable struct LDLSolver{Tf<:AbstractFloat,Ti<:Integer} <: LinearSolver
    # QDLDL Factorization
    F::QDLDL.QDLDLFactorisation{Tf,Ti}
    A_sparse::SparseMatrixCSC{Tf,Ti}
    # Allocate memory
    AtoPAPt::Vector{Ti}
    Pr::Vector{Ti}
    Pc::Vector{Ti}
    Pv::Vector{Tf}
    num_entries::Vector{Ti}
    Pr_triu::Vector{Ti}
    Pv_triu::Vector{Tf}
end

function LDLSolver(A::SparseMatrixCSC{Tf,Ti}, F::QDLDL.QDLDLFactorisation{Tf,Ti}) where {Tf<:AbstractFloat,Ti<:Integer}
    AtoPAPt = zeros(Ti,nnz(A))
    Pr = zeros(Ti, nnz(A))
    Pc = zeros(Ti, size(A, 1) + 1)
    Pv = zeros(Tf, nnz(A))
    num_entries = zeros(Ti,size(A, 2))
    Pr_triu = zeros(Ti, nnz(F.workspace.triuA))
    Pv_triu = zeros(Tf, nnz(F.workspace.triuA))
    return LDLSolver{Tf,Ti}(F, copy(A), AtoPAPt, Pr, Pc, Pv, num_entries, Pr_triu, Pv_triu)
end

function factorize!(s::LDLSolver{Tf,Ti}, A::SparseMatrixCSC{Tf,Ti}) where {Tf<:AbstractFloat, Ti<:Integer}
    # Reset the pre-allocated fields
    s.AtoPAPt .= 0
	s.Pr .= 0
    s.Pc .= 0
	s.Pv .= 0.0
    s.num_entries .= 0
	s.Pr_triu .= 0
	s.Pv_triu .= 0.0

    # Triangularize the matrix with the allocation-free method.
    A = _permute_symmetricAF(A, s.AtoPAPt, s.F.iperm, s.Pr, s.Pc, s.Pv,
		s.num_entries, s.Pr_triu, s.Pv_triu)  #returns an upper triangular matrix

    # Update the workspace, triuA is the only field we need to update
    s.F.workspace.triuA.nzval .= A.nzval

    # factor the matrix
    QDLDL.refactor!(s.F)

    # return nothing
end

# the main function without extra argument checks
# following the book: Timothy Davis - Direct Methods for Sparse Linear Systems
function _permute_symmetricAF(
        A::SparseMatrixCSC{Tf, Ti},
        AtoPAPt::AbstractVector{Ti},
        iperm::AbstractVector{Ti},
		Pr::AbstractVector{Ti},
        Pc::AbstractVector{Ti},
        Pv::AbstractVector{Tf},
		num_entries::Vector{Ti},
		Pr_triu::AbstractVector{Ti},
		Pv_triu::AbstractVector{Tf},
        ) where {Tf <: AbstractFloat, Ti <: Integer}

    # 1. count number of entries that each column of P will have
    n = size(A, 2)
    Ar = A.rowval
    Ac = A.colptr
    Av = A.nzval

    # count the number of upper-triangle entries in columns of P, keeping in mind the row permutation
    for colA = 1:n
        colP = iperm[colA]
        # loop over entries of A in column A...
        for row_idx = Ac[colA]:Ac[colA+1]-1
            rowA = Ar[row_idx]
            rowP = iperm[rowA]
            # ...and check if entry is upper triangular
            if rowA <= colA
                # determine to which column the entry belongs after permutation
                col_idx = max(rowP, colP)
                num_entries[col_idx] += one(Ti)
            end
        end
    end
    # 2. calculate permuted Pc = P.colptr from number of entries
    Pc[1] = one(Ti)
    @inbounds for k = 1:n
        Pc[k + 1] = Pc[k] + num_entries[k]

        # reuse this vector memory to keep track of free entries in rowval
        num_entries[k] = Pc[k]
    end
    # use alias
    row_starts = num_entries

    # 3. permute the row entries and position of corresponding nzval
    for colA = 1:n
        colP = iperm[colA]
        # loop over rows of A and determine where each row entry of A should be stored
        for rowA_idx = Ac[colA]:Ac[colA+1]-1
            rowA = Ar[rowA_idx]
            # check if upper triangular
            if rowA <= colA
                rowP = iperm[rowA]
                # determine column to store the entry
                col_idx = max(colP, rowP)

                # find next free location in rowval (this results in unordered columns in the rowval)
                rowP_idx = row_starts[col_idx]

                # store rowval and nzval
                Pr[rowP_idx] = min(colP, rowP)
                Pv[rowP_idx] = Av[rowA_idx]

                #record this into the mapping vector
                AtoPAPt[rowA_idx] = rowP_idx

                # increment next free location
                row_starts[col_idx] += 1
            end
        end
    end
	nz_new = Pc[end] - 1
	for i = 1:nz_new
		Pr_triu[i] = Pr[i]
		Pv_triu[i] = Pv[i]
	end
    P = SparseMatrixCSC{Tf, Ti}(n, n, Pc, Pr_triu, Pv_triu)

    return P
end

"""
    LDL solver
"""
function ldl_solver(A::SparseMatrixCSC{T,Int}) where T
    LDLSolver(A, qdldl(A))
end

ldl_solver(A::Array{T, 2}) where T = ldl_solver(sparse(A))

function linear_solve!(solver::LDLSolver{Tv,Ti}, x::Vector{Tv}, A::SparseMatrixCSC{Tv,Ti}, b::Vector{Tv};
    reg=0.0, fact::Bool = true) where {Tv<:AbstractFloat,Ti<:Integer}
    fact && factorize!(solver, A) # factorize
    x .= b
    QDLDL.solve!(solver.F, x) # solve
end

function linear_solve!(solver::LDLSolver{Tv,Ti}, x::Vector{Tv}, A::AbstractMatrix{Tv}, b::Vector{Tv};
    reg=0.0, fact::Bool = true) where {Tv<:AbstractFloat,Ti<:Integer}

    # fill sparse_matrix
    n, m = size(A)
    for i = 1:n
        for j = 1:m
            solver.A_sparse[i, j] = A[i, j]
        end
    end

    linear_solve!(solver, x, solver.A_sparse, b, reg=reg, fact=fact)
end

function linear_solve!(s::LDLSolver{T}, x::Matrix{T}, A::Matrix{T},
    b::Matrix{T};
    reg::T = 0.0,
    fact::Bool = true) where T

    fill!(x, 0.0)
    n, m = size(x)
    r_idx = 1:n
    fact && factorize!(s, A)

    x .= b
    for j = 1:m
        xv = @views x[r_idx, j]
        QDLDL.solve!(solver.F, xv)
    end
end


text:
module ens_forecast

export da_cycles

using Statistics
using LinearAlgebra
using Random

using Distributions
using ProgressMeter
using PyCall

struct Forecast_Info
    errs
    errs_fcst
    crps
    crps_fcst
    spread
    spread_fcst
    Q_hist
    P_hist
    analyses
    trues
    model_errs
    inflation_hist
end

xskillscore = pyimport("xskillscore")
xarray = pyimport("xarray")

function make_psd(A, tol=1e-6)
    L, Q = eigen(A)
    L[L .< 0] .= tol
    return Symmetric(Q*diagm(0=>L)*inv(Q))
end

function da_cycles(; x0::AbstractVector{float_type},
                     ensembles::AbstractVector{<:AbstractMatrix{float_type}},
                     models::AbstractVector{<:Function}, model_true::Function,
                     obs_ops::AbstractVector{<:AbstractMatrix}, H_true=I,
                     mappings::Union{Matrix{AbstractArray}, Nothing}=nothing,
                     model_errs::AbstractVector{<:Union{AbstractMatrix{float_type}, Nothing}},
                     model_errs_prescribed::AbstractVector{<:Union{AbstractMatrix{float_type}, Nothing}},
                     integrators::AbstractVector{<:Function}, integrator_true::Function,
                     da_method::Function, localization::AbstractMatrix{float_type},
                     ens_sizes::AbstractVector{int_type}, Δt::float_type, window::int_type,
                     n_cycles::int_type, outfreq::int_type,
                     model_sizes::AbstractVector{int_type}, R::Symmetric{float_type},
                     ens_errs::Union{AbstractVector{<:AbstractMatrix{float_type}}, Nothing}=nothing,
                     ρ::float_type,
                     Q_p::Union{AbstractVector{<:AbstractMatrix{float_type}}, Nothing}=nothing,
                     ρ_all::float_type=0.01, all_orders::Bool=true,
                     combine_forecasts::Bool=true, gen_ensembles::Bool=false,
                     assimilate_obs::Bool=true, save_Q_hist::Bool=false,
                     save_P_hist::Bool=false, save_analyses::Bool=false, save_trues::Bool=false,
                     prev_analyses::Union{AbstractArray{float_type}, Nothing}=nothing,
                     leads::int_type=1, ref_model::int_type=1) where {float_type<:AbstractFloat, int_type<:Integer}
    n_models = length(models)
    obs_err_dist = MvNormal(R)
    R_inv = inv(R)

    x_true = x0

    errs = Array{float_type}(undef, n_cycles, model_sizes[ref_model])
    errs_fcst = Array{float_type}(undef, n_cycles, model_sizes[ref_model])
    crps = Array{float_type}(undef, n_cycles)
    crps_fcst = Array{float_type}(undef, n_cycles)
    Q_hist = Array{float_type}(undef, n_models, n_cycles)
    if save_Q_hist
        Q_hist = Array{Matrix{float_type}}(undef, n_models, n_cycles)
    end
    if save_P_hist
        P_hist = Array{Matrix{float_type}}(undef, n_models, n_cycles)
    else
        P_hist = nothing
    end
    spread = Array{float_type}(undef, n_cycles)
    spread_fcst = Array{float_type}(undef, n_cycles)
    inflation_hist = Array{float_type}(undef, n_cycles)

    inflation_all = ones(leads)

    if save_analyses
        analyses = Array{float_type}(undef, n_cycles, model_sizes[ref_model],
                                     all_orders ? sum(ens_sizes) : ens_sizes[ref_model])
    else
        analyses = nothing
    end

    if save_trues
        trues = Array{float_type}(undef, n_cycles, model_sizes[ref_model])
    else
        trues = nothing
    end

    model_errs_leads = Array{AbstractMatrix{float_type}}(undef, n_models, leads)

    for model=1:n_models
        for lead=1:leads
            model_errs_leads[model, lead] = model_errs[model]*lead^2
        end
    end

    if all_orders
        orders = []
        for i=1:n_models
            order = Array(1:n_models)
            replace!(order, 1=>i, i=>1)
            append!(orders, [order])
        end
    else
        order = Array(1:n_models)
        replace!(order, 1=>ref_model, ref_model=>1)
        orders = [order]
    end

    if mappings === nothing
        if (~all(model_sizes[1] .== model_sizes))
            error("Must specify mappings")
        else
            mappings = Matrix{AbstractArray}(undef, n_models, n_models)
            for m1=1:n_models
                for m2=1:n_models
                    mappings[m1, m2] = I(model_sizes[1])
                end
            end
        end
    end

    t = 0.0

    @showprogress for cycle=1:n_cycles
        y = H_true*x_true + rand(obs_err_dist)

        lead = mod(cycle, leads)

        for model=1:n_models
            model_size = model_sizes[model]
            E = ensembles[model]

            H_m = obs_ops[model]

            x_m = mean(E, dims=2)
            innovation = y - H_m*x_m
            P_p = Symmetric(cov(E'))

            if save_P_hist
                P_hist[model, cycle] = P_p
            end

            C = innovation*innovation' - R - H_m*P_p*H_m'
            if rank(H_m) >= model_size
                Q_est = pinv(H_m)*C*pinv(H_m)'
            else
                A = Array{float_type}(undef, size(R)[1]^2, length(Q_p))
                for p=1:length(Q_p)
                    A[:, p] = vec(H_m*Q_p[p]*H_m')
                end
                q = A \ vec(C)
                Q_est = sum([q[p]*Q_p[p] for p=1:length(Q_p)])
            end

            Q = Symmetric(ρ*Q_est + (1 - ρ)*model_errs_leads[model, lead + 1])

            if !isposdef(Q)
                Q = make_psd(Q)
            end

            if save_Q_hist
                Q_hist[model, cycle] = Q
            else
                Q_hist[model, cycle] = tr(Q)
            end
            model_errs_leads[model, lead + 1] = Q

            E += rand(MvNormal(model_errs_leads[model, lead + 1]), ens_sizes[model])

            ensembles[model] = E
        end

        ensembles_new = similar(ensembles)
        # Iterative multi-model data assimilation
        if combine_forecasts
            for (i, order) in enumerate(orders)
                for model=2:n_models
                    # Posterior ensemble of the previous model is used as the prior
                    # ensemble for the next model
                    if model == 2
                        E = ensembles[order[model-1]]
                    else
                        E = ensembles_new[i]
                    end

                    E_model = ensembles[order[model]]

                    H_model = mappings[order[1], order[model]]

                    P_f = cov(E_model')
                    if localization !== nothing
                        P_f = (mappings[ref_model, order[model]]*localization*mappings[ref_model, order[model]]').*P_f
                    else
                        P_f = Diagonal(diagm(diag(P_f)))
                    end
                    P_f_inv = Symmetric(inv(P_f))

                    if localization !== nothing
                        localization_m = mappings[ref_model, order[model-1]]*localization*mappings[ref_model, order[model-1]]'
                    else
                        localization_m = nothing
                    end

                    # Assimilate the forecast of each ensemble member of the current
                    # model as if it were an observation
                    E = da_method(E=E, R=Symmetric(Matrix(P_f)), R_inv=P_f_inv, H=H_model,
                                  y=mean(E_model, dims=2)[:, 1],
                                  localization=localization_m)

                    ensembles_new[i] = E
                end
            end
        end

        if (n_models > 1) & (combine_forecasts)
            ensembles = ensembles_new
        end

        if all_orders & (~all(model_sizes[1] .== model_sizes))
            error("Not implemented")
        end

        if (~all_orders) & (~all(ens_sizes[1] .== ens_sizes))
            error("Ensemble sizes must be the same")
        end

        if all_orders
            E_all = hcat([mappings[model, ref_model]*ensembles[model] for model=1:n_models]...)
        else
            E_all = ensembles[1]
        end

	    if (n_models > 1)
            H = obs_ops[ref_model]

            x_m = mean(E_all, dims=2)
            innovation = y - H*x_m

            P_e = innovation*innovation'
            P_f = Symmetric(cov(E_all'))
            λ = tr(P_e - R)/tr(H*P_f*H')
            λ = max(λ, 0)

            inflation_all[lead + 1] = ρ_all*λ + (1 - ρ_all)*inflation_all[lead + 1]
	        inflation_hist[cycle] = inflation_all[lead + 1]

            E_all = x_m .+ sqrt(inflation_all[lead + 1])*(E_all .- x_m)
        end

        errs_fcst[cycle, :] = mean(E_all, dims=2) - pinv(obs_ops[ref_model])*H_true*x_true

	    E_corr_fcst_array = xarray.DataArray(data=E_all, dims=["dim", "member"])
        crps_fcst[cycle] = xskillscore.crps_ensemble(pinv(obs_ops[ref_model])*H_true*x_true, E_corr_fcst_array).values[1]
        spread_fcst[cycle] = mean(std(E_all, dims=2))

        if assimilate_obs & (mod(cycle, leads) == 0)
            E_a = da_method(E=E_all, R=R, R_inv=R_inv, H=obs_ops[ref_model],
                            y=y, localization=localization)

            E_corr_array = xarray.DataArray(data=E_a, dims=["dim", "member"])
            crps[cycle] = xskillscore.crps_ensemble(pinv(obs_ops[ref_model])*H_true*x_true, E_corr_array).values[1]

            spread[cycle] = mean(std(E_a, dims=2))

            errs[cycle, :] = mean(E_a, dims=2) - pinv(obs_ops[ref_model])*H_true*x_true
        else
            E_a = E_all
        end

        if save_analyses
            analyses[cycle, :, :] = E_a
        end

        if save_trues
            trues[cycle, :] = x_true
        end

        for model=1:n_models
            if gen_ensembles & (mod(cycle, leads) == 0)
                E = mappings[ref_model, model]*pinv(obs_ops[ref_model])*H_true*x_true .+ rand(MvNormal(ens_errs[model]), ens_sizes[model])
            elseif (prev_analyses !== nothing) & (mod(cycle, leads) == 0)
                E = mappings[ref_model, model]*prev_analyses[cycle, :, [0; cumsum(ens_sizes)][model]+1:[0; cumsum(ens_sizes)][model+1]]
            else
                if all_orders
                    E = mappings[ref_model, model]*E_a[:, [0; cumsum(ens_sizes)][model]+1:[0; cumsum(ens_sizes)][model+1]]
                else
                    E = mappings[ref_model, model]*E_a
                end
            end

            if model_errs_prescribed[model] === nothing
                pert = zeros(model_sizes[model])
            else
                pert = rand(MvNormal(model_errs_prescribed[model]))
            end
            Threads.@threads for i=1:ens_sizes[model]
                integration = integrators[model](models[model], E[:, i], t,
                                                 t + window*outfreq*Δt, Δt, inplace=false)
                E[:, i] = integration[end, :] + pert
            end

            ensembles[model] = E
        end

        x_true = integrator_true(model_true, x_true, t, t + window*outfreq*Δt, Δt)

        t += window*outfreq*Δt
    end

    return Forecast_Info(errs, errs_fcst, crps, crps_fcst, spread, spread_fcst, Q_hist,
                         P_hist, analyses, trues, model_errs_leads, inflation_hist)
end

end

max_stars_repo_name: eviatarbach/mmda

max_issues_repo_name: eviatarbach/mmda

max_forks_repo_name: eviatarbach/mmda


text:
abstract type Writer{C,T} end

mutable struct TypedBuffer{T<:NTuple{<:Any,Array}}
    size::Int # The current number of data entries in the buffer.
    maxsize::Int # The maximum number of data entries in the buffer.
    data::T # Tuple that contains a fixed number of data vectors with length `maxsize`.
end

@generated function TypedBuffer(collector::Collector{C}, maxsize::Integer) where C
    buffer_datatypes = C.types
    n_datavectors = length(buffer_datatypes)
    T = Tuple{(Vector{buffer_datatypes[i]} for i in 1:n_datavectors)...}

    return :(TypedBuffer{$T}(0, maxsize, ($([:(Vector{$(buffer_datatypes[i])}(undef, maxsize))
                                             for i in 1:n_datavectors]...),)))
end

"""
    collectdata(writer, args...)

Collect data with the collector in `writer`.

When the the `writer`'s local buffer is full, automatically send the buffered data to the master
process to write it to disk.

# Arguments
* `writer::Writer`: the Writer to use.
* `args...`: arguments to pass on to the collector's functions.
"""
@generated function collectdata(writer::Writer{C,T}, args...) where {C,T}
    n_calls = length(T.parameters)
    ex = :()
    for call_idx in 1:n_calls
        ex = :($(ex.args...);
               @inbounds data[$call_idx][size] = calls[$call_idx](report_args...))
    end

    return quote
        collector = writer.collector
        if collector.condition(args...) == true
            collector.prepare(args...)

            # Hoist from the loop.
            data = writer.localbuffer.data
            calls = collector.calls

            for report_args in collector.iter(args...)
                ensureroom(writer)
                size = (writer.localbuffer.size += 1)
                $ex
            end

            collector.finish(args...)
        end
        return nothing
    end
end

"""
    flushdata(writer)

Send the `writer`'s locally buffered data to the master process, which will write it to disk.

The flushdata function has to be called right before the simulation function returns. This ensures
that any remaining locally buffered data is send to the master proces before the local buffer is
closed.

# Arguments
* `writer::Writer`: the Writer whose buffer to flush.
"""
function flushdata(writer::Writer)
    if myid() == 1 # Make re-using the buffer safe by forcing a copy when the caller is the master.
        put!(writer.remotequeue, deepcopy(writer.localbuffer))
    else
        put!(writer.remotequeue, writer.localbuffer)
    end
    writer.localbuffer.size = 0
    return nothing
end

function ensureroom(writer::Writer)
    if writer.localbuffer.size >= writer.localbuffer.maxsize
        flushdata(writer)
    end
    return nothing
end


text:
const default_p0 = 1e5
const default_κ = 2/7
const default_g = 9.81

function dimensional_π(z, p::Params)
    return z*p.Π
end

function dimensional_p(z, p::Params; p0 = default_p0, κ = default_κ)
    return p0*(z/p.π0)^(1/κ)
end

function dimensional_r(r, p::Params)
    return get_Rd(p)*r
end

function dimensional_pseudoz(z, p::Params; g = default_g)
    θ0 = p.θ0*p.S*p.Π
    cp = p.π0*p.Π
    H = cp*θ0/g
    return H*(1 - z/p.π0)
end

function dimensional_q(q, z, p::Params; p0 = default_p0, κ = default_κ, g = default_g)
    πp = dimensional_π(z,p)/dimensional_p(z, p; p0 = p0, κ = κ)
    return p.S*p.f*g*κ*πp*q
end

function dimensional_u(u, p::Params)
    return get_Rd(p)*p.f*u
end

function redimensionalize_q!(q, d::Domain, p::Params; p0 = default_p0, κ = default_κ, g = default_g)
    for I = CartesianIndices(d)
        x, y, z = d[I]
        q[I] = dimensional_q(q[I], z, p; p0 = p0, κ = κ, g = g)
    end
    return q
end

function redimensionalize_u!(u, d::Domain, p::Params)
    for I = CartesianIndices(d)
        u[I] = dimensional_u(u[I], p)
    end
    return u
end


text:
# These are the tests for deprecated features, they should be deleted along with them

@testset "Trie: path iterator" begin
    t = Trie{Int}()
    t["rob"] = 27
    t["roger"] = 52
    t["kevin"] = Int8(11)
    t0 = t
    t1 = t0.children['r']
    t2 = t1.children['o']
    t3 = t2.children['b']
    @test collect(path(t, "b")) == [t0]
    @test collect(path(t, "rob")) == [t0, t1, t2, t3]
    @test collect(path(t, "robb")) == [t0, t1, t2, t3]
    @test collect(path(t, "ro")) == [t0, t1, t2]
    @test collect(path(t, "roa")) == [t0, t1, t2]
end


text:
using SimJulia

struct TestException <: Exception end

function test_callback(ev::AbstractEvent)
  println("Hi I timed out at $(now(ev))")
end

function test_callback_exception(ev::Event)
  throw(TestException())
end

sim = Simulation()
@callback test_callback(timeout(sim, 1))
@callback test_callback(timeout(sim, 3))
run(sim, 2)
sim = Simulation()
try
  run(sim, Event(sim))
catch exc
  println("$exc has been thrown!")
end
sim = Simulation(3)
start_time = now(sim)
@callback test_callback(timeout(sim, 1))
run(sim, start_time+2)
println(now(sim)-start_time)
sim = Simulation()
start_time = now(sim)
ev = Event(sim)
@callback test_callback_exception(ev)
succeed(ev)
try
  run(sim)
catch exc
  println("$exc has been thrown after $(now(sim)-start_time)!")
end


text:
module GridGraphs

using Test
using Gridap
using Gridap.Helpers

export GridGraph
export FullGridGraph
export cell_nfaces
export nface_cells
export connections
export test_grid_graph
export test_full_grid_graph
import Base: ndims

# Interfaces

"""
Abstract type that provides aggregation info of cell to nfaces and nface to cells
"""
abstract type GridGraph end

function ndims(::GridGraph)
  @abstractmethod
end

function cell_nfaces(g::GridGraph,n::Integer)::IndexCellVector
  @abstractmethod
end

function nface_cells(g::GridGraph,n::Integer)::IndexCellVector
  @abstractmethod
end

function connections(g::GridGraph,from::Integer,to::Integer)
  nd = ndims(g)
  if from == nd
    return cell_nfaces(g,to)
  elseif to == nd
    return nface_cells(g,from)
  else
    @unreachable
  end
end

"""
Abstract type that provides aggregation info of nfaces to mfaces for all n,m
"""
abstract type FullGridGraph <: GridGraph end

function connections(g::FullGridGraph,from::Integer,to::Integer)::IndexCellVector
  @abstractmethod
end

function cell_nfaces(g::FullGridGraph,n::Integer)
  from = ndims(g)
  connections(g,from,n)
end

function nface_cells(g::FullGridGraph,n::Integer)
  to = ndims(g)
  connections(g,n,to)
end

# Testers

function test_grid_graph(g::GridGraph,nd::Integer)
  _nd = ndims(g)
  @test isa(_nd,Integer)
  @test _nd == nd
  for d in 0:nd
    ca = cell_nfaces(g,d)
    cb = connections(g,nd,d)
    @test isa(ca,IndexCellVector{<:Integer})
    @test isa(cb,IndexCellVector{<:Integer})
    @test ca == cb
    ca = nface_cells(g,d)
    cb = connections(g,d,nd)
    @test isa(ca,IndexCellVector{<:Integer})
    @test isa(cb,IndexCellVector{<:Integer})
    @test ca == cb
  end
end

function test_full_grid_graph(g::FullGridGraph,nd::Integer)
  test_grid_graph(g,nd)
  for i in 0:nd
    for j in 0:nd
      ca = connections(g,i,j)
      @test isa(ca,IndexCellVector{<:Integer})
    end
  end
end

# Concrete implementations

struct GridGraphFromData <: GridGraph
  primal::Vector{IndexCellVector}
  dual::Vector{IndexCellVector}
end

function ndims(g::GridGraphFromData)
  n = length(g.primal)
  m = length(g.dual)
  @assert n == m
  n-1
end

function cell_nfaces(g::GridGraphFromData,n::Integer)
  g.primal[n+1]
end

function nface_cells(g::GridGraphFromData,n::Integer)
  g.dual[n+1]
end

struct FullGridGraphFromData <: FullGridGraph
  data::Array{IndexCellArray,2}
end

function ndims(g::FullGridGraphFromData)
  n,m = size(g.data)
  @assert n == m
  n-1
end

function connections(g::FullGridGraphFromData,from::Integer,to::Integer)
  g.data[from+1,to+1]
end

end # module


text:
module Utils

struct CompletedProcess
    stdout::String
    stderr::String
    proc::Base.Process
end

function open_stdin_new(f, cmd)
    out = IOBuffer()
    err = IOBuffer()
    cmd = pipeline(cmd; stderr = err, stdout = out)
    proc = open(cmd, write = true) do proc
        f(proc)
        return proc
    end
    completed = CompletedProcess(String(take!(out)), String(take!(err)), proc)
end

function open_stdin_old(f, cmd)
    inp = Pipe()
    out = Pipe()
    err = Pipe()
    cmd = pipeline(cmd; stdin = inp, stdout = out, stderr = err)
    proc = run(cmd; wait = false)
    close(out.in)
    close(err.in)
    outstr = Ref{String}()
    errstr = Ref{String}()
    @sync begin
        @async outstr[] = read(out, String)
        @async errstr[] = read(err, String)
        try
            f(inp)
        finally
            close(inp)
        end
        wait(proc)
    end
    return CompletedProcess(outstr[], errstr[], proc)
end

if VERSION < v"1.3-"
    open_stdin(args...) = open_stdin_old(args...)
else
    open_stdin(args...) = open_stdin_new(args...)
end

function exec(code)
    julia = Base.julia_cmd()
    script = "include_string(Main, read(stdin, String))"
    cmd = `$julia --startup-file=no -e $script`
    setup = Base.load_path_setup_code()
    cmd = ignorestatus(cmd)
    completed = open_stdin(cmd) do input
        write(input, setup)
        println(input)
        write(input, code)
        close(input)
    end
    @debug(
        "Done `exec(code)`",
        code = Text(code),
        stdout = Text(completed.stdout),
        stderr = Text(completed.stderr),
    )
    return completed
end

Base.success(c::CompletedProcess) = success(c.proc)

end  # module


text:
using DynamicPolynomials
const  MPoly = DynamicPolynomials.MultivariatePolynomials
using GroupsCore
using PermutationGroups

struct VariablePermutation <: SymbolicWedderburn.ByPermutations end

function SymbolicWedderburn.action(a::VariablePermutation, g::PermutationGroups.AbstractPerm, m::Monomial)
    v = variables(m)
    return m(v => SymbolicWedderburn.action(a, g, v))
end

function SymbolicWedderburn.action(::VariablePermutation, g::PermutationGroups.AbstractPerm, v::AbstractVector)
    return map(i -> v[i^g], eachindex(v))
end

abstract type OnMonomials <: SymbolicWedderburn.ByLinearTransformation end

function SymbolicWedderburn.action(
    a::Union{VariablePermutation,OnMonomials},
    el::GroupElement,
    term::MPoly.AbstractTerm,
)
    return MPoly.coefficient(term) * SymbolicWedderburn.action(a, el, MPoly.monomial(term))
end
function SymbolicWedderburn.action(
    a::Union{VariablePermutation,OnMonomials},
    el::GroupElement,
    poly::MPoly.AbstractPolynomial,
)
    return sum([SymbolicWedderburn.action(a, el, term) for term in MPoly.terms(poly)])
end

function SymbolicWedderburn.decompose(
    k::MPoly.AbstractPolynomialLike,
    hom::SymbolicWedderburn.InducedActionHomomorphism,
)
    # correct only if basis(hom) == monomials
    I = SymbolicWedderburn._int_type(hom)
    indcs = I[hom[mono] for mono in MPoly.monomials(k)]
    coeffs = MPoly.coefficients(k)

    return indcs, coeffs
end


text:
using QMTK
import QMTK: RealSpaceType, UnRandomizedError
using Compat.Test

@test Float64 <: RealSpaceType{Float64}
@test (Float64 <: RealSpaceType{Float32}) == false

@testset "Constructors" begin

# setUp
space = RealSpace()
@test space.data < space.max
@test space.data >= space.min
@test typeof(space.data) <: Float64
@test israndomized(space) == true

space = RealSpace(2, 2, min=-1, max=1)
@test size(space.data) == (2, 2)
@test all(space.data .< 1)
@test all(space.data .>= -1)
@test typeof(space.data) <: Array{Float64, 2}
@test israndomized(space) == true

space = RealSpace(Float32, 2, 2, min=-1, max=1)
@test size(space.data) == (2, 2)
@test all(space.data .< 1)
@test all(space.data .>= -1)
@test typeof(space.data) <: Array{Float32, 2}

end

@testset "methods" begin

    space = RealSpace()
    space = reset!(space)
    @test space.data == space.min
    @test israndomized(space) == false

    data = acquire(space)
    @test space.data === data

    @test_throws UnRandomizedError shake!(space)
    space = randomize!(space)
    @test israndomized(space) == true

    space = RealSpace(2, 2)
    space = reset!(space)
    @test all(space.data .== space.min)
    @test israndomized(space) == false

    data = acquire(space)
    @test space.data !== data

    @test_throws UnRandomizedError shake!(space)
    space = randomize!(space)
    @test israndomized(space) == true

    # TODO: Test distribution after binning is implemented
    # @test shake!(space)
end

@testset "Iterator" begin

    space = RealSpace(min=1, max=2.1)
    @test collect(1:0.2:2.1) == collect(traverse(space, step=0.2))

end


text:
@testset "Index creation" begin
    t = Table(a = [11,12,12], b = [2,3,1], c = ["a", "b", "c"])

    @test NoIndex(t) === NoIndex()

    @test UniqueIndex{(:b,)}(t) === UniqueIndex{(:b,)}()

    @test (SortIndex{(:a,)}(t)::SortIndex{(:a,)}).order == [1, 2, 3]
    @test (SortIndex{(:a,:b)}(t)::SortIndex{(:a,:b)}).order == [1, 3, 2]

    @test (UniqueSortIndex{(:b,)}(t)::UniqueSortIndex{(:b,)}).order == [3, 1, 2]

    @test (HashIndex{(:a,)}(t)::HashIndex{(:a,)}).dict == Dict((11,) => [1], (12,) => [2, 3])
    @test (HashIndex{(:a,:b)}(t)::HashIndex{(:a,:b)}).dict == Dict((11, 2) => [1], (12, 3) => [2], (12, 1) => [3])

    @test (UniqueHashIndex{(:c,)}(t)::UniqueHashIndex{(:c,)}).dict == Dict(("a",) => 1, ("b",) => 2, ("c",) => 3)

    t2 = accelerate(t, SortIndex{(:b,)})
    @test accelerators(t2) isa Tuple{SortIndex{(:b,)}}
    @test accelerators(t2)[1].order == [3, 1, 2]
end


text:
# --------------------------------------------------------------------------- #
# Binary transformators
# --------------------------------------------------------------------------- #

function eval(block::B) where {B <: AbstractBinaryTransformer}
    error("not implemented yet")
end

function eval(block::B, time) where {B <: AbstractBinaryTransformer}
    return eval(block)
end

mutable struct Plus{T} <: AbstractBinaryTransformer
    x1::InputPin{T}
    x2::InputPin{T}
    y::OutputPin{T}

    function Plus{T}() where {T}
        x1 = InputPin{T}()
        x2 = InputPin{T}()
        y  = OutputPin{T}()

        obj = new()

        obj.x1 = x1
        obj.x2 = x2
        obj.y  = y

        x1.containing_object = obj
        x2.containing_object = obj
        y.containing_object  = obj

        return obj
    end
end
function Plus() where {T}
    return Plus{Float64}()
end

function eval(block::Plus)
    val = block.x1.value + block.x2.value
    block.y.value = val
    return val
end

mutable struct Minus{T} <: AbstractBinaryTransformer
    x1::InputPin{T}
    x2::InputPin{T}
    y::OutputPin{T}

    function Minus{T}() where {T}
        x1 = InputPin{T}()
        x2 = InputPin{T}()
        y  = OutputPin{T}()

        obj = new()

        obj.x1 = x1
        obj.x2 = x2
        obj.y  = y

        x1.containing_object = obj
        x2.containing_object = obj
        y.containing_object  = obj

        return obj
    end
end
function Minus() where {T}
    return Minus{Float64}()
end

function eval(block::Minus)
    val = block.x1.value - block.x2.value
    block.y.value = val
    return val
end

mutable struct Multiply{T} <: AbstractBinaryTransformer
    x1::InputPin{T}
    x2::InputPin{T}
    y::OutputPin{T}

    function Multiply{T}() where {T}
        x1 = InputPin{T}()
        x2 = InputPin{T}()
        y  = OutputPin{T}()

        obj = new()

        obj.x1 = x1
        obj.x2 = x2
        obj.y  = y

        x1.containing_object = obj
        x2.containing_object = obj
        y.containing_object  = obj

        return obj
    end
end
function Multiply() where {T}
    return Multiply{Float64}()
end

function eval(block::Multiply)
    val = block.x1.value * block.x2.value
    block.y.value = val
    return val
end

mutable struct Divide{T} <: AbstractBinaryTransformer
    x1::InputPin{T}
    x2::InputPin{T}
    y::OutputPin{T}

    function Divide{T}() where {T}
        x1 = InputPin{T}()
        x2 = InputPin{T}()
        y  = OutputPin{T}()

        obj = new()

        obj.x1 = x1
        obj.x2 = x2
        obj.y  = y

        x1.containing_object = obj
        x2.containing_object = obj
        y.containing_object  = obj

        return obj
    end
end
function Divide() where {T}
    return Divide{Float64}()
end

function eval(block::Divide)
    val = block.x1.value / block.x2.value
    block.y.value = val
    return val
end

max_stars_repo_name: HomoModelicus/julia

max_issues_repo_name: HomoModelicus/julia

max_forks_repo_name: HomoModelicus/julia


text:
using JLD
using Printf

result = load("NonlinElliptic2d_data_freq600s6.jld")["result"]

arr_kernel = ["Matern5half", "Matern7half", "Matern9half"]

# arr_kernel = ["Matern5half"]
arr_h = [0.02,0.01,0.005,0.0025]
# arr_ρ = [2.0, 3.0, 4.0]
arr_ρ = [4.0]

data = Dict()

for kernel in arr_kernel
    data[kernel] = Dict()
    for ρ in arr_ρ
        data[kernel][ρ] = Dict()
        L2err = zeros(length(arr_h))
        Linferr = zeros(length(arr_h))
        time = zeros(length(arr_h))
        for idx in 1:length(arr_h)
            L2err[idx] = result[("kernel", kernel)][("h",arr_h[idx])][("rho", ρ)]["L2"]
            Linferr[idx] = result[("kernel", kernel)][("h",arr_h[idx])][("rho", ρ)]["Linf"]
            time[idx] = result[("kernel", kernel)][("h",arr_h[idx])][("rho", ρ)]["time"]
        end
        data[kernel][ρ]["L2"] = L2err
        data[kernel][ρ]["Linf"] = Linferr
        data[kernel][ρ]["time"] = time
    end
end

function linreg1d(x,y)
    ymean = sum(y)/length(y)
    xmean = sum(x)/length(x)
    β = sum((y.-ymean).*(x.-xmean))/sum((x.-xmean).*(x.-xmean))
    b = y .- β*x
    return β, b
end

using PyCall
using PyPlot
fsize = 15.0
tsize = 15.0
tdir = "in"
major = 5.0
minor = 3.0
lwidth = 0.8
lhandle = 2.0
plt.style.use("default")
rcParams = PyDict(matplotlib["rcParams"])
# rcParams["text.usetex"] = true
# I am not able to make tex font work. Using the anaconda pyplot environment leads to MKL errors, while the julia python environment cannot support tex font.
# need to ask Daniel how to have good plots
# rcParams["text.latex.preamble"] = [raw"\usepackage{amsmath}"]
rcParams["font.size"] = fsize
rcParams["legend.fontsize"] = tsize
rcParams["xtick.direction"] = tdir
rcParams["ytick.direction"] = tdir
rcParams["xtick.major.size"] = major
rcParams["xtick.minor.size"] = minor
rcParams["ytick.major.size"] = 5.0
rcParams["ytick.minor.size"] = 3.0
rcParams["axes.linewidth"] = lwidth
rcParams["legend.handlelength"] = lhandle

# @show result

arr_N = (1.0./arr_h.+1).^2

plotL2 = plt.figure()
ax = plotL2.add_subplot(111)
for kernel in arr_kernel
    for ρ in arr_ρ
        X = arr_N
        Y = data[kernel][ρ]["L2"]
        slope, b = linreg1d(log.(X),log.(Y))
        kernelname = kernel[1:end-4] * "/2"
        ax.plot(X, Y, "-s", label = kernelname*L", $\rho=$"*"$ρ, slope $(@sprintf("%.2f", slope))")
    end
end
plt.xscale("log")
plt.yscale("log")
plt.xlabel(L"$N_{\mathrm{domain}}}$")
plt.ylabel(L"$L^{2}$ error")
plt.title(L"Solution Accuracy $L^{2}$")
plt.legend()
# plt.legend(bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0.)
display(plotL2)
plotL2.savefig("L2err_Cholesky_elliptic.pdf", bbox_inches="tight",dpi=100,pad_inches=0.15)

plotLinf = plt.figure()
ax = plotLinf.add_subplot(111)
for kernel in arr_kernel
    for ρ in arr_ρ
        X = arr_N
        Y = data[kernel][ρ]["Linf"]
        kernelname = kernel[1:end-4] * "/2"
        slope, b = linreg1d(log.(X),log.(Y))
        ax.plot(X, Y, "-s", label = kernelname*L", $\rho=$"*"$ρ, slope $(@sprintf("%.2f", slope))")
    end
end
plt.xscale("log")
plt.yscale("log")
plt.xlabel(L"$N_{\mathrm{domain}}}$")
plt.ylabel(L"$L^{\infty}$ error")
plt.title(L"Solution Accuracy $L^{\infty}$")
plt.legend()
# plt.legend(bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0.)
display(plotLinf)

plottime = plt.figure()
ax = plottime.add_subplot(111)
for kernel in arr_kernel
    for ρ in arr_ρ
        X = arr_N
        Y = data[kernel][ρ]["time"]
        slope, b = linreg1d(log.(X),log.(Y))
        kernelname = kernel[1:end-4] * "/2"
        ax.plot(X, Y, "-s", label = kernelname*L", $\rho=$"*"$ρ, slope $(@sprintf("%.2f", slope))")
    end
end
plt.xscale("log")
plt.yscale("log")
plt.xlabel(L"$N_{\mathrm{domain}}}$")
plt.ylabel("Time (s)")
plt.title("Computational Efficiency")
plt.legend()
# plt.legend(bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0.)
display(plottime)
plottime.savefig("time_Cholesky_elliptic.pdf", bbox_inches="tight",dpi=100,pad_inches=0.15)

plottradeoff = plt.figure()
ax = plottradeoff.add_subplot(111)
for kernel in arr_kernel
    for ρ in arr_ρ
        X = data[kernel][ρ]["time"]
        Y = data[kernel][ρ]["L2"]
        slope, b = linreg1d(log.(X),log.(Y))
        kernelname = kernel[1:end-4] * "/2"
        ax.plot(X, Y, "-s", label = kernelname*L", $\rho=$"*"$ρ, slope $(@sprintf("%.2f", slope))")
    end
end
plt.xscale("log")
plt.yscale("log")
plt.xlabel("Time")
plt.ylabel(L"$L^{2}$ error")
plt.title("Speed–Accuracy Tradeoff")
plt.legend()
# plt.legend(bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0.)
display(plottradeoff)


text:
@recipe function dummy(curve::ROCData)
    xlim := (0,1)
    ylim := (0,1)
    xlab := "false positive rate"
    ylab := "true positive rate"
    title --> "Receiver Operator Characteristic"
    @series begin
        color --> :black
        linestyle --> :dash
        label := ""
        [0, 1], [0, 1]
    end
    @series begin
        curve.FPR, curve.TPR
    end
end


text:
# MIT license
# Copyright (c) Microsoft Corporation. All rights reserved.
# See LICENSE in the project root for full license information.

using DelimitedFiles: readdlm # used in sourcefile_to_catalog
using StringEncodings
using StaticArrays
using Unitful
import Unitful: Length, Temperature, Quantity, Units

"""
    generate_jls(sourcenames::Vector{<:AbstractString}, mainfile::AbstractString, jldir::AbstractString, sourcedir::AbstractString; test::Bool = false)

Generates .jl files in `jldir`: a `mainfile` and several catalog files.

Each catalog file is a module representing a distinct glass catalog (e.g. NIKON, SCHOTT), generated from corresponding
AGF files in `sourcedir`. These are then included and exported in `mainfile`.

In order to avoid re-definition of constants `AGF_GLASS_NAMES` and `AGF_GLASSES` during testing, we have an optional
`test` argument. If `true`, we generate a .jl file that defines glasses with `GlassType.TEST` to avoid namespace/ID
clashes.
"""
function generate_jls(
    sourcenames::Vector{<:AbstractString},
    mainfile::AbstractString,
    jldir::AbstractString,
    sourcedir::AbstractString;
    test::Bool = false
)
    glasstype = test ? "TEST" : "AGF"
    id = 1
    catalogfiles = []
    glassnames = []

    # generate several catalog files (.jl)
    for catalogname in sourcenames
        # parse the sourcefile (.agf) into a catalog (native Julia dictionary)
        sourcefile = joinpath(sourcedir, "$(catalogname).agf")
        catalog = sourcefile_to_catalog(sourcefile)

        # parse the catalog into a module string and write it to a catalog file (.jl)
        id, modstring = catalog_to_modstring(id, catalogname, catalog, glasstype)
        push!(catalogfiles, "$(catalogname).jl")
        catalogpath = joinpath(jldir, catalogfiles[end])
        @info "Writing $catalogpath"
        open(catalogpath, "w") do io
            write(io, modstring)
        end

        # track glass names for lookup purposes
        append!(glassnames, ["$(catalogname).$(glassname)" for glassname in keys(catalog)])
    end

    # generate the parent main file (.jl)
    agfstrings = [
        "export $(join(sourcenames, ", "))",
        "",
        ["include(\"$(catalogfile)\")" for catalogfile in catalogfiles]...,
        "",
        "const $(glasstype)_GLASS_NAMES = [$(join(repr.(glassnames), ", "))]",
        "const $(glasstype)_GLASSES = [$(join(glassnames, ", "))]",
        ""
    ]
    @info "Writing $mainfile"
    open(mainfile, "w") do io
        write(io, join(agfstrings, "\n"))
    end
end

"""
Parse a `sourcefile` (.agf) into a native Dict, `catalogdict`, where each `kvp = (glassname, glassinfo)` is a glass.

| 1  | 2        | 3         | 4     | 5     | 6                     | 7             | 8         | 9         | 10    | 11    |
|:---|:---------|:----------|:------|:------|:----------------------|:--------------|:----------|:----------|:------|:------|
| NM | raw name | dispform  | ?     |  Nd   | Vd                    | [exclude_sub  | status    | meltfreq] |
| ED | TCE      | ?         | p     |  ΔPgF | [ignore_thermal_exp]  |
| CD | C1       | C2        | C3    |  C4   | C5                    | C6            | C7        | C8        | C9    | C10   |
| TD | D₀       | D₁        | D₂    |  E₀   | E₁                    | λₜₖ           | temp      |
| OD | relcost  | CR        | FR    |  SR   | AR                    | PR
| LD | λmin     | λmax      |
| IT | T1       | T2        | T3    |
"""
function sourcefile_to_catalog(sourcefile::AbstractString)
    catalogdict = Dict{String,Dict{String}}()

    # store persistent variables between loops
    glassname = ""
    rowbuffer = []

    # `rowspecs` provides tuples (keys, defaultvalues) which describe how to parse each tokenized row
    nullspec = ("", [])
    rowspecs = Dict(
        "CC" => nullspec,
        "GC" => nullspec,
        "NM" => (
            "raw_name dispform _        Nd   Vd   exclude_sub status meltfreq",
            [nothing, NaN,     nothing, NaN, NaN, 0,          0,     0       ]
        ),
        "ED" => (
            "TCE  _        p    ΔPgF ignore_thermal_exp",
            [NaN, nothing, NaN, NaN, 0                 ]
        ),
        "CD" => (
            join([string('C', i) for i in 1:10], ' '),
            repeat([NaN], 10)
        ),
        "TD" => (
            "D₀ D₁ D₂ E₀ E₁ λₜₖ temp",
            [0, 0, 0, 0, 0, 0, 20  ]
        ),
        "OD" => (
            "relcost CR FR SR AR PR",
            repeat([-1], 6)
        ),
        "LD" => (
            "λmin λmax",
            [NaN, NaN ]
        ),
        "IT" => nullspec,
    )

    # call process_rowbuffer!() when rowbuffer is a vector corresponding to a full AGF row
    # e.g. rowbuffer = ["OD", -1.0, 5.0, 5.0, 5.0, -1.0, -1.0]
    # the function will parse the row data into catalogdict and flush the row buffer
    function process_rowbuffer!()
        # an empty rowbuffer gets passed at the beginning of each file
        if isempty(rowbuffer) return end
        token = rowbuffer[1]

        # NM and IT need special treatment
        if token == "NM"
            glassname = make_valid_name(string(rowbuffer[2]))
            catalogdict[glassname] = Dict{String,Any}()
            catalogdict[glassname]["transmission"] = Vector{SVector{3,Float64}}(undef, 0)
        elseif token == "IT"
            if rowbuffer[2] != 0 && length(rowbuffer) == 4
                push!(catalogdict[glassname]["transmission"], SVector{3,Float64}(rowbuffer[2:4]))
            end
        end

        if token ∉ keys(rowspecs)
            @error "Unrecognised row token: $token" rowbuffer
            return
        end
        rowspec = rowspecs[token]

        # parse row buffer into catalogdict, according to rowspec
        for (i, (key, defaultvalue)) in enumerate(zip(split(rowspec[1]), rowspec[2]))
            value = length(rowbuffer) < i + 1 || rowbuffer[i + 1] == "-" ? defaultvalue : rowbuffer[i + 1]
            catalogdict[glassname][key] = value
        end

        # flush rowbuffer
        rowbuffer = []
    end

    is_utf8 = isvalid(readuntil(sourcefile, " "))
    # use DelimitedFiles.readdlm to parse the source file conveniently (with type inference)
    for line in eachrow(readdlm(sourcefile))
        for item in line
            if !is_utf8
                item = decode(Vector{UInt8}(item), "UTF-16")
                _item = tryparse(Float64, item)
                item = _item === nothing ? item : _item
            end

            if item == "" # eol
                break
            elseif item ∈ keys(rowspecs) # process buffer when a token is reached, instead of at eol (see Issue #106)
                process_rowbuffer!()
            end

            push!(rowbuffer, item)
        end
    end
    process_rowbuffer!()

    return catalogdict
end

"""
Make a valid Julia variable name from an arbitrary string.
"""
function make_valid_name(name::AbstractString)
    # remove invalid characters
    name = replace(name, "*" => "_STAR")
    name = replace(name, r"""[ ,.:;?!()&-]""" => "_")
    # cant have module names which are just numbers so add a _ to the start
    if tryparse(Int, "$(name[1])") !== nothing
        name = "_" * name
    end
    return name
end

"""
Convert a `catalog` dict into a `modstring` which can be written to a Julia source file.
"""
function catalog_to_modstring(
    start_id::Integer, catalogname::AbstractString, catalog::Dict{<:AbstractString}, glasstype::AbstractString
)
    id = start_id
    isCI = haskey(ENV, "CI")

    modstrings = [
        "module $catalogname",
        "using ..GlassCat: Glass, GlassID, $glasstype",
        "export $(join(keys(catalog), ", "))",
        ""
    ]
    for (glassname, glassinfo) in catalog
        argstring = glassinfo_to_argstring(glassinfo, id, glasstype)
        push!(modstrings, "const $glassname = Glass($argstring)")
        id += 1
    end
    append!(modstrings, ["end # module", ""]) # last "" is for \n at EOF

    return id, join(modstrings, "\n")
end

"""
Convert a `glassinfo` dict into a `docstring` to be prepended to a `Glass` const.
"""
function glassinfo_to_docstring(
    glassinfo::Dict{<:AbstractString},
    id::Integer,
    catalogname::AbstractString,
    glassname::AbstractString,
    glasstype::AbstractString
)
    raw_name = glassinfo["raw_name"] == glassname ? "" : " ($(glassinfo["raw_name"]))"
    pad(str, padding=25) = rpad(str, padding)
    getinfo(key, default=0.0) = get(glassinfo, key, default)

    return join([
        "\"\"\"    $catalogname.$glassname$raw_name",
        "```",
        "$(pad("ID:"))$glasstype:$id",
        "$(pad("RI @ 587nm:"))$(getinfo("Nd"))",
        "$(pad("Abbe Number:"))$(getinfo("Vd"))",
        "$(pad("ΔPgF:"))$(getinfo("ΔPgF"))",
        "$(pad("TCE (÷1e-6):"))$(getinfo("TCE"))",
        "$(pad("Density:"))$(getinfo("p"))g/m³",
        "$(pad("Valid wavelengths:"))$(getinfo("λmin"))μm to $(getinfo("λmax"))μm",
        "$(pad("Reference Temp:"))$(getinfo("temp", 20.0))°C",
        "```",
        "\"\"\""
    ], "\n")
end

"""
Convert a `glassinfo` dict into an `argstring` to be passed into a `Glass` constructor.
"""
function glassinfo_to_argstring(glassinfo::Dict{<:AbstractString}, id::Integer, glasstype::AbstractString)
    argstrings = []
    for fn in string.(fieldnames(Glass))
        if fn == "ID"
            push!(argstrings, "GlassID($glasstype, $id)")
        elseif fn in ["D₀", "D₁", "D₂", "E₀", "E₁", "λₜₖ"]
            push!(argstrings, repr(get(glassinfo, fn, 0.0)))
        elseif fn == "temp"
            push!(argstrings, repr(get(glassinfo, fn, 20.0)))
        elseif fn == "transmission"
            v = get(glassinfo, "transmission", nothing)
            if isnothing(v)
                push!(argstrings, repr(nothing))
            else
                str = join(["($(join(a, ", ")))" for a in v], ", ")
                push!(argstrings, "[$str]")
            end
        elseif fn == "transmissionN"
            continue
        else
            push!(argstrings, repr(get(glassinfo, fn, NaN)))
        end
    end
    return join(argstrings, ", ")
end


text:
abstract type AbstractFloat8 <: AbstractFloat end
primitive type Float8 <: AbstractFloat8 8 end        # standard 3 exp version
primitive type Float8_4 <: AbstractFloat8 8 end      # version with 4 exp bits

# conversions
Float8(x::UInt8) = reinterpret(Float8,x)
Float8_4(x::UInt8) = reinterpret(Float8_4,x)
UInt8(x::T) where {T<:AbstractFloat8} = reinterpret(UInt8,x)
bitstring(x::AbstractFloat8) = bitstring(reinterpret(UInt8,x))
Float8(x::T) where {T<:Union{Int16,Int32,Int64,Float64,Float16}} = Float8(Float32(x))
Float8_4(x::T) where {T<:Union{Int16,Int32,Int64,Float64,Float16}} = Float8_4(Float32(x))
(::Type{T})(x::AbstractFloat8) where {T<:Union{Int16,Int32,Int64,Float64,Float16}} = T(Float32(x))

# masks, UInt8s with 1s for the respective parts
sign_mask(::Type{T}) where {T<:AbstractFloat8} = 0x80
exponent_mask(::Type{Float8}) = 0x70
exponent_mask(::Type{Float8_4}) = 0x78
significand_mask(::Type{Float8}) = 0x0f
significand_mask(::Type{Float8_4}) = 0x07

sign_mask(::Type{Float32}) =            0x8000_0000
exponent_mask(::Type{Float32}) =        0x7f80_0000
significand_mask(::Type{Float32}) =     0x007f_ffff
n_exponent_bits(::Type{Float32}) =      8
n_significant_bits(::Type{Float32}) =   23

# number of exp/sig bits
n_exponent_bits(::Type{Float8}) = 3
n_exponent_bits(::Type{Float8_4}) = 4
n_significant_bits(::Type{Float8}) = 4
n_significant_bits(::Type{Float8_4}) = 3
bias(::Type{Float8}) = 3
bias(::Type{Float8_4}) = 7

eps(::Type{Float8}) = Float8(0x02)
eps(::Type{Float8_4}) = Float8_4(0x20)

# define inifinities and nan
inf8(::Type{Float8}) = Float8(0x70)
inf8(::Type{Float8_4}) = Float8_4(0x78)

nan8(::Type{Float8}) = Float8(0x78)
nan8(::Type{Float8_4}) = Float8_4(0x7c)

const NaN8 = nan8(Float8)
const Inf8 = inf8(Float8)

const NaN8_4 = nan8(Float8_4)
const Inf8_4 = inf8(Float8_4)

typemin(::Type{T}) where {T<:AbstractFloat8} = -inf8(T)
typemax(::Type{T}) where {T<:AbstractFloat8} = inf8(T)

# smallest non-subnormal number, largest representable number
floatmin(::Type{Float8}) = Float8(0x10)
floatmin(::Type{Float8_4}) = Float8_4(0x08)
floatmax(::Type{Float8}) = Float8(0x6f)
floatmax(::Type{Float8_4}) = Float8_4(0x77)

# one and zero element
one(::Type{Float8}) = Float8(0x30)
one(::Type{Float8_4}) = Float8_4(0x38)
zero(::Type{T}) where {T<:AbstractFloat8} = Float8(0x00)

one(x::AbstractFloat8) = one(typeof(x))
zero(x::AbstractFloat8) = zero(typeof(x))

# is positive zero 0x00 or negative zero 0x80
iszero(x::AbstractFloat8) = reinterpret(UInt8, x) == 0x00 || reinterpret(UInt8,x) == 0x80
-(x::T) where {T<:AbstractFloat8} = reinterpret(T, reinterpret(UInt8, x) ⊻ 0x80)
Bool(x::AbstractFloat8) = iszero(x) ? false : isone(x) ? true : throw(InexactError(:Bool, Bool, x))

abs(x::T) where {T<:AbstractFloat8} = reinterpret(T, reinterpret(UInt8, x) & 0x7f)
isnan(x::T) where {T<:AbstractFloat8} = reinterpret(UInt8,x) & 0x7f > exponent_mask(T)
isfinite(x::T) where {T<:AbstractFloat8} = reinterpret(UInt8,x) & exponent_mask(T) != exponent_mask(T)

precision(::Type{Float8}) = 5
precision(::Type{Float8_4}) = 4

signbit(x::AbstractFloat8) = UInt8(x) > 0x7f

function sign(x::T) where {T<:AbstractFloat8}
    if isnan(x) || iszero(x)
        return x
    elseif signbit(x)
        return -one(T)
    else
        return one(T)
    end
end

first_sig_bit_mask(::Type{Float8}) = 0x00000008
first_sig_bit_mask(::Type{Float8_4}) = 0x00000004

sig_bit_shift(::Type{Float8}) = 19          # 23 significand bits for Float32 - 4 significand bits for Float8
sig_bit_shift(::Type{Float8_4}) = 20        # 23 significand bits for Float32 - 3 significand bits for Float8_4

bias_difference(::Type{Float8}) = 0x0000007c        # = 124, 127 for Float32 minus 3 for Float 8
bias_difference(::Type{Float8_4}) = 0x00000078      # = 120, 127 for Float32 minus 7 for Float 8_4

exp_bits_all_one(::Type{Float8}) = 0x00000007
exp_bits_all_one(::Type{Float8_4}) = 0x0000000f

round(x::T, r::RoundingMode{:ToZero}) where {T<:AbstractFloat8} = T(round(Float32(x), r))
round(x::T, r::RoundingMode{:Down}) where {T<:AbstractFloat8} = T(round(Float32(x), r))
round(x::T, r::RoundingMode{:Up}) where {T<:AbstractFloat8} = T(round(Float32(x), r))
round(x::T, r::RoundingMode{:Nearest}) where {T<:AbstractFloat8} = T(round(Float32(x), r))

function ==(x::AbstractFloat8, y::AbstractFloat8)
    if isnan(x) || isnan(y)     # Alternatively, For Float16: (ix|iy)&0x7fff > 0x7c00
        return false
    end
    if iszero(x) && iszero(y)   # For Float16: (ix|iy)&0x7fff == 0x0000
        return true
    end
    return reinterpret(UInt8,x) == reinterpret(UInt8,y)
end

for op in (:<, :<=, :isless)
    @eval ($op)(a::T, b::T) where {T<:AbstractFloat8} = ($op)(Float32(a), Float32(b))
end

for op in (:+, :-, :*, :/, :\, :^)
    @eval ($op)(a::Float8, b::Float8) = Float8(($op)(Float32(a), Float32(b)))
    @eval ($op)(a::Float8_4, b::Float8_4) = Float8_4(($op)(Float32(a), Float32(b)))
end

for func in (:sin,:cos,:tan,:asin,:acos,:atan,:sinh,:cosh,:tanh,:asinh,:acosh,
             :atanh,:exp,:exp2,:exp10,:expm1,:log,:log2,:log10,:sqrt,:cbrt,:lgamma,:log1p)
    @eval begin
        $func(a::Float8) = Float8($func(Float32(a)))
        $func(a::Float8_4) = Float8_4($func(Float32(a)))
    end
end

for func in (:atan,:hypot)
    @eval begin
        $func(a::Float8,b::Float8) = Float8($func(Float32(a),Float32(b)))
        $func(a::Float8_4,b::Float8_4) = Float8_4($func(Float32(a),Float32(b)))
    end
end

function Base.show(io::IO,x::Float8)
    if isnan(x)
        print(io,"NaN8")
    elseif isinf(x)
        if UInt8(x) > 0x80  # is negative?
            print(io,"-Inf8")
        else
            print(io,"Inf8")
        end
    else
        io2 = IOBuffer()
        print(io2,Float32(x))
        f = String(take!(io2))
        print(io,"Float8("*f*")")
    end
end

function Base.show(io::IO,x::Float8_4)
    if isnan(x)
        print(io,"NaN8_4")
    elseif isinf(x)
        if UInt8(x) > 0x80  # is negative?
            print(io,"-Inf8_4")
        else
            print(io,"Inf8_4")
        end
    else
        io2 = IOBuffer()
        print(io2,Float32(x))
        f = String(take!(io2))
        print(io,"Float8_4("*f*")")
    end
end

function nextfloat(x::T) where {T<:AbstractFloat8}
    if isnan(x) || x == inf8(T)
        return x
    elseif x == -zero(T)
        return T(0x01)
    elseif UInt8(x) < 0x80  # positive numbers
        return T(UInt8(x)+0x1)
    else                    # negative numbers
        return T(UInt8(x)-0x1)
    end
end

function prevfloat(x::T) where {T<:AbstractFloat8}
    if isnan(x) || x == -inf8(T)
        return x
    elseif x == zero(T)
        return T(0x81)
    elseif UInt8(x) < 0x80
        return T(UInt8(x)-0x1)
    else
        return T(UInt8(x)+0x1)
    end
end

promote_rule(::Type{Float8}, ::Type{Float64}) = Float64
promote_rule(::Type{Float8}, ::Type{Float32}) = Float32
promote_rule(::Type{Float8}, ::Type{Float16}) = Float16

promote_rule(::Type{Float8}, ::Type{Int64}) = Float8
promote_rule(::Type{Float8}, ::Type{Int32}) = Float8
promote_rule(::Type{Float8}, ::Type{Int16}) = Float8

promote_rule(::Type{Float8}, ::Type{Bool}) = Float8

promote_rule(::Type{Float8_4}, ::Type{Float64}) = Float64
promote_rule(::Type{Float8_4}, ::Type{Float32}) = Float32
promote_rule(::Type{Float8_4}, ::Type{Float16}) = Float16

promote_rule(::Type{Float8_4}, ::Type{Int64}) = Float8
promote_rule(::Type{Float8_4}, ::Type{Int32}) = Float8
promote_rule(::Type{Float8_4}, ::Type{Int16}) = Float8

promote_rule(::Type{Float8_4}, ::Type{Bool}) = Float8


text:
function substitute_underscores!(underscores_to_gensyms, meta_level, other)
    other
end
function substitute_underscores!(underscores_to_gensyms, meta_level, maybe_argument::Symbol)
    if meta_level < 1 && all(isequal('_'), string(maybe_argument))
        if !haskey(underscores_to_gensyms, maybe_argument)
            underscores_to_gensyms[maybe_argument] = gensym(maybe_argument)
        end
        underscores_to_gensyms[maybe_argument]
    else
        maybe_argument
    end
end
function substitute_underscores!(underscores_to_gensyms, meta_level, code::Expr)
    # have to do this the old fashioned way because _ has a special meaning in MacroTools
    expanded_code = if code.head === :macrocall && length(code.args) === 3
        name, location, body = code.args
        if name === Symbol("@_")
            anonymous(location, body)
        elseif name === Symbol("@>")
            make_chain(location, body)
        else
            code
        end
    else
        code
    end
    head = expanded_code.head
    new_meta_level = if head == :quote
        meta_level + 1
    elseif head == :$
        meta_level - 1
    else
        meta_level
    end
    Expr(
        head,
        (
            substitute_underscores!(underscores_to_gensyms, new_meta_level, code) for
            code in expanded_code.args
        )...,
    )
end

function anonymous(location, other)
    other
end

function anonymous(location, body::Expr)
    underscores_to_gensyms = Dict{Symbol,Symbol}()
    meta_level = 0
    substituted_body = substitute_underscores!(underscores_to_gensyms, meta_level, body)
    if length(underscores_to_gensyms) == 0
        body
    else
        Expr(
            :function,
            Expr(
                :call,
                gensym("@_"),
                Iterators.map(value, sort!(collect(underscores_to_gensyms), by = first))...,
            ),
            Expr(:block, location, substituted_body),
        )
    end
end

"""
    macro _(body)

Terser function syntax. The arguments are inside the `body`; the first argument is `_`, the second argument is `__`, etc. Will `@inline`.

```jldoctest anonymous
julia> using LightQuery

julia> using Test: @inferred

julia> @inferred (@_ _ + 1)(1)
2

julia> @inferred map((@_ __ - _), (1, 2), (2, 1))
(1, -1)
```

If there are no `_` arguments, read as is.

```jldoctest anonymous
julia> (@_ x -> x + 1)(1)
2
```
"""
macro _(body)
    anonymous(__source__, body) |> esc
end
export @_

function link(location, object, call::Expr)
    underscores_to_gensyms = Dict{Symbol,Symbol}()
    meta_level = 0
    body = substitute_underscores!(underscores_to_gensyms, meta_level, call)
    if length(underscores_to_gensyms) == 0
        Expr(:block, location, Expr(:call, call, object))
    elseif length(underscores_to_gensyms) == 1 && haskey(underscores_to_gensyms, :_)
        Expr(:let, Expr(:(=), underscores_to_gensyms[:_], object), Expr(:block, location, body))
    else
        throw(ArgumentError("Chain segments must contain single underscores only"))
    end

end
function link(location, object, call)
    Expr(:call, call, object)
end

function make_chain(location, maybe_chain)
    if @capture maybe_chain object_ |> call_
        link(location, make_chain(location, object), call)
    else
        maybe_chain
    end
end

"""
    macro >(body)

If body is in the form `object_ |> call_`, call [`@_`](@ref) on `call`, and recur on `object`.

```jldoctest chain
julia> using LightQuery

julia> @> 0 |> _ - 1 |> abs |> Base.abs
1
```

You can't include multiple arguments.

```jldoctest chain
julia> @> _ |> _ + __
ERROR: LoadError: ArgumentError: Chain segments must contain single underscores only
[...]
```

You can nest chains:

```jldoctest chain
julia> @> 1 |> (@> _ + 1 |> _ + 1)
3
```

Handles interpolations seamlessly:

```jldoctest chain
julia> @> 1 |> :(_ + \$_)
:(_ + 1)
```
"""
macro >(body)
    make_chain(__source__, body) |> esc
end
export @>


text:
# # Example: Linear regression

# ## Defining the linear regression model

# In this example, we fit a Bayesian linear regression model with the
# canonical link function.

# Suppose that we are given a matrix of features `X` and a column vector of
# labels `y`. `X` has `n` rows and `p` columns. `y` has `n` elements. We assume
# that our observation vector `y` is a realization of a random variable `Y`.
# We define `μ` (mu) as the expected value of `Y`, i.e. `μ := E[Y]`. Our model
# comprises three components:
#
# 1. The probability distribution of `Y`: for linear regression, we assume that each `Yᵢ` follows a normal distribution with mean `μᵢ` and variance `σ²`.
# 2. The systematic component, which consists of linear predictor `η` (eta), which we define as `η := α + Xβ`, where `α` is the scalar intercept and `β` is the column vector of `p` coefficients.
# 3. The link function `g`, which provides the following relationship: `g(E[Y]) = g(μ) = η = Xβ`. It follows that `μ = g⁻¹(η)`, where `g⁻¹` denotes the inverse of `g`. For linear regression, the canonical link function is the identity function. Therefore, when using the canonical link function, `μ = g⁻¹(η) = η`.
#
# In this model, the parameters that we want to estimate are `α`, `β`, and `σ`.
# We need to select prior distributions for these parameters. For `α`, we choose
# a normal distribution with zero mean and unit variance. For each `βᵢ`,
# we choose a normal distribution with zero mean and unit variance. Here, `βᵢ`
# denotes the `i`th component of `β`. For `σ`, we will choose a half-normal
# distribution with unit variance.

# We define this model using Soss:

using Soss
using Random

model = @model X begin
    p = size(X, 2) # number of features
    α ~ Normal(0, 1) # intercept
    β ~ Normal(0, 1) |> iid(p) # coefficients
    σ ~ HalfNormal(1) # dispersion
    η = α .+ X * β # linear predictor
    μ = η # `μ = g⁻¹(η) = η`
    y ~ For(eachindex(μ)) do j
        Normal(μ[j], σ) # `Yᵢ ~ Normal(mean=μᵢ, variance=σ²)`
    end
end;

# In Soss, models are _first-class_ and _function-like_, and applying a model to its arguments gives a _joint distribution_.

# Just a few of the things we can do in Soss:
# - Sample from the forward model
# - Condition a joint distribution on a subset of parameters
# - Have arbitrary Julia values (yes, even other models) as inputs or outputs of a model
# - Build a new model for the _predictive_ distribution, for assigning parameters to particular values

# ## Sampling from the forward model

# First, create some fake data:

X = randn(6,2)

# Now, sample from the forward model:

forward_sample = rand(model(X=X))

# The `pairs` function can make this a little easier to read:

pairs(forward_sample)

# ## Use MCMC to sample from the posterior distribution

# First, generate some fake data:

num_rows = 1_000
num_features = 2
X = randn(num_rows, num_features)

# Pick the true values for our coefficients `β`:

β_true = [2.0, -1.0]

# We also need to pick a true value for the intercept `α`:

α_true = 1.0

# And we also need to pick a true value for the dispersion parameter `σ`

σ_true = 0.5

# Now, generate the true labels:

η_true = α_true .+ X * β_true
μ_true = η_true
noise = randn(num_rows) .* σ_true
y_true = μ_true .+ noise

# Now we use MCMC (specifically, the No-U-turn sampler) to sample from the
# posterior distribution:

posterior = dynamicHMC(model(X=X) | (y=y_true,))

# Often, the posterior distributions are easier to work with in terms of
# `particles` (built using [MonteCarloMeasurements.jl](https://github.com/baggepinnen/MonteCarloMeasurements.jl)):

particles(posterior)

# Again, the `pairs` function can make this a little easier to read:

pairs(particles(posterior))

# Compare the posterior distributions on `σ`, `α`, and `β` to the true values:

@show σ_true; @show α_true; @show β_true;

# We did a pretty good job at recovering the true parameter values!

# ## Construct the posterior predictive distribution

# For model diagnostics and prediction, we need the posterior _predictive distribution_:

posterior_predictive = predictive(model, :β)

# This requires `X` and `β` as inputs, so we can do something like this to do a _posterior predictive check (PPC)_

y_ppc = [rand(posterior_predictive(;X=X, p...)).y for p in posterior]

# We can compare the posterior predictive distribution on `y` to the true values of `y`:

y_true - particles(y_ppc)

# These play a role similar to that of residuals in a non-Bayesian approach (there's plenty more detail to go into, but that's for another time).

# ## So, what's really happening here?

# Under the hood, `rand` and `logdensity` specify different ways of "running" the model.

 # `rand`  turns each `v ~ dist` into `v = rand(dist)`, finally outputting the `NamedTuple` of all values it has seen.

# `logdensity` steps through the same program, but instead accumulates a log-density. It begins by initializing `_ℓ = 0.0`. Then at each step, it turns `v ~ dist` into `_ℓ += logdensity(dist, v)`, before finally returning `_ℓ`.

# Note that I said "turns into" instead of "interprets". Soss uses [`GG.jl`](https://github.com/thautwarm/GG.jl) to generate specialized code for a given model, inference primitive (like `rand` and `logdensity`), and type of data.

# This idea can be used in much more complex ways. `weightedSample` is a sort of hybrid between `rand` and `logdensity`. For data that are provided, it increments a `_ℓ` using `logdensity`. Unknown values are sampled using `rand`.

# TODO: Fix weightedSample
ℓ, proposal = weightedSample(model(X=X) | (y=y_true,));

# `ℓ`:

ℓ

# `proposal.β`:

proposal.β

# Again, there's no runtime check needed for this. Each of these is compiled the first time it is called, so future calls are very fast. Functions like this are great to use in tight loops.


text:
# test saving and loading of trees

using Test
using IncrementalInference
using MetaGraphs
using Graphs

##

@testset "Test loading and saving of Bayes (Junction) tree" begin
##

fg = generateGraph_Kaess(graphinit=false)
tree = buildTreeReset!(fg)

# save and load tree
saveTree(tree)
tree2 = loadTree()

# perform a few spot checks to see that the trees are similar
@test length(tree.cliques) == length(tree2.cliques)
@test getEliminationOrder(tree) == getEliminationOrder(tree2)

for (clid,cl) in tree.cliques
  fsyms = getFrontals(cl)
  cl2 = getClique(tree2, fsyms[1])
  fsyms2 = getFrontals(cl2)
  @test fsyms == fsyms2
  @test getCliqSeparatorVarIds(cl) == getCliqSeparatorVarIds(cl2)
  @test typeof(cl) == typeof(cl2)
end

##
end

@testset "Test loading and saving of Bayes (Junction) tree" begin
##

fg = generateGraph_Kaess(graphinit=false)
tree = buildTreeReset!(fg)

# save and load tree as array
filepath = saveTree([tree;deepcopy(tree)])
trees = loadTree(filepath)

# perform a few spot checks to see that the trees are similar
@test length(tree.cliques) == length(trees[1].cliques)
@test getEliminationOrder(tree) == getEliminationOrder(trees[1])

for (clid,cl) in tree.cliques
  fsyms = getFrontals(cl)
  cl2 = getClique(trees[1], fsyms[1])
  fsyms2 = getFrontals(cl2)
  @test fsyms == fsyms2
  @test getCliqSeparatorVarIds(cl) == getCliqSeparatorVarIds(cl2)
  @test typeof(cl) == typeof(cl2)
end

##
end


text:
"""
    distance_periodic(basis::AbstractMatrix, v1::AbstractVector, v2::AbstractVector)

Return the distance between two vectors in a given crystal basis. This function considers all
possible neighbors given the cell boundaries, and returns the distance between the closest pair.
"""
function distance_periodic(
    basis::AbstractMatrix,
    v1::AbstractVector,
    v2::AbstractVector,
)
    # Vectors must be the same dimensionality
    @assert length(v1) == length(v2) "Vectors are of different lengths."
    # Truncate the vector components so they're between 0 and 1
    (v1t, v2t) = map(v -> v - floor.(v), (v1, v2))
    # Get the difference vector
    dred = (v1 - v2) - round.(v1t - v2t)
    # Transform it to Cartesian coordinates and get the norm
    return norm(basis*dred)
end

function distance_periodic(b::AbstractMatrix, a1::AtomPosition, a2::AtomPosition)
    return distance_periodic(b, coord(a1), coord(a2))
end

function distance_periodic(l::AtomList)
    # Matrix with all the distances
    M = zeros(Float64, natom(l), natom(l))
    for a in 1:natom
        # Distances on diagonal should be zero/don't need to be calculated
        for b in 1:(a-1)
            dist = distance_periodic(basis(l), coord(l[a]), coord(l[b]))
            # Assign off-diagonal elements symmetrically
            M[a,b] = dist
            M[b,a] = dist
        end
    end
    return M
end

"""
    delaunay_periodic(l::AtomList) -> BitMatrix

Generates the Delaunay graph of a set of points in a periodic basis.
"""
function delaunay_periodic(l::AtomList{D}) where D
    # Initialize the adjacency matrix
    M = BitMatrix(zeros(Bool, length(points), length(points)))
    # Throw an exception if there is no basis given
    basis(l) == zeros(BasisVectors{D}) && error("No basis vectors were provided.")

    return M
end

"""
    packing_efficiency(xtal::AbstractCrystal)

Determines the packing efficiency of a crystal.
"""
function packing_efficiency(xtal::Crystal)

end

max_stars_repo_name: brainandforce/Xtal

max_issues_repo_name: brainandforce/Xtal

max_forks_repo_name: brainandforce/Xtal


text:
using StatsBase: mean
using .NonadiabaticDistributions: NonadiabaticDistribution
using ComponentArrays: ComponentVector

export Ehrenfest

"""
    Ehrenfest{T} <: AbstractEhrenfest

Ehrenfest molecular dynamics. Classical molecular dynamics where the force is derived
by averaging contributions from multiple electronic states.

```jldoctest
Simulation{Ehrenfest}(Atoms(:H), DoubleWell())

# output

Simulation{Ehrenfest{Float64}}:
  Atoms{1, Float64}([:H], UInt8[0x01], [1837.4715941070515])
  DoubleWell{Int64, Int64, Int64, Int64}
  mass: Int64 1
  ω: Int64 1
  γ: Int64 1
  Δ: Int64 1
```
"""
struct Ehrenfest{T} <: AbstractEhrenfest
    density_propagator::Matrix{Complex{T}}
    function Ehrenfest{T}(n_states::Integer) where {T}
        density_propagator = zeros(n_states, n_states)
        new{T}(density_propagator)
    end
end

function Simulation{Ehrenfest}(atoms::Atoms{S,T}, model::Model; kwargs...) where {S,T}
    Simulation(atoms, model, Ehrenfest{T}(NQCModels.nstates(model)); kwargs...)
end

function DynamicsMethods.DynamicsVariables(
    sim::AbstractSimulation{<:AbstractEhrenfest}, v, r, electronic::NonadiabaticDistribution
)
    σ = NonadiabaticDistributions.initialise_adiabatic_density_matrix(electronic, sim.calculator, r)
    return ComponentVector(v=v, r=r, σreal=σ, σimag=zero(σ))
end

function acceleration!(dv, v, r, sim::AbstractSimulation{<:Ehrenfest}, t, σ)
    dv .= zero(eltype(dv))
    for I in eachindex(dv)
        for J in eachindex(σ)
            dv[I] -= sim.calculator.adiabatic_derivative[I][J] * real(σ[J])
        end
    end
    DynamicsUtils.divide_by_mass!(dv, sim.atoms.masses)
    return nothing
end

function Estimators.adiabatic_population(::AbstractSimulation{<:Ehrenfest}, u)
    σ = DynamicsUtils.get_quantum_subsystem(u)
    return real.(diag(σ))
end

function Estimators.diabatic_population(sim::AbstractSimulation{<:AbstractEhrenfest}, u)
    r = DynamicsUtils.get_positions(u)
    U = NonadiabaticDistributions.evaluate_transformation(sim.calculator, r)

    σ = DynamicsUtils.get_quantum_subsystem(u)

    return real.(diag(U * σ * U'))
end

function DynamicsUtils.classical_hamiltonian(sim::Simulation{<:Ehrenfest}, u)
    kinetic = DynamicsUtils.classical_kinetic_energy(sim, DynamicsUtils.get_velocities(u))

    Calculators.evaluate_potential!(sim.calculator, DynamicsUtils.get_positions(u))
    Calculators.eigen!(sim.calculator)
    potential = sum(diag(DynamicsUtils.get_quantum_subsystem(u)) .* sim.calculator.eigen.values)

    return kinetic + potential
end