text:
using TSSM
using Combinatorics
using LinearAlgebra

function sign(p::Array{Int,1})
    n = length(p)
    k = 0
    for i=1:n
        for j=i+1:n
            if p[i]>p[j]
                k += 1
            end
        end
    end
   (-1)^k
end

struct MultiFor
    k::Array{Int,1}
end

function Base.iterate(MF::MultiFor, k::Array{Int,1}=Int[])
    if k==MF.k
        return nothing
    end
    if k==Int[]
        k = ones(Int, length(MF.k))
        return(copy(k), k)
    end
    for i=1:length(k)
        if k[i]<MF.k[i]
            k[i] += 1
            for j = 1:i-1
                k[j] = 1
            end
            return (copy(k), k)
        end
    end
end

function find_xxx(v::Vector{Int}, lena::Int, slater_exchange, slater_indices; must_be_included::Int=0)
    res = Tuple{Int,Int,Float64}[]
    for i=1:lena
        for j=1:lena
            if must_be_included==0 || ((must_be_included in slater_indices[i]) || (must_be_included in slater_indices[i]))
                if v==slater_exchange[i,j][1]
                    push!(res, (i,j, slater_exchange[i,j][2]))
                end
            end
        end
    end
    res
end

function add_xxx!(d::Dict{Tuple{Int,Int}, Float64}, v::Tuple{Int,Int,Float64})
    key = (v[1],v[2])
    f = get!(d, key, 0.0) + v[3]
    if f==0
        delete!(d, key)
    else
        d[key] = f
    end
    d
end

function add_xxx!(d::Dict{Tuple{Int,Int}, Float64}, v::Vector{Tuple{Int,Int,Float64}})
    for x in v
        add_xxx!(d, x)
    end
    d
end

function init_mctdhf_combinatorics(f::Int, N::Int)
    slater_indices = collect(Combinatorics.Combinations(1:N, f))

    density_rules = [Tuple{Int,Int,Int}[] for j=1:N, l=1:N]
    for J = MultiFor([N for j=2:f])
        for j=1:N
            for l=j:N
                J1 = vcat(j,J)
                J2 = vcat(l,J)
                #j1 = findfirst(slater_indices, sort(J1))
                j1 = something(findfirst(isequal(sort(J1)), slater_indices), 0)
                if j1>0
                    #j2 = findfirst(slater_indices, sort(J2))
                    j2 = something(findfirst(isequal(sort(J2)), slater_indices), 0)
                    if j2>0
                        s1 = sign(J1)
                        s2 = sign(J2)
                        t = (j1, j2, s1*s2)
                        push!(density_rules[j,l], t)
                        if j!=l
                            push!(density_rules[l,j], t)
                        end
                    end
                end
            end
        end
    end

    density2_rules = [Tuple{Int,Int,Int}[] for j=1:N, l=1:N, p=1:N, q=1:N]
    if f==2
        for j=1:N
            for l=j:N
                for p=1:N
                    for q=1:N
                        J1 = vcat(j, p)
                        J2 = vcat(l, q)
                        #j1 = findfirst(slater_indices, sort(J1))
                        j1 = something(findfirst(isequal(sort(J1)), slater_indices), 0)
                        if j1>0
                            #j2 = findfirst(slater_indices, sort(J2))
                            j2 = something(findfirst(isequal(sort(J2)), slater_indices), 0)
                            if j2>0
                                s1 = sign(J1)
                                s2 = sign(J2)
                                t = (j1, j2, s1*s2)
                                push!(density2_rules[j,l,p,q], t)
                                if j!=l
                                    push!(density2_rules[l,j,q,p], t)
                                end
                            end
                        end
                    end
                end
            end
        end
    else
    for J = MultiFor([N for j=3:f])
        for j=1:N
            for l=j:N
                for p=1:N
                    for q=1:N
                        J1 = vcat(j, p, J)
                        J2 = vcat(l, q, J)
                        #j1 = findfirst(slater_indices, sort(J1))
                        j1 = something(findfirst(isequal(sort(J1)), slater_indices), 0)
                        if j1>0
                            #j2 = findfirst(slater_indices, sort(J2))
                            j2 = something(findfirst(isequal(sort(J2)), slater_indices), 0)
                            if j2>0
                                s1 = sign(J1)
                                s2 = sign(J2)
                                t = (j1, j2, s1*s2)
                                push!(density2_rules[j,l,p,q], t)
                                if j!=l
                                    push!(density2_rules[l,j,q,p], t)
                                end
                            end
                        end
                    end
                end
            end
        end
    end
    end

    lena = binomial(N,f)
    slater_exchange = [(Int[], 1) for j=1:lena, l=1:lena]
    # entry = (exchange_index_pairs, sign)
    for j=1:lena
        for l=1:lena
            u = symdiff(slater_indices[j], slater_indices[l])
            if length(u)==2
                #i1 = findfirst(slater_indices[j], u[1])
                i1 = something(findfirst(isequal(u[1]), slater_indices[j]), 0)
                #i2 = findfirst(slater_indices[l], u[2])
                i2 = something(findfirst(isequal(u[2]), slater_indices[l]), 0)
                s = (-1)^(i1+i2)
                slater_exchange[j,l] = (u, s)
            elseif length(u)==4
                u = [u[1], u[3], u[2], u[4]]
                #i1 = findfirst(slater_indices[j], u[1])
                i1 = something(findfirst(isequal(u[1]), slater_indices[j]), 0)
                #i2 = findfirst(slater_indices[l], u[2])
                i2 = something(findfirst(isequal(u[2]), slater_indices[l]), 0)
                #i3 = findfirst(slater_indices[j], u[3])
                i3 = something(findfirst(isequal(u[3]), slater_indices[j]), 0)
                #i4 = findfirst(slater_indices[l], u[4])
                i4 = something(findfirst(isequal(u[4]), slater_indices[l]), 0)
                s = (-1)^(i1+i2+i3+i4)
                slater_exchange[j,l] = (u, s)
            end
        end
    end

    res = [Dict{Tuple{Int,Int},Float64}([]) for p=1:N, q=1:N]
    for p=1:N
        for q=1:N
            if p==q
                for j=1:lena
                    if p in slater_indices[j]
                       add_xxx!(res[p,q], (j,j, 1.0))
                    end
                end
            end
            v = find_xxx([q,p], lena, slater_exchange, slater_indices)
            add_xxx!(res[p,q], v)
        end
    end
    slater1_rules = [[(key[1],key[2],val) for (key,val) in res[p,q]] for p=1:N, q=1:N]

    res = [Dict{Tuple{Int,Int},Float64}([]) for p=1:N, q=1:N, r=1:N, s=1:N]
    for p=1:N
        for q=1:N
            for r=1:N
                for s=1:N
                    if p==q && r==s && p!=r
                        for j=1:lena
                            if p in slater_indices[j] && r in slater_indices[j]
                                add_xxx!(res[p,q,r,s], (j,j, 0.5))
                            end
                        end
                    end
                    if p==s && r==q && p!=r
                        for j=1:lena
                            if p in slater_indices[j] && r in slater_indices[j]
                                add_xxx!(res[p,q,r,s], (j,j, -0.5))
                            end
                        end
                    end
                    if p==q && r!=s
                        v = find_xxx([s,r], lena, slater_exchange, slater_indices, must_be_included=p)
                        add_xxx!(res[p,q,r,s], v)
                    end
                    if q==r && s!=p
                        v0 = find_xxx([s,p], lena, slater_exchange, slater_indices, must_be_included=q)
                        v = [(j,k,-sigma) for (j,k, sigma) in v0]
                        add_xxx!(res[p,q,r,s], v)
                    end
                    v = find_xxx([s,r,q,p], lena, slater_exchange, slater_indices)
                    add_xxx!(res[p,q,r,s], v)
                    v0 = find_xxx([s,p,q,r], lena, slater_exchange, slater_indices)
                    v = [(j,k,-sigma) for (j,k, sigma) in v0]
                    add_xxx!(res[p,q,r,s], v)
                end
            end
        end
    end
    slater2_rules = [[(key[1],key[2],val) for (key,val) in res[p,q,r,s]] for p=1:N, q=1:N, r=1:N, s=1:N]

    orthogonalization_rules = [Tuple{Int,Int,Int}[] for p=1:N, q=1:N]
    for p=1:N
        for j=1:lena
            #k = findfirst(slater_indices[j], p)
            k = something(findfirst(isequal(p), slater_indices[j]), 0)
            if k>0
                for q=1:p-1
                    J = copy(slater_indices[j])
                    J[k] = q
                    J1 = sort(J)
                    #l = findfirst(slater_indices, J1)
                    l = something(findfirst(isequal(J1), slater_indices), 0)
                    if l>0
                        s = sign(J)*sign(J1)
                        push!(orthogonalization_rules[p,q], (j,l, s))
                    end
                end
            end
        end
    end

    slater_indices, density_rules, density2_rules, slater_exchange, slater1_rules, slater2_rules, orthogonalization_rules
end

