text:
# Projector function for TT. Goes into the mode-coupling matrix.
function Ξ_TT(𝐖::SpectralVector{T, AA},
              w3j²₀₀::WignerSymbolVector{T, Int},
              ℓ₁::Int, ℓ₂::Int) where {T, AA}
    Ξ = zero(T)
    ℓ₃_start = max(firstindex(w3j²₀₀), firstindex(𝐖))
    ℓ₃_end = min(lastindex(w3j²₀₀), lastindex(𝐖))
    @inbounds @simd for ℓ₃ ∈ ℓ₃_start:ℓ₃_end
        Ξ += (2ℓ₃ + 1) * w3j²₀₀[ℓ₃] * 𝐖[ℓ₃]
    end
    return Ξ / (4π)
end

# Projector function for EE. Goes into the mode-coupling matrix.
# Note that w3j² refers to the square of ( ℓ ℓ₂ ℓ₃ 0 -2 2 )
function Ξ_EE(𝐖::SpectralVector{T, AA},
              w3j²₂₂::WignerSymbolVector{T, Int},
              ℓ₁::Int, ℓ₂::Int) where {T, AA}
    Ξ = zero(T)
    ℓ₃_start = max(firstindex(w3j²₂₂), firstindex(𝐖))
    ℓ₃_end = min(lastindex(w3j²₂₂), lastindex(𝐖))
    if isodd(ℓ₁ + ℓ₂ + ℓ₃_start)
        ℓ₃_start += 1
    end
    @inbounds @simd for ℓ₃ ∈ ℓ₃_start:2:ℓ₃_end
        Ξ += (2ℓ₃ + 1) * w3j²₂₂[ℓ₃] * 𝐖[ℓ₃]
    end
    return Ξ / (4π)
end

# Projector function for EE. Goes into the mode-coupling matrix.
# Note that w3j² refers to the square of ( ℓ ℓ₂ ℓ₃ 0 -2 2 )
function Ξ_EB(𝐖::SpectralVector{T, AA},
              w3j²₂₂::WignerSymbolVector{T, Int},
              ℓ₁::Int, ℓ₂::Int) where {T, AA}
    Ξ = zero(T)
    ℓ₃_start = max(firstindex(w3j²₂₂), firstindex(𝐖))
    ℓ₃_end = min(lastindex(w3j²₂₂), lastindex(𝐖))
    if iseven(ℓ₁ + ℓ₂ + ℓ₃_start)
        ℓ₃_start += 1
    end
    @inbounds @simd for ℓ₃ ∈ ℓ₃_start:2:ℓ₃_end
        Ξ += (2ℓ₃ + 1) * w3j²₂₂[ℓ₃] * 𝐖[ℓ₃]
    end
    return Ξ / (4π)
end

# Projector function for TE. Goes into the mode-coupling matrix.
# Note that w3j₀₀₂₂ refers to ( ℓ ℓ₂ ℓ₃ 0 0 0 ) × ( ℓ ℓ₂ ℓ₃ 0 -2 2 )
function Ξ_TE(𝐖::SpectralVector{T, AA},
              w3j₀₀₂₂::WignerSymbolVector{T, Int},
              ℓ₁::Int, ℓ₂::Int) where {T, AA}
    Ξ = zero(T)
    ℓ₃_start = max(firstindex(w3j₀₀₂₂), firstindex(𝐖))
    ℓ₃_end = min(lastindex(w3j₀₀₂₂), lastindex(𝐖))
    if isodd(ℓ₁ + ℓ₂ + ℓ₃_start)
        ℓ₃_start += 1
    end
    @inbounds @simd for ℓ₃ ∈ ℓ₃_start:2:ℓ₃_end
        Ξ += (2ℓ₃ + 1) * w3j₀₀₂₂[ℓ₃] * 𝐖[ℓ₃]
    end
    return Ξ / (4π)
end

# use a view of a memory buffer and fill with wigner 3j
function fill_3j!(buffer::Array{T,N}, ℓ₁, ℓ₂, m₁, m₂) where {T,N}
    w = WignerF(T, ℓ₁, ℓ₂, m₁, m₂)  # set up the wigner recurrence
    buffer_view = uview(buffer, 1:length(w.nₘᵢₙ:w.nₘₐₓ))  # preallocated buffer
    w3j = WignerSymbolVector(buffer_view, w.nₘᵢₙ:w.nₘₐₓ)
    wigner3j_f!(w, w3j)  # deposit symbols into buffer
    return w3j
end

# inner MCM loop TT
function inner_mcm⁰⁰!(𝐌::SpectralArray{T,2},
                      Vᵢⱼ::SpectralVector{T}) where {T}
    @assert axes(𝐌, 1) == axes(𝐌, 2)
    lmin, lmax = first(axes(𝐌, 1)), last(axes(𝐌, 1))
    thread_buffers = get_thread_buffers(T, 2lmax+1)

    @qthreads for ℓ₁ in lmin:lmax
        buffer = thread_buffers[Threads.threadid()]
        for ℓ₂ in ℓ₁:lmax
            w3j²₀₀ = fill_3j!(buffer, ℓ₁, ℓ₂, 0, 0)
            w3j²₀₀.symbols .= w3j²₀₀.symbols .^ 2  # square the symbols
            Ξ = Ξ_TT(Vᵢⱼ, w3j²₀₀, ℓ₁, ℓ₂)
            𝐌[ℓ₁, ℓ₂] = (2ℓ₂ + 1) * Ξ
            𝐌[ℓ₂, ℓ₁] = (2ℓ₁ + 1) * Ξ
        end
    end
    return 𝐌
end

# inner MCM loop TE and TB
function inner_mcm⁰²!(𝐌::SpectralArray{T,2}, Vᵢⱼ::SpectralVector{T}) where {T}
    @assert axes(𝐌, 1) == axes(𝐌, 2)
    lmin, lmax = first(axes(𝐌, 1)), last(axes(𝐌, 1))
    thread_buffers_0 = get_thread_buffers(T, 2lmax+1)
    thread_buffers_2 = get_thread_buffers(T, 2lmax+1)
    @qthreads for ℓ₁ in lmin:lmax
        tid = Threads.threadid()
        buffer0 = thread_buffers_0[tid]
        buffer2 = thread_buffers_2[tid]
        for ℓ₂ in ℓ₁:lmax
            w3j₀₀ = fill_3j!(buffer0, ℓ₁, ℓ₂, 0, 0)
            w3j₂₂ = fill_3j!(buffer2, ℓ₁, ℓ₂, -2, 2)
            w3j₀₀₂₂ = w3j₀₀
            w3j₀₀₂₂.symbols .*= w3j₂₂.symbols
            Ξ = Ξ_TE(Vᵢⱼ, w3j₀₀₂₂, ℓ₁, ℓ₂)
            𝐌[ℓ₁, ℓ₂] = (2ℓ₂ + 1) * Ξ
            𝐌[ℓ₂, ℓ₁] = (2ℓ₁ + 1) * Ξ
        end
    end
    return 𝐌
end

