text:
function compute_OOOO!(I::IntegralHelper{T,Chonky,<:AbstractRestrictedOrbitals}, aoints::IntegralHelper{T,SparseERI, AtomicOrbitals}) where T<:AbstractFloat

    idxs = aoints["ERI"].indexes
    eri_vals = aoints["ERI"].data

    core = Options.get("drop_occ")
    ndocc = I.molecule.Nα
    nbf = size(I.orbitals.C,1)
    o = (1+core):ndocc
    Co = I.orbitals.C[:,o]
    if eltype(I.orbitals.C) !== T
        Co = T.(Co)
    end

    # Create a partially-contracted dense array from SparseERI
    iνρσ = zeros(T, length(o), nbf, nbf, nbf)
    @sync for i = 1:length(o)
        Threads.@spawn begin
        @inbounds begin
        @fastmath begin
        for z = eachindex(eri_vals)
            V = eri_vals[z]
            μ,ν,ρ,σ = idxs[z] .+ 1
            γμν = μ !== ν
            γρσ = ρ !== σ
            γab = Fermi.index2(μ,ν) !== Fermi.index2(ρ, σ)
            Vμ = V*Co[μ,i]
            Vν = V*Co[ν,i]
            Vρ = V*Co[ρ,i]
            Vσ = V*Co[σ,i]

            if γab && γμν && γρσ
                iνρσ[i,ν,ρ,σ] += Vμ
                iνρσ[i,ν,σ,ρ] += Vμ
                iνρσ[i,μ,ρ,σ] += Vν
                iνρσ[i,μ,σ,ρ] += Vν
                iνρσ[i,σ,μ,ν] += Vρ
                iνρσ[i,σ,ν,μ] += Vρ
                iνρσ[i,ρ,μ,ν] += Vσ
                iνρσ[i,ρ,ν,μ] += Vσ

            elseif γab && γμν
                iνρσ[i,ν,ρ,σ] += Vμ
                iνρσ[i,μ,ρ,σ] += Vν
                iνρσ[i,σ,μ,ν] += Vρ
                iνρσ[i,σ,ν,μ] += Vρ

            elseif γab && γρσ
                iνρσ[i,ν,ρ,σ] += Vμ
                iνρσ[i,ν,σ,ρ] += Vμ
                iνρσ[i,σ,μ,ν] += Vρ
                iνρσ[i,ρ,μ,ν] += Vσ

            elseif γμν && γρσ
                iνρσ[i,ν,ρ,σ] += Vμ
                iνρσ[i,ν,σ,ρ] += Vμ
                iνρσ[i,μ,ρ,σ] += Vν
                iνρσ[i,μ,σ,ρ] += Vν
            elseif γab
                iνρσ[i,ν,ρ,σ] += Vμ
                iνρσ[i,σ,μ,ν] += Vρ
            else
                iνρσ[i,ν,ρ,σ] += Vμ
            end
        end
        end
        end
        end
    end

    @tensoropt (μ=>100x, ν=>100x, ρ=>100x, σ=>100x, i=>10x, j=>10x, k=>10x, l=>10x, a=>80x, b=>80x, c=>80x, d=>80) begin
        OOOO[i,j,k,l] :=  iνρσ[i, ν, ρ, σ]*Co[ν, j]*Co[ρ, k]*Co[σ, l]
    end
    I["OOOO"] = OOOO
end

function compute_OOOV!(I::IntegralHelper{T,Chonky,<:AbstractRestrictedOrbitals}, aoints::IntegralHelper{T,SparseERI, AtomicOrbitals}) where T<:AbstractFloat

    eri_vals = aoints["ERI"].data
    idxs = aoints["ERI"].indexes

    inac = Options.get("drop_vir")
    core = Options.get("drop_occ")
    ndocc = I.molecule.Nα
    nbf = size(I.orbitals.C,1)
    o = (1+core):ndocc
    v = (ndocc+1):(nbf - inac)
    Cv = I.orbitals.C[:,v]
    Co = I.orbitals.C[:,o]

    # Create a partially-contracted dense array from SparseERI
    iνρσ = zeros(T, length(o), nbf, nbf, nbf)
    @sync for i = 1:length(o)
        Threads.@spawn begin
        @inbounds begin
        @fastmath begin
        for z = eachindex(eri_vals)
            V = eri_vals[z]
            μ,ν,ρ,σ = idxs[z] .+ 1
            γμν = μ !== ν
            γρσ = ρ !== σ
            γab = Fermi.index2(μ,ν) !== Fermi.index2(ρ, σ)
            Vμ = V*Co[μ,i]
            Vν = V*Co[ν,i]
            Vρ = V*Co[ρ,i]
            Vσ = V*Co[σ,i]

            if γab && γμν && γρσ
                iνρσ[i,ν,ρ,σ] += Vμ
                iνρσ[i,ν,σ,ρ] += Vμ
                iνρσ[i,μ,ρ,σ] += Vν
                iνρσ[i,μ,σ,ρ] += Vν
                iνρσ[i,σ,μ,ν] += Vρ
                iνρσ[i,σ,ν,μ] += Vρ
                iνρσ[i,ρ,μ,ν] += Vσ
                iνρσ[i,ρ,ν,μ] += Vσ

            elseif γab && γμν
                iνρσ[i,ν,ρ,σ] += Vμ
                iνρσ[i,μ,ρ,σ] += Vν
                iνρσ[i,σ,μ,ν] += Vρ
                iνρσ[i,σ,ν,μ] += Vρ

            elseif γab && γρσ
                iνρσ[i,ν,ρ,σ] += Vμ
                iνρσ[i,ν,σ,ρ] += Vμ
                iνρσ[i,σ,μ,ν] += Vρ
                iνρσ[i,ρ,μ,ν] += Vσ

            elseif γμν && γρσ
                iνρσ[i,ν,ρ,σ] += Vμ
                iνρσ[i,ν,σ,ρ] += Vμ
                iνρσ[i,μ,ρ,σ] += Vν
                iνρσ[i,μ,σ,ρ] += Vν
            elseif γab
                iνρσ[i,ν,ρ,σ] += Vμ
                iνρσ[i,σ,μ,ν] += Vρ
            else
                iνρσ[i,ν,ρ,σ] += Vμ
            end
        end
        end
        end
        end
    end

    @tensoropt (μ=>100x, ν=>100x, ρ=>100x, σ=>100x, i=>10x, j=>10x, k=>10x, l=>10x, a=>80x, b=>80x, c=>80x, d=>80) begin
        OOOV[i,j,k,a] :=  iνρσ[i, ν, ρ, σ]*Co[ν, j]*Co[ρ, k]*Cv[σ, a]
    end
    I["OOOV"] = OOOV
end

function compute_OOVV!(I::IntegralHelper{T,Chonky,<:AbstractRestrictedOrbitals}, aoints::IntegralHelper{T,SparseERI, AtomicOrbitals}) where T<:AbstractFloat

    eri_vals = aoints["ERI"].data
    idxs = aoints["ERI"].indexes

    inac = Options.get("drop_vir")
    core = Options.get("drop_occ")
    ndocc = I.molecule.Nα
    nbf = size(I.orbitals.C,1)
    o = (1+core):ndocc
    v = (ndocc+1):(nbf - inac)
    Cv = I.orbitals.C[:,v]
    Co = I.orbitals.C[:,o]
    if eltype(I.orbitals.C) !== T
        Cv = T.(Cv)
        Co = T.(Co)
    end

    # Create a partially-contracted dense array from SparseERI
    iνρσ = zeros(T, length(o), nbf, nbf, nbf)
    @sync for i = 1:length(o)
        Threads.@spawn begin
        @inbounds begin
        @fastmath begin
        for z = eachindex(eri_vals)
            V = eri_vals[z]
            μ,ν,ρ,σ = idxs[z] .+ 1
            γμν = μ !== ν
            γρσ = ρ !== σ
            γab = Fermi.index2(μ,ν) !== Fermi.index2(ρ, σ)
            Vμ = V*Co[μ,i]
            Vν = V*Co[ν,i]
            Vρ = V*Co[ρ,i]
            Vσ = V*Co[σ,i]

            if γab && γμν && γρσ
                iνρσ[i,ν,ρ,σ] += Vμ
                iνρσ[i,ν,σ,ρ] += Vμ
                iνρσ[i,μ,ρ,σ] += Vν
                iνρσ[i,μ,σ,ρ] += Vν
                iνρσ[i,σ,μ,ν] += Vρ
                iνρσ[i,σ,ν,μ] += Vρ
                iνρσ[i,ρ,μ,ν] += Vσ
                iνρσ[i,ρ,ν,μ] += Vσ

            elseif γab && γμν
                iνρσ[i,ν,ρ,σ] += Vμ
                iνρσ[i,μ,ρ,σ] += Vν
                iνρσ[i,σ,μ,ν] += Vρ
                iνρσ[i,σ,ν,μ] += Vρ

            elseif γab && γρσ
                iνρσ[i,ν,ρ,σ] += Vμ
                iνρσ[i,ν,σ,ρ] += Vμ
                iνρσ[i,σ,μ,ν] += Vρ
                iνρσ[i,ρ,μ,ν] += Vσ

            elseif γμν && γρσ
                iνρσ[i,ν,ρ,σ] += Vμ
                iνρσ[i,ν,σ,ρ] += Vμ
                iνρσ[i,μ,ρ,σ] += Vν
                iνρσ[i,μ,σ,ρ] += Vν
            elseif γab
                iνρσ[i,ν,ρ,σ] += Vμ
                iνρσ[i,σ,μ,ν] += Vρ
            else
                iνρσ[i,ν,ρ,σ] += Vμ
            end
        end
        end
        end
        end
    end

    @tensoropt (μ=>100x, ν=>100x, ρ=>100x, σ=>100x, i=>10x, j=>10x, k=>10x, l=>10x, a=>80x, b=>80x, c=>80x, d=>80) begin
        OOVV[i,j,a,b] :=  iνρσ[i, ν, ρ, σ]*Co[ν, j]*Cv[ρ, a]*Cv[σ, b]
    end
    I["OOVV"] = OOVV
end

function compute_OVOV!(I::IntegralHelper{T,Chonky,<:AbstractRestrictedOrbitals}, aoints::IntegralHelper{T,SparseERI, AtomicOrbitals}) where T<:AbstractFloat

    eri_vals = aoints["ERI"].data
    idxs = aoints["ERI"].indexes

    inac = Options.get("drop_vir")
    core = Options.get("drop_occ")
    ndocc = I.molecule.Nα
    nbf = size(I.orbitals.C,1)
    o = (1+core):ndocc
    v = (ndocc+1):(nbf - inac)
    Cv = I.orbitals.C[:,v]
    Co = I.orbitals.C[:,o]
    if eltype(I.orbitals.C) !== T
        Cv = T.(Cv)
        Co = T.(Co)
    end

    # Create a partially-contracted dense array from SparseERI
    iνρσ = zeros(T, length(o), nbf, nbf, nbf)
    @sync for i = 1:length(o)
        Threads.@spawn begin
        @inbounds begin
        @fastmath begin
        for z = eachindex(eri_vals)
            V = eri_vals[z]
            μ,ν,ρ,σ = idxs[z] .+ 1
            γμν = μ !== ν
            γρσ = ρ !== σ
            γab = Fermi.index2(μ,ν) !== Fermi.index2(ρ, σ)
            Vμ = V*Co[μ,i]
            Vν = V*Co[ν,i]
            Vρ = V*Co[ρ,i]
            Vσ = V*Co[σ,i]

            if γab && γμν && γρσ
                iνρσ[i,ν,ρ,σ] += Vμ
                iνρσ[i,ν,σ,ρ] += Vμ
                iνρσ[i,μ,ρ,σ] += Vν
                iνρσ[i,μ,σ,ρ] += Vν
                iνρσ[i,σ,μ,ν] += Vρ
                iνρσ[i,σ,ν,μ] += Vρ
                iνρσ[i,ρ,μ,ν] += Vσ
                iνρσ[i,ρ,ν,μ] += Vσ

            elseif γab && γμν
                iνρσ[i,ν,ρ,σ] += Vμ
                iνρσ[i,μ,ρ,σ] += Vν
                iνρσ[i,σ,μ,ν] += Vρ
                iνρσ[i,σ,ν,μ] += Vρ

            elseif γab && γρσ
                iνρσ[i,ν,ρ,σ] += Vμ
                iνρσ[i,ν,σ,ρ] += Vμ
                iνρσ[i,σ,μ,ν] += Vρ
                iνρσ[i,ρ,μ,ν] += Vσ

            elseif γμν && γρσ
                iνρσ[i,ν,ρ,σ] += Vμ
                iνρσ[i,ν,σ,ρ] += Vμ
                iνρσ[i,μ,ρ,σ] += Vν
                iνρσ[i,μ,σ,ρ] += Vν
            elseif γab
                iνρσ[i,ν,ρ,σ] += Vμ
                iνρσ[i,σ,μ,ν] += Vρ
            else
                iνρσ[i,ν,ρ,σ] += Vμ
            end
        end
    end
