text:
using DifferentialEquations

"""
Calculate the Froude number.

"""
function froude(Q::Real, xs::CrossSection)
    g = 9.806
    Fr = Q / (depth(xs) * width(xs) * sqrt(g * depth(xs)))
end

"""
Ordinary differential equation describing the Gradually-Varied-Flow model.

"""
function dydx(y, p, x)
    i = trunc(Int, ceil(x / (p[3] / length(p[1]))))
    i = i > 0 ? i : 1
    # TODO: Change this dynamically depending on whether the flow is subcritical or supercritical
    pm = -1 # integrate upstream
    Q = p[1]
    xs = p[2][i]
    xs.Yₘ = y * (xs.r + 1) / xs.r
    S0 = xs.S0
    Sf = xs.n^2 * Q^2 / (width(xs)^2 * depth(xs)^(10/3))
    Fr = froude(Q, xs)
    pm * (S0 - Sf) / (1 - Fr^2)
end

"""
Calculate a water surface profile by solving the Gradually-Varied-Flow equation.

- `Q`: discharge
- `ybc`: downstream boundary condition for depth
- `z0`: bed elevation at downstream cross section
- `S0`: bed slope for reach
- `n`: Manning's roughness coefficient
- `x`: downstream distance for each cross section
- `b`: bottom width for each cross section
- `r`: channel geometry coefficient

"""
function gvf(Q, ybc, w, S0, n, x, wbf, ybf, r)
    c = [Dingman(wbf[i], ybf[i], ybc, r[i], S0[i], n) for i in 1:length(x)]
    prob = DiscreteProblem(dydx, ybc, (x[1], x[end]), (Q, c, x[end]))
    h = try
        sol = solve(prob, Tsit5(), abstol=1e-2, saveat=x)
        sol.u
    catch
        zeros(length(x)) .- 9999.
    end
    h
end


text:
module BenchLookup
using BenchmarkTools
using SymSpellChecker
using SymSpellChecker: VerbosityALL

suite = BenchmarkGroup()

d = SymSpell()

push!(d, "a", 5)
push!(d, "bc", 2)
s = "abc"

suite["simple_lookup"] = @benchmarkable lookup($d, $s, max_edit_distance = 1)

d2 = SymSpell(joinpath(@__DIR__, "..", "assets", "frequency_dictionary_en_82_765.txt"))
s1 = "quintesestnial"
s2 = "goophr"

suite[s1] = @benchmarkable lookup($d2, $s1)
suite[s2] = @benchmarkable lookup($d2, $s2, verbosity = VerbosityALL)

end # module

BenchLookup.suite

ext: jl


text:
"Parameters for right-hand side function of Variational Partitioned Runge-Kutta methods."
const ParametersVPRKpVariational = AbstractParametersVPRK{:vprk_pvariational}

"Variational partitioned Runge-Kutta integrator."
struct IntegratorVPRKpVariational{DT, TT, D, S,
                PT  <: ParametersVPRK{DT,TT},
                PPT <: ParametersVPRKpVariational{DT,TT},
                ST  <: NonlinearSolver{DT},
                PST <: NonlinearSolver{DT},
                IT <: InitialGuessIODE{TT}} <: AbstractIntegratorVPRKwProjection{DT,TT,D,S}

    params::PT
    pparams::PPT
    solver::ST
    projector::PST
    iguess::IT
    caches::CacheDict{PPT}

    function IntegratorVPRKpVariational(params::ParametersVPRK{DT,TT,D,S},
                    pparams::ParametersVPRKpVariational{DT,TT,D,S},
                    solver::ST, projector::PST, iguess::IT, caches) where {DT,TT,D,S,ST,PST,IT}
        new{DT, TT, D, S, typeof(params), typeof(pparams), ST, PST, IT}(params, pparams, solver, projector, iguess, caches)
    end

    function IntegratorVPRKpVariational{DT,D}(equations::NamedTuple, tableau::TableauVPRK{TT}, Δt::TT; R=[1,1]) where {DT, TT, D}
        # get number of stages
        S = tableau.s

        # create params
        params = ParametersVPRK{DT,D}(equations, tableau, Δt)

        # create projector params
        R  = convert(Vector{TT}, R)
        R1 = [R[1], zero(TT)]
        R2 = [zero(TT), R[2]]
        pparams = ParametersVPRKpVariational{DT,D}(equations, tableau, Δt, NamedTuple{(:R,:R1,:R2)}((R,R1,R2)))

        # create cache dict
        caches = CacheDict(pparams)

        # create nonlinear solver
        solver = create_nonlinear_solver(DT, D*S, params, caches)

        # create projector
        projector = create_nonlinear_solver(DT, 2*D, pparams, caches)

        # create initial guess
        iguess = InitialGuessIODE(get_config(:ig_extrapolation), equations[:v̄], equations[:f̄], Δt)

        # create integrator
        IntegratorVPRKpVariational(params, pparams, solver, projector, iguess, caches)
    end

    function IntegratorVPRKpVariational(equation::Union{IODE{DT},LODE{DT}}, tableau, Δt; kwargs...) where {DT}
        IntegratorVPRKpVariational{DT, ndims(equation)}(get_functions(equation), tableau, Δt; kwargs...)
    end
end

function Integrators.initialize!(int::IntegratorVPRKpVariational{DT}, sol::AtomicSolutionPODE{DT},
                                 cache::IntegratorCacheVPRK{DT}=int.caches[DT]) where {DT}
    sol.t̄ = sol.t - timestep(int)

    equation(int, :v̄)(sol.t, sol.q, sol.v)
    equation(int, :f̄)(sol.t, sol.q, sol.v, sol.f)

    initialize!(int.iguess, sol.t, sol.q, sol.p, sol.v, sol.f,
                            sol.t̄, sol.q̄, sol.p̄, sol.v̄, sol.f̄)

    # initialise projector
    equation(int, :g)(sol.t, sol.q, cache.λ, cache.G[1])
    # cache.U[1] .= cache.λ
end

function initial_guess!(int::IntegratorVPRKpVariational{DT}, sol::AtomicSolutionPODE{DT},
                        cache::IntegratorCacheVPRK{DT}=int.caches[DT]) where {DT}
    for i in eachstage(int)
        evaluate!(int.iguess, sol.q̄, sol.p̄, sol.v̄, sol.f̄,
                              sol.q, sol.p, sol.v, sol.f,
                              cache.q̃, cache.ṽ,
                              tableau(int).q.c[i])

        for k in eachdim(int)
            int.solver.x[ndims(int)*(i-1)+k] = cache.ṽ[k]
        end
    end
end

function initial_guess_projection!(int::IntegratorVPRKpVariational, sol::AtomicSolutionPODE)
    offset_q = 0
    offset_λ = ndims(int)

    for k in eachdim(int)
        int.projector.x[offset_q+k] = sol.q[k]
        int.projector.x[offset_λ+k] = 0
    end
end

function compute_projection!(
                x::Vector{ST}, q::SolutionVector{ST}, p::SolutionVector{ST}, λ::SolutionVector{ST},
                U::Vector{Vector{ST}}, G::Vector{Vector{ST}},
                params::ParametersVPRKpVariational{DT,TT,D,S}) where {ST,DT,TT,D,S}

    @assert D == length(q) == length(p) == length(λ)
    @assert D == length(U[1]) == length(U[2])
    @assert D == length(G[1]) == length(G[2])

    # copy x to q, λ
    for k in 1:D
        q[k] = x[0*D+k]
        λ[k] = x[1*D+k]
    end

    # compute u=λ and g=∇ϑ(q)⋅λ
    U[1] .= 0
    U[2] .= λ
    # U[1] .= λ
    # U[2] .= 0

    params.equ[:g](params.t̄, q, λ, G[1])
    G[2] .= 0
    # G[1] .= 0
    # params.equ.g(params.t̄, q, λ, G[2])

    # compute p=ϑ(q)
    params.equ[:ϑ](params.t̄, q, λ, p)
end

"Compute stages of projected variational partitioned Runge-Kutta methods."
function Integrators.function_stages!(x::Vector{ST}, b::Vector{ST},
                params::ParametersVPRKpVariational{DT,TT,D,S},
                caches::CacheDict) where {ST,DT,TT,D,S}

    @assert length(x) == length(b)

    # get cache for internal stages
    cache = caches[ST]

    compute_projection!(x, cache.q̃, cache.p̃, cache.λ, cache.U, cache.G, params)

    # # compute b = - [q̄-q-U]
    for k in 1:D
        b[0*D+k] = - (cache.q̃[k] - params.q̄[k]) + params.Δt * params.pparams[:R][2] * cache.U[2][k]
        # b[0*D+k] = - (cache.q[k] - params.q̄[k])
    end

    # compute b = - [p̄-p-G]
    for k in 1:D
        b[1*D+k] = - (cache.p̃[k] - params.p̄[k])
        # b[1*D+k] = - (cache.p[k] - params.p̄[k]) + params.Δt * params.R[2] * cache.G[2][k]
    end
end

function Integrators.integrate_step!(int::IntegratorVPRKpVariational{DT,TT}, sol::AtomicSolutionPODE{DT,TT},
                                     cache::IntegratorCacheVPRK{DT}=int.caches[DT]) where {DT,TT}
    # add perturbation for next time step to solution
    # (same vector field as previous time step)
    project_solution!(int, sol, int.pparams.pparams[:R1], cache)

    # update nonlinear solver parameters from cache
    update_params!(int.params, sol)

    # compute initial guess
    initial_guess!(int, sol, cache)

    # reset solution
    reset!(sol, timestep(int))

    # call nonlinear solver
    solve!(int.solver)

    # print solver status
    print_solver_status(int.solver.status, int.solver.params)

    # check if solution contains NaNs or error bounds are violated
    check_solver_status(int.solver.status, int.solver.params)

    # compute vector fields at internal stages
    compute_stages!(int.solver.x, cache.Q, cache.V, cache.P, cache.F, int.params)

    # compute unprojected solution
    update_solution!(int, sol, cache)

    # set time and solution for projection solver
    update_params!(int.pparams, sol)

    # set initial guess for projection
    initial_guess_projection!(int, sol)

    # call projection solver
    solve!(int.projector)

    # print solver status
    print_solver_status(int.projector.status, int.projector.params)

    # check if solution contains NaNs or error bounds are violated
    check_solver_status(int.projector.status, int.projector.params)

    # compute projection vector fields
    compute_projection!(int.projector.x, cache.q̃, cache.p̃, cache.λ,
                                         cache.U, cache.G, int.pparams)

    # add projection to solution
    project_solution!(int, sol, int.pparams.pparams[:R2], cache)

    # copy solution to initial guess
    update_vector_fields!(int.iguess, sol.t, sol.q, sol.p, sol.v, sol.f)
end


text:
@testset "moinput" begin
    x = [rand(5) for _ in 1:4]
    mgpi = MOInput(x, 3)

    @test length(mgpi) == 12
    @test size(mgpi) == (12,)
    @test size(mgpi, 1) == 12
    @test size(mgpi, 2) == 1
    @test lastindex(mgpi) == 12
    @test firstindex(mgpi) == 1
    @test iterate(mgpi) == (mgpi[1], 1)
    @test iterate(mgpi, 2) == (mgpi[3], 3)
    @test_throws BoundsError mgpi[0]

    @test mgpi[2] == (x[2], 1)
    @test mgpi[5] == (x[1], 2)
    @test mgpi[7] == (x[3], 2)
    @test all([(x_, i) for i in 1:3 for x_ in x] .== mgpi)
end


text:
mutable struct OptimStruct

    lossfunc::String
    method::String
    normalizeenergy::Bool
    normalizestress::Bool
    eta::Vector{Float64}
    kappa::Vector{Int64}
    weightouter::Vector{Float64}
    etaspace::Matrix{Float64}
    kappalist::Matrix{Int64}

end

function setoptim(lossfunc=nothing, method=nothing, normalizeenergy=nothing, normalizestress=nothing,
     eta=nothing, kappa=nothing, weightouter=nothing, etaspace=nothing, kappalist=nothing)

    if lossfunc === nothing
        lossfunc = "energyforce"
    end
    if method === nothing
        method = "lsq"
    end
    if normalizeenergy === nothing
        normalizeenergy = false
    end
    if normalizestress === nothing
        normalizestress = false
    end
    if eta === nothing
        eta = []
    end
    if kappa === nothing
        kappa = []
    end
    if weightouter === nothing
        weightouter = [1.0, 1.0, 1.0]
    end
    if etaspace === nothing
        etaspace = reshape([0.0 1.0],(1,2))
    end
    if kappalist === nothing
        kappalist = reshape([0],(1,1))
    end

    optim = OptimStruct(lossfunc, method, normalizeenergy, normalizestress, eta, kappa, weightouter, etaspace, kappalist)

    return optim
end

function getoptim(optim::OptimStruct)

    lossfunc = optim.lossfunc
    method = optim.method
    normalizeenergy = optim.normalizeenergy
    normalizestress = optim.normalizestress
    eta = optim.eta
    kappa = optim.kappa
    weightouter = optim.weightouter
    etaspace = optim.etaspace
    kappalist = optim.kappalist

    return lossfunc, method, normalizeenergy, normalizestress, eta, kappa, weightouter, etaspace, kappalist
end

max_stars_repo_name: cesmix-mit/MDP

max_issues_repo_name: cesmix-mit/MDP

