text:
let
    # Test core NRLMSISE - Test cases match Daniel Brodo's C implementation
    input  = SatelliteDynamics.NRLMSISE_Input[SatelliteDynamics.NRLMSISE_Input() for i in 1:17]
    output = SatelliteDynamics.NRLMSISE_Output[SatelliteDynamics.NRLMSISE_Output() for i in 1:17]
    flags  = SatelliteDynamics.NRLMSISE_Flags()

    # Set input values
    aph = ones(Float64, 7).*100.0

    flags.switches[1] = 0
    for i in 2:24
        flags.switches[i] = 1
    end

    for i in 1:17
        input[i].doy   = 172
        input[i].year  = 0
        input[i].sec   = 29000
        input[i].alt   = 400
        input[i].g_lat = 60
        input[i].g_lon = -70
        input[i].lst   = 16
        input[i].f107A = 150
        input[i].f107  = 150
        input[i].ap    = 4
    end

    input[2].doy       = 81
    input[3].sec       = 75000
    input[3].alt       = 1000
    input[4].alt       = 100
    input[11].alt      = 0
    input[12].alt      = 10
    input[13].alt      = 30
    input[14].alt      = 50
    input[15].alt      = 70
    input[17].alt      = 100
    input[5].g_lat     = 0
    input[6].g_lon     = 0
    input[7].lst       = 4
    input[8].f107A     = 70
    input[9].f107      = 180
    input[10].ap       = 40
    input[16].ap_array = aph
    input[17].ap_array = aph

    # Run model
    SatelliteDynamics.gtd7!(input[1], flags, output[1])

    # Evalute 1 to 15
    for i in 1:15
        SatelliteDynamics.gtd7!(input[i], flags, output[i])
    end

    # Evaluate 16 and 17
    for i in 16:17
        SatelliteDynamics.gtd7!(input[i], flags, output[i])
    end

    # Test Case 1
    @test isapprox(output[1].t[1], 1250.54, atol=1.0e-2)
    @test isapprox(output[1].t[2], 1241.42, atol=1.0e-2)
    @test isapprox(output[1].d[1], 6.6651769e+05, rtol=1.0e-6)
    @test isapprox(output[1].d[2], 1.1388056e+08, rtol=1.0e-6)
    @test isapprox(output[1].d[3], 1.9982109e+07, rtol=1.0e-6)
    @test isapprox(output[1].d[4], 4.0227636e+05, rtol=1.0e-6)
    @test isapprox(output[1].d[5], 3.5574650e+03, rtol=1.0e-6)
    @test isapprox(output[1].d[7], 3.4753124e+04, rtol=1.0e-6)
    @test isapprox(output[1].d[8], 4.0959133e+06, rtol=1.0e-6)
    @test isapprox(output[1].d[9], 2.6672732e+04, rtol=1.0e-6)
    @test isapprox(output[1].d[6], 4.0747135e-15, rtol=1.0e-6)

    # Test Case 2
    @test isapprox(output[2].t[1], 1166.75, atol=1.0e-2)
    @test isapprox(output[2].t[2], 1161.71, atol=1.0e-2)
    @test isapprox(output[2].d[1], 3.4072932e+06, rtol=1.0e-6)
    @test isapprox(output[2].d[2], 1.5863334e+08, rtol=1.0e-6)
    @test isapprox(output[2].d[3], 1.3911174e+07, rtol=1.0e-6)
    @test isapprox(output[2].d[4], 3.2625595e+05, rtol=1.0e-6)
    @test isapprox(output[2].d[5], 1.5596182e+03, rtol=1.0e-6)
    @test isapprox(output[2].d[7], 4.8542085e+04, rtol=1.0e-6)
    @test isapprox(output[2].d[8], 4.3809667e+06, rtol=1.0e-6)
    @test isapprox(output[2].d[9], 6.9566820e+03, rtol=1.0e-6)
    @test isapprox(output[2].d[6], 5.0018457e-15, rtol=1.0e-6)

    # Test Case 3
    @test isapprox(output[3].t[1], 1239.89, atol=1.0e-2)
    @test isapprox(output[3].t[2], 1239.89, atol=1.0e-2)
    @test isapprox(output[3].d[1], 1.1237672e+05, rtol=1.0e-6)
    @test isapprox(output[3].d[2], 6.9341301e+04, rtol=1.0e-6)
    @test isapprox(output[3].d[3], 4.2471052e+01, rtol=1.0e-6)
    @test isapprox(output[3].d[4], 1.3227501e-01, rtol=1.0e-6)
    @test isapprox(output[3].d[5], 2.6188484e-05, rtol=1.0e-6)
    @test isapprox(output[3].d[7], 2.0167499e+04, rtol=1.0e-6)
    @test isapprox(output[3].d[8], 5.7412559e+03, rtol=1.0e-6)
    @test isapprox(output[3].d[9], 2.3743942e+04, rtol=1.0e-6)
    @test isapprox(output[3].d[6], 2.7567723e-18, rtol=1.0e-6)

    # Test Case 4
    @test isapprox(output[4].t[1], 1027.32, atol=1.0e-2)
    @test isapprox(output[4].t[2], 206.89, atol=1.0e-2)
    @test isapprox(output[4].d[1], 5.4115544e+07, rtol=1.0e-6)
    @test isapprox(output[4].d[2], 1.9188934e+11, rtol=1.0e-6)
    @test isapprox(output[4].d[3], 6.1158256e+12, rtol=1.0e-6)
    @test isapprox(output[4].d[4], 1.2252011e+12, rtol=1.0e-6)
    @test isapprox(output[4].d[5], 6.0232120e+10, rtol=1.0e-6)
    @test isapprox(output[4].d[7], 1.0598797e+07, rtol=1.0e-6)
    @test isapprox(output[4].d[8], 2.6157367e+05, rtol=1.0e-6)
    @test isapprox(output[4].d[9], 2.8198794e-42, rtol=1.0e-6)
    @test isapprox(output[4].d[6], 3.5844263e-10, rtol=1.0e-6)

    # Test Case 5
    @test isapprox(output[5].t[1], 1212.40, atol=1.0e-2)
    @test isapprox(output[5].t[2], 1208.14, atol=1.0e-2)
    @test isapprox(output[5].d[1], 1.8511225e+06, rtol=1.0e-6)
    @test isapprox(output[5].d[2], 1.4765548e+08, rtol=1.0e-6)
    @test isapprox(output[5].d[3], 1.5793562e+07, rtol=1.0e-6)
    @test isapprox(output[5].d[4], 2.6337950e+05, rtol=1.0e-6)
    @test isapprox(output[5].d[5], 1.5887814e+03, rtol=1.0e-6)
    @test isapprox(output[5].d[7], 5.8161668e+04, rtol=1.0e-6)
    @test isapprox(output[5].d[8], 5.4789845e+06, rtol=1.0e-6)
    @test isapprox(output[5].d[9], 1.2644459e+03, rtol=1.0e-6)
    @test isapprox(output[5].d[6], 4.8096302e-15, rtol=1.0e-6)

    # Test Case 6
    @test isapprox(output[6].t[1], 1220.15, atol=1.0e-2)
    @test isapprox(output[6].t[2], 1212.71, atol=1.0e-2)
    @test isapprox(output[6].d[1], 8.6730952e+05, rtol=1.0e-6)
    @test isapprox(output[6].d[2], 1.2788618e+08, rtol=1.0e-6)
    @test isapprox(output[6].d[3], 1.8225766e+07, rtol=1.0e-6)
    @test isapprox(output[6].d[4], 2.9222142e+05, rtol=1.0e-6)
    @test isapprox(output[6].d[5], 2.4029624e+03, rtol=1.0e-6)
    @test isapprox(output[6].d[7], 3.6863892e+04, rtol=1.0e-6)
    @test isapprox(output[6].d[8], 3.8972755e+06, rtol=1.0e-6)
    @test isapprox(output[6].d[9], 2.6672732e+04, rtol=1.0e-6)
    @test isapprox(output[6].d[6], 4.3558656e-15, rtol=1.0e-6)

    # Test Case 7
    @test isapprox(output[7].t[1], 1116.39, atol=1.0e-2)
    @test isapprox(output[7].t[2], 1113.00, atol=1.0e-2)
    @test isapprox(output[7].d[1], 5.7762512e+05, rtol=1.0e-6)
    @test isapprox(output[7].d[2], 6.9791387e+07, rtol=1.0e-6)
    @test isapprox(output[7].d[3], 1.2368136e+07, rtol=1.0e-6)
    @test isapprox(output[7].d[4], 2.4928677e+05, rtol=1.0e-6)
    @test isapprox(output[7].d[5], 1.4057387e+03, rtol=1.0e-6)
    @test isapprox(output[7].d[7], 5.2919856e+04, rtol=1.0e-6)
    @test isapprox(output[7].d[8], 1.0698141e+06, rtol=1.0e-6)
    @test isapprox(output[7].d[9], 2.6672732e+04, rtol=1.0e-6)
    @test isapprox(output[7].d[6], 2.4706514e-15, rtol=1.0e-6)

    # Test Case 8
    @test isapprox(output[8].t[1], 1031.25, atol=1.0e-2)
    @test isapprox(output[8].t[2], 1024.85, atol=1.0e-2)
    @test isapprox(output[8].d[1], 3.7403041e+05, rtol=1.0e-6)
    @test isapprox(output[8].d[2], 4.7827201e+07, rtol=1.0e-6)
    @test isapprox(output[8].d[3], 5.2403800e+06, rtol=1.0e-6)
    @test isapprox(output[8].d[4], 1.7598746e+05, rtol=1.0e-6)
    @test isapprox(output[8].d[5], 5.5016488e+02, rtol=1.0e-6)
    @test isapprox(output[8].d[7], 8.8967757e+04, rtol=1.0e-6)
    @test isapprox(output[8].d[8], 1.9797408e+06, rtol=1.0e-6)
    @test isapprox(output[8].d[9], 9.1218149e+03, rtol=1.0e-6)
    @test isapprox(output[8].d[6], 1.5718887e-15, rtol=1.0e-6)

    # Test Case 9
    @test isapprox(output[9].t[1], 1306.05, atol=1.0e-2)
    @test isapprox(output[9].t[2], 1293.37, atol=1.0e-2)
    @test isapprox(output[9].d[1], 6.7483388e+05, rtol=1.0e-6)
    @test isapprox(output[9].d[2], 1.2453153e+08, rtol=1.0e-6)
    @test isapprox(output[9].d[3], 2.3690095e+07, rtol=1.0e-6)
    @test isapprox(output[9].d[4], 4.9115832e+05, rtol=1.0e-6)
    @test isapprox(output[9].d[5], 4.5787811e+03, rtol=1.0e-6)
    @test isapprox(output[9].d[7], 3.2445948e+04, rtol=1.0e-6)
    @test isapprox(output[9].d[8], 5.3708331e+06, rtol=1.0e-6)
    @test isapprox(output[9].d[9], 2.6672732e+04, rtol=1.0e-6)
    @test isapprox(output[9].d[6], 4.5644202e-15, rtol=1.0e-6)

    # Test Case 10
    @test isapprox(output[10].t[1], 1361.87, atol=1.0e-2)
    @test isapprox(output[10].t[2], 1347.39, atol=1.0e-2)
    @test isapprox(output[10].d[1], 5.5286008e+05, rtol=1.0e-6)
    @test isapprox(output[10].d[2], 1.1980413e+08, rtol=1.0e-6)
    @test isapprox(output[10].d[3], 3.4957978e+07, rtol=1.0e-6)
    @test isapprox(output[10].d[4], 9.3396184e+05, rtol=1.0e-6)
    @test isapprox(output[10].d[5], 1.0962548e+04, rtol=1.0e-6)
    @test isapprox(output[10].d[7], 2.6864279e+04, rtol=1.0e-6)
    @test isapprox(output[10].d[8], 4.8899742e+06, rtol=1.0e-6)
    @test isapprox(output[10].d[9], 2.8054448e+04, rtol=1.0e-6)
    @test isapprox(output[10].d[6], 4.9745431e-15, rtol=1.0e-6)

    # Test Case 11
    @test isapprox(output[11].t[1], 1027.32, atol=1.0e-2)
    @test isapprox(output[11].t[2], 281.46, atol=1.0e-2)
    @test isapprox(output[11].d[1], 1.3754876e+14, rtol=1.0e-6)
    @test isapprox(output[11].d[2], 0.0000000e+00, rtol=1.0e-6)
    @test isapprox(output[11].d[3], 2.0496870e+19, rtol=1.0e-6)
    @test isapprox(output[11].d[4], 5.4986954e+18, rtol=1.0e-6)
    @test isapprox(output[11].d[5], 2.4517332e+17, rtol=1.0e-6)
    @test isapprox(output[11].d[7], 0.0000000e+00, rtol=1.0e-6)
    @test isapprox(output[11].d[8], 0.0000000e+00, rtol=1.0e-6)
    @test isapprox(output[11].d[9], 0.0000000e+00, rtol=1.0e-6)
    @test isapprox(output[11].d[6], 1.2610657e-03, rtol=1.0e-6)

    # Test Case 12
    @test isapprox(output[12].t[1], 1027.32, atol=1.0e-2)
    @test isapprox(output[12].t[2], 227.42, atol=1.0e-2)
    @test isapprox(output[12].d[1], 4.4274426e+13, rtol=1.0e-6)
    @test isapprox(output[12].d[2], 0.0000000e+00, rtol=1.0e-6)
    @test isapprox(output[12].d[3], 6.5975672e+18, rtol=1.0e-6)
    @test isapprox(output[12].d[4], 1.7699293e+18, rtol=1.0e-6)
    @test isapprox(output[12].d[5], 7.8916800e+16, rtol=1.0e-6)
    @test isapprox(output[12].d[7], 0.0000000e+00, rtol=1.0e-6)
    @test isapprox(output[12].d[8], 0.0000000e+00, rtol=1.0e-6)
    @test isapprox(output[12].d[9], 0.0000000e+00, rtol=1.0e-6)
    @test isapprox(output[12].d[6], 4.0591394e-04, rtol=1.0e-6)

    # Test Case 13
    @test isapprox(output[13].t[1], 1027.32, atol=1.0e-2)
    @test isapprox(output[13].t[2], 237.44, atol=1.0e-2)
    @test isapprox(output[13].d[1], 2.1278288e+12, rtol=1.0e-6)
    @test isapprox(output[13].d[2], 0.0000000e+00, rtol=1.0e-6)
    @test isapprox(output[13].d[3], 3.1707906e+17, rtol=1.0e-6)
    @test isapprox(output[13].d[4], 8.5062798e+16, rtol=1.0e-6)
    @test isapprox(output[13].d[5], 3.7927411e+15, rtol=1.0e-6)
    @test isapprox(output[13].d[7], 0.0000000e+00, rtol=1.0e-6)
    @test isapprox(output[13].d[8], 0.0000000e+00, rtol=1.0e-6)
    @test isapprox(output[13].d[9], 0.0000000e+00, rtol=1.0e-6)
    @test isapprox(output[13].d[6], 1.9508222e-05, rtol=1.0e-6)

    # Test Case 14
    @test isapprox(output[14].t[1], 1027.32, atol=1.0e-2)
    @test isapprox(output[14].t[2], 279.56, atol=1.0e-2)
    @test isapprox(output[14].d[1], 1.4121835e+11, rtol=1.0e-6)
    @test isapprox(output[14].d[2], 0.0000000e+00, rtol=1.0e-6)
    @test isapprox(output[14].d[3], 2.1043696e+16, rtol=1.0e-6)
    @test isapprox(output[14].d[4], 5.6453924e+15, rtol=1.0e-6)
    @test isapprox(output[14].d[5], 2.5171417e+14, rtol=1.0e-6)
    @test isapprox(output[14].d[7], 0.0000000e+00, rtol=1.0e-6)
    @test isapprox(output[14].d[8], 0.0000000e+00, rtol=1.0e-6)
    @test isapprox(output[14].d[9], 0.0000000e+00, rtol=1.0e-6)
    @test isapprox(output[14].d[6], 1.2947090e-06, rtol=1.0e-6)

    # Test Case 15
    @test isapprox(output[15].t[1], 1027.32, atol=1.0e-2)
    @test isapprox(output[15].t[2], 219.07, atol=1.0e-2)
    @test isapprox(output[15].d[1], 1.2548844e+10, rtol=1.0e-6)
    @test isapprox(output[15].d[2], 0.0000000e+00, rtol=1.0e-6)
    @test isapprox(output[15].d[3], 1.8745328e+15, rtol=1.0e-6)
    @test isapprox(output[15].d[4], 4.9230510e+14, rtol=1.0e-6)
    @test isapprox(output[15].d[5], 2.2396854e+13, rtol=1.0e-6)
    @test isapprox(output[15].d[7], 0.0000000e+00, rtol=1.0e-6)
    @test isapprox(output[15].d[8], 0.0000000e+00, rtol=1.0e-6)
    @test isapprox(output[15].d[9], 0.0000000e+00, rtol=1.0e-6)
    @test isapprox(output[15].d[6], 1.1476677e-07, rtol=1.0e-6)
