text:
using Printf
using CUDA
using CUDA.CUSPARSE
using FiniteDiff
using ForwardDiff
using KernelAbstractions
using LinearAlgebra
using Random
using SparseArrays
using Test
using TimerOutputs
using ExaPF
import ExaPF: PowerSystem, AutoDiff

const PS = PowerSystem

@testset "Compute reduced Hessian on CPU" begin
    @testset "Case $case" for case in ["case9.m", "case30.m"]
        if has_cuda_gpu()
            ITERATORS = zip([CPU(), CUDADevice()], [Vector, CuVector])
        else
            ITERATORS = zip([CPU()], [Vector])
        end
        RTOL = 1e-6
        ATOL = 1e-6
        @testset "Device $device" for (device, T) in ITERATORS
            ##################################################
            # Initialization
            ##################################################
            datafile = joinpath(dirname(@__FILE__), "..", "..", "data", case)
            tolerance = 1e-8
            pf = PS.PowerNetwork(datafile)
            polar = PolarForm(pf, device)
            cpu_polar = PolarForm(pf, CPU())
            pv = pf.pv ; npv = length(pv)
            pq = pf.pq ; npq = length(pq)
            ref = pf.ref ; nref = length(ref)
            nbus = pf.nbus
            ngen = get(polar, PS.NumberOfGenerators())

            pv2gen = polar.indexing.index_pv_to_gen
            ref2gen = polar.indexing.index_ref_to_gen
            gen2bus = polar.indexing.index_generators
            cache = ExaPF.get(polar, ExaPF.PhysicalState())
            cpu_cache = ExaPF.get(cpu_polar, ExaPF.PhysicalState())
            ExaPF.init_buffer!(polar, cache)
            ExaPF.init_buffer!(cpu_polar, cpu_cache)

            xk = ExaPF.initial(polar, State())
            u = ExaPF.initial(polar, Control())
            nx = length(xk) ; nu = length(u)

            jx = AutoDiff.Jacobian(polar, ExaPF.power_balance, State())
            ju = AutoDiff.Jacobian(polar, ExaPF.power_balance, Control())
            ∂obj = ExaPF.AdjointStackObjective(polar)
            pbm = AutoDiff.TapeMemory(ExaPF.active_power_generation, ∂obj, nothing)

            cpu_jx = AutoDiff.Jacobian(cpu_polar, ExaPF.power_balance, State())
            cpu_ju = AutoDiff.Jacobian(cpu_polar, ExaPF.power_balance, Control())
            cpu_∂obj = ExaPF.AdjointStackObjective(cpu_polar)

            ##################################################
            # Step 1: computation of first-order adjoint
            ##################################################
            conv = powerflow(polar, jx, cache, NewtonRaphson())
            cpu_conv = powerflow(cpu_polar, cpu_jx, cpu_cache, NewtonRaphson())
            ExaPF.update!(polar, PS.Generators(), PS.ActivePower(), cache)
            ExaPF.update!(cpu_polar, PS.Generators(), PS.ActivePower(), cpu_cache)
            @test conv.has_converged
            @test cpu_conv.has_converged
            # No need to recompute ∇gₓ
            ∇gₓ = jx.J
            ∇gᵤ = AutoDiff.jacobian!(polar, ju, cache)
            # test jacobian wrt x
            ∇gᵥ = AutoDiff.jacobian!(polar, jx, cache)
            @test isequal(∇gₓ, ∇gᵥ)

            # Fetch values found by Newton-Raphson algorithm
            vm = cache.vmag
            va = cache.vang
            pg = cache.pg
            # State & Control
            x = [va[pv] ; va[pq] ; vm[pq]]
            u = [vm[ref]; vm[pv]]
            # Test with Matpower's Jacobian
            V = vm .* exp.(im * va)
            Ybus = pf.Ybus
            Jₓ = ExaPF.matpower_jacobian(polar, State(), ExaPF.power_balance, V)
            @test isapprox(∇gₓ, Jₓ )
            # Hessian vector product
            ExaPF.gradient_objective!(polar, pbm, cache)
            ∇fₓ = ∂obj.∇fₓ
            ∇fᵤ = ∂obj.∇fᵤ
            λ  = -(Array(∇gₓ')) \ Array(∇fₓ)
            grad_adjoint = Array(∇fᵤ) + Array(∇gᵤ)' * λ

            ##################################################
            # Step 2: computation of Hessian of powerflow g
            ##################################################
            ## w.r.t. xx
            function jac_diff(x)
                vm_ = copy(vm)
                va_ = copy(va)
                va_[pv] = x[1:npv]
                va_[pq] = x[npv+1:npv+npq]
                vm_[pq] = x[npv+npq+1:end]
                V = vm_ .* exp.(im * va_)
                Jx = ExaPF.matpower_jacobian(polar, State(), ExaPF.power_balance, V)
                return Jx' * λ
            end

            # Evaluate Hessian-vector product (full ∇²gₓₓ is a 3rd dimension tensor)
            ∇²gλ = ExaPF.matpower_hessian(cpu_polar, ExaPF.power_balance, cpu_cache, λ)
            H_fd = FiniteDiff.finite_difference_jacobian(jac_diff, x)
            @test isapprox(∇²gλ.xx, H_fd, rtol=RTOL)
            ybus_re, ybus_im = ExaPF.Spmat{T{Int}, T{Float64}}(Ybus)
            pbus = T(real(pf.sbus))
            qbus = T(imag(pf.sbus))
            dev_pv = T(pf.pv)
            dev_pq = T(pf.pq)
            dev_ref = T(pf.ref)
            F = zeros(Float64, npv + 2*npq)
            nx = size(∇²gλ.xx, 1)
            nu = size(∇²gλ.uu, 1)

            # Hessian-vector product using forward over adjoint AD
            HessianAD = AutoDiff.Hessian(polar, ExaPF.power_balance)

            tgt = rand(nx + nu)
            projp = zeros(nx + nu)
            dev_tgt = T(tgt)
            dev_projp = T(projp)
            # set tangents only for x direction
            dev_λ = T(λ)
            tgt[nx+1:end] .= 0.0
            dev_tgt = T(tgt)
            AutoDiff.adj_hessian_prod!(polar, HessianAD, dev_projp, cache, dev_λ, dev_tgt)
            projp = Array(dev_projp)
            @test isapprox(projp[1:nx], ∇²gλ.xx * tgt[1:nx])

            tgt = rand(nx + nu)
            # set tangents only for u direction
            tgt[1:nx] .= 0.0
            dev_tgt = T(tgt)
            AutoDiff.adj_hessian_prod!(polar, HessianAD, dev_projp, cache, dev_λ, dev_tgt)
            projp = Array(dev_projp)
            # (we use absolute tolerance as Huu is equal to 0 for case9)
            @test isapprox(projp[nx+1:end], ∇²gλ.uu * tgt[nx+1:end], atol=ATOL)

            # check cross terms ux
            tgt = rand(nx + nu)
            tgt .= 1.0
            # Build full Hessian
            H = [
                ∇²gλ.xx ∇²gλ.xu';
                ∇²gλ.xu ∇²gλ.uu
            ]
            dev_tgt = T(tgt)
            AutoDiff.adj_hessian_prod!(polar, HessianAD, dev_projp, cache, dev_λ, dev_tgt)
            projp = Array(dev_projp)
            @test isapprox(projp, H * tgt)

            ## w.r.t. uu
            function jac_u_diff(u)
                vm_ = copy(vm)
                va_ = copy(va)
                vm_[ref] = u[1:nref]
                vm_[pv] = u[nref+1:end]
                V = vm_ .* exp.(im * va_)
                Ju = ExaPF.matpower_jacobian(polar, Control(), ExaPF.power_balance, V)
                return Ju' * λ
            end

            Hᵤᵤ_fd = FiniteDiff.finite_difference_jacobian(jac_u_diff, u)

            ## w.r.t. xu
            function jac_xu_diff(x)
                vm_ = copy(vm)
                va_ = copy(va)
                va_[pv] = x[1:npv]
                va_[pq] = x[npv+1:npv+npq]
                vm_[pq] = x[npv+npq+1:end]
                V = vm_ .* exp.(im * va_)
                Ju = ExaPF.matpower_jacobian(polar, Control(), ExaPF.power_balance, V)[:, 1:nref+npv]
                return Ju' * λ
            end

            Hₓᵤ_fd = FiniteDiff.finite_difference_jacobian(jac_xu_diff, x)
            @test isapprox(∇²gλ.xu[1:nref+npv, :], Hₓᵤ_fd, rtol=RTOL)

            ##################################################
            # Step 3: include constraints in Hessian
            ##################################################
            # h1 (state)      : xl <= x <= xu
            # h2 (by-product) : yl <= y <= yu
            # Test sequential evaluation of Hessian
            x = [va[pv] ; va[pq] ; vm[pq]]
            u = [vm[ref]; vm[pv]; pg[pv2gen]]

            μ = rand(ngen)
            ∂₂Q = ExaPF.matpower_hessian(cpu_polar, ExaPF.reactive_power_constraints, cpu_cache, μ)
            function jac_x(z)
                x_ = z[1:nx]
                u_ = z[1+nx:end]
                # Transfer control
                ExaPF.transfer!(polar, cache, u_)
                # Transfer state (manually)
                cache.vang[pv] .= x_[1:npv]
                cache.vang[pq] .= x_[npv+1:npv+npq]
                cache.vmag[pq] .= x_[npv+npq+1:end]
                ExaPF.update!(polar, PS.Generators(), PS.ActivePower(), cache)
                Jx = ExaPF.matpower_jacobian(polar, ExaPF.reactive_power_constraints, State(), cache)
                Ju = ExaPF.matpower_jacobian(polar, ExaPF.reactive_power_constraints, Control(), cache)
                return [Jx Ju]' * μ
            end

            H_fd = FiniteDiff.finite_difference_jacobian(jac_x, [x; u])
            @test isapprox(∂₂Q.uu, H_fd[nx+1:end, nx+1:end], rtol=RTOL)
            @test isapprox(∂₂Q.xx, H_fd[1:nx, 1:nx], rtol=RTOL)
            @test isapprox(∂₂Q.xu, H_fd[nx+1:end, 1:nx], rtol=RTOL)

            # Test with AutoDiff.Hessian
            hess_reactive = AutoDiff.Hessian(polar, ExaPF.reactive_power_constraints)

            # XX
            tgt = rand(nx + nu)
            projp = zeros(nx + nu)
            tgt[nx+1:end] .= 0.0
            dev_tgt = T(tgt)
            dev_projp = T(projp)
            dev_μ = T(μ)
            AutoDiff.adj_hessian_prod!(polar, hess_reactive, dev_projp, cache, dev_μ, dev_tgt)
            projp = Array(dev_projp)
            @test isapprox(projp[1:nx], ∂₂Q.xx * tgt[1:nx])

            # UU
            tgt = rand(nx + nu)
            tgt[1:nx] .= 0.0
            dev_tgt = T(tgt)
            dev_projp .= 0.0
            AutoDiff.adj_hessian_prod!(polar, hess_reactive, dev_projp, cache, dev_μ, dev_tgt)
            projp = Array(dev_projp)
            @test isapprox(projp[1+nx:end], ∂₂Q.uu * tgt[1+nx:end])

            # XU
            tgt = rand(nx + nu)
            dev_tgt = T(tgt)
            dev_projp .= 0.0
            AutoDiff.adj_hessian_prod!(polar, hess_reactive, dev_projp, cache, dev_μ, dev_tgt)
            H = [
                ∂₂Q.xx ∂₂Q.xu' ;
                ∂₂Q.xu ∂₂Q.uu
            ]
            projp = Array(dev_projp)
            @test isapprox(projp, H * tgt)

            # Hessian w.r.t. Line-flow
            hess_lineflow = AutoDiff.Hessian(polar, ExaPF.flow_constraints)
            ncons = ExaPF.size_constraint(polar, ExaPF.flow_constraints)
            μ = rand(ncons)
            function flow_jac_x(z)
                x_ = z[1:nx]
                u_ = z[1+nx:end]
                # Transfer control
                ExaPF.transfer!(polar, cache, u_)
                # Transfer state (manually)
                cache.vang[pv] .= x_[1:npv]
                cache.vang[pq] .= x_[npv+1:npv+npq]
                cache.vmag[pq] .= x_[npv+npq+1:end]
                V = cache.vmag .* exp.(im .* cache.vang)
                Jx = ExaPF.matpower_jacobian(polar, State(), ExaPF.flow_constraints, V)
                Ju = ExaPF.matpower_jacobian(polar, Control(), ExaPF.flow_constraints, V)
                return [Jx Ju]' * μ
            end

            H_fd = FiniteDiff.finite_difference_jacobian(flow_jac_x, [x; u])
            tgt = rand(nx + nu)
            projp = zeros(nx + nu)
            dev_tgt = T(tgt)
            dev_projp = T(projp)
            dev_μ = T(μ)
            AutoDiff.adj_hessian_prod!(polar, hess_lineflow, dev_projp, cache, dev_μ, dev_tgt)
            projp = Array(dev_projp)
            @test isapprox(projp, H_fd * tgt, rtol=RTOL)

            # Hessian w.r.t. Active-power generation
            hess_pg = AutoDiff.Hessian(cpu_polar, ExaPF.active_power_constraints)
            μ = rand(1)
            # TODO: No GPU implementation so far
            AutoDiff.adj_hessian_prod!(cpu_polar, hess_pg, projp, cache, μ, tgt)
            ∂₂P = ExaPF.matpower_hessian(cpu_polar, ExaPF.active_power_constraints, cpu_cache, μ)
            H = [
                ∂₂P.xx ∂₂P.xu' ;
                ∂₂P.xu ∂₂P.uu
            ]
            @test isapprox(projp, (H * tgt))
        end
    end
end


text:
using MendelIterHardThreshold
result = IterHardThreshold("gwas 1 Control.txt")
# result = IterHardThreshold("test_control.txt")

ext: jl


text:
using Test
using WorldBankData

us_gnp_data = Any[  Dict("total" => 23, "per_page" => "25000", "pages" => 1, "page" => 1),
                  Any[ Dict("date" => "2012", "value" => 52340, "indicator" => Dict("id" => "NY.GNP.PCAP.CD", "value" => "GNI per capita, Atlas method (current US\$)"), "country" => Dict("id" => "US", "value" => "United States"), "decimal" => "0"),
                    Dict("date" => "2011", "value" => 50650, "indicator" => Dict("id" => "NY.GNP.PCAP.CD", "value" => "GNI per capita, Atlas method (current US\$)"), "country" => Dict("id" => "US", "value" => "United States"), "decimal" => "0"),
                    Dict("date" => "2010", "value" => 48960, "indicator" => Dict("id" => "NY.GNP.PCAP.CD", "value" => "GNI per capita, Atlas method (current US\$)"), "country" => Dict("id" => "US", "value" => "United States"), "decimal" => "0"),
                    Dict("date" => "2009", "value" => 48040, "indicator" => Dict("id" => "NY.GNP.PCAP.CD", "value" => "GNI per capita, Atlas method (current US\$)"), "country" => Dict("id" => "US", "value" => "United States"), "decimal" => "0"),
                    Dict("date" => "2008", "value" => 49350, "indicator" => Dict("id" => "NY.GNP.PCAP.CD", "value" => "GNI per capita, Atlas method (current US\$)"), "country" => Dict("id" => "US", "value" => "United States"), "decimal" => "0"),
                    Dict("date" => "2007", "value" => 48640, "indicator" => Dict("id" => "NY.GNP.PCAP.CD", "value" => "GNI per capita, Atlas method (current US\$)"), "country" => Dict("id" => "US", "value" => "United States"), "decimal" => "0"),
                    Dict("date" => "2006", "value" => 48080, "indicator" => Dict("id" => "NY.GNP.PCAP.CD", "value" => "GNI per capita, Atlas method (current US\$)"), "country" => Dict("id" => "US", "value" => "United States"), "decimal" => "0")
                  ]
                ]

us_gnp = WorldBankData.parse_wdi("NY.GNP.PCAP.CD", us_gnp_data[2], 2006, 2012)

@test us_gnp[!, :year] == Float64[2012, 2011, 2010, 2009, 2008, 2007, 2006]
@test us_gnp[!, :NY_GNP_PCAP_CD] == Float64[52340, 50650, 48960, 48040, 49350, 48640, 48080]


text:
# mtdg_inference.jl

export ParamsMTDg, PriorMTDg, ModMTDg,
  sim_mtdg, symmetricDirPrior_mtdg, transTensor_mtdg,
  counttrans_mtdg, mcmc!, timemod!, llik_MTDg,
  TankReduction, forecDist_MTDg;

mutable struct ParamsMTDg
    lλ::Vector{Float64} # will be a prob. vector of length L + 1
    ζ::Vector{Int} # will be length TT - L
    lQ0::Vector{Float64}
    lQ::Vector{Matrix{Float64}} # organized so first index is now and lag 1 is the next index

    ParamsMTDg(lλ, ζ, lQ0, lQ) = new(deepcopy(lλ), deepcopy(ζ), deepcopy(lQ0), deepcopy(lQ))
end

mutable struct PriorMTDg
  λ::Union{Vector{Float64}, SparseDirMix, SBMprior}
  Q0::Vector{Float64}
  Q::Union{Vector{Matrix{Float64}}, Vector{Vector{SparseDirMix}}, Vector{Vector{SBMprior}}}
  PriorMTDg(λ, Q0, Q) = new(deepcopy(λ), deepcopy(Q0), deepcopy(Q))
end

mutable struct ModMTDg
  L::Int # maximal order
  K::Int # number of states
  TT::Int
  S::Vector{Int}
  prior::PriorMTDg
  state::ParamsMTDg
  iter::Int

  ModMTDg(L, K, TT, S, prior, state) = new(L, K, TT, deepcopy(S), deepcopy(prior), deepcopy(state), 0)
end

"""
    sim_mtdg(TT, nburn, L, K, λ, Q0, Q)
"""
function sim_mtdg(TT::Int, nburn::Int, L::Int, K::Int,
  λ::Vector{Float64}, Q0::Vector{Float64}, Q::Vector{Matrix{Float64}})

  Nsim = nburn + TT

  ζ = [ StatsBase.sample(Weights(λ)) - 1 for i in 1:(Nsim-L) ]

  S = Vector{Int}(undef, Nsim)
  S[1:L] = StatsBase.sample(1:K, L)

  for tt in (L+1):(Nsim)
    i = tt - L
    if ζ[i] == 0
        pvec = deepcopy(Q0)
    else
        Slagrev_now = S[range(tt-1, step=-1, length=L)]
        pvec = deepcopy( Q[ζ[i]][:, Slagrev_now[ζ[i]] ] )
    end
    S[tt] = StatsBase.sample(Weights( pvec ))
  end

  S[(nburn+1):(Nsim)], ζ[(nburn+1):(Nsim-L)]
end

"""
    symmetricPrior_mtdg(size_λ, size_Q0, size_Q, L, K)
"""
function symmetricDirPrior_mtdg(size_λ::Float64, size_Q0::Float64, size_Q::Float64,
  L::Int, K::Int)

  α0_λ = fill( size_λ / float(L), L )
  α0_Q0 = fill( size_Q0 / float(K), K )
  a0_Q = size_Q / float(K)
  α0_Q = [ fill( a0_Q, (K, K) ) for ℓ=1:L ]

  (α0_λ, α0_Q0, α0_Q)
end

"""
    transTensor_mtdg(L, K, λ, Q0, Q)

Calculate full transition tensor from λ and Q.
"""
function transTensor_mtdg(L::Int, K::Int, λ::Vector{Float64},
    Q0::Vector{Float64}, Q::Vector{Matrix{Float64}})

    froms, nfroms = create_froms(K, L) # in this case, ordered by now, lag1, lag2, etc.
    Ωmat = zeros(Float64, (K, nfroms))
    for i in 1:nfroms
        for k in 1:K
            Ωmat[k,i] += λ[1] * Q0[k]
            for ℓ in 1:L
                Ωmat[k,i] += λ[ℓ+1] .* Q[ℓ][k,froms[i][ℓ]]
            end
        end
    end
    reshape(Ωmat, fill(K,L+1)...)
end

"""
    rpost_lλ_mtdg(α0_λ, ζ, L)
"""
function rpost_lλ_mtdg(α0_λ::Vector{Float64}, ζ::Vector{Int}, L::Int)

    Nζ = StatsBase.counts(ζ, 0:L)
    α1_λ = α0_λ .+ Nζ
    lλ_out = SparseProbVec.rDirichlet(α1_λ, logout=true)

    lλ_out
end
function rpost_lλ_mtdg(prior_λ::SparseDirMix, ζ::Vector{Int}, L::Int)

    Nζ = StatsBase.counts(ζ, 0:L)
    α1_λ = prior_λ.α .+ Nζ
    d = SparseProbVec.SparseDirMix(α1_λ, prior_λ.β)
    lλ_out = SparseProbVec.rand(d, logout=true)

    lλ_out
end
function rpost_lλ_mtdg(prior_λ::SBMprior, ζ::Vector{Int}, L::Int)

    Nζ = StatsBase.counts(ζ, 0:L)
    post_lλ = SparseProbVec.SBM_multinom_post(prior_λ, Nζ)
    lλ_out = SparseProbVec.rand(post_lλ, logout=true)

    lλ_out
end

"""
    counttrans_mtdg(S, TT, ζ, L, K)

    ### Example
    ```julia
    L = 2
    K = 3
    TT = 12
    S = [1,2,1,3,3,1,2,1,3,2,1,1]
    ζ = [1,2,0,0,1,2,1,2,1,2]
      counttrans_mtdg(S, TT, ζ, L, K)
    ```
"""
function counttrans_mtdg(S::Vector{Int}, TT::Int, ζ::Vector{Int},
  L::Int, K::Int)

  ## initialize
  N0_out = zeros(Int, K)
  N_out = [ zeros(Int, (K,K)) for ℓ = 1:L ]

  ## pass through data and add counts
  for tt in (L+1):(TT)
    if ζ[tt-L] == 0
        N0_out[ S[tt] ] += 1
    else
        Slagrev_now = deepcopy( S[range(tt-1, step=-1, length=L)] )
        from = deepcopy( Slagrev_now[ ζ[tt-L] ] )
        N_out[ζ[tt-L]][ S[tt], from ] += 1
    end
  end

  N0_out, N_out
end

"""
    rpost_lQ_mtdg(S, TT, prior_Q0, Prior_Q, ζ, L, K)
"""
function rpost_lQ_mtdg(S::Vector{Int}, TT::Int,
    prior_Q0::Vector{Float64}, prior_Q::Vector{Matrix{Float64}},
    ζ::Vector{Int}, L::Int, K::Int)

    N0, N = counttrans_mtdg(S, TT, ζ, L, K)

    ## initialize
    α0_Q0 = deepcopy(prior_Q0)
    α0_Q = deepcopy(prior_Q)
    lQ_out = [ Matrix{Float64}(undef, K, K) for ℓ in 1:L ]

    α1_Q0 = α0_Q0 .+ N0
    α1_Q = [ α0_Q[ℓ] .+ N[ℓ] for ℓ in 1:L ]

    lQ0_out = SparseProbVec.rDirichlet(α1_Q0, logout=true)
    for ℓ in 1:L
        for j in 1:K
            lQ_out[ℓ][:,j] = SparseProbVec.rDirichlet(α1_Q[ℓ][:,j], logout=true)
        end
    end

    lQ0_out, lQ_out
end
function rpost_lQ_mtdg(S::Vector{Int}, TT::Int,
    prior_Q0::Vector{Float64}, prior_Q::Vector{Vector{SparseDirMix}},
    ζ::Vector{Int}, L::Int, K::Int)

    N0, N = counttrans_mtdg(S, TT, ζ, L, K)

    ## initialize
    α0_Q0 = deepcopy(prior_Q0)
    lQ_out = [ Matrix{Float64}(undef, K, K) for ℓ in 1:L ]

    α1_Q0 = α0_Q0 .+ N0
    lQ0_out = SparseProbVec.rDirichlet(α1_Q0, logout=true)

    for ℓ in 1:L
        for j in 1:K
            α1 = prior_Q[ℓ][j].α .+ N[ℓ][:,j]
            d = SparseProbVec.SparseDirMix(α1, prior_Q[ℓ][j].β)
            lQ_out[ℓ][:,j] = SparseProbVec.rand(d, logout=true)
        end
    end

    lQ0_out, lQ_out
end
function rpost_lQ_mtdg(S::Vector{Int}, TT::Int,
    prior_Q0::Vector{Float64}, prior_Q::Vector{Vector{SBMprior}},
    ζ::Vector{Int}, L::Int, K::Int)

    N0, N = counttrans_mtdg(S, TT, ζ, L, K)

    ## initialize
    α0_Q0 = deepcopy(prior_Q0)
    lQ_out = [ Matrix{Float64}(undef, K, K) for ℓ in 1:L ]

    α1_Q0 = α0_Q0 .+ N0
    lQ0_out = SparseProbVec.rDirichlet(α1_Q0, logout=true)

    for ℓ in 1:L
        for j in 1:K
            d = SparseProbVec.SBM_multinom_post(prior_Q[ℓ][j], N[ℓ][:,j])
            lQ_out[ℓ][:,j] = SparseProbVec.rand(d, logout=true)
        end
    end

    lQ0_out, lQ_out
end

"""
    rpost_ζ_mtdg(S, TT, lλ, lQ0, lQ, L, K)
"""
function rpost_ζ_mtdg(S::Vector{Int}, TT::Int,
  lλ::Vector{Float64}, lQ0::Vector{Float64}, lQ::Vector{Matrix{Float64}},
  L::Int, K::Int)

  ζ_out = Vector{Int}(undef, TT-L)

  for i in 1:(TT-L)
      tt = i + L
      Slagrev_now = S[range(tt-1, step=-1, length=L)]
      lp = Vector{Float64}(undef, L+1)

      lp[1] = lλ[1] + lQ0[S[tt]]
      for j in 1:L
          lp[j] = lλ[j+1] + lQ[j][ append!([deepcopy(S[tt])], deepcopy(Slagrev_now[j]))... ]
      end

      w = exp.( lp .- maximum(lp) )
      ζ_out[i] = StatsBase.sample(Weights(w)) - 1 # so that the values are 0:L
  end

  ζ_out
end

"""
    rpost_ζ_mtdg_marg(S, ζ_old, lλ, prior_Q0, prior_Q, TT, L, K)

Collapsed conditinal updates for ζ marginalizing over Q
"""
function rpost_ζ_mtdg_marg(S::Vector{Int}, ζ_old::Vector{Int},
    prior_Q0::Vector{Float64}, # always Dirichlet on intercept
    prior_Q::Union{Vector{Matrix{Float64}}, Vector{Vector{SparseDirMix}}, Vector{Vector{SBMprior}}},
    lλ::Vector{Float64},
    TT::Int, L::Int, K::Int)

    prior_Q_Dir_flag = typeof(prior_Q) <: Vector{Matrix{Float64}}

    α0_Q0 = deepcopy(prior_Q0)

    ζ_out = deepcopy(ζ_old)
    N0_now, N_now = counttrans_mtdg(S, TT, ζ_out, L, K) # in each N, rows are tos, cols are froms

    for i in 1:(TT-L)  # i indexes ζ, tt indexes S
        tt = i + L
        Slagrev_now = S[range(tt-1, step=-1, length=L)]

        N0init = deepcopy(N0_now)
        Ninit = deepcopy(N_now)
        if ζ_out[i] == 0
            N0init[ S[tt] ] -= 1
        else
            Ninit[ζ_out[i]][ S[tt], Slagrev_now[ ζ_out[i] ] ] -= 1
        end
        α1_Q0 = α0_Q0 .+ N0init
        if prior_Q_Dir_flag
            α1_Q = deepcopy(prior_Q)
            for ℓ in 1:L
                α1_Q[ℓ] += Ninit[ℓ]
            end
        end

        eSt = [1.0*(ii==S[tt]) for ii in 1:K]

        lmvbn_without = Vector{Float64}(undef, L+1)
        lmvbn_with = Vector{Float64}(undef, L+1)

        lmvbn_without[1] = SparseProbVec.lmvbeta( α1_Q0 ) # Dirichlet prior on Q0
        lmvbn_with[1] = SparseProbVec.lmvbeta( α1_Q0 + eSt ) # Dirichlet prior on Q0

        if prior_Q_Dir_flag
            for ℓ in 1:L
                lmvbn_without[ℓ+1] = SparseProbVec.lmvbeta( α1_Q[ℓ][:,Slagrev_now[ℓ]] )
                lmvbn_with[ℓ+1] = SparseProbVec.lmvbeta( α1_Q[ℓ][:,Slagrev_now[ℓ]] + eSt )
            end
        else # assumes a matching method for another prior type
            for ℓ in 1:L
                lmvbn_without[ℓ+1] = SparseProbVec.logMarginal( prior_Q[ℓ][Slagrev_now[ℓ]],
                        Ninit[ℓ][:,Slagrev_now[ℓ]] )
                lmvbn_with[ℓ+1] = SparseProbVec.logMarginal( prior_Q[ℓ][Slagrev_now[ℓ]],
                        Ninit[ℓ][:,Slagrev_now[ℓ]] + Int64.(eSt) )
            end
        end

        lw = lλ .+ sum(lmvbn_without)
        for ℓ in 0:L
            lw[ℓ+1] -= lmvbn_without[ℓ+1]
            lw[ℓ+1] += lmvbn_with[ℓ+1]
        end

        w = exp.( lw .- maximum(lw) )
        ζ_out[i] = StatsBase.sample(Weights( w )) - 1 # 0 refers to intercept (no lag dependence)
        N0_now = deepcopy(N0init)
        N_now = deepcopy(Ninit)
        if ζ_out[i] == 0
            N0_now[ S[tt] ] += 1
        else
            N_now[ζ_out[i]][ S[tt], Slagrev_now[ ζ_out[i] ] ] += 1
        end

    end

    ζ_out
end

"""
    MetropIndep_λζ(S::Vector{Int}, lλ_old::Vector{Float64}, ζ_old::,
        prior_λ, prior_Q0, prior_Q,
        TT::Int, L::Int, K::Int)

    Independence Metropolis step for λ and ζ.
"""
function MetropIndep_λζ(S::Vector{Int}, lλ_old::Vector{Float64}, ζ_old::Vector{Int},
    prior_λ::Union{Vector{Float64}, SparseDirMix, SBMprior},
    prior_Q0::Vector{Float64},
    prior_Q::Union{Vector{Matrix{Float64}}, Vector{Vector{SparseDirMix}}, Vector{Vector{SBMprior}}},
    TT::Int, L::Int, K::Int)

  if typeof(prior_λ) == Vector{Float64}
      lλ_cand = SparseProbVec.rDirichlet(prior_λ, logout=true)
  elseif typeof(prior_λ) == SBMprior
      lλ_cand = SparseProbVec.rand(prior_λ, logout=true)
  elseif typeof(prior_λ) == SparseDirMix
      lλ_cand = SparseProbVec.rand(prior_λ, logout=true)
  end

  λ_cand = exp.(lλ_cand)
  ζ_cand = [ StatsBase.sample( Weights(λ_cand) ) - 1 for i in 1:(TT-L) ]

  N0_cand, N_cand = counttrans_mtdg(S, TT, ζ_cand, L, K) # in each N, rows are tos, cols are froms
  N0_old, N_old = counttrans_mtdg(S, TT, ζ_old, L, K) # in each N, rows are tos, cols are froms

  ll_cand = SparseProbVec.lmvbeta(prior_Q0 .+ N0_cand) # can ignore denominators
  ll_old = SparseProbVec.lmvbeta(prior_Q0 .+ N0_old)

  if typeof(prior_Q) == Vector{Matrix{Float64}}
      for ℓ in 1:L
          for k in 1:K
              ll_cand += SparseProbVec.lmvbeta(prior_Q[ℓ][:,k] .+ N_cand[ℓ][:,k]) # can ignore denominators
              ll_old += SparseProbVec.lmvbeta(prior_Q[ℓ][:,k] .+ N_old[ℓ][:,k])
          end
      end
  else # assumes the logMarginal method exists for the type of prior
      for ℓ in 1:L
          for k in 1:K
              ll_cand += SparseProbVec.logMarginal(prior_Q[ℓ][k], N_cand[ℓ][:,k]) # can ignore denominators
              ll_old += SparseProbVec.logMarginal(prior_Q[ℓ][k], N_old[ℓ][:,k])
          end
      end
  end

  lu = log(rand())
  if lu < (ll_cand - ll_old)
      lλ_out = lλ_cand
      ζ_out = ζ_cand
  else
      lλ_out = lλ_old
      ζ_out = ζ_old
  end

  (lλ_out, ζ_out)
end

"""
mcmc!(model::ModMTDg, n_keep::Int, save::Bool=true,
    report_filename::String="out_progress.txt", thin::Int=1, jmpstart_iter::Int=25,
    report_freq::Int=1000;
    monitor::Vector{Symbol}=[:lλ, :lQ0, :lQ])
"""
function mcmc!(model::ModMTDg, n_keep::Int; save::Bool=true,
    report_filename::String="out_progress.txt", thin::Int=1, jmpstart_iter::Int=25,
    report_freq::Int=1000,
    monitor::Vector{Symbol}=[:lλ, :lQ0, :lQ])

    ## output files
    report_file = open(report_filename, "a+")
    write(report_file, "Commencing MCMC at $(Dates.now()) at iteration $(model.iter) for $(n_keep * thin) iterations.\n")
    close(report_file)

    if save
        sims = postSimsInit(monitor, n_keep, model.state)
    end

    ## sampling
    for i in 1:n_keep
        for j in 1:thin

            jmpstart = (model.iter % jmpstart_iter == 0)

            if jmpstart

                model.state.lλ, model.state.ζ = MetropIndep_λζ(model.S,
                    model.state.lλ, model.state.ζ, model.prior.λ,
                    model.prior.Q0, model.prior.Q,
                    model.TT, model.L, model.K)

            else

                model.state.ζ = rpost_ζ_mtdg_marg(model.S, model.state.ζ,
                    model.prior.Q0, model.prior.Q, model.state.lλ,
                    model.TT, model.L, model.K)

                model.state.lλ = rpost_lλ_mtdg(model.prior.λ, model.state.ζ, model.L)

            end

            model.state.lQ0, model.state.lQ = rpost_lQ_mtdg(model.S, model.TT,
                    model.prior.Q0, model.prior.Q,
                    model.state.ζ, model.L, model.K)

            model.iter += 1
            if model.iter % report_freq == 0
                report_file = open(report_filename, "a+")
                write(report_file, "Iter $(model.iter) at $(Dates.now())\n")
                llik_now = llik_MTDg(model.S, model.state.lλ,
                    model.state.lQ0, model.state.lQ)
                write(report_file, "Log-likelihood $(llik_now)\n")
                write(report_file, "Current lag weights: $( round.(exp.(model.state.lλ), digits=3) )\n\n")
                close(report_file)
            end
        end

        if save
            for field in monitor
              sims[i][field] = deepcopy(getfield(model.state, field))
            end
            sims[i][:llik] = llik_MTDg(model.S, model.state.lλ,
                model.state.lQ0, model.state.lQ)
        end

    end

    if save
        return sims
    else
        return model.iter
    end

end

"""
    llik_MTDg(S::Vector{Int}, lλ::Vector{Float64},
        lQ0::Vector{Float64}, lQ::Vector{Matrix{Float64}})

Computes the (conditional on intial values) log-likelihood for a time series.

### Example
```julia
    lλ = log.([0.0, 0.0, 0.0, 1.0])
    lQ0 = log.([0.2, 0.8])
    lQ = [ log.(reshape([0.7, 0.3, 0.25, 0.75], (2,2))),
           log.(reshape([0.7, 0.3, 0.25, 0.75], (2,2))),
           log.(reshape([0.7, 0.3, 0.25, 0.75], (2,2))) ]
    S = [1,1,2,1,1,2,2,1,2,2,1,2,1,2,1,1,1,1,2,1,2]
    llik_MTDg(S, lλ, lQ0, lQ)
```
"""
function llik_MTDg(S::Vector{Int}, lλ::Vector{Float64},
    lQ0::Vector{Float64}, lQ::Vector{Matrix{Float64}})

    TT = length(S)
    L = length(lλ) - 1

    n = TT - L

    lpt = Matrix{Float64}(undef, n, L+1)

    for i in 1:n
        tt = i + L
        Slagrev_now = S[range(tt-1, step=-1, length=L)]

        lpt[i,1] = lλ[1] + lQ0[S[tt]]
        for ℓ in 1:L
            lpt[i,ℓ+1] = lλ[ℓ+1] + lQ[ℓ][S[tt], Slagrev_now[ℓ]]
        end
    end

    lptvec = [ SparseProbVec.logsumexp(lpt[i,:]) for i in 1:n ]
    llik = sum(lptvec)

    return llik
end

"""
    TankReduction(λ, Q0, Q)

    Computes the unique Tank Reduction of the MTDg with an intercept.

    ### Example
    ```julia
    lam = [0.4, 0.5, 0.1]
    Q0 = [0.1, 0.2, 0.7]
    Q = [ [0.4 0.4 0.6; 0.3 0.1 0.2; 0.3 0.5 0.2],
          [0.2 0.5 0.2; 0.5 0.2 0.1; 0.3 0.3 0.7] ]

    TankReduction(lam, Q0, Q)
    ```
"""
function TankReduction(λ::Vector{Float64},
    Q0::Vector{Float64}, Q::Vector{Matrix{Float64}})

    # here each Q is organized with tos as rows and froms as columns

    K = length(Q0)
    L = length(λ) - 1

    Z0 = λ[1] .* Q0
    Z = [ λ[ℓ+1] .* Q[ℓ] for ℓ = 1:L ]

    a = [ [ minimum( Z[ℓ][k,:] ) for k = 1:K ] for ℓ = 1:L ]

    ZZ0 = deepcopy(Z0)
    ZZ = deepcopy(Z)
    for ℓ = 1:L
        ZZ0 += a[ℓ]
        for k = 1:K
            ZZ[ℓ][:,k] -= a[ℓ]
        end
    end

    λr = vcat( sum(ZZ0), [ sum(ZZ[ℓ][:,1]) for ℓ = 1:L ]  )
    Q0r = deepcopy(ZZ0)
    Qr = deepcopy(ZZ)

    if λr[1] > 0.0
        Q0r ./= λr[1]
    end

    for ℓ = 1:L
        if λr[ℓ+1] > 0.0
            Qr[ℓ] ./= λr[ℓ+1]
        end
    end

    λr, Q0r, Qr
end

"""
    forecDist_MTDg(Slag::Vector{Int}, λ::Vector{Float64}, Q::Matrix{Float64})

Computes the forecast distribution given lagged values and parameters.

### Example
```julia
    lam = [0.2, 0.4, 0.3, 0.1]
    Q0 = [0.5, 0.5]
    Q = [ reshape([0.7, 0.3, 0.25, 0.75], (2,2)) for ell = 1:3 ]
    Slagrev = [1,1,2]
    forecDist_MTDg(Slagrev, lam, Q0, Q)
```
"""
function forecDist_MTDg(Slagrev::Vector{Int}, λ::Vector{Float64},
    Q0::Vector{Float64}, Q::Vector{Matrix{Float64}})
    K = length(Q0)
    L = length(λ) - 1
    @assert length(Slagrev) == L
    w = zeros(Float64, K)
    for k in 1:K
        w[k] += λ[1] * Q0[k]
        for ℓ in 1:L
            w[k] += λ[ℓ+1] * Q[ℓ][k, Slagrev[ℓ]]
        end
    end
    w
end

max_stars_repo_name: mheiner/MTD

max_issues_repo_name: mheiner/MTD

max_forks_repo_name: mheiner/MTD


text:
# Autogenerated wrapper script for Clipper_jll for x86_64-linux-musl-cxx03
export libcclipper

JLLWrappers.@generate_wrapper_header("Clipper")
JLLWrappers.@declare_library_product(libcclipper, "libcclipper.so")
function __init__()
    JLLWrappers.@generate_init_header()
    JLLWrappers.@init_library_product(
        libcclipper,
        "lib/libcclipper.so",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@generate_init_footer()
end  # __init__()


text:
#=== TypeAliases ===#

# Opaque types
typealias CL_platform_id        Ptr{Void}
typealias CL_device_id          Ptr{Void}
typealias CL_context            Ptr{Void}
typealias CL_command_queue      Ptr{Void}
typealias CL_mem                Ptr{Void}
typealias CL_program            Ptr{Void}
typealias CL_kernel             Ptr{Void}
typealias CL_event              Ptr{Void}
typealias CL_sampler            Ptr{Void}

# Scalar types
typealias CL_char   Int8
typealias CL_uchar  UInt8
typealias CL_short  Int16
typealias CL_ushort UInt16
typealias CL_int    Int32
typealias CL_uint   UInt32
typealias CL_long   Int64
typealias CL_ulong  UInt64

typealias CL_half   Float16
typealias CL_float  Float32
typealias CL_double Float64

typealias CL_bool                       CL_uint
typealias CL_bitfield                   CL_ulong
typealias CL_device_type                CL_bitfield
typealias CL_platform_info              CL_uint
typealias CL_device_info                CL_uint
typealias CL_device_fp_config           CL_bitfield
typealias CL_device_mem_cache_type      CL_uint
typealias CL_device_local_mem_type      CL_uint
typealias CL_device_exec_capabilities   CL_bitfield
typealias CL_device_svm_capabilities    CL_bitfield
typealias CL_command_queue_properties   CL_bitfield
typealias CL_device_partition_property  Cssize_t #intptr_t
typealias CL_device_affinity_domain     CL_bitfield

typealias CL_context_properties         Cssize_t #intptr_t
typealias CL_context_info               CL_uint
typealias CL_queue_properties           CL_bitfield
typealias CL_command_queue_info         CL_uint
typealias CL_channel_order              CL_uint
typealias CL_channel_type               CL_uint
typealias CL_mem_flags                  CL_bitfield
typealias CL_svm_mem_flags              CL_bitfield
typealias CL_mem_object_type            CL_uint
typealias CL_mem_info                   CL_uint
typealias CL_mem_migration_flags        CL_bitfield
typealias CL_image_info                 CL_uint
typealias CL_buffer_create_type         CL_uint
typealias CL_addressing_mode            CL_uint
typealias CL_filter_mode                CL_uint
typealias CL_sampler_info               CL_uint
typealias CL_map_flags                  CL_bitfield
typealias CL_pipe_properties            Cssize_t #intptr_t
typealias CL_pipe_info                  CL_uint
typealias CL_program_info               CL_uint
typealias CL_program_build_info         CL_uint
typealias CL_build_status               CL_int
typealias CL_kernel_info                CL_uint
typealias CL_kernel_work_group_info     CL_uint
typealias CL_event_info                 CL_uint
typealias CL_command_type               CL_uint
typealias CL_profiling_info             CL_uint
typealias CL_sampler_properties         CL_bitfield
typealias CL_kernel_exec_info           CL_uint

# Scalar OpenGL types ! We should get these from OpenGL.jl

typealias GL_uint                       UInt32
typealias GL_int                        Int32

typealias GL_enum                       GL_uint

# interop types

typealias CL_GL_object_type             CL_uint
typealias CL_GL_texture_info            CL_uint
typealias CL_GL_platform_info           CL_uint
typealias CL_gl_context_info            CL_uint

typealias GL_sync                       Ptr{Void}

#=== Image Types ===#

immutable CL_image_format
    image_channel_order::CL_channel_order
    image_channel_data_type::CL_channel_type
end

immutable CL_image_desc
    image_type::CL_mem_object_type
    image_width::Csize_t
    image_depth::Csize_t
    image_array_size::Csize_t
    image_row_pitch::Csize_t
    image_slice_pitch::Csize_t
    num_mip_levels::CL_uint
    num_samples::CL_uint
    buffer::CL_mem
end

immutable CL_buffer_region
    origin::Csize_t
    size::Csize_t
end

#=== Conversion Functions ===#

cl_char(x)     = Int8(x)
cl_uchar(x)    = UInt8(x)
cl_short(x)    = Int16(x)
cl_ushort(x)   = UInt16(x)
cl_int(x)      = Int32(x)
cl_uint(x)     = UInt32(x)
cl_long(x)     = Int64(x)
cl_ulong(x)    = UInt64(x)

cl_half(x)     = UInt16(x)
cl_float(x)    = Float32(x)
cl_double(x)   = Float64(x)

cl_bool(x)     = x != 0 ? cl_uint(1) : cl_uint(0)
cl_bitfield(x) = cl_ulong(x)

cl_command_queue_properties(x) = cl_ulong(x)
cl_device_type(x)              = cl_bitfield(x)
cl_platform_info(x)            = cl_uint(x)
cl_device_info(x)              = cl_uint(x)
cl_device_fp_config(x)         = cl_bitfield(x)
cl_device_mem_cache_type(x)    = cl_uint(x)
cl_device_local_mem_type(x)    = cl_uint(x)
cl_device_exec_capabilities(x) = cl_bitfield(x)
cl_device_svm_capabilities(x)  = cl_bitfield(x)

cl_context_properties(x)       = convert(CL_context_properties, x)
cl_context_info(x)             = cl_uint(x)
cl_queue_properties(x)         = cl_bitfield(x)
cl_command_queue_info(x)       = cl_uint(x)
cl_channel_order(x)            = cl_uint(x)
cl_channel_type(x)             = cl_uint(x)
cl_mem_flags(x)                = cl_bitfield(x)
cl_svm_mem_flags(x)            = cl_bitfield(x)
cl_mem_object_type(x)          = cl_uint(x)
cl_mem_info(x)                 = cl_uint(x)
cl_image_info(x)               = cl_uint(x)
cl_buffer_create_type(x)       = cl_uint(x)
cl_addressing_mode(x)          = cl_uint(x)
cl_filter_mode(x)              = cl_uint(x)
cl_sampler_info(x)             = cl_uint(x)
cl_map_flags(x)                = cl_bitfield(x)
cl_pipe_properties(x)          = convert(CL_pipe_properties, x)
cl_pipe_info(x)                = cl_uint(x)
cl_program_info(x)             = cl_uint(x)
cl_program_build_info(x)       = cl_uint(x)
cl_build_status(x)             = cl_int(x)
cl_kernel_info(x)              = cl_uint(x)
cl_kernel_work_group_info(x)   = cl_uint(x)
cl_event_info(x)               = cl_uint(x)
cl_command_type(x)             = cl_uint(x)
cl_profiling_info(x)           = cl_uint(x)
cl_sampler_properties(x)       = cl_bitfield(x)
cl_kernel_exec(x)              = cl_uint(x)
cl_platform_id(x)              = Ptr{Void}(x)


text:
# ODEInterface.jl Algorithms

abstract type ODEInterfaceAlgorithm <: DiffEqBase.AbstractODEAlgorithm end
abstract type ODEInterfaceImplicitAlgorithm <: ODEInterfaceAlgorithm end
abstract type ODEInterfaceExplicitAlgroithm <: ODEInterfaceAlgorithm end
"""
dopri5: Hairer's classic implementation of the Dormand-Prince 4/5 method.
"""
struct dopri5 <: ODEInterfaceExplicitAlgroithm end
"""
dop853: Explicit Runge-Kutta 8(5,3) by Dormand-Prince.
"""
struct dop853 <: ODEInterfaceExplicitAlgroithm end
"""
odex: GBS extrapolation-algorithm based on the midpoint rule.
"""
struct odex <: ODEInterfaceExplicitAlgroithm end
"""
seulex: Extrapolation-algorithm based on the linear implicit Euler method.
"""
struct seulex{T} <: ODEInterfaceImplicitAlgorithm
    jac_lower::T
    jac_upper::T
end
"""
radau: Implicit Runge-Kutta (Radau IIA) of variable order between 5 and 13.
"""
struct radau{T} <: ODEInterfaceImplicitAlgorithm
    jac_lower::T
    jac_upper::T
end
"""
radau5: Implicit Runge-Kutta method (Radau IIA) of order 5.
"""
struct radau5{T} <: ODEInterfaceImplicitAlgorithm
    jac_lower::T
    jac_upper::T
end
"""
rodas: Rosenbrock 4(3) method.
"""
struct rodas{T} <: ODEInterfaceImplicitAlgorithm
    jac_lower::T
    jac_upper::T
end
"""
ddeabm: Adams-Bashforth-Moulton Predictor-Corrector method (order between 1 and 12)
"""
struct ddeabm <: ODEInterfaceExplicitAlgroithm end
"""
ddebdf: Backward Differentiation Formula (orders between 1 and 5)
"""
struct ddebdf{T} <: ODEInterfaceImplicitAlgorithm
    jac_lower::T
    jac_upper::T
end

seulex(;jac_lower=nothing,jac_upper=nothing) = seulex(jac_lower,jac_upper)
radau(;jac_lower=nothing,jac_upper=nothing) = radau(jac_lower,jac_upper)
radau5(;jac_lower=nothing,jac_upper=nothing) = radau5(jac_lower,jac_upper)
rodas(;jac_lower=nothing,jac_upper=nothing) = rodas(jac_lower,jac_upper)
ddebdf(;jac_lower=nothing,jac_upper=nothing) = ddebdf(jac_lower,jac_upper)


text:
module GalacticOptim

using DiffEqBase, Requires
using DiffResults, ForwardDiff, Zygote
using Optim

include("problem.jl")
include("solve.jl")
include("function.jl")

export OptimizationProblem, OptimizationFunction
export solve

export BBO

end # module


text:
function get_refraction(obj::AtmosphericRefractionModel, arg0::jdouble)
    return jcall(obj, "getRefraction", jdouble, (jdouble,), arg0)
end


text:
include("structure.jl")
using LinearAlgebra

function initialize(file)
# PART OF YAML DATA EXTRACTION
    data = YAML.load_file(file)
    raw = data["structure"]
    s = Structure(
        hcat(raw["lattice"]...)',
        raw["species"],
        hcat(raw["fcoords"]...)',
        data["replicate"],
        data["composition"]
    )
    sc = supercell(s)
# Now, the shortest distance is being evualated by periodic boundary conditions
# and stored in 'vec_mat'
    rg = [1, 0, -1]
    shifts = (collect(Iterators.product(rg, rg, rg)))
    shifts = reshape(shifts, :, 1)
    cart_shifts = []
    distance_matrix = zeros(size(sc.fcoords)[1], size(sc.fcoords)[1])
    vec_mat = zeros(size(sc.fcoords)[1], size(sc.fcoords)[1], 3)
    min_dist = 1e-6
    for i in shifts
        push!(cart_shifts, collect(i))
    end

    for i in 1:size(sc.fcoords)[1]
        for j in 1:size(sc.fcoords)[1]
            if i == j
                continue
            else
            d_vecs = []
            for sh = 1:length(cart_shifts)
                push!(d_vecs, sc.lattice*reshape(sc.fcoords[j, :] + cart_shifts[sh] - sc.fcoords[i, :], (3, 1)))
            end
            for k = 1:length(d_vecs)
                for l = 1:length(d_vecs[k])
                    d_vecs[k][l] = round(d_vecs[k][l], digits = 5)
                end
            end
            min_vec = [Inf, Inf, Inf]
                for k in d_vecs
                    if norm(k) < norm(min_vec) && norm(k) > min_dist
                        vec_mat[i, j, :] = k
                        min_vec = k
                    end
                end
            end
        end
    end
# The distance matrix is being created out of the shortest distances
    for i = 1:size(vec_mat)[1]
        for j = 1:size(vec_mat)[1]
            distance_matrix[i, j] = round(norm(vec_mat[i, j, :]), digits = 4) #changed round(digits 3!)
        end
    end
# Structures from the functions below are created here.
    distances = (sort!(unique(distance_matrix)))
    distance_dict = Dict(dist => index-1 for (index, dist) in enumerate(distances))
    shell_matrix = map(num -> distance_dict[num], distance_matrix)
    shells = coord_shells(shell_matrix)
    triplet_shell_matrix = triplet_shell(shell_matrix)
    map!(t -> cutoff(t), triplet_shell_matrix, triplet_shell_matrix) #induces cutoff
    mole_fracs, elem_dict = mole_fractions(sc)
    pair_prefactors = calc_pair_pref(mole_fracs, shells, shell_matrix)
    trip_prefactors = calc_trip_prefs(triplet_shell_matrix, shells)
# Finally, the two arrays for SRO value storage
    alpha = ones(length(mole_fracs), length(mole_fracs), length(shells))
    beta = ones(length(mole_fracs), length(mole_fracs), length(mole_fracs), length(shells), length(shells), length(shells))
# DONE :-)
    return Backpack(sc, elem_dict, mole_fracs, shell_matrix, triplet_shell_matrix, alpha, beta, pair_prefactors, trip_prefactors)
end

function cutoff((s1, s2, s3), cutoff=30)
# Simple cutoff help function
    any([s1 >= cutoff, s2 >= cutoff, s3 >= cutoff ]) && return (0,0,0)
    any([s1 == 0, s2 == 0, s3 == 0]) && return (0,0,0)
    return (s1, s2, s3)
end

function triplet_shell(shell_matrix::Array, cutoff = Inf)
# constructs a 3D-matrix containing a tuple of all three shells between
# atoms i, j and k
    sm = shell_matrix
    triplet_shell_matrix = [(sm[i, j], sm[j, k], sm[i, k]) for i=1:size(sm)[1], j=1:size(sm)[1], k=1:size(sm)[1]]
    return triplet_shell_matrix
end

function coord_shells(shell_matrix::Array)
# returns a dict with the shell number as key and the number of appearing atoms
# (with pbc) as values
    shell_dict = Dict(i => 0.0 for i = 1:maximum(shell_matrix))
    for i = 1:size(shell_matrix)[1], j = i:size(shell_matrix)[1]
        shell_matrix[i, j] == 0 && continue
        shell_dict[shell_matrix[i, j]] += 2/size(shell_matrix)[1]
    end
    return shell_dict
end

function calc_trip_prefs(triplet_shell_matrix::Array, M::Dict)
# returns a dictionary with all occurring triplets as keys and their corresponding
# prefactors
    tt = unique(triplet_shell_matrix) #all triplet types
    triplet_prefactors = Dict(map(t-> trip_pref_helper(t, M, size(triplet_shell_matrix)[1]), tt))
    return triplet_prefactors
end

function trip_pref_helper(s::Tuple, M::Dict, sc_size)
# Calculates the fixed part ν and rules out cases with a zero
    if any(i -> i == 0, s) == true
        return (0,0,0) => nothing
    else
        return s => 1/(3*sc_size)*(1/(M[s[1]]*M[s[2]])+1/(M[s[1]]*M[s[3]])+1/(M[s[2]]*M[s[3]]))
    end
end

function calc_pair_pref(mole_fractions, shells, shell_matrix)
# returns an 3D-array with all prefactors. Prefactor for a pair between elements
# 1 and 3 - 4 shells apart- is found at pair_prefactors[1, 3, 4] (and [3, 1, 4])
    pair_prefactors = zeros(length(mole_fractions), length(mole_fractions), length(shells))
    for i = 1:length(shells)
        for (idx1, val1) in enumerate(mole_fractions)
            if val1[1][2] != 0
                for (idx2, val2) in enumerate(mole_fractions)
                    if val2[1][2] != 0
                        #Factor 2 because there is no double counting later on
                        pair_prefactors[val1[1][2], val2[1][2], i] = 2/(val1[2]*val2[2]*shells[i]*size(shell_matrix)[1])
                    end
                end
            end
        end
    end
    return pair_prefactors
end

function mole_fractions(s::Structure)
# Returns mole fraction and an dictionary with all elements in it. The values
# in element_dict are the keys for mole_fractions
    mf = Dict()
    elem_dict = Dict()
    seq = s.species
    elem_num = 1
    for (idx1, val1) in enumerate(s.composition)
        sub_sum = (sum(values(val1[2])))
        for (idx2, val2) in enumerate(val1[2])
            if val2[1] == "Vac" #Vacancy exception: returns 0
                mf[(idx1, 0)] = val2[2]/sub_sum
                elem_dict[val2[1]] = (idx1, 0)
                i = 1
                j = 1
            else
                mf[(idx1, elem_num)] = val2[2]/sub_sum #key is sublattice + element number
                elem_dict[val2[1]] = (idx1, elem_num) #key above is the value here
                i = 1
                j = 1
                elem_num += 1
            end
            for i = 1:length(seq)
                if seq[i] == val1[1] && j <= val2[2]
                    seq[i] = val2[1]
                    j += 1
                end
            end
            i += 1
        end
    end
    return mf, elem_dict
end

max_stars_repo_name: ulipoto/WCSRO

max_issues_repo_name: ulipoto/WCSRO

max_forks_repo_name: ulipoto/WCSRO


text:
using ClimateMachine
const clima_dir = dirname(dirname(pathof(ClimateMachine)));

if parse(Bool, get(ENV, "CLIMATEMACHINE_PLOT_EDMF_COMPARISON", "false"))
    plot_dir = joinpath(clima_dir, "output", "sbl_edmf", "pycles_comparison")
else
    plot_dir = nothing
end

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

data_file = Dataset(joinpath(PyCLES_output_dataset_path, "Gabls.nc"), "r")

#! format: off
best_mse = OrderedDict()
best_mse["prog_ρ"] = 9.3809207150296822e-03
best_mse["prog_ρu_1"] = 6.7269975116623837e+03
best_mse["prog_ρu_2"] = 6.8630628605220889e-01
#! format: on

computed_mse = compute_mse(
    solver_config.dg.grid,
    solver_config.dg.balance_law,
    time_data,
    dons_arr,
    data_file,
    "Gabls",
    best_mse,
    1800,
    plot_dir,
)

@testset "SBL Anelastic Solution Quality Assurance (QA) tests" begin
    #! format: off
    test_mse(computed_mse, best_mse, "prog_ρ")
    test_mse(computed_mse, best_mse, "prog_ρu_1")
    test_mse(computed_mse, best_mse, "prog_ρu_2")
    #! format: on
end


text:
# set the working directory to the directory in which this file is contained
cd(@__DIR__)
println("Installing and compiling dependencies...")
println("This will take several minutes...")
println("If prompted to install Anaconda, press Y to accept...")
# load the package manager
using Pkg
Pkg.add("IJulia")
Pkg.precompile()
# activate the project environment
Pkg.activate("")
# initialize project environment
Pkg.instantiate()
Pkg.update()
Pkg.precompile()


text:
using GuidingCenterOrbits

# Dummy print test
@test True

# Write your own  test here


text:
@testset "partial" begin

@testset "create_∂ and apply_∂!, 1D" begin
    N = (10,)
    M = prod(N)

    Fu = rand(N...)
    Gv = similar(Fu)
    gv = zeros(M)

    Mask = similar(Fu)  # masking array
    for nw = 1:1
        sub = Vector{Int}(undef, 1)
        sub′ = Vector{Int}(undef, 1)

        Nw = N[nw]
        ∆w⁻¹ = rand(Nw)

        for ns = (-1,1), isbloch = (true,false)  # (backward,forward difference), (Bloch,symmetric)
            ∂ws = spzeros(M,M)

            for ind = 1:M
                sub .= CartesianIndices(N)[ind].I  # subscripts of diagonal entry
                indw = sub[nw]
                ∆wᵢ⁻¹ = ∆w⁻¹[indw]
                ∂ws[ind,ind] = -ns * ∆wᵢ⁻¹  # diagonal entries

                # Calculate the column index of the off-diagonal entry in the row `ind`.
                sub′ .= sub  # subscripts of off-diagonal entry
                if ns == 1  # forward difference
                    if sub′[nw] == Nw
                        sub′[nw] = 1
                    else
                        sub′[nw] += 1
                    end
                else  # backward difference
                    if sub′[nw] == 1
                        sub′[nw] = Nw
                    else
                        sub′[nw] -= 1
                    end
                end

                ind′ = LinearIndices(N)[sub′...]
                ∂ws[ind, ind′] += ns * ∆wᵢ⁻¹  # off-diagonal entry
            end

            if !isbloch  # symmetry boundary
                # Initialize the masking array.
                Mask .= 1
                Mask[Base.setindex(axes(Mask), 1, nw)...] = 0  # not .= 0 unlike code for 2D and 3D
                if ns == 1  # forward difference
                    # The input fields at the symmetry boundary should be zero, so apply the
                    # mask to the input field.
                    ∂ws = ∂ws * spdiagm(0=>Mask[:])
                else  # backward difference
                    # The output fields at the symmetry boundary should be zero, so apply
                    # mask to the output field.
                    ∂ws = spdiagm(0=>Mask[:]) * ∂ws
                end
            end

            # Test create_∂.
            @test create_∂(nw, ns==1, [N...], ∆w⁻¹, isbloch) == ∂ws

            # Test apply_∂!.
            fu = Fu[:]
            mul!(gv, ∂ws, fu)
            apply_∂!(Gv, Fu, Val(:(=)), nw, ns==1, ∆w⁻¹, isbloch)
            @test Gv[:] ≈ gv

            # print("matrix: "); @btime mul!($gv, $∂ws, $fu)
            # print("matrix-free: "); @btime apply_∂!($Gv, $Fu, $nw, $ns==1, $∆w⁻¹)
            # println()
        end
    end
end  # @testset "create_∂ and apply_∂!, 1D"

@testset "create_∂ and apply_∂!, 2D" begin
    N = (8,9)
    M = prod(N)

    Fu = rand(N...)
    Gv = similar(Fu)
    gv = zeros(M)

    Mask = similar(Fu)  # masking array
    for nw = 1:2
        sub = Vector{Int}(undef, 2)
        sub′ = Vector{Int}(undef, 2)

        Nw = N[nw]
        ∆w⁻¹ = rand(Nw)

        for ns = (-1,1), isbloch = (true,false)  # (backward,forward difference), (Bloch,symmetric)
            ∂ws = spzeros(M,M)

            for ind = 1:M
                sub .= CartesianIndices(N)[ind].I  # subscripts of diagonal entry
                indw = sub[nw]
                ∆wᵢ⁻¹ = ∆w⁻¹[indw]
                ∂ws[ind,ind] = -ns * ∆wᵢ⁻¹  # diagonal entries

                # Calculate the column index of the off-diagonal entry in the row `ind`.
                sub′ .= sub  # subscripts of off-diagonal entry
                if ns == 1  # forward difference
                    if sub′[nw] == Nw
                        sub′[nw] = 1
                    else
                        sub′[nw] += 1
                    end
                else  # backward difference
                    if sub′[nw] == 1
                        sub′[nw] = Nw
                    else
                        sub′[nw] -= 1
                    end
                end

                ind′ = LinearIndices(N)[sub′...]
                ∂ws[ind, ind′] += ns * ∆wᵢ⁻¹  # off-diagonal entry
            end

            if !isbloch  # symmetry boundary
                # Initialize the masking array.
                Mask .= 1
                Mask[Base.setindex(axes(Mask), 1, nw)...] .= 0
                if ns == 1  # forward difference
                    # The input fields at the symmetry boundary should be zero, so apply the
                    # mask to the input field.
                    ∂ws = ∂ws * spdiagm(0=>Mask[:])
                else  # backward difference
                    # The output fields at the symmetry boundary should be zero, so apply
                    # mask to the output field.
                    ∂ws = spdiagm(0=>Mask[:]) * ∂ws
                end
            end

            # Test create_∂.
            @test create_∂(nw, ns==1, [N...], ∆w⁻¹, isbloch) == ∂ws

            # Test apply_∂!.
            fu = Fu[:]
            mul!(gv, ∂ws, fu)
            apply_∂!(Gv, Fu, Val(:(=)), nw, ns==1, ∆w⁻¹, isbloch)
            @test Gv[:] ≈ gv

            # print("matrix: "); @btime mul!($gv, $∂ws, $fu)
            # print("matrix-free: "); @btime apply_∂!($Gv, $Fu, $nw, $ns==1, $∆w⁻¹)
            # println()
        end
    end
end  # @testset "create_∂ and apply_∂!, 2D"

@testset "create_∂ and apply_∂!, 3D" begin
    N = (8,9,10)
    M = prod(N)

    Fu = rand(N...)
    Gv = similar(Fu)
    gv = zeros(M)

    Mask = similar(Fu)  # masking array
    for nw = 1:3
        sub = Vector{Int}(undef, 3)
        sub′ = Vector{Int}(undef, 3)

        Nw = N[nw]
        ∆w⁻¹ = rand(Nw)

        for ns = (-1,1), isbloch = (true,false)  # (backward,forward difference), (Bloch,symmetric)
            ∂ws = spzeros(M,M)

            for ind = 1:M
                sub .= CartesianIndices(N)[ind].I  # subscripts of diagonal entry
                indw = sub[nw]
                ∆wᵢ⁻¹ = ∆w⁻¹[indw]
                ∂ws[ind,ind] = -ns * ∆wᵢ⁻¹  # diagonal entries

                # Calculate the column index of the off-diagonal entry in the row `ind`.
                sub′ .= sub  # subscripts of off-diagonal entry
                if ns == 1  # forward difference
                    if sub′[nw] == Nw
                        sub′[nw] = 1
                    else
                        sub′[nw] += 1
                    end
                else  # backward difference
                    if sub′[nw] == 1
                        sub′[nw] = Nw
                    else
                        sub′[nw] -= 1
                    end
                end

                ind′ = LinearIndices(N)[sub′...]
                ∂ws[ind, ind′] += ns * ∆wᵢ⁻¹  # off-diagonal entry
            end

            if !isbloch  # symmetry boundary
                # Initialize the masking array.
                Mask .= 1
                Mask[Base.setindex(axes(Mask), 1, nw)...] .= 0
                if ns == 1  # forward difference
                    # The input fields at the symmetry boundary should be zero, so apply the
                    # mask to the input field.
                    ∂ws = ∂ws * spdiagm(0=>Mask[:])
                else  # backward difference
                    # The output fields at the symmetry boundary should be zero, so apply
                    # mask to the output field.
                    ∂ws = spdiagm(0=>Mask[:]) * ∂ws
                end
            end

            # Test create_∂.
            @test create_∂(nw, ns==1, [N...], ∆w⁻¹, isbloch) == ∂ws

            # Test apply_∂!.
            fu = Fu[:]
            mul!(gv, ∂ws, fu)
            apply_∂!(Gv, Fu, Val(:(=)), nw, ns==1, ∆w⁻¹, isbloch)
            @test Gv[:] ≈ gv

            # print("matrix: "); @btime mul!($gv, $∂ws, $fu)
            # print("matrix-free: "); @btime apply_∂!($Gv, $Fu, $nw, $ns==1, $∆w⁻¹)
            # println()
        end
    end
end  # @testset "create_∂ and apply_∂!, 3D"

end  # @testset "differential"


text:
#
# PointSpreadFunctions.jl --
#
# Modeling and fitting of Point Spread Functions (PSF) for Julia.
#

module PointSpreadFunctions

export
    AbstractPSF,
    AiryPSF,
    CauchyPSF,
    GaussianPSF,
    MoffatPSF,
    getfwhm,
    findzeros

import Base: getindex, isvalid
using Base: @propagate_inbounds

using OptimPack

using SpecialFunctions

function fit end
function immutable end

include("models.jl")
include("fitting.jl")

end # module


text:
"""
A collection of functions for solving and analyzing school-choice problems.
See `examples/Tutorial.jl` for usage examples.

Some of the algorithms in this module are optimized for preference lists given
in "permutation" order, and others in "rating" order. Which order is used can be confirmed
by inspecting each function call. If the input calls for `students`, then the function
expects "rating" order, where if column `i` of students is `[3, 1, 2]`, this implies that
school 1 is `i`'s 3rd choice, school 2 is her 1st choice, and school 3 is her 2nd choice.
If the input calls for `students_inv`, then the expected input is `[2, 3, 1]`, which means
the same. These can be trivially related via `invperm()`.
"""
module DeferredAcceptance

using StatsBase
using LinearAlgebra
using Random
using Combinatorics
using Polyhedra

export singletiebreaking, multipletiebreaking, hybridtiebreaking, welfaretiebreaking, choiceaugmented
export deferredacceptance, nonatomicdeferredacceptance, nonatomicdeferredacceptance_iid, TTC, TTC_match, RSD,
       nonatomicsecant, nonatomictatonnement
export isstable, ismarketclearing, DA_rank_dist
export assignmentfrompreflists, demandfrompreflists, demandfromMNL_iid, demandfromMNL_singlescore,
       demandfrommixedMNL_tscores

include("demandfunctions.jl")
include("discretematch.jl")
include("nonatomic.jl")
include("tiebreakers.jl")

"""
    argsort(vec)

Associate each item in vec with its (descending) rank. Convenience wrapper of `sortperm()`.
"""
function argsort(vec::AbstractArray{<:Real, 1})::AbstractArray{Int, 1}
    return invperm(sortperm(vec))
end

end

max_stars_repo_name: maxkapur/DeferredAcceptance

max_issues_repo_name: maxkapur/DeferredAcceptance

max_forks_repo_name: maxkapur/DeferredAcceptance


text:
module LinAlgTools

include("gaussian_elimination.jl")

end


text:
using DelimitedFiles

# write_to_smat
function writeToSMAT(M,filename)
    #if unweighted
    if size(M,2) == 2
        M = hcat(M,ones(Int,size(M,1)))
    end
    m = maximum(M)+1
    nz = size(M,1)
    open(filename, "w") do f
        write(f, "$m   $m   $nz \n")
        writedlm(f,M)
    end
end
function findin_index(x,y)
     indices_in_y = zeros(Int64,length(x))
     already_exist = findall((in)(y), x)
     donot_exist = setdiff(1:length(x),already_exist)
     funcmap = i -> indices_in_y[findall(x.==y[i])] .= i
     lookfor_indices = findall((in)(x), y)
     map(funcmap,lookfor_indices)
     return indices_in_y
end

function run_script()
    M = readdlm("playerinfo.txt",header=true)
    names = M[1][:,3].*" ".*M[1][:,2]
    name_ids = M[1][:,1]
    names_data = hcat(name_ids,names)

    allfiles = readdir("pitcherinfo")
    myedge_list = Array{Int64}(undef,0,2)
    for i = 1:length(allfiles)
        fileid = allfiles[i]
        Mi = readdlm(joinpath("pitcherinfo",fileid))
        # 1(skip) - 20 - 1(skip) - 20
        # the crawler needs to change to not crawl "pitcher" every 20 lines
        interestingids = findall(Mi[:,1].!="Pitcher")
        curfullnames = Mi[interestingids,1].*" ".*Mi[interestingids,2]
        indexnumbers = findin_index(curfullnames,names_data[:,2])
        connectto = indexnumbers[findall(indexnumbers.!=0)]
        curplayerid = fileid[1:findfirst("_",fileid)[1]-1]
        curplayerid_inlist = findfirst(name_ids.==curplayerid)
        new_connections = hcat(ones(Int64,length(connectto)).*curplayerid_inlist,connectto)
        myedge_list = vcat(myedge_list,new_connections)
    end
myedge_list.-=1
return myedge_list,names_data
end

max_stars_repo_name: empet/more-graphs

max_issues_repo_name: empet/more-graphs

max_forks_repo_name: empet/more-graphs


text:
using CUSPARSE
using CUDArt
using Base.Test

m = 25
n = 35
k = 10
blockdim = 5

@testset "mv" begin
    for elty in [Float32,Float64,Complex64,Complex128]
        A = sparse(rand(elty,m,m))
        x = rand(elty,m)
        y = rand(elty,m)
        alpha = rand(elty)
        beta = rand(elty)
        @testset "mv_symm" begin
            A_s = A + A.'
            @testset "csr" begin
                d_x = CudaArray(x)
                d_y = CudaArray(y)
                d_A = Symmetric(CudaSparseMatrixCSR(A_s))
                d_y = CUSPARSE.mv('N',alpha,d_A,d_x,beta,d_y,'O')
                h_y = to_host(d_y)
                z = alpha * A_s * x + beta * y
                @test z ≈ h_y
                x_  = rand(elty,n)
                d_x = CudaArray(x_)
                @test_throws DimensionMismatch CUSPARSE.mv('T',alpha,d_A,d_x,beta,d_y,'O')
                @test_throws DimensionMismatch CUSPARSE.mv('N',alpha,d_A,d_y,beta,d_x,'O')
            end
            @testset "csc" begin
                d_x = CudaArray(x)
                d_y = CudaArray(y)
                d_A = Symmetric(CudaSparseMatrixCSC(A_s))
                d_y = CUSPARSE.mv('N',alpha,d_A,d_x,beta,d_y,'O')
                h_y = to_host(d_y)
                z = alpha * A_s * x + beta * y
                @test z ≈ h_y
                x_  = rand(elty,n)
                d_x = CudaArray(x_)
                @test_throws DimensionMismatch CUSPARSE.mv('T',alpha,d_A,d_x,beta,d_y,'O')
                @test_throws DimensionMismatch CUSPARSE.mv('N',alpha,d_A,d_y,beta,d_x,'O')
            end
        end

        @testset "mv_herm" begin
            A_h = A + A'
            @testset "csr" begin
                d_x = CudaArray(x)
                d_y = CudaArray(y)
                d_A = Hermitian(CudaSparseMatrixCSR(A_h))
                d_y = CUSPARSE.mv('N',alpha,d_A,d_x,beta,d_y,'O')
                h_y = to_host(d_y)
                z = alpha * A_h * x + beta * y
                @test z ≈ h_y
                x_  = rand(elty,n)
                d_x = CudaArray(x_)
                @test_throws DimensionMismatch CUSPARSE.mv('T',alpha,d_A,d_x,beta,d_y,'O')
                @test_throws DimensionMismatch CUSPARSE.mv('N',alpha,d_A,d_y,beta,d_x,'O')
            end
            @testset "csc" begin
                d_x = CudaArray(x)
                d_y = CudaArray(y)
                d_A = Hermitian(CudaSparseMatrixCSC(A_h))
                d_y = CUSPARSE.mv('N',alpha,d_A,d_x,beta,d_y,'O')
                h_y = to_host(d_y)
                z = alpha * A_h * x + beta * y
                @test z ≈ h_y
                x_  = rand(elty,n)
                d_x = CudaArray(x_)
                @test_throws DimensionMismatch CUSPARSE.mv('T',alpha,d_A,d_x,beta,d_y,'O')
                @test_throws DimensionMismatch CUSPARSE.mv('N',alpha,d_A,d_y,beta,d_x,'O')
            end
        end
        A = sparse(rand(elty,m,n))
        x = rand(elty,n)
        y = rand(elty,m)
        alpha = rand(elty)
        beta = rand(elty)
        @testset "mv" begin
            @testset "csr" begin
                d_x = CudaArray(x)
                d_y = CudaArray(y)
                d_A = CudaSparseMatrixCSR(A)
                @test_throws DimensionMismatch CUSPARSE.mv('T',alpha,d_A,d_x,beta,d_y,'O')
                @test_throws DimensionMismatch CUSPARSE.mv('N',alpha,d_A,d_y,beta,d_x,'O')
                d_z = CUSPARSE.mv('N',alpha,d_A,d_x,beta,d_y,'O')
                h_z = to_host(d_z)
                z = alpha * A * x + beta * y
                @test z ≈ h_z
                d_z = CUSPARSE.mv('N',d_A,d_x,beta,d_y,'O')
                h_z = to_host(d_z)
                z = A * x + beta * y
                @test z ≈ h_z
                d_z = CUSPARSE.mv('N',d_A,d_x,d_y,'O')
                h_z = to_host(d_z)
                z = A * x + y
                @test z ≈ h_z
                d_z = CUSPARSE.mv('N',alpha,d_A,d_x,d_y,'O')
                h_z = to_host(d_z)
                z = alpha * A * x + y
                @test z ≈ h_z
                d_z = CUSPARSE.mv('N',alpha,d_A,d_x,'O')
                h_z = to_host(d_z)
                z = alpha * A * x
                @test z ≈ h_z
                d_z = CUSPARSE.mv('N',d_A,d_x,'O')
                h_z = to_host(d_z)
                z = A * x
                @test z ≈ h_z
                d_z = d_A*d_x
                h_z = to_host(d_z)
                z = A * x
                @test z ≈ h_z
            end
            @testset "csc" begin
                d_x = CudaArray(x)
                d_y = CudaArray(y)
                d_A = CudaSparseMatrixCSC(A)
                @test_throws DimensionMismatch CUSPARSE.mv('T',alpha,d_A,d_x,beta,d_y,'O')
                @test_throws DimensionMismatch CUSPARSE.mv('N',alpha,d_A,d_y,beta,d_x,'O')
                d_z = CUSPARSE.mv('N',alpha,d_A,d_x,beta,d_y,'O')
                h_z = to_host(d_z)
                z = alpha * A * x + beta * y
                @test z ≈ h_z
                d_z = CUSPARSE.mv('N',d_A,d_x,beta,d_y,'O')
                h_z = to_host(d_z)
                z = A * x + beta * y
                @test z ≈ h_z
                d_z = CUSPARSE.mv('N',d_A,d_x,d_y,'O')
                h_z = to_host(d_z)
                z = A * x + y
                @test z ≈ h_z
                d_z = CUSPARSE.mv('N',alpha,d_A,d_x,d_y,'O')
                h_z = to_host(d_z)
                z = alpha * A * x + y
                @test z ≈ h_z
                d_z = CUSPARSE.mv('N',alpha,d_A,d_x,'O')
                h_z = to_host(d_z)
                z = alpha * A * x
                @test z ≈ h_z
                d_z = CUSPARSE.mv('N',d_A,d_x,'O')
                h_z = to_host(d_z)
                z = A * x
                @test z ≈ h_z
                d_z = d_A*d_x
                h_z = to_host(d_z)
                z = A * x
                @test z ≈ h_z
            end
            @testset "bsr" begin
                d_x = CudaArray(x)
                d_y = CudaArray(y)
                d_A = CudaSparseMatrixCSR(A)
                d_A = CUSPARSE.switch2bsr(d_A,convert(Cint,blockdim))
                @test_throws DimensionMismatch CUSPARSE.mv('T',alpha,d_A,d_x,beta,d_y,'O')
                @test_throws DimensionMismatch CUSPARSE.mv('N',alpha,d_A,d_y,beta,d_x,'O')
                d_z = CUSPARSE.mv('N',alpha,d_A,d_x,beta,d_y,'O')
                h_z = to_host(d_z)
                z = alpha * A * x + beta * y
                @test z ≈ h_z
                d_z = CUSPARSE.mv('N',d_A,d_x,beta,d_y,'O')
                h_z = to_host(d_z)
                z = A * x + beta * y
                @test z ≈ h_z
                d_z = CUSPARSE.mv('N',d_A,d_x,d_y,'O')
                h_z = to_host(d_z)
                z = A * x + y
                @test z ≈ h_z
                d_z = CUSPARSE.mv('N',alpha,d_A,d_x,d_y,'O')
                h_z = to_host(d_z)
                z = alpha * A * x + y
                @test z ≈ h_z
                d_z = CUSPARSE.mv('N',alpha,d_A,d_x,'O')
                h_z = to_host(d_z)
                z = alpha * A * x
                @test z ≈ h_z
                d_z = CUSPARSE.mv('N',d_A,d_x,'O')
                h_z = to_host(d_z)
                z = A * x
                @test z ≈ h_z
                d_z = d_A*d_x
                h_z = to_host(d_z)
                z = A * x
                @test z ≈ h_z
                @test_throws DimensionMismatch CUSPARSE.mv('T',alpha,d_A,d_x,beta,d_y,'O')
                @test_throws DimensionMismatch CUSPARSE.mv('N',alpha,d_A,d_y,beta,d_x,'O')
            end
            @testset "hyb" begin
                d_x = CudaArray(x)
                d_y = CudaArray(y)
                d_A = CudaSparseMatrixCSR(A)
                d_A = CUSPARSE.switch2hyb(d_A)
                d_z = CUSPARSE.mv('N',alpha,d_A,d_x,beta,d_y,'O')
                h_z = to_host(d_z)
                z = alpha * A * x + beta * y
                @test z ≈ h_z
                d_z = CUSPARSE.mv('N',d_A,d_x,beta,d_y,'O')
                h_z = to_host(d_z)
                z = A * x + beta * y
                @test z ≈ h_z
                d_z = CUSPARSE.mv('N',d_A,d_x,d_y,'O')
                h_z = to_host(d_z)
                z = A * x + y
                @test z ≈ h_z
                d_z = CUSPARSE.mv('N',alpha,d_A,d_x,d_y,'O')
                h_z = to_host(d_z)
                z = alpha * A * x + y
                @test z ≈ h_z
                d_z = CUSPARSE.mv('N',alpha,d_A,d_x,'O')
                h_z = to_host(d_z)
                z = alpha * A * x
                @test z ≈ h_z
                d_z = CUSPARSE.mv('N',d_A,d_x,'O')
                h_z = to_host(d_z)
                z = A * x
                @test z ≈ h_z
                d_z = d_A*d_x
                h_z = to_host(d_z)
                z = A * x
                @test z ≈ h_z
                @test_throws DimensionMismatch CUSPARSE.mv('T',alpha,d_A,d_x,beta,d_y,'O')
                @test_throws DimensionMismatch CUSPARSE.mv('N',alpha,d_A,d_y,beta,d_x,'O')
            end
        end
    end
end


text:
using KitBase, DataFrames
using KitBase.CSV
cd(@__DIR__)

###
# parameters
###
file = "data/pdfs.csv"
kn = [1e-3, 1e-3]
dx = [1e-2, 1e-2]

###
# workflow
###
f = open(file)
data = []
for line in eachline(f)
    a = split(line, ",")
    b = [parse(Float64, a[i]) for i = 1:length(a)]
    push!(data, b)
end
pdfs = data[3:end];

vs = VSpace1D(minimum(data[1]), maximum(data[1]), length(data[1]), data[1], data[1][2:end] .- data[1][1:end - 1], data[2])
println(vs.u)
#println(vs.weights)
ws = [moments_conserve(pdfs[i], vs.u, vs.weights) for i in axes(pdfs, 1)]
wa = zeros(9, 3)
for i in 1:9
    wa[i, :] .= ws[i]
end

w = [ws[3], ws[6]]
sw = [(ws[2] - ws[1]) / dx[1], (ws[5] - ws[4]) / dx[2]]

function f_ns(vs, w, sw, kn, γ = 3, K = 0, μ = ref_vhs_vis(kn, 1, 0.5), ω = 0.81)
    μ = ref_vhs_vis(kn, 1, 0.5)
    prim = conserve_prim(w, γ)
    Mu, Mxi, _, _1 = gauss_moments(prim, K)
    a = pdf_slope(prim, sw, K)
    swt = -prim[1] .* moments_conserve_slope(a, Mu, Mxi, 1)
    A = pdf_slope(prim, swt, K)
    τ = vhs_collision_time(prim, μ, ω)
    fr = chapman_enskog(vs.u, prim, a, A, τ)

    return fr, τ
end

fr1, τ1 = f_ns(vs, w[1], sw[1], kn[1])
fr2, τ2 = f_ns(vs, w[2], sw[2], kn[2])

###
# Write
###
df0 = DataFrame(density = wa[:, 1], velocity = wa[:, 2], energy = wa[:, 3])
CSV.write("data/w.csv", df0)

df1 = DataFrame(gradient1 = sw[1], gradient2 = sw[2])
CSV.write("data/dw.csv", df1)

df2 = DataFrame(fns1 = fr1, fns2 = fr2)
CSV.write("data/fns.csv", df2)

df3 = DataFrame(tau = [τ1, τ2], kn = kn, dx = dx)
CSV.write("data/paras.csv", df3)

max_stars_repo_name: CSMMLab/Flowmachine

max_issues_repo_name: CSMMLab/Flowmachine

max_forks_repo_name: CSMMLab/Flowmachine


text:
module auxiliaryFunctions

import LinearAlgebra: zeros, Matrix, svd, pinv, transpose, det, norm, I, rank
import HomotopyContinuation: solve, randn, differentiate, solutions, real_solutions, System, @polyvar
import DynamicPolynomials: Polynomial, Term
import Combinatorics: binomial, powerset, multiexponents
import Suppressor: @suppress

export affineVeronese,
	projVeronese,
	addNoise,
	calculateMeanDistanceToVariety,
	findEqListOfDegrees,
	comparisonOfMethods,
	makeCombinations,
	fillUpWithZeros,
	sampsonDistance,
	weightedGradientDescent,
	cleanUp

function addNoise(points, variance)
	points = [[entry+randn(Float64) for entry in point] for point in points]
	return(points)
end

#Take the filledzero list of degrees and return, if the polynomial generators are minimal.
function cleanUp(intermediateValues, listOfDegrees, varlength)
	totalexponents = vcat(map(i -> collect(multiexponents(varlength,-i)), -listOfDegrees[end][1]:0)...)
	array=[]
	indecesOfBasis = []
	for q in 1:length(listOfDegrees)
		currentsum = listOfDegrees[q][1]==listOfDegrees[1][1] ? 0 : sum([en[2] for en in listOfDegrees[1:q-1]])
		for i in 1:listOfDegrees[q][2]
			push!(array,intermediateValues[i+currentsum])
			push!(indecesOfBasis, length(array))
		end
		currentexponents = vcat(map(i -> collect(multiexponents(varlength,-i)), -listOfDegrees[q][1]:0)...)
		for val in 1:sum(totalexponents[1])-sum(currentexponents[1])
			indexarray = [[findfirst(t->t==adder+entry, totalexponents) for entry in currentexponents] for adder in collect(multiexponents(varlength,val)) ]
			for i in 1:listOfDegrees[q][2]
				for j in 1:length(indexarray)
					helper = [0. for _ in 1:length(totalexponents)]
					for k in 1:length(indexarray[j])
						helper[indexarray[j][k]] = intermediateValues[currentsum+i][length(totalexponents)-length(collect(multiexponents(varlength,listOfDegrees[q][1])))+k-1]
					end
					push!(array,helper)
				end
			end
		end
	end
	#Return true if the rank does not drop (ideal is minimal) and return the reduced ideal generators. Else return false
	if rank(hcat(array...))==length(array)
		redarray = reduce([array[ind] for ind in indecesOfBasis], array)
		return(true, [ar./norm(ar) for ar in redarray])
	else
		return(false, [])
	end
end

#Calculate f mod I for f not in I
function reduce(truebasis, array)
	for i in 1:length(truebasis)
		element = Base.copy(truebasis[i])
		fil = sort(filter(t->element!=t, array), rev=true)
		indeces = [findfirst(t->norm(t)>1e-15, el) for el in fil]

		#Gauss elimination algorithm
		for j in 1:length(fil)
			filind = findfirst(t->norm(t)>1e-15, fil[j])
			if norm(element[filind])>1e-15
				c1 = fil[j][filind]
				c2 = element[filind]
				element = element - fil[j]./c1.*c2
			end
			truebasis[i] = element
		end
	end
	return(truebasis)
end

function jacobianProd(veronese, var)
	d = length(var)
	dMatrix = Array{Polynomial, 2}(undef, length(veronese), length(veronese))
	dArray = [collect(differentiate(poly, var) for poly in veronese)][1]
	for i in 1:size(dMatrix)[1]
		for j in 1:size(dMatrix)[2]
			dMatrix[i,j] = Polynomial(differentiate(veronese[i], var)'*differentiate(veronese[j], var))
		end
	end
	return(dMatrix)
end

function vandermonde(n, d, array, proj=false)
	exponents = vcat(map(i -> collect(multiexponents(n,-i)), -d:0)...)
	if (proj == true)
		exponents = []
		for k in n:-1:0
			exp = collect(multiexponents(length(array[1])-1,k))
			[append!(entry,n-k) for entry in exp]
			append!(exponents,exp)
		end
	end
	output = Array{Float64,2}(undef, length(array), length(exponents))
	for i in 1:length(array)
		for j in 1:length(exponents)
			prod = 1.0
			for k in 1:length(exponents[j])
				if(exponents[j][k]!=0)
					prod = prod*((array[i][k])^(exponents[j][k]))
				end
			end
			output[i,j]=prod
		end
	end
	return(output)
end

function evaluationOfMatrix(gamma, Z, var)
	output = zeros(size(gamma)[1], size(gamma)[2])
	for iter in 1:length(Z)
		helper = Array{Float64, 2}(undef, size(gamma)[1], size(gamma)[2])
		for i in 1:size(output)[1]
			for j in 1:size(output)[2]
				helper[i,j] = (gamma[i,j])(var=>Z[iter])
			end
		end
		output = output + helper
	end
	return(output./length(Z))
end

function affineVeronese(n, var)
	exponents = vcat(map(i -> collect(multiexponents(length(var),-i)), -n:0)...)
	output = [prod(var.^exp) for exp in exponents]
	return(output)
end

function projVeronese(n,var)
	exponents = vcat(map(i -> collect(multiexponents(length(var)-1,-i)), -n:0)...)
	output = [prod(var[1:length(var)-1].^exp)*var[length(var)]^(n-sum(exp)) for exp in exponents]
	return(output)
end

function calculateMeanDistanceToVariety(points, equations, var)
	#HomotopyContinuation: https://www.juliahomotopycontinuation.org/examples/critical-points/
	@polyvar u[1:length(var)]
	if (length(var)-length(equations) == 1)
		d = [differentiate(equation, var) for equation in equations]
		matrix = Array{Polynomial, 2}(undef, length(var), length(var))
		matrix[1,:] = var-u
		for i in 1:length(d)
			matrix[i+1,:] = d[i]
		end
		systemArray = [det(matrix)]
		append!(systemArray, equations)
		systemArray = filter!(t->t!=0,systemArray)
		F_u = System(systemArray, variables = var, parameters = u)
	elseif (length(var)-length(equations) > 1)
		d = [differentiate(equation, var) for equation in equations]
		matrix = Array{Polynomial, 2}(undef, length(equations)+1, length(var))
		matrix[1,:] = var-u
		for i in 1:length(d)
			matrix[i+1,:] = d[i]
		end
		binomialsets = filter(p->length(p)==length(equations)+1, collect(powerset(1:length(var))))
		saverMatrix=Array{Polynomial,2}(undef,length(equations)+1,length(equations)+1)
		systemArray=Array{Polynomial,1}(undef,0)
		for entry in binomialsets
			for i in 1:length(entry)
				saverMatrix[:,i] = matrix[:, entry[i]]
			end
			push!(systemArray,det(saverMatrix))
		end
		append!(systemArray, equations)
		systemArray = filter!(t->t!=0,systemArray)
		F_u = System(systemArray, variables = var, parameters = u)
	else
		throw(error("The method is not yet supported for non-complete intersections!"))
	end
	p = randn(ComplexF64, length(points[1]))
	#@suppress begin
		try
			result_p = solve(F_u, target_parameters = p)
			realSolutions = solve(
								F_u,
								solutions(result_p);
								start_parameters =  p,
								target_parameters = points,
								transform_result = (r,u) -> minimum([norm(rel-u) for rel in solutions(r)])
			)
			realSolutions = filter(p -> p!=Inf, realSolutions)
			if !isempty(realSolutions)
				return(sum(realSolutions)/length(realSolutions))
			else
				return("dim>0")
			end
		catch e
			return("dim>0")
		end
	#end
end

#=
 Output is a good estimate for the starting vector of our Least Square Iteration
 =#
function comparisonOfMethods(n,points,numEq,tau)
	timer = round(Int64, time() * 1000)
	@polyvar var[1:length(points[1])]
	projPoints = [vcat(point,[1]) for point in points]
	@polyvar projVar[1:length(projPoints[1])]

	veroneseProj = projVeronese(n,projVar)
	veroProdProj = sum([projVeronese(n,point)*projVeronese(n,point)' for point in projPoints])/length(projPoints)
	jacoProdProj = evaluationOfMatrix(jacobianProd(veroneseProj,projVar), projPoints, projVar)
	svdSingular = svd(pinv(jacoProdProj)*veroProdProj)
	firstS = [entry / maximum(svdSingular.S) for entry in svdSingular.S]
	smallestS = firstS[length(firstS)-numEq]*tau
	numberOfSmallSingularValues = length(filter(p-> p<=smallestS, firstS))
	firstV = [svdSingular.V[:,i] for i in (length(veroneseProj)-numberOfSmallSingularValues+1):length(veroneseProj)]
	timer2 = round(Int64, time() * 1000)
	try
		Vandermonde = vandermonde(length(var),n,points,false)
		svdVander = svd(Vandermonde)
		secondS = [entry / maximum(svdVander.S) for entry in svdVander.S]
		secondV = svdVander.V[:,(length(veroneseProj)-numberOfSmallSingularValues)+1:length(veroneseProj)]
		secondV = [secondV[:,i] for i in 1:size(secondV)[2]]
		return(firstV, secondV)
	catch e
		return(firstV, firstV)
		prinln("Error caught",e)
	end
end

function weightedGradientDescent(points, n, var, curw0, nEq, maxIter, saverArray, zeroEntries)

	w0Matrix = Array{Float64,2}(undef,length(curw0[1]),length(curw0))
	for i in 1:length(curw0)
		w0Matrix[:,i] = curw0[i]
	end
	function zeroMatrix(w0Matrix)
		mat = zeros(Float64,length(curw0[1]),length(curw0))
		for zero in zeroEntries
			mat[zero[2],zero[1]] = w0Matrix[zero[2],zero[1]]
		end
		return(mat)
	end

	lossFct = w0Matrix -> sum([(projVeronese(n,vcat(points[j],1))'*w0Matrix)*saverArray[j]*w0Matrix'*projVeronese(n,vcat(points[j],1)) for j in 1:length(points)])/length(points)+100*sum([entry.^2 for entry in zeroMatrix(w0Matrix)])
	curLoss = lossFct(w0Matrix)
	dLossFct = w0Matrix -> 2*sum([(projVeronese(n,vcat(points[j],1))*projVeronese(n,vcat(points[j],1))')*w0Matrix*saverArray[j] for j in 1:length(points)])./length(points)+200*zeroMatrix(w0Matrix)
	i = 1
	while  i < maxIter && norm(dLossFct(w0Matrix)) > 1e-2
		w0Matrix, curLoss = backtracking_line_search(w0Matrix, dLossFct, lossFct)
		i=i+1
	end
	return([w0Matrix[:,i]./norm(w0Matrix[:,i]) for i in 1:size(w0Matrix)[2]], curLoss)
end

function backtracking_line_search(w0Matrix, dLossFct, lossFct; r=1e-3, s=0.7)
	dLoss = dLossFct(w0Matrix)./norm(dLossFct(w0Matrix))
	α = 0; t = 0.1; β = 10000
	while norm(β-α)>1e-15
		if lossFct(w0Matrix)-lossFct(w0Matrix - t * dLoss) < r*t*norm(dLoss)^2
			β = t
			t = (α+β)/2
		elseif s*norm(dLoss'*dLoss)<norm(dLoss'*dLossFct(w0Matrix - t*dLoss)./norm(dLossFct(w0Matrix - t*dLoss)))
			α = t
			t = (α+β)/2
		else
			break
		end
	end
	return(w0Matrix - t * dLoss, lossFct(w0Matrix - t * dLoss))
end

function sampsonDistance(points, nEq, n, var, startValues)
	@polyvar zed[1:length(points[1])+1]
	veronese = projVeronese(n, zed)
	Qstart = [start'*veronese for start in startValues]
	matrix = Array{Polynomial,2}(undef, nEq, length(points[1]))
	for i in 1:nEq, j in 1:nEq
		matrix[i,j] = differentiate(Qstart[i],zed)'*differentiate(Qstart[j],zed)
	end
	#TODO Whatever needs to be done here.
	helper = Array{Float64,2}(undef, nEq, nEq)
	saverArray = []

	for point in points
		for i in 1:nEq
			for j in 1:nEq
				helper[i,j] = matrix[i,j](zed=>vcat(point,1))
			end
		end
		append!(saverArray,[pinv(helper)])
	end
	return(saverArray)
end

function makeCombinations(values)
	saver = []
	for entry in values
		combis = collect(powerset(entry[2],entry[1],entry[1]))
		append!(saver, [combis])
	end
	return(cartesianUnion(saver))
end

function cartesianUnion(combinations)

	product = 1
	for comb in combinations
		product = product*length(comb)
	end
	output = Array{Any, 1}(undef, product)
	for i in 1:length(output)
		output[i]=[]
	end
	number = [1 for i in 1:length(combinations)]
	max = [length(comb) for comb in combinations]
	a = []
	j = 1
	while number[1]<=length(combinations[1])
		a = []
		for i in 1:length(combinations)
			append!(a,combinations[i][number[i]])
		end
		output[j] = [element for element in a]
		number[length(number)] = number[length(number)]+1
		number = checkNumber(number, max)
		j = j+1
	end
	return(output)
end

function checkNumber(number, max)
	for i in length(number):-1:2
		if(number[i]>max[i])
			number[i-1]=number[i-1]+1
			number[i]=1
		end
	end
	return(number)
end

function fillUpWithZeros(combination, n, numEq,d)
	zeroEntries = []

	for i in 1:length(combination)
		output = Array{Float64, 1}(undef, length(combination[end]))
		output[1:length(output)-length(combination[i])] .= 0
		output[length(output)-length(combination[i])+1:end] = combination[i]
		append!(zeroEntries,[[i,j] for j in 1:length(output)-length(combination[i])])
		combination[i] = [out for out in output]
	end
	return(combination, zeroEntries)
end

function findEqListOfDegrees(listOfDegrees)
	sort!(listOfDegrees)
	degreeList = []
	for x in listOfDegrees
		index = 0
		for entry in degreeList
			if entry[1]==x
				entry[2]=entry[2]+1
				index = 1
				break
			end
		end
		if(index==0)
			append!(degreeList, [[x,1]])
		end
	end
	return(	degreeList)
end

function findAllCompositions(amount, number)
	saverList = []
	for j in 1:number
		if(number%j==0)
			append!(saverList,[[j]])
		end
	end

	product = 1
	for k in (amount-1):-1:1
		helperList = []
		for entry in saverList
			append!(helperList, findDivisors(number/prod(entry), entry))
		end
		saverList = helperList
	end

	saverList = filter(p-> prod(p)==number,saverList)
	for list in saverList
		sort!(list)
	end
	saverSet = Set(saverList)
	saverList = collect(saverSet)
	return(saverList)
end

function findDivisors(number, currentList)
	if(number <= 0)
		return([])
	end
	saverList = []
	for k in 1:number
		if number%k==0
			helper = [list for list in currentList]
			append!(helper,[k])
			append!(saverList,[helper])
		end
	end
	return(saverList)
end

end


text:
using PyPlot
using Statistics
using LinearAlgebra

""" Function for results of Fig. A.1 """
function RPE_SPE_corr(;results_fit_path =  joinpath(dirname(pathof(SinergiafMRI_datafit)),"mcmc_rl_fit/projects/fmri/someresults/fits/"),
                        algo_rpe = "sarsa_lambda",
                        algo_spe = "forward_learner",
                        fitType = "pop",
                        subject_of_interest = 3,
                        )

    save_path = makehomesavepath("RPE_SPE_corr")
    mkdir(save_path)

    data_path = "./mcmc_rl_fit/projects/fmri/data/SARSPEICZVG_all_fMRI.csv"

    dataFmriSubj = 0
    rpe_pop = 0
    surprise_trial_idx = 0
    all_RPE_surprise = Array{Float64,1}()
    all_SPE_surprise = Array{Float64,1}()
    all_RPE = Array{Float64,1}()
    all_SPE = Array{Float64,1}()
    all_RPE_nonsurprise = Array{Float64,1}()
    all_SPE_nonsurprise = Array{Float64,1}()
    corrs = []
    corrs_non_surprise = []
    fig = 0
    axes = 0

    n_actions_persubject = []
    n_surprise_trials_persubj = []

    n_actions_subject = 0
    n_actions_total = 0
    n_surprise_trials_subj = 0

    fig, axes = subplots(nrows=1, ncols=3, figsize=(12,3))

    for i_subj in 1:21

        df_rpe = getSubjSignal(subjID = i_subj,
                        results_fit_path = results_fit_path,
                        signalname = "RPE",
                        fitType = fitType,
                        algoname = algo_rpe)

        df_spe = getSubjSignal(subjID = i_subj,
                        results_fit_path = results_fit_path,
                        signalname = "SPE",
                        # signalname = "Surprise",
                        fitType = fitType,
                        algoname = algo_spe)

        correlation = cor(df_rpe[:,1], df_spe[:,1])
        push!(corrs, correlation)

        (dataFmriSubj, stats) = RL_Fit.load_SARSPEI(data_path; subject_id = i_subj)

        surprise_trial_idx = findall(dataFmriSubj[:,8] .== 1) .-1
        # NOTE:
        # Shifting -1 is correct.
        # In the csv files the surprise indicator is placed on the landing state
        # But to calculate the signals we take the tuple (s,a,s'), ie the column of the csv file
        # corresponding to the current state. So the RPE calculated on the surprise trial is
        # placed at -1 the index of the surprise trial.

        all_RPE = vcat(all_RPE, df_rpe[:,1])
        all_SPE = vcat(all_SPE, df_spe[:,1])

        all_RPE_surprise= vcat(all_RPE_surprise, df_rpe[:,1][surprise_trial_idx])
        all_SPE_surprise = vcat(all_SPE_surprise, df_spe[:,1][surprise_trial_idx])

        n_actions_subject = length(df_rpe[:,1])
        n_surprise_trials_subj = length(surprise_trial_idx)

        push!(n_actions_persubject, n_actions_subject)
        push!(n_surprise_trials_persubj, n_surprise_trials_subj)

        non_surprise = setdiff(1:n_actions_subject, surprise_trial_idx)

        corr_non_surprise = cor(df_rpe[:,1][non_surprise], df_spe[:,1][non_surprise])

        push!(corrs_non_surprise, corr_non_surprise)

        all_RPE_nonsurprise= vcat(all_RPE_nonsurprise, df_rpe[:,1][non_surprise])
        all_SPE_nonsurprise = vcat(all_SPE_nonsurprise, df_spe[:,1][non_surprise])

        n_actions_total += length(df_rpe[:,1])
        if i_subj == subject_of_interest
            axes[1].scatter(df_rpe[:,1], df_spe[:,1])
            axes[1].set_xlabel("RPE")
            axes[1].set_ylabel("SPE")
            axes[1].set_title("subj: $fitType, Cor(RPE, SPE) = $(round(correlation,digits=3)) [$(round(corr_non_surprise,digits=3))], #S-trials:$n_surprise_trials_subj/$n_actions_subject")
            axes[1].scatter(df_rpe[:,1][surprise_trial_idx], df_spe[:,1][surprise_trial_idx], color="r")
            axes[1].grid()
            axes[1].set_xlim((-1, +1) )
            axes[1].set_ylim((0., +1) )
        end

    end
    fig.suptitle("SPE/RPE correlations. Param: $fitType. #S-trials: $(length(all_SPE_surprise))/$n_actions_total", y=1.)

    axes[2].hist(all_RPE_surprise, alpha=0.9, color="r", edgecolor="black", density=true)#,normed=true)
    axes[3].hist(all_SPE_surprise, alpha=0.9, color="r", edgecolor="black", density=true)#,normed=true)
    axes[2].set_title("Histogram of RPE at surprise trials")
    axes[3].set_title("Histogram of SPE at surprise trials")

    axes[3].hist(all_SPE_nonsurprise, color="blue", alpha=0.5, edgecolor="black", density=true)#,normed=true)

    fig.tight_layout()

    fig.savefig(joinpath(save_path, "Hist_RPE_SPE.pdf"))   # save the figure to file

    abs_corrs = abs.(corrs)
    abs_corrs_non_surprise = abs.(corrs_non_surprise)

    println("mean abs(corr): $(round(mean(abs_corrs),digits=3))")
    println("std corr: $(round(std(abs_corrs),digits=3))")
    println("minimum corr: $(round(minimum(corrs),digits=3))")
    println("maximum corr: $(round(maximum(corrs),digits=3))")

    println("mean abs(corr_non_surprise): $(round(mean(abs_corrs_non_surprise),digits=3))")
    println("std corr_non_surprise: $(round(std(abs_corrs_non_surprise),digits=3))")
    println("minimum corr_non_surprise: $(round(minimum(corrs_non_surprise),digits=3))")
    println("maximum corr_non_surprise: $(round(maximum(corrs_non_surprise),digits=3))")

    all_RPE_surprise, all_SPE_surprise, corrs, corrs_non_surprise,
    n_actions_persubject, n_surprise_trials_persubj
end

""" Function for results of Fig. A.2 """
function signals_corr(;
                        results_fit_path =  joinpath(dirname(pathof(SinergiafMRI_datafit)),"mcmc_rl_fit/projects/fmri/someresults/fits/"),
                         algoname1 = "hybrid_actorcritic_FWD",
                          algoname2 = "hybrid_actorcritic_FWD",
                        signalname1 = "RPE",
                        signalname2 = "SPE",
                        fitType = "pop",
                        subject_of_interest = 3)

    save_path = makehomesavepath("signals_corr")
    mkdir(save_path)

    data_path = "./mcmc_rl_fit/projects/fmri/data/SARSPEICZVG_all_fMRI.csv"

    dataFmriSubj = 0
    rpe_pop = 0
    surprise_trial_idx = 0
    all_signal1_surprise = Array{Float64,1}()
    all_signal2_surprise = Array{Float64,1}()
    all_signal1 = Array{Float64,1}()
    all_signal2 = Array{Float64,1}()
    all_signal1_nonsurprise = Array{Float64,1}()
    all_signal2_nonsurprise = Array{Float64,1}()

    all_signal1_surprise_DM = Array{Float64,1}()
    all_signal2_surprise_DM = Array{Float64,1}()
    all_signal1_nonsurprise_DM = Array{Float64,1}()
    all_signal2_nonsurprise_DM = Array{Float64,1}()

    corrs = []
    corrs_non_surprise = []
    fig = 0
    axes = 0

    n_actions_persubject = []
    n_surprise_trials_persubj = []

    n_actions_subject = 0
    n_actions_total = 0
    n_surprise_trials_subj = 0

    fig, axes = subplots(nrows=1, ncols=3, figsize=(9,3))

    for i_subj in 1:21

        df_signal1 = getSubjSignal(subjID = i_subj,
                        results_fit_path = results_fit_path,
                        signalname = signalname1,
                        fitType = fitType,
                        algoname = algoname1)

        df_signal2 = getSubjSignal(subjID = i_subj,
                        results_fit_path = results_fit_path,
                        signalname = signalname2,
                        # signalname = "Surprise",
                        fitType = fitType,
                        algoname = algoname2)

        # --- Signals the way they are
        signal1 = df_signal1[:,1]
        signal2 = df_signal2[:,1]

        # --- Mean center everyone
        signal1_DM = meancenter(df_signal1[:,1])
        signal2_DM = meancenter(df_signal2[:,1])

        correlation = cor(signal1_DM, signal2_DM)
        @show correlation
        push!(corrs, correlation)
        all_signal1 = vcat(all_signal1, signal1)
        all_signal2 = vcat(all_signal2, signal2)

        (dataFmriSubj, stats) = RL_Fit.load_SARSPEI(data_path; subject_id = i_subj)

        surprise_trial_idx = findall(dataFmriSubj[:,8] .== 1) .-1
        # NOTE:
        # Shifting -1 is correct.
        # In the csv files the surprise indicator is placed on the landing state
        # But to calculate the signals we take the tuple (s,a,s'), ie the column of the csv file
        # corresponding to the current state. So the RPE calculated on the surprise trial is
        # placed at -1 the index of the surprise trial.

        # --- Signals the way they are
        signal1surprise = signal1[surprise_trial_idx]
        signal2surprise = signal2[surprise_trial_idx]
        # --- Demeaned
        signal1surprise_DM = signal1_DM[surprise_trial_idx]
        signal2surprise_DM = signal2_DM[surprise_trial_idx]

        all_signal1_surprise= vcat(all_signal1_surprise, signal1surprise)
        all_signal2_surprise = vcat(all_signal2_surprise, signal2surprise)

        all_signal1_surprise_DM= vcat(all_signal1_surprise_DM, signal1surprise_DM)
        all_signal2_surprise_DM = vcat(all_signal2_surprise_DM, signal2surprise_DM)

        n_actions_subject = length(signal1)
        n_surprise_trials_subj = length(surprise_trial_idx)

        push!(n_actions_persubject, n_actions_subject)
        push!(n_surprise_trials_persubj, n_surprise_trials_subj)
        non_surprise = setdiff(1:n_actions_subject, surprise_trial_idx)

        # --- Signals the way they are
        signal1nonsurprise = signal1[non_surprise]
        signal2nonsurprise = signal2[non_surprise]
        # --- Demeaned
        signal1nonsurprise_DM = signal1_DM[non_surprise]
        signal2nonsurprise_DM = signal2_DM[non_surprise]

        correlation_non_surprise = cor(signal1nonsurprise_DM, signal2nonsurprise_DM)
        push!(corrs_non_surprise, correlation_non_surprise)

        all_signal1_nonsurprise= vcat(all_signal1_nonsurprise, signal1nonsurprise)
        all_signal2_nonsurprise = vcat(all_signal2_nonsurprise, signal2nonsurprise)

        all_signal1_nonsurprise_DM= vcat(all_signal1_nonsurprise_DM, signal1nonsurprise_DM)
        all_signal2_nonsurprise_DM = vcat(all_signal2_nonsurprise_DM, signal2nonsurprise_DM)

        n_actions_total += length(signal1)
        if i_subj == subject_of_interest
            axes[1].scatter(signal1, signal2)
            axes[1].set_xlabel(signalname1)
            axes[1].set_ylabel(signalname2)
            axes[1].set_title("subj: $fitType, Cor($signalname1, $signalname2) = $(round(correlation,digits=3)) [$(round(correlation_non_surprise,digits=3))], #S-trials:$n_surprise_trials_subj/$n_actions_subject")
            axes[1].scatter(signal1surprise, signal2surprise, color="r")
            axes[1].grid()
            axes[1].set_xlim((-1, +1) )
            axes[1].set_ylim((0., +1) )
        end

    end
    fig.suptitle("$signalname2/$signalname1 correlations. Param: $fitType. #S-trials: $(length(all_signal2_surprise))/$n_actions_total", y=1.)

    nbins = 15

    axes[2].hist(all_signal1_surprise, alpha=0.9, color="r", edgecolor="black", density=true) #, normed=true)
    axes[3].hist(all_signal2_surprise, alpha=0.9, color="r", edgecolor="black", density=true, bins=nbins)#, normed=true)

    axes[2].set_title("Histogram of $signalname1 at surprise trials")
    axes[3].set_title("Histogram of $signalname2 at surprise trials")

    axes[3].hist(all_signal2_nonsurprise, color="blue", alpha=0.5, edgecolor="black", density=true, bins=nbins)#, normed=true)

    fig.tight_layout()
    fig.savefig(joinpath(save_path, "Hist_$signalname1$signalname2.pdf"))   # save the figure to file

    abs_corrs = abs.(corrs)
    abs_corrs_non_surprise = abs.(corrs_non_surprise)

    println("mean corr: $(round(mean(corrs),digits=3))")
    println("std  corr: $(round(std(corrs),digits=3))")
    println("mean abs(corr): $(round(mean(abs_corrs),digits=3))")
    println("std  abs(corr): $(round(std(abs_corrs),digits=3))")
    println("minimum corr: $(round(minimum(corrs),digits=3))")
    println("maximum corr: $(round(maximum(corrs),digits=3))")

    println("mean corr_non_surprise: $(round(mean(corrs_non_surprise),digits=3))")
    println("std  corr_non_surprise: $(round(std(corrs_non_surprise),digits=3))")
    println("mean abs(corr_non_surprise): $(round(mean(abs_corrs_non_surprise),digits=3))")
    println("std  abs(corr_non_surprise): $(round(std(abs_corrs_non_surprise),digits=3))")
    println("minimum corr_non_surprise: $(round(minimum(corrs_non_surprise),digits=3))")
    println("maximum corr_non_surprise: $(round(maximum(corrs_non_surprise),digits=3))")

    all_signal1_surprise, all_signal2_surprise, corrs, corrs_non_surprise,
    n_actions_persubject, n_surprise_trials_persubj
end

function cosinesimilarity(a,b)
    sum(a.*b) / (sqrt(sum(a.^2)*sum(b.^2)))
end

function meancenter(a)
    a = a .- mean(a)
end


text:
using CUDA.CUSPARSE
using LinearAlgebra, SparseArrays

@testset "opnorm" for T in [Float32, Float64, ComplexF32, ComplexF64]
    S = sprand(T, 10, 10, 0.1)
    dS_csc = CuSparseMatrixCSC(S)
    dS_csr = CuSparseMatrixCSR(S)
    @test opnorm(S, Inf) ≈ opnorm(dS_csc, Inf)
    @test opnorm(S, Inf) ≈ opnorm(dS_csr, Inf)
end