mutable struct MCTDHF1D <: TimeSplittingSpectralMethodComplex1D{Float64}
    m::Schroedinger1D
    f::Int # number of electrons
    N::Int # number of orbitals
    lena::Int # number of (independent) coefficients
    spins::Array{Int, 1}
    spin_restricted::Bool

    slater_indices    :: Array{Array{Int64,1},1}
    density_rules     :: Array{Array{Tuple{Int64,Int64,Int64},1},2}
    density2_rules    :: Array{Array{Tuple{Int64,Int64,Int64},1},4}
    slater_exchange   :: Array{Tuple{Array{Int64,1},Int64},2}
    slater1_rules     :: Array{Array{Tuple{Int64,Int64,Float64},1},2}
    slater2_rules     :: Array{Array{Tuple{Int64,Int64,Float64},1},4}
    orthogonalization_rules :: Array{Array{Tuple{Int64,Int64,Int64},1},2}

    Vee               :: Array{Float64,2}
    density_matrix    :: Array{Complex{Float64},2}
    density2_tensor   :: Array{Complex{Float64},4}
    u_pq              :: Array{Complex{Float64},1}
    u_pqs	      :: Array{Complex{Float64},1}

    k1
    k2
    k3
    k4

    function MCTDHF1D(f::Integer, N::Integer,
                      nx::Integer, xmin::Real, xmax::Real; spins::Array{Int,1}=ones(Int, N),
                      spin_restricted::Bool=true,
                      potential1::Function=TSSM.none_1D, potential1_t::Function=TSSM.none_2D ,potential2::Function=TSSM.none_2D)
        m = Schroedinger1D(nx, xmin, xmax, potential=potential1, potential_t = potential1_t)
        lena = binomial(N,f)
        if spin_restricted
            spins = [(-1)^(k+1) for k=1:N]
        end
        slater_indices, density_rules, density2_rules, slater_exchange, slater1_rules, slater2_rules, orthogonalization_rules =
              init_mctdhf_combinatorics(f, N)
        mm = new(m, f, N, lena, spins, spin_restricted, slater_indices, density_rules,
                density2_rules, slater_exchange, slater1_rules, slater2_rules, orthogonalization_rules,
                zeros(Float64, nx, nx),
                zeros(Complex{Float64},N,N),  zeros(Complex{Float64},N,N,N,N),
        zeros(Complex{Float64}, nx), zeros(Complex{Float64}, nx), nothing, nothing, nothing, nothing )
        set_potential2!(mm, potential2)
        mm
    end
end

function set_potential2!(m::MCTDHF1D, V::Function)
    xx = get_nodes(m.m)
    n = get_nx(m.m)
    for ix=1:n
        x = xx[ix]
        for iy=1:n
            y = xx[iy]
            m.Vee[ix, iy] = V(x,y)
        end
    end
end

set_potential1!(m::MCTDHF1D, V::Function) = set_potential!(m.m, V)
set_potential1_t!(m::MCTDHF1D, V::Function) = set_potential_t!(m.m, V)

mutable struct Orbital
    phi::WfSchroedinger1D
    spin::Int
end

mutable struct WfMCTDHF1D <: WaveFunctionComplex1D{Float64}
    a::Array{Complex{Float64}, 1}
    o::Array{Orbital, 1}
    m::MCTDHF1D
    function WfMCTDHF1D(m::MCTDHF1D)
        a = zeros(Complex{Float64}, m.lena)
        if m.spin_restricted
            w = [WfSchroedinger1D(m.m) for j=1:m.N]
            o = [Orbital(w[div(j+1,2)], m.spins[j]) for j=1:m.N]
        else
            o = [Orbital(WfSchroedinger1D(m.m), m.spins[j]) for j=1:m.N]
        end
        new(a, o, m)
    end
end

function TSSM.wave_function(m::MCTDHF1D )
    WfMCTDHF1D(m)
end

LinearAlgebra.norm(o::Orbital) = norm(o.phi)
TSSM.inner_product(o1::Orbital, o2::Orbital) = (o1.spin==o2.spin ? inner_product(o1.phi, o2.phi) : 0.0im)
TSSM.potential_matrix_element(o1::Orbital, o2::Orbital) = (o1.spin==o2.spin ? potential_matrix_element(o1.phi, o2.phi) : 0.0im)
TSSM.kinetic_matrix_element(o1::Orbital, o2::Orbital) = (o1.spin==o2.spin ? kinetic_matrix_element(o1.phi, o2.phi) : 0.0im)

function LinearAlgebra.scale!(o::Orbital, f::Number)
    scale!(o.phi, f)
end

function LinearAlgebra.axpy!(o1::Orbital, o2::Orbital, f::Number)
    axpy!(o1.phi, o2.phi, f)
end

function LinearAlgebra.scale!(psi::WfMCTDHF1D, f::Number)
    st = psi.m.spin_restricted ? 2 : 1
    for j=1:st:psi.m.N
        scale!(psi.o[j], f)
    end
    psi.a[:] *= f
end

function LinearAlgebra.axpy!(psi1::WfMCTDHF1D, psi2::WfMCTDHF1D, f::Number)
    st = psi1.m.spin_restricted ? 2 : 1
    for j=1:st:psi1.m.N
        axpy!(psi1.o[j], psi2.o[j], f)
    end
    psi1.a[:] += f*psi2.a[:]
end

function Base.copy!(psi1::WfMCTDHF1D, psi2::WfMCTDHF1D)
    st = psi1.m.spin_restricted ? 2 : 1
    for j=1:st:psi1.m.N
        copy!(psi1.o[j].phi, psi2.o[j].phi)
    end
    for j=1:psi1.m.N
        psi1.o[j].spin = psi2.o[j].spin
    end
    psi1.a[:] = psi2.a[:]
end

using HDF5

function TSSM.save(psi::WfMCTDHF1D, filename::AbstractString)
    st = psi.m.spin_restricted ? 2 : 1
    for k=1:st:psi.m.N
       save(psi.o[k].phi, filename, string("orbital_",k, "_real"),
             string("orbital_", k, "_imag", ), append=(k>1))
    end
    h5open(filename,"r+") do file
        file["coefficients_real"] =  real(psi.a)
        file["coefficients_imag"] =  imag(psi.a)
        file["spins"] = Cint[psi.o[k].spin for k=1:psi.m.N]
        attrs(file)["number_of_particles"] = psi.m.f
        attrs(file)["number_of_orbitals"] = psi.m.N
        attrs(file)["spin_restricted"] = psi.m.spin_restricted ? 1 : 0
    end
    filename
end

function TSSM.load!(psi::WfMCTDHF1D, filename::AbstractString)
    st = psi.m.spin_restricted ? 2 : 1
    for k=1:st:psi.m.N
       load!(psi.o[k].phi, filename, string("orbital_",k, "_real"),
             string("orbital_", k, "_imag", ))
    end
    h5open(filename,"r") do file
        psi.a = read(file["coefficients_real"])+ 1im*read(file["coefficients_imag"])
    end
    psi
end

function gen_density_matrix(psi::WfMCTDHF1D)
    m = psi.m
    N = m.N
    rho = m.density_matrix
    rho[:,:] .= 0.0
    for j=1:N
        for l=j:N
            for (u, v, s) in m.density_rules[j,l]
                rho[j,l] += s * conj(psi.a[u]) * psi.a[v]
                if l!=j
                   rho[l,j] = conj(rho[j,l])
                end
            end
        end
    end
    nothing
end

function gen_density2_tensor(psi::WfMCTDHF1D; mult_inverse_density_matrix::Bool=true)
    m = psi.m
    N = m.N
    rho = m.density2_tensor
    rho[:,:,:,:] .= 0.0
    for j=1:N
        for l=j:N
            for p=1:N
                for q=1:N
                    for (u, v, s) in m.density2_rules[j,l,p,q]
                        rho[j,l,p,q] += s * conj(psi.a[u]) * psi.a[v]
                        if l!=j
                           rho[l,j,q,p] = conj(rho[j,l,p,q])
                        end
                    end
                end
            end
        end
    end
    if mult_inverse_density_matrix
        X = bunchkaufman(m.density_matrix)
        for p=1:N
            for q=1:N
                rho[:,:,p,q] = X \ rho[:,:,p,q]
            end
        end
    end
    nothing
end

LinearAlgebra.norm(psi::WfMCTDHF1D) = norm(psi.a)
#Note, only correct if orbitals are ortonormal

function TSSM.inner_product(psi1::WfMCTDHF1D, psi2::WfMCTDHF1D)
    m = psi1.m
    if m ≠ psi2.m
        error("psi1 and psi2 must belong to the same method")
    end
    ip = zeros(Complex{Float64}, m.N,m.N)
    st = m.spin_restricted ? 2 : 1
    for j=1:st:m.N
        for l=1:st:m.N
            ip[j,l] = inner_product(psi1.o[j], psi2.o[l])
            if m.spin_restricted && j+1<=m.N && l+1<=m.N
                ip[j+1,l+1] = ip[j,l]
            end
        end
    end

    ps = [(p, sign(p)) for p in (Combinatorics.permutations(1:m.f))]
    d = 0.0im
    for j=1:m.lena
        J = m.slater_indices[j]
        for l=1:m.lena
            L = m.slater_indices[l]
            aa = conj(psi1.a[j])*psi2.a[l]
            for (p,s) in ps
                L1 = getindex(L, p)
                d = d + s*aa*prod([ip[J[k],L1[k]] for k=1:m.f])
            end
        end
    end
    d
end

