text:
function getTableauVSPARKMidpointProjection(name, q::CoefficientsRK{T}, p::CoefficientsRK{T}, d=[]; R∞=1) where {T}
    @assert q.s == p.s
    s = q.s

    g = getCoefficientsGLRK(1)
    α̃ = g.a
    β = g.b
    γ = g.c

    α  = zeros(T, s, 1)
    α .= 0.5

    q_ã = zeros(T, 1, s)
    for i in 1:s
        q_ã[1,i] = q.b[i] / β[1] * ( β[1] - α[i,1] )
    end

    p_ã = zeros(T, 1, s)
    for i in 1:s
        p_ã[1,i] = p.b[i] / β[1] * ( β[1] - α[i,1] )
    end

    β  = zero(g.b)
    β .= α̃[1,1] * (1 + R∞)

    d = ones(T, 1)

    ω  = zeros(T, 1, 2)
    ω .= [0.0  1.0]
    δ  = zeros(T, 0, 1)

    return TableauVSPARKprimary(name, min(q.o, p.o),
                        q.a, p.a, α, α,
                        q_ã, p_ã, α̃, α̃,
                        q.b, p.b, β, β,
                        q.c, p.c, γ, d,
                        ω, δ)
end

"Tableau for Gauss-Legendre method with s stages and midpoint projection."
function getTableauVSPARKGLRKpMidpoint(s)
    glrk = getCoefficientsGLRK(s)
    getTableauVSPARKMidpointProjection(Symbol("vsparkglrk", s, "pMidpoint"), glrk, glrk; R∞=(-1)^s)
end

function getTableauVSPARKSymmetricProjection(name, q::CoefficientsRK{T}, p::CoefficientsRK{T}, d=[]; R∞=1) where {T}

    @assert q.s == p.s

    o = min(q.o, p.o)

    a_q = q.a
    a_p = p.a

    α_q = zeros(T, q.s, 2)
    α_q[:,1] .= 0.5

    α_p = zeros(T, p.s, 2)
    α_p[:,1] .= 0.5

    a_q̃ = Array(transpose(hcat(zero(q.b), q.b)))
    a_p̃ = Array(transpose(hcat(zero(p.b), p.b)))

    α_q̃ = [[0.0  0.0]
           [0.5  R∞*0.5]]
    α_p̃ = [[0.0  0.0]
           [0.5  R∞*0.5]]

    b_q = q.b
    b_p = p.b
    β_q = [0.5, R∞*0.5]
    β_p = [0.5, R∞*0.5]

    c_q = q.c
    c_p = p.c
    c_λ = [ 0.0, 1.0]
    d_λ = [ 0.5, 0.5]

    ω_λ  = zeros(T, 1, 3)
    ω_λ .= [0.5 R∞*0.5 0.0]

    δ_λ  = zeros(T, 1, 2)
    δ_λ .= [-1.0 +1.0]

    if length(d) == 0
        return TableauVSPARKprimary(name, o,
                            a_q, a_p, α_q, α_p,
                            a_q̃, a_p̃, α_q̃, α_p̃,
                            b_q, b_p, β_q, β_p,
                            c_q, c_p, c_λ, d_λ,
                            ω_λ, δ_λ)
    else
        @assert length(d) == q.s == p.s

        return TableauVSPARKprimary(name, o,
                            a_q, a_p, α_q, α_p,
                            a_q̃, a_p̃, α_q̃, α_p̃,
                            b_q, b_p, β_q, β_p,
                            c_q, c_p, c_λ, d_λ,
                            ω_λ, δ_λ, d)
    end

end

"Tableau for Gauss-Lobatto IIIA-IIIB method with two stages and symmetric projection."
function getTableauVSPARKLobIIIAIIIB2pSymmetric()
    d = [+1.0, -1.0]
    getTableauVSPARKSymmetricProjection(:LobIIIAIIIB2pSymmetric, getCoefficientsLobIIIA2(), getCoefficientsLobIIIB2(), d; R∞=-1)
end

"Tableau for Gauss-Lobatto IIIA-IIIB method with three stages and symmetric projection."
function getTableauVSPARKLobIIIAIIIB3pSymmetric()
    d = [+0.5, -1.0, +0.5]
    getTableauVSPARKSymmetricProjection(:LobIIIAIIIB3pSymmetric, getCoefficientsLobIIIA3(), getCoefficientsLobIIIB3(), d; R∞=+1)
end

"Tableau for Gauss-Legendre method with s stages and symplectic projection."
function getTableauVSPARKGLRKpSymmetric(s)
    glrk = getCoefficientsGLRK(s)
    getTableauVSPARKSymmetricProjection(Symbol("vpglrk", s, "pSymmetric"), glrk, glrk; R∞=(-1)^s)
end

function getTableauVSPARKSymplecticProjection(name, q::CoefficientsRK{T}, p::CoefficientsRK{T}, d=[]; R∞=+1) where {T}

    la = getCoefficientsLobIIIA2()
    lb = getCoefficientsLobIIIB2()

    @assert q.s == p.s
    @assert q.b == p.b
    @assert q.c == p.c

    @assert la.s == lb.s
    @assert la.b == lb.b
#    @assert la.c == lb.c

    s = q.s
    σ = la.s
    ρ = 0

    a_q = q.a
    a_p = p.a
    b_q = q.b
    b_p = p.b

    # β_q = la.b
    # β_p = lb.b
   β_q = [0.5, R∞*0.5]
   β_p = [0.5, R∞*0.5]

    α_q = zeros(T, s, 2)
    α_q[:,1] .= 0.5

    α_p = zeros(T, s, 2)
    α_p[:,1] .= 0.5

#    α_q = zeros(T, s, σ)
#    α_p = zeros(T, s, σ)
#    for i in 1:s
#        for j in 1:σ
#            α_q[i,j] = #q.b[i] / β[1] * ( β[1] - α[i,1] )
#            α_p[i,j] = #p.b[i] / β[1] * ( β[1] - α[i,1] )
#        end
#    end

    a_q̃ = zeros(T, σ, s)
    a_p̃ = zeros(T, σ, s)
    for i in 1:σ
        for j in 1:s
            a_q̃[i,j] = b_q[j] / β_p[i] * (β_p[i] - α_p[j,i])
            a_p̃[i,j] = b_p[j] / β_q[i] * (β_q[i] - α_q[j,i])
        end
    end

    α_q̃ = la.a
    α_p̃ = lb.a

    c_q = q.c
    c_p = p.c
    c_λ = la.c
    d_λ = [ 0.5, 0.5]

    ω_λ = [0.5 0.5 0.0
            0.0 0.0 1.0]

    δ_λ = zeros(T, ρ, σ)

    if length(d) == 0
        return TableauVSPARKprimary(name, min(q.o, p.o),
                            a_q, a_p, α_q, α_p,
                            a_q̃, a_p̃, α_q̃, α_p̃,
                            b_q, b_p, β_q, β_p,
                            c_q, c_p, c_λ, d_λ,
                            ω_λ, δ_λ)
    else
        @assert length(d) == s

        return TableauVSPARKprimary(name, o,
                            a_q, a_p, α_q, α_p,
                            a_q̃, a_p̃, α_q̃, α_p̃,
                            b_q, b_p, β_q, β_p,
                            c_q, c_p, c_λ, d_λ,
                            ω_λ, δ_λ, d)
    end

end

function getTableauVSPARKGLRKpSymplectic(s)
    glrk = getCoefficientsGLRK(s)
    getTableauVSPARKSymplecticProjection(Symbol("VSPARK", s, "pSymplectic"), glrk, glrk; R∞=(-1)^s)
end


text:
using LowRankModels, DataFrames, Random, StatsBase

Random.seed!(0)

println("censored data example")
# boolean example with only entries greater than threshold t observed
# ie, censored data
# example with only entries greater than threshold t observed
m,n,k,ktrue = 100,100,1,1
A = rand(m,ktrue)*rand(ktrue,n)
B = round.(Int, ktrue*rand(m,n) .>= A) # Bernoulli samples with probability proportional to A
losses = fill(QuadLoss(),n)
r = QuadReg(.1)
obs = Array{Tuple{Int,Int}}(undef,0)
for i=1:m
    for j=1:n
        if B[i,j] == 1
            push!(obs, (i,j))
        end
    end
end
(train_observed_features, train_observed_examples,
    test_observed_features,  test_observed_examples) =
        get_train_and_test(obs, m, n, .2)

train_glrm = GLRM(B,losses,r,r,k, observed_features=train_observed_features, observed_examples=train_observed_examples,)
test_glrm = GLRM(B,losses,r,r,k, observed_features=test_observed_features, observed_examples=test_observed_examples)

function censored_regularization_path(train_glrm::GLRM, test_glrm::GLRM; params=Params(), reg_params=exp10.(range(2,stop=-2,length=5)),
                                         holdout_proportion=.1, verbose=true,
                                         ch::ConvergenceHistory=ConvergenceHistory("reg_path"))
    m,n = size(train_glrm.A)
    ntrain = sum(map(length, train_glrm.observed_features))
    ntest = sum(map(length, test_glrm.observed_features))
    train_error = Array{Float64}(undef,length(reg_params))
    test_error = Array{Float64}(undef,length(reg_params))
    solution = Array{Tuple{Float64,Float64}}(undef, length(reg_params))
    train_time = Array{Float64}(undef, length(reg_params))
    for iparam=1:length(reg_params)
        reg_param = reg_params[iparam]
        # evaluate train and test error
        if verbose println("fitting train GLRM for reg_param $reg_param") end
        mul!(train_glrm.rx, reg_param)
        mul!(train_glrm.ry, reg_param)
        train_glrm.X, train_glrm.Y = randn(train_glrm.k,m), randn(train_glrm.k,n)
        X, Y, ch = fit!(train_glrm; params=params, ch=ch, verbose=verbose)
        train_time[iparam] = ch.times[end]
        if verbose println("computing train and test error for reg_param $reg_param:") end
        train_error[iparam] = objective(train_glrm, X, Y, include_regularization=false) / ntrain
        if verbose println("\ttrain error: $(train_error[iparam])") end
        test_error[iparam] = objective(test_glrm, X, Y, include_regularization=false) / ntest
        if verbose println("\ttest error:  $(test_error[iparam])") end
        solution[iparam] = (sum(X)+sum(Y), sum(abs.(X))+sum(abs.(Y)))
        if verbose println("\tsum of solution, one norm of solution:  $(solution[iparam])") end
    end
    return train_error, test_error, train_time, reg_params, solution
end

(train_error, test_error, train_time, reg_params, solution) =
    censored_regularization_path(train_glrm, test_glrm,
                                 params=ProxGradParams(1, max_iter=50, abs_tol=.001, min_stepsize=.1),
                                 reg_params=exp10.(range(2, stop=-2, length=3)) )

