text:
module GAMSFiles

using DataStructures, OffsetArrays

export parsegams, getwithkey, sexpr

include("types.jl")
include("consts.jl")
include("io.jl")
include("lexer.jl")
include("parser.jl")

"""
    modex = parsegams(Module, modname, gams)
    modex = parsegams(Module, filename)

Return an expression defining a module which, when evaluated (`mod = eval(modex)`),
contains a (non-exported) function `f(x)` evaluating the objective function defined in the
corresponding GAMS file.
"""
function parsegams(::Type{Module}, modname::Symbol, gams::Dict{String,Any})
    parseconsts!(gams)
    vars = gams["variables"]
    # Create the computational part of the function body
    preexprs, bodyexprs, solved = parseassignments(gams)
    # "uncheck" the solved flag for initialized variables
    for (v, vinfo) in vars
        if isinitialized(vinfo)
            solved[getname(v)] = false
        end
    end
    # Create the return statement
    retvalstr = gams["minimizing"]
    push!(bodyexprs, :(return $(Symbol(retvalstr))))
    # For the solved variables, allocate as needed. For the unsolved variables,
    # store them in declaration order for destructuring
    uvars = []
    for (v, vinfo) in vars
        if solved[getname(v)]
            if v isa GArray && getname(v) != gams["minimizing"]
                vsym = Symbol(getname(v))
                allocex = allocate_expr(v.indices, gams["sets"])
                pushfirst!(bodyexprs, :($vsym = $allocex))
            end
        else
            push!(uvars, v)
        end
    end
    # As needed, destructure the input vector
    varnames = Vector{Symbol}(undef, length(uvars))
    if length(uvars) == 1 && (v = uvars[1]) isa GArray
        varnames[1] = Symbol(getname(v))
        xin = varnames[1]
        xaxes = getaxes(v.indices, gams["sets"])
    else
        for (i, v) in enumerate(uvars)
            v isa GText || error("destructured inputs require scalar args, got ", v, ".\nPerhaps there is an unused variable in the file?")
            varnames[i] = Symbol(getname(v))
        end
        xin = gensym("x")
        xaxes = (Base.OneTo(length(varnames)),)
        pushfirst!(bodyexprs, Expr(:(=), Expr(:tuple, varnames...), xin))
    end
    szcheck = :(@assert(axes($xin) == $xaxes))
    body = Expr(:block, szcheck, bodyexprs...)
    # Set up the constant expressions
    initexprs = Expr[]
    for p in ("parameters", "tables")
        if haskey(gams, p)
            for (k, v) in gams[p]
                v = v isa Ref ? v[] : v
                push!(initexprs, Expr(:const, :($(Symbol(getname(k))) = $v)))
            end
        end
    end
    append!(initexprs, preexprs)
    modex = quote
        module $modname
        using OffsetArrays
        $gamsfuncs
        $(Expr(:block, initexprs...))
        function objective($xin)
            $body
        end
        end  # module
    end
    # Initialization information
    x0 = fillr(NaN, xaxes)
    lo, up = fillr(-Inf, xaxes), fillr(Inf, xaxes)
    for (i, v) in enumerate(uvars)
        vinfo = vars[v]
        for (prop, val) in vinfo.assignments
            inds = nothing
            if isa(prop, GText)
                inds = (i,)
            elseif isa(prop, GArray)
                inds = map(x->x.val, prop.indices)
            end
            if isa(prop, Union{GText,GArray})
                c = val.val
                if getname(prop) ∈ ("l", "fx")
                    x0[inds...] = c
                elseif getname(prop) == "lo"
                    lo[inds...] = c
                elseif getname(prop) == "up"
                    up[inds...] = c
                end
            end
        end
    end
    return Expr(:toplevel, modex.args[2]), (x0, lo, up)
end
function parsegams(::Type{Module}, filename::AbstractString, gams::Dict{String,Any})
    bname, _ = splitext(basename(filename))
    return parsegams(Module, Symbol(bname), gams)
end
function parsegams(::Type{Module}, filename::AbstractString)
    gams = parsegams(filename)
    bname, _ = splitext(basename(filename))
    return parsegams(Module, Symbol(bname), gams)
end

function parseconsts!(gams::Dict{String,Any})
    sets = get(gams, "sets", nothing)
    if haskey(gams, "parameters")
        params = gams["parameters"]
        for (key, val) in params
            c = allocate(key, sets)
            if !isempty(val)
                if key isa GText
                    c[] = parse(Float64, val)
                else
                    lines = strip.(split(val, '\n'))
                    for line in lines
                        ln = split(line)
                        @assert(length(ln) == 2)
                        c[parse(Int, ln[1])] = parse(Float64, ln[2])
                    end
                end
            end
            params[key] = c
        end
    end
    if haskey(gams, "tables")
        tables = gams["tables"]
        for (key, val) in tables
            @assert(key isa GArray)
            c = allocate(key, sets)
            lines = split(val, '\n')
            firstrow = true
            cols = Int[]
            for line in lines
                isempty(strip(line)) && continue
                if line[1] == '+'
                    firstrow = true
                end
                if firstrow
                    empty!(cols)
                    # Identify the locations of the column ends (labels are right-justified)
                    j = 2
                    while j <= length(line)
                        i, j = bracket_text(line, j)
                        if i <= length(line)
                            push!(cols, parse(Int, line[i:j-1]))
                        end
                    end
                    firstrow = false
                    continue
                end
                i, j = bracket_text(line, 1)
                row = parse(Int, line[i:j-1])
                for k = 1:length(cols)
                    i, j = bracket_text(line, j)
                    c[row,cols[k]] = parse(Float64, line[i:j-1])
                end
            end
            tables[key] = c
        end
    end
    gams
end

function bracket_text(line, i)
    while i <= length(line) && isspace(line[i])
        i = nextind(line, i)
    end
    j = i
    while j <= length(line) && !isspace(line[j]) && (j == i || line[j] != '-')  # sometimes columns have no gap, use the sign as an indicator
        j = nextind(line, j)
    end
    i, j
end

# Process assignments that can be resolved to a constant
function parseassignments!(exprs, asgn::Vector{<:Pair}, params::Dict, sets)
    for (k, v) in asgn
        x, success = evalconst(v)
        if success
            if k isa GArray && all(x->isa(x, GNumber), k.indices)
                # This can be evaluated ahead of time
                c = params[k]
                inds = map(x->x.val, k.indices)
                c[inds...] = x
            elseif k isa GText
                params[k][] = x
            else
                lhs = convert(Expr, k)
                push!(exprs, assignexpr(lhs, x, sets))
            end
        else
            push!(exprs, assignexpr(convert(Expr, k), convert(Expr, v), sets))
        end
    end
    return exprs
end

function parseassignments(gams)
    vars = getvars(gams["variables"])
    sets = get(gams, "sets", Dict{String,Any}())
    preexprs, bodyexprs = Expr[], Expr[]
    if haskey(gams, "parameters") && haskey(gams, "assignments")
        parseassignments!(preexprs, gams["assignments"], gams["parameters"], sets)
    end
    # Anyone using a solver approach (general constrained optimization) might want to
    # skip the next part (see below).
    if haskey(gams, "equations")
        solved = eqs2assigns!(bodyexprs, gams["equations"], vars, sets, gams["minimizing"])
    else
        solved = nothing
    end
    return preexprs, bodyexprs, solved
end

"""
    solved = eqs2assigns!(bodyexprs, equations, vars, sets, solvevar)

Convert `equations`, containing a list of relations, into Julia expressions
corresponding to assignments, storing then in `bodyexprs`. `solvevar` is the variable you
are trying to solve for (e.g., the objective value), and only those entries in `equations`
needed to solve for it are processed. Any processed entry in `equations` is removed.

Returns `solved`, a dictionary indicating the status of all variables listed in `vars`.
`sets` contains the ranges defined as "sets" in the GAMS file.

This is designed to convert simple constraints into assignment statements,
to allow an objective to be evaluated. For example, if you're minimizing `objval`
(`solvevar == "objval"`) it converts the GAMS equation
    eq1..  x + objval =e= 2
into
    objval = 2 - x
and the GAMS equations
    eq1(i)..  2*x(i) - y =g= 0
    eq0..    5*y - objval =e= 0
into
    y = Inf
    for i = irange
        y = min(y, 2*x[i])
    end
    objval = 5*y
"""
function eqs2assigns!(bodyexprs, equations, vars, sets, solvevar)
    # Determine equation(s) that can be solved for needed variables
    eqs = []   # equations in reverse order of evaluation
    neededvars = Set([solvevar])
    solved = Dict([v=>isinitialized(vinfo) for (v,vinfo) in vars])
    todelete_eqs = []  # it's not nice to modify containers that you're...
    todelete_nvs = []  # ...iterating over, so these store items to be killed...
    toadd_nvs    = []  # ...or added, and are processed after the loop is finished.
    while !isempty(equations) && !isempty(neededvars)
        anysucceeded = false
        empty!(todelete_nvs)
        for var in neededvars
            empty!(todelete_eqs)
            for (key, eq) in equations
                eqsol, succeeded = solvefor(eq, var)
                deparen!(eqsol)
                if succeeded
                    push!(eqs, key=>eqsol)
                    solved[var] = true
                    push!(todelete_nvs, var)
                    add_dependents!(toadd_nvs, eqsol.args[2], solved)
                    push!(todelete_eqs, key=>eq)
                    anysucceeded = true
                    # Don't break because we may have multiple constraint equations
                end
            end
            for pr in todelete_eqs
                deletefrom!(equations, pr)
            end
        end
        for var in todelete_nvs
            delete!(neededvars, var)
        end
        for var in toadd_nvs
            push!(neededvars, var)
        end
        anysucceeded || break  # triggered when the only remaining variables are inputs
    end
    @assert(!isempty(eqs))
    sentinels = Dict()
    for (key, eq) in reverse(eqs)
        @assert(eq isa GCall && length(eq.args)==2)
        lhs, rhs, op = eq.args[1], eq.args[2], eqops[getname(eq)]
        eqj = assignexpr(convert(Expr, lhs), convert(Expr, rhs), key, sets, op)
        eqj = replace_reductions!(eqj, sets)
        push!(bodyexprs, eqj)
        sentinels[lhs] = op
    end
    for (key, op) in sentinels
        op == :(=) && continue
        val = op == :> ? -Inf : Inf
        if key isa GText
            init = Expr(:(=), Symbol(key.text), val)
        elseif key isa GArray
            init = Expr(:call, :fill!, Symbol(key.name), val)
        else
            error("unexpected key ", key)
        end
        pushfirst!(bodyexprs, init)
    end
    return solved
end

function deletefrom!(a::Vector, pr::Pair)
    j = 0
    for i = 1:length(a)
        if a[i] == pr
            j = i
        end
    end
    if j != 0
        deleteat!(a, j)
    end
    a
end
deletefrom!(a::AbstractDict, pr::Pair) = delete!(a, first(pr))

function solvefor(eq, solvevar)
    @assert((eq isa GCall) & haskey(eqops, getname(eq)))
    lhs, rhs, op = eq.args[1], eq.args[2], getname(eq)
    onleft, onright = hasvar(lhs, solvevar), hasvar(rhs, solvevar)
    # Make sure exactly one has it
    onleft != onright || return eq, false
    if onright
        lhs, rhs, op = rhs, lhs, flipop(op)
    end
    if lhs isa Parens && length(lhs.args) == 1
        lhs = lhs.args[1]
    end
    if lhs isa Union{GText,GArray} && getname(lhs) == solvevar
        return GCall(op, [lhs, rhs]), true
    end
    if (lhs isa GCall && getname(lhs) ∈ ("+", "-"))
        if length(lhs.args) == 1
            # unary -
            @assert(getname(lhs) == "-")
            neweq = GCall(flipop(op), lhs.args, unaryminus(rhs))
        else
            @assert(length(lhs.args)==2)
            arg1, arg2 = lhs.args
            has1, has2 = hasvar(arg1, solvevar), hasvar(arg2, solvevar)
            has1 != has2 || return eq, false
            if has1
                neweq = GCall(op, [arg1, GCall(flipsign(getname(lhs)), [rhs, arg2])])
            else
                if getname(lhs) == "+"
                    neweq = GCall(op, [arg2, GCall(flipsign(getname(lhs)), [rhs, arg1])])
                else
                    neweq = GCall(flipop(op), [arg2, GCall("-", [arg1, rhs])])
                end
            end
        end
        return solvefor(neweq, solvevar)
    end
    return eq, false
end

function hasvar(expr, varname::AbstractString)
    expr isa GText && return expr.text == varname
    expr isa GArray && return expr.name == varname
    expr isa GNumber && return false
    expr isa Union{GCall,Parens} || error(expr, " unexpected")
    hv = false
    for i = 1:length(expr.args)
        hv |= hasvar(expr.args[i], varname)
        hv && break
    end
    return hv
end

function add_dependents!(neededvars, rhs, solved)
    if rhs isa Union{GCall,Parens}
        for a in rhs.args
            if !get(solved, getname(a), true)
                push!(neededvars, getname(a))
            end
            add_dependents!(neededvars, a, solved)
        end
    else
        if !get(solved, getname(rhs), true)
            push!(neededvars, getname(rhs))
        end
    end
    neededvars
end

function isinitialized(vinfo)
    for (prop, val) in vinfo.assignments
        isa(prop, GText) && prop.text ∈ ("l", "fx") && return true
    end
    return false
end
isinitialized(::Nothing) = false

function assignexpr(lhs, rhs, sets::Dict, op = :(=))
    if op == :(=)
        body = Expr(:(=), lhs, rhs)
    elseif op == :> || op == :<
        body = Expr(:(=), lhs, Expr(:call, op == :> ? :max : :min, lhs, rhs))
    else
        error("operator $op not recognized")
    end
    if lhs isa Expr && lhs.head == :ref
        for j = 2:length(lhs.args)
            s = lhs.args[j]
            if s isa Symbol
                rng = sets[string(s)]
                body = quote
                    for $s = $rng
                        $body
                    end
                end
            end
        end
    end
    return body
end

function assignexpr(lhs, rhs, eqname::AbstractLex, sets::Dict, op = :(=))
    if op == :(=)
        body = Expr(:(=), lhs, rhs)
    elseif op == :> || op == :<
        body = Expr(:(=), lhs, Expr(:call, op == :> ? :max : :min, lhs, rhs))
    else
        error("operator $op not recognized")
    end
    if eqname isa GArray
        for s in eqname.indices
            if !isa(s, GNumber)
                rng = sets[getname(s)]
                itersym = Symbol(getname(s))
                body = quote
                    for $itersym = $rng
                        $body
                    end
                end
            end
        end
    end
    return body