function TSSM.distance(psi1::WfMCTDHF1D, psi2::WfMCTDHF1D; expanded::Bool=false)
    m = psi1.m
    if m ≠ psi2.m
        error("psi1 and psi2 must belong to the same method")
    end
    if expanded
        n2 = norm(psi1)^2+norm(psi2)^2-2*real(inner_product(psi1,psi2))
        if n2<0.0
            return sqrt(complex(n2))
        else
            return sqrt(n2)
        end
    else
        n2 = norm(psi1.a-psi2.a).^2
        st = m.spin_restricted ? 2 : 1
        for j=1:st:m.N
            n2 += distance(psi1.o[j].phi, psi2.o[j].phi)^2
        end
        return sqrt(n2)
    end
end

function TSSM.to_real_space!(psi::WfMCTDHF1D)
    st = psi.m.spin_restricted ? 2 : 1
    for j=1:st:psi.m.N
        to_real_space!(psi.o[j].phi)
    end
end

function TSSM.to_frequency_space!(psi::WfMCTDHF1D)
    st = psi.m.spin_restricted ? 2 : 1
    for j=1:st:psi.m.N
        to_frequency_space!(psi.o[j].phi)
    end
end

function set_zero!(psi::WfMCTDHF1D)
    st = psi.m.spin_restricted ? 2 : 1
    for j=1:st:psi.m.N
        get_data(psi.o[j].phi, true)[:] = 0.0
    end
    psi.a[:] = 0.0
end

function TSSM.set!(psi::WfMCTDHF1D, x::Number)
    st = psi.m.spin_restricted ? 2 : 1
    for j=1:st:psi.m.N
        set!(psi.o[j].phi, x)
    end
    psi.a[:] .= x
end

function gen_rhs1!(rhs::WfMCTDHF1D, psi::WfMCTDHF1D; include_kinetic_part::Bool=false,
                   include_one_particle_potential_part::Bool=true)
    if !(include_kinetic_part||include_one_particle_potential_part)
        return # nothing to do
    end
    m = rhs.m
    st = m.spin_restricted ? 2 : 1
    if m ≠ psi.m
        error("rhs and psi must belong to the same method")
    end
    for q=1:st:m.N
        if include_kinetic_part
            add_apply_A!(psi.o[q].phi, rhs.o[q].phi, 1im)
        end
        if include_one_particle_potential_part
            add_apply_B!(psi.o[q].phi, rhs.o[q].phi, 1im)
        end
        h = inner_product(psi.o[q], rhs.o[q])
        for (j,l,f) in m.slater1_rules[q,q]
            rhs.a[j] += h*f*psi.a[l]
        end
        if m.spin_restricted && q+1<=m.N
            for (j,l,f) in m.slater1_rules[q+1,q+1]
                rhs.a[j] += h*f*psi.a[l]
            end
        end
        for p=1:st:q-1
            if psi.o[p].spin==psi.o[q].spin
                h = inner_product(psi.o[p], rhs.o[q])
                for (j,l,f) in m.slater1_rules[q,p]
                    rhs.a[j] += h*f*psi.a[l]
                    rhs.a[l] += conj(h)*f*psi.a[j]
                end
                if m.spin_restricted &&  q+1<=m.N && p+1<=m.N
                    for (j,l,f) in m.slater1_rules[q+1,p+1]
                        rhs.a[j] += h*f*psi.a[l]
                        rhs.a[l] += conj(h)*f*psi.a[j]
                    end
                end
            end
        end
    end
end

function project_out_orbitals!(rhs:: WfMCTDHF1D, psi::WfMCTDHF1D)
    m = psi.m
    st = m.spin_restricted ? 2 : 1
    c = zeros(Complex{Float64},m.N)
    for p = 1:st:m.N
        for q = 1:st:m.N
            if psi.o[p].spin==psi.o[q].spin
                c[q] = inner_product(psi.o[q], rhs.o[p])
            end
        end
        for q = 1:st:m.N
            if psi.o[p].spin==psi.o[q].spin
                axpy!(rhs.o[p], psi.o[q], -c[q])
            end
        end
    end
end

function gen_rhs2!(rhs::WfMCTDHF1D, psi::WfMCTDHF1D)
    m = rhs.m
    if m ≠ psi.m
        error("rhs and psi must belong to the same method")
    end
    n = get_nx(m.m)
    dx = (get_xmax(m.m)-get_xmin(m.m))/n
    u_pq = m.u_pq
    u_pqs = m.u_pqs
    to_real_space!(psi)
    to_real_space!(rhs)
    st = m.spin_restricted ? 2 : 1
    for p=1:st:m.N
        for q=1:st:p
            if psi.o[p].spin==psi.o[q].spin
                u_pq[:] = m.Vee * (conj(get_data(psi.o[p].phi, true)).*get_data(psi.o[q].phi, true))
                for s=1:st:m.N
                    u_pqs[:] = u_pq .* get_data(psi.o[s].phi, true)
                    for r=1:st:m.N
                        if psi.o[r].spin==psi.o[s].spin

                            u = get_data(rhs.o[r].phi, true)
                            u[:] += (m.density2_tensor[r,s,p,q]*(m.f-1)*dx) * u_pqs
                            if m.spin_restricted && s+1<=m.N && r+1<=m.N
                                u[:] += (m.density2_tensor[r+1,s+1,p,q]*(m.f-1)*dx) * u_pqs
                            end

                            h = dot(get_data(psi.o[r].phi, true), u_pqs) * dx^2
                            for (j,l,f) in m.slater2_rules[q,p,s,r]
                                rhs.a[j] += h*f*psi.a[l]
                            end
                            if m.spin_restricted
                                if q+1<=m.N && p+1<=m.N
                                    for (j,l,f) in m.slater2_rules[q+1,p+1,s,r]
                                        rhs.a[j] += h*f*psi.a[l]
                                    end
                                end
                                if s+1<=m.N && r+1<=m.N
                                    for (j,l,f) in m.slater2_rules[q,p,s+1,r+1]
                                        rhs.a[j] += h*f*psi.a[l]
                                    end
                                end
                                if q+1<=m.N && p+1<=m.N && s+1<=m.N && r+1<=m.N
                                    for (j,l,f) in m.slater2_rules[q+1,p+1,s+1,r+1]
                                        rhs.a[j] += h*f*psi.a[l]
                                    end
                                end
                            end
                            if p!=q
                                h = conj(h)
                                for (j,l,f) in m.slater2_rules[p,q,r,s]
                                    rhs.a[j] += h*f*psi.a[l]
                                end
                                if m.spin_restricted
                                    if q+1<=m.N && p+1<=m.N
                                        for (j,l,f) in m.slater2_rules[p+1,q+1,r,s]
                                            rhs.a[j] += h*f*psi.a[l]
                                        end
                                    end
                                    if s+1<=m.N && r+1<=m.N
                                        for (j,l,f) in m.slater2_rules[p,q,r+1,s+1]
                                            rhs.a[j] += h*f*psi.a[l]
                                        end
                                    end
                                    if q+1<=m.N && p+1<=m.N && s+1<=m.N && r+1<=m.N
                                        for (j,l,f) in m.slater2_rules[p+1,q+1,r+1,s+1]
                                            rhs.a[j] += h*f*psi.a[l]
                                        end
                                    end
                                end
                            end
                        end
                    end
                end
                if p!=q
                    u_pq = conj(u_pq)
                    for s=1:st:m.N
                        u_pqs[:] = u_pq .* get_data(psi.o[s].phi, true)
                        for r=1:st:m.N
                            if psi.o[r].spin==psi.o[s].spin
                                u = get_data(rhs.o[r].phi, true)
                                u[:] += (m.density2_tensor[r,s,q,p]*(m.f-1)*dx) * u_pqs
                                if m.spin_restricted && s+1<=m.N && r+1<=m.N
                                    u[:] += (m.density2_tensor[r+1,s+1,q,p]*(m.f-1)*dx) * u_pqs
                                end
                            end
                        end
                    end
                end
            end
        end
    end
end

function gen_rhs!(rhs::WfMCTDHF1D, psi::WfMCTDHF1D; include_kinetic_part::Bool=false,
                   include_one_particle_potential_part::Bool=true )
    if rhs.m ≠ psi.m
        error("rhs and psi must belong to the same method")
    end
    orthonormalize_orbitals!(psi)
    gen_density_matrix(psi)
    gen_density2_tensor(psi)
    set!(rhs, 0.0)
    gen_rhs1!(rhs, psi, include_kinetic_part=include_kinetic_part,
                        include_one_particle_potential_part=include_one_particle_potential_part)
    gen_rhs2!(rhs, psi)
    project_out_orbitals!(rhs, psi)
    scale!(rhs, -1im)
end

function orthonormalize_orbitals!(psi::WfMCTDHF1D)
    m = psi.m
    g = zeros(Complex{Float64}, m.N)
    st = m.spin_restricted ? 2 : 1

    for p = 1:st:m.N
        for q=1:st:p-1
            g[q] = inner_product(psi.o[q], psi.o[p])
            if m.spin_restricted && q+1<=m.N
                g[q+1] = g[q]
            end
        end
        a1 = zeros(Complex{Float64}, m.lena)
        for q=1:st:p-1
            if psi.o[p].spin==psi.o[q].spin
                axpy!(psi.o[p], psi.o[q], -g[q])
                for (j, l, s) in m.orthogonalization_rules[p,q]
                    a1[l] += s*g[q]*psi.a[j]
                end
            end
        end
        f = norm(psi.o[p])
        scale!(psi.o[p],1/f)
        for j=1:m.lena
            if p in m.slater_indices[j]
                a1[j] += f*psi.a[j]
            else
                a1[j] += psi.a[j]
            end
        end
        psi.a[:] = a1

        if m.spin_restricted && p+1<=m.N
            a1 = zeros(Complex{Float64}, m.lena)
            for q=1:st:p-1
                if psi.o[p].spin==psi.o[q].spin
                    if q+1<=m.N
                        for (j, l, s) in m.orthogonalization_rules[p+1,q+1]
                            a1[l] += s*g[q+1]*psi.a[j]
                        end
                    end
                end
            end
            for j=1:m.lena
                if p+1 in m.slater_indices[j]
                    a1[j] += f*psi.a[j]
                else
                    a1[j] += psi.a[j]
                end
            end
            psi.a[:] = a1
        end
    end
    psi