max_forks_repo_name: cesmix-mit/MDP


text:
struct EnumerableJoin{T,TKey,TI,SO,SI,OKS<:Function,IKS<:Function,RS<:Function} <: Enumerable
    outer::SO
    inner::SI
    outerKeySelector::OKS
    innerKeySelector::IKS
    resultSelector::RS
end

Base.eltype(::Type{EnumerableJoin{T,TKeyOuter,TI,SO,SI,OKS,IKS,RS}}) where {T,TKeyOuter,TI,SO,SI,OKS,IKS,RS} = T

function join(outer::Enumerable, inner::Enumerable, f_outerKeySelector::Function, outerKeySelector::Expr, f_innerKeySelector::Function, innerKeySelector::Expr, f_resultSelector::Function, resultSelector::Expr)
    TO = eltype(outer)
    TI = eltype(inner)
    TKeyOuter = Base._return_type(f_outerKeySelector, Tuple{TO,})
    TKeyInner = Base._return_type(f_innerKeySelector, Tuple{TI,})

    if TKeyOuter!=TKeyInner
        error("The keys in the join clause have different types, $TKeyOuter and $TKeyInner.")
    end

    SO = typeof(outer)
    SI = typeof(inner)

    T = Base._return_type(f_resultSelector, Tuple{TO,TI})

    OKS = typeof(f_outerKeySelector)
    IKS = typeof(f_innerKeySelector)
    RS = typeof(f_resultSelector)

    return EnumerableJoin{T,TKeyOuter,TI,SO,SI,OKS,IKS,RS}(outer,inner,f_outerKeySelector,f_innerKeySelector,f_resultSelector)
end

function Base.iterate(iter::EnumerableJoin{T,TKeyOuter,TI,SO,SI,OKS,IKS,RS}) where {T,TKeyOuter,TI,SO,SI,OKS,IKS,RS}
    results = Array{T}(undef, 0)

    inner_dict = OrderedDict{TKeyOuter,Array{TI,1}}()
    for i in iter.inner
        key = iter.innerKeySelector(i)
        if !haskey(inner_dict, key)
            inner_dict[key] = Array{TI}(undef, 0)
        end
        push!(inner_dict[key], i)
    end

    for i in iter.outer
        outerKey = iter.outerKeySelector(i)
        if haskey(inner_dict,outerKey)
            for j in inner_dict[outerKey]
                push!(results, iter.resultSelector(i,j))
            end
        end
    end

    if length(results)==0
        return nothing
    end

    return results[1], (results, 2)
end

function Base.iterate(iter::EnumerableJoin{T,TKeyOuter,TI,SO,SI,OKS,IKS,RS}, state) where {T,TKeyOuter,TI,SO,SI,OKS,IKS,RS}
    if state[2]>length(state[1])
        return nothing
    else
        return state[1][state[2]], (state[1], state[2]+1)
    end
end


text:
# This file was generated by the Julia Swagger Code Generator
# Do not modify this file directly. Modify the swagger specification instead.

include("model_CloudError.jl")
include("model_ErrorAdditionalInfo.jl")
include("model_ErrorResponse.jl")
include("model_ParameterDefinitions.jl")
include("model_ParameterDefinitionsValue.jl")
include("model_ParameterValues.jl")
include("model_ParameterValuesValue.jl")
include("model_PolicyDefinitionGroup.jl")
include("model_PolicyDefinitionReference.jl")
include("model_PolicySetDefinition.jl")
include("model_PolicySetDefinitionListResult.jl")
include("model_PolicySetDefinitionProperties.jl")


text:
using TomoForward
using XfromProjections

# using MKLSparse # uncomment if you've installed MKLSparse, it will boost the performance

# test slice by slice
img = zeros(128, 128, 128)
img[40:70, 40:60, 50:70] .= 1.0

nslice = size(img, 3)

nangles = 90
detcount = Int(floor(size(img,1)*1.4))
proj_geom = ProjGeom(1.0, detcount, LinRange(0,pi,nangles+1)[1:nangles])

isdefined_A = @isdefined A
if isdefined_A == false
    A = fp_op_parallel2d_strip(proj_geom, size(img, 1), size(img, 2))
end

p = zeros(nangles, nslice, detcount)
for i=1:nslice
    p[:,i,:] = reshape(A * vec(img[:,:,i]), nangles, detcount);
end

p = permutedims(p, [2,3,1])

u0 = zeros(size(img));
@time u = recon2d_stack_tv_primaldual!(u0, A, p, 20, 0.1);

# using PyPlot
# imshow(u[:,:,60])

max_stars_repo_name: JuliaTomo/DynamicTomographyWithFlow

max_issues_repo_name: JuliaTomo/DynamicTomographyWithFlow

max_forks_repo_name: JuliaTomo/DynamicTomographyWithFlow


text:
module MaxHelpingHandRainbowSpinnerColorAreaController

using ..Ahorn, Maple

@mapdef Entity "MaxHelpingHand/RainbowSpinnerColorAreaController" RainbowSpinnerColorAreaController(x::Integer, y::Integer, width::Integer=Maple.defaultBlockWidth, height::Integer=Maple.defaultBlockHeight,
    colors::String="89E5AE,88E0E0,87A9DD,9887DB,D088E2", gradientSize::Number=280.0, loopColors::Bool=false, centerX::Number=0.0, centerY::Number=0.0, gradientSpeed::Number=50.0)

const placements = Ahorn.PlacementDict(
    "Rainbow Spinner Color Area Controller\n(max480's Helping Hand)" => Ahorn.EntityPlacement(
        RainbowSpinnerColorAreaController,
        "rectangle"
    )
)

Ahorn.minimumSize(entity::RainbowSpinnerColorAreaController) = 8, 8
Ahorn.resizable(entity::RainbowSpinnerColorAreaController) = true, true

Ahorn.selection(entity::RainbowSpinnerColorAreaController) = Ahorn.getEntityRectangle(entity)

function Ahorn.render(ctx::Ahorn.Cairo.CairoContext, entity::RainbowSpinnerColorAreaController, room::Maple.Room)
    width = Int(get(entity.data, "width", 32))
    height = Int(get(entity.data, "height", 32))

    Ahorn.drawRectangle(ctx, 0, 0, width, height, (0.4, 0.4, 1.0, 0.4), (0.4, 0.4, 1.0, 1.0))
end

end

max_stars_repo_name: swoolcock/MaxHelpingHand

max_issues_repo_name: swoolcock/MaxHelpingHand

max_forks_repo_name: swoolcock/MaxHelpingHand


text:
function res_con(model::ContactModel, env::Environment{<:World,LinearizedCone}, z, θ, κ)
	q0, q1, u1, w1, μ, h = unpack_θ(model, θ)
	q2, γ1, b1, ψ1, s1, η1, s2 = unpack_z(model, env, z)

	# @warn "define residual order"
	[s1 .- ϕ_func(model, env, q2);
	 s2 .- (μ[1] * γ1 .- E_func(model, env) * b1);
	 γ1 .* s1 .- κ;
	 b1 .* η1 .- κ;
	 ψ1 .* s2 .- κ]
end

function rz_approx!(s, rz, z, θ)
	model = s.model
	env = s.env
	nq = model.dim.q
	nc = model.dim.c
	nb = model.dim.c * friction_dim(env)

	rz .= 0.0

	q0, q1, u1, w1, μ, h = unpack_θ(model, θ)
	q2, γ1, b1, ψ1, s1, η1, s2 = unpack_z(model, env, z)

	k = kinematics(model, q2)
	λ1 = contact_forces(model, env, γ1, b1, q2, k)
	vT = velocity_stack(model, env, q1, q2, k, h)

	# Dynamics
	idyn = index_dyn(model, env, quat = false)
	iimp = index_imp(model, env, quat = false)
	imdp = index_mdp(model, env, quat = false)
	ifri = index_fri(model, env, quat = false)
	ibimp = index_bimp(model, env, quat = false)
	ibmdp = index_bmdp(model, env, quat = false)
	ibfri = index_bfri(model, env, quat = false)

	iq2 = index_q2(model, env, quat = false)
	iγ1 = index_γ1(model, env, quat = false)
	ib1 = index_b1(model, env, quat = false)
	iψ1 = index_ψ1(model, env, quat = false)
	is1 = index_s1(model, env, quat = false)
	iη1 = index_η1(model, env, quat = false)
	is2 = index_s2(model, env, quat = false)

	rz[idyn, iq2] = model.dyn.∂q2(h, q0, q1, u1, w1, λ1, q2) + s.con.dJ(λ1, q2)
	rz[idyn, [iq2; iγ1; ib1]] += s.con.dλ1(q2)' * s.con.dcf(γ1, b1, q2, k)

	# Maximum dissipation
	rz[imdp, [iψ1; iη1]] = s.con.mdψη(vT, ψ1, η1)
	rz[imdp, iq2] += s.con.mdvs(vT, ψ1, η1) * s.con.vsq2(q1, q2, k, h)

	# Other constraints
	s.con.rcz(view(rz, collect([iimp; ifri; ibimp; ibmdp; ibfri]), :), z, θ)
end

function rθ_approx!(s, rθ, z, θ)
	model = s.model
	env = s.env
	nq = model.dim.q
	nc = model.dim.c
	nb = model.dim.c * friction_dim(env)

	rθ .= 0.0

	q0, q1, u1, w1, μ, h = unpack_θ(model, θ)
	q2, γ1, b1, ψ1, s1, η1, s2 = unpack_z(model, env, z)

	k = kinematics(model, q2)
	λ1 = contact_forces(model, env, γ1, b1, q2, k)
	vT = velocity_stack(model, env, q1, q2, k, h)

	idyn = index_dyn(model, env, quat = false)
	iimp = index_imp(model, env, quat = false)
	imdp = index_mdp(model, env, quat = false)
	ifri = index_fri(model, env, quat = false)
	ibimp = index_bimp(model, env, quat = false)
	ibmdp = index_bmdp(model, env, quat = false)
	ibfri = index_bfri(model, env, quat = false)

	iq0 = index_q0(model)
	iq1 = index_q1(model)
	iu1 = index_u1(model)
	iw1 = index_w1(model)
	iμ  = index_μ(model)
	ih  = index_h(model)

	# Dynamics
	idx = collect([iq0; iq1; iu1; iw1; ih])
	rθ[idyn, idx] = model.dyn.dθ(h, q0, q1, u1, w1, λ1, q2)

	# Maximum dissipation
	idx = collect([iq1; ih])
	rθ[imdp, idx] = s.con.vsq1h(q1, q2, k, h)

	# Other constraints
	s.con.rcθ(view(rθ, collect([iimp; ifri; ibimp; ibmdp; ibfri]), :), z, θ)
end


text:
module Mill

using Flux
using MLDataPattern
using SparseArrays
using Statistics
using Combinatorics
using Zygote
using HierarchicalUtils
using Zygote: @adjoint
using LinearAlgebra
import Base.reduce

MLDataPattern.nobs(::Missing) = nothing

const VecOrRange = Union{UnitRange{Int},AbstractVector{Int}}

"""
	catobs(xs...)

	concatenates all observations from all xs together
"""
function catobs end;

include("bags.jl")
export AlignedBags, ScatteredBags, length2bags

include("util.jl")
include("threadfuns.jl")

include("datanodes/datanode.jl")
export AbstractNode, AbstractProductNode, AbstractBagNode
export ArrayNode, BagNode, WeightedBagNode, ProductNode, LazyNode, IdentityModel
export NGramMatrix, NGramIterator
export catobs, removeinstances

include("aggregations/aggregation.jl")
# agg. types exported in aggregation.jl
export AggregationFunction, Aggregation

include("modelnodes/modelnode.jl")
export AbstractMillModel, ArrayModel, BagModel, ProductModel, LazyModel, IdentityModel, identity_model
export reflectinmodel

include("conv.jl")
export bagconv, BagConv

include("bagchain.jl")
export BagChain

include("replacein.jl")
export replacein, findin

include("hierarchical_utils.jl")
export printtree

Base.show(io::IO, ::T) where T <: Union{AbstractNode, AbstractMillModel, AggregationFunction} = show(io, Base.typename(T))
Base.show(io::IO, ::MIME"text/plain", n::Union{AbstractNode, AbstractMillModel}) = HierarchicalUtils.printtree(io, n; htrunc=3)
Base.getindex(n::Union{AbstractNode, AbstractMillModel}, i::AbstractString) = HierarchicalUtils.walk(n, i)

include("partialeval.jl")
const _emptyismissing = Ref(false)

function emptyismissing(a)
    _emptyismissing[] = a
end

const _terseprint = Ref(true)

function terseprint(a)
    _terseprint[] = a
end

function base_show_terse(io::IO, x::Type{T}) where {T<:Union{AbstractNode,AbstractMillModel}}
    if hasproperty(x, :body) && !hasproperty(x.body, :name) && hasproperty(x.body, :body)
        print(io, "$(x.body.body.name){…}")
        return
    elseif hasproperty(x, :body) && !hasproperty(x, :name)
        print(io, "$(x.body.name){…}")
        return
    else
        print(io, "$(x.name){…}")
        return
    end
end