# inner MCM loop for spin 2, called "EE" in Planck notation
function inner_mcm⁺⁺!(𝐌::SpectralArray{T,2}, Vᵢⱼ::SpectralVector{T}) where {T}
    @assert axes(𝐌, 1) == axes(𝐌, 2)
    lmin, lmax = first(axes(𝐌, 1)), last(axes(𝐌, 1))
    thread_buffers = get_thread_buffers(T, 2lmax+1)

    @qthreads for ℓ₁ in lmin:lmax
        buffer = thread_buffers[Threads.threadid()]
        for ℓ₂ in ℓ₁:lmax
            w3j²₂₂ = fill_3j!(buffer, ℓ₁, ℓ₂, -2, 2)
            w3j²₂₂.symbols .= w3j²₂₂.symbols .^ 2  # square the symbols
            Ξ = Ξ_EE(Vᵢⱼ, w3j²₂₂, ℓ₁, ℓ₂)
            𝐌[ℓ₁, ℓ₂] = (2ℓ₂ + 1) * Ξ
            𝐌[ℓ₂, ℓ₁] = (2ℓ₁ + 1) * Ξ
        end
    end
    return 𝐌
end

# inner MCM loop for spin 2
function inner_mcm⁻⁻!(𝐌::SpectralArray{T,2}, Vᵢⱼ::SpectralVector{T}) where {T}
    @assert axes(𝐌, 1) == axes(𝐌, 2)
    lmin, lmax = first(axes(𝐌, 1)), last(axes(𝐌, 1))
    thread_buffers = get_thread_buffers(T, 2lmax+1)

    @qthreads for ℓ₁ in lmin:lmax
        buffer = thread_buffers[Threads.threadid()]
        for ℓ₂ in ℓ₁:lmax
            w3j²₂₂ = fill_3j!(buffer, ℓ₁, ℓ₂, -2, 2)
            w3j²₂₂.symbols .= w3j²₂₂.symbols .^ 2  # square the symbols
            Ξ = Ξ_EB(Vᵢⱼ, w3j²₂₂, ℓ₁, ℓ₂)
            𝐌[ℓ₁, ℓ₂] = (2ℓ₂ + 1) * Ξ
            𝐌[ℓ₂, ℓ₁] = (2ℓ₁ + 1) * Ξ
        end
    end
    return 𝐌
end

@doc raw"""
    mcm(spec::Symbol, alm₁::Alm{T}, alm₂::Alm{T}; lmax=nothing)

Compute the mode-coupling matrix. See the [Spectral Analysis](@ref)
section in the documentation for examples. These are used by applying the
linear solve operator `\` to a `SpectralArray{T,1}`.

Choices for `spec`:

* `:TT`, identical to `M⁰⁰`
* `:TE`, identical to `:ET`, `:TB`, `:BT`, `:M⁰²`, `:M²⁰`
* `:EE_BB`, returns coupling matrix for stacked EE and BB vectors
* `:EB_BE`, returns coupling matrix for stacked EB and BE vectors
* `:M⁺⁺`, sub-block of spin-2 mode-coupling matrices
* `:M⁻⁻`, sub-block of spin-2 mode-coupling matrices

# Arguments:
- `spec::Symbol`: cross-spectrum of the mode-coupling matrix
- `alm₁::Alm{T}`: first mask's spherical harmonic coefficients
- `alm₂::Alm{T}`: second mask's spherical harmonic coefficients

# Keywords
- `lmin=0`: minimum multiple for mode-coupling matrix
- `lmax=nothing`: maximum multipole for mode-coupling matrix

# Returns:
- the mode coupling matrix. for single symbols, this returns a
    `SpectralArray{T,2}`. if spec is `:EE_BB` or `:EB_BE`, returns a
    `BlockSpectralMatrix{T}` with 2×2 blocks.
"""
function mcm(spec::Symbol, alm₁::Alm{Complex{T}}, alm₂::Alm{Complex{T}};
             lmin=0, lmax=nothing) where T
    if isnothing(lmax)  # use alm lmax if an lmax is not specified
        lmax = min(alm₁.lmax, alm₂.lmax)
    end
    Vᵢⱼ = SpectralVector(alm2cl(alm₁, alm₂)[1:(lmax+1)])  # zero-indexed
    if spec ∈ (:TT, :M⁰⁰)
        𝐌 = spectralzeros(lmin:lmax, lmin:lmax)
        return inner_mcm⁰⁰!(𝐌, Vᵢⱼ)
    elseif spec ∈ (:TE, :ET, :TB, :BT, :M⁰², :M²⁰)
        𝐌 = spectralzeros(lmin:lmax, lmin:lmax)
        return inner_mcm⁰²!(𝐌, Vᵢⱼ)
    elseif spec == :M⁺⁺
        𝐌 = spectralzeros(lmin:lmax, lmin:lmax)
        return inner_mcm⁺⁺!(𝐌, Vᵢⱼ)
    elseif spec == :M⁻⁻
        𝐌 = spectralzeros(lmin:lmax, lmin:lmax)
        return inner_mcm⁻⁻!(𝐌, Vᵢⱼ)
    elseif spec == :EE_BB
        𝐌⁺⁺ = spectralzeros(lmin:lmax, lmin:lmax)
        𝐌⁻⁻ = spectralzeros(lmin:lmax, lmin:lmax)
        inner_mcm⁺⁺!(𝐌⁺⁺, Vᵢⱼ)
        inner_mcm⁻⁻!(𝐌⁻⁻, Vᵢⱼ)
        return [ 𝐌⁺⁺  𝐌⁻⁻;
                 𝐌⁻⁻  𝐌⁺⁺ ]
    elseif spec == :EB_BE
        𝐌⁺⁺ = spectralzeros(lmin:lmax, lmin:lmax)
        𝐌⁻⁻ = spectralzeros(lmin:lmax, lmin:lmax)
        inner_mcm⁺⁺!(𝐌⁺⁺, Vᵢⱼ)
        inner_mcm⁻⁻!(𝐌⁻⁻, Vᵢⱼ)
        return [ 𝐌⁺⁺   (-𝐌⁻⁻);
                (-𝐌⁻⁻)   𝐌⁺⁺ ]
    end
    throw(ArgumentError("$(spec) not a valid spectrum."))
end

function mcm(spec::Tuple{Symbol,Symbol}, alm₁::Alm{Complex{T}}, alm₂::Alm{Complex{T}};
             lmin=0, lmax=nothing) where T
    if isnothing(lmax)  # use alm lmax if an lmax is not specified
        lmax = min(alm₁.lmax, alm₂.lmax)
    end
    Vᵢⱼ = SpectralVector(alm2cl(alm₁, alm₂)[1:(lmax+1)])  # zero-indexed
    if spec == (:EE_BB, :EB_BE)
        𝐌⁺⁺ = spectralzeros(lmin:lmax, lmin:lmax)
        𝐌⁻⁻ = spectralzeros(lmin:lmax, lmin:lmax)
        inner_mcm⁺⁺!(𝐌⁺⁺, Vᵢⱼ)
        inner_mcm⁻⁻!(𝐌⁻⁻, Vᵢⱼ)
        EE_BB = [ 𝐌⁺⁺  𝐌⁻⁻;
                  𝐌⁻⁻  𝐌⁺⁺ ]
        EB_BE = [ 𝐌⁺⁺   (-𝐌⁻⁻);
                 (-𝐌⁻⁻)   𝐌⁺⁺ ]
        return EE_BB, EB_BE
    end
end

# convenience function
mcm(spec::Symbol, m₁::HealpixMap, m₂::HealpixMap; lmin=0, lmax=nothing) =
    mcm(spec, map2alm(m₁), map2alm(m₂); lmin=lmin, lmax=lmax)