end

function LinearAlgebra.normalize!(psi::WfMCTDHF1D)
    orthonormalize_orbitals!(psi)
    psi.a[:] /= norm(psi)
end

function potential_energy_1(psi::WfMCTDHF1D)
    m = psi.m
    V = 0.0
    st = m.spin_restricted ? 2 : 1
    for p=1:st:m.N
        h = potential_matrix_element(psi.o[p], psi.o[p])
        for (j,l,f) in m.slater1_rules[p,p]
            V += real(h*f*psi.a[j]*conj(psi.a[l]))
        end
        if m.spin_restricted && p+1<=m.N+1
            for (j,l,f) in m.slater1_rules[p+1,p+1]
                V += real(h*f*psi.a[j]*conj(psi.a[l]))
            end
        end
        for q=1:st:p-1
            if psi.o[p].spin==psi.o[q].spin
                h = potential_matrix_element(psi.o[p], psi.o[q])
                for (j,l,f) in m.slater1_rules[p,q]
                    V += 2*real(h*f*psi.a[j]*conj(psi.a[l]))
                end
                if m.spin_restricted &&  q+1<=m.N && p+1<=m.N+1
                    for (j,l,f) in m.slater1_rules[p+1,q+1]
                        V += 2*real(h*f*psi.a[j]*conj(psi.a[l]))
                    end
                end
            end
        end
    end
    real(V)
end

function TSSM.kinetic_energy(psi::WfMCTDHF1D)
    m = psi.m
    V = 0
    st = m.spin_restricted ? 2 : 1
    for p=1:st:m.N
        h = kinetic_matrix_element(psi.o[p], psi.o[p])
        for (j,l,f) in m.slater1_rules[p,p]
            V += real(h*f*psi.a[j]*conj(psi.a[l]))
        end
        if m.spin_restricted && p+1<=m.N+1
            for (j,l,f) in m.slater1_rules[p+1,p+1]
                V += real(h*f*psi.a[j]*conj(psi.a[l]))
            end
        end
        for q=1:st:p-1
            if psi.o[p].spin==psi.o[q].spin
                h = kinetic_matrix_element(psi.o[p], psi.o[q])
                for (j,l,f) in m.slater1_rules[p,q]
                    V += 2*real(h*f*psi.a[j]*conj(psi.a[l]))
                end
                if m.spin_restricted &&  q+1<=m.N && p+1<=m.N+1
                    for (j,l,f) in m.slater1_rules[p+1,q+1]
                        V += 2*real(h*f*psi.a[j]*conj(psi.a[l]))
                    end
                end
            end
        end
    end
    V
end

function potential_energy_2(psi::WfMCTDHF1D)
    m = psi.m
    V = 0
    n = get_nx(m.m)
    dx = (get_xmax(m.m)-get_xmin(m.m))/n
    u_pq = m.u_pq
    u_pqs = m.u_pqs
    to_real_space!(psi)
    st = m.spin_restricted ? 2 : 1
    for p=1:st:m.N
        for q=1:st:p
            if psi.o[p].spin==psi.o[q].spin
                u_pq[:] = m.Vee * (conj(get_data(psi.o[p].phi, true)).*get_data(psi.o[q].phi, true))
                for s=1:st:m.N
                    u_pqs[:] = u_pq .* get_data(psi.o[s].phi, true)
                    for r=1:st:m.N
                        if psi.o[r].spin==psi.o[s].spin
                            h = dot(get_data(psi.o[r].phi, true), u_pqs)
                            for (j,l,f) in m.slater2_rules[q,p,s,r]
                                V += h*f*conj(psi.a[j])*psi.a[l]
                            end
                            if m.spin_restricted
                                if q+1<=m.N && p+1<=m.N
                                    for (j,l,f) in m.slater2_rules[q+1,p+1,s,r]
                                        V += h*f*conj(psi.a[j])*psi.a[l]
                                    end
                                end
                                if s+1<=m.N && r+1<=m.N
                                    for (j,l,f) in m.slater2_rules[q,p,s+1,r+1]
                                        V += h*f*conj(psi.a[j])*psi.a[l]
                                    end
                                end
                                if q+1<=m.N && p+1<=m.N && s+1<=m.N && r+1<=m.N
                                    for (j,l,f) in m.slater2_rules[q+1,p+1,s+1,r+1]
                                        V += h*f*conj(psi.a[j])*psi.a[l]
                                    end
                                end
                            end
                            if p!=q
                                h = conj(h)
                                for (j,l,f) in m.slater2_rules[p,q,r,s]
                                    V += h*f*conj(psi.a[j])*psi.a[l]
                                end
                                if m.spin_restricted
                                    if q+1<=m.N && p+1<=m.N
                                        for (j,l,f) in m.slater2_rules[p+1,q+1,r,s]
                                            V += h*f*conj(psi.a[j])*psi.a[l]
                                        end
                                    end
                                    if s+1<=m.N && r+1<=m.N
                                        for (j,l,f) in m.slater2_rules[p,q,r+1,s+1]
                                            V += h*f*conj(psi.a[j])*psi.a[l]
                                        end
                                    end
                                    if q+1<=m.N && p+1<=m.N && s+1<=m.N && r+1<=m.N
                                        for (j,l,f) in m.slater2_rules[p+1,q+1,r+1,s+1]
                                            V += h*f*conj(psi.a[j])*psi.a[l]
                                        end
                                    end
                                end
                            end
                        end
                    end
                end
            end
        end
    end
    V *= dx^2
    real(V)
end

TSSM.potential_energy(psi::WfMCTDHF1D) = potential_energy_1(psi) + potential_energy_2(psi)

function TSSM.get_energy_expectation_deviation(psi::WfMCTDHF1D; tmp::WfMCTDHF1D=psi.m.k1)
    gen_rhs!(tmp, psi, include_kinetic_part=true, include_one_particle_potential_part=true )
    E = -imag(dot(psi.a, tmp.a))
    dev = norm(tmp.a+1im*E*psi.a)
    E, dev
end

TSSM.interaction_energy(psi::WfMCTDHF1D) = 0.0 #dummy

function TSSM.set_time!(psi::WfMCTDHF1D, t::Number)
   st = psi.m.spin_restricted ? 2 : 1
   for j=1:st:psi.m.N
       set_time!(psi.o[j].phi, t)
   end
end

TSSM.get_time(psi::WfMCTDHF1D) = get_time(psi.o[1].phi)
TSSM.set_propagate_time_together_with_A!(m::MCTDHF1D, flag::Bool) = set_propagate_time_together_with_A!(m.m, flag)
TSSM.get_propagate_time_together_with_A(m::MCTDHF1D) = get_propagate_time_together_with_A(m.m)

function TSSM.propagate_A!(psi::WfMCTDHF1D, dt::Number)
    st = psi.m.spin_restricted ? 2 : 1
    for j=1:st:psi.m.N
        propagate_A!(psi.o[j].phi, dt)
    end
end

function TSSM.propagate_B!(psi::WfMCTDHF1D, dt::Number)
    st = psi.m.spin_restricted ? 2 : 1
    for j=1:st:psi.m.N
        propagate_B!(psi.o[j].phi, dt)
    end
end

function TSSM.imaginary_time_propagate_A!(psi::WfMCTDHF1D, dt::Real)
    st = psi.m.spin_restricted ? 2 : 1
    for j=1:st:psi.m.N
        imaginary_time_propagate_A!(psi.o[j].phi, dt)
    end
end

function TSSM.imaginary_time_propagate_B!(psi::WfMCTDHF1D, dt::Real)
    st = psi.m.spin_restricted ? 2 : 1
    for j=1:st:psi.m.N
        imaginary_time_propagate_B!(psi.o[j].phi, dt)
    end
end

function TSSM.add_apply_A!(this::WfMCTDHF1D, other::WfMCTDHF1D, coefficient::Number=1.0)
    if this.m ≠ other.m
        error("this and other must belong to the same method")
    end
    st = this.m.spin_restricted ? 2 : 1
    for j=1:st:this.m.N
        add_apply_A!(this.o[j].phi, other.o[j].phi, coefficient)
    end
end

function TSSM.add_phi_A!(this::WfMCTDHF1D, other::WfMCTDHF1D, dt::Number, n::Integer, coefficient::Number=1.0)
    if this.m ≠ other.m
        error("this and other must belong to the same method")
    end
    st = this.m.spin_restricted ? 2 : 1
    for j=1:st:this.m.N
        add_phi_A!(this.o[j].phi, other.o[j].phi, dt, n, coefficient)
    end
    other.a[:] += coefficient/factorial(n)*this.a[:]