end

let
    # Date and conversions
    epc = Epoch(2018, 1, 1, 12, 0, 0)
    oe = [R_EARTH + 500e3, 0.001, 97.8, 75, 25, 45]
    eci  = sOSCtoCART(oe, use_degrees=true)
    ecef = sECItoECEF(epc, eci)
    geod = sECEFtoGEOD(ecef, use_degrees=true)

    rho = density_nrlmsise00(epc, geod, use_degrees=true)

    @test 0.0 <= rho <= 1.0e-13
end


text:
using UniqueVectors: UniqueVector
using SparseArrays: spzeros, spdiagm
using LinearAlgebra: I, eigen, svd, diagm
using IterativeSolvers: lobpcg

export PSOModel
"""
    PSOModel{T, U}(K::U, G::U, C::U, P::U, Q::U, ports::UniqueVector{T}) where {T, U}
    PSOModel(K::U, G::U, C::U, P::U, Q::U, ports::AbstractVector{T}) where {T, U}

A Positive Second Order model, representing the equations `(K + G∂ₜ + C∂²ₜ)Φ = Px`,
`y = Qᵀ∂ₜΦ`.
"""
struct PSOModel{T, U<:AbstractMatrix{Float64}} <: AdmittanceModel{T}
    K::U
    G::U
    C::U
    P::U
    Q::U
    ports::UniqueVector{T}
    function PSOModel{T, U}(K::U, G::U, C::U, P::U, Q::U,
        ports::UniqueVector{T}) where {T, U}
        l = size(K, 1)
        @assert size(P) == (l, length(ports))
        @assert size(Q) == (l, length(ports))
        @assert all([size(y) == (l, l) for y in [K, G, C]])
        return new{T, U}(K, G, C, P, Q, ports)
    end
end

PSOModel(K::U, G::U, C::U, P::U, Q::U, ports::AbstractVector{T}) where {T, U} =
    PSOModel{T, U}(K, G, C, P, Q, UniqueVector(ports))

get_Y(pso::PSOModel) = [pso.K, pso.G, pso.C]

get_P(pso::PSOModel) = pso.P

get_Q(pso::PSOModel) = pso.Q

get_ports(pso::PSOModel) = pso.ports

function partial_copy(pso::PSOModel{T, U};
    Y::Union{Vector{V}, Nothing}=nothing,
    P::Union{V, Nothing}=nothing,
    Q::Union{V, Nothing}=nothing,
    ports::Union{AbstractVector{W}, Nothing}=nothing) where {T, U, V, W}
    Y = Y == nothing ? get_Y(pso) : Y
    P = P == nothing ? get_P(pso) : P
    Q = Q == nothing ? get_Q(pso) : Q
    ports = ports == nothing ? get_ports(pso) : ports
    return PSOModel(Y..., P, Q, ports)
end

compatible(psos::AbstractVector{PSOModel{T, U}}) where {T, U} = true

function canonical_gauge(pso::PSOModel)
    F = qr(pso.P)
    (m, n) = size(pso.P)
    transform = collect(F.Q * Matrix(I,m,m)) * [transpose(inv(collect(F.R))) zeros(n, m-n);
        zeros(m-n, n) Matrix(I, m-n, m-n)]
    return apply_transform(pso, transform)
end

# NOTE we cannot write lossy_modes_sparse because KrylovKit.geneigsolve
# solves Av = λBv but requires that B is positive semidefinite and A is
# symmetric or Hermitian. For lossy modes, we can make B positive semidefinite
# and A nonsymmetric or A symmetric and B not positive semidefinite. If
# KrylovKit gains the ability to solve either type of eigenproblem we can
# write lossy_modes_sparse.

export lossy_modes_dense
"""
    lossy_modes_dense(pso::PSOModel; min_freq::Real=0,
        max_freq::Union{Real, Nothing}=nothing, real_tol::Real=Inf, imag_tol::Real=Inf)

Use a dense eigenvalue solver to find the decaying modes of the PSOModel. Each mode consists
of an eigenvalue `λ` and a spatial distribution vector `v`. The frequency of the mode is
`imag(λ)/(2π)` and the decay rate is `-2*real(λ)/(2π)`. `min_freq` and `max_freq` give the
frequency range in which to find eigenvalues. Passivity is enforced using an inverse power
method solver, and `real_tol` and `imag_tol` are the tolerances used in that method to
determine convergence of the real and imaginary parts of the eigenvalue.
"""
function lossy_modes_dense(pso::PSOModel; min_freq::Real=0,
    max_freq::Union{Real, Nothing}=nothing, real_tol::Real=Inf, imag_tol::Real=Inf)
    K, G, C = get_Y(pso)
    Z = zero(K)
    M = collect([C Z; Z I]) \ collect([-G -K; I Z])
    M_scale = maximum(abs.(M))
    # for some reason eigen(Y \ X) seems to perform better than eigen(X, Y)
    values, vectors = eigen(M/M_scale)
    values *= M_scale
    freqs = imag.(values)/(2π)
    if max_freq == nothing
        max_freq = maximum(freqs)
    end
    inds = findall(max_freq .>= freqs .>= min_freq)
    values, vectors = values[inds], vectors[:, inds]
    # enforce passivity using inverse power method
    pred(x, y) = ((real(y) <= 0) && (abs(real(x-y)) < real_tol) &&
        (abs(imag(x-y)) < imag_tol))
    for i in 1:length(values)
        if real(values[i]) > 0
            λ, v = inv_power_eigen(M, v0=vectors[:, i], λ0=values[i], pred=pred)
            values[i] = λ
            vectors[:, i] = v
        end
    end
    # the second half of the eigenvector is the spatial part
    vectors = vectors[size(K, 1)+1:end, :]
    # normalize the columns
    return values, vectors./sqrt.(sum(abs2.(vectors), dims=1))
end