df = DataFrame(train_error = train_error, test_error = test_error,
                   train_time = train_time, reg_param = reg_params, solution_1norm = [s[2] for s in solution])
println(df)

println(solution)

ext: jl


text:
using MAT
using ReachabilityAnalysis
using SparseArrays
using LinearAlgebra
using StructuralDynamicsODESolvers
using LazySets.Arrays

using Plots
using LaTeXStrings
using Plots.PlotMeasures

# utils for plots
__toL(x, digits) = L"%$(round(x, digits=digits))"
xticklatex(vec, digits) = (vec, __toL.(vec, Ref(digits)))

LazySets.set_rtol(Float64, eps(Float64))
LazySets.set_ztol(Float64, eps(Float64))
LazySets.set_atol(Float64, eps(Float64))

include("../krylov_concept.jl")

include("lambs_modelo.jl")

# cases:
#    - "constant" (Heaviside)
#    - "rickler" (step-load approximation of the rickler wavelet)
#
function sol_numerica_lamb(p::FEMProblem; case="heaviside", δt=nothing, δtfac=1.0,
                           alg=Trapezoidal, X0=nothing, kwargs...)

    if isnothing(δt)
        δt = step_size_lamb(p, δtfac)
    end

    Tmax = 1.0 # seconds
    NSTEPS = round(Int, Tmax/δt)
    tsp = range(0.0, Tmax, length=NSTEPS+1)

    m = size(p.s.K, 1)
    B = Matrix(1.0I, m, m)
    X = Universe(m)

    H = Heaviside
    if case == "constant" || case == "heaviside"
        Fval = -1e6
        Fval = [p.b * Fval for ti in tsp]

    elseif case == "ricker"
        ancho_escalon = 0.05 # 50 ms de ancho
        # por cuestion de simetria se usa 2e6 / 2 respecto del articulo de Accurate Wave Prop
        fr(t) = -1e6 * (H(ancho_escalon*3 - t) - 3*H(ancho_escalon*2 - t) + 3*H(ancho_escalon - t))
        Fval = [p.b * fr(ti) for ti in tsp] # "step-loading" en el articulo

    else
        throw(ArgumentError("case $case not defined"))
    end
    sys = SecondOrderConstrainedLinearControlContinuousSystem(sparse(p.s.M), p.s.C, p.s.K, B, Universe(m), Fval)

    if isnothing(X0)
        X0 = (zeros(m), zeros(m))
    end
    prob = @ivp(sys, u(0) ∈ X0)
    alg = alg(Δt=δt)
    sol = solve(prob, alg, NSTEPS=NSTEPS)

    # al final no invertimos
    # invertimos el signo de todos los desplazamientos para comparar con el articulo Accurate Wave prop
    #sol.U .= -1 .* sol.U
    # lo mismo con las velocidades
    #sol.U′ .= -1 .* sol.U′

    return sol
end

function sol_numerica_lamb_statistics(p::FEMProblem; case="ricker", δt=nothing, δtfac=1.0,
                                      alg=Trapezoidal, X0=nothing,
                                      sampler, idx, vecsamples,
                                      kwargs...)
    if isnothing(δt)
        δt = step_size_lamb(p, δtfac)
    end

    Tmax = 1.0 # seconds
    NSTEPS = round(Int, Tmax/δt)
    tsp = range(0.0, Tmax, length=NSTEPS+1)

    m = size(p.s.K, 1)
    B = Matrix(1.0I, m, m)
    X = Universe(m)

    H = Heaviside
    if case == "constant" || case == "heaviside"
        Fval = -1e6
        Fval = [p.b * Fval for ti in tsp]

    elseif case == "ricker"
        ancho_escalon = 0.05 # 50 ms de ancho
        # por cuestion de simetria se usa 2e6 / 2 respecto del articulo de Accurate Wave Prop
        fr(t) = -1e6 * (H(ancho_escalon*3 - t) - 3*H(ancho_escalon*2 - t) + 3*H(ancho_escalon - t))
        Fval = [p.F * fr(ti) for ti in tsp] # "step-loading" en el articulo

    else
        throw(ArgumentError("case $case not defined"))
    end
    sys = SecondOrderConstrainedLinearControlContinuousSystem(sparse(p.s.M), p.s.C, p.s.K, B, Universe(m), Fval)

    if isnothing(X0)
        X0 = (zeros(m), zeros(m))
    end
    prob = @ivp(sys, u(0) ∈ X0)
    alg = alg(Δt=δt)

    sol_statistics = StructuralDynamicsODESolvers._solve_statistics(alg, prob, NSTEPS, sampler, idx, vecsamples)

    # invertimos el signo de todos los desplazamientos para comparar con el articulo Accurate Wave prop
    #sol.U .= -1 .* sol.U

    return sol_statistics
end

function plot_nodo(sol; nodeidx=3, direction="Horizontal", method="Newmark")

    fig = plot(xlab=L"\textrm{Time (s)}",
               ylab=L"\textrm{%$direction displacement} (\times 10^{-5} \textrm{m})",
               legend=:topright,
               legendfontsize=25,
               tickfont=font(20, "Times"),
               guidefontsize=10,
               xguidefont=font(20, "Times"),
               yguidefont=font(20, "Times"),
               xtick = xticklatex([0, 0.2, 0.4, 0.6, 0.8, 1.0], 2),
               ytick = ([-4e-5, -2e-5, 0, 2e-5, 4e-5, 6e-5], [L"-4", L"-2", L"0", L"2", L"4", L"6"]),
               xlims=(0.0, 1.0),
               ylims=(-5e-5, 8e-5),
               bottom_margin=10mm, left_margin=6mm, right_margin=12mm, top_margin=3mm, size=(900, 600))

    plot!(fig, sol, vars=(0, nodeidx), seriestype=:path, markershape=:none, lw=2.0, lab=L"\textrm{%$method}")
    fig
end

function plot_lambs(p::FEMProblem; nodeidx=3, direction="Horizontal", case="heaviside", kwargs...)

    fig = plot(xlab=L"\textrm{Time (s)}",
               ylab=L"\textrm{%$direction displacement} (\times 10^{-5} \textrm{m})",
               legend=:topright,
               legendfontsize=25,
               tickfont=font(20, "Times"),
               guidefontsize=10,
               xguidefont=font(20, "Times"),
               yguidefont=font(20, "Times"),
               xtick = xticklatex([0, 0.2, 0.4, 0.6, 0.8, 1.0], 2),
               ytick = ([-4e-5, -2e-5, 0, 2e-5, 4e-5, 6e-5], [L"-4", L"-2", L"0", L"2", L"4", L"6"]),
               xlims=(0.0, 1.0),
               ylims=(-5e-5, 8e-5),
               bottom_margin=10mm, left_margin=6mm, right_margin=12mm, top_margin=3mm, size=(900, 600))

    # resolver numericas
    sol_newmark = sol_numerica_lamb(p; nodeidx=nodeidx, case=case, alg=Trapezoidal, kwargs...)
    sol_bathe = sol_numerica_lamb(p; nodeidx=nodeidx, case=case, alg=Bathe, kwargs...)

    # resolver flowpipe
    out = _sol_reachability_lamb(p; case=case, nodeidx=nodeidx, kwargs...)

    plot!(fig, out, vars=(0, nodeidx), lc=:blue, c=:blue, lab=L"\textrm{Flowpipe}")
    plot!(fig, sol_newmark, vars=(0, nodeidx), seriestype=:path, markershape=:none, lw=2.0, lab=L"\textrm{Newmark}")
    plot!(fig, sol_bathe, vars=(0, nodeidx), seriestype=:path, markershape=:none, lw=2.0, lab=L"\textrm{Bathe}")
    fig
end