function mcm(spec::Symbol, f₁::CovField, f₂::CovField; lmin=0, lmax=nothing)
    mask1 = (string(spec)[1] == 'T') ? f₁.maskT : f₁.maskP
    mask2 = (string(spec)[2] == 'T') ? f₂.maskT : f₂.maskP
    mcm(spec, map2alm(mask1), map2alm(mask2); lmin=lmin, lmax=lmax)
end

"""Scale a map."""
function scale!(m::HealpixMap, s::Number)
    m .*= s
end
function scale!(m::PolarizedHealpixMap, sT::Number, sP::Number)
    m.i .*= sT
    m.q .*= sP
    m.u .*= sP
end
scale!(m::PolarizedHealpixMap, s::Number) = scale!(m, s, s)

"""
    mask!(m::HealpixMap, mask)
    mask!(m::PolarizedHealpixMap, maskT, maskP)

Mask a map or polarized map in place.

# Arguments:
- `m::Union{HealpixMap,PolarizedHealpixMap}`: map or polarized map to mask
- `maskT::HealpixMap`: mask for first map's intensity
- `maskP::HealpixMap`: mask for first map's polarization
"""
function mask!(m::HealpixMap, mask)
    m .*= mask
    return m
end
function mask!(m::PolarizedHealpixMap, maskT, maskP)
    m.i .*= maskT
    m.q .*= maskP
    m.u .*= maskP
    return m
end
mask!(m::PolarizedHealpixMap, mask) = mask!(m, mask, mask)

"""
    master(map₁::PolarizedHealpixMap, maskT₁::HealpixMap, maskP₁::HealpixMap,
           map₂::PolarizedHealpixMap, maskT₂::HealpixMap, maskP₂::HealpixMap; already_masked=false)

Perform a mode-decoupling calculation for two polarized maps, along with masks to apply.
Returns spectra for ``TT``, ``TE``, ``ET``, ``EE``, ``EB``, ``BE``, and ``BB``.

# Arguments:
- `map₁::PolarizedHealpixMap`: the first IQU map
- `maskT₁::HealpixMap`: mask for first map's intensity
- `maskP₁::HealpixMap`: mask for first map's polarization
- `map₂::PolarizedHealpixMap`: the second IQU map
- `maskT₂::HealpixMap`: mask for second map's intensity
- `maskP₂::HealpixMap`: mask for second map's polarization

# Keywords
- `already_masked::Bool=false`: are the input maps already multiplied with the masks?
- `lmin::Int=0`: minimum multipole

# Returns:
- `Dict{Symbol,SpectralVector}`: spectra `Dict`, indexed with `:TT`, `:TE`, `:ET`, etc.
"""
function master(map₁::PolarizedHealpixMap, maskT₁::HealpixMap, maskP₁::HealpixMap,
                map₂::PolarizedHealpixMap, maskT₂::HealpixMap, maskP₂::HealpixMap;
                already_masked::Bool=false, lmin::Int=0)
    if already_masked
        maskedmap₁, maskedmap₂ = map₁, map₂
    else
        maskedmap₁ = deepcopy(map₁)
        maskedmap₂ = deepcopy(map₂)
        mask!(maskedmap₁, maskT₁, maskP₁)
        mask!(maskedmap₂, maskT₂, maskP₂)
    end
    return maskedalm2spectra(map2alm(maskedmap₁), map2alm(maskT₁), map2alm(maskP₁),
                             map2alm(maskedmap₂), map2alm(maskT₂), map2alm(maskP₂);
                             lmin=lmin)
end

"""Construct a NamedTuple with T,E,B names for the alms."""
function name_alms(alms::Vector)
    return (T=alms[1], E=alms[2], B=alms[3])
end

"""Compute spectra from alms of masked maps and alms of the masks themselves."""
function maskedalm2spectra(maskedmap₁vec::Vector{A}, maskT₁::A, maskP₁::A,
                           maskedmap₂vec::Vector{A}, maskT₂::A, maskP₂::A;
                           lmin=0) where {CT, A <: Alm{CT}}
    ## add TEB names
    maskedmap₁ = name_alms(maskedmap₁vec)
    maskedmap₂ = name_alms(maskedmap₂vec)
    spectra = Dict{Symbol, SpectralVector}()

    ## spectra that are independent
    for (X, Y) in ((:T,:T), (:T,:E), (:E,:T), (:T,:B), (:B,:T))
        spec = Symbol(X, Y)  # join X and Y

        ## select temp or pol mask
        maskX = (X == :T) ? maskT₁ : maskP₁
        maskY = (Y == :T) ? maskT₂ : maskP₂

        ## compute mcm
        M = mcm(spec, maskX, maskY; lmin=lmin)
        pCl = SpectralVector(alm2cl(maskedmap₁[X], maskedmap₂[Y]))[IdentityRange(lmin:end)]
        Cl = M \ pCl
        spectra[spec] = Cl  # store the result
    end

    M_EE_BB, M_EB_BE = mcm((:EE_BB, :EB_BE), maskP₁, maskP₂; lmin=lmin)
    ## EE and BB have to be decoupled together
    pCl_EE = SpectralVector(alm2cl(maskedmap₁[:E], maskedmap₂[:E]))[IdentityRange(lmin:end)]
    pCl_BB = SpectralVector(alm2cl(maskedmap₁[:B], maskedmap₂[:B]))[IdentityRange(lmin:end)]
    ## apply the 2×2 block mode-coupling matrix to the stacked EE and BB spectra
    @spectra Cl_EE, Cl_BB = M_EE_BB \ [pCl_EE; pCl_BB]
    spectra[:EE] = Cl_EE
    spectra[:BB] = Cl_BB

    ## EB and BE have to be decoupled together
    pCl_EB = SpectralVector(alm2cl(maskedmap₁[:E], maskedmap₂[:B]))[IdentityRange(lmin:end)]
    pCl_BE = SpectralVector(alm2cl(maskedmap₁[:B], maskedmap₂[:E]))[IdentityRange(lmin:end)]
    ## apply the 2×2 block mode-coupling matrix to the stacked EB and BE spectra
    @spectra Cl_EB, Cl_BE = M_EB_BE \ [pCl_EB; pCl_BE]
    spectra[:EB] = Cl_EB
    spectra[:BE] = Cl_BE

    return spectra
end

max_stars_repo_name: xzackli/PlanckCov

max_issues_repo_name: xzackli/PlanckCov


text:
## ----------  Work

function solve_work!(m :: Model, sol :: Solution)
    for t = workspan(m) : -1 : 1
        solve_work_period!(m, sol, t);
    end
end

function solve_work_period!(m :: Model, sol :: Solution, t :: Integer)
    set_solution!(sol, t, solve_work_period())
end

# stub +++++
function solve_work_period()
    kPrimeFct = nothing;
    cFct = nothing;
    return WorkSolution(kPrimeFct, cFct)
end

# function interpolate_kprime(kGridV :: AbstractVector{F}, kPrimeV :: AbstractVector{F}) where F
#     interp = LinearInterpolation(kGridV, kPrimeV);
#     return interp
# end

# -----------------

ext: jl


text:
function obtain_raw_data(base_dir::String, sha::String)
    data_path = joinpath(base_dir,"raw_input_data")
    ispath(data_path) && rm(data_path, recursive = true)
    mkpath(data_path)

    if Sys.iswindows()
        POWERSYSTEMSTESTDATA_URL = "https://github.com/GridMod/RTS-GMLC/archive/$(sha).zip"
    else
        POWERSYSTEMSTESTDATA_URL = "https://github.com/GridMod/RTS-GMLC/archive/$(sha).tar.gz"
    end

    tempfilename= Base.download(POWERSYSTEMSTESTDATA_URL)
    Sys.iswindows() && unzip_windows(tempfilename, base_dir)
    Sys.islinux() && unzip_unix(tempfilename, base_dir)
    Sys.isapple() && unzip_unix(tempfilename, base_dir)

    mv(joinpath(base_dir, "RTS-GMLC-$(sha)"), data_path, force=true)

    return data_path