end

flipop(op::Symbol) = op == :(=) ? op : (op == :< ? :> : (op == :> ? :< : error(op, " not recognized")))
flipop(op::String) = op == "=e=" ? op : (op == "=l=" ? "=g=" : (op == "=g=" ? "=l=" : error(op, " not recognized")))

flipsign(op::String) = op == "+" ? "-" : (op == "-" ? "+" : error(op, " is not a sign operator"))

function unaryminus(x)
    # A double-negative is a positive
    if isa(x, GCall) && getname(x) == "-" && length(x.args) == 1
        return x.args[1]
    end
    return GCall("-", x)
end

function allvars(gams)
    totvars = Dict{String,Any}()
    if haskey(gams, "variables")
        getvars!(totvars, gams["variables"])
    end
    if haskey(gams, "parameters")
        getvars!(totvars, gams["parameters"])
    end
    if haskey(gams, "tables")
        getvars!(totvars, gams["tables"])
    end
    return totvars
end

function getvars!(vars, pairiter)
    for (key, val) in pairiter
        vars[getname(key)] = val
    end
    return vars
end
getvars(pairiter) = getvars!(Dict{String,Any}(), pairiter)

function evalconst(expr::AbstractLex)
    isa(expr, GNumber) && return expr.val, true
    if isa(expr, GCall)
        # Apply evalconst to the args of the call
        arg_sxs = map(evalconst, expr.args)
        if all(x->x[2], arg_sxs)
            # we were able to evaluate all the args
            args = map(x->x[1], arg_sxs)
            f = getfield(GamsFuncs, Symbol(expr.name))
            return f(args...), true
        end
    elseif isa(expr, Parens)
        if length(expr.args) == 1
            return evalconst(expr.args[1])
        end
    end
    return NaN, false
end

function getaxes(setnames, sets)
    axs = []
    for name in setnames
        s = sets[string(name)]
        push!(axs, s isa UnitRange ? s : Base.OneTo(length(s)))
    end
    (axs...,)
end

allocate(array::GArray, sets) = allocate(array.indices, sets)
allocate(x::GText, sets) = Ref(0.0)   # a ref so that it can be passed in updatable form

function allocate(setnames, sets)
    axs = getaxes(setnames, sets)
    axs isa Tuple{} && return 0.0
    axs isa Tuple{Base.OneTo{Int},Vararg{Base.OneTo{Int}}} &&
        return Array{Float64}(undef, length.(axs))
    return OffsetArray{Float64}(undef, axs)
end

function allocate_expr(setnames, sets)
    axs = getaxes(setnames, sets)
    axs isa Tuple{} && return 0.0
    axs isa Tuple{Base.OneTo{Int},Vararg{Base.OneTo{Int}}} &&
        return :(Array{Float64}(undef, $(length.(axs))))
    return :(OffsetArray{Float64}(undef, $axs))
end

function splitws(str; rmsemicolon::Bool=false)
    iend = lastindex(str)
    if rmsemicolon && str[iend] == ';'
        iend = prevind(str, iend)
    end
    m = match(r"\s", str)
    m == nothing && return str, ""
    return str[1:m.offset-1], strip(str[m.offset:iend])
end

function replace_reductions!(ex::Expr, sets)
    for i = 1:length(ex.args)
        ex.args[i] = replace_reductions!(ex.args[i], sets)
    end
    if ex.head == :call && ex.args[1] ∈ (:sum, :prod, :smin, :smax)
        @assert(length(ex.args) == 3)
        vars, val = ex.args[2], ex.args[3]
        index_exprs = Expr[]
        if vars isa Symbol
            rng = sets[String(vars)]
            push!(index_exprs, :($vars = $rng))
        else
            for v in vars
                rng = sets[String(v)]
                push!(index_exprs, :($v = $rng))
            end
        end
        return Expr(:call, gamsf2jf[ex.args[1]], Expr(:comprehension, Expr(:generator, val, index_exprs...)))
    end
    return ex
end
replace_reductions!(arg, sets) = arg

fillr(val, inds::Tuple{}) = Ref(val)
fillr(val, inds::Tuple{Base.OneTo,Vararg{Base.OneTo}}) = fill(val, length.(inds))
fillr(val, inds) = fill(val, inds)

end # module


text:
# Returns an array of positive prime numbers less than or equal to lim
function sieve(lim :: Int)
    if lim < 2 return [] end
    limi :: Int = (lim - 1) ÷ 2 # calculate the required array size
    isprime :: Array{Bool} = trues(limi)
    llimi :: Int = (isqrt(lim) - 1) ÷ 2 # and calculate maximum root prime index
    result :: Array{Int} = [2]  #Initial array
    for i in 1:limi
        if isprime[i]
            p = i + i + 1 # 2i + 1
            if i <= llimi
                for j = (p*p-1)>>>1:p:limi # quick shift/divide in case LLVM doesn't optimize divide by 2 away
                    isprime[j] = false
                end
            end
            push!(result, p)
        end
    end
    return result
end

ext: jl

max_stars_repo_name: ethansaxenian/RosettaDecode

max_issues_repo_name: ethansaxenian/RosettaDecode

max_forks_repo_name: ethansaxenian/RosettaDecode


text:
include("MatAcoustFluidModule.jl")
include("FEMMAcoustModule.jl")
include("FEMMAcoustSurfModule.jl")
include("AlgoAcoustModule.jl")


text:
"""
    init_cg_params(CGParamStruct, NetSettings)

Takes a Conjugate Gradient Parameter struct as argument, Initialize the Conjugate
Gradient Parameters.
"""
function init_cg_params(CGParamStruct, NetSettings)
    @unpack dim_rbm = NetSettings
    r = zeros(Complex{Float64}, dim_rbm)
    dir = SharedArray(zeros(Complex{Float64}, dim_rbm))
    dw = zeros(Complex{Float64}, dim_rbm)
    Sdir_w = SharedArray(zeros(Complex{Float64}, dim_rbm, nworkers()))
    Sdir = zeros(Complex{Float64}, dim_rbm)
    res = 0.0
    res_prev = 1.0
    return CGParamStruct(r, dir, dw, Sdir_w, Sdir, res, res_prev)
end

"""
    distributed_o_dir!(Sdir_w, VarDervs, dir)

Multiply the part of adjoint(O)*O*dir that is sampled by processor `myid() - 1`.
Views are used to prevent allocation-expensive copies upon slicing. `Sdir_w` is
updated in-place.
"""
function distributed_o_dir!(Sdir_w, VarDervs, dir)
    @fastmath mul!(view(Sdir_w, :, myid() - 1),  # view to ensure writing to Sdir-slice, instead of a copy
        view(VarDervs.o_tot, :, :, myid() - 1),
        BLAS.gemv('C', view(VarDervs.o_tot, :, :, myid() - 1), dir))
end

"""
    LinearAlgebra.mul!(Sdir, Sdir_w, VarDervs, dir)

Definition of 'VarDervs * dir' multiplication, so that calculation of entire
S-matrix is prevented: S * dir = adjoint(OptParams.O) * [OptParams.O * dir], see
(Carleo, Troyer, 2017). Note that the VARDERIVS type stores adjoint(O) instead
of O itself. This function is called directly from `my_cg!` for calculating
matrix-vector products. This is done in parallel, after which the parallel
results are summed and scaled. The result is stored in-place in `Sdir`.
See also `init_vardervs` and `my_cg!`.
"""
function LinearAlgebra.mul!(Sdir::Array{Complex{Float64}, 1},
        Sdir_w::SharedArray{Complex{Float64}, 2},
        VarDervs::VARDERIVS,
        dir::SharedArray{Complex{Float64}, 1},
        NetSettings)
    @unpack mc_trials = NetSettings
    if nworkers() > 30
        @sync for i=1:nworkers()
            @async remotecall_wait(distributed_o_dir!,
                workers()[i],
                Sdir_w,
                VarDervs,
                dir)
        end
    else
        for i=1:nworkers()
            @fastmath mul!(view(Sdir_w, :, i),  # view to ensure writing to Sdir-slice, instead of a copy
                view(VarDervs.o_tot, :, :, i),
                BLAS.gemv('C', view(VarDervs.o_tot, :, :, i), dir))
        end
    end
    Sdir .= dropdims(sum(Sdir_w, dims=2), dims=2)
    # overwrites Sdir with adjoint(O) * (O * dir)
    lmul!(1 / (mc_trials * nworkers()), Sdir)
    # overwrites Sdir with scaled Sdir
    axpy!(VarDervs.reg, dir, Sdir)
    # overwrites Sdir with regularized Sdir.
end

""" --- NOT USED ---
    diag_precond!(VarDervs, CGP, NetSettings)

Takes the variational derivatives and calculates only the diagonal of the
S-matrix to be used for preconditioning in the CG algorithm.
See also the matrix-free version of `my_cg!`.

"""
function diag_precond!(VarDervs::VARDERIVS, CGP::CGPARAMS, NetSettings)
    @unpack dim_rbm, mc_trials = NetSettings
    for i=1:dim_rbm
        for j=1:nworkers()
            CGP.m[i] += norm(view(VarDervs.o_tot, i, :, j))^2
        end
    end
    CGP.m ./= (mc_trials * nworkers())
end

"""
    my_cg!(CGP, VarDervs, OptParams, lrate, NetSettings[, tol] )

Calculates the change in network parameters iteratively by conjugate gradient.
Matrix-vector products are calculated without an explictly built S-matrix.
`OptParams.dw` is updated in-place.
"""
function my_cg!(CGP::CGPARAMS, VarDervs::VARDERIVS, OptParams::OPTPARAMSMFREE,
    lrate::Float64, NetSettings; tol=1e-3)
    @unpack dim_rbm = NetSettings
    i_max = dim_rbm

    fill!(OptParams.dw_tot, zero(eltype(OptParams.dw_tot)))
    fill!(CGP.dir, zero(eltype(CGP.dir)))
    CGP.res_prev = 1.0

    CGP.r .= -lrate .* OptParams.f
    res_0 = norm(CGP.r)
    CGP.res = res_0

    i = 0
    while i < i_max && CGP.res / res_0 > tol
        β = CGP.res^2 / CGP.res_prev^2
        CGP.dir .= CGP.r .+ β .* CGP.dir
        mul!(CGP.Sdir, CGP.Sdir_w, VarDervs, CGP.dir, NetSettings)
        α = CGP.res^2 / dot(CGP.dir, CGP.Sdir)
        OptParams.dw_tot .+= α .* CGP.dir
        CGP.r .-= α .* CGP.Sdir
        CGP.res_prev = CGP.res
        CGP.res = norm(CGP.r)
        i += 1
    end
    return i
end

"""
    my_cg!(CGP, VarDervs, OptParams, lrate, NetSettings[, tol] )

Calculates the change in network parameters iteratively by conjugate gradient.
Matrix-vector products are calculated with the explicitly built S-matrix.
`OptParams.dw` is updated in-place.
"""
function my_cg!(CGP::CGPARAMS,
        S::Array{Complex{Float64}, 2},
        OptParams::OPTPARAMS,
        lrate::Float64,
        NetSettings;
        tol=1e-3)
    @unpack dim_rbm = NetSettings
    i_max = dim_rbm

    fill!(OptParams.dw_tot, zero(eltype(OptParams.dw_tot)))
    fill!(CGP.dir, zero(eltype(CGP.dir)))
    CGP.res_prev = 1.0

    CGP.r .= -lrate .* OptParams.f
    res_0 = norm(CGP.r)
    CGP.res = res_0

    i = 0
    while i < i_max && (CGP.res / res_0) > tol
        β = CGP.res^2 / CGP.res_prev^2
        CGP.dir .= CGP.r .+ β .* CGP.dir
        mul!(CGP.Sdir, S, CGP.dir)
        α = CGP.res^2 / dot(CGP.dir, CGP.Sdir)
        OptParams.dw_tot .+= α .* CGP.dir
        CGP.r .-= α .* CGP.Sdir
        CGP.res_prev = CGP.res
        CGP.res = norm(CGP.r)
        i += 1
    end
end

# Supporting Base functions for VARDERIVS type.
Base.size(VarDervs::VARDERIVS) = (size(VarDervs.o_tot, 1), size(VarDervs.o_tot, 1)) # should return size of S
Base.size(VarDervs::VARDERIVS, i::Int) = size(VarDervs.o_tot, 1)  # S is square
Base.eltype(VarDervs::VARDERIVS) = eltype(VarDervs.o_tot)


text:
# Shortcuts for now
export reals, bools, integers, unit

bools(θ) = unit(θ) > 0.5
bound(θ, a, b, c) = unit(θ, T) * b - a + a
reals(θ, T = Float64) = (unit(θ, T) - 0.5) * 100.0
integers(θ) = rand(Int64)

const ℝ = reals
const 𝔹 = bools
const ℤ = integers

function choose(branches, θ)
  return 1
end


text:
"""
    conv_bn(kernelsize, inplanes, outplanes;
            stride = 1, pad = 0, usebias = true, rev = false)

Create a convolution + batch normalization pair with ReLU activation.

# Arguments
- `kernelsize`: size of the convolution kernel (tuple)
- `inplanes`: number of input feature maps
- `outplanes`: number of output feature maps
- `stride`: stride of the convolution kernel
- `pad`: padding of the convolution kernel
- `usebias`: set to `true` to use a bias in the convolution layer
- `rev`: set to `true` to place the batch norm before the convolution
"""
function conv_bn(kernelsize, inplanes, outplanes;
                 stride = 1, pad = 0, usebias = true, rev = false)
  conv_layer = []
  if usebias
    push!(conv_layer, Conv(kernelsize, Int(inplanes) => Int(outplanes),
                           stride = stride, pad = pad, init = Flux.kaiming_normal))
  else
    push!(conv_layer, Conv(kernelsize, Int(inplanes) => Int(outplanes),
                           stride = stride,
                           pad = pad,
                           init = Flux.kaiming_normal,
                           bias = Flux.Zeros()))
  end

  if rev
    push!(conv_layer, BatchNorm(Int(inplanes), relu))
    return reverse(conv_layer)
  end

  push!(conv_layer, BatchNorm(Int(outplanes), relu))
  return conv_layer
end

