text:
module LatCFT

using PyPlot
using PlotTools
using EDutils

function local_op_in_basis(lop::Vector{Array}, bvecsL::Vector{Array}, bvecsR::Vector{Array})
    @assert length(lop) == ndims(bvecsL[1])
    @assert length(lop) == ndims(bvecsR[1])

    op_new = zeros(eltype(bvecsL[1]), (length(bvecsL),length(bvecsR)))
    for j in 1:length(bvecsR)
        v = bvecsR[j]
        opv = EDutils.apply_local_op(vec(v), size(v), lop)
        for k in 1:length(bvecsL)
            w = bvecsL[k]
            op_new[k,j] = (vec(w)' * vec(opv))[1]
        end
    end
    op_new
end

function svgens(n::Int, Hcomp, normfac::Float64=1.0)
    N = length(Hcomp)
    d = size(Hcomp[1][1], 1) #assume all sites have same local dimension
    ns_max = div(maximum(map(ndims, Hcomp[1])), 2) #assume range of terms is the same for all sites

    L = Array[]
    for j in 1:N
        hs = Hcomp[j]
        Ln = 0.0
        for h in hs
            ns = div(ndims(h), 2)
            size(h, 1) == 0 && continue
            #Note: Recall that kron() argument ordering is the reverse of col-major index ordering
            hr = reshape(kron(eye(d^(ns_max-ns)), reshape(h, (d^ns, d^ns))), ntuple(x->d, 2ns_max))
            pos = j + ns/2 - 0.5
            Ln += cis(n * pos * 2π/N) * hr
        end
        push!(L, Ln)
    end
    L
end

function svgens{M}(n::Int, Hcomp::NTuple{M,Array}, Nsites::Int, normfac::Float64=1.0)
    hs = NTuple{M,Array}[Hcomp for n in 1:Nsites]
    svgens(n, hs, normfac)
end

eV_site_decomp(eV::Matrix, d::Int, N::Int) = Array[reshape(eV[:,j], ntuple(x->d,N)) for j in 1:size(eV,2)]

function svgens_in_basis{M}(n::Int, Hcomp::NTuple{M,Matrix}, Nsites::Int, eVL::Matrix, eVR::Matrix, d::Int)
    eVL = eV_site_decomp(eVL, d, Nsites)
    eVR = eV_site_decomp(eVR, d, Nsites)

    Hcomp_r = ntuple(m->reshape(Hcomp[m], ntuple(x->d, 2m)), M)

    Ln = svgens(n, Hcomp_r, Nsites)

    local_op_in_basis(Ln, eVL, eVR)
end

function Hns_ed(N::Number, evd, h_comp, d::Int, maxn::Int=2)
    en_pbc, mom_pbc, eV = evd[1:3]

    #note: shifting spectrum of H should not be necessary, since the modulation should cancel any identity contribution
    Matrix[svgens_in_basis(+n, h_comp, N, eV, eV, d) for n in 1:maxn]
end

function Hns_ed!(Hn_data, ev_data, h_comp, d::Int, maxn::Int=2; overwrite::Bool=true)
    for (N, evd) in ev_data
        !overwrite && haskey(Hn_data, N) && continue

        en_pbc, mom_pbc, eVR = evd[1:3]
        eVL = length(evd) >= 4 ? evd[4] : eVR

        Hn_data[N] = Hns_ed(N, evd, h_comp, d, maxn)
    end
    Hn_data
end

moms_Tev(Tev, N::Number, mom_res::Int) = angle(Tev.^mom_res)/π * N/(2mom_res)

function guess_Ts(H2)
    iT1 = indmax(abs(H2[:,1]))
    iT2 = indmax(abs(H2[1,:]))
    iT1,iT2
end

function normalising_factor(en, H2)
    iT1, iT2 = guess_Ts(H2)
    4.0 / (en[iT1] - en[1] + en[iT2] - en[1])
end

function normalising_factors{T,S1,S2}(ev_data::Dict{T,S1}, Hn_data::Dict{T,S2})
    fs = Dict{T,Float64}()
    for (k, Hns) in Hn_data
        fs[k] = normalising_factor(ev_data[k][1], Hns[2])
    end
    fs
end

function normalise_Hn!(Hn_data, ev_data)
    for (N, Hns) in Hn_data
        en = ev_data[N][1]
        fac = normalising_factor(en, Hns[2])
        @show N, fac
        for Hn in Hns
            scale!(Hn, fac)
        end
    end
end

function estimate_c(Hns)
    H2 = Hns[2]
    2 * (H2' * H2)[1,1]
end

function estimate_c_filterT(Hns)
    H2 = Hns[2]
    iT1, iT2 = guess_Ts(H2)
    2 * abs2(H2[iT1,1])
end

function plot_spec(ev_num::Int, N::Int, edata, op::Matrix; mom_res::Int=1, norm_fac=1.0, E0=nothing, c="r")
    en, Tev = edata[1:2]
    E0 = (E0 == nothing ? en[1] : E0)
    en = norm_fac * (en - E0)

    spins = moms_Tev(Tev, N, mom_res)
    ols = op[:,ev_num]

    plot(spins, en, "kx", markeredgewidth=1.0)
    scatter(spins, en, s=300*log10(100 * abs(ols)), alpha=0.8, edgecolors="None", color=c)
    scatter(spins[ev_num], en[ev_num], s=200, alpha=0.8, edgecolors=c, facecolors="none")
    xlabel(L"s")
    ylabel(L"\Delta")

    ols
end

function ols_low_high(op::Matrix, en::Vector, moms::Vector, tol::Float64=0.01; max_mom::Float64=Inf)
    lowen = Vector{Bool}[en .< en[j] for j in 1:length(en)]
    ols_lower = Float64[vecnorm(op[lowen[j],j]) for j in 1:length(en)]
    ols_higher = Float64[vecnorm(op[!lowen[j],j]) for j in 1:length(en)]
    inds = find(ol-> ol .<= tol, ols_lower)
    mominds = find(p-> abs(p) - sqrt(eps()) < max_mom, moms)
    inds = intersect(inds, mominds)
    ols_lower, ols_higher, inds
end

function qps(Hdata, en, moms, tol::Float64=0.01; max_mom::Float64=Inf)
    H1 = Hdata[1]
    op = (H1 + H1')/2
    ols_lower, ols_higher, qpinds = ols_low_high(op, en, moms, tol, max_mom=max_mom - 1.0)
    qpinds, ols_lower, ols_higher
end

function ps(Hdata, en, moms, tol::Float64=0.01; max_mom::Float64=Inf)
    H1,H2 = Hdata[1:2]

    ols_lower1, ols_higher1, inds1 = ols_low_high((H1 + H1')/2, en, moms, tol, max_mom=max_mom - 2.0)
    ols_lower2, ols_higher2, inds2 = ols_low_high((H2 + H2')/2, en, moms, tol, max_mom=max_mom - 2.0)
    pinds = intersect(inds1, inds2)
    ols_lower = ols_lower1 + ols_lower2
    ols_higher = ols_higher1 + ols_higher2

    pinds, ols_lower, ols_higher
end

function plot_qps(edata, Hdata, N::Int, tol::Float64=0.01; mom_res::Int=1, norm_fac=1.0, E0=nothing, max_mom::Float64=Inf, c::String="r", mrk="o")
    en, Tev = edata[1:2]
    E0 = (E0 == nothing ? en[1] : E0)
    en = norm_fac * (en - E0)

    spins = moms_Tev(Tev, N, mom_res)

    qpinds, ols = qps(Hdata, en, spins, tol, max_mom=max_mom)[1:2]

    plot(spins, en, "kx", markeredgewidth=0.7)
    scatter(spins, en, s=Float64[j in qpinds ? 100.0 : 0.0 for j in 1:length(ols)], alpha=0.8, edgecolors="None", color=c, zorder=10, marker=mrk)
    xlabel(L"s")
    ylabel(L"\Delta")

    qpinds, ols
end

function plot_ps(edata, Hdata, N::Int, tol::Float64=0.01; mom_res::Int=1, norm_fac=1.0, E0=nothing, max_mom::Float64=Inf, min_en::Float64=0.0, max_en::Float64=Inf, c::String="r", mrk="o", mark_pris::Bool=true)
    en, Tev = edata[1:2]
    E0 = (E0 == nothing ? en[1] : E0)
    en = norm_fac * (en - E0)

    spins = moms_Tev(Tev, N, mom_res)

    pinds, ols = ps(Hdata, en, spins, tol, max_mom=max_mom)[1:2]

    ind1 = findfirst(e-> e > min_en - sqrt(eps()), en)
    ind2 = findlast(e-> e < max_en + sqrt(eps()), en)
    plot(spins[ind1:ind2], en[ind1:ind2], "kx", markeredgewidth=0.7)
    if mark_pris
        scatter(spins[ind1:ind2], en[ind1:ind2], s=Float64[j in pinds ? 100.0 : 0.0 for j in ind1:ind2], alpha=0.8, edgecolors="None", color=c, zorder=10, marker=mrk)
    end
    xlabel(L"s")
    ylabel(L"\Delta")

    pinds = pinds[(pinds .>= ind1) & (pinds .<= ind2)]
    pinds, ols[ind1:ind2]
end

function plot_scaling(ev_data, Hs_data, normfacs, f, tol::Float64=0.01; Ns=sort(collect(keys(Hs_data))), plot_ratio::Bool=false, mom_res::Int=1, max_mom::Float64=Inf, min_err::Float64=0.0, min_en::Float64=0.0, max_en::Float64=Inf, fmt::String="o", Npow::Float64=1.0, fit::Bool=true, fit_Ns=Ns, col=nothing)
    data = Dict{Int, Any}()
    maxplot = 0
    for N in Ns
        en, Tev = ev_data[N][1:2]
        en -= en[1]
        en *= normfacs[N]
        spins = moms_Tev(Tev, N, mom_res)
        inds, ols_lower, ols_higher = f(Hs_data[N], en, spins, tol, max_mom=max_mom)
        ols_lower = ols_lower[inds]
        ols_higher = ols_higher[inds]
        flt = (abs(ols_lower) .> min_err) & (spins[inds] .> -sqrt(eps())) & (en[inds] .> min_en-sqrt(eps())) & (en[inds] .< max_en+sqrt(eps()))
        data[N] = (ols_lower[flt], ols_higher[flt], en[inds][flt], spins[inds][flt])
        maxplot = max(maxplot, countnz(flt))
    end

    xs = 1.0 ./ collect(Ns).^Npow

    res = []
    if plot_ratio
        for j in 1:maxplot
            ys = Float64[j <= length(data[N][1]) ? data[N][2][j] / data[N][1][j] : NaN for N in Ns]
            !all(isnan(ys)) && plot(xs, ys, fmt)
        end
        ylabel("Ratio of higher to lower energy overlaps")
    else
        for j in 1:maxplot
            yslo = Float64[j <= length(data[N][1]) ? data[N][1][j] : NaN for N in Ns]
            en = Float64[j <= length(data[N][1]) ? data[N][3][j] : NaN for N in Ns][end]
            mom = Float64[j <= length(data[N][1]) ? data[N][4][j] : NaN for N in Ns][end]
            if !all(isnan(yslo))
                push!(res, (xs, yslo, en, mom))

                if fit
                    flt = !isnan(yslo)
                    flt_Ns = Bool[N in fit_Ns for N in Ns[flt]]
                    a, b = linreg(xs[flt][flt_Ns], yslo[flt][flt_Ns])
                    @show a, b
                    fitx = linspace(0, maximum(xs))
                    plot(fitx, a + fitx * b, "-", color=col)
                end

                if !isnan(en + mom)
                    lbl = "\$ E\\approx $(round(real(en),1)), \\;p= $(round(Int, real(mom))) \$"
                else
                    lbl = ""
                end
                plot(xs, yslo, fmt, label=lbl, color=col)
            end
        end
        ylabel("Norm of overlap vector")
    end

    xlabel("\$1/N^{$(Npow)}\$")

    res
end

function towers(towerf, Hdata, pris::Vector{Int}, threshold::Float64=0.0)
    ols = zeros(Complex128, (size(Hdata[1],1), length(pris)))

    for j in 1:length(pris)
        ols[:,j] = towerf(Hdata, pris[j])
    end

    towers = zeros(Int, size(ols,1))
    olsabs = abs(ols)
    for j in 1:size(ols,1)
        towers[j] = indmax(olsabs[j,:])
        if olsabs[j, towers[j]] < threshold
            towers[j] = 0
        end
    end

    towers, ols
end

function H1_tower(Hdata, ev_num::Int)
    H1 = Hdata[1]

    op = expm(1.0im * (H1+H1'))
    ols = op[:,ev_num]
    ols[ev_num] = 0.0

    ols
end

function plot_H1_towers(edata, Hdata, norm_fac, N::Number, qpris::Vector{Int}; E0=nothing, cs=nothing, cdef=nothing, mom_res::Int=1, sz::Int=20, szp::Int=40, pex_yrange=0.12, pex_xspace=0.12, pex_rounddigits=6, maxDelta::Float64=Inf, tower_tol::Float64=1e-6, qpmarker="D", qpris_markflt=trues(length(qpris)))
    en, Tev = edata[1:2]
    E0 = (E0 == nothing ? en[1] : E0)
    en = norm_fac * (en - E0)
    spins = moms_Tev(Tev, N, mom_res)

    tows, ols = towers(H1_tower, Hdata, qpris, tower_tol)
    if cs == nothing
        cs_rest = nothing
    else
        cs_mod = [cdef, cs...]
        cs_rest = [cs_mod[ind+1] for ind in tows]
    end

    xs_p, ys_p = PlotTools.pointexclusion1D(spins[qpris], en[qpris], pex_yrange, pex_xspace, pex_rounddigits)
    flt = ys_p[qpris_markflt] .< maxDelta
    cs_flt = cs == nothing ? nothing : cs[qpris_markflt][flt]

    scatter(xs_p[qpris_markflt][flt], ys_p[qpris_markflt][flt], edgecolors=cs_flt, marker=qpmarker, s=szp, zorder=20, facecolors="none")

    flt = ys_p[!qpris_markflt] .< maxDelta
    cs_flt = cs == nothing ? nothing : cs[!qpris_markflt][flt]
    scatter(xs_p[!qpris_markflt][flt], ys_p[!qpris_markflt][flt], c=cs_flt, marker="o", s=sz, zorder=10)

    notqpris = trues(en)
    notqpris[qpris] = false
    xs, ys = PlotTools.pointexclusion1D(spins[notqpris], en[notqpris], pex_yrange, pex_xspace, pex_rounddigits)
    flt = ys .< maxDelta
    cs_rest_flt = cs_rest == nothing ? nothing : cs_rest[notqpris][flt]
    scatter(xs[flt], ys[flt], c=cs_rest_flt, marker="o", s=sz, zorder=10)
end

function H1H2_tower(Hdata, ev_num::Int)
    H1, H2 = Hdata[1:2]

    op = expm(1.0im * (H1+H1'+H2+H2'))
    ols = op[:,ev_num]
    ols[ev_num] = 0.0

    ols
end

function plot_H1H2_towers(edata, Hdata, norm_fac, N::Number, pris::Vector{Int}; E0=nothing, cs=nothing, mom_res::Int=1, sz::Int=20, szp::Int=40, pex_yrange=0.12, pex_xspace=0.12, pex_rounddigits=6, maxDelta::Float64=Inf)
    en, Tev = edata[1:2]
    E0 = (E0 == nothing ? en[1] : E0)
    en = norm_fac * (en - E0)
    spins = moms_Tev(Tev, N, mom_res)

    tows, ols = towers(H1H2_tower, Hdata, pris)
    if cs == nothing
        cs_rest = nothing
    else
        cs_rest = [cs[ind] for ind in tows]
    end

    xs_p, ys_p = PlotTools.pointexclusion1D(spins[pris], en[pris], pex_yrange, pex_xspace, pex_rounddigits)
    flt = ys_p .< maxDelta
    cs_flt = cs == nothing ? nothing : cs[flt]
    scatter(xs_p[flt], ys_p[flt], c=cs_flt, marker="D", s=szp, zorder=10)

    notpris = trues(en)
    notpris[pris] = false
    xs, ys = PlotTools.pointexclusion1D(spins[notpris], en[notpris], pex_yrange, pex_xspace, pex_rounddigits)
    flt = ys .< maxDelta
    cs_rest_flt = cs_rest == nothing ? nothing : cs_rest[notpris][flt]
    scatter(xs[flt], ys[flt], c=cs_rest_flt, marker="o", s=sz, zorder=10)
end

end

max_stars_repo_name: amilsted/spinchain-virasoro-demo

max_issues_repo_name: amilsted/spinchain-virasoro-demo

max_forks_repo_name: amilsted/spinchain-virasoro-demo


text:
"""
    prim_mst(g, distmx=weights(g))

Return a vector of edges representing the minimum spanning tree of a connected, undirected graph `g` with optional
distance matrix `distmx` using [Prim's algorithm](https://en.wikipedia.org/wiki/Prim%27s_algorithm).
Return a vector of edges.
"""
function prim_mst end
@traitfn function prim_mst(g::AG::(!IsDirected),
    distmx::AbstractMatrix{T}=weights(g)) where {T <: Real, U, AG <: AbstractGraph{U}}

    nvg = nv(g)

    pq = PriorityQueue{U, T}()
    finished = zeros(Bool, nvg)
    wt = fill(typemax(T), nvg) #Faster access time
    parents = zeros(U, nv(g))

    pq[1] = typemin(T)
    wt[1] = typemin(T)

    while !isempty(pq)
        v = dequeue!(pq)
        finished[v] = true

        for u in neighbors(g, v)
            finished[u] && continue

            if wt[u] > distmx[u, v]
                wt[u] = distmx[u, v]
                pq[u] = wt[u]
                parents[u] = v
            end
        end
    end

    return [Edge{U}(parents[v], v) for v in vertices(g) if parents[v] != 0]
end

ext: jl


text:
using Test
using MathOptInterface
const MOI = MathOptInterface
const MOIT = MOI.Test
const MOIU = MOI.Utilities
const MOIB = MOI.Bridges

import KNITRO

# Default configuration.
const config = MOIT.TestConfig(atol=1e-5, rtol=1e-8,
                               optimal_status=MOI.LOCALLY_SOLVED,
                               query=false,
                               infeas_certificates=false, # Do not ask for infeasibility certificates.
                               modify_lhs=false)

const config_noduals = MOIT.TestConfig(atol=1e-5, rtol=1e-8,
                                       optimal_status=MOI.LOCALLY_SOLVED,
                                       query=false,
                                       duals=false,
                                       infeas_certificates=false,
                                       modify_lhs=false)

@testset "MOI utils" begin
    @testset "SolverName" begin
        optimizer = KNITRO.Optimizer()
        @test MOI.get(optimizer, MOI.SolverName()) == "Knitro"
        KNITRO.free(optimizer)
    end
    @testset "supports_default_copy_to" begin
        optimizer = KNITRO.Optimizer()
        @test MOIU.supports_default_copy_to(optimizer, false)
        # Use `@test !...` if names are not supported
        @test MOIU.supports_default_copy_to(optimizer, true)
        KNITRO.free(optimizer)
    end
    @testset "MOI.Silent" begin
        optimizer = KNITRO.Optimizer()
        @test MOI.supports(optimizer, MOI.Silent())
        @test MOI.get(optimizer, MOI.Silent()) == false
        MOI.set(optimizer, MOI.Silent(), true)
        @test MOI.get(optimizer, MOI.Silent()) == true
        KNITRO.free(optimizer)
    end
    @testset "MOI.TimeLimitSec" begin
        optimizer = KNITRO.Optimizer()
        @test MOI.supports(optimizer, MOI.TimeLimitSec())
        # TimeLimitSec is set to 1e8 by default in Knitro.
        @test MOI.get(optimizer, MOI.TimeLimitSec()) == 1e8
        my_time_limit = 10.
        MOI.set(optimizer, MOI.TimeLimitSec(), my_time_limit)
        @test MOI.get(optimizer, MOI.TimeLimitSec()) == my_time_limit
        KNITRO.free(optimizer)
    end
    @testset "MOI.RawAttribute" begin
        # Test special RawAttributes
        optimizer = KNITRO.Optimizer()
        option_file = joinpath(dirname(pathof(KNITRO)),"..", "examples", "knitro.opt")
        MOI.set(optimizer, MOI.RawParameter("option_file"), option_file)
        tuner_file = joinpath(dirname(pathof(KNITRO)),"..", "examples", "tuner-fixed.opt")
        MOI.set(optimizer, MOI.RawParameter("tuner_file"), tuner_file)
        KNITRO.free(optimizer)
    end
end

const OPTIMIZER = KNITRO.Optimizer()
MOI.set(OPTIMIZER, MOI.RawParameter("outlev"), 0)
MOI.empty!(OPTIMIZER)

# Build bridge optimizer.
const BRIDGED = MOIB.full_bridge_optimizer(OPTIMIZER, Float64)

@testset "MOI Linear tests" begin
    exclude = ["linear1",
               "linear2", # DualObjectivevalue not supported
               "linear4", # ConstraintSet not supported
               "linear5",
               "linear6", # ConstraintSet not supported
               "linear7", # Delete not allowed
               "linear8a", # Behavior in infeasible case doesn't match test.
               "linear8b", # Investigate
               "linear8c", # Problem catching infeasibility ray
               "linear10", # Delete not allowed
               "linear11", # problem accessing constraint function
               "linear12", # Same as above.
               "linear14", # Delete not allowed
               "linear15", # DualObjectivevalue not supported
               ]
    MOIT.contlineartest(BRIDGED, config, exclude)
    # Test linear2 and linear15 without querying the dual solution.
    MOIT.linear15test(BRIDGED, config_noduals)
    MOIT.linear2test(BRIDGED, config_noduals)
end

@testset "MOI QP/QCQP tests" begin
    # Exclude QP tests as Knitro does not support update in objective
    # function after a solve.
    exclude = String["qp2", "qp3"]
    MOIT.contquadratictest(BRIDGED, config, exclude)
end

@testset "MOI complementarity tests" begin
    MOIT.test_qp_complementarity_constraint(BRIDGED, config)
    MOIT.test_linear_mixed_complementarity(BRIDGED, config)
end

#= @testset "MOI SOCP tests" begin =#
#=     # TODO: DualObjectivevalue not supported =#
#=     # Presolve must be switch off to get proper dual variables. =#
#=     config2 = MOIT.TestConfig(atol=1e-4, rtol=1e-4, infeas_certificates=false, =#
#=                               optimal_status=MOI.LOCALLY_SOLVED, query=false) =#
#=     # Behavior in infeasible case doesn't match test. =#
#=     exclude = ["lin4"] =#
#=     BRIDGED2 = MOIB.full_bridge_optimizer(KNITRO.Optimizer(outlev=0, presolve=0), Float64) =#
#=     MOIT.lintest(BRIDGED2, config2) =#
#= end =#

@testset "MOI NLP tests" begin
    MOIT.nlptest(OPTIMIZER, config)
end

@testset "MOI MILP test" begin
    MOIT.knapsacktest(OPTIMIZER, config)
end

KNITRO.free(OPTIMIZER)


text:
# Indicator constraints

MOI.supports_constraint(::Optimizer, ::Type{<:MOI.VectorAffineFunction}, ::Type{<:MOI.Indicator{MOI.ACTIVATE_ON_ONE, <:MOI.LessThan}}) = true

function MOI.add_constraint(o::Optimizer, func::MOI.VectorAffineFunction{T}, set::MOI.Indicator{MOI.ACTIVATE_ON_ONE, LT}) where {T<:Real, LT<:MOI.LessThan}
    allow_modification(o)
    first_index_terms  = [v.scalar_term for v in func.terms if v.output_index == 1]
    scalar_index_terms = [v.scalar_term for v in func.terms if v.output_index != 1]
    length(first_index_terms) == 1 || error("There should be exactly one term in output_index 1, found $(length(first_index_terms))")
    y = VarRef(first_index_terms[1].variable.value)
    x = [VarRef(vi.variable.value) for vi in scalar_index_terms]
    a = [vi.coefficient for vi in scalar_index_terms]
    b = func.constants[2]
    # a^T x + b <= c ===> a^T <= c - b

    cr = add_indicator_constraint(o.inner, y, x, a, MOI.constant(set.set) - b)
    ci = CI{MOI.VectorAffineFunction{T}, MOI.Indicator{MOI.ACTIVATE_ON_ONE, LT}}(cr.val)
    register!(o, ci)
    register!(o, cons(o, ci), cr)
    return ci
end

function MOI.delete(o::Optimizer, ci::CI{MOI.VectorAffineFunction{T}, MOI.Indicator{MOI.ACTIVATE_ON_ONE, LT}}) where {T<:Real, LT<:MOI.LessThan}
    _throw_if_invalid(o, ci)
    allow_modification(o)
    delete!(o.constypes[MOI.VectorAffineFunction{T}, MOI.Indicator{MOI.ACTIVATE_ON_ONE, LT}], ConsRef(ci.value))
    delete!(o.reference, cons(o, ci))
    delete(o.inner, ConsRef(ci.value))
    return nothing
end

function MOI.get(o::Optimizer, ::MOI.ConstraintFunction, ci::CI{MOI.VectorAffineFunction{T}, MOI.Indicator{MOI.ACTIVATE_ON_ONE, LT}}) where {T<:Real, LT<:MOI.LessThan}
    _throw_if_invalid(o, ci)
    indicator_cons = cons(o, ci)::Ptr{SCIP_CONS}
    bin_var = SCIPgetBinaryVarIndicator(indicator_cons)::Ptr{SCIP_VAR}
    slack_var = SCIPgetSlackVarIndicator(indicator_cons)::Ptr{SCIP_VAR}
    linear_cons = SCIPgetLinearConsIndicator(indicator_cons)::Ptr{SCIP_CONS}

    nvars::Int = SCIPgetNVarsLinear(o, linear_cons)
    vars = unsafe_wrap(Array{Ptr{SCIP_VAR}}, SCIPgetVarsLinear(o, linear_cons), nvars)
    vals = unsafe_wrap(Array{Float64}, SCIPgetValsLinear(o, linear_cons), nvars)
    aff_terms = [AFF_TERM(vals[i], VI(ref(o, vars[i]).val)) for i=1:nvars
                 if vars[i] != slack_var]

    ind_terms = [VEC_TERM(1, AFF_TERM(1.0, VI(ref(o, bin_var).val)))]
    vec_terms = [VEC_TERM(2, term) for term in aff_terms]

    return VAF(vcat(ind_terms, vec_terms), [0.0, 0.0])
end

function MOI.get(o::Optimizer, ::MOI.ConstraintSet, ci::CI{MOI.VectorAffineFunction{T}, MOI.Indicator{MOI.ACTIVATE_ON_ONE, LT}}) where {T<:Real, LT<:MOI.LessThan}
    _throw_if_invalid(o, ci)
    indicator_cons = cons(o, ci)::Ptr{SCIP_CONS}
    linear_cons = SCIPgetLinearConsIndicator(indicator_cons)::Ptr{SCIP_CONS}

    lhs = SCIPgetLhsLinear(o, linear_cons)
    rhs = SCIPgetRhsLinear(o, linear_cons)
    lhs == -SCIPinfinity(o) || error("Have lower bound on indicator constraint!")

    return MOI.Indicator{MOI.ACTIVATE_ON_ONE}(MOI.LessThan(rhs))
end

function MOI.get(o::Optimizer, ::MOI.ListOfConstraintIndices{F, S}) where {F <: MOI.VectorAffineFunction{Float64}, S <: MOI.Indicator{MOI.ACTIVATE_ON_ONE, MOI.LessThan{Float64}}}
    return Vector{F, S}(o.constypes[F,S])
end


text:
#!/usr/bin/julia

__precompile__(true)

module FunctionWrappers

using Compat

# Used to bypass NULL check
@inline function assume(v::Bool)
    Base.llvmcall(("declare void @llvm.assume(i1)",
                   """
                   %v = trunc i8 %0 to i1
                   call void @llvm.assume(i1 %v)
                   ret void
                   """), Cvoid, Tuple{Bool}, v)
end

Base.@pure is_singleton(@nospecialize(T)) = isdefined(T, :instance)
# Base.@pure get_instance(@nospecialize(T)) = Base.getfield(T, :instance)

@inline function convert_ret(::Type{Ret}, ret) where Ret
    # Only treat `Cvoid` as ignoring return value.
    # Treating all singleton as ignoring return value is also possible as shown in the
    # commented out implementation but it doesn't seem necessary.
    # The stricter rule may help catching errors and can be more easily changed later.
    Ret === Cvoid && return
    # is_singleton(Ret) && return get_instance(Ret)
    return convert(Ret, ret)
end

# Call wrapper since `cfunction` does not support non-function
# or closures
struct CallWrapper{Ret} <: Function end
(::CallWrapper{Ret})(f, args...) where Ret = convert_ret(Ret, f(args...))

# Specialized wrapper for
for nargs in 0:128
    @eval (::CallWrapper{Ret})(f, $((Symbol("arg", i) for i in 1:nargs)...)) where {Ret} =
        convert_ret(Ret, f($((Symbol("arg", i) for i in 1:nargs)...)))
end

# Convert return type and generates cfunction signatures
Base.@pure map_rettype(T) =
    (isbitstype(T) || T === Any || is_singleton(T)) ? T : Ref{T}
Base.@pure function map_cfunc_argtype(T)
    if is_singleton(T)
        return Ref{T}
    end
    return (isbitstype(T) || T === Any) ? T : Ref{T}
end
Base.@pure function map_argtype(T)
    if is_singleton(T)
        return Any
    end
    return (isbitstype(T) || T === Any) ? T : Any
end
Base.@pure get_cfunc_argtype(Obj, Args) =
    Tuple{Ref{Obj}, (map_cfunc_argtype(Arg) for Arg in Args.parameters)...}

if isdefined(Base, Symbol("@cfunction"))
    @generated function gen_fptr(::Type{Ret}, ::Type{Args}, ::Type{objT}) where {Ret,Args,objT}
        quote
            @cfunction($(CallWrapper{Ret}()), $(map_rettype(Ret)),
                       ($(get_cfunc_argtype(objT, Args).parameters...),))
        end
    end
else
    function gen_fptr(::Type{Ret}, ::Type{Args}, ::Type{objT}) where {Ret,Args,objT}
        cfunction(CallWrapper{Ret}(), map_rettype(Ret), get_cfunc_argtype(objT, Args))
    end
end

mutable struct FunctionWrapper{Ret,Args<:Tuple}
    ptr::Ptr{Cvoid}
    objptr::Ptr{Cvoid}
    obj
    objT
    function (::Type{FunctionWrapper{Ret,Args}})(obj::objT) where {Ret,Args,objT}
        objref = Base.cconvert(Ref{objT}, obj)
        new{Ret,Args}(gen_fptr(Ret, Args, objT),
                      Base.unsafe_convert(Ref{objT}, objref), objref, objT)
    end
    (::Type{FunctionWrapper{Ret,Args}})(obj::FunctionWrapper{Ret,Args}) where {Ret,Args} = obj
end

Base.convert(::Type{T}, obj) where {T<:FunctionWrapper} = T(obj)
Base.convert(::Type{T}, obj::T) where {T<:FunctionWrapper} = obj

@noinline function reinit_wrapper(f::FunctionWrapper{Ret,Args}) where {Ret,Args}
    objref = f.obj
    objT = f.objT
    ptr = gen_fptr(Ret, Args, objT)::Ptr{Cvoid}
    f.ptr = ptr
    f.objptr = Base.unsafe_convert(Ref{objT}, objref)
    return ptr
end

@generated function do_ccall(f::FunctionWrapper{Ret,Args}, args) where {Ret,Args}
    # Has to be generated since the arguments type of `ccall` does not allow
    # anything other than tuple (i.e. `@pure` function doesn't work).
    quote
        $(Expr(:meta, :inline))
        ptr = f.ptr
        if ptr == C_NULL
            # For precompile support
            ptr = reinit_wrapper(f)
        end
        assume(ptr != C_NULL)
        objptr = f.objptr
        ccall(ptr, $(map_rettype(Ret)),
              (Ptr{Cvoid}, $((map_argtype(Arg) for Arg in Args.parameters)...)),
              objptr, $((:(convert($(Args.parameters[i]), args[$i]))
                         for i in 1:length(Args.parameters))...))
    end
end

@inline (f::FunctionWrapper)(args...) = do_ccall(f, args)

# Testing only
const identityAnyAny = FunctionWrapper{Any,Tuple{Any}}(identity)

end


text:
function getDeltaKernel{T <: FloatingPoint}(::Type{T})
    nType = T == Float64 ? "double" : "float"

    return "
        #if defined(cl_khr_fp64)  // Khronos extension available?
        #pragma OPENCL EXTENSION cl_khr_fp64 : enable
        #elif defined(cl_amd_fp64)  // AMD extension available?
        #pragma OPENCL EXTENSION cl_amd_fp64 : enable
        #endif
        #define number $nType

        #define R1(x,y) row1[y*D2 + x]
        #define R2(x,y) row2[y*D2 + x]
        #define R3(x,y) row3[y*D2 + x]
        #define R4(x,y) row4[y*D2 + x]
        #define R5(x,y) row5[y*D2 + x]
        #define R6(x,y) row6[y*D2 + x]

        #define Out(x,y) out[y*D2 + x]

        __kernel void delta(
                      __global const number *row1,
                      __global const number *row2,
                      __global const number *row3,
                      __global const number *row4,
                      __global const number *row5,
                      __global const number *row6,
                      __global number *out,
                      const number r1,
                      const number r2,
                      const number r3,
                      const number r4,
                      const number r5,
                      const number r6,
                      const int D1,
                      const int D2) {

        int i = get_global_id(0);
        int j = get_global_id(1);

        Out(i,j) = r1 * R1(i,j) + r2 * R2(i,j) + r3 * R3(i,j) + r4 * R4(i,j) + r5 * R5(i,j) + r6 * R6(i,j);
    }
"
end

function deltaCL!{T <: FloatingPoint}(
    row1_buff, row2_buff, row3_buff, row4_buff, row5_buff, row6_buff,
    out_buff,
    r1 :: Real, r2 :: Real, r3 :: Real , r4 :: Real, r5 :: Real, r6 :: Real,
    d1 :: Int64, d2 :: Int64,
    ctx, queue, program, :: Type{T})

    k = Kernel(program, "delta")

    call(queue, k, (d1,d2), nothing, row1_buff, row2_buff, row3_buff, row4_buff, row5_buff, row6_buff, out_buff, convert(T, r1), convert(T, r2), convert(T, r3), convert(T, r4), convert(T, r5), convert(T, r6), int32(d1), int32(d2))
end


text:
# Anderson-Darling test

export OneSampleADTest, KSampleADTest

abstract ADTest <: HypothesisTest

## ONE SAMPLE AD-TEST
### http://www.itl.nist.gov/div898/handbook/eda/section3/eda35e.htm

function adstats{T<:Real}(x::AbstractVector{T}, d::UnivariateDistribution)
    n = length(x)
    y = sort(x)
    μ = mean(y)
    σ = std(y)
    A² = convert(typeof(μ), -n)
    for i = 1:n
        zi = (y[i] - μ)/σ
        zni1 = (y[n - i + 1] - μ)/σ
        lcdfz = logcdf(d, zi)
        lccdfz = logccdf(d, zni1)
        A² -= (2*i - 1)/n * (lcdfz + lccdfz)
    end
    (n, μ, σ, A²)
end

immutable OneSampleADTest <: ADTest
    n::Int      # number of observations
    μ::Float64  # sample mean
    σ::Float64  # sample std
    A²::Float64 # Anderson-Darling test statistic
end

function OneSampleADTest{T<:Real}(x::AbstractVector{T}, d::UnivariateDistribution)
    OneSampleADTest(adstats(x, d)...)
end

testname(::OneSampleADTest) = "One sample Anderson-Darling test"

function show_params(io::IO, x::OneSampleADTest, ident="")
    println(io, ident, "number of observations:   $(x.n)")
    println(io, ident, "sample mean:              $(x.μ)")
    println(io, ident, "sample SD:                $(x.σ)")
    println(io, ident, "A² statistic:             $(x.A²)")
end

### Ralph B. D'Agostino, Goodness-of-Fit-Techniques, CRC Press, Jun 2, 1986
### https://books.google.com/books?id=1BSEaGVBj5QC&pg=PA97, p.127
function pvalue(x::OneSampleADTest)
    z = x.A²*(1.0 + .75/x.n + 2.25/x.n/x.n)

    if z < .200
        1.0 - exp(-13.436+101.14z-223.73z^2)
    elseif .200 < z < .340
        1.0 - exp(-8.318+42.796z-59.938z^2)
    elseif .340 < z < .600
        exp(0.9177-4.279z-1.38z^2)
    elseif z < 153.467
        exp(1.2937-5.709z+0.0186z^2)
    else
        0.0
    end
end

## K-SAMPLE ANDERSON DARLING TEST
### k-Sample Anderson-Darling Tests, F. W. Scholz; M. A. Stephens, Journal of the American Statistical Association, Vol. 82, No. 399. (Sep., 1987), pp. 918-924.

immutable KSampleADTest <: ADTest
    k::Int        # number of samples
    n::Int        # number of observations
    σ::Float64   # variance A²k
    A²k::Float64 # Anderson-Darling test statistic
end

function KSampleADTest{T<:Real}(xs::AbstractVector{T}...; modified=true)
    KSampleADTest(a2_ksample(xs, modified)...)
end

testname(::KSampleADTest) = "k-sample Anderson-Darling test"

function show_params(io::IO, x::KSampleADTest, ident="")
    println(io, ident, "number of samples:        $(x.k)")
    println(io, ident, "number of observations:   $(x.n)")
    println(io, ident, "SD of A²k:               $(x.σ)")
    println(io, ident, "A²k statistic:           $(x.A²k)")
end

function pvalue(x::KSampleADTest)
    m = x.k - 1
    Tk = (x.A²k - m) / x.σ
    sig = [0.25, 0.1, 0.05, 0.025, 0.01]
    b0 = [0.675, 1.281, 1.645, 1.96, 2.326]
    b1 = [-0.245, 0.25, 0.678, 1.149, 1.822]
    b2 = [-0.105, -0.305, -0.362, -0.391, -0.396]
    tm = b0 + b1 / sqrt(m) + b2 / m

    # Fit a quadratic polynomial to tm and log(sig)
    # Adapted from code by Paulo José Salz Jabardo
    A = zeros(eltype(tm), 5, 3)
    A[:,1] = 1.0
    @inbounds for i = 1:2, k = 1:5
        A[k,i+1] = A[k,i] * tm[k]
    end
    f = A \ log(sig)

    exp(f[1] + f[2]*Tk + f[3]*Tk^2)
end

function a2_ksample(samples, modified=true)
    k = length(samples)
    k < 2 && error("Need at least two samples")

    n = map(length, samples)
    Z = sort(vcat(samples...))
    N = length(Z)
    Z⁺ = unique(Z)
    L = length(Z⁺)

    L < 2 && error("Need more then 1 observation")
    minimum(n) == 0 && error("One of the samples is empty")

    fij = zeros(Int, k, L)
    for i in 1:k
        for s in samples[i]
            fij[i, searchsortedfirst(Z⁺, s)] += 1
        end
    end

    A²k = 0.
    if modified
        for i in 1:k
            inner = 0.
            Mij = 0.
            Bj = 0.
            for j = 1:L
                lj = sum(fij[:,j])
                Mij += fij[i, j]
                Bj += lj
                Maij = Mij - fij[i, j]/2.
                Baj = Bj - lj/2.
                inner += lj/N * (N*Maij-n[i]*Baj)^2 / (Baj*(N-Baj) - N*lj/4.)
            end
            A²k += inner / n[i]
        end
        A²k *= (N - 1.) / N
    else
        for i in 1:k
            inner = 0.
            Mij = 0.
            Bj = 0.
            for j = 1:L-1
                lj = sum(fij[:,j])
                Mij += fij[i, j]
                Bj += lj
                inner += lj/N * (N*Mij-n[i]*Bj)^2 / (Bj*(N-Bj))
            end
            A²k += inner / n[i]
        end
    end

    H = sum(map(i->1./i, n))
    h = sum(1./(1:N-1))
    g = 0.
    for i in 1:N-2
        for j in i+1:N-1
            g += 1. / ((N - i) * j)
        end
    end

    a = (4*g - 6)*(k - 1) + (10 - 6*g)*H
    b = (2*g - 4)*k^2 + 8*h*k + (2*g - 14*h - 4)*H - 8*h + 4*g - 6
    c = (6*h + 2*g - 2)*k^2 + (4*h - 4*g + 6)*k + (2*h - 6)*H + 4*h
    d = (2*h + 6)*k^2 - 4*h*k
    σ²n = (a*N^3 + b*N^2 + c*N + d) / ((N - 1.) * (N - 2.) * (N - 3.))

    (k, N, sqrt(σ²n), A²k)
end


text:
## Type Hierarchy ##
abstract AbstractConverter
abstract FrameByFrameConverter <: AbstractConverter
abstract TrajectoryConverter <: AbstractConverter

# vc performs voice conversion based on specified converter.
function vc(c::FrameByFrameConverter,
            fm::AbstractMatrix{Float64} # feature matrix
            )
    # Split src feature matrix to power and spectral features
    power, src =  fm[1,:], fm[2:end,:]

    T = size(src, 2)
    converted = similar(fm)

    # Perform feature mapping for each time
    for t=1:T
        @inbounds converted[2:end,t] = fvconvert(c, src[:,t])
    end

    # keep original power
    # note that it is assumed that 0-th coef. represents power coef.
    converted[1,:] = power

    converted
end

# vc performs trajectory-based feature conversion. To reduce computational
# complexiy in practice, we split the input sequence to a set of sub-sequences
# and perform trajectory-based conversion for each sub-sequence.
function vc(c::TrajectoryConverter,
            fm::AbstractMatrix{Float64};  # feature matrix
            )
    # Split src feature matrix to power and spectral features
    power, src =  fm[1,:], fm[2:end,:]

    # D = length/2 + 1: order of static spectral features + power
    D, T = size(src, 1)>>1 + 1, size(src, 2)
    converted = Array(eltype(fm), D, T)

    # Perform Trajectory-based mapping
    count = 0
    limit = length(c) # maximum length of sub-sequence
    while true
        b = count * limit + 1
        e = (count+1) * limit
        if e > T
            e = T
        end
        phrase = src[:,b:e]
        converted[2:end,b:e] = fvconvert(c, phrase)

        if e == T
            break
        end
        count += 1
    end

    # keep original power
    converted[1,:] = power

    converted
end


text:
import Base.Sort: Forward, Ordering, Algorithm, lt
using Distributed

using StatsBase

function getmedians(x, n)
    q,r = divrem(length(x), n+1)

    if q == 0
        return x
    end
    buckets = [q for _ in 1:n+1]
    for i in 1:r
        buckets[i] += 1
    end
    pop!(buckets)
    x[cumsum(buckets)]
end

function sortandsample_array(ord, xs, nsamples, presorted=false)
    r = sample(1:length(xs), min(length(xs), nsamples),
               replace=false, ordered=true)
    if !presorted
        sorted = sort(xs, order=ord)
        chunk = tochunk(sorted)
        samples = sorted[r]
    else
        chunk = nothing # avoid communicating metadata if already sorted
        samples = chunk[r]
    end
    (chunk, samples)
end

function batchedsplitmerge(chunks, splitters, batchsize, start_proc=1;
                           merge=merge_sorted,
                           by=identity,
                           sub=getindex,
                           order=default_ord)

    if batchsize >= length(chunks)
        if isempty(splitters)
            # this means we need to combine all chunks
            # into a single chunk.
            return [collect_merge(merge, chunks)]
        else
            # split each chunk by the splitters, then merge
            # corresponding portions from all the chunks
            return splitmerge(chunks, splitters, merge, by, sub, order)
        end
    end

    # group chunks into batches:
    # e.g. if length(chunks) = 128, and batchsize = 8
    # then we will have q = 16, r=0 and length(b) = 16
    q, r = divrem(length(chunks), batchsize)
    b = [batchsize for _ in 1:q]
    r != 0 && push!(b, r) # add the remaining as a batch
    # create ranges that can extract a batch each
    batch_ranges = map(UnitRange, cumsum(vcat(1, b[1:end-1])), cumsum(b))
    # make the batches of chunks
    batches = map(range->chunks[range], batch_ranges)

    # splitmerge each batch
    topsplits, lowersplits = recursive_splitters(order, splitters, length(chunks), batchsize)

    sorted_batches = map(batches) do b
        isempty(topsplits) ?
            [collect_merge(merge, b)] :
            splitmerge(b, topsplits, merge, by, sub, order)
    end

    range_groups = transpose_vecvec(sorted_batches)

    chunks = []
    p = start_proc
    for i = 1:length(range_groups)
        s = lowersplits[i]
        group = range_groups[i]
        if !isempty(s)
            # recursively call batchedsplitmerge with a smaller batch
            cs = batchedsplitmerge(group, s, batchsize; merge = merge, by=by, sub=sub, order=order)
            append!(chunks, cs)
        else
            push!(chunks, collect_merge(merge, group))
        end
    end
    return chunks
end

function collect_merge(merge, group)
    #delayed((xs...) -> treereduce(merge, Any[xs...]))(group...)
    t = treereduce(delayed(merge), group)
end

# Given sorted chunks, splits each chunk according to splitters
# then merges corresponding splits together to form length(splitters) + 1 sorted chunks
# these chunks will be in turn sorted
function splitmerge(chunks, splitters, merge, by, sub, ord)
    c1 = map(c->splitchunk(c, splitters, by, sub, ord), chunks)
    map(cs->collect_merge(merge, cs), transpose_vecvec(c1))
end

function splitchunk(c, splitters, by, sub, ord)
    function getbetween(xs, lo, hi, ord)
        i = searchsortedlast(xs, lo, order=ord)
        j = searchsortedlast(xs, hi, order=ord)
        (i+1):j
    end

    function getgt(xs, lo, ord)
        i = searchsortedlast(xs, lo, order=ord)
        (i+1):length(xs)
    end

    function getlt(xs, hi, ord)
        j = searchsortedlast(xs, hi, order=ord)
        1:j
    end

    between = map((hi, lo) -> delayed(x->sub(x, getbetween(by(x), hi, lo, ord)))(c),
                  splitters[1:end-1], splitters[2:end])
    hi = splitters[1]
    lo = splitters[end]

    a = delayed(x->sub(x, getlt(by(x), hi, ord)))(c)
    b = between
    c = delayed(x->sub(x, getgt(by(x), lo, ord)))(c)
    return vcat(a, b, c)

   #[delayed(c->sub(c, getlt(by(c), hi, ord)))(c);
   # between; delayed(c->sub(c, getgt(by(c), lo, ord)))(c)]
end

# transpose a vector of vectors
function transpose_vecvec(xs)
    map(1:length(xs[1])) do i
        map(x->x[i], xs)
    end
end

const use_shared_array = Ref(true)
function _promote_array(x::AbstractArray{T}, y::AbstractArray{S}) where {T,S}
    Q = promote_type(T,S)
    ok = (isa(x, Array) || isa(x, SharedArray)) && (isa(y, Array) || isa(y, SharedArray))
    if ok && isbitstype(Q) && use_shared_array[] && Distributed.check_same_host([workers()..., 1])
        return SharedArray{Q}(length(x)+length(y), pids=Distributed.procs())
    else
        return similar(x, Q, length(x)+length(y))
    end
end

function merge_sorted(ord::Ordering, x::AbstractArray, y::AbstractArray)
    z = _promote_array(x, y)
    _merge_sorted(ord, z, x, y)
end

function _merge_sorted(ord::Ordering, z::AbstractArray{T}, x::AbstractArray{T}, y::AbstractArray{S}) where {T, S}
    n = length(x) + length(y)
    i = 1; j = 1; k = 1
    len_x = length(x)
    len_y = length(y)
    @inbounds while i <= len_x && j <= len_y
        if lt(ord, x[i], y[j])
            z[k] = x[i]
            i += 1
        else
            z[k] = y[j]
            j += 1
        end
        k += 1
    end
    remaining, m = i <= len_x ? (x, i) : (y, j)
    @inbounds while k <= n
        z[k] = remaining[m]
        k += 1
        m += 1
    end
    z
end

"""
split the splitters themselves into batches.

Args:
- `ord` -- `Sorting.Ordering` object
- `splitters` -- the `nchunks-1` splitters
- `batchsize` -- batch size

Returns:
A Tuple{Vector, Vector{Vector}} -- the coarse splitters which
create `batchsize` splits, finer splitters within those batches
which create a total of `nchunks` splits.

```julia
julia> Dagger.recursive_splitters(Dagger.default_ord,
            [10,20,30,40,50,60], 5,3)
([30], Any[[10, 20], [40, 50, 60]])
```

The first value `[30]` represents a coarse split that cuts the dataset
from -Inf-30, and 30-Inf. Each part is further recursively split using
the next set of splitters

"""
function recursive_splitters(ord, splitters, nchunks, batchsize)

    # final number of chunks
    noutchunks = length(splitters) + 1
    # chunks per batch
    perbatch = ceil(Int, nchunks / batchsize)
    root = getmedians(splitters, perbatch-1)

    subsplits = []
    i = 1
    for c in root
        j = findlast(x->lt(ord, x, c), splitters)
        if j===nothing
            j = length(splitters)
        end
        push!(subsplits, splitters[i:j])
        i = j+2
    end
    push!(subsplits, splitters[i:end])
    root, subsplits
end

arrayorvcat(x::AbstractArray,y::AbstractArray) = vcat(x,y)
arrayorvcat(x,y) = [x,y]

const default_ord = Base.Sort.ord(isless, identity, false, Forward)

"""
`dsort_chunks(cs, [nchunks, nsamples]; options...)`

Sort contents of chunks (`cs`) and return a new set of chunks
such that the chunks when concatenated return a sorted collection.
Each chunk in turn is sorted.

Args:
- `nchunks` -- the number of chunks to produce, regardless of how many chunks were given as input
- `nsamples` -- the number of elements to sample from each chunk to guess the splitters (`nchunks-1` splitters) each chunk will be delimited by the splitter.
- `merge` -- a function to merge two sorted collections.
- `sub` -- a function to get a subset of the collection takes (collection, range) (defaults to `getindex`)
- `order` -- `Base.Sort.Ordering` to be used for sorting
- `batchsize` -- number of chunks to split and merge at a time (e.g. if there are 128 input chunks and 128 output chunks, and batchsize is 8, then we first sort among batches of 8 chunks -- giving 16 batches. Then we sort among the first chunk of the first 8 batches (all elements less than the first splitter), then go on to the first 8 chunks of the second 8 batches, and so on...
- `chunks_presorted` -- is each chunk in the input already sorted?
- `sortandsample` -- a function to sort a chunk, then sample N elements to infer the splitters. It takes 3 arguments: (collection, N, presorted). presorted is a boolean which is true if the chunk is already sorted.
Returns a tuple of `(chunk, samples)` where `chunk` is the `Dagger.Chunk` object. `chunk` can be nothing if no change to the initial array was made (e.g. it was already sorted)
- `affinities` -- a list of processes where the output chunks should go. If the length is not equal to `nchunks` then affinities array is cycled through.
"""
function dsort_chunks(cs, nchunks=length(cs), nsamples=2000;
                      merge = merge_sorted,
                      by=identity,
                      sub=getindex,
                      order=default_ord,
                      batchsize=max(2, nworkers()),
                      splitters=nothing,
                      chunks_presorted=false,
                      sortandsample = (x,ns, presorted)->sortandsample_array(order, x,ns, presorted),
                      affinities=workers(),
                     )
    if splitters !== nothing
        # this means splitters are given beforehand
        nsamples = 0 # no samples needed
    end

    # first sort each chunk and sample nsamples elements from each chunk
    cs1 = map(c->delayed(sortandsample)(c, nsamples, chunks_presorted), cs)

    batchsize = max(2, batchsize)
    # collect the samples
    xs = collect(treereduce(delayed(vcat), cs1))
    if length(cs1) == 1
        xs = [xs]
    end

    # get nchunks-1 splitters that equally split the samples into nchunks pieces
    if splitters === nothing
        samples = reduce((a,b)->merge_sorted(order, a, b), map(x->x[2], xs))
        splitters = getmedians(samples, nchunks-1)
    end

                       # if first(x) === nothing this means that
                       # the chunk was already sorted, so just
                       # use the input chunk as-is
    cs2 = map((x,c) -> x === nothing ? c : x, map(first, xs), cs)

    # main sort routine. At this point:
    # we know the splitters we want to use,
    # and each chunk is already sorted
    batchedsplitmerge(cs2,
                      splitters,
                      batchsize;
        merge=merge, by=by, sub=sub, order=order)
    #=
    for (w, c) in zip(Iterators.cycle(affinities), cs)
        propagate_affinity!(c, Dagger.OSProc(w) => 1)
    end
    =#
end

function propagate_affinity!(c, aff)
    if !isa(c, Thunk)
        return
    end
    if c.affinity !== nothing
        push!(c.affinity, aff)
    else
        c.affinity = [aff]
    end

    for t in c.inputs
        propagate_affinity!(t, aff)
    end
end

function Base.sort(v::ArrayOp;
               lt=Base.isless,
               by=identity,
               nchunks=nothing,
               rev::Bool=false,
               batchsize=max(2, nworkers()),
               nsamples=2000,
               order::Ordering=default_ord)
    v1 = compute(v)
    ord = Base.Sort.ord(lt,by,rev,order)
    nchunks = nchunks === nothing ? length(v1.chunks) : nchunks
    cs = dsort_chunks(v1.chunks, nchunks, nsamples,
                      order=ord, merge=(x,y)->merge_sorted(ord, x,y))
    foreach(persist!, cs)
    t=delayed((xs...)->[xs...]; meta=true)(cs...)
    # `compute(t)` only computes references to materialized version of `cs`
    # we don't want the scheduler to think that `cs`s' job is done
    # so we call persist! above
    chunks = compute(t)
    dmn = ArrayDomain((1:sum(length(domain(c)) for c in chunks),))
    DArray(eltype(v1), dmn, map(domain, chunks), chunks)
end


text:
function Dateneinlesen(Dimension)
    x10 = -1.5; # kleinst möglicher x-Wert
    x20 = -1.5; # kleinst möglicher y-Wert
    N = Dimension; # Anzahl der betrachteten Parallelen Strahlen
    JJ = Dimension+1; # Anzahl der betrachteten Winkelschritte
    nx1 = Dimension; # Anzahl der Spalten
    nx2 = Dimension; # Anzahl der Zeilen
    x11 = x10+nx1; # größt möglicher x-Wert
    x21 = x20+nx2; # größt möglicher y-Wert
    verschx1 = 0.25; # Verschiebung x damit Strahlen nicht auf Kante liegen
    verschx2 = 0.25; # Verschiebung y damit Strahlen nicht auf Kante liegen
    a = -3.5; # für spätere Aufgabe
    b = 3.5; # für spätere Aufgabe
    laengenmatrix=zeros(N*JJ,size(nx1*nx2,1));
    theta1 = 0.0000001; #Referenzwinkel
    counter = 1;
    return x10,x20,N,JJ,nx1,nx2,x11,x21,verschx1,verschx2,a,b,laengenmatrix,theta1,counter;

end


text:
using BenchmarkTools, Random, Statistics, WaveFD

_nthreads = [2^i for i in 0:floor(Int,log2(Sys.CPU_THREADS))]
if Sys.CPU_THREADS ∉ _nthreads
    push!(_nthreads, Sys.CPU_THREADS)
end

const SUITE = BenchmarkGroup()

z0,y0,x0,dz,dy,dx,nt,dt = 0.0,0.0,0.0,10.0,10.0,10.0,3000,0.001

n_2D = (z=parse(Int,get(ENV,"2D_NZ","501")), x=parse(Int,get(ENV,"2D_NX","1001")))
n_3D = (z=parse(Int,get(ENV,"3D_NZ","101")), y=parse(Int,get(ENV,"3D_NY","201")), x=parse(Int,get(ENV,"3D_NX","301")))

nb_2D = (z=parse(Int,get(ENV,"2D_NBZ","$(n_2D.z)")), x=parse(Int,get(ENV,"2D_NBX","8")))
nb_3D = (z=parse(Int,get(ENV,"3D_NBZ","$(n_3D.z)")), y=parse(Int,get(ENV,"3D_NBY","8")), x=parse(Int,get(ENV,"3D_NBX","8")))

@info "size 2D: $n_2D, use ENV[\"2D_NZ\"], ENV[\"2D_NX\"] to customize"
@info "size 3D: $n_3D, use ENV[\"3D_NZ\"], ENV[\"3D_NY\"], ENV[\"3D_NX\"] to customize"
@info "cache block size 2D: $nb_2D, use ENV[\"2D_NBZ\"], ENV[\"2D_NBX\"] to customize"
@info "cache block size 3D: $nb_3D, use ENV[\"3D_NBZ\"], ENV[\"3D_NBY\"], ENV[\"3D_NBX\"] to customize"

rz,rx = dz*ones(n_2D.x),dx*[0:n_2D.x-1;]

iz,ix,c = WaveFD.hickscoeffs(dz, dx, z0, x0, n_2D.z, n_2D.x, rz, rx)
blocks = WaveFD.source_blocking(n_2D.z, n_2D.x, 4, 4, iz, ix, c)
field = zeros(n_2D.z, n_2D.x)
data = rand(nt, n_2D.x)
it = 1

SUITE["2D Utils"] = BenchmarkGroup()
SUITE["2D Utils"]["source_blocking"] = @benchmarkable WaveFD.source_blocking($(n_2D.z), $(n_2D.x), $(nb_2D.z), $(nb_2D).x, $iz, $ix, $c);
SUITE["2D Utils"]["injectdata!"] = @benchmarkable WaveFD.injectdata!($field, $blocks, $data, $it)

dtmod = 0.0002
ntmod = nt * round(Int,dt / dtmod)
h = WaveFD.interpfilters(dtmod, dt, 0, WaveFD.LangC(), Sys.CPU_THREADS)
data_mod = rand(ntmod, n_2D.x)
SUITE["2D Utils"]["interpadjoint!"] = @benchmarkable WaveFD.interpadjoint!($h, $data_mod, $data)
SUITE["2D Utils"]["interpforward!"] = @benchmarkable WaveFD.interpforward!($h, $data, $data_mod)

compressor = WaveFD.Compressor(Float32, Float32, UInt32, (n_2D.z,n_2D.x), (32,32), 1e-2, 1024, false)
field2 = rand(Float32,n_2D.z,n_2D.x)
SUITE["2D compression"] = BenchmarkGroup()
SUITE["2D compression"]["write"] = @benchmarkable WaveFD.compressedwrite(io, $compressor, 1, $field2) setup=(open($compressor); io=open(tempname(),"w")) teardown=(close($compressor); close(io))
SUITE["2D compression"]["read"] = @benchmarkable WaveFD.compressedread!(io, $compressor, 1, $field2) setup=(open($compressor); tfile=tempname(); _io=open(tfile,"w"); WaveFD.compressedwrite(_io, $compressor, 1, $field2); close(_io); io=open(tfile)) teardown=(close($compressor); close(io))

rng2 = (10:n_2D.z-10,10:n_2D.x)
_field2 = view(field2, rng2...)
_compressor = WaveFD.Compressor(Float32, Float32, UInt32, size(_field2), (32,32), 1e-2, 1024, true)
SUITE["2D compression, interior"] = BenchmarkGroup()
SUITE["2D compression, interior"]["write"] = @benchmarkable WaveFD.compressedwrite(io, $_compressor, 1, $field2, $rng2) setup=(open($_compressor); io=open(tempname(),"w")) teardown=(close($_compressor); close(io))
SUITE["2D compression, interior"]["read"] = @benchmarkable WaveFD.compressedread!(io, $_compressor, 1, $field2, $rng2) setup=(open($_compressor); tfile=tempname(); _io=open(tfile,"w"); WaveFD.compressedwrite(_io, $_compressor, 1, $field2, rng2); close(_io); io=open(tfile)) teardown=(close($_compressor); close(io))

SUITE["2DAcoIsoDenQ_DEO2_FDTD"] = BenchmarkGroup([Dict("ncells"=>n_2D.z*n_2D.x,"nbz"=>nb_2D.z,"nbx"=>nb_2D.x,"nthreads"=>_nthreads)])
function p2diso(nthreads,nz,nx,nbz,nbx)
    p = WaveFD.Prop2DAcoIsoDenQ_DEO2_FDTD(freesurface=false, nz=nz, nx=nx, nbz=nbz, nbx=nbx, dz=10.0, dx=10.0, dt=0.001, nthreads=nthreads)
    v,b,pcur,pold = WaveFD.V(p),WaveFD.B(p),WaveFD.PCur(p),WaveFD.POld(p)
    v .= 1500
    b .= 1
    rand!(pcur)
    rand!(pold)
    p
end
for nthreads in _nthreads
    SUITE["2DAcoIsoDenQ_DEO2_FDTD"]["$nthreads threads"] = @benchmarkable WaveFD.propagateforward!(p) setup=(p=p2diso($nthreads,$(n_2D.z),$(n_2D.x),$(nb_2D.z),$(nb_2D.x))) teardown=(free(p))
end

function fields2diso(p::WaveFD.Prop2DAcoIsoDenQ_DEO2_FDTD)
    nz,nx=size(p)
    δm = Dict("v"=>rand(Float32,nz,nx))
    fields = Dict("pspace"=>rand(Float32,nz,nx))
    δm,fields
end
SUITE["2DAcoIsoDenQ_DEO2_FDTD"]["imaging condition, standard"] = @benchmarkable WaveFD.adjointBornAccumulation!(p,WaveFD.Prop2DAcoIsoDenQ_DEO2_FDTD_Model_V(),ic,δm,fields) setup=(p=p2diso(Sys.CPU_THREADS,$(n_2D.z),$(n_2D.x),$(nb_2D.z),$(nb_2D.x)); ic=WaveFD.ImagingConditionStandard(); (δm,fields)=fields2diso(p)) teardown=(free(p))
SUITE["2DAcoIsoDenQ_DEO2_FDTD"]["imaging condition, wave field separation FWI"] = @benchmarkable WaveFD.adjointBornAccumulation!(p,WaveFD.Prop2DAcoIsoDenQ_DEO2_FDTD_Model_V(),ic,δm,fields) setup=(p=p2diso(Sys.CPU_THREADS,$(n_2D.z),$(n_2D.x),$(nb_2D.z),$(nb_2D.x)); ic=WaveFD.ImagingConditionWaveFieldSeparationFWI(); (δm,fields)=fields2diso(p)) teardown=(free(p))
SUITE["2DAcoIsoDenQ_DEO2_FDTD"]["imaging condition, wave field separation RTM"] = @benchmarkable WaveFD.adjointBornAccumulation!(p,WaveFD.Prop2DAcoIsoDenQ_DEO2_FDTD_Model_V(),ic,δm,fields) setup=(p=p2diso(Sys.CPU_THREADS,$(n_2D.z),$(n_2D.x),$(nb_2D.z),$(nb_2D.x));  ic=WaveFD.ImagingConditionWaveFieldSeparationRTM(); (δm,fields)=fields2diso(p)) teardown=(free(p))

SUITE["2DAcoVTIDenQ_DEO2_FDTD"] = BenchmarkGroup([Dict("ncells"=>n_2D.z*n_2D.x,"nbz"=>nb_2D.z,"nbx"=>nb_2D.x,"nthreads"=>_nthreads)])
function p2dvti(nthreads,nz,nx,nbz,nbx)
    p = WaveFD.Prop2DAcoVTIDenQ_DEO2_FDTD(freesurface=false, nz=nz, nx=nx, nbz=nbz, nbx=nbx, dz=10.0, dx=10.0, dt=0.001, nthreads=nthreads)
    v,b,ϵ,η,f,pcur,pold = WaveFD.V(p),WaveFD.B(p),WaveFD.Eps(p),WaveFD.Eta(p),WaveFD.F(p),WaveFD.PCur(p),WaveFD.POld(p)
    rand!(pcur)
    rand!(pold)
    v .= 1500
    b .= 1
    ϵ .= 0.2
    η .= 0.0
    f .= 0.85
    p
end
for nthreads in _nthreads
    SUITE["2DAcoVTIDenQ_DEO2_FDTD"]["$nthreads threads"] = @benchmarkable WaveFD.propagateforward!(p) setup=(p=p2dvti($nthreads,$(n_2D.z),$(n_2D.x),$(nb_2D.z),$(nb_2D.x))) teardown=(free(p))
end

function fields2dvti(p::WaveFD.Prop2DAcoVTIDenQ_DEO2_FDTD)
    nz,nx=size(p)
    δm = Dict("v"=>rand(Float32,nz,nx))
    fields = Dict("pspace"=>rand(Float32,nz,nx),"mspace"=>rand(Float32,nz,nx))
    δm,fields
end
SUITE["2DAcoVTIDenQ_DEO2_FDTD"]["imaging condition, standard"] = @benchmarkable WaveFD.adjointBornAccumulation!(p,WaveFD.Prop2DAcoVTIDenQ_DEO2_FDTD_Model_V(),ic,δm,fields) setup=(p=p2dvti(Sys.CPU_THREADS,$(n_2D.z),$(n_2D.x),$(nb_2D.z),$(nb_2D.x)); ic=WaveFD.ImagingConditionStandard(); (δm,fields)=fields2dvti(p)) teardown=(free(p))
SUITE["2DAcoVTIDenQ_DEO2_FDTD"]["imaging condition, wave field separation FWI"] = @benchmarkable WaveFD.adjointBornAccumulation!(p,WaveFD.Prop2DAcoVTIDenQ_DEO2_FDTD_Model_V(),ic,δm,fields) setup=(p=p2dvti(Sys.CPU_THREADS,$(n_2D.z),$(n_2D.x),$(nb_2D.z),$(nb_2D.x)); ic=WaveFD.ImagingConditionWaveFieldSeparationFWI(); (δm,fields)=fields2dvti(p)) teardown=(free(p))
SUITE["2DAcoVTIDenQ_DEO2_FDTD"]["imaging condition, wave field separation RTM"] = @benchmarkable WaveFD.adjointBornAccumulation!(p,WaveFD.Prop2DAcoVTIDenQ_DEO2_FDTD_Model_V(),ic,δm,fields) setup=(p=p2dvti(Sys.CPU_THREADS,$(n_2D.z),$(n_2D.x),$(nb_2D.z),$(nb_2D.x)); ic=WaveFD.ImagingConditionWaveFieldSeparationRTM(); (δm,fields)=fields2dvti(p)) teardown=(free(p))

SUITE["2DAcoTTIDenQ_DEO2_FDTD"] = BenchmarkGroup([Dict("ncells"=>n_2D.z*n_2D.x,"nbz"=>nb_2D.z,"nbx"=>nb_2D.x,"nthreads"=>_nthreads)])
function p2dtti(nthreads,nz,nx,nbz,nbx)
    p = WaveFD.Prop2DAcoTTIDenQ_DEO2_FDTD(freesurface=false, nz=nz, nx=nx, nbz=nbz, nbx=nbx, dz=10.0, dx=10.0, dt=0.001, nthreads=nthreads)
    v,b,ϵ,η,f,sinθ,cosθ,pcur,pold = WaveFD.V(p),WaveFD.B(p),WaveFD.Eps(p),WaveFD.Eta(p),WaveFD.F(p),WaveFD.SinTheta(p),WaveFD.CosTheta(p),WaveFD.PCur(p),WaveFD.POld(p)
    v .= 1500
    b .= 1
    ϵ .= 0.2
    η .= 0.0
    f .= 0.85
    cosθ .= cos(pi/4)
    sinθ .= sin(pi/4)
    rand!(pcur)
    rand!(pold)
    p
end
for nthreads in _nthreads
    SUITE["2DAcoTTIDenQ_DEO2_FDTD"]["$nthreads threads"] = @benchmarkable WaveFD.propagateforward!(p) setup=(p=p2dtti($nthreads,$(n_2D.z),$(n_2D.x),$(nb_2D.z),$(nb_2D.x))) teardown=(free(p))
end

function fields2dtti(p::WaveFD.Prop2DAcoTTIDenQ_DEO2_FDTD)
    nz,nx = size(p)
    δm = Dict("v"=>rand(Float32,nz,nx))
    fields = Dict("pspace"=>rand(Float32,nz,nx),"mspace"=>rand(Float32,nz,nx))
    δm,fields
end
SUITE["2DAcoTTIDenQ_DEO2_FDTD"]["imaging condition, standard"] = @benchmarkable WaveFD.adjointBornAccumulation!(p,WaveFD.Prop2DAcoTTIDenQ_DEO2_FDTD_Model_V(),ic,δm,fields) setup=(p=p2dtti(Sys.CPU_THREADS,$(n_2D.z),$(n_2D.x),$(nb_2D.z),$(nb_2D.x)); ic=WaveFD.ImagingConditionStandard(); (δm,fields)=fields2dtti(p)) teardown=(free(p))
SUITE["2DAcoTTIDenQ_DEO2_FDTD"]["imaging condition, wave field separation FWI"] = @benchmarkable WaveFD.adjointBornAccumulation!(p,WaveFD.Prop2DAcoTTIDenQ_DEO2_FDTD_Model_V(),ic,δm,fields) setup=(p=p2dtti(Sys.CPU_THREADS,$(n_2D.z),$(n_2D.x),$(nb_2D.z),$(nb_2D.x)); ic=WaveFD.ImagingConditionWaveFieldSeparationFWI(); (δm,fields)=fields2dtti(p)) teardown=(free(p))
SUITE["2DAcoTTIDenQ_DEO2_FDTD"]["imaging condition, wave field separation RTM"] = @benchmarkable WaveFD.adjointBornAccumulation!(p,WaveFD.Prop2DAcoTTIDenQ_DEO2_FDTD_Model_V(),ic,δm,fields) setup=(p=p2dtti(Sys.CPU_THREADS,$(n_2D.z),$(n_2D.x),$(nb_2D.z),$(nb_2D.x)); ic=WaveFD.ImagingConditionWaveFieldSeparationRTM(); (δm,fields)=fields2dtti(p)) teardown=(free(p))

nz,ny,nx = 501,301,1001

rz = [dz for iy = 1:ny, ix = 1:nx][:]
ry = [(iy-1)*dy for iy = 1:ny, ix = 1:nx][:]
rx = [(ix-1)*dx for iy = 1:ny, ix = 1:nx][:]

iz,iy,ix,c = WaveFD.hickscoeffs(dz, dy, dx, z0, y0, x0, nz, ny, nx, rz, ry, rx)
blocks = WaveFD.source_blocking(nz, ny, nx, 4, 4, 4, iz, iy, ix, c)
field = zeros(nz, ny, nx)
data = rand(nt, nx*ny)
it = 1

SUITE["3D Utils"] = BenchmarkGroup()
SUITE["3D Utils"]["source_blocking"] = @benchmarkable WaveFD.source_blocking($nz, $ny, $nx, 256, 8, 8, $iz, $iy, $ix, $c);
SUITE["3D Utils"]["injectdata!"] = @benchmarkable WaveFD.injectdata!($field, $blocks, $data, $it)

compressor = WaveFD.Compressor(Float32, Float32, UInt32, (n_3D.z,n_3D.y,n_3D.x), (32,32,32), 1e-2, 1024, false)
field3 = rand(Float32,n_3D.z,n_3D.y,n_3D.x)
SUITE["3D compression"] = BenchmarkGroup()
SUITE["3D compression"]["write"] = @benchmarkable WaveFD.compressedwrite(io, $compressor, 1, $field3) setup=(open($compressor); io=open(tempname(), "w")) teardown=(close($compressor); close(io))
SUITE["3D compression"]["read"] = @benchmarkable WaveFD.compressedread!(io, $compressor, 1, $field3) setup=(open($compressor); tfile=tempname(); _io=open(tfile,"w"); WaveFD.compressedwrite(_io, $compressor, 1, $field3); close(_io); io=open(tfile)) teardown=(close($compressor); close(io))

rng3 = (10:n_3D.z-10,10:n_3D.y,10:n_3D.x)
_field3 = view(field3, rng3...)
_compressor = WaveFD.Compressor(Float32, Float32, UInt32, size(_field3), (32,32,32), 1e-2, 1024, true)
SUITE["3D compression, interior"] = BenchmarkGroup()
SUITE["3D compression, interior"]["write"] = @benchmarkable WaveFD.compressedwrite(io, $_compressor, 1, $field3, $rng3) setup=(open($_compressor); io=open(tempname(),"w")) teardown=(close($_compressor); close(io))
SUITE["3D compression, interior"]["read"] = @benchmarkable WaveFD.compressedread!(io, $_compressor, 1, $field3, $rng3) setup=(open($_compressor); tfile=tempname(); _io=open(tfile,"w"); WaveFD.compressedwrite(_io, $_compressor, 1, $field3, rng3); close(_io); io=open(tfile)) teardown=(close($_compressor); close(io))

SUITE["3DAcoIsoDenQ_DEO2_FDTD"] = BenchmarkGroup([Dict("ncells"=>n_3D.z*n_3D.y*n_3D.x,"nbz"=>nb_3D.z,"nby"=>nb_3D.y,"nbx"=>nb_3D.x,"nthreads"=>_nthreads)])
function p3diso(nthreads,nz,ny,nx,nbz,nby,nbx)
    p = WaveFD.Prop3DAcoIsoDenQ_DEO2_FDTD(freesurface=false, nz=nz, ny=ny, nx=nx, nbz=nbz, nby=nby, nbx=nbx, dz=10.0, dy=10.0, dx=10.0, dt=0.001, nthreads=nthreads)
    v,b,pcur,pold = WaveFD.V(p), WaveFD.B(p),WaveFD.PCur(p),WaveFD.POld(p)
    v .= 1500
    b .= 1
    rand!(pcur)
    rand!(pold)
    p
end
for nthreads in _nthreads
    SUITE["3DAcoIsoDenQ_DEO2_FDTD"]["$nthreads threads"] = @benchmarkable WaveFD.propagateforward!(p) setup=(p=p3diso($nthreads,$(n_3D.z),$(n_3D.y),$(n_3D.x),$(nb_3D.z),$(nb_3D.y),$(nb_3D.x))) teardown=(free(p)) seconds=15
end

function fields3diso(p::WaveFD.Prop3DAcoIsoDenQ_DEO2_FDTD)
    nz,ny,nx=size(p)
    δm = Dict("v"=>rand(Float32,nz,ny,nx))
    fields = Dict("pspace"=>rand(Float32,nz,ny,nx))
    δm,fields
end
SUITE["3DAcoIsoDenQ_DEO2_FDTD"]["imaging condition, standard"] = @benchmarkable WaveFD.adjointBornAccumulation!(p,WaveFD.Prop3DAcoIsoDenQ_DEO2_FDTD_Model_V(),ic,δm,fields) setup=(p=p3diso(Sys.CPU_THREADS,$(n_3D.z),$(n_3D.y),$(n_3D.x),$(nb_3D.z),$(nb_3D.y),$(nb_3D.x)); ic=WaveFD.ImagingConditionStandard(); (δm,fields)=fields3diso(p)) teardown=(free(p))
SUITE["3DAcoIsoDenQ_DEO2_FDTD"]["imaging condition, wave field separation FWI"] = @benchmarkable WaveFD.adjointBornAccumulation!(p,WaveFD.Prop3DAcoIsoDenQ_DEO2_FDTD_Model_V(),ic,δm,fields) setup=(p=p3diso(Sys.CPU_THREADS,$(n_3D.z),$(n_3D.y),$(n_3D.x),$(nb_3D.z),$(nb_3D.y),$(nb_3D.x)); ic=WaveFD.ImagingConditionWaveFieldSeparationFWI(); x=zeros(Float32,$(n_3D.z),$(n_3D.y),$(n_3D.x)); (δm,fields)=fields3diso(p)) teardown=(free(p))
SUITE["3DAcoIsoDenQ_DEO2_FDTD"]["imaging condition, wave field separation RTM"] = @benchmarkable WaveFD.adjointBornAccumulation!(p,WaveFD.Prop3DAcoIsoDenQ_DEO2_FDTD_Model_V(),ic,δm,fields) setup=(p=p3diso(Sys.CPU_THREADS,$(n_3D.z),$(n_3D.y),$(n_3D.x),$(nb_3D.z),$(nb_3D.y),$(nb_3D.x)); ic=WaveFD.ImagingConditionWaveFieldSeparationRTM(); x=zeros(Float32,$(n_3D.z),$(n_3D.y),$(n_3D.x)); (δm,fields)=fields3diso(p)) teardown=(free(p))

SUITE["3DAcoVTIDenQ_DEO2_FDTD"] = BenchmarkGroup([Dict("ncells"=>n_3D.z*n_3D.y*n_3D.x,"nbz"=>nb_3D.z,"nby"=>nb_3D.y,"nbx"=>nb_3D.x,"nthreads"=>_nthreads)])
function p3dvti(nthreads,nz,ny,nx,nbz,nby,nbx)
    p = WaveFD.Prop3DAcoVTIDenQ_DEO2_FDTD(freesurface=false, nz=nz, ny=ny, nx=nx, nbz=nbz, nby=nby, nbx=nbx, dz=10.0, dy=10.0, dx=10.0, dt=0.001, nthreads=nthreads)
    v,b,ϵ,η,f,pcur,pold = WaveFD.V(p),WaveFD.B(p),WaveFD.Eps(p),WaveFD.Eta(p),WaveFD.F(p),WaveFD.PCur(p),WaveFD.POld(p)
    rand!(pcur)
    rand!(pold)
    v .= 1500
    b .= 1
    ϵ .= 0.2
    η .= 0.0
    f .= 0.85
    p
end
for nthreads in _nthreads
    SUITE["3DAcoVTIDenQ_DEO2_FDTD"]["$nthreads threads"] = @benchmarkable WaveFD.propagateforward!(p) setup=(p=p3dvti($nthreads,$(n_3D.z),$(n_3D.y),$(n_3D.x),$(nb_3D.z),$(nb_3D.y),$(nb_3D.x))) teardown=(free(p)) seconds=15
end

function fields3dvti(p::WaveFD.Prop3DAcoVTIDenQ_DEO2_FDTD)
    nz,ny,nx=size(p)
    δm = Dict("v"=>rand(Float32,nz,ny,nx))
    fields = Dict("pspace"=>rand(Float32,nz,ny,nx),"mspace"=>rand(Float32,nz,ny,nx))
    δm,fields
end
SUITE["3DAcoVTIDenQ_DEO2_FDTD"]["imaging condition, standard"] = @benchmarkable WaveFD.adjointBornAccumulation!(p,WaveFD.Prop3DAcoVTIDenQ_DEO2_FDTD_Model_V(),ic,δm,fields) setup=(p=p3dvti(Sys.CPU_THREADS,$(n_3D.z),$(n_3D.y),$(n_3D.x),$(nb_3D.z),$(nb_3D.y),$(nb_3D.x)); ic=WaveFD.ImagingConditionStandard(); (δm,fields)=fields3dvti(p)) teardown=(free(p))
SUITE["3DAcoVTIDenQ_DEO2_FDTD"]["imaging condition, wave field separation FWI"] = @benchmarkable WaveFD.adjointBornAccumulation!(p,WaveFD.Prop3DAcoVTIDenQ_DEO2_FDTD_Model_V(),ic,δm,fields) setup=(p=p3dvti(Sys.CPU_THREADS,$(n_3D.z),$(n_3D.y),$(n_3D.x),$(nb_3D.z),$(nb_3D.y),$(nb_3D.x)); ic=WaveFD.ImagingConditionWaveFieldSeparationFWI(); (δm,fields)=fields3dvti(p)) teardown=(free(p))
SUITE["3DAcoVTIDenQ_DEO2_FDTD"]["imaging condition, wave field separation RTM"] = @benchmarkable WaveFD.adjointBornAccumulation!(p,WaveFD.Prop3DAcoVTIDenQ_DEO2_FDTD_Model_V(),ic,δm,fields) setup=(p=p3dvti(Sys.CPU_THREADS,$(n_3D.z),$(n_3D.y),$(n_3D.x),$(nb_3D.z),$(nb_3D.y),$(nb_3D.x)); ic=WaveFD.ImagingConditionWaveFieldSeparationRTM(); (δm,fields)=fields3dvti(p)) teardown=(free(p))

SUITE["3DAcoTTIDenQ_DEO2_FDTD"] = BenchmarkGroup([Dict("ncells"=>n_3D.z*n_3D.y*n_3D.x,"nbz"=>nb_3D.z,"nby"=>nb_3D.y,"nbx"=>nb_3D.x,"nthreads"=>_nthreads)])
function p3dtti(nthreads,nz,ny,nx,nbz,nby,nbx)
    p = WaveFD.Prop3DAcoTTIDenQ_DEO2_FDTD(freesurface=false, nz=nz, ny=ny, nx=nx, nbz=nbz, nby=nby, nbx=nbx, dz=10.0, dx=10.0, dt=0.001, nthreads=nthreads)
    v,b,ϵ,η,f,sinθ,cosθ,sinϕ,cosϕ,pcur,pold = WaveFD.V(p),WaveFD.B(p),WaveFD.Eps(p),WaveFD.Eta(p),WaveFD.F(p),WaveFD.SinTheta(p),WaveFD.CosTheta(p),WaveFD.SinPhi(p),WaveFD.CosPhi(p),WaveFD.PCur(p),WaveFD.POld(p)
    v .= 1500
    b .= 1
    ϵ .= 0.2
    η .= 0.0
    f .= 0.85
    cosθ .= cos(pi/4)
    sinθ .= sin(pi/4)
    cosϕ .= cos(pi/8)
    sinϕ .= sin(pi/8)
    rand!(pcur)
    rand!(pold)
    p
end
for nthreads in _nthreads
    SUITE["3DAcoTTIDenQ_DEO2_FDTD"]["$nthreads threads"] = @benchmarkable WaveFD.propagateforward!(p) setup=(p=p3dtti($nthreads,$(n_3D.z),$(n_3D.y),$(n_3D.x),$(nb_3D.z),$(nb_3D.y),$(nb_3D.x))) teardown=(free(p)) seconds=15
end

function fields3dtti(p::WaveFD.Prop3DAcoTTIDenQ_DEO2_FDTD)
    nz,ny,nx = size(p)
    δm = Dict("v"=>rand(Float32,nz,ny,nx))
    fields = Dict("pspace"=>rand(Float32,nz,ny,nx),"mspace"=>rand(Float32,nz,ny,nx))
    δm,fields
end
SUITE["3DAcoTTIDenQ_DEO2_FDTD"]["imaging condition, standard"] = @benchmarkable WaveFD.adjointBornAccumulation!(p,WaveFD.Prop3DAcoTTIDenQ_DEO2_FDTD_Model_V(),ic,δm,fields) setup=(p=p3dtti(Sys.CPU_THREADS,$(n_3D.z),$(n_3D.y),$(n_3D.x),$(nb_3D.z),$(nb_3D.y),$(nb_3D.x)); ic=WaveFD.ImagingConditionStandard(); (δm,fields)=fields3dtti(p)) teardown=(free(p))
SUITE["3DAcoTTIDenQ_DEO2_FDTD"]["imaging condition, wave field separation FWI"] = @benchmarkable WaveFD.adjointBornAccumulation!(p,WaveFD.Prop3DAcoTTIDenQ_DEO2_FDTD_Model_V(),ic,δm,fields) setup=(p=p3dtti(Sys.CPU_THREADS,$(n_3D.z),$(n_3D.y),$(n_3D.x),$(nb_3D.z),$(nb_3D.y),$(nb_3D.x)); ic=WaveFD.ImagingConditionWaveFieldSeparationFWI(); (δm,fields)=fields3dtti(p)) teardown=(free(p))
SUITE["3DAcoTTIDenQ_DEO2_FDTD"]["imaging condition, wave field separation RTM"] = @benchmarkable WaveFD.adjointBornAccumulation!(p,WaveFD.Prop3DAcoTTIDenQ_DEO2_FDTD_Model_V(),ic,δm,fields) setup=(p=p3dtti(Sys.CPU_THREADS,$(n_3D.z),$(n_3D.y),$(n_3D.x),$(nb_3D.z),$(nb_3D.y),$(nb_3D.x)); ic=WaveFD.ImagingConditionWaveFieldSeparationRTM(); (δm,fields)=fields3dtti(p)) teardown=(free(p))

include(joinpath(pkgdir(WaveFD), "benchmark", "mcells_per_second.jl"))

SUITE

max_stars_repo_name: ChevronETC/WaveFD

max_issues_repo_name: ChevronETC/WaveFD

max_forks_repo_name: ChevronETC/WaveFD


text:
const TwoCorrelatedOU = CorrelatedOU+CorrelatedOU
const TwoOU = OU+OU
@testset "expectedperiodogram" begin
    store = allocate_memory_EI_FGH(CorrelatedOU,100,1)
    storeadd = allocate_memory_EI_FGH(TwoCorrelatedOU,100,1)
    storeuni = allocate_memory_EI_FGH(OU,100,1)
    storeadduni = allocate_memory_EI_FGH(TwoOU,100,1)
    badkernel = ones(50) # kernel should be length 200
    @testset "Function" begin
        @test_throws DimensionMismatch WhittleLikelihoodInference.mult_allocated_kernel!(store.funcmemory.allocatedarray,badkernel)
        @test_throws DimensionMismatch WhittleLikelihoodInference.mult_allocated_kernel!(storeuni.funcmemory.allocatedarray,badkernel)
        @test length(EI(OU(1.0,1.0),10,1/2)) == 10
        @test EI!(store,CorrelatedOU(1.0,1.0,0.5)) == nothing
        @test EI!(storeuni,OU(1.0,1.0)) == nothing
        @test EI!(storeadd,TwoCorrelatedOU([1.0,1.0,0.5,1.0,1.0,0.5])) == nothing
        @test EI!(storeadduni,TwoOU(ones(4))) == nothing
    end
    @testset "Gradient" begin
        @test_throws DimensionMismatch WhittleLikelihoodInference.mult_allocated_kernel!(store.gradmemory.allocatedarray,badkernel)
        @test_throws DimensionMismatch WhittleLikelihoodInference.mult_allocated_kernel!(storeuni.gradmemory.allocatedarray,badkernel)
        @test grad_EI!(store,CorrelatedOU(1.0,1.0,0.5)) == nothing
        @test grad_EI!(storeuni,OU(1.0,1.0)) == nothing
        @test grad_EI!(storeadd,TwoCorrelatedOU([1.0,1.0,0.5,1.0,1.0,0.5])) == nothing
        @test grad_EI!(storeadduni,TwoOU(ones(4))) == nothing
    end
    @testset "Hessian" begin
        @test_throws DimensionMismatch WhittleLikelihoodInference.mult_allocated_kernel!(store.hessmemory.allocatedarray,badkernel)
        @test_throws DimensionMismatch WhittleLikelihoodInference.mult_allocated_kernel!(storeuni.hessmemory.allocatedarray,badkernel)
        @test hess_EI!(store,CorrelatedOU(1.0,1.0,0.5)) == nothing
        @test hess_EI!(storeuni,OU(1.0,1.0)) == nothing
        @test hess_EI!(storeadd,TwoCorrelatedOU([1.0,1.0,0.5,1.0,1.0,0.5])) == nothing
        @test hess_EI!(storeadduni,TwoOU(ones(4))) == nothing
    end
end


text:
#=

  The BUGS Book
  Example 3.5.1 Heart transplants: learning from data

  ~/jags/bugs_book_3_5_1.R
  """
  model {
    yT ~ dbin(pT, nT)
    pT ~ dunif(0, 1)
    for (i in 1:8) {
      sP[i] ~ dexp(theta)
    }
    theta ~ dgamma(0.001, 0.001)
    surv.t <- pT/theta # expected survival with transplant
    Is <- surv.t - 2
  }

  Data: list(yT=8, nT=10, sP=c(2,3,4,4,6,7,10,12))

  Output:
         Mean      SD  Naive SE Time-series SE
Is     3.1442 2.28250 0.0080699      0.0085034
pT     0.7500 0.12014 0.0004248      0.0005856
surv.t 5.1442 2.28250 0.0080699      0.0085034
theta  0.1667 0.05899 0.0002086      0.0002086
  """

  Note that most value are fixed/derived.

  Summary Statistics
  parameters      mean       std   naive_se      mcse       ess      rhat   ess_per_sec
      Symbol   Float64   Float64    Float64   Float64   Float64   Float64       Float64

       theta    0.1473    0.0000     0.0000    0.0000   20.5530    0.9999        6.1813
          pT    0.8623    0.0000     0.0000    0.0000   20.5530    0.9999        6.1813
          yT    8.0000    0.0000     0.0000    0.0000       NaN       NaN           NaN
      surv_t    5.8559    0.0000     0.0000    0.0000       NaN       NaN           NaN
          Is    3.8559    0.0000     0.0000    0.0000   20.5530    0.9999        6.1813

  Cf ~/webppl/bugs_book_3_5_1.wppl

=#

using Turing, StatsPlots, DataFrames
include("jl_utils.jl")

@model function bugs_book_3_5_1(sPs)
    n = length(sPs)
    nT = 10

    theta ~ Gamma(0.001,1/0.001)
    pT ~ Uniform(0,1)
    yT ~ Binomial(nT,pT)

    surv_t ~ Dirac(pT/theta) # expected survival with transplant
    Is ~ Dirac(surv_t - 2.0)

    true ~ Dirac(yT == 8)
    sPs ~ filldist(Exponential(theta),n)

end

sPs = [2,3,4,4,6,7,10,12]
model = bugs_book_3_5_1(sPs)

num_chains = 4

# chns = sample(model, Prior(), 10_000)
# chns = sample(model, MH(), 100_000)
# chns = sample(model, PG(15), 10_000)
# chns = sample(model, SMC(10_000), 10_000)
# chns = sample(model, IS(), 10_000)

# chns = sample(model, NUTS(), 10_000)
chns = sample(model, HMC(0.1,10), 10_000)

display(chns)
# display(plot(chns))

max_stars_repo_name: tias/hakank

max_issues_repo_name: tias/hakank

max_forks_repo_name: tias/hakank


text:
# Script that allows the user to sample from a beta model with fitnesses time varying,
# according to one of the dgps described in the paper, and then filter with the
# SD beta model

#add packages to LOAD path
push!(LOAD_PATH,"./DynNets/src/")

using  Utilities, DynNets, StaticNets, StatsBase

#%% Sample DGP

N = 10 # number of nodes
NTV = round.(Int,N/2) # Number of nodes with time varying fitnesses
T = 250 #
dynType = "SIN"
minDeg, maxDeg = degb = [3,8]
unifDeg = round(Int,(maxDeg-  minDeg)/2)

#targetStatic = true=> use a single static estimate for targeting in SD estimate
#targetStatic = false=> use a the mean of single snapshots  estimates for targeting in SD estimate
targetStatic = true
# Define time varying fitneesse
dynFitDgp, indsTVnodes =  StaticNets.dgpDynamic(StaticNets.fooErgmDirBin1,dynType,N,T;
                                  NTV = NTV, degIOUncMeans =unifDeg*ones(2N),
                                 degb = degb  )

# Sample time series of networks from the dynamical fitnesses
dynDegsSam = zeros(size(dynFitDgp))
for t=1:T
    matSamp = StaticNets.sampl(StaticNets.ErgmDirBin1(dynFitDgp[:,t]),1;  parGroupsIO = dynFitDgp[:,t])
    dynDegsSam[:,t] = [sumSq(matSamp,2);sumSq(matSamp,1)]
end
estFitSS_T =  StaticNets.estimate( StaticNets.SnapSeqNetDirBin1(dynDegsSam); identPost = true,identIter= false,targetErr = 0.001 )
for t=1:T
    dynFitDgp[:,t] =  StaticNets.identify(StaticNets.ErgmDirBin1(dynDegsSam[:,t]),dynFitDgp[:,t];idType = "firstZero" )
    estFitSS_T[:,t] =  StaticNets.identify(StaticNets.ErgmDirBin1(dynDegsSam[:,t]),estFitSS_T[:,t];idType = "firstZero" )
end

degsIO_T = dynDegsSam
indsGroups = [Int.(1:2N),Int.(1:2N)] #[Int.(1:2N),ones(2N)] #
modGasDirBin1_3Npars = DynNets.SdErgmDirBin1(degsIO_T, [zeros(2N), zeros(2N) ,zeros(2N)],
                                            indsGroups, "FISHER-DIAG")
rmseSS =sqrt.(meanSq((dynFitDgp - estFitSS_T).^2,2))
if targetStatic
    estFitStatic = DynNets.estimateSnapSeq(modGasDirBin1_3Npars, degsIO_T = mean(degsIO_T,dims=2) )
    estFitSS_T = repeat(estFitStatic, 1, 5)
end

#%% Estimate SD beta model ---Can take a while (~30 min for N=10)
@time parGas_3Npars, ~  =  DynNets.estimateTarg(modGasDirBin1_3Npars;SSest =estFitSS_T)
gasFiltFit,~ = DynNets.score_driven_filter_or_dgp(modGasDirBin1_3Npars,[parGas_3Npars[1];parGas_3Npars[2];parGas_3Npars[3]])
rmseGas =sqrt.(meanSq((dynFitDgp - gasFiltFit).^2,2))

# plot the resulting filtered path
indsPlot = [ 5]
for Nplot =indsPlot
    plot(1:T,dynFitDgp[Nplot,:])
    plot(1:T,filFitGas[Nplot,:,:],".r")
    plot(1:T,estFitSS[Nplot,:,:],".b")
end

#%%Add Confidence bands accounting for parameter uncertainty
using ForwardDiff

# TO BE ADJUSTED TO THE CASE WITHOUT PARAMETERS RESTRICTIONS1!!
# NOW IT HANDLES ONLY THE RESTRICTED CASE, AND MIGHT CONTAIN BUGS

Ttrain = T# round(Int, T/2) #70 106 #
#Estimate gas model on  windows of length Ttrain
modGasDirBin1_N_p_2 = SdErgmDirBin1(degsIO_T[:,1:Ttrain],"FISHER-DIAG", "ONE_PAR_ALL")
gasParEstOnTrain,~ = estimateTarg(modGasDirBin1_N_p_2;SSest = estFitSS_T[:,1:Ttrain] )

#f(x) = DynNets.score_driven_filter_or_dgp(modGasDirBin1_N_p_2, x;groupsInds = modGasDirBin1_N_p_2.groupsInds)[2]
GBA = 1
allpar0 = array_2_vec_all_par(modGasDirBin1_N_p_2, gasParEstOnTrain)
BA_re_hat = allpar0[end-1:end]
function UnrestrAB(vecRePar::Array{<:Real,1}) #restrict a vector of only A and B
      diag_B_Re = vecRePar[1:GBA]
      diag_A_Re = vecRePar[GBA+1:end]
      diag_B_Un = log.(diag_B_Re ./ (1 .- diag_B_Re ))
      diag_A_Un = log.(diag_A_Re)
      vecUnPar =  [diag_B_Un; diag_A_Un]
      return vecUnPar
end
function loglike(Model::DynNets.SdErgmDirBin1, degs_t::Array{<:Real,1},
                        f_t::Array{<:Real,1})

    thetas_mat_t_exp, exp_mat_t = StaticNets.expMatrix2(StaticNets.fooErgmBin1,f_t)
    exp_deg_t = sum(exp_mat_t,dims = 2)

    loglike_t = sum(f_t.*degs_t) -  sum(UpperTriangular(log.(1 .+ thetas_mat_t_exp))) #  sum(log.(1 + thetas_mat_t_exp))
    return loglike_t
 end
function logLike_t(Model, obsT, vReGasPar)
 groupsInds = Model.groupsInds
      N2,T = size(obsT);N = round(Int,N2/2)
      NGW,GBA,ABgroupsIndNodesIO,indTvNodesIO = NumberOfGroupsAndABindNodes(Model, groupsInds)
      # Organize parameters of the GAS update equation
      WGroupsIO = vReGasPar[1:NGW]
      #    StaticNets.expMatrix2(StaticNets.fooErgmDirBin1,WGroupsIO )
      W_allIO = WGroupsIO[groupsInds[1]]
      BgasGroups  = vReGasPar[NGW+1:NGW+GBA]
      AgasGroups  = vReGasPar[NGW+GBA+1:NGW+2GBA]
      #distribute nodes among groups with equal parameters
      AgasIO = AgasGroups[ABgroupsIndNodesIO[indTvNodesIO]]
      BgasIO = BgasGroups[ABgroupsIndNodesIO[indTvNodesIO]]
      WgasIO   = W_allIO[indTvNodesIO]
      UMallNodesIO = W_allIO
      UMallNodesIO  =  WgasIO ./ (1 .- BgasIO)
      ftotIO_t =  identify(Model,UMallNodesIO)
      I_tm1 = UniformScaling(N)
      loglike_t = zero(Real)
      for t=1:T-1
            degsIO_t = obsT[:,t] # vector of in and out degrees
            ftotIO_tp1,loglike_t = predict_score_driven_par(Model,N,degsIO_t,ftotIO_t,I_tm1,
                                                  indTvNodesIO,WgasIO,BgasIO,AgasIO)
            ftotIO_t = ftotIO_tp1
      end
      return  loglike_t::T where T <:Real
end
function logLike_T(Model, obsT, vReGasPar)
      groupsInds = Model.groupsInds
      N2,T = size(obsT);N = round(Int,N2/2)
      NGW,GBA,ABgroupsIndNodesIO,indTvNodesIO = NumberOfGroupsAndABindNodes(Model, groupsInds)
      # Organize parameters of the GAS update equation
      WGroupsIO = vReGasPar[1:NGW]
      #    StaticNets.expMatrix2(StaticNets.fooErgmDirBin1,WGroupsIO )
      W_allIO = WGroupsIO[groupsInds[1]]
      BgasGroups  = vReGasPar[NGW+1:NGW+GBA]
      AgasGroups  = vReGasPar[NGW+GBA+1:NGW+2GBA]
      #distribute nodes among groups with equal parameters
      AgasIO = AgasGroups[ABgroupsIndNodesIO[indTvNodesIO]]
      BgasIO = BgasGroups[ABgroupsIndNodesIO[indTvNodesIO]]
      WgasIO   = W_allIO[indTvNodesIO]
      # start values equal the unconditional mean, but  constant ones remain equal to the unconditional mean, hence initialize as:
      UMallNodesIO = W_allIO
      UMallNodesIO  =  WgasIO ./ (1 .- BgasIO)
      ftotIO_t =  identify(Model,UMallNodesIO)
      I_tm1 = UniformScaling(N)
      loglike_T = zero(Real)
      for t=1:T-1
            degsIO_t = obsT[:,t] # vector of in and out degrees
            ftotIO_tp1,loglike_t = predict_score_driven_par(Model,N,degsIO_t,ftotIO_t,I_tm1,
                                                  indTvNodesIO,WgasIO,BgasIO,AgasIO)
            ftotIO_t = ftotIO_tp1
            loglike_T += loglike_t
      end
      return  loglike_T::T where T <:Real
end

obsT = modGasDirBin1_N_p_2.obsT
T_train = size(obsT)[2]
f_T(x) = logLike_T(modGasDirBin1_N_p_2::SdErgmDirBin1, obsT, x)
# @time ForwardDiff.hessian(f_t, allpar0)
#@time FiniteDiff.finite_difference_hessian(f_T, allpar0)

covA0 =  ForwardDiff.hessian(f_T, allpar0)./T_train
# estimate outer product of scores
vec_of_f_t =[x -> logLike_t(modGasDirBin1_N_p_2::SdErgmDirBin1, obsT[:, 1:t], x) for t in 1:Ttrain]
# run once to precompile the functions
[f(allpar0) for f in vec_of_f_t]
vec_g_t = [ForwardDiff.gradient(f, allpar0) for f in vec_of_f_t[2:end]]
global covB0 = zeros(size(covA0))
 for g_t in vec_g_t
      global covB0 += g_t * g_t'
 end
 global covB0 = covB0./ T_train
#%%
diagCorrect = 1e-3
covMatHat = pinv(covA0) * covB0 * pinv(covA0)
      covMatHat = Symmetric(covMatHat)

covMatHat_corr = (covMatHat .+ Diagonal(diagCorrect.*ones(size(covMatHat)[1])))

parDistr = MvNormal(allpar0, diag(covMatHat_corr))

samp = rand(parDistr, 2000)
indsToKeep = (0 .< samp[end-1, :] .< 1) .& (samp[end, :] .>0)
sum(indsToKeep)
#hist(samp[end,:])
N_samp_par = 1000
allpar_sample = samp[:, indsToKeep][:, 1:N_samp_par]

#%%
ftotIO_0 = meanSq(estimateSnapSeq(modGasDirBin1_N_p_2; degsIO_T = obsT[:, 1:5]), 2)
ftot_T_est, _ = score_driven_filter_or_dgp(modGasDirBin1_N_p_2, allpar0,ftotIO_0 = ftotIO_0)
ftot_T, _ = score_driven_filter_or_dgp(modGasDirBin1_N_p_2, allpar_sample[:,3],  ftotIO_0 = ftotIO_0)
plot(ftot_T_est')
#plot(ftot_T')

Nterms = 2*N

global gasFiltPar_conf = zeros(N_samp_par, Nterms, T_train)
      for i= 1:N_samp_par
       # @show(i)
       global gasFiltPar_conf[i,:,:], _ = score_driven_filter_or_dgp(modGasDirBin1_N_p_2, allpar_sample[:,i],  ftotIO_0 = ftotIO_0)
      end

indsToRem = [any(isnan.(gasFiltPar_conf[i, : , :])) for i in 1:N_samp_par ]
sum(indsToRem)
gasFiltPar_conf = gasFiltPar_conf[ .!indsToRem, : , : ]

using Statistics
# Compute confidence bands
quant_vals = [0.95, 0.05]
 conf_band = zeros(length(quant_vals),T_train, Nterms)
 for t=1:T_train
     for k=1:Nterms
      filt_t = gasFiltPar_conf[:,k,t]
      conf_band[:,t,k] = Statistics.quantile(filt_t[.!isnan.(filt_t)],quant_vals)
    end
 end

figure()
      legTex = ["DGP";"SDERGM"; "95%"]
      node = 11
      plot(1:T_train,ftot_T_est[node, :],"r")
      plt.fill_between(1:T_train, conf_band[1, :, node], y2 =conf_band[2,:,node],color =(0.9, 0.2 , 0.2, 0.1)  )#, color='b', alpha=.1)

max_stars_repo_name: domenicodigangi/SDERGM

max_issues_repo_name: domenicodigangi/SDERGM

max_forks_repo_name: domenicodigangi/SDERGM


text:
# Rational Quadratic Isotropic Covariance Function

"""
    RQIso <: Isotropic{SqEuclidean}

Isotropic Rational Quadratic kernel (covariance)
```math
k(x,x') = σ²(1 + (x - x')ᵀ(x - x')/(2αℓ²))^{-α}
```
with length scale ``ℓ``, signal standard deviation ``σ``, and shape parameter ``α``.
"""
mutable struct RQIso{T<:Real} <: Isotropic{SqEuclidean}
    "Squared length scale"
    ℓ2::T
    "Signal variance"
    σ2::T
    "Shape parameter"
    α::T
    "Priors for kernel parameters"
    priors::Array
end

"""
Rational Quadratic Isotropic Covariance Function

    RQIso(ll:T, lσ::T, lα::T)

# Arguments
  - `ll::Real`: length scale (given on log scale)
  - `lσ::Real`: signal standard deviation (given on log scale)
  - `lα::Real`: shape parameter (given on log scale)
"""
RQIso(ll::T, lσ::T, lα::T) where T = RQIso{T}(exp(2 * ll), exp(2 * lσ), exp(lα), [])

function set_params!(rq::RQIso, hyp::AbstractVector)
    length(hyp) == 3 || throw(ArgumentError("Rational Quadratic function has three parameters"))
    rq.ℓ2, rq.σ2, rq.α = exp(2 * hyp[1]), exp(2 * hyp[2]), exp(hyp[3])
end

get_params(rq::RQIso{T}) where T = T[log(rq.ℓ2) / 2, log(rq.σ2) / 2, log(rq.α)]
get_param_names(rq::RQIso) = [:ll, :lσ, :lα]
num_params(rq::RQIso) = 3

cov(rq::RQIso, r::Number) = rq.σ2 * (1 + r / (2 * rq.α * rq.ℓ2))^(-rq.α)

@inline dk_dll(rq::RQIso, r::Real) =
    (s = r / rq.ℓ2; rq.σ2 * s * (1 + s / (2 * rq.α))^(-rq.α - 1)) # dK_d(log ℓ)dK_dℓ
@inline function dk_dlα(rq::RQIso, r::Real)
    s = r / rq.ℓ2
    part = 1 + s / (2 * rq.α)
    rq.σ2 * part^(-rq.α) * (s / (2 * part) - rq.α * log(part))  # dK_d(log α)
end
@inline function dk_dθp(rq::RQIso, r::Real, p::Int)
    if p==1
        return dk_dll(rq, r)
    elseif p==2
        return dk_dlσ(rq, r)
    elseif p==3
        return dk_dlα(rq, r)
    else
        return NaN
    end
end


text:
module GenTabularTMazeDpi

using Curiosity
using Curiosity: DPI, HordeDPI
using Curiosity: OpenWorld
using MinimalRLCore
const TTMU = Curiosity.TabularTMazeUtils
include("../experiments/tabular_tmaze.jl")
using Random
using StatsBase
using GVFHordes
using Statistics
using JLD2

using ProgressMeter

StatsBase.sample(p::GVFHordes.GVFParamFuncs.FunctionalPolicy, s, actions) =
    sample(Weights([p.func(;state_t = s, action_t = a) for a in actions]))

StatsBase.sample(rng, p::GVFHordes.GVFParamFuncs.AbstractPolicy, s, actions) =
    sample(Weights([get(p;state_t = s, action_t = a) for a in actions]))

function gen_dpi(num_states=100000, seed=1; exploring_starts=false, h_kwargs...)
    Random.seed!(seed)

    parsed = TabularTMazeExperiment.default_args()
    parsed["exploring_starts"] = exploring_starts
    parsed["demon_learner"] = "Q"

    dpis = DPI[]

    state_filter_func =  Curiosity.FeatureSubset(identity, 1:1)

    for gvf_i ∈ 1:4
        cumulant_schedule = TTMU.get_cumulant_schedule(parsed)
        exploring_starts = parsed["exploring_starts"]
        env = TabularTMaze(exploring_starts, cumulant_schedule)

        ss = Vector{Float64}[]
        as = Int[]
        policy = TTMU.GoalPolicy(gvf_i)

        while length(ss) < num_states
            s = MinimalRLCore.start!(env)
            push!(ss, state_filter_func(s))
            a = policy(s)

            for t ∈ 1:rand(1:15)
                s, r, term = MinimalRLCore.step!(env, a)
                a = policy(s)
                push!(ss, state_filter_func(s))
                if term
                    break
                end
            end
        end

         push!(dpis, DPI(ss, policy; h_kwargs...))
    end

    return HordeDPI(dpis, state_filter_func)
end

function save_dpi(num_states, seed; exploring_starts=false, h_kwargs...)
    hdpi = gen_dpi(num_states, seed; exploring_starts=exploring_starts, h_kwargs...)
    @save "./src/data/dpi/ttmaze_exploring_starts_$(exploring_starts).jld2" hdpi
end

end

max_stars_repo_name: MatthewMcLeod/curiosity

max_issues_repo_name: MatthewMcLeod/curiosity

max_forks_repo_name: MatthewMcLeod/curiosity


text:
# trying to see range bearing AprilTag detections

using Images, ImageDraw, ImageView, FileIO
using YAML, JLD, HDF5
using Colors

using AprilTags

include(joinpath(dirname(@__FILE__),"racecarUtils.jl"))

cfg = loadConfig()
# some default, probably not right
# cw, ch = 330.4173, 196.32587
# focal = 340.97913
# tagsize = 0.172

cw, ch = cfg[:intrinsics][:cx], cfg[:intrinsics][:cy]
focal = cfg[:intrinsics][:fx]
tagsize = 0.172

# datafolder = ENV["HOME"]*"/data/racecar/smallloop/"
datafolder = ENV["HOME"]*"/data/racecar/thursday/"
@load datafolder*"tag_det_per_pose.jld" tag_bag
# tag_bag[0]

imgfolder = "newCalibrationTest3"
# camcount = readdlm(datafolder*"$(imgfolder)/cam-count.csv",',')
camlookup = Dict{Int, String}()
for i in 1:length(tag_bag)
  camlookup[i-1] = "camera_image$(5*(i-1)).jpeg"
  # camlookup[camcount[i,1]] = strip(camcount[i,2])
end

IMGS = []
psid = 1
for psid in 0:(length(tag_bag)-1)
  img = load(datafolder*"$(imgfolder)/$(camlookup[psid])")
  push!(IMGS, deepcopy(img))
  for (tid, vals) in tag_bag[psid]
    drawTagLine!(IMGS[psid+1], vals)
  end
end

IMGS[1]
IMGS[20]
IMGS[30]
IMGS[40]
IMGS[50]
IMGS[59]
IMGS[60]
IMGS[61]
IMGS[62]
IMGS[63]
IMGS[64]
IMGS[70]
IMGS[71]
IMGS[72]
IMGS[73]
IMGS[74]
IMGS[75]
IMGS[76]
IMGS[77]
IMGS[78]
IMGS[79]
IMGS[87]
IMGS[88]
IMGS[90]
IMGS[99]
IMGS[100]
IMGS[101]
IMGS[102]
IMGS[109]
IMGS[110]
IMGS[119]
IMGS[120]
IMGS[121]
IMGS[122]
IMGS[123]
IMGS[124]
IMGS[125]
IMGS[126]
IMGS[127]
IMGS[128]
IMGS[129]
IMGS[130]
IMGS[131]
IMGS[132]
IMGS[139]
IMGS[140]
IMGS[149]
IMGS[150]
IMGS[159]
IMGS[160]
IMGS[163]

tag_bag[60][18][:bearing]
tag_bag[61][18][:bearing]
tag_bag[62][18][:bearing]
tag_bag[63][18][:bearing]

tag_bag[77]

tag_bag[86]
tag_bag[87]
tag_bag[88]

println("$(tag_bag[98][0][:bearing]), $(tag_bag[98][0][:range])")
println("$(tag_bag[99][0][:bearing]), $(tag_bag[99][0][:range])")
println("$(tag_bag[100][0][:bearing]), $(tag_bag[100][0][:range])")
println("$(tag_bag[101][0][:bearing]), $(tag_bag[101][0][:range])")
println("$(tag_bag[102][0][:bearing]), $(tag_bag[102][0][:range])")

println("$(tag_bag[98][1][:bearing]), $(tag_bag[98][1][:range])")
println("$(tag_bag[99][1][:bearing]), $(tag_bag[99][1][:range])")
println("$(tag_bag[100][1][:bearing]), $(tag_bag[100][1][:range])")
println("$(tag_bag[101][1][:bearing]), $(tag_bag[101][1][:range])")
println("$(tag_bag[102][1][:bearing]), $(tag_bag[102][1][:range])")

tag_bag[109]
tag_bag[110]

tag_bag[123]

IMGS[125]
tag_bag[124]

tag_bag[125]

# pose :x1
psid = 1

pose_tags = tag_bag[psid]
img = load(datafolder*"images/$(camlookup[psid])")
push!(IMGS, deepcopy(img))

drawTagLine!(IMGS[2], pose_tags[0])
drawTagLine!(IMGS[2], pose_tags[1])
drawTagLine!(IMGS[2], pose_tags[16])
drawTagLine!(IMGS[2], pose_tags[584])

IMGS[2]

# pose :x10
psid = 10

pose_tags = tag_bag[psid]
img = load(datafolder*"images/$(camlookup[psid])")
push!(IMGS, deepcopy(img))

drawTagLine!(IMGS[end], pose_tags[0])
drawTagLine!(IMGS[end], pose_tags[1])
drawTagLine!(IMGS[end], pose_tags[16])

IMGS[end]

# pose :x25
psid = 25

pose_tags = tag_bag[psid]
img = load(datafolder*"images/$(camlookup[psid])")
push!(IMGS, deepcopy(img))

drawTagLine!(IMGS[end], pose_tags[8])
drawTagLine!(IMGS[end], pose_tags[9])
drawTagLine!(IMGS[end], pose_tags[10])
drawTagLine!(IMGS[end], pose_tags[16])

IMGS[end]

# pose :x25
psid = 26

pose_tags = tag_bag[psid]
img = load(datafolder*"images/$(camlookup[psid])")
push!(IMGS, deepcopy(img))

drawTagLine!(IMGS[end], pose_tags[9])
drawTagLine!(IMGS[end], pose_tags[10])
drawTagLine!(IMGS[end], pose_tags[16])

IMGS[end]

using TransformUtils

# redetect tags

detector = AprilTagDetector()

tags = detector(img)

pose = homography_to_pose(tags[5].H, -focal, focal, cw, ch)

cV = pose[1:3,4]

bV = cfg[:bRc]*cV

bearing = atan2(-bV[1], bV[3])
# tag_det[id][:range] = [norm(tag_det[id][:pos][[1;3]]);]

foreach(tag->drawTagBox!(img,tag, width = 5, drawReticle = false), tags)

img

#


text:
# genetics.jl

# This script defines the "alogirthm" used to solve the problem in the form of genes
# Genes are defined in a tree structure, referencing their parents.

###### Example Genetics ######

# Create All Genes
# Define all Proteins Generated by Genes
# Proteins are all functions
# Note: protein lifetime length is stored in the gene for that protein

println("Genetics Importing Begun")

include("output.jl")

function root_protein(molecule)

	output = 0
	backreg = 0
	println(molecule)
	println("Replicating Molecule")
	return ([molecule], backreg, output)
end

root = Gene(nothing, root_protein, 1.0,  0.0)

function dmg_protein(molecule)

	output = 1
	backreg = 0
	println(molecule)
	println("Destroying Molecule")
	return ([], backreg, output)
end

destroyMoleculeGene = Gene(root, dmg_protein, 0.5, 5.0)

function duplicate_protein(molecule)

	output = 0
	backreg = 0
	println(molecule)
	println("Duplicate Molecule")
	return ([molecule, molecule], backreg, output)
end

dupMolGene = Gene(root, duplicate_protein, 0.5, 5.0)

# Finally set list of genes/proteins and their initial abundances
pa1 = ProteinAbundance(root, 10)
pa2 = ProteinAbundance(destroyMoleculeGene, 20)
pa3 = ProteinAbundance(dupMolGene, 70)
proteinDistribution = [pa1, pa2, pa3]

println("Genetics Importing Finished")

max_stars_repo_name: voutasaurus/metalbolicalgorithm

max_issues_repo_name: voutasaurus/metalbolicAlgorithm

max_forks_repo_name: voutasaurus/metalbolicAlgorithm


text:
# Sturm_Drang.jl
# Sturm and Urge - Calculation of electronic densities of states for conjugated polymer chains, with off-diagonal disorder parameterised by statistical mechanics.
# https://github.com/jarvist/Teclo

println("Sturm und Drang: DoS by Sturm sequences")

push!(LOAD_PATH,"./") # Temporary version of modules in PWD
using Sturm
using ApproxFunFromFile

#Bolztmann's constant in units of eV - thereby all the potentials (of functional form or tabulated data) in units eV
kB=8.6173324E-5

N=10^6 #6 # Length of tridiagonal Hamiltonian constructed; larger value -> better statistics, slower runtime

disorder=0.0 # Energetic disorder eV, Gaussian form, for the site energies of the polymer

function BoltzmannDoS(U, modelJ, PREFIX="")
    println("\nBoltzmannDoS with: $(PREFIX) ")

    # Outputfiles, streamed to with C-like @printf
    DoSfile=open("$(PREFIX)_DoS.dat","w+")
    @printf(DoSfile,"#T sigma Cumulative-DoS DoS-this-bin\n#Trying plotting like: gnuplot> p \"DoS.dat\" u 2:4 \n")
    onsetfile=open("$(PREFIX)_onset.dat","w+")
    potfile=open("$(PREFIX)_potential.dat","w+")
    @printf(potfile,"# Try plotting like: gnuplot> p \"potential.dat\" u 1:2, \"\" u 1:(0.1*\$3) w l \n")

    # The @sync @parallel magic makes this forloop execute in parallel, at least on Linux.
    @sync @parallel for T=300.0:100:300.0 #T=100.0:100:400 #:0.1:1
        B=1/(T*kB) #300K * k_B in eV

        Z,epsilon=quadgk(theta -> exp(-U(theta)*B),0.0,360.0) # Now using Julia language (>0.4) built in quadgk numeric integration
        Z=Z/30 # to improve statistics on rejection sampling of thetas NB: A TOTAL HACK
        # Calculation partition function Z; particular to this potential energy surface + temperature
        println("Integration of Z via quadgk method, estimated upper bound on absolute error: ",epsilon)
        println("Partition function for Z(E0=",E0,",T=",T,")=",Z)

    # Following checks the partition function code, outputting p(robability) as a fn(theta) for varying P
        @printf(potfile,"# Potential and probability density at T=%f",T)
        @printf(potfile,"# theta U(theta) probability(theta)")
        for theta = 0.0:1.0:360.0
#       println("Partition function Z=",Z)
            p=exp(-U(theta)*B)/Z
#        println(t," ",p)
            @printf(potfile,"%f %f %f\n",theta,U(theta),p)
        end

        # generates separate (D)iagonal and (E)-offdiagonal terms of Tight Binding Hamiltonian
        D,E=randH(5.0,disorder, modelJ, B,Z,U,N)

#        showcompact(diagm(E.^0.5,-1)+diagm(D)+diagm(E.^0.5,1)) # Show dense form of matrix. First expands to a dense form, so only use for N<~10^4 !

        @printf("Calculated with E0= %f Z= %f\n",E0,Z)
#       println("STURM sequence method...")
        pDoS=0
        pold=0
        onset=false
        for sigma=3.2:0.01:6.8 # Energy window (eV) over which to bin eigenvalues
            # Sturm sequence returns number of eigenvalues of Hamiltonian {D,E} below sigma
            pDoS=sturm(D,E,sigma)
            @printf(".") # Progress indicator - marching dots ('......')
            @printf(DoSfile,"%f %f %f %f\n",T,sigma,pDoS, pDoS-pold)
            if (pDoS-pold>10.0 && onset==false)
                @printf(onsetfile,"%f %f %f\n",T,sigma,pDoS-pold)
                onset=true
            end
            pold=pDoS
        end
        @printf("\n") # New line after all DoS calculated
    end
    close(DoSfile)
    close(onsetfile)
    close(potfile)
end

# Model setup
J0=0.8
modelJ(theta) = J0*cos(theta*pi/180.0).^2

# Size of potential from Raos P3HT ForceField paper (Moreno et al. J.Phys.Chem.B 2010)
# 'full' potential energy Fig 4.a. puts a barrier at 90 degress of ~3.0 kCal / mol = 126 meV
E0=0.126

BoltzmannDoS(x->0.0 , modelJ, "Flat") # Flat pot energy distro; normal modelJ=J0*cos^2 form
BoltzmannDoS(x->0.0, theta->J0+0.0*(theta), "FlatWhite") # Flat pot energy distro, constant modelJ=J0 form

BoltzmannDoS(x->0.0, theta->rand(length(theta))+0.0*(theta), "FlatWhiteNoise") # Wigner semi-circle ?
BoltzmannDoS(x->0.0, theta->randn(length(theta))+0.0*(theta), "FlatWhiteNoiseNormal") #

U(theta)=( E0 * sin(theta*pi/180.0)^2 ) #P3HT like
BoltzmannDoS(U,modelJ,"P3HT")
U2(theta)=( E0 * (-sin(theta*pi/180.0)^2 - sin(2*theta*pi/180.0)^2 ) ) # PFO like
BoltzmannDoS(U2,modelJ,"PFO")
# PFO forcefield: See Figure 5.7, Page 213: https://dx.doi.org/10.6084/m9.figshare.91370.v1
U3,raw=ApproxFunVandermonde("INDT-modred-eV.dat",25) # Use Vandermonde interpolation to load an ApproxFun function
# Using 'U2' to avoid trying to redefine type of U, which Julia does not like.
BoltzmannDoS(U3,modelJ,"INDT-sampled")

end

# This code was used for debugging + profiling

#println("Elements...(offdiag^2, diag))");
#println([E;D])

#println("Full square matrix H");
H=diagm(E.^0.5,-1)+diagm(D)+diagm(E.^0.5,1) #build full matrix from diagonal elements; for comparison
println(H)

println("Eigenvalues")
println(eigvals(H))
println("Min Eigenvalue")
println(eigmin(H))

println("I time a single histogram (Es<sigma) point:")
@time sturm(D,E,5.0)
println("I time eigvals(H)")
@time eigvals(H)

max_issues_repo_name: jarvist/Teclo

max_forks_repo_name: jarvist/Teclo


text:
"""
    struct ContinuousSpace{T<:Number} <: AbstractContinuousSpace
        low::T
        high::T
    end

The lower bound and upper bound are specifed by `low` and `high`.
"""
struct ContinuousSpace{T<:Number} <: AbstractContinuousSpace
    low::T
    high::T
end

eltype(s::ContinuousSpace{T}) where T = T
sample(s::ContinuousSpace) = s.low + rand() * (s.high - s.low)
in(x, s::ContinuousSpace) = s.low ≤ x ≤ s.high


text:
###
### This file contains operations defined on Hamiltonians
###
function Base.getindex(ham::AbstractHamiltonian{T}, address1, address2) where T
    # calculate the matrix element when only two bitstring addresses are given
    # this is NOT used for the QMC algorithm and is currenlty not used either
    # for building the matrix for conventional diagonalisation.
    # Only used for verifying matrix.
    # This will be slow and inefficient. Avoid using for larger Hamiltonians!
    address1 == address2 && return diagonal_element(ham, address1) # diagonal
    for (add,val) in offdiagonals(ham, address2) # off-diag column as iterator
        add == address1 && return val # found address1
    end
    return zero(T) # address1 not found
end

function Base.:*(h::AbstractHamiltonian{E}, v::AbstractDVec{K,V}) where {E, K, V}
    T = promote_type(E, V) # allow for type promotion
    w = empty(v, T) # allocate new vector; non-mutating version
    for (key, val) in pairs(v)
        w[key] += diagonal_element(h, key)*val
        off = offdiagonals(h, key)
        for (add, elem) in off
            w[add] += elem*val
        end
    end
    return w
end

# three argument version
function LinearAlgebra.mul!(w::AbstractDVec, h::AbstractHamiltonian, v::AbstractDVec)
    empty!(w)
    for (key,val) in pairs(v)
        w[key] += diagonal_element(h, key)*val
        for (add,elem) in offdiagonals(h, key)
            w[add] += elem*val
        end
    end
    return w
end

"""
    dot(x, H::AbstractHamiltonian, v)

Evaluate `x⋅H(v)` minimizing memory allocations.
"""
function LinearAlgebra.dot(x::AbstractDVec, LO::AbstractHamiltonian, v::AbstractDVec)
    return dot(LOStructure(LO), x, LO, v)
end
# specialised method for UniformProjector
function LinearAlgebra.dot(::UniformProjector, LO::AbstractHamiltonian{T}, v::AbstractDVec{K,T2}) where {K, T, T2}
    result = zero(promote_type(T,T2))
    for (key,val) in pairs(v)
        result += diagonal_element(LO, key) * val
        for (add,elem) in offdiagonals(LO, key)
            result += elem * val
        end
    end
    return result
end

LinearAlgebra.dot(::AdjointUnknown, x, LO::AbstractHamiltonian, v) = dot_from_right(x,LO,v)
# default for LOs without special structure: keep order

function LinearAlgebra.dot(::LOStructure, x, LO::AbstractHamiltonian, v)
    if length(x) < length(v)
        return conj(dot_from_right(v, LO', x)) # turn args around to execute faster
    else
        return dot_from_right(x,LO,v) # original order
    end
end

"""
    Hamiltonians.dot_from_right(x, LO, v)
Internal function evaluates the 3-argument `dot()` function in order from right
to left.
"""
function dot_from_right(
    x::AbstractDVec{K,T}, LO::AbstractHamiltonian{U}, v::AbstractDVec{K,V}
) where {K,T,U,V}
    result = zero(promote_type(T, U, V))
    for (key, val) in pairs(v)
        result += conj(x[key]) * diagonal_element(LO, key) * val
        for (add, elem) in offdiagonals(LO, key)
            result += conj(x[add]) * elem * val
        end
    end
    return result
end

LinearAlgebra.adjoint(op::AbstractHamiltonian) = adjoint(LOStructure(op), op)

"""
    adjoint(::LOStructure, op::AbstractHamiltonian)

Represent the adjoint of an `AbstractHamiltonian`. Extend this method to define custom
adjoints.
"""
function LinearAlgebra.adjoint(::S, op) where {S<:LOStructure}
    error(
        "`adjoint()` not defined for `AbstractHamiltonian`s with `LOStructure` `$(S)`. ",
        " Is your Hamiltonian hermitian?"
    )
end

LinearAlgebra.adjoint(::IsHermitian, op) = op # adjoint is known