function generar_figura_malla2_singleton_position()

    # =================
    # setear parametros
    # =================
    p = load_lambs(mesh=2);
    @show m = size(p.s.M, 1)
    nodeidx = 3
    δtfac = 1.0
    δt = step_size_lamb(p, δtfac) # 0.001
    case = "heaviside"
    condicion_inicial = "singleton"
    mk = 140
    tol = 1e-7

    # ====================
    # resolver numericas
    # ====================
    println("Resolviendo con Newmark... ")
    @time sol_newmark = sol_numerica_lamb(p; nodeidx=nodeidx, case=case, alg=Trapezoidal, δt=δt, condicion_inicial=condicion_inicial)

    println("Resolviendo con Bathe... ")
    @time sol_bathe = sol_numerica_lamb(p; nodeidx=nodeidx, case=case, alg=Bathe, δt=δt, condicion_inicial=condicion_inicial)

    # =================
    # resolver flowpipe
    # =================
    println("Resolviendo Flowpipe... ")
    @time out = _sol_reachability_lamb(p; case=case, nodeidx=nodeidx, δt=δt, condicion_inicial=condicion_inicial, mk=mk, tol=tol)

    # =================
    # graficar
    # =================
    fig = plot(xlab=L"\textrm{Time (s)}",
               ylab=L"\textrm{Displacement} (\times 10^{-5} \textrm{m})",
               legend=:bottomright,
               legendfontsize=15,
               tickfont=font(25, "Times"),
               guidefontsize=15,
               xguidefont=font(25, "Times"),
               yguidefont=font(25, "Times"),
               xtick = xticklatex([0, 0.2, 0.4, 0.6, 0.8, 1.0], 2),
               ytick = ([-1e-5, 0, 1e-5, 2e-5, 3e-5, 4e-5], [L"-1", L"0", L"1", L"2", L"3", L"4"]),
               xlims=(0.0, 1.0),
               ylims=(-1.2e-5, 4.2e-5),
               bottom_margin=10mm, left_margin=6mm, right_margin=12mm, top_margin=3mm, size=(900, 600))

    plot!(fig, out, vars=(0, nodeidx), c=:lightblue, lc=:black, lw=0.5, alpha=1., lab=L"\textrm{Flowpipe}")

    # grafico continuo
    plot!(fig, sol_newmark, vars=(0, nodeidx), seriestype=:path, marker=:none, lw=2.0, color=:red, lab="")
    plot!(fig, sol_bathe, vars=(0, nodeidx), seriestype=:path, marker=:none, lw=2.0, color=:green, lab="")

    # con markers salteados
    paso = 15
    npoints = length(sol_newmark.U)
    plot!(fig, sol_newmark.t[1:paso:npoints], [sol_newmark.U[i][nodeidx] for i in 1:paso:npoints], vars=(0, nodeidx), linetype=:scatter, marker=:circle, markersize=8, lw=2.0, color=:red, lab=L"\textrm{Newmark}")
    plot!(fig, sol_bathe.t[1:paso:npoints], [sol_bathe.U[i][nodeidx] for i in 1:paso:npoints], vars=(0, nodeidx), linetype=:scatter, marker=:utriangle, markersize=8, lw=2.0, color=:green, lab=L"\textrm{Bathe}")

    # grafico bounding box donde hago zoom
    tzoom_left = 0.40
    tzoom_right = 0.50
    yzoom_bottom = 2e-5
    yzoom_top = 4e-5

    PBotLeft  = [tzoom_left, yzoom_bottom]
    PBotRight = [tzoom_right, yzoom_bottom]
    PTopLeft  = [tzoom_left, yzoom_top]
    PTopRight = [tzoom_right, yzoom_top]

    L1 = LineSegment(PBotLeft, PBotRight)
    L2 = LineSegment(PBotRight, PTopRight)
    L3 = LineSegment(PTopRight, PTopLeft)
    L4 = LineSegment(PTopLeft, PBotLeft)

    plot!(fig, L1, color=:black, style=:dash, lw=3.0, seriestype=:solid, marker=:none, alpha=1.)
    plot!(fig, L2, color=:black, style=:dash, lw=3.0, seriestype=:solid, marker=:none, alpha=1.)
    plot!(fig, L3, color=:black, style=:dash, lw=3.0, seriestype=:solid, marker=:none, alpha=1.)
    plot!(fig, L4, color=:black, style=:dash, lw=3.0, seriestype=:solid, marker=:none, alpha=1.)

    savefig(fig, "malla2_singleton_pos_nodeidx_3.pdf")

    # =============================
    # graficar parte zoomeada
    # =============================

    # como todos los solvers usan el mismo step size, puedo usar Newmark para encontrar
    # los indices donde hago zoom
    a = findfirst(>(0), sol_newmark.t .- tzoom_left)
    b = findfirst(>(0), sol_newmark.t .- tzoom_right)

    fig = plot(xlab=L"\textrm{Time (s)}",
               ylab=L"\textrm{Displacement} (\times 10^{-5} \textrm{m})",
               legend=:bottomright,
               legendfontsize=15,
               tickfont=font(25, "Times"),
               guidefontsize=15,
               xguidefont=font(25, "Times"),
               yguidefont=font(25, "Times"),
               xtick = xticklatex([0.4, 0.425, 0.45, 0.475, 0.5], 3),
               ytick = ([2.5e-5, 2.75e-5, 3.0e-5, 3.25e-5, 3.5e-5], [L"2.5", L"2.75", L"3.0", L"3.25", L"3.5"]),
               xlims=(0.395, 0.505),
               ylims=(2.4e-5, 3.7e-5),
               bottom_margin=10mm, left_margin=6mm, right_margin=12mm, top_margin=3mm, size=(900, 600))

    plot!(fig, out[a:b], vars=(0, nodeidx), c=:lightblue, lc=:black, lw=0.5, alpha=1., lab=L"\textrm{Flowpipe}")

    plot!(fig, sol_newmark.t[a:b], [sol_newmark.U[i][nodeidx] for i in a:b], vars=(0, nodeidx), seriestype=:path, markershape=:none, lw=2.0, color=:red, lab="", alpha=1.)
    plot!(fig, sol_bathe.t[a:b], [sol_bathe.U[i][nodeidx] for i in a:b], seriestype=:path, markershape=:none, lw=2.0, color=:green, lab="", alpha=1.)

    # con markers salteados
    paso = 2
    plot!(fig, sol_newmark.t[a:paso:b], [sol_newmark.U[i][nodeidx] for i in a:paso:b], vars=(0, nodeidx), linetype=:scatter, marker=:circle, markersize=10, lw=2.0, color=:red, lab=L"\textrm{Newmark}")
    plot!(fig, sol_bathe.t[a:paso:b], [sol_bathe.U[i][nodeidx] for i in a:paso:b], vars=(0, nodeidx), linetype=:scatter, marker=:utriangle, markersize=10, lw=2.0, color=:green, lab=L"\textrm{Bathe}")

    savefig(fig, "malla2_singleton_pos_nodeidx_3_zoom.pdf")

    return nothing
end

function generar_figura_malla2_singleton_velocity()

    # =================
    # setear parametros
    # =================
    p = load_lambs(mesh=2);
    @show m = size(p.s.M, 1)
    nodeidx = 3
    δtfac = 1.0
    δt = step_size_lamb(p, δtfac) # 0.001
    case = "heaviside"
    condicion_inicial = "singleton"
    mk = 140
    tol = 1e-7

    # ====================
    # resolver numericas
    # ====================
    println("Resolviendo con Newmark... ")
    @time sol_newmark = sol_numerica_lamb(p; nodeidx=nodeidx, case=case, alg=Trapezoidal, δt=δt, condicion_inicial=condicion_inicial)

    println("Resolviendo con Bathe... ")
    @time sol_bathe = sol_numerica_lamb(p; nodeidx=nodeidx, case=case, alg=Bathe, δt=δt, condicion_inicial=condicion_inicial)

    # =================
    # resolver flowpipe
    # =================
    println("Resolviendo Flowpipe... ")
    @time out = _sol_reachability_lamb(p; case=case, nodeidx=nodeidx+m, δt=δt, condicion_inicial=condicion_inicial, mk=mk, tol=tol)

    # =================
    # graficar
    # =================
    fig = plot(xlab=L"\textrm{Time (s)}",
               ylab=L"\textrm{Velocity} (\times 10^{-4} \textrm{m/s})",
               legend=:topright,
               legendfontsize=15,
               tickfont=font(25, "Times"),
               guidefontsize=15,
               xguidefont=font(25, "Times"),
               yguidefont=font(25, "Times"),
               xtick = xticklatex([0, 0.2, 0.4, 0.6, 0.8, 1.0], 2),
               ytick = ([-1e-4, 0, 1e-4, 2e-4, 3e-4, 4e-4, 5e-4], [L"-1", L"0", L"1", L"2", L"3", L"4", L"5"]),
               xlims=(-0.02, 1.0),
               ylims=(-2.1e-4, 5.5e-4),
               bottom_margin=10mm, left_margin=6mm, right_margin=12mm, top_margin=3mm, size=(900, 600))

    plot!(fig, out, vars=(0, nodeidx+m), c=:lightblue, lc=:black, lab=L"\textrm{Flowpipe}", alpha=1.)

    # grafico continuo
    plot!(fig, sol_newmark.t, [sol_newmark.U′[i][nodeidx] for i in 1:length(sol_newmark.U′)], vars=(0, nodeidx), seriestype=:path, markershape=:none, lw=2.0, lab="")
    plot!(fig, sol_bathe.t, [sol_bathe.U′[i][nodeidx] for i in 1:length(sol_bathe.U′)], seriestype=:path, markershape=:none, lw=2.0, lab="")

    # con markers salteados
    paso = 15
    npoints = length(sol_newmark.U)
    plot!(fig, sol_newmark.t[1:paso:npoints], [sol_newmark.U′[i][nodeidx] for i in 1:paso:npoints], vars=(0, nodeidx), linetype=:scatter, marker=:circle, markersize=10, lw=2.0, c=:red, lab=L"\textrm{Newmark}")
    plot!(fig, sol_bathe.t[1:paso:npoints], [sol_bathe.U′[i][nodeidx] for i in 1:paso:npoints], linetype=:scatter, marker=:utriangle, markersize=10, lw=2.0, c=:green, lab=L"\textrm{Bathe}")

    # grafico bounding box donde hago zoom
    tzoom_left = 0.8
    tzoom_right = 0.99
    yzoom_bottom = -1e-4
    yzoom_top = 1e-4

    PBotLeft  = [tzoom_left, yzoom_bottom]
    PBotRight = [tzoom_right, yzoom_bottom]
    PTopLeft  = [tzoom_left, yzoom_top]
    PTopRight = [tzoom_right, yzoom_top]

    L1 = LineSegment(PBotLeft, PBotRight)
    L2 = LineSegment(PBotRight, PTopRight)
    L3 = LineSegment(PTopRight, PTopLeft)
    L4 = LineSegment(PTopLeft, PBotLeft)

    plot!(fig, L1, color=:black, style=:dash, lw=3.0, seriestype=:solid, marker=:none, alpha=1.)
    plot!(fig, L2, color=:black, style=:dash, lw=3.0, seriestype=:solid, marker=:none, alpha=1.)
    plot!(fig, L3, color=:black, style=:dash, lw=3.0, seriestype=:solid, marker=:none, alpha=1.)
    plot!(fig, L4, color=:black, style=:dash, lw=3.0, seriestype=:solid, marker=:none, alpha=1.)

    savefig(fig, "malla2_singleton_vel_nodeidx_3.pdf")

    # =============================
    # graficar parte zoomeada
    # =============================

    # como todos los solvers usan el mismo step size, puedo usar Newmark para encontrar
    # los indices donde hago zoom
    a = findfirst(>(0), sol_newmark.t .- tzoom_left)
    b = findfirst(>(0), sol_newmark.t .- tzoom_right)

    fig = plot(xlab=L"\textrm{Time (s)}",
               ylab=L"\textrm{Velocity} (\times 10^{-4} \textrm{m/s})",
               legend=:topright,
               legendfontsize=15,
               tickfont=font(25, "Times"),
               guidefontsize=15,
               xguidefont=font(25, "Times"),
               yguidefont=font(25, "Times"),
               xtick = xticklatex([0.8, 0.9, 1.0], 3),
               ytick = ([-1.0e-4, -0.5e-4, 0, 0.5e-4, 1e-4], [L"-1.0", L"-0.5", L"0.0", L"0.5", L"1.0"]),
               xlims=(0.79, 1.01),
               ylims=(-1.05e-4, 1.05e-4),
               bottom_margin=10mm, left_margin=6mm, right_margin=12mm, top_margin=3mm, size=(900, 600))

    plot!(fig, out[a:b-1], vars=(0, nodeidx+m), lc=:black, c=:lightblue, lw=0.5, alpha=1., lab=L"\textrm{Flowpipe}")
    plot!(fig, sol_bathe.t[a:b], [sol_bathe.U′[i][nodeidx] for i in a:b], seriestype=:path, markershape=:none, c=:red, lw=2.0, lab="", alpha=1.)
    plot!(fig, sol_newmark.t[a:b], [sol_newmark.U′[i][nodeidx] for i in a:b], vars=(0, nodeidx), seriestype=:path, markershape=:none, c=:green, lw=2.0, lab="", alpha=1.)

    # con markers salteados
    paso = 2
    plot!(fig, sol_newmark.t[a:paso:b], [sol_newmark.U′[i][nodeidx] for i in a:paso:b], vars=(0, nodeidx), linetype=:scatter, marker=:circle, markersize=10, lw=2.0, color=:red, alpha=1., lab=L"\textrm{Newmark}")
    plot!(fig, sol_bathe.t[a:paso:b], [sol_bathe.U′[i][nodeidx] for i in a:paso:b], vars=(0, nodeidx), linetype=:scatter, marker=:utriangle, markersize=10, lw=2.0, color=:green, alpha=1., lab=L"\textrm{Bathe}")

    savefig(fig, "malla2_singleton_vel_nodeidx_3_zoom.pdf")

    return nothing