end
    end
    end

    @tensoropt (μ=>100x, ν=>100x, ρ=>100x, σ=>100x, i=>10x, j=>10x, k=>10x, l=>10x, a=>80x, b=>80x, c=>80x, d=>80) begin
        OVOV[i,a,j,b] :=  iνρσ[i, ν, ρ, σ]*Cv[ν, a]*Co[ρ, j]*Cv[σ, b]
    end
    I["OVOV"] = OVOV
end

function compute_OVVV!(I::IntegralHelper{T,Chonky,<:AbstractRestrictedOrbitals}, aoints::IntegralHelper{T,SparseERI, AtomicOrbitals}) where T<:AbstractFloat

    eri_vals = aoints["ERI"].data
    idxs = aoints["ERI"].indexes

    inac = Options.get("drop_vir")
    core = Options.get("drop_occ")
    ndocc = I.molecule.Nα
    nbf = size(I.orbitals.C,1)
    o = (1+core):ndocc
    v = (ndocc+1):(nbf - inac)
    Cv = I.orbitals.C[:,v]
    Co = I.orbitals.C[:,o]
    if eltype(I.orbitals.C) !== T
        Cv = T.(Cv)
        Co = T.(Co)
    end

    # Create a partially-contracted dense array from SparseERI
    iνρσ = zeros(T, length(o), nbf, nbf, nbf)
    @sync for i = 1:length(o)
        Threads.@spawn begin
        @inbounds begin
        @fastmath begin
        for z = eachindex(eri_vals)
            V = eri_vals[z]
            μ,ν,ρ,σ = idxs[z] .+ 1
            γμν = μ !== ν
            γρσ = ρ !== σ
            γab = Fermi.index2(μ,ν) !== Fermi.index2(ρ, σ)
            Vμ = V*Co[μ,i]
            Vν = V*Co[ν,i]
            Vρ = V*Co[ρ,i]
            Vσ = V*Co[σ,i]

            if γab && γμν && γρσ
                iνρσ[i,ν,ρ,σ] += Vμ
                iνρσ[i,ν,σ,ρ] += Vμ
                iνρσ[i,μ,ρ,σ] += Vν
                iνρσ[i,μ,σ,ρ] += Vν
                iνρσ[i,σ,μ,ν] += Vρ
                iνρσ[i,σ,ν,μ] += Vρ
                iνρσ[i,ρ,μ,ν] += Vσ
                iνρσ[i,ρ,ν,μ] += Vσ

            elseif γab && γμν
                iνρσ[i,ν,ρ,σ] += Vμ
                iνρσ[i,μ,ρ,σ] += Vν
                iνρσ[i,σ,μ,ν] += Vρ
                iνρσ[i,σ,ν,μ] += Vρ

            elseif γab && γρσ
                iνρσ[i,ν,ρ,σ] += Vμ
                iνρσ[i,ν,σ,ρ] += Vμ
                iνρσ[i,σ,μ,ν] += Vρ
                iνρσ[i,ρ,μ,ν] += Vσ

            elseif γμν && γρσ
                iνρσ[i,ν,ρ,σ] += Vμ
                iνρσ[i,ν,σ,ρ] += Vμ
                iνρσ[i,μ,ρ,σ] += Vν
                iνρσ[i,μ,σ,ρ] += Vν
            elseif γab
                iνρσ[i,ν,ρ,σ] += Vμ
                iνρσ[i,σ,μ,ν] += Vρ
            else
                iνρσ[i,ν,ρ,σ] += Vμ
            end
        end
        end
        end
        end
    end

    @tensoropt (μ=>100x, ν=>100x, ρ=>100x, σ=>100x, i=>10x, j=>10x, k=>10x, l=>10x, a=>80x, b=>80x, c=>80x, d=>80) begin
        OVVV[i,a,b,c] := iνρσ[i, ν, ρ, σ]*Cv[ν, a]*Cv[ρ, b]*Cv[σ, c]
    end
    I["OVVV"] = OVVV
end

function compute_VVVV!(I::IntegralHelper{T,Chonky,<:AbstractRestrictedOrbitals}, aoints::IntegralHelper{T,SparseERI, AtomicOrbitals}) where T<:AbstractFloat

    eri_vals = aoints["ERI"].data
    idxs = aoints["ERI"].indexes

    inac = Options.get("drop_vir")
    ndocc = I.molecule.Nα
    nbf = size(I.orbitals.C,1)
    v = (ndocc+1):(nbf - inac)
    Cv = I.orbitals.C[:,v]
    if eltype(I.orbitals.C) !== T
        Cv = T.(Cv)
    end

    # Create a partially-contracted dense array from SparseERI
    aνρσ = zeros(T, length(v), nbf, nbf, nbf)
    @sync for a = 1:length(v)
        Threads.@spawn begin
        @inbounds begin
        @fastmath begin
        for z = eachindex(eri_vals)
            V = eri_vals[z]
            μ,ν,ρ,σ = idxs[z] .+ 1
            γμν = μ !== ν
            γρσ = ρ !== σ
            γab = Fermi.index2(μ,ν) !== Fermi.index2(ρ, σ)
            Vμ = V*Cv[μ,a]
            Vν = V*Cv[ν,a]
            Vρ = V*Cv[ρ,a]
            Vσ = V*Cv[σ,a]

            if γab && γμν && γρσ
                aνρσ[a,ν,ρ,σ] += Vμ
                aνρσ[a,ν,σ,ρ] += Vμ
                aνρσ[a,μ,ρ,σ] += Vν
                aνρσ[a,μ,σ,ρ] += Vν
                aνρσ[a,σ,μ,ν] += Vρ
                aνρσ[a,σ,ν,μ] += Vρ
                aνρσ[a,ρ,μ,ν] += Vσ
                aνρσ[a,ρ,ν,μ] += Vσ

            elseif γab && γμν
                aνρσ[a,ν,ρ,σ] += Vμ
                aνρσ[a,μ,ρ,σ] += Vν
                aνρσ[a,σ,μ,ν] += Vρ
                aνρσ[a,σ,ν,μ] += Vρ

            elseif γab && γρσ
                aνρσ[a,ν,ρ,σ] += Vμ
                aνρσ[a,ν,σ,ρ] += Vμ
                aνρσ[a,σ,μ,ν] += Vρ
                aνρσ[a,ρ,μ,ν] += Vσ

            elseif γμν && γρσ
                aνρσ[a,ν,ρ,σ] += Vμ
                aνρσ[a,ν,σ,ρ] += Vμ
                aνρσ[a,μ,ρ,σ] += Vν
                aνρσ[a,μ,σ,ρ] += Vν
            elseif γab
                aνρσ[a,ν,ρ,σ] += Vμ
                aνρσ[a,σ,μ,ν] += Vρ
            else
                aνρσ[a,ν,ρ,σ] += Vμ
            end
        end
        end
        end
        end
    end
    @tensoropt (μ=>100x, ν=>100x, ρ=>100x, σ=>100x, i=>10x, j=>10x, k=>10x, l=>10x, a=>80x, b=>80x, c=>80x, d=>80) begin
        VVVV[a,b,c,d] :=  aνρσ[a, ν, ρ, σ]*Cv[ν, b]*Cv[ρ, c]*Cv[σ, d]
    end
    I["VVVV"] = VVVV
end

max_stars_repo_name: FermiQC/Fermi

max_issues_repo_name: FermiQC/Fermi

max_forks_repo_name: FermiQC/Fermi


text:
# common packages to load
# using Gadfly (fancy plotting)
# using PyPlot (simple plotting)
# using Distributions (statistics/probability)

ext: jl


text:
"""
    @Array(sub, size)

Homogenous array of elements.
"""
macro Array(sub, size...)
    if all(iscontextfree, size)
        Expr(:call, GlobalRef(Constructs, :SizedArray), sub, size...)
    else
        Expr(:call, GlobalRef(Constructs, :ContextualArray), sub, map(x -> iscontextfree(x) ? x : Expr(:quote, x), size)...)
    end
end

"""
    BaseArray{T, N} <: Wrapper{T, Array{T, N}}

Abstract base type for Array wrapper.
"""
abstract type BaseArray{T, N} <: Wrapper{T, Array{T, N}} end

"""
    SizedArray{T, TSubCon<:Construct{T}, N} <: Wrapper{T, Array{T, N}}

Homogenous array of elements.
"""
struct SizedArray{T, TSubCon<:Construct{T}, N} <: BaseArray{T, N}
    subcon::TSubCon
    size::NTuple{N, UInt}
end

SizedArray(subcon::Construct, size::Vararg{Integer, N}) where {N} = SizedArray(subcon, convert(NTuple{N, UInt}, size))
SizedArray(type::Type, size::Vararg{Integer, N}) where {N} = SizedArray(Construct(type), size...)

subcon(wrapper::SizedArray) = wrapper.subcon
function deserialize(array::SizedArray{T, TSubCon, N}, s::IO; contextkw...) where {T, TSubCon, N}
    result = Array{T}(undef, array.size)
    for i in eachindex(result)
        result[i] = deserialize(array.subcon, s; contextkw...)
    end
    result
end
function serialize(array::SizedArray{T, TSubCon, N}, s::IO, obj::Array{T, N}; contextkw...) where {T, TSubCon, N}
    actualsize = size(obj)
    if actualsize != array.size
        throw(ValidationError("expected $(array.size) elements, found $actualsize."))
    end
    bytecount = 0
    for i in eachindex(obj)
        bytecount += serialize(array.subcon, s, obj[i]; contextkw...)
    end
    bytecount
end
estimatesize(array::SizedArray; contextkw...) = prod(array.size) * estimatesize(array.subcon; contextkw...)

struct ContextualArray{T, TSubCon<:Construct{T}, N} <: BaseArray{T, N}
    subcon::TSubCon
    size::NTuple{N, Any}
end

ContextualArray(subcon::Construct, size::Vararg{Any, N}) where {N} = ContextualArray(subcon, size)
ContextualArray(type::Type, size::Vararg{Any, N}) where {N} = ContextualArray(Construct(type), size...)

subcon(wrapper::ContextualArray) = wrapper.subcon

Construct(::Type{Array{T, N}}, size::Vararg{Integer, N}) where {T, N} = SizedArray(T, size...)


text:
# this is the interface
macro keep(t::Symbol,
    by::Union{Vector{Symbol}, Nothing},
    sort::Union{Vector{Symbol}, Nothing},
    arguments::Union{Vector{Symbol},Union{Expr, Nothing}},
    filter::Union{Expr, Nothing},
    use::Union{String, Nothing},
    options::Union{Dict{String,Any}, Nothing})
    error("""\n
        The syntax is:
            keep <varlist>
        or
            keep if <condition>
    """)
end

# this is `keep <varlist>`
macro keep(t::Symbol, by::Nothing, sort::Nothing, arguments::Vector{Symbol}, filter::Nothing, use::Nothing, options::Nothing)
    return esc(
        quote
            Douglass.@keep_var!($t, $arguments)
        end
    )
end

# this is `keep if <condition>`
macro keep(t::Symbol, by::Union{Vector{Symbol}, Nothing},
    sort::Union{Vector{Symbol}, Nothing},
    arguments::Nothing,
    filter::Expr,
    use::Nothing,
    options::Nothing)
    return esc(
        quote
            Douglass.@keep_if!($t, $by, $sort, $arguments, $filter, $use, $options)
        end
    )
end

# @keep_var! <varlist>
macro keep_var!(t::Symbol,varlist::Vector{Symbol})
    return esc(
        quote
            # check that all variables are present
            Douglass.@assert_vars_present($t, $varlist)
            # keep them
            select!($t, $varlist )
        end
    )
end

# @keep_if! <filter>
# This is the vector-valued function
# macro keep_if!(t::Symbol, filter::Expr)
#     esc(
#         quote
#             # check that filter expands to a Vector of Union{Bool, Missing}
#             Douglass.@assert_filter($t, $filter)
#             # drop it like it's hot
#             keepme = Douglass.@with($t, $filter)
#             filter!(r -> keepme[DataFrames.row(r)] , $t)
#         end
#     )
# end

