text:
function gmmresults()
    # example of GMM: draws from N(0,1). We estimate mean and variance.
    y = randn(1000,1)
    # 3 moment conditions implied by drawing from N(0,σ²):
    # mean = 0, variance = constant, skew = 0
    moments = θ -> [y.-θ[1] (y.^2.0).-θ[2] (y.-θ[1]).^3.0]
    # first round consistent
    W = Matrix{Float64}(I,3,3)
    θ = [0.0, 1.0]
    θhat, objvalue, D, ms, converged = gmm(moments, θ, W)
    # second round efficient
    Ωinv = inv(cov(ms))
    gmmresults(moments, θhat, Ωinv, "GMM example, two step")
    # CUE
    gmmresults(moments, θhat)
    return
end

function gmmresults(moments, θ, weight="", title="", names="", efficient=true)
    n,g = size(moments(θ))
    CUE = false
    if weight !="" # if weight provided, use it
        θhat, objvalue, D, ms, converged = gmm(moments, θ, weight)
    else # do CUE
        θhat, objvalue, D, ms, weight, converged = gmm(moments, θ)
        CUE=true
    end
    k,g = size(D)
    # estimate variance
    efficient ? V = inv(D*weight*D')/n : V = V*D*weight*NeweyWest(ms)*weight*D'*V/n
    if names==""
        names = 1:k
        names = names'
    end
    se = sqrt.(diag(V))
    t = θhat ./ se
    p = 2.0 .- 2.0*cdf.(Ref(TDist(n-k)),abs.(t))
    PrintDivider()
    if title !="" printstyled(title, color=:yellow); println() end
    print("GMM Estimation Results    Convergence: ")
    printstyled(converged, color=:green)
    CUE ? message=" (CUE)" : message=" (user weight matrix)"
    printstyled(message, color=:cyan)
    println()
    println("Observations: ", n)
    println("Hansen-Sargan statistic: ", round(n*objvalue, digits=5))
    if g > k
        println("Hansen-Sargan p-value: ", round(1.0 - cdf(Chisq(g-k),n*objvalue), digits=5))
    end
    a =[θhat se t p]
    println("")
    PrintEstimationResults(a, names)
    println()
    PrintDivider()
    return θhat, objvalue, V, D, weight, converged
end

max_stars_repo_name: mcreel/EconometricsNotes

max_issues_repo_name: mcreel/EconometricsNotes

max_forks_repo_name: mcreel/EconometricsNotes


text:
# QuEST_jl/src/base/unitaries.jl
#

function compactUnitary(qureg       ::QuEST_Types.Qureg,
                        targetQubit ::Integer,
                        α           ::Complex{Qreal},
                        β           ::Complex{Qreal})   ::Nothing

    alpha = QuEST_Types.Complex(real(α),imag(α))
    beta  = QuEST_Types.Complex(real(β),imag(β))

    ccall(:compactUnitary, Cvoid,
          (QuEST_Types.Qureg, Cint,        QuEST_Types.Complex, QuEST_Types.Complex),
          qureg,              targetQubit, alpha,               beta)
end

function controlledCompactUnitary(qureg          ::QuEST_Types.Qureg,
                                  controlQubit   ::Integer,
                                  targetQubit    ::Integer,
                                  α              ::Complex{Qreal},
                                  β              ::Complex{Qreal} ) ::Nothing

    alpha = QuEST_Types.Complex(real(α),imag(α))
    beta  = QuEST_Types.Complex(real(β),imag(β))

    ccall(:controlledCompactUnitary,
          Cvoid,
          (QuEST_Types.Qureg, Cint, Cint, QuEST_Types.Complex, QuEST_Types.Complex),
          qureg,
          Cint(controlQubit),
          Cint(targetQubit),
          alpha,
          beta)
    return nothing
end

function controlledMultiQubitUnitary(qureg   ::QuEST_Types.Qureg,
                                     ctrl    ::Integer,
                                     targs   ::Vector{QubitIdx},
                                     u       ::QuEST_Types.ComplexMatrixN) ::Nothing
    @assert 0 ≤ ctrl < qureg.numQubitsRepresented
    @assert all( tidx -> 0 ≤ targs[tidx] < qureg.numQubitsRepresented,
                 1:length(targs) )
    @assert ctrl ∉ targs

    @assert u.numQubits < qureg.numQubitsRepresented

    ccall(:controlledMultiQubitUnitary, Cvoid,
          (QuEST_Types.Qureg, Cint, Ptr{Cint}, Cint,          QuEST_Types.ComplexMatrixN),
          qureg,              ctrl, targs,     length(targs), u)
end

function controlledNot(qureg         ::QuEST_Types.Qureg,
                       controlQubit  ::Integer,
                       targetQubit   ::Integer)         ::Nothing

    @assert 0 ≤ controlQubit < qureg.numQubitsRepresented
    @assert 0 ≤ targetQubit  < qureg.numQubitsRepresented
    @assert controlQubit != targetQubit

    ccall(:controlledNot, Cvoid,
          (QuEST_Types.Qureg, QubitIdx,      QubitIdx),
          qureg,              controlQubit,  targetQubit)
end

function controlledPauliY(qureg         ::QuEST_Types.Qureg,
                          controlQubit  ::Integer,
                          targetQubit   ::Integer)         ::Nothing

    @assert 0 ≤ controlQubit < qureg.numQubitsRepresented
    @assert 0 ≤ targetQubit  < qureg.numQubitsRepresented
    @assert controlQubit != targetQubit

    ccall(:controlledPauliY,  Cvoid,
          (QuEST_Types.Qureg, Cint,          Cint),
          qureg,              controlQubit,  targetQubit)
end

function controlledPhaseFlip(qureg         ::QuEST_Types.Qureg,
                             controlQubit  ::Integer,
                             targetQubit   ::Integer)   ::Nothing

    @assert 0 ≤ controlQubit < qureg.numQubitsRepresented
    @assert 0 ≤ targetQubit  < qureg.numQubitsRepresented
    @assert controlQubit != targetQubit

    ccall(:controlledPhaseFlip, Cvoid,
          (QuEST_Types.Qureg, Cint,         Cint),
          qureg,              controlQubit, targetQubit)
end

function controlledPhaseShift(qureg    ::QuEST_Types.Qureg,
                              idQubit1 ::Integer,
                              idQubit2 ::Integer,
                              angle    ::Qreal)   ::Nothing

    @assert 0 ≤ idQubit1 < qureg.numQubitsRepresented
    @assert 0 ≤ idQubit2  < qureg.numQubitsRepresented
    @assert idQubit1 != idQubit2

    ccall(:controlledPhaseShift,  Cvoid,
          (QuEST_Types.Qureg, Cint,     Cint,      Qreal),
          qureg,              idQubit1, idQubit2,  angle)
end

function controlledRotateAroundAxis(qureg        ::QuEST_Types.Qureg,
                                    controlQubit ::Integer,
                                    targetQubit  ::Integer,
                                    angle        ::Qreal,
                                    axis         ::Union{NTuple{3,Qreal},Vector{Qreal}})  ::Nothing

    @assert 0 ≤ controlQubit < qureg.numQubitsRepresented
    @assert 0 ≤ targetQubit  < qureg.numQubitsRepresented
    @assert controlQubit != targetQubit

    @assert length(axis) == 3

    q_axis = QuEST_Types.Vector(axis[1], axis[2], axis[3])

    ccall(:controlledRotateAroundAxis,  Cvoid,
          (QuEST_Types.Qureg, Cint,           Cint,         Qreal, QuEST_Types.Vector),
          qureg,              controlQubit,   targetQubit,  angle, q_axis)
end

function controlledRotateX(qureg         ::QuEST_Types.Qureg,
                           controlQubit  ::Integer,
                           targetQubit   ::Integer,
                           angle         ::Qreal)  ::Nothing

    @assert 0 ≤ controlQubit < qureg.numQubitsRepresented
    @assert 0 ≤ targetQubit  < qureg.numQubitsRepresented
    @assert controlQubit != targetQubit

    ccall(:controlledRotateX, Cvoid,
          (QuEST_Types.Qureg, Cint,          Cint,          Qreal),
          qureg,              controlQubit,  targetQubit,   angle)
end

function controlledRotateY(qureg         ::QuEST_Types.Qureg,
                           controlQubit  ::Integer,
                           targetQubit   ::Integer,
                           angle         ::Qreal)  ::Nothing

    @assert 0 ≤ controlQubit < qureg.numQubitsRepresented
    @assert 0 ≤ targetQubit  < qureg.numQubitsRepresented
    @assert controlQubit != targetQubit

    ccall(:controlledRotateY,    Cvoid,
          (QuEST_Types.Qureg, Cint,          Cint,        Qreal),
          qureg,              controlQubit,  targetQubit, angle)
end

function controlledRotateZ(qureg         ::QuEST_Types.Qureg,
    controlQubit  ::Integer,
    targetQubit   ::Integer,
    angle         ::Qreal)  ::Nothing

    ccall(:controlledRotateZ,
          Cvoid,
          (QuEST_Types.Qureg, Cint, Cint, Qreal),
          qureg,
          Cint(controlQubit),
          Cint(targetQubit),
          angle)
return nothing
end

function _quest_mtx_2(U ::Matrix{Complex{Qreal}}) ::QuEST_Types.ComplexMatrix2
    @assert size(U) == (2,2)
    u = QuEST_Types.ComplexMatrix2(
        ( (real(U[1,1]), real(U[1,2])), (real(U[2,1]), real(U[2,2])) ),
        ( (imag(U[1,1]), imag(U[1,2])), (imag(U[2,1]), imag(U[2,2])) )  )
    return u
end

function _quest_mtx_4(U ::Matrix{Complex{Qreal}}) ::QuEST_Types.ComplexMatrix4
    @assert size(U) == (4,4)
    u = QuEST_Types.ComplexMatrix4( ( ( real(U[1,1]), real(U[1,2]), real(U[1,3]), real(U[1,4]) ),
                          ( real(U[2,1]), real(U[2,2]), real(U[2,3]), real(U[2,4]) ),
                          ( real(U[3,1]), real(U[3,2]), real(U[3,3]), real(U[3,4]) ),
                          ( real(U[4,1]), real(U[4,2]), real(U[4,3]), real(U[4,4]) ) ),
                        ( ( imag(U[1,1]), imag(U[1,2]), imag(U[1,3]), imag(U[1,4]) ),
                          ( imag(U[2,1]), imag(U[2,2]), imag(U[2,3]), imag(U[2,4]) ),
                          ( imag(U[3,1]), imag(U[3,2]), imag(U[3,3]), imag(U[3,4]) ),
                          ( imag(U[4,1]), imag(U[4,2]), imag(U[4,3]), imag(U[4,4]) ) )  )
end

function controlledTwoQubitUnitary(qureg         ::QuEST_Types.Qureg,
                                   controlQubit  ::Integer,
                                   targetQubit1  ::Integer,
                                   targetQubit2  ::Integer,
                                   U             ::Matrix{Complex{Qreal}}) ::Nothing

    @assert 0 ≤ controlQubit < qureg.numQubitsRepresented
    @assert 0 ≤ targetQubit1 < qureg.numQubitsRepresented
    @assert 0 ≤ targetQubit2 < qureg.numQubitsRepresented
    @assert targetQubit1 != targetQubit2
    @assert controlQubit ∉ [targetQubit1,targetQubit2]

    @assert size(U) == (4, 4)
    u = _quest_mtx_4(U)
    ccall(:controlledTwoQubitUnitary,
          Cvoid,
          (QuEST_Types.Qureg, Cint, Cint, Cint, QuEST_Types.ComplexMatrix4),
          qureg,
          Cint(controlQubit),
          Cint(targetQubit1),
          Cint(targetQubit2),
          u)
return nothing
end

function controlledUnitary(qureg        ::QuEST_Types.Qureg,
                           controlQubit ::Integer,
                           targetQubit  ::Integer,
                           U            ::Matrix{Complex{Qreal}}) ::Nothing

    @assert 0 ≤ controlQubit < qureg.numQubitsRepresented
    @assert 0 ≤ targetQubit  < qureg.numQubitsRepresented
    @assert controlQubit != targetQubit

    u = _quest_mtx_2(U)

    ccall(:controlledUnitary,  Cvoid,
          (QuEST_Types.Qureg, Cint,          Cint,          QuEST_Types.ComplexMatrix2),
          qureg,              controlQubit,  targetQubit,   u)
end

function hadamard(qureg       ::QuEST_Types.Qureg,
                  targetQubit ::Integer) ::Nothing

    ccall(:hadamard, Cvoid,
          (QuEST_Types.Qureg, Cint),
          qureg,              targetQubit)
end

function multiControlledMultiQubitUnitary(qureg  ::QuEST_Types.Qureg,
                                          ctrls  ::Vector{QubitIdx},
                                          targs  ::Vector{QubitIdx},
                                          u      ::QuEST_Types.ComplexMatrixN) ::Nothing

    ccall(:multiControlledMultiQubitUnitary, Cvoid,
          (QuEST_Types.Qureg, Ptr{Cint}, Cint,          Ptr{Cint}, Cint,          QuEST_Types.ComplexMatrixN),
          qureg,              ctrls,     length(ctrls), targs,     length(targs), u)
end

function multiControlledPhaseFlip(qureg         ::QuEST_Types.Qureg,
                                  controlQubits ::Vector{QubitIdx}) ::Nothing

    ccall(:multiControlledPhaseFlip, Cvoid,
          (QuEST_Types.Qureg,     Ptr{Cint},     Cint),
          qureg,                  controlQubits, length(controlQubits))
end

function multiControlledPhaseShift(qureg         ::QuEST_Types.Qureg,
                                   controlQubits ::Vector{QubitIdx},
                                   angle         ::Qreal)         ::Nothing

    ccall(:multiControlledPhaseShift, Cvoid,
          (QuEST_Types.Qureg, Ptr{Cint},     Cint,                  Qreal),
          qureg,              controlQubits, length(controlQubits), angle)
end

function multiControlledTwoQubitUnitary(qureg           ::QuEST_Types.Qureg,
                                        controlQubits   ::Vector{QubitIdx},
                                        targetQubit1    ::Integer,
                                        targetQubit2    ::Integer,
                                        U               ::Matrix{Complex{Qreal}}) ::Nothing

    @assert 0 ≤ length(controlQubits) ≤ qureg.numQubitsRepresented -2
    @assert 0 ≤ targetQubit1 < qureg.numQubitsRepresented
    @assert 0 ≤ targetQubit2 < qureg.numQubitsRepresented
    @assert targetQubit1 != targetQubit2
    @assert targetQubit1 ∉ controlQubits
    @assert targetQubit2 ∉ controlQubits

    u = _quest_mtx_4(U)

    ccall(:multiControlledTwoQubitUnitary, Cvoid,
          (QuEST_Types.Qureg, Ptr{Cint},     Cint,                  Cint,          Cint,           QuEST_Types.ComplexMatrix4),
          qureg,              controlQubits, length(controlQubits), targetQubit1,  targetQubit2,   u)
end

function multiControlledUnitary(qureg         ::QuEST_Types.Qureg,
                                controlQubits ::Vector{QubitIdx},
                                targetQubit   ::Integer,
                                U             ::Matrix{Complex{Qreal}}) ::Nothing

    u = _quest_mtx_2(U)

    ccall(:multiControlledUnitary,  Cvoid,
          (QuEST_Types.Qureg, Ptr{Cint},     Cint,                  Cint,          QuEST_Types.ComplexMatrix2),
          qureg,              controlQubits, length(controlQubits), targetQubit,   u)
end

function multiQubitUnitary(qureg ::QuEST_Types.Qureg,
                           targs ::Vector{QubitIdx},
                           u     ::QuEST_Types.ComplexMatrixN)  ::Nothing

    @assert u.numQubits ≤ qureg.numQubitsRepresented

    ccall(:multiQubitUnitary, Cvoid,
          (QuEST_Types.Qureg, Ptr{Cint}, Cint,           QuEST_Types.ComplexMatrixN),
          qureg,              targs,     length(targs),  u)
end

function multiRotatePauli(qureg         ::QuEST_Types.Qureg,
                          targetQubits  ::Vector{QubitIdx},
                          targetPaulis  ::Vector{QuEST_Types.pauliOpType},
                          angle         ::Qreal)          ::Nothing

    @assert length(targetQubits) == length(targetPaulis)
    #@assert all( σ -> 0 ≤ σ ≤ 3,   targetPaulis )

    ccall(:multiRotatePauli, Cvoid,
          (QuEST_Types.Qureg, Ptr{Cint},     Ptr{QuEST_Types.pauliOpType}, Cint,                 Qreal),
          qureg,              targetQubits,  targetPaulis,                 length(targetPaulis), angle)
end

function multiRotateZ(qureg      ::QuEST_Types.Qureg,
                      qubits     ::Vector{QubitIdx},
                      angle      ::Qreal)         ::Nothing
    ccall(:multiRotateZ, Cvoid,
          (QuEST_Types.Qureg, Ptr{Cint}, Cint,           Qreal),
          qureg,              qubits,    length(qubits), angle)
end

function multiStateControlledUnitary(qureg           ::QuEST_Types.Qureg,
                                     controlQubits   ::Vector{QubitIdx},
                                     controlState    ::Vector{Cint},
                                     targetQubit     ::Integer,
                                     U               ::Matrix{Complex{Qreal}}) ::Nothing

    @assert length(controlQubits) == length(controlState)

    u = _quest_mtx_2(U)

    ccall(:multiStateControlledUnitary, Cvoid,
          (QuEST_Types.Qureg, Ptr{Cint},      Ptr{Cint},    Cint,                  Cint,        QuEST_Types.ComplexMatrix2),
          qureg,              controlQubits,  controlState, length(controlQubits), targetQubit, u)
end

function pauliX(qureg       ::QuEST_Types.Qureg,
                targetQubit ::Integer)         :: Nothing

    ccall(:pauliX, Cvoid,
          (QuEST_Types.Qureg, Cint),
          qureg,              targetQubit)
end

function pauliY(qureg       ::QuEST_Types.Qureg,
                targetQubit ::Integer)         :: Nothing

    ccall(:pauliY, Cvoid,
          (QuEST_Types.Qureg, Cint),
          qureg,              targetQubit)
end
function pauliZ(qureg ::QuEST_Types.Qureg, targetQubit ::Integer):: Nothing

    ccall(:pauliZ,  Cvoid,
          (QuEST_Types.Qureg, Cint),
          qureg,              targetQubit)
end

function phaseShift(qureg       ::QuEST_Types.Qureg,
                    targetQubit ::Integer,
                    angle       ::Qreal) ::Nothing

    ccall(:phaseShift, Cvoid,
          (QuEST_Types.Qureg, Cint,        Qreal),
          qureg,              targetQubit, angle)
end

function rotateAroundAxis(qureg         ::QuEST_Types.Qureg,
                          rotQubit      ::Integer,
                          angle         ::Qreal,
                          axis          ::Union{NTuple{3,Qreal},Vector{Qreal}}) ::Nothing

    @assert length(axis) == 3

    q_axis = QuEST_Types.Vector(axis[1], axis[2], axis[3])

    ccall(:rotateAroundAxis,  Cvoid,
          (QuEST_Types.Qureg, Cint,     Qreal, QuEST_Types.Vector),
          qureg,              rotQubit, angle, q_axis)
end

function rotateX(qureg    ::QuEST_Types.Qureg,
                 rotQubit ::Integer,
                 angle    ::Qreal)           ::Nothing

    ccall(:rotateX,  Cvoid,
          (QuEST_Types.Qureg, Cint,      Qreal),
          qureg,              rotQubit,  angle)
end

function rotateY(qureg    ::QuEST_Types.Qureg,
                 rotQubit ::Integer,
                 angle    ::Qreal)           ::Nothing

    ccall(:rotateY, Cvoid,
          (QuEST_Types.Qureg, Cint,      Qreal),
          qureg,              rotQubit,  angle)
end

function rotateZ(qureg    ::QuEST_Types.Qureg,
                 rotQubit ::Integer,
                 angle    ::Qreal)           ::Nothing

    ccall(:rotateZ, Cvoid,
          (QuEST_Types.Qureg, Cint,     Qreal),
          qureg,              rotQubit, angle)
end

function sGate(qureg       ::QuEST_Types.Qureg,
               targetQubit ::Integer)         ::Nothing

    ccall(:sGate, Cvoid,
          (QuEST_Types.Qureg, Cint),
          qureg,              targetQubit)
end

function sqrtSwapGate(qureg  ::QuEST_Types.Qureg,
                      qubit1 ::Integer,
                      qubit2 ::Integer)         ::Nothing

    ccall(:sqrtSwapGate, Cvoid,
          (QuEST_Types.Qureg, Cint,   Cint),
          qureg,              qubit1, qubit2)
end

function swapGate(qureg  ::QuEST_Types.Qureg,
                  qubit1 ::Integer,
                  qubit2 ::Integer)         ::Nothing

    ccall(:swapGate,  Cvoid,
          (QuEST_Types.Qureg, Cint,   Cint),
          qureg,              qubit1, qubit2)
end

function tGate(qureg       ::QuEST_Types.Qureg,
               targetQubit ::Integer)         ::Nothing

    ccall(:tGate, Cvoid,
          (QuEST_Types.Qureg, Cint),
          qureg,              targetQubit)
end

function twoQubitUnitary(qureg           ::QuEST_Types.Qureg,
                         targetQubit1    ::Integer,
                         targetQubit2    ::Integer,
                         U               ::Matrix{Complex{Qreal}}) ::Nothing

    u = _quest_mtx_4(U)

    ccall(:twoQubitUnitary, Cvoid,
          (QuEST_Types.Qureg, Cint,          Cint,          QuEST_Types.ComplexMatrix4),
          qureg,              targetQubit1,  targetQubit2,  u)
end

function unitary(qureg           ::QuEST_Types.Qureg,
                 targetQubit     ::Integer,
                 U               ::Matrix{Complex{Qreal}}) ::Nothing

    u = _quest_mtx_2(U)

    ccall(:unitary, Cvoid,
          (QuEST_Types.Qureg, Cint,           QuEST_Types.ComplexMatrix2),
          qureg,              targetQubit,    u)
end

#EOF

ext: jl


text:
module CompAlgebra

include("AlgebraicNets.jl")
include("MathFormulas.jl")

end


text:
#=
# TODO: Infer dimensionality of user function from context type
mutable struct RevRelaxMeta{T} end
Cassette.@context RelaxCtx

macro scalar_relax_rule(f, rev)
  esc(quote
    function Cassette.overdub(::RelaxCtx{RelaxMeta{T}}, ::typeof($f), x::VRev{T,F}) where {T<:Number,F}
        r = y -> x.rev(_val(x) ∩ $rev)
        return VRev{T,typeof(r)}(($f)(x), r)
    end
  end)
end

@norev_scalar(RelaxCtx, identity)
@norev_scalar(RelaxCtx, one)
@norev_scalar(RelaxCtx, zero)
@norev_scalar(RelaxCtx, transpose)

@scalar_set_rule(sqrt, y^2)
@scalar_set_rule(real, y)

@scalar_set_rule(acos, cos(y))
@scalar_set_rule(acot, cot(y))
@scalar_set_rule(acsc, csc(y))
@scalar_set_rule(asec, sec(y))
@scalar_set_rule(asin, sin(y))
@scalar_set_rule(atan, tan(y))

@scalar_set_rule(acosd, cosd(y))
@scalar_set_rule(acotd, cotd(y))
@scalar_set_rule(acscd, cscd(y))
@scalar_set_rule(asecd, secd(y))
@scalar_set_rule(asind, sind(y))
@scalar_set_rule(atand, tand(y))

@scalar_set_rule(asinh, sinh(y))
@scalar_set_rule(acosh, cosh(y))
@scalar_set_rule(atanh, tanh(y))
@scalar_set_rule(asech, sech(y))
@scalar_set_rule(acsch, csch(y))
@scalar_set_rule(acoth, coth(y))

@scalar_set_rule(sinh, asinh(y))
@scalar_set_rule(tanh, atanh(y))

# has point discontinuity (& unbounded)
@scalar_set_rule(csch, acsch(y))
@scalar_set_rule(coth, acoth(y))

function Cassette.overdub(::RelaxCtx{RelaxMeta{T}}, ::typeof(zeros), n::Int) where {T<:Number}
    return zeros(VRev(T), n)
end
function Cassette.overdub(::RelaxCtx{RelaxMeta{T}}, ::typeof(zeros), dims...) where {T<:Number}
    return zeros(VRev(T), dims...)
end

function Cassette.overdub(ctx::RelaxCtx{RelaxMeta{T}}, ::typeof(hcat), A...) where {T<:Number}
    vA = hcat(A...)
    sz = size(vA)
    vR = zeros(VRev, sz...)
    vR[:] = vA[:]
    return vR
end
function Cassette.overdub(ctx::RelaxCtx{RelaxMeta{T}}, ::typeof(vcat), A...) where {T<:Number}
    vA = vcat(A...)
    sz = size(vA)
    vR = zeros(VRev, sz...)
    vR[:] = vA[:]
    return vR
end

# Complex(), Real(), hypot, fma, muladd, rem2pi, mod, deg2rad, rad2deg, ldexp

#=

@scalar_rule cotd(x) -(π / oftype(x, 180)) * (1 + Ω ^ 2)
@scalar_rule cscd(x) -(π / oftype(x, 180)) * Ω * cotd(x)
@scalar_rule csch(x) -(coth(x)) * Ω
@scalar_rule sec(x) Ω * tan(x)
@scalar_rule secd(x) (π / oftype(x, 180)) * Ω * tand(x)
@scalar_rule sech(x) -(tanh(x)) * Ω

@scalar_rule acot(x) -(inv(1 + x ^ 2))
@scalar_rule acsc(x) -(inv(x ^ 2 * sqrt(1 - x ^ -2)))
@scalar_rule acsc(x::Real) -(inv(abs(x) * sqrt(x ^ 2 - 1)))
@scalar_rule asec(x) inv(x ^ 2 * sqrt(1 - x ^ -2))
@scalar_rule asec(x::Real) inv(abs(x) * sqrt(x ^ 2 - 1))

@scalar_rule cosd(x) -(π / oftype(x, 180)) * sind(x)
@scalar_rule cospi(x) -π * sinpi(x)
@scalar_rule sind(x) (π / oftype(x, 180)) * cosd(x)
@scalar_rule sinpi(x) π * cospi(x)
@scalar_rule tand(x) (π / oftype(x, 180)) * (1 + Ω ^ 2)

@scalar_rule sinc(x) cosc(x)

@scalar_rule round(x) zero(x)
@scalar_rule floor(x) zero(x)
@scalar_rule ceil(x) zero(x)

=#

#@scalar_set_rule(imag(x::Real) Zero()
=#

max_stars_repo_name: PSORLab/RSActivationFunctions

max_issues_repo_name: PSORLab/RSActivationFunctions

max_forks_repo_name: PSORLab/RSActivationFunctions


text:
# ------------------------------------------------------------------
# Licensed under the MIT License. See LICENSE in the project root.
# ------------------------------------------------------------------

"""
    UniformWeighting()

A weighting method that assigns uniform weights to points in spatial object.
"""
struct UniformWeighting <: WeightingMethod end

weight(object, method::UniformWeighting) =
  weight(domain(object), method)

weight(domain::Domain, ::UniformWeighting) =
  GeoWeights(domain, ones(nelements(domain)))


text:
import DelayEmbeddings:
    Dataset

import StateSpaceReconstruction:
    Embeddings,
    minima_and_stepsizes

import StaticArrays:
    SVector,
    MVector

import ..TransferOperators:
    TransferOperatorEstimatorRectangularBinning

import Base.Cartesian

"""
    bin_origins_subdivided_boxcovering(axisminima, stepsizes,
            n_intervals_eachaxis) -> Array{Float64, 2}

Consider a (single) hyperrectangular box covering a point cloud. We're given the
minimal along each coordinate axis (`axisminima`), a set of `stepsizes` along
each axis (which might be different for each axis), and a set of
intervals ()`n_intervals_per_axis`) to which divide each axis into.

Together, this information gives a complete box covering of the points. This
function returns the origins of all the boxes that the original (single) big
hyperrectangle gets subdivided into, given the `stepsizes` and
`n_intervals_per_axis`.

The origins are returned as an array where each column is the origin of one
box.
"""
function bin_origins_subdivided_boxcovering(axisminima, stepsizes, n_intervals_eachaxis)
    D = length(axisminima)
    coords = zeros(Float64, D, prod(n_intervals_eachaxis))

    n_origins_generated = 0
    for I in Iterators.product([1:i for i in n_intervals_eachaxis]...,)
        n_origins_generated += 1
        for dim in 1:D
            coords[dim, n_origins_generated] = axisminima[dim] + stepsizes[dim]*(I[dim] - 1)
        end
    end
    coords
end

export bin_origins_subdivided_boxcovering

"""
    boxorigins(::StateSpaceReconstruction.AbstractEmbedding, ϵ) -> Array{Float64, 2}

Consider the box covering of the embedding induced by the binning scheme `ϵ`.
This functions returns the origin of each of those bins (always counting
from minima to maxima along each axis, with boxes extending 1% past the extrema
in both directions along each axis).
"""
function boxorigins end

function boxorigins(points::AbstractArray{T, 2}, ϵ) where T
    # Make sure that the array contains points as columns.
    if size(points, 1) > size(points, 2)
        throw(DomainError("The dimension of the dataset exceeds the number of points."))
    end

    D = size(points, 1)
    n_pts = size(points, 2)

    axisminima = zeros(Float64, D)
    top = zeros(Float64, D)

    for i = 1:D
        axisminima[i] = minimum(points[i, :])
        top[i] = maximum(points[i, :])
    end
    axisminima = axisminima - (top - axisminima) / 100
    top = top + (top - axisminima) / 100

    stepsizes = Vector{Float64}(undef, D)
    if typeof(ϵ) <: Float64
        stepsizes = [ϵ for i in 1:D]
    elseif typeof(ϵ) == Vector{Float64}
        stepsizes = ϵ
    elseif typeof(ϵ) <: Int
        stepsizes = (top - axisminima) / ϵ
    elseif typeof(ϵ) == Vector{Int}
        stepsizes = (top - axisminima) ./ ϵ
    end
    n_intervals_eachaxis = floor.(Int, (top .- axisminima) ./ stepsizes) .+ 1

    bin_origins_subdivided_boxcovering(axisminima,
        stepsizes,
        n_intervals_eachaxis)
end

function boxorigins(E::Embeddings.AbstractEmbedding, ϵ)
    boxorigins(E.points, ϵ)
end

function boxorigins(D::Dataset, ϵ)
    boxorigins(transpose(Matrix(D)), ϵ)
end

function boxorigins(pts::Vector{Vector{T}}, ϵ) where T
    c = boxorigins(hcat(pts...,), ϵ)
    [c[:, i] for i = 1:maximum(size(c))]
end

function boxorigins(pts::Vector{SVector{D, T}}, ϵ) where {D, T}
    c = boxorigins(Array(hcat(pts...,)), ϵ)
    [c[:, i] for i = 1:maximum(size(c))]
end

function boxorigins(pts::Vector{MVector{D, T}}, ϵ) where {D, T}
    c = boxorigins(Array(hcat(pts...,)), ϵ)
    [c[:, i] for i = 1:maximum(size(c))]
end

export boxorigins

######################################################
# Some function barriers to reduce memory allocations
######################################################
function fill_edges!(edges, a::Float64, b::Float64, c::Float64, d::Float64)
    edges[1] = a
    edges[2] = b
    edges[3] = c
    edges[4] = d
end

function find_lengths_intersections!(sorted_edges, edges, permutation,
        temp, tmp_reset, T, intersecting_lengths, c)

    # partialsort reduces allocations by half
    partialsortperm!(permutation, edges, 1:2, rev = true)
    sort!(edges, rev = true)

    T .= temp[permutation]

    if T[1] != T[2]
        intersecting_lengths[c] = edges[2] - edges[3]
    else
        intersecting_lengths[c] = 0.0
    end

    temp .= tmp_reset
end

######################################################
"""
    μ_of_bₐ_induced_by_bins_from_binningscheme_ϵⱼ(
        ϵF::Vector{Float64},
        bin_origin_bₐ::Vector{Float64},
        ϵⱼ::Vector{Float64},
        binorigins_visitedbins_ϵⱼ::Array{Float64, 2},
        μϵⱼ::Vector{Float64}) -> Float64

## Returns

The measure ``\\mu(b_a)`` of the bin ``b_a`` in the partition
given by the binning scheme `ϵF` induced by the elements of another
partition given by the binning scheme `ϵⱼ` overlapping with ``b_a``.

## Explanation

Consider a rectangular binning ``B_{\\epsilon_{F}}`` of some points, given by the
partition scheme `ϵF`. Next, consider a separate rectanglar binning
``B_{\\epsilon_{j}}`` of the same points, but formed from the partition scheme `ϵⱼ`.

Both partitions are now made up of rectangular boxes, but differ in that the
boxes have different sizes. Thus, an invariant measure computed over
``B_{\\epsilon_{F}}`` will in general be different from an invariant measure
computed over ``B_{\\epsilon_{j}}`` (the first partition may be coarse than
the second, for example). However, for a given box ``b_a \\in B_{\\epsilon_{F}}``,
there will always be some boxes ``b_k \\in B_{\\epsilon_{j}}`` overlapping with it.
Each of the ``b_k \\in B_{\\epsilon_{j}}`` have a measure ``\\mu_k`` associated
with it. We want to know what the measure of ``b_a`` would be from the perspective
of ``B_{\\epsilon_{j}}``.

This function finds the measure of the target bin ``b_a \\in B_{\\epsilon_{F}}``
induced by the bins ``\\{b_k \\in B_{\\epsilon_{j}} : b_k \\cap b_a \\neq \\emptyset \\}``
(that is, by the ``b_k`` having have some intersection with ``b_a``).

## Arguments

- `ϵF`: Edge lengths of boxes in the ``B_{\\epsilon_{F}}`` rectangular partition.
- `bin_origin_bₐ`: Origin of the a-th bin in ``B_{\\epsilon_{F}}``.
- `ϵⱼ`: Edge lengths of boxes in the ``B_{\\epsilon_{j}}`` rectangular partition.
- `binorigins_visitedbins_ϵⱼ`: The origin of each visited bin in
    ``P_{\\epsilon_{j}}``. Array where columns are the coordinates of the
    bin origins.
- `μϵⱼ`: the measures of the visited bins in ``B_{\\epsilon_{j}}``.
"""
function μ_of_bₐ_induced_by_bins_from_binningscheme_ϵⱼ(
            ϵF::Vector{Float64},
            bin_origin_bₐ::Vector{Float64},
            ϵⱼ::Vector{Float64},
            binorigins_visitedbins_ϵⱼ::Array{Float64, 2},
            μϵⱼ::Vector{Float64})

    dim = length(ϵF)
    volume_ϵⱼ = prod(ϵⱼ)
    n_visited_bins_ϵⱼ = length(μϵⱼ)
    binning_target_bin1 = falses(n_visited_bins_ϵⱼ)
    binning_target_bin2 = zeros(n_visited_bins_ϵⱼ)
    edges = zeros(Float64, 4)
    sorted_edges = zeros(Float64, 4)
    permutation = zeros(Int, 4)
    temp::Vector{Int} = [1, 1, 2, 2]
    tmp_reset::Vector{Int} = [1, 1, 2, 2]
    T::Vector{Number} = zeros(Number, size(temp))
    intersecting_lengths = zeros(Float64, dim)

    for b = 1:n_visited_bins_ϵⱼ
        for c = 1:dim
            fill_edges!(edges,
               bin_origin_bₐ[c] + ϵF[c],
               bin_origin_bₐ[c],
               binorigins_visitedbins_ϵⱼ[c, b] + ϵⱼ[c],
               binorigins_visitedbins_ϵⱼ[c, b])

            find_lengths_intersections!(
                sorted_edges, edges, permutation,
                temp, tmp_reset, T,
                intersecting_lengths, c)
        end

        intersecting_volume = prod(intersecting_lengths)
        if intersecting_volume > 0
            binning_target_bin1[b] = true
            binning_target_bin2[b] = intersecting_volume
        end
    end
    ρₐ = transpose(μϵⱼ[binning_target_bin1]) *
        binning_target_bin2[binning_target_bin1]/volume_ϵⱼ
end

export μ_of_bₐ_induced_by_bins_from_binningscheme_ϵⱼ

"""
    μ_allbins_ϵF_induced_by_binningscheme_ϵⱼ(ϵF,
        binorigins_ϵF::Array{Float64, 2},
        ϵⱼ,
        binorigins_visitedbins_ϵⱼ::Array{Float64, 2},
        μϵⱼ::Vector{Float64})  -> Vector{Float64}

## Returns
The measure ``\\mu(b_a)`` of the bin ``b_a`` in the partition
given by the binning scheme `ϵF` induced by the elements of another
partition given by the binning scheme `ϵⱼ` overlapping with ``b_a``.

## Explanation
Consider a rectangular binning ``B_{\\epsilon_{F}}`` of some points, given by the
partition scheme `ϵF`. Next, consider a separate rectanglar binning
``B_{\\epsilon_{j}}`` of the same points, but formed from the partition scheme `ϵⱼ`.

Both partitions are now made up of rectangular boxes, but differ in that the
boxes have different sizes. Thus, an invariant measure computed over
``B_{\\epsilon_{F}}`` will in general be different from an invariant measure
computed over ``B_{\\epsilon_{j}}`` (the first partition may be coarse than
the second, for example). However, for a given box ``b_a \\in B_{\\epsilon_{F}}``,
there will always be some boxes ``b_k \\in B_{\\epsilon_{j}}`` overlapping with it.
Each of the ``b_k \\in B_{\\epsilon_{j}}`` have a measure ``\\mu_k`` associated
with it. We want to know what the measure of ``b_a`` would be from the perspective
of ``B_{\\epsilon_{j}}``.

This function finds the measure of the target bin ``b_a \\in B_{\\epsilon_{F}}``
induced by the bins ``\\{b_k \\in B_{\\epsilon_{j}} : b_k \\cap b_a \\neq \\emptyset \\}``
(that is, by the ``b_k`` having have some intersection with ``b_a``).

## Arguments

- `ϵF`: Edge lengths along each coordinate axis in ``B_{\\epsilon_{F}}``.
- `binorigins_ϵF`: Coordinates of all bins ``b_i \\in B_{\\epsilon_{F}}``.
- `ϵⱼ`: Edge lengths along each coordinate axis in ``B_{\\epsilon_{j}}``.
- `binorigins_visitedbins_ϵⱼ`: Coordinates of the origin of each visited bin
    in ``P_{\\epsilon_{j}}``. Each column of this array is a bin origin.
- `μϵⱼ`: the measures of the visited bins in ``B_{\\epsilon_{j}}``.
"""
function μ_allbins_ϵF_induced_by_binningscheme_ϵⱼ(
            points,
            ϵF,
            binorigins_ϵF::Array{Float64, 2},
            ϵⱼ,
            binorigins_visitedbins_ϵⱼ::Array{Float64, 2},
            μϵⱼ::Vector{Float64})

    ϵF = minima_and_stepsizes(points, ϵF)[2]
    ϵⱼ = minima_and_stepsizes(points, ϵⱼ)[2]
    n = size(binorigins_ϵF, 2)
    ρ = zeros(Float64, n)

    # Loop over the bins bᵢ in the partition induced by the binning scheme ϵF,
    # and find the measure of those bins in terms of the measure of the
    # bins of the other partition.
    for i = 1:n
        bₐ = binorigins_ϵF[:, i]
        ρ[i] = μ_of_bₐ_induced_by_bins_from_binningscheme_ϵⱼ(
                    ϵF,
                    bₐ,
                    ϵⱼ,
                    binorigins_visitedbins_ϵⱼ, # only need to check the visited bins in the other partition
                    μϵⱼ)
    end

    return ρ
end

export μ_allbins_ϵF_induced_by_binningscheme_ϵⱼ

"""
    μϵF_induced_by_ϵj(points, ϵF, ϵⱼ)

## Returns
The measures ``\\{\\mu(b_f)\\}`` of the bins ``b_f \\in B_{\\epsilon_{F}}``
(``B_{\\epsilon_{F}}`` is a partition given by the binning scheme `ϵF`)
induced by the bins ``b_j \\in B_{\\epsilon_{j}}`` (``B_{\\epsilon_{j}}``
is a partition given by the binning scheme `ϵⱼ`) which have different
measures ``\\{\\mu(b_j)\\}`` associated with them. The partitions may
have a different number of elements, and
``\\{\\mu(b_f)\\} = \\{\\mu(b_j)\\}`` only if
``B_{\\epsilon_{F}} = B_{\\epsilon_{j}}``.

This gives us a way of comparing the measures obtained at different
resolutions of a system.

## Explanation

Consider a cloud of `points` and two separate binning schemes, `ϵF` and `ϵⱼ`.
The binning schemes dictate how to create a rectangular partition covering the
points.

Create the partition ``B_{\\epsilon_{F}}`` from `ϵF` and a separate partition
``B_{\\epsilon_{j}}`` from `ϵⱼ`. Then, compute an invariant measure over each
partition separately.

The boxes in ``B_{\\epsilon_{F}}`` are have sizes different from
the boxes in ``B_{\\epsilon_{j}}``. Thus, the invariant measure associated
with each of the partitions convey something about the dynamics at different
scales. We may want to know how the dynamics exposed by
``B_{\\epsilon_{j}}`` is expressed in terms of another coarse-graining
``B_{\\epsilon_{F}}``.

This function computes the measure of the bins ``b_f \\in B_{\\epsilon_{F}}``
induced by the bins ``b_j \\in B_{\\epsilon_{j}}``. That is, for each ``b_f``,
find the set of bins
``\\{b_k \\in B_{\\epsilon_{j}} : b_k \\cap b_f \\neq \\emptyset \\}``.

Taking this further, we can decide on a final rectangular grid (resolution)
and compute the measure over that grid as an average over multiple separate
partitions (finer or coarser, depending on the specific application).

## Arguments
- **`points`**: The points representing observations from the system from
    which to compute an invariant measure.
- **`ϵF`**: Edge lengths along each coordinate axis in the partition ``B_{\\epsilon_{F}}``.
- **`ϵⱼ`**: Edge lengths along each coordinate axis in the partition ``B_{\\epsilon_{j}}``.
"""
function μϵF_induced_by_ϵj(points, ϵF, ϵⱼ)
    ϵF = minima_and_stepsizes(points, ϵF)[2]
    ϵⱼ = minima_and_stepsizes(points, ϵⱼ)[2]

    # All possible bins in the final partition given by ϵF
    allbins_ϵF = boxorigins(points, ϵF)

    # The unique bins visited by the orbit in the partition given by ϵⱼ
    visitedbins_ϵⱼ = unique(assign_coordinate_labels(points, ϵⱼ), dims = 2)

    # The transfer operator in the partition given by ϵⱼ
    TOϵⱼ = TransferOperatorEstimatorRectangularBinning(points, ϵⱼ)

    # The measure of the of the visited bins in the partition given by ϵⱼ
    μϵⱼ = invariantmeasure(TOϵⱼ).dist

    # The induced measure at resolution ϵF. This does not filter
    # zero entries, so be careful when comparing with the distribution
    # obtained using invariantmeasure(::TransferOperatorRectangularBinning)
    μϵF = μ_allbins_ϵF_induced_by_binningscheme_ϵⱼ(
        points,
        ϵF,
        allbins_ϵF,
        ϵⱼ,
        visitedbins_ϵⱼ,
        μϵⱼ)
end

μϵF_induced_by_ϵj(E::Embeddings.AbstractEmbedding, ϵF, ϵⱼ) =
    μϵF_induced_by_ϵj(E.points, ϵF, ϵⱼ)

export μϵF_induced_by_ϵj

"""
    induced_measure(points, ϵF, ϵⱼ)

Express the invariant measure of a rectangular box covering of `points`, at
the resolution given by the binning scheme `ϵⱼ` as the invariant measure at
a separate resolution given by the binning scheme `ϵF`.

See the documentation for [`μϵF_induced_by_ϵⱼ`](@ref) for more details.

## Arguments
- **`points`**: The points representing observations from the system from
    which to compute an invariant measure.
- **`ϵF`**: Edge lengths along each coordinate axis in the partition ``B_{\\epsilon_{F}}``.
- **`ϵⱼ`**: Edge lengths along each coordinate axis in the partition ``B_{\\epsilon_{j}}``.
"""
function induced_measure(points, ϵF, ϵⱼ)
    μϵF_induced_by_ϵj(points,  ϵF, ϵⱼ)
end

export induced_measure


text:
using LAMMPS

function lammpspotential(x, t, a, b, c, pbc, unitstyle, pair_style, pair_coeff)

    lmp = LMP()

    # handle boundary conditions
    bcs = ""
    for i = 1:3
        if pbc[i] == 1
            bcs = bcs * " p"
        else
            bcs = bcs * " f"
        end
    end

    command(lmp, "units " * unitstyle)
    command(lmp, "boundary " * bcs)
    command(lmp, "atom_style atomic")
    command(lmp, "atom_modify map array")
    command(lmp, "box tilt large")

    # create simulation box
    xlo = 0.0; ylo = 0.0; zlo = 0.0;
    xhi = a[1]; yhi = b[2]; zhi = c[3];
    xy = b[1]; xz = c[1]; yz = c[2];
    if (abs(xy) + abs(xz) + abs(yz)) <= 1e-10
        command(lmp, "region simbox block $xlo $xhi $ylo $yhi $zlo $zhi units box")
    else
        command(lmp, "region simbox prism $xlo $xhi $ylo $yhi $zlo $zhi $xy $xz $yz units box")
    end
    command(lmp, "create_box 1 simbox")

    # pair potential
    command(lmp, "pair_style " * pair_style)
    for i = 1:length(pair_coeff)
        command(lmp, "pair_coeff " * pair_coeff[i])
    end

    # set mass
    ntypes = length(unique(t))
    for i = 1:ntypes
        itype = Int32(i)
        command(lmp, "mass $itype 1.0")
    end

    # create atom types and atom positions
    n = size(x,2)
    id = Int32.(Array(1:(n)))
    v = 0.0*x
    image = zeros(Int32, n)
    bexpand = 0
    LAMMPS.API.lammps_create_atoms(lmp, n, id, Int32.(t), x, v, image, bexpand)

    # run lammps
    command(lmp, "run 0")

    # potential energy
    etot = LAMMPS.API.lammps_get_thermo(lmp, "pe")
    pe = Float64(etot)*n

    # extract forces
    forces = extract_atom(lmp, "f")
    id = extract_atom(lmp, "id")
    pos = extract_atom(lmp, "x")
    type = extract_atom(lmp, "type")

    return pe, forces, id, pos, type

    # pos = extract_atom(lmp, "x")
    # t = extract_atom(lmp, "type")
    # i = extract_atom(lmp, "id")
    # nlocal = extract_global(lmp, "nlocal")
    # ntypes = extract_global(lmp, "ntypes")
end

max_stars_repo_name: cesmix-mit/MDP

max_issues_repo_name: cesmix-mit/MDP

max_forks_repo_name: cesmix-mit/MDP


text:
# using SymPy

# x, y = @vars x y
# u = x*(1-x)*y*(1-y)
# v = x*(1-x)*y*(1-y)^2

# p = x*(1-x)*y*(1-y)

# ux = diff(u,x)
# uy = diff(u,y)
# vx = diff(v,x)
# vy = diff(v,y)
# px = diff(p,x)
# py = diff(p,y)
# f = u*ux + v*uy - (diff(ux,x)+diff(uy,y)) + px
# g = u*vx + v*vy - (diff(vx,x)+diff(vy,y)) + py
# h = diff(u, x) + diff(v, y)
# println(replace(replace(sympy.julia_code(simplify(f)), ".^"=>"^"), ".*"=>"*"))
# println(replace(replace(sympy.julia_code(simplify(g)), ".^"=>"^"), ".*"=>"*"))
# println(replace(replace(sympy.julia_code(simplify(h)), ".^"=>"^"), ".*"=>"*"))

function u_exact(x,y)
    x*(1-x)*y*(1-y)
end

function v_exact(x,y)
    x*(1-x)*y*(1-y)
end

function p_exact(x,y)
    x*(1-x)*y*(1-y)
end

function t_exact(x,y)
    x*(1-x)*y*(1-y)
end

function k_exact(x,y)
    1 + x^2 + x / (1 + y^2)
end

function ffunc_(x, y)
    x*y*(1 - x)*(1 - y)*(-x*y*(1 - x) + x*(1 - x)*(1 - y)) + x*y*(1 - x)*(1 - y)*(-x*y*(1 - y) + y*(1 - x)*(1 - y)) - x*y*(1 - y) + 0.02*x*(1 - x) + y*(1 - x)*(1 - y) + 0.02*y*(1 - y)
end

function gfunc_(x, y)
    x*y*(1 - x)*(1 - y)*(-x*y*(1 - x) + x*(1 - x)*(1 - y)) + x*y*(1 - x)*(1 - y)*(-x*y*(1 - y) + y*(1 - x)*(1 - y)) - x*y*(1 - x) + x*(1 - x)*(1 - y) + 0.02*x*(1 - x) + 0.02*y*(1 - y)
end

function hfunc_(x,y)
    -x*y*(1 - x) - x*y*(1 - y) + x*(1 - x)*(1 - y) + y*(1 - x)*(1 - y)
end

function heat_source_func(x, y)
    # x^2*y*(x - 1)^2*(y - 1)*(2*y - 1) + x*y^2*(x - 1)*(2*x - 1)*(y - 1)^2 - 2.0*x*(x - 1) - 2.0*y*(y - 1) # for k=1.0
    (2*x^2*y*(x - 1)*(2*y - 1) + x*y*(x - 1)*(y - 1)*(y^2 + 1)^2*(x*(x - 1)*(2*y - 1) + y*(2*x - 1)*(y - 1)) - (y^2 + 1)*(2*x*(x - 1)*(x + (x^2 + 1)*(y^2 + 1)) + 2*y*(x + (x^2 + 1)*(y^2 + 1))*(y - 1) + y*(2*x - 1)*(y - 1)*(2*x*(y^2 + 1) + 1)))/(y^2 + 1)^2
end

using LinearAlgebra
using MAT
using AdFem
using PyPlot
using SparseArrays

m = 20
n = 20
h = 1/n
nu = 0.01

F1 = compute_fem_source_term1(eval_f_on_gauss_pts(ffunc_, m, n, h), m, n, h)
F2 = compute_fem_source_term1(eval_f_on_gauss_pts(gfunc_, m, n, h), m, n, h)
H = h^2*eval_f_on_fvm_pts(hfunc_, m, n, h)
B = constant(compute_interaction_matrix(m, n, h))

# compute F
Laplace = nu * constant(compute_fem_laplace_matrix1(m, n, h))
heat_source = eval_f_on_gauss_pts(heat_source_func, m, n, h)
heat_source = constant(compute_fem_source_term1(heat_source, m, n, h))

kgauss = eval_f_on_gauss_pts(k_exact, m, n, h)
LaplaceK = constant(compute_fem_laplace_matrix1(kgauss, m, n, h))

# The temperature term is
# T_xx + T_yy + q = u T_x + v T_y
#
function compute_residual(S)
    u, v, p, T = S[1:(m+1)*(n+1)],
        S[(m+1)*(n+1)+1:2(m+1)*(n+1)],
        S[2(m+1)*(n+1)+1:2(m+1)*(n+1)+m*n],
        S[2(m+1)*(n+1)+m*n+1:end]
    G = eval_grad_on_gauss_pts([u;v], m, n, h)
    ugauss = fem_to_gauss_points(u, m, n, h)
    vgauss = fem_to_gauss_points(v, m, n, h)
    ux, uy, vx, vy = G[:,1,1], G[:,1,2], G[:,2,1], G[:,2,2]

    interaction = compute_interaction_term(p, m, n, h) # julia kernel needed
    f1 = compute_fem_source_term1(ugauss.*ux, m, n, h)
    f2 = compute_fem_source_term1(vgauss.*uy, m, n, h)
    f3 = -interaction[1:(m+1)*(n+1)]
    f4 = Laplace*u
    f5 = -F1
    F = f1 + f2 + f3 + f4 + f5

    g1 = compute_fem_source_term1(ugauss.*vx, m, n, h)
    g2 = compute_fem_source_term1(vgauss.*vy, m, n, h)
    g3 = -interaction[(m+1)*(n+1)+1:end]
    g4 = Laplace*v
    g5 = -F2
    G = g1 + g2 + g3 + g4 + g5

    H0 = -B * [u;v] + H

    # T_xx + T_yy = u T_x + v T_y - heat_source
    T0 = LaplaceK * T + compute_fem_advection_matrix1(ugauss,vgauss, m, n, h) * T - heat_source

    R = [F;G;H0;T0]
    return R
end

function compute_jacobian(S)
    u, v, p, T = S[1:(m+1)*(n+1)],
        S[(m+1)*(n+1)+1:2(m+1)*(n+1)],
        S[2(m+1)*(n+1)+1:2(m+1)*(n+1)+m*n],
        S[2(m+1)*(n+1)+m*n+1:end]

    G = eval_grad_on_gauss_pts([u;v], m, n, h)
    ugauss = fem_to_gauss_points(u, m, n, h)
    vgauss = fem_to_gauss_points(v, m, n, h)
    ux, uy, vx, vy = G[:,1,1], G[:,1,2], G[:,2,1], G[:,2,2]

    M1 = constant(compute_fem_mass_matrix1(ux, m, n, h))
    M2 = constant(compute_fem_advection_matrix1(constant(ugauss), constant(vgauss), m, n, h)) # a julia kernel needed
    M3 = Laplace
    Fu = M1 + M2 + M3

    Fv = constant(compute_fem_mass_matrix1(uy, m, n, h))

    N1 = constant(compute_fem_mass_matrix1(vy, m, n, h))
    N2 = constant(compute_fem_advection_matrix1(constant(ugauss), constant(vgauss), m, n, h))
    N3 = Laplace
    Gv = N1 + N2 + N3

    Gu = constant(compute_fem_mass_matrix1(vx, m, n, h))

    # M = constant(compute_fem_laplace_matrix1(m, n, h)) +
        #  compute_fem_advection_matrix1(ugauss,vgauss, m, n, h)

    M = LaplaceK + constant(compute_fem_advection_matrix1(ugauss,vgauss, m, n, h))

    gradT = eval_grad_on_gauss_pts1(T, m, n, h)
    Tx, Ty = gradT[:,1], gradT[:,2]
    DU_TX = constant(compute_fem_mass_matrix1(Tx, m, n, h))       # (m+1)*(n+1), (m+1)*(n+1)
    DV_TY = constant(compute_fem_mass_matrix1(Ty, m, n, h))       # (m+1)*(n+1), (m+1)*(n+1)

    J0 = [Fu Fv                         # Jacobian of [u; v]
          Gu Gv]

    J1 = [J0 -B'                        # Jacobian of [u; v; p]
          -B spdiag(zeros(size(B,1)))]

    N = 2*(m+1)*(n+1) + m*n                                     # length of [u; v; p]
    J = [J1 SparseTensor(spzeros(N,(m+1)*(n+1)))                # Jacobian of [u; v; p; T]
        DU_TX DV_TY SparseTensor(spzeros((m+1)*(n+1), m*n)) M]
end

NT = 5    # number of iterations for Newton's method, 10 => 5

bd = bcnode("all", m, n, h)
bd = [bd; bd .+ (m+1)*(n+1);
    2*(m+1)*(n+1)+(n-1)*m+1:2*(m+1)*(n+1)+(n-1)*m+2;
    (2*(m+1)*(n+1)+m*n) .+ bd]
# only apply Dirichlet to velocity; set left bottom two points to zero to fix rank deficient problem for pressure

function solve_steady_cavityflow_one_step(S)
    residual = compute_residual(S)
    J = compute_jacobian(S)

    J, _ = fem_impose_Dirichlet_boundary_condition1(J, bd, m, n, h)
    residual = scatter_update(residual, bd, zeros(length(bd)))    # residual[bd] .= 0.0 in Tensorflow syntax

    d = J\residual
    residual_norm = norm(residual)
    op = tf.print("residual norm", residual_norm)
    d = bind(d, op)
    S_new = S - d
    return S_new
end

function condition(i, S_arr)
    i <= NT + 1
end

function body(i, S_arr)
    S = read(S_arr, i-1)
    op = tf.print("i=",i)
    i = bind(i, op)
    S_new = solve_steady_cavityflow_one_step(S)
    S_arr = write(S_arr, i, S_new)
    return i+1, S_arr
end

xy = fem_nodes(m, n, h)
x, y = xy[:,1], xy[:,2]
u0 = @. u_exact(x,y)
v0 = @. v_exact(x,y)
t0 = @. t_exact(x,y)

xy = fvm_nodes(m, n, h)
x, y = xy[:,1], xy[:,2]
p0 = @. p_exact(x,y)

S_arr = TensorArray(NT+1)
S_arr = write(S_arr, 1, zeros(m*n+3*(m+1)*(n+1)))

i = constant(2, dtype=Int32)

_, S = while_loop(condition, body, [i, S_arr])
S = set_shape(stack(S), (NT+1, 2*(m+1)*(n+1)+m*n+(m+1)*(n+1)))

sess = Session(); init(sess)
output = run(sess, S)

matwrite("CHTcoupled_data.mat",
    Dict(
        "V"=>output[end, :]
    ))

figure(figsize=(25,10))
subplot(241)
title("u velocity")
visualize_scalar_on_fem_points(output[NT+1, 1:(m+1)*(n+1)], m, n, h)
subplot(245)
visualize_scalar_on_fem_points(u0, m, n, h)

subplot(242)
title("v velocity")
visualize_scalar_on_fem_points(output[NT+1, (m+1)*(n+1)+1:2*(m+1)*(n+1)], m, n, h)
subplot(246)
visualize_scalar_on_fem_points(v0, m, n, h)

subplot(243)
visualize_scalar_on_fvm_points(output[NT+1, 2*(m+1)*(n+1)+1:2*(m+1)*(n+1)+m*n], m, n, h)
title("pressure")
subplot(247)
visualize_scalar_on_fvm_points(p0, m, n, h)
title("")

subplot(244)
title("temperature")
visualize_scalar_on_fem_points(output[NT+1, 2*(m+1)*(n+1)+m*n+1:end], m, n, h)
subplot(248)
visualize_scalar_on_fem_points(t0, m, n, h)

tight_layout()

# final_u=output[NT+1, 1:(1+m)*(1+n)]
# final_v=output[NT+1, (1+m)*(1+n)+1:2*(m+1)*(n+1)]
# final_p=output[NT+1, 2*(m+1)*(n+1)+1:end]

# u1 = final_u[50*101+1: 50*101+101]
# u2 = final_u[51:101:end]

# v1 = final_v[50*101+1: 50*101+101]
# v2 = final_v[51:101:end]
# xx = 0:0.01:1

# figure();plot(xx, u1);
# savefig("u_horizontal.png")

# figure();plot(xx, u2);
# savefig("u_vertical.png")

# figure();plot(xx, v1);
# savefig("v_horizontal.png")

# figure();plot(xx, v2);
# savefig("v_vertical.png")

# p1 = final_p[49*100+1: 49*100+100]
# p2 = final_p[50*100+1: 50*100+100]
# p3 = 0.5 * (p1 .+ p2)

# p4 = final_p[50:100:end]
# p5 = final_p[51:100:end]
# p6 = 0.5 * (p4 .+ p5)

# xx = 0.005:0.01:1

# figure();plot(xx, p3);
# savefig("p_horizontal.png")

# figure();plot(xx, p6);
# savefig("p_vertical.png")


text:
###include("Crystal1563.jl")
###include("DFToutMod.jl")
#using XMLDict

#######################
module DOS
"""
DOS
"""

using LinearAlgebra
using Plots
using Base.Threads
using ..CrystalMod:get_grid
using ..TB:calc_energy_fft
using ..TB:tb_crys
using ..TB:tb_crys_kspace
using ..CrystalMod:crystal
using ..CrystalMod:orbital_index
using ..TB:summarize_orb
using ..TB:get_energy_electron_density_kspace
using ..ThreeBodyTB:convert_energy
using ..ThreeBodyTB:convert_dos
using ..ThreeBodyTB:global_energy_units
using ..BandStruct:no_display

function get_projtype(tbc, ptype=missing)

    ind2orb, orb2ind, etotal, nval = orbital_index(tbc.crys)

    if ismissing(ptype)
        if length(Set(tbc.crys.types)) != 1
            ptype=:atomic
        else
            ptype=:orbs
        end
        println("Projection type: ", ptype)
    end

    names = []
    PROJ = []
    pwan=[]

    if ptype == :atomic || ptype == "atomic" || ptype == :atoms || ptype == "atoms" || ptype == :atom || ptype == "atom" || ptype == :Atomic || ptype == "Atomic"
        for ti in Set(tbc.crys.stypes)

            proj_inds = Int64[]
            for n = 1:tbc.tb.nwan
                at,t, orb = ind2orb[n]
                sorb = summarize_orb(orb)
                if t == ti
                    push!(proj_inds, n)
                end
            end

            push!(names, String(ti))
            push!(PROJ, proj_inds)
            push!(pwan, length(proj_inds))
        end

        println("PROJ")
        for a in zip(names, PROJ, pwan)
            println(a)
        end

    else

        for o in [:s, :p, :d, :f]
            proj_inds = Int64[]
            for n = 1:tbc.tb.nwan
                at,t, orb = ind2orb[n]
                sorb = summarize_orb(orb)
                if o == sorb
                    push!(proj_inds, n)
                end
            end
            if length(proj_inds) > 0
                push!(names, String(o))
                push!(PROJ, proj_inds)
                push!(pwan, length(proj_inds))
            end

        end
    end

    return names, PROJ, pwan

end

"""
    function projection(tbc::tb_crys, vects, sk3, grid; ptype=missing)

Figures out the projections.
`ptype` can be `:atomic` or `:orbs` for atom projection or orbital projection (:s,:p,:d)
Default is to choose `:atomic` except for elemental systems.
"""
function projection(tbc::tb_crys, vects, sk3, grid; ptype=missing)

    names, PROJ, pwan = get_projtype(tbc, ptype)

    nk = size(vects)[1]

    temp = zeros(Complex{Float64}, tbc.tb.nwan, tbc.tb.nwan)
    proj = zeros(nk, tbc.tb.nwan, length(PROJ))

#    sk5 = zeros(Complex{Float64}, tbc.tb.nwan, tbc.tb.nwan)
    sk = zeros(Complex{Float64}, tbc.tb.nwan, tbc.tb.nwan)
    c=0
#=
    for k1 = 1:grid[1]
        for k2 = 1:grid[2]
            for k3 = 1:grid[3]
                c += 1
                sk5[:,:] = ( 0.5 * (sk3[:, :, k1, k2, k3] + sk3[:, :, k1, k2, k3]'))^0.5
                sv = sk5 * vects[c,:,:]
                for (pind, proj_inds) in enumerate(PROJ)
                    for p in proj_inds
                        proj[c,:, pind] += real(conj(sv[p,:]) .* sv[p,:])
                    end
#                    if c == 1
#                        println("pind ", pind, " $proj_inds ", proj[c,:,pind])
#                        println("check ", 0.5*sum(vects[c,:,:]' * (sk3[:, :, k1, k2, k3] + sk3[:, :, k1, k2, k3]') * vects[c,:,:]))
#                    end
                end
            end
        end
    end
=#
#    for k1 = 1:grid[1]
#        for k2 = 1:grid[2]
#            for k3 = 1:grid[3]
#                c += 1
                for (pind, proj_inds) in enumerate(PROJ)
                    for p in proj_inds
                        for a = 1:tbc.tb.nwan
                            for j = 1:tbc.tb.nwan
#                                t = vects[c,p,a]*conj(vects[c,j,a])
#                                proj[c,a, pind] += 0.5*real( (t*sk[j,p]  + conj(t)* conj(sk[j,p])))

                                t = vects[:,p,a].*conj(vects[:,j,a])
                                @threads for c = 1:grid[1]*grid[2]*grid[3]

                                    k3 = mod(c-1 , grid[3])+1
                                    k2 = 1 + mod((c-1) ÷ grid[3], grid[2])
                                    k1 = 1 + (c-1) ÷ (grid[2]*grid[3])

#                                    sk[:,:] = ( 0.5 * ( (@view sk3[:, :, k1, k2, k3]) + (@view sk3[:, :, k1, k2, k3])'))

                                    proj[c,a, pind] += 0.5*real( (t[c]*sk3[j,p,k1,k2,k3]  + conj(t[c])* conj(sk3[j,p, k1,k2,k3])))

                            end
                        end
                    end
                end

#                    if c == 1
#                        println("pind ", pind, " $proj_inds ", proj[c,:,pind])
#                        println("check ", 0.5*sum(vects[c,:,:]' * (sk3[:, :, k1, k2, k3] + sk3[:, :, k1, k2, k3]') * vects[c,:,:]))
#                    end
#                end
#            end
#        end
    end

    #    for (pind, proj_inds) in enumerate(PROJ)
#        println("SUM $pind", sum(proj[:,:,pind], dims=1) / prod(grid))
#    end

                    #                    for p in proj_inds
#                        for a = 1:tbc.tb.nwan
#                            for b = 1:tbc.tb.nwan
#                                #                                temp[a,b] = vects[c, a,p]'*sk3[a, b, k1, k2, k3]* vects[c, b,p]
#                                temp[a,b] = vects[c, p,a]'*sk3[a, b, k1, k2, k3]* vects[c, p,b]
#                            end
#                        end
#                        temp = temp + conj(temp)
#                        proj[c,:, pind] += 0.5*sum(real(temp), dims=1)[:]
#                    end
#                    if k1 == 1 && k2 == 1 && k3 == 1
#                        println("inds, ",proj_inds)
#                        println(proj[c,:, pind])
#                    end
#
#                end
#            end
#        end

#    end

    return proj, names, pwan

end

function projection(tbcK::tb_crys_kspace, vects, SK; ptype=missing)

    names, PROJ, pwan = get_projtype(tbcK, ptype)

    nk = size(vects)[1]

    temp = zeros(Complex{Float64}, tbcK.tb.nwan, tbcK.tb.nwan)
    proj = zeros(nk, tbcK.tb.nwan, length(PROJ))

#    sk5 = zeros(Complex{Float64}, tbcK.tb.nwan, tbcK.tb.nwan)
    sk = zeros(Complex{Float64}, tbcK.tb.nwan, tbcK.tb.nwan)
    c=0

    for k = 1:nk
        c += 1
        sk[:,:] = ( 0.5 * (SK[ :, :,k] + SK[ :, :, k]'))
        for (pind, proj_inds) in enumerate(PROJ)
            for p in proj_inds
                for a = 1:tbcK.tb.nwan
                    for j = 1:tbcK.tb.nwan
                        t = vects[c,p,a]*conj(vects[c,j,a])
                        proj[c,a, pind] += 0.5*real( (t*sk[j,p]  + conj(t)* conj(sk[j,p])))
                    end
                end
            end
        end
    end

    return proj, names, pwan

end

"""
    function gaussian_dos(tbc::tb_crys; grid=missing, smearing=0.02, npts=missing, proj_type=missing, do_display=true)

Simple Gaussian DOS, mostly for testing.

- `npts` is number of energies
- `proj_type` can be `"none"`, `"atomic"`, or `"orbital"`
- `do_display=false` will suppress the actual plot

The combination of `smearing` and `grid` are important to get converged results.

See also `dos`

`return energies, dos, projected_dos, pdos_names`
"""
function gaussian_dos(tbc::tb_crys; grid=missing, smearing=0.02, npts=missing, proj_type=missing, do_display=true)

    if ismissing(grid)
        grid = get_grid(tbc.crys)
        grid = Int64.(round.(grid * 1.4))
        println("grid $grid")
    end

    @time etot, efermi, vals, vects, hk3, sk3 = calc_energy_fft(tbc, grid=grid, smearing=smearing, return_more_info=true)

#    println("precheck ", sum(vects[1,:,:]' * sk3[:,:,1,1,1] * vects[1,:,:]))

    #prelim
    vals = vals .- efermi

    nk = size(vals)[1]

    vmin = minimum(vals)
    vmax = min(maximum(vals), 5.0)
    r = vmax - vmin

    if ismissing(npts)
        npts = Int64(round(r * 100 ))
    end

    energies = collect(vmin - r*0.02 : r*1.04 / npts    : vmax + r*0.02 + 1e-7)

    dos = zeros(length(energies))

    if ismissing(proj_type) ||  proj_type != "none"
        do_proj=true
        proj, names, pwan =  projection(tbc, vects, sk3, grid, ptype=proj_type)
        nproj = size(proj)[3]

        pdos = zeros(length(energies),nproj)

    else
        do_proj=false
        nproj=0
        pdos=missing
        names=missing
    end

    for (c,e) in enumerate(energies)
        dos[c] = sum(exp.( -0.5 * (vals[:,:] .- e).^2 / smearing^2 ) )
    end
    dos = dos / smearing / (2.0*pi)^0.5 / nk

    if do_proj
        for i = 1:nproj
            for (c,e) in enumerate(energies)
                pdos[c, i] = sum(proj[:,:,i] .*  exp.( -0.5 * (vals[:,:] .- e).^2 / smearing^2 ) )
            end
        end
        pdos = pdos / smearing / (2.0*pi)^0.5 / nk
    end

    println("Int DOS " , sum(dos) * (energies[2]-energies[1]) )

    ind = energies .< 0

    println("Int DOS occ " , sum(dos[ind]) * (energies[2]-energies[1]) )
    for p in 1:nproj
        println("Int pDOS occ $p : " , sum(pdos[ind, p]) * (energies[2]-energies[1]) )
    end

    energies = convert_energy(energies)
    dos = convert_dos(dos)
    pdos = convert_dos(pdos)

    plot_dos(energies, dos, pdos, names, do_display=do_display)

    return energies, dos, pdos, names

end

function gaussian_dos(tbcK::tb_crys_kspace; smearing=0.03, npts=missing, proj_type=missing, do_display=true)

    @time bandenergy, eden, vects, vals, efermi, error_flag = get_energy_electron_density_kspace(tbcK.tb, tbcK.nelec, smearing=0.01)

#    etot, efermi, vals, vects, hk3, sk3 = calc_energy_fft(tbc, grid=grid, smearing=smearing, return_more_info=true)

#    println("precheck ", sum(vects[1,:,:]' * sk3[:,:,1,1,1] * vects[1,:,:]))

    #prelim
    vals = vals .- efermi

    nk = size(vals)[1]

    vmin = minimum(vals) - 0.1
    vmax = maximum(vals) + 0.1

    println("v $vmin $vmax ")

#    vmax = min(maximum(vals), 5.0)
    r = vmax - vmin

    if ismissing(npts)
        npts = Int64(round(r * 100 ))
    end

    energies = collect(vmin - r*0.02 : r*1.04 / npts    : vmax + r*0.02 + 1e-7)

    dos = zeros(length(energies))

    if ismissing(proj_type) ||  proj_type != "none"
        do_proj=true
        proj, names, pwan =  projection(tbcK, vects, tbcK.tb.Sk, ptype=proj_type)
        nproj = size(proj)[3]

        pdos = zeros(length(energies),nproj)

    else
        do_proj=false
        nproj=0
        pdos=missing
        names=missing
    end

    W = repeat(tbcK.tb.kweights, 1, tbcK.tb.nwan)

    for (c,e) in enumerate(energies)
        dos[c] = sum(exp.( -0.5 * (vals[:,:] .- e).^2 / smearing^2 ) .* W )
    end
    dos = dos / smearing / (2.0*pi)^0.5 / 2

    if do_proj
        for i = 1:nproj
            for (c,e) in enumerate(energies)
                pdos[c, i] = sum(proj[:,:,i] .*  exp.( -0.5 * (vals[:,:] .- e).^2 / smearing^2 ) .* W)
            end
        end
        pdos = pdos / smearing / (2.0*pi)^0.5 / 2
    end

    println("Int DOS " , sum(dos) * (energies[2]-energies[1]) )

    ind = energies .< 0

    println("Int DOS occ " , sum(dos[ind]) * (energies[2]-energies[1]) )
    for p in 1:nproj
        println("Int pDOS occ $p : " , sum(pdos[ind, p]) * (energies[2]-energies[1]) )
    end

    energies = convert_energy(energies)
    dos = convert_dos(dos)
    pdos = convert_dos(pdos)

    plot_dos(energies, dos, pdos, names, do_display=do_display)

    return energies, dos, pdos, names

end

"""
    function setup_tetra(grid)

Setup simple tetrahedron method
based on tetra.f90 in QE. Simple tetra method.

"""
function setup_tetra(grid)

    ntetra = prod(grid) * 6

    tetra_map = zeros(Int64, 4, ntetra)

    kpts = zeros(Int64, prod(grid), 3)
    kpts_dict = Dict()

    c=0
    for k1 = 0:grid[1]-1
        for k2 = 0:grid[2]-1
            for k3 = 0:grid[3]-1
                c+=1
                kpts[c,:] = [k1+1, k2+1, k3+1]
                kpts_dict[ [k1+1, k2+1, k3+1]] = c

            end
        end
    end

    c=1

    for k1 = 0:grid[1]-1
        for k2 = 0:grid[2]-1
            for k3 = 0:grid[3]-1

                k1a = (k1) % grid[1] + 1
                k2a = (k2) % grid[2] + 1
                k3a = (k3) % grid[3] + 1

                k1b = (k1+1) % grid[1] + 1
                k2b = (k2+1) % grid[2] + 1
                k3b = (k3+1) % grid[3] + 1

                #corners of cube (eight)
                n1 = kpts_dict[[k1a, k2a, k3a]]
                n2 = kpts_dict[[k1a, k2a, k3b]]
                n3 = kpts_dict[[k1a, k2b, k3a]]
                n4 = kpts_dict[[k1a, k2b, k3b]]
                n5 = kpts_dict[[k1b, k2a, k3a]]
                n6 = kpts_dict[[k1b, k2a, k3b]]
                n7 = kpts_dict[[k1b, k2b, k3a]]
                n8 = kpts_dict[[k1b, k2b, k3b]]

                #six tetra per cube
                tetra_map[1, c] = n1
                tetra_map[2, c] = n2
                tetra_map[3, c] = n3
                tetra_map[4, c] = n6

                tetra_map[1, c+1] = n2
                tetra_map[2, c+1] = n3
                tetra_map[3, c+1] = n4
                tetra_map[4, c+1] = n6

                tetra_map[1, c+2] = n1
                tetra_map[2, c+2] = n3
                tetra_map[3, c+2] = n5
                tetra_map[4, c+2] = n6

                tetra_map[1, c+3] = n3
                tetra_map[2, c+3] = n4
                tetra_map[3, c+3] = n6
                tetra_map[4, c+3] = n8

                tetra_map[1, c+4] = n3
                tetra_map[2, c+4] = n6
                tetra_map[3, c+4] = n7
                tetra_map[4, c+4] = n8

                tetra_map[1, c+5] = n3
                tetra_map[2, c+5] = n5
                tetra_map[3, c+5] = n6
                tetra_map[4, c+5] = n7

                c+=6

            end
        end
    end

    return kpts, kpts_dict, tetra_map

end

"""
    function dos(tbc::tb_crys; grid=missing, npts=missing, proj_type=missing, do_display=true)

DOS, using tetrahedral integration

- `grid` is the k-point grid. Defaults to 1.6 times the default grid for energy integration.
- `npts` is number of energies
- `proj_type` can be `"none"`, `"atomic"`, or `"orbital"`. Defaults to `atomic` if more than one atom type.
- `do_display=false` will suppress plotting

`return energies, dos, projected_dos, pdos_names`
"""
function dos(tbc::tb_crys; grid=missing, npts=missing, proj_type=missing, do_display=true)

    if ismissing(grid)
        grid = get_grid(tbc.crys)
        grid = Int64.(round.(grid * 1.4))
        println("grid $grid")
    end

    @time etot, efermi, vals, vects,hk3, sk3 = calc_energy_fft(tbc, grid=grid, return_more_info=true)

#    println("vects ", size(vects))
#    println(vects[1,:,:])
#    println("vals ", size(vals))
#    println(vals[1, :])
#    println("sk ", size(sk3))
#    println(sk3[:,:,1,1,1])

    begin

    vals = vals .- efermi

    nk = size(vals)[1]

    vmin = minimum(vals)
    vmax = maximum(vals)
    vmax = min(maximum(vals), 5.0)
    r = vmax - vmin

    if ismissing(npts)
        npts = Int64(round(r * 100 ))
    end

    energies = collect(vmin - r*0.01 : r*1.02 / npts    : vmax + r*0.01 + 1e-7)

    dos = zeros(size(energies))
    dos_id = zeros( length(energies), nthreads())

    if ismissing(proj_type) ||  (proj_type != "none"   && proj_type != :none)
        do_proj=true
        #println("Projection")
        proj, names, pwan =  projection(tbc, vects, sk3, grid, ptype=proj_type)
        nproj = size(proj)[3]

        pdos = zeros(length(energies),nproj)
        pdos_id = zeros(length(energies),nproj, nthreads())

    else
        do_proj=false
        nproj=0
        pdos=missing
        names=missing
    end

        #println("setup tetra")
        kpts, kpts_dict, tetra_map = setup_tetra(grid)

    #    e=zeros(4)

    ntetra = nk*6
    norm = Float64(1/ntetra)

    #    ex=zeros(4)
#    px=zeros(4)
#    ktet = zeros(Int64, 4)

    eps = 1e-7

    #println("calc dos tetra")

    end

    range = 1:length(energies)

    @threads for nt = 1: ntetra
        id = threadid()
        #id = 1
        ex=zeros(4)
        px=zeros(4)
        e=zeros(4, tbc.tb.nwan)
        p=zeros(4, tbc.tb.nwan, nproj)
        ktet = zeros(Int64, 4)
        wt = zeros(4)

        ktet[:] .= tetra_map[:, nt]

        e[:,:] .= vals[ktet, :]

        if do_proj
            p[:,:, :] .= proj[ktet,:, :]
            f14 = 0.0
            f24 = 0.0
            f35 = 0.0
            G = 0.0
            f13 = 0.0
            f31 = 0.0
            f14 = 0.0
            f41 =  0.0
            f23 = 0.0
            f32 = 0.0
            f24 = 0.0
            f42 = 0.0

            f12 =  0.0
            f21 = 0.0

        end

        for nband = 1:tbc.tb.nwan

            ex[:] = e[:,nband]

            if do_proj==false  ####

                sort!(ex)

                inds = (ex[1] .< energies .< ex[4] )

                for c in range[inds]

                    en = energies[c]

                    if en < ex[4] && en >= ex[3]

                        dos_id[c, id] += 3.0 * (ex[4] - en)^2 / (ex[4] - ex[1] + eps) / (ex[4] - ex[2] + eps) / (ex[4] - ex[3] + eps)

                    elseif en < ex[3] && en >= ex[2]

                        dos_id[c, id] += 1.0 / (ex[3] - ex[1] + eps) / (ex[4] - ex[1] + eps) * (3.0 * (ex[2] - ex[1]) + 6.0 * (en - ex[2]) - 3.0 * (ex[3] - ex[1] + ex[4] - ex[2]) / (ex[3]-ex[2] + eps) / (ex[4] - ex[2] + eps) * (en -ex[2])^2)

                    elseif en < ex[2] && en >= ex[1]

                        dos_id[c, id] += 3.0 * (en-ex[1])^2 / (ex[2] - ex[1] + eps) / (ex[3] - ex[1] + eps) / (ex[4] - ex[1] + eps)

                    end

                end

            elseif do_proj == true ####

                perm = sortperm(ex)
                ex[:] = ex[perm]

                inds = (ex[1] .< energies .< ex[4] )

                for c in range[inds]
#                for c in 1:length(energies)
                    en = energies[c]

                    if en < ex[4] && en >= ex[1]

                        if en < ex[4] && en >= ex[3]

                            f14 = (en-ex[4])/(ex[1]-ex[4] - eps)
                            f24 = (en-ex[4])/(ex[2]-ex[4] - eps)
                            f34 = (en-ex[4])/(ex[3]-ex[4] - eps)

                            G = 3.0 * f14 * f24 * f34 / (ex[4] - en + eps)

                            @inbounds wt[1] = f14 / 3.0
                            @inbounds wt[2] = f24 / 3.0
                            @inbounds wt[3] = f34 / 3.0
                            @inbounds wt[4] = (3.0 - f14 - f24 - f34) / 3.0

                        elseif en < ex[3] && en >= ex[2]

                            f13 = (en-ex[3])/(ex[1]-ex[3] - eps)
                            f31 = 1.0 - f13
                            f14 = (en-ex[4])/(ex[1]-ex[4] - eps)
                            f41 = 1.0 - f14
                            f23 = (en-ex[3])/(ex[2]-ex[3] - eps)
                            f32 = 1.0 - f23
                            f24 = (en-ex[4])/(ex[2]-ex[4] - eps)
                            f42 = 1.0 - f24

                            G   =  3.0 * (f23*f31 + f32*f24)

                            @inbounds wt[1]  =  f14 / 3.0 + f13*f31*f23 / (G + eps)
                            @inbounds wt[2]  =  f23 / 3.0 + f24*f24*f32 / (G + eps)
                            @inbounds wt[3]  =  f32 / 3.0 + f31*f31*f23 / (G + eps)
                            @inbounds wt[4]  =  f41 / 3.0 + f42*f24*f32 / (G + eps)

                            G   =  G / (ex[4]-ex[1] + eps)

                        elseif en < ex[2] && en >= ex[1]

                            f12 = (en-ex[2])/(ex[1]-ex[2] - eps)
                            f21 = 1.0 - f12
                            f13 = (en-ex[3])/(ex[1]-ex[3] - eps)
                            f31 = 1.0 - f13
                            f14 = (en-ex[4])/(ex[1]-ex[4] - eps)
                            f41 = 1.0 - f14

                            G  =  3.0 * f21 * f31 * f41 / (en-ex[1] + eps)

                            @inbounds wt[1] =  (f12 + f13 + f14) / 3.0
                            @inbounds wt[2] =  f21  / 3.0
                            @inbounds wt[3] =  f31  / 3.0
                            @inbounds wt[4] =  f41  / 3.0

                        end

                        @inbounds dos_id[c, id] += sum(wt) * G

                        for projind = 1:nproj
                            px[:] = @view p[perm,nband, projind]
                            @inbounds pdos_id[c,projind, id] += ( px' * wt )  * G
                        end

                    end

                end

            end #### if do_proj

        end

    end

    dos[:] = sum(dos_id, dims=2) * norm
    dos = max.(dos, 0.0)

    #println("Int DOS " , sum(dos) * (energies[2]-energies[1]) )
    correct =  tbc.tb.nwan / ( sum(dos) * (energies[2]-energies[1]) )

    dos = dos * correct

    if nproj > 0
        pdos = max.(pdos, 0.0)
        pdos[:,:] = sum(pdos_id, dims=3) * norm
        pdos = pdos * correct

        #println("Int PDOS " , sum(dos, dims=1) * (energies[2]-energies[1]) )

        t = sum(pdos, dims=1) * (energies[2]-energies[1])
        fix = pwan' ./ t
        pdos = pdos .* fix
        pdos2 = zeros(size(pdos))

    else
        pdos2=missing
    end

    #simple smoothing
    dos2 = zeros(size(dos))

    for i = 3:(length(dos)-2)
        dos2[i] =  0.125*dos[i-2] + 0.25*dos[i-1] + 0.25*dos[i] + 0.25*dos[i+1] + 0.125*dos[i+2]
        if nproj>0
            #            pdos2[i,:] = 0.25*pdos[i-1,:] + 0.5*pdos[i,:] + 0.25*pdos[i+1,:]
            pdos2[i,:] = 0.125*pdos[i-2,:] + 0.25*pdos[i-1,:] + 0.25*pdos[i,:] + 0.25*pdos[i+1,:] + 0.125*pdos[i+2,:]
        end
    end

    if nproj > 0
        pdos2 = max.(pdos2, 0.0)
        pdos2 = convert_dos(pdos2)
    end

    energies = convert_energy(energies)
    dos2 = convert_dos(dos2)

    plot_dos(energies, dos2, pdos2, names, do_display=do_display)

    return energies, dos2, pdos2, names

end

"""
    function plot_dos(energies, dos, pdos, names; filename=missing, do_display=true)

Does the actual DOS plotting, called by `dos` or `gaussian_dos`
"""
function plot_dos(energies, dos, pdos, names; filename=missing, do_display=true)

    if ismissing(names)
        plot(legend=false, grid=false, framestyle=:box)
    else
        plot(legend=true, grid=false, framestyle=:box)
    end

    plot!(energies, dos, color="black", lw=4, label="Total", legend=:topleft, xtickfontsize=12,ytickfontsize=12, legendfontsize=12)

    if !ismissing(names)
        colors = ["blue", "orange", "green", "magenta", "cyan", "red", "yellow"]
        for i in 1:size(pdos)[2]
            color = colors[i%7+1]
#            println("i $i $color", names[i])
            plot!(energies, pdos[:,i], color=color, lw=3, label=names[i])
        end
    end

    if global_energy_units == "eV"
        ylabel!("DOS  ( 1 / eV )", guidefontsize=16)
        xlabel!("Energy - E_F ( eV )", guidefontsize=16)
    else
        ylabel!("DOS  ( 1 / Ryd. )", guidefontsize=16)
        xlabel!("Energy - E_F ( Ryd. )", guidefontsize=16)
    end

    xl1 = minimum(energies)
    xl2 = min(10.0, maximum(energies))
#    xl2 = ( maximum(energies))

    xlims!(xl1, xl2)

    if do_display && ! no_display
        display(plot!([0,0], [0, maximum(dos) * 1.1], color="black", linestyle=:dash, label=""))
    else

        plot!([0,0], [0, maximum(dos) * 1.1], color="black", linestyle=:dash, label="")

    end

    if !ismissing(filename)
        savefig(filename)
    end

end

end #end module


text:
include("Graph_class.jl")
using Random

function readGraph(fileName, graphType) # Read graph from a file | graphType = [weighted, unweighted]
    # Initialized
    n = 0
    origin = Dict{Int32, Int32}()
    label = Dict{Int32, Int32}()
    edge = Set{Tuple{Int32, Int32, Float64}}()

    getid(x :: Int32) = haskey(label, x) ? label[x] : label[x] = n += 1

    open(fileName) do f1
        for line in eachline(f1)
            # Read origin data from file
            buf = split(line)
            u = parse(Int32, buf[1])
            v = parse(Int32, buf[2])
            if graphType == "weighted"
                w = parse(Float64, buf[3])
            else
                w = 1.0
            end
            if u == v
                continue
            end
            # Label the node
            u1 = getid(u)
            v1 = getid(v)
            origin[u1] = u
            origin[v1] = v
            # Store the edge
            if u1 > v1
                u1, v1 = v1, u1
            end
            push!(edge, (u1, v1, w)) # each edge stores only once
        end
    end

    # Store data into the struct Graph
    m = length(edge)
    V = Array{Int32, 1}(undef, n)
    E = Array{Tuple{Int32, Int32, Int32, Float64}, 1}(undef, m)

    for i = 1 : n
        V[i] = origin[i]
    end

    ID = 0
    for (u, v, w) in edge
        ID = ID + 1
        E[ID] = (ID, u, v, w)
    end

    return Graph(n, m, V, E)
end

function BFS(st, g) # Search from the node st
    q = []
    h = Array{Int8, 1}(undef, size(g, 1))
    fill!(h, 0)
    push!(q, st)
    h[st] = 1
    front = 1
    rear = 1
    while front <= rear
        u = q[front]
        front = front + 1
        for v in g[u]
            if h[v] == 0
                h[v] = 1
                rear = rear + 1
                push!(q, v)
            end
        end
    end
    return q
end

function getLLC(G :: Graph) # Find the largest link clique
    # Create the Link Table
    g = Array{Array{Int32, 1}, 1}(undef, G.n)
    for i = 1 : G.n
        g[i] = []
    end
    for (ID, u, v, w) in G.E
        push!(g[u], v)
        push!(g[v], u)
    end
    # Find the node Set of LLC
    n2 = 0
    S = Array{Int32, 1}(undef, G.n)
    visited = Array{Int8, 1}(undef, G.n)
    fill!(visited, 0)
    for i = 1 : G.n
        if visited[i] == 0
            nodeSet = BFS(i, g)
            for x in nodeSet
                visited[x] = 1
            end
            tn = size(nodeSet, 1)
            if (tn > n2)
                n2 = tn
                for j = 1 : n2
                    S[j] = nodeSet[j]
                end
            end
        end
    end
    # Store the result
    label = Array{Int32, 1}(undef, G.n)
    fill!(label, 0)
    V2 = Array{Int32, 1}(undef, n2)
    for i = 1 : n2
        label[S[i]] = i
        V2[i] = G.V[S[i]]
    end
    E2 = []
    nID = 0
    for (ID, u, v, w) in G.E
        if (label[u] > 0) && (label[v] > 0)
            nID += 1
            push!(E2, (nID, label[u], label[v], w))
        end
    end
    m2 = size(E2, 1)

    #G2 = Graph(n2, m2, V2, E2)
    #g_new = Array{Array{Int32, 1}, 1}(undef, G2.n)
    #for i = 1 : G2.n
    #    g_new[i] = []
    #end
    #for (ID, u, v, w) in G2.E
    #    push!(g_new[u], v)
    #end

    #edgelist = BFS(1, g_new);
    #println("The size of travelled list is ", size(edgelist, 1))

    return Graph(n2, m2, V2, E2)
end

### New
### !!! NOTE: there exists only one edge between two nodes in the graph

## (1) Weighted cascade model to assign edge weights
function assignWeightedCascade(G :: Graph)
    g = Array{Array{Int32, 1}, 1}(undef, G.n)
    for i = 1 : G.n
        g[i] = []
    end
    for (ID, u, v, w) in G.E
        push!(g[v], u) # in-edges for v
    end
    p1 = Array{Float64, 1}(undef, G.n) # store edge probabilities
    for i = 1 : G.n
        if size(g[i], 1) > 0
            p1[i] = 1 / size(g[i], 1)
        else
            p1[i] = 0
        end
    end

    E2 = []
    for (ID, u, v, w) in G.E
        push!(E2, (ID, u, v, p1[v])) #reassign edge weights according to weighted cascade model
    end

    return Graph(G.n, G.m, G.V, E2)
end

## (2) trialency model to assign edge weights
function assignTrivalency(G :: Graph)
    p2 = Trivalency(G.m)
    E2 = []
    idx = 1
    for (ID, u, v, w) in G.E
        push!(E2, (ID, u, v, p2[idx])) # reassign edge weights according to trialency model
        idx += 1
    end

    return Graph(G.n, G.m, G.V, E2)
end

"""
(3) Independent Cascade Model
Input: the graph with proabilities assigned

Reference:
  [1] David Kempe, Jon Kleinberg, and Eva Tardos.
      Influential nodes in a diffusion model for social networks.
      In Automata, Languages and Programming, 2005.
"""

# Rounds that takes care of the nodes "Activated" & "Acknowledge"
function cascadeIC_one_round(g :: Array{Array{Tuple{Int32, Float64}, 1}, 1}, A, activated_nodes_of_last_round, delta)
    activated_nodes_of_this_round = Set{Int32}()
    akg_nodes_of_this_round = Set{Int32}()
    tmp = 0
    for s in activated_nodes_of_last_round
        #println(s)
        for nb in g[s]
            if nb[1] in A
                continue
            end
            tmp = rand(1)[1]
            if tmp < delta * nb[2]
                push!(activated_nodes_of_this_round, nb[1])
            elseif delta*nb[2]<= tmp && tmp < nb[2]
                push!(akg_nodes_of_this_round, nb[1])
            end
        end
    end

    #activated_nodes_of_this_round = toarray(activated_nodes_of_this_round) # convert to vector
    #A = vcat(A, activated_nodes_of_this_round) # add to original vector

    return activated_nodes_of_this_round, akg_nodes_of_this_round
end

# seeds should be a vector iter: number of iterations
# return a vector of probabilites of getting activated, and time
function cascadeIC(G :: Graph, seeds, iter = 1000, delta = 0.5)
    p4 = zeros(Float64, G.n)

    # Create the Link Table
    g = Array{Array{Tuple{Int32, Float64}, 1}, 1}(undef, G.n)
    for i = 1 : G.n
        g[i] = []
    end
    for (_, u, v, w) in G.E
        push!(g[u], (v, w))
    end

    T = time()

    A, B, activated_nodes_of_one_round = Set{Int32}(), Set{Int32}(), Set{Int32}()

    for i = 1 : iter
        A = Set{Int32}(seeds)
        activated_nodes_of_one_round = Set{Int32}(seeds)
        B = Set{Int32}()
        #println("A is ", A) A test
        while length(activated_nodes_of_one_round) != 0
            (activated_nodes_of_one_round, akg_nodes_of_this_round) = cascadeIC_one_round(g, A, activated_nodes_of_one_round, delta)
            union!(A, activated_nodes_of_one_round)
            union!(B, akg_nodes_of_this_round)
        end

        union!(A, B)
        C = toarray(A)
        for node in C
            p4[node] += 1
        end

    end

    p4 = p4 ./ iter

    T = time() - T

    return T, p4
end


text:
#****************************************************************************
#                     Molecular Dynamics Potentials [MDP]
#                            CESMIX-MIT Project
#
#  Contributing authors: Ngoc-Cuong Nguyen [cuongng@mit.edu, exapde@gmail.com]
#****************************************************************************

function prebox(pbc::Array{Int64,1}, rmax::Array{Float64,1}, a::Array{Float64,1}, b::Array{Float64,1}, c=nothing)

if c===nothing
    B2R, R2B = domainmaps(a, b);

    v,w = boundingbox(pbc, rmax, a, b);
    vr = B2R*v; wr = B2R*w;

    pimages = periodicimages(pbc, a, b);
else
    B2R, R2B = domainmaps(a, b, c);

    v,w = boundingbox(pbc, rmax, a, b, c);
    vr = B2R*v; wr = B2R*w;

    pimages = periodicimages(pbc, a, b, c);
end

return B2R, R2B, v, w, vr, wr, pimages

end


text:
using Documenter

makedocs(
    sitename = "BAWJ",
    format = Documenter.HTML(),
    pages = [
        "Business Applications with Julia" => "index.md",
        "Blog" => "blog.md",
        "Section 1" => [
            "1 - Intro Notebook" => "chapter1.md",
            "2 - Overview Julia" => "chapter2.md"
        ],
        "Section 2 - Building the application" => [
            "3 - The design" => "chapter3.md",
            "4 - Implementing the design" => "chapter4.md",
            "5 - Creating modules" => "chapter5.md",
            "6 - Testing the application" => "chapter6.md"
        ],
        "Section 3 - Modules" => [
            "7 - Writing test software" => "chapter7.md",
            "8 - Documenting the modules" => "chapter8.md"
        ],
        "Section 4 - SSH enabled containers" => [
            "9 - Introduction" => "chapter9.md",
            "10 - Create SSH enabled Containers" => "chapter10.md",
            "11 - Raspberry Pi" => "chapter11.md"
        ],
        "Section 5 Up and running" => [
            "12 - Running functions remote" => "chapter12.md",
            "13 - Running the containers" => "chapter13.md",
            "14 - Logging" => "chapter14.md",
            "15 - Create a Package" => "chapter15.md",
            "19 - Activity diagrams" => "chapter19.md",
            "20 - Using Remote Channels (old stuff)" => "chapter20.md"
        ],

        "Appendix" => "appendix.md",
        "References" => "references.md"
    ]
)

max_stars_repo_name: PallHaraldsson/BAWJ

max_issues_repo_name: PallHaraldsson/BAWJ

max_forks_repo_name: PallHaraldsson/BAWJ


text:
using Test
using Pkg.Artifacts
using NCDatasets

using RRTMGP
using RRTMGP.Device: array_type, array_device
using RRTMGP.LookUpTables
using RRTMGP.Vmrs
using RRTMGP.AtmosphericStates
using RRTMGP.Optics
using RRTMGP.Sources
using RRTMGP.Fluxes
using RRTMGP.BCs
using RRTMGP.RTE
using RRTMGP.RTESolver

import CLIMAParameters as CP

include(joinpath(pkgdir(RRTMGP), "parameters", "create_parameters.jl"))
# overriding some parameters to match with RRTMGP FORTRAN code
overrides = (; grav = 9.80665, molmass_dryair = 0.028964, molmass_water = 0.018016)
param_set = create_insolation_parameters(FT, overrides)

include("reference_files.jl")
include("read_rfmip_clear_sky.jl")
#---------------------------------------------------------------
function sw_rfmip(
    ::Type{OPC},
    ::Type{SRC},
    ::Type{VMR},
    ::Type{FT},
    ::Type{I},
    ::Type{DA},
) where {FT <: AbstractFloat, I <: Int, DA, OPC, SRC, VMR}
    opc = Symbol(OPC)
    sw_file = get_ref_filename(:lookup_tables, :clearsky, λ = :sw) # sw lookup tables
    sw_input_file = get_ref_filename(:atmos_state, :clearsky)      # clear-sky atmos state
    # reference data files for comparison
    flux_up_file = get_ref_filename(:comparison, :clearsky, λ = :sw, flux_up_dn = :flux_up, opc = :TwoStream)
    flux_dn_file = get_ref_filename(:comparison, :clearsky, λ = :sw, flux_up_dn = :flux_dn, opc = :TwoStream)

    FTA1D = DA{FT, 1}
    FTA2D = DA{FT, 2}
    max_threads = Int(256)
    exp_no = 1

    # reading shortwave lookup data
    ds_sw = Dataset(sw_file, "r")
    lookup_sw, idx_gases = LookUpSW(ds_sw, I, FT, DA)
    close(ds_sw)
    # reading rfmip data to atmospheric state
    ds_sw_in = Dataset(sw_input_file, "r")

    (as, _, sfc_alb_direct, zenith, toa_flux, usecol) =
        setup_rfmip_as(ds_sw_in, idx_gases, exp_no, lookup_sw, FT, DA, VMR, max_threads)
    close(ds_sw_in)

    ncol, nlay, ngpt = as.ncol, as.nlay, lookup_sw.n_gpt
    nlev = nlay + 1
    op = OPC(FT, ncol, nlay, DA)            # allocating optical properties object
    src_sw = SRC(FT, DA, nlay, ncol)        # allocating longwave source function object
    #src_sw = source_func_shortwave(FT, ncol, nlay, opc, DA)        # allocating longwave source function object

    # setting up boundary conditions
    inc_flux_diffuse = nothing
    sfc_alb_diffuse = FTA2D(deepcopy(sfc_alb_direct))
    bcs_sw = SwBCs{FT, FTA1D, Nothing, FTA2D}(zenith, toa_flux, sfc_alb_direct, inc_flux_diffuse, sfc_alb_diffuse)
    fluxb_sw = FluxSW(ncol, nlay, FT, DA) # flux storage for bandwise calculations
    flux_sw = FluxSW(ncol, nlay, FT, DA)  # shortwave fluxes for band calculations

    # initializing RTE solver
    slv = Solver(as, op, nothing, src_sw, nothing, bcs_sw, nothing, fluxb_sw, nothing, flux_sw)
    #--------------------------------------------------
    solve_sw!(slv, max_threads, lookup_sw)

    for i in 1:10
        @time solve_sw!(slv, max_threads, lookup_sw)
    end

    # reading comparison data
    flip_ind = nlev:-1:1

    ds_flux_up = Dataset(flux_up_file, "r")
    comp_flux_up = ds_flux_up["rsu"][:][flip_ind, :, exp_no]
    close(ds_flux_up)

    ds_flux_dn = Dataset(flux_dn_file, "r")
    comp_flux_dn = ds_flux_dn["rsd"][:][flip_ind, :, exp_no]
    close(ds_flux_dn)

    flux_up = Array(slv.flux_sw.flux_up)
    flux_dn = Array(slv.flux_sw.flux_dn)

    for i in 1:ncol
        if usecol[i] == 0
            flux_up[:, i] .= FT(0)
            flux_dn[:, i] .= FT(0)
        end
    end

    max_err_flux_up = maximum(abs.(flux_up .- comp_flux_up))
    max_err_flux_dn = maximum(abs.(flux_dn .- comp_flux_dn))

    println("=======================================")
    println("Clear-sky shortwave test, opc  = $opc")
    println("max_err_flux_up = $max_err_flux_up")
    println("max_err_flux_dn = $max_err_flux_dn")

    toler = FT(0.001)
    @test maximum(abs.(flux_up .- comp_flux_up)) ≤ toler
    @test maximum(abs.(flux_dn .- comp_flux_dn)) ≤ toler
    return nothing
end

sw_rfmip(TwoStream, SourceSW2Str, VmrGM, Float64, Int, array_type()) # two-stream solver should be used for the short-wave problem
#sw_rfmip(OneScalar, VmrGM, Float64, Int, array_type()) # this only computes flux_dn_dir

max_stars_repo_name: climate-machine/JRRTMGP

max_issues_repo_name: climate-machine/JRRTMGP

max_forks_repo_name: climate-machine/JRRTMGP


text:
#InspectDR: Top-level interface for Cairo layer
#-------------------------------------------------------------------------------

#Render entire plot
#-------------------------------------------------------------------------------
function render(canvas::PCanvas2D, plot::Plot2D)
	#Render annotation/axes
	render(canvas, plot.annotation, plot.layout)
	if plot.layout.legend.enabled
		legend_render(canvas, plot)
	end

	grid = gridlines(plot.axes, canvas.ext)
	render_grid(canvas, plot.layout, grid)
	#TODO: render axes first once drawing is multi-threaded.
#	render_axes

	#Plot actual data
Cairo.save(canvas.ctx)
	setclip(canvas.ctx, canvas.graphbb)
	render(canvas, plot.display_data)

	#Plot secondary annotation:
	render(canvas, plot.markers, plot.axes)
	render(canvas, plot.atext, plot.axes)

Cairo.restore(canvas.ctx)

	#Re-render axis over data:
	render_axes(canvas, plot.layout, grid)
end

#Render entire plot within provided bounding box:
function render(ctx::CairoContext, plot::Plot2D, bb::BoundingBox)
	graphbb = graphbounds(bb, plot.layout, plot.axes)
	update_ddata(plot) #Also computes new extents
	canvas = PCanvas2D(ctx, bb, graphbb, getextents_xfrm(plot))
	render(canvas, plot)
end

#Last line


text:
# Regression Metrics

"""
    mae(y_pred, y_true)

Mean Absolute Error. Calculated as `sum(|y_true .- y_pred|) / length(y_true)` based on provided `y_pred` and `y_true`.
"""
function mae(y_pred, y_true)
    @assert length(y_true) == length(y_pred)
    return sum(abs.(y_true .- y_pred)) / length(y_true)
end

"""
    mse(y_pred, y_true)

Mean Squared Error. Calculated as `sum((y_true .- y_pred).^2) / length(y_true)` based on provided `y_pred` and `y_true`.
"""
function mse(y_pred, y_true)
    @assert length(y_true) == length(y_pred)
    return sum((y_true .- y_pred).^2) / length(y_true)
end

"""
    male(y_pred, y_true)

Mean Absolute Logarithmic Error. Calculated as `sum(|log.(y_true) .- log.(y_pred)|) / length(y_true)` based on provided `y_pred` and `y_true`.
"""
function male(y_pred, y_true)
    @assert length(y_true) == length(y_pred)
    epsilon = eps(eltype(y_pred)(1.0))
    return sum(abs.(log.(y_true .+ epsilon) .- log.(y_pred .+ epsilon)))  / length(y_true)
end

"""
    msle(y_pred, y_true)

Mean Absolute Logarithmic Error.
Calculated as `sum((log.(y_true) .- log.(y_pred)).^2) / length(y_true)` based on provided `y_pred` and `y_true`.
"""
function msle(y_pred, y_true)
    @assert length(y_true) == length(y_pred)
    epsilon = eps(eltype(y_pred)(1.0))
    return sum((log.(y_true .+ epsilon) .- log.(y_pred .+ epsilon)).^2)  / length(y_true)
end

"""
    r2_score(y_pred, y_true)

Calculates the r2 (Coefficient of Determination) score for the provided `y_pred` and `y_true`.
Best possible score is 1.0 and it can be negative (because the model can be arbitrarily worse).
A constant model that always predicts the expected value of y, disregarding the input features, would get a r2_score of `0.0`.
"""
function r2_score(y_pred, y_true)
    @assert length(y_true) == length(y_pred)
    ss_res = sum((y_true .- y_pred).^2)
    mean = sum(y_true) / length(y_true)
    ss_total = sum((y_true .- mean).^2)
    return 1 - ss_res/(ss_total + eps(eltype(y_pred)))
end

"""
    adjusted_r2_score(y_pred, y_true, n)

Modified version of `r2_score` that has been adjusted for the number of predictors in the model. Here the argument `n` is for the number of predictors(or independent variables in X).

See also: [`r2_score`](@ref)
"""
function adjusted_r2_score(y_pred, y_true, n)		# n -> number of predictors(independent variables in X)
    @assert length(y_true) == length(y_pred)
    score = r2_score(y_pred, y_true)
    epsilon = eps(eltype(y_pred)(1.0))
    return 1 - ((1 - score) * (length(y_true) -1)) / abs(( length(y_true) - n - 1 + epsilon))
end