"""
    cat_channels(x, y)

Concatenate `x` and `y` along the channel dimension (third dimension).
Equivalent to `cat(x, y; dims=3)`.
Convenient binary reduction operator for use with `Parallel`.
"""
cat_channels(x, y) = cat(x, y; dims = 3)

"""
    weights(model)

Load the pre-trained weights for `model` using the stored artifacts.
"""
weights(model) = BSON.load(joinpath(@artifact_str(model), "$model.bson"), @__MODULE__)[:weights]
pretrain_error(model) = throw(ArgumentError("No pre-trained weights available for $model."))

function _maybe_big_show(io, model)
  if isdefined(Flux, :_big_show)
    if isnothing(get(io, :typeinfo, nothing)) # e.g. top level in REPL
      Flux._big_show(io, model)
    else
      show(io, model)
    end
  else
    show(io, model)
  end
end


text:
"""
Description

        The fred() method is a wrapper to download financial and economic time series data from the
        St. Louis Federal Reserve (FRED).

Usage

        DGS = fred("DGS10")
        CPI = fred()

Method Signature(s)

      fred(series::ASCIIString="CPIAUCNS")

Details

      The fred method takes a string argument that corresponds to a series code from the St. Louis Federal
      Reserve (FRED) database. It returns the data in the TimeSeries.TimeArray data structure.  When no argument
      is provided, the default data set is the Consumer Price Index for All Urban Consumers: All Items (CPIAUCNS).

References

      https://research.stlouisfed.org/fred2

See Also

      yahoo() which is a wrapper from downloading financial time series for stocks from Yahoo Finance.

"""

function fred(series::ASCIIString="CPIAUCNS")
    resp = get("http://research.stlouisfed.org/fred2/series/$series/downloaddata/$series.csv")
    timearray(resp, m=series)
end


text:
struct DefaultWeightingFunction
    b::Float64
end
function (dwf::DefaultWeightingFunction)(_, y)
    if y ≤ 0
        return 0.0
    elseif 0 < y < dwf.b
        return y / dwf.b
    else
        return 1.0
    end
end

const inv2π = 1.0 / (2π)

struct Binormal
    sigma::Float64
end
function (bi::Binormal)(x, y)
    invsigma2 = 1 / abs2(bi.sigma)
    return inv2π * invsigma2 * exp(-(abs2(x) + abs2(y)) * 0.5 * invsigma2)
end

"""
    PersistenceImage

`PersistenceImage` provides a vectorization method for persistence diagrams. Each point in
the diagram is first transformed into `birth`, `persistence` coordinates. Then, it is
weighted by a weighting function and widened by a distribution (default: gaussian with σ=1).
Once all the points are transformed, their distributions are summed together and discretized
into an image.

The weighting ensures points near the diagonal have a small contribution. This ensures this
representation of the diagram is stable.

Once a `PersistenceImage` is constructed (see below), it can called like a function to
transform a diagram to an image.

Infinite intervals in the diagram are ignored.

# Constructors

    PersistenceImage(ylims, xlims; kwargs...)

Create an image ranging from `ylims[1]` to `ylims[2]` in the ``y`` direction and
equivalently for the ``x`` direction.

    PersistenceImage(diagrams; zero_start=true, margin=0.1, kwargs...)

Learn the ``x`` and ``y`` ranges from diagrams, ensuring all diagrams will fully fit in the
image. Limits are increased by the `margin`. If `zero_start` is true, set the minimum `y`
value to 0. If all intervals in diagrams have the same birth (e.g. in the zeroth dimension),
a single column image is produced.

## Keyword Arguments

* `size`: integer or tuple of two integers. Determines the size of the array containing the
  image. Defaults to 5.

* `distribution`: A function or callable object used to smear each interval in diagram.  Has
  to be callable with two `Float64`s as input and should return a `Float64`. Defaults to a
  normal distribution.

* `sigma`: The width of the normal distribution mentioned above. Only applicable when
  `distribution` is unset. Defaults to twice the size of each pixel.

* `weight`: A function or callable object used as the weighting function. Has to be callable
  with two `Float64`s as input and should return a `Float64`. Should equal 0.0 for x=0, but
  this is not enforced. Defaults to function that is zero at ``y=0``, and increases linearly
  to 1 until `slope_end` is reached.

* `slope_end`: the relative ``y`` value at which the default weight function stops
  increasing. Defaults to 1.0.

# Example

```jldoctest
julia> diag_1 = PersistenceDiagram([(0, 1), (0, 1.5), (1, 2)]);

julia> diag_2 = PersistenceDiagram([(1, 2), (1, 1.5)]);

julia> image = PersistenceImage([diag_1, diag_2])
5×5 PersistenceImage(
  distribution = PersistenceDiagrams.Binormal(0.5499999999999999),
  weight = PersistenceDiagrams.DefaultWeightingFunction(1.65),
)

julia> image(diag_1)
5×5 Matrix{Float64}:
 0.156707  0.164263  0.160452  0.149968  0.133353
 0.344223  0.355089  0.338991  0.308795  0.268592
 0.571181  0.577527  0.535069  0.47036   0.396099
 0.723147  0.714873  0.639138  0.536823  0.432264
 0.700791  0.677237  0.582904  0.46433   0.352962

```

# Reference

Adams, H., Emerson, T., Kirby, M., Neville, R., Peterson, C., Shipman, P., ... &
Ziegelmeier, L. (2017). Persistence images: A stable vector representation of persistent
homology. [The Journal of Machine Learning Research, 18(1), 218-252]
(http://www.jmlr.org/papers/volume18/16-337/16-337.pdf).
"""
struct PersistenceImage{X<:AbstractVector{Float64},Y<:AbstractVector{Float64},D,W}
    ys::Y
    xs::X
    distribution::D
    weight::W
end

function PersistenceImage(
    ys::AbstractArray{Float64},
    xs::AbstractArray{Float64};
    sigma=nothing,
    distribution=nothing,
    weight=nothing,
    slope_end=nothing,
)
    if !issorted(xs) || !issorted(ys)
        throw(ArgumentError("`xs` and `ys` must be increasing"))
    end
    if !isnothing(sigma) && !isnothing(distribution)
        throw(ArgumentError("`sigma` and `distribution` can't be set at the same time"))
    elseif !isnothing(sigma)
        if sigma > 0
            distribution = Binormal(sigma)
        else
            throw(ArgumentError("`sigma` must be positive"))
        end
    elseif !isnothing(distribution)
        distribution = distribution
    else
        w = max((last(ys) - first(ys)) / length(ys), (last(xs) - first(xs)) / length(xs))
        distribution = Binormal(2 * w)
    end
    if !isnothing(weight) && !isnothing(slope_end)
        throw(ArgumentError("`weight` and `slope_end` can't be set at the same time"))
    elseif !isnothing(weight)
        weight = weight
    elseif !isnothing(slope_end)
        if 0 < slope_end ≤ 1
            weight = DefaultWeightingFunction(slope_end * last(ys))
        else
            throw(ArgumentError("`0 < slope_end ≤ 1` does not hold"))
        end
    else
        weight = DefaultWeightingFunction(last(ys))
    end

    return PersistenceImage(ys, xs, distribution, weight)
end
function PersistenceImage(ylims::Tuple, xlims::Tuple; size=5, kwargs...)
    s = length(size) == 1 ? (size, size) : size
    ys = range(ylims[1], ylims[2]; length=s[1] + 1)
    xs = range(xlims[1], xlims[2]; length=s[2] + 1)

    return PersistenceImage(ys, xs; kwargs...)
end
function PersistenceImage(
    diagrams;
    zero_start=true,
    margin=0.1,
    size=5,
    sigma=nothing,
    distribution=nothing,
    kwargs...,
)
    if margin < 0
        throw(ArgumentError("`margin` must be non-negative"))
    end
    ysize, xsize = length(size) == 1 ? (size, size) : size

    finite = Iterators.filter(isfinite, Iterators.flatten(diagrams))
    # Generator used for Julia 1.0 compatibility.
    min_persistence, max_persistence = extrema(persistence(int) for int in finite)
    min_birth, max_birth = extrema(birth(int) for int in finite)
    if min_birth == max_birth
        xsize = 1
    end

    if zero_start
        min_persistence = 0.0
    end
    ywidth = max_persistence - min_persistence
    min_persistence = max(0.0, min_persistence - margin * ywidth)
    max_persistence = max_persistence + margin * ywidth

    xwidth = max_birth - min_birth
    min_birth = min_birth - margin * xwidth
    max_birth = max_birth + margin * xwidth

    ylims = (min_persistence, max_persistence)
    xlims = (min_birth, max_birth)

    return PersistenceImage(
        ylims, xlims; size=(ysize, xsize), sigma=sigma, distribution=distribution, kwargs...
    )
end

output_size(pi::PersistenceImage) = (length(pi.ys) - 1) * (length(pi.xs) - 1)

function Base.show(io::IO, pi::PersistenceImage)
    return print(io, "$(length(pi.ys) - 1)×$(length(pi.xs) - 1) PersistenceImage")
end
function Base.show(io::IO, ::MIME"text/plain", pi::PersistenceImage)
    println(io, pi, "(")
    println(io, "  distribution = ", pi.distribution, ",")
    println(io, "  weight = ", pi.weight, ",")
    return print(io, ")")
end

function (pi::PersistenceImage)(diagram::PersistenceDiagram)
    n = length(pi.ys) - 1
    m = length(pi.xs) - 1
    result = zeros(n, m)
    distribution_buffer = zeros(n + 1, m + 1)

    @inbounds for interval in Iterators.filter(isfinite, diagram)
        b = birth(interval)
        p = persistence(interval)
        w = pi.weight(b, p)

        for j in 1:(m + 1), i in 1:(n + 1)
            x = pi.xs[j]
            y = pi.ys[i]
            distribution_buffer[i, j] = w * pi.distribution(x - b, y - p) * 0.25
        end

        for j in 1:m, i in 1:n
            result[i, j] +=
                distribution_buffer[i, j] +
                distribution_buffer[i + 1, j] +
                distribution_buffer[i, j + 1] +
                distribution_buffer[i + 1, j + 1]
        end
    end
    return result
end


text:
using Combinatorics

"""
one_body_to_matrix and two_body_to_matrix for testing purposes
returns the 2^L×2^L matrices of a_k^†a_l and a_k^† a_l^† a_m a_n
"""
function one_body_to_matrix(k,l,L)
    A = zeros(Float64,2^L,2^L)
    for i in 0:2^L-1
        i_bin = digits(i,base=2,pad=L)
        for j in 0:2^L-1
            j_bin = digits(j,base=2,pad=L) #corresponding index as d-tuples
            if i_bin[k]==1 && j_bin[l]==1 && j-2^(l-1) == i-2^(k-1)
                A[i+1,j+1] = (-1)^(sum(i_bin[1:(k-1)])+sum(j_bin[1:(l-1)]))
            end
        end
    end
    return A
end

function two_body_to_matrix(k,l,m,n,L)
    A = zeros(Float64,2^L,2^L)
    for i in 0:2^L-1
        i_bin = digits(i,base=2,pad=L)
        for j in 0:2^L-1
            j_bin = digits(j,base=2,pad=L) #corresponding index as d-tuples
            if i_bin[k]==1 && i_bin[l]==1 && j_bin[m]==1 && j_bin[n]==1 && j-2^(m-1)-2^(n-1) == i-2^(k-1)-2^(l-1)
                A[i+1,j+1] = (-1)^(sum(i_bin[1:(k-1)])+sum(i_bin[1:(l-1)])+sum(j_bin[1:(m-1)])+sum(j_bin[1:(n-1)]))
            end
        end
    end
    return A
end

#naive implementation