end


text:
# init file for "make devrepl"
using Revise
using Plots
unicodeplots()
using JuliaFormatter
using QuantumControlBase.TestUtils: test
using LiveServer: LiveServer, serve, servedocs as _servedocs
include(joinpath(@__DIR__, "clean.jl"))

servedocs(; kwargs...) = _servedocs(; skip_dirs=["docs/src/examples"], kwargs...)

REPL_MESSAGE = """
*******************************************************************************
DEVELOPMENT REPL

Revise, JuliaFormatter, LiveServer, Plots with unicode backend are active.

* `help()` – Show this message
* `include("test/runtests.jl")` – Run the entire test suite
* `include("test/generate_example_tests.jl")` – Convert all examples to tests
* `include("test/examples/simple_state_to_state.jl")` –
  Run an individual example test (after converting examples)!
* `test()` – Run the entire test suite in a subprocess with coverage
* `test(genhtml=true)` – Generate an HTML coverage report
* `include("docs/make.jl")` – Generate the documentation
* `format(".")` – Apply code formatting to all files
* `servedocs([port=8000, verbose=false])` –
  Build and serve the documentation. Automatically recompile and redisplay on
  changes
* `clean()` – Clean up build/doc/testing artifacts
* `distclean()` – Restore to a clean checkout state
*******************************************************************************
"""

"""Show help"""
help() = println(REPL_MESSAGE)

ext: jl


text:
using ReservoirComputing

const approx_res_size = 10
const radius = 0.99
const sparsity = 0.1
const alpha = 1.0
const beta = 1*10^(-1)
const extended_states = false
const input_weight = 0.1
const gates_weight = 0.8
const nla_type = NLADefault()
const activation = tanh
const predict_len = 10

train = ones(Float64, 2, 10)

W = init_reservoir_givensp(approx_res_size, radius, sparsity)
W_in = irrational_sign_input(approx_res_size, size(train, 1), input_weight)

gruesn = GRUESN(W, train, W_in,
        extended_states = extended_states,
        gates_weight = gates_weight,
        nla_type = nla_type,
        activation = activation,
        alpha = alpha)

@test isequal(approx_res_size, gruesn.res_size)
@test isequal(gruesn.in_size, size(train, 1))
@test isequal(train, gruesn.train_data)
@test isequal(alpha, gruesn.alpha)
@test isequal(nla_type, gruesn.nla_type)
@test isequal(activation, gruesn.activation)
@test size(gruesn.W) == (gruesn.res_size, gruesn.res_size)
@test size(gruesn.W_in) == (gruesn.res_size, gruesn.in_size)
@test isequal(gruesn.extended_states, extended_states)

#test train
W_out = ESNtrain(gruesn, beta)
@test size(W_out) == (gruesn.in_size, gruesn.res_size)
#test predict
output = GRUESNpredict(gruesn, predict_len, W_out)
@test size(output) == (gruesn.in_size, predict_len)


text:
using Plots, LaTeXStrings, Plots.Measures
using ReachabilityAnalysis: dim
using MAT

(@isdefined TARGET_FOLDER) ? nothing : TARGET_FOLDER = ""

# load the model
include("Heat1D_Model.jl")

# plot the center
# flowpipe
δ = 1e-5
file = joinpath(@__DIR__, "heat1d_99.mat")

println("Solving with set propagation (BOX)")
@time sol = solve_heat1d(NSTEPS=30_000, δ=δ, file=file, model=StepIntersect(Forward()));

# numerical solutions
println("Solving with Euler's method (sol +)")
@time sol₊ = solve_heat1d_implicit_euler(file=file, δ=δ, tol=0.1, NSTEPS=30_000);

println("Solving with Euler's method (sol -)")
@time sol₋ = solve_heat1d_implicit_euler(file=file, δ=δ, tol=-0.1, NSTEPS=30_000);

node = 51 # 1 a 99 (nodo medio = 51)

sold₊ = displacements(sol₊)
sold₋ = displacements(sol₋)

tdom = times(sol₊)

sol₊_num(node) = [x[node] for x in sold₊[1:end-1]];
sol₋_num(node) = [x[node] for x in sold₋[1:end-1]];

sol₊_node = sol₊_num(node)
sol₋_node = sol₋_num(node);

# temperature evolution at the central node
function plot_node_vs_t()

    fig = plot(xlab=L"\textrm{Time (s)}", ylab=L"\textrm{Temperature at the center}",
               legendfontsize=18, legend=:topright,
               tickfont=font(20, "Times"), guidefontsize=20,
               xguidefont=font(30, "Times"),
               yguidefont=font(30, "Times"),
               #title=L"\textrm{SDOF with } \alpha = %$α",
               xtick=[0.0, 0.1, 0.2, 0.3], # , ytick=[-1.0, 0.0, 1.0],
               xlims=(0.0, 0.33), ylims=(0.0, 0.2, 0.4, 0.6, 0.8),
               bottom_margin=6mm, left_margin=6mm, right_margin=4mm, top_margin=3mm, size=(900, 600))

    #    fig = plot(xlab="Time (s)", ylab="Temperature at the center")

    # \delta = 10^{-5}
    plot!(fig, sol[1:5:end], vars=(0, node), alpha=.1, lc=:blue, lab=L"\textrm{Flowpipe}")

    plot!(fig, tdom[1:350:end], sol₊_node[1:350:end], linestyle=:solid, marker=:utriangle, lab=L"\textrm{Backward Euler}", color=:red)
    plot!(fig, tdom[1:350:end], sol₋_node[1:350:end], lab="", marker=:utriangle, linestyle=:solid, color=:red)

    # Analytic solution
    sol₊_analytic = evolucion_temp_teo(node/100., 1.1, tdom)
    sol₋_analytic = evolucion_temp_teo(node/100., 0.9, tdom)

    plot!(fig, tdom, sol₊_analytic, c=:magenta, linestyle = :solid, lab=L"\textrm{Analytic (PDE)}", lw=1.2)
    plot!(fig, tdom, sol₋_analytic, c=:magenta, linestyle = :solid, lab="")

    #xtick=([0.010, 0.015, 0.020, 0.025, 0.030], [L"0.010", L"0.015", L"0.020", L"0.025", L"0.030"])
    lens!(fig, [0.01, 0.03], [0.6, 0.85], inset = (1, bbox(0.58, 0.33, 0.4, 0.4)),
           subplot=2,
           tickfont=font(18, "Times"),
           xticks=([0.01, 0.020, 0.030], [L"0.01", L"0.02", L"0.03"]),
           yticks=([0.60, 0.65, 0.70, 0.75, 0.80, 0.85], ["", L"0.65", "", L"0.75", "", L"0.85"]))
    fig
end

fig = plot_node_vs_t()
savefig(fig, joinpath(TARGET_FOLDER, "temp_center.pdf"))

# temperature profile for different instants of time
function plot_espacial()

    fig = plot(xlab=L"\textrm{Position}", ylab=L"\textrm{Temperature}",
               legendfontsize=18, legend=:top,
               tickfont=font(20, "Times"), guidefontsize=20,
               xguidefont=font(30, "Times"),
               yguidefont=font(30, "Times"),
               ylims=(-0.05, 1.25),
               xlims=(-0.02, 1.02),
               #title=L"\textrm{SDOF with } \alpha = %$α",
               xtick=[0.0, 0.25, 0.50, 0.75, 1.0], # , ytick=[-1.0, 0.0, 1.0],
               ytick=[0.0, 0.25, 0.50, 0.75, 1.0, 1.25],
               bottom_margin=6mm, left_margin=6mm, right_margin=4mm, top_margin=3mm, size=(900, 600))

    paso = 5
    ancho_linea = 0.75
    col_linea = :blue
    marker_col = :blue

    tiempo_perfil = 0.0
    col = :blue

    fp = flowpipe_snapshot(sol, tiempo_perfil)
    xs, Tteo₋ = perfil_temp_teo(101, 0.9, tiempo_perfil)
    xs, Tteo₊ = perfil_temp_teo(101, 1.1, tiempo_perfil)
    plot!(fig, fp, lw=.1, c=col)
    plot!(fig, xs, Tteo₊, lc=col_linea, lw=ancho_linea, linestyle=:solid, lab="")
    plot!(fig, xs[1:paso:end], Tteo₊[1:paso:end], seriestype=:scatter, marker=:utriangle, c=marker_col, lab=L"t=0")
    plot!(fig, xs, Tteo₋, lc=col_linea, lw=ancho_linea, linestyle=:solid, lab="")
    plot!(fig, xs[1:paso:end], Tteo₋[1:paso:end], seriestype=:scatter, marker=:utriangle, c=marker_col, lab="")

    tiempo_perfil = 0.03
    col = :green
    col_linea = :green
    marker_col = :green

    fp = flowpipe_snapshot(sol, tiempo_perfil)
    xs, Tteo₋ = perfil_temp_teo(101, 0.9, tiempo_perfil)
    xs, Tteo₊ = perfil_temp_teo(101, 1.1, tiempo_perfil)
    plot!(fig, fp, lw=.1, c=col)
    plot!(fig, xs, Tteo₊, lc=col_linea, lw=ancho_linea, linestyle=:solid, lab="")
    plot!(fig, xs[1:paso:end], Tteo₊[1:paso:end], marker=:utriangle, c=marker_col, seriestype=:scatter, lab=L"t=0.03")
    plot!(fig, xs, Tteo₋, lc=col_linea, lw=ancho_linea, linestyle=:solid, lab="")
    plot!(fig, xs[1:paso:end], Tteo₋[1:paso:end], marker=:utriangle, c=marker_col, seriestype=:scatter, lab="")

    tiempo_perfil = 0.1
    col = :yellow
    col_linea = :orange
    marker_col = :orange

    fp = flowpipe_snapshot(sol, tiempo_perfil)
    xs, Tteo₋ = perfil_temp_teo(101, 0.9, tiempo_perfil)
    xs, Tteo₊ = perfil_temp_teo(101, 1.1, tiempo_perfil)
    plot!(fig, fp, lw=.1, c=col)
    plot!(fig, xs, Tteo₊, lc=col_linea, lw=ancho_linea, linestyle=:solid, lab="")
    plot!(fig, xs[1:paso:end], Tteo₊[1:paso:end], marker=:utriangle, c=marker_col, seriestype=:scatter, lab=L"t=0.1")
    plot!(fig, xs, Tteo₋, lc=col_linea, lw=ancho_linea, linestyle=:solid, lab="")
    plot!(fig, xs[1:paso:end], Tteo₋[1:paso:end], marker=:utriangle, c=marker_col, seriestype=:scatter, lw=0.0, lab="")

    return fig