function base_show_full(io::IO, x::Type{T}) where {T<:Union{AbstractNode,AbstractMillModel}}
    # basically copied from the Julia sourcecode, seems it's one of most robust fixes to Pevňákoviny
    # specifically function show(io::IO, @nospecialize(x::Type))
    if x isa DataType
        Base.show_datatype(io, x)
        return
    elseif x isa Union
        if x.a isa DataType && Core.Compiler.typename(x.a) === Core.Compiler.typename(DenseArray)
            T2, N = x.a.parameters
            if x == StridedArray{T2,N}
                print(io, "StridedArray")
                Base.show_delim_array(io, (T2,N), '{', ',', '}', false)
                return
            elseif x == StridedVecOrMat{T2}
                print(io, "StridedVecOrMat")
                Base.show_delim_array(io, (T2,), '{', ',', '}', false)
                return
            elseif StridedArray{T2,N} <: x
                print(io, "Union")
                Base.show_delim_array(io, vcat(StridedArray{T2,N}, Base.uniontypes(Core.Compiler.typesubtract(x, StridedArray{T2,N}))), '{', ',', '}', false)
                return
            end
        end
        print(io, "Union")
        Base.show_delim_array(io, Base.uniontypes(x), '{', ',', '}', false)
        return
    end

    # this type assert is behaving obscurely. When in Mill, it does not assert that LazyNode{T<:Symbol,D} where D is UnionAll, but in debugging using Debugger, it does
    # x::UnionAll
    if Base.print_without_params(x)
        return show(io, Base.unwrap_unionall(x).name)
    end

    if x.var.name === :_ || Base.io_has_tvar_name(io, x.var.name, x)
        counter = 1
        while true
            newname = Symbol(x.var.name, counter)
            if !Base.io_has_tvar_name(io, newname, x)
                newtv = TypeVar(newname, x.var.lb, x.var.ub)
                x = UnionAll(newtv, x{newtv})
                break
            end
            counter += 1
        end
    end

    show(IOContext(io, :unionall_env => x.var), x.body)
    print(io, " where ")
    show(io, x.var)
end

function Base.show(io::IO, x::Type{T}) where {T<:Union{AbstractNode,AbstractMillModel}}
    if _terseprint[]
        return base_show_terse(io, x)
    else
        return base_show_full(io, x)
    end
end

end


text:
# This file contains the base definitions,
# Plus mass-action kinetics.

import Base.==,Base.hash,Base.reverse

const IntVec = AbstractArray{Unsigned,1}
const StrVec = AbstractArray{String,1}

"""
    Reaction(
        reactants::Vector
        stoichr::Vector
        products::Vector
        stoichp::Vector
        kf
        kr,
    )

Defines a reaction by the given reactants and stoichiometry `stoichr`,
the given products and stoichiometry `stoichp`, and the forward and backward
rate constants `kf` and `kr`.

Note: the default definition of a reaction is reversible.
This simplifies analyses for most reactions.
Note that no truly irreversible reactions exist in nature,
And in fact true irreversible reactions _can't_ exist because
this would violate microscopic reversibility.
However, conceptually it can sometimes be convenient to
think of a reaction as being purely unidirectional,
in which case one can set kf or kr to 0.0, which is treated
as a special case.
"""
mutable struct Reaction
    reactants::IntVec # vector of reactants
    stoichr::IntVec   # vector of stoichiometric coefficients of reactants
    products::IntVec  # vector of products
    stoichp::IntVec   # vector of stoichiometric coefficients of products
    kf::Number        # rate constant, forward. Zero means no forward reaction.
    kr::Number        # rate constant, backward. Zero means no backward reaction.
end

struct ReactionNetwork
    n_species::Integer
    reactions::Vector{Reaction}
    names::StrVec     # Species names
end

ReactionNetwork(ns::Integer,r::Vector{Reaction},names=String[]) =
    ReactionNetwork(ns,r,names)

ReactionNetwork(r::Vector{Reaction},names=String[]) =
    ReactionNetwork(get_nspecies(r),r,names)

function get_nspecies(reactions::Vector{Reaction})
    ns = 0
    for r in reactions
        ns = maximum(cat(dims=1,ns,r.reactants,r.products))
    end
    return ns
end

function show_reactants(io::IO, reactants::IntVec, stoich::IntVec, names::StrVec)
    strs = String[]

    for si in 1:length(reactants)
        s = reactants[si]

        cstr = (stoich[si] == 1) ? "" : "$(stoich[si])"

        push!(strs, isempty(names) ? "$cstr(S$s) " : "$cstr$(names[s]) ")
    end

    print(io, join(strs, "+ "))
end

function Base.show(io::IO, r::Reaction, names=String[])
    showdict = Dict((true, true ) => ("↮ ",""),
                    (true, false) => ("← ","[k-:$(r.kr)]"),
                    (false,true ) => ("→ ","[k+:$(r.kf)]"),
                    (false,false) => ("⇋ ","[k+:$(r.kf) k-:$(r.kr)]"))
    eq_sym, k_str = showdict[(r.kf==0, r.kr==0)]

    # Show left-hand side
    show_reactants(io, r.reactants, r.stoichr, names)

    print(io, eq_sym)

    # Show right-hand side
    show_reactants(io, r.products, r.stoichp, names)

    print(io, k_str)
end

function Base.show(io::IO, rn::ReactionNetwork)
    println(io, "Reaction Network: ")
    for r in rn.reactions
        show(io, r, rn.names)
        println(io)
    end
end

function reaction_current_1side(reactants::IntVec, stoich::IntVec, z)
    j = 1

    for si = 1 : length(reactants)
        s = reactants[si]
        c = stoich[si]

        j *= z[s] ^ c
    end

    return j
end

function reaction_current(ri::Reaction, z)
    jf = ri.kf*reaction_current_1side(ri.reactants, ri.stoichr, z)
    jr = ri.kr*reaction_current_1side(ri.products,  ri.stoichp, z)
    return jf,jr
end

function mass_action_1side!(dz, reactants::IntVec, stoich::IntVec, j, pos::Bool)
    for si = 1 : length(reactants)
        s = reactants[si]
        c = stoich[si]

        dz[s] += pos ? c*j : -c*j
    end
end

function mass_action!(dz, ri::Reaction, z)
    jf,jr = reaction_current(ri, z)
    j=jf-jr

    # calculate change in concentration
    mass_action_1side!(dz, ri.reactants, ri.stoichr, j, false)

    mass_action_1side!(dz, ri.products,  ri.stoichp, j, true)
end

function precompute_concpowers(z, n_powers)
    z_powers = zeros(length(z), n_powers)

    z_powers[:,1] = z
    for c = 2 : n_powers
        z_powers[:,c] = z_powers[:,c-1] .* z
    end

    return z_powers
end

# evolve concentration of reactants based on mass-action kinetics.
# re: list of reactions
# z: concentration of reactants
# returns: dz, instantanous change in z over time
function mass_action(reactions::Vector{Reaction}, z)
    dz = zero(z)

    for ri in reactions
        mass_action!(dz, ri, z)
    end

    return dz
end

function jacobian(rn::ReactionNetwork)
    jm = zeros(rn.n_species,rn.n_species)

    cm = mass_action(rn.reactions,zeros(rn.n_species))

    for i=1:rn.n_species
        z = zeros(rn.n_species)
        z[i]=1.0
        jm[:,i] = mass_action(rn.reactions,z)
    end

    return cm,jm
end

reverse(r::Reaction) = Reaction(r.products,r.stoichp,r.reactants,r.stoichr,r.kr,r.kf,r.names)

# Function to determine if a reaction is redundant;
# i.e. A + 2B -> 2B + A is redundant.
# Note that for this to work, the reaction must be in normal form,
# Which means the reactants and products must be a sorted list.
isredundant(r::Reaction) = (r.products == r.reactants) && (r.stoichp == r.stoichr )

function test_equality(r::Reaction,p::Reaction)
    return (r.stoichr == p.stoichr) &&
           (r.stoichp == p.stoichp) &&
           (r.reactants == p.reactants) &&
           (r.products  == p.products) &&
           (r.kf == p.kf) &&
           (r.kr == p.kr)
end

# Two reactions are equal if they are equal in either
# the forward or backward directions
function ==(r::Reaction,p::Reaction)
    return test_equality(r,p) || test_equality(reverse(r),p)
end

function hash1(r::Reaction)
    return foldr(hash, [r.reactants,r.stoichr,r.products,r.stoichp,r.kf,r.kr,r.names]; init=zero(UInt64))
end
function hash(r::Reaction)
    h1 = hash1(r)
    h2 = hash1(reverse(r))
    return h1+h2
end


text:
importall Base

abstract Engine
type Nuts <: Engine
  max_depth::Int64
end
Nuts(;max_depth::Number=10) = Nuts(max_depth)

type Static <: Engine
  int_time::Float64
end
Static(;int_time::Number=2 * pi) = Static(int_time)

abstract Metric
type unit_e <: Metric
end
type dense_e <: Metric
end
type diag_e <: Metric
end

abstract Algorithm

type Hmc <: Algorithm
  engine::Engine
  metric::Metric
  stepsize::Float64
  stepsize_jitter::Float64
end
Hmc(;engine::Engine=Nuts(), metric::DataType=diag_e,
  stepsize::Number=1.0, stepsize_jitter::Number=1.0) =
    Hmc(engine, metric(), stepsize, stepsize_jitter)
Hmc(engine::Engine) = Hmc(engine, diag_e(), 1.0, 1.0)

type Adapt
  engaged::Bool
  gamma::Float64
  delta::Float64
  kappa::Float64
  t0::Float64
  init_buffer::Int64
  term_buffer::Int64
  window::Int64
end
Adapt(;engaged::Bool=true, gamma::Number=0.05, delta::Number=0.8,
  kappa::Number=0.75, t0::Number=10.0,
  init_buffer::Number=75, term_buffer::Number=50, window::Number=25) =
    Adapt(engaged, gamma, delta, kappa, t0, init_buffer, term_buffer, window)

type Sample <: Methods
  num_samples::Int64
  num_warmup::Int64
  save_warmup::Bool
  thin::Int64
  adapt::Adapt
  algorithm::Algorithm
end
Sample(;num_samples::Number=1000, num_warmup::Number=1000,
  save_warmup::Bool=false, thin::Number=1,
  adapt::Adapt=Adapt(), algorithm::Algorithm=Hmc()) =
    Sample(num_samples, num_warmup, save_warmup, thin, adapt, algorithm)

function sample_show(io::IO, s::Sample, compact::Bool)
  if compact
    println("Sample($(s.num_samples), $(s.num_warmup), $(s.save_warmup), $(s.thin), $(s.adapt), $(s.algorithm))")
  else
    println("  method =                  Sample()")
    println("    num_samples =             ", s.num_samples)
    println("    num_warmup =              ", s.num_warmup)
    println("    save_warmup =             ", s.save_warmup)
    println("    thin =                    ", s.thin)
    if isa(s.algorithm, Hmc)
      println("    algorithm =               HMC()")
      if isa(s.algorithm.engine, Nuts)
        println("      engine =                  NUTS()")
        println("        max_depth =               ", s.algorithm.engine.max_depth)
      elseif isa(s.algorithm.engine, Static)
        println("      engine =                  Static()")
        println("        int_time =                ", s.algorithm.engine.int_time)
      end
    else
      println("    algorithm =               Unknown")
    end
    println("      metric =                  ", typeof(s.algorithm.metric))
    println("      stepsize =                ", s.algorithm.stepsize)
    println("      stepsize_jitter =         ", s.algorithm.stepsize_jitter)
    println("    adapt =                   Adapt()")
    println("      gamma =                   ", s.adapt.gamma)
    println("      delta =                   ", s.adapt.delta)
    println("      kappa =                   ", s.adapt.kappa)
    println("      t0 =                      ", s.adapt.t0)
    println("      init_buffer =             ", s.adapt.init_buffer)
    println("      term_buffer =             ", s.adapt.term_buffer)
    println("      window =                  ", s.adapt.window)
  end
end

show(io::IO, s::Sample) = sample_show(io, s, false)
showcompact(io::IO, s::Sample) = sample_show(io, s, true)


text:
using SparseArrays,LinearAlgebra
n = 10000
L = sprand(Float64,n,n,0.3)
u0 = sprand(Float64,n,0.3)
u1 = similar(u0)

BLAS.set_num_threads(1)
@time mul!(u1,L,u0);

BLAS.set_num_threads(4)
@time mul!(u1,L,u0);

@time mul!(u1,L,u0);
u0_d = rand(Float64,n)

u1_d = similar(u0_d)

@time mul!(u1_d,L,u0_d);


text:
"
Author: Daniel O'Connor
Email: dt.oconnor@outlook.com

Based off the work published in
https://journals.aps.org/prb/abstract/10.1103/PhysRevB.66.094203
Quantum annealing by the path-integral Monte Carlo method:
The two-dimensional random Ising model
Roman Martoňák, Giuseppe E. Santoro, and Erio Tosatti
Phys. Rev. B 66, 094203 – Published 13 September 2002
"

