text:
"""
 Tests for https://arxiv.org/abs/1303.5103 with test problems as in the paper.
 DRI1, RI1, RI3, RI5, RI6, RDI1WM, RDI2WM, RDI3WM, RDI4WM
"""

import Statistics # for mean values of trajectories
import LinearAlgebra # for the normn
using StochasticDiffEq
using Test
using Random
#using DiffEqGPU

#using Plots

function generate_weak_solutions(prob, alg, dts, numtraj; ensemblealg=EnsembleThreads())
  sols = []
  for i in 1:length(dts)
    sol = solve(prob,alg;ensemblealg=ensemblealg,dt=dts[i],save_start=false,save_everystep=false,weak_timeseries_errors=false,weak_dense_errors=false,trajectories=Int(numtraj))
    println(i)
    push!(sols,sol)
  end
  return sols
end

function prob_func(prob, i, repeat)
    remake(prob,seed=seeds[i])
end

"""
 Test Scalar SDEs (oop)
"""

@info "Scalar oop noise"

numtraj = Int(2e5) # in the paper they use 1e9
u₀ = 0.0
f(u,p,t) = 1//2*u+sqrt(u^2+1)
g(u,p,t) = sqrt(u^2+1)
dts = 1 .//2 .^(4:-1:1)
tspan = (0.0,2.0) # 2.0 in paper

h1(z) = z^3-6*z^2+8*z
#analytical_sol(t) = E(f(X(t))) = E(h1(arsinh(X(t))) = t^3-3*t^2+2*t
#analytical_sol(2) = 0 and analytical_sol(1)=0

seed = 100
Random.seed!(seed)
seeds = rand(UInt, numtraj)

prob = SDEProblem(f,g,u₀,tspan)
ensemble_prob = EnsembleProblem(prob;
        output_func = (sol,i) -> (h1(asinh(sol[end])),false),
        prob_func = prob_func
        )
_solutions = @time generate_weak_solutions(ensemble_prob, DRI1(), dts, numtraj, ensemblealg=EnsembleThreads())

errors = [LinearAlgebra.norm(Statistics.mean(sol.u)) for sol in _solutions]
m = log(errors[end]/errors[1])/log(dts[end]/dts[1])
@test -(m-2) < 0.3

#convergence_plot = plot(dts, errors, xaxis=:log, yaxis=:log)
#savefig(convergence_plot, "DRI1-CPU-final-"*string(numtraj)*".pdf")
println("DRI1:", m)

numtraj = Int(5e5)
seed = 100
Random.seed!(seed)
seeds = rand(UInt, numtraj)

_solutions = @time generate_weak_solutions(ensemble_prob, RI1(), dts, numtraj, ensemblealg=EnsembleThreads())

errors = [LinearAlgebra.norm(Statistics.mean(sol.u)) for sol in _solutions]
m = log(errors[end]/errors[1])/log(dts[end]/dts[1])
@test -(m-2) < 0.3

println("RI1:", m)

numtraj = Int(5e5)
seed = 100
Random.seed!(seed)
seeds = rand(UInt, numtraj)

_solutions = @time generate_weak_solutions(ensemble_prob, RI3(), dts, numtraj, ensemblealg=EnsembleThreads())

errors = [LinearAlgebra.norm(Statistics.mean(sol.u)) for sol in _solutions]
m = log(errors[end]/errors[1])/log(dts[end]/dts[1])
@test -(m-2) < 0.3

println("RI3:", m)

numtraj = Int(5e5)
seed = 100
Random.seed!(seed)
seeds = rand(UInt, numtraj)

_solutions = @time generate_weak_solutions(ensemble_prob, RI5(), dts, numtraj, ensemblealg=EnsembleThreads())

errors = [LinearAlgebra.norm(Statistics.mean(sol.u)) for sol in _solutions]
m = log(errors[end]/errors[1])/log(dts[end]/dts[1])
@test -(m-2) < 0.3

println("RI5:", m)

numtraj = Int(7e5)
seed = 100
Random.seed!(seed)
seeds = rand(UInt, numtraj)

_solutions = @time generate_weak_solutions(ensemble_prob, RI6(), dts, numtraj, ensemblealg=EnsembleThreads())

errors = [LinearAlgebra.norm(Statistics.mean(sol.u)) for sol in _solutions]
m = log(errors[end]/errors[1])/log(dts[end]/dts[1])
@test -(m-2) < 0.3

println("RI6:", m)

numtraj = Int(1e3)
seed = 100
Random.seed!(seed)
seeds = rand(UInt, numtraj)

_solutions = @time generate_weak_solutions(ensemble_prob, RDI1WM(), dts, numtraj, ensemblealg=EnsembleThreads())

errors = [LinearAlgebra.norm(Statistics.mean(sol.u)) for sol in _solutions]
m = log(errors[end]/errors[1])/log(dts[end]/dts[1])
@test -(m-1) < 0.3

println("RDI1WM:", m)

numtraj = Int(7e5)
seed = 100
Random.seed!(seed)
seeds = rand(UInt, numtraj)

_solutions = @time generate_weak_solutions(ensemble_prob, RDI2WM(), dts, numtraj, ensemblealg=EnsembleThreads())

errors = [LinearAlgebra.norm(Statistics.mean(sol.u)) for sol in _solutions]
m = log(errors[end]/errors[1])/log(dts[end]/dts[1])
@test -(m-2) < 0.3

println("RDI2WM:", m)

_solutions = @time generate_weak_solutions(ensemble_prob, RDI3WM(), dts, numtraj, ensemblealg=EnsembleThreads())

errors = [LinearAlgebra.norm(Statistics.mean(sol.u)) for sol in _solutions]
m = log(errors[end]/errors[1])/log(dts[end]/dts[1])
@test -(m-2) < 0.3

numtraj = Int(3e5)
seed = 100
Random.seed!(seed)
seeds = rand(UInt, numtraj)

println("RDI3WM:", m)

_solutions = @time generate_weak_solutions(ensemble_prob, RDI4WM(), dts, numtraj, ensemblealg=EnsembleThreads())

errors = [LinearAlgebra.norm(Statistics.mean(sol.u)) for sol in _solutions]
m = log(errors[end]/errors[1])/log(dts[end]/dts[1])
@test -(m-2) < 0.3

println("RDI4WM:", m)

"""
 Test Scalar SDEs (iip)
"""

@info "Scalar iip noise"