end

fig = plot_espacial()
savefig(fig, joinpath(TARGET_FOLDER, "temp_barra.pdf"))

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:
mutable struct OptimizerStats
    detailed_stats::Bool
    objective_value::Float64
    termination_status::Int
    primal_status::Int
    dual_status::Int
    solver_solve_time::Float64
    result_count::Int
    has_values::Bool
    has_duals::Bool
    # Candidate solution
    objective_bound::Union{Missing, Float64}
    # Use missing instead of nothing so that CSV writting doesn't fail
    dual_objective_value::Union{Missing, Float64}
    # Work counters
    solve_time::Float64
    barrier_iterations::Union{Missing, Int}
    simplex_iterations::Union{Missing, Int}
    node_count::Union{Missing, Int}
    timed_solve_time::Float64
    timed_calculate_aux_variables::Float64
    timed_calculate_dual_variables::Float64
    solve_bytes_alloc::Union{Missing, Float64}
    sec_in_gc::Union{Missing, Float64}
end

function OptimizerStats()
    return OptimizerStats(
        false,
        NaN,
        -1,
        -1,
        -1,
        NaN,
        -1,
        false,
        false,
        missing,
        missing,
        NaN,
        missing,
        missing,
        missing,
        NaN,
        0,
        0,
        missing,
        missing,
    )
end

"""
Construct OptimizerStats from a vector that was serialized to HDF5.
"""
function OptimizerStats(data::Vector{Float64})
    vals = Vector(undef, length(data))
    to_missing = Set((
        :objective_bound,
        :dual_objective_value,
        :barrier_iterations,
        :simplex_iterations,
        :node_count,
        :solve_bytes_alloc,
        :sec_in_gc,
    ))
    for (i, name) in enumerate(fieldnames(OptimizerStats))
        if name in to_missing && isnan(data[i])
            vals[i] = missing
        else
            vals[i] = data[i]
        end
    end
    return OptimizerStats(vals...)
end

"""
Convert OptimizerStats to a matrix of floats that can be serialized to HDF5.
"""
function to_matrix(stats::T) where {T <: OptimizerStats}
    field_values = Matrix{Float64}(undef, fieldcount(T), 1)
    for (ix, field) in enumerate(fieldnames(T))
        value = getfield(stats, field)
        field_values[ix] = ismissing(value) ? NaN : value
    end
    return field_values
end

function to_dataframe(stats::OptimizerStats)
    df = DataFrames.DataFrame([to_namedtuple(stats)])
    return df
end

function to_dict(stats::OptimizerStats)
    data = Dict()
    for field in fieldnames(typeof(stats))
        data[String(field)] = getfield(stats, field)
    end

    return data
end

function get_column_names(::Type{OptimizerStats})
    return collect(string.(fieldnames(OptimizerStats)))
end


text:
using VectorizationBase: OffsetPrecalc, StaticBool, Bit, static, NativeTypes, Index, gep_quote, VectorIndex,
    AbstractMask, NativeTypesExceptBit, AbstractSIMDVector, IndexNoUnroll, AbstractStridedPointer, AbstractSIMD
using VectorizationBase: contiguous_batch_size, contiguous_axis, val_stride_rank, bytestrides, offsets, memory_reference,
    vmaximum, fmap, FloatingTypes, IntegerIndex, LazyMulAdd

LoopVectorization.check_args(::Type{T}, ::Type{T}) where T<:Tropical = true
LoopVectorization.check_type(::Type{Tropical{T}}) where {T} = LoopVectorization.check_type(T)

@inline Base.FastMath.add_fast(a::Tropical, b::Tropical) = Tropical(Base.FastMath.max_fast(content(a), content(b)))

@inline function VectorizationBase._vstore!(
    ptr::AbstractStridedPointer, vu::Tropical{<:VecUnroll{Nm1,W}}, u::Unroll{AU,F,N,AV,W}, a::A, s::S, nt::NT, si::StaticInt{RS}
) where {A<:StaticBool,S<:StaticBool,NT<:StaticBool,RS,AU,F,N,AV,W,Nm1}
    VectorizationBase._vstore!(notropical(ptr), content(vu), u, a, s, nt, si)
end

@inline function VectorizationBase._vstore!(ptr::AbstractStridedPointer, vu::Tropical{<:VecUnroll{Nm1,W}}, u::Unroll{AU,F,N,AV,W}, m::AbstractMask{W}, a::A, s::S, nt::NT, si::StaticInt{RS}
) where {A<:StaticBool,S<:StaticBool,NT<:StaticBool,RS,AU,F,N,AV,W,Nm1}
    VectorizationBase._vstore!(notropical(ptr), content(vu), u, m, a, s, nt, si)
end

@inline function VectorizationBase._vstore!(
    g::G, ptr::AbstractStridedPointer{T,D,C}, vu::Tropical{<:VecUnroll{U,W}}, u::Unroll{AU,F,N,AV,1,M,X,I}, a::A, s::S, nt::NT, si::StaticInt{RS}
) where {T,D,C,U,AU,F,N,W,M,I,G<:Function,AV,A<:StaticBool, S<:StaticBool, NT<:StaticBool, RS,X}
    VectorizationBase._vstore!(g, notropical(ptr), content(vu), u, a, s, nt, si)
end
@inline function VectorizationBase.__vstore!(
    f::F, ptr::Ptr{Tropical{T}}, v::Tropical{T}, i::IntegerIndex, a::A, s::S, nt::NT, si::StaticInt{RS}
) where {T<:NativeTypesExceptBit, F<:Function,A<:StaticBool,S<:StaticBool,NT<:StaticBool,RS}
    VectorizationBase.__vstore!(f, Ptr{T}(ptr), content(v), i, a, s, nt, si)
end
@inline function VectorizationBase.__vstore!(
        ptr::Ptr{Tropical{T}}, v::Tropical{VT}, i::I, m::AbstractMask{W}, a::A, s::S, nt::NT, si::StaticInt{RS}
    ) where {W, T <: NativeTypesExceptBit, VT <: Vec, I <: Index, A <: StaticBool, S <: StaticBool, NT <: StaticBool, RS}
    VectorizationBase.__vstore!(Ptr{T}(ptr), content(v), i, m, a, s, nt, si)
end
@inline function VectorizationBase.__vstore!(
    ptr::Ptr{Tropical{T}}, v::Tropical{VT}, i::VectorIndex{W}, a::A, s::S, nt::NT, si::StaticInt{RS}
) where {T,VT<:Vec,W,S<:StaticBool,A<:StaticBool,NT<:StaticBool,RS}
    VectorizationBase.__vstore!(Ptr{T}(ptr), content(v), i, a, s, nt, si)
end

@inline function VectorizationBase.__vload(ptr::Ptr{Tropical{T}}, i::I, m::AbstractMask, a::A, si::StaticInt{RS})  where {A <: StaticBool, T <: NativeTypes, I <: Index, RS}
    Tropical(VectorizationBase.__vload(Ptr{T}(ptr), i, m, a, si))
end
@inline function VectorizationBase.__vload(ptr::Ptr{Tropical{T}}, i::I, a::A, si::StaticInt{RS}) where {A <: StaticBool, T <: NativeTypes, I <: Index, RS}
    Tropical(VectorizationBase.__vload(Ptr{T}(ptr), i, a, si))
end

@inline function VectorizationBase._vbroadcast(a::Union{Val{W},StaticInt{W}}, s::Tropical{T}, si::StaticInt{RS}) where {W,T,RS}
    Tropical(VectorizationBase._vbroadcast(a, content(s), si))
end

@inline function notropical(ptr::VectorizationBase.StridedPointer{Tropical{T},N,C,B,R,X,O}) where {T,N,C,B,R,X,O}
    stridedpointer(Ptr{T}(ptr.p), ptr.si, StaticInt{B}())
end

@inline function notropical(ptr::OffsetPrecalc{<:Tropical})
    VectorizationBase.OffsetPrecalc(notropical(ptr.ptr), ptr.precalc)
end