module PIMC
    include("tools.jl")
    include("SA.jl")
    using ..SA
    using ..SamplingTools: hj_info, default_temp_range, generate_nbr_dict,
                            generate_nbr, pimc_aggregation, add_noise,
                            add_cross_talk, hj_info!, generate_nbr!
    using Random
    using Distributed
    using LinearAlgebra
    using SharedArrays
    function anneal!(spin_lattice::Array{Float64, 2},
                    nbs_idxs::Array{Int64, 2},
                    nbs_vals::Array{Float64, 2},
                    a_schedule::Array{Float64, 1},
                    # must be same size or will crash from @inbounds
                    b_schedule::Array{Float64, 1},
                    temp::Float64)
        maxnb = size(nbs_idxs, 1)
        nspins = size(spin_lattice, 1)
        nslices = size(spin_lattice, 2)
        perm = collect(1:nspins)
        istep = 1
        a_field = 0.0
        b_field = 0.0
        teff = Float64(nslices) * temp
        jperp = 0.0
        islice = 0
        idx = 1
        ispin = 1
        spin_nb = 1
        deg = 1
        ediff =  0.0

        @fastmath @inbounds begin
            for istep in eachindex(a_schedule)
                a_field = a_schedule[istep]
                b_field = b_schedule[istep]
                jperp = -0.5 * teff * log(tanh(a_field/teff))
                # perform local moves
                for islice = 1:nslices
                    randperm!(perm)
                    for idx in eachindex(perm)
                        ispin = perm[idx]
                        ediff = 0.0
                        # add z components
                        for deg = 1:maxnb
                            spin_nb = nbs_idxs[deg, ispin]
                            if spin_nb === ispin
                                ediff += -2.0*b_field*nbs_vals[deg, ispin]*spin_lattice[ispin, islice]
                            elseif spin_nb !== 0
                                ediff += -2.0*b_field*spin_lattice[ispin, islice]*(nbs_vals[deg, ispin]*spin_lattice[spin_nb, islice])
                            end
                        end
                        # add x components
                        # periodic boundaries
                        if islice == 1
                            tleft = nslices
                            tright = 2
                        elseif islice == nslices
                            tleft = nslices-1
                            tright = 1
                        else
                            tleft = islice-1
                            tright = islice+1
                        end
                        ediff += 2.0*spin_lattice[ispin, islice]*(jperp*spin_lattice[ispin, tleft])
                        ediff += 2.0*spin_lattice[ispin, islice]*(jperp*spin_lattice[ispin, tright])
                        if ediff <= 0.0
                            spin_lattice[ispin, islice] *= -1.0
                        elseif exp(-1.0 * ediff / teff) > rand()
                            spin_lattice[ispin, islice] *= -1.0
                        end
                    end
                end
                # perform global moves
                randperm!(perm)
                for idx in eachindex(perm)
                    ispin = perm[idx]
                    ediff = 0.0
                    for islice = 1:nslices
                        for deg = 1:maxnb
                            spin_nb = nbs_idxs[deg, ispin]
                            if spin_nb === ispin
                                ediff += -2.0*b_field*nbs_vals[deg, ispin]*spin_lattice[ispin, islice]
                            elseif spin_nb !== 0
                                ediff += -2.0*b_field*spin_lattice[ispin, islice]*(nbs_vals[deg, ispin]*spin_lattice[spin_nb, islice])
                            end
                        end
                    end
                    if ediff <= 0.0
                        spin_lattice[ispin, :] *= -1.0
                    elseif exp(-1.0 * ediff / teff) > rand()
                        spin_lattice[ispin, :] *= -1.0
                    end
                end
            end
        end
        # return spin_vector
    end

    function annealiq!(spin_lattice::Array{Float64, 2},
                    nbs_idxs::Array{Int64, 2},
                    nbs_vals::Array{Float64, 2},
                    a_schedule::Array{Float64, 2},
                    # must be same size or will crash from @inbounds
                    b_schedule::Array{Float64, 2},
                    temp::Float64)
        maxnb = size(nbs_idxs, 1)
        nspins = size(spin_lattice, 1)
        nslices = size(spin_lattice, 2)
        perm = collect(1:nspins)
        num_sweeps = size(a_schedule, 2)
        istep = 1
        a_field = 0.0
        b_field = 0.0
        teff = Float64(nslices) * temp
        jperp = 0.0
        islice = 0
        idx = 1
        ispin = 1
        spin_nb = 1
        deg = 1
        ediff =  0.0

        @fastmath @inbounds begin
            for istep = 1:num_sweeps
                # perform local moves
                for islice = 1:nslices
                    randperm!(perm)
                    for idx in eachindex(perm)
                        ispin = perm[idx]
                        a_field = a_schedule[ispin, istep]
                        b_field = b_schedule[ispin, istep]
                        jperp = -0.5 * teff * log(tanh(a_field/teff))
                        ediff = 0.0
                        # add z components
                        for deg = 1:maxnb
                            spin_nb = nbs_idxs[deg, ispin]
                            if spin_nb === ispin
                                ediff += -2.0*b_field*nbs_vals[deg, ispin]*spin_lattice[ispin, islice]
                            elseif spin_nb !== 0
                                ediff += -2.0*b_field*spin_lattice[ispin, islice]*(nbs_vals[deg, ispin]*spin_lattice[spin_nb, islice])
                            end
                        end
                        # add x components
                        # periodic boundaries
                        if islice == 1
                            tleft = nslices
                            tright = 2
                        elseif islice == nslices
                            tleft = nslices-1
                            tright = 1
                        else
                            tleft = islice-1
                            tright = islice+1
                        end
                        ediff += 2.0*spin_lattice[ispin, islice]*(jperp*spin_lattice[ispin, tleft])
                        ediff += 2.0*spin_lattice[ispin, islice]*(jperp*spin_lattice[ispin, tright])
                        if ediff <= 0.0
                            spin_lattice[ispin, islice] *= -1.0
                        elseif exp(-1.0 * ediff / teff) > rand()
                            spin_lattice[ispin, islice] *= -1.0
                        end
                    end
                end
                # perform global moves
                randperm!(perm)
                for idx in eachindex(perm)
                    ispin = perm[idx]
                    b_field = b_schedule[ispin, istep]
                    ediff = 0.0
                    for islice = 1:nslices
                        for deg = 1:maxnb
                            spin_nb = nbs_idxs[deg, ispin]
                            if spin_nb === ispin
                                ediff += -2.0*b_field*nbs_vals[deg, ispin]*spin_lattice[ispin, islice]
                            elseif spin_nb !== 0
                                ediff += -2.0*b_field*spin_lattice[ispin, islice]*(nbs_vals[deg, ispin]*spin_lattice[spin_nb, islice])
                            end
                        end
                    end
                    if ediff <= 0.0
                        spin_lattice[ispin, :] *= -1.0
                    elseif exp(-1.0 * ediff / teff) > rand()
                        spin_lattice[ispin, :] *= -1.0
                    end
                end
            end
        end
        # return spin_vector
    end

    # NOT AVAILABLE YET
    function anneal_bath!(spin_lattice::Array{Float64, 2},
                    nbs_idxs::Array{Int64, 2},
                    nbs_vals::Array{Float64, 2},
                    a_schedule::Array{Float64, 1},
                    # must be same size or will crash from @inbounds
                    b_schedule::Array{Float64, 1},
                    temp::Float64,
                    alpha::Float64)
        maxnb = size(nbs_idxs, 1)
        nspins = size(spin_lattice, 1)
        nslices = size(spin_lattice, 2)
        perm = collect(1:nspins)
        istep = 1
        a_field = 0.0
        b_field = 0.0
        teff = Float64(nslices) * temp
        jperp = 0.0
        islice = 0
        idx = 1
        ispin = 1
        spin_val = 1.0
        spin_nb = 1
        deg = 1
        ediff =  0.0
        lookuptable = -(alpha * pi /(nslices^2)) ./ (sin.(pi/nslices * collect(1:(nslices-1)))).^2
        cslice = 1

        @fastmath @inbounds begin
            for istep in eachindex(a_schedule)
                a_field = a_schedule[istep]
                b_field = b_schedule[istep]
                jperp = -0.5 * teff * log(tanh(a_field/teff))
                for islice = 1:nslices
                    randperm!(perm)
                    for idx in eachindex(perm)
                        ispin = perm[idx]
                        spin_val = spin_lattice[ispin, islice]
                        ediff = 0.0
                        # add z components
                        for deg = 1:maxnb
                            spin_nb = nbs_idxs[deg, ispin]
                            if spin_nb === ispin
                                ediff += -2.0*b_field*nbs_vals[deg, ispin]*spin_val
                            elseif spin_nb !== 0
                                ediff += -2.0*b_field*spin_val*(nbs_vals[deg, ispin]*spin_lattice[spin_nb, islice])
                            end
                        end
                        # add x components
                        # periodic boundaries
                        if islice == 1
                            tleft = nslices
                            tright = 2
                        elseif islice == nslices
                            tleft = nslices-1
                            tright = 1
                        else
                            tleft = islice-1
                            tright = islice+1
                        end
                        ediff += 2.0*spin_val*(jperp*spin_lattice[ispin, tleft])
                        ediff += 2.0*spin_val*(jperp*spin_lattice[ispin, tright])
                        # add depolar term
                        for cslice = 1:(nslices-1)
                            ediff += lookuptable[cslice] * spin_val * spin_lattice[ispin, (((cslice-1) + islice) % nslices) + 1]
                        end
                        if ediff <= 0.0
                            spin_lattice[ispin, islice] *= -1.0
                        elseif exp(-1.0 * ediff / teff) > rand()
                            spin_lattice[ispin, islice] *= -1.0
                        end
                    end
                end
            end
        end
        # return spin_vector
    end

    function sample(h::Dict, J::Dict; num_reads=10000, temp=0.1,
                    a_sch=1.0 .- collect(LinRange(0, 1, 1000)),
                    b_sch=collect(LinRange(0, 1, 1000)), nslices=64,
                    add_xtalk=false, noisy=false, pre_anneal=false,
                    add_bath=false, coupling_const=0.001,
                    initial_state=Array{Float64, 1}())
        # forward anneal usage only
        independent_qubit_sechdules = false
        nspins = length(h)
        if isa(a_sch, Dict) | isa(b_sch, Dict)
            # if one input is not a dict but one if, convert to dict of lists
            if !isa(a_sch, Dict)
                schedule = deepcopy(a_sch)
                a_sch = Dict([(i, schedule) for i = 1:nspins])
            elseif !isa(b_sch, Dict)
                schedule = deepcopy(b_sch)
                b_sch = Dict([(i, schedule) for i = 1:nspins])
            end

            # check consistent inputs
            @assert length(a_sch) == nspins && length(b_sch) == nspins
            num_steps = length(a_sch[1])
            @assert all([length(a_sch[i]) == num_steps for i = 1:nspins])
            @assert all([length(b_sch[i]) == num_steps for i = 1:nspins])

            independent_qubit_sechdules = true
            a_sch_all = zeros(Float64, nspins, num_steps)
            b_sch_all = zeros(Float64, nspins, num_steps)
            for i = 1:nspins
                a_sch_all[i, :] = a_sch[i]
                b_sch_all[i, :] = b_sch[i]
            end
        end

        if length(initial_state) == length(h)
            @assert length(h) == length(initial_state)
            spin_array = zeros(Float64, nspins, nslices)
            for islice = 1:nslices
                spin_array[:, islice] = deepcopy(initial_state)
            end
            spin_vectors = [deepcopy(spin_array) for i =1:num_reads]
        else
            if pre_anneal
                result = SA.sample(h, J, num_reads=nslices*num_reads, return_raw=true)
                spin_vectors = [Array{Float64, 2}(hcat(result[1+(i-1)*nslices:i*nslices]...)) for i =1:num_reads]
            else
                spin_vectors = rand([-1.0, 1.0], nspins, nslices, num_reads)
                spin_vectors = [spin_vectors[:, :, i] for i =1:num_reads]
            end
        end
        JM, M, degs = hj_info(h, J)
        Js = copy(JM)
        nbs_idxs, nbs_vals = generate_nbr(nspins, M, maximum(degs) + 1)
        if add_bath
            if !noisy & !add_xtalk
                @inbounds for i = 1:num_reads
                    anneal_bath!(spin_vectors[i], nbs_idxs, nbs_vals, a_sch, b_sch, temp, coupling_const)
                end
            elseif noisy & add_xtalk
                @inbounds for i = 1:num_reads
                    hc, jc = add_cross_talk(add_noise(h, J)...)
                    hj_info!(Js, M, degs, hc, jc)
                    generate_nbr!(nbs_idxs, nbs_vals, nspins, M, maximum(degs) + 1)
                    anneal_bath!(spin_vectors[i], nbs_idxs, nbs_vals, a_sch, b_sch, temp, coupling_const)
                end
            elseif add_xtalk & !noisy
                hc, jc = add_cross_talk(h, J)
                hj_info!(Js, M, degs, hc, jc)
                generate_nbr!(nbs_idxs, nbs_vals, nspins, M, maximum(degs) + 1)
                @inbounds for i = 1:num_reads
                    anneal_bath!(spin_vectors[i], nbs_idxs, nbs_vals, a_sch, b_sch, temp, coupling_const)
                end
            else
                @inbounds for i = 1:num_reads
                    hc, jc = add_noise(h, J)
                    hj_info!(Js, M, degs, hc, jc)
                    generate_nbr!(nbs_idxs, nbs_vals, nspins, M, maximum(degs) + 1)
                    anneal_bath!(spin_vectors[i], nbs_idxs, nbs_vals, a_sch, b_sch, temp, coupling_const)
                end
            end
        elseif independent_qubit_sechdules
            if !noisy & !add_xtalk
                @inbounds for i = 1:num_reads
                    annealiq!(spin_vectors[i], nbs_idxs, nbs_vals, a_sch_all, b_sch_all, temp)
                end
            elseif noisy & add_xtalk
                @inbounds for i = 1:num_reads
                    hc, jc = add_cross_talk(add_noise(h, J)...)
                    hj_info!(Js, M, degs, hc, jc)
                    generate_nbr!(nbs_idxs, nbs_vals, nspins, M, maximum(degs) + 1)
                    annealiq!(spin_vectors[i], nbs_idxs, nbs_vals, a_sch_all, b_sch_all, temp)
                end
            elseif add_xtalk & !noisy
                hc, jc = add_cross_talk(h, J)
                hj_info!(Js, M, degs, hc, jc)
                generate_nbr!(nbs_idxs, nbs_vals, nspins, M, maximum(degs) + 1)
                @inbounds for i = 1:num_reads
                    annealiq!(spin_vectors[i], nbs_idxs, nbs_vals, a_sch_all, b_sch_all, temp)
                end
            else
                @inbounds for i = 1:num_reads
                    hc, jc = add_noise(h, J)
                    hj_info!(Js, M, degs, hc, jc)
                    generate_nbr!(nbs_idxs, nbs_vals, nspins, M, maximum(degs) + 1)
                    annealiq!(spin_vectors[i], nbs_idxs, nbs_vals, a_sch_all, b_sch_all, temp)
                end
            end
        else
            if !noisy & !add_xtalk
                @inbounds for i = 1:num_reads
                    anneal!(spin_vectors[i], nbs_idxs, nbs_vals, a_sch, b_sch, temp)
                end
            elseif noisy & add_xtalk
                @inbounds for i = 1:num_reads
                    hc, jc = add_cross_talk(add_noise(h, J)...)
                    hj_info!(Js, M, degs, hc, jc)
                    generate_nbr!(nbs_idxs, nbs_vals, nspins, M, maximum(degs) + 1)
                    anneal!(spin_vectors[i], nbs_idxs, nbs_vals, a_sch, b_sch, temp)
                end
            elseif add_xtalk & !noisy
                hc, jc = add_cross_talk(h, J)
                hj_info!(Js, M, degs, hc, jc)
                generate_nbr!(nbs_idxs, nbs_vals, nspins, M, maximum(degs) + 1)
                @inbounds for i = 1:num_reads
                    anneal!(spin_vectors[i], nbs_idxs, nbs_vals, a_sch, b_sch, temp)
                end
            else
                @inbounds for i = 1:num_reads
                    hc, jc = add_noise(h, J)
                    hj_info!(Js, M, degs, hc, jc)
                    generate_nbr!(nbs_idxs, nbs_vals, nspins, M, maximum(degs) + 1)
                    anneal!(spin_vectors[i], nbs_idxs, nbs_vals, a_sch, b_sch, temp)
                end
            end
        end
        return pimc_aggregation(spin_vectors, num_reads, collect(UpperTriangular(JM)))
    end

    export sample, anneal!, annealiq!