end

function unzip_unix(filename, directory)
    @assert success(`tar -xvf $filename -C $directory`) "Unable to extract $filename to $directory"
end

function unzip_windows(filename, directory)
    home = (Base.VERSION < v"0.7-") ? JULIA_HOME : Sys.BINDIR
    @assert success(`$home/7z x $filename -y -o$directory`) "Unable to extract $filename to $directory"
end


text:
struct KoopmanProblem{X,Y,U,C,PR,B,TR,TS,P,O}
    x::X
    y::Y
    b::U
    inds::C
    prob::PR
    basis::B
    train::TR
    test::TS
    alg::P
    options::O
    eval_expression::Bool
end
struct KoopmanSolution{O,IN,S, E, F, A, P}
    k::O
    inds::IN
    sets::S
    error::E
    folds::F
    alg::A
    options::P
end

select_by(::Val, sol::KoopmanSolution) = begin
    @unpack k, error = sol
    i = argmin(error)
    return k[i], error[i]
end

select_by(::Val{:kfold}, sol::KoopmanSolution) = begin
    @unpack k, folds, error  = sol
    size(k, 1) <= 1 && return select_by(1, sol)
    i = argmin(mean(folds, dims = 1)[1,:])
    return k[i], error[i]
end

## Apply the problem to get the operator

# DMD-Like
function CommonSolve.init(prob::AbstractDataDrivenProblem{N,C,P}, alg::AbstractKoopmanAlgorithm, args...; kwargs...) where {N,C,P}
    # Build a basis
    s_x = size(prob.X,1)
    s_u = size(prob.U,1)

    x = [Symbolics.variable(:x, i) for i in 1:s_x]
    u = [Symbolics.variable(:u, i) for i in 1:s_u]
    t = Symbolics.variable(:t)

    b = Basis([x; u], x, controls = u, iv = t)

    init(prob, b, alg, args...; kwargs...)
end

# All (g(E))DMD like
function CommonSolve.init(prob::AbstractDiscreteProb{N,C}, b::AbstractBasis, alg::A, args...; B = [], eval_expression = false,  kwargs...) where {N,C, A <: AbstractKoopmanAlgorithm}
    @is_applicable prob

    @unpack X,p,t,U = prob

    x = b(X[:,1:end-1], p, t[1:end-1], U[:,1:end-1])
    y = b(X[:, 2:end], p, t[2:end], U[:, 2:end])

    if !isempty(controls(b))
        inds = .! is_dependent(map(eq->Num(eq.rhs),equations(b)), Num.(controls(b)))[1,:]
    else
        inds = ones(Bool, length(b))
    end

    options = DataDrivenCommonOptions(alg, N; kwargs...)

    @unpack sampler = options

    train, test = sampler(prob)

    return KoopmanProblem(
        x, y, B, inds, prob, b, train, test, alg, options, eval_expression
    )
end

function CommonSolve.init(prob::AbstracContProb{N,C}, b::AbstractBasis, alg::A, args...; B = [], eval_expression = false,  kwargs...) where {N,C, A <: AbstractKoopmanAlgorithm}
    @is_applicable prob

    @unpack DX,X,p,t,U = prob

    x = b(prob)

    y = similar(x)

    J = jacobian(b)

    if !isempty(U)
        # Find the indexes of the control states
        inds = .! is_dependent(map(eq->Num(eq.rhs),equations(b)), Num.(controls(b)))[1,:]
        for i in 1:length(prob)
           y[:, i] .= J(X[:, i], p, t[i], U[:, i])*DX[:, i]
        end
    else
        inds = ones(Bool, length(b))
        for i in 1:length(prob)
            y[:, i] .= J(X[:, i], p, t[i])*DX[:, i]
        end
    end

    options = DataDrivenCommonOptions(alg, N; kwargs...)

    @unpack sampler = options

    # Right now just ignore this
    #train , test = nothing, nothing
    train, test = sampler(prob)

    return KoopmanProblem(
        x, y, B, inds, prob, b, train, test, alg, options, eval_expression
    )
end

function derive_operator(alg, x, y, b, z, inds)
    if all(inds)
        K, B = alg(x,y)
        Q = y[inds, :]*x'
        P = x*x'
        C = z / y[inds, :]
    elseif isempty(b)
        K, B =  alg(x[inds, :], y[inds, :], x[.! inds, :])
        Q = y[inds, :]*x'
        P = x*x'
        C = z / y[inds, :]
    else
        K, B =  alg(x[inds, :], y[inds, :], x[.! inds, :], b)
        Q = y[inds, :]*x'
        P = x*x'
        C = z / y[inds, :]
    end
    return (K, B, C, P, Q,)
end

function operator_error(f, g)
    (x,y,K,B,C,P,Q,inds) -> begin
        k_ = Matrix(K)
        isempty(B) && return g(f(k_*x, C, y))
        return g(f(k_*x[inds, :]+B*x[.! inds, :], C, y))
    end
end

function CommonSolve.solve!(k::KoopmanProblem)

    @unpack x, y, b, inds, prob, basis, train, test, alg, options, eval_expression = k
    @unpack normalize, denoise, sampler, maxiter, abstol, reltol, verbose, progress,f,g,digits,kwargs = options

    z = get_target(prob)
    xₜ = x[:, test]
    zₜ = z[:, test]

    trainerror = zeros(eltype(z), length(train), length(train))
    testerror = zeros(eltype(z), length(train))
    fg = operator_error(f, g)
    ops = []

    for (i,t) in enumerate(train)

        op = derive_operator(alg, x[:, t], y[:, t], b, z[:, t], inds)

        push!(ops, op)

        testerror[i] = fg(xₜ, zₜ, op..., inds)

        for (j, tt) in enumerate(train)
            trainerror[i, j] = fg(x[:, tt], z[:, tt], op..., inds)
        end
    end

    sol = KoopmanSolution(ops, inds, (train,test),testerror, trainerror, alg, options)

    return DataDrivenSolution(prob, sol, basis, alg; eval_expression = eval_expression, digits = digits, kwargs...)
end


text:
using DirectDetections, Distributions, Plots

@named X = DirectDetections.Planet(
    Priors(
        a = TruncatedNormal(16, 3, 1, 30),
        e = Beta(1.2, 5),
        τ = Uniform(0., 1),
        ω = VonMises(0,1),
        i = VonMises(0,1),
        Ω = VonMises(0,1),
        H = Uniform(0, 100)
    ),
)

##
using DirectImages
cd(@__DIR__)
imagesc = map(centered, readfits("image-examples-1.fits",:))
pwd()
imshow2([
    imagesc[1]
    imagesc[2]
    imagesc[3]
    imagesc[4]
    imagesc[5]
], cmap=:magma, clims=(-1.0, 6.0))
##

system_images = DirectDetections.Images(
    (band=:H, image=imagesc[1], platescale=10.0, epoch=times[1], contrast=contrasts[1]),
    (band=:H, image=imagesc[2], platescale=10.0, epoch=times[2], contrast=contrasts[2]),
    (band=:H, image=imagesc[3], platescale=10.0, epoch=times[3], contrast=contrasts[3]),
    (band=:H, image=imagesc[4], platescale=10.0, epoch=times[4], contrast=contrasts[4]),
    (band=:H, image=imagesc[5], platescale=10.0, epoch=times[5], contrast=contrasts[5]),
)
##