export lossless_modes_dense
"""
    lossless_modes_dense(pso::PSOModel; min_freq::Real=0,
        max_freq::Union{Real, Nothing}=nothing)

Use a dense svd solver to find the modes of the PSOModel neglecting loss. Each mode consists
of an eigenvalue `λ` and an eigenvector `v`. The frequency of the mode is `imag(λ)/(2π)`
and the decay rate is `-2*real(λ)/(2π)`. `min_freq` and `max_freq` give the
frequency range in which to find eigenvalues.
"""
function lossless_modes_dense(pso::PSOModel; min_freq::Real=0,
    max_freq::Union{Real, Nothing}=nothing)
    K, _, C = get_Y(pso)
    K = .5 * (transpose(K) + K)
    C = .5 * (transpose(C) + C)
    sqrt_C, sqrt_K = sqrt(collect(C)), sqrt(collect(K))
    M = sqrt_C \ sqrt_K
    scale = maximum(abs.(M))
    U, S, Vt = svd(M/scale)
    freqs = S * scale/(2π)
    vectors = real.(sqrt_C \ U)
    if max_freq == nothing
        max_freq = maximum(freqs)
    end
    inds = findall(max_freq .>= freqs .>= min_freq)
    freqs, vectors = freqs[inds], vectors[:, inds]
    inds = sortperm(freqs)
    return freqs[inds] * 2π * 1im, vectors[:, inds]
end

export lossless_modes_sparse
"""
    lossless_modes_sparse(pso::PSOModel; num_modes=1, maxiter=200)

Use a sparse generalized eigenvalue solver to find the `num_modes` lowest frequency modes
of the PSOModel neglecting loss. Each mode consists of an eigenvalue `λ` and an eigenvector
`v`. The frequency of the mode is `imag(λ)/(2π)` and the decay rate is `-2*real(λ)/(2π)`.
"""
function lossless_modes_sparse(pso::PSOModel; num_modes=1, maxiter=200)
    K, _, C = get_Y(pso)
    K = (K + transpose(K))/2
    C = (C + transpose(C))/2
    K_scale, C_scale = maximum(abs.(K)), maximum(abs.(C))
    result = lobpcg(K/K_scale, C/C_scale, false, num_modes, maxiter=maxiter)
    values = result.λ * (K_scale/C_scale)
    vectors = result.X
    return real.(sqrt.(values)) * 1im, vectors
end

function port_matrix(circuit::Circuit{T}, port_edges::Vector{Tuple{T, T}}) where T
    coord_matrix = coordinate_matrix(circuit)
    if length(port_edges) > 0
        port_indices = [(findfirst(isequal(p[1]), circuit.vertices),
                         findfirst(isequal(p[2]), circuit.vertices)) for p in port_edges]
        return hcat([coord_matrix[:, i1] - coord_matrix[:, i2]
            for (i1, i2) in port_indices]...)
    else
        return spzeros(size(coord_matrix, 1), 0)
    end
end

circuit_to_pso_matrix(m::AbstractMatrix{<:Real}) =
    (-m + spdiagm(0=>sum(m, dims=2)[:,1]))[2:end, 2:end]

"""
    PSOModel(circuit::Circuit{T}, port_edges::Vector{Tuple{T, T}},
        port_names::AbstractVector{S}) where {T, S}

Create a PSOModel for a circuit with ports on given edges. The spanning tree where the
first vertex is chosen as ground and all other vertices are neighbors of ground is used.
"""
function PSOModel(circuit::Circuit{T}, port_edges::Vector{Tuple{T, T}},
    port_names::AbstractVector{S}) where {T, S}
    Y = [circuit_to_pso_matrix(m) for m in matrices(circuit)]
    P = port_matrix(circuit, port_edges)
    return PSOModel(Y..., P, P, port_names)
end

coordinate_transform(m::AbstractMatrix{<:Real}) = m[:,2:end] .- m[:,1]

function coordinate_transform(coord_matrix_from::AbstractMatrix{<:Real},
    coord_matrix_to::AbstractMatrix{<:Real})
    return coordinate_transform(coord_matrix_to)/coordinate_transform(coord_matrix_from)
end

"""
    PSOModel(circuit::Circuit{T}, port_edges::Vector{Tuple{T, T}},
        port_names::AbstractVector{S}, tree::SpanningTree{T}) where {T, S}

Create a PSOModel for a circuit with ports on given edges and using the given spanning tree.
"""
function PSOModel(circuit::Circuit{T}, port_edges::Vector{Tuple{T, T}},
    port_names::AbstractVector{S}, tree::SpanningTree{T}) where {T, S}
    pso = PSOModel(circuit, port_edges, port_names)
    transform = transpose(coordinate_transform(coordinate_matrix(circuit, tree)))
    return apply_transform(pso, transform)
end

ext: jl


text:
using Revise
import Optim
using ADCME
mpi_init()

r = mpi_rank()
θ0 = Variable(rand(2))
θ = mpi_bcast(θ0)
a = 1.
b = 100.
x, y = θ[1], θ[2]
fxy = (a-x)^2 + b*(y-x^2)^2
local_loss = fxy * (r+1)
loss = mpi_sum(local_loss)
g = gradients(loss, θ0)

sess = Session(); init(sess)

_f_eval = 0
_g_eval = 0
function f(x0)
    global _f_eval += 1
    L = run(sess, loss, θ0=>x0)
    L
end

function g!(G, x0)
    global _g_eval += 1
    gv = run(sess, g, θ0=>x0)
    G[:] = gv
end

x0 = rand(2)
mpi_optimize(f, g!, x0)

mpi_finalize()


text:
using Distributions
using LinearAlgebra
using Plots

#This file is simply for visualizing the behavior of various kernel functions