end

# ========================================
# Metodo de reachability
# ========================================

function _sol_reachability_lamb(p::FEMProblem; δt=nothing, nodeidx=[3, 4],
                                case="ricker", δtfac=1.0, alg=Trapezoidal,
                                mk=150, tol=1e-12,
                                Eplus_optim=true,
                                kwargs...)

    if case == "heaviside"
        return _sol_reachability_lamb_heaviside(p; δt=δt, nodeidx=nodeidx, δtfac=δtfac, alg=alg, mk=mk, tol=tol, Eplus_optim=Eplus_optim, kwargs...)
    else
        error("ver lambs_krylov_old.jl")
    end
end

function _sol_reachability_lamb_heaviside(p::FEMProblem;
                                          δtfac=1.0,
                                          δt=step_size_lamb(p, δtfac),
                                          nodeidx=[3, 4],
                                          alg=Trapezoidal,
                                          mk=150, tol=1e-12,
                                          Eplus_optim=true,
                                          condicion_inicial="singleton",

                                          # para condicion inicial "distribuida"
                                          radio_desplazamientos = nothing, # 3.5e-7
                                          radio_velocidades = nothing, # 5e-6

                                          kwargs...)

    # ===================
    # Definir la carga
    # ===================

    # @assert δt <= δtnom

    Tmax = 1.0 # seconds
    NSTEPS = round(Int, Tmax/δt)

    H = Heaviside

    ancho_escalon = 0.05 # 50 ms de ancho
    # por cuestion de simetria se usa 2e6 / 2 respecto del articulo de Accurate Wave Prop
    Δ = ancho_escalon
    α = -1e6
    fload(t) = α * H(t)

    # ==============================
    # Homogeneizar
    # ==============================
    m = size(p.s.K, 1)
    Shom = _homogeneize(p)
    Aext = Shom.A
    @assert size(Aext, 1) == 2m+1

    # =====================================
    # Preparacion de la condicion inicial
    # =====================================

    if condicion_inicial == "singleton"
        X0_init = Singleton(zeros(2m+1))
        X0_init.element[end] = α

    elseif condicion_inicial == "distribuida"
        X0_init = Hyperrectangle(zeros(2m+1), zeros(2m+1))
        if isnothing(radio_desplazamientos) || isnothing(radio_velocidades)
            error("debe especificar `radio_desplazamientos` y `radio_velocidades`")
        end
        X0_init.center[end] = α # fuerza
        X0_init.radius[1:m] .= radio_desplazamientos # desplazamientos iniciales
        X0_init.radius[m+1:2m] .= radio_velocidades # velocidades iniciales

    else
        error("condicion_inicial = $condicion_inicial no existe")
    end

    # ==============================
    # Propagacion a tiempo continuo
    # ==============================
    Aᵀδ = copy(transpose(Aext .* δt))
    num_nodes = isa(nodeidx, Int) ? 1 : length(nodeidx)

    # alg = LGG09(δ=δt, n=2m+1, vars=nodeidx)
    dirs = Vector{LazySets.Arrays.SingleEntryVector{Float64}}()
    for j in 1:num_nodes
        node = nodeidx[j]
        push!(dirs, SingleEntryVector(node, 2m+1, 1.0))  # 0, 0, 1, ... , 0, 0, 0, 0, 0, 0
        push!(dirs, SingleEntryVector(node, 2m+1, -1.0))
    end
    cdirs = CustomDirections(dirs)

    println("Calculando el factor de bloating...")

    Eplus = _build_Eplus_optim(Aext, X0_init, δt; m=mk, tol=tol, cutoff=eps(Float64))

    println("Propagando a tiempo continuo...")

    if condicion_inicial == "singleton"
        ΦX0_init = _expv(Aext, X0_init.element, 1, δt, m=mk, tol=tol) |> Singleton
        Ω0 = CH(X0_init, ΦX0_init ⊕ Eplus)

        Mout = Matrix{Float64}(undef, 2*num_nodes, NSTEPS)
        for j in 1:num_nodes
            node = nodeidx[j]
            d = zeros(2m+1)
            d[node] = 1.0
            println("Propagando indice $node ...")
            idx = (2(j-1) + 1):2j
            _reach_homog_krylov_LGG09_modif!(view(Mout, idx, :), Ω0, Aᵀδ, d, NSTEPS; m=mk, tol=tol)
        end
        F = [TemplateReachSet(cdirs, view(Mout, :, k), (0.0 .. δt) + (k-1) * δt) for k in 1:NSTEPS]

    elseif condicion_inicial == "distribuida"
        # Ω0 = CH(X0_init, ΦX0_init * X0_init ⊕ Eplus)

        Mout = Matrix{Float64}(undef, 2*num_nodes, NSTEPS)
        for j in 1:num_nodes
            node = nodeidx[j]

            d = zeros(2m+1)
            d[node] = 1.0

            println("Propagando indice $node ...")
            idx = (2(j-1) + 1):2j
            _reach_homog_krylov_LGG09_modif!(view(Mout, idx, :), X0_init, Aext, δt, Eplus, d, NSTEPS; m=mk, tol=tol)
        end
        F = [TemplateReachSet(cdirs, view(Mout, :, k), (0.0 .. δt) + (k-1) * δt) for k in 1:NSTEPS]

    else
        error("condicion_inicial = $condicion_inicial no existe")
    end

    # niego para ajustar el sistema de coordenadas con el articulo
    proyectar = get(kwargs, :proyectar, false)
    if proyectar
        # OJO si uso esto, el flowpipe tiene una sola variable 1
        @assert length(nodeidx) == 1
        out = [ReachSet(Interval(-Mout[2, k], Mout[1, k]), (0.0 .. δt) + (k-1) * δt) for k in 1:size(Mout, 2)]
    else
        #α = -1.0
        α = 1.0
        out = [ReachSet(α * set(R), tspan(R)) for R in F]
    end

    return out
end

# computes the support functions of Ω0 = CH(X0, Φ*X0 ⊕ Eplus)
# along +d and -d for NSTEPS
#
# for k = 0, 1, ..., NSTEPS-1
#
# ρ(d, Φ^k Ω0) = ρ((Φ^T)^k d, CH(X0, Φ * X0 ⊕ Eplus))
#              = max(ρ((Φ^T)^k d, X0), ρ((Φ^T)^k d, Φ * X0) + ρ((Φ^T)^k d, Eplus))
#              = max(ρ((Φ^T)^k d, X0), ρ((Φ^T)^(k+1) d, X0) + ρ((Φ^T)^k d, Eplus))
#              = max(ρ(vk, X0), ρ((Φ^T)^k * vk, X0) + ρ(vk, Eplus)), vk := (Φ^T)^k d
#
function _reach_homog_krylov_LGG09_modif!(out, X0, A, δ, Eplus,
                                          d::AbstractVector, NSTEPS;
                                          hermitian=false, m=min(30, size(Aᵀδ, 1)), tol=1e-7)

    @assert size(out, 1) == 2
    @assert size(out, 2) == NSTEPS

    Aᵀδ = transpose(A) .* δ

    # initialization of the krylov subspace
    # esta parte genera una base de
    # {Aᵀδ * d, (Aᵀδ)^2 * d, (Aᵀδ)^3 * d, ..., (Aᵀδ)^m * d}
    TA, Tb = eltype(Aᵀδ), eltype(d)
    T = promote_type(TA, Tb)
    Ks = KrylovSubspace{T, real(T)}(length(d), m)
    arnoldi!(Ks, Aᵀδ, d; m=m, ishermitian=hermitian, tol=tol)

    # rᵢ stores is the cache for each vector: (Φᵀ)^i d and (Φᵀ)^{i+1} d
    rᵢ = deepcopy(d)
    rᵢ₊₁ = similar(d)

    #
    # TODO se puede optimizar el loop ya que el ρ(rᵢ₊₁, X0) se repite en ρ(rᵢ, X0)
    #
    @inbounds for i in 1:NSTEPS
        expv!(rᵢ₊₁, i*1.0, Ks) # rᵢ₊₁ <- exp(A^T δ * i) * d
        out[1, i] = max(ρ(rᵢ, X0), ρ(rᵢ₊₁, X0) + ρ(rᵢ, Eplus))
        out[2, i] = max(ρ(-rᵢ, X0), ρ(-rᵢ₊₁, X0) + ρ(-rᵢ, Eplus))
        copyto!(rᵢ, rᵢ₊₁)
    end
    return out
end

# ===========================
# estadisticas con numericas
# ===========================

using Sobol
using StructuralDynamicsODESolvers: SolutionExtrema

sol_numericas_random_statistics(p; m, vecsamples::Int, idx, radio_desplazamientos, radio_velocidades) = sol_numericas_random_statistics(p; m=m,
    vecsamples=[vecsamples], idx=idx, radio_desplazamientos=radio_desplazamientos, radio_velocidades=radio_velocidades)

# ejemplo: vecsamples = [10, 20, 100]
function sol_numericas_random_statistics(p; m, vecsamples::Vector{Int}, idx, δt, radio_desplazamientos, radio_velocidades)
    #X0rand_desplazamientos = [[2(rand() - 0.5) * 5e-7 for _ in 1:825] for i in 1:NTRAJ]
    #X0rand_velocidades = [[2(rand() - 0.5) * 5e-7 for _ in 1:825] for i in 1:NTRAJ]
    #remake = [(X0d, X0v) for (X0d, X0v) in zip(X0rand_desplazamientos, X0rand_velocidades)]

    H0 = Hyperrectangle(zeros(2m), vcat(fill(radio_desplazamientos, m), fill(radio_velocidades, m)))

    sampleado = :vertices

    if sampleado == :sobol
        # para samplear dentro de H0
        sampler = SobolSeq(H0)
    elseif sampleado == :vertices
        # para samplear en los vertices
        sampler = VertexSampler(H0)
    end

    @time out = sol_numerica_lamb_statistics(p; case="heaviside", alg=Trapezoidal, δt=δt, X0=nothing, sampler=sampler, idx=idx, vecsamples=vecsamples)
end

Sobol.SobolSeq(H::Hyperrectangle) = SobolSeq(low(H), high(H))

struct VertexSampler{T}
    H::T
end

function StructuralDynamicsODESolvers._next!(S::VertexSampler, unew)
    c = S.H.center
    r = S.H.radius
    for i in 1:length(unew)
        a = ifelse(rand() > 0.5, 1, -1)
        unew[i] = c[i] + a*r[i]
    end
end

StructuralDynamicsODESolvers._next!(s::ScaledSobolSeq, unew) = next!(s, unew)