end

max_stars_repo_name: dtoconnor/Annealing-Samplers

max_issues_repo_name: dtoconnor/Annealing-Samplers

max_forks_repo_name: dtoconnor/Annealing-Samplers


text:
module VCFTools

using CodecZlib
using Distributions
using HypothesisTests
using ProgressMeter
using DelimitedFiles
using Dates
import GeneticVariation.VCF

export conformgt_by_id, conformgt_by_pos,
    convert_gt, copy_gt!,
    filter_genotype, gtstats,
    nrecords, nsamples, openvcf,
    convert_ht, copy_ht!,
    filter, filter_header, convert_ds,
    mask_gt, geno_ismissing,
    copy_gt_as_is!, copy_ht_as_is!

# package code goes here
include("gtstats.jl")
include("conformgt.jl")
include("convert.jl")
include("filter.jl")

end # module


text:
module TestVariableVectorize

using Test

using MathOptInterface
const MOI = MathOptInterface

function runtests()
    for name in names(@__MODULE__; all = true)
        if startswith("$(name)", "test_")
            @testset "$(name)" begin
                getfield(@__MODULE__, name)()
            end
        end
    end
    return
end

include("../utilities.jl")

function test_get_scalar_constraint()
    mock = MOI.Utilities.MockOptimizer(
        MOI.Utilities.UniversalFallback(MOI.Utilities.Model{Float64}()),
    )
    bridged_mock = MOI.Bridges.Variable.Vectorize{Float64}(mock)
    x, cx = MOI.add_constrained_variable(bridged_mock, MOI.GreaterThan(1.0))
    func = 2.0 * x
    set = MOI.GreaterThan(5.0)
    err =
        MOI.ScalarFunctionConstantNotZero{Float64,typeof(func),typeof(set)}(1.0)
    @test_throws err MOI.add_constraint(bridged_mock, func + 1.0, set)

    c = MOI.add_constraint(bridged_mock, func, set)
    @test MOI.get(bridged_mock, MOI.ConstraintFunction(), c) ≈ func
    @test MOI.get(bridged_mock, MOI.ConstraintSet(), c) == set
    MOI.set(bridged_mock, MOI.ConstraintName(), c, "c")

    @testset "Mock model" begin
        MOI.set(
            mock,
            MOI.VariableName(),
            MOI.get(mock, MOI.ListOfVariableIndices()),
            ["y"],
        )
        MOI.set(
            mock,
            MOI.ConstraintName(),
            MOI.get(
                mock,
                MOI.ListOfConstraintIndices{
                    MOI.VectorOfVariables,
                    MOI.Nonnegatives,
                }(),
            ),
            ["cy"],
        )
        s = """
        variables: y
        cy: [y] in MathOptInterface.Nonnegatives(1)
        c: 2.0y >= 3.0
        """
        model = MOI.Utilities.Model{Float64}()
        MOI.Utilities.loadfromstring!(model, s)
        MOI.Test.util_test_models_equal(mock, model, ["y"], ["cy", "c"])
    end
    @testset "Bridged model" begin
        MOI.set(bridged_mock, MOI.VariableName(), x, "x")
        s = """
        variables: x
        x >= 1.0
        c: 2.0x >= 5.0
        """
        model = MOI.Utilities.Model{Float64}()
        MOI.Utilities.loadfromstring!(model, s)
        MOI.Test.util_test_models_equal(
            bridged_mock,
            model,
            ["x"],
            ["c"],
            [("x", MOI.GreaterThan{Float64}(1.0))],
        )
    end
end

function test_exp3_with_add_constrained_variable_y()
    mock = MOI.Utilities.MockOptimizer(
        MOI.Utilities.UniversalFallback(MOI.Utilities.Model{Float64}()),
    )
    bridged_mock = MOI.Bridges.Variable.Vectorize{Float64}(mock)
    mock.optimize! =
        (mock::MOI.Utilities.MockOptimizer) -> MOI.Utilities.mock_optimize!(
            mock,
            [log(5), 0.0],
            (MOI.ScalarAffineFunction{Float64}, MOI.LessThan{Float64}) =>
                [0.0],
            (MOI.VectorAffineFunction{Float64}, MOI.ExponentialCone) =>
                [[-1.0, log(5) - 1, 1 / 5]],
        )

    MOI.empty!(bridged_mock)
    x = MOI.add_variable(bridged_mock)
    @test MOI.get(bridged_mock, MOI.NumberOfVariables()) == 1
    xc = MOI.add_constraint(bridged_mock, 2.0x, MOI.LessThan(4.0))
    y, yc = MOI.add_constrained_variable(bridged_mock, MOI.LessThan(5.0))
    @test yc.value == y.value == -1
    @test MOI.get(bridged_mock, MOI.NumberOfVariables()) == 2
    @test length(MOI.get(bridged_mock, MOI.ListOfVariableIndices())) == 2
    @test Set(MOI.get(bridged_mock, MOI.ListOfVariableIndices())) == Set([x, y])
    ec = MOI.add_constraint(
        bridged_mock,
        MOI.Utilities.operate(vcat, Float64, x, 1.0, y),
        MOI.ExponentialCone(),
    )

    MOI.optimize!(bridged_mock)
    @test MOI.get(bridged_mock, MOI.VariablePrimal(), x) ≈ log(5)
    @test MOI.get(bridged_mock, MOI.VariablePrimal(), y) ≈ 5.0
    @test MOI.get(bridged_mock, MOI.ConstraintPrimal(), xc) ≈ 2log(5)
    @test MOI.get(bridged_mock, MOI.ConstraintPrimal(), yc) ≈ 5
    @test MOI.get(bridged_mock, MOI.ConstraintPrimal(), ec) ≈ [log(5), 1.0, 5.0]
    @test MOI.get(bridged_mock, MOI.ConstraintDual(), xc) ≈ 0.0
    @test MOI.get(bridged_mock, MOI.ConstraintDual(), yc) ≈ -1 / 5
    @test MOI.get(bridged_mock, MOI.ConstraintDual(), ec) ≈
          [-1.0, log(5) - 1, 1 / 5]

    err = ErrorException(
        "Cannot add two `VariableIndex`-in-`MathOptInterface.LessThan{Float64}`" *
        " on the same variable MathOptInterface.VariableIndex(-1).",
    )
    @test_throws err MOI.add_constraint(bridged_mock, y, MOI.LessThan(4.0))

    cis = MOI.get(
        bridged_mock,
        MOI.ListOfConstraintIndices{
            MOI.VectorAffineFunction{Float64},
            MOI.ExponentialCone,
        }(),
    )
    @test length(cis) == 1

    err = ArgumentError(
        "Variable bridge of type `$(MathOptInterface.Bridges.Variable.VectorizeBridge{Float64,MathOptInterface.Nonpositives})`" *
        " does not support accessing the attribute `MathOptInterface.Test.UnknownVariableAttribute()`.",
    )
    @test_throws err MOI.get(
        bridged_mock,
        MOI.Test.UnknownVariableAttribute(),
        y,
    )
    ci = MOI.ConstraintIndex{MOI.VariableIndex,MOI.LessThan{Float64}}(y.value)
    attr = MOI.ConstraintSet()
    err = MOI.SetAttributeNotAllowed(
        attr,
        "The variable `MathOptInterface.VariableIndex(12345676)` is bridged by the `VectorizeBridge`.",
    )
    @test_throws err MOI.set(bridged_mock, attr, ci, MOI.LessThan(4.0))

    change = MOI.MultirowChange(y, [(3, 0.0)])
    message =
        "The change $change" *
        " contains variables bridged into a function with nonzero constant."
    err = MOI.ModifyConstraintNotAllowed(cis[1], change, message)
    @test_throws err MOI.modify(bridged_mock, cis[1], change)

    change = MOI.ScalarCoefficientChange(y, 0.0)
    attr = MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}()
    message =
        "The change MathOptInterface.ScalarCoefficientChange{Float64}(MathOptInterface.VariableIndex(-1), 0.0)" *
        " contains variables bridged into a function with nonzero constant."
    err = MOI.ModifyObjectiveNotAllowed(change, message)
    @test_throws err MOI.modify(bridged_mock, attr, change)

    MOI.set(bridged_mock, MOI.VariableName(), x, "x")
    MOI.set(bridged_mock, MOI.ConstraintName(), xc, "xc")
    MOI.set(bridged_mock, MOI.ConstraintName(), ec, "ec")
    z = MOI.get(mock, MOI.ListOfVariableIndices())[2]

    MOI.set(mock, MOI.VariableName(), z, "z")
    MOI.set(
        mock,
        MOI.ConstraintName(),
        MOI.get(
            mock,
            MOI.ListOfConstraintIndices{
                MOI.VectorOfVariables,
                MOI.Nonpositives,
            }(),
        ),
        ["zc"],
    )
    s = """
    variables: x, z
    zc: [z] in MathOptInterface.Nonpositives(1)
    xc: 2.0x <= 4.0
    ec: [x, 1.0, z + 5.0] in MathOptInterface.ExponentialCone()
    """
    model = MOI.Utilities.Model{Float64}()
    MOI.Utilities.loadfromstring!(model, s)
    MOI.Test.util_test_models_equal(mock, model, ["x", "z"], ["zc", "xc", "ec"])

    MOI.set(bridged_mock, MOI.VariableName(), y, "y")
    s = """
    variables: x, y
    y <= 5.0
    xc: 2.0x <= 4.0
    ec: [x, 1.0, y] in MathOptInterface.ExponentialCone()
    """
    model = MOI.Utilities.Model{Float64}()
    MOI.Utilities.loadfromstring!(model, s)
    MOI.Test.util_test_models_equal(
        bridged_mock,
        model,
        ["x", "y"],
        ["xc", "ec"],
        [("y", MOI.LessThan{Float64}(5.0))],
    )

    @test MOI.supports(
        bridged_mock,
        MOI.VariablePrimalStart(),
        MOI.VariableIndex,
    )
    @test MOI.supports(
        bridged_mock,
        MOI.VariablePrimalStart(),
        typeof(MOI.Bridges.bridge(bridged_mock, y)),
    )
    MOI.set(bridged_mock, MOI.VariablePrimalStart(), y, 1.0)
    @test MOI.get(mock, MOI.VariablePrimalStart(), z) == -4
    @test MOI.get(bridged_mock, MOI.VariablePrimalStart(), y) == 1

    _test_delete_bridged_variable(
        bridged_mock,
        y,
        MOI.LessThan{Float64},
        2,
        ((MOI.VectorOfVariables, MOI.Nonpositives, 0),),
    )
    return
end

end  # module

TestVariableVectorize.runtests()


text:
using SpecialFunctions