@inline function VectorizationBase._vload(ptr::AbstractStridedPointer{Tropical{T}}, u::Unroll, ::A, ::StaticInt{RS}) where {T,A<:StaticBool,RS}
    res = VectorizationBase._vload(notropical(ptr), u, A(), StaticInt{RS}())
    Tropical(res)
end

@inline function VectorizationBase._vload(ptr::AbstractStridedPointer{Tropical{T}}, u::Unroll, m::AbstractMask, ::A, ::StaticInt{RS}) where {T,A<:StaticBool,RS}
    res = VectorizationBase._vload(notropical(ptr), u, m, A(), StaticInt{RS}())
    Tropical(res)
end

@generated function VectorizationBase.zero_vecunroll(::StaticInt{N}, ::StaticInt{W}, ::Type{Tropical{T}}, ::StaticInt{RS}) where {N,W,T,RS}
    quote
        $(Expr(:meta,:inline))
        t = Base.Cartesian.@ntuple $N n -> VectorizationBase._vbroadcast(StaticInt{$W}(), $(zero(Tropical{T}).n), StaticInt{$RS}())
        Tropical(VecUnroll(t))
    end
end

@inline function Base.promote_rule(::Type{Tropical{T1}}, ::Type{Tropical{T2}}) where {T1<:VecUnroll,T2<:Vec}
    Tropical{promote_rule(T1, T2)}
end

@inline function VectorizationBase._vzero(::StaticInt{W}, ::Type{T}, ::StaticInt{RS}) where {W,FT,T<:Tropical{FT},RS}
    Tropical(VectorizationBase._vbroadcast(StaticInt{W}(), zero(T).n, StaticInt{RS}()))
end

# `gep` is a shorthand for "get element pointer"
@inline function VectorizationBase._gep(ptr::Ptr{Tropical{T}}, ::StaticInt{N}, ::StaticInt{RS}) where {N, T <: NativeTypes, RS}
    Ptr{Tropical{T}}(VectorizationBase._gep(Ptr{T}(ptr), StaticInt{N}(), StaticInt{RS}()))
end
@inline function VectorizationBase._gep(ptr::Ptr{Tropical{T}}, i::I, ::StaticInt{RS}) where {I <: IntegerIndex, T <: NativeTypes, RS}
    Ptr{Tropical{T}}(VectorizationBase._gep(Ptr{T}(ptr), i, StaticInt{RS}()))
end
@inline function VectorizationBase._gep(ptr::Ptr{Tropical{T}}, i::LazyMulAdd{M,O,I}, ::StaticInt{RS}) where {T <: NativeTypes, I <: Integer, O, M, RS}
    Ptr{Tropical{T}}(VectorizationBase._gep(Ptr{T}(ptr), i, StaticInt{RS}()))
end

for TP in [:NativeTypes, :AbstractSIMD]
    @eval @inline function Base.fma(x::Tropical{V}, y::Tropical{V}, z::Tropical{V}) where {V<:$TP}
        Tropical(Base.FastMath.max_fast(content(z), Base.FastMath.add_fast(content(x), content(y))))
    end
    @eval @inline function Base.fma(::StaticInt{N}, y::Tropical{V}, z::Tropical{V}) where {N,V<:$TP}
        Base.FastMath.add_fast(Base.FastMath.mul_fast(StaticInt{N}(), y), z)
    end

    for f ∈ [:(Base.:(*)), :(Base.FastMath.mul_fast)]
        @eval begin
            @inline $f(::StaticInt{0}, vx::Tropical{T}) where {T<:$TP} = zero(Tropical{T})
            @inline $f(::StaticInt{1}, vx::Tropical{T}) where {T<:$TP} = vx
            @inline $f(vx::Tropical{T}, ::StaticInt{0}) where {T<:$TP} = zero(Tropical{T})
            @inline $f(vx::Tropical{T}, ::StaticInt{1}) where {T<:$TP} = vx
        end
    end
    for f ∈ [:(Base.:(+)), :(Base.FastMath.add_fast)]
        @eval begin
            @inline $f(::StaticInt{0}, vx::Tropical{T}) where {T<:$TP} = vx
            @inline $f(vx::Tropical{T}, ::StaticInt{0}) where {T<:$TP} = vx
        end
    end
end
# julia 1.5 patch
@inline function VectorizationBase.VecUnroll(data::Tuple{T,Vararg{T,N}}) where {N,T<:Tropical}
    Tropical(VecUnroll(map(content, data)))
end

@inline LoopVectorization.vecmemaybe(x::Tropical) = x
@inline function VectorizationBase.collapse_add(vu::Tropical{VecUnroll{N,W,T,V}}) where {N,W,T,V}
    Tropical(VectorizationBase.collapse_max(content(vu)))
end
@inline function VectorizationBase.contract_add(vu::Tropical{VecUnroll{N,W,T,V}}, ::StaticInt{K}) where {N,W,T,V,K}
    Tropical(VectorizationBase.contract_max(content(vu), StaticInt{K}()))
end
@inline function VectorizationBase.reduced_add(x::Tropical, y::Tropical)
    Tropical(VectorizationBase.reduced_max(content(x), content(y)))
end

@inline function VectorizationBase.ifelse(f::F, m::AbstractMask, v1::Tropical, v2::Tropical, v3::Tropical) where {F}
    Tropical(VectorizationBase.ifelse(m, content(f(v1, v2, v3)), content(v3)))
end

@inline VectorizationBase.vsum(x::Tropical{<:AbstractSIMD}) = Tropical(VectorizationBase.vmaximum(content(x)))

# Overwrite the `mul!` in LinearAlgebra (also changes the behavior of `*` in Base)!
using Octavian
const XTranspose{T} = Transpose{T, <:AbstractVecOrMat{T}}
for TA in [:AbstractMatrix, :XTranspose]
    for TB in [:AbstractMatrix, :XTranspose]
        @eval function LinearAlgebra.mul!(o::AbstractMatrix{T}, a::$TA{T}, b::$TB{T}, α::Number, β::Number) where {T<:Tropical{<:NativeTypes}}
            α = _convert_to_static(T, α)
            β = _convert_to_static(T, β)
            Octavian.matmul!(o, a, b, α, β)
        end
    end
end
# NOTE: benchmark shows, the type instability here can be optimized by the compiler
# so you do not need to worry about the overheads.
@inline _convert_to_static(::Type{T}, α::TropicalTypes) where T<:TropicalTypes = α
@inline function _convert_to_static(::Type{T}, α::Number) where T<:TropicalTypes
    if iszero(α)
        return StaticInt{0}()
    elseif isone(α)
        return StaticInt{1}()
    else
        throw(ArgumentError("$α is not a valid tropical number."))
    end
end

using Octavian: zstridedpointer, preserve_buffer, matmul_sizes, matmul_params, dontpack, matmul_st_pack_dispatcher!,
    loopmul!, inlineloopmul!, maybeinline, matmul_only_β!, One, Zero, ArrayInterface, block_sizes, __matmul!

# a patch to allow tropical types
@inline function Octavian._matmul!(C::AbstractMatrix{T}, A, B, α, β, nthread, MKN) where {T<:Tropical{<:NativeTypes}}
    M, K, N = MKN === nothing ? matmul_sizes(C, A, B) : MKN
    if M * N == 0
        return
    elseif K == 0
        matmul_only_β!(C, β)
        return
    end
    W = pick_vector_width(T)
    pA = zstridedpointer(A); pB = zstridedpointer(B); pC = zstridedpointer(C);
    Cb = preserve_buffer(C); Ab = preserve_buffer(A); Bb = preserve_buffer(B);
    mᵣ, nᵣ = matmul_params(Val(T))
    GC.@preserve Cb Ab Bb begin
        if maybeinline(M, N, T, ArrayInterface.is_column_major(A)) # check MUST be compile-time resolvable
            inlineloopmul!(pC, pA, pB, One(), Zero(), M, K, N)
            return
        else
            (nᵣ ≥ N) && @goto LOOPMUL
            if (Sys.ARCH === :x86_64) || (Sys.ARCH === :i686)
                (M*K*N < (StaticInt{4_096}() * W)) && @goto LOOPMUL
            else
                (M*K*N < (StaticInt{32_000}() * W)) && @goto LOOPMUL
            end
            __matmul!(pC, pA, pB, α, β, M, K, N, nthread)
            return
            @label LOOPMUL
            loopmul!(pC, pA, pB, α, β, M, K, N)
            return
        end
    end
end

@inline function Octavian._matmul_serial!(
    C::AbstractMatrix{T}, A::AbstractMatrix, B::AbstractMatrix, α, β, MKN
) where {T<:Tropical}
    M, K, N = MKN === nothing ? matmul_sizes(C, A, B) : MKN
    if M * N == 0
        return
    elseif K == 0
        matmul_only_β!(C, β)
        return
    end
    pA = zstridedpointer(A); pB = zstridedpointer(B); pC = zstridedpointer(C);
    Cb = preserve_buffer(C); Ab = preserve_buffer(A); Bb = preserve_buffer(B);
    Mc, Kc, Nc = block_sizes(Val(T)); mᵣ, nᵣ = matmul_params(Val(T));
    GC.@preserve Cb Ab Bb begin
        if maybeinline(M, N, T, ArrayInterface.is_column_major(A)) # check MUST be compile-time resolvable
            inlineloopmul!(pC, pA, pB, One(), Zero(), M, K, N)
            return
        elseif (nᵣ ≥ N) || dontpack(pA, M, K, Mc, Kc, T)
            loopmul!(pC, pA, pB, α, β, M, K, N)
            return
        else
            matmul_st_pack_dispatcher!(pC, pA, pB, α, β, M, K, N)
            return
        end
    end