# version with `by`
# this effectively creates a new dummy variable by group and drops if it's 1
macro keep_if!(t::Symbol, varlist_by::Vector{Symbol},
    varlist_sort::Union{Vector{Symbol}, Nothing},
    arguments::Nothing,
    filter::Expr,
    use::Nothing,
    options::Union{Expr, Nothing})

    # replace non-indexed variables with _n
    Douglass.ref_quotenodes!(filter)

    assigned_var = gensym()
    assigned_var_qn = QuoteNode(assigned_var)

    return esc(
        quote
            # sort, if we need to, (first by-variables, then sort-variables)
            if !isnothing($(varlist_sort)) && !isempty($(varlist_sort))
                sort!($t, $varlist_sort)
            end

            #determine type of resulting column from the type of the first element
            $t[!,$(assigned_var_qn)] = trues(size($t,1))

            # this is the function that maps every sub-df into its transformed df
            my_f = _df -> Douglass.@with _df begin
                # define _N
                _N = size(_df, 1)
                # fill the new variable, row by row
                for _n in 1:size(_df,1)
                    if (isnothing($filter) ? false : $filter)  # if condition is not satisfied, leave with missing
                        $(assigned_var_qn)[_n] = false
                    end
                end
                _df
            end
            t2 = combine(my_f, groupby($t,$varlist_by, sort=false, skipmissing = false ))
            filter!(r -> !t2[DataFrames.row(r),$(assigned_var_qn)] , $t)
            select!($t, Not($(assigned_var_qn)))
        end
    )
end

# version without `by`
macro keep_if!(t::Symbol, by::Nothing,
    varlist_sort::Union{Vector{Symbol}, Nothing},
    arguments::Nothing,
    filter::Expr,
    use::Nothing,
    options::Union{Expr, Nothing})

    # replace non-indexed variables with _n
    Douglass.ref_quotenodes!(filter)

    return esc(
        quote
            # sort, if we need to, (first by-variables, then sort-variables)
            if !isnothing($(varlist_sort)) && !isempty($(varlist_sort))
                sort!($t, $varlist_sort)
            end

            Douglass.@with $t begin
                keepme = BitArray{1}(trues(size($t,1)))
                # define _N
                _N = size($t, 1)
                for _n in 1:size($t,1)
                    if (isnothing($filter) ? false : $filter)  # if condition is not satisfied, leave with missing
                        keepme[_n] = 0
                    end
                end
                filter!(r -> !keepme[DataFrames.row(r)] , $t)
            end

        end
    )
end


text:
using Test

for testscen in 1:2
    valdir, scenario, use_permafrost, use_seaice = get_scenario(testscen)
    println(scenario)

    m = test_page_model()
    include("../src/components/CH4cycle.jl")

    addch4cycle(m, use_permafrost)

    set_param!(m, :CH4Cycle, :e_globalCH4emissions, readpagedata(m, "test/validationdata/$valdir/e_globalCH4emissions.csv"))
    set_param!(m, :CH4Cycle, :rtl_g_landtemperature, readpagedata(m, "test/validationdata/$valdir/rtl_g_landtemperature.csv"))
    set_param!(m, :CH4Cycle, :permtce_permafrostemissions, readpagedata(m, "test/validationdata/$valdir/perm_tot_ce_ch4.csv"))
    set_param!(m, :CH4Cycle, :y_year_0, 2015.)

    p = load_parameters(m)
    p["y_year"] = Mimi.dim_keys(m.md, :time)
    set_leftover_params!(m, p)

    # running Model
    run(m)

    conc = m[:CH4Cycle,  :c_CH4concentration]
    conc_compare = readpagedata(m, "test/validationdata/$valdir/c_ch4concentration.csv")

    @test conc ≈ conc_compare rtol = 1e-4
end


text:
__precompile__()

module PhyloNetworks

    # stdlib (standard libraries)
    using Dates
    using Distributed
    using LinearAlgebra # for LowerTriangular, logdet, diag, mul!
    using Printf: @printf, @sprintf
    using Random
    using Statistics: mean, quantile, median

    # other libraries, indicate compatible version in Project.toml
    using BioSequences
    using BioSymbols
    using Combinatorics: combinations
    using CSV
    using DataFrames
    using DataStructures # for updateInCycle with priority queue
    using Distributions #for RateVariationAcrossSites
    using GLM # for the lm function
    using NLopt # for branch lengths optimization
    using StaticArrays
    using StatsBase # sample, coef etc.
    using StatsFuns # logsumexp, logaddexp, various cdf
    using StatsModels # re-exported by GLM. for ModelFrame ModelMatrix Formula etc

    import Base: show
    import GLM: ftest

    const DEBUGC = false # even more debug messages
    global CHECKNET = false # for debugging only

    export ftest
    export
        ## Network Definition
        HybridNetwork,
        DataCF,
        Quartet,
        readTopology,
        readTopologyLevel1,
        tipLabels,
        writeTopology,
        writeSubTree!,
        hybridlambdaformat,
        deleteleaf!,
        printEdges,
        printNodes,
        sorttaxa!,
        ## SNAQ
        readTrees2CF,
        countquartetsintrees,
        readTableCF,
        readTableCF!,
        writeTableCF,
        readInputTrees,
        readNexusTrees,
        summarizeDataCF,
        snaq!,
        readSnaqNetwork,
        topologyMaxQPseudolik!,
        topologyQPseudolik!,
        ## Network Manipulation
        # getParent, getParents, getMajorParentEdge, getMinorParentEdge, getChildren,
        # functions above: first rename them throughout to be consistent with other packages, like:
        # parent child parents children parentmajor parentminor ancestor sibling offspring
        rootatnode!,
        rootonedge!,
        directEdges!,
        preorder!,
        cladewiseorder!,
        fittedQuartetCF,
        rotate!,
        setLength!,
        setGamma!,
        mapAllelesCFtable,
        deleteHybridThreshold!,
        displayedTrees,
        majorTree,
        minorTreeAt,
        displayedNetworkAt!,
        hardwiredClusters,
        hardwiredCluster,
        hardwiredCluster!,
        hardwiredClusterDistance,
        biconnectedComponents,
        blobDecomposition!,
        blobDecomposition,
        ## Network Bootstrap
        treeEdgesBootstrap,
        hybridDetection,
        summarizeHFdf,
        hybridBootstrapSupport,
        bootsnaq,
        readBootstrapTrees,
        writeMultiTopology,
        readMultiTopologyLevel1,
        readMultiTopology,
        hybridatnode!,
        undirectedOtherNetworks,
        ## Network Calibration
        getNodeAges,
        pairwiseTaxonDistanceMatrix,
        calibrateFromPairwiseDistances!,
        ## Network PCM
        phyloNetworklm,
        PhyloNetworkLinearModel,
        simulate,
        TraitSimulation,
        ParamsBM,
        ParamsMultiBM,
        ShiftNet,
        shiftHybrid,
        getShiftEdgeNumber,
        getShiftValue,
        sharedPathMatrix,
        descendenceMatrix,
        regressorShift,
        regressorHybrid,
        ancestralStateReconstruction,
        ReconstructedStates,
        sigma2_estim,
        mu_estim,
        lambda_estim,
        expectations,
        expectationsPlot,
        predint,
        predintPlot,
        vcv,
        ## Discrete Trait PCM
        parsimonySoftwired,
        parsimonyGF,
        maxParsimonyNet,
        TraitSubstitutionModel,
        EqualRatesSubstitutionModel,
        BinaryTraitSubstitutionModel,
        TwoBinaryTraitSubstitutionModel,
        JC69, HKY85,
        nstates,
        Q,
        getlabels,
        nparams,
        RateVariationAcrossSites,
        randomTrait,
        randomTrait!,
        fitdiscrete,
        readfastatodna,
        stationary,
        empiricalDNAfrequencies,
        nni!,
        mapindividuals,
        phyLiNC!,
        # neighbor joining
        nj

    include("types.jl")
    include("auxiliary.jl")
    include("update.jl")
    include("undo.jl")
    include("addHybrid_snaq.jl")
    include("addHybrid.jl")
    include("deleteHybrid.jl")
    include("moves_snaq.jl")
    include("moves_semidirected.jl")
    include("readwrite.jl")
    include("readData.jl")
    include("snaq_optimization.jl")
    include("pseudolik.jl")
    include("descriptive.jl")
    include("manipulateNet.jl")
    include("bootstrap.jl")
    include("multipleAlleles.jl")
    include("compareNetworks.jl")
    include("traits.jl")
    include("parsimony.jl")
    include("pairwiseDistanceLS.jl")
    include("interop.jl")
    include("substitutionModels.jl")
    include("biconnectedComponents.jl")
    include("traitsLikDiscrete.jl")
    include("deprecated.jl")
    include("nj.jl")
    include("phyLiNCoptimization.jl")

end #module


text:
include("Utils.jl")
include("Kmeans.jl")
module Algorithms
  using Utils
  using K_means
export nearest_neighbor, gram_schmidt, k_means