function generar_figura_malla2_envolventes_position(; vecsamples = nothing)

    # =================
    # setear parametros
    # =================
    p = load_lambs(mesh=2);
    @show m = size(p.s.M, 1)
    nodeidx = 3
    δt = step_size_lamb(p, 1.0) # 0.001
    case = "heaviside"
    condicion_inicial = "distribuida"
    mk = 140
    tol = 1e-7

    radio_desplazamientos = 3.5e-7
    radio_velocidades = 5e-6

    # ====================
    # correr simulaciones
    # ====================
    out_extrema_malla2 = nothing
    if !isnothing(vecsamples)
        out_extrema_malla2 = sol_numericas_random_statistics(p, m=m, idx=nodeidx, vecsamples=vecsamples, δt=δt, radio_desplazamientos=radio_desplazamientos, radio_velocidades=radio_velocidades)
    end

    # ====================
    # calcular flowpipe
    # ====================
    @time out_flowpipe = _sol_reachability_lamb(p, case=case, condicion_inicial=condicion_inicial, nodeidx=nodeidx, mk=mk, tol=tol, δt=δt, proyectar=true, radio_desplazamientos=radio_desplazamientos, radio_velocidades=radio_velocidades)

    # ====================
    # graficar flowpipe
    # ====================
    # q = length(vecsamples)
    # c = distinguishable_colors(q);
    fig = plot()

    plot!(fig, out_flowpipe, vars=(0, 1), lw=0.0, lc=:lightblue, c=:lightblue, lab=L"\textrm{Flowpipe}")

    if !isnothing(vecsamples)
        [plot!(fig, out_extrema_malla2.t, out_extrema_malla2.hist[k][:U_min], lw=2.0, lab="$k") for (i, k) in enumerate(vecsamples)]
        [plot!(fig, out_extrema_malla2.t, out_extrema_malla2.hist[k][:U_max], lw=2.0, lab="") for (i, k) in enumerate(vecsamples)]
    end
    savefig(fig, "malla2_distribuida_pos_nodeidx_3.pdf")

    out_flowpipe, out_extrema_malla2, fig
end

function generar_figura_malla2_envolventes_velocity(; vecsamples = nothing)

    # =================
    # setear parametros
    # =================
    p = load_lambs(mesh=2);
    @show m = size(p.s.M, 1)
    nodeidx = 3
    δt = step_size_lamb(p, 1.0)
    case = "heaviside"
    condicion_inicial = "distribuida"
    mk = 140
    tol = 1e-7

    radio_desplazamientos = 3.5e-7
    radio_velocidades = 5e-6

    # ====================
    # correr simulaciones
    # ====================
    println("Corriendo simulaciones...")
    out_extrema_malla2 = nothing
    if !isnothing(vecsamples)
        out_extrema_malla2 = sol_numericas_random_statistics(p, m=m, idx=nodeidx, vecsamples=vecsamples, δt=δt, radio_desplazamientos=radio_desplazamientos, radio_velocidades=radio_velocidades)
    end

    # ====================
    # calcular flowpipe
    # ====================
    println("Calculando el flowpipe...")
    @time out_flowpipe = _sol_reachability_lamb(p, case=case, condicion_inicial=condicion_inicial, nodeidx=nodeidx+m, mk=mk, tol=tol, δt=δt, proyectar=true, radio_desplazamientos=radio_desplazamientos, radio_velocidades=radio_velocidades)

    # ====================
    # graficar flowpipe
    # ====================
    # q = length(vecsamples)
    # c = distinguishable_colors(q);
    fig = plot()

    plot!(fig, out_flowpipe, vars=(0, 1), lw=0.0, lc=:lightblue, c=:lightblue, lab=L"\textrm{Flowpipe}")
    [plot!(fig, out_extrema_malla2.t, out_extrema_malla2.hist[k][:Udot_min], lw=2.0, lab="$k") for (i, k) in enumerate(vecsamples)]
    [plot!(fig, out_extrema_malla2.t, out_extrema_malla2.hist[k][:Udot_max], lw=2.0, lab="") for (i, k) in enumerate(vecsamples)]

    savefig(fig, "malla2_distribuida_vel_nodeidx_3.pdf")

    out_flowpipe, out_extrema_malla2, fig
end

function _norma_inf_envolvente_numericas(sol, vecsamples)
    sol_norma_inf = Vector{Float64}()
    for k in vecsamples
        Udotmin_k = sol.hist[k][:Udot_min] .|> abs
        Udotmax_k = sol.hist[k][:Udot_max] .|> abs

        # norma infinito de las envolventes numericas
        sol_norma_inf_k = maximum(max.(Udotmin_k, Udotmax_k))
        push!(sol_norma_inf, sol_norma_inf_k)
    end
    return sol_norma_inf
end

function _norma_1_envolvente_numericas(sol, vecsamples)
    sol_norma_1 = Vector{Float64}()
    for k in vecsamples
        Udotmin_k = sol.hist[k][:Udot_min] .|> abs
        Udotmax_k = sol.hist[k][:Udot_max] .|> abs

        # norma infinito de las envolventes numericas
        sol_norma_1_k = sum(max(Udotmin_k[i], Udotmax_k[i]) for i in 1:length(Udotmin_k))
        push!(sol_norma_1, sol_norma_1_k)
    end
    δt = out_extrema_malla2_vel.alg.Δt
    return sol_norma_1 * δt
end

function generar_figura_malla2_envolvente_velocidad()

    results = matread("results.mat")

    vecsamples = "N" .* string.(results["vecsamples"])

    times = results["out_extrema_malla2_vel_t"]
    vmin = [results["out_extrema_malla2_vel_hist"][k]["NUdot_min"] for k in vecsamples]
    vmax = [results["out_extrema_malla2_vel_hist"][k]["NUdot_max"] for k in vecsamples];

    xx = [Interval(a, b) for (a, b) in zip(results["out_flowpipe_vel_lo"], results["out_flowpipe_vel_hi"])];
    tt = [IntervalArithmetic.Interval(a, b) for (a, b) in zip(results["out_flowpipe_vel_tstart"], results["out_flowpipe_vel_tend"])];
    F = [ReachSet(a, b) for (a, b) in zip(xx, tt)];

    cols = distinguishable_colors(length(vecsamples))

    fig = plot(xlab=L"\textrm{Time (s)}",
               ylab=L"\textrm{Velocity} (\times 10^{-4} \textrm{m/s})",
               legend=:outerright,
               legendfontsize=20,
               tickfont=font(20, "Times"),
               guidefontsize=10,
               xguidefont=font(20, "Times"),
               yguidefont=font(20, "Times"),
               xtick = xticklatex([0, 0.2, 0.4, 0.6, 0.8, 1.0], 2),
               ytick = ([-10e-4, -5e-4, 0, 5e-4, 10e-4], [L"-10", L"-5", L"0", L"5", L"10"]),
               xlims=(0.0, 1.0),
               ylims=(-1.3e-3, 1.2e-3),
               bottom_margin=10mm, left_margin=6mm, right_margin=12mm, top_margin=3mm, size=(900, 600))

    labs = [L"10^0", L"10^1", L"10^2", L"10^3", L"10^4", L"10^5"]

    plot!(fig, F, vars=(0, 1), lw=1., lc=:lightblue, c=:lightblue, lab=L"\textrm{Flowpipe}")
    [plot!(fig, times, vmin[i], lw=2.0, c=cols[i], lab=labs[i], ls=:solid) for i in 1:length(vmin)]
    [plot!(fig, times, vmax[i], lw=2.0, c=cols[i], lab="", ls=:solid) for i in 1:length(vmin)]

    savefig(fig, "fig_vel.pdf")

    fig
end

max_stars_repo_name: JuliaReach/SetPropagation-FEM-Examples

max_issues_repo_name: JuliaReach/SetPropagation-FEM-Examples

max_forks_repo_name: JuliaReach/SetPropagation-FEM-Examples