"""
h: two-dimensional array of the connected sites (kinetic+external potential)
V: four-dimensional array of the connected sites (interaction)
"""
function sc_to_mat(h,V)
    L = size(h,1)
    A = zeros(2^L,2^L)
    for i in 1:L
        for j in i:L
            for k in j:L
                for l in k:L
                    if !isapprox(h[i,j],0.0,atol=1e-10)
                        H = one_body_to_matrix(i,j,L)
                        A += h[i,j]*(H+H')
                    end
                    if !isapprox(V[i,j,k,l],0.0,atol=1e-10)
                        H = two_body_to_matrix(i,j,k,l,L)
                        A += V[i,j,k,l]*(H+H')
                    end
                end
            end
        end
    end
    return A
end

#odd index = spin up, even index = spin down
#returns TT of a half-filled state
function half_filling(N)
    tt_vec = Array{Array{Float64,3},1}(undef,2N)
    for i in 1:N
        tt_vec[2i-1] = zeros(2,1,1)
        tt_vec[2i-1][1,1,1] = 1.0
        tt_vec[2i] = zeros(2,1,1)
        tt_vec[2i][2,1,1] = 1.0
    end
    return ttvector{Float64}(tt_vec,2*ones(Int,2N),ones(Int,2N+1),zeros(2N))
end

#auxiliary functions for a_p^†a_q
function mpo_core_id(;T=Float64)
    out = zeros(T,2,2,1,1)
    out[1,1,1,1] = 1.0
    out[2,2,1,1] = 1.0
    return out
end

function mpo_core_ferm_sign(;T=Float64)
    out = zeros(T,2,2,1,1)
    out[1,1,1,1] = 1.0
    out[2,2,1,1] = -1.0
    return out
end

function mpo_core_creation(;T=Float64)
    out = zeros(T,2,2,1,1)
    out[2,1,1,1] = 1.0
    return out
end

function mpo_core_annihilation(;T=Float64)
    out = zeros(T,2,2,1,1)
    out[1,2,1,1] = 1.0
    return out
end

#returns bosonic or fermionic MPO of a_p^†a_q
function one_body_mpo(p,q,L;fermion=true,T=Float64)
    H = Array{Array{T,4},1}(undef,L)
    if p == q
        for i in 1:L
            H[i] = mpo_core_id(;T=T)
        end
        H[p][1,1,1,1] = 0.0
    else
        H[p] = mpo_core_creation(;T=T)
        H[q] = mpo_core_annihilation(;T=T)
        for i in 1:min(p,q)-1
            H[i] = mpo_core_id(;T=T)
        end
        for i in max(p,q)+1:L
            H[i] = mpo_core_id(;T=T)
        end
        for i in min(p,q)+1:max(p,q)-1
            if fermion
                H[i] = mpo_core_ferm_sign(;T=T)
            else
                H[i] = mpo_core_id(;T=T)
            end
        end
    end
    return ttoperator{T}(H,2*ones(Int64,L),ones(Int64,L+1),zeros(Int64,L))
end

#returns bosonic or fermionic MPO of a_k^† a^†_l a_m a_n
#assume k<l,m<n
function two_body_mpo(k,l,m,n,L;T=Float64)
    H = Array{Array{T,4},1}(undef,L)
    if l == m
        A = one_body_mpo(k,m,L;T=T)
        B = one_body_mpo(l,n,L;T=T)
        C = one_body_mpo(k,n,L;T=T)
        return  (-1.0*A)*B + C
    else
        A = one_body_mpo(k,m,L;T=T)
        B = one_body_mpo(l,n,L;T=T)
        return -1.0*(A*B)
    end
end

"""
returns an MPO version of
H = Σ_{ij} h_ij (a_i† a_j + c.c.) + Σ_{ijkl} V_{ijkl} (a_i†a_j†a_ka_l + c.c.)
"""
#assuming diagonal terms are divided by 2 in the h and V matrix
function hV_to_mpo(h,V;tol=1e-8)
    L = size(h,1)
    i_list = findall(!iszero,h)
    if length(i_list)>0
        i = i_list[1]
        A = h[i]*(one_body_mpo(i[1],i[2],L)+one_body_mpo(i[2],i[1],L))
        for i in i_list[2:end]
            H = one_body_mpo(i[1],i[2],L)
            G = one_body_mpo(i[2],i[1],L)
            A = A+h[i]*(G+H)
        end
        A = tt_rounding(A,tol=tol)
    end
    for i in findall(!iszero,V)
        H = two_body_mpo(i[1],i[2],i[3],i[4],L)
        G = two_body_mpo(i[4],i[3],i[2],i[1],L)
        A = tt_rounding(A+V[i]*(H+G),tol=tol)
    end
    return A
end

"""
Examples of standard Hamiltonians
"""

#returns the h,V for the hubbard model (assuming natural labelling of sites)
function hubbard_1D(L;t=1,U=1, pbc=false)
    h = zeros(2L,2L)
    V = zeros(2L,2L,2L,2L)
    for i in 1:L-1
        h[2i,2i+2]=-t
        h[2i-1,2i+1]=-t
        V[2i-1,2i,2i-1,2i]=-U #because of the anticommutation rules
    end
    V[2L-1,2L,2L-1,2L]=-U
    if pbc
        h[2L-1,1]=-t
        h[2L,2]=-t
    end
    return h,V
end

#switch sites j and k in a one-dimensional spin chain model
function site_switch(j::Integer,k::Integer,L::Integer)
    out = collect(1:2L)
    out[2j-1],out[2j] = 2k-1,2k
    out[2k-1],out[2k] = 2j-1,2j
    return out
end

#2D Hubbard with cylindrical boundary conditions along L
function hubbard_2D(w,L;t=1,U=1,w_pbc = false,L_pbc=true)
    h = zeros(2w*L,2w*L)
    V = zeros(2w*L,2w*L,2w*L,2w*L)
    for i in 0:w-2
        for j in 0:L-2
            V[2(j*w+i)+1,2(j*w+i)+2,2(j*w+i)+1,2(j*w+i)+2] = -U
            h[2(j*w+i)+1,2((j+1)*w+i)+1] = -t #right connection
            h[2(j*w+i)+1,2(j*w+i+1)+1] = -t #bottom connection
            h[2(j*w+i)+2,2((j+1)*w+i)+2] = -t #right connection
            h[2(j*w+i)+2,2(j*w+i+1)+2] = -t #bottom connection
        end
        #j=L-1
        V[2((L-1)*w+i)+1,2((L-1)*w+i)+2,2((L-1)*w+i)+1,2((L-1)*w+i)+2] = -U
        h[2((L-1)*w+i)+1,2((L-1)*w+i+1)+1] = -t #bottom connection
        h[2((L-1)*w+i)+2,2((L-1)*w+i+1)+2] = -t #bottom connection
        if L_pbc && L>2
            h[2(L*w+i-w)+1,2i+1] =-t #right connection for j=L-1
            h[2(L*w+i-w)+2,2i+2] =-t #right connection for j=L-1
        end
    end
    #i = w-1
    for j in 0:L-2
        V[2(j*w+w-1)+1,2(j*w+w-1)+2,2(j*w+w-1)+1,2(j*w+w-1)+2] = -U
        h[2(j*w+w-1)+1,2((j+1)*w+w-1)+1] = -t #right connection
        h[2(j*w+w-1)+2,2((j+1)*w+w-1)+2] = -t #right connection
        if w_pbc && w>2
            h[2(j*w+w-1)+1,2j*w+1] = -t #top connection
            h[2(j*w+w-1)+2,2j*w+2] = -t #top connection
        end
    end
    #i = w-1, j=L-1
    V[2L*w-1,2L*w,2L*w-1,2L*w] = -U
    if L_pbc && L>2
        h[2(L*w-1)+1,2(w-1)+1] = -t #right connection
        h[2(L*w-1)+2,2(w-1)+2] = -t #right connection
    end
    if w_pbc && w>2
        h[2(L*w-1)+1,2(L-1)*w+1] = -t #right connection
        h[2(L*w-1)+2,2(L-1)*w+2] = -t #right connection
    end
    return h,V
end

"""
PPP Hamiltonian for cyclic polyene C_NH_N (ref. G. Fano, F. Ortolani † and L. Ziosi, The density matrix renormalization group method. Application to the PPP model of a cyclic polyene chain, J. Chem. Phys. 1998)
H = β ∑_{<μ,ν>,σ} a^†_μσ a^†_νσ + c.c. + 0.5 ∑_μ,ν γ_μν (n_μ-1)(n_ν-1)

β = -2.5eV, γ_μν = 1/(γ0^-1+d_μν), γ0 = 10.84eV
d_μν = b sin(π/N*(μ-ν [N]))/sin(π/N), b = 1.4 A

Ground-state energy = -12.72 eV (N=6)
"""
function PPP_C_NH_N(N;β=-2.5/27.2113845,b=1.4*1.8897259886,γ0=10.84/27.2113845,order=collect(1:2N),tol=1e-8)
    @assert isperm(order) "Ordering given is not a permutation."
    h = zeros(2N,2N)
    γ = sum(1/(1/γ0+b*sin(k/N*pi)/sin(pi/N)) for k in 1:N)
    for i in 1:N-1
        h[2i-1,2i+1] = β
        h[2i,2i+2] = β
    end
    h[2N,2] = β
    h[2N-1,1] = β
    H_tto = hV_to_mpo(h[order,order],zeros(2N,2N,2N,2N))
    σ = invperm(order)
    id = -1.0*id_tto(2N)
    for i in 1:N
        Hi = (one_body_mpo(σ[2i],σ[2i],2N)+one_body_mpo(σ[2i-1],σ[2i-1],2N))+id
        for j in 1:N
            Hj = (one_body_mpo(σ[2j],σ[2j],2N)+one_body_mpo(σ[2j-1],σ[2j-1],2N)) + id
            Htemp = Hi*Hj
            γij = 1/(1/γ0+b*sin(pi/N*mod(i-j,N))/sin(pi/N))
            H_tto = H_tto + 0.5*γij*Htemp
            H_tto = tt_rounding(H_tto,tol=tol)
        end
    end
    return H_tto
end

max_stars_repo_name: msdupuy/Tensor-Train-Julia

max_issues_repo_name: msdupuy/Tensor-Train-Julia

max_forks_repo_name: msdupuy/Tensor-Train-Julia


text:
using Base.Test
include("testsystem.jl")
srand(11)
n = 20
Y, X = sample(n, M)
kf = KalmanFilter(Y, M)

@test eltype(kf) === GaussianDistributions.Gaussian{Float64,Float64}
@test length(kf) == n
est = collect(kf)
@test eltype(est) == eltype(kf)

f(t, x, P, Ppred, ll, K) = ll
import Kalman: mappedreturntype
Kalman.mappedreturntype(_, ::Type{typeof(f)}) = Float64
kf2 = Kalman.MappedKalmanFilter(zip(Base.Iterators.countfrom(1),Y), M, f)

@test eltype(kf2) === Float64
@test length(kf2) == n
ll = sum(kf2)
@test typeof(ll) == eltype(kf2)


text:
using Test
using CommonLH

@testset "checkLH" begin
	@testset "validate" begin
		@test_throws ErrorException validate(1.2, Float64)

		validate([1.2, 2.3], Float64)
		validate([1.2, 2.3], Float64, sz = (2,), lb = 1.0)
		validate([1.2 2.1; 2.2 3.2], Float64, sz = (2,2))
		# checkLH.validate(1.2, Float64, sz = ())
		validate([1.2, 2.3], [], lb = 1.1)

		@test_throws DimensionMismatch validate([1.2, 2.3, 3.4], Float64, sz = (2,))
		@test_throws ErrorException  validate([2.1, 3.2], Float64, ub = 3.1)
	end

	@testset "validate_scalar" begin
		validate_scalar(1.1, Float64, lb = 0.9, ub = 1.2)
		validate_scalar(1, Integer, lb = 0.9, ub = 1.3)
		@test_throws ErrorException validate_scalar(1.1, [], lb = 1.2, ub = 1.3)
		@test_throws ErrorException validate_scalar(1.1, Integer, lb = 0.9, ub = 1.3)
		@test_throws ErrorException validate_scalar([1.1, 2.2], [], lb = 1.2, ub = 1.3)
	end
end

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


text:
module SelfEnergy2DSymmetric

using ..BoundaryConditions
using ..Domains
using ..Points
using SparseArrays

import ..Symmetric, ..Unsymmetric
import LinearAlgebra: I
import ..SelfEnergy

function SelfEnergy{Symmetric}(domain::LatticeDomain{2,Symmetric,Cartesian}, αx_half::Vector, αy_half::Vector)

    Nx = length(αx_half)-1
    Ny = length(αy_half)-1
    N = Nx*Ny

    Ix = sparse(I,Nx,Nx)
    Iy = sparse(I,Ny,Ny)
    a = domain.shape.a
    αx_half⁻¹ = 1 ./αx_half
    αy_half⁻¹ = 1 ./αy_half
    dx = domain.dx
    dy = domain.dy

    bc1 = domain.boundary.bcs[1]
    if typeof(bc1)<:DirichletBC{1}
        Σ0L = kron(Iy,sparse([1],[1],[-αx_half⁻¹[1]/dx^2],Nx,Nx))
        Σ1L = kron(Iy,sparse([1],[1],[+1/dx],Nx,Nx))
        fL = FΣ2D(dx,0)
    elseif typeof(bc1)<:NeumannBC{1}
        Σ0L = kron(Iy,sparse([1],[1],[+αx_half⁻¹[1]/dx^2],Nx,Nx))
        Σ1L = kron(Iy,sparse([1],[1],[-1/dx],Nx,Nx))
        fL = FΣ2D(dx,0)
    elseif typeof(bc1)<:FloquetBC{1}
        # Σ0L = sparse([li],[ri],[α_half⁻¹[ri+1]/dx^2],N,N)
        # Σ1L = sparse([li],[ri],[-1/dx],N,N)
        # fL = (ω,ka,_...) -> ffl(ka,a)
        throw("not yet")
    elseif typeof(bc1)<:MatchedBC{1}
        Σ0L = kron(Iy,sparse([1],[1],[+αx_half⁻¹[1]/dx^2],Nx,Nx))
        Σ1L = kron(Iy,sparse([1],[1],[-1/dx],Nx,Nx))
        if isempty(bc1.in)
            fL = FΣ2D(dx,+1)
        else
            fL = FΣ2D(dx,-1)
        end
    else
        Σ0L = spzeros(N,N)
        Σ1L = spzeros(N,N)
        fL = FΣ2D(dx,0)
    end

    bc2 = domain.boundary.bcs[2]
    if typeof(bc2)<:DirichletBC{2}
        Σ0R = kron(Iy,sparse([Nx],[Nx],[-αx_half⁻¹[Nx+1]/dx^2],Nx,Nx))
        Σ1R = kron(Iy,sparse([Nx],[Nx],[-1/dx],Nx,Nx))
        fR = FΣ2D(dx,0)
    elseif typeof(bc2)<:NeumannBC{2}
        Σ0R = kron(Iy,sparse([Nx],[Nx],[+αx_half⁻¹[Nx+1]/dx^2],Nx,Nx))
        Σ1R = kron(Iy,sparse([Nx],[Nx],[+1/dx],Nx,Nx))
        fR = FΣ2D(dx,0)
    elseif typeof(bc2)<:FloquetBC{2}
        # Σ0R = sparse([ri],[1],[-α_half⁻¹[li]/dx^2],N,N)
        # Σ1R = sparse([ri],[1],[+1/dx],N,N)
        # fR = (ω,ka,_...) -> ffr(ka,a)
        throw("not yet")
    elseif typeof(bc2)<:MatchedBC{2}
        Σ0R = kron(Iy,sparse([Nx],[Nx],[+αx_half⁻¹[Nx+1]/dx^2],Nx,Nx))
        Σ1R = kron(Iy,sparse([Nx],[Nx],[+1/dx],Nx,Nx))
        if isempty(bc2.in)
            fR = FΣ2D(dx,+1)
        else
            fR = FΣ2D(dx,-1)
        end
    else
        Σ0R = spzeros(N,N)
        Σ1R = spzeros(N,N)
        fR = FΣ2D(dx,0)
    end

    bc3 = domain.boundary.bcs[3]
    if typeof(bc3)<:DirichletBC{3}
        Σ0B = kron(sparse([1],[1],[-αy_half⁻¹[1]/dy^2],Ny,Ny),Ix)
        Σ1B = kron(sparse([1],[1],[+1/dy],Ny,Ny),Ix)
        fB = FΣ2D(dy,0)
    elseif typeof(bc3)<:NeumannBC{3}
        Σ0B = kron(sparse([1],[1],[+αy_half⁻¹[1]/dy^2],Ny,Ny),Ix)
        Σ1B = kron(sparse([1],[1],[-1/dy],Ny,Ny),Ix)
        fB = FΣ2D(dy,0)
    elseif typeof(bc3)<:FloquetBC{3}
        # Σ0L = sparse([li],[ri],[α_half⁻¹[ri+1]/dx^2],N,N)
        # Σ1L = sparse([li],[ri],[-1/dx],N,N)
        # fL = (ω,ka,_...) -> ffl(ka,a)
        throw("not yet")
    elseif typeof(bc3)<:MatchedBC{3}
        Σ0B = kron(sparse([1],[1],[+αy_half⁻¹[1]/dy^2],Ny,Ny),Ix)
        Σ1B = kron(sparse([1],[1],[-1/dy],Ny,Ny),Ix)
        if isempty(bc3.in)
            fB = FΣ2D(dy,+1)
        else
            fB = FΣ2D(dy,-1)
        end
    else
        Σ0B = spzeros(N,N)
        Σ1B = spzeros(N,N)
        fB = FΣ2D(dy,0)
    end

    bc4 = domain.boundary.bcs[4]
    if typeof(bc4)<:DirichletBC{4}
        Σ0T = kron(sparse([Ny],[Ny],[-αy_half⁻¹[Ny+1]/dy^2],Ny,Ny),Ix)
        Σ1T = kron(sparse([Ny],[Ny],[-1/dy],Ny,Ny),Ix)
        fT = FΣ2D(dy,0)
    elseif typeof(bc4)<:NeumannBC{4}
        Σ0T = kron(sparse([Ny],[Ny],[+αy_half⁻¹[Ny+1]/dy^2],Ny,Ny),Ix)
        Σ1T = kron(sparse([Ny],[Ny],[+1/dy],Ny,Ny),Ix)
        fT = FΣ2D(dy,0)
    elseif typeof(bc4)<:FloquetBC{4}
        # Σ0R = sparse([ri],[1],[-α_half⁻¹[li]/dx^2],N,N)
        # Σ1R = sparse([ri],[1],[+1/dx],N,N)
        # fR = (ω,ka,_...) -> ffr(ka,a)
        throw("not yet")
    elseif typeof(bc4)<:MatchedBC{4}
        Σ0T = kron(sparse([Ny],[Ny],[+αy_half⁻¹[Ny+1]/dy^2],Ny,Ny),Ix)
        Σ1T = kron(sparse([Ny],[Ny],[+1/dy],Ny,Ny),Ix)
        if isempty(bc4.in)
            fT = FΣ2D(dy,+1)
        else
            fT = FΣ2D(dy,-1)
        end
    else
        Σ0T = spzeros(N,N)
        Σ1T = spzeros(N,N)
        fT = FΣ2D(dy,0)
    end

    Σ2 = spzeros(N,N)

    fs = (fL,fR,fB,fT)

    return SelfEnergy{2,Symmetric,4,typeof(fs)}((-Σ0L,-Σ0R,-Σ0B,-Σ0T),(Σ1L,Σ1R,Σ1B,Σ1T),Σ2,fs)
end

struct FΣ2D <: Function
    dx::Float64
    sign::Int
end
function (fm::FΣ2D)(ω::Number,ka,ky=0,kz=0)
    if iszero(fm.sign)
        return complex(1.0,0)
    else
        k̂ = sqrt(ω^2 - ky^2 - kz^2)*fm.dx/2
        abs(k̂)<1 || throw(DomainError(fm.dx,"lattice spacing $(fm.dx) too large. for |ω|=$(abs(ω)), ky=$ky, kz=$kz, should be < $(abs(1/sqrt(ω^2 - ky^2 - kz^2)))"))
        kx = 2asin(k̂)/fm.dx
        return exp(fm.sign*1im*kx*fm.dx)
    end
end
function (fm::FΣ2D)(ω::Matrix,ka,ky,kz)
    if iszero(fm.sign)
        return one(convert(Matrix{ComplexF64},ω))
    else
        kx::Matrix{ComplexF64} = 2asin(sqrt(ω^2 - one(ω)*ky^2 - one(ω)*kz^2)*fm.dx/2)/fm.dx
        return exp(fm.sign*1im*kx*fm.dx)
    end
end

end

using .SelfEnergy2DSymmetric


text:
module BhAbs_Solver

#Julia libs
using LinearAlgebra
using LinearAlgebra.BLAS
using DifferentialEquations
using SparseArrays
using Dates
using Parameters
using HDF5
using DelimitedFiles
using ProgressMeter

#Import structures
include("./structs/Operators.jl")
include("./structs/Parameters.jl")

#Operators
include("./operators/differential_operator.jl")
include("./operators/position_operator.jl")
include("./operators/absorption_operator.jl")
include("./operators/BC_operator.jl")
include("./operators/potential_operator.jl")

#Simulation helpers
include("./sim_utils/get_gaussian_pulse.jl")    # change to gaussian pulse

include("./sim_utils/setup_operators.jl")
include("./sim_utils/setup_folders.jl")

include("./sim_utils/save_parameters_to_file.jl")
include("./sim_utils/get_parameters_from_file.jl")

include("./sim_utils/print_parameters_to_screen.jl")

include("./sim_utils/time_to_seconds.jl")
include("./sim_utils/exit_message.jl")
include("./sim_utils/welcome_message.jl")

#Solvers
include("./solvers/KG_Solver.jl")

#Data analysis aids
include("./data_read_utils/get_field_energy.jl")
include("./data_read_utils/get_field_point.jl")
include("./data_read_utils/get_field_slice.jl")
include("./data_read_utils/get_fields.jl")
include("./data_read_utils/get_time_variables.jl")
include("./data_read_utils/get_space_variables.jl")

#Visualization aids
#include("./visualization/sim_visualizer.jl")

#Template_runs
include("./template_runs/single_bh_runs.jl")
include("./template_runs/binary_bh_runs.jl")
include("./template_runs/resume_simulation.jl")

export run_single_BH_simulation , run_binary_BH_simulation , resume_simulation_from_file
export get_time_variables , get_field_energy , get_field_point, get_field_slice, get_fields

end


text:
using Plots
using Serialization
function make_fig()
    julia_set = deserialize("./data/julia_data")
    frames = size(julia_set)[end]
    gr()
    for frame in 1:frames
        img = heatmap(julia_set[:,:,frame],legend=:none,axis=nothing, size=size(julia_set[:,:,frame]))
        if frame < 10
            savefig(img,"./imgs/julia_set_0$frame.png")
        elseif frame < 100
            savefig(img,"./imgs/julia_set_$frame.png")
        else
            println("If more than 100 frames code needs to be modified here")
            break
        end
    end
    run(`ffmpeg -r 10 -i ./imgs/julia_set_%02d.png julia_video.gif`)
end

make_fig()

max_stars_repo_name: ofbennett/julia-set-fractal

max_issues_repo_name: ofbennett/julia-set-fractal

max_forks_repo_name: ofbennett/julia-set-fractal


text:
@testset "swap" begin

    a=Int(1)
    b=Int(2)

    @swap(a,b)

    @test a == 2
    @test b == 1

    a=rand(5)
    b=rand(8)

    @swap(a,b)

    @test length(a) == 8
    @test length(b) == 5

end;

max_stars_repo_name: vincent-picaud/DirectConvolution

max_issues_repo_name: vincent-picaud/DirectConvolution

max_forks_repo_name: vincent-picaud/DirectConvolution


text:
@inline function ExaTronKernels.dnrm2(n::Int,x::CuDeviceArray{Float64,1},incx::Int)
    tx = threadIdx().x

    v = 0.0
    if tx <= n  # No check on ty so that each warp has v.
        @inbounds v = x[tx]*x[tx]
    end

    # shfl_down_sync() will automatically sync threads in a warp.

    offset = 16
    while offset > 0
        v += CUDA.shfl_down_sync(0xffffffff, v, offset)
        offset >>= 1
    end
    v = sqrt(v)
    v = CUDA.shfl_sync(0xffffffff, v, 1)

    return v
end


text:
using Zlib
zliblength(str) = length(Zlib.compress(str,9,false,true))
using Blosc
blosclzlength(s) = length(Blosc.compress(convert(Vector{Uint8}, s), clevel=9, cname=:blosclz))
lz4length(s) = length(Blosc.compress(convert(Vector{Uint8}, s), clevel=9, cname=:lz4))
lz4hclength(s) = length(Blosc.compress(convert(Vector{Uint8}, s), clevel=9, cname=:lz4hc))
bzliblength(s) = length(Blosc.compress(convert(Vector{Uint8}, s), clevel=9, cname=:zlib))

function report(name, func, input)
  tic()
  len = func(input)
  t = toq()
  @printf("%s, time = %.3e seconds, compression ratio = %.3f\n", name, t, length(input)/len)
end

for exponent in 1:7
  n = 10^exponent
  input = Uint8[1:n];
  println("\nInput of length 10^$exponent")
  report("zlib         ", (input) -> zliblength(input), input)
  report("zlib in blosc", (input) -> lz4hclength(input), input)
  report("lz4hc        ", (input) -> bzliblength(input), input)
  report("lz4          ", (input) -> lz4length(input), input)
  report("blosclz      ", (input) -> blosclzlength(input), input)
end


text:
LL = 20
SDE_probs = zeros(LL)
for i in 1:LL
    dt = 1//2^(i+1)
    SDE_probs[i] = dt
end

function sample_SDE_path(level, ω; TEST=false)
    # solve on finest grid
    SDE_prob_coarse = SDEProblem(f, g, u₀, (0.0, 1.0), dt=SDE_probs[level + one(level)])
    sol1 = solve(SDE_prob_coarse, EM(), save_noise = true) # TO DO, dt needs to be varied
    Qf = sol1.u[1:2:end]

    # compute difference when not on coarsest grid
    dQ = copy(Qf)
    if level != Level(0)
        SDE_prob_fine = SDEProblem(f, g, u₀, (0.0, 1.0), dt=SDE_probs[level])
        sol2 = solve(SDE_prob_fine, EM())
        Qc = sol2.u
        dQ -= Qc
    end

    if TEST
        return dQ, Qf, Qc
    else
        return dQ, Qf
    end
end

# define the problem
a, b, c = sample_SDE_path(7, 1., TEST=true)

println("succes")

p2 = plot()
plot!(p2, b, label="fine")
plot!(p2, c, label="coarse")


text:
using Millboard
using Test

Millboard.set_table_mode(:markdown)
Base.eval(:(have_color = false))

board = []
@test """
""" == table(board) |> string

board = [11 12 13; 21 22 23; 31 32 33]
@test """
|   |  1 |  2 |  3 |
|---|----|----|----|
| 1 | 11 | 12 | 13 |
| 2 | 21 | 22 | 23 |
| 3 | 31 | 32 | 33 |""" == table(board) |> string

Base.eval(:(have_color = true))


text:
using CodedComputing
using Random, MPI, HDF5, H5Sparse, LinearAlgebra, SparseArrays
using Dates
using Test

"""

Return an array composed of the PCA computed iterates.
"""
function load_pca_iterates(outputfile::AbstractString, outputdataset::AbstractString)
    @test HDF5.ishdf5(outputfile)
    h5open(outputfile, "r") do fid
        @test outputdataset in keys(fid)
        if "iterates" in keys(fid) && length(size(fid["iterates"])) == 3 && size(fid["iterates"], 3) > 0
            @test fid[outputdataset][:, :] ≈ fid["iterates"][:, :, end]
            return [fid["iterates"][:, :, i] for i in 1:size(fid["iterates"], 3)]
        else # return only the final iterate if intermediate iterates aren't stored
            return [fid[outputdataset][:, :]]
        end
    end
end

function load_logreg_iterates(outputfile::AbstractString, outputdataset::AbstractString)
    @test HDF5.ishdf5(outputfile)
    h5open(outputfile, "r") do fid
        @test outputdataset in keys(fid)
        if "iterates" in keys(fid) && length(size(fid["iterates"])) == 2 && size(fid["iterates"], 2) > 0
            @test fid[outputdataset][:] ≈ fid["iterates"][:, end]
            return [fid["iterates"][:, i] for i in 1:size(fid["iterates"], 2)]
        else # return only the final iterate if intermediate iterates aren't stored
            return [fid[outputdataset][:]]
        end
    end
end

function test_pca_iterates(;X::AbstractMatrix, niterations::Integer, ncomponents::Integer,
                            ev::Real, outputfile::AbstractString, outputdataset::AbstractString, atol=1e-2)
    dimension, nsamples = size(X)
    Vs = load_pca_iterates(outputfile, outputdataset)
    fs = [explained_variance(X, V) for V in Vs]
    # @test length(Vs) == niterations
    # @test all((V)->size(V)==(dimension, ncomponents), Vs)
    @test size(Vs[end]) == (dimension, ncomponents)
    @test Vs[end]'*Vs[end] ≈ I atol=1e-6
    @test fs[end] < ev || isapprox(fs[end], ev, atol=atol)
    return Vs, fs
end

# @testset "latency.jl" begin
#     kernel = "../src/latency/kernel.jl"
#     nwait = 1
#     nworkers = 3
#     niterations = 10
#     timeout = 0.1
#     nbytes = 100
#     nrows = 100
#     ncols = 100
#     ncomponents = 3
#     density = 0.1
#     outputfile = tempname()
#     mpiexec(cmd -> run(```
#         $cmd -n $(nworkers+1) julia --project $kernel $outputfile
#             --niterations $niterations
#             --nbytes $nbytes
#             --nrows $nrows
#             --ncols $ncols
#             --ncomponents $ncomponents
#             --density $density
#             --nwait $nwait
#             --timeout $timeout
#     ```))
#     df = df_from_latency_file(outputfile)
#     @test all(diff(df.timestamp) .>= timeout)
# end

@time @testset "Linalg.jl" begin include("linalg_test.jl") end
@time @testset "concurrentcircularbuffer.jl" begin include("concurrentcircularbuffer_test.jl") end
@time @testset "partitioning.jl" begin include("partitioning_test.jl") end
@time @testset "eventdriven.jl" begin include("eventdriven_test.jl") end
@time @testset "profiling.jl" begin include("profiling_test.jl") end
@time @testset "loadbalancing.jl" begin include("loadbalancing_test.jl") end
@time @testset "treegradient.jl" begin include("treegradient_test.jl") end
@time @testset "mul.jl" begin include("mul_test.jl") end
@time @testset "logreg.jl" begin include("logreg_test.jl") end
@time @testset "pca.jl" begin include("pca_test.jl") end


text:
"""
    PG(n_particles::Int, n_iters::Int)

Particle Gibbs sampler.

Usage:

```julia
PG(100, 100)
```

Example:

```julia
# Define a simple Normal model with unknown mean and variance.
@model gdemo(x) = begin
  s ~ InverseGamma(2,3)
  m ~ Normal(0,sqrt.(s))
  x[1] ~ Normal(m, sqrt.(s))
  x[2] ~ Normal(m, sqrt.(s))
  return s, m
end

sample(gdemo([1.5, 2]), PG(100, 100))
```
"""
mutable struct PG <: InferenceAlgorithm
  n_particles           ::    Int         # number of particles used
  n_iters               ::    Int         # number of iterations
  resampler             ::    Function    # function to resample
  space                 ::    Set         # sampling space, emtpy means all
  gid                   ::    Int         # group ID
  PG(n1::Int, n2::Int) = new(n1, n2, resampleSystematic, Set(), 0)
  PG(n1::Int, n2::Int, resampler::Function, space::Set, gid::Int) = new(n1, n2, resampler, space, gid)
  function PG(n1::Int, n2::Int, space...)
    space = isa(space, Symbol) ? Set([space]) : Set(space)
    new(n1, n2, resampleSystematic, space, 0)
  end
  PG(alg::PG, new_gid::Int) = new(alg.n_particles, alg.n_iters, alg.resampler, alg.space, new_gid)
end

const CSMC = PG # type alias of PG as Conditional SMC

Sampler(alg::PG) = begin
  info = Dict{Symbol, Any}()
  info[:logevidence] = []
  Sampler(alg, info)
end

step(model::Function, spl::Sampler{PG}, vi::VarInfo, _::Bool) = step(model, spl, vi)

step(model::Function, spl::Sampler{PG}, vi::VarInfo) = begin
  particles = ParticleContainer{Trace}(model)

  vi.num_produce = 0;  # Reset num_produce before new sweep\.
  ref_particle = isempty(vi) ?
                 nothing :
                 forkr(Trace(model, spl, vi))

  set_retained_vns_del_by_spl!(vi, spl)
  resetlogp!(vi)

  if ref_particle == nothing
    push!(particles, spl.alg.n_particles, spl, vi)
  else
    push!(particles, spl.alg.n_particles-1, spl, vi)
    push!(particles, ref_particle)
  end

  while consume(particles) != Val{:done}
    resample!(particles, spl.alg.resampler, ref_particle)
  end

  ## pick a particle to be retained.
  Ws, _ = weights(particles)
  indx = randcat(Ws)
  push!(spl.info[:logevidence], particles.logE)

  particles[indx].vi
end

sample(model::Function, alg::PG;
       save_state=false,         # flag for state saving
       resume_from=nothing,      # chain to continue
       reuse_spl_n=0             # flag for spl re-using
      ) = begin

  spl = reuse_spl_n > 0 ?
        resume_from.info[:spl] :
        Sampler(alg)

  @assert typeof(spl.alg) == typeof(alg) "[Turing] alg type mismatch; please use resume() to re-use spl"

  n = reuse_spl_n > 0 ?
      reuse_spl_n :
      alg.n_iters
  samples = Vector{Sample}()

  ## custom resampling function for pgibbs
  ## re-inserts reteined particle after each resampling step
  time_total = zero(Float64)

  vi = resume_from == nothing ?
       VarInfo() :
       resume_from.info[:vi]

  pm = nothing
  if PROGRESS spl.info[:progress] = ProgressMeter.Progress(n, 1, "[PG] Sampling...", 0) end

  for i = 1:n
    time_elapsed = @elapsed vi = step(model, spl, vi)
    push!(samples, Sample(vi))
    samples[i].value[:elapsed] = time_elapsed

    time_total += time_elapsed

    if PROGRESS  && spl.alg.gid == 0
      ProgressMeter.next!(spl.info[:progress])
    end
  end

  @info("[PG] Finished with")
  @info("  Running time    = $time_total;")

  loge = exp.(mean(spl.info[:logevidence]))
  if resume_from != nothing   # concat samples
    pushfirst!(samples, resume_from.value2...)
    pre_loge = resume_from.weight
    # Calculate new log-evidence
    pre_n = length(resume_from.value2)
    loge = exp.((log(pre_loge) * pre_n + log(loge) * n) / (pre_n + n))
  end
  c = Chain(loge, samples)       # wrap the result by Chain

  if save_state               # save state
    save!(c, spl, model, vi)
  end

  c
end

assume(spl::Sampler{T}, dist::Distribution, vn::VarName, _::VarInfo) where T<:Union{PG,SMC} = begin
  vi = current_trace().vi
  if isempty(spl.alg.space) || vn.sym in spl.alg.space
    if ~haskey(vi, vn)
      r = rand(dist)
      push!(vi, vn, r, dist, spl.alg.gid)
      spl.info[:cache_updated] = CACHERESET # sanity flag mask for getidcs and getranges
    elseif is_flagged(vi, vn, "del")
      unset_flag!(vi, vn, "del")
      r = rand(dist)
      vi[vn] = vectorize(dist, r)
      setgid!(vi, spl.alg.gid, vn)
      setorder!(vi, vn, vi.num_produce)
    else
      updategid!(vi, vn, spl)
      r = vi[vn]
    end
  else # vn belongs to other sampler <=> conditionning on vn
    if haskey(vi, vn)
      r = vi[vn]
    else
      r = rand(dist)
      push!(vi, vn, r, dist, -1)
    end
    acclogp!(vi, logpdf_with_trans(dist, r, istrans(vi, vn)))
  end
  r, zero(Real)
end

assume(spl::Sampler{A}, dists::Vector{D}, vn::VarName, var::Any, vi::VarInfo) where {A<:Union{PG,SMC},D<:Distribution} =
  error("[Turing] PG and SMC doesn't support vectorizing assume statement")

observe(spl::Sampler{T}, dist::Distribution, value, vi) where T<:Union{PG,SMC} = begin
  produce(logpdf(dist, value))
  zero(Real)
end

observe(spl::Sampler{A}, ds::Vector{D}, value::Any, vi::VarInfo) where {A<:Union{PG,SMC},D<:Distribution} =
  error("[Turing] PG and SMC doesn't support vectorizing observe statement")


text:
using Kinetic, Flux
using KitBase.JLD2, KitBase.WriteVTK
using Flux: @epochs
import PyPlot as plt

cd(@__DIR__)
include("../common.jl")

begin
    set = Setup(
        case = "cylinder",
        space = "2d2f2v",
        boundary = ["maxwell", "extra", "mirror", "mirror"],
        limiter = "minmod",
        cfl = 0.5,
        maxTime = 10.0, # time
    )
    ps = CSpace2D(1.0, 6.0, 60, 0.0, π, 50, 1, 1)
    vs = VSpace2D(-10.0, 10.0, 48, -10.0, 10.0, 48)
    #gas = Gas(Kn = 1e-3, Ma = 5.0, K = 1.0)
    gas = Gas(Kn = 1e-2, Ma = 5.0, K = 1.0)

    prim0 = [1.0, 0.0, 0.0, 1.0]
    prim1 = [1.0, gas.Ma * sound_speed(1.0, gas.γ), 0.0, 1.0]
    fw = (args...) -> prim_conserve(prim1, gas.γ)
    ff = function(args...)
        prim = conserve_prim(fw(args...), gas.γ)
        h = maxwellian(vs.u, vs.v, prim)
        b = h .* gas.K / 2 / prim[end]
        return h, b
    end
    bc = function(x, y)
        if abs(x^2 + y^2 - 1) < 1e-3
            return prim0
        else
            return prim1
        end
    end
    ib = IB2F(fw, ff, bc)

    ks = SolverSet(set, ps, vs, gas, ib)
end

#@load "kn3ref.jld2" ctr
@load "kn2ref.jld2" ctr

rg_ref = zeros(ks.ps.nr, ks.ps.nθ)
@inbounds Threads.@threads for j = 1:ks.ps.nθ
    rg_ref[1, j] = 1
    for i = 2:ks.ps.nr
        swx1 = (ctr[i+1, j].w - ctr[i-1, j].w) / (1e-6 + ks.ps.x[i+1, j] - ks.ps.x[i-1, j])
        swy1 = (ctr[i+1, j].w - ctr[i-1, j].w) / (1e-6 + ks.ps.y[i+1, j] - ks.ps.y[i-1, j])
        swx2 = (ctr[i, j+1].w - ctr[i, j-1].w) / (1e-6 + ks.ps.x[i, j+1] - ks.ps.x[i, j-1])
        swy2 = (ctr[i, j+1].w - ctr[i, j-1].w) / (1e-6 + ks.ps.y[i, j+1] - ks.ps.y[i, j-1])
        swx = (swx1 + swx2) ./ 2
        swy = (swy1 + swy2) ./ 2

        rg_ref[i, j] = judge_regime(ks, ctr[i, j].h, ctr[i, j].prim, swx, swy)
    end
end

rg_kngll = zeros(ks.ps.nr, ks.ps.nθ)
@inbounds Threads.@threads for j = 1:ks.ps.nθ
    rg_kngll[1, j] = 1
    for i = 2:ks.ps.nr
        swx1 = (ctr[i+1, j].w - ctr[i-1, j].w) / (1e-6 + ks.ps.x[i+1, j] - ks.ps.x[i-1, j])
        swy1 = (ctr[i+1, j].w - ctr[i-1, j].w) / (1e-6 + ks.ps.y[i+1, j] - ks.ps.y[i-1, j])
        swx2 = (ctr[i, j+1].w - ctr[i, j-1].w) / (1e-6 + ks.ps.x[i, j+1] - ks.ps.x[i, j-1])
        swy2 = (ctr[i, j+1].w - ctr[i, j-1].w) / (1e-6 + ks.ps.y[i, j+1] - ks.ps.y[i, j-1])
        swx = (swx1 + swx2) ./ 2
        swy = (swy1 + swy2) ./ 2
        sw = @. sqrt(swx^2 + swy^2)

        L = abs(ctr[i, j].w[1] / sw[1])
        ℓ = (1 / ctr[i, j].prim[end])^ks.gas.ω / ctr[i, j].prim[1] * sqrt(ctr[i, j].prim[end]) * ks.gas.Kn
        rg_kngll[i, j] = ifelse(ℓ / L > 0.05, 1, 0)
        #rg_kngll[i, j] = ifelse(ℓ / L > 0.01, 1, 0)
    end
end

begin
    plt.close("all")
    fig = plt.figure("contour", figsize=(8, 4))
    plt.contourf(
        ps.x[1:ks.ps.nr, 1:ks.ps.nθ],
        ps.y[1:ks.ps.nr, 1:ks.ps.nθ],
        #rg_ref,
        rg_kngll,
        #rg_nn,
        levels = 20,
        cmap = plt.ColorMap("inferno"),
    )
    plt.colorbar()
    plt.xlabel("x")
    plt.ylabel("y")
    plt.tight_layout()
    plt.xlim(-6, 6)
    plt.ylim(0, 6)
    plt.display(fig)
end
#fig.savefig("cylinder_rgref_kn3.pdf")
#fig.savefig("cylinder_rgkngll_kn2_c001.pdf")

@load "nn_pro.jld2" nn

global X = Float32.([[1.0, 0.0, 0.0, 1.0]; zeros(4); 1e-4])
global Y = Float32(0.0)
for j = 1:ks.ps.nθ, i = 2:ks.ps.nr
    swx1 = (ctr[i+1, j].w - ctr[i-1, j].w) / (ks.ps.x[i+1, j] - ks.ps.x[i-1, j])
    swy1 = (ctr[i+1, j].w - ctr[i-1, j].w) / (ks.ps.y[i+1, j] - ks.ps.y[i-1, j])
    swx2 = (ctr[i, j+1].w - ctr[i, j-1].w) / (ks.ps.x[i, j+1] - ks.ps.x[i, j-1])
    swy2 = (ctr[i, j+1].w - ctr[i, j-1].w) / (ks.ps.y[i, j+1] - ks.ps.y[i, j-1])
    swx = (swx1 + swx2) ./ 2
    swy = (swy1 + swy2) ./ 2

    x, y = regime_data(ks, ctr[i, j].w, swx, swy, ctr[i, j].h)
    global X = hcat(X, x)
    global Y = hcat(Y, y)
end

data = Flux.Data.DataLoader((X, Y), shuffle = true)
ps = params(nn)
loss(x, y) = Flux.binarycrossentropy(nn(x), y)
cb = () -> println("loss: $(loss(X, Y))")
opt = ADAM()

@epochs 20 Flux.train!(loss, ps, data, opt, cb = Flux.throttle(cb, 1))
accuracy(nn, X, Y)

rg_nn = zeros(ks.ps.nr, ks.ps.nθ)
@inbounds Threads.@threads for j = 1:ks.ps.nθ
    rg_nn[1, j] = 1
    for i = 2:ks.ps.nr
        swx1 = (ctr[i+1, j].w - ctr[i-1, j].w) / (1e-6 + ks.ps.x[i+1, j] - ks.ps.x[i-1, j])
        swy1 = (ctr[i+1, j].w - ctr[i-1, j].w) / (1e-6 + ks.ps.y[i+1, j] - ks.ps.y[i-1, j])
        swx2 = (ctr[i, j+1].w - ctr[i, j-1].w) / (1e-6 + ks.ps.x[i, j+1] - ks.ps.x[i, j-1])
        swy2 = (ctr[i, j+1].w - ctr[i, j-1].w) / (1e-6 + ks.ps.y[i, j+1] - ks.ps.y[i, j-1])
        swx = (swx1 + swx2) ./ 2
        swy = (swy1 + swy2) ./ 2
        sw = sqrt.(swx.^2 + swy.^2)
        tau = vhs_collision_time(ctr[i, j].prim, ks.gas.μᵣ, ks.gas.ω)

        rg_nn[i, j] = nn([ctr[i, j].w; sw; tau])[1] |> round |> Int
    end
end

begin
    plt.close("all")
    fig = plt.figure("contour", figsize=(8, 4))
    plt.contourf(
        ks.ps.x[1:ks.ps.nr, 1:ks.ps.nθ],
        ks.ps.y[1:ks.ps.nr, 1:ks.ps.nθ],
        rg_nn,
        levels = 20,
        cmap = plt.ColorMap("inferno"),
    )
    plt.colorbar()
    plt.xlabel("x")
    plt.ylabel("y")
    plt.tight_layout()
    plt.xlim(-6, 6)
    plt.ylim(0, 6)
    plt.display(fig)
end
#fig.savefig("cylinder_rgnn_kn3.pdf")
fig.savefig("cylinder_rgnn_kn2.pdf")

max_stars_repo_name: CSMMLab/Flowmachine

max_issues_repo_name: CSMMLab/Flowmachine

max_forks_repo_name: CSMMLab/Flowmachine


text:
using Random
using Ripserer
using Test

using Ripserer: chain_element_type, coefficient, index

using Ripserer: ReducedMatrix, record!, commit!, clear_buffer!
using Ripserer: WorkingChain, nonheap_push!, repair!

cofacet_type(::Type{<:A}) where {D,T,I,A<:Simplex{D,T,I}} = Simplex{D + 1,T,I}
facet_type(::Type{<:A}) where {D,T,I,A<:Simplex{D,T,I}} = Simplex{D - 1,T,I}

@testset "ReducedMatrix" begin
    for S in (Simplex{2,Int,Int}, Simplex{1,Int,Int32}), T in (Mod{3}, Rational{Int})
        C = cofacet_type(S)
        CE = chain_element_type(C, T)
        SE = chain_element_type(S, T)

        columns = C.([3, 10, 6, 8], 1)
        colelems = CE.(columns)

        fwd = Base.Order.Forward
        rev = Base.Order.Reverse

        @testset "ReducedMatrix with $S and $T" begin
            @testset "a fresh ReducedMatrix is empty" begin
                matrix = ReducedMatrix{C,SE}(fwd)
                @test length(matrix) == 0
                for col in Iterators.flatten((columns, colelems))
                    @test isempty(matrix[col])
                end
            end

            @testset "is empty after committing nothing" begin
                matrix = ReducedMatrix{C,SE}(fwd)
                commit!(matrix, columns[1], T(2))

                @test length(matrix) == 0
                for col in Iterators.flatten((columns, colelems))
                    @test isempty(matrix[col])
                end
            end

            @testset "is still empty after recording values" begin
                matrix = ReducedMatrix{C,SE}(rev)
                vals = [SE(S(1, 1)), SE(S(2, 1)), SE(S(3, 1)), SE(S(4, 1))]
                record!(matrix, vals, one(T))
                record!(matrix, S(5, 1))

                @test length(matrix) == 0
                for col in Iterators.flatten((columns, colelems))
                    @test isempty(matrix[col])
                end
            end

            @testset "commiting adds a column and changes length" begin
                matrix = ReducedMatrix{C,SE}(fwd)
                vals = [SE(S(1, 1)), SE(S(2, 1)), SE(S(3, 1)), SE(S(4, 1)), SE(S(-1, 1))]
                record!(matrix, vals, one(T))
                record!(matrix, S(4, 1))

                commit!(matrix, columns[1], T(2))
                clear_buffer!(matrix)
                commit!(matrix, columns[2], T(2))

                @test length(matrix) == 1
                @test index.(simplex.(collect(matrix[columns[1]]))) == [4, 3, 2]
                @test index.(simplex.(collect(matrix[colelems[1]]))) == [4, 3, 2]

                @test coefficient.(collect(matrix[columns[1]])) == [T(4), T(2), T(2)]
                @test coefficient.(collect(matrix[colelems[1]])) == [T(4), T(2), T(2)]

                for col in Iterators.flatten((columns[2:end], colelems[2:end]))
                    @test isempty(matrix[col])
                end
            end

            @testset "clear_buffer! undos changes" begin
                matrix = ReducedMatrix{C,SE}(rev)
                vals = [
                    SE(S(1, 1)),
                    SE(S(2, 1)),
                    SE(S(3, 1)),
                    SE(S(4, 1)),
                    SE(S(4, 1)),
                    SE(S(-1, 1)),
                ]
                record!(matrix, vals, one(T))
                clear_buffer!(matrix)
                commit!(matrix, columns[2], T(2))

                record!(matrix, vals, one(T))
                commit!(matrix, columns[1], T(2))

                @test length(matrix) == 1
                @test index.(simplex.(collect(matrix[columns[1]]))) == [2, 3, 4]
                @test index.(simplex.(collect(matrix[colelems[1]]))) == [2, 3, 4]

                @test coefficient.(collect(matrix[columns[1]])) == [T(2), T(2), T(4)]
                @test coefficient.(collect(matrix[colelems[1]])) == [T(2), T(2), T(4)]

                for col in Iterators.flatten((columns[2:end], colelems[2:end]))
                    @test isempty(matrix[col])
                end
            end

            @testset "committing values that sum to 0 does not create a column" begin
                matrix = ReducedMatrix{C,SE}(fwd)
                vals = [
                    SE(S(1, 1)),
                    SE(S(-2, 1)),
                    SE(S(4, 1)),
                    SE(S(-4, 1)),
                    SE(S(2, 1)),
                    SE(S(-1, 1)),
                ]
                record!(matrix, vals, one(T))
                commit!(matrix, columns[1], T(2))

                @test length(matrix) == 0
                for col in Iterators.flatten((columns, colelems))
                    @test isempty(matrix[col])
                end
            end

            @testset "committing multiple times creates multiple columns" begin
                matrix = ReducedMatrix{C,SE}(rev)
                vals_1 = [
                    SE(S(-1, 1)), SE(S(4, 1)), SE(S(-2, 1)), SE(S(-3, 1)), SE(S(-4, 1))
                ]
                record!(matrix, vals_1, -one(T))
                commit!(matrix, columns[1], T(2))
                clear_buffer!(matrix)

                vals_2 = [SE(S(4, 1)), SE(S(1, 1)), SE(S(5, 1)), SE(S(6, 1)), SE(S(-1, 1))]
                record!(matrix, vals_2, one(T))
                commit!(matrix, columns[2], T(-1))
                clear_buffer!(matrix)

                @test length(matrix) == 2

                @test index.(simplex.(collect(matrix[columns[1]]))) == [1, 2, 3]
                @test index.(simplex.(collect(matrix[columns[2]]))) == [4, 5, 6]

                for col in Iterators.flatten((columns[3:end], colelems[3:end]))
                    @test isempty(matrix[col])
                end
            end
        end
    end
end

@testset "WorkingChain" begin
    for S in (Simplex{2,Int,Int}, Simplex{1,Int,Int32}), T in (Mod{3}, Rational{Int})
        SE = chain_element_type(S, T)

        elements = SE.(S.([1, -7, 2, 3, 4, 7, 5, 6, -1], [1, 7, 1, 1, 4, 7, 5, 6, 1]))
        unq_elements = SE.(S.([7, 2, 3, 4, 5, 6], [7, 1, 1, 4, 5, 6]))

        fwd = Base.Order.Forward
        rev = Base.Order.Reverse

        @testset "a fresh WorkingChain is empty and pop! yields nothing" begin
            working_boundary = WorkingChain{SE}(fwd)

            @test isempty(working_boundary)
            @test pop!(working_boundary) ≡ nothing
            @test_throws BoundsError first(working_boundary)
        end

        @testset "pushing elements and popping finds the lowest simplex (cohomology)" begin
            working_boundary = WorkingChain{SE}(fwd)
            for e in elements
                push!(working_boundary, e)
            end

            @test pop!(working_boundary) == SE(S(3, 1))
        end

        @testset "the same happens for nonheap_push! with repair! (homology)" begin
            working_boundary = WorkingChain{SE}(rev)
            for e in elements
                nonheap_push!(working_boundary, e)
            end
            repair!(working_boundary)

            @test pop!(working_boundary) == SE(S(6, 6))
        end

        @testset "adding inverses removes elements" begin
            working_boundary = WorkingChain{SE}(fwd)
            for e in unq_elements
                nonheap_push!(working_boundary, e)
            end
            repair!(working_boundary)
            for e in unq_elements[1:2:end]
                push!(working_boundary, -e)
            end

            for e in sort(unq_elements[2:2:end]; order=fwd)
                @test pop!(working_boundary) == e
            end
            @test pop!(working_boundary) ≡ nothing
        end
    end
end


text:
abstract type AbstractSingleValuedUtility end
abstract type AbstractIsoElastic <: AbstractSingleValuedUtility end

dutil(c, ps::AbstractIsoElastic)=c^(-rra(ps))
invdutil(m, ps::AbstractIsoElastic)=m^(-1/rra(ps))

"
   IsoElastic

IsoElastic specification for a utility with one good.

The parametrization is either
```math
U(c) = \frac{c^{1-\alpha}-1}{1-\alpha}
```

or

```math
U(c) = \frac{c^{1-\alpha}}{1-\alpha}
```

(depending on whether IsoElasticVar2 is passed to `util`.)

The field `elast` corresponds to α.
"
struct IsoElastic <: AbstractIsoElastic
    elast::Float64
    function IsoElastic(elast)
        elast == 1 && error("Please use log utility specification")
        new(elast)
    end
end
struct LogUtility <: AbstractIsoElastic end
function IsoElastic(;a::Float64)
    a == 1 && return LogUtility()
    return(IsoElastic(a))
end

# Methods for iso-elastic types
struct IsoElasticVar2 end

util(c, ps::IsoElastic, ::Type{IsoElasticVar2})=c^(1-ps.elast)/(1-ps.elast)
util(c, ps::IsoElastic)=util(c,ps,IsoElasticVar2)-1/(1-ps.elast)
invutil(u, ps::IsoElastic)=(u*(1-ps.elast)+1)^(1/(1-ps.elast))

util(c, ::LogUtility)=log(c)
invutil(u, ::LogUtility)=exp(u)

rra(ps::IsoElastic)=ps.elast
rra(::LogUtility)=1.

" Constant absolute risk aversion specification "
struct CARA <: AbstractSingleValuedUtility
    a::Float64
end

util(c, ps::CARA)=(1 - exp(-ps.a * c))/ps.a
invutil(u, ps::CARA)=(-1/ps.a)*log(1 - ps.a*u)
dutil(c, ps::CARA)=exp(-ps.a * c)
invdutil(m, ps::CARA)=-log(m)/ps.a

" Risk neutral specification "
struct RiskNeutral <: AbstractIsoElastic end

util(c, ps::RiskNeutral)=c
invutil(u, ps::RiskNeutral)=u
dutil(c, ps::RiskNeutral)=one(c)
invdutil(c, ps::RiskNeutral)=error("Tried to invert constant marginal utility from risk neutral agent")


text:
using Revise, BeamPropagation, Distributions, StatsBase, StaticArrays, PhysicalConstants.CODATA2018, Plots, BenchmarkTools

const h = PlanckConstant.val
const ħ = h / 2π
const λ = @with_unit 626 "nm"
const k = 2π / λ
const m = @with_unit 57 "u"
const Δv = ħ * k / m

vz_μ = @with_unit 125 "m/s"
vz_σ = @with_unit 25 "m/s"
vxy_μ = @with_unit 0 "m/s"
vxy_σ = @with_unit 25 "m/s"
exit_radius = @with_unit 4 "mm"

const r = (Normal(0, exit_radius/2), Normal(0, exit_radius/2), Normal(0, 0))
const v = (Normal(vxy_μ, vxy_σ), Normal(vxy_μ, vxy_σ), Normal(vz_μ, vz_σ))
const a = (Normal(0, 0), Normal(0, 0), Normal(0, 0))

VBRs = Weights([
    0.9457,     # to 000
    0.0447,     # 100
    3.9e-3,     # 0200
    2.7e-3,     # 200
    9.9e-4,     # 0220
    11.3e-4,    # 0110
    3.9e-4,     # 1200
    1.5e-4,     # 300
    1.3e-4,     # 1220
    0.7e-4,     # 110 (only N=1, assuming roughly 2/3 to 1/3 rotational branching)
    0.4e-4,     # 110, N=2
    5.7e-5,     # 220
    4.3e-5      # other states
])

addressed = [
    true,   # to 000
    true,   # 100
    true,   # 0200
    true,   # 200
    true,   # 0220
    true,   # 0110
    true,   # 1200
    true,   # 300
    true,   # 1220
    true,   # 110
    false,   # 110, N=2
    false,   # 220
    false   # other states
]

transverse = [
    false,  # to 000
    false,  # 100
    false,  # 0200
    false,  # 200
    false,  # 0220
    false,  # 0110
    false,  # 1200
    false,  # 300
    true,   # 1220
    true,   # 110
    false,   # 110, N=2
    false,   # 220
    false   # other states
]

@inline function random_unit3Dvector()
    θ = rand(Uniform(0, 2π))
    z = rand(Uniform(-1, 1))
    return @SVector [sqrt(1-z^2)*cos(θ), sqrt(1-z^2)*sin(θ), z]
end

@inline function transverse_on(z)
    if 0.175 < z < 0.20     # 1st transverse region 17.5 - 20 cm after cell
        return true
    elseif 0.45 < z < 0.50  # 2nd transverse region 45 - 50 cm after cell
        return true
    elseif 0.56 < z < 0.61  # 2nd transverse region 56 - 61 cm after cell
        return true
    else
        return false
    end
end

const detect_rad  = @with_unit 0.5 "cm"
const detect_zloc = @with_unit 70 "cm"
const detect_zlen = @with_unit 0.5 "cm"
const dead_rad = @with_unit 1 "cm"
const dead_len = detect_zloc + detect_zlen

n = Int64(40e6)
scattering_rate = @with_unit 1.2 "MHz"
save_every      = 100
delete_every    = 15
dt              = 1 / scattering_rate
max_steps       = Int64(3.5e4)

@inline function simple_prop(r, v)
    dist_detect = detect_zloc - r[3]
    x_final = r[1] + v[1] * dist_detect / v[3]
    y_final = r[2] + v[2] * dist_detect / v[3]
    return sqrt(x_final^2 + y_final^2)
end
@inline discard(r, v) = (simple_prop(r, v) > dead_rad) || (r[3] > dead_len)
@inline is_detectable(r) = v_perp(r) < detect_rad && (detect_zloc + detect_zlen > r[3] > detect_zloc)

@inline function save(i, r, v, a, state, s)
    s.vzs[i] = v[3]
    s.detectable[i] = simple_prop(r, v) < detect_rad
    return nothing
end

@inline function f(i, r, v, a, state, dt, p, s)
    if is_detectable(r) && ~s.detectable[i]
        s.detectable[i] = true
    end
    if p.addressed[state] & (~p.transverse[state] | transverse_on(r[3]))
        state′ = sample(1:13, p.VBRs)
        v′ = @SVector [v[1], v[2], v[3] - Δv]
        v′ += Δv .* random_unit3Dvector()
        s.photons[i] += 1
    else
        state′ = state
        v′ = v
    end
    s.vzs[i] = v′[3]
    s.states[i] = state′
    return (state′, v′, a)
end

vzs           = zeros(Float64, n)
photons       = zeros(Int64, n)
detectable    = zeros(Bool, n)
states        = ones(Int64, n)

p = @params (VBRs, transverse, addressed)
s = @params (vzs, photons, detectable, states)

@time s₀, sₜ = propagate!(n, r, v, a, f, save, discard, delete_every, dt, max_steps, p, s)

bright = [addressed[final_state] for final_state in sₜ.states]

barhist(s₀.vzs[s₀.detectable], bins=0:25/4:225, alpha=0.5, xlim=[0,225])
barhist!(sₜ.vzs[sₜ.detectable], bins=0:25/4:225, alpha=0.5)
barhist!(sₜ.vzs[sₜ.detectable .& bright], bins=0:25/4:225, alpha=0.5)

sum(abs.(s₀.vzs[s₀.detectable]) .< 10) / n
sum(abs.(sₜ.vzs[sₜ.detectable]) .< 10) / n
sum(abs.(sₜ.vzs[sₜ.detectable .& bright]) .< 10) / n

v_class = 0 .< sₜ.vzs .< 100
histogram2d(s₀.vzs[sₜ.detectable .& bright .& v_class], sₜ.vzs[sₜ.detectable .& bright.& v_class], bins=20)


text:
## posun homografii

base_path_bikes = "/Users/MrTrololord/Downloads/ox_affine/images/bikes/"
# base_path_bikes = "/BIRL/mount/ox_affine/images/leuven/"

img = load(joinpath(base_path_bikes, "img1.png"))
k = 4
imgw = load(joinpath(base_path_bikes, "img$(k).png"))

H = LAP_julia.load_H(joinpath(base_path_bikes, "H1to$(k)p"))
out = LAP_julia.make_flow_from_H(H, size(img))

flip_out = reverse(out, dims = 1)
flip_out = real.(flip_out) .- imag(flip_out) .* 1im

showflow(imresize(flip_out, ratio=0.1))

img, imgw_large = LAP_julia.pad_images(img, imgw)
warped_img = warp_img(img, real.(out), imag.(out), border_strat=:zeros)
# viewimg(warped_img)
imgw_large

# show warp:
maximum(LAP_julia.vec_len.(out))
showflow(imresize(out, ratio=0.1))
showflow(imresize(reverse(out, dims = 1).*(-1), ratio=0.1))
showflow(imresize(reverse(out, dims = 1), ratio=0.1))

img
imgw

#reverse
imgww = warp_img(imgw, -real.(out), -imag.(out), border_strat=:zeros)
img

## posun mym alg

viewimg(img) = colorview(Gray, img)

img, imgw = Gray.(img), Gray.(imgw)
img, imgw = Float64.(Gray.(img)), Float64.(Gray.(imgw))
# img, imgw = imresize(img, (287, 410)), imresize(imgw, (287, 410))

method = sparse_pflap_psnr
timer = TimerOutput(string(method));
method_kwargs = Dict(:display => false, :timer => timer, :match_source_histogram => true, :point_count => 200)
flow_est, source_reg, timer, time_in_secs = test_registration_alg(method, img, imgw, out.*(-1), timer=method_kwargs[:timer], method_kwargs=method_kwargs);

method = pflap
timer = TimerOutput(string(method));
method_kwargs = Dict(:display => false, :timer => timer, :match_source_histogram => true, :prefilter => true)
flow_est, source_reg, timer, time_in_secs = test_registration_alg(method, img, imgw, out.*(-1), timer=method_kwargs[:timer], method_kwargs=method_kwargs);

source_reg_zeros = warp_img(imgw, -real.(flow_est), -imag.(flow_est), border_strat=:zeros)
img
viewimg(source_reg_zeros)
imgw_large

# with landmarks
imgshow(source_reg_zeros, origin_bot=true)
addpoints(transformed_source_inds)

imgshow(imgw, origin_bot=true)
addpoints(source_inds)

# ground  truth
warped_img

showflow(imresize(flow_est, ratio=0.1))
showflow(imresize(out, ratio=0.1))
# showflow(imresize(view(flow_est.*(-1), 1:50, 1:50), ratio=0.4))

imgoverlay(source_reg_zeros, imgw, figtitle="prekryv s moji registraci")
imgoverlay(warped_img, imgw, figtitle="prekryv s posunem z homografie")

## see points

imgshow(source_reg, figtitle = "Source Registered")
imgshow(img, figtitle = "Target")

showflow(flow.*(-1), figtitle = "Truth Flow")
showflow(flow_est, figtitle = "Estimated Flow")

inds, mag = find_edge_points(img, debug=true, spacing = 20)
imgshow(mag)
addpoints(inds)

### get max displacement

ox_folder = "/Users/MrTrololord/Downloads/ox_affine/images/"

folders = ["bikes", "leuven", "trees", "ubc"]
# folder = "trees"
for folder in folders
    k = 2
    H = LAP_julia.load_H(joinpath(ox_folder, folder, "H1to$(k)p"))
    out = LAP_julia.make_flow_from_H(H, size(img))
    cur_max = maximum(LAP_julia.vec_len.(out))
    @info cur_max
    max_disp = cur_max
    min_max_disp = cur_max

    for k in 3:6
        H = LAP_julia.load_H(joinpath(ox_folder, folder, "H1to$(k)p"))
        out = LAP_julia.make_flow_from_H(H, size(img))
        cur_max = maximum(LAP_julia.vec_len.(out))
        @info cur_max
        max_disp = max(cur_max, max_disp)
        min_max_disp = min(cur_max, min_max_disp)
    end

    @info folder max_disp min_max_disp
end

### get example images

ox_folder = "/Users/MrTrololord/Downloads/ox_affine/images/"
plots_folder = "/Users/MrTrololord/Google_Drive/cvut/bakalarka/plots"

folders = ["bikes", "leuven", "trees", "ubc"]
method = sparse_pflap

for folder in folders
    k = 3
    img = load(joinpath(ox_folder, folder, "img1.png"))
    imgw = load(joinpath(ox_folder, folder, "img$(k).png"))
    H = LAP_julia.load_H(joinpath(ox_folder, folder, "H1to$(k)p"))
    ground_flow = LAP_julia.make_flow_from_H(H, size(img))
    img_save_path = joinpath(plots_folder, "$(folder)_target.png")
    save(img_save_path, img)
    imgw_save_path = joinpath(plots_folder, "$(folder)_source.png")
    save(imgw_save_path, imgw)

    flip_out = reverse(ground_flow, dims = 1)
    flip_out = real.(flip_out) .- imag(flip_out) .* 1im

    showflow(flip_out, figtitle="Ground Truth Deformation")
    ground_flow_save_path = joinpath(plots_folder, "$(folder)_ground_truth_flow.png")
    savefig(ground_flow_save_path)

    timer = TimerOutput(string(method));
    method_kwargs = Dict(:display => false, :timer => timer, :match_source_histogram => true)
    img, imgw = Gray.(img), Gray.(imgw)
    flow_est, source_reg, timer, time_in_secs = test_registration_alg(method, img, imgw, ground_flow.*(-1), timer=method_kwargs[:timer], method_kwargs=method_kwargs);

    flip_flow_est = reverse(flow_est, dims = 1)
    flip_flow_est = real.(flip_flow_est) .- imag(flip_flow_est) .* 1im

    showflow(flip_flow_est, figtitle="Estimated Deformation")
    ground_flow_save_path = joinpath(plots_folder, "$(folder)_estimated_flow_with_$(string(method)).png")
    savefig(ground_flow_save_path)

    @info folder
end

------------------------------------------------------------
                                              Time
                                      ----------------------
           Tot / % measured:               18.1s / 97.8%

Section                       ncalls     time   %tot     avg
------------------------------------------------------------
sparse_pflap_psnr                  1    17.7s   100%   17.7s
  single filter pyramid level      8    17.5s  98.7%   2.18s
    sparse lap                    15    16.5s  93.4%   1.10s
      filtering                   15    15.8s  89.0%   1.05s
      prepare A and b             15    768ms  4.34%  51.2ms
        window sum part 1         45    473ms  2.67%  10.5ms
        window sum part 2         30    295ms  1.66%  9.82ms
      solve linear systems        15    483μs  0.00%  32.2μs
      calculate flow              15   94.3μs  0.00%  6.29μs
    image interpolation           15    612ms  3.46%  40.8ms
    flow interpolation            15    274ms  1.55%  18.3ms
    filter inds                   15    283μs  0.00%  18.8μs
  setup                            1    236ms  1.33%   236ms
    find edge points               1    153ms  0.86%   153ms
    hist match                     1   65.1ms  0.37%  65.1ms
------------------------------------------------------------

------------------------------------------------------------
                                              Time
                                      ----------------------
           Tot / % measured:               44.1s / 100%

Section                       ncalls     time   %tot     avg
------------------------------------------------------------
pflap                              1    43.9s   100%   43.9s
  single filter pyramid level      8    43.8s   100%   5.47s
    lap                           16    19.2s  43.7%   1.20s
      filtering                   16    13.9s  31.6%   866ms
      prepare A and b             16    2.00s  4.56%   125ms
        window sum part 1         48    1.07s  2.44%  22.3ms
        window sum part 2         32    491ms  1.12%  15.4ms
      solve linear systems        16    1.75s  3.98%   109ms
      calculate flow              16    330ms  0.75%  20.6ms
    prefiltering                  40    14.0s  31.9%   350ms
    smoothing                     16    7.98s  18.2%   499ms
    inpainting                    16    1.92s  4.37%   120ms
      replicating borders         16    293ms  0.67%  18.3ms
    image interpolation           16    493ms  1.12%  30.8ms
  setup                            1   90.9ms  0.21%  90.9ms
    hist match                     1   82.7ms  0.19%  82.7ms
------------------------------------------------------------


text:
using PyCall
using Base.Test

enums = pyimport("hearthstone.enums")

CardClass = enums[:CardClass]
CardType = enums[:CardType]
GameTag = enums[:GameTag]

@test 3 == CardType[:HERO]

# CardClass
# (0,"INVALID")
# (1,"DEATHKNIGHT")
# (2,"DRUID")
# (3,"HUNTER")
# (4,"MAGE")
# (5,"PALADIN")
# (6,"PRIEST")
# (7,"ROGUE")
# (8,"SHAMAN")
# (9,"WARLOCK")
# (10,"WARRIOR")
# (11,"DREAM")
# (12,"NEUTRAL")

# CardType
# (0,:INVALID)
# (1,:GAME)
# (2,:PLAYER)
# (3,:HERO)
# (4,:MINION)
# (5,:SPELL)
# (6,:ENCHANTMENT)
# (7,:WEAPON)
# (8,:ITEM)
# (9,:TOKEN)
# (10,:HERO_POWER)

cards = pyimport("fireplace.cards")
cards[:filter](collectible=true) # [fireplace.__init__]: Initializing card database
                                 # [fireplace.__init__]: Merged 2528 cards
db = cards[:db]

karazhan = db["KAR_037t"] # PyObject <KAR_037t: 'Secrets of Karazhan'>
@test "Secrets of Karazhan" == karazhan[:name]
@test CardType[:ENCHANTMENT] == karazhan[:type]
@test CardClass[:NEUTRAL] == karazhan[:card_class]
@test !karazhan[:collectible]

# [x for x in a if x[2][:name]=="Alexstrasza" ]
# EX1_561: 'Alexstrasza'
alex = db["EX1_561"]
@test "Alexstrasza" == alex[:name]
@test CardType[:MINION] == alex[:type]
@test CardClass[:NEUTRAL] == alex[:card_class]
@test alex[:collectible]


text:
type StochasticDeployment <: DeploymentModel
    m::JuMP.Model
    x::Vector{JuMP.Variable}
end
deployment(m::StochasticDeployment) = [round(Int,x) for x in JuMP.getvalue(m.x)]

function StochasticDeployment(p::DeploymentProblem; nperiods=params.nperiods, tol=params.δ,
    solver=GurobiSolver(OutputFlag=0))

    nperiods = min(length(p.train), nperiods)
    demand = p.demand[p.train,:]
    I = 1:p.nlocations
    J = 1:p.nregions
    T = 1:nperiods

    m = JuMP.Model(solver=solver)
    JuMP.@variable(m, x[1:p.nlocations] >= 0, Int)
    JuMP.@variable(m, y[1:p.nlocations,1:p.nregions,1:nperiods] >= 0, Int)
    JuMP.@variable(m, z[1:p.nregions,1:nperiods] >= 0, Int)

    JuMP.@objective(m, Min, sum(z[j,t] for j=J, t=T) + tol*sum(y[i,j,t] for i=I, j=J, t=T))

    JuMP.@constraint(m, sum(x[i] for i=I) <= p.nambulances)

    for j in J # coverage over all regions
        JuMP.@constraint(m, sum(x[i] for i in filter(i->p.coverage[j,i], I)) >= 1)
    end

    # flow constraints at each station
    for i in I, t in T
        JuMP.@expression(m, outflow, sum(y[i,j,t] for j in filter(j->p.coverage[j,i], J)))
        JuMP.@constraint(m, x[i] >= outflow)
    end

    # shortfall from satisfying demand/calls
    for j in J, t in T
        JuMP.@expression(m, inflow, sum(y[i,j,t] for i in filter(i->p.coverage[j,i], I)))
        JuMP.@constraint(m, z[j,t] >= demand[t,j] - inflow)
    end

    StochasticDeployment(m, x)
end

solve(model::StochasticDeployment) = JuMP.solve(model.m)
solve(model::StochasticDeployment, p::DeploymentProblem) = JuMP.solve(model.m)


text:
# Julia wrapper for header: /usr/local/include/libavfilter/buffersrc.h
# Automatically generated using Clang.jl wrap_c, version 0.0.0

export
    av_buffersrc_buffer,
    av_buffersrc_write_frame,
    av_buffersrc_add_frame

function av_buffersrc_buffer(ctx,buf)
    ccall((:av_buffersrc_buffer,libavfilter),Cint,(Ptr{AVFilterContext},Ptr{AVFilterBufferRef}),ctx,buf)
end

function av_buffersrc_write_frame(ctx,frame)
    ccall((:av_buffersrc_write_frame,libavfilter),Cint,(Ptr{AVFilterContext},Ptr{AVFrame}),ctx,frame)
end

function av_buffersrc_add_frame(ctx,frame)
    ccall((:av_buffersrc_add_frame,libavfilter),Cint,(Ptr{AVFilterContext},Ptr{AVFrame}),ctx,frame)
end