function gram_schmidt(vector_list)
  epsilon = 1e-6
  k = length(vector_list)
  qs = []
  for vector in vector_list
    q_i = vector
    # orthogonalization with previous vectors
    for j = 1:length(qs)
      q_i -= (qs[j]'vector).*qs[j]
    end
    # check for dependence, quit if q_i ~= 0
    if norm(q_i) <= epsilon
      break
    end
    # normalize and store the result
    q_i = q_i/norm(q_i)
    push!(qs,q_i)
  end
  return qs
end

function nearest_neighbor(query, vector_list)
  """
  args:
    query: vector query
    vector_list: list of vectors

  prints the index of the closest vector, the vector itself, and
  the distance to the minimum vector

  returns: the closest vector to the query
  """
  min_dist, min_vector, min_index = Inf, nothing, nothing
  for (index, vector) in enumerate(vector_list)
    dist_to_vector = dist(query, vector)
    if dist_to_vector < min_dist
      min_dist = dist_to_vector
      min_vector = vector
      min_index = index
    end
  end
  println("Index: $(min_index) \t Vector: $(min_vector) \t Distance: $(min_dist)")
  return min_vector
end
end

max_stars_repo_name: reesepathak/mma-functions

max_issues_repo_name: reesepathak/mma-functions

max_forks_repo_name: reesepathak/mma-functions


text:
#in module Units

#addition
(+)(x::Physical{T1, m, kg, s, A, K, mol, cd}, y::Physical{T2, m, kg, s, A, K, mol, cd}) where {T1, T2, m, kg, s, A, K, mol, cd} =
    Physical(x.value + y.value, m, kg, s, A, K, mol, cd)

#negate
(-)(x::Physical{T, m, kg, s, A, K, mol, cd}) where {T, m, kg, s, A, K, mol, cd} =
    Physical(x.value - y.value, m, kg, s, A, K, mol, cd)

#substration
(-)(x::Physical{T1, m, kg, s, A, K, mol, cd}, y::Physical{T2, m, kg, s, A, K, mol, cd}) where {T1, T2, m, kg, s, A, K, mol, cd} =
    Physical(x.value - y.value, m, kg, s, A, K, mol, cd)

#multiplication with pure number
(*)(k, x::Physical{T, m, kg, s, A, K, mol, cd}) where {T, m, kg, s, A, K, mol, cd} =
    Physical{T, m, kg, s, A, K, mol, cd}(k * x.value)

(*)(x::Physical{T, m, kg, s, A, K, mol, cd}, k) where {T, m, kg, s, A, K, mol, cd} =
    Physical{T, m, kg, s, A, K, mol, cd}(x.value * k)

#general multiplication
(*)(x::Physical{T1, m1, kg1, s1, A1, K1, mol1, cd1}, y::Physical{T2, m2, kg2, s2, A2, K2, mol2, cd2}) where {T1, T2, m1, kg1, s1, A1, K1, mol1, cd1, m2, kg2, s2, A2, K2, mol2, cd2} =
    Physical(x.value * y.value, m1+m2, kg1+kg2, s1+s2, A1+A2, K1+K2, mol1+mol2, cd1+cd2)

#division with a pure number
(/)(k, x::Physical{T, m, kg, s, A, K, mol, cd}) where {T, m, kg, s, A, K, mol, cd} =
    Physical{T, -m, -kg, -s, -A, -K, -mol, -cd}(k / x.value)

(/)(x::Physical{T, m, kg, s, A, K, mol, cd}, k) where {T, m, kg, s, A, K, mol, cd} =
    Physical{T, m, kg, s, A, K, mol, cd}(x.value * k)

#general division
(/)(x::Physical{T1, m1, kg1, s1, A1, K1, mol1, cd1}, y::Physical{T2, m2, kg2, s2, A2, K2, mol2, cd2}) where {T1, T2, m1, kg1, s1, A1, K1, mol1, cd1, m2, kg2, s2, A2, K2, mol2, cd2} =
    Physical(x.value / y.value, m1-m2, kg1-kg2, s1-s2, A1-A2, K1-K2, mol1-mol2, cd1-cd2)

#power and square root
(^)(x::Physical{T, m, kg, s, A, K, mol, cd}, k) where {T, m, kg, s, A, K, mol, cd} =
    Physical(x.value ^ k, m*k, kg*k, s*k, A*k, K*k, mol*k, cd*k)

sqrt(x::Physical{T, m, kg, s, A, K, mol, cd}) where {T, m, kg, s, A, K, mol, cd} = let k = 1//2
    Physical(sqrt(x.value), m*k, kg*k, s*k, A*k, K*k, mol*k, cd*k)
end

#exponent and logrithm
#warning: all units are dropped. Probably I should come up with a better solution later
(^)(x, y::Physical) = x.value ^ k
(^)(x::Physical, y::Physical) = x ^ y.value
exp(x::Physical) = exp(x.value)

log(x::Physical) = log(x.value)
log2(x::Physical) = log2(x.value)
log10(x::Physical) = log10(x.value)
log(b, m) = log(m) - log(b)

#handling ambigious
(^)(x::Physical{T, m, kg, s, A, K, mol, cd}, k::Integer) where {T, m, kg, s, A, K, mol, cd} =
    Physical(x.value ^ k, m*k, kg*k, s*k, A*k, K*k, mol*k, cd*k)

(+)(x::Physical, y::Physical) = begin
    warn("Type mismatch while adding $x and $y")
    x.value + y.value
end

(-)(x::Physical, y::Physical) = begin
    warn("Type mismatch while substracting $x and $y")
    x.value + y.value
end


text:
# ******************************************************************************
# PACKAGES AND INCLUDES
# ******************************************************************************
using Plots
using ImplicitEquations

include("project2.jl")
include("helpers.jl")
include("simple.jl")

# ******************************************************************************
# PATH STORING FUNCTIONS
# ******************************************************************************
function path_hooke_jeeves(f, g, c, x, num_evals, ρ1, ρ2, γ, δ)
    x1_arr = Float64[]; x2_arr = Float64[]; y_arr = Float64[]; max_c_arr = Float64[];
    push!(x1_arr, x[1]); push!(x2_arr, x[2]); push!(y_arr, f(x))
    evals = c(x); max_c = maximum([evals[1], evals[2], 0]); push!(max_c_arr, max_c)
    α, y, n, ϵ = 1, f(x) + penalty(c, x, ρ1, ρ2, δ), length(x), 1e-9
    while α > ϵ && count(f, g, c) < (num_evals-8)
        evals = c(x); max_c = maximum([evals[1], evals[2], 0]); push!(max_c_arr, max_c)
        x, y, α = hooke_jeeves(x -> f(x) + penalty(c, x, ρ1, ρ2, δ), x, y, α, n, 0.5)
        push!(x1_arr, x[1]);push!(x2_arr, x[2]); push!(y_arr, f(x))
        ρ1 *= γ; ρ2 *= γ
    end
    return x, x1_arr, x2_arr, y_arr, max_c_arr
end

function path_mesh_adaptive(f, g, c, x, num_evals, ρ1, ρ2, γ, δ)
    x1_arr = Float64[]; x2_arr = Float64[]; y_arr = Float64[]; max_c_arr = Float64[];
    push!(x1_arr, x[1]); push!(x2_arr, x[2]); push!(y_arr, f(x))
    evals = c(x); max_c = maximum([evals[1], evals[2], 0]); push!(max_c_arr, max_c)
    α, y, n, ϵ = 1, f(x) + penalty(c, x, ρ1, ρ2, δ), length(x), 1e-9
    while α > ϵ && count(f, g, c) < (num_evals-8)
        evals = c(x); max_c = maximum([evals[1], evals[2], 0]); push!(max_c_arr, max_c)
        x, y, α = mesh_adaptive_direct_search(x -> f(x) + penalty(c, x, ρ1, ρ2, δ), x, y, α, n)
        push!(x1_arr, x[1]); push!(x2_arr, x[2]); push!(y_arr, f(x))
        ρ1 *= γ; ρ2 *= γ
    end
    return x, x1_arr, x2_arr, y_arr, max_c_arr
end

# ******************************************************************************
# PLOTTING HELPER FUNCTIONS
# ******************************************************************************
function scatter_helper(x1_arr, x2_arr, x0, x_best, color)
    scatter!([x1_arr], [x2_arr], color=color, label = "", markersize = 3, outline = color)
    scatter!([x0[1]], [x0[2]], color=color,  marker = :rect,label = "", markersize = 5)
    scatter!([x_best[1]], [x_best[2]], color=color, marker = :star4, label = "", markersize = 7)
end

function plot_simple1_contour()
    x1 = -3.0:0.01:3.0
    x2 = -3.0:0.01:3.0
    plot_f(x1, x2) = begin
        -x1 * x2 + 2.0 / (3.0 * sqrt(3.0))
    end
    X = repeat(reshape(x1, 1, :), length(x2), 1)
    Y = repeat(x2, 1, length(x1))
    Z = map(plot_f, X, Y)
    contour(x1, x2, Z, levels=31, fill=false, c=cgrad(:viridis, rev = true), linewidth=2.0, colorbar=true)

    c1(x1,x2) = -x1 - x2
    c2(x1,x2) = x1 + x2^2 - 1

    plot!((Lt(c1,0)) & (Lt(c2,0)), fill=(0.65,"#495C6F"), lims=[-3,3],
        label="Feasible Region")
end

function plot_simple2_contour()
    x1 = -3.0:0.01:3.0
    x2 = -3.0:0.01:3.0
    plot_f(x1, x2) = begin
        (1.0 - x1)^2 + 100.0 * (x2 - x1^2)^2
    end
    X = repeat(reshape(x1, 1, :), length(x2), 1)
    Y = repeat(x2, 1, length(x1))
    Z = map(plot_f, X, Y)
    h = [minimum(Z), 2.5, 10, 25, 100, 250, 500, 750,1000, 1500, 2500, 3500, 5000, 6500, 8000, 10000,maximum(Z)]
    contour(x1, x2, Z, levels=h, fill=false, c=cgrad(:viridis, rev = true), linewidth=2.0, colorbar=true)

    c1(x1,x2) = (x1-1)^3 - x2 + 1
    c2(x1,x2) = x1 + x2 - 2

    plot!((Lt(c1,0)) & (Lt(c2,0)), fill=(0.65,"#495C6F"), lims=[-3,3],
        label="Feasible Region")
end

# ******************************************************************************
# PATH PLOTTING FUNCTIONS
# ******************************************************************************
function plot_simple1_hooke_jeeves(f, g, c, n)

    plot_simple1_contour()

    x01 = [2.5, -2.5]; x02 = [2.5, 2.5]; x03 = [-2.0, 2.0]
    x1_best, x11_arr, x21_arr, _, _ =
        path_hooke_jeeves(f, g, c, x01, n, 1.01, 0, 1.001, 0.5);
    x2_best, x12_arr, x22_arr, _, _ =
        path_hooke_jeeves(f, g, c, x02, n, 1.01, 0, 1.001, 0.5);
    x3_best, x13_arr, x23_arr, _, _ =
        path_hooke_jeeves(f, g, c, x03, n, 1.01, 0, 1.001, 0.5);

    # Plot first path
    plot!(x11_arr,x21_arr,color=:magenta, label = "x0: (2.5, -2.5)",legend = :bottomleft,
        title = "Hooke-Jeeves Method to Minimize Simple1", titlefontsize = 12,
        xlabel = "x1", ylabel = "x2", dpi=300, colorbar = true)
    scatter_helper(x11_arr, x21_arr, x01, x1_best, :magenta)
    # Plot second path
    plot!(x12_arr,x22_arr,color=:red, label = "x0: (2.5, 2.5)")
    scatter_helper(x12_arr, x22_arr, x02, x2_best, :red)
    # Plot third path
    plot!(x13_arr,x23_arr,color=:darkorange, label = "x0: (-2.0, 2.0)")
    scatter_helper(x13_arr, x23_arr, x03, x3_best, :darkorange)

    png("simple1_hooke_jeeves.png")
end

function plot_simple1_mesh_adaptive(f, g, c, n)
    # Plot contour
    plot_simple1_contour()

    # Store paths
    x01 = [2.5, -2.5]; x02 = [2.5, 2.5]; x03 = [-2.0, 2.0]
    x1_best, x11_arr, x21_arr, _, _ =
        path_mesh_adaptive(f, g, c, x01, n, 1.01, 0, 1.001, 0.5);
    x2_best, x12_arr, x22_arr, _, _ =
        path_mesh_adaptive(f, g, c, x02, n, 1.01, 0, 1.001, 0.5);
    x3_best, x13_arr, x23_arr, _, _ =
        path_mesh_adaptive(f, g, c, x03, n, 1.01, 0, 1.001, 0.5);

    # Plot first path
    plot!(x11_arr,x21_arr,color=:magenta, label = "x0: (2.5, -2.5)",legend = :bottomleft,
        title = "Mesh Adaptive Direct Search Method to Minimize Simple1", titlefontsize = 12,
        xlabel = "x1", ylabel = "x2", dpi=300, colorbar = true)
    scatter_helper(x11_arr, x21_arr, x01, x1_best, :magenta)
    # Plot second path
    plot!(x12_arr,x22_arr,color=:red, label = "x0: (2.5, 2.5)")
    scatter_helper(x12_arr, x22_arr, x02, x2_best, :red)
    # Plot third path
    plot!(x13_arr,x23_arr,color=:darkorange, label = "x0: (-2.0, 2.0)")
    scatter_helper(x13_arr, x23_arr, x03, x3_best, :darkorange)

    png("simple1_mesh_adaptive.png")
end

function plot_simple2_hooke_jeeves(f, g, c, n)
    # Plot contour
    plot_simple2_contour()

    # Store paths
    x01 = [2.5, -2.5]; x02 = [2.5, 2.5]; x03 = [-2.0, 2.0]
    x1_best, x11_arr, x21_arr, _, _ =
        path_hooke_jeeves(f, g, c, x01, n, 100, 1, 1.1, 0.75);
    x2_best, x12_arr, x22_arr, _, _ =
        path_hooke_jeeves(f, g, c, x02, n, 100, 1, 1.1, 0.75);
    x3_best, x13_arr, x23_arr, _, _ =
        path_hooke_jeeves(f, g, c, x03, n, 100, 1, 1.1, 0.75);

    # Plot first path
    plot!(x11_arr,x21_arr,color=:magenta, label = "x0: (2.5, -2.5)",legend = :bottomleft,
        title = "Hooke-Jeeves Method to Minimize Simple2", titlefontsize = 12,
        xlabel = "x1", ylabel = "x2", dpi=300, colorbar = true)
    scatter_helper(x11_arr, x21_arr, x01, x1_best, :magenta)
    # Plot second path
    plot!(x12_arr,x22_arr,color=:red, label = "x0: (2.5, 2.5)")
    scatter_helper(x12_arr, x22_arr, x02, x2_best, :red)
    # Plot third path
    plot!(x13_arr,x23_arr,color=:darkorange, label = "x0: (-2.0, 2.0)")
    scatter_helper(x13_arr, x23_arr, x03, x3_best, :darkorange)

    png("simple2_hooke_jeeves.png")
end

function plot_simple2_mesh_adaptive(f, g, c, n)
    # Plot contour
    plot_simple2_contour()

    # Store paths
    x01 = [0.0, -2.0]; x02 = [1.0, -1.0]; x03 = [2.0, 0.0]
    x1_best, x11_arr, x21_arr, _, _ =
        path_mesh_adaptive(f, g, c, x01, n, 100, 1, 1.01, 1.0);
    x2_best, x12_arr, x22_arr, _, _ =
        path_mesh_adaptive(f, g, c, x02, n, 100, 1, 1.01, 1.0);
    x3_best, x13_arr, x23_arr, _, _ =
        path_mesh_adaptive(f, g, c, x03, n, 100, 1, 1.01, 1.0);

    # Plot first path
    plot!(x11_arr,x21_arr,color=:magenta, label = "x0: (0.0, -2.0)",legend = :bottomleft,
        title = "Mesh Adaptive Direct Search Method to Minimize Simple2", titlefontsize = 12,
        xlabel = "x1", ylabel = "x2", dpi=300, colorbar = true)
    scatter_helper(x11_arr, x21_arr, x01, x1_best, :magenta)
    # Plot second path
    plot!(x12_arr,x22_arr,color=:red, label = "x0: (1.0, -1.0)")
    scatter_helper(x12_arr, x22_arr, x02, x2_best, :red)
    # Plot third path
    plot!(x13_arr,x23_arr,color=:darkorange, label = "x0: (2.0, 0.0)")
    scatter_helper(x13_arr, x23_arr, x03, x3_best, :darkorange)

    png("simple2_mesh_adaptive.png")
end

function plot_obj_fun_max_con_hooke_jeeves(f, g, c, n)

    x01 = [2.5, -2.5]; x02 = [2.5, 2.5]; x03 = [-2.0, 2.0]
    _, _, _, y1_arr, max_c1_arr =
        path_hooke_jeeves(f, g, c, x01, n, 100, 1, 1.1, 0.75);
    _, _, _, y2_arr, max_c2_arr =
        path_hooke_jeeves(f, g, c, x02, n, 100, 1, 1.1, 0.75);
    _, _, _, y3_arr, max_c3_arr =
        path_hooke_jeeves(f, g, c, x03, n, 100, 1, 1.1, 0.75);

    x_lim = minimum([length(y1_arr), length(y2_arr), length(y3_arr), 20])
    x_vals = 1:x_lim

    plot(x_vals, y1_arr[1:x_lim], lw=3, color = "#B2BfCC", label="x0: (2.5, -2.5)",legend = :topright,
        title="Objective Function vs Iteration for Hooke-Jeeves",
        xlabel="Iteration",ylabel="Maximum Constraint Violation",dpi=300)
    plot!(x_vals, y2_arr[1:x_lim], lw = 3, color = "#495C6F", label = "x0: (2.5, 2.5)")
    plot!(x_vals, y3_arr[1:x_lim], lw = 3, color = "#1E262E", label = "x0: (-2.0, 2.0)")

    png("objective_function_hooke_jeeves.png")

    plot(x_vals, max_c1_arr[1:x_lim], lw=3, color = "#B2BfCC", label="x0: (2.5, -2.5)",legend = :topright,
        title="Maximum Constraint Violation for Hooke-Jeeves",
        xlabel="Iteration",ylabel="Maximum Constraint Violation",dpi=300)
    plot!(x_vals, max_c2_arr[1:x_lim], lw = 3, color = "#495C6F", label = "x0: (2.5, 2.5)")
    plot!(x_vals, max_c3_arr[1:x_lim], lw = 3, color = "#1E262E", label = "x0: (-2.0, 2.0)")

    png("constraint_violation_hooke_jeeves.png")
end

function plot_obj_fun_max_con_mesh_adaptive(f, g, c, n)

    x01 = [0.0, -2.0]; x02 = [1.0, -1.0]; x03 = [2.0, 0.0]
    _, _, _, y1_arr, max_c1_arr =
        path_mesh_adaptive(f, g, c, x01, n, 100, 1, 1.01, 1.0);
    _, _, _, y2_arr, max_c2_arr =
        path_mesh_adaptive(f, g, c, x02, n, 100, 1, 1.01, 1.0);
    _, _, _, y3_arr, max_c3_arr =
        path_mesh_adaptive(f, g, c, x03, n, 100, 1, 1.01, 1.0);

    x_lim = minimum([length(y1_arr), length(y2_arr), length(y3_arr), 20])
    x_vals = 1:x_lim

    plot(x_vals, y1_arr[1:x_lim], lw=3, color = "#B2BfCC", label="x0: (0.0, -2.0)",legend = :topright,
        title="Objective Function for Mesh Adaptive",
        xlabel="Iteration",ylabel="Maximum Constraint Violation",dpi=300)
    plot!(x_vals, y2_arr[1:x_lim], lw = 3, color = "#495C6F", label = "x0: (1.0, -1.0)")
    plot!(x_vals, y3_arr[1:x_lim], lw = 3, color = "#1E262E", label = "x0: (2.0, 0.0)")

    png("objective_function_mesh_adaptive.png")

    plot(x_vals, max_c1_arr[1:x_lim], lw=3, color = "#B2BfCC", label="x0: (0.0, -2.0)",legend = :topright,
        title="Maximum Constraint Violation for Mesh Adaptive",
        xlabel="Iteration", ylabel="Maximum Constraint Violation",dpi=300)
    plot!(x_vals, max_c2_arr[1:x_lim], lw = 3, color = "#495C6F", label = "x0: (1.0, -1.0)")
    plot!(x_vals, max_c3_arr[1:x_lim], lw = 3, color = "#1E262E", label = "x0: (2.0, 0.0)")

    png("constraint_violation_mesh_adaptive.png")
end

n = 100000
# Path plots
plot_simple1_hooke_jeeves(simple1, simple1_gradient, simple1_constraints, n)
plot_simple1_mesh_adaptive(simple1, simple1_gradient, simple1_constraints, n)
plot_simple2_hooke_jeeves(simple2, simple2_gradient, simple2_constraints, n)
plot_simple2_mesh_adaptive(simple2, simple2_gradient, simple2_constraints, n)

# Objective function and maximum constraint plots
plot_obj_fun_max_con_hooke_jeeves(simple2, simple2_gradient, simple2_constraints, n)
plot_obj_fun_max_con_mesh_adaptive(simple2, simple2_gradient, simple2_constraints, n)


text:
struct LinearCombination{T, As<:Tuple{Vararg{LinearMap}}} <: LinearMap{T}
    maps::As
    function LinearCombination{T, As}(maps::As) where {T, As}
        N = length(maps)
        sz = size(maps[1])
        for n in 1:N
            size(maps[n]) == sz || throw(DimensionMismatch("LinearCombination"))
            promote_type(T, eltype(maps[n])) == T || throw(InexactError())
        end
        new{T, As}(maps)
    end
end

LinearCombination{T}(maps::As) where {T, As} = LinearCombination{T, As}(maps)

# basic methods
Base.size(A::LinearCombination) = size(A.maps[1])
LinearAlgebra.issymmetric(A::LinearCombination) = all(issymmetric, A.maps) # sufficient but not necessary
LinearAlgebra.ishermitian(A::LinearCombination) = all(ishermitian, A.maps) # sufficient but not necessary
LinearAlgebra.isposdef(A::LinearCombination) = all(isposdef, A.maps) # sufficient but not necessary

# adding linear maps
"""
    A::LinearMap + B::LinearMap

Construct a `LinearCombination <: LinearMap`, a (lazy) representation of the sum
of the two operators. Sums of `LinearMap`/`LinearCombination` objects and
`LinearMap`/`LinearCombination` objects are reduced to a single `LinearCombination`.
In sums of `LinearMap`s and `AbstractMatrix`/`UniformScaling` objects, the latter
get promoted to `LinearMap`s automatically.

# Examples
```jldoctest; setup=(using LinearAlgebra, LinearMaps)
julia> CS = LinearMap{Int}(cumsum, 3)::LinearMaps.FunctionMap;

julia> LinearMap(ones(Int, 3, 3)) + CS + I + rand(3, 3);
```
"""
function Base.:(+)(A₁::LinearMap, A₂::LinearMap)
    size(A₁) == size(A₂) || throw(DimensionMismatch("+"))
    T = promote_type(eltype(A₁), eltype(A₂))
    return LinearCombination{T}(tuple(A₁, A₂))
end
function Base.:(+)(A₁::LinearMap, A₂::LinearCombination)
    size(A₁) == size(A₂) || throw(DimensionMismatch("+"))
    T = promote_type(eltype(A₁), eltype(A₂))
    return LinearCombination{T}(tuple(A₁, A₂.maps...))
end
Base.:(+)(A₁::LinearCombination, A₂::LinearMap) = +(A₂, A₁)
function Base.:(+)(A₁::LinearCombination, A₂::LinearCombination)
    size(A₁) == size(A₂) || throw(DimensionMismatch("+"))
    T = promote_type(eltype(A₁), eltype(A₂))
    return LinearCombination{T}(tuple(A₁.maps..., A₂.maps...))
end
Base.:(-)(A₁::LinearMap, A₂::LinearMap) = +(A₁, -A₂)

# comparison of LinearCombination objects, sufficient but not necessary
Base.:(==)(A::LinearCombination, B::LinearCombination) = (eltype(A) == eltype(B) && A.maps == B.maps)

# special transposition behavior
LinearAlgebra.transpose(A::LinearCombination) = LinearCombination{eltype(A)}(map(transpose, A.maps))
LinearAlgebra.adjoint(A::LinearCombination)   = LinearCombination{eltype(A)}(map(adjoint, A.maps))

# multiplication with vectors
if VERSION < v"1.3.0-alpha.115"

function A_mul_B!(y::AbstractVector, A::LinearCombination, x::AbstractVector)
    # no size checking, will be done by individual maps
    A_mul_B!(y, A.maps[1], x)
    l = length(A.maps)
    if l>1
        z = similar(y)
        for n in 2:l
            A_mul_B!(z, A.maps[n], x)
            y .+= z
        end
    end
    return y
end

else # 5-arg mul! is available for matrices

# map types that have an allocation-free 5-arg mul! implementation
const FreeMap = Union{MatrixMap,UniformScalingMap}

function A_mul_B!(y::AbstractVector, A::LinearCombination{T,As}, x::AbstractVector) where {T, As<:Tuple{Vararg{FreeMap}}}
    # no size checking, will be done by individual maps
    A_mul_B!(y, A.maps[1], x)
    for n in 2:length(A.maps)
        mul!(y, A.maps[n], x, true, true)
    end
    return y
end
function A_mul_B!(y::AbstractVector, A::LinearCombination, x::AbstractVector)
    # no size checking, will be done by individual maps
    A_mul_B!(y, A.maps[1], x)
    l = length(A.maps)
    if l>1
        z = similar(y)
        for n in 2:l
            An = A.maps[n]
            if An isa FreeMap
                mul!(y, An, x, true, true)
            else
                A_mul_B!(z, A.maps[n], x)
                y .+= z
            end
        end
    end
    return y
end

function LinearAlgebra.mul!(y::AbstractVector, A::LinearCombination{T,As}, x::AbstractVector, α::Number=true, β::Number=false) where {T, As<:Tuple{Vararg{FreeMap}}}
    length(y) == size(A, 1) || throw(DimensionMismatch("mul!"))
    if isone(α)
        iszero(β) && (A_mul_B!(y, A, x); return y)
        !isone(β) && rmul!(y, β)
    elseif iszero(α)
        iszero(β) && (fill!(y, zero(eltype(y))); return y)
        isone(β) && return y
        # β != 0, 1
        rmul!(y, β)
        return y
    else # α != 0, 1
        if iszero(β)
            A_mul_B!(y, A, x)
            rmul!(y, α)
            return y
        elseif !isone(β)
            rmul!(y, β)
        end # β-cases
    end # α-cases

    for An in A.maps
        mul!(y, An, x, α, true)
    end
    return y
end

end # VERSION

At_mul_B!(y::AbstractVector, A::LinearCombination, x::AbstractVector) = A_mul_B!(y, transpose(A), x)

Ac_mul_B!(y::AbstractVector, A::LinearCombination, x::AbstractVector) = A_mul_B!(y, adjoint(A), x)


text:
function Subtract()
    return Subtract(())
end

function value(obj::Subtract, arg0::jdouble, arg1::jdouble)
    return jcall(obj, "value", jdouble, (jdouble, jdouble), arg0, arg1)
end


text:
struct UnitCell{D}
    n     :: Int
    sites :: Vector{NTuple{D, Float64}}
    as    :: Vector{NTuple{D, Float64}}

    function UnitCell{D}(n     :: Int,
                         sites :: Vector,
                         as    :: Vector) where{D}
        length(sites) == n || error("number of sites don't match!")
        length(as) == D || error("dimension of unitcell vector don't match")
        new{D}(n, sites, as)
    end
end

UnitCell{1}(n::Int, sites::Vector{Float64}, a::Float64) =
    UnitCell{1}(n, [(site,) for site in sites], [(a,)])

@inline coordsum(A::NTuple{N}, B::NTuple{N}) where {N} = A .+ B

struct QLattice{D}
    unitc :: UnitCell{D}
    sizes :: NTuple{D, Int}
    bcs   :: NTuple{D, Symbol}
end

QLattice(uc::UnitCell{1}, lx::Int, bc::Symbol) where{D} =
    QLattice{1}(uc, (lx,), (bc,))

@inline dimension(lattice::QLattice{D}) where{D} = D

"return a changed boundary version of `lattice` in the `d`th direction"
function changeboundary(lattice::QLattice, d::Int, boundary::Symbol)
    D = dimension(lattice)
    1 <= d <= D || error("can't change boundary!")
    QLattice(lattice.unitc,
             lattice.sizes,
             Tuple([lattice.bcs[1:d-1]..., boundary, lattice.bcs[d+1:D]...]))
end

"return a chaged size version of `lattice` in the `d`th direction"
function changesize(lattice::QLattice, d::Int, dim::Int)
    D = dimension(lattice)
    1 <= d <= D || error("can't change size!")
    QLattice(lattice.unitc,
             Tuple([lattice.sizes[1:d-1]..., dim, lattice.sizes[d+1:D]...]),
             lattice.bcs)
end

"""
    sitelinearindex(lattice, ucidx, x_uc)

for a given lattice `lattice`, unite cell site `ucidx`, and a
corrdinate of unitcell `x_uc` which can be any integer number, find
the site linear index of that site if allowed according to the
boundary condition!

"""
function sitelinearindex(lattice::QLattice{D},
                         ucidx::Int,
                         x_uc::NTuple{D, Int}) where {D}

    x_uc_vec = zeros(Int, D)
    crossings = zeros(Int, D)
    for i in 1:D
        if lattice.bcs[i] == :OBC
            if 1 <= x_uc[i] <= lattice.sizes[i]
                x_uc_vec[i] = x_uc[i]
            else
                return nothing, crossings
            end

        elseif lattice.bcs[i] in [:PBC, :APBC]
            x_uc_vec[i] = mod(x_uc[i] - 1, lattice.sizes[i]) + 1
            crossings[i] = fld(x_uc[i] - 1, lattice.sizes[i])

        elseif lattice.bcs[i] == :INF
            if i < D
                error("only the last dimensions can have INF boundary conition!")
            end
            x_uc_vec[i] = x_uc[i]

        else
            error("Exhaustive check!")
        end
    end
    index = ucidx + lattice.unitc.n *
        sum((x_uc_vec .- 1) .* [1, cumprod([lattice.sizes...])[1:D-1]...])

    return index, Tuple(crossings)
end

abstract type AbstractQType end
struct SpinType <: AbstractQType
    d :: Int
end

# because different species of fermions should be different for
# anticommutation relations
struct Fermion{Name} <: AbstractQType
    Fermion{Name}() where {Name} = new{typeassert(Name, Symbol)}()
end

abstract type QuantumOperator end
abstract type FermionOp end

struct FCreate{F} <: FermionOp
    FCreate{F}() where {F} = new{typeassert(F, Fermion)}()
end

struct FAnnihilate{F} <: FermionOp
    FAnnihilate{F}() where {F}= new{typeassert(F, Fermion)}()
end

function fermionoperators(f::Fermion)
    FCreate{f}(), FAnnihilate{f}()
end
abstract type AbstractQInteraction{T, N} end

struct QAmpTerm{O, N, T}
    amp :: T
    ops :: NTuple{N, O}
end

abstract type AbstractQModelInteraction{D, N, T} end
struct QModelInteraction{D, N, T} <: AbstractQModelInteraction{D, N, T}
    amp     :: T
    ucidxs  :: NTuple{N, Int}
    offsets :: NTuple{N, NTuple{D, Int}}
    terms   :: Vector{QAmpTerm{Symbol,N,T}}
end

# struct SymQModelInteraction{D, N, T} <: AbstractQModelInteraction{D, N, T}
#     amp     :: T
#     ucidxs  :: NTuple{N, Int}
#     offsets :: NTuple{N, NTuple{D, Int}}
#     terms   :: Vector{NTuple{N, SymMatrix{T}}}
# end

struct FermionQModelInteraction{D, N, T} <: AbstractQModelInteraction{D, N, T}
    amp     :: T
    ucidxs  :: NTuple{N, Int}
    offsets :: NTuple{N, NTuple{D, Int}}
    terms   :: Vector{NTuple{N, FermionOp}}
end

function supportrange(inter::AbstractQModelInteraction{D, N, T}) where{D,N,T}
    [(minimum([offset[d] for offset in inter.offsets]),
      maximum([offset[d] for offset in inter.offsets])) for d in 1:D]
end

support(::AbstractQModelInteraction{D, N, T}) where{D, N, T} = N
eltype(::AbstractQModelInteraction{D, N, T}) where{D, N, T} = T

# This is the generic linear QTerm
struct QInteraction{T, N} <: AbstractQInteraction{T, N}
    amp   :: T
    sites :: NTuple{N, Int}
    terms :: Vector{NTuple{N, Matrix{T}}}
end

support(::AbstractQInteraction{T,N}) where{T, N} = N
removehead(A::QInteraction{T, N}) where {T, N} =
    QInteraction{T, N-1}(A.amp, A.sites[2:N], [term[2:N] for term in A.terms])

struct QTerm{OP, N}
    sites :: NTuple{N, Int}
    ops   :: NTuple{N, OP}
end

eltype(term::QTerm{OP}) where {OP} = eltype(OP)
support(::QTerm{OP, N}) where {OP, N} = N
removehead(A::QTerm{OP, N}) where {OP, N} =
    QTerm{OP, N-1}(A.sites[2:N], A.ops[2:N])

# function *(a::T, term::QTerm) where {T<:Number}
#     T == eltype(eltype(term)) || error("oops, $T vs $(eltype(eltype(term)))")
#     N = support(term)
#     QTerm{OP, N}(term.sites, [op for op in term.ops[1:N-1]]..., a*term.ops[N])
# end

abstract type AbstractQModel{Q, D} end
struct UnitCellQModel{Q<:AbstractQType, D, J<:AbstractQModelInteraction} <: AbstractQModel{Q, D}
    qtype    :: Q
    lattice  :: QLattice{D}
    symmetry :: Type{<:AbstractCharge}
    inters   :: Vector{J}
end

@inline dimension(::AbstractQModel{Q, D}) where {Q, D} = D
@inline numofsites(m::UnitCellQModel) = prod(m.lattice.sizes)*m.lattice.unitc.n

function changeboundary(model::UnitCellQModel, d::Int, boundary::Symbol)
    UnitCellQModel(model.qtype,
                   changeboundary(model.lattice, d, boundary),
                   model.inters)
end

function changesize(model::UnitCellQModel, d::Int, dim::Int)
    UnitCellQModel(model.qtype,
                   changesize(model.lattice, d, dim),
                   model.inters)
end

function largestxrange(model::UnitCellQModel)
    D = dimension(model)
    ranges = [supportrange(inter)[D] for inter in model.inters]
    maximum([r[2] for r in ranges]) - minimum([r[1] for r in ranges])
end

max_stars_repo_name: amiragha/mps

max_issues_repo_name: amiragha/mps

max_forks_repo_name: amiragha/mps


text:
const mutated = nothing
# TODO: define a singleton on Python side

_gen_doc(stem, jl) = """
    __$(stem)__(self, other)

Python's $(_pymodelref("__$(stem)__")).
Default implementation is roughly: `broadcast($jl, self, other)`.
"""

_gen_rdoc(stem) = """
    __r$(stem)__(self, other)

Python's $(_pymodelref("__r$(stem)__")).
Default implementation: `PyBase.__$(stem)__(other, self)`.
"""

_gen_idoc(stem, jl) = """
    __i$(stem)__(self, other)

Python's $(_pymodelref("__i$(stem)__")).
If `self` is mutated, `PyBase.mutated` must be returned.
Default implementation:

```julia
function __i$(stem)__(self, other)
    broadcast!($jl, self, self, other)
    return PyBase.mutated
end
```
"""

for (stem, jl) in [
        (:add, +),
        (:sub, -),
        (:mul, *),
        (:truediv, /),
        (:floordiv, ÷),
        (:mod, Base.mod),
        (:pow, ^),
        (:lshift, <<),
        (:rshift, >>),
        (:and, &),
        (:xor, Base.xor),
        (:or, |),
        ]
    py = Symbol(:__, stem, :__)
    @eval $py(a, b) = _wrap(broadcast($jl, a, b))
    @eval $(Symbol(:__r, stem, :__))(a, b) = $py(b, a)
    @eval function $(Symbol(:__i, stem, :__))(a, b)
        broadcast!($jl, a, a, b)
        return mutated
    end
    doc = _gen_doc(stem, jl)
    rdoc = _gen_rdoc(stem)
    idoc = _gen_idoc(stem, jl)
    @eval @doc $doc $py
    @eval @doc $rdoc $(Symbol(:__r, stem, :__))
    @eval @doc $idoc $(Symbol(:__i, stem, :__))
end

__pow__(a, b, modulo) = _wrap(broadcast(powermod, a, b, modulo))

__divmod__(a, b) = (_wrap(broadcast(÷, a, b)),
                    _wrap(broadcast(Base.mod, a, b)))
__rdivmod__(a, b) = __divmod__(b, a)

__matmul__(a, b) = _wrap(a * b)
__rmatmul__(a, b) = __matmul__(b, a)

function __imatmul__(y, a, b)
    if applicable(mul!, y, a, b)
        mul!(y, a, b)
        return mutated
    else
        return _wrap(a * b)
    end
end

for (stem, jl) in [
        (:neg, -),
        (:pos, +),
        (:abs, abs),
        (:invert, ~),
        (:complex, ComplexF64),
        (:int, Int128),
        (:float, Float64),
        # TODO: Use BigFloat/BigInt if mpmath is installed?
        (:trunc, trunc),
        (:floor, floor),
        (:ceil, ceil),
        ]
    py = Symbol(:__, stem, :__)
    @eval $py(a) = _wrap($jl(a))
    doc = """
        __$(stem)__(self)

    Python's $(_pymodelref("__$(stem)__")).
    Default implementation is roughly: `$jl(self)`.
    """
    @eval @doc $doc $py
end

"""
    __round__(self, [ndigits = nothing])

Python's $(_pymodelref("__round__")).
Default implementation is roughly:

```julia
__round__(self) = round(self)
__round__(self, ndigits) = round(self, digits=ndigits)
```
"""
__round__(self) = _wrap(round(self))
__round__(self, ::Nothing) = __round__(self)
__round__(self, ndigits) = _wrap(round(self, digits=ndigits))


text:
#
# This function builds the grid of the 3D density function and fills an array of
# mutable structures of type MutableGMD3DGrid, containing the position of the atoms of
# grid, the closest atom to that position, and distance.
#

function gmd3D_setgrid(mysim,solute;dmax=7.0,step=0.5,dmin=1.0)

  xmin = [ mysim.atoms[solute[1]].coor[1],
           mysim.atoms[solute[1]].coor[2],
           mysim.atoms[solute[1]].coor[3] ]
  xmax = [ mysim.atoms[solute[1]].coor[1],
           mysim.atoms[solute[1]].coor[2],
           mysim.atoms[solute[1]].coor[3] ]

  for i in 2:length(solute)
    xmin[1] = min(xmin[1],mysim.atoms[solute[i]].coor[1])
    xmin[2] = min(xmin[2],mysim.atoms[solute[i]].coor[2])
    xmin[3] = min(xmin[3],mysim.atoms[solute[i]].coor[3])
    xmax[1] = max(xmax[1],mysim.atoms[solute[i]].coor[1])
    xmax[2] = max(xmax[2],mysim.atoms[solute[i]].coor[2])
    xmax[3] = max(xmax[3],mysim.atoms[solute[i]].coor[3])
  end

  xmin[1] = xmin[1] - dmax
  xmin[2] = xmin[2] - dmax
  xmin[3] = xmin[3] - dmax
  xmax[1] = xmax[1] + dmax
  xmax[2] = xmax[2] + dmax
  xmax[3] = xmax[3] + dmax

  nx = (xmax[1]-xmin[1])/step+1
  ny = (xmax[2]-xmin[2])/step+1
  nz = (xmax[3]-xmin[3])/step+1

  # Counting the number of grid points

  dmin2 = dmin^2
  dmax2 = dmax^2
  ngrid = 0
  for ix in 1:nx
    x = xmin[1] + step*(ix-1)
    for iy in 1:ny
      y = xmin[2] + step*(iy-1)
      for iz in 1:nz
        z = xmin[3] + step*(iz-1)
        for iat in 1:length(solute)
          d2 = ( mysim.atoms[solute[iat]].coor[1] - x )^2 +
               ( mysim.atoms[solute[iat]].coor[2] - y )^2 +
               ( mysim.atoms[solute[iat]].coor[3] - z )^2
          if d2 <= dmax2 && d2 > dmin2
            ngrid = ngrid + 1
            break
          end
        end
      end
    end
  end

  # Initializing grid structure

  grid = Vector{MutableGMD3DGrid}(undef,ngrid)

  # Saving the data for each grid point

  igrid = 0
  for ix in 1:nx
    x = xmin[1] + step*(ix-1)
    for iy in 1:ny
      y = xmin[2] + step*(iy-1)
      for iz in 1:nz
        z = xmin[3] + step*(iz-1)
        keep = false
        for iat in 1:length(solute)
          d2 = ( mysim.atoms[solute[iat]].coor[1] - x )^2 +
               ( mysim.atoms[solute[iat]].coor[2] - y )^2 +
               ( mysim.atoms[solute[iat]].coor[3] - z )^2
          if d2 <= dmax2 && d2 > dmin2
            if keep == false
              keep = true
              igrid = igrid + 1
              grid[igrid] = MutableGMD3DGrid()
              grid[igrid].x = [ x, y, z ]
              grid[igrid].atom = solute[iat]
              grid[igrid].dmin = sqrt(d2)
            end
            if d2 <= grid[igrid].dmin
              grid[igrid].atom = solute[iat]
              grid[igrid].dmin = sqrt(d2)
            end
          end
        end
      end
    end
  end

  return grid

end


text:
immutable Maton
    n::Int
    inis::Set{Int}
    finis::Set{Int}
    over::Vector{Pair{Int,Int}}
    by::Vector{eltype(ABET)}
end

immutable Ducer
    n::Int
    inis::Set{Int}
    finis::Set{Int}
    over::Vector{Pair{Int,Int}}
    by::Vector{Pair{eltype(ABET),eltype(ABET)}}
end

include("maton.jl")

include("ducer.jl")

mac(x::Word) = wm(x.w)

mac(x::BotS) = botm()
mac(x::TopS) = topm()

mac(x::AltS) = alt(mac(x.a), mac(x.b))
mac(x::CatS) = cat(mac(x.a), mac(x.b))
mac(x::IntS) = int(mac(x.a), mac(x.b))

mac(x::StarS) = star(mac(x.a))
mac(x::PlusS) = plus(mac(x.a))

mac(x::BotR) = botd()
mac(x::TopR) = topd()

mac(x::AltR) = alt(mac(x.a), mac(x.b))
mac(x::CatR) = cat(mac(x.a), mac(x.b))

mac(x::StarR) = star(mac(x.a))
mac(x::PlusR) = plus(mac(x.a))


text:
using Test

@testset "FinPoly" begin
  include("FinPolynomials.jl")
end

@testset "SumProdPoly" begin
  include("SumProdPolynomials.jl")
end


text:
# using SuperResolution

mutable struct EvalSR
    _selectIDX::Array{Tuple{Int64,Int64},1}
    _modus::Int64
    wModel::Array{SuperResolution.ModelSR,1}
    RefModel::Array{SuperResolution.ModelSR,1}
    distKL::Array{SuperResolution.DistanceSR,1}
    wkldSR::Array{SuperResolution.clusWKLD,1}
    gapSR::Array{SuperResolution.clusGAP,1}

    max_mean_idx::Array{Int64,1}
    max_median_idx::Array{Int64,1}
    max_max_idx::Array{Int64,1}
    max_minDelta_idx::Array{Int64,1}

    diff_WK_max_idx::Array{Int64,1}
    diff_GAP_max_idx::Array{Int64,1}

    N_observations::Array{Int64,1}
    K_ground_truth::Array{Int64,1}

    K_KL_estimated_by_mean::Array{Int64,1}
    K_KL_estimated_by_median::Array{Int64,1}
    K_KL_estimated_by_max::Array{Int64,1}
    K_KL_estimated_by_minDelta::Array{Int64,1}
    K_KL_estimated_by_95quantile::Array{Int64,1}

    K_GAP_estimated::Array{Int64,1}

    corrDetecRate_KL_by_mean::Float64
    corrDetecRate_KL_by_median::Float64
    corrDetecRate_KL_by_max::Float64
    corrDetecRate_KL_by_minDelta::Float64
    corrDetecRate_KL_by_95quantile::Float64

    corrDetecRate_GAP::Float64

    VALID_DF::DataFrame
    error_DF::DataFrame

    locID_DF::Array{DataFrame,1}
    evo_DF::Array{DataFrame,1}

    function EvalSR(SEG::SuperResolution.SegmentSR,RESTR_SEG::Array{Tuple{Int64,Int64},1},modINPUT=5)#RESTR_SEG::Int64)
        obj=new()
        obj._selectIDX=RESTR_SEG
        obj._modus=modINPUT
        print("*** modINPUT = $(modINPUT) ***\n")
        NUM_SEG=length(RESTR_SEG)
        #obj._selectIDX=RESTR_SEG[1:RESTR_SEG]
        obj.wModel=Array{SuperResolution.ModelSR,1}(undef,NUM_SEG)
        obj.RefModel=Array{SuperResolution.ModelSR,1}(undef,NUM_SEG)
        obj.distKL=Array{SuperResolution.DistanceSR,1}(undef,NUM_SEG)
        obj.wkldSR=Array{SuperResolution.clusWKLD,1}(undef,NUM_SEG)
        obj.gapSR=Array{SuperResolution.clusGAP,1}(undef,NUM_SEG)

        obj.diff_WK_max_idx=Array{Int64,1}(undef,NUM_SEG)
        obj.diff_GAP_max_idx=Array{Int64,1}(undef,NUM_SEG)

        obj.K_ground_truth=Array{Int64,1}(undef,NUM_SEG)
        obj.N_observations=Array{Int64,1}(undef,NUM_SEG)

        select_mode=[   :K_KL_estimated_by_mean,                    #   modINPUT=1
                        :K_KL_estimated_by_median,                  #   modINPUT=2
                        :K_KL_estimated_by_max,                     #   modINPUT=3
                        :K_KL_estimated_by_minDelta,                #   modINPUT=4
                        :K_KL_estimated_by_95quantile]              #   modINPUT=5

        @time(for i=1:NUM_SEG
                print("\n>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>\n")
                print(">>> Starting Segment $(i)/$(NUM_SEG) with index $(SEG.idx[i])\n")
                print("\n>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>\n")
                print("(wModel:)\n")
                obj.wModel[i]=SuperResolution.ModelSR(SEG,RESTR_SEG[i],1)     #for segment create all models  WKLD
                # print("(RefModel:)\n")
                #obj.RefModel[i]=SuperResolution.ModelSR(SEG,RESTR_SEG[i],2)   #for segment create all models KLD
                # obj.RefModel[i]=obj.wModel[i]   #temporary workaround to save computation time

                obj.wModel[i].wkld[1:(end-1)].=obj.wModel[i].wkld[2:end]
                obj.wModel[i].weight[1:(end-1)].=obj.wModel[i].weight[2:end]
                obj.wModel[i].geom[1:(end-1)].=obj.wModel[i].geom[2:end]

                obj.wkldSR[i]=SuperResolution.clusWKLD(obj.wModel[i],1)    #cluster WKLD
                # obj.gapSR[i]=SuperResolution.clusGAP(obj.wModel[i],obj.RefModel[i],1)  #Gap statistic
        end)
        obj.K_ground_truth=map(i->SEG.K[CartesianIndex(RESTR_SEG[i])],1:NUM_SEG)
        obj.N_observations=map(i->SEG.N[CartesianIndex(RESTR_SEG[i])],1:NUM_SEG)

        obj.max_mean_idx=map(i->obj.wkldSR[i].max_mean_idx[1],1:length(obj.wkldSR))
        obj.max_median_idx=map(i->obj.wkldSR[i].max_median_idx[1],1:length(obj.wkldSR))
        obj.max_max_idx=map(i->obj.wkldSR[i].max_max_idx[1],1:length(obj.wkldSR))
        obj.max_minDelta_idx=map(i->obj.wkldSR[i].max_minDelta_idx[1],1:length(obj.wkldSR))
        # obj.diff_WK_max_idx=map(i->obj.gapSR[i].diff_WK_max_idx[1],1:length(obj.gapSR))

        #obj.K_KL_estimated_by_mean=map(i->obj.wModel[i].K[CartesianIndex(RESTR_SEG[i])],1:NUM_SEG)
        obj.K_KL_estimated_by_mean=map(i->size(obj.wModel[i].wkld[obj.max_mean_idx[i]],1)-1,1:NUM_SEG)
        obj.K_KL_estimated_by_median=map(i->size(obj.wModel[i].wkld[obj.max_median_idx[i]],1)-1,1:NUM_SEG)
        obj.K_KL_estimated_by_max=map(i->size(obj.wModel[i].wkld[obj.max_max_idx[i]],1)-1,1:NUM_SEG)
        obj.K_KL_estimated_by_minDelta=map(i->size(obj.wModel[i].wkld[obj.max_minDelta_idx[i]],1)-1,1:NUM_SEG)
        obj.K_KL_estimated_by_95quantile=map(i->length( obj.wkldSR[i].wkldMinDelta[ .!(obj.wkldSR[i].wkldMinDelta .< quantile(obj.wkldSR[i].wkldMinDelta, 0.95))] )+1,1:length(obj.wkldSR))  # +1 to account for indexing

        # obj.K_KL_estimated_by_95quantile=map(i->length( obj.wkldSR[i].wkldMin[ .!(obj.wkldSR[i].wkldMin .< quantile(obj.wkldSR[i].wkldMin, 0.95))] )+1,1:length(obj.wkldSR))  # +2 to account for indexing and +1 to account for floor quantile to closest existing wlkd value lower than the quantile
        # obj.K_GAP_estimated=map(i->size(obj.wModel[i].wkld[obj.diff_WK_max_idx[i]],1),1:NUM_SEG)
        #obj.corrDetecRate_GAP=sum(abs.(obj.K_ground_truth-obj.diff_WK_max_idx))/sum(obj.K_ground_truth)
        obj.corrDetecRate_KL_by_mean=mean(exp.((-abs.(obj.K_ground_truth-obj.K_KL_estimated_by_mean))./obj.K_ground_truth))
        obj.corrDetecRate_KL_by_median=mean(exp.((-abs.(obj.K_ground_truth-obj.K_KL_estimated_by_median))./obj.K_ground_truth))
        obj.corrDetecRate_KL_by_max=mean(exp.((-abs.(obj.K_ground_truth-obj.K_KL_estimated_by_max))./obj.K_ground_truth))
        obj.corrDetecRate_KL_by_minDelta=mean(exp.((-abs.(obj.K_ground_truth-obj.K_KL_estimated_by_minDelta))./obj.K_ground_truth))
        obj.corrDetecRate_KL_by_95quantile=mean(exp.((-abs.(obj.K_ground_truth-obj.K_KL_estimated_by_95quantile))./obj.K_ground_truth))
        # obj.corrDetecRate_GAP=mean(exp.((-abs.(obj.K_ground_truth-obj.K_GAP_estimated))./obj.K_ground_truth))
        obj.VALID_DF=DataFrame( idx=obj._selectIDX,
                                N_observations=obj.N_observations,
                                K_ground_truth=obj.K_ground_truth,
                                K_KL_estimated_by_mean=obj.K_KL_estimated_by_mean,
                                K_KL_estimated_by_median=obj.K_KL_estimated_by_median,
                                K_KL_estimated_by_max=obj.K_KL_estimated_by_max,
                                K_KL_estimated_by_minDelta=obj.K_KL_estimated_by_minDelta,
                                K_KL_estimated_by_95quantile=obj.K_KL_estimated_by_95quantile,
                                corrDetecRate_KL_by_mean=obj.corrDetecRate_KL_by_mean,
                                corrDetecRate_KL_by_median=obj.corrDetecRate_KL_by_median,
                                corrDetecRate_KL_by_max=obj.corrDetecRate_KL_by_max,
                                corrDetecRate_KL_by_minDelta=obj.corrDetecRate_KL_by_minDelta,
                                corrDetecRate_KL_by_95quantile=obj.corrDetecRate_KL_by_95quantile,
                                dev_by_mean=(obj.K_ground_truth-obj.K_KL_estimated_by_mean),
                                dev_by_median=(obj.K_ground_truth-obj.K_KL_estimated_by_median),
                                dev_by_max=(obj.K_ground_truth-obj.K_KL_estimated_by_max),
                                dev_by_minDelta=(obj.K_ground_truth-obj.K_KL_estimated_by_minDelta),
                                dev_by_95quantile=(obj.K_ground_truth-obj.K_KL_estimated_by_95quantile) )

        print("\n<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<\n")
        print("<<< K_ground_truth $(obj.K_ground_truth)\n")
        print("<<< N_observations $(obj.N_observations)\n")
        print("------------------------------------------------\n")
        # print("<<< K_KL_estimated_by_mean $(obj.K_KL_estimated_by_mean)\n")
        # print("<<< KL(mean) correct detection rate $(obj.corrDetecRate_KL_by_mean)\n")
        # print("<<< deviation from ground truth $(obj.K_ground_truth-obj.K_KL_estimated_by_mean)\n")
        # print("------------------------------------------------\n")
        print("<<< K_KL_estimated_by_median $(obj.K_KL_estimated_by_median)\n")
        print("<<< KL(median) correct detection rate $(obj.corrDetecRate_KL_by_median)\n")
        print("<<< deviation from ground truth $(obj.K_ground_truth-obj.K_KL_estimated_by_median)\n")
        print("------------------------------------------------\n")
        # print("<<< K_KL_estimated_by_max $(obj.K_KL_estimated_by_max)\n")
        # print("<<< KL(max) correct detection rate $(obj.corrDetecRate_KL_by_max)\n")
        # print("<<< deviation from ground truth $(obj.K_ground_truth-obj.K_KL_estimated_by_max)\n")
        # print("------------------------------------------------\n")
        print("<<< K_KL_estimated_by_minDelta $(obj.K_KL_estimated_by_minDelta)\n")
        print("<<< KL(minDelta) correct detection rate $(obj.corrDetecRate_KL_by_minDelta)\n")
        print("<<< deviation from ground truth $(obj.K_ground_truth-obj.K_KL_estimated_by_minDelta)\n")
        print("------------------------------------------------\n")
        print("<<< K_KL_estimated_by_95quantile $(obj.K_KL_estimated_by_95quantile)\n")
        print("<<< KL(95quantile) correct detection rate $(obj.corrDetecRate_KL_by_95quantile)\n")
        print("<<< deviation from ground truth $(obj.K_ground_truth-obj.K_KL_estimated_by_95quantile)\n")
        print("------------------------------------------------\n")
        # print("<<< K_GAP_estimated $(obj.K_GAP_estimated)\n")
        # print("<<< GAP correct detection rate $(obj.corrDetecRate_GAP)\n")
        # print("<<< deviation from ground truth $(obj.K_ground_truth-obj.K_GAP_estimated)\n")
        # print("<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<\n")
        obj.locID_DF=map(k->DataFrame(  mol_ID=obj.wModel[k].mol_ID,
                                        est_ID=obj.wModel[k].member[end-getproperty(obj, select_mode[modINPUT])[k]+1],
                                        x=obj.wModel[k].loc[1][:,1],
                                        y=obj.wModel[k].loc[1][:,2],
                                        # xC=vcat(map(i->repeat([obj.wModel[k].loc[end-obj.K_KL_estimated_by_95quantile[k]+1][i,:][1]],sum(obj.wModel[k].member[end-obj.K_KL_estimated_by_95quantile[k]+1].==unique(obj.wModel[k].member[end-obj.K_KL_estimated_by_95quantile[k]+1])[i])),1:length(unique(obj.wModel[k].member[end-obj.K_KL_estimated_by_95quantile[k]+1])))...),
                                        # yC=vcat(map(i->repeat([obj.wModel[k].loc[end-obj.K_KL_estimated_by_95quantile[k]+1][i,:][2]],sum(obj.wModel[k].member[end-obj.K_KL_estimated_by_95quantile[k]+1].==unique(obj.wModel[k].member[end-obj.K_KL_estimated_by_95quantile[k]+1])[i])),1:length(unique(obj.wModel[k].member[end-obj.K_KL_estimated_by_95quantile[k]+1])))...),
                                        x_clus_ID=vcat(map(i->obj.wModel[k].loc[end-getproperty(obj, select_mode[modINPUT])[k]+1][:,1][(1:length(obj.wModel[k].loc[end-getproperty(obj, select_mode[modINPUT])[k]+1][:,1]))[obj.wModel[k].clus_ID[end-getproperty(obj, select_mode[modINPUT])[k]+1].==i]],obj.wModel[k].member[end-getproperty(obj, select_mode[modINPUT])[k]+1])...),
                                        y_clus_ID=vcat(map(i->obj.wModel[k].loc[end-getproperty(obj, select_mode[modINPUT])[k]+1][:,2][(1:length(obj.wModel[k].loc[end-getproperty(obj, select_mode[modINPUT])[k]+1][:,2]))[obj.wModel[k].clus_ID[end-getproperty(obj, select_mode[modINPUT])[k]+1].==i]],obj.wModel[k].member[end-getproperty(obj, select_mode[modINPUT])[k]+1])...),
                                        x_mol_ID=vcat(map(i->mean(obj.wModel[k].loc[1][:,1][obj.wModel[k].mol_ID.==i]),obj.wModel[k].mol_ID)...),
                                        y_mol_ID=vcat(map(i->mean(obj.wModel[k].loc[1][:,2][obj.wModel[k].mol_ID.==i]),obj.wModel[k].mol_ID)...),
                                        intra_dist_clus_ID=map(j->pairwise(Euclidean(), transpose(hcat( obj.wModel[k].loc[1][:,1],obj.wModel[k].loc[1][:,2])),
                                                                                transpose(hcat( vcat(map(i->obj.wModel[k].loc[end-getproperty(obj, select_mode[modINPUT])[k]+1][:,1][(1:length(obj.wModel[k].loc[end-getproperty(obj, select_mode[modINPUT])[k]+1][:,1]))[obj.wModel[k].clus_ID[end-getproperty(obj, select_mode[modINPUT])[k]+1].==i]],obj.wModel[k].member[end-getproperty(obj, select_mode[modINPUT])[k]+1])...),
                                                                                vcat(map(i->obj.wModel[k].loc[end-getproperty(obj, select_mode[modINPUT])[k]+1][:,2][(1:length(obj.wModel[k].loc[end-getproperty(obj, select_mode[modINPUT])[k]+1][:,2]))[obj.wModel[k].clus_ID[end-getproperty(obj, select_mode[modINPUT])[k]+1].==i]],obj.wModel[k].member[end-getproperty(obj, select_mode[modINPUT])[k]+1])...))),
                                                        dims=2)[j,j],1:length(obj.wModel[k].loc[1][:,1])),
                                        ),1:length(obj._selectIDX))

        obj.evo_DF=map(k->DataFrame(    K=map(i->length(unique(obj.wModel[k].member[i])),1:length(obj.wModel[k].member)),
                                        median_intra_dist_clus_ID=map(l->
                                                                    median(map(j->
                                                                        pairwise(Euclidean(), transpose(hcat( obj.wModel[k].loc[1][:,1],obj.wModel[k].loc[1][:,2])),
                                                                                transpose(hcat( vcat(map(i->obj.wModel[k].loc[l][:,1][(1:length(obj.wModel[k].loc[l][:,1]))[obj.wModel[k].clus_ID[l].==i]],obj.wModel[k].member[l])...),
                                                                                vcat(map(i->obj.wModel[k].loc[l][:,2][(1:length(obj.wModel[k].loc[l][:,2]))[obj.wModel[k].clus_ID[l].==i]],obj.wModel[k].member[l])...))),
                                                                                dims=2)[j,j],
                                                                            1:length(obj.wModel[k].loc[1][:,1]))),
                                                                    1:length(obj.wModel[k].loc))
                                        ),1:length(obj._selectIDX))

        obj.error_DF=DataFrame(     idx=obj._selectIDX,
                                    clus_vs_mol_MSE=map(k->(1/length(obj.locID_DF[k].x_clus_ID)) * sum( (obj.locID_DF[k].x_clus_ID-obj.locID_DF[k].x_mol_ID).^2 + (obj.locID_DF[k].y_clus_ID-obj.locID_DF[k].y_mol_ID).^2 ),1:length(obj.locID_DF)),
                                    clus_vs_mol_MAE=map(k->(1/length(obj.locID_DF[k].x_clus_ID)) * sum(sqrt.( (obj.locID_DF[k].x_clus_ID-obj.locID_DF[k].x_mol_ID).^2 + (obj.locID_DF[k].y_clus_ID-obj.locID_DF[k].y_mol_ID).^2 )),1:length(obj.locID_DF))
        )
        print("------------------------------------------------\n")
        print("<<< mean clus_vs_mol_MSE $(mean(obj.error_DF.clus_vs_mol_MSE))\n")
        print("<<< median clus_vs_mol_MSE $(median(obj.error_DF.clus_vs_mol_MSE))\n")
        print("<<< sample standard deviation $(Statistics.std(obj.error_DF.clus_vs_mol_MSE))\n")
        print("------------------------------------------------\n")
        return obj
    end
end
export EvalSR

# map(i-> EVAL.wModel[1].loc[end-EVAL.K_KL_estimated_by_95quantile[1]+1][:,1][collect(1:length(EVAL.wModel[1].loc[end-EVAL.K_KL_estimated_by_95quantile[1]+1][:,1]))[EVAL.wModel[1].clus_ID[end-EVAL.K_KL_estimated_by_95quantile[1]+1].==i]],EVAL.wModel[1].member[end-EVAL.K_KL_estimated_by_95quantile[1]+1])

max_stars_repo_name: PGSch/SuperResolution

max_issues_repo_name: PGSch/SuperResolution

max_forks_repo_name: PGSch/SuperResolution


text:
function hessexpr(ex_f::Expr, xs::Symbol)
    ex_g = gradexpr(ex_f, xs)
    ex_h = gradexpr(ex_g, xs)
    return ex_h
end

function hessexpr(ex_f::Expr, xs::Vector{Symbol})
    n = length(xs)
    exs_g = gradexpr(ex_f, xs)
    exs_h = Array(Any, n, n)
    for i in 1:n
        for j in 1:n
            exs_h[i, j] = gradexpr(exs_g[i], xs[j])
        end
    end
    return exs_h
end


text:
using UnicodePlots, Dates

shiftprintln(shift) = x -> print("\u001b[$(shift)G\u001b[0K$(x)\u001b[1B")

"""
    show_dashboard(current_epoch, n_of_epochs, start_time, accuracy_curve,
	 	target_accuracy, best_accuracy=max(accuracy_curve...);
	 	height=min(20, displaysize(stdout)[1]),
		width=min(80, displaysize(stdout)[2]))

Show information about the current state of the training process.
"""
function show_dashboard(current_epoch, n_of_epochs, start_time,
	accuracy_curve, target_accuracy, best_accuracy=max(accuracy_curve...);
	height=min(20, displaysize(stdout)[1]), width=min(80, displaysize(stdout)[2]))

	plot_height = height - 3
	plot_width = width - 37
	elapsed_time = now() - start_time
	estimated_time = (elapsed_time * (n_of_epochs - current_epoch)) ÷ current_epoch

	p = lineplot(accuracy_curve, xlabel="Epoch", ylabel="Accuracy", color=:green,
		xlim=(1, n_of_epochs), ylim=(0, target_accuracy), width=plot_width, height=plot_height)
	lines!(p, 0, best_accuracy, n_of_epochs, best_accuracy, color=:blue)
	lines!(p, 0, target_accuracy, n_of_epochs, target_accuracy, color=:magenta)

	show(p)
	print("\r\u001b[$(plot_height + 2)A")
	shiftedprintln = shiftprintln(plot_width + 21)

	shiftedprintln("Elapsed Time:")
	shiftedprintln(Dates.CompoundPeriod(elapsed_time) |> canonicalize)
	shiftedprintln("ETA:")
	shiftedprintln(Dates.CompoundPeriod(estimated_time) |> canonicalize)
	print("\u001b[2B")
	shiftedprintln("\u001b[32;1mLast Accuracy:\u001b[0m")
	shiftedprintln("\u001b[32m$(accuracy_curve[end])\u001b[0m")
	shiftedprintln("\u001b[34;1mBest Accuracy:\u001b[0m")
	shiftedprintln("\u001b[34m$(best_accuracy)\u001b[0m")
	shiftedprintln("\u001b[35;1mTarget Accuracy:\u001b[0m")
	shiftedprintln("\u001b[35m$(target_accuracy)\u001b[0m")
	print("\r\u001b[$(plot_height - 9)B\n")
end

max_stars_repo_name: gabriel-ss/music-autotagger

max_issues_repo_name: gabriel-ss/music-autotagger

max_forks_repo_name: gabriel-ss/music-autotagger


text:
#!/usr/bin/julia --color=yes

using ArgParse
using JSON

function bench(f, simulate=false)
    fp = joinpath(JSON_DATA_DIR, string(f, ".json"))
    if !isfile(fp)
        println("Downloading benchmark file...")
        download(DATA_SOURCES[f], fp)
    end
    gc()  # run gc so it doesn't affect benchmarks
    t = if args["parse"]["parse-file"]
        tic()
        JSON.parsefile(fp)
        toq()
    else
        data = readstring(fp)
        tic()
        JSON.Parser.parse(data)
        toq()
    end

    if !simulate
        print_with_color(:yellow, " [Bench$FLAGS] ")
        println(f, " ", t, " seconds")
    end
    t
end

const JSON_DATA_DIR = joinpath(dirname(dirname(@__FILE__)), "data")
const s = ArgParseSettings(description="Benchmark JSON.jl")

const DATA_SOURCES = Dict(
    "canada" => "https://raw.githubusercontent.com/miloyip/nativejson-benchmark/v1.0.0/data/canada.json",
    "citm_catalog" => "https://raw.githubusercontent.com/miloyip/nativejson-benchmark/v1.0.0/data/citm_catalog.json",
    "citylots" => "https://raw.githubusercontent.com/zemirco/sf-city-lots-json/master/citylots.json",
    "twitter" => "https://raw.githubusercontent.com/miloyip/nativejson-benchmark/v1.0.0/data/twitter.json")

@add_arg_table s begin
    "parse"
        action = :command
        help = "Run a JSON parser benchmark"
    "list"
        action = :command
        help = "List available JSON files for use"
end

@add_arg_table s["parse"] begin
    "--include-compile", "-c"
        help = "If set, include the compile time in measurements"
        action = :store_true
    "--parse-file", "-f"
        help = "If set, measure JSON.parsefile, hence including IO time"
        action = :store_true
    "file"
        help = "The JSON file to benchmark (leave out to benchmark all)"
        required = false
end

const args = parse_args(ARGS, s)

if args["%COMMAND%"] == "parse"
    const FLAGS = string(
            args["parse"]["include-compile"] ? "C" : "",
            args["parse"]["parse-file"] ? "F" : "")

    if args["parse"]["file"] ≠ nothing
        const file = args["parse"]["file"]

        if !args["parse"]["include-compile"]
            bench(file, true)
        end
        bench(file)
    else
        times = 1.0
        if args["parse"]["include-compile"]
            error("Option --include-compile can only be used for single file.")
        end
        for k in sort(collect(keys(DATA_SOURCES)))
            bench(k, true)  # warm up compiler
        end
        for k in sort(collect(keys(DATA_SOURCES)))
            times *= bench(k)  # do benchmark
        end
        print_with_color(:yellow, " [Bench$FLAGS] ")
        println("Total (G.M.) ", times^(1/length(DATA_SOURCES)), " seconds")
    end
elseif args["%COMMAND%"] == "list"
    println("Available benchmarks are:")
    for k in sort(collect(keys(DATA_SOURCES)))
        println(" • $k")
    end
end