@testset "math" begin
    @testset "log10" begin
        @test testf(a->log10.(a), Float32[100])
    end

    @testset "pow" begin
        for T in (Float16, Float32, Float64, ComplexF32, ComplexF64)
            range = (T<:Integer) ? (T(5):T(10)) : T
            @test testf((x,y)->x.^y, rand(Float32, 1), rand(range, 1))
            @test testf((x,y)->x.^y, rand(Float32, 1), -rand(range, 1))
        end
    end

    @testset "isinf" begin
      for x in (Inf32, Inf, NaN32, NaN)
        @test testf(x->isinf.(x), [x])
      end
    end

    @testset "isnan" begin
      for x in (Inf32, Inf, NaN32, NaN)
        @test testf(x->isnan.(x), [x])
      end
    end

    for op in (exp, angle, exp2, exp10,)
        @testset "$op" begin
            for T in (Float16, Float32, Float64)
                @test testf(x->op.(x), rand(T, 1))
                @test testf(x->op.(x), -rand(T, 1))
            end
        end
    end

    for op in (expm1,)
        @testset "$op" begin
            # FIXME: add expm1(::Float16) to Base
            for T in (Float32, Float64)
                @test testf(x->op.(x), rand(T, 1))
                @test testf(x->op.(x), -rand(T, 1))
            end
        end
    end

    for op in (exp, abs, abs2, angle, log)
        @testset "Complex - $op" begin
            for T in (ComplexF16, ComplexF32, ComplexF64)
                @test testf(x->op.(x), rand(T, 1))
                @test testf(x->op.(x), -rand(T, 1))
            end

        end
    end
    @testset "mod and rem" begin
        for T in (Float16, Float32, Float64)
            @test testf(a->rem.(a, T(2)), T[0, 1, 1.5, 2, -1])
            @test testf(a->rem.(a, T(2), RoundNearest), T[0, 1, 1.5, 2, -1])
            @test testf(a->mod.(a, T(2)), T[0, 1, 1.5, 2, -1])
        end
    end

    @testset "rsqrt" begin
        # GPUCompiler.jl#173: a CUDA-only device function fails to validate
        function kernel(a)
            a[] = CUDA.rsqrt(a[])
            return
        end

        # make sure this test uses an actual device function
        @test_throws ErrorException kernel(ones(1))

        for T in (Float16, Float32)
            a = CuArray{T}([4])
            @cuda kernel(a)
            @test Array(a) == [0.5]
        end
    end

    @testset "fma" begin
        for T in (Float16, Float32, Float64)
            @test testf((x,y,z)->fma.(x,y,z), rand(T, 1), rand(T, 1), rand(T, 1))
            @test testf((x,y,z)->fma.(x,y,z), rand(T, 1), -rand(T, 1), -rand(T, 1))
        end
    end

    # something from SpecialFunctions.jl
    @testset "erf" begin
        @test testf(a->SpecialFunctions.erf.(a), Float32[1.0])
    end

    @testset "exp" begin
        # JuliaGPU/CUDA.jl#1085: exp uses Base.sincos performing a global CPU load
        @test testf(x->exp.(x), [1e7im])
    end
end


text:
"""
    FPTInfo{S,T}

The struct
```
struct FPTInfo{S,T}
    condCoord::NTuple{N,S}
    upCrossing::NTuple{N,Bool}
    autoRenewed::NTuple{N,Bool}
    reset::NTuple{N,T}
end
```
serves as a container for the information regarding first passage time
observations. `condCoord` is an NTuple of coordinates that are conditioned on
the first passage time nature of the observations. `upCrossing` indicates
whether observations of the corresponding coordinate are up-crossings or
down-crossings. `autoRenewed` indicates whether process starts from the
renewed state (i.e. normally the process is unconstrained until it hits level
`reset` for the first time, however `autoRenewed` process is constrained on the
first passage time from the very beginnig). `reset` level is the level that
needs to be reached before the process starts to be conditioned on the first
passage time.
"""
struct FPTInfo{S,T,N}
    condCoord::NTuple{N,S}
    upCrossing::NTuple{N,Bool}
    autoRenewed::NTuple{N,Bool}
    reset::NTuple{N,T}

    FPTInfo(condCoord::NTuple{N,S}, upCrossing::NTuple{N,Bool},
            reset::NTuple{N,T},
            autoRenewed::NTuple{N,Bool} = Tuple(fill(false,length(condCoord)))
            ) where {S,T,N} = new{S,T,N}(condCoord, upCrossing,
                                         autoRenewed, reset)
end

"""
    checkSingleCoordFpt(XXᵒ, c, cidx, fpt)

Verify whether coordinate `c` (with index number `cidx`) of path `XXᵒ`.yy
adheres to the first passage time observation scheme specified by the object
`fpt`.
"""
function checkSingleCoordFpt(XXᵒ, c, cidx, fpt)
    k = length(XXᵒ.yy)
    thrsd = XXᵒ.yy[end][c]
    renewed = fpt.autoRenewed[cidx]
    s = fpt.upCrossing[cidx] ? 1 : -1
    for i in 1:k
        if !renewed && (s*XXᵒ.yy[i][c] <= s*fpt.reset[cidx])
            renewed = true
        elseif renewed && (s*XXᵒ.yy[i][c] > s*thrsd)
            return false
        end
    end
    return renewed
end

"""
    checkFpt(::PartObs, XXᵒ, fpt)

First passage time constrains are automatically satisfied for the partially
observed scheme
"""
checkFpt(::PartObs, XXᵒ, fpt) = true

"""
    checkFpt(::FPT, XXᵒ, fpt)

Verify whether path `XXᵒ`.yy adheres to the first passage time observation
scheme specified by the object `fpt`.
"""
function checkFpt(::FPT, XXᵒ, fpt)
    for (cidx, c) in enumerate(fpt.condCoord)
        if !checkSingleCoordFpt(XXᵒ, c, cidx, fpt)
            return false
        end
    end
    return true
end

"""
    check_full_path_fpt(::PartObs, ::Any, ::Any, ::Any)

First passage time constrains are automatically satisfied for the partially
observed scheme
"""
check_full_path_fpt(::PartObs, ::Any, ::Any, ::Any) = true

"""
    check_full_path_fpt(::PartObs, XXᵒ, m, fpt)

Verify whether all paths in the range `iRange`, i.e. `XXᵒ`[i].yy, i in `iRange`
adhere to the first passage time observation scheme specified by the object
`fpt`
"""
function check_full_path_fpt(::FPT, XXᵒ, iRange, fpt)
    for i in iRange
        if !checkFpt(FPT(), XXᵒ[i], fpt[i])
            return false
        end
    end
    return true
end


text:
## test_spectral.jl : test spectral method for GRF generation

@testset "spectral decomposition   " begin

## 1d Mat\'ern ##
cov = CovarianceFunction(1,Matern(0.3,1))
pts = 0:0.01:1
grf = GaussianRandomField(cov,Spectral(),pts)
@test isa(grf,GaussianRandomField)
@test isa(grf.cov,CovarianceFunction)
@test isa(grf.cov.cov,Matern)
@test ndims(grf.cov) == 1
@test isa(grf,GaussianRandomField{C,Spectral} where {C})
@test length(grf.pts[1]) == length(pts)
@test length(sample(grf)) == length(pts)

## 2d Exponential ##
cov = CovarianceFunction(2,Exponential(1))
pts1 = 0:0.05:1
pts2 = 0:0.05:1
grf = GaussianRandomField(cov,Spectral(),pts1,pts2)
@test isa(grf,GaussianRandomField)
@test isa(grf.cov,CovarianceFunction)
@test isa(grf.cov.cov,Exponential)
@test ndims(grf.cov) == 2
@test isa(grf,GaussianRandomField{C,Spectral} where {C})
@test length(grf.pts) == 2
@test length(grf.pts[1]) == length(pts1)
@test length(grf.pts[2]) == length(pts2)
@test size(sample(grf),1) == length(pts1)
@test size(sample(grf),2) == length(pts2)

## 3d Gaussian ##
cov = CovarianceFunction(3,Matern(1,2.5))
pts1 = 0:0.1:1
pts2 = 0:0.1:1
pts3 = 0:0.1:1
grf = GaussianRandomField(cov,Spectral(),pts1,pts2,pts3)
@test isa(grf,GaussianRandomField)
@test isa(grf.cov,CovarianceFunction)
@test isa(grf.cov.cov,Matern)
@test ndims(grf.cov) == 3
@test isa(grf,GaussianRandomField{C,Spectral} where {C})
@test length(grf.pts) == 3
@test length(grf.pts[1]) == length(pts1)
@test length(grf.pts[2]) == length(pts2)
@test length(grf.pts[3]) == length(pts3)
@test size(sample(grf),1) == length(pts1)
@test size(sample(grf),2) == length(pts2)
@test size(sample(grf),3) == length(pts3)

## non-SPD covariance matrix ##
cov = CovarianceFunction(2,SquaredExponential(0.5))
pts1 = 0:0.05:1
pts2 = 0:0.05:1
@suppress grf = GaussianRandomField(cov,Spectral(),pts1,pts2)
@test isa(grf,GaussianRandomField)
@test isa(grf.cov,CovarianceFunction)
@test isa(grf.cov.cov,SquaredExponential)
@test ndims(grf.cov) == 2
@test isa(grf,GaussianRandomField{C,Spectral} where {C})
@test length(grf.pts) == 2
@test length(grf.pts[1]) == length(pts1)
@test length(grf.pts[2]) == length(pts2)
@test size(sample(grf),1) == length(pts1)
@test size(sample(grf),2) == length(pts2)

## test non-equidistant structured grid  ##
cov = CovarianceFunction(2,Matern(1.,2.5,σ=1.,p=2))
pts1 = sin.(-pi/2:0.2:pi/2)
pts2 = sin.(-pi/2:0.1:pi/2)
grf = GaussianRandomField(cov,Spectral(),pts1,pts2)
@test isa(grf,GaussianRandomField)
@test isa(grf.cov,CovarianceFunction)
@test isa(grf.cov.cov,Matern)
@test ndims(grf.cov) == 2
@test isa(grf,GaussianRandomField{C,Spectral} where {C})
@test length(grf.pts) == 2
@test length(grf.pts[1]) == length(pts1)
@test length(grf.pts[2]) == length(pts2)
@test size(sample(grf),1) == length(pts1)
@test size(sample(grf),2) == length(pts2)

## test anisotropic
cov = CovarianceFunction(2,AnisotropicExponential([1000 0; 0 1000]))
pts1 = range(0,stop = 10,length = 64)
pts2 = range(0,stop = 10,length = 64)
grf = GaussianRandomField(cov,Spectral(),pts1,pts2)
@test isa(grf,GaussianRandomField)
@test isa(grf.cov,CovarianceFunction)
@test isa(grf.cov.cov,AnisotropicExponential)
@test ndims(grf.cov) == 2
@test isa(grf,GaussianRandomField{C,Spectral} where {C})
@test length(grf.pts) == 2
@test length(grf.pts[1]) == length(pts1)
@test length(grf.pts[2]) == length(pts2)
@test size(sample(grf),1) == length(pts1)
@test size(sample(grf),2) == length(pts2)

end


text:
@testset "imfill" begin
    #2 Dimensional case
    image = Bool.([0 0 0 0 0 0 1 0 0 0
         	   0 1 1 1 1 1 0 0 0 0
         	   0 1 0 0 1 1 0 0 0 0
         	   0 1 1 1 0 1 0 0 0 0
         	   0 1 1 1 1 1 0 0 0 0
         	   0 0 0 0 0 0 0 1 1 1
         	   0 0 0 0 0 0 0 1 0 1
         	   0 0 0 0 0 0 0 1 1 1
	       	   0 0 0 0 0 0 0 0 0 0
	           0 0 0 0 0 0 0 0 0 0])
    answer = Bool.([0 0 0 0 0 0 1 0 0 0
         	    0 1 1 1 1 1 0 0 0 0
         	    0 1 1 1 1 1 0 0 0 0
         	    0 1 1 1 1 1 0 0 0 0
         	    0 1 1 1 1 1 0 0 0 0
         	    0 0 0 0 0 0 0 1 1 1
         	    0 0 0 0 0 0 0 1 1 1
         	    0 0 0 0 0 0 0 1 1 1
		    0 0 0 0 0 0 0 0 0 0
	            0 0 0 0 0 0 0 0 0 0])
	@test imfill(.!image, (0,2)) == (.!answer)

	#3 Dimensional case
	A = zeros(Bool,3,size(image)[1],size(image)[2])
	A[1,:,:] = image
	A[2,:,:] = image
	A[3,:,:] = image
	B = zeros(Bool,3,size(answer)[1],size(answer)[2])
	B[1,:,:] = answer
	B[2,:,:] = answer
	B[3,:,:] = answer
	@test imfill(.!A, (0,8)) == (.!B)

	# Complete white image and complete black image case
	img = zeros(Bool,10,10)
	@test imfill(img,(0,10)) == img
	@test imfill(.!img,(0,99)) == .!img
	# Miscellaneous case
	img[5,5] = true
	@test imfill(.!img,(0,10)) == .!img
	@test imfill(.!img,(0,99)) == zeros(Bool,10,10)
end