u₀ = [0.0]
f1!(du,u,p,t) = @.(du = 1//2*u+sqrt(u^2 +1))
g1!(du,u,p,t) = @.(du = sqrt(u^2 +1))
dts = 1 .//2 .^(4:-1:1)
tspan = (0.0,2.0)

h1(z) = z^3-6*z^2+8*z

prob = SDEProblem(f1!,g1!,u₀,tspan)
ensemble_prob = EnsembleProblem(prob;
        output_func = (sol,i) -> (h1(asinh(sol[end][1])),false),
        prob_func = prob_func
        )

numtraj = Int(2e5)
seed = 100
Random.seed!(seed)
seeds = rand(UInt, numtraj)

_solutions = @time generate_weak_solutions(ensemble_prob, DRI1(), dts, numtraj, ensemblealg=EnsembleThreads())

errors = [LinearAlgebra.norm(Statistics.mean(sol.u)) for sol in _solutions]
m = log(errors[end]/errors[1])/log(dts[end]/dts[1])
@test -(m-2) < 0.3

println("DRI1:", m)

numtraj = Int(5e5)
seed = 100
Random.seed!(seed)
seeds = rand(UInt, numtraj)

_solutions = @time generate_weak_solutions(ensemble_prob, RI1(), dts, numtraj, ensemblealg=EnsembleThreads())

errors = [LinearAlgebra.norm(Statistics.mean(sol.u)) for sol in _solutions]
m = log(errors[end]/errors[1])/log(dts[end]/dts[1])
@test -(m-2) < 0.3

println("RI1:", m)

numtraj = Int(5e5)
seed = 100
Random.seed!(seed)
seeds = rand(UInt, numtraj)

_solutions = @time generate_weak_solutions(ensemble_prob, RI3(), dts, numtraj, ensemblealg=EnsembleThreads())

errors = [LinearAlgebra.norm(Statistics.mean(sol.u)) for sol in _solutions]
m = log(errors[end]/errors[1])/log(dts[end]/dts[1])
@test -(m-2) < 0.3

println("RI3:", m)

numtraj = Int(5e5)
seed = 100
Random.seed!(seed)
seeds = rand(UInt, numtraj)

_solutions = @time generate_weak_solutions(ensemble_prob, RI5(), dts, numtraj, ensemblealg=EnsembleThreads())

errors = [LinearAlgebra.norm(Statistics.mean(sol.u)) for sol in _solutions]
m = log(errors[end]/errors[1])/log(dts[end]/dts[1])
@test -(m-2) < 0.3

println("RI5:", m)

numtraj = Int(7e5)
seed = 100
Random.seed!(seed)
seeds = rand(UInt, numtraj)

_solutions = @time generate_weak_solutions(ensemble_prob, RI6(), dts, numtraj, ensemblealg=EnsembleThreads())

errors = [LinearAlgebra.norm(Statistics.mean(sol.u)) for sol in _solutions]
m = log(errors[end]/errors[1])/log(dts[end]/dts[1])
@test -(m-2) < 0.3

println("RI6:", m)

numtraj = Int(1e3)
seed = 100
Random.seed!(seed)
seeds = rand(UInt, numtraj)

_solutions = @time generate_weak_solutions(ensemble_prob, RDI1WM(), dts, numtraj, ensemblealg=EnsembleThreads())

errors = [LinearAlgebra.norm(Statistics.mean(sol.u)) for sol in _solutions]
m = log(errors[end]/errors[1])/log(dts[end]/dts[1])
@test -(m-1) < 0.3

println("RDI1WM:", m)

numtraj = Int(7e5)
seed = 100
Random.seed!(seed)
seeds = rand(UInt, numtraj)

_solutions = @time generate_weak_solutions(ensemble_prob, RDI2WM(), dts, numtraj, ensemblealg=EnsembleThreads())

errors = [LinearAlgebra.norm(Statistics.mean(sol.u)) for sol in _solutions]
m = log(errors[end]/errors[1])/log(dts[end]/dts[1])
@test -(m-2) < 0.3

println("RDI2WM:", m)

numtraj = Int(3e5)
seed = 100
Random.seed!(seed)
seeds = rand(UInt, numtraj)

_solutions = @time generate_weak_solutions(ensemble_prob, RDI3WM(), dts, numtraj, ensemblealg=EnsembleThreads())

errors = [LinearAlgebra.norm(Statistics.mean(sol.u)) for sol in _solutions]
m = log(errors[end]/errors[1])/log(dts[end]/dts[1])
@test -(m-2) < 0.3

println("RDI3WM:", m)

_solutions = @time generate_weak_solutions(ensemble_prob, RDI4WM(), dts, numtraj, ensemblealg=EnsembleThreads())

errors = [LinearAlgebra.norm(Statistics.mean(sol.u)) for sol in _solutions]
m = log(errors[end]/errors[1])/log(dts[end]/dts[1])
@test -(m-2) < 0.3

println("RDI4WM:", m)

"""
 Test non-commutative noise SDEs (iip)
"""

@info "Non-commutative noise"

u₀ = [1.0,1.0]
function f2!(du,u,p,t)
  du[1] = -273//512*u[1]
  du[2] = -1//160*u[1]-(-785//512+sqrt(2)/8)*u[2]
end
function g2!(du,u,p,t)
  du[1,1] = 1//4*u[1]
  du[1,2] = 1//16*u[1]
  du[2,1] = (1-2*sqrt(2))/4*u[1]
  du[2,2] = 1//10*u[1]+1//16*u[2]
end
dts = 1 .//2 .^(3:-1:0)
tspan = (0.0,3.0)

h2(z) = z^2 # but apply it only to u[1]

prob = SDEProblem(f2!,g2!,u₀,tspan,noise_rate_prototype=zeros(2,2))
ensemble_prob = EnsembleProblem(prob;
        output_func = (sol,i) -> (h2(sol[end][1]),false),
        prob_func = prob_func
        )

numtraj = Int(1e5)
seed = 100
Random.seed!(seed)
seeds = rand(UInt, numtraj)

_solutions = @time generate_weak_solutions(ensemble_prob, DRI1(), dts, numtraj, ensemblealg=EnsembleThreads())

errors = [LinearAlgebra.norm(Statistics.mean(sol.u)-exp(-3.0)) for sol in _solutions]
m = log(errors[end]/errors[1])/log(dts[end]/dts[1])
@test -(m-2) < 0.3

println("DRI1:", m)

numtraj = Int(2e5)
seed = 100
Random.seed!(seed)
seeds = rand(UInt, numtraj)
_solutions = @time generate_weak_solutions(ensemble_prob, RI1(), dts, numtraj, ensemblealg=EnsembleThreads())

errors = [LinearAlgebra.norm(Statistics.mean(sol.u)-exp(-3.0)) for sol in _solutions]
m = log(errors[end]/errors[1])/log(dts[end]/dts[1])
@test -(m-2) < 0.3

println("RI1:", m)

numtraj = Int(2e5)
seed = 100
Random.seed!(seed)
seeds = rand(UInt, numtraj)
_solutions = @time generate_weak_solutions(ensemble_prob, RI3(), dts, numtraj, ensemblealg=EnsembleThreads())

errors = [LinearAlgebra.norm(Statistics.mean(sol.u)-exp(-3.0)) for sol in _solutions]
m = log(errors[end]/errors[1])/log(dts[end]/dts[1])
@test -(m-2) < 0.3

println("RI3:", m)

numtraj = Int(2e5)
seed = 100
Random.seed!(seed)
seeds = rand(UInt, numtraj)
_solutions = @time generate_weak_solutions(ensemble_prob, RI5(), dts, numtraj, ensemblealg=EnsembleThreads())

errors = [LinearAlgebra.norm(Statistics.mean(sol.u)-exp(-3.0)) for sol in _solutions]
m = log(errors[end]/errors[1])/log(dts[end]/dts[1])
@test -(m-2) < 0.3

println("RI5:", m)

numtraj = Int(2e5)
seed = 100
Random.seed!(seed)
seeds = rand(UInt, numtraj)
_solutions = @time generate_weak_solutions(ensemble_prob, RI6(), dts, numtraj, ensemblealg=EnsembleThreads())

errors = [LinearAlgebra.norm(Statistics.mean(sol.u)-exp(-3.0)) for sol in _solutions]
m = log(errors[end]/errors[1])/log(dts[end]/dts[1])
@test -(m-2) < 0.3

println("RI6:", m)

numtraj = Int(1e3)
seed = 10
Random.seed!(seed)
seeds = rand(UInt, numtraj)

_solutions = @time generate_weak_solutions(ensemble_prob, RDI1WM(), dts, numtraj, ensemblealg=EnsembleThreads())

errors = [LinearAlgebra.norm(Statistics.mean(sol.u)-exp(-3.0)) for sol in _solutions]
m = log(errors[end]/errors[1])/log(dts[end]/dts[1])
@test -(m-1) < 0.3

println("RDI1WM:", m)

numtraj = Int(1e5)
seed = 10
Random.seed!(seed)
seeds = rand(UInt, numtraj)

_solutions = @time generate_weak_solutions(ensemble_prob, RDI2WM(), dts, numtraj, ensemblealg=EnsembleThreads())

errors = [LinearAlgebra.norm(Statistics.mean(sol.u)-exp(-3.0)) for sol in _solutions]
m = log(errors[end]/errors[1])/log(dts[end]/dts[1])
@test -(m-2) < 0.3

println("RDI2WM:", m)

_solutions = @time generate_weak_solutions(ensemble_prob, RDI3WM(), dts, numtraj, ensemblealg=EnsembleThreads())

errors = [LinearAlgebra.norm(Statistics.mean(sol.u)-exp(-3.0)) for sol in _solutions]
m = log(errors[end]/errors[1])/log(dts[end]/dts[1])
@test -(m-2) < 0.3

println("RDI3WM:", m)

_solutions = @time generate_weak_solutions(ensemble_prob, RDI4WM(), dts, numtraj, ensemblealg=EnsembleThreads())

errors = [LinearAlgebra.norm(Statistics.mean(sol.u)-exp(-3.0)) for sol in _solutions]
m = log(errors[end]/errors[1])/log(dts[end]/dts[1])
@test -(m-2) < 0.3

println("RDI4WM:", m)

"""
 Test Diagonal noise SDEs (iip), SIAM Journal on Numerical Analysis, 47 (2009), pp. 1713–1738
"""

@info "Diagonal noise"

u₀ = [0.1,0.1]
function f3!(du,u,p,t)
  du[1] = 3//2*u[1]
  du[2] = 3//2*u[2]
end
function g3!(du,u,p,t)
  du[1] = 1//10*u[1]
  du[2] = 1//10*u[1]
end
dts = 1 .//2 .^(3:-1:0)
tspan = (0.0,1.0)

h3(z) = z^2 # == 1//10**exp(3//2*t) if h3(z) = z and  == 1//100**exp(301//100*t) if h3(z) = z^2 )

prob = SDEProblem(f3!,g3!,u₀,tspan)
ensemble_prob = EnsembleProblem(prob;
        output_func = (sol,i) -> (h3(sol[end][1]),false),
        prob_func = prob_func
        )

numtraj = Int(5e4)
seed = 100
Random.seed!(seed)
seeds = rand(UInt, numtraj)

_solutions = @time generate_weak_solutions(ensemble_prob, DRI1(), dts, numtraj, ensemblealg=EnsembleThreads())

errors = [LinearAlgebra.norm(Statistics.mean(sol.u)-1//100*exp(301//100)) for sol in _solutions]
m = log(errors[end]/errors[1])/log(dts[end]/dts[1])
@test -(m-2) < 0.3

println("DRI1:", m)

_solutions = @time generate_weak_solutions(ensemble_prob, RI1(), dts, numtraj, ensemblealg=EnsembleThreads())

errors = [LinearAlgebra.norm(Statistics.mean(sol.u)-1//100*exp(301//100)) for sol in _solutions]
m = log(errors[end]/errors[1])/log(dts[end]/dts[1])
@test -(m-2) < 0.3

println("RI1:", m)

_solutions = @time generate_weak_solutions(ensemble_prob, RI3(), dts, numtraj, ensemblealg=EnsembleThreads())

errors = [LinearAlgebra.norm(Statistics.mean(sol.u)-1//100*exp(301//100)) for sol in _solutions]
m = log(errors[end]/errors[1])/log(dts[end]/dts[1])
@test -(m-2) < 0.3

println("RI3:", m)

_solutions = @time generate_weak_solutions(ensemble_prob, RI5(), dts, numtraj, ensemblealg=EnsembleThreads())

errors = [LinearAlgebra.norm(Statistics.mean(sol.u)-1//100*exp(301//100)) for sol in _solutions]
m = log(errors[end]/errors[1])/log(dts[end]/dts[1])
@test -(m-2) < 0.3

println("RI5:", m)

numtraj = Int(1e5)
seed = 70
Random.seed!(seed)
seeds = rand(UInt, numtraj)

_solutions = @time generate_weak_solutions(ensemble_prob, RI6(), dts, numtraj, ensemblealg=EnsembleThreads())

errors = [LinearAlgebra.norm(Statistics.mean(sol.u)-1//100*exp(301//100)) for sol in _solutions]
m = log(errors[end]/errors[1])/log(dts[end]/dts[1])
@test -(m-2) < 0.45

println("RI6:", m)

numtraj = Int(1e3)
seed = 100
Random.seed!(seed)
seeds = rand(UInt, numtraj)

_solutions = @time generate_weak_solutions(ensemble_prob, RDI1WM(), dts, numtraj, ensemblealg=EnsembleThreads())

errors = [LinearAlgebra.norm(Statistics.mean(sol.u)-1//100*exp(301//100)) for sol in _solutions]
m = log(errors[end]/errors[1])/log(dts[end]/dts[1])
@test -(m-1) < 0.45

println("RDI1WM:", m)

numtraj = Int(1e5)
seed = 70
Random.seed!(seed)
seeds = rand(UInt, numtraj)

_solutions = @time generate_weak_solutions(ensemble_prob, RDI2WM(), dts, numtraj, ensemblealg=EnsembleThreads())

errors = [LinearAlgebra.norm(Statistics.mean(sol.u)-1//100*exp(301//100)) for sol in _solutions]
m = log(errors[end]/errors[1])/log(dts[end]/dts[1])
@test -(m-2) < 0.45

println("RDI2WM:", m)

numtraj = Int(5e4)
seed = 100
Random.seed!(seed)
seeds = rand(UInt, numtraj)

_solutions = @time generate_weak_solutions(ensemble_prob, RDI3WM(), dts, numtraj, ensemblealg=EnsembleThreads())

errors = [LinearAlgebra.norm(Statistics.mean(sol.u)-1//100*exp(301//100)) for sol in _solutions]
m = log(errors[end]/errors[1])/log(dts[end]/dts[1])
@test -(m-2) < 0.3

println("RDI3WM:", m)

_solutions = @time generate_weak_solutions(ensemble_prob, RDI4WM(), dts, numtraj, ensemblealg=EnsembleThreads())

errors = [LinearAlgebra.norm(Statistics.mean(sol.u)-1//100*exp(301//100)) for sol in _solutions]
m = log(errors[end]/errors[1])/log(dts[end]/dts[1])
@test -(m-2) < 0.3

println("RDI4WM:", m)


text:
using MarketData
FactCheck.setstyle(:compact)
FactCheck.onlystats(true)

facts("helper method") do

    context("single-letter methods work") do
        @fact o(AAPL).colnames --> ["Open"]
        @fact h(AAPL).colnames --> ["High"]
        @fact l(AAPL).colnames --> ["Low"]
        @fact c(AAPL).colnames --> ["Close"]
        @fact v(AAPL).colnames --> ["Volume"]
    end
end

ext: jl


text:
import HCubature: hcubature

"""
$(TYPEDEF)

`CGLiouvillian` defines the Liouvillian operator corresponding to the CGME.

# Fields

$(FIELDS)
"""
struct CGLiouvillian <: AbstractLiouvillian
    "CGME kernels"
    kernels::Any
    "close system unitary"
    unitary::Any
    "absolute error tolerance for integration"
    atol::Float64
    "relative error tolerance for integration"
    rtol::Float64
    "cache matrix for inplace unitary"
    Ut1::Union{Matrix,MMatrix}
    "cache matrix for inplace unitary"
    Ut2::Union{Matrix,MMatrix}
    "cache matrix for integration"
    Ut::Union{Matrix,MMatrix}
end

function CGLiouvillian(kernels, U, atol, rtol)
    m_size = size(kernels[1][2])
    Λ = m_size[1] <= 10 ? zeros(MMatrix{m_size[1],m_size[2],ComplexF64}) :
        zeros(ComplexF64, m_size[1], m_size[2])
    unitary = isinplace(U) ? U.func : (cache, t) -> cache .= U(t)
    CGLiouvillian(kernels, unitary, atol, rtol, similar(Λ), similar(Λ), Λ)
end

function (CG::CGLiouvillian)(du, u, p, t::Real)
    tf = p.tf
    for (inds, coupling, cfun, Ta) in CG.kernels
        lower_bound = t < Ta / 2 ? [-t, -t] : [-Ta, -Ta] / 2
        upper_bound = t + Ta / 2 > tf ? [tf, tf] .- t : [Ta, Ta] / 2
        for (i, j) in inds
            function integrand(x)
                Ut = CG.Ut
                Ut1 = CG.Ut1
                Ut2 = CG.Ut2
                CG.unitary(Ut, t)
                Ut2 .= CG.unitary(Ut2, t + x[2]) * Ut'
                Ut1 .= CG.unitary(Ut1, t + x[1]) * Ut'
                At1 = Ut .= Ut1' * coupling[i](p(t + x[1])) * Ut1
                At2 = Ut1 .= Ut2' * coupling[j](p(t + x[2])) * Ut2
                cfun[i, j](t + x[2], t + x[1]) *
                (At1 * u * At2 - 0.5 * (At2 * At1 * u + u * At2 * At1)) / Ta
            end
            cg_res, err = hcubature(
                integrand,
                lower_bound,
                upper_bound,
                rtol=CG.rtol,
                atol=CG.atol,
            )
            axpy!(1.0, cg_res, du)
        end
    end
end


text:
# Type constructors e.g. Tensor{2, 3}(arg)

# Tensor from function
@generated function (S::Union{Type{Tensor{order, dim}}, Type{SymmetricTensor{order, dim}}})(f::Function) where {order, dim}
    TensorType = get_base(get_type(S))
    if order == 1
        exp = tensor_create(TensorType, (i) -> :(f($i)))
    elseif order == 2
        exp = tensor_create(TensorType, (i,j) -> :(f($i, $j)))
    elseif order == 4
        exp = tensor_create(TensorType, (i,j,k,l) -> :(f($i, $j, $k, $l)))
    end
    quote
        $(Expr(:meta, :inline))
        @inbounds return $TensorType($exp)
    end
end

# Applies the function f to all indices f(1), f(2), ... f(n_independent_components)
@generated function apply_all(S::Union{Type{Tensor{order, dim}}, Type{SymmetricTensor{order, dim}}}, f::Function) where {order, dim}
    TensorType = get_base(get_type(S))
    exp = tensor_create_linear(TensorType, (i) -> :(f($i)))
    quote
        $(Expr(:meta, :inline))
        @inbounds return $TensorType($exp)
    end
end

@inline function apply_all(S::Union{Tensor{order, dim}, SymmetricTensor{order, dim}}, f::Function) where {order, dim}
    apply_all(get_base(typeof(S)), f)
end

# Tensor from AbstractArray
function Tensor{order, dim}(data::AbstractArray) where {order, dim}
    N = n_components(Tensor{order, dim})
    length(data) != n_components(Tensor{order, dim}) && throw(ArgumentError("wrong number of elements, expected $N, got $(length(data))"))
    return apply_all(Tensor{order, dim}, @inline function(i) @inbounds data[i]; end)
end

# SymmetricTensor from AbstractArray
@generated function SymmetricTensor{order, dim}(data::AbstractArray) where {order, dim}
    N = n_components(Tensor{order,dim})
    expN = Expr(:tuple, [:(data[$i]) for i in 1:N]...)
    M = n_components(SymmetricTensor{order,dim})
    expM = Expr(:tuple, [:(data[$i]) for i in 1:M]...)
    return quote
        L = length(data)
        if L != $N && L != $M
            throw(ArgumentError("wrong number of vector elements, expected $($N) or $($M), got $L"))
        end
        if L == $M
            @inbounds return SymmetricTensor{order, dim}($expM)
        end
        @inbounds S = Tensor{order, dim}($expN)
        return convert(SymmetricTensor{order, dim}, S)
    end
end

# Tensor from heterogeneous Tuple
function Tensor{order, dim}(data::Tuple) where {order, dim}
    return Tensor{order,dim}(promote(data...))
end

# SymmetricTensor from heterogeneous Tuple
function SymmetricTensor{order, dim}(data::Tuple) where {order, dim}
    return SymmetricTensor{order,dim}(promote(data...))
end

# one (identity tensor)
for TensorType in (SymmetricTensor, Tensor)
    @eval begin
        @inline Base.one(::Type{$(TensorType){order, dim}}) where {order, dim} = one($TensorType{order, dim, Float64})
        @inline Base.one(::Type{$(TensorType){order, dim, T, M}}) where {order, dim, T, M} = one($TensorType{order, dim, T})
        @inline Base.one(::$TensorType{order, dim, T}) where {order, dim, T} = one($TensorType{order, dim, T})

        @generated function Base.one(S::Type{$(TensorType){order, dim, T}}) where {order, dim, T}
            !(order in (2,4)) && throw(ArgumentError("`one` only defined for order 2 and 4"))
            δ = (i,j) -> i == j ? :(o) : :(z)
            ReturnTensor = get_base(get_type(S))
            if order == 2
                f = (i,j) -> :($(δ(i,j)))
            elseif order == 4 && $TensorType == Tensor
                f = (i,j,k,l) -> :($(δ(i,k)) * $(δ(j,l)))
            else # order == 4 && TensorType == SymmetricTensor
                f = (i,j,k,l) -> :(($(δ(i,k)) * $(δ(j,l)) + $(δ(i,l))* $(δ(j,k))) / 2)
            end
            exp = tensor_create(ReturnTensor, f)
            return quote
                $(Expr(:meta, :inline))
                o = one(T)
                z = zero(T)
                $ReturnTensor($exp)
            end
        end
    end
end

# zero, one, rand
for (op, el) in ((:zero, :(zero(T))), (:ones, :(one(T))), (:rand, :(()->rand(T))), (:randn,:(()->randn(T))))
for TensorType in (SymmetricTensor, Tensor)
    @eval begin
        @inline Base.$op(::Type{$TensorType{order, dim}}) where {order, dim} = $op($TensorType{order, dim, Float64})
        @inline Base.$op(::Type{$TensorType{order, dim, T, N}}) where {order, dim, T, N} = $op($TensorType{order, dim, T})
        @inline Base.$op(::Type{$TensorType{order, dim, T}}) where {order, dim, T} = fill($el, $TensorType{order, dim})
    end
end
@eval @inline Base.$op(S::Type{Vec{dim}}) where {dim} = $op(Vec{dim, Float64})
@eval @inline Base.$op(t::AllTensors) = $op(typeof(t))
end

@inline Base.fill(el::Number, S::Type{T}) where {T <: Union{Tensor, SymmetricTensor}} = apply_all(get_base(T), i -> el)
@inline Base.fill(f::Function, S::Type{T}) where {T <: Union{Tensor, SymmetricTensor}} = apply_all(get_base(T), i -> f())

# Array with zero/ones
@inline Base.zeros(::Type{T}, dims::Int...) where {T <: Union{Tensor, SymmetricTensor}} = fill(zero(T), (dims))
@inline Base.ones(::Type{T}, dims::Int...) where {T <: Union{Tensor, SymmetricTensor}} = fill(one(T), (dims))

# diagm
@generated function LinearAlgebra.diagm(S::Type{T}, v::Union{AbstractVector, Tuple}) where {T <: SecondOrderTensor}
    TensorType = get_base(get_type(S))
    ET = eltype(get_type(S)) == Any ? eltype(v) : eltype(get_type(S)) # lol
    f = (i,j) -> i == j ? :($ET(v[$i])) : :(o)
    exp = tensor_create(TensorType, f)
    return quote
        $(Expr(:meta, :inline))
        o = zero($ET)
        @inbounds return $TensorType($exp)
    end
end
@inline LinearAlgebra.diagm(::Type{Tensor{2, dim}}, v::T) where {dim, T<:Number} = v * one(Tensor{2, dim, T})
@inline LinearAlgebra.diagm(::Type{SymmetricTensor{2, dim}}, v::T) where {dim, T<:Number} = v * one(SymmetricTensor{2, dim, T})

"""
    basevec(::Type{Vec{dim, T}})
    basevec(::Type{Vec{dim, T}}, i)
    basevec(::Vec{dim, T})
    basevec(::Vec{dim, T}, i)

Return a tuple with the base vectors corresponding to the dimension `dim` and type
`T`. An optional integer `i` can be used to extract the i:th base vector.
The alias `eᵢ` can also be used, written `e\\_i<TAB>`.

# Examples
```jldoctest
julia> eᵢ(Vec{2, Float64})
([1.0, 0.0], [0.0, 1.0])

julia> eᵢ(Vec{2, Float64}, 2)
2-element Vec{2, Float64}:
 0.0
 1.0
```
"""
@inline function basevec(::Type{Vec{1, T}}) where {T}
    o = one(T)
    return (Vec{1, T}((o,)), )
end
@inline function basevec(::Type{Vec{2, T}}) where {T}
    o = one(T)
    z = zero(T)
    return (Vec{2, T}((o, z)),
            Vec{2, T}((z, o)))
end
@inline function basevec(::Type{Vec{3, T}}) where {T}
    o = one(T)
    z = zero(T)
    return (Vec{3, T}((o, z, z)),
            Vec{3, T}((z, o, z)),
            Vec{3, T}((z, z, o)))
end

@inline basevec(::Type{Vec{dim}}) where {dim} = basevec(Vec{dim, Float64})
@inline basevec(::Type{Vec{dim, T}}, i::Int) where {dim, T} = basevec(Vec{dim, T})[i]
@inline basevec(::Type{Vec{dim}}, i::Int) where {dim} = basevec(Vec{dim, Float64})[i]
@inline basevec(v::Vec{dim, T}) where {dim, T} = basevec(typeof(v))
@inline basevec(v::Vec{dim, T}, i::Int) where {dim, T} = basevec(typeof(v), i)

const eᵢ = basevec


text:
heat_eq_dir = joinpath(output_root,"HeatEquation")

@testset "∂_t T = K ΔT + 1, T = 0 ∈ ∂Ω, diffusion-explicit" begin
  dd = DomainDecomp(gm=1)
  dss = DomainSubSet(gm=true)
  K, maxiter, Δt = 1.0, 1000, 0.005
  grid = Grid(0.0, 1.0, 10)
  q = StateVec(( (:T, dss), ), grid, dd)
  tmp = StateVec(( (:ΔT, dss), ), grid, dd)
  rhs = deepcopy(q)
  for i in 1:maxiter
    for k in over_elems_real(grid)
      tmp[:ΔT, k] = Δ_z(q[:T, Cut(k)], grid)
      rhs[:T, k] = K*tmp[:ΔT, k] + 1
    end
    for k in over_elems(grid)
      q[:T, k] += Δt*rhs[:T, k]
    end
    apply_Dirichlet!(q, :T, grid, 0.0, Zmax())
    apply_Dirichlet!(q, :T, grid, 0.0, Zmin())
  end
  sol_analtyic = 1/2*grid.zc - grid.zc.^2/2
  sol_error = [abs(q[:T, k] - sol_analtyic[k]) for k in over_elems_real(grid)]
  @test all(sol_error .< grid.Δz)
  export_plots && export_state(q, grid, heat_eq_dir, "T.csv")
  export_plots && plot_state(q, grid, heat_eq_dir, :T, filename="T_unit_source.png")
end

@testset "∂_t T = K ΔT, T={1, 0} ∈ {z_{min}, z_{max}}, diffusion-explicit" begin
  dd = DomainDecomp(gm=1)
  dss = DomainSubSet(gm=true)
  K, maxiter, Δt = 1.0, 1000, 0.002
  grid = Grid(0.0, 1.0, 10)
  q = StateVec(( (:T, dss), ), grid, dd)
  tmp = StateVec(( (:ΔT, dss), ), grid, dd)
  rhs = deepcopy(q)
  for i in 1:maxiter
    for k in over_elems_real(grid)
      tmp[:ΔT, k] = Δ_z(q[:T, Cut(k)], grid)
      rhs[:T, k] = K*tmp[:ΔT, k]
    end
    for k in over_elems(grid)
      q[:T, k] += Δt*rhs[:T, k]
    end
    apply_Dirichlet!(q, :T, grid, 1.0, Zmin())
    apply_Dirichlet!(q, :T, grid, 0.0, Zmax())
  end
  sol_analtyic = 1 .- grid.zc
  sol_error = [abs(q[:T, k] - sol_analtyic[k]) for k in over_elems_real(grid)]
  @test all(sol_error .< grid.Δz)
  export_plots && plot_state(q, grid, heat_eq_dir, :T, filename="T1_diffExp.png")
end

@testset "∂_t T = K ΔT, T = {1, 0} ∈ {z_{min}, z_{max}}, diffusion implicit" begin
  dd = DomainDecomp(gm=1)
  dss = DomainSubSet(gm=true)
  K, maxiter, Δt = 1.0, 1000, 0.1
  grid = Grid(0.0, 1.0, 10)
  k_star1 = first_interior(grid, Zmin())
  k_star2 = first_interior(grid, Zmax())
  q = StateVec(( (:T, dss), (:a, dss), ), grid, dd)
  tmp = StateVec(( (:ΔT, dss), (:a_tau, dss), (:ρ_0, dss), (:K, dss), (:a, dss), ), grid, dd)
  rhs = deepcopy(q)
  for k in over_elems(grid)
    q[:a, k] = 1.0
    tmp[:ρ_0, k] = 1.0
    tmp[:a_tau, k] = 1.0
    tmp[:K, k] = 1.0
  end
  for i in 1:maxiter
    for k in over_elems_real(grid)
      rhs[:T, k] = 0.0
      if k==k_star1
        gv = 2*1 - q[:T, k]
        grad_T = ∇_z_flux([gv, 0], grid)
        bc_surf = ∇_z_flux([grad_T, 0], grid)
        rhs[:T, k] += bc_surf
      end
      if k==k_star2
        gv = - q[:T, k]
        grad_T = ∇_z_flux([0, gv], grid)
        bc_surf = ∇_z_flux([0, grad_T], grid)
        rhs[:T, k] += bc_surf
      end
    end
    solve_tdma!(q, rhs, tmp, :T, :ρ_0, :a, :a_tau, :K, grid, Δt)
    extrap!(q, :T, grid)
  end
  sol_analtyic = 1 .- grid.zc
  sol_numerical = [q[:T, k] for k in over_elems_real(grid)]
  sol_error = [abs(q[:T, k] - sol_analtyic[k]) for k in over_elems_real(grid)]
  @test all(sol_error .< grid.Δz)
  export_plots && plot_state(q, grid, heat_eq_dir, :T, filename= "T1_diffImp.png")
end

@testset "∂_t T = K ΔT, -K∇T = F ∈ z_{min}, T = 0 ∈ z_{max}, explicit Euler" begin
  dd = DomainDecomp(gm=1)
  dss = DomainSubSet(gm=true)
  K, maxiter, Δt = 1.0, 10, 0.001, 10
  grid = Grid(0.0, 1.0, 10)
  k_star1 = first_interior(grid, Zmin())
  k_star2 = first_interior(grid, Zmax())
  q = StateVec(( (:T1, dss), (:T2, dss), ), grid, dd)
  tmp = StateVec(( (:ΔT1, dss), (:ΔT2, dss), ), grid, dd)
  rhs = deepcopy(q)
  for i in 1:maxiter
    q_out = -K*1
    q_in = -q_out
    for k in over_elems_real(grid)
      tmp[:ΔT1, k] = Δ_z(q[:T1, Cut(k)], grid)
      tmp[:ΔT2, k] = Δ_z(q[:T2, Cut(k)], grid)
      rhs[:T1, k] = K*tmp[:ΔT1, k]
      rhs[:T2, k] = K*tmp[:ΔT2, k]
      if k==k_star1 && i>1
        rhs[:T2, k] += -∇_z_flux([q_in, 0.0], grid)
      end
    end
    for k in over_elems(grid)
      q[:T1, k] += Δt*rhs[:T1, k]
      q[:T2, k] += Δt*rhs[:T2, k]
    end
    apply_Neumann!(q   , :T1, grid, q_out, Zmin())
    apply_Dirichlet!(q , :T1, grid, 0.0  , Zmax())
    apply_Neumann!(q   , :T2, grid, 0.0  , Zmin())
    apply_Dirichlet!(q , :T2, grid, 0.0  , Zmax())
    @test all([q[:T1, k] ≈ q[:T2, k] for k in over_elems_real(grid)])
  end
end

@testset "∂_t T = ∇ • (K(z) ∇T) + 1, T = 0 ∈ ∂Ω, explicit Euler" begin
  dd = DomainDecomp(gm=1)
  dss = DomainSubSet(gm=true)
  K, maxiter, Δt = 1.0, 1000, 0.001
  grid = Grid(0.0, 1.0, 10)
  unknowns = ( (:T, dss), )
  vars = ( (:ΔT, dss), (:K_thermal, dss) )
  q = StateVec(unknowns, grid, dd)
  tmp = StateVec(vars, grid, dd)
  rhs = deepcopy(q)
  cond_thermal(z) = z > .5 ? 1 : .1
  for k in over_elems(grid)
    tmp[:K_thermal, k] = cond_thermal(grid.zc[k])
  end
  for i in 1:maxiter
    for k in over_elems_real(grid)
      tmp[:ΔT, k] = Δ_z(q[:T, Cut(k)], grid, tmp[:K_thermal, Cut(k)])
      rhs[:T, k] = K*tmp[:ΔT, k] + 1
    end
    for k in over_elems(grid)
      q[:T, k] += Δt*rhs[:T, k]
    end
    apply_Dirichlet!(q, :T, grid, 0.0, Zmax())
    apply_Dirichlet!(q, :T, grid, 0.0, Zmin())
  end
end

@testset "∂_t T = ∇ • (K(z) ∇T) + 1, T = {0, 0} ∈ {z_{min}, z_{max}}, comparison" begin
  dd = DomainDecomp(gm=1)
  dss = DomainSubSet(gm=true)
  VS, TS, F, K = 1.0, 1.0, 1.0, 1.0
  maxiter, Δt = 5000, 0.005
  grid = Grid(0.0, 1.0, 10)
  k_star1 = first_interior(grid, Zmin())
  k_star2 = first_interior(grid, Zmax())
  q = StateVec(( (:T_explicit_surf, dss), (:T_explicit_vol, dss), (:T_implicit, dss), (:a, dss), ), grid, dd)
  tmp = StateVec(( (:ΔT, dss), (:a_tau, dss), (:ρ_0, dss), (:K_thermal, dss), (:a, dss), ), grid, dd)
  rhs = deepcopy(q)
  for k in over_elems(grid)
    q[:a, k] = 1.0
    tmp[:ρ_0, k] = 1.0
    tmp[:a_tau, k] = 1.0
  end
  cond_thermal(z) = z > .5 ? 1 : .1
  for k in over_elems(grid)
    tmp[:K_thermal, k] = cond_thermal(grid.zc[k])
  end
  for i in 1:maxiter
    for k in over_elems_real(grid)
      tmp[:ΔT, k] = Δ_z(q[:T_explicit_surf, Cut(k)], grid, tmp[:K_thermal, Cut(k)])
      rhs[:T_explicit_surf, k] = tmp[:ΔT, k] + VS
      tmp[:ΔT, k] = Δ_z(q[:T_explicit_vol, Cut(k)], grid, tmp[:K_thermal, Cut(k)])
      rhs[:T_explicit_vol, k] = tmp[:ΔT, k] + VS
      rhs[:T_implicit, k] = VS
      k==k_star1 && (rhs[:T_implicit, k] += bc_source(q, grid, tmp, :T_implicit, :ρ_0, :a, :K_thermal, Zmin(), TS, Dirichlet(), DiffusionAbsorbed()))
      k==k_star2 && (rhs[:T_implicit, k] += bc_source(q, grid, tmp, :T_implicit, :ρ_0, :a, :K_thermal, Zmax(), TS, Dirichlet(), DiffusionAbsorbed()))
      k==k_star1 && (rhs[:T_explicit_vol, k] += bc_source(q, grid, tmp, :T_explicit_vol, :ρ_0, :a, :K_thermal, Zmin(), TS, Dirichlet(), DiffusionAbsorbed()))
      k==k_star2 && (rhs[:T_explicit_vol, k] += bc_source(q, grid, tmp, :T_explicit_vol, :ρ_0, :a, :K_thermal, Zmax(), TS, Dirichlet(), DiffusionAbsorbed()))
    end

    for k in over_elems_real(grid)
      q[:T_explicit_surf, k] += Δt*rhs[:T_explicit_surf, k]
      q[:T_explicit_vol, k] += Δt*rhs[:T_explicit_vol, k]
    end
    solve_tdma!(q, rhs, tmp, :T_implicit, :ρ_0, :a, :a_tau, :K_thermal, grid, Δt)

    apply_Dirichlet!(q, :T_explicit_surf, grid, TS, Zmax())
    apply_Dirichlet!(q, :T_explicit_surf, grid, TS, Zmin())

    assign_ghost!(q, :T_explicit_vol, grid, 0.0, Zmin())
    assign_ghost!(q, :T_explicit_vol, grid, 0.0, Zmax())

    assign_ghost!(q, :T_implicit, grid, 0.0, Zmin())
    assign_ghost!(q, :T_implicit, grid, 0.0, Zmax())
  end
  extrap!(q, :T_explicit_vol, grid, TS, Zmin())
  extrap!(q, :T_explicit_vol, grid, TS, Zmax())
  extrap!(q, :T_implicit, grid, TS, Zmin())
  extrap!(q, :T_implicit, grid, TS, Zmax())

  @test q[:T_explicit_vol, Dual(k_star1)][1] ≈ TS
  @test q[:T_explicit_vol, Dual(k_star2)][2] ≈ TS
  @test all([q[:T_explicit_vol, k] ≈ q[:T_explicit_surf, k] for k in over_elems_real(grid)])
  @test all([q[:T_implicit, k] ≈ q[:T_explicit_surf, k] for k in over_elems_real(grid)])

  export_plots && plot_state(q, grid, heat_eq_dir, :T_explicit_surf, filename="T_varK_explicit_surf.png")
  export_plots && plot_state(q, grid, heat_eq_dir, :T_explicit_vol, filename="T_varK_explicit_vol.png")
  export_plots && plot_state(q, grid, heat_eq_dir, :T_implicit, filename="T_varK_implicit.png")
end

@testset "∂_t T = ∇ • (K(z) ∇T) + VS, -K(z) ∇T = 1 ∈ z_{min}, T = 0 ∈ z_{max}, comparison" begin
  dd = DomainDecomp(gm=1)
  dss = DomainSubSet(gm=true)
  VS, TS, F = 1.0, 0.0, 2.0
  IC = 0.0
  maxiter, Δt = 10000, 0.0001
  grid = Grid(0.0, 1.0, 20)
  k_star1 = first_interior(grid, Zmin())
  k_star2 = first_interior(grid, Zmax())
  q = StateVec(( (:T_explicit_surf, dss), (:T_explicit_vol, dss), (:T_implicit, dss), (:a, dss), ), grid, dd)
  tmp = StateVec(( (:ΔT, dss), (:a_tau, dss), (:ρ_0, dss), (:K_thermal, dss), ), grid, dd)
  rhs = deepcopy(q)
  for k in over_elems(grid)
    q[:a, k] = 1.0
    tmp[:ρ_0, k] = 1.0
    tmp[:a_tau, k] = 1.0
  end
  cond_thermal(z) = 5+3*cos(6*π*z)
  for k in over_elems(grid)
    tmp[:K_thermal, k] = cond_thermal(grid.zc[k])
    q[:T_explicit_surf, k] = IC
    q[:T_explicit_vol, k] = IC
    q[:T_implicit, k] = IC
  end
  K_surf = tmp[:K_thermal, Dual(k_star1)][1]
  apply_Neumann!(q, :T_explicit_surf, grid, F/K_surf, Zmin())
  apply_Dirichlet!(q, :T_explicit_surf, grid, TS, Zmax())
  for i in 1:maxiter
    assign_ghost!(q, :T_explicit_vol, grid, 0.0, Zmin())
    assign_ghost!(q, :T_explicit_vol, grid, 0.0, Zmax())
    for k in over_elems_real(grid)
      tmp[:ΔT, k] = Δ_z_dual(q[:T_explicit_surf, Cut(k)], grid, tmp[:K_thermal, Dual(k)])
      rhs[:T_explicit_surf, k] = tmp[:ΔT, k] + VS

      tmp[:ΔT, k] = Δ_z_dual(q[:T_explicit_vol, Cut(k)], grid, tmp[:K_thermal, Dual(k)])
      rhs[:T_explicit_vol, k] = tmp[:ΔT, k] + VS
      k==k_star1 && (rhs[:T_explicit_vol, k] += bc_source(q, grid, tmp, :T_explicit_vol, :ρ_0, :a, :K_thermal, Zmin(), F, Neumann(), DiffusionAbsorbed()))
      k==k_star2 && (rhs[:T_explicit_vol, k] += bc_source(q, grid, tmp, :T_explicit_vol, :ρ_0, :a, :K_thermal, Zmax(), TS, Dirichlet(), DiffusionAbsorbed()))

      rhs[:T_implicit, k] = VS

      k==k_star1 && (rhs[:T_implicit, k] += bc_source(q, grid, tmp, :T_implicit, :ρ_0, :a, :K_thermal, Zmin(), F, Neumann(), DiffusionAbsorbed()))
      k==k_star2 && (rhs[:T_implicit, k] += bc_source(q, grid, tmp, :T_implicit, :ρ_0, :a, :K_thermal, Zmax(), TS, Dirichlet(), DiffusionAbsorbed()))
    end

    for k in over_elems_real(grid)
      q[:T_explicit_surf, k] += Δt*rhs[:T_explicit_surf, k]
      q[:T_explicit_vol, k] += Δt*rhs[:T_explicit_vol, k]
    end
    solve_tdma!(q, rhs, tmp, :T_implicit, :ρ_0, :a, :a_tau, :K_thermal, grid, Δt)

    apply_Neumann!(q, :T_explicit_surf, grid, F/K_surf, Zmin())
    apply_Dirichlet!(q, :T_explicit_surf, grid, TS, Zmax())

    apply_Neumann!(q, :T_explicit_vol, grid, F/K_surf, Zmin())
    apply_Dirichlet!(q, :T_explicit_vol, grid, TS, Zmax())

    apply_Neumann!(q, :T_implicit, grid, F/K_surf, Zmin())
    apply_Dirichlet!(q, :T_implicit, grid, TS, Zmax())
  end

  apply_Neumann!(q, :T_explicit_vol, grid, F/K_surf, Zmin())
  extrap!(q, :T_explicit_vol, grid, TS, Zmax())

  apply_Neumann!(q, :T_implicit, grid, F/K_surf, Zmin())
  extrap!(q, :T_implicit, grid, TS, Zmax())

  @test q[:T_explicit_vol, Dual(k_star2)][2] ≈ TS
  @test q[:T_implicit, Dual(k_star2)][2] ≈ TS

  export_plots && plot_state(q, grid, heat_eq_dir, :T_explicit_surf, filename="T_cosK_explicit_surf.png")
  export_plots && plot_state(q, grid, heat_eq_dir, :T_explicit_vol, filename="T_cosK_explicit_vol.png")
  export_plots && plot_state(q, grid, heat_eq_dir, :T_implicit, filename="T_cosK_implicit.png")

  @test all([abs(q[:T_explicit_vol, k] - q[:T_explicit_surf, k]) < 10*Δt for k in over_elems(grid)])
  @test all([abs(q[:T_explicit_vol, k] - q[:T_explicit_surf, k]) < 10*Δt for k in over_elems(grid)])
  @test all([abs(q[:T_implicit, k]     - q[:T_explicit_surf, k]) < 10*Δt for k in over_elems(grid)])
end

max_stars_repo_name: akshaysridhar/CLIMA

max_issues_repo_name: akshaysridhar/CLIMA

max_forks_repo_name: akshaysridhar/CLIMA


text:
# Pose3Pose3 evaluation functions

export Pose3Pose3, PackedPose3Pose3

# ------------------------------------
"""
$(TYPEDEF)

Rigid transform factor between two Pose3 compliant variables.
"""
struct Pose3Pose3{T <: IIF.SamplableBelief} <: IIF.AbstractManifoldMinimize
  z::T
end
# convenience and default constructor
Pose3Pose3() = Pose3Pose3(MvNormal(zeros(6),LinearAlgebra.diagm([0.01*ones(3);0.0001*ones(3)])))

DFG.getManifold(::Pose3Pose3) = Manifolds.SpecialEuclidean(3)

Pose3Pose3(::UniformScaling) = Pose3Pose3() # MvNormal(zeros(3),LinearAlgebra.diagm([1.0;1.0;1.0])) )

function getSample(cf::CalcFactor{<:Pose3Pose3}, N::Int=1)

  Xc = rand(cf.factor.z)
  #NOTE be carefull using this as template as it will not work in general for all manifolds
  M = getManifold(Pose3)
  ϵ = getPointIdentity(Pose3)
  X = hat(M, ϵ, Xc)
  # return a vector
  return X
end

function (cf::CalcFactor{<:Pose3Pose3})(X, p, q)
    M = getManifold(Pose3)
    q̂ = Manifolds.compose(M, p, exp(M, identity_element(M, p), X)) #for groups
    #TODO allocalte for vee! see Manifolds #412, fix for AD
    Xc = zeros(6)
    vee!(M, Xc, q, log(M, q, q̂))
    return Xc
end

# TODO Deprecate

# """
# $(TYPEDEF)
# """
# mutable struct PP3REUSE
#   wTi::SE3
#   wTj::SE3
#   iTi::SE3
#   PP3REUSE() = new(SE3(0),SE3(0),SE3(0))
# end

# function fastpose3pose3residual!( reusethrid::PP3REUSE,
#                                   meas,
#                                   wXi,
#                                   wXj  )
#   #
#   reusethrid.wTi.t[1:3] = wXi[1:3]
#   TransformUtils.convert!(reusethrid.wTi.R, Euler(wXi[4],wXi[5],wXi[6]))
#   reusethrid.wTj.t[1:3] = wXj[1:3]
#   TransformUtils.convert!(reusethrid.wTj.R, Euler(wXj[4],wXj[5],wXj[6]))

#   # TODO -- convert to in place convert! functions, many speed-ups possible here
#   jTi = SE3( matrix(reusethrid.wTj)\matrix(reusethrid.wTi) )
#   # also wasted memory here, should operate directly on iTi and not be assigning new memory
#   reusethrid.iTi = (SE3(meas[1:3],Euler(meas[4:6]...)) * jTi)
#   return veeEuler(reusethrid.iTi)
# end

# """
# $(TYPEDEF)

# Rigid transform factor between two Pose3 compliant variables.
# """
# mutable struct Pose3Pose3{T <: IIF.SamplableBelief} <: AbstractRelativeRoots
#     Zij::T
#     reuse::Vector{PP3REUSE}
#     Pose3Pose3{T}() where T = new{T}()
#     Pose3Pose3{T}(s::T) where {T <: SamplableBelief} = new{T}(s, PP3REUSE[PP3REUSE() for i in 1:Threads.nthreads()]  )
# end
# Pose3Pose3(z::T=MvNormal(zeros(6),LinearAlgebra.diagm([0.01*ones(3);0.0001*ones(3)]))) where {T <: IIF.SamplableBelief} = Pose3Pose3{T}(z)

# function getSample(cf::CalcFactor{<:Pose3Pose3}, N::Int=1)
#   return (rand(cf.factor.Zij, N), )
# end
# function (cf::CalcFactor{<:Pose3Pose3})(meas,
#                                         wXi,
#                                         wXj  )
#   #
#   reusethrid = cf.factor.reuse[Threads.threadid()]
#   return fastpose3pose3residual!(reusethrid, meas, wXi, wXj)
# end

#TODO Serialization

"""
$(TYPEDEF)

Serialization type for `Pose3Pose3`.
"""
mutable struct PackedPose3Pose3 <: IncrementalInference.PackedInferenceType
  Zij::String
  PackedPose3Pose3() = new()
  PackedPose3Pose3(x::AbstractString) = new(x)
end
function convert(::Type{Pose3Pose3}, packed::PackedPose3Pose3)
  return Pose3Pose3( convert(SamplableBelief, packed.Zij) )
end
function convert(::Type{PackedPose3Pose3}, obj::Pose3Pose3)
  return PackedPose3Pose3( convert(PackedSamplableBelief, obj.z) )
end

#


text:
@testset verbose = true "History" begin
    history = TemporalHistory([0.2, 0.8, 1.1], ["a", "b", "c"], 0.0, 2.0);
    @test duration(history) == 2.
    @test nb_events(history) == 3
    @test nb_events(history, 1.0, 2.0) == 1
    @test has_events(history)
    @test !has_events(history, 1.5, 2.0)
    push!(history, 1.7, "d")
    @test has_events(history, 1.5, 2.0)
end


text:
using TransferEntropy
using CausalityToolsBase

include("wrapper_transferentropy_regular.jl")
include("wrapper_transferentropy_cond.jl")


text:
function levelset_normal(levelset, p::V, invjac) where {V<:AbstractVector}
    g = vec(gradient(levelset, p))
    n = invjac .* g
    return n / norm(n)
end

function normalize_normals!(normals)
    dim,npts = size(normals)
    for i = 1:npts
        n = normals[:,i]
        normals[:,i] .= n/norm(n)
    end
end

function levelset_normal(levelset, points::M, invjac) where {M<:AbstractMatrix}
    npts = size(points)[2]
    if npts == 0
        return []
    else
        g = hcat([gradient(levelset, points[:, i])' for i = 1:npts]...)
        normals = diagm(invjac) * g
        normalize_normals!(normals)
        return normals
    end
end

function levelset_coefficients(distancefunc,mesh)
    nodalcoordinates = nodal_coordinates(mesh)
    return distancefunc(nodalcoordinates)
end

function reference_cell_area()
    return 4.0
end

function reference_face_normals()
    n1 = [0.,-1.]
    n2 = [1.,0.]
    n3 = [0.,1.]
    n4 = [-1.,0.]
    return [n1,n2,n3,n4]
end

function reference_bottom_face_midpoint()
    [0.0, -1.0]
end

function reference_right_face_midpoint()
    [1.0, 0.0]
end

function reference_top_face_midpoint()
    [0.0, 1.0]
end

function reference_left_face_midpoint()
    [-1.0, 0.0]
end

function reference_face(faceid)
    if faceid == 1
        return (2, -1.0)
    elseif faceid == 2
        return (1, +1.0)
    elseif faceid == 3
        return (2, +1.0)
    elseif faceid == 4
        return (1, -1.0)
    else
        error("Expected faceid ∈ {1,2,3,4}, got faceid = $faceid")
    end
end

function reference_face_midpoints()
    [
        reference_bottom_face_midpoint(),
        reference_right_face_midpoint(),
        reference_top_face_midpoint(),
        reference_left_face_midpoint(),
    ]
end

function extend_to_face(points, faceid)
    dir, coordval = reference_face(faceid)
    @assert dir == 1 || dir == 2
    flipdir = dir == 1 ? 2 : 1
    return extend([coordval], flipdir, points)
end

function opposite_face(faceid)
    if faceid == 1
        return 3
    elseif faceid == 2
        return 4
    elseif faceid == 3
        return 1
    elseif faceid == 4
        return 2
    else
        error("Expected faceid ∈ {1,2,3,4} got faceid = $faceid")
    end
end

function tangents(normals)
    rot = [
        0.0 -1.0
        1.0 0.0
    ]
    return rot * normals
end

function scale_area(normals,invjac)
    t = tangents(normals)
    den = sqrt.((t .^ 2)' * (invjac .^ 2))
    return 1.0 ./ den
end

function scale_area(cellmap::CellMap, normals)
    invjac = inverse_jacobian(cellmap)
    return scale_area(normals,invjac)
end

function cell_sign_to_row(s)
    (s == -1 || s == +1) || error("Use ±1 to index into rows (i.e. phase), got index = $s")
    row = s == +1 ? 1 : 2
    return row
end

function cell_couple_sign_to_row(s1,s2)
    (s1 == -1 || s1 == +1) || error("Expected sign ∈ {-1,1}, got sign $s1")
    (s2 == -1 || s2 == +1) || error("Expected sign ∈ {-1,1}, got sign $s2")
    if s1 == +1 && s2 == +1
        return 1
    elseif s1 == +1 && s2 == -1
        return 2
    elseif s1 == -1 && s2 == +1
        return 3
    elseif s1 == -1 && s2 == -1
        return 4
    end
end

function dimension(basis::TensorProductBasis{dim}) where {dim}
    return dim
end

function number_of_basis_functions(basis::TensorProductBasis{dim,T,NF}) where {dim,T,NF}

    return NF
end

function interpolation_points(basis::TensorProductBasis{D,T,N}) where {D,T,N}
    return basis.points
end

function number_of_symmetric_degrees_of_freedom(dim)
    if dim == 2
        return 3
    else
        error("Current support for dim = 2, got dim = $dim")
    end
end


text:
"""
# PairStrings.jl

Embed `key => value` Pairs in a string.

```
julia> using PairStrings

julia> pairs_str("Foo :FOO => 1 Bar :BAR => 2")
2-element Array{Pair{Symbol,Int64},1}:
 :FOO => 1
 :BAR => 2

julia> pairs"Foo :FOO => 1 Bar :BAR => 2"
2-element Array{Pair{Symbol,Int64},1}:
 :FOO => 1
 :BAR => 2
 ```

e.g. Define a pin assignment map for a Raspberry Pi GPIO header.

```julia
julia> pins = pairs\"\"\"
                               R'Pi GPIO Pins
                              ┌──────────────┐
                              │     1  2     │
                              │ 3V3 ■  ■ 5V  │
    :RTC_I2C => (2, 3)   (SDA)│   2 ■  ■ 5V  │      AVR_ISP_5V
                         (SCL)│   3 ■  ■ GND │
                              │   4 ■  ■ 14  │(TX) :CONSOLE_RX => 14
                              │ GND ■  ■ 15  │(RX) :CONSOLE_TX => 15
                              │  17 ■  ■ 18  │
                              │  27 ■  ■ GND │
                              │  22 ■  ■ 23  │     :AVR_ISP_RESET => 23
                              │ 3V3 ■  ■ 24  │     :AVR_ISP_MOSI  => 24
    :SENSOR_MOSI => 10  (MOSI)│  10 ■  ■ GND │      AVR_ISP_GND
    :SENSOR_MISO =>  9  (MISO)│   9 ■  ■ 25  │     :AVR_ISP_MISO  => 25
    :SENSOR_SCLK => 11  (SCLK)│  11 ■  ■  8  │     :AVR_ISP_CLK   =>  8
                              │ GND ■  ■  7  │
                         (IDD)│   0 ■  ■  1  │(IDC)
    :SENSOR_CS   =>  5        │   5 ■  ■ GND │
                              │   6 ■  ■ 12  │
                              │  13 ■  ■ GND │
                              │  19 ■  ■ 16  │     :LED_ERROR  => 16
                              │  26 ■  ■ 20  │     :LED_POWER  => 20
                              │ GND ■  ■ 21  │     :LED_STATUS => 21
                              │    39  40    │
                              └──────────────┘
\"\"\"
14-element Array{Pair{Symbol,B} where B,1}:
       :RTC_I2C => (2, 3)
    :CONSOLE_RX => 14
    :CONSOLE_TX => 15
 :AVR_ISP_RESET => 23
  :AVR_ISP_MOSI => 24
   :SENSOR_MOSI => 10
   :SENSOR_MISO => 9
  :AVR_ISP_MISO => 25
   :SENSOR_SCLK => 11
   :AVR_ISP_CLK => 8
     :SENSOR_CS => 5
     :LED_ERROR => 16
     :LED_POWER => 20
    :LED_STATUS => 21

julia> using PiGPIOMEM

julia> gpio = Dict(k => GPIOPin(v) for (k, v) in pins)

julia> set_output_mode(gpio[:LED_POWER])

julia> gpio[:LED_POWER][] = true
```
"""
module PairStrings

export @pairs_str, @docpairs_str

"""
Match `key => value`,
      `key => "value",
      `key => (value)` or
      `key => [value]`
"""
const pair_pattern = r"""
    [^\s]+                  # Pair key. Any non-whitespace characters.

    \s*                     # Optional whitespace.
    =>                      # Pair operator.
    \s*                     # Optional whitespace.

    (                       # Pair value options:
         \[   [^\]]*   \]   #   `[ ... ]`
     |   \(    [^)]*   \)   #   `( ... )`
     |    "    [^"]*    "   #   `" ... "`
     |        [^\s]+        #   Any non-whitespace characters.
    )
"""x

"""
    match_pair(@__MODULE__, string, [i=1]) => Pair, next_i

Search for `a => b` in `string` (starting at index `i`).
"""
function match_pair(mod, s, i=1)
    m = match(pair_pattern, s, i)
    if m == nothing
        m, i
    else
        pair = include_string(mod, m.match)
        i = m.offset + length(m.match)
        pair, i
    end
end

"""
    pairs(@__MODULE__, string) => Vector{Pair}

Extract `Vector{Pair}` from a string.
"""
function pairs(mod, s::AbstractString)
    result = []
    pair, i = match_pair(mod, s)
    while pair != nothing
        push!(result, pair)
        pair, i = match_pair(mod, s, i)
    end
    [x for x in result]
end

"""
    pairs"Foo :a => 1 Bar :b => 2" => [:a => 1, :b => 2]

Extract `Vector{Pair}` from a string.
"""
macro pairs_str(s)
    :(pairs($(esc(:(@__MODULE__))), $(esc(s))))
end

"""
    docpairs"Foo :a => 1 Bar :b => 2"
        => ("Foo :a => 1 Bar :b => 2",
            [:a => 1, :b => 2])

Extract `Vector{Pair}` from a string (and also return the original string).
"""
macro docpairs_str(s)
    :(($(esc(s)), pairs($(esc(:(@__MODULE__))), $(esc(s)))))
end

"""
    prefix_filter(prefix, Dict(pairs))
    prefix_filter("A_", Dict(pairs":A_X => 1 :B_X => 2 :A_Y => 3"))
        => Dict(:X => 1, Y => 3)

Filter a Dict of pairs based on a prefix.
Remove the prefix from the output keys.
"""
function prefix_filter(prefix, dict::Dict{K,V}) where {K, V}
    prefix = string(prefix)
    dict = filter(p->startswith(string(p[1]), prefix), dict)
    Dict(K(string(k)[1+length(prefix):end]) => v for (k, v) in dict)
end

end # module PairStrings


text:
const olm_mcts_c = 1e5 # exploration parameter for MCTS
const olm_mcts_iterations = 300 # exploration parameter for MCTS

mutable struct MctsNode
  x::AbstractArray{Float64, 1}
  u::AbstractArray{Float64, 1}
  q::Float64
  N::Int
  Ns::Int

  MctsNode() = new(Float64[], Float64[], 0.0, 1, 1)
  MctsNode(x, u, r) = new(x, u, r, 1, 1)
end

function plan_mcts(x::AbstractArray{Float64, 1}, agent::Agent, world::World,
                   reward::Function, ctrl_d::Discretization, depth::Int)
  x = copy(x)
  x[1] = mod(x[1], world.road.path.S[end])
  x = copy(x)
  value = MctsNode()
  value.x = x
  node = Tree(value)

  num_iterations = round(Int, olm_mcts_iterations / 4 * depth)

  visited = Set{MctsNode}()
  for i in 1:num_iterations
    simulate_mcts(node, agent, world, reward,
                  ctrl_d, visited, depth)
  end

  us = Float64[]
  qs = -Inf
  for a in 1:length(node.next)
    if node.next[a].value.q > qs
      qs = node.next[a].value.q
      us = node.next[a].value.u
    end
  end

  return us
end

function simulate_mcts(node::Tree, agent::Agent, world::World,
                       reward::Function, ctrl_d::Discretization,
                       visited::Set{MctsNode}, depth::Int)
  if depth <= 0
    return 0.0
  end

  if !(node.value in visited)
    push!(visited, node.value)

    la_len = ctrl_d.thr[end] * ctrl_d.pt[end] # number of actions to survey
    node.next = Array{Tree, 1}(undef, la_len)
    for la in 0:(la_len - 1)
      u = dis.ls2x(ctrl_d, la)
      agent.custom = u
      nx = copy(node.value.x)
      sim.advance!(sim.default_dynamics!, nx, Pair(agent, world), 0.0, olm_dt,
                   olm_h)
      nx[1] = mod(nx[1], world.road.path.S[end])

      r = reward(node.value.x, u, nx, agent, world)

      value = MctsNode(nx, u, r)
      node.next[la + 1] = Tree(value)
    end

    return rollout(node.value.x, agent, world, reward, ctrl_d, depth)
  end

  # find the best action for value and exploration
  as = -1
  vs = -Inf
  for a in 1:length(node.next)
    v = node.next[a].value.q + olm_mcts_c * sqrt(log(node.value.Ns) /
                                                 node.next[a].value.N)
    if v > vs
      vs = v
      as = a
    end
  end

  # transition to the new state deterministically
  x = node.value.x
  u = node.next[as].value.u
  nx = node.next[as].value.x
  r = reward(x, u, nx, agent, world)

  q = r + olm_gamma * simulate_mcts(node.next[as], agent, world, reward,
                                    ctrl_d, visited, depth - 1)

  # update Qsa for fuck's sake ---------------------------------------------- #
  #dq = (q - node.next[as].value.q)
  #node.next[as].value.q += dq / node.next[as].value.N

  # this works best BY FAR
  node.next[as].value.q = max(q, node.next[as].value.q)

  #alpha = 0.5
  #Qsa = node.next[as].value.q
  #Qsa = Qsa * (1 - alpha) + q * alpha
  #node.next[as].value.q = Qsa
  # ------------------------------------------------------------------------- #

  node.next[as].value.N += 1
  node.value.Ns += 1

  return q
end

function rollout(x::AbstractArray{Float64, 1}, agent::Agent,
                 world::World, reward::Function,
                 ctrl_d::Discretization, depth::Int)
  if depth <= 0
    return 0.0
  end

  # choose middle action, generally reasonable
  ad = fill(0, ctrl_d.dim)
  for i in 1:length(ad)
    ad[i] = div(ctrl_d.pt[i], 2)
  end
  u = dis.xd2x(ctrl_d, ad)
  agent.custom = u
  nx = copy(x)
  sim.advance!(sim.default_dynamics!, nx, Pair(agent, world), 0.0, olm_dt,
               olm_h)
  r = reward(x, u, nx, agent, world)

  return r + olm_gamma * rollout(nx, agent, world, reward, ctrl_d, depth - 1)
end

function controller_mcts!(u::AbstractArray{Float64},
                          x::AbstractArray{Float64},
                          dx::AbstractArray{Float64},
                          agent_world::Pair{Agent, World}, t::Float64)
  agent = agent_world.first
  u[1] = agent.custom[1]
  u[2] = agent.custom[2]
end


text:
using LinearAlgebra
using SparseArrays
using Plots

"This routine solves the periodic advection-diffusion equation using Forward Euler
in time and 2nd order finite differences in space."

m = 100 # number of points

# define spatial grid
x = LinRange(-1,1,m+2) # x_0, x_1, ..., x_m, x_{m+1} = x_0
xint = x[1:end-1]
h = x[2]-x[1]

a = 1.0 # advection speed
ϵ = h/2 # diffusion coefficient
dt = h # timestep
T = 2.0 # final time

# initial condition and forcing
u0(x) = sin(pi*x)
f(x,t) = 0.0

# approximation of second derivative: largest eig = O(1/h^2)
A = (1/h^2) * spdiagm(0=>2*ones(m+1),-1=>-ones(m),1=>-ones(m))
A[1,end] = -1/h^2
A[end,1] = -1/h^2

# approximation of first derivative: largest eig = O(1/h)
Q = (1/(2*h)) * spdiagm(1=>ones(m),-1=>-ones(m))
Q[1,end] = -1/(2*h)
Q[end,1] = 1/(2*h)

u = u0.(xint)
Nsteps = ceil(Int,T/dt)
dt = T / Nsteps

interval = 10
@gif for k = 1:Nsteps
    tk = k*dt
    u .= u + dt * (f.(xint,tk) - (a*Q + ϵ*A)*u)
    if k % interval==0
        plot(xint,u,linewidth=2,label="Solution",ylims=(-1.5,1.5))
        println("on timestep $k out of $Nsteps.")
    end
end every interval

maxλ = Float64[]
ϵvec = .0:.0001:.025
for ϵ = ϵvec
    append!(maxλ,maximum(abs.(eigvals(Matrix(I - dt*(Q + ϵ*A))))))
end
scatter(ϵvec,maxλ,leg=false)
plot!(title="Largest eigenvalue magnitude with ϵ")


text:
######################### Initialization #######################
function checkPRatio(Pratio)
    if Pratio < 0
        throw(ArgumentError("Pratio must be positive"))
    end
    if Pratio > 1
        Pratio = 1/Pratio
    end

    return Pratio
end

function initializeShockTubeFDM(nCells=100; domainLength=1, Pratio=10)
    Pratio = checkPRatio(Pratio)

    # Create arrays to store data (cell # = position in array)
    dx = Array{Float64, 1}(undef, nCells)
    U = Array{Float64, 1}(undef, nCells)
    P = Array{Float64, 1}(undef, nCells)
    T = Array{Float64, 1}(undef, nCells)

    # Apply initial conditions (Fig. 1 in Henry's paper)
    for i in 1:nCells
        if i <= (nCells/2)
            T[i] = 0.00348432
            P[i] = 1
        else
            T[i] = 0.00278746
            P[i] = Pratio
        end
        U[i] = 0

        # dx = Distance between cell centers i and i+1
        dx[i] = domainLength / nCells
    end

    return dx, P, T, U
end

function initializeShockTube_StructuredFVM(nCells=100; domainLength=1, Pratio=10, silent=true)
    if !silent
        println("Meshing shock tube, $nCells cells")
    end
    dx, P, T, U = initializeShockTubeFDM(nCells, domainLength=domainLength, Pratio=Pratio)

    cellPrimitives = Array{Float64, 2}(undef, nCells, 3)
    for c in 1:nCells
        cellPrimitives[c,:] = [ P[c], T[c], U[c] ]
    end

    return dx, cellPrimitives
end
#
# Wrapper for FDM initialization function, adding a mesh definition suitable for FVM and vector-format velocity
function initializeShockTubeFVM(nCells=100; domainLength=1, Pratio=10, silent=true)
    if !silent
        println("Meshing shock tube, $nCells cells")
    end
    dx, P, T, U = initializeShockTubeFDM(nCells, domainLength=domainLength, Pratio=Pratio)
    U = []

    #Shock tube dimensions
    h = 0.1
    w = 0.1

    cells = []
    faces = []
    fAVecs = []
    fCenters = []
    boundaryFaces = [ [nCells,], [nCells+1,] ]
    cVols = []
    cCenters = []

    fAVec = [h*w, 0, 0]
    cV = h*w*domainLength/nCells
    dx = dx[1]
    for i in 1:nCells
        # Modify natural face numbering to have boundary faces numbered last
        if i == 1
            push!(cells, [nCells, 1])
            push!(faces, [i, i+1])
        elseif i == nCells
            push!(cells, [nCells-1, nCells+1])
        else
            push!(cells, [i-1, i])
            push!(faces, [i, i+1])
        end

        push!(U, [0.0, 0.0, 0.0] )
        push!(cVols, cV)
        push!(fAVecs, fAVec)
        push!(cCenters, [ (i-0.5)*dx, 0.0, 0.0 ])

        # Face centers also adjusted to have boundaries last
        if i != nCells
            push!(fCenters, [ i*dx, 0.0, 0.0 ])
        else
            # Left boundary face
            push!(fCenters, [ 0.0, 0.0, 0.0 ])
        end
    end

    # Last face
    push!(fAVecs, fAVec)
    # Boundary faces
    push!(faces, [-1,1])
    push!(faces, [nCells, -1])
    push!(fCenters, [ nCells*dx, 0.0, 0.0 ])

    # Returns in mesh format
    mesh = [ cells, cVols, cCenters, faces, fAVecs, fCenters, boundaryFaces ]
    return mesh, P, T, U
end

# Wrapper for FDM initialization function, adding a mesh definition suitable for FVM and vector-format velocity
function initializeShockTube3DFVM(nCells=100; domainLength=1, Pratio=10, silent=true)
    if !silent
        println("Meshing shock tube, $nCells cells")
    end
    dx, P, T, U = initializeShockTubeFDM(nCells, domainLength=domainLength, Pratio=Pratio)

    cellPrimitives = Array{Float64, 2}(undef, nCells, 5)
    for c in 1:nCells
        cellPrimitives[c,:] = [ P[c], T[c], 0.0, 0.0, 0.0 ]
    end

    #Shock tube dimensions
    h = 0.1
    w = 0.1

    cells = []
    faces = []
    fAVecs = []
    fCenters = []
    boundaryFaces = [ [nCells,], [nCells+1,], [] ]
    cVols = []
    cCenters = []
    cellSizes = zeros(nCells, 3)

    fAVec = [h*w, 0, 0]
    cV = h*w*domainLength/nCells
    dx = dx[1]
    for i in 1:nCells
        # Modify natural face numbering to have boundary faces numbered last
        if i == 1
            push!(cells, [nCells, 1])
            push!(faces, [i, i+1])
            push!(fAVecs, fAVec)
        elseif i == nCells
            push!(cells, [nCells-1, nCells+1])
        else
            push!(cells, [i-1, i])
            push!(faces, [i, i+1])
            push!(fAVecs, fAVec)
        end

        push!(cVols, cV)
        push!(cCenters, [ (i-0.5)*dx, 0.0, 0.0 ])

        # Face centers also adjusted to have boundaries last
        if i != nCells
            push!(fCenters, [ i*dx, 0.0, 0.0 ])
        end

        cellSizes[i, :] = [dx, h, w]
    end

    # Boundary faces
    push!(faces, [1,-1])
    push!(fAVecs, -fAVec)
    push!(fCenters, [ 0.0, 0.0, 0.0 ])

    push!(faces, [nCells, -1])
    push!(fAVecs, fAVec)
    push!(fCenters, [ nCells*dx, 0.0, 0.0 ])

    # Add the empty faces on the sides
    for c in 1:nCells
        push!(faces, [c, -1])
        push!(faces, [c, -1])
        push!(faces, [c, -1])
        push!(faces, [c, -1])
        push!(fCenters, [ (c-0.5)*dx, 0.05, 0.0 ])
        push!(fCenters, [ (c-0.5)*dx, -0.05, 0.0 ])
        push!(fCenters, [ (c-0.5)*dx, 0.0, 0.05 ])
        push!(fCenters, [ (c-0.5)*dx, 0.0, -0.05 ])
        push!(boundaryFaces[3], nCells+1+c)
        push!(boundaryFaces[3], nCells+2+c)
        push!(boundaryFaces[3], nCells+3+c)
        push!(boundaryFaces[3], nCells+4+c)
        push!(cells[c], nCells+1+c)
        push!(cells[c], nCells+2+c)
        push!(cells[c], nCells+3+c)
        push!(cells[c], nCells+4+c)
        push!(fAVecs, [0, h*w, 0])
        push!(fAVecs, [0, -h*w, 0])
        push!(fAVecs, [0, 0, h*w])
        push!(fAVecs, [0, 0, -h*w])
    end

    # Returns in mesh format
    mesh = Mesh(cells, cVols, cCenters, cellSizes, faces, fAVecs, fCenters, boundaryFaces)
    return mesh, cellPrimitives
end

############################ Plotting ############################
function plotShockTubeResults_Plotly(P, U, T, rho)
    plots = []
    xAxis = Array{Float64, 1}(undef, nCells)
    for i in 1:nCells
        xAxis[i] = i/nCells - 1/(2*nCells)
    end

    pPlot = plot(xAxis, P, label="P (Pa)", title="Pressure", xlabel="x (m)", bottom_margin=15mm, left_margin=10mm)
    rhoPlot = plot(xAxis, rho, label="rho (kg/m3)", title="Density", xlabel="x (m)", bottom_margin=15mm, left_margin=10mm)
    uPlot = plot(xAxis, U, label="Velocity (m/s)", title="Velocity", xlabel="x (m)", bottom_margin=15mm, left_margin=10mm)
    TPlot = plot(xAxis, T, label="T (K)", title="Temperature", xlabel="x (m)", bottom_margin=15mm, left_margin=10mm)
    plots = [pPlot, rhoPlot, uPlot, TPlot]
    plot(plots..., layout=(2, 2), size=(1720, 880), window_title="Euler1D_Draft_Henry", legend=false)
    gui()
end

function plotShockTubeResults_PyPlot(P, U, T, rho)
    plots = []
    xAxis = Array{Float64, 1}(undef, nCells)
    for i in 1:nCells
        xAxis[i] = i/nCells - 1/(2*nCells)
    end

    pPlot = plot(xAxis, P, label="P (Pa)", title="Pressure", xlabel="x (m)")
    rhoPlot = plot(xAxis, rho, label="rho (kg/m3)", title="Density", xlabel="x (m)")
    uPlot = plot(xAxis, U, label="Velocity (m/s)", title="Velocity", xlabel="x (m)")
    TPlot = plot(xAxis, T, label="T (K)", title="Temperature", xlabel="x (m)")
    plots = [pPlot, rhoPlot, uPlot, TPlot]
    plot(plots..., layout=(2, 2), size=(860, 600), window_title="Euler1D_Draft_Henry", legend=false)
    gui()
end

max_stars_repo_name: henrystoldt/fvCFD


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

          Solving the HJB with a state constraint using
                the explicit method

               Translated from matlab code on Ben Moll's website:
                              http://www.princeton.edu/~moll/HACTproject.htm

==============================================================================#

using LinearAlgebra, SparseArrays, Plots

γ = 2 # parameter from CRRA utility
r = 0.035 # the interest rate, exogenous in this version
ρ = 0.05 # the discount rate

z1=.1
z2 =.2
z= [z1 z2]

λ1 = 1.5
λ2 = 1.0

λ = [λ1 λ2]

H = 500 # number of points in the grid space
amin=-0.02
amax= 3

a = LinRange(amin,amax,H)
da = (amax-amin)/(H-1)

aa = [a a]
zz = ones(H,1)*z

maxit = 100 # the maximum number of iterations we allow the finite differencing algorithm
crit = 10^(-6) # our critical value
Δ = 1000

dVf=zeros(H,2)
dVb = zeros(H,2)
c= zeros(H,2)

# Create the matrix that captures the systems dynamic from λ process
Aswitch = [-sparse(I,H,H)*λ[1] sparse(I,H,H)*λ[1]; sparse(I,H,H)*λ[2] -sparse(I,H,H)*λ[2]]

# Inital Guess
v0=zeros(H,2)
v0[:,1] = (z[1] .+ r.*a).^(1-γ)/(1-γ)/ρ
v0[:,2] = (z[2] .+ r.*a).^(1-γ)/(1-γ)/ρ
global v=v0

dist=[]
V_n=[]

# The finite differeing loop

for n in 1:maxit
    global V=v
    push!(V_n,V)
    # Forward differencing
    dVf[1:H-1,:]=(V[2:H,:]-V[1:H-1,:])/da
    dVf[H,:] = (z .+ r.*amax).^(-γ) # impose state constraint at the max, just in case
    # Backward differencing
    dVb[2:H,:]=(V[2:H,:]-V[1:H-1,:])/da
    dVb[1,:]= (z .+r.*amin).^(-γ)

    # Find consumption and savings with the forward difference
    cf = dVf.^(-1/γ)
    sf = zz +r.*aa -cf

    # Find consumption and savings with the backward difference
    cb= dVb.^(-1/γ)
    sb = zz +r.*aa -cb

    # Find consumption and savings at steady state
    c0 = zz + r.*aa
    dV0 = c0.^(-γ)

    # Now implement the upwind scheme in order to select the best differencing method
    If = sf .> 0
    Ib = sb .< 0
    I0 = (1 .- If - Ib)

    # State constraint at amin is automatically implemented

    dV_Upwind = dVf.*If + dVb.*Ib + dV0.*I0

    global c = dV_Upwind.^(-1/γ)
    u = c.^(1-γ)/(1-γ)
    # Construct matrix for the evolution of the system
    X = -min.(sb,0)/da
    Y= -max.(sf,0)/da + min.(sb,0)/da
    Z = max.(sf,0)/da

    A1 = spdiagm(0=>Y[:,1], -1 => X[2:H,1], 1=> Z[1:H-1,1])
    A2 = spdiagm(0=>Y[:,2], -1 => X[2:H,2], 1=> Z[1:H-1,2])
    AA = [A1 spzeros(H,H); spzeros(H,H) A2]

    global A = AA + Aswitch

    B = (ρ +1/Δ)*sparse(I,2*H,2*H) -A
    u_stacked = [u[:,1]; u[:,2]]
    V_stacked = [V[:,1]; V[:,2]]

    b = u_stacked + V_stacked/Δ
    V_stacked = B\b # Solves the system of equations

    V = [V_stacked[1:H] V_stacked[H+1:2*H]]

    V_change = V - v

    global v = V

    push!(dist, findmax(abs.(V_change))[1])

    if dist[n] < crit
        println("Value Function converged, Iteration=")
        println(n)
        break
    end
end

#====================================================
        Fokker-Planck Equation
====================================================#

AT = A' # we need the transpose of our A matrix
b =zeros(2*H,1)

# We need to fix one value to prevent the matrix from being singular
i_fix =1
b[i_fix] = .1
row = [zeros(1,i_fix-1) 1 zeros(1,2*H-i_fix)]
AT[i_fix,:] = row

# Using this solve the linear system
gg = AT\b
g_sum = gg'*ones(2*H,1)*da
gg = gg./g_sum # we need gg to sum to one

g = [gg[1:H] gg[H+1:2*H]]

# Graphs
adot = zz + r.*aa -c
plot(a, adot, xlabel="a", ylabel="\$s_{i}(a)\$", xlims=(amin,amax))
plot!(a, zeros(H,1), label="")

S = g[:,1]'*a*da + g[:,2]'*a*da

# Approximation at the borrowing constraint
u1=(z1+r*amin)^(-γ)
u2 = c[1,2]^(-γ)
u11 = -γ*(z1+r*amin)^(-γ-1)
ν = sqrt(-2*((ρ-r)*u1 + λ1*(u1-u2))/u11)
s_approx = -ν*(a.-amin).^(1/2)

amax1=1

plot(a, adot[:,1], xlabel="a", ylabel="\$s_{i}(a)\$", xlims=(amin,amax1), label="\$s_{1}(a)\$")
plot!(a, adot[:,2], label="\$s_{2}(a)\$")
plot!(a, zeros(H,1), line=:dash, label="", color=:black)
plot!(a, s_approx, label="Approximatation")

plot(a,g[:,1],xlabel="a", ylabel="\$g_{i}(a)\$", label="\$g_{1}(a)\$", xlims=(amin,amax))
plot!(a,g[:,2],label="\$g_{2}(a)\$")


text:
const digits = "123456789ABCDEFGHJKLMNPQRSTUVWXYZabcdefghijkmnopqrstuvwxyz"
function encodebase58(b::Vector{<:Integer})
    out = Vector{Char}(34)
    for j in endof(out):-1:1
        local c::Int = 0
        for i in eachindex(b)
            c = c * 256 + b[i]
            b[i] = c ÷ 58
            c %= 58
        end
        out[j] = digits[c + 1]
    end
    local i = 1
    while i < endof(out) && out[i] == '1'
        i += 1
    end
    return join(out[i:end])
end

const COINVER = 0x00
function pubpoint2address(x::Vector{UInt8}, y::Vector{UInt8})
    c = vcat(0x04, x, y)
    c = vcat(COINVER, digest("ripemd160", sha256(c)))
    d = sha256(sha256(c))
    return encodebase58(vcat(c, d[1:4]))
end
pubpoint2address(x::AbstractString, y::AbstractString) =
    pubpoint2address(hex2bytes(x), hex2bytes(y))

max_stars_repo_name: ethansaxenian/RosettaDecode

max_issues_repo_name: ethansaxenian/RosettaDecode

max_forks_repo_name: ethansaxenian/RosettaDecode


text:
using JuliaDFT
using Test

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

max_stars_repo_name: mailhexu/JuliaDFT

max_issues_repo_name: mailhexu/JuliaDFT

max_forks_repo_name: mailhexu/JuliaDFT


text:
function make_mIHTvar(r) # r = number of traits
    n = 1000  # number of samples
    p = 10000 # number of SNPs
    k = 10    # number of causal SNPs

    # set random seed for reproducibility
    Random.seed!(2021)

    # simulate `.bed` file with no missing data
    x = simulate_random_snparray(undef, n, p)
    xla = SnpLinAlg{Float64}(x, model=ADDITIVE_MODEL, center=true, scale=true)

    # intercept is the only nongenetic covariate
    z = ones(n, 1)

    # simulate response y, true model b, and the correct non-0 positions of b
    Y, true_Σ, true_b, correct_position = simulate_random_response(xla, k, r);

    # everything is transposed for multivariate analysis!
    Yt = Matrix(Y')
    Zt = Matrix(z')
    v = MendelIHT.mIHTVariable(Transpose(xla), Zt, Yt, k, trues(1))

    return v, Transpose(xla), Yt, Zt, k
end

@testset "update_support!" begin
    idx = BitVector([false, false, true, true, false])
    b = rand(5, 5)
    MendelIHT.update_support!(idx, b)
    @test all(idx)

    b[:, 1] .= 0.0
    MendelIHT.update_support!(idx, b)
    @test idx[1] == false
    @test all(idx[2:5] .== true)

    b[3, 3] = 0.0
    MendelIHT.update_support!(idx, b)
    @test idx[1] == false
    @test all(idx[2:5] .== true)

    b[:, 3] .= 0.0
    MendelIHT.update_support!(idx, b)
    @test idx[1] == false
    @test idx[2] == true
    @test idx[3] == false
    @test idx[4] == true
    @test idx[5] == true
end

@testset "benchmarks for multivariate function" begin

    # r = 2 traits
    v, xla, Yt, Zt, k = make_mIHTvar(2)
    zkeep = trues(1)
    MendelIHT.init_iht_indices!(v, false, trues(size(xla, 2)), false)
    @btime MendelIHT.score!($v)           # 9.304 ms (45 allocations: 2.33 KiB)
    @btime MendelIHT.loglikelihood($v)    # 3.377 μs (2 allocations: 208 bytes)
    @btime MendelIHT.update_xb!($v)       # 187.489 μs (8 allocations: 656 bytes)
    @btime MendelIHT.update_μ!($v)        # 2.337 μs (0 allocations: 0 bytes)
    @btime MendelIHT._choose!($v)         # 16.698 μs (3 allocations: 208 bytes)
    @btime MendelIHT.solve_Σ!($v)         # 2.889 μs (2 allocations: 48 bytes)
    # @btime MendelIHT.iht_stepsize(v) # cannot be done because BenchmarkTools require multiple evaluations of the same function call, which makes Γ not pd
    @btime MendelIHT.vectorize!($(v.full_b), $(v.B), $(v.C), $zkeep) # 10.299 μs (0 allocations: 0 bytes)
    @btime MendelIHT.unvectorize!($(v.full_b), $(v.B), $(v.C), $zkeep) # 10.548 μs (0 allocations: 0 bytes)
    @btime MendelIHT.update_support!($(v.idx), $(v.B)) # 23.204 μs (0 allocations: 0 bytes)

    # r = 10 traits
    v, xla, Yt, Zt, k = make_mIHTvar(10)
    MendelIHT.init_iht_indices!(v, false, trues(size(xla, 2)), false)
    @btime MendelIHT.score!($v)           # 43.252 ms (221 allocations: 11.58 KiB)
    @btime MendelIHT.loglikelihood($v)    # 12.208 μs (2 allocations: 1.03 KiB)
    @btime MendelIHT.update_xb!($v)       # 219.192 μs (8 allocations: 656 bytes)
    @btime MendelIHT.update_μ!($v)        # 1.963 μs (0 allocations: 0 bytes)
    @btime MendelIHT._choose!($v)         # 68.161 μs (5 allocations: 432 bytes)
    @btime MendelIHT.solve_Σ!($v)         # 12.623 μs (2 allocations: 48 bytes)
    # @btime MendelIHT.iht_stepsize(v) # cannot be done because BenchmarkTools require multiple evaluations of the same function call, which makes Γ not pd
    @btime MendelIHT.vectorize!($(v.full_b), $(v.B), $(v.C), $zkeep) # 23.712 μs (0 allocations: 0 bytes)
    @btime MendelIHT.unvectorize!($(v.full_b), $(v.B), $(v.C), $zkeep) # 23.761 μs (0 allocations: 0 bytes)
    @btime MendelIHT.update_support!($(v.idx), $(v.B)) # 50.047 μs (0 allocations: 0 bytes)
end

@testset "multivariate fit_iht SnpLinAlg" begin
    n = 1000  # number of samples
    p = 10000 # number of SNPs
    k = 10    # number of causal SNPs
    r = 2     # number of traits

    # set random seed for reproducibility
    Random.seed!(2021)

    # simulate `.bed` file with no missing data
    x = simulate_random_snparray(undef, n, p)
    xla = SnpLinAlg{Float64}(x, model=ADDITIVE_MODEL, center=true, scale=true)

    # intercept is the only nongenetic covariate
    z = ones(n, 1)
    intercepts = [10.0 1.0] # each trait have different intercept

    # simulate response y, true model b, and the correct non-0 positions of b
    Y, true_Σ, true_b, correct_position = simulate_random_response(xla, k, r, Zu=z*intercepts, overlap=2)
    correct_snps = [x[1] for x in correct_position] # causal snps
    Yt = Matrix(Y'); # in MendelIHT, multivariate traits should be rows

    # Y and X should be transposed
    @test_throws DimensionMismatch fit_iht(Yt, xla, k=12)
    @test_throws DimensionMismatch fit_iht(Y, Transpose(xla), k=12)

    # no debias
    @time result = fit_iht(Yt, Transpose(xla), k=12, debias=false)
    @test size(result.beta) == (r, p)
    @test result.k == 12
    @test result.traits == 2
    @test result.iter ≥ 5
    @test all(result.σg .> 0)
    # @test all(result.beta[1, correct_snps] - true_b[correct_snps, 1] .< 0.15) # estimates are close to truth
    # @test all(result.beta[2, correct_snps] - true_b[correct_snps, 2] .< 0.15) # estimates are close to truth

    # yes debias
    @time result2 = fit_iht(Yt, Transpose(xla), k=12, debias=true)
    @test size(result2.beta) == (r, p)
    @test result2.k == 12
    @test result2.traits == 2
    @test result2.iter ≥ 5
    @test all(result2.σg .> 0)
    # @test all(result2.beta[1, correct_snps] - true_b[correct_snps, 1] .< 0.15) # estimates are close to truth
    # @test all(result2.beta[2, correct_snps] - true_b[correct_snps, 2] .< 0.15) # estimates are close to truth
end

@testset "initialze beta" begin
    n = 1000  # number of samples
    p = 10000 # number of SNPs
    k = 10    # number of causal SNPs
    r = 2     # number of traits

    # set random seed for reproducibility
    Random.seed!(2021)

    # simulate `.bed` file with no missing data
    x = simulate_random_snparray(undef, n, p)
    xla = SnpLinAlg{Float64}(x, model=ADDITIVE_MODEL, center=true, scale=true)

    # intercept is the only nongenetic covariate
    z = ones(n, 1)
    intercepts = [10.0 1.0] # each trait have different intercept

    # simulate response y, true model b, and the correct non-0 positions of b
    Y, true_Σ, true_b, correct_position = simulate_random_response(xla, k, r, Zu=z*intercepts, overlap=2)
    correct_snps = [x[1] for x in correct_position] # causal snps
    Yt = Matrix(Y'); # in MendelIHT, multivariate traits should be rows

    @time result = fit_iht(Yt, Transpose(xla), k=12, init_beta=true)
    @test size(result.beta) == (r, p)
    @test result.k == 12
    @test result.traits == 2
    @test result.iter ≥ 5
    @test all(result.σg .> 0)
end


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

using JuliaFEM
using Test

abstract type HeatProblem <: AbstractProblem

end

function HeatProblem(dim::Int=1, elements=[])
    return Problem{HeatProblem}(dim, elements)
end

function get_unknown_field_name(::Type{P}) where P<:HeatProblem
    return "temperature"
end

function get_unknown_field_type(::Type{P}) where P<:HeatProblem
    return Float64
end

""" Calculate a potential Π = Wint - Wext of system. """
function get_potential_energy(problem::Problem{HeatProblem}, element::Element{Quad4}, ip::IntegrationPoint, time::Number; variation=nothing)
    k = element("temperature thermal conductivity", ip, time)
    f = element("temperature load", ip, time)
    T = element("temperature", ip, time, variation)
    c = element("temperature nonlinearity coefficient", ip, time)
    gradT = element("temperature", ip, time, Val{:grad}, variation)
    Wint = (k + c*T) * 1/2*vecdot(gradT, gradT)
    Wext = f*T
    W = Wint - Wext
    J = get_jacobian(element, ip, time)
    return W*det(J)
end

function get_potential_energy(problem::Problem{HeatProblem}, element::Element{Seg2}, ip::IntegrationPoint, time::Number; variation=nothing)
    T = element("temperature", ip, time, variation)[1]
    T_ext = element("temperature external", ip, time)[1]
    coeff = element("temperature coefficient", ip, time)[1]
    q0 = coeff*(T_ext^4 - T^4)
    Wint = 0.0
    Wext = q0*T
    W = Wint - Wext
    J = get_jacobian(element, ip, time)
    return W*norm(J)
end

function test_potential_energy_method()

    # create model -- start
    element = Quad4([1, 2, 3, 4])
    element["geometry"] = Vector[[0.0,0.0], [1.0,0.0], [1.0,1.0], [0.0,1.0]]
    element["temperature thermal conductivity"] = 6.0
    element["temperature load"] = [0.0, 0.0, 0.0, 0.0]
    element["temperature nodal load"] = [3.0, 3.0, 0.0, 0.0]
    element["temperature nonlinearity coefficient"] = 6.0
    element["temperature"] = (0.0 => zeros(Float64, 4))
    problem = HeatProblem()
    push!(problem, element)
    # create model -- end

    solve!(problem, [1, 2], 0.0)
    temp = element("temperature", [0.0, -1.0], 0.0)
    err = temp - 2/3
    @info("error: $err")
    @test isapprox(err, 0.0)
end

function test_potential_energy_method_2()

    # create model -- start
    N = Vector[[0.0, 0.0], [1.0, 0.0], [1.0, 1.0], [0.0, 1.0]]
    element1 = Quad4([1, 2, 3, 4])
    element1["geometry"] = Vector[N[1], N[2], N[3], N[4]]
    element1["temperature thermal conductivity"] = 6.0
    element1["temperature load"] = [0.0, 0.0, 0.0, 0.0]
    element1["temperature nonlinearity coefficient"] = [0.0, 0.0, 0.0, 0.0]
    element1["temperature"] = (0.0 => zeros(Float64, 4))

    element2 = Seg2([1, 2])
    element2["geometry"] = Vector[N[1], N[2]]
    element2["temperature coefficient"] = 3.0e-8 # ~ 5.7e-8 * 0.5
    element2["temperature external"] = 100.0
    element2["temperature"] = (0.0 => zeros(Float64, 2))
    # create model -- end

    problem = HeatProblem()
    push!(problem, element1)
    push!(problem, element2)
    solve!(problem, [1, 2], 0.0)

    temp = element1("temperature", [0.0, -1.0], 0.0)
    err = temp - 0.5
    @info("error: $err")
    @test isapprox(err, 0.0, atol=1.0e-6)

    # @test isapprox(temp, 2.93509690572300E+00)  # tested using Code Aster
end


text:
# Struct for use in linear system solver
struct AOᴸ
  ψ::InfiniteCanonicalMPS
  H::InfiniteMPOMatrix
  n::Int
end

function (A::AOᴸ)(x)
  ψ = A.ψ
  H = A.H
  ψᴴ = dag(ψ)
  ψ′ = ψᴴ'
  n = A.n
  N = nsites(ψ)
  #@assert n == N

  l = linkinds(only, ψ.AL)
  l′ = linkinds(only, ψ′.AL)
  r = linkinds(only, ψ.AR)
  r′ = linkinds(only, ψ′.AR)
  δˡ(n) = δ(l[n], l′[n])
  δʳ(n) = δ(dag(r[n]), prime(r[n]))
  xT = translatecell(translator(ψ), x, -1)
  for j in (2 - N):1
    xT = xT * H[j][n, n] * ψ.AL[j] * ψ′.AL[j]
  end
  xR = x * ψ.C[1] * (ψ′.C[1] * δʳ(1)) * denseblocks(δˡ(1))
  return xT - xR
end

function apply_local_left_transfer_matrix(
  Lstart::Vector{ITensor}, H::InfiniteMPOMatrix, ψ::InfiniteCanonicalMPS, n_1::Int64
)
  dₕ = length(Lstart)
  ψ′ = dag(ψ)'

  Ltarget = Vector{ITensor}(undef, size(H[n_1])[1])
  for j in 1:dₕ
    init = false
    for k in reverse(j:dₕ)
      if !isempty(H[n_1][k, j]) && isassigned(Lstart, k) && !isempty(Lstart[k])
        if isassigned(Ltarget, j) && init
          Ltarget[j] += Lstart[k] * ψ.AL[n_1] * ψ′.AL[n_1] * H[n_1][k, j]
        else
          Ltarget[j] = Lstart[k] * ψ.AL[n_1] * ψ′.AL[n_1] * H[n_1][k, j]
          init = true
        end
      end
    end
  end
  return Ltarget
end

# apply the left transfer matrix at position n1 to the vector Lstart considering it at position m, adding to Ltarget
function apply_local_left_transfer_matrix(
  Lstart::ITensor,
  m::Int64,
  H::InfiniteMPOMatrix,
  ψ::InfiniteCanonicalMPS,
  n_1::Int64;
  reset=true,
)
  Ltarget = Vector{ITensor}(undef, size(H[n_1])[1])
  for j in 1:m
    if !isempty(H[n_1][m, j])
      Ltarget[j] = Lstart * ψ.AL[n_1] * H[n_1][m, j] * dag(prime(ψ.AL[n_1]))
    end
  end
  return Ltarget
end

#apply the left transfer matrix n1:n1+nsites(ψ)-1
function apply_left_transfer_matrix(
  Lstart::ITensor, m::Int64, H::InfiniteMPOMatrix, ψ::InfiniteCanonicalMPS, n_1::Int64
)
  Ltarget = apply_local_left_transfer_matrix(Lstart, m, H, ψ, n_1)
  for j in 1:(nsites(ψ) - 1)
    Ltarget = apply_local_left_transfer_matrix(Ltarget, H, ψ, n_1 + j)
  end
  return Ltarget
end

# Also input C bond matrices to help compute the right fixed points
# of ψ (R ≈ C * dag(C))
function left_environment(H::InfiniteMPOMatrix, ψ::InfiniteCanonicalMPS; tol=1e-10)
  N = nsites(H)
  @assert N == nsites(ψ)

  # Do the 1-site case first
  ψ′ = dag(ψ)'

  l = linkinds(only, ψ.AL)
  l′ = linkinds(only, ψ′.AL)
  r = linkinds(only, ψ.AR)
  r′ = linkinds(only, ψ′.AR)
  s = siteinds(only, ψ)
  δʳ(n) = δ(dag(r[n]), prime(r[n]))
  δˡ(n) = δ(l[n], l′[n])
  δˢ(n) = δ(dag(s[n]), prime(s[n]))

  eₗ = [0.0]
  dₕ = size(H[1])[1]
  Ls = [Vector{ITensor}(undef, dₕ) for j in 1:N]
  #Building the L vector for n_1 = 1
  # TM is 2 3 ... N 1
  Ls[1][end] = δˡ(1)   # exact by construction
  localR = ψ.C[1] * δʳ(1) * ψ′.C[1] #to revise
  for n in reverse(1:(dₕ - 1))
    temp_Ls = apply_left_transfer_matrix(
      translatecell(translator(ψ), Ls[1][n + 1], -1), n + 1, H, ψ, 2 - N
    )
    for j in 1:n
      if isassigned(temp_Ls, j)
        if isassigned(Ls[1], j)
          Ls[1][j] += temp_Ls[j]
        else
          Ls[1][j] = temp_Ls[j]
        end
      end
    end
    if !isempty(H[1][n, n])
      λ = H[1][n, n][1, 1]
      δiag = δˢ(1)
      #@assert norm(H[1][n, n] - λ * δiag) == 0 "Non identity diagonal not implemented in MPO"
      @assert abs(λ) <= 1 "Diverging term"
      eₗ[1] = (Ls[1][n] * localR)[]
      Ls[1][n] += -(eₗ[1] * denseblocks(δˡ(1)))
      if λ == 1
        A = AOᴸ(ψ, H, n)
        Ls[1][n], info = linsolve(A, Ls[1][n], 1, -1; tol=tol)
      else
        println("Not implemented")
        flush(stdout)
        flush(stderr)
      end
    end
  end
  for n in 2:N
    Ls[n] = apply_local_left_transfer_matrix(Ls[n - 1], H, ψ, n)
  end
  return CelledVector(Ls), eₗ[1]
end

# Struct for use in linear system solver
struct AOᴿ
  ψ::InfiniteCanonicalMPS
  H::InfiniteMPOMatrix
  n::Int
end

function (A::AOᴿ)(x)
  ψ = A.ψ
  H = A.H
  ψᴴ = dag(ψ)
  ψ′ = ψᴴ'
  n = A.n
  N = nsites(ψ)
  #@assert n == N

  l = linkinds(only, ψ.AL)
  l′ = linkinds(only, ψ′.AL)
  r = linkinds(only, ψ.AR)
  r′ = linkinds(only, ψ′.AR)
  δˡ(n) = δ(l[n], l′[n])
  δʳ(n) = δ(dag(r[n]), prime(r[n]))
  xT = translatecell(translator(ψ), x, 1)
  for j in reverse(1:N)
    xT = xT * ψ.AR[j] * H[j][n, n] * ψ′.AR[j]
  end
  xR = x * ψ.C[0] * (ψ′.C[0] * δˡ(0)) * denseblocks(δʳ(0))
  return xT - xR
end

# apply the left transfer matrix at position n1 to the vector Lstart, replacing Ltarget
function apply_local_right_transfer_matrix!(
  Lstart::Vector{ITensor}, H::InfiniteMPOMatrix, ψ::InfiniteCanonicalMPS, n_1::Int64
)
  dₕ = length(Lstart)
  ψ′ = dag(ψ)'
  for j in reverse(1:dₕ)
    init = false
    for k in reverse(1:j)
      if !isempty(H[n_1][j, k]) && isassigned(Lstart, k) && !isempty(Lstart[k])
        if isassigned(Lstart, j) && init
          Lstart[j] += Lstart[k] * ψ.AR[n_1] * H[n_1][j, k] * ψ′.AR[n_1]
        else
          Lstart[j] = Lstart[k] * ψ.AR[n_1] * H[n_1][j, k] * ψ′.AR[n_1]
          init = true
        end
      end
    end
  end
end

function apply_local_right_transfer_matrix(
  Lstart::Vector{ITensor}, H::InfiniteMPOMatrix, ψ::InfiniteCanonicalMPS, n_1::Int64
)
  dₕ = length(Lstart)
  ψ′ = dag(ψ)'
  Ltarget = Vector{ITensor}(undef, size(H[n_1])[1])
  for j in reverse(1:dₕ)
    init = false
    for k in reverse(1:j)
      if !isempty(H[n_1][j, k]) && isassigned(Lstart, k) && !isempty(Lstart[k])
        if isassigned(Ltarget, j) && init
          Ltarget[j] += Lstart[k] * ψ.AR[n_1] * H[n_1][j, k] * ψ′.AR[n_1]
        else
          Ltarget[j] = Lstart[k] * ψ.AR[n_1] * H[n_1][j, k] * ψ′.AR[n_1]
          init = true
        end
      end
    end
  end
  return Ltarget
end

# apply the left transfer matrix at position n1 to the vector Lstart considering it at position m, adding to Ltarget
function apply_local_right_transfer_matrix(
  Lstart::ITensor,
  m::Int64,
  H::InfiniteMPOMatrix,
  ψ::InfiniteCanonicalMPS,
  n_1::Int64;
  reset=true,
)
  dₕ = size(H[n_1])[1]
  ψ′ = dag(prime(ψ.AR[n_1]))
  Ltarget = Vector{ITensor}(undef, dₕ)
  for j in m:dₕ
    if !isempty(H[n_1][j, m])
      Ltarget[j] = Lstart * ψ.AR[n_1] * H[n_1][j, m] * ψ′ #TODO optimize
    end
  end
  return Ltarget
end

#apply the right transfer matrix n1:n1+nsites(ψ)-1
function apply_right_transfer_matrix(
  Lstart::ITensor, m::Int64, H::InfiniteMPOMatrix, ψ::InfiniteCanonicalMPS, n_1::Int64
)
  Ltarget = apply_local_right_transfer_matrix(Lstart, m, H, ψ, n_1)
  for j in 1:(nsites(ψ) - 1)
    Ltarget = apply_local_right_transfer_matrix(Ltarget, H, ψ, n_1 - j)
  end
  return Ltarget
end

function right_environment(H::InfiniteMPOMatrix, ψ::InfiniteCanonicalMPS; tol=1e-10)
  N = nsites(H)
  @assert N == nsites(ψ)

  l = linkinds(only, ψ.AL)
  r = linkinds(only, ψ.AR)
  s = siteinds(only, ψ)
  δʳ(n) = δ(dag(r[n]), prime(r[n]))
  δˡ(n) = δ(l[n], dag(prime(l[n])))
  δˢ(n) = δ(dag(s[n]), prime(s[n]))

  eᵣ = [0.0]
  dₕ = size(H[1])[1]
  Rs = [Vector{ITensor}(undef, dₕ) for j in 1:N]
  #Building the L vector for n_1 = 1
  # TM is 2-N 3-N ... 0
  Rs[1][1] = δʳ(0)
  localL = ψ.C[0] * δˡ(0) * dag(prime(ψ.C[0]))
  for n in 2:dₕ
    temp_Rs = apply_right_transfer_matrix(
      translatecell(translator(ψ), Rs[1][n - 1], 1), n - 1, H, ψ, N
    )
    for j in n:dₕ
      if isassigned(temp_Rs, j)
        if isassigned(Rs[1], j)
          Rs[1][j] += temp_Rs[j]
        else
          Rs[1][j] = temp_Rs[j]
        end
      end
    end
    if !isempty(H[1][n, n])
      λ = H[1][n, n][1, 1]
      δiag = δˢ(1)
      #@assert norm(H[1][n, n] - λ * δiag) == 0 "Non identity diagonal not implemented in MPO"
      @assert abs(λ) <= 1 "Diverging term"
      eᵣ[1] = (localL * Rs[1][n])[]
      Rs[1][n] += -(eᵣ[1] * denseblocks(δʳ(0)))
      if λ == 1
        A = AOᴿ(ψ, H, n)
        Rs[1][n], info = linsolve(A, Rs[1][n], 1, -1; tol=tol)
      else
        println("Not yet implemented")
        flush(stdout)
        flush(stderr)
      end
    end
  end
  if N > 1
    Rs[N] = apply_local_right_transfer_matrix(
      translatecell(translator(ψ), Rs[1], 1), H, ψ, N
    )
    for n in reverse(2:(N - 1))
      Rs[n] = apply_local_right_transfer_matrix(Rs[n + 1], H, ψ, n)
    end
  end
  return CelledVector(Rs), eᵣ[1]
end

function vumps(H::InfiniteMPOMatrix, ψ::InfiniteMPS; kwargs...)
  return vumps(H, orthogonalize(ψ, :); kwargs...)
end

struct H⁰
  L::Vector{ITensor}
  R::Vector{ITensor}
end

function (H::H⁰)(x)
  L = H.L
  R = H.R
  dₕ = length(L)
  result = L[1] * x * R[1]
  for j in 2:dₕ
    result += L[j] * x * R[j]
  end
  return noprime(result)
end

struct H¹
  L::Vector{ITensor}
  R::Vector{ITensor}
  T::Matrix{ITensor}
end

function (H::H¹)(x)
  L = H.L
  R = H.R
  T = H.T
  dₕ = length(L)
  result = ITensor(prime(inds(x)))
  for i in 1:dₕ
    for j in 1:dₕ
      if !isempty(T[i, j])
        result += L[i] * x * T[i, j] * R[j]
      end
    end
  end
  return noprime(result)
end

function tdvp_iteration_sequential(
  solver::Function,
  H::InfiniteMPOMatrix,
  ψ::InfiniteCanonicalMPS;
  (ϵᴸ!)=fill(1e-15, nsites(ψ)),
  (ϵᴿ!)=fill(1e-15, nsites(ψ)),
  time_step,
  solver_tol=(x -> x / 100),
  eager=true,
)
  ψ = copy(ψ)
  ϵᵖʳᵉˢ = max(maximum(ϵᴸ!), maximum(ϵᴿ!))
  _solver_tol = solver_tol(ϵᵖʳᵉˢ)
  N = nsites(ψ)

  C̃ = InfiniteMPS(Vector{ITensor}(undef, N))
  Ãᶜ = InfiniteMPS(Vector{ITensor}(undef, N))
  Ãᴸ = InfiniteMPS(Vector{ITensor}(undef, N))
  Ãᴿ = InfiniteMPS(Vector{ITensor}(undef, N))

  eL = zeros(N)
  eR = zeros(N)
  for n in 1:N
    L, eL[n] = left_environment(H, ψ; tol=_solver_tol) #TODO currently computing two many of them
    R, eR[n] = right_environment(H, ψ; tol=_solver_tol) #TODO currently computing two many of them
    if N == 1
      # 0-site effective Hamiltonian
      E0, C̃[n], info0 = solver(H⁰(L[1], R[2]), time_step, ψ.C[1], _solver_tol, eager)
      # 1-site effective Hamiltonian
      E1, Ãᶜ[n], info1 = solver(
        H¹(L[0], R[2], H[1]), time_step, ψ.AL[1] * ψ.C[1], _solver_tol, eager
      )
      Ãᴸ[1] = ortho_polar(Ãᶜ[1], C̃[1])
      Ãᴿ[1] = ortho_polar(Ãᶜ[1], C̃[0])
      ψ.AL[1] = Ãᴸ[1]
      ψ.AR[1] = Ãᴿ[1]
      ψ.C[1] = C̃[1]
    else
      # 0-site effective Hamiltonian
      E0, C̃[n], info0 = solver(H⁰(L[n], R[n + 1]), time_step, ψ.C[n], _solver_tol, eager)
      E0′, C̃[n - 1], info0′ = solver(
        H⁰(L[n - 1], R[n]), time_step, ψ.C[n - 1], _solver_tol, eager
      )
      # 1-site effective Hamiltonian
      E1, Ãᶜ[n], info1 = solver(
        H¹(L[n - 1], R[n + 1], H[n]), time_step, ψ.AL[n] * ψ.C[n], _solver_tol, eager
      )
      Ãᴸ[n] = ortho_polar(Ãᶜ[n], C̃[n])
      Ãᴿ[n] = ortho_polar(Ãᶜ[n], C̃[n - 1])
      ψ.AL[n] = Ãᴸ[n]
      ψ.AR[n] = Ãᴿ[n]
      ψ.C[n] = C̃[n]
      ψ.C[n - 1] = C̃[n - 1]
    end
  end
  for n in 1:N
    ϵᴸ![n] = norm(Ãᶜ[n] - Ãᴸ[n] * C̃[n])
    ϵᴿ![n] = norm(Ãᶜ[n] - C̃[n - 1] * Ãᴿ[n])
  end
  return ψ, (eL / N, eR / N)
end

function tdvp_iteration_parallel(
  solver::Function,
  H::InfiniteMPOMatrix,
  ψ::InfiniteCanonicalMPS;
  (ϵᴸ!)=fill(1e-15, nsites(ψ)),
  (ϵᴿ!)=fill(1e-15, nsites(ψ)),
  time_step,
  solver_tol=(x -> x / 100),
  eager=true,
)
  ψ = copy(ψ)
  ϵᵖʳᵉˢ = max(maximum(ϵᴸ!), maximum(ϵᴿ!))
  _solver_tol = solver_tol(ϵᵖʳᵉˢ)
  N = nsites(ψ)

  C̃ = InfiniteMPS(Vector{ITensor}(undef, N))
  Ãᶜ = InfiniteMPS(Vector{ITensor}(undef, N))
  Ãᴸ = InfiniteMPS(Vector{ITensor}(undef, N))
  Ãᴿ = InfiniteMPS(Vector{ITensor}(undef, N))

  eL = zeros(1)
  eR = zeros(1)
  L, eL[1] = left_environment(H, ψ; tol=_solver_tol) #TODO currently computing two many of them
  R, eR[1] = right_environment(H, ψ; tol=_solver_tol) #TODO currently computing two many of them
  for n in 1:N
    if N == 1
      # 0-site effective Hamiltonian
      E0, C̃[n], info0 = solver(H⁰(L[1], R[2]), time_step, ψ.C[1], _solver_tol, eager)
      # 1-site effective Hamiltonian
      E1, Ãᶜ[n], info1 = solver(
        H¹(L[0], R[2], H[1]), time_step, ψ.AL[1] * ψ.C[1], _solver_tol, eager
      )
      Ãᴸ[1] = ortho_polar(Ãᶜ[1], C̃[1])
      Ãᴿ[1] = ortho_polar(Ãᶜ[1], C̃[0])
      ψ.AL[1] = Ãᴸ[1]
      ψ.AR[1] = Ãᴿ[1]
      ψ.C[1] = C̃[1]
    else
      # 0-site effective Hamiltonian
      for n in 1:N
        E0, C̃[n], info0 = solver(H⁰(L[n], R[n + 1]), time_step, ψ.C[n], _solver_tol, eager)
        E1, Ãᶜ[n], info1 = solver(
          H¹(L[n - 1], R[n + 1], H[n]), time_step, ψ.AL[n] * ψ.C[n], _solver_tol, eager
        )
      end
      # 1-site effective Hamiltonian
      for n in 1:N
        Ãᴸ[n] = ortho_polar(Ãᶜ[n], C̃[n])
        Ãᴿ[n] = ortho_polar(Ãᶜ[n], C̃[n - 1])
        ψ.AL[n] = Ãᴸ[n]
        ψ.AR[n] = Ãᴿ[n]
        ψ.C[n] = C̃[n]
      end
    end
  end
  for n in 1:N
    ϵᴸ![n] = norm(Ãᶜ[n] - Ãᴸ[n] * C̃[n])
    ϵᴿ![n] = norm(Ãᶜ[n] - C̃[n - 1] * Ãᴿ[n])
  end
  return ψ, (eL / N, eR / N)
end


text:
using GR
y=cumsum(randn(100_000))
p=plot(y)
display(p)


text:
"""
    base(v::AbstractVertex)

Return base vertex
"""
function base end

"""
    OutputsVertex

Decorates an AbstractVertex with output edges.
"""
struct OutputsVertex <: AbstractVertex
    base::AbstractVertex
    outs::AbstractVector{AbstractVertex}
end
OutputsVertex(v::AbstractVertex) = OutputsVertex(v, AbstractVertex[])
init!(v::OutputsVertex, p::AbstractVertex) = foreach(in -> push!(outputs(in), p), inputs(v))
clone(v::OutputsVertex, ins::AbstractVertex...;cf=clone) = OutputsVertex(cf(base(v), ins...,cf=cf))

base(v::OutputsVertex) = v.base
(v::OutputsVertex)(x...) = base(v)(x...)

inputs(v::OutputsVertex) = inputs(base(v))
outputs(v::OutputsVertex) = v.outs

"""
    InputSizeVertex

Vertex with an (immutable) size.

Intended use is for wrapping an InputVertex in conjuntion with mutation
"""
struct InputSizeVertex <: AbstractVertex
    base::AbstractVertex
    size::Integer

    function InputSizeVertex(b::OutputsVertex, size::Integer)
        this = new(b, size)
        init!(b, this)
        return this
    end
end
InputSizeVertex(name, size::Integer) = InputSizeVertex(InputVertex(name), size)
InputSizeVertex(b::AbstractVertex, size::Integer) = InputSizeVertex(OutputsVertex(b), size)
clone(v::InputSizeVertex, ins::AbstractVertex...;cf=clone) = InputSizeVertex(cf(base(v), ins...,cf=cf), cf(v.size,cf=cf))

base(v::InputSizeVertex)::AbstractVertex = v.base
(v::InputSizeVertex)(x...) = base(v)(x...)

inputs(v::InputSizeVertex) = inputs(base(v))
outputs(v::InputSizeVertex) = outputs(base(v))

"""
    MutationTrait

Base type for traits relevant when mutating.
"""
abstract type MutationTrait end
# For convenience as 99% of all traits are immutable. Don't forget to implement for stateful traits or else there will be pain!
clone(t::MutationTrait, cf=nothing) = t

"""
    MutationSizeTrait
Base type for mutation traits relevant to size
"""
abstract type MutationSizeTrait <: MutationTrait end

"""
    SizeTransparent
Base type for mutation traits which are transparent w.r.t size, i.e size changes propagate both forwards and backwards.
"""
abstract type SizeTransparent <: MutationSizeTrait end
"""
    SizeStack
Transparent size trait type where inputs are stacked, i.e output size is the sum of all input sizes
"""
struct SizeStack <: SizeTransparent end
"""
    SizeInvariant
Transparent size trait type where all input sizes must be equal to the output size, e.g. elementwise operations (including broadcasted).
"""
struct SizeInvariant <: SizeTransparent end
"""
    SizeAbsorb
Size trait type for which size changes are absorbed, i.e they do not propagate forward.

Note that size changes do propagate backward as changing the input size of a vertex requires that the output size of its input is also changed and vice versa.
"""
struct SizeAbsorb <: MutationSizeTrait end
"""
    Immutable

Trait for vertices which are immutable. Typically inputs and outputs as those are fixed to the surroundings (e.g a data set).
"""
struct Immutable <: MutationTrait end
trait(v::AbstractVertex) = Immutable()

"""
    DecoratingTrait <: MutationTrait

Avbstract trait which wraps another trait. The wrapped trait of a `DecoratingTrait t` is accessible through `base(t)`.
"""
abstract type DecoratingTrait <: MutationTrait end

struct NamedTrait <: DecoratingTrait
    base::MutationTrait
    name
end
base(t::NamedTrait) = t.base
clone(t::NamedTrait; cf=clone) = NamedTrait(cf(base(t), cf=cf), cf(t.name, cf=cf))

struct SizeChangeLogger <: DecoratingTrait
    level::LogLevel
    infostr::InfoStr
    base::MutationTrait
end
SizeChangeLogger(base::MutationTrait) = SizeChangeLogger(FullInfoStr(), base)
SizeChangeLogger(infostr::InfoStr, base::MutationTrait) = SizeChangeLogger(Logging.Info, infostr, base)
base(t::SizeChangeLogger) = t.base
infostr(t::SizeChangeLogger, v::AbstractVertex) = infostr(t.infostr, v)
clone(t::SizeChangeLogger; cf=clone) = SizeChangeLogger(cf(t.level, cf=cf), cf(t.infostr,cf=cf), cf(base(t), cf=cf))

struct SizeChangeValidation <: DecoratingTrait
    base::MutationTrait
end
base(t::SizeChangeValidation) = t.base
clone(t::SizeChangeValidation; cf=clone) = SizeChangeValidation(cf(base(t), cf=cf))

"""
    MutationVertex

Vertex which may be subject to mutation.

Scope is mutations which affect the input and output sizes as such changes needs to be propagated to the neighbouring vertices.

The member op describes the type of mutation, e.g if individual inputs/outputs are to be pruned vs just changing the size without selecting any particular inputs/outputs.

The member trait describes the nature of the vertex itself, for example if size changes
are absorbed (e.g changing an nin x nout matrix to an nin - Δ x nout matrix) or if they
propagate to neighbouring vertices (and if so, how).
"""
struct MutationVertex <: AbstractVertex
    base::AbstractVertex
    op::MutationOp
    trait::MutationTrait

    function MutationVertex(b::OutputsVertex, s::MutationOp, t::MutationTrait)
        this = new(b, s, t)
        init!(b, this)
        return this
    end
end
MutationVertex(b::AbstractVertex, s::MutationOp, t::MutationTrait) = MutationVertex(OutputsVertex(b), s, t)

clone(v::MutationVertex, ins::AbstractVertex...; cf=clone) = MutationVertex(
cf(base(v), ins...,cf=cf),
cf(v.op, cf=cf),
cf(v.trait, cf=cf))

base(v::MutationVertex) = v.base
op(v::MutationVertex) = v.op
trait(v::MutationVertex) = v.trait
(v::MutationVertex)(x...) = base(v)(x...)

inputs(v::MutationVertex)  = inputs(base(v))
outputs(v::MutationVertex) = outputs(base(v))

# Stuff for displaying information about vertices

show_less(io::IO, v::InputSizeVertex) = show_less(io, base(v))
show_less(io::IO, v::MutationVertex) = print(io, name(v))
show_less(io::IO, v::OutputsVertex) = show_less(io, base(v))

function show(io::IO, v::OutputsVertex)
     show(io, base(v))
     print(io, ", outputs=")
     show(io, outputs(v))
 end

# Stuff for logging

name(v::InputSizeVertex) = name(base(v))
name(v::OutputsVertex) = name(base(v))
name(v::MutationVertex) = name(trait(v), v)
name(t::MutationTrait, v) = summary(v) * "::" * summary(t)
name(t::NamedTrait, v) = t.name
name(t::DecoratingTrait, v) = name(base(t), v)

struct MutationTraitInfoStr <: InfoStr  end
struct MutationSizeTraitInfoStr <: InfoStr  end
struct NinInfoStr <: InfoStr  end
struct NoutInfoStr <: InfoStr  end
SizeInfoStr() = ComposedInfoStr(PrefixedInfoStr("nin=", BracketInfoStr(NinInfoStr())), PrefixedInfoStr("nout=", BracketInfoStr(NoutInfoStr())))

struct OutputsInfoStr <: InfoStr
    infostr::InfoStr
end
OutputsInfoStr() =BracketInfoStr(OutputsInfoStr(NameInfoStr()))

NameAndIOInfoStr() = push!(NameAndInputsInfoStr(), PrefixedInfoStr("outputs=", OutputsInfoStr()))

FullInfoStr() = push!(NameAndIOInfoStr(), SizeInfoStr(), MutationSizeTraitInfoStr())

infostr(::NinInfoStr, v::AbstractVertex) = "unknown"
infostr(::NoutInfoStr, v::AbstractVertex) = "unknown"
infostr(::OutputsInfoStr, v::AbstractVertex) = "unknown"

infostr(i::MutationTraitInfoStr, v::AbstractVertex) = infostr(i, trait(v))
infostr(::MutationTraitInfoStr, t::MutationTrait) = replace(string(t), "\"" => "")
infostr(i::MutationSizeTraitInfoStr, v::AbstractVertex) = infostr(i, trait(v))
infostr(i::MutationSizeTraitInfoStr, t::DecoratingTrait) = infostr(i, base(t))
infostr(::MutationSizeTraitInfoStr, t::MutationSizeTrait) = string(t)
infostr(::MutationSizeTraitInfoStr, t::Immutable) = string(t)
infostr(::NinInfoStr, v::MutationVertex) = join(string.(nin(v)), ", ")
infostr(::NinInfoStr, v::InputSizeVertex) = "N/A"
infostr(::NoutInfoStr, v::MutationVertex) = string(nout(v))
infostr(::NoutInfoStr, v::InputSizeVertex) = string(nout(v))
infostr(i::OutputsInfoStr, v::InputSizeVertex) = infostr(i, base(v))
infostr(i::OutputsInfoStr, v::MutationVertex) = infostr(i, base(v))
infostr(i::OutputsInfoStr, v::OutputsVertex) = join(infostr.(i.infostr, outputs(v)), ", ")