text:
function save_posterior_samples!(sp_num::Int64, Pr::Priors, D::InputData, Posterior::PosteriorSample,
                                 resid::Residuals, fixed_effects::FixedEffects,
                                 genetic_effects::GeneticEffects, Factors::LatentFactors,
                                 interaction_effects::InteractionEffects)
    #save posteriors. Re-scale samples back to original variances.
    sp = Pr.sp;
    VY = D.Var_Y';
    Lambda = broadcast(*, Factors.Lambda, sqrt.(VY'));     #re-scale by Y variances
    G_h2 = Factors.h2;
    U = genetic_effects.U;
    delta = Factors.delta;
    genetic_ps = genetic_effects.ps ./ VY';
    resid_ps = resid.ps ./ VY';

    # save factor samples
    Lambda = Lambda[:,1:Factors.k];

    Posterior.Lambda[sp_num] = copy(Lambda);
    Posterior.delta[sp_num] = copy(delta);
    Posterior.G_h2[sp_num] = copy(G_h2);
    Posterior.U[sp_num] = copy(U);

    Posterior.no_f[sp_num] = Factors.k;

    Posterior.ps[:,sp_num] = copy(genetic_ps);
    Posterior.resid_ps[:,sp_num] = copy(resid_ps);

    # save B,U,W
    Posterior.B = Posterior.B + (broadcast(*, fixed_effects.B, sqrt.(VY'))  ./ sp);
    Posterior.d = Posterior.d + (broadcast(*, genetic_effects.d, sqrt.(VY')) ./ sp);
    Posterior.W = Posterior.W + (broadcast(*, interaction_effects.W, sqrt.(VY')) ./ sp);

end

	function PosteriorMeans(Posterior::PosteriorSample, D::InputData, Pr::Priors)
    kmax = convert(Int64, maximum(Posterior.no_f));

    Lambda_est = zeros(D.p, kmax);
    G_s = zeros(D.p, D.p, Pr.sp);
    P_s = zeros(D.p, D.p, Pr.sp);
    E_s = zeros(D.p, D.p, Pr.sp);
    G_est = zeros(D.p, D.p);
    P_est = zeros(D.p, D.p);
    E_est = zeros(D.p, D.p);

    factor_h2s_est = zeros(kmax);
    for j=1:Pr.sp
        Lj  = Posterior.Lambda[j];
        h2j = Posterior.G_h2[j];

        Pj = Lj * Lj'                   + Diagonal(1. ./ (Posterior.ps[:,j])) + Diagonal(1. ./(Posterior.resid_ps[:,j]));
        Gj = Lj * Diagonal(     h2j) * Lj' + Diagonal(1. ./ (Posterior.ps[:,j]));
        Ej = Lj * Diagonal(1 .- h2j) * Lj'                                    + Diagonal(1. ./(Posterior.resid_ps[:,j]));

        P_s[:, :, j] = copy(Pj);
        G_s[:, :, j] = copy(Gj);
        E_s[:, :, j] = copy(Ej);

        P_est = P_est + Pj./Pr.sp;
        G_est = G_est + Gj./Pr.sp;
        E_est = E_est + Ej./Pr.sp;

        k_local = size(h2j, 1)
        Lambda_est[:, 1:k_local]  = Lambda_est[:, 1:k_local] + Lj./Pr.sp;
        factor_h2s_est[1:k_local] = factor_h2s_est[1:k_local] + h2j./Pr.sp;

    end

    posterior_mean = Dict("G" => G_est,
      "P" => P_est,
      "E" => E_est,
      "Gs" => G_s,
      "Ps" => P_s,
      "Es" => E_s,
      "Lambda" => Lambda_est,
      "F_h2s" => factor_h2s_est)
    posterior_mean
end


text:
# HeuristicFrictionModels.jl

# Notes:
# - For the routines in this section it is expected that any heuristic model
#   provided implement two routines.  These routines are as follows:
#   #1: dzdt = heuristic_state_equation(mdl, t, nrm, pos, vel, z)
#   #2: F = heuristic_force_equation(mdl, t, nrm, pos, vel, z, dzdt)
#
# The inputs to these functions are as follows:
#   - mdl: A type based upon HeuristicFrictionModel
#   - t: A numeric value defining the time point at which to evaluate the model.
#   - nrm: The compressive normal force.
#   - pos: The relative position of the two contacting bodies at time t.
#   - vel: The relative (sliding) velocity between the two contacting bodies.
#   - z: An N-element state vector at time t.
#   - dzdt: An N-element vector of the state variable time derivatives.

"""
Evaluates a heuristic-type friction model at the given state.

The output of the function is a named tuple containing the following fields.
- f: The friction force.
- dzdt: The time derivative(s) of the state variable(s).
"""
function friction(
    mdl::HeuristicFrictionModel,
    t::T,
    nrm::T,
    pos::T,
    vel::T,
    z::Array{T}
) where T <: Number

    # Evaluate the state equation of the heuristic system
    dzdt = heuristic_state_equation(mdl, t, nrm, pos, vel, z)

    # Evaluate the friction force of the heuristic system
    F = heuristic_force_equation(mdl, t, nrm, pos, vel, z, dzdt)

    # Output
    return (f = F, dzdt = dzdt)
end

"""
Evaluates a heuristic-type friction model by integrating the differential
equations describing the model state.  The user can specify that the solver
choose the time step by only supplying a start and end simulation time, or the
user can specify the time points the solver returns by specifying more than 2
time points.  Notice, if supplying time points at which the solver should
return the results, the time points do not need to be evenly spaced; however,
they do need to be either monotonically increasing or decreasing.

Control over the ODE solver tolerances and maximum allowable time step are
exposed to the calling code via the varargs 'p' input parameter.  The following
inputs are currently recognized.
- reltol: The relative ODE solver tolerance. The defaults is 1e-8.
- abstol: The absolute ODE solver tolerance.  The default is 1e-6.
- dtmax: The maximum allowable step size.  The default is 1e-3.

The output of the function is a named tuple containing the following fields.
- t: The solution time points.
- f: The friction force.
- z: The state variable(s).
- dzdt: The time derivative(s) of the state variable(s).
"""
function friction(
    mdl::HeuristicFrictionModel,
    t::Array{T},
    nrm::Function,
    pos::Function,
    vel::Function,
    z0::Array{T};
    p...
) where T <: Number

    function diffeq!(du_, u_, p_, t_)
        x = pos(t_)
        v = vel(t_)
        N = nrm(t_)
        dzdt = heuristic_state_equation(mdl, t_, N, x, v, u_)
        for i in 1:length(dzdt)
            du_[i] = dzdt[i]
        end
    end

    args = extract_options(p)

    prob = ODEProblem(diffeq!, z0, [first(t), last(t)])

    if length(t) == 2
        sol = solve(
            prob,
            alg_hints = [:stiff],
            reltol = args.reltol,
            abstol = args.abstol,
            dtmax = args.dtmax
        )
    elseif length(t) > 2
        sol = solve(
            prob,
            alg_hints = [:stiff],
            saveat = t,
            reltol = args.reltol,
            abstol = args.abstol,
            dtmax = args.dtmax
        )
    else
        error("No time range for the solution has been given.  Ensure at least two values are provided in the time array input.")
    end

    z = Array(sol)
    npts = length(sol.t)
    nargs = length(z0)
    dzdt = zeros(T, nargs, npts)
    F = zeros(T, npts)
    for i in 1:npts
        x = pos(sol.t[i])
        v = vel(sol.t[i])
        N = nrm(sol.t[i])
        dzdt[:,i] = heuristic_state_equation(mdl, sol.t[i], N, x, v, z[:,i])
        F[i] = heuristic_force_equation(mdl, sol.t[i], N, x, v, z[:,i], dzdt[:,i])
    end

    return (t = sol.t, f = F, z = z, dzdt = dzdt)
end

"""
Evaluates a heuristic-type friction model, that is not dependent upon position,
by integrating the differential equations describing the model state.  The user
can specify that the solver choose the time step by only supplying a start and
end simulation time, or the user can specify the time points the solver returns
by specifying more than 2 time points.  Notice, if supplying time points at
which the solver should return the results, the time points do not need to be
evenly spaced; however, they do need to be either monotonically increasing or
decreasing.

Control over the ODE solver tolerances and maximum allowable time step are
exposed to the calling code via the varargs 'p' input parameter.  The following
inputs are currently recognized.
- reltol: The relative ODE solver tolerance. The defaults is 1e-8.
- abstol: The absolute ODE solver tolerance.  The default is 1e-6.
- dtmax: The maximum allowable step size.  The default is 1e-3.

The output of the function is a named tuple containing the following fields.
- t: The solution time points.
- f: The friction force.
- z: The state variable(s).
- dzdt: The time derivative(s) of the state variable(s).
"""
function friction(
    mdl::HeuristicFrictionModel,
    t::Array{T},
    nrm::Function,
    vel::Function,
    z0::Array{T};
    p...
) where T <: Number

    pos(ti) = zero(T)
    return friction(mdl, t, nrm, pos, vel, z0; p...)
end

"""
Evaluates a heuristic friction model by integrating the state equations using
the supplied data.  If the solver requires points lying between provided data
points linear interpolation will be used.  The user can specify that the solver
choose the time step by only supplying a start and end simulation time, or the
user can specify the time points the solver returns by specifying more than 2
time points.  Notice, if supplying time points at which the solver should return
the results, the time points do not need to be evenly spaced; however, they do
need to be either monotonically increasing or decreasing.

Control over the ODE solver tolerances and maximum allowable time step are
exposed to the calling code via the varargs 'p' input parameter.  The following
inputs are currently recognized.
- reltol: The relative ODE solver tolerance. The defaults is 1e-8.
- abstol: The absolute ODE solver tolerance.  The default is 1e-6.
- dtmax: The maximum allowable step size.  The default is 1e-3.

The output of the function is a named tuple containing the following fields.
- t: The solution time points.
- f: The friction force.
- z: The state variable(s).
- dzdt: The time derivative(s) of the state variable(s).
"""
function friction(
    mdl::HeuristicFrictionModel,
    t::Array{T},
    nrm::Array{T},
    pos::Array{T},
    vel::Array{T},
    z0::Array{T};
    p...
) where T <: Number

    normal_interp = LinearInterpolation(
        t, nrm,
        extrapolation_bc = Line()
    )
    position_interp = LinearInterpolation(
        t, pos,
        extrapolation_bc = Line()
    )
    velocity_interp = LinearInterpolation(
        t, vel,
        extrapolation_bc = Line()
    )

    N(ti) = normal_interp(ti)
    X(ti) = position_interp(ti)
    V(ti) = velocity_interp(ti)

    return friction(mdl, t, N, X, V, z0; p...)
end

"""
Evaluates a heuristic friction model that doesn't depend upon position data by
integrating the state equations using the supplied data.  If the solver requires
points lying between provided data points linear interpolation will be used.
The user can specify that the solver choose the time step by only supplying a
start and end simulation time, or the user can specify the time points the
solver returns by specifying more than 2 time points.  Notice, if supplying time
points at which the solver should return the results, the time points do not
need to be evenly spaced; however, they do need to be either monotonically
increasing or decreasing.

Control over the ODE solver tolerances and maximum allowable time step are
exposed to the calling code via the varargs 'p' input parameter.  The following
inputs are currently recognized.
- reltol: The relative ODE solver tolerance. The defaults is 1e-8.
- abstol: The absolute ODE solver tolerance.  The default is 1e-6.
- dtmax: The maximum allowable step size.  The default is 1e-3.

The output of the function is a named tuple containing the following fields.
- t: The solution time points.
- f: The friction force.
- z: The state variable(s).
- dzdt: The time derivative(s) of the state variable(s).
"""
function friction(
    mdl::HeuristicFrictionModel,
    t::Array{T},
    nrm::Array{T},
    vel::Array{T},
    z0::Array{T};
    p...
) where T <: Number

    normal_interp = LinearInterpolation(
        t, nrm,
        extrapolation_bc = Line()
    )
    velocity_interp = LinearInterpolation(
        t, vel,
        extrapolation_bc = Line()
    )

    N(ti) = normal_interp(ti)
    X(ti) = zero(T)
    V(ti) = velocity_interp(ti)

    return friction(mdl, t, N, X, V, z0; p...)
end


text:
export irToJulia

function erasedToJuliaTy(e::ErasedType)
    ! = erasedToJuliaTy
    @match e begin
        ERArrow(a, r) => :Function
        ERNom(name) =>
            @match name begin
                :i64 => :Int64
                :i32 => :Int32
                :i16 => :Int16
                :i8 => :Int8
                :f16 => :Float16
                :f32 => :Float32
                :f64 => :Float64
                :bool => :Bool
                :str => :String
                :char => :Char
                _ => :Any
            end
        ERTuple(xs) => :Tuple
        ERType(t) => Symbol
        ERAny => :Any
    end
end

function irToJulia(decl, ln::LineNumberNode)
    a = MLFS.HM.VerboseUN[]
    MLFS.HM.VerboseUN[] = true
    r = irToJulia_(decl, ln)
    MLFS.HM.VerboseUN[] = a
    r
end

function irToJulia_(decl::IR.Decl, ln::LineNumberNode)
    @match decl begin
        IR.Perform(e) => irToJulia_(e, ln)
        IR.Assign(s, t, e) =>
            let t = erasedToJuliaTy(tyErase(t))
                :($s :: $t = $(irToJulia_(e, ln)))
            end
    end
end

function irToJulia_(expr::IR.Expr, ln::LineNumberNode)
    @match expr begin
        IR.Expr(ln, t, e) =>
            let e = irToJulia_(e, ln)
                if t === nothing
                    e
                else
                    t = erasedToJuliaTy(tyErase(t))
                    :($e :: $t)
                end
            end
    end
end

function irToJulia_(expr::IR.ExprImpl, ln::LineNumberNode)
    @switch expr begin
    @case IR.ETypeVal(t)
        QuoteNode(Symbol(string(t)))

    @case IR.EExt(ex)
        ex

    @case IR.EVar(s)
        s
    @case IR.ELet(decls, expr)
        Expr(:block,
            ln,
            map(x -> irToJulia_(x, ln), decls)...,
            irToJulia_(expr, ln)
        )
    @case IR.EITE(a, b, c)
        Expr(:if,
            irToJulia_(a, ln),
            irToJulia_(b, ln),
            irToJulia_(c, ln),
        )
    @case IR.EFun(s, d)
        Expr(
            :function,
            Expr(:tuple, s),
            Expr(:block,
                ln,
                irToJulia_(d, ln)))
    @case IR.EApp(f, a)
        Expr(:call, irToJulia_(f, ln), irToJulia_(a, ln))
    @case IR.ETup(xs)
        Expr(:tuple, irToJulia_.(xs, [ln])...)
    @case IR.EInt(i, _) || IR.EFloat(i, _) ||
          IR.EStr(i) || IR.EChar(i) ||
          IR.EBool(i)
        i
    @case IR.EIm()
        error("should be eliminated")
    end
end


text:
"""
Checkerboard initialization: Assaad four site version for square lattice
"""
function find_four_site_hopping_corners(l::Lattice)
  pc = Int(floor(l.sites/4))

  A_corners = zeros(Int, pc)
  idx = 1
  tolinidx = reshape(collect(1:l.sites),l.L,l.L)
  for x in 1:2:l.L
      for y in 1:2:l.L
          A_corners[idx] = tolinidx[y,x]
          idx += 1
      end
  end
  B_corners = l.neighbors[1,l.neighbors[2,A_corners]]
  return A_corners, B_corners
end

function build_four_site_hopping_matrix_exp(mc::AbstractDQMC{CBAssaad}, corners::Tuple{Array{Int64,1},Array{Int64,1}}, prefac::Float64=0.5)
  p = mc.p
  l = mc.l

  pc = Int(floor(l.sites/4))
  chkr_hop_4site = Array{SparseMatrixCSC, 3}(undef,pc,2,2) # i, group (A, B), hopping flavor (tx, ty)
  chkr_hop_4site_inv = Array{SparseMatrixCSC, 3}(undef,pc,2,2)

  # build explicitly based on analytic form (only without mag field)
  for tflv in 1:2
    for g in 1:2
      for c in 1:pc

        chkr_hop_4site[c,g,tflv] = sparse(1.0I, l.sites, l.sites)
        chkr_hop_4site_inv[c,g,tflv] = sparse(1.0I, l.sites, l.sites)

        # Tijmn where i ist bottom left and i<j<m<n.
        i = corners[g][c]
        j = l.neighbors[1,i]
        m = l.neighbors[2,i]
        n = l.neighbors[2,j]

        inds = sort!([corners[g][c], l.neighbors[1,i], l.neighbors[2,i], l.neighbors[2,j]])

        th = l.t[1,tflv]
        tv = l.t[2,tflv]
        fac = -prefac * p.delta_tau

        chkr_hop_4site[c,g,tflv][i,i] = chkr_hop_4site[c,g,tflv][j,j] = chkr_hop_4site[c,g,tflv][m,m] = chkr_hop_4site[c,g,tflv][n,n] = cosh(fac * -th) * cosh(fac * -tv) #
        chkr_hop_4site[c,g,tflv][i,n] = chkr_hop_4site[c,g,tflv][j,m] = chkr_hop_4site[c,g,tflv][m,j] = chkr_hop_4site[c,g,tflv][n,i] = sinh(fac * -th) * sinh(fac * -tv)
        chkr_hop_4site[c,g,tflv][i,j] = chkr_hop_4site[c,g,tflv][j,i] = chkr_hop_4site[c,g,tflv][m,n] = chkr_hop_4site[c,g,tflv][n,m] = cosh(fac * -th) * sinh(fac * -tv) #
        chkr_hop_4site[c,g,tflv][i,m] = chkr_hop_4site[c,g,tflv][j,n] = chkr_hop_4site[c,g,tflv][m,i] = chkr_hop_4site[c,g,tflv][n,j] = sinh(fac * -th) * cosh(fac * -tv)

        chkr_hop_4site_inv[c,g,tflv][i,i] = chkr_hop_4site_inv[c,g,tflv][j,j] = chkr_hop_4site_inv[c,g,tflv][m,m] = chkr_hop_4site_inv[c,g,tflv][n,n] = cosh(-fac * -th) * cosh(-fac * -tv)
        chkr_hop_4site_inv[c,g,tflv][i,n] = chkr_hop_4site_inv[c,g,tflv][j,m] = chkr_hop_4site_inv[c,g,tflv][m,j] = chkr_hop_4site_inv[c,g,tflv][n,i] = sinh(-fac * -th) * sinh(-fac * -tv)
        chkr_hop_4site_inv[c,g,tflv][i,j] = chkr_hop_4site_inv[c,g,tflv][j,i] = chkr_hop_4site_inv[c,g,tflv][m,n] = chkr_hop_4site_inv[c,g,tflv][n,m] = cosh(-fac * -th) * sinh(-fac * -tv)
        chkr_hop_4site_inv[c,g,tflv][i,m] = chkr_hop_4site_inv[c,g,tflv][j,n] = chkr_hop_4site_inv[c,g,tflv][m,i] = chkr_hop_4site_inv[c,g,tflv][n,j] = sinh(-fac * -th) * cosh(-fac * -tv)

      end
    end
  end

  return chkr_hop_4site, chkr_hop_4site_inv
end

function init_checkerboard_matrices(mc::AbstractDQMC{CBAssaad})
  p = mc.p
  l = mc.l

  println("Initializing hopping exponentials (Checkerboard)")
  pc = Int(floor(l.sites/4))

  corners = find_four_site_hopping_corners(l)

  chkr_hop_4site_half, chkr_hop_4site_half_inv = build_four_site_hopping_matrix_exp(mc, corners, 0.5)
  chkr_hop_4site, chkr_hop_4site_inv = build_four_site_hopping_matrix_exp(mc, corners, 1.)

  eTx_A_half = foldl(*,chkr_hop_4site_half[:,1,1])
  eTx_B_half = foldl(*,chkr_hop_4site_half[:,2,1])
  eTy_A_half = foldl(*,chkr_hop_4site_half[:,1,2])
  eTy_B_half = foldl(*,chkr_hop_4site_half[:,2,2])
  eTx_A = foldl(*,chkr_hop_4site[:,1,1])
  eTx_B = foldl(*,chkr_hop_4site[:,2,1])
  eTy_A = foldl(*,chkr_hop_4site[:,1,2])
  eTy_B = foldl(*,chkr_hop_4site[:,2,2])

  eTx_A_half_inv = foldl(*,chkr_hop_4site_half_inv[:,1,1])
  eTx_B_half_inv = foldl(*,chkr_hop_4site_half_inv[:,2,1])
  eTy_A_half_inv = foldl(*,chkr_hop_4site_half_inv[:,1,2])
  eTy_B_half_inv = foldl(*,chkr_hop_4site_half_inv[:,2,2])
  eTx_A_inv = foldl(*,chkr_hop_4site_inv[:,1,1])
  eTx_B_inv = foldl(*,chkr_hop_4site_inv[:,2,1])
  eTy_A_inv = foldl(*,chkr_hop_4site_inv[:,1,2])
  eTy_B_inv = foldl(*,chkr_hop_4site_inv[:,2,2])

  if p.opdim == 3
    eT_A_half = cat(eTx_A_half,eTy_A_half,eTx_A_half,eTy_A_half, dims=(1,2))
    eT_B_half = cat(eTx_B_half,eTy_B_half,eTx_B_half,eTy_B_half, dims=(1,2))
    eT_A_half_inv = cat(eTx_A_half_inv,eTy_A_half_inv,eTx_A_half_inv,eTy_A_half_inv, dims=(1,2))
    eT_B_half_inv = cat(eTx_B_half_inv,eTy_B_half_inv,eTx_B_half_inv,eTy_B_half_inv, dims=(1,2))
    eT_A = cat(eTx_A,eTy_A,eTx_A,eTy_A, dims=(1,2))
    eT_B = cat(eTx_B,eTy_B,eTx_B,eTy_B, dims=(1,2))
    eT_A_inv = cat(eTx_A_inv,eTy_A_inv,eTx_A_inv,eTy_A_inv, dims=(1,2))
    eT_B_inv = cat(eTx_B_inv,eTy_B_inv,eTx_B_inv,eTy_B_inv, dims=(1,2))
  else
    eT_A_half = cat(eTx_A_half,eTy_A_half, dims=(1,2))
    eT_B_half = cat(eTx_B_half,eTy_B_half, dims=(1,2))
    eT_A_half_inv = cat(eTx_A_half_inv,eTy_A_half_inv, dims=(1,2))
    eT_B_half_inv = cat(eTx_B_half_inv,eTy_B_half_inv, dims=(1,2))
    eT_A = cat(eTx_A,eTy_A, dims=(1,2))
    eT_B = cat(eTx_B,eTy_B, dims=(1,2))
    eT_A_inv = cat(eTx_A_inv,eTy_A_inv, dims=(1,2))
    eT_B_inv = cat(eTx_B_inv,eTy_B_inv, dims=(1,2))
  end

  l.chkr_hop_half = [eT_A_half, eT_B_half]
  l.chkr_hop_half_inv = [eT_A_half_inv, eT_B_half_inv]
  l.chkr_hop_half_dagger = [adjoint(eT_A_half), adjoint(eT_B_half)]
  l.chkr_hop = [eT_A, eT_B]
  l.chkr_hop_inv = [eT_A_inv, eT_B_inv]
  l.chkr_hop_dagger = [adjoint(eT_A), adjoint(eT_B)]

  muv = vcat(fill(p.mu1,l.sites),fill(p.mu2,l.sites))
  muv = repeat(muv, outer=[Int(p.flv/2)])
  l.chkr_mu_half = sparse(Diagonal(exp.(-0.5*p.delta_tau * -muv)))
  l.chkr_mu_half_inv = sparse(Diagonal(exp.(0.5*p.delta_tau * -muv)))
  l.chkr_mu = sparse(Diagonal(exp.(-p.delta_tau * -muv)))
  l.chkr_mu_inv = sparse(Diagonal(exp.(p.delta_tau * -muv)))

  l.n_groups = 2

  hop_mat_exp_chkr = l.chkr_hop_half[1] * l.chkr_hop_half[2] * sqrt.(l.chkr_mu)
  r = effreldiff(l.hopping_matrix_exp,hop_mat_exp_chkr)
  r[findall(x->x==zero(x), hop_mat_exp_chkr)] .= 0
  println("Checkerboard - exact (abs):\t\t", maximum(absdiff(l.hopping_matrix_exp,hop_mat_exp_chkr)))
  # println("Checkerboard - exact (eff rel):\t\t", maximum(r))
end

#### WITH ARTIFICIAL B-FIELD

function build_four_site_hopping_matrix_Bfield(mc::AbstractDQMC{CBAssaad}, corner::Int,f::Int,s::Int)
  l = mc.l
  H = heltype(mc)

  sites_clockwise = [corner, l.neighbors[1,corner], l.neighbors[1,l.neighbors[2,corner]], l.neighbors[2,corner]]
  shift_sites_clockwise = circshift(sites_clockwise,-1)

  h = l.t[1,f]
  v = l.t[2,f]

  hop = -1 * repeat(H[v, h, v, h], outer=2)

  @inbounds for k in 1:4
    i = sites_clockwise[k]
    j = shift_sites_clockwise[k]

    hop[k] *= exp(im * l.peirls[s,f][i,j]) # i,j = trg,src
    hop[k+4] *= exp(im * l.peirls[s,f][j,i]) # j,i = trg,src
  end

  # return T only filled with T_ijkl hoppings
  return sparse([sites_clockwise; shift_sites_clockwise],[shift_sites_clockwise; sites_clockwise],hop,l.sites,l.sites)
end

function build_four_site_hopping_matrix_exp_Bfield(mc::AbstractDQMC{CBAssaad}, corners::Tuple{Array{Int64,1},Array{Int64,1}}, prefac::Float64=0.5)
  p = mc.p
  l = mc.l
  H = heltype(mc)

  B = zeros(2,2) # rowidx = spin up,down, colidx = flavor
  if p.Bfield
    B[1,1] = B[2,2] = 2 * pi / l.sites
    B[1,2] = B[2,1] = - 2 * pi / l.sites
  end

  pc = Int(floor(l.sites/4))

  chkr_hop_4site = Array{SparseMatrixCSC{H, Int}, 4}(undef, pc,2,2,2) # i, group (A, B), spin (up, down), flavor (x, y)
  chkr_hop_4site_inv = Array{SparseMatrixCSC{H, Int}, 4}(undef, pc,2,2,2)

  # build numerically (due to mag field)
  for f in 1:2
    for s in 1:2
      for g in 1:2
        for c in 1:pc

          fac = -prefac * p.delta_tau

          chkr_hop_4site[c,g,s,f] = sparse(rem_eff_zeros!(exp(fac * Matrix(build_four_site_hopping_matrix_Bfield(mc,corners[g][c],f,s)))))
          chkr_hop_4site_inv[c,g,s,f] = sparse(rem_eff_zeros!(exp(-fac * Matrix(build_four_site_hopping_matrix_Bfield(mc,corners[g][c],f,s)))))

        end
      end
    end
  end

  return chkr_hop_4site, chkr_hop_4site_inv
end

function init_checkerboard_matrices_Bfield(mc::AbstractDQMC{CBAssaad})
  p = mc.p
  l = mc.l
  H = heltype(mc)

  println("Initializing hopping exponentials (Bfield, Checkerboard)")
  # pc = Int(floor(l.sites/4))

  corners = find_four_site_hopping_corners(l)

  chkr_hop_4site_half, chkr_hop_4site_half_inv = build_four_site_hopping_matrix_exp_Bfield(mc, corners, 0.5)
  chkr_hop_4site, chkr_hop_4site_inv = build_four_site_hopping_matrix_exp_Bfield(mc, corners, 1.)

  eT_half = Array{SparseMatrixCSC{H, Int}, 3}(undef,2,2,2) # group (A, B), spin (up, down), flavor (x, y)
  eT_half_inv = Array{SparseMatrixCSC{H, Int}, 3}(undef,2,2,2)
  eT = Array{SparseMatrixCSC{H, Int}, 3}(undef,2,2,2)
  eT_inv = Array{SparseMatrixCSC{H, Int}, 3}(undef,2,2,2)

  for g in 1:2
    for f in 1:2
      for s in 1:2
        eT_half[g,s,f] = foldl(*,chkr_hop_4site_half[:,g,s,f])
        eT_half_inv[g,s,f] = foldl(*,chkr_hop_4site_half_inv[:,g,s,f])
        eT[g,s,f] = foldl(*,chkr_hop_4site[:,g,s,f])
        eT_inv[g,s,f] = foldl(*,chkr_hop_4site_inv[:,g,s,f])
      end
    end
  end

  if p.opdim == 3
    eT_A_half = cat(eT_half[1,1,1], eT_half[1,2,2], eT_half[1,2,1], eT_half[1,1,2], dims=(1,2))
    eT_B_half = cat(eT_half[2,1,1], eT_half[2,2,2], eT_half[2,2,1], eT_half[2,1,2], dims=(1,2))
    eT_A_half_inv = cat(eT_half_inv[1,1,1], eT_half_inv[1,2,2], eT_half_inv[1,2,1], eT_half_inv[1,1,2], dims=(1,2))
    eT_B_half_inv = cat(eT_half_inv[2,1,1], eT_half_inv[2,2,2], eT_half_inv[2,2,1], eT_half_inv[2,1,2], dims=(1,2))
    eT_A = cat(eT[1,1,1], eT[1,2,2], eT[1,2,1], eT[1,1,2], dims=(1,2))
    eT_B = cat(eT[2,1,1], eT[2,2,2], eT[2,2,1], eT[2,1,2], dims=(1,2))
    eT_A_inv = cat(eT_inv[1,1,1], eT_inv[1,2,2], eT_inv[1,2,1], eT_inv[1,1,2], dims=(1,2))
    eT_B_inv = cat(eT_inv[2,1,1], eT_inv[2,2,2], eT_inv[2,2,1], eT_inv[2,1,2], dims=(1,2))
  else # O(2) and O(1) model
    eT_A_half = cat(eT_half[1,1,1], eT_half[1,2,2], dims=(1,2))
    eT_B_half = cat(eT_half[2,1,1], eT_half[2,2,2], dims=(1,2))
    eT_A_half_inv = cat(eT_half_inv[1,1,1], eT_half_inv[1,2,2], dims=(1,2))
    eT_B_half_inv = cat(eT_half_inv[2,1,1], eT_half_inv[2,2,2], dims=(1,2))
    eT_A = cat(eT[1,1,1], eT[1,2,2], dims=(1,2))
    eT_B = cat(eT[2,1,1], eT[2,2,2], dims=(1,2))
    eT_A_inv = cat(eT_inv[1,1,1], eT_inv[1,2,2], dims=(1,2))
    eT_B_inv = cat(eT_inv[2,1,1], eT_inv[2,2,2], dims=(1,2))
  end

  l.chkr_hop_half = [eT_A_half, eT_B_half]
  l.chkr_hop_half_inv = [eT_A_half_inv, eT_B_half_inv]
  l.chkr_hop_half_dagger = [adjoint(eT_A_half), adjoint(eT_B_half)]
  l.chkr_hop = [eT_A, eT_B]
  l.chkr_hop_inv = [eT_A_inv, eT_B_inv]
  l.chkr_hop_dagger = [adjoint(eT_A), adjoint(eT_B)]

  muv = vcat(fill(p.mu1,l.sites),fill(p.mu2,l.sites))
  muv = repeat(muv, outer=[Int(p.flv/2)])
  l.chkr_mu_half = sparse(Diagonal(exp.(-0.5*p.delta_tau * -muv)))
  l.chkr_mu_half_inv = sparse(Diagonal(exp.(0.5*p.delta_tau * -muv)))
  l.chkr_mu = sparse(Diagonal(exp.(-p.delta_tau * -muv)))
  l.chkr_mu_inv = sparse(Diagonal(exp.(p.delta_tau * -muv)))

  l.n_groups = 2

  hop_mat_exp_chkr = l.chkr_hop_half[1] * l.chkr_hop_half[2] * sqrt.(l.chkr_mu)
  r = effreldiff(l.hopping_matrix_exp,hop_mat_exp_chkr)
  r[findall(x -> x == zero(x), hop_mat_exp_chkr)] .= 0
  println("Checkerboard (Bfield) - exact (abs):\t\t", maximum(absdiff(l.hopping_matrix_exp,hop_mat_exp_chkr)))
  # println("Checkerboard (Bfield) - exact (eff rel):\t\t", maximum(r))
end

max_stars_repo_name: crstnbr/dqmc

max_issues_repo_name: crstnbr/dqmc

max_forks_repo_name: crstnbr/dqmc


text:
import ConstructionBase

export ParameterizedMeasure

abstract type ParameterizedMeasure{N} <: AbstractMeasure end

function Base.getproperty(μ::ParameterizedMeasure{N}, prop::Symbol) where {N}
    return getproperty(getfield(μ, :par), prop)
end

function Base.propertynames(μ::ParameterizedMeasure{N}) where {N}
    return N
end

function Base.show(io::IO, μ::ParameterizedMeasure{()})
    io = IOContext(io, :compact => true)
    print(io, nameof(typeof(μ)), "()")
end

function Base.show(io::IO, μ::ParameterizedMeasure{N}) where {N}
    io = IOContext(io, :compact => true)
    print(io, nameof(typeof(μ)))
    print(io, getfield(μ, :par))
end

# Allow things like
#
# julia> Normal{(:μ,)}(2)
# Normal(μ = 2,)
#
function (::Type{P})(args...) where {N,P<:ParameterizedMeasure{N}}
    C = constructorof(P)
    return C(NamedTuple{N}(args...))
end

(::Type{P})(; kwargs...) where {P<:ParameterizedMeasure} = P(NamedTuple(kwargs))

function ConstructionBase.setproperties(
    d::P,
    nt::NamedTuple
) where {P<:ParameterizedMeasure}
    return constructorof(P)(merge(params(d), nt))
end

###############################################################################
# params

export params

params(μ::ParameterizedMeasure) = getfield(μ, :par)

function params(μ::AbstractMeasure, constraints::NamedTuple{C}) where {C}
    NamedTuple{paramnames(μ, constraints)}(params(μ))
end

params(μ::AbstractMeasure) = NamedTuple()

###############################################################################
# paramnames

export paramnames

paramnames(μ) = paramnames(typeof(μ))

paramnames(::Type{PM}) where {N,PM<:ParameterizedMeasure{N}} = N

paramnames(μ::AbstractMeasure) = propertynames(μ)

params(::Type{PM}) where {N,PM<:ParameterizedMeasure{N}} = N

function paramnames(μ, constraints::NamedTuple{N}) where {N}
    tuple((k for k in paramnames(μ) if k ∉ N)...)
end

###############################################################################
# kernelfactor

function kernelfactor(::Type{P}) where {N,P<:ParameterizedMeasure{N}}
    (constructorof(P), N)
end

function kernelfactor(::P) where {N,P<:ParameterizedMeasure{N}}
    (constructorof(P), N)
end