@named HD82134 = System(
    Priors(
        μ = TruncatedNormal(2.0, 0.1, 0, Inf),
        plx = TruncatedNormal(45., 0.02, 0, Inf),
    ),
    system_images,
    X,
)

chain = DirectDetections.hmc(
    HD82134, .60,
    MCMCThreads(),
    num_chains=4,
    adaptation =  1_000,
    iterations = 10_000,
    tree_depth =     12,
)

##
@show snr = mean(chain["X[H]"]) / std(chain["X[H]"])
##
DirectDetections.MCMCChains.gelmandiag(chain)
##
plot(chain)
##
plotmodel(chain, alpha=0.05,  lims=1000)
##
savefig("../docs/src/assets/images-model-plot.svg")

##
using PairPlots
table = (;
a=         chain["X[a]"],
H=         chain["X[H]"],
e=         chain["X[e]"],
i=rad2deg.(chain["X[i]"]),
Ω=rad2deg.(chain["X[Ω]"]),
ω=rad2deg.(chain["X[ω]"]),
τ=         chain["X[τ]"],
)
labels=[
"a",
"H",
"e",
"i",
"\\Omega",
"\\omega",
"\\tau",
]
units = [
"(au)",
"(arb.)",
"",
"(\\degree)",
"(\\degree)",
"(\\degree)",
"",
]
corner(table, labels, units)
##
cd(@__DIR__)
savefig("../docs/src/assets/images-corner-plot.svg")

## This block generates the images we used for the example
using Random
using DirectOrbits
Random.seed!(4321)

truth = (;
    a = 16,
    τ = 0.75,
    ω = 0.1,
    Ω = 0.0,
    e = 0.25,
    i = 0.1,
    μ = 2.0,
    plx = 45.
)
truth_elements = KeplerianElements(truth)
times = sort(rand(Uniform(0, period(truth_elements)/2), 5,))
points = hcat(
    raoff.(truth_elements, times),
    decoff.(truth_elements, times)
)

# Create synthetic images at each time with those points
using DirectImages, ImageFiltering
images_contrasts = map(eachrow(points)) do (ra,dec)
    x = -ra
    y = dec
    img = zeros(201,201)
    r = imgsep(img)
    img[r .< 2] .= NaN

    img = map(zip(img, r)) do (px,r)
        # px + 2000randn()/0.5r
        # px + 900randn()/0.5r
        px + 200randn()/0.5r
    end

    img = centered(img)

    # img = imfilter(img, Kernel.gaussian(5), NA())
    img_for_contrast = imfilter(img, Kernel.gaussian(5), "replicate")
    contrast = contrast_interp(img_for_contrast)

    img[round(Int,x/10), round(Int,y/10)] += 200

    img = imfilter(img, Kernel.gaussian(5), "replicate")

    img, contrast
end
images = [img for (img,contrast) in images_contrasts]
contrasts = [contrast for (img,contrast) in images_contrasts]
writefits("image-examples-1.fits", images...)

##

imshow(
    reduce(hcat, snrmap.(images, contrasts)),
    cmap=:magma,
    clims=(-1.0, 6.0)
)


text:
@testset "epoch" begin
    d = collect(1:100)
    evt = DataFrame(:latency=>(50))

    ep = τ -> Unfold.epoch(d,evt,τ,1)[1][1,:,1]
    # check time delays
    @test ep((0,10.)) ≈ collect(50:60.)
    @test ep((-10,10.)) ≈ collect(40:60.)
    @test ep((-10,0.)) ≈ collect(40:50.)
    @test ep((5,15)) ≈ collect(55:65.)
    @test ep((-15,-5)) ≈ collect(35:45.)

    # check corner cases (sample doesnt end on sampling rate)
    @test ep((0.6,2)) ≈ collect(51:52.)
    @test ep((0.2,2)) ≈ collect(50:52.)

    # test sampling frequencies

    ep = τ -> Unfold.epoch(d,evt,τ,2)[1][1,:,1]
    @test ep((-1.0, 2)) ≈ collect(48:54.)

    ep = τ -> Unfold.epoch(d,evt,τ,0.5)[1][1,:,1]
    @test ep((-4.0, 8)) ≈ collect(48:54.)

# rounding bug when latency was .5 -> bug #78
d = zeros((1, 1270528))
evt = DataFrame(:latency=>(181603.5))
ep = τ -> Unfold.epoch(d,evt,τ,256.)[1][1,:,1]
ep((-0.1, 0.8))
end

max_stars_repo_name: unfoldtoolbox/jlunfold

max_issues_repo_name: unfoldtoolbox/jlunfold

max_forks_repo_name: unfoldtoolbox/jlunfold


text:
using Test, BifurcationKit, ForwardDiff, RecursiveArrayTools, LinearAlgebra
const BK = BifurcationKit
N = 10
M = 5
δ = 1e-8
####################################################################################################
# test the jacobian of the multiple shooting functional using Linear flow
# TODO do example with A matrix and exp(At)
vf(x, p) = x
flow(x, p, t) = exp(t) .* x
dflow(x, p, dx, t) = (t = t, u = flow(x, p, t), du = exp(t) .* dx)
section(x, T) = dot(x[1:N], ones(N))
section(x::BorderedArray, T) = section(vec(x.u[:,:]), T)
par = nothing

fl = BK.Flow(vf, flow, dflow)

probSh = BK.ShootingProblem(M, fl,
	LinRange(0, 1, M+1) |> diff,
	section, false)

poguess = VectorOfArray([rand(N) for ii=1:M])
	po = BorderedArray(poguess, 1.)

dpoguess = VectorOfArray([rand(N) for ii=1:M])
	dpo = BorderedArray(dpoguess, 2.)

# use of AbstractArray structure
pov = vcat(vec(po.u), po.p)
dpov = vcat(vec(dpo.u), dpo.p)
resv = probSh(pov, par)

dresv = probSh(pov, par, dpov; δ = δ)
resad = ForwardDiff.derivative(t -> probSh(pov .+ t .* dpov, par), 0.)
@test norm(resad[1:end-1] - dresv[1:end-1], Inf) < 1e-14

# use of BorderedArray structure
res = probSh(po, par)
@test norm(vec(res.u[:,:]) - resv[1:end-1] , Inf) ≈ 0
@test norm(res.p - resv[end], Inf) ≈ 0

dres = probSh(po, par, dpo; δ = δ)
@test norm(dres.p - dresv[end], Inf) ≈ 0
@test norm(vec(dres.u[:,:]) - dresv[1:end-1], Inf) ≈ 0
####################################################################################################
# test the jacobian of the multiple shooting functional using nonLinear flow
vf(x, p) = x.^2
flow(x, p, t) = x ./ (1 .- t .* x)
dflow(x, p, dx, t) = (t = t, u = flow(x, p, t), du = dx ./ (1 .- t .* x).^2)

fl = BK.Flow(vf, flow, dflow)

probSh = BK.ShootingProblem(M, fl,
	LinRange(0,1,M+1) |> diff ,
	section, false)

poguess = VectorOfArray([rand(N) for ii=1:M])
	po = BorderedArray(poguess, 1.)

dpoguess = VectorOfArray([rand(N) for ii=1:M])
	dpo = BorderedArray(dpoguess, 2.)