end # function

function Octavian._matmul!(y::AbstractVector{T}, A::AbstractMatrix, x::AbstractVector, α, β, MKN, contig_axis) where {T<:Tropical}
  @tturbo for m ∈ indices((A,y),1)
    yₘ = zero(T)
    for n ∈ indices((A,x),(2,1))
      yₘ += A[m,n]*x[n]
    end
    y[m] = α * yₘ + β * y[m]
  end
  return y
end
function Octavian._matmul_serial!(y::AbstractVector{T}, A::AbstractMatrix, x::AbstractVector, α, β, MKN) where {T<:Tropical}
  @turbo for m ∈ indices((A,y),1)
    yₘ = zero(T)
    for n ∈ indices((A,x),(2,1))
      yₘ += A[m,n]*x[n]
    end
    y[m] = α * yₘ + β * y[m]
  end
  return y
end

Octavian.matmul_params(::Val{T}) where {T <: Tropical} = LoopVectorization.matmul_params()
@inline Octavian.incrementp(A::AbstractStridedPointer{<:Tropical,3}, a::Ptr) = VectorizationBase.increment_ptr(A, a, (Zero(), Zero(), One()))
@inline Octavian.increment2(B::AbstractStridedPointer{<:Tropical,2}, b::Ptr, ::StaticInt{nᵣ}) where {nᵣ} = VectorizationBase.increment_ptr(B, b, (Zero(), StaticInt{nᵣ}()))
@inline Octavian.increment1(C::AbstractStridedPointer{<:Tropical,2}, c::Ptr, ::StaticInt{mᵣW}) where {mᵣW} = VectorizationBase.increment_ptr(C, c, (StaticInt{mᵣW}(), Zero()))


text:
mutable struct Region
    style::Int32
    triclinic::Int32
    interior::Int32                     # 1 for interior, 0 for exterior
    scaleflag::Int32                    # 1 for lattice, 0 for box
    bboxflag::Int32                # 1 if bounding box is computable
    varshape::Int32                # 1 if region shape changes over time
    dynamic::Int32                 # 1 if position/orient changes over time
    moveflag::Int32
    rotateflag::Int32    # 1 if position/orientation changes
    openflag::Int32                # 1 if any face is open

    r::Vector{Float64}                   # distance between particle & surf, r > 0.0
    delx::Vector{Float64}
    dely::Vector{Float64}
    delz::Vector{Float64}    # vector from surface pt to particle
    radius::Vector{Float64}              # curvature of region at contact point
    iwall::Vector{Int32}                # unique id of wall for storing shear history
    openfaces::Vector{Int32};#[6];           # flags for which faces are open
    tri::Vector{Int32}#[12*3];    # 3 corner pts of 12 triangles (2 per face)
    theta::Float64 # orientation
    rprev::Float64 # speed of time-dependent radius, if applicable

    face::Array{Float64,2}#[6*3];   # unit normals of 6 box faces
    corners::Array{Float64,2}#[8*3];                     # 8 corner points
    faces::Array{Float64,2};#[6*4*3];   # 4 corner pts of 6 prism faces
    boxhi::Vector{Float64}#[3];   # orthogonal box global bounds
    boxlo::Vector{Float64}#[3];   # orthogonal box global bounds
    boxtilt::Vector{Float64}#[3]; # 3 tilt factors
    clo::Vector{Float64}#[3]; # opposite corners of prism
    chi::Vector{Float64}#[3]; # opposite corners of prism
    scale::Vector{Float64}#[3];
    extent_lo::Vector{Float64}#[3];
    extent_hi::Vector{Float64}#[3];
    a::Vector{Float64}#[3];  # edge vectors of region
    b::Vector{Float64}#[3];  # edge vectors of region
    c::Vector{Float64}#[3];  # edge vectors of region
    h::Vector{Float64}#[6];
    h_inv::Vector{Float64}#[6];
    dx::Vector{Float64}#[3];    # current displacement
    v::Vector{Float64}#[3];                 # translational velocity
    rpoint::Vector{Float64}#[3];            # current origin of rotation axis
    omega::Vector{Float64}#[3];             # angular velocity
    xcenter::Vector{Float64}#[3];    # translated/rotated center of cylinder/sphere (only used if varshape)
    prev::Vector{Float64}#[5];       # stores displacement (X3), angle and if

    Region() = new()
end

function initregion(boxhi, boxtilt=[0.0; 0.0; 0.0])

    reg = Region();

    reg.boxlo = [0.0; 0.0; 0.0];
    reg.boxhi = boxhi;
    reg.boxtilt = boxtilt;

    if (reg.boxtilt[1] == 0.0 && reg.boxtilt[2] == 0.0 && reg.boxtilt[3] == 0.0)
        reg.triclinic = 0;
    else
        reg.triclinic = 1;
    end

    return reg
end

function printout(reg::Region)
    display("boxlo: "); display(reg.boxlo);
    display("boxhi: "); display(reg.boxhi);
    display("boxtilt: "); display(reg.boxtilt);
    display("clo: "); display(reg.clo);
    display("chi: "); display(reg.chi);
    display("extent_lo: "); display(reg.extent_lo);
    display("extent_hi: "); display(reg.extent_hi);
    display("a: "); display(reg.a);
    display("b: "); display(reg.b);
    display("c: "); display(reg.c);
    display("scale: "); display(reg.scale);
    display("h: "); display(reg.h);
    display("h_inv: "); display(reg.h_inv);
    display("dx: "); display(reg.dx);
    display("v: "); display(reg.v);
    display("rpoint: "); display(reg.rpoint);
    display("omega: "); display(reg.omega);
    display("xcenter: "); display(reg.xcenter);
    display("prev: "); display(reg.prev);
end

max_stars_repo_name: cesmix-mit/MDP

max_issues_repo_name: cesmix-mit/MDP

max_forks_repo_name: cesmix-mit/MDP


text:
using QuasiGeostrophy
using LinearAlgebra, FFTW, BenchmarkTools, Plots

a = 0; b = 22.0
Δt = 1e-1 / 2
Ω = S¹(a,b) # Domain
Nx = 2^6;
fourier_grid = FourierGrid(Nx, Ω) # Grid
x = fourier_grid.grid[1]
kx = fourier_grid.wavenumbers[1]
f = @. sin( 4 * 2π * x / (b-a)) + 0im
f̂ = copy(f)
FFTW.set_num_threads(Threads.nthreads())
P = plan_fft(f) # typeof(iP) <: AbstractFFTs.ScaledPlan
iP = plan_ifft(f)
mul!(f̂, P, f)
∂x = FourierOperator(im .* kx)
argmax(real.((-∂x^2 - ∂x^4).op))

function closure_rhs(∂x, P, iP)
    function closed_rhs(state)
        f̂ = state[1]
        ff = convolve(f̂, f̂, P, iP)
        return [ ∂x( -0.5 .* ff - ∂x(f̂))]
    end
end

ks_rhs = closure_rhs(∂x, P, iP)

function closure_lhs(∂x, Δt)
    function lhs()
        q1_lhs = 1 + Δt * ∂x^4
        return [q1_lhs]
    end
end

ks_lhs = closure_lhs(∂x, Δt)
operators = ks_lhs()
lhs_operator = [inv(operator) for operator in operators]

function imex_step!(state, lhs_operator, rhs, Δt)
    rhs_state = rhs(state)
    current_state = copy(state)
    for i in eachindex(state)
        tmp = lhs_operator[i](current_state[i] .+ Δt .* rhs_state[i])
        @. state[i] =  tmp
    end
end
##
f = @. sin( 8π * x / (b-a)) + 0im
f̂ = copy(f)
mul!(f̂, P, f)
state = [f̂]
nsteps = 40*200
total_array = zeros(length(f), nsteps)
##
gr(size = (400,400))
for i in 1:nsteps
    imex_step!(state, lhs_operator, ks_rhs, Δt)
    u = state[1]
    total_array[:, i] .= real.(iP * u)
    if (i%100) == 0
        mul!(f, iP, f̂)
        p1 = plot(x, real.(f), ylims = (-2.0,2.0), label = "real space", xlabel = "position", ylabel  = "amplitude")
        p2 = scatter(log.(abs.(f̂))[1:div(length(f̂),2)+1], label = "spectrum", ylims = (-7, 3), xlabel = "wavenumber", ylabel = "spectral amplitude")
        display(plot(p1,p2))
        sleep(0.1)
    end
end
##
contourf(total_array, linewidth = 0, color = :thermometer)

max_stars_repo_name: upiterbarg/QuasiGeostrophy

max_issues_repo_name: sandreza/QuasiGeostrophy

max_forks_repo_name: upiterbarg/QuasiGeostrophy


text:
using VectorSphericalWaves
n_max = 40
functions_tobe_tested = [VectorSphericalWaves.γ_mn, VectorSphericalWaves.γ_mn_dash]
for f in functions_tobe_tested
    println("now testing <<$f>>")
    for n = 1:n_max
        for m = -n:n
            try
                f(m, n);
            catch e
                println("here is the arguments that caused error: n=$n, m=$m")
                println(e)
                break
            end
        end
    end
end