function GaussianKernel(x, y, θ::Float64=1.0)
    x = vec(x)
    y = vec(y)
    M = 1/sqrt(2*pi)*exp.(-θ*.5*(x.-y').^2)
    return M
end

function generateSample(n, θ)
    x = collect(0:1/n:1-1/n)
    K = GaussianKernel(x, x, θ) + UniformScaling(1e-8) #add diagonal terms to ensure PD-ness
    C = cholesky(K)
    z = C.L*(randn(n))
    return x, z
end

x1, z1 = generateSample(10, 50.0)
x2, z2 = generateSample(50, 50.0)
x3, z3 = generateSample(150, 50.0)
x4, z4 = generateSample(30, 50.0)

display(plot(x1, z1, seriestype = :scatter))
display(plot!(x2, z2, seriestype = :scatter))
display(plot!(x3, z3, seriestype = :scatter))
display(plot!(x4, z4, seriestype = :scatter))

max_stars_repo_name: dbindel/btg

max_issues_repo_name: dbindel/btg

max_forks_repo_name: dbindel/btg


text:
using Pkg
Pkg.activate(@__DIR__)
pkg"instantiate"
pkg"precompile"

using Distributed
@everywhere using Pkg
@everywhere Pkg.activate(@__DIR__)

# using Revise

using ArgParse
using Caesar, RoME, DistributedFactorGraphs
using Dates
using Glob
using Interpolations

include(joinpath(@__DIR__, "CommonUtils.jl"))
include(joinpath(@__DIR__, "Plotting.jl"))

pargs = if !isdefined(Main, :parsed_args)
  parse_commandline()
else
  parsed_args
end

pargs["reportDir"] = if !haskey(pargs, "reportDir") && isdefined(Main, :fg)
  getLogPath(fg)
else
  # "/tmp/caesar/2020-02-22T02:21:20.777/fg_after_x781.tar.gz"
  # "/tmp/caesar/2020-02-23T01:43:32.222/fg_after_x1391.tar.gz"
  pargs["reportDir"]
end

# @show pargs["reportDir"]
# @show splitpath(pargs["reportDir"])

# load the factor graph needed
fg = if !isdefined(Main, :fg)
  println("going to load fg")
  fg = LightDFG{SolverParams}(params=SolverParams())
  @show pathElem = splitpath(pargs["reportDir"])
  @show getSolverParams(fg).logpath = joinpath(pathElem[1:end-1]...)
  loadDFG(pargs["reportDir"], Main, fg)
  ensureAllInitialized!(fg)
  fg
else
  fg
end

wtdsh = if !isdefined(Main, :wtdsh)
  # load wtdsh and dashboard
  @show wtFile = string(joinLogPath(fg, "dash.json"))
  wtdshA = readdlm(joinLogPath(fg, "dash.csv"), ',')
  # make array of array
  wtdsh = []
  for i in 1:size(wtdshA,1)
    push!(wtdsh, wtdshA[i,:])
  end
  wtdsh
else
  wtdsh
end

## draw fg
if pargs["drawFG"]
  drawGraph(fg, filepath=joinpath(getLogPath(fg),"fg.pdf"), show=false, engine="sfdp")
end

##  Draw trajectory & Analyze solve run

plb = plotSandsharkFromDFG(fg, drawTriads=false)
plb |> PDF(joinpath(getLogPath(fg),"traj.pdf"))
# pla = drawPosesLandmarksAndOdo(fg, ppbrDict, navkeys, X, Y, lblkeys, lblX, lblY)

plc = plotFrontendTiming(wtdsh)
plc |> PDF(joinpath(getLogPath(fg),"timing.pdf"),75cm,25cm)
# accumulated real time slack
wtt = (x->x[6]).(wtdsh)
wtt[end]

# set all to solvable=1

map(x->setSolvable!(fg, x, 1), setdiff(ls(fg), ls(fg, r"drt_\d")))
map(x->setSolvable!(fg, x, 0), union(ls(fg, r"drt\d"),lsf(fg, r"drt"))  )
map(x->setSolvable!(fg, x, 1 ),[lsf(fg, Pose2Pose2);
                                lsf(fg, Pose2Point2Range);
                                lsfPriors(fg)] )

# Check how long not solvable tail is?
# map(x->(x,isSolvable(getVariable(fg,x))), getLastPoses(fg, filterLabel=r"x\d", number=15))

# after all initialized
plb = plotSandsharkFromDFG(fg, drawTriads=false, lbls=false)

## add reference layer
posesyms = ls(fg, r"x\d") |> sortDFG
filter!(x->isInitialized(fg, x), posesyms)
filter!(x->getSolverData(getVariable(fg, x), :lbl) != nothing, posesyms)
XX = (x->(getSolverData(getVariable(fg, x), :lbl).val[1,1])).(posesyms)
YY = (x->(getSolverData(getVariable(fg, x), :lbl).val[2,1])).(posesyms)
pl = Gadfly.layer(x=XX, y=YY, Geom.path, Theme(default_color=colorant"magenta"))
union!(plb.layers, pl)
plb |> PDF(joinLogPath(fg,"traj_ref.pdf"))

XXlbl = deepcopy(XX)
YYlbl = deepcopy(YY)

## plot the DEAD RECKON THREAD

drt_data, TTm, XXm, YYm, XXf, YYf = loadResultsDRT(fg)

mask = YYf .< -32
XXfm = XXf[mask]
YYfm = YYf[mask]
# only start does not match
TTmm = sum(mask)<length(TTm) ? TTm[end-sum(mask)+1:end] : TTm

# plot DRT
pl = Gadfly.plot(x=XXm[1:10:end], y=YYm[1:10:end], Geom.point, Theme(default_color=colorant"khaki"))
pl |> PDF(joinLogPath(fg,"drt.pdf"))
pl = Gadfly.plot(x=XXf[1:10:end], y=YYf[1:10:end], Geom.point, Theme(default_color=colorant"green"))
pl |> PDF(joinLogPath(fg,"drtf.pdf"))

pl = Gadfly.layer(x=XXfm, y=YYfm, Geom.path, Theme(default_color=colorant"green"))
union!(plb.layers, pl)
pl = Gadfly.layer(x=XXm[mask], y=YYm[mask], Geom.path, Theme(default_color=colorant"khaki"))
union!(plb.layers, pl)
plb |> PDF(joinLogPath(fg,"traj_ref_drt.pdf"))

odo_data = readdlm(joinLogPath(fg, "ODO.csv"), ',')

XX = Float64.(odo_data[:,3])
YY = Float64.(odo_data[:,4])

# remove those near zero
mask = 40.0 .< (XX.^2 + YY.^2)

XXf = XX #[mask]
YYf = YY #[mask]

pl = Gadfly.plot(x=XXf[1:10:end], y=YYf[1:10:end], Geom.point, Theme(default_color=colorant"black"))
pl |> PDF(joinLogPath(fg,"odo.pdf"))

pl = Gadfly.layer(x=XXf, y=YYf, Geom.path, Theme(default_color=colorant"black"))
union!(plb.layers, pl)
plb |> PDF(joinLogPath(fg,"traj_ref_drt_odo.pdf"))

odo_data = readdlm(joinLogPath(fg, "RAWODO.csv"), ',')

DX = Float64.(odo_data[:,3:5])

nXYT = devAccumulateOdoPose2(DX, getFactorType(getFactor(fg, :x0f1)).Z.μ)

# remove those near zero for better visualization
mask = 40.0 .< (nXYT[:,1].^2 + nXYT[:,2].^2)

XXf = nXYT[:,1][mask]
YYf = nXYT[:,2][mask]

# pl = Gadfly.plot(x=XXf, y=YYf, Geom.path, Theme(default_color=colorant"black"))
pl = plotTrajectoryArrayPose2(nXYT)
pl |> PDF(joinLogPath(fg,"dirodo.pdf"))

pl = Gadfly.layer(x=XXf, y=YYf, Geom.path, Theme(default_color=colorant"black"))
union!(plb.layers, pl)
plb |> PDF(joinLogPath(fg,"traj_ref_drt_dirodo.pdf"))

## summarize errors

posesyms = ls(fg, r"x\d") |> sortDFG
filter!(x->isInitialized(fg, x), posesyms)
filter!(x->getSolverData(getVariable(fg, x), :lbl) != nothing, posesyms)
Xlbl = (x->(getSolverData(getVariable(fg, x), :lbl).val[1,1])).(posesyms)
Ylbl = (x->(getSolverData(getVariable(fg, x), :lbl).val[2,1])).(posesyms)

# pose PPE
Xppe = (x->(getVariablePPE(getVariable(fg, x)).suggested[1])).(posesyms)
Yppe = (x->(getVariablePPE(getVariable(fg, x)).suggested[2])).(posesyms)

# timestamp
ts = (x->getTimestamp(getVariable(fg, x))).(posesyms) .|> datetime2unix
T0 = ts[1]
ts .-= ts[1]
ts[end]

# filter on timestamps
# remove 420-450s
mask = 420 .< ts .< 450
imask = xor.(mask, true)

# from DRT
drtt = datetime2unix.(TTmm)
drtt .-= drtt[1]

itpx = LinearInterpolation(drtt, XXfm)
itpy = LinearInterpolation(drtt, YYfm)

itplblx = LinearInterpolation(ts[imask], Xlbl[imask])
itplbly = LinearInterpolation(ts[imask], Ylbl[imask])

plx = Gadfly.plot(
  Gadfly.layer(x=ts[1:end-1], y=Xppe[1:end-1], Geom.line, Theme(default_color=colorant"black")),
  Gadfly.layer(x=ts[1:end-1], y=itpx.(ts[1:end-1]), Geom.line, Theme(default_color=colorant"red")),
  Gadfly.layer(x=ts[1:end-1], y=itplblx.(ts[1:end-1]), Geom.line, Theme(default_color=colorant"magenta")),
  Guide.xlabel("time [s]"), Guide.ylabel("East [m]"),
)

ply = Gadfly.plot(
  Gadfly.layer(x=ts[1:end-1], y=Yppe[1:end-1], Geom.line, Theme(default_color=colorant"black")),
  Gadfly.layer(x=ts[1:end-1], y=itpy.(ts[1:end-1]), Geom.line, Theme(default_color=colorant"red")),
  Gadfly.layer(x=ts[1:end-1], y=itplbly.(ts[1:end-1]), Geom.line, Theme(default_color=colorant"magenta")),
  Guide.xlabel(""), Guide.ylabel("North [m]"),
)

plx |> PDF(joinLogPath(fg, "drtOverlayx.pdf"),12cm,5cm)
ply |> PDF(joinLogPath(fg, "drtOverlayy.pdf"),12cm,5cm)

drtErr = (Xppe[1:end-1]-itpx.(ts[1:end-1])).^2 + (Yppe[1:end-1]-itpy.(ts[1:end-1])).^2 .|> sqrt

ple = Gadfly.plot(
  Gadfly.layer(x=ts[1:end-1], y=drtErr, Geom.line, Theme(default_color=colorant"black")),
)

errX = Xppe[1:end-1]-itpx.(ts[1:end-1])
errY = Yppe[1:end-1]-itpy.(ts[1:end-1])
plhx = Gadfly.plot(
  Gadfly.layer(x=errX, Geom.histogram(density=true)),
  Guide.xlabel("x [m]"), Guide.ylabel("density"), Guide.title("DRT error, σ=$(round(std(errX),digits=3))"),
) # , Theme(default_color=colorant"black")
plhy = Gadfly.plot(
  Gadfly.layer(x=errY, Geom.histogram(density=true)),
  Guide.xlabel("y [m]"), Guide.title("DRT error, σ=$(round(std(errY),digits=3))"),
) # , Theme(default_color=colorant"black")

plhx |> PDF(joinLogPath(fg, "drtErrHistx.pdf"),6cm,5cm)
plhy |> PDF(joinLogPath(fg, "drtErrHisty.pdf"),6cm,5cm)

## plot densities

include(joinpath(@__DIR__, "PrepMakieBackground.jl"))

## more conventional

include(joinpath(@__DIR__, "MakiePlotsFG.jl"))

# nvsyms = ls(fg, r"x\d") |> length

# how to suppress window and simply export
pl, Z = plotVariableBeliefs(fg, r"x\d", sortVars=true, fade=10, fadeFloor=0.2, resolution=(1920,1080), colormap=:blues)
# pl |> typeof |> fieldnames
pl

Base.rm(joinLogPath(fg,"fgBeliefs.png"), force=true)
Makie.save(joinLogPath(fg,"fgBeliefs.png"), pl)

addLinesBelief!(fg, pl, TTm, drt=false, ppe=true, ref=true)

Base.rm(joinLogPath(fg,"fgBeliefsLines.png"), force=true)
Makie.save(joinLogPath(fg,"fgBeliefsLines.png"), pl)

# pargs["plotSeriesBeliefs"] = 3
0

# plot a series of frames
if 0 <= pargs["plotSeriesBeliefs"]

pattern = "fg_x"
ext = ".tar.gz"
files = glob("$(pattern)*$ext", getLogPath(fg))
indiv = splitpath.(files) .|> x->x[end]
presort = (x->x[length(pattern)+1:end-length(ext)]).(indiv)
sorted = parse.(Int, presort) |> sortperm
# reduce the number of frames, if requested
indiv = pargs["plotSeriesBeliefs"] == 0 ? indiv[sorted] : indiv[sorted[1:pargs["plotSeriesBeliefs"]]]
# get range from last file in glob
fgl = LightDFG{SolverParams}(params=SolverParams())
loadDFG(joinLogPath(fg, indiv[end]), Main, fgl)
drt_data, TTm, XXm, YYm, XXf, YYf = loadResultsDRT(fg)
minmax = getRangeCartesian(fgl,r"x\d",digits=-1)
xmin=minmax[1,1]-10;xmax=minmax[1,2]+10;ymin=minmax[2,1]-10;ymax=minmax[2,2]+10;
Base.mkpath(joinLogPath(fg, "frames"))
Base.mkpath(joinLogPath(fg, "lines"))
global frame = 0
for ind in indiv
  global frame += 1
  println("frame $frame of $(length(indiv))")
  fname = split(ind, '.')[1]
  if pargs["skip"] && isfile(joinLogPath(fg,"frames/$fname.png")) && isfile(joinLogPath(fg,"lines/$fname.png"))
    @info "skip $fname since its already there"
    continue
  end
  fgl = LightDFG{SolverParams}(params=SolverParams())
  loadDFG(joinLogPath(fg, ind), Main, fgl)
  nvars = length(ls(fgl, r"x\d"))
  # if nvars < 15
  #   continue
  # end
  try
    pl, Z = plotVariableBeliefs(fgl, r"x\d", solvable=1, sortVars=true, fade=minimum([15; nvars]), fadeFloor=0.2,
                              xmin=xmin,xmax=xmax,ymin=ymin,ymax=ymax,
                              resolution=(1920, 1080) )
    #
    Base.rm(joinLogPath(fg,"frames/$fname.png"), force=true)
    Makie.save(joinLogPath(fg,"frames/$fname.png"), pl)

    # draw lines
    addLinesBelief!(fgl, pl, TTm)
    Base.rm(joinLogPath(fg,"lines/$fname.png"), force=true)
    Makie.save(joinLogPath(fg,"lines/$fname.png"), pl)
  catch ex
    @error ex
  end
  GC.gc()
end

end

## Look at factor reports separately
if pargs["reportRanges"]
  reportFactors(fg, Pose2Point2Range, show=false)
end

if pargs["reportPoses"]
  reportFactors(fg, Pose2Pose2, show=false)
end

0

## Develop ranges alongside

if pargs["drawAllRanges"]

  mkpath(joinLogPath(fg, "ranges"))

  files = glob("fg_x*.tar.gz", getLogPath(fg))
  indiv = splitpath.(files) .|> x->x[end]

  # vsym = ls(fg, r"x\d") |> sortDFG

  PP2 = ls(fg, r"x\d")
  fcPP2R = ls(fg, Pose2Point2Range)

  # j=10
  for ind in indiv
    fname = split(ind, '.')[1]
    fgl = LightDFG{SolverParams}(params=SolverParams())
    loadDFG(joinLogPath(fg, ind), Main, fgl)
    vsym = ls(fgl, r"x\d") |> sortDFG
    filter!(x->isSolved(fgl, x), vsym)
    j = length(vsym)
    if j < 10
      continue
    end
  # for j in 10:length(vsym)
    for i in 1:10
      global latestPP2R = intersect(union(map(v->ls(fgl,v),vsym[j-i:j])...), fcPP2R)
      0 < length(latestPP2R) ? break : nothing
    end
    if length(latestPP2R) == 0
      continue
    end
    rangePose = intersect(ls(fgl, latestPP2R[1]), PP2)[1]
    hdl = []
    plotFactor(fgl, latestPP2R[1], hdl=hdl)
    hdl[4].coord = Coord.Cartesian(xmin=50, xmax=300)
    hdl[4] |> PNG(joinLogPath(fg, "ranges/$rangePose.png"))
  end

end

0

## Plot reference poses

#
# fg, dashboard = main(logSpeed=0.25, lcm=LCMLog(joinpath(ENV["HOME"],"data","sandshark","lcmlog","lcmlog-2019-11-26.01")) ) # bad ranges

# find . -type f -name '*.pdf' -print0 |
#   while IFS= read -r -d '' file
#     do convert -verbose -density 500 -quality 100 "${file}" "png/${file%.*}.png"
#   done

# ls nobytes/ | sed 's/bt_//g' | sed 's/.pdf//g' | xargs -L1


text:
# Use baremodule to shave off a few KB from the serialized `.ji` file
baremodule CUTENSOR_jll
using Base
using Base: UUID
using LazyArtifacts
import JLLWrappers

JLLWrappers.@generate_main_file_header("CUTENSOR")
JLLWrappers.@generate_main_file("CUTENSOR", UUID("35b6c64b-1ee1-5834-92a3-3f624899209a"))
end  # module CUTENSOR_jll


text:
module ContourWrapper

    export contoursw, levels, lines, coordinates

    using Contour

    function contoursw(args...)
        contours(args...)
    end
end


text:
# by lazarusA & Julius Krumbiegel
# using GLMakie # HIDE
using Makie, CairoMakie, HTTP, CSV, DataFrames, DataFramesMeta, Suppressor
using Images, ColorSchemes, Colors, Statistics
using Lazy: @>
CairoMakie.activate!()
let
    data_url = "https://bit.ly/3kmlGn2"
    @suppress begin
        global astronauts
        astronauts = CSV.File(HTTP.download(data_url)) |> DataFrame
    end
    offhr, rPts, α = 3.5, 45, 35
    astro = @> begin
        astronauts
        @orderby(:year_of_mission)
        @where(:total_eva_hrs .> 0.0)
        @transform(ntotEVA = :total_eva_hrs/ maximum(:total_eva_hrs))
        @transform(θ = LinRange(0, 2π - 2π/length(:name), length(:name)))
        @transform(align = tuple.(ifelse.(π/2 .< :θ .< 3π/2, ^(:right), ^(:left)), ^(:center)))
        @transform(texttheta = ifelse.(π/2 .< :θ .< 3π/2, :θ .+ π, :θ))
        @transform(evaM = log10.((60 * :eva_hrs_mission/median(60 * :eva_hrs_mission)) .+ offhr))
        @transform(xM = rPts * :evaM .* cos.(:θ), yM = rPts * :evaM .* sin.(:θ))
        @transform(xMnm = rPts * (:evaM .- :total_number_of_missions .* :evaM./α) .* cos.(:θ))
        @transform(yMnm = rPts * (:evaM .- :total_number_of_missions .* :evaM./α) .* sin.(:θ))
    end
    valYear = @> begin astro
            @where([true; :year_of_mission[2:end] .!= :year_of_mission[1:end-1]])
        end
    vehicles = @> begin astro
            @where([true; :ascend_shuttle[2:end] .!= :ascend_shuttle[1:end-1]])
        end
    tierra = "https://climate.nasa.gov/system/internal_resources/details/original/309_ImageWall5_768px-60.jpg"
    @suppress begin
        global imgEarth
        tierra = HTTP.download(tierra)
        imgEarth = load(tierra)
    end
    function getPoints(xi,yi,xf,yf)
        xyos = []
        for i in 1:length(xo)
            push!(xyos, [xi[i], yi[i]])
            push!(xyos, [xf[i], yf[i]])
        end
        xyos
    end
    x, y = astro.xM, astro.yM # end points
    xs, ys = astro.xMnm, astro.yMnm # short lines starts
    xo = yo = zeros(length(x)) # origin
    xnb, ynb = 90 * cos.(astro.θ), 90 * sin.(astro.θ)
    xne, yne = 100 * cos.(astro.θ), 100 * sin.(astro.θ)
    ps, colorp = 0.5 .+ 3astro.ntotEVA, astro.total_eva_hrs ; # point size, color palette
    gridLines = LinRange(log10(offhr), maximum(astro.evaM), 6)
    horas = (10 .^ gridLines .- offhr)*median(60*astro.eva_hrs_mission)/60
    xg = [rPts * gl .* cos.(astro.θ) for gl in gridLines]
    yg = [rPts * gl .* sin.(astro.θ) for gl in gridLines]
    # in order to use linesegments (faster to plot)
    xyos = getPoints(xo,yo,xs,ys)
    xys  = getPoints(xs,ys,x,y)
    xybe = getPoints(xnb,ynb,xne,yne)
    cpDo = repeat(colorp, inner=2)
    psDo = repeat(ps, inner=2);

    with_theme(theme_black()) do
        fig = Figure(resolution = (1200, 1200))
        cmap = :rainbow2
        ax = CairoMakie.Axis(fig[1, 1],
            title = "ASTRONAUTS' EXTRAVEHICULAR ACTIVITIES",
            autolimitaspect = 1)
        hidespines!(ax)
        hidedecorations!(ax)
        image!(-20..20, -17..17, rotr90(imgEarth))
        text!(astro.name, position = @.(Point(cos(astro.θ), sin(astro.θ)) * 85),
            rotation = astro.texttheta, textsize = 6, align = astro.align)
        text!(string.(valYear.year_of_mission),
            position = @.(Point(cos(valYear.θ), sin(valYear.θ)) * 65),
            rotation = valYear.texttheta, textsize = 10, align = valYear.align)
        text!(vehicles.ascend_shuttle,
            position = @.(Point(cos(vehicles.θ), sin(vehicles.θ)) * 73),
            rotation = vehicles.texttheta, textsize = 6, align = vehicles.align)

        pltobj = scatter!(ax, astro[:,:xM], astro[:,:yM], color = colorp,
            colormap = cmap, markersize = 3*ps, strokewidth = 0)

        linesegments!(Point2f0.(xyos), color = cpDo, linewidth = psDo/2, colormap = (cmap, 0.15))
        linesegments!(Point2f0.(xys), color = cpDo, linewidth = psDo/2,colormap = cmap)
        linesegments!(Point2f0.(xybe), color = cpDo, linewidth = psDo/3, colormap = (cmap,0.5))

        Colorbar(fig[1,1], pltobj,
            label = "Total duration of all extravehicular activities in hours",
            tellheight = false, tellwidth = false, ticklabelsize = 12, flipaxis = true,
            vertical = false, ticksize=15, tickalign = 1, width = Relative(1.5/4),
            halign = :right, valign = :bottom, labelsize = 12)

        for (indx, gl) in enumerate(gridLines)
            xg, yg = rPts * gl .* cos.(astro.θ), rPts * gl .* sin.(astro.θ)
            hrs = Int64(round(horas[indx], digits = 0))
            lines!(xg, yg, linewidth = 0.5, linestyle = :dash, color = :white)
            text!(string.(hrs), position = (xg[1] + 0.5, y[1]+0.5),
                color = "#FFDD33", textsize = 14)
        end
        lines!([rPts*gridLines[1], rPts*gridLines[end]],[0,0],linestyle = :dash,
            linewidth = 2, color =  "#FFDD33")
        text!("evaM (hrs)", position = (47, -3.5), color = "#FFDD33", textsize = 16)
        text!("evaM ≡ Duration of extravehicular \n activities during the mission in hours",
            position = (rPts*gridLines[end-2], 90), color = "#FFDD33", textsize = 16)
        text!("using Makie", position = (-99, -94), textsize = 18, color = :white)
        text!("Visualization by @LazarusAlon and Julius Krumbiegel ",
            position = (-99, -97), textsize = 12, color = "#61AFEF")
        text!("Data - Astronaut Database - Mariya Stavnichuk and Tatsuya Corlett",
            position = (-99, -99), textsize = 10)
        limits!(ax, -100, 100, -100, 100)
        save(joinpath(@__DIR__, "output", "theme_dark_astronauts.png"), fig, px_per_unit = 2) # HIDE
    end

end
using Pkg # HIDE
Pkg.status(["CairoMakie", "HTTP", "CSV", "DataFrames", "DataFramesMeta", # HIDE
    "Suppressor", "Images", "ColorSchemes","Colors", "Statistics", "Lazy" ]) # HIDE

max_stars_repo_name: MisterBiggs/BeautifulMakie

max_issues_repo_name: MisterBiggs/BeautifulMakie

max_forks_repo_name: MisterBiggs/BeautifulMakie


text:
type SCMMulti
    pathAS::Real          # standard deviation of Laplace
    nPaths::Integer      # number of paths

    SCMMulti(; pathAS=2.0, nPaths=3) = new(pathAS,nPaths)
end

function generate_channel(chan::SCMMulti, nAntennas; nCoherence=1,nBatches=1)
    h = zeros(Complex128,nAntennas, nCoherence, nBatches)
    t = zeros(Complex128,nAntennas, nBatches)
    for i in 1:nBatches
        gains = rand(chan.nPaths)
        gains = gains ./ sum(gains,1)
        angles = (rand(chan.nPaths)-0.5)*180 # path center: uniform [-90,90] deg

        (h[:,:,i],t[:,i]) = scm_channel(angles, gains, nAntennas; nCoherence=nCoherence,AS=chan.pathAS)
    end
    return (h,t)
end

max_stars_repo_name: tum-msv/learning-mmse-est


text:
# Subplot geometries

abstract SubplotGeometry <: Gadfly.GeometryElement

immutable SubplotLayer
    statistic::Gadfly.StatisticElement
    geom::Gadfly.GeometryElement

    function SubplotLayer(geom::Gadfly.GeometryElement=Geom.nil(),
                          statistic::Gadfly.StatisticElement=Stat.nil())
        new(statistic, geom)
    end
end

# Adding elements to subplots in a generic way.

function add_subplot_element(subplot::SubplotGeometry, arg::Function)
    add_subplot_element(subplot, arg())
end

function add_subplot_element(subplot::SubplotGeometry, arg::SubplotLayer)
    push!(subplot.layers, arg)
end

function add_subplot_element(subplot::SubplotGeometry, arg::Gadfly.GeometryElement)
    push!(subplot.layers, SubplotLayer(arg))
end

function add_subplot_element(subplot::SubplotGeometry, arg::Gadfly.StatisticElement)
    push!(subplot.statistics, arg)
end

function add_subplot_element(subplot::SubplotGeometry, arg::Gadfly.GuideElement)
    push!(subplot.guides, arg)
end

function add_subplot_element{T <: Gadfly.Element}(subplot::SubplotGeometry,
                                                  arg::Type{T})
    add_subplot_element(subplot, arg())
end

function add_subplot_element(subplot::SubplotGeometry, arg)
    error("Subplots do not support elements of type $(typeof(arg))")
end

immutable SubplotGrid <: SubplotGeometry
    layers::Vector{SubplotLayer}
    statistics::Vector{Gadfly.StatisticElement}
    guides::Vector{Gadfly.GuideElement}
    free_x_axis::Bool
    free_y_axis::Bool

    # Current plot has no way of passing existing aesthetics. It always produces
    # these using scales.
    function SubplotGrid(elements::Gadfly.ElementOrFunction...;
                         free_x_axis=false, free_y_axis=false)
        subplot = new(SubplotLayer[], Gadfly.StatisticElement[],
                      Gadfly.GuideElement[], free_x_axis, free_y_axis)

        for element in elements
            add_subplot_element(subplot, element)
        end

        # TODO: Handle default guides and statistics
        subplot
    end
end

const subplot_grid = SubplotGrid

function element_aesthetics(geom::SubplotGrid)
    vars = [:xgroup, :ygroup]
    for layer in geom.layers
        append!(vars, element_aesthetics(layer.geom))
    end
    vars
end

# Render a subplot grid geometry, which consists of rendering and arranging
# many smaller plots.
function render(geom::SubplotGrid, theme::Gadfly.Theme,
                superplot_aes::Gadfly.Aesthetics)
    if superplot_aes.xgroup === nothing && superplot_aes.ygroup === nothing
        error("Geom.subplot_grid requires \"xgroup\" and/or \"ygroup\" to be bound.")
    end

    # partition the each aesthetic into a matrix of aesthetics
    aes_grid = Gadfly.aes_by_xy_group(superplot_aes)
    n, m = size(aes_grid)

    coord = Coord.cartesian()
    scales = Dict{Symbol, Gadfly.ScaleElement}()
    plot_stats = Gadfly.StatisticElement[stat for stat in geom.statistics]
    layer_stats = Gadfly.StatisticElement[typeof(layer.statistic) == Stat.nil ?
                       Geom.default_statistic(layer.geom) : layer.statistic
                   for layer in geom.layers]

    layer_aes_grid = Array(Array{Gadfly.Aesthetics, 1}, n, m)
    for i in 1:n, j in 1:m
        layer_aes = fill(copy(aes_grid[i, j]), length(geom.layers))

        for (layer_stat, aes) in zip(layer_stats, layer_aes)
            Stat.apply_statistics(Gadfly.StatisticElement[layer_stat],
                                  scales, coord, aes)
        end

        plot_aes = cat(layer_aes...)
        Stat.apply_statistics(plot_stats, scales, coord, plot_aes)

        aes_grid[i, j] = plot_aes
        layer_aes_grid[i, j] = layer_aes
    end

    # apply geom-wide statistics
    geom_aes = cat(aes_grid...)
    geom_stats = Gadfly.StatisticElement[]

    if !geom.free_x_axis
        push!(geom_stats, Stat.xticks())
    end

    if !geom.free_y_axis
        push!(geom_stats, Stat.yticks())
    end

    Stat.apply_statistics(geom_stats, scales, coord, geom_aes)

    # if either axis is on a free scale, we need to apply row/column-wise
    # tick statistics.
    if (geom.free_x_axis)
        for j in 1:m
            col_aes = cat([aes_grid[i, j] for i in 1:n]...)
            Stat.apply_statistic(Stat.xticks(), scales, coord, col_aes)
            for i in 1:n
                aes_grid[i, j] = cat(aes_grid[i, j], col_aes)
            end
        end
    end

    if (geom.free_y_axis)
        for i in 1:n
            row_aes = cat([aes_grid[i, j] for j in 1:m]...)
            Stat.apply_statistic(Stat.yticks(), scales, coord, row_aes)
            for j in 1:m
                aes_grid[i, j] = cat(aes_grid[i, j], row_aes)
            end
        end
    end

    for i in 1:n, j in 1:m
        Gadfly.inherit!(aes_grid[i, j], geom_aes)
    end

    canvas_grid = Array(Canvas, n, m)

    xtitle = "x"
    for v in [:x, :xmin, :xmax]
        if haskey(superplot_aes.titles, v)
            xtitle = superplot_aes.titles[v]
            break
        end
    end

    ytitle = "y"
    for v in [:y, :ymin, :ymax]
        if haskey(superplot_aes.titles, v)
            ytitle = superplot_aes.titles[v]
            break
        end
    end

    xlabels = superplot_aes.xgroup_label(1.0:m)
    ylabels = superplot_aes.ygroup_label(1.0:n)

    for i in 1:n, j in 1:m
        p = Plot()
        p.theme = theme
        for layer in geom.layers
            plot_layer = Gadfly.Layer()
            plot_layer.statistic = layer.statistic
            plot_layer.geom = layer.geom
            push!(p.layers, plot_layer)
        end
        guides = Gadfly.GuideElement[guide for guide in geom.guides]

        # default guides
        push!(guides, Guide.background())

        if i == n
            push!(guides, Guide.xticks())
            if !is(superplot_aes.xgroup, nothing)
                push!(guides, Guide.xlabel(xlabels[j]))
            end
        else
            push!(guides, Guide.xticks(label=false))
        end

        if j == 1
            push!(guides, Guide.yticks())
            if !is(superplot_aes.ygroup, nothing)
                push!(guides, Guide.ylabel(ylabels[i]))
            end
        else
            push!(guides, Guide.yticks(label=false))
        end

        canvas_grid[i, j] =
            Gadfly.render_prepared(
                            p, aes_grid[i, j], layer_aes_grid[i, j],
                            layer_stats,
                            Dict{Symbol, Gadfly.ScaleElement}(),
                            plot_stats,
                            guides,
                            preserve_plot_canvas_size=true)

        canvas_grid[i, j] = pad(canvas_grid[i, j], 3.0mm)
    end

    gridstack(canvas_grid, 0w, 0h, halign=hright)
end


text:
# Configurations and contact forces
struct NewtonResidualConfigurationForce{T,vq2,vu1,vγ1,vb1,vd,vI,vq0,vq1} <: NewtonResidual
    r::Vector{T}                           # residual

    q2::Vector{vq2}                    # rsd objective views
    u1::Vector{vu1}                    # rsd objective views
    γ1::Vector{vγ1}                    # rsd objective views
    b1::Vector{vb1}                    # rsd objective views

    rd::Vector{vd}                         # rsd dynamics lagrange multiplier views
    rI::Vector{vI}                         # rsd dynamics -I views [q2, γ1, b1]

    q0::Vector{vq0}                        # rsd dynamics q0 views
    q1::Vector{vq1}                        # rsd dynamics q1 views
end

function NewtonResidualConfigurationForce(model::Model, env::Environment, H::Int)

    nq = model.nq # configuration
    nu = model.nu # control
    nc = model.nc # contact
    nb = nc * friction_dim(env) # linear friction
    nd = nq + nc + nb # implicit dynamics constraint
    nr = nq + nu + nc + nb# + nd # size of a one-time-step block

    off = 0
    iu = SizedVector{nu}(off .+ (1:nu)); off += nu # index of the control u1
    iγ = SizedVector{nc}(off .+ (1:nc)); off += nc # index of the impact γ1
    ib = SizedVector{nb}(off .+ (1:nb)); off += nb # index of the linear friction b1
    iq = SizedVector{nq}(off .+ (1:nq)); off += nq # index of the configuration q2

    iz = vcat(iq, iγ, ib) # index of the IP solver solution [q2, γ1, b1]

    iν = SizedVector{nd}(1:nd); #off += nd # index of the dynamics lagrange multiplier ν1

    r = zeros(H * (nr + nd))

    u1  = [view(r, (t - 1) * nr .+ iu) for t = 1:H]
    γ1  = [view(r, (t - 1) * nr .+ iγ) for t = 1:H]
    b1  = [view(r, (t - 1) * nr .+ ib) for t = 1:H]
    q2  = [view(r, (t - 1) * nr .+ iq) for t = 1:H]
    rI  = [view(r, (t - 1) * nr .+ iz) for t = 1:H]

    rd  = [view(r, H * nr + (t - 1) * nd .+ iν) for t = 1:H]

    q0  = [view(r, (t - 3) * nr .+ iq) for t = 3:H]
    q1  = [view(r, (t - 2) * nr .+ iq) for t = 2:H]

    T = eltype(r)

    return NewtonResidualConfigurationForce{T, eltype.((q2, u1, γ1, b1))...,eltype.((rd, rI, q0, q1))...}(
        r, q2, u1, γ1, b1, rd, rI, q0, q1)
end

# Configurations
struct NewtonResidualConfiguration{T,vq2,vu1,vd,vI,vq0,vq1} <: NewtonResidual
    r::Vector{T}                           # residual

    q2::Vector{vq2}                    # rsd objective views
    u1::Vector{vu1}                    # rsd objective views

    rd::Vector{vd}                         # rsd dynamics lagrange multiplier views
    rI::Vector{vI}                         # rsd dynamics -I views [q2]
    q0::Vector{vq0}                        # rsd dynamics q0 views
    q1::Vector{vq1}                        # rsd dynamics q1 views
end

function NewtonResidualConfiguration(model::Model, env::Environment, H::Int)
    nq = model.nq # configuration
    nu = model.nu # control
    nc = model.nc # contact
    nd = nq    # implicit dynamics constraint
    nr = nq + nu # size of a one-time-step block

    off = 0
    iu = SizedVector{nu}(off .+ (1:nu)); off += nu # index of the control u1
    iq = SizedVector{nq}(off .+ (1:nq)); off += nq # index of the configuration q2
    iz = copy(iq) # index of the IP solver solution [q2]

    iν = SizedVector{nd}(1:nd) # index of the dynamics lagrange multiplier ν1

    r = zeros(H * (nr + nd))

    u1  = [view(r, (t - 1) * nr .+ iu) for t = 1:H]
    q2  = [view(r, (t - 1) * nr .+ iq) for t = 1:H]
    rI  = [view(r, (t - 1) * nr .+ iz) for t = 1:H]

    rd  = [view(r, H * nr + (t - 1) * nd .+ iν) for t = 1:H]

    q0  = [view(r, (t - 3) * nr .+ iq) for t = 3:H]
    q1  = [view(r, (t - 2) * nr .+ iq) for t = 2:H]

    T = eltype(r)

    return NewtonResidualConfiguration{T, eltype.((q2, u1))...,eltype.((rd, rI, q0, q1))...}(
        r, q2, u1, rd, rI, q0, q1)
end

function NewtonResidual(model::Model, env::Environment, H::Int;
    mode = :configurationforce)

    if mode == :configurationforce
        return NewtonResidualConfigurationForce(model, env, H)
    elseif mode == :configuration
        NewtonResidualConfiguration(model, env, H)
    else
        @error "mode not implemented"
    end
end

function residual!(res::NewtonResidual, core::Newton,
    ν::Vector{Nν}, im_traj::ImplicitTrajectory, traj::ContactTraj, ref_traj::ContactTraj, window::Vector{Int}) where Nν

    # unpack
    opts = core.opts
    obj = core.obj
    res.r .= 0.0

    # Objective
    gradient!(res, obj, core, traj, ref_traj)

    for (i, t) in enumerate(window[1:end-2])
        # Lagrangian
        i >= 3 && mul!(res.q2[i-2], transpose(im_traj.δq0[t]), ν[i], 1.0, 1.0)
        i >= 2 && mul!(res.q2[i-1], transpose(im_traj.δq1[t]), ν[i], 1.0, 1.0)
        mul!(res.u1[i], transpose(im_traj.δu1[t]), ν[i], 1.0, 1.0)

        # Implicit dynamics
        res.rd[i] .+= im_traj.d[t]

        # Minus Identity term #∇qk1, ∇γk, ∇bk
        res.rI[i] .-= ν[i]
    end

    return nothing
end

function update_traj!(traj_cand::ContactTraj, traj::ContactTraj,
        ν_cand::Vector, ν::Vector, Δ::NewtonResidualConfigurationForce, α::T) where T

    H = traj_cand.H

    for t = 1:H
        traj_cand.q[t+2] .= traj.q[t+2] .- α .* Δ.q2[t]
        traj_cand.u[t] .= traj.u[t] .- α .* Δ.u1[t]
        traj_cand.γ[t] .= traj.γ[t] .- α .* Δ.γ1[t]
        traj_cand.b[t] .= traj.b[t] .- α .* Δ.b1[t]

        ν_cand[t] .= ν[t] .- α .* Δ.rd[t]
    end

    update_z!(traj_cand)
    update_θ!(traj_cand)

    return nothing
end

function update_traj!(traj_cand::ContactTraj, traj::ContactTraj,
        ν_cand::Vector, ν::Vector, Δ::NewtonResidualConfiguration, α::T) where T

    H = traj_cand.H

    for t = 1:H
        traj_cand.q[t+2] .= traj.q[t+2] .- α .* Δ.q2[t]
        traj_cand.u[t] .= traj.u[t] .- α .* Δ.u1[t]

        ν_cand[t] .= ν[t] .- α .* Δ.rd[t]
    end

    update_z!(traj_cand)
    update_θ!(traj_cand)

    return nothing
end

function gradient!(res::NewtonResidualConfigurationForce{T,vq2,vu1,vγ1,vb1,vd,vI,vq0,vq1}, obj::TrackingObjective{Q,U,C,B}, core::Newton{T,nq,nu,nw,nc,nb,nz,nθ,nν,NJ,NR,NI,O,LS}, traj::ContactTraj{T,nq,nu,nw,nc,nb,nz,nθ}, ref_traj::ContactTraj{T,nq,nu,nw,nc,nb,nz,nθ}) where {T,vq2,vu1,vγ1,vb1,vd,vI,vq0,vq1,Q,U,C,B,nq,nu,nw,nc,nb,nz,nθ,nν,NJ,NR,NI,O,LS}
    for t = 1:traj.H
        # Cost function
        delta!(core.Δq[t], traj.q[t+2], ref_traj.q[t+2])
        delta!(core.Δu[t], traj.u[t], ref_traj.u[t])
        delta!(core.Δγ[t], traj.γ[t], ref_traj.γ[t])
        delta!(core.Δb[t], traj.b[t], ref_traj.b[t])

        # res.q2[t] .+= obj.q[t] * core.Δq[t]
        mul!(res.q2[t], obj.q[t], core.Δq[t], 1.0, 1.0)
        # res.u1[t] .+= obj.u[t] * core.Δu[t]
        mul!(res.u1[t], obj.u[t], core.Δu[t], 1.0, 1.0)
        # res.γ1[t] .+= obj.γ[t] * core.Δγ[t]
        mul!(res.γ1[t], obj.γ[t], core.Δγ[t], 1.0, 1.0)
        # res.b1[t] .+= obj.b[t] * core.Δb[t]
        mul!(res.b1[t], obj.b[t], core.Δb[t], 1.0, 1.0)
    end
end

function gradient!(res::NewtonResidualConfiguration{T,vq2,vu1,vd,vI,vq0,vq1}, obj::TrackingObjective{Q,U,C,B}, core::Newton{T,nq,nu,nw,nc,nb,nz,nθ,nν,NJ,NR,NI,O,LS}, traj::ContactTraj{T,nq,nu,nw,nc,nb,nz,nθ}, ref_traj::ContactTraj{T,nq,nu,nw,nc,nb,nz,nθ}) where {T,vq2,vu1,vd,vI,vq0,vq1,Q,U,C,B,nq,nu,nw,nc,nb,nz,nθ,nν,NJ,NR,NI,O,LS}
    for t = 1:traj.H
        # Cost function
        delta!(core.Δq[t], traj.q[t+2], ref_traj.q[t+2])
        delta!(core.Δu[t], traj.u[t], ref_traj.u[t])

        # res.q2[t] .+= obj.q[t] * core.Δq[t]
        # res.u1[t] .+= obj.u[t] * core.Δu[t]
        mul!(res.q2[t], obj.q[t], core.Δq[t], 1.0, 1.0)
        mul!(res.u1[t], obj.u[t], core.Δu[t], 1.0, 1.0)
    end
end

function gradient!(res::NewtonResidualConfigurationForce{T,vq2,vu1,vγ1,vb1,vd,vI,vq0,vq1}, obj::TrackingVelocityObjective{Q,V,U,C,B}, core::Newton{T,nq,nu,nw,nc,nb,nz,nθ,nν,NJ,NR,NI,O,LS}, traj::ContactTraj{T,nq,nu,nw,nc,nb,nz,nθ}, ref_traj::ContactTraj{T,nq,nu,nw,nc,nb,nz,nθ}) where {T,vq2,vu1,vγ1,vb1,vd,vI,vq0,vq1,Q,V,U,C,B,nq,nu,nw,nc,nb,nz,nθ,nν,NJ,NR,NI,O,LS}
    for t = 1:traj.H
        # Cost function
        delta!(core.Δq[t], traj.q[t+2], ref_traj.q[t+2])
        delta!(core.Δu[t], traj.u[t], ref_traj.u[t])
        delta!(core.Δγ[t], traj.γ[t], ref_traj.γ[t])
        delta!(core.Δb[t], traj.b[t], ref_traj.b[t])

        # res.q2[t] .+= obj.q[t] * core.Δq[t]
        mul!(res.q2[t], obj.q[t], core.Δq[t], 1.0, 1.0)
        # res.u1[t] .+= obj.u[t] * core.Δu[t]
        mul!(res.u1[t], obj.u[t], core.Δu[t], 1.0, 1.0)
        # res.γ1[t] .+= obj.γ[t] * core.Δγ[t]
        mul!(res.γ1[t], obj.γ[t], core.Δγ[t], 1.0, 1.0)
        # res.b1[t] .+= obj.b[t] * core.Δb[t]
        mul!(res.b1[t], obj.b[t], core.Δb[t], 1.0, 1.0)

        # velocity
        # res.q2[t] .+= obj.v[t] * (traj.q[t+2] - traj.q[t+1])
        mul!(res.q2[t], obj.v[t], traj.q[t+2], 1.0, 1.0)
        mul!(res.q2[t], obj.v[t], traj.q[t+1], -1.0, 1.0)

        t == 1 && continue
        # res.q2[t-1] .-= obj.v[t] * (traj.q[t+2] - traj.q[t+1])
        mul!(res.q2[t-1], obj.v[t], traj.q[t+2], -1.0, 1.0)
        mul!(res.q2[t-1], obj.v[t], traj.q[t+1], 1.0, 1.0)
    end
end

function gradient!(res::NewtonResidualConfiguration{T,vq2,vu1,vd,vI,vq0,vq1}, obj::TrackingVelocityObjective{Q,V,U,C,B}, core::Newton{T,nq,nu,nw,nc,nb,nz,nθ,nν,NJ,NR,NI,O,LS}, traj::ContactTraj{T,nq,nu,nw,nc,nb,nz,nθ}, ref_traj::ContactTraj{T,nq,nu,nw,nc,nb,nz,nθ}) where {T,vq2,vu1,vd,vI,vq0,vq1,Q,V,U,C,B,nq,nu,nw,nc,nb,nz,nθ,nν,NJ,NR,NI,O,LS}
    for t = 1:traj.H
        # Cost function
        delta!(core.Δq[t], traj.q[t+2], ref_traj.q[t+2])
        delta!(core.Δu[t], traj.u[t], ref_traj.u[t])

        mul!(res.q2[t], obj.q[t], core.Δq[t], 1.0, 1.0)
        mul!(res.u1[t], obj.u[t], core.Δu[t], 1.0, 1.0)

        # velocity
        mul!(res.q2[t], obj.v[t], traj.q[t+2], 1.0, 1.0)
        mul!(res.q2[t], obj.v[t], traj.q[t+1], -1.0, 1.0)
        mul!(res.q2[t], obj.v[t], obj.v_target[t], 1.0, 1.0)

        t == 1 && continue
        mul!(res.q2[t-1], obj.v[t], traj.q[t+2], -1.0, 1.0)
        mul!(res.q2[t-1], obj.v[t], traj.q[t+1], 1.0, 1.0)
        mul!(res.q2[t-1], obj.v[t], obj.v_target[t], -1.0, 1.0)
    end
end


text:
using ZonalFlow
using Test

@testset "NL equalities for 2x2" begin
    lx = 4.0*Float64(pi);
    ly = 2.0*Float64(pi);
    nx = 2;
    ny = 2;
    Ω = 2.0*Float64(pi)

    for θ in [0.0]
        for τ in [10.0,20.0]
            for Ξ in [0.2]

                @info "Parameters: θ = $θ, τ = $τ, Ξ = $Ξ"

                β = 2.0*Ω*cos(θ*Float64(pi))
                ζ0 = ic_pert_eqm(lx,ly,nx,ny,Ξ*Ω); # one ic for all

                sol1 = nl(lx,ly,nx,ny,Ξ*Ω,β,τ/Ω,ic=ζ0,dt=0.001,t_end=500.0,savefreq=10);
                P1,O1 = zonalenergy(lx,ly,nx,ny,sol1.u)
                A1 = meanvorticity(lx,ly,nx,ny,sol1.u)
                E1,Z1 = energy(lx,ly,nx,ny,sol1.u)

                @info "Do NL, QL and CE2 match?"
                Λ = 0
                sol2 = gql(lx,ly,nx,ny,Λ,Ξ*Ω,β,τ/Ω,dt=0.001,ic=ζ0,savefreq=10);
                sol3 = gce2(lx,ly,nx,ny,Λ,Ξ*Ω,β,τ/Ω,ic=ζ0,dt=0.001,t_end=500.0,poscheck=false,savefreq=10);
                sol4 = gce2(lx,ly,nx,ny,Λ,Ξ*Ω,β,τ/Ω,ic=ζ0,dt=0.001,t_end=500.0,savefreq=10,poscheck=true);
                @test sol1.u[end][:,1:Λ+1] ≈ sol2.u[end][:,1:Λ+1] atol=1e-6
                @test sol1.u[end][:,1:Λ+1] ≈ sol3.u[end].x[1] atol=1e-6
                @test sol1.u[end][:,1:Λ+1] ≈ sol4.u[end].x[1] atol=1e-6
                P2,O2 = zonalenergy(lx,ly,nx,ny,sol2.u);
                P3,O3 = zonalenergy(lx,ly,nx,ny,Λ,sol3.u);
                P4,O4 = zonalenergy(lx,ly,nx,ny,Λ,sol4.u);
                @test P1[end,:] ≈ P2[end,:] atol=1e-6
                @test P1[end,:] ≈ P3[end,:] atol=1e-6
                @test P1[end,:] ≈ P4[end,:] atol=1e-6
                A2 = meanvorticity(lx,ly,nx,ny,sol2.u)
                A3 = meanvorticity(lx,ly,nx,ny,sol3.u)
                A4 = meanvorticity(lx,ly,nx,ny,sol4.u)
                @test A1[end,:] ≈ A2[end,:] atol=1e-6
                @test A1[end,:] ≈ A3[end,:] atol=1e-6
                @test A1[end,:] ≈ A4[end,:] atol=1e-6
                E2,Z2 = energy(lx,ly,nx,ny,sol2.u)
                E3,Z3 = energy(lx,ly,nx,ny,Λ,sol3.u)
                E4,Z4 = energy(lx,ly,nx,ny,Λ,sol4.u)
                @test E1[end] ≈ E2[end] atol=1e-6
                @test E1[end] ≈ E3[end] atol=1e-6
                @test E1[end] ≈ E4[end] atol=1e-6

                @info "Do NL, GQl(M) and GCE2(M) match?"
                Λ = 1
                sol2 = gql(lx,ly,nx,ny,Λ,Ξ*Ω,β,τ/Ω,dt=0.001,ic=ζ0,t_end=500.0,savefreq=10);
                sol3 = gce2(lx,ly,nx,ny,Λ,Ξ*Ω,β,τ/Ω,ic=ζ0,dt=0.001,t_end=500.0,poscheck=false,savefreq=10);
                sol4 = gce2(lx,ly,nx,ny,Λ,Ξ*Ω,β,τ/Ω,ic=ζ0,dt=0.001,t_end=500.0,savefreq=10,poscheck=true);
                @test sol1.u[end][:,1:Λ+1] ≈ sol2.u[end][:,1:Λ+1] atol=1e-6
                @test sol1.u[end][:,1:Λ+1] ≈ sol3.u[end].x[1] atol=1e-6
                @test sol1.u[end][:,1:Λ+1] ≈ sol4.u[end].x[1] atol=1e-6
                P2,O2 = zonalenergy(lx,ly,nx,ny,sol2.u);
                P3,O3 = zonalenergy(lx,ly,nx,ny,Λ,sol3.u);
                P4,O4 = zonalenergy(lx,ly,nx,ny,Λ,sol4.u);
                @test P1[end,:] ≈ P2[end,:] atol=1e-6
                @test P1[end,:] ≈ P3[end,:] atol=1e-6
                @test P1[end,:] ≈ P4[end,:] atol=1e-6
                A2 = meanvorticity(lx,ly,nx,ny,sol2.u)
                A3 = meanvorticity(lx,ly,nx,ny,sol3.u)
                A4 = meanvorticity(lx,ly,nx,ny,sol4.u)
                @test A1[end,:] ≈ A2[end,:] atol=1e-6
                @test A1[end,:] ≈ A3[end,:] atol=1e-6
                @test A1[end,:] ≈ A4[end,:] atol=1e-6
                E2,Z2 = energy(lx,ly,nx,ny,sol2.u)
                E3,Z3 = energy(lx,ly,nx,ny,Λ,sol3.u)
                E4,Z4 = energy(lx,ly,nx,ny,Λ,sol4.u)
                @test E1[end] ≈ E2[end] atol=1e-6
                @test E1[end] ≈ E3[end] atol=1e-6
                @test E1[end] ≈ E4[end] atol=1e-6

                @info "Do default parameters work?"
                Λ = 0
                sol2 = gql(lx,ly,nx,ny,Λ,Ξ*Ω,β,τ/Ω,ic=ζ0);
                sol3 = gce2(lx,ly,nx,ny,Λ,Ξ*Ω,β,τ/Ω,ic=ζ0);
                sol4 = gce2(lx,ly,nx,ny,Λ,Ξ*Ω,β,τ/Ω,ic=ζ0,poscheck=true);
                @test sol1.u[end][:,1:Λ+1] ≈ sol2.u[end][:,1:Λ+1] atol=1e-6
                @test sol1.u[end][:,1:Λ+1] ≈ sol3.u[end].x[1] atol=1e-6
                @test sol1.u[end][:,1:Λ+1] ≈ sol4.u[end].x[1] atol=1e-6
                P2,O2 = zonalenergy(lx,ly,nx,ny,sol2.u);
                P3,O3 = zonalenergy(lx,ly,nx,ny,Λ,sol3.u);
                P4,O4 = zonalenergy(lx,ly,nx,ny,Λ,sol4.u);
                @test P1[end,:] ≈ P2[end,:] atol=1e-6
                @test P1[end,:] ≈ P3[end,:] atol=1e-6
                @test P1[end,:] ≈ P4[end,:] atol=1e-6
                A2 = meanvorticity(lx,ly,nx,ny,sol2.u)
                A3 = meanvorticity(lx,ly,nx,ny,sol3.u)
                A4 = meanvorticity(lx,ly,nx,ny,sol4.u)
                @test A1[end,:] ≈ A2[end,:] atol=1e-6
                @test A1[end,:] ≈ A3[end,:] atol=1e-6
                @test A1[end,:] ≈ A4[end,:] atol=1e-6
                E2,Z2 = energy(lx,ly,nx,ny,sol2.u)
                E3,Z3 = energy(lx,ly,nx,ny,Λ,sol3.u)
                E4,Z4 = energy(lx,ly,nx,ny,Λ,sol4.u)
                @test E1[end] ≈ E2[end] atol=1e-6
                @test E1[end] ≈ E3[end] atol=1e-6
                @test E1[end] ≈ E4[end] atol=1e-6

                Λ = 1
                sol2 = gql(lx,ly,nx,ny,Λ,Ξ*Ω,β,τ/Ω,ic=ζ0);
                sol3 = gce2(lx,ly,nx,ny,Λ,Ξ*Ω,β,τ/Ω,ic=ζ0);
                sol4 = gce2(lx,ly,nx,ny,Λ,Ξ*Ω,β,τ/Ω,ic=ζ0,poscheck=true);
                @test sol1.u[end][:,1:Λ+1] ≈ sol2.u[end][:,1:Λ+1] atol=1e-6
                @test sol1.u[end][:,1:Λ+1] ≈ sol3.u[end].x[1] atol=1e-6
                @test sol1.u[end][:,1:Λ+1] ≈ sol4.u[end].x[1] atol=1e-6
                P2,O2 = zonalenergy(lx,ly,nx,ny,sol2.u);
                P3,O3 = zonalenergy(lx,ly,nx,ny,Λ,sol3.u);
                P4,O4 = zonalenergy(lx,ly,nx,ny,Λ,sol4.u);
                @test P1[end,:] ≈ P2[end,:] atol=1e-6
                @test P1[end,:] ≈ P3[end,:] atol=1e-6
                @test P1[end,:] ≈ P4[end,:] atol=1e-6
                A2 = meanvorticity(lx,ly,nx,ny,sol2.u)
                A3 = meanvorticity(lx,ly,nx,ny,sol3.u)
                A4 = meanvorticity(lx,ly,nx,ny,sol4.u)
                @test A1[end,:] ≈ A2[end,:] atol=1e-6
                @test A1[end,:] ≈ A3[end,:] atol=1e-6
                @test A1[end,:] ≈ A4[end,:] atol=1e-6
                E2,Z2 = energy(lx,ly,nx,ny,sol2.u)
                E3,Z3 = energy(lx,ly,nx,ny,Λ,sol3.u)
                E4,Z4 = energy(lx,ly,nx,ny,Λ,sol4.u)
                @test E1[end] ≈ E2[end] atol=1e-6
                @test E1[end] ≈ E3[end] atol=1e-6
                @test E1[end] ≈ E4[end] atol=1e-6
            end
        end
    end
end


text:
module Xtal

using StaticArrays
using LinearAlgebra
using Printf
using ComputedFieldTypes

const ELEMENTS =
(
    "H",  "He", "Li", "Be", "B",  "C",  "N",  "O",  "F",  "Ne",
    "Na", "Mg", "Al", "Si", "P",  "S",  "Cl", "Ar", "K",  "Ca",
    "Sc", "Ti", "V",  "Cr", "Mn", "Fe", "Co", "Ni", "Cu", "Zn",
    "Ga", "Ge", "As", "Se", "Br", "Kr", "Rb", "Sr", "Y",  "Zr",
    "Nb", "Mo", "Tc", "Ru", "Rh", "Pd", "Ag", "Cd", "In", "Sn",
    "Sb", "Te", "I",  "Xe", "Cs", "Ba", "La", "Ce", "Pr", "Nd",
    "Pm", "Sm", "Eu", "Gd", "Tb", "Dy", "Ho", "Er", "Tm", "Yb",
    "Lu", "Hf", "Ta", "W",  "Re", "Os", "Ir", "Pt", "Au", "Hg",
    "Tl", "Pb", "Bi", "Po", "As", "Rn", "Fr", "Ra", "Ac", "Th",
    "Pa", "U",  "Np", "Pu", "Am", "Cm", "Bk", "Cf", "Es", "Fm",
    "Md", "No", "Lr", "Rf", "Db", "Sg", "Bh", "Hs", "Mt", "Ds",
    "Rg", "Cn", "Nh", "Fl", "Mc", "Lv", "Ts", "Og"
)

# TODO: can we modify this to make it unnecessary?
const ELEMENT_LOOKUP = Dict{String, Int}([(ELEMENTS[n] => n) for n in 1:length(ELEMENTS)])

# Matrices used to reduce conventional cells to corresponding primitive cells
# These can be accessed by calling the centering letter as a symbol
# So if you want the face-centered reduction matrix, use REDUCTION_MATRIX_3D[:F]
const REDUCTION_MATRIX_3D =
(
    # Transpose these matrices in your head, please
    A = SMatrix{3,3,Float64}( 1,    0,    0,
                              0,  1/2,  1/2,
                              0, -1/2,  1/2),
    B = SMatrix{3,3,Float64}( 1/2,  0,  1/2,
                                0,  1,    0,
                             -1/2,  0,  1/2),
    C = SMatrix{3,3,Float64}( 1/2,  1/2,  0,
                             -1/2,  1/2,  0,
                                0,    0,  1),
    F = SMatrix{3,3,Float64}( 1/2,  1/2,    0,
                              1/2,    0,  1/2,
                                0,  1/2,  1/2),
    I = SMatrix{3,3,Float64}(-1/2,  1/2,  1/2,
                              1/2, -1/2,  1/2,
                              1/2,  1/2, -1/2),
    P = SMatrix{3,3,Float64}(LinearAlgebra.I)
)

"""
    TOL_DEF

Default tolerance for discrepancies in floating point values.
"""
const TOL_DEF = 1e-8

const BOHR2ANG = 0.529177

# the value of 2m/ħ^2 as used in VASP
const CVASP = 0.262465831

"""
    _allsame(itr)

Returns `true` if all the elements of an iterator are identical.
"""
_allsame(itr) = all(x -> x == first(itr), itr)

"""
    _selfdot(v)

Computes the dot product of a vector with itself.
"""
_selfdot(v) = dot(v,v)

# Methods used in vector operations that go beyond what's available in LinearAlgebra
include("vectors.jl")
# Abstract types used in type tree
include("types.jl")
export AbstractLattice, AbstractCrystal, AbstractCrystalData, AbstractRealSpaceData,
       AbstractReciprocalSpaceData, AbstractKPoints, AbstractDensityOfStates, AbsractPotential,
       AbstractPseudopotential
# Methods and structs for working with crystal lattices
include("lattices.jl")
export BasisVectors, RealLattice, ReciprocalLattice
export dual, prim, conv, cell_lengths, cell_volume, lattice2D, lattice3D, maxHKLindex
# Methods and structs for working with atomic positions
include("atoms.jl")
export AtomPosition, AtomList
export atomname, atomicno, coord, natom, basis, cartesian, reduce_coords
# Methods and structs for working with crystal data
include("crystals.jl")
export Crystal, CrystalWithDatasets
export natom_template
# Methods and structs for working with different types of data associated with crystals
include("data.jl")
export RealSpaceDataGrid, KPointGrid, KPointList, BandAtKPoint, BandStructure, HKLData,
       ReciprocalWavefunction, DensityOfStates, ProjectedDensityOfStates, FatBands, AtomicData,
       SphericalComponents
export grid, gridsize, nkpt, nband, fermi, energies, nelectrons
include("potentials.jl")
export XCFunctional, HGHPseudopotential, zatom, zion
# Methods and structs for working with different file formats
include("filetypes.jl")
export readXYZ, readXSF3D, writeXSF, readCPcoeff, readCPgeo, readCPcell, read_abinit_density,
       read_abinit_potential, readWAVECAR, readDOSCAR, readPROCAR, readHGH
# Show methods for pretty printing this module's structs
include("show.jl")
# Precompilation directives
include("precompile.jl")

end # end of module

max_stars_repo_name: brainandforce/Xtal

max_issues_repo_name: brainandforce/Xtal

max_forks_repo_name: brainandforce/Xtal