# use of AbstractArray structure
pov = vcat(vec(po.u), po.p)
dpov = vcat(vec(dpo.u), dpo.p)
resv = probSh(pov, par)

dresv = probSh(pov, par, dpov; δ = δ)
resad = ForwardDiff.derivative(t -> probSh(pov .+ t .* dpov, par), 0.)
@test norm(resad[1:end-1] - dresv[1:end-1], Inf) < 1e-14

# use of BorderedArray structure
res = probSh(po, par)
@test norm(vec(res.u[:,:]) - resv[1:end-1] , Inf) ≈ 0
@test norm(res.p - resv[end], Inf) ≈ 0

dres = probSh(po, par, dpo; δ = δ)
@test norm(dres.p - dresv[end], Inf) ≈ 0
@test norm(vec(dres.u[:,:]) - dresv[1:end-1], Inf) ≈ 0
####################################################################################################
# test the hyperplane projections for Poincare Shooting
M = 3
normals = [rand(50) for ii=1:M]
for ii=1:M
	normals[ii] /= norm(normals[ii])
end
centers = [rand(50) for ii=1:M]

hyper = BK.SectionPS(normals, centers)

x = 1:50 |>collect .|> Float64 |> vec
x = rand(50)
xb = BK.R(hyper, x, 1)
# test
for ii=1:M
	x2 = BK.E(hyper, xb, ii)
	# test that x2 in Sigma2
	@test dot(x2 - centers[ii], normals[ii]) < 1e-14
	# that we have Rk∘Ek = Id and Ek∘Rk = IdΣ
	@test BK.R(hyper, x2, ii) ≈ xb
end

# test of the derivatives of E and R
dx = rand(50)
_out1 = (BK.R(hyper, x .+ δ .* dx, 1) - BK.R(hyper, x, 1)) ./ δ
_out2 = zero(_out1)
_out2 = BK.dR!(hyper, _out2, dx, 1)

@test norm(_out2 - _out1, Inf) < 1e-5

dx = rand(49)
_out1 = ForwardDiff.derivative(t -> BK.E(hyper, xb .+ t .* dx,1), 0.)
_out2 = BK.dE(hyper, dx, 1)
@test norm(_out2 - _out1, Inf) < 1e-12


text:
module EMLreboiler
	using DanaTypes
	using DotPlusInheritance
	using Reexport
	@reexport using ...stage_separators.EMLtank
	import EMLtypes.length
	include("reboiler/thermosyphon.jl")
	include("reboiler/reboilerSteady.jl")
	include("reboiler/reboilerReact.jl")
	include("reboiler/reboiler.jl")
end


text:
# [scr/vector_to_state.jl]

"""
Build a state from a state vector by recursively moving through the
elements of a more complex example state (matrix, cell array, or struct),
placing the relevant numerical values from the vector into the
appropriate places. This is used in odehybrid to convert a state vector
into the original complex setup of states. The reverse function is
state_to_vector.

state = vector_to_state(vector, state);

Inputs:

    vector  A vector containing the numerical values from the object
    state   A numeric, cell array, or struct array type representing the
            sizes and types for the values in vector

Outputs:

    state   A numeric, cell array, or struct array type
    count   (Internal use only)

Example:

    mutable struct TestStruc
        a
        bcd
    end
    Base.:(==)(a::TestStruc, b::TestStruc) = (Base.:(==)(a.a, b.a) && Base.:(==)(a.bcd, b.bcd)) # Custom

    x = [ [1 3; 4 2], TestStruc([5; 6], [7:9; 10]),  3.14 ]
    v = state_to_vector(x)
    x2 = vector_to_state(v, x)
    x .== x2[1]

See also: state_to_vector.

Online doc: http://www.anuncommonlab.com/doc/odehybrid/vector_to_state.html

Copyright 2014 An Uncommon Lab
"""
function vector_to_state(vector, state)

    count = 0

    return vector_to_state(vector, state, count)
end

function vector_to_state(vector, state, count)
    """
        Internal function that converts vector to a struct state
    """

    state_copy = deepcopy(state) # So we can setproperty! without modifying original

	# Should be a struct - convert and store in each field.
    fields = fieldnames(typeof(state_copy))
    for k = 1:length(fields)
        temp, count = vector_to_state(vector, getproperty(state_copy, fields[k]), count)
        setproperty!(state_copy, fields[k], temp)
    end

    return state_copy, count
end

function vector_to_state(vector, state::Tuple, count)
    """
        Internal function that converts vector to a tuple state
    """

    result = zeros(length(state))
    # First convert to an array, and then to a tuple
    for k = 1:length(vector)
        result[k], temp = vector_to_state(vector[k], state[k])
    end

    count += length(state)
    return tuple(result...), count
end

function vector_to_state(vector, state::Char, count)
    """
        Internal function that converts vector to a Char state
    """

    state = Char(vector[count .+ 1])
    count = count + 1

    return state, count
end

function vector_to_state(vector, state::Array{Char}, count)
    """
        Internal function that converts vector to a Char array state
    """

    state[:] = vector[count .+ range(1, length(state), step = 1)]
    count = count + length(state);

    # state = string(state...) # Return State as a string rather than char array
    return state, count
end

function vector_to_state(vector, state::Number, count)
    """
        Internal function that converts vector to a Number state
    """

    state = float(vector[count + 1])
    count = count + 1

    return state, count
end

function vector_to_state(vector, state::Array{<:Number}, count)
    """
        Internal function that converts vector to a Number array state
    """

    state_copy = deepcopy(state)
    state_copy[:] = float.(vector[count .+ range(1, length(state_copy), step = 1)])
    count = count + length(state_copy);

    return state_copy, count
end

function vector_to_state(vector, state::Array, count)
    """
        Internal function that converts vector to a vector of structs state
    """

    state_copy = deepcopy(state) # So we can setproperty! without modifying original

	# If it's a struct, convert and store in each field.
    for n = 1:length(state_copy)
        temp, count = vector_to_state(vector, state_copy[n], count)
        state_copy[n] = temp
    end

    return state_copy, count
end


text:
using WAVI
function MISMIP_PLUS_Ice0()
    #Grid and boundary conditions
    nx = 80
    ny = 10
    nσ = 4
    x0 = 0.0
    y0 = -40000.0
    dx = 8000.0
    dy = 8000.0
    h_mask=trues(nx,ny)
    u_iszero = falses(nx+1,ny); u_iszero[1,:].=true
    v_iszero=falses(nx,ny+1); v_iszero[:,1].=true; v_iszero[:,end].=true
    grid = Grid(nx = nx,
                ny = ny,
                nσ = nσ,
                x0 = x0,
                y0 = y0,
                dx = dx,
                dy = dy,
                h_mask = h_mask,
                u_iszero = u_iszero,
                v_iszero = v_iszero)

    #Bed
    bed = WAVI.mismip_plus_bed #function definition

    #solver parameters
    maxiter_picard = 1
    solver_params = SolverParams(maxiter_picard = maxiter_picard)

    #Physical parameters
    default_thickness = 100.0 #set the initial condition this way
    accumulation_rate = 0.3
    params = Params(default_thickness = default_thickness,
                    accumulation_rate = accumulation_rate)

    #make the model
    model = Model(grid = grid,
                     bed_elevation = bed,
                     params = params,
                     solver_params = solver_params)

    #timestepping parameters
    dt = 0.5
    end_time = 10000.
    timestepping_params = TimesteppingParams(dt = dt, end_time = end_time)

    simulation = Simulation(model = model,
                        timestepping_params = timestepping_params)

    #perform the simulation
    run_simulation!(simulation)
    return simulation