text:
"""
Bubble积分
qpp = k1+k2; qfs = k3-k2; qex = k1-k3
"""
struct Bubble{T, P}
    lval :: Float64
    #这里面的顺序是α, β, b_1, b_2, k_n, k_1, k_2
    #k1，k2, b1, b2会用来决定qpp,qfs等数值
    V :: Array{Float64, 7}
end


text:
using Plots, Printf, DelimitedFiles, Random, Statistics, ProgressMeter
using MDToolbox
using CUDA

using BSON: @save, @load
@load "afm_generate_test_data.bson" afms models models_template qs myseed sigma_noise nframe

configs = [AfmizeConfig(10.0 * (pi / 180),
    r,
    MDToolbox.Point2D(-250, -200),
    MDToolbox.Point2D(250, 200),
    MDToolbox.Point2D(6.25, 6.25),
    MDToolbox.defaultParameters())
    for r in [10.0, 15.0, 20.0, 25.0, 30.0]]

for iatom = 1:models_template.natom
    push!(models_template.mass, MDToolbox.defaultParameters()[models_template.resname[iatom]])
end

afms_d = cu.(afms)
models_template_d = gpu32(models_template)
qs_d = cu(qs)

# single frame
r = MDToolbox.getafmposterior_gpu(afms_d[1], models_template_d, qs_d, configs)

# all frames with parallel map
#r = pmap(x -> getafmposterior(x, models_template, qs, configs), afms[1:20])

#@save "afm_analyze_test_data.bson" r configs

max_stars_repo_name: matsunagalab/notebooks

max_issues_repo_name: matsunagalab/notebooks

max_forks_repo_name: matsunagalab/notebooks


text:
 module Factorial2

 function factorial2(number::Int64)
     fact = one(number)
     if number%2 == 0
         for m in 1:number
             if m%2 == 0
                 fact *= m
             end
         end
     elseif number%2 == 1
         for m in 1:number
             if m%2 == 1
                 fact *= m
             end
         end
     end
     return fact
 end

 end # module


text:
# Single decision #
# ========================== #
mutable struct DecisionBridge{T} <: MOIB.Variable.AbstractBridge
    decision::Decision{T}
    variable::MOI.VariableIndex

    function DecisionBridge(decision::Decision{T},
                            variable::MOI.VariableIndex) where T
        return new{T}(decision, variable)
    end
end

function MOIB.Variable.bridge_constrained_variable(::Type{DecisionBridge{T}},
                                                   model::MOI.ModelLike,
                                                   set::SingleDecisionSet{T,NoSpecifiedConstraint}) where T
    variable = MOI.add_variable(model)
    return DecisionBridge(set.decision, variable)
end

function MOIB.Variable.bridge_constrained_variable(::Type{DecisionBridge{T}},
                                                   model::MOI.ModelLike,
                                                   set::SingleDecisionSet{T,S}) where {T, S}
    variable, constraint = MOI.add_constrained_variable(model, set.constraint)
    return DecisionBridge(set.decision, variable)
end

function MOIB.Variable.supports_constrained_variable(
    ::Type{DecisionBridge{T}}, ::Type{SingleDecisionSet{T,S}}) where {T, S}
    return true
end

function MOIB.added_constrained_variable_types(::Type{<:DecisionBridge})
    return Tuple{DataType}[]
end
function MOIB.added_constraint_types(::Type{<:DecisionBridge})
    return Tuple{DataType, DataType}[]
end

# Attributes, Bridge acting as a model
MOI.get(bridge::DecisionBridge, ::MOI.NumberOfVariables) = 1

function MOI.get(bridge::DecisionBridge, ::MOI.ListOfVariableIndices)
    return [bridge.variable]
end

# References
function MOI.delete(::MOI.ModelLike, bridge::DecisionBridge)
    MOI.delete(bridge.variable)
end

# Attributes, Bridge acting as a constraint
function MOI.get(::MOI.ModelLike, ::MOI.ConstraintSet,
                 bridge::DecisionBridge)
    return SingleDecisionSet(bridge.decision)
end

function MOI.get(model::MOI.ModelLike, ::MOI.ConstraintPrimal,
                 bridge::DecisionBridge{T}) where T
    if bridge.decision.state == Taken
        return bridge.decision.value
    end
    return MOI.get(model, MOI.VariablePrimal(), bridge.variable)
end

function MOI.get(model::MOI.ModelLike, attr::Union{MOI.VariablePrimal, MOI.VariablePrimalStart},
                 bridge::DecisionBridge{T}) where T
    if bridge.decision.state == Taken
        return bridge.decision.value
    end
    return MOI.get(model, attr, bridge.variable)
end

function MOI.set(model::MOI.ModelLike, attr::MOI.VariablePrimalStart,
                 bridge::DecisionBridge{T}, value) where T
    MOI.set(model, attr, bridge.variable, value)
end

function MOI.set(model::MOI.ModelLike, attr::MOI.ConstraintSet,
                 bridge::DecisionBridge{T}, new_set::SingleDecisionSet{T}) where T
    bridge.decision = new_set.decision
end

function MOIB.bridged_function(bridge::DecisionBridge{T}) where T
    if bridge.decision.state == NotTaken
        # Use mapped variable as standard MOI variable.
        return MOI.ScalarAffineFunction{T}(MOI.SingleVariable(bridge.variable))
    else
        # Give mapped variable as dummy decision to be held by objective/constraint bridges
        # in order to unbridge properly. Give the value of the taken decision in the
        # decision part constant given to the objective/constraint bridges
        return MOI.ScalarAffineFunction{T}([MOI.ScalarAffineTerm(one(T), bridge.variable)],
                                        bridge.decision.value)
    end
end

function MOIB.Variable.unbridged_map(bridge::DecisionBridge, vi::MOI.VariableIndex)
    return (bridge.variable => MOI.SingleVariable(vi),)
end

# Single known decision #
# ========================== #
mutable struct KnownBridge{T} <: MOIB.Variable.AbstractBridge
    known::Decision{T}

    function KnownBridge(known::Decision{T}) where T
        return new{T}(known)
    end
end

function MOIB.Variable.bridge_constrained_variable(::Type{KnownBridge{T}},
                                                   model::MOI.ModelLike,
                                                   set::SingleKnownSet{T}) where T
    return KnownBridge(set.known)
end

function MOIB.Variable.supports_constrained_variable(
    ::Type{KnownBridge{T}}, ::Type{SingleKnownSet{T}}) where T
    return true
end

function MOIB.added_constrained_variable_types(::Type{<:KnownBridge})
    return Tuple{DataType}[]
end
function MOIB.added_constraint_types(::Type{<:KnownBridge})
    return Tuple{DataType, DataType}[]
end

# Attributes, Bridge acting as a model
MOI.get(bridge::KnownBridge, ::MOI.NumberOfVariables) = 0
function MOI.get(bridge::KnownBridge, ::MOI.ListOfVariableIndices)
    return MOI.VariableIndex[]
end

# References
function MOI.delete(::MOI.ModelLike, bridge::KnownBridge)
    # Nothing to do
    return nothing
end

# Attributes, Bridge acting as a constraint
function MOI.get(::MOI.ModelLike, ::MOI.ConstraintSet,
                 bridge::KnownBridge)
    return SingleKnownSet(bridge.known)
end

function MOI.get(model::MOI.ModelLike, ::MOI.ConstraintPrimal,
                 bridge::KnownBridge{T}) where T
    return bridge.known.value
end

function MOI.get(model::MOI.ModelLike, attr::Union{MOI.VariablePrimal, MOI.VariablePrimalStart},
                 bridge::KnownBridge{T}) where T
    return bridge.known.value
end

function MOI.set(model::MOI.ModelLike, attr::MOI.ConstraintSet,
                 bridge::KnownBridge{T}, new_set::SingleKnownSet{T}) where T
    bridge.known = new_set.known
    return nothing
end

function MOIB.bridged_function(bridge::KnownBridge{T}) where T
    # Give the value of the known decision in the
    # known part constant given to the objective/constraint bridges
    return convert(MOI.ScalarAffineFunction{T}, bridge.known.value)
end

function MOIB.Variable.unbridged_map(bridge::KnownBridge, vi::MOI.VariableIndex)
    # Ignore unbridging without errors
    return (vi => MOI.SingleVariable(vi),)
end

# Multiple decisions #
# ========================== #
struct DecisionsBridge{T} <: MOIB.Variable.AbstractBridge
    decisions::Vector{Decision{T}}
    variables::Vector{MOI.VariableIndex}
end

function MOIB.Variable.bridge_constrained_variable(::Type{DecisionsBridge{T}},
                                                   model::MOI.ModelLike,
                                                   set::MultipleDecisionSet{T,NoSpecifiedConstraint}) where T
    variables = MOI.add_variables(model, length(set.decisions))
    return DecisionsBridge(set.decisions, variables)
end

function MOIB.Variable.bridge_constrained_variable(::Type{DecisionsBridge{T}},
                                                   model::MOI.ModelLike,
                                                   set::MultipleDecisionSet{T,S}) where {T, S}
    variables, constraints = MOI.add_constrained_variables(model, set.constraint)
    return DecisionsBridge(set.decisions, variables)
end

function MOIB.Variable.supports_constrained_variable(
    ::Type{DecisionsBridge{T}}, ::Type{MultipleDecisionSet{T,S}}) where {T, S}
    return true
end

function MOIB.added_constrained_variable_types(::Type{<:DecisionsBridge})
    return Tuple{DataType}[]
end
function MOIB.added_constraint_types(::Type{<:DecisionsBridge})
    return Tuple{DataType, DataType}[]
end

# Attributes, Bridge acting as a model
MOI.get(bridge::DecisionsBridge, ::MOI.NumberOfVariables) = length(bridge.variables)

function MOI.get(bridge::DecisionsBridge, ::MOI.ListOfVariableIndices)
    return bridge.variables
end

# References
function MOI.delete(model::MOI.ModelLike, bridge::DecisionsBridge)
    MOI.delete(model, bridge.variables)
end

function MOI.delete(model::MOI.ModelLike, bridge::DecisionsBridge, i::MOIB.Variable.IndexInVector)
    MOI.delete(model, bridge.variables[i.value])
    deleteat!(bridge.variables, i.value)
end

# Attributes, Bridge acting as a constraint
function MOI.get(::MOI.ModelLike, ::MOI.ConstraintSet,
                 bridge::DecisionsBridge)
    return MultipleDecisionSet(bridge.decisions)
end

function MOI.get(model::MOI.ModelLike, ::MOI.ConstraintPrimal,
                 bridge::DecisionsBridge{T}) where T
    return [decision.state == Taken ? decision_value.(bridge.decisions) :
            MOI.get(model,
                    MOI.VariablePrimal(),
                    bridge.variables[i]) for (i,decision) in bridge.decisions]
end

function MOI.get(model::MOI.ModelLike,
                 attr::Union{MOI.VariablePrimal, MOI.VariablePrimalStart},
                 bridge::DecisionsBridge, i::MOIB.Variable.IndexInVector)
    decision = bridge.decisions[i.value]
    if decision.state == Taken
        return decision.value
    else
        return MOI.get(model, attr, bridge.variables[i.value])
    end
end

function MOI.set(model::MOI.ModelLike, attr::MOI.VariablePrimalStart,
                 bridge::DecisionsBridge{T}, val, i::MOIB.Variable.IndexInVector) where T
    MOI.set(model, attr, bridge.variables[i.value], val)
end

function MOI.set(model::MOI.ModelLike, attr::MOI.ConstraintSet,
                 bridge::DecisionsBridge{T}, new_set::MultipleDecisionSet{T}) where T
    bridge.decisions .= new_set.decisions
end

function MOIB.bridged_function(bridge::DecisionsBridge{T}, i::MOIB.Variable.IndexInVector) where T
    decision = bridge.decisions[i.value]
    if decision.state == NotTaken
        # Use mapped variable as standard MOI variable.
        return MOI.ScalarAffineFunction{T}(MOI.SingleVariable(bridge.variables[i.value]))
    else
        # Give mapped variable as dummy decision to be held by objective/constraint bridges
        # in order to unbridge properly. Give the value of the taken decision in the
        # decision part constant given to the objective/constraint bridges
        return MOI.ScalarAffineFunction{T}([MOI.ScalarAffineTerm(one(T), bridge.variables[i.value])],
                                           decision.value)
    end
end
function MOIB.Variable.unbridged_map(bridge::DecisionsBridge, vi::MOI.VariableIndex, i::MOIB.Variable.IndexInVector)
    return (bridge.variables[i.value] => MOI.SingleVariable(vi),)
end

# Multiple knowns #
# ========================== #
struct KnownsBridge{T} <: MOIB.Variable.AbstractBridge
    knowns::Vector{Decision{T}}
end

function MOIB.Variable.bridge_constrained_variable(::Type{KnownsBridge{T}},
                                                   model::MOI.ModelLike,
                                                   set::MultipleKnownSet{T}) where T
    return KnownsBridge(set.knowns)
end

function MOIB.Variable.supports_constrained_variable(
    ::Type{KnownsBridge{T}}, ::Type{MultipleKnownSet{T}}) where T
    return true
end

function MOIB.added_constrained_variable_types(::Type{<:KnownsBridge})
    return Tuple{DataType}[]
end
function MOIB.added_constraint_types(::Type{<:KnownsBridge})
    return Tuple{DataType, DataType}[]
end

# Attributes, Bridge acting as a model
MOI.get(bridge::KnownsBridge, ::MOI.NumberOfVariables) = 0
function MOI.get(bridge::KnownsBridge, ::MOI.ListOfVariableIndices)
    return MOI.VariableIndex[]