end

#@time simulation = MISMIP_PLUS_test();


text:
# TODO: switch to keyword argument h = ...
# TODO: switch API to dispatch, e.g., `example_mesh(Scramjet(), h = .1)`

# predefined domains
Base.@kwdef struct RectangularDomain{Tv, Ti}
    xlims::SVector{2,Tv} = SVector{2}(-1., 1.)
    ylims::SVector{2,Tv} = SVector{2}(-1., 1.)
    segment_markers::SVector{4,Ti} = SVector{4}(1, 2, 3, 4)
end

SquareDomain(xlims::SVector{2, Tv} = SVector{2}(-1., 1.),
             segment_markers::SVector{4, Ti} = SVector{4}(1, 2, 3, 4)) where {Tv, Ti} =
    RectangularDomain(xlims, xlims, segment_markers)

Base.@kwdef struct RectangularDomainWithHole{Tv, Ti}
    xlims::SVector{2,Tv} = SVector{2}(-1., 1.)
    ylims::SVector{2,Tv} = SVector{2}(-1., 1.)
    hole_xlims::SVector{2,Tv} = SVector{2}(-.1, .1)
    hole_ylims::SVector{2,Tv} = SVector{2}(-.1, .1)
    segment_markers::SVector{8,Ti} = SVector{8}(1, 1, 1, 1, 2, 2, 2, 2)
end

Base.@kwdef struct Scramjet{Ti}
    # segment marker legend: 1 = wall, 2 = inflow, 3 = outflow
    segment_markers::SVector{8,Ti} = SVector{8}(1, 3, 1, 2, 1, 1, 1, 1)
end