end

# References
function MOI.delete(model::MOI.ModelLike, bridge::KnownsBridge)
    # Nothing to do
    return nothing
end

function MOI.delete(model::MOI.ModelLike, bridge::KnownsBridge, i::MOIB.Variable.IndexInVector)
    # Nothing to do
    return nothing
end

# Attributes, Bridge acting as a constraint
function MOI.get(::MOI.ModelLike, ::MOI.ConstraintSet,
                 bridge::KnownsBridge)
    return MultipleKnownSet(bridge.knowns)
end

function MOI.get(model::MOI.ModelLike, ::MOI.ConstraintPrimal,
                 bridge::KnownsBridge{T}) where T
    return map(bridge.knowns) do known
        known.value
    end
end

function MOI.get(model::MOI.ModelLike,
                 attr::Union{MOI.VariablePrimal, MOI.VariablePrimalStart},
                 bridge::KnownsBridge, i::MOIB.Variable.IndexInVector)
    return bridge.knowns[i.value].value
end

function MOI.set(model::MOI.ModelLike, attr::MOI.ConstraintSet,
                 bridge::KnownsBridge{T}, new_set::MultipleDecisionSet{T}) where T
    bridge.knowns .= new_set.knowns
end

function MOIB.bridged_function(bridge::KnownsBridge{T}, i::MOIB.Variable.IndexInVector) where T
    known = bridge.knowns[i.value]
    # Give the value of the known decision in the
    # known part constant given to the objective/constraint bridges
    return convert(MOI.ScalarAffineFunction{T}, known.value)
end
function MOIB.Variable.unbridged_map(bridge::KnownsBridge, vi::MOI.VariableIndex, i::MOIB.Variable.IndexInVector)
    # Ignore unbridging without errors
    return (vi => MOI.SingleVariable(vi),)
end

# Modifications #
# ========================== #
# To enable certain modifications (decision coefficient changes and decision state updates)
# a few methods must be added to MathOptInterface
function MOIB.is_bridged(b::MOIB.AbstractBridgeOptimizer,
                         change::Union{DecisionCoefficientChange, DecisionMultirowChange})
    return MOIB.is_bridged(b, change.decision)
end

function MOIB.is_bridged(b::MOIB.AbstractBridgeOptimizer,
                         change::Union{DecisionStateChange, DecisionsStateChange,
                                       KnownCoefficientChange, KnownMultirowChange,
                                       KnownValueChange, KnownValuesChange})
    # These modifications should not be handled by the variable bridges
    return false
end

function MOIB.modify_bridged_change(b::MOIB.AbstractBridgeOptimizer, obj,
                                    change::DecisionCoefficientChange)
    f = MOIB.bridged_variable_function(b, change.decision)
    # Continue modification with mapped variable
    MOI.modify(b, obj, DecisionCoefficientChange(only(f.terms).variable_index, change.new_coefficient))
    return nothing
end

function MOIB.modify_bridged_change(b::MOIB.AbstractBridgeOptimizer, obj,
                                    change::DecisionMultirowChange)
    f = MOIB.bridged_variable_function(b, change.decision)
    # Continue modification with mapped variable
    MOI.modify(b, obj, DecisionMultirowChange(only(f.terms).variable_index, change.new_coefficients))
    return nothing
end


text:
using GradientSketching
using Test, Random, LinearAlgebra

@time @testset "Projection" begin
    # vector gradient with vector sketches
    h = zeros(2)
    project!(h, 1, [1, 0])
    project!(h, 2, [0, 1])
    @test h ≈ [1, 2]

    y = randn(2)
    Q = qr(randn(2, 2)).Q # orthonormal matrix
    Qy = Q'*y
    project!(h, Qy[1], Q[:, 1])
    project!(h, Qy[2], Q[:, 2])
    @test h ≈ y

    # vector gradient with coordinate sketches
    h = zeros(2)
    project!(h, 10, 1)
    @test h[1] == 10
    @test h[2] == 0
    project!(h, 20, 2)
    @test h[1] == 10
    @test h[2] == 20

    # matrix gradient with vector sketches
    h = zeros(2, 2)
    project!(h, [1, 2], [1, 0])
    project!(h, [3, 4], [0, 1])
    @test h ≈ [1 2; 3 4]

    Y = randn(2, 2)
    QY = Q'*Y
    project!(h, QY[1, :], Q[:, 1])
    project!(h, QY[2, :], Q[:, 2])
    @test h ≈ Y

    # vector gradient with matrix sketches
    h = zeros(2)
    project!(h, [1, 2], [1 0;0 1])
    @test h ≈ [1, 2]

    S = randn(2, 2)
    Sy = S'*y
    project!(h, Sy, S)
    @test h ≈ y

    # matrix gradient with matrix sketches
    h = zeros(2, 2)
    project!(h, [1 2; 3 4], [1 0;0 1])
    @test h ≈ [1 2; 3 4]

    SY = S'*Y
    project!(h, SY, S)
    @test h ≈ Y

    # vector of arrays gradient with vector sketches
    h = [zeros(2, 2), zeros(2, 2)]
    project!(h, ones(2, 2), [1, 0])
    project!(h, 2.0.*ones(2, 2), [0, 1])
    @test h[1] ≈ ones(2, 2)
    @test h[2] ≈ 2.0.*ones(2, 2)

    # same as above but saving to a view
    for hi in h
        hi .= 0
    end
    project!(view(h, :), ones(2, 2), [1, 0])
    project!(view(h, :), 2.0.*ones(2, 2), [0, 1])
    @test h[1] ≈ ones(2, 2)
    @test h[2] ≈ 2.0.*ones(2, 2)

    # vector of arrays gradient with coordinate sketches
    h = [zeros(2, 2), zeros(2, 2)]
    v1 = fill(10, 2, 2)
    v2 = fill(20, 2, 2)
    project!(h, v1, 1)
    @test h[1] ≈ v1
    @test h[2] ≈ zeros(2, 2)
    project!(h, v2, 2)
    @test h[1] ≈ v1
    @test h[2] ≈ v2

    # vector of arrays gradient with matrix sketches
    h = [zeros(2, 2), zeros(2, 2)]
    project!(h, [ones(2, 2), 2.0.*ones(2, 2)], [1 0;0 1])
    @test h[1] ≈ ones(2, 2)
    @test h[2] ≈ 2.0.*ones(2, 2)

    S = randn(2, 2)
    S∇ = [
        S[1, 1] .* ones(2, 2) .+ S[2, 1] .* ones(2, 2),
        S[1, 2] .* ones(2, 2) .+ S[2, 2] .* ones(2, 2),
    ]
    project!(h, S∇, S)
    @test h[1] ≈ ones(2, 2)
    @test h[2] ≈ ones(2, 2)

    # same as the previous, but saving the result to a view
    for hi in h
        hi .= 0
    end
    project!(view(h, :), S∇, S)
    @test h[1] ≈ ones(2, 2)
    @test h[2] ≈ ones(2, 2)

    # try for many random matrices to ensure there are no problem with singular StS matrices
    h = ones(Float64, 5)
    S = zeros(Float64, 5, 10)
    for _ in 1:100
        S .= randn(5, 10)
        h .= 1
        project!(h, S'*h, S)
        @test h ≈ ones(Float64, 5)
    end
end

@time @testset "BiasSEGA" begin

    ### Constructors
    sega = BiasSEGA((10, 5))
    @test eltype(sega) == Float64
    @test size(sega) == (10, 5)

    sega = BiasSEGA{Float32}((10, 5))
    @test eltype(sega) == Float32
    @test size(sega) == (10, 5)

    sega = BiasSEGA(3)
    @test eltype(sega) == Float64
    @test size(sega) == (3, )

    sega = BiasSEGA{Float32}(3)
    @test eltype(sega) == Float32
    @test size(sega) == (3, )

    sega = BiasSEGA(zeros(5))
    @test eltype(sega) == Float64
    @test size(sega) == (5, )

    ### Vector sketches
    sega = BiasSEGA(2)
    ∇ = zeros(2)

    # Project onto the first dimension
    project!(sega, 1, [1, 0])
    @test gradient(sega) ≈ [1, 0]

    # Project onto the second dimension
    project!(sega, 1, [0, 1])
    @test gradient(sega) ≈ [1, 1]

    ### 1-dimensional matrix sketches for 2-dimensional gradients
    sega = BiasSEGA((2, 3))
    ∇ = zeros(2, 3)
    correct = zeros(2, 3)

    # Project onto the first dimension
    correct[1, :] .= 1
    project!(sega, ones(1, 3), [1 0]')
    @test gradient(sega) ≈ correct

    # Project onto the second dimension
    correct[2, :] .= 1
    project!(sega, ones(1, 3), [0 1]')
    @test gradient(sega) ≈ correct

    # Vector sketches for 2-dimensional gradients
    sega = BiasSEGA((2, 3))
    ∇ = zeros(2, 3)
    correct = zeros(2, 3)

    # Project onto the first dimension
    correct[1, :] .= 1
    project!(sega, ones(3), [1, 0])
    @test gradient(sega) ≈ correct

    # Project onto the second dimension
    correct[2, :] .= 1
    project!(sega, ones(3), [0, 1])
    @test gradient(sega) ≈ correct

    ### Matrix sketches
    sega = BiasSEGA((2, 3))
    ∇ = zeros(2, 3)
    correct = ones(2, 3)
    S = [1.0 2; 3 4]
    project!(sega, S'*correct, S)
    @test gradient(sega) ≈ correct

    ### Vector sketches with approximate projection
    sega = BiasSEGA((2, 3))
    ∇ = zeros(2, 3)
    correct = zeros(2, 3)

    # Project onto the first dimension
    correct[1, :] .= 1
    projecta!(sega, ones(3), [1, 0])
    @test gradient(sega) ≈ correct

    # Project onto the second dimension
    correct[2, :] .= 1
    projecta!(sega, ones(3), [0, 1])
    @test gradient(sega) ≈ correct

    ### Matrix sketches for 1-dimensional gradient with approximate projection
    # Test that we get arbitrarily close after a very large number of iterations
    sega = BiasSEGA(2)
    ∇ = zeros(2)
    correct = ones(2)
    S = [1 2; 3 4]
    projecta!(sega, S'*correct, S, γ=10000)
    @test gradient(sega) ≈ correct

    ### Matrix sketches for 2-dimensional gradient with approximate projection
    # Test that we get arbitrarily close after a very large number of iterations
    sega = BiasSEGA((2, 3))
    ∇ = zeros(2, 3)
    correct = ones(2, 3)
    S = [1 2; 3 4]
    projecta!(sega, S'*correct, S, γ=10000)
    @test gradient(sega) ≈ correct
end

@time @testset "SEGA" begin

    ### Constructors
    sega = SEGA((10, 5))
    @test eltype(sega) == Float64
    @test size(sega) == (10, 5)

    sega = SEGA{Float32}((10, 5))
    @test eltype(sega) == Float32
    @test size(sega) == (10, 5)

    sega = SEGA(3)
    @test eltype(sega) == Float64
    @test size(sega) == (3, )

    sega = SEGA{Float32}(3)
    @test eltype(sega) == Float32
    @test size(sega) == (3, )

    sega = SEGA(zeros(5))
    @test eltype(sega) == Float64
    @test size(sega) == (5, )

    sega = SEGA(zeros(5), zeros(5), zeros(5))
    @test eltype(sega) == Float64
    @test size(sega) == (5, )

    ### Matrix sketches
    sega = SEGA((2, 3))
    ∇ = zeros(2, 3)
    correct = ones(2, 3)
    S = [1.0 2; 3 4]
    project!(sega, 1, S'*correct, S)
    unbias!(sega)
    @test gradient(sega) ≈ correct

    ### Test SEGA un-biasing
    h = ones(2)
    sega = SEGA(h)

    project!(sega, 1/2, 0, 1)
    g = unbias!(sega)
    @test g ≈ [-1, 1]

    project!(sega, 1/2, 0, 2)
    g = unbias!(sega)
    @test g ≈ [0, -1]

    ### Same as previous, but for a vector of arrays gradient
    h = [ones(2) for _ in 1:2]
    sega = SEGA(h)

    project!(sega, 1/2, zeros(2), 1)
    g = unbias!(sega)
    @test g ≈ [fill(-1, 2), ones(2)]

    project!(sega, 1/2, zeros(2), 2)
    g = unbias!(sega)
    @test g ≈ [zeros(2), fill(-1, 2)]

    ### Test SEGA un-biasing when minimizing f(x) = x'*Q*x ./ 2 .+ c'*x
    # over many realizations of the starting point and current gradient estimate
    Random.seed!(123)
    Q = Symmetric(randn(2, 2))
    c = randn(2)
    f = (x) -> x'*Q*x ./ 2 .+ c'*x
    ∇ = (x) -> Q*x .+ c
    x = randn(2)
    h = zeros(2)
    s = zeros(2)
    avg = zeros(2)
    nsamples = 100000
    for i in 1:nsamples
        x .= randn(2)
        s .= randn(2)
        h .= randn(2)
        sega = SEGA(h)
        project!(sega, 1/2, dot(s, ∇(x)), s)
        g = unbias!(sega)
        avg .+= g .- ∇(x)
    end
    avg ./= nsamples * length(avg)
    @test avg ≈ zeros(2) atol=1e-2
end