function triangulate_domain(domain::RectangularDomain; h = .1)
    @unpack xlims,ylims,segment_markers = domain
    triin=Triangulate.TriangulateIO()
    triin.pointlist=Matrix{Cdouble}([xlims[1] ylims[1];
                                     xlims[2] ylims[1];
                                     xlims[2] ylims[2];
                                     xlims[1] ylims[2];
                                    ]')
    triin.segmentlist=Matrix{Cint}([1 2; 2 3; 3 4; 4 1;]')
    triin.segmentmarkerlist=Vector{Int32}(segment_markers)
    triout = triangulate(triin, h^2)
    return triout
end

# Todo: allow for general polygon with polygonal hole (using counterclockwise ordering)
function triangulate_domain(domain::RectangularDomainWithHole; h = .1)
    @unpack xlims, ylims, hole_xlims, hole_ylims, segment_markers  = domain
    triin=Triangulate.TriangulateIO()
    triin.pointlist=Matrix{Cdouble}([xlims[1] ylims[1];
                                    xlims[2] ylims[1];
                                    xlims[2] ylims[2];
                                    xlims[1] ylims[2];
                                    hole_xlims[1] hole_ylims[1];
                                    hole_xlims[2] hole_ylims[1];
                                    hole_xlims[2] hole_ylims[2];
                                    hole_xlims[1] hole_ylims[2];
                                    ]')
    triin.segmentlist = Matrix{Cint}([1 2; 2 3; 3 4; 4 1; 5 6; 6 7; 7 8; 8 5; ]')
    triin.segmentmarkerlist = Vector{Int32}(segment_markers)
    avg(x) = sum(x) / length(x)
    triin.holelist = [avg(hole_xlims) avg(hole_ylims)]'
    triout = triangulate(triin, h^2)
    return triout
end

function triangulate_domain(domain::Scramjet; h = .1)
    triin=Triangulate.TriangulateIO()
    triin.pointlist=Matrix{Cdouble}([0.0 0.0;
                                     8.0 0.0;
                                     8.0 0.8;
                                     0.0 2.0;
                                     2.0 0.7;
                                     4.0 0.2;
                                     7.0 0.6;
                                     6.0 0.7;
                                    ]')
    triin.segmentlist=Matrix{Cint}([1 2; 2 3; 3 4; 4 1; 5 6; 6 7; 7 8; 8 5;]')
    triin.segmentmarkerlist=Vector{Int32}(domain.segment_markers)

    hole_x = sum(triin.pointlist[1,5:8]) / length(triin.pointlist[1,5:8])
    hole_y = sum(triin.pointlist[2,5:8]) / length(triin.pointlist[2,5:8])
    triin.holelist=[hole_x hole_y]'

    triout = triangulate(triin, h^2)
    return triout
end


text:
"""
    isBinaryOr(b::AbstractBoolVar, x::AbstractBoolVar, y::AbstractBoolVar)

Is Or constraint, states that `b <=> x or y`.
"""
struct isBinaryOr <: Constraint
    b::AbstractBoolVar
    x::AbstractBoolVar
    y::AbstractBoolVar
    active::StateObject{Bool}

    function isBinaryOr(b::AbstractBoolVar, x::AbstractBoolVar, y::AbstractBoolVar, trailer)
        constraint = new(b, x, y, StateObject{Bool}(true, trailer))
        addOnDomainChange!(b, constraint)
        addOnDomainChange!(x, constraint)
        addOnDomainChange!(y, constraint)
        return constraint
    end
end

"""
    propagate!(constraint::isBinaryOr, toPropagate::Set{Constraint}, prunedDomains::CPModification)

`isBinaryOr` propagation function.
"""
function propagate!(constraint::isBinaryOr, toPropagate::Set{Constraint}, prunedDomains::CPModification)
    if isbound(constraint.b)
        if !assignedValue(constraint.b)
            prunedX = remove!(constraint.x.domain, true)
            prunedY = remove!(constraint.y.domain, true)
            if !isempty(prunedX)
                addToPrunedDomains!(prunedDomains, constraint.x, prunedX)
                triggerDomainChange!(toPropagate, constraint.x)
            end
            if !isempty(prunedY)
                addToPrunedDomains!(prunedDomains, constraint.y, prunedY)
                triggerDomainChange!(toPropagate, constraint.y)
            end
            setValue!(constraint.active, false)
            return !isempty(constraint.x.domain) && !isempty(constraint.y.domain)
        else # b is true
            if isbound(constraint.x) && !assignedValue(constraint.x)
                prunedY = remove!(constraint.y.domain, false)
                if !isempty(prunedY)
                    addToPrunedDomains!(prunedDomains, constraint.y, prunedY)
                    triggerDomainChange!(toPropagate, constraint.y)
                end
                setValue!(constraint.active, false)
                return !isempty(constraint.y.domain)
            elseif isbound(constraint.y) && !assignedValue(constraint.y)
                prunedX = remove!(constraint.x.domain, false)
                if !isempty(prunedX)
                    addToPrunedDomains!(prunedDomains, constraint.x, prunedX)
                    triggerDomainChange!(toPropagate, constraint.x)
                end
                setValue!(constraint.active, false)
                return !isempty(constraint.x.domain)
            elseif (isbound(constraint.x) && assignedValue(constraint.x)) || (isbound(constraint.y) && assignedValue(constraint.y))
                setValue!(constraint.active, false)
            end
        end
    else # b not bound
        if (isbound(constraint.x) && assignedValue(constraint.x)) || (isbound(constraint.y) && assignedValue(constraint.y))
            prunedB = remove!(constraint.b.domain, false)
            addToPrunedDomains!(prunedDomains, constraint.b, prunedB)
            triggerDomainChange!(toPropagate, constraint.b)
            setValue!(constraint.active, false)
            return true
        elseif (isbound(constraint.x) && !assignedValue(constraint.x)) && (isbound(constraint.y) && !assignedValue(constraint.y))
            prunedB = remove!(constraint.b.domain, true)
            addToPrunedDomains!(prunedDomains, constraint.b, prunedB)
            triggerDomainChange!(toPropagate, constraint.b)
            setValue!(constraint.active, false)
            return true
        end
    end

    return true
end

variablesArray(constraint::isBinaryOr) = [constraint.b, constraint.x, constraint.y]

function Base.show(io::IO, ::MIME"text/plain", con::isBinaryOr)
    println(io, typeof(con), ": ", con.b.id, " ≡ ", con.x.id, " ∨ ", con.y.id, ", active = ", con.active)
    print(io, "   ", con.x)
end

function Base.show(io::IO, con::isBinaryOr)
    print(io, typeof(con), ": ", con.b.id, " ≡ ", con.x.id, " ∨ ", con.y.id)
end


text:
import SpecialFunctions: erf

# Construct custom localiser for the GaAs area
function construct_localiser_AlGaAsSiO2H(lattice, atoms; σ=30)
    @assert lattice == Diagonal(diag(lattice))
    @assert maximum(lattice) == lattice[3, 3]  # 3rd axis is the elongated one
    @assert length(atoms) == 6                 # Two types of atoms
    @assert atoms[1][1].symbol == :Ga          # Gallium symbol
    @assert atoms[2][1].symbol == :As          # Arsenic symbol
    @assert σ ≥ 15

    # The slabs are usually from 0.0 through to around 0.333
    # Using this we find first and last Ga and As atom
    last_z  = sort([r[3] for iat in (1, 2) for r in atoms[iat][2]], by=z -> abs(z - 0.3))[1]
    first_z = sort([r[3] for iat in (1, 2) for r in atoms[iat][2]],
                   by=z -> min(abs(z - 0.0), abs(z - 1.0)))[1]
    first_z > 0.75 && (first_z -= 1.0)

    # https://en.wikipedia.org/wiki/Atomic_radii_of_the_elements_%28data_page%29
    atomic_radii = Dict(:Ga => 1.36 * DFTK.units.Å, :As => 1.14 * DFTK.units.Å)
    max_at = max(atomic_radii[:Ga], atomic_radii[:As])
    last_z  += max_at / lattice[3, 3]  # Convert to fractional coords
    first_z -= max_at / lattice[3, 3]

    @assert -0.25 < first_z < 0.25
    @assert 0.25 < last_z < 0.75
    periodise(x) = ifelse(x < 0.75, x + 1, x)

    first_z = periodise(first_z)
    last_z = periodise(last_z)
    r -> max(0, + 0.5 * erf(σ * (periodise(r[3]) - first_z))
                - 0.5 * erf(σ * (periodise(r[3]) - last_z)))
end

config[@__DIR__] = (
    directory=@__DIR__,
    case="AlGaAsSiO2H",
    n_repeat=10,
    kwargs_scf=(kspacing=0.15, smearing=Smearing.Gaussian(), temperature=0.001),
    firstmixkey="Custom",
    #
    construct_localiser=construct_localiser_AlGaAsSiO2H,
    mixings=Dict(
        "Custom"      => (α, localiser) -> HybridMixing(α=α, kTF=1.0, εr=14, localisation=localiser),
        "HybridGaAs"  => (α, localiser) -> HybridMixing(α=α, kTF=1.0, εr=14, localisation=identity),
        "HybridVac"   => (α, localiser) -> HybridMixing(α=α, kTF=1.0, εr=1,  localisation=identity),
        "Kerker"      => (α, localiser) -> KerkerMixing(α=α, kTF=1.0),
        "None"        => (α, localiser) -> SimpleMixing(α=α),
        "RestaGaAs"   => (α, localiser) -> RestaMixing( α=α, kTF=1.0, εr=14),
    ),
)


text:
using FractionalCalculus
using Test

@testset "Test Caputo Direct Fractional Derivative" begin
    @test isapprox(fracdiff(x->x, 0.5, 0, 1, 1e-8, Caputo_Direct())[1], 2/sqrt(pi); atol = 1e-5)
    @test isapprox(fracdiff(x->x^5, 0.5, 0, 3.2, 1e-8, Caputo_Direct())[1], 4.300306216488329e2; atol = 1e-5)
    @test isapprox(fracdiff(x->x^5, 0.25, 0, 3.2, 1e-8, Caputo_Direct())[1], 382.1228113951680; atol = 1e-5)
    @test isapprox(fracdiff(x->(x-3)^3+(x+2)^2+5, 0.5, 0, 0.83, 1e-8, Caputo_Direct())[1], 23.899930581118699; atol = 1e-5)
    @test isapprox(fracdiff(x->(x-3)^3+(x+2)^2+5, 0.25, 0, 0.83, 1e-8, Caputo_Direct())[1], 22.96354626572943; atol = 1e-5)
    @test isapprox(fracdiff(x->cos(x), 0.5, 0, 0.34, 1e-8, Caputo_Direct())[1], -0.147174774607683; atol = 1e-5)
    @test isapprox(fracdiff(x->cos(x), 0.25, 0, 0.34, 1e-8, Caputo_Direct())[1], -0.093074344902431; atol = 1e-5)
    @test isapprox(fracdiff(x->exp(-x^2), 0.5, 0, 2.3, 1e-8, Caputo_Direct())[1], -0.505891017154289; atol = 1e-5)
    @test isapprox(fracdiff(x->exp(-x^2), 0.25, 0, 2.3, 1e-8, Caputo_Direct())[1], -0.762927553656252; atol = 1e-5)
end

@testset "Test Caputo Piecewise Fractional Derivative" begin
    @test isapprox(fracdiff(x->x, 0.5, 1, 1e-4, Caputo_Piecewise()), 2/sqrt(pi); atol = 1e-5)
    @test isapprox(fracdiff(x->x^5, 0.5, 3.2, 1e-4, Caputo_Piecewise()), 4.300306216488329e2; atol = 1e-5)
    @test isapprox(fracdiff(cos, 0.5, 0.34, 1e-4, Caputo_Piecewise()), -0.147174774607683; atol = 1e-5)
    @test isapprox(fracdiff(x->exp(-x^2), 0.25, 2.3, 1e-3, Caputo_Piecewise()), -0.762927553656252; atol = 1e-5)
end

@testset "Test Caputo Diethelm Fractional Derivative" begin
    @test isapprox(fracdiff(x->x, 0.5, 1, 1e-4, Caputo_Diethelm()), 2/sqrt(pi); atol = 1e-5)
    @test isapprox(fracdiff(x->x^5, 0.5, 3.2, 1e-5, Caputo_Diethelm()), 4.300306216488329e2; atol = 1e-5)
end

@testset "Test Grunwald-Letnikov Fractional Derivative" begin
    @test isapprox(fracdiff(x->x, 0.5, 0, 0.5, GL_Direct())[1], 0.7978845583186518; atol = 1e-5)
    @test isapprox(fracdiff(x->x^5, 0.5, 0, 3.2, GL_Direct())[1], 4.300306216488329e2; atol = 1e-5)
end

@testset "Test RLDiff Matrix Fractional deriavtive" begin
    @test isapprox(fracdiff(x->x, 0.5, 1, 0.0001, RLDiff_Matrix())[end], 2/sqrt(pi); atol = 1e-4)
end