text:
export Nonlinear, Unscented, Sampling, Extended

abstract type ApproximationMethod end
abstract type Unscented <: ApproximationMethod end
abstract type Sampling <: ApproximationMethod end
abstract type Extended <: ApproximationMethod end

"""
Description:

    Nonlinear node modeling a nonlinear relation. Updates for
    the nonlinear node are computed through the unscented transform (by default),
    importance sampling, or local linear approximation.

    For more details see "On Approximate Nonlinear Gaussian Message Passing on
    Factor Graphs", Petersen et al. 2018.

    f(out, in1) = δ(out - g(in1))

Interfaces:

    1. out
    2. in1

Construction:

    Nonlinear{T}(out, in1; g=g, id=:my_node)
    Nonlinear{T}(out, in1; g=g, g_inv=g_inv, id=:my_node)
    Nonlinear{T}(out, in1, in2, ...; g=g, id=:my_node)
    Nonlinear{T}(out, in1, in2, ...; g=g, g_inv=(g_inv_in1, g_inv_in2, ...), id=:my_node)
    Nonlinear{T}(out, in1, in2, ...; g=g, g_inv=(g_inv_in1, nothing, ...), id=:my_node)

    where T encodes the approximation method: Unscented, Sampling, or Extended.
"""
mutable struct Nonlinear{T<:ApproximationMethod} <: DeltaFactor
    id::Symbol
    interfaces::Array{Interface,1}
    i::Dict{Symbol, Interface}

    g::Function # Vector function that expresses the output as a function of the inputs
    g_inv::Union{Function, Nothing, Vector} # Inverse of g with respect to individual inbounds (optional)
    alpha::Union{Float64, Nothing} # Spread parameter for unscented transform
    dims::Union{Vector, Nothing} # Dimensions of messages on (all) interfaces
    n_samples::Union{Int64, Nothing} # Number of samples for sampling

    function Nonlinear{T}(id::Symbol,
                          g::Function,
                          g_inv::Union{Function, Nothing, Vector},
                          alpha::Union{Float64, Nothing},
                          dims::Union{Vector, Nothing},
                          n_samples::Union{Int64, Nothing},
                          out::Variable,
                          args::Vararg) where T<:ApproximationMethod

        n_args = length(args)
        for i=1:n_args
            @ensureVariables(args[i])
        end
        self = new(id, Vector{Interface}(undef, n_args+1), Dict{Symbol,Interface}(), g, g_inv, alpha, dims, n_samples)
        ForneyLab.addNode!(currentGraph(), self)
        self.i[:out] = self.interfaces[1] = associate!(Interface(self), out)
        for k = 1:n_args
            self.i[:in*k] = self.interfaces[k+1] = associate!(Interface(self), args[k])
        end

        return self
    end
end

function Nonlinear{Unscented}(out::Variable, args::Vararg; g::Function, g_inv=nothing, alpha=nothing, dims=nothing, id=ForneyLab.generateId(Nonlinear{Unscented}))
    return Nonlinear{Unscented}(id, g, g_inv, alpha, dims, nothing, out, args...)
end

function Nonlinear{Sampling}(out::Variable, args::Vararg; g::Function, dims=nothing, n_samples=nothing, id=ForneyLab.generateId(Nonlinear{Sampling}))
    return Nonlinear{Sampling}(id, g, nothing, nothing, dims, n_samples, out, args...)
end

function Nonlinear{Extended}(out::Variable, args::Vararg; g::Function, g_inv=nothing, dims=nothing, id=ForneyLab.generateId(Nonlinear{Extended}))
    return Nonlinear{Extended}(id, g, g_inv, nothing, dims, nothing, out, args...)
end

# A breaker message is required if interface is partnered with a Nonlinear{Sampling} inbound and there are multiple inbounds
function requiresBreaker(interface::Interface, partner_interface::Interface, partner_node::Nonlinear{Sampling})
    backward = (partner_interface != partner_node.i[:out]) # Interface is partnered with an inbound
    multi_in = isMultiIn(partner_node) # Boolean to indicate a nonlinear node with multiple stochastic inbounds

    return backward && multi_in
end

# A breaker message is required if interface is partnered with a Nonlinear{Unscented/Extended} inbound and no inverse is available
function requiresBreaker(interface::Interface, partner_interface::Interface, partner_node::Nonlinear{T}) where T<:Union{Unscented, Extended}
    backward = (partner_interface != partner_node.i[:out]) # Interface is partnered with an inbound
    multi_in = isMultiIn(partner_node) # Boolean to indicate a nonlinear node with multiple stochastic inbounds
    inx = findfirst(isequal(partner_interface), partner_node.interfaces) - 1 # Find number of inbound interface; 0 for outbound
    undefined_inverse = (partner_node.g_inv === nothing) || (multi_in && (inx > 0) && (partner_node.g_inv[inx] === nothing)) # (Inbound-specific) inverse is undefined

    return backward && undefined_inverse
end

function breakerParameters(interface::Interface, partner_interface::Interface, partner_node::Nonlinear)
    requiresBreaker(interface, partner_interface, partner_node) || error("Breaker dimensions requested for non-breaker interface: $(interface)")

    if partner_node.dims === nothing
        dims = ()
    else
        inx = findfirst(isequal(partner_interface), partner_node.interfaces) # Find interface index
        dims = partner_node.dims[inx] # Extract dimensionality from node.dims vector
    end

    return (Message{GaussianMeanVariance, variateType(dims)}, dims)
end

slug(::Type{Nonlinear{T}}) where T<:ApproximationMethod = "g{$(removePrefix(T))}"

"""
Determine whether there are multiple stochastic inbound edges
"""
function isMultiIn(node::Nonlinear)
    stochastic_in_count = 0
    for iface in node.interfaces[2:end]
        if !isDeltaConstrained(iface.partner)
            stochastic_in_count += 1
        end
    end

    return stochastic_in_count > 1
end


text:
function select_dates(obj::DatesSelector, arg0::AbsoluteDate, arg1::AbsoluteDate)
    return jcall(obj, "selectDates", List, (AbsoluteDate, AbsoluteDate), arg0, arg1)
end

ext: jl


text:
using DataFrames
using LinearAlgebra
using PhysicalConstants.CODATA2018
using Unitful

import Unitful:
    nm, μm, mm, cm, ns, μs, ms, ps, s

"""
    set_units(df::DataFrame)

Return a dataframe with the same structure as the input but with length and
time units set.
"""
function set_units(df::DataFrame)
    ## mm columns, assumes structure as 2021-10-05, should be generalised
    len_cols = [:r1, :r2, :r1x, :r2x, :x1, :x2, :xb1, :xb2, :xr1, :xr2,
        :xs, :xt1, :xt2, :y1, :y2, :yb1, :yb2, :yr1, :yr2, :ys, :yt1, :yt2, :z1,
        :z2, :zb1, :zb2, :zr1, :zr2, :zs, :zt1, :zt2, :rt1, :rt2]

    # ns cols:
    time_cols = [:t1, :t2, :ta1, :ta2, :tr1, :tr2]

    units_df = transform(df, len_cols .=> (x -> x * mm), renamecols=false)
    transform!(units_df, time_cols .=> (x -> x * ns), renamecols=false)
    return units_df
end

"""
	radial_correction(xb::Vector{<:Real}, yb::Vector{<:Real},zb::Vector{<:Real},
	                            r::Vector{<:Real})
Computes the x, y position of the interaction from a radius
of interaction and the x, y position from the SiPM barycentre
"""
function radial_correction(xb::Vector{<:Real}, yb::Vector{<:Real},
    zb::Vector{<:Real}, r::Vector{<:Real})
    ϕ = atan.(yb, xb)
    x = r .* cos.(ϕ)
    y = r .* sin.(ϕ)
    return x, y, zb
end

"""
    predict_interaction_radius
Given a calibration funcion and the physical limits of the Detector
and function return the radius of interaction (maxr - DOI).
"""
function predict_interaction_radius(cal_func::Function,
        minr::Real, maxr::Real, bias::Real=zero(Real))
  function prediction(sigmas::Vector{<:Real})
    rpred = cal_func.(sigmas) .+ bias
    rpred[rpred .< minr] .= minr
    rpred[rpred .> maxr] .= maxr
    return rpred
  end
end

"""
    time_of_flight(df::DataFrame, ri::Symbol, rsipm::Real, n::Real)

Calculate the time of flight between an interaction radius (ri) and the
detection radius (rsipm).
"""
function time_of_flight(df::DataFrame, ri::Symbol, rsipm::Real, n::Real=1.6)
    cc       = SpeedOfLightInVacuum / n
    path_len = rsipm * mm .- df[!, ri]
    return uconvert.(ps, path_len ./ cc)
end

"""
    time_of_flight(df::DataFrame, source::Vector{Symbol},
                    interaction::Vector{Symbol}, n::Real)

Time of flight between two points
"""
function time_of_flight(df::DataFrame, source_xyz::Vector{Symbol},
        int_xyz::Vector{Symbol}, n::Real=1.0)
    cc       = SpeedOfLightInVacuum / n
    pos_df   = combine(df, source_xyz => ByRow(vcat) => :spos, int_xyz => ByRow(vcat) => :ipos)
    path_len = combine(pos_df, [:ipos, :spos] => ByRow((x, y) -> norm(x .- y)) => :path)
    return uconvert.(ps, path_len[!, :path] ./ cc)
end

"""
    interaction_time(df::DataFrame, ri::Symbol, td::Symbol, rsipm::Real, n::Real)

Calculate the interaciton time given radius of interaction (ri) and detection time.
"""
function interaction_time(df::DataFrame, ri::Symbol, td::Symbol,
        rsipm::Real, n::Real=1.69)
    flight_time = time_of_flight(df, ri, rsipm, n)
    return uconvert.(ps, df[!, td]) .- flight_time
end

"""
    CRT(df::DataFrame, times::Vector{Symbol}, radii::Vector{Symbol}, det_rad::Real)

Calculate the difference between the calculated and expected interaction
time differences.
"""
#TODO: Protections for which symbols valid?
function CRT(df::DataFrame, times::Vector{Symbol},
        radii::Vector{Symbol}, det_rad::Real, n::Real=1.69)::Vector{<:Real}
    int1_flight = time_of_flight(df, [:xs, :ys, :zs], [:xt1, :yt1, :zt1])
    int2_flight = time_of_flight(df, [:xs, :ys, :zs], [:xt2, :yt2, :zt2])
    true_dt     = int2_flight ./ ps .- int1_flight ./ ps

    int1_det    = interaction_time(df, radii[1], times[1], det_rad, n)
    int2_det    = interaction_time(df, radii[2], times[2], det_rad, n)
    calc_dt     = int2_det ./ ps .- int1_det ./ ps

    return calc_dt - true_dt
end

max_stars_repo_name: jjgomezcadenas/ATools

max_issues_repo_name: jjgomezcadenas/ATools

max_forks_repo_name: jjgomezcadenas/ATools


text:
using SolidStateDetectors
using Unitful
using DataFrames
using ProgressMeter

T = Float32

function CreateSSDEvent(df::SubDataFrame, offset::CartesianPoint{T})
    pos = Matrix(df[:, [:X, :Y, :Z]])
    starting_positions = Vector{CartesianPoint{T}}()
    for i in 1:size(pos, 1)
        push!(starting_positions, CartesianPoint{T}(CartesianPoint{T}(pos[i, :]/1000) + offset))
    end
    energy_deps = convert(Array{T}, df[:, :E]) * u"keV"

    return Event(starting_positions, energy_deps)
end

function GetClusteredGeant4Events(gdf::GroupedDataFrame{DataFrame}, offset::CartesianPoint{T},
    radius = 0.2u"mm")::DetectorHitEvents
    events = Vector{Event}()
    for i in length(gdf)
        event = GenerateSSDEvent(gdf[i], offset)
        push!(events, event)
    end
    events_clustered = SolidStateDetectors.cluster_detector_hits(events, radius)

    return events_clustered
end

function DriftGeant4Events(gdf::GroupedDataFrame{DataFrame}, sim::Simulation, offset::CartesianPoint{T}; time_step = 0.1u"ns", stepLimiter::Integer=Inf)
    events = Vector{Event}()
    @info "Simulating events"
    @showprogress for i in 1:min(stepLimiter, length(gdf))
        event = CreateSSDEvent(gdf[i], offset)
        DriftEvent(event, sim, time_step)
        SimulateWaveform(event, sim, time_step)
        push!(events, event)
    end
    return events
end

function DriftEvent(event::Event, simulation::Simulation, time_step)
    drift_charges!(event, simulation, Δt = time_step)
end

function SimulateWaveform(event::Event, simulation::Simulation, time_step = 0.1u"ns")
    simulate!(event, simulation, Δt = time_step)
end

max_stars_repo_name: rjtayl/NDP

max_issues_repo_name: rjtayl/NDP

max_forks_repo_name: rjtayl/NDP


text:
import Tulip

using SparseArrays

@testset "LinearSolver" begin
    # TODO: Tulip constructor

    Aref = [
        1.0 0.0 1.1 1.2 1.3 1.4;
        0.0 1.0 2.1 2.2 2.3 2.4;
        1.0 0.0 3.1 3.2 3.3 3.4;
        0.0 0.0 1.0 0.0 1.0 0.0;
        0.0 0.0 0.0 1.0 0.0 1.0
    ]

    m0, n0, n, R = 3, 2, 4, 2
    M, N = m0 + R, n0 + n
    B0 = Aref[1:3, 1:2]
    B  = Aref[1:3, 3:end]
    blocks = [1, 2, 1, 2]

    A = UnitBlockAngularMatrix(B0, B, R, blocks)

    ls = UnitBlockAngularFactor(A)

    θ = ones(N)
    regP = 1e-2 .* ones(N)
    regD = 1e-2 .* ones(M)
    Tulip.KKT.update!(ls, θ, regP, regD)

    dx = zeros(N)
    dy = zeros(M)
    ξp = ones(M)
    ξd = ones(N)
    Tulip.KKT.solve!(dx, dy, ls, ξp, ξd)

    lsref = Tulip.KKT.Dense_SymPosDef(Aref)
    Tulip.KKT.update!(lsref, θ, regP, regD)
    dxref = zeros(N)
    dyref = zeros(M)
    Tulip.KKT.solve!(dxref, dyref, lsref, ξp, ξd)

    @test dx ≈ dxref
    @test dy ≈ dyref

end

@testset "Optimize" begin
    Aref = sparse([
        1.0 -1.0  0.0  0.0 1.1 1.2 1.3 1.4;
        0.0  0.0  1.0 -1.0 2.1 2.2 2.3 2.4;
        0.0  0.0  0.0  0.0 1.0 0.0 1.0 0.0;
        0.0  0.0  0.0  0.0 0.0 1.0 0.0 1.0
    ])
    m0, n0, n, R = 2, 4, 4, 2

    obj = [10000, 10000, 10000, 10000, 10, 20, 100, 200.0]
    lvar = zeros(8)
    uvar = fill(Inf, 8)

    lcon = ones(4)
    ucon = ones(4)
    connames = fill("", 4)
    varnames = fill("", 8)

    m = Tulip.Model{Float64}()

    m.params.OutputLevel = 1
    m.params.Presolve = 0

    m.params.MatrixOptions = Tulip.TLA.MatrixOptions(UnitBlockAngularMatrix; m0=m0, n0=n0, n=n, R=R)
    m.params.KKTOptions = Tulip.KKT.SolverOptions(UnitBlockAngularFactor)
    Tulip.load_problem!(m.pbdata, "Test", true, obj, 0.0, Aref, lcon, ucon, lvar, uvar, connames, varnames)
    Tulip.optimize!(m)
end


text:
type BoonModel <: DischargeModel
    M
    β
    H
    Q
end

BoonModel(N) = BoonModel(2,zeros(2),zeros(2,N),zeros(N))

function boon(h,β)
    β[1]*abs(h[1]).^β[2]*(h[1]-h[2])
end

function boon(H::Matrix{Float64},β)
    Q = zeros(size(H,2))
    for i in 1:length(Q)
        Q[i] = boon(H[:,i],β)
    end
    Q
end

function boon_gradient!(storage,h,β)
    storage[1] = abs(h[1]).^β[2]*(h[1]-h[2])
    storage[2] = β[1]*β[2]*abs(h[1]).^(β[2]-1)*(h[1]-h[2])
    storage
end

boon_gradient(h,β) = boon_gradient!(zeros(2),h,β)

function boon_gradient(H::Matrix{Float64},β)
    storage = zeros(2,size(H,2))
    for i in 1:size(H,2)
        storage[:,i] = boon_gradient(H[:,i],β)
    end
    storage
end

function ssr(H,Q,β)
    0.5*sum(abs2,Q-boon(H,β))
end

function gssr!(storage,H,Q,β)
    rs = Q-boon(H,β)
    gs = boon_gradient(H,β)

    storage[1] = dot(rs,vec(gs[1,:]))
    storage[2] = dot(rs,vec(gs[2,:]))
    storage
end

function StatsBase.fit(M::BoonModel,H::Matrix{Float64},Q::Vector{Float64})
    β = optimize(β->ssr(H,Q,β),M.β).minimizer
    BoonModel(2,β,H,Q)
end

function StatsBase.fit!(M::BoonModel,H::Matrix{Float64},Q::Vector{Float64})
    β = optimize(β->ssr(H,Q,β),M.β,ftol=1e-32).minimizer
    M.β = β
    M.H = H
    M.Q = Q
    M
end

#(β,g) -> gssr!(g,H,Q,β)
StatsBase.coef(M::BoonModel) = M.β
StatsBase.nobs(M::BoonModel) = length(M.Q)
StatsBase.model_response(M::BoonModel) = M.Q

function StatsBase.predict(M::BoonModel,H)
    Q = zeros(size(H,2))
    for i in 1:length(Q)
        Q[i] = boon(H[:,i],M.β)
    end
    Q
end
StatsBase.predict(M::BoonModel) = predict(M,M.H)

StatsBase.residuals(M::BoonModel) = M.Q-predict(M)
StatsBase.residuals(M::BoonModel,H,Q) = Q-predict(M,H)


text:
istrait(::Type{T}) where T = false
istrait(::Type{Trait}) where Trait <: AbstractTrait = true

isleaftrait(::Type{T}) where T = false
isleaftrait(::Type{Trait}) where Trait <: AbstractTrait = isempty(subtraits(Trait))
isleaftrait(::Type{AbstractTrait{T}}) where T = false
isleaftrait(::Type{NullTrait{T}}) where T = false

subtraits(::Type{T}) where T = DataType[]
subtraits(::Type{Trait}) where Trait <: AbstractTrait = Vector{UnionAll}(subtypes(Trait))

function hastrait(x,::Type{T}) where T
    info(T, " is not a trait")
    return false
end
hastrait(x::T,::Type{Trait}) where {T, Trait <: AbstractTrait} = traitdispatch(Trait{T}) != NullTrait{T}
hastrait(x::T,::Type{AbstractTrait{T}}) where T = false
hastrait(x::T,::Type{NullTrait{T}}) where T = NullTrait{T} ∈ traits(x)
hastraits(x,traits...) = all(map(trait -> hastrait(x,trait),[traits...]))
hastraits(x::T) where T = !hastrait(x,NullTrait{T})

function implementstrait(x,::Type{T}) where T
    info(T, " is not a trait")
    return false
end
implementstrait(::Type{T},::Type{Trait}) where {T, Trait <: AbstractTrait} = traitdispatch(Trait{T}) != NullTrait{T}
implementstrait(::Type{T},::Type{AbstractTrait{T}}) where T = false
implementstrait(::Type{T},::Type{NullTrait{T}}) where T = NullTrait{T} ∈ traits(x)
implementstraits(x,traits...) = all(map(trait -> hastrait(x,trait),[traits...]))
implementstraits(x::Type{T}) where T = !implementstrait(x,NullTrait{T})

notraits(x) = !hastraits(x)
notraits(x::Type{T}) where T = !implementstraits(x)

traits() = subtraits(AbstractTrait)
function traits(x::T) where T
    traitlist = _traits(x,AbstractTrait)
    if isempty(traitlist)
        return [NullTrait{T}]
    else
        return traitlist
    end
end

function _traits(x::T,::Type{Trait}) where {T, Trait <: AbstractTrait}
    traitlist = Vector{UnionAll}()
    if hastrait(x,Trait)
        push!(traitlist,Trait)
    end
    subtraits = Vector{UnionAll}(subtypes(Trait))
    if isempty(subtraits)
        return traitlist
    else
        for subtrait in subtraits
            append!(traitlist,_traits(x,subtrait))
        end
    end
    return traitlist
end


text:
"""
Сниппет для перевода dataframe в json

https://github.com/JuliaData/DataFrames.jl/issues/873
"""

using JSON, DataFrames, Query

function df2json(df::DataFrame)
  len = length(df[:, 1])
  indices = names(df)
  jsonarray = [Dict([string(index) => (isna(df[!, index][i]) ? nothing : df[!, index][i])
                                                    for index in indices])
                                                          for i in 1:len]
  return JSON.json(jsonarray)
end

function writejson(path::String, df::DataFrame)
  f = open(path,"w")
  write(f,df2json(df))
  close(f)
end

max_stars_repo_name: kuznetsovin/julia-first-app

max_issues_repo_name: kuznetsovin/julia-first-app

max_forks_repo_name: kuznetsovin/julia-first-app


text:
macro hessexpr(ex_f, xs)
    return esc(hessexpr(ex_f, xs))
end


text:
module GroupsCore

import Random
import Markdown

abstract type Group end
abstract type GroupElement end

export Group, GroupElement
export commutator, gens, hasgens, isfiniteorder, ngens, order
export istrivial
# export one!, inv!, mul!, conj!, commutator!, div_left!, div_right!

include("exceptions.jl")
include("groups.jl")
include("group_elements.jl")

include("extensions.jl")

end


text:
# Autogenerated wrapper script for Birch_Standard_jll for aarch64-linux-musl-cxx11
export libbirch_standard, libbirch_standard_debug

using boost_jll
using CompilerSupportLibraries_jll
using Eigen_jll
using LibBirch_jll
using LibYAML_jll
JLLWrappers.@generate_wrapper_header("Birch_Standard")
JLLWrappers.@declare_library_product(libbirch_standard, "libbirch-standard.so.0")
JLLWrappers.@declare_library_product(libbirch_standard_debug, "libbirch-standard-debug.so.0")
function __init__()
    JLLWrappers.@generate_init_header(boost_jll, CompilerSupportLibraries_jll, Eigen_jll, LibBirch_jll, LibYAML_jll)
    JLLWrappers.@init_library_product(
        libbirch_standard,
        "lib/libbirch-standard.so",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libbirch_standard_debug,
        "lib/libbirch-standard-debug.so",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@generate_init_footer()
end  # __init__()


text:
import RestrictedBoltzmannMachines as RBMs
using Test: @test, @testset
using LinearAlgebra: Diagonal, logdet

@testset "linalg" begin
    A = randn(2,2) + Diagonal([10, 10])
    B = randn(2,2)
    C = randn(2,2)
    D = randn(2,2) + Diagonal([10, 10])
    @test RBMs.block_matrix_logdet(A, B, C, D) ≈ logdet([A B; C D])
    @test RBMs.block_matrix_invert(A, B, C, D) ≈ inv([A B; C D])

    A = randn(3,3) + Diagonal([10, 10, 10])
    B = randn(3,2)
    C = randn(2,3)
    D = randn(2,2) + Diagonal([10, 10])
    @test RBMs.block_matrix_logdet(A, B, C, D) ≈ logdet([A B; C D])
    @test RBMs.block_matrix_invert(A, B, C, D) ≈ inv([A B; C D])

    A = randn(2,2) + Diagonal([10, 10])
    B = randn(2,3)
    C = randn(3,2)
    D = randn(3,3) + Diagonal([10, 10, 10])
    @test RBMs.block_matrix_logdet(A, B, C, D) ≈ logdet([A B; C D])
    @test RBMs.block_matrix_invert(A, B, C, D) ≈ inv([A B; C D])
end


text:
# Goal: Implement Conways Game of Life in Julia
# author: Jay Jagtap

#=
Rules for Game of Life

 Now Neighbors           Next
#   --- ---------           --------------
#   1   0,1             ->  0  # Lonely
#   1   4,5,6,7,8       ->  0  # Overcrowded
#   1   2,3             ->  1  # Lives
#   0   3               ->  1  # It takes three to give birth!
#   0   0,1,2,4,5,6,7,8 ->  0  # Barren

=#

function nextGen(A)
  #=
    Function returns next Generation of Life based on above mentioned Rules
    param A: 2D matrix of previous Generation
    return next_Gen: 2D matrix derived from A(previous Generation)
  =#

  # Initializing next_Gen with all Zero Values
  next_Gen = zeros(Int, size(A,1), size(A,2))

  for i in 1:size(next_Gen, 1)
    for j in 1:size(next_Gen, 2)
      alive_count = get_no_of_alive_neighbors(A, i, j)
      if A[i,j] == 0
        if alive_count == 3
            next_Gen[i,j] = 1
        end
      else
        if alive_count == 2 || alive_count == 3
            next_Gen[i,j] = 1
        end
      end
    end
  end
  return next_Gen
end

function get_no_of_alive_neighbors(A, i, j)
    #=
    param A: 2D matrix
    i: row index
    j: column index
    return alive_count: Number of 1's(alive cells) in neighborhood
    =#

    # Collecting Neighborhood co-ordinates
    neighbor_coord = []
      for k in 1:3
        if k < 2
          push!(neighbor_coord , (i-1,j-1))
          push!(neighbor_coord , (i-1,j))
          push!(neighbor_coord , (i-1,j+1))
        elseif k > 2
          push!(neighbor_coord , (i+1,j-1))
          push!(neighbor_coord , (i+1,j))
          push!(neighbor_coord , (i+1,j+1))
        else
          push!(neighbor_coord , (i,j-1))
          push!(neighbor_coord , (i,j+1))
        end
      end

    # Deriving alive count from neighbors
    alive = 0
    for coord in neighbor_coord
        if (coord[1] >= 1 && coord[1] <= size(A,1) ) && (coord[2] >= 1 && coord[2] <= size(A,2))
            if A[coord[1] , coord[2]] == 1
                alive=alive+1
            end
        end
    end
    return alive
end

function main()
  initial_Gen =
  [  0 0 0 1 1 1 0 0;
     0 1 0 1 0 1 1 1;
     0 1 0 1 0 1 0 1;
     0 1 0 1 0 1 1 1;
     0 1 0 1 0 1 0 1;
     0 1 0 1 0 1 0 1;
     1 1 1 1 1 0 0 0;
     1 1 1 1 1 1 1 1]

 no_of_iterations = 20
 println("Generation: initial_Gen")
 display(initial_Gen)
 println()
 for i in 1:no_of_iterations
    println("Generation: " , i)
    current_Gen = nextGen(initial_Gen)
    display(current_Gen)
    println()
    initial_Gen = current_Gen
 end

end

main()


text:
include("mountain_car_rtaa_planner.jl")
include("mountain_car_cmax_planner.jl")
include("mountain_car_planner_utils.jl")


text:
module IE

using Hecke

import Base.==

mutable struct FactoredIdeal
  exp::Vector{Int}
  idl::NfAbsOrdIdl
  norm::Int
  pred::Int
  function FactoredIdeal(n::Int)
    r = new()
    r.exp = zeros(Int, n)
    r.norm = 1
    return r
  end
end

function ==(a::FactoredIdeal, b::FactoredIdeal)
  return a.exp == b.exp
end

function Base.copy(a::FactoredIdeal)
  b = FactoredIdeal(length(a.exp))
  b.exp = copy(a.exp)
  b.norm = a.norm
  return b
end

mutable struct IdealEnum
  FB::Vector{NfAbsOrdIdl}
  cur::Vector{FactoredIdeal}
  lim::Vector{Int}
  mi::Int
  mi_idx::Vector{Int}

  function IdealEnum(FB::Vector{<: NfAbsOrdIdl})
    r = new()
    r.FB = FB
    r.cur = [FactoredIdeal(length(FB))]
    r.cur[1].idl = ideal(order(FB[1]), 1)
    r.lim = ones(Int, length(FB))
    r.mi = minimum(norm(x) for x = FB)
    r.mi_idx = find(x -> r.mi == norm(x), FB)
    return r
  end
end

function Base.show(io::IO, IE::IdealEnum)
  println(io, "enumerate ideals generated by ", length(IE.FB), " elements.")
end

function norm(I::FactoredIdeal)
  return I.norm
end

function Base.iterate(IE::IdealEnum, s::Int = 1)
  if s == length(IE.cur) || norm(IE.cur[s+1]) > IE.mi
    id = copy(IE.mi_idx)
    for j = id
      expand!(IE, j)
    end
    mi = minimum(norm(IE.FB[i])^IE.lim[i] for i=1:length(IE.FB))
    idx = find(i -> norm(IE.FB[i])^IE.lim[i] == mi, 1:length(IE.FB))
    IE.mi = mi
    IE.mi_idx = idx
  end
  return ideal(IE, s), s + 1
end

function expand!(IE::IdealEnum, i::Int)
  v = IE.lim[i]
  IE.lim[i] += 1
  N = []
  for ik = 1:length(IE.cur)
    k = IE.cur[ik]
    if k.exp[i] == 0
      kk = copy(k)
      kk.exp[i] = v
      kk.pred = ik
      kk.norm *= norm(IE.FB[i])^v
      push!(N, kk)
    end
  end
  i = 1
  j = 1
  while i <= length(IE.cur) && j <= length(N)
    if norm(IE.cur[i]) > norm(N[j])
      insert!(IE.cur, i, N[j])
      j += 1
    else
      i += 1
    end
  end
  while j <= length(N)
    push!(IE.cur, N[j])
    j += 1
  end
end

function Hecke.ideal(IE::IdealEnum, s::Int)
  if isdefined(IE.cur[s], :idl)
    return IE.cur[s].idl
  end
  I = ideal(IE, IE.cur[s].pred)
  i = findfirst(x -> IE.cur[s].exp[x] != IE.cur[IE.cur[s].pred].exp[x], 1:length(IE.FB))
  IE.cur[s].idl = I*IE.FB[i]^IE.cur[s].exp[i]
  return IE.cur[s].idl
end

end


text:
using Transducers

nums = [parse(Int, x) for x in readlines("1a.txt")]

increasing((x,y)) = x<y

nums |> Partition(3,1) |> Map(sum) |> Partition(2,1) |> Filter(increasing) |> collect |> size


text:
########################################################################
# Sum of atoms.
########################################################################

"""
    SumAtom(n, atoms::Vector{AbstractAtom})

Create an atom as the sum of individual atoms in a vector.
"""
mutable struct SumAtom{T1<:Int64, T2<:Tuple{Vararg{AbstractAtom}}} <: AbstractAtom
    n::T1
    atoms::T2
    function SumAtom(atoms::Tuple{Vararg{AbstractAtom}})
        n = checklength(atoms, "all atoms must be the same length")
        new{Int64, Tuple{Vararg{AbstractAtom}}}(n, atoms)
    end
end

Base.length(a::SumAtom) = a.n
Base.vec(a::SumAtom) = mapreduce(vec, hcat, a.atoms)
Base.:(*)(M::AbstractLinearOp, a::SumAtom) = mapreduce(a->M*a, hcat, a.atoms)
LinearAlgebra.mul!(Ma::AbstractMatrix{Float64}, M::LinearOp, a::SumAtom) = pmap(i->mul!(view(Ma,:,i), M, a.atoms[i]), collect(1:length(a.atoms)))

########################################################################
# Sum of atomic sets.
########################################################################

"""
    SumAtomicSet(A::Vector{AbstractAtomicSet})

Create an atomic set from the Minkowski sum of atomic sets.
"""
struct SumAtomicSet{T1<:Int64, T2<:Tuple{Vararg{AbstractAtomicSet}}} <: AbstractAtomicSet
    n::T1
    sets::T2
    function SumAtomicSet(sets::Tuple{Vararg{AbstractAtomicSet}})
        n = checklength(sets, "all atomic sets must be the same length")
        new{Int64, Tuple{Vararg{AbstractAtomicSet}}}(n, sets)
    end
end

"""
    sum(A, B)
    +(A, B)

Create the sum of the atomic sets A and B.
"""
function Base.:(+)(A::AbstractAtomicSet, B::AbstractAtomicSet)
    return SumAtomicSet(tuple(A, B))
end
function Base.:(+)(A::SumAtomicSet, B::AbstractAtomicSet)
    return SumAtomicSet(tuple(A.sets..., B))
end
function Base.:(+)(A::AbstractAtomicSet, B::SumAtomicSet)
    return SumAtomicSet(tuple(A, B.sets...))
end
function Base.:(+)(A::SumAtomicSet, B::SumAtomicSet)
    return SumAtomicSet(tuple(A.sets..., B.sets...))
end

"""
    support(A::SumAtomicSet, z)

Support function of a sum of atomic sets.
"""
function support(A::SumAtomicSet, z::Vector{Float64})
    return mapreduce(Ai->support(Ai, z), +, A.sets)
end

"""
    expose(A::SumAtomicSet, z)

Expose atoms in the sum of atomic sets. For two atomic sets
A and B, and a vector z,

    expose((A+B),z) = expose(A,z) + expose(B,z)

"""
function expose!(A::SumAtomicSet, z::Vector{Float64}, a::SumAtom)
    k = length(A.sets)
    pmap(i->expose!(A.sets[i], z, a.atoms[i]), collect(1:k))
    return nothing
end

function expose(A::SumAtomicSet, z::Vector{Float64})
    atoms = pmap(Ai->expose(Ai, z), A.sets)
    return SumAtom(atoms)
end

Base.length(A::SumAtomicSet) = A.n
atom_name(A::SumAtomicSet) = "Sum of atomic sets"
atom_description(A::SumAtomicSet) = "A₁ + A₂ + ⋯"
atom_parameters(A::SumAtomicSet) = "$(length(A.sets)) sets; n = $(A.n)"

########################################################################
# Face of sum of atomic sets.
########################################################################
mutable struct SumFace{T1<:Int64, T2<:Tuple{Vararg{AbstractFace}}} <: AbstractFace
    n::T1
    faces::T2
    function SumFace(faces::Tuple{Vararg{AbstractFace}})
        n = checklength(faces, "all faces must be the same length")
        return new{Int64, Tuple{Vararg{AbstractFace}}}(n, faces)
    end
end

"""
    face(A1 + ... + Ak, z)

Return collection of faces of the atomic sets `Ai`'s exposed by the vector `z`.
"""
function face!(A::SumAtomicSet, z::Vector{Float64}, F::SumFace)
    k = length(A.sets)
    pmap(i->face!(A.sets[i], z, F.faces[i]), collect(1:k))
    return nothing
end

function face(A::SumAtomicSet, z::Vector{Float64})
    faces = pmap(Ai->face(Ai, z), A.sets)
    return SumFace(faces)
end

"""
Given a face and a set of weights, reveal the corresponding
point on the face.
"""
Base.:(*)(F::SumFace, c::Vector) = vec(F)*c
Base.:(*)(M::AbstractLinearOp, F::SumFace) = mapreduce(Fi->M*Fi, hcat, F.faces)
Base.:(*)(λ::Real, F::SumFace) = mapreduce(Fi->λ*Fi, hcat, F.faces)
Base.length(F::SumFace) = F.n
rank(F::SumFace) = map(rank, F.faces)
vec(F::SumFace) = mapreduce(vec, hcat, F.faces)
face_name(F::SumFace) = "Face of the sum of atomic sets"
face_parameters(F::SumFace) = "rank = $(rank(F)); n = $(length(F))"


text:
using Pkg
using FredData, HTTP, DataFrames, JSON, Dates;
using Gadfly, Cairo, Compose, VegaLite, VegaDatasets
using CSV

cd(@__DIR__)

# Add Fred API key:
my_api_key = "-------------------------- USE YOUR FRED API KEY HERE --------------------------"
api_fred = Fred(my_api_key);

# Specific date for mapping
date_last = "2020-11-01"; date_last_txt = "2020 November";
date_baseline = "2020-01-01"; date_baseline_txt = "2020 January";

us10m = dataset("us-10m") # Guide map tool

# (1) GDP and Weekly economic index by Lewis-Mertens-Stock =====================================================
include("srcfigure-macro-output.jl");
# (2) Wage and labor force participation =======================================================================
include("figure-macro-labor-wagelabour.jl");
# (3) Unemployment and part time emp for econ reason ==========================================================
include("figure-macro-labor-unemp.jl");
# (4) Retail sales, excl. food services =======================================================================
include("figure-macro-retail.jl");
# (5) Maps - state level labor force particiaption and unemployment since 2020 January ========================
include("figure-macro-labor-maps.jl");
# (6) State-level employment - Texas & Nevada =================================================================
include("figure-macro-labor-state-industrytxnv.jl");
# (7) Census - Household Pulse Survey plotting calculation ====================================================
include("figure-micro-census-foodinseq.jl");


text:
"""
    CUDA.CUDNN

High level interface to cuDNN functions. See
https://github.com/JuliaGPU/CUDA.jl/blob/master/lib/state/README.md
for a design overview.
"""
module CUDNN

using ..APIUtils

using ..CUDA
using ..CUDA: CUstream, libraryPropertyType
using ..CUDA: libcudnn, @retry_reclaim, isdebug, initialize_context

using CEnum: @cenum

# core library
include("libcudnn_common.jl")
include("error.jl")
include("libcudnn.jl")
include("libcudnn_deprecated.jl")

# low-level wrappers
include("util.jl")
include("base.jl")
include("descriptors.jl")
include("tensor.jl")
include("inplace.jl")
include("optensor.jl")
include("reduce.jl")
include("convolution.jl")
include("pooling.jl")
include("activation.jl")
include("softmax.jl")
include("dropout.jl")
include("rnn.jl")
include("multiheadattn.jl")
include("normalization.jl")

function math_mode(mode=CUDA.math_mode())
    if mode == CUDA.PEDANTIC_MATH
        # don't use tensor cores.
        # on A100, only use them for TF32
        CUDNN_DEFAULT_MATH
    elseif mode == CUDA.DEFAULT_MATH
        # allow tensor core usage
        CUDNN_TENSOR_OP_MATH
    elseif mode == CUDA.FAST_MATH
        # also downcast inputs
        CUDNN_TENSOR_OP_MATH_ALLOW_CONVERSION
    end
end

# cache for created, but unused handles
const idle_handles = HandleCache{CuContext,cudnnHandle_t}()

function handle()
    cuda = CUDA.active_state()

    # every task maintains library state per device
    LibraryState = @NamedTuple{handle::cudnnHandle_t, stream::CuStream}
    states = get!(task_local_storage(), :CUDNN) do
        Dict{CuContext,LibraryState}()
    end::Dict{CuContext,LibraryState}

    # get library state
    @noinline function new_state(cuda)
        new_handle = pop!(idle_handles, cuda.context) do
            cudnnCreate()
        end

        finalizer(current_task()) do task
            push!(idle_handles, cuda.context, new_handle) do
                context!(cuda.context; skip_destroyed=true) do
                    cudnnDestroy(new_handle)
                end
            end
        end

        cudnnSetStream(new_handle, cuda.stream)

        (; handle=new_handle, cuda.stream)
    end
    state = get!(states, cuda.context) do
        new_state(cuda)
    end

    # update stream
    @noinline function update_stream(cuda, state)
        cudnnSetStream(state.handle, cuda.stream)
        (; state.handle, cuda.stream)
    end
    if state.stream != cuda.stream
        states[cuda.context] = state = update_stream(cuda, state)
    end

    return state.handle
end

## logging

const log_messages = []
const log_lock = ReentrantLock()
const log_cond = Ref{Any}()    # root

function log_message(sev, udata, dbg_ptr, ptr)
    dbg = unsafe_load(dbg_ptr)

    # find the length of the message, as denoted by two null terminators
    len = 0
    while true && len < 10000
        if unsafe_load(ptr, len+1) == 0 && unsafe_load(ptr, len+2) == 0
            break
        end
        len += 1
    end
    str = unsafe_string(ptr, len)   # XXX: can this yield?

    # print asynchronously
    @spinlock log_lock begin
        push!(log_messages, (; sev, dbg, str))
    end
    ccall(:uv_async_send, Cint, (Ptr{Cvoid},), udata)

    return
end

function _log_message(sev, dbg, str)
    lines = split(str, '\0')
    msg = join(lines, '\n')
    if sev == CUDNN_SEV_INFO
        @debug msg
    elseif sev == CUDNN_SEV_WARNING
        @warn msg
    elseif sev == CUDNN_SEV_ERROR
        @error msg
    elseif sev == CUDNN_SEV_FATAL
        error(msg)
    end
    return
end

function __runtime_init__()
    if version() < v"8.0"
        @warn "This version of CUDA.jl only supports CUDNN 8.0 or higher"
    end

    # register a log callback
    if (isdebug(:init, CUDNN) || Base.JLOptions().debug_level >= 2) &&
       version() >= v"8.2"  # NVIDIA bug #3256123
        log_cond[] = Base.AsyncCondition() do async_cond
            message =  @lock log_lock popfirst!(log_messages)
            _log_message(message...)
        end

        callback = @cfunction(log_message, Nothing,
                              (cudnnSeverity_t, Ptr{Cvoid}, Ptr{cudnnDebug_t}, Ptr{UInt8}))
        cudnnSetCallback(typemax(UInt32), log_cond[], callback)
    end
end

end


text:
module ROCKET

using Distributions
using CoordinateDescent, ProximalBase

export
  teInference

include("utils.jl")
include("inference.jl")

end


text:
module ChromNet

include("util.jl")
include("groupgm.jl")

end


text:
@testset "Cosine Similarity" begin
    for (point_a, point_b, similarity_expected) in [
        ([1, 0], [0, 1], 0),
        ([1, 0], [1, 1], √2 / 2),
        ([1, 0, 0], [1, 0, 0], 1),
        ([1, 0, 0], [0, 0, 1], 0),
        ([1, 0, 0], [1, 1, 1], √3 / 3),
    ]
        @test cosine_similarity(point_a, point_b) ≈ similarity_expected
    end
end

@testset "Angle between vectors" begin
    for (point_a, point_b, angle_expected) in [
        ([1, 0], [1, 0], 0),
        ([1, 0], [0, 1], π / 2),
        ([1, 0], [1, 1], π / 4),
        ([1, 0, 0], [1, 0, 1], π / 4),
    ]
        @test angle_between(point_a, point_b) ≈ angle_expected
    end
end

@testset "Distance between points" begin
    for (point_a, point_b, distance_expected) in [
        ([0], [0], 0),
        ([0], [1], 1),
        ([0], [-1], 1),
        ([0, 0], [0, 1], 1),
        ([0, 0], [1, 1], √2),
        ([0, 0, 0], [1, 1, 1], √3),
    ]
        @test distance(point_a, point_b) ≈ distance_expected
    end
end

@testset "Distance from point to line" begin
    for (point, line, distance_expected) in [
        ([0, 0], Line([0, 0], [1, 0]), 0),
        ([1, 0], Line([0, 0], [1, 0]), 0),
        ([1, 1], Line([0, 0], [1, 0]), 1),
        ([1, 1], Line([0, 0], [1, 1]), 0),
        ([1, 1], Line([0, 0], [-1, -1]), 0),
        ([8, 7], Line([0, 0], [1, 0]), 7),
        ([20, -3], Line([0, 0], [1, 0]), 3),
        ([20, -3, 1], Line([0, 0, 0], [1, 0, 0]), √10),
    ]
        @test distance(point, line) ≈ distance_expected
    end
end


text:
padded_axes(axs, it::InterpolationType) = (ax->padded_axis(ax, it)).(axs)
padded_axes(axs::NTuple{N,AbstractUnitRange}, it::NTuple{N,InterpolationType}) where N =
    padded_axis.(axs, it)

padded_similar(::Type{TC}, inds::Tuple{Vararg{Base.OneTo{Int}}}) where TC = Array{TC}(undef, length.(inds))
padded_similar(::Type{TC}, inds) where TC = OffsetArray{TC}(undef, inds)

# Narrow ax by the amount that axpad is larger
padinset(ax::AbstractUnitRange, axpad) = 2*first(ax)-first(axpad):2*last(ax)-last(axpad)
function padinset(ax::Base.OneTo, axpad::Base.OneTo)
    # We don't have any types that pad asymmetrically. Therefore if they both start at 1,
    # they must be the same
    @assert ax == axpad
    return ax
end

ct!(coefs, indscp, A, indsA) = copyto!(coefs, CartesianIndices(indscp), A, CartesianIndices(indsA))

copy_with_padding(A, it) = copy_with_padding(eltype(A), A, it)
function copy_with_padding(::Type{TC}, A, it::DimSpec{InterpolationType}) where {TC}
    indsA = axes(A)
    indspad = padded_axes(indsA, it)
    coefs = padded_similar(TC, indspad)
    if indspad == indsA
        coefs = copyto!(coefs, A)
    else
        fill!(coefs, zero(TC))
        ct!(coefs, indsA, A, indsA)
    end
    coefs
end

prefilter!(::Type{TWeights}, A::AbstractArray, ::BSpline{D}, ::GridType) where {TWeights,D<:Union{Constant,Linear}} = A

function prefilter(
    ::Type{TWeights}, ::Type{TC}, A::AbstractArray,
    it::Union{BSpline,Tuple{Vararg{Union{BSpline,NoInterp}}}}
    ) where {TWeights,TC}
    ret = copy_with_padding(TC, A, it)
    prefilter!(TWeights, ret, it)
end

function prefilter(
    ::Type{TWeights}, ::Type{TC}, A::AbstractArray,
    it::Union{BSpline,Tuple{Vararg{Union{BSpline,NoInterp}}}},
    λ::Real, k::Int
    ) where {TWeights,TC}
    ret = copy_with_padding(TC, A, it)
    prefilter!(TWeights, ret, it, λ, k)
end

function prefilter!(
    ::Type{TWeights}, ret::TCoefs, it::BSpline
    ) where {TWeights,TCoefs<:AbstractArray}
    sz = size(ret)
    first = true
    for dim in 1:ndims(ret)
        M, b = prefiltering_system(TWeights, eltype(TCoefs), sz[dim], degree(it))
        A_ldiv_B_md!(popwrapper(ret), M, popwrapper(ret), dim, b)
    end
    ret
end

function diffop(::Type{TWeights}, n::Int, k::Int) where {TWeights}
    D = spdiagm(0 => ones(TWeights,n))
    for i in 1:k
        D = diff(D; dims=1)
    end
    ## TODO: Normalize by n?
    D' * D
end
diffop(n::Int, k::Int) = diffop(Float64, n, k)
### TODO: add compiled constructor for most common operators of order k=1,2

function prefilter!(
    ::Type{TWeights}, ret::TCoefs, it::BSpline, λ::Real, k::Int
    ) where {TWeights,TCoefs<:AbstractArray}

    sz = size(ret)

    # Test if regularizing is allowed
    fallback = if ndims(ret) > 1
        @warn "Smooth BSpline only available for Vectors, fallback to non-smoothed"
        true
    elseif λ < 0
        @warn "Smooth BSpline require a non-negative λ, fallback to non-smoothed: $(λ)"
        true
    elseif λ == 0
        true
    else
        false
    end

    if fallback
        return prefilter!(TWeights, ret, it)
    end

    M, b = prefiltering_system(TWeights, eltype(TCoefs), sz[1], degree(it))
    ## Solve with regularization
    # Convert to dense Matrix because `*(Woodbury{Adjoint}, Woodbury)` is not defined
    Mt = Matrix(M)'
    Q = diffop(TWeights, sz[1], k)
    K = cholesky(Mt * Matrix(M) + λ * Q)
    B = Mt * popwrapper(ret)
    ldiv!(popwrapper(ret), K, B)

    ret
end

function prefilter!(
    ::Type{TWeights}, ret::TCoefs, its::Tuple{Vararg{Union{BSpline,NoInterp}}}
    ) where {TWeights,TCoefs<:AbstractArray}
    sz = size(ret)
    for dim in 1:ndims(ret)
        it = iextract(its, dim)
        if it != NoInterp
            M, b = prefiltering_system(TWeights, eltype(TCoefs), sz[dim], degree(it))
            if M != nothing
                A_ldiv_B_md!(popwrapper(ret), M, popwrapper(ret), dim, b)
            end
        end
    end
    ret
end

prefiltering_system(::Any, ::Any, ::Any, ::Any) = nothing, nothing

popwrapper(A) = A
popwrapper(A::OffsetArray) = A.parent

"""
    M, b = prefiltering_system{T,TC,GT<:GridType,D<:Degree}m(::T, ::Type{TC}, n::Int, ::Type{D}, ::Type{GT})

Given element types (`T`, `TC`) and interpolation scheme
(`GT`, `D`) as well the number of rows in the data input (`n`), compute the
system used to prefilter spline coefficients. Boundary conditions determine the
values on the first and last rows.

Some of these boundary conditions require that these rows have off-tridiagonal
elements (e.g the `[1,3]` element of the matrix). To maintain the efficiency of
solving tridiagonal systems, the [Woodbury matrix identity](https://en.wikipedia.org/wiki/Woodbury_matrix_identity)
is used to add additional elements off the main 3 diagonals.

The filtered coefficients are given by solving the equation system

    M * c = v + b

where `c` are the sought coefficients, and `v` are the data points.
"""
prefiltering_system

"""
    dl, d, du = inner_system_diags{T,IT}(::Type{T}, n::Int, ::Type{IT})

Helper function to generate the prefiltering equation system: generates the diagonals
for a `n`-by-`n` tridiagonal matrix with eltype `T` corresponding to the interpolation
type `IT`.

`dl`, `d`, and `du` are intended to be used e.g. as in `M = Tridiagonal(dl, d, du)`
"""
inner_system_diags


text:
"""
    exampleregistry()

Create a [`Registry`](#) of mathematical functions for testing purposes.

## Examples

{cell}
```julia
using FeatureRegistries
registry = FeatureRegistries.exampleregistry()
```
"""
function exampleregistry()

    registry = Registry(
        (;
            id = Field(String, name = "ID"),
            category = Field(
                Symbol,
                name = "Category",
                description = "The category of mathematical function",
                optional = true,
            ),
            instance = Field(
                Function,
                name = "Function instance",
                description = "The Julia function object"
            ),
        ),
        name = "Mathematical functions",
        loadfn = row -> row.instance,
        description= """
        Contains mathematical functions. `load`ing an entry will return the
        corresponding function.
        """
    )
    push!(registry, (id = "sin", instance = sin, category = :trigonometric))
    push!(registry, (id = "tan", instance = tan, category = :trigonometric))
    push!(registry, (id = "cos", instance = cos, category = :trigonometric))
    push!(registry, (id = "log2", instance = log2, category = :logarithmic))
    push!(registry, (id = "log", instance = log, category = :logarithmic))
    push!(registry, (id = "log1p", instance = log1p, category = :logarithmic))

    registry
end


text:
matrixdata =
"""
"random symmetric matrix"
function randsym(::Type{T}, n) where T
 A = zeros(T, n, n)
  for j = 1:n
      for i = j:n
          A[i,j] = randn()
      end
  end
  A = A + tril(A, -1)'
  return A
end
randsym(n) = randsym(Float64, n)
include_generator(FunctionName, "randsym", randsym)
include_generator(Group, :random, randsym)
include_generator(Group, :symmetric, randsym)
"""
open(joinpath(MatrixDepot.user_dir(), "generator.jl"), "w") do f
    write(f, matrixdata)
end

# include the just written user file
MatrixDepot.init(ignoredb=true)

n = rand(1:8)
@test matrixdepot("randsym", n) !== nothing
@test mdinfo("randsym") != nothing
@test "randsym" in MatrixDepot.mdlist(:random)
@test "randsym" in MatrixDepot.mdlist(:symmetric)

import MatrixDepot: include_generator, Group
@test_throws ArgumentError include_generator(Group, :lkjasj, sin)


text:
function read_libsvm(fname::ASCIIString, dimension = :col)
    label = Float64[]
    mI = Int64[]
    mJ = Int64[]
    mV = Float64[]
    fi = open(fname, "r")
    cnt = 1
    for line in eachline(fi)
        line = split(strip(line), " ")
        push!(label, parse(Float64, line[1]))
        line = line[2:end]
        for itm in line
            itm = split(itm, ":")
            push!(mI, parse(Int, itm[1]) + 1)
            push!(mJ, cnt)
            push!(mV, parse(Float64, itm[2]))
        end
        cnt += 1
    end
    close(fi)

    if dimension == :col
        (sparse(mI,mJ,mV), label)
    else
        (sparse(mJ,mI,mV), label)
    end
end


text:
"""
CTPG_train(dynamics_plant::Function, dynamics_controller::Function, cost_running::Function, cost_terminal::Function, cost_regularisor::Function, policy_NN, scenario;
solve_alg = Tsit5(), sense_alg = InterpolatingAdjoint(autojacvec = ZygoteVJP()), ensemble_alg = EnsembleThreads(), opt_1 = ADAM(0.01), opt_2 = LBFGS(), maxiters_1 = 100, maxiters_2 = 100, progress_plot = true, solve_kwargs...)

`CTPG_train()` provides a high-level interface for optimisation of the neural networks inside an ODE-represented dynamics based on Continuous-Time Policy Gradient (CTPG) methods that belong to the adjoint sensitivity analysis techniques. The code implemented and the default values for keyword arguments are specified considering training of a neural controller as the main application. In the context herein, a neural controller refers to a dynamic controller that incorporates neural-network-represented components at some points in its mathematical description.

The code utilises the functionalities provided by the [DiffEqFlux.jl](https://github.com/SciML/DiffEqFlux.jl) and [DiffEqSensitivity.jl](https://github.com/SciML/DiffEqSensitivity.jl) packages, and the Automatic Differentiation (AD) capabilities provided by the [Zygote.jl](https://github.com/FluxML/Zygote.jl) package that is integrated in DiffEqFlux.jl. `CTPG_train()` presumes the consistency of the functions provided as its input arguments with the AD tool, hence, the dynamics and cost functions should maintain their transparence against AD tools.

The optimisation (training) problem minimises the cost function defined over deterministic samples of the initial plant state `x₀` and the reference `r` by performing ensemble simulation based on parallelised computation.

The signals are defined as described below:

- `t`: time
- `x`: plant state
- `y`: plant output (= sensor output)
- `x_c`: controller state
- `u`: plant input (= controller output)
- `r`: exogenous reference
- `x_aug`: augmented forward dynamics state (= `[x; x_c; ∫cost_running]`)
- `p_NN`: neural network parameter

The arguments should be provided as explained below:

- `dynamics_plant`: Describes the dynamics of the plant to be controlled. Input arguments `x` and `u` should be of Vector type.
- `dynamics_controller`: Describes the dynamics of the controller that includes neural networks components. Input arguments `x_c`, `y`, `r`, and `p_NN` should be of Vector type.
- `dynamics_sensor`: Describes the dynamics of the sensor that measures output variables fed to the controller. Input arguments `x` should be of Vector type:
- `cost_running`: Describes the running cost defined as the integrand of the Lagrange-form continuous functional. Input arguments `x`, `y`, `u`, and `r` should be of Vector type.
- `cost_terminal`: Describes the terminal cost defined as the Mayer-form problem cost function. Defines a Bolza-form problem along with `cost_running`. Input arguments `x_f` and `r` should be of Vector type.
- `cost_regularisor`: Describes the regularisation term appended to the cost (loss) function. Input argument `p_NN` should be of Vector type.
- `policy_NN`: The neural networks entering into the controller dynamics. DiffEqFlux-based FastChain is recommended for its construction.
- `scenario`: Contains the parameters related with the ensemble-based training scenarios.
    - `ensemble`: A vector of the initial plant state `x₀` and the reference `r` constituting the trajectory realisations.
    - `t_span`: Time span for forward-pass integration
    - `t_save`: Array of time points to be saved while solving ODE. Typically defined as `t_save = t_span[1]:Δt_save:t_span[2]`
    - `dim_x`: `length(x)`
    - `dim_x_c`: `length(x_c)`

The keyword arguments should be provided as explained below:

- `solve_alg`: The algorithm used for solving ODEs. Default value is `Tsit5()`
- `sense_alg`: The algorithm used for adjoint sensitivity analysis. Default value is `InterpolatingAdjoint(autojacvec = ZygoteVJP())`, because the control problems usually render the `BacksolveAdjoint()` unstable. The vjp choice `autojacvec = ReverseDiffVJP(true)` is usually faster than `ZygoteVJP()`, when the ODE function does not have any branching inside. Please refer to the [DiffEqFlux documentation](https://diffeqflux.sciml.ai/dev/ControllingAdjoints/) for further details.
- `ensemble_alg`: The algorithm used for handling ensemble of ODEs. Default value is `EnsembleThreads()` for multi-threaded computation in CPU.
- `opt_1`: The algorithm used for the first phase of optimisation which rapidly delivers the parameter to a favourable region around a local minimum. Default value is `ADAM(0.01)`.
- `opt_2`: The algorithm used for the second phase of opitmisaiton. Defalut value is `LBFGS()` which refines the result of the first phase to find a more precise minimum. Please refer to the [DiffEqFlux documentation](https://diffeqflux.sciml.ai/dev/sciml_train/) for further details about two-phase composition of optimisers.
- `maxiters_1`: The maximum number of iterations allowed for the first phase of optimisation with `opt_1`. Defalut value is `100`.
- `maxiters_2`: The maximum number of iterations allowed for the second phase of optimisation with `opt_2`. Defalut value is `100`.
- `progress_plot`: The indicator to plot the state history for a nominal condition among the ensemble during the learning process. Default value is `true`.
- `i_nominal`: The index to select the case to plot using `progress_plot` during optimisation process from the `ensemble` defined in `scenario`. Defalut value is `nothing`.
- `p_NN_0`: Initial value of the NN parameters supplied by the user to bypass random initialisation of `p_NN` or to continue optimisation from the previous result. Defalut value is `nothing`.
- `solve_kwargs...`: Additional keyword arguments that are passed onto the ODE solver.

`CTPG_train()` returns the following outputs:

- `result`: The final result of parameter optimisation.
- `fwd_ensemble_sol`: The ensemble solution of forward simulation using the final neural network parameters.
- `loss_history`: The history of loss function evaluated at each iteration.
"""

function CTPG_train(dynamics_plant::Function, dynamics_controller::Function, dynamics_sensor::Function, cost_running::Function, cost_terminal::Function, cost_regularisor::Function, policy_NN, scenario; solve_alg = Tsit5(), sense_alg = InterpolatingAdjoint(autojacvec = ZygoteVJP()), ensemble_alg = EnsembleThreads(), opt_1 = ADAM(0.01), opt_2 = LBFGS(), maxiters_1 = 100, maxiters_2 = 100, progress_plot = true, i_nominal = nothing, p_NN_0 = nothing, solve_kwargs...)

    # scenario parameters
    @unpack ensemble, t_span, t_save, dim_x, dim_x_c = scenario
    dim_ensemble = length(ensemble)
    mean_factor  = Float32(1 / dim_ensemble)
    dim_t_save   = length(t_save)
    if isnothing(i_nominal)
        i_nominal = max(round(Int, dim_ensemble / 2), 1)
    end

    # NN parameters initialisation
    if isnothing(p_NN_0)
        p_NN_0 = initial_params(policy_NN)
    end

    # augmented dynamics
    function fwd_dynamics(r)
        return function (x_aug, p_NN, t)
            x   = x_aug[1:dim_x]
            x_c = x_aug[dim_x+1:end-1]
            # ∫cost_running = x_aug[end]

            y            = dynamics_sensor(t, x)
            (dx_c, u, _) = dynamics_controller(t, x_c, y, r, p_NN, policy_NN)
            dx           = dynamics_plant(t, x, u)

            return [dx; dx_c; cost_running(t, x, y, u, r)]
        end
    end

    # ODE problem construction
    prob_base = ODEProblem(fwd_dynamics(ensemble[i_nominal].r), [ensemble[i_nominal].x₀; zeros(Float32, dim_x_c + 1)], t_span, p_NN_0)

    function generate_probs(p_NN)
        return function (prob, i, repeat)
            remake(prob, f = fwd_dynamics(ensemble[i].r), u0 = [ensemble[i].x₀; zeros(Float32, dim_x_c + 1)], p = p_NN)
        end
    end

    if dim_ensemble == 1
        prob_mtk = modelingtoolkitize(prob_base)
        prob_base = ODEProblem(prob_mtk, [], t_span, jac = true)
        ensemble_alg = EnsembleSerial()
    end

    # loss function definition
    function loss(p_NN)
        ensemble_prob = EnsembleProblem(prob_base, prob_func = generate_probs(p_NN))

        fwd_ensemble_sol_full = solve(ensemble_prob, solve_alg, ensemble_alg, saveat = t_save, trajectories = dim_ensemble, sensealg = sense_alg; solve_kwargs...)

        # version 1: mean(sol[end])
        # sol_length = Float32.([max(1,length(fwd_ensemble_sol_full[i])) for i in 1:dim_ensemble])
        # sol_length_ratio = maximum(sol_length)./sol_length

        # loss_val = mean([fwd_ensemble_sol_full[i][end][end] + cost_terminal(fwd_ensemble_sol_full[i][end][1:dim_x], ensemble[i].r) for i in 1:dim_ensemble] .* sol_length_ratio) + cost_regularisor(p_NN)

        # version 2: mean(Array[end])
        # loss_val = mean([(Array(fwd_ensemble_sol_full[i])[end,end] + cost_terminal(Array(fwd_ensemble_sol_full[i])[1:dim_x,end], ensemble[i].r)) for i in 1:dim_ensemble] .* sol_length_ratio) + cost_regularisor(p_NN)

        # version 3: scalar operation (best for code robustness as it can handle divergent case in ensemble)
        loss_val = 0.0f0
        for i in 1:dim_ensemble
            fwd_sol = Array(fwd_ensemble_sol_full[i])

            if size(fwd_sol,2) > dim_t_save
                fwd_sol = fwd_sol[:,1:dim_t_save]
            end

            if size(fwd_sol,2) == dim_t_save
                x_aug_f = fwd_sol[:,end]
                loss_val += (x_aug_f[end] + cost_terminal(x_aug_f[1:dim_x], ensemble[i].r)) * mean_factor
            else
                loss_val += 1000.0f0
            end
        end
        loss_val += cost_regularisor(p_NN)

        return loss_val, fwd_ensemble_sol_full
    end

    # learning progress callback setup
    loss_history = fill(NaN32, maxiters_1 + maxiters_2 + 2)
    iterator_learning = 1
    cb_progress = function (p_NN_val, loss_val, fwd_ensemble_sol_full; plot_val = progress_plot)
        @show (loss_val, iterator_learning);
        loss_history[iterator_learning] = loss_val
        if plot_val
            fwd_ensemble_sol = Array(fwd_ensemble_sol_full[i_nominal])
            display(scatter(fwd_ensemble_sol[1:dim_x, :]', label = :false, plot_title = "System State: Learning Iteration $(iterator_learning)", layout = (dim_x, 1), size = (700, 200 * dim_x)))
        end
        iterator_learning += 1
        return false
    end

    # NN training
    if maxiters_1 <= 0
        result_coarse = (; u = p_NN_0)
    else
        result_coarse = DiffEqFlux.sciml_train(loss, p_NN_0, opt_1; cb = cb_progress, maxiters = maxiters_1)
    end

    if maxiters_2 <= 0
        result = result_coarse
    else
        result = DiffEqFlux.sciml_train(loss, result_coarse.u, opt_2; cb = cb_progress, maxiters = maxiters_2)
    end

    # Forward solution for optimised p_NN
    function eval_IO(r, p_NN)
        # Evaluation of input u, output y, and NN output y_NN at each t
        return function (t, x_aug)
            x = x_aug[1:dim_x]
            x_c = x_aug[dim_x+1:end-1]

            y = dynamics_sensor(t, x)
            (_, u, y_NN) = dynamics_controller(t, x_c, y, r, p_NN, policy_NN)

            return (; u = u, y = y, y_NN = y_NN)
        end
    end

    function generate_savedata(p_NN)
        return function (sol, i)
            IO = eval_IO(ensemble[i].r, p_NN).(sol.t, sol.u)
            return ((; sol = sol, u = [u for (u, y, y_NN) in IO], y = [y for (u, y, y_NN) in IO], y_NN = [y_NN for (u, y, y_NN) in IO]), false)
        end
    end

    fwd_ensemble_sol = solve(EnsembleProblem(prob_base, prob_func = generate_probs(result.u), output_func = generate_savedata(result.u)), solve_alg, ensemble_alg, saveat = t_save, trajectories = dim_ensemble, sensealg = sense_alg; solve_kwargs...)

    loss_history = filter(!isnan, loss_history)

    return result, fwd_ensemble_sol, loss_history
end


text:
module expmtdata
using ArgCheck

mutable struct ExperimentData{MT <: AbstractMatrix}
    YsRaw::Vector{MT}
    Ys::Vector{MT}
    Us::Vector{MT}
    ix_lkp
    function ExperimentData(YsRaw, Ys, Us, ix_lkp)
        MT = unique([typeof(x) for x in YsRaw])[1]
        YsRaw = [y[2:end,:] for y in YsRaw]  # note that Y, U take 2:end, 1:end-1.
        YsRaw = convert(Vector{MT}, YsRaw)
        Ys = convert(Vector{MT}, Ys)
        Us = convert(Vector{MT}, Us)
        new{MT}(YsRaw, Ys, Us, ix_lkp)
    end
end

"""
    get_data(s::ExperimentData, ix, splittype, tasktype)

Convenience utility for accessing data stored in an ExperimentData struct.
Specify the index of the target task, and then select from:

splittype:
* **:all**        - return the concatentation of all training/validation/test data.
* **:trainvalid** - return the concatentation of all training/validation data.
* **:split**      - return individual (3x) outputs for training/validation/test data.
* **:test**       - return only the test data
* **:train**      - return only the train data
* **:valid**      - return only the validation data.

tasktype:
* **:stl**   - single task model. Return train/validation/test data from this task's data.
* **:pool**  - pooled/cohort model. Here, training and validation data are from the
         complement of the selected index, returned in individual wrappers.

Note that in all cases, the output will be (a) Dict(s) containing the following fields:

* **:Y**    - the observation matrix (each column is an observation).
* **:U**    - the input matrix (each column is a datapoint).
* **:Yraw** - the raw data before standardisation and another manipulation. (Possibly legacy?)

Othe kwargs:

* `concat`  - By default, each boundary encountered between files will result in
a separate Dict, so the return values will be a vector of Dicts. However, for
more basic models (such as linear regression) with no assumption of temporal
continuity, it may be simpler to operate on a standard input and output data
matrix. Setting `concat = true` will return just a single Dict in an array with
all the data. Choosing `simplify=true` will further remove the array, returning
only Dicts.
* `stratified`  - (STL only) stratify the validation/test sets across files in
each style. By default, the test set will come at the end of the concatenation
of all files. Stratifying will mean there are L test sets from each of L files.
For the pooled dataset, the test set is partially stratified, that is, it is
stratified over the *types* (i.e. a % of each style), but not over the *files*
within the types. Given that our goal is MTL, this seems appropriate.
* `split`  - The train/validation/test split as a simplicial 3-dim vector.
* `simplify` - See `concat`. Used without `concat` this option does nothing.
"""
function get_data(s::ExperimentData, ix::Int, splittype::Symbol, tasktype::Symbol;
        concat::Bool=false, stratified=false, split=[0.7,0.1,0.2], simplify::Bool=false)
    @argcheck splittype ∈ [:all, :trainvalid, :split, :test, :train, :valid]
    @argcheck tasktype ∈ [:stl, :pool, :pooled]
    @assert !(stratified && splittype != :stl) "stratified only available for STL. Pooled is 'semi-stratified'(!)"
    if tasktype == :stl
        get_data_stl(s, ix, splittype; concat=concat, stratified=stratified,
            split=split, simplify=simplify)
    else
        splitpool = split[1:2]./sum(split[1:2])
        get_data_pooled(s, ix, splittype; concat=concat, split=splitpool, simplify=simplify)
    end
end

function get_data_stl(s::ExperimentData, ix::Int, splittype::Symbol;
        concat::Bool=false, stratified=false, split=[0.7,0.1,0.2], simplify=false)
    @argcheck splittype ∈ [:all, :trainvalid, :split, :test, :train, :valid]
    ixs = s.ix_lkp[ix]
    stratified = false

    # Get STL data (needed for everything)
    cYsraw = s.YsRaw[ixs]
    cYs    = s.Ys[ixs]
    cUs    = s.Us[ixs]

    if splittype == :all
        if concat
            cYsraw, cYs, cUs = _concat(cYsraw, cYs, cUs, simplify);
        end
        return _create_y_u_raw_dict(cYs, cUs, cYsraw)
    end

    Ns     = [size(y, 2) for y in cYs]
    !stratified && begin; cYsraw, cYs, cUs = _concat(cYsraw, cYs, cUs, false); end

    train, valid, test = create_data_split(cYs, cUs, cYsraw, split);
    if !stratified && !concat
        train, valid, test = _unconcatDicts(train, valid, test, Ns)
    elseif concat
        train, valid, test = _concatDicts(train, valid, test)
        if simplify
            rmv(x) = (length(x) == 1) ? x[1] : x
            train, valid, test = rmv(train), rmv(valid), rmv(test)
        end
    end

    if splittype == :split
        return train, valid, test
    elseif splittype == :test
        return test
    elseif splittype == :train
        return train
    elseif splittype == :valid
        return valid
    elseif splittype == :trainvalid
        return _concatDict([train, valid])
    else
        error("Unreachable error")
    end
end

function get_data_pooled(s::ExperimentData, ix::Int, splittype::Symbol;
        concat::Bool=false, split=[0.875, 0.125], simplify=false)
    @argcheck splittype ∈ [:all, :split, :test, :train, :valid]
    @argcheck length(split) == 2

    if splittype == :all
        L = length(s.ix_lkp)
        all = vcat([get_data_stl(s, i, :all; concat=concat) for i in 1:L]...)
        rmv = !simplify ? identity : (x-> (length(x) == 1) ? x[1] : x)
        return !concat ? all : _concatDicts(all)
    end

    # Get STL data (needed for everything)
    test = get_data_stl(s, ix, :all; concat=concat, stratified=false, simplify=simplify)

    if splittype == :test
        return test
    end

    train = Dict[]
    valid = Dict[]

    for i in setdiff(1:length(s.ix_lkp), ix)
        _train, _valid, _test = get_data_stl(s, i, :split; concat=concat,
            stratified=false, split=vcat(split, 0))
        train = vcat(train, _train)
        valid = vcat(valid, _valid)
    end

    if concat
        train, valid = _concatDicts(train), _concatDicts(valid)
        if simplify
            rmv(x) = (length(x) == 1) ? x[1] : x
            train, valid, test = rmv(train), rmv(valid), rmv(test)
        end
    end

    if splittype == :split
        return train, valid, test
    elseif splittype == :train
        return train
    elseif splittype == :valid
        return valid
    else
        error("Unreachable error")
    end

end

"""
    create_data_split(Y, U, Yraw, split=[0.7,0.1,0.2])

Create a partition of the data into train/validation/test components. The
default size of these partitions is 70%/10%/20% respectively. This can be
switched up by supplying the argument `split=[a, b, c]` s.t. sum(a,b,c) = 1.

It is assumed that *columns* of Y and U contain datapoints, and that *rows* of
Yraw contain datapoints. The output will be a Dict containing the following fields:

* **:Y**    - the observation matrix (each column is an observation).
* **:U**    - the input matrix (each column is a datapoint).
* **:Yrawv - the raw data before standardisation and another manipulation. (Possibly legacy?)
"""
function create_data_split(Y::Matrix{T}, U::Matrix{T}, Yraw::Matrix{T}, split=[0.7,0.1,0.2]) where T
    N = size(Y,2)
    @argcheck size(Y, 2) == size(U, 2)
    @argcheck sum(split) == 1
    int_split = zeros(Int, 3)
    int_split[1] = Int(round(N*split[1]))
    int_split[2] = Int(round(N*split[2]))
    int_split[3] = N - sum(int_split[1:2])
    split = vcat(0, cumsum(int_split), N) .+1
    train, valid, test = (Dict(:Y=>Y[:,split[i]:split[i+1]-1],
                               :U=>U[:,split[i]:split[i+1]-1],
                               :Yraw=>Yraw[split[i]:split[i+1]-1,:]) for i in 1:3)
    return train, valid, test
end

function create_data_split(Y::Vector{MT}, U::Vector{MT}, Yraw::Vector{MT},
    split=[0.7,0.1,0.2]) where MT <: AbstractMatrix
    @argcheck length(Y) == length(U) == length(Yraw)
    train, valid, test = Dict[], Dict[], Dict[]
    for i = 1:length(Y)
        _train, _valid, _test = create_data_split(Y[i], U[i], Yraw[i], split);
        push!(train, _train)
        push!(valid, _valid)
        push!(test, _test)
    end
    return train, valid, test
end

#= -----------------------------------------------------------------------
                Utilities for utilities for accessing data  :S
    This whole thing got way out of hand. I'm sure this can be tidied up =>
    is a function of setting out without a plan, under the assumption that
    "it wouldn't take long". You know what I'm talking about.
   ----------------------------------------------------------------------- =#

function _concat(Ysraw::Vector{MT}, Ys::Vector{MT}, Us::Vector{MT}, simplify::Bool) where MT <: AbstractMatrix
    v = simplify ? identity : x -> [x]
    return v(reduce(vcat, Ysraw)), v(reduce(hcat, Ys)), v(reduce(hcat, Us))
end

function _concat(trainsplit::Vector{D}, simplify::Bool) where D <: Dict
    v = simplify ? identity : x -> [x]
    Ysraw = reduce(vcat, [x[:Yraw] for x in trainsplit])
    Ys = reduce(hcat, [x[:Y] for x in trainsplit])
    Us = reduce(hcat, [x[:U] for x in trainsplit])
    return v(Ysraw), v(Ys), v(Us)
end

function _concatDicts(trainsplit::Vector{D}) where D <: Dict
    Ysraw, Ys, Us = _concat(trainsplit, true)
    return Dict(:Y=>Ys, :U=>Us, :Yraw=>Ysraw)
end

function _concatDicts(train::Vector{D}, valid::Vector{D}, test::Vector{D}) where D <: Dict
    train = _concatDicts(train)
    valid = _concatDicts(valid)
    test  = _concatDicts(test)
    return train, valid, test
end

function _unconcat(Yraw::AbstractMatrix{T}, Y::AbstractMatrix{T}, U::AbstractMatrix{T},
    breaks::Vector{I}) where {T <: Real, I <: Int}

    N = size(Y,2)
    @argcheck size(Y, 2) == size(U, 2) == size(Yraw, 1) == sum(breaks)
    nsplit = length(breaks)
    split = cumsum([0; breaks]) .+ 1
    Ys = [Y[:,split[i]:split[i+1]-1] for i in 1:nsplit]
    Us = [U[:,split[i]:split[i+1]-1] for i in 1:nsplit]
    Ysraw = [Yraw[split[i]:split[i+1]-1, :] for i in 1:nsplit]
    return Ysraw, Ys, Us
end

function _unconcatDict(d::Dict, breaks::Vector{I}) where I <: Int
    Ysraw, Ys, Us = _unconcat(d[:Yraw], d[:Y], d[:U], breaks)
    return [Dict(:Y=>y, :U=>u, :Yraw=>yraw) for (y,u,yraw) in zip(Ys, Us, Ysraw)]
end

# TODO: DRY fail here.

function _unconcatDicts(ds::Vector{D}, breaks::Vector{I}) where {D <: Dict, I <: Int}
    N = length(ds)
    Ls = [size(d[:Yraw], 1) for d in ds]
    @assert sum(Ls) == sum(breaks)
    bs  = copy(breaks)
    cbs = cumsum(breaks)
    out = Dict[]
    for nn in 1:N
        ix = something(findlast(cbs .<= Ls[nn]), 0)
        (ix > 0 && ix < length(bs) && cbs[ix] != Ls[nn]) && (ix += 1)
        ix = max(ix, 1)
        cbreaks = bs[1:ix]
        cbreaks[end] = Ls[nn] - (ix > 1 ? cbs[ix-1] : 0)
        bs = bs[ix:end]
        bs[1] -= cbreaks[end]
        cbs = cbs[ix:end] .- Ls[nn]

        out = vcat(out, _unconcatDict(ds[nn], cbreaks))
    end
    return out
end

function _unconcatDicts(train::Vector{D}, valid::Vector{D}, test::Vector{D},
    breaks::Vector{I}) where {D <: Dict, I <: Int}
    Ls = [sum([size(d[:Y], 2) for d in data]) for data in [train, valid, test]]
    bs  = copy(breaks)
    cbs = cumsum(breaks)
    out = []
    for (nn, data) in enumerate([train, valid, test])
        ix = something(findlast(cbs .<= Ls[nn]), 0)
        (ix > 0 && ix < length(bs) && cbs[ix] != Ls[nn]) && (ix += 1)
        ix = max(ix, 1)
        cbreaks = bs[1:ix]
        cbreaks[end] = Ls[nn] - (ix > 1 ? cbs[ix-1] : 0)
        bs = bs[ix:end]
        bs[1] -= cbreaks[end]
        cbs = cbs[ix:end] .- Ls[nn]
        push!(out, _unconcatDicts(data, cbreaks))
    end
    return out[1], out[2], out[3]
end

function _create_y_u_raw_dict(Ys::Vector{MT}, Us::Vector{MT}, Ysraw::Vector{MT}
    ) where MT <: AbstractMatrix
    [Dict(:Y=>y, :U=>u, :Yraw=>yraw) for (y,u,yraw) in zip(Ys, Us, Ysraw)]
end

function _create_y_u_raw_dict(Ys::AbstractMatrix, Us::AbstractMatrix, Ysraw::AbstractMatrix)
    [Dict(:Y=>Ys, :U=>Us, :Yraw=>Ysraw)]
end

function restrict_expmtdata(expmtdata, N)
    # restrict training set size to N chunks of 64 per style.
    Ys, test_Ys = [], []
    Ysraw, test_Ysraw = [], []
    Us, test_Us = [], []
    num_styles = length(expmtdata.ix_lkp)
    lkp, test_lkp = [[] for i in 1:num_styles], [[] for i in 1:num_styles]

    for s in 1:num_styles
        Nrem = N
        for i in expmtdata.ix_lkp[s]
            y, yraw, u = expmtdata.Ys[i], expmtdata.YsRaw[i], expmtdata.Us[i]
            n = size(y, 2) ÷ 64
            ncur = min(n, Nrem)
            if ncur == 0
                push!(test_Ys, y)
                push!(test_Ysraw, yraw)
                push!(test_Us, u)
                push!(test_lkp[s], i)
            elseif ncur == n
                push!(Ys, y)
                push!(Ysraw, yraw)
                push!(Us, u)
                push!(lkp[s], i)
            else
                push!(Ys, y[:, 1:64*ncur])
                push!(Ysraw, yraw[1:64*ncur, :])
                push!(Us, u[:, 1:64*ncur])
                push!(test_Ys, y[:, 64*ncur+1:end])
                push!(test_Ysraw, yraw[64*ncur+1:end, :])
                push!(test_Us, u[:, 64*ncur+1:end])
                push!(test_lkp[s], i)
                push!(lkp[s], i)
            end
            Nrem -= ncur
        end
    end

    return (ExperimentData(Ysraw, Ys, Us, lkp),
        ExperimentData(test_Ysraw, test_Ys, test_Us, test_lkp))
end
end

max_stars_repo_name: ornithos/pytorch-mtds-mocap

max_issues_repo_name: ornithos/pytorch-mtds-mocap

max_forks_repo_name: ornithos/pytorch-mtds-mocap


text:
using LabelledArrays
using OrdinaryDiffEq
import OrdinaryDiffEq: solve
using ModelingToolkit
using Petri
import Petri: stripnullterms, fluxes, odefunc, mk_function
using Test
using Plots

N(x) = sum(x)
function solver1(f, u0, p, t)
    prob = ODEProblem(f,u0,(0.0,365.0),p)
    soln = solve(prob,Tsit5())
    return prob, soln
end

# the following approach works
@variables S,E,I,R, β,γ,μ
sir = Petri.Model([S,I,R],[(S+I, 2I), (I,R)])

u0 = @LArray [100.0, 1, 0] (:S, :I, :R)
p = @LArray [0.35, 0.05] (:μ, :β)
t = (0, 365.0)

@info "SIR Solving"
fex = odefunc(sir, :sir)
f = eval(fex)
@time prob, soln = solver1(f, u0, p, t)
@time prob, soln = solver1(f, u0, p, t)

@info "successfully ran global eval passed to function"

function makesolve(m, name, u0, p, t)
    u0 = @LArray [100.0, 1, 0] (:S, :I, :R)
    p = @LArray [0.35, 0.05] (:μ, :β)
    fex = odefunc(m, name)
    f = eval(fex)
    prob, soln = Base.invokelatest(solver1, f, u0, p, (0, 365.0))
end

@time prob2, soln2 = makesolve(sir, :sir, u0, p, t)
@time prob2, soln2 = makesolve(sir, :sir, u0, p, t)
@info "makesolve approach worked but was slow."

norm(x)= sqrt(sum(map(y->y*y, x)))

@show  norm(soln(360.0) - soln2(360.0))

@info "Running GG version"

function ggsolve(m, name, u0, p, t)
    u0 = @LArray [100.0, 1, 0] (:S, :I, :R)
    p = @LArray [0.35, 0.05] (:μ, :β)
    fex = odefunc(m, name)
    f = mk_function(m)
    du = similar(u0)
    f(du, u0, p, 0)
    prob = ODEProblem(f, u0, t, p)
    soln = solve(prob, Tsit5())
    return prob, soln
end

@time prob3, soln3 = ggsolve(sir, :sir, u0, p, t)
@time prob3, soln3 = ggsolve(sir, :sir, u0, p, t)
@info "gg approach worked and is fast."

@show  norm(soln(360.0) - soln3(360.0))


text:
using FITSIO.Libcfitsio

function writehealpix(filename, pixels, nside, ordering, coordsys)
    if eltype(pixels) == Float32
        tform = "1E"
    elseif eltype(pixels) == Float64
        tform = "1D"
    end

    file = fits_create_file("!"*filename)
    try
        fits_create_img(file, Int16, Int[])
        fits_write_date(file)
        fits_movabs_hdu(file, 1)
        fits_create_binary_tbl(file, length(pixels), [("SIGNAL", tform, "")], "BINTABLE")
        fits_write_key(file, "PIXTYPE",  "HEALPIX", "HEALPIX pixelization")
        fits_write_key(file, "ORDERING", ordering,  "Pixel ordering scheme (either RING or NESTED)")
        fits_write_key(file, "NSIDE",    nside,     "Resolution parameter for HEALPIX")
        fits_write_key(file, "COORDSYS", coordsys,  "Pixelization coordinate system")
        fits_write_comment(file, "G = galactic, E = ecliptic, C = celestial = equatorial")
        fits_write_col(file, 1, 1, 1, pixels)
    finally
        fits_close_file(file)
    end
end

function readhealpix(filename)
    file = fits_open_file(filename)
    try
        hdutype = fits_movabs_hdu(file, 2)
        tform, tform_comment = fits_read_key_str(file, "TFORM1")
        if tform == "1E"
            T = Float32
        elseif tform == "1D"
            T = Float64
        end

        naxes, naxis_comment = fits_read_key_lng(file, "NAXIS")
        naxis, nfound = fits_read_keys_lng(file, "NAXIS", 1, naxes)
        nside, nside_comment = fits_read_key_lng(file, "NSIDE")
        npix = 12*nside*nside

        ordering, ordering_comment = fits_read_key_str(file, "ORDERING")
        coordsys, coordsys_comment = fits_read_key_str(file, "COORDSYS")

        pixels = zeros(T, npix)
        fits_read_col(file, 1, 1, 1, pixels)

        return pixels, nside, ordering, coordsys

    finally
        fits_close_file(file)
    end
end

# test reading/writing Healpix maps as FITS binary tables using the Libcfitsio interface
for T in (Float32, Float64)
    nside = 4
    npix = 12*nside*nside
    pixels = rand(T, npix)
    ordering = "NESTED"
    coordsys = "G"

    filename = tempname() * ".fits"
    try
        writehealpix(filename, pixels, nside, ordering, coordsys)
        @test readhealpix(filename) == (pixels, nside, ordering, coordsys)
    finally
        ispath(filename) && rm(filename)
    end
end


text:
using StatsBase

function perturbation(f, M, a, b; λ::Real=1.0, maxitter::Int=10, ϵ::Real=1e-3)
        feval = (a, b) -> f(optimaltransport(M, a, b, λ=λ, maxitter=maxitter, ϵ=ϵ))
        ∇af, ∇bf = Zygote.gradient(feval, a, b)
        return ∇af .- mean(∇af), ∇bf .- mean(∇bf)
end

function pert_utility(M, a, b; λ::Real=1.0, maxitter::Int=10, ϵ::Real=1e-3)
   return perturbation(Q -> utility(Q, M), M, a, b, λ=λ, maxitter=maxitter, ϵ=ϵ)
end

function pert_entropy(M, a, b; λ::Real=1.0, maxitter::Int=10, ϵ::Real=1e-3)
   return perturbation(Q -> entropy(Q), M, a, b, λ=λ, maxitter=maxitter, ϵ=ϵ)
end

max_stars_repo_name: PoisotLab/OTSIN

max_issues_repo_name: PoisotLab/OTSIN

max_forks_repo_name: PoisotLab/OTSIN


text:
### A Pluto.jl notebook ###
# v0.16.1

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : missing
        el
    end
end

# ╔═╡ 8a2604ac-1f9b-11ec-30a2-0998ea216cac
begin
    import Pkg
    Pkg.activate("/tmp/timeseriespkg/")
    Pkg.add([
        Pkg.PackageSpec(name="FluxArchitectures"),
        Pkg.PackageSpec(name="Plots", version="1"),
        Pkg.PackageSpec(name="DataFrames", version="1"),
        Pkg.PackageSpec(name="CSV", version="0.9"),
    ])
	Pkg.add(["PlutoUI"])
    using FluxArchitectures, Random, Plots, DataFrames, Dates, CSV, PlutoUI
end

# ╔═╡ e9e23536-7553-4d61-bfa9-e3a76c6d6737
begin
	df = CSV.read("/home/kris/Downloads/opendata_covid19_tests_total.csv.1", DataFrame, )
	dropmissing!(df)
end

# ╔═╡ aeedcfe8-1e98-49af-8ea1-55864638f6bf
cases = df[:, [:DailyCases, :TotalCases, :PerPopulation]] |> Matrix{Float32}

# ╔═╡ 3085780c-1977-4683-80e9-c1151b2d804a
begin
	poollength = 21; horizon = 14; datalength = 550;
	#input, target = get_data(:traffic, poollength, datalength, horizon)
	input, target = prepare_data(cases, poollength, datalength, horizon)
end

# ╔═╡ 2c9e7f3a-a913-4c3e-a70b-4e5b676bdcfa
input;

# ╔═╡ be7755f1-cfb2-4d42-940c-fe522390ad87
begin
	inputsize = size(input, 1)
	convlayersize = 4
	recurlayersize = 6
	skiplength = 120

	model = LSTnet(inputsize, convlayersize, recurlayersize, poollength, skiplength,
        init=Flux.zeros32, initW=Flux.zeros32)
end

# ╔═╡ bc612d5c-9092-489e-ab79-f62f0a814ae1
loss(x, y) = Flux.mse(model(x), y')

# ╔═╡ 7d37133c-b84f-4087-ae2b-3af5203cd028
cbb = function ()
    Flux.reset!(model)
    pred = model(input)' |> cpu
    Flux.reset!(model)
	# @show losss = Flux.mse(pred, target)
    p1 = plot(pred, label="Predict")
    p1 = plot!(target, label="Data", title="Loss $(loss(input, target))")
    display(plot(p1))
end

# ╔═╡ 6bbf14b9-c88f-44ab-ad8f-8de25174b998
begin
	@show start_loss = loss(input, target)
	for _ = 1:20
		Flux.train!(loss, Flux.params(model),Iterators.repeated((input, target), 20), ADAM(0.01))
	end
	@show final_loss = loss(input, target)
end

# ╔═╡ fbd02e0b-c2de-4a39-9896-84346e5f12f8
Threads.nthreads()

# ╔═╡ 7958c8df-79fa-44ef-b27b-0859ee47be08
loss(input, target)

# ╔═╡ 06f33530-06bb-4a75-954c-a48115e07bb4
@bind range_i Slider(1:500, show_value=true)

# ╔═╡ 648d121f-d8f7-4b7f-b416-5450976c2861
prediction = model(input)[range_i:end]

# ╔═╡ 56a5241b-4b0d-4f76-9e0c-180e1823c1f3
truth = target[range_i:end]

# ╔═╡ dd097505-8353-4dfd-8cef-63a7c59473af
plot([prediction, truth], labels=["prediction" "target"])

# ╔═╡ 3a48bb35-baf3-4583-bc1c-4da64eb7dd6e
plot(cases[:,1])

# ╔═╡ Cell order:
# ╠═8a2604ac-1f9b-11ec-30a2-0998ea216cac
# ╠═e9e23536-7553-4d61-bfa9-e3a76c6d6737
# ╠═aeedcfe8-1e98-49af-8ea1-55864638f6bf
# ╠═3085780c-1977-4683-80e9-c1151b2d804a
# ╠═2c9e7f3a-a913-4c3e-a70b-4e5b676bdcfa
# ╠═be7755f1-cfb2-4d42-940c-fe522390ad87
# ╠═bc612d5c-9092-489e-ab79-f62f0a814ae1
# ╠═7d37133c-b84f-4087-ae2b-3af5203cd028
# ╠═6bbf14b9-c88f-44ab-ad8f-8de25174b998
# ╠═fbd02e0b-c2de-4a39-9896-84346e5f12f8
# ╠═7958c8df-79fa-44ef-b27b-0859ee47be08
# ╠═06f33530-06bb-4a75-954c-a48115e07bb4
# ╠═648d121f-d8f7-4b7f-b416-5450976c2861
# ╠═56a5241b-4b0d-4f76-9e0c-180e1823c1f3
# ╠═dd097505-8353-4dfd-8cef-63a7c59473af
# ╠═3a48bb35-baf3-4583-bc1c-4da64eb7dd6e


text:
using OrdinaryDiffEq, Test, Random
using DiffEqProblemLibrary.ODEProblemLibrary: importodeproblems; importodeproblems()
import DiffEqProblemLibrary.ODEProblemLibrary: prob_ode_linear
Random.seed!(100)

@testset "Tstops Tests on the Interval [0, 1]" begin
  prob = prob_ode_linear

  sol =solve(prob,Tsit5(),dt=1//2^(6),tstops=[1/2])
  @test 1//2 ∈ sol.t

  sol =solve(prob,RK4(),dt=1//3,tstops=[1/2], adaptive=false)
  @test sol.t == [0,1/3,1/2,1/3+1/2,1]

  sol =solve(prob,RK4(),dt=1//3,tstops=[1/2],d_discontinuities=[-1/2,1/2,3/2], adaptive=false)
  @test sol.t == [0,1/3,1/2,1/3+1/2,1]

  # TODO
  integrator = init(prob,RK4(),tstops=[1/5,1/4,1/3,1/2,3/4], adaptive=false)

  sol =solve(prob,RK4(),tstops=[1/5,1/4,1/3,1/2,3/4], adaptive=false)
  @test sol.t == [0,1/5,1/4,1/3,1/2,3/4,1]

  sol =solve(prob,RK4(),tstops=[0,1/5,1/4,1/3,1/2,3/4,1], adaptive=false)
  @test sol.t == [0,1/5,1/4,1/3,1/2,3/4,1]

  sol = solve(prob,RK4(),tstops=0:1//16:1, adaptive=false)
  @test sol.t == collect(0:1//16:1)

  sol = solve(prob,RK4(),tstops=range(0,stop=1,length=100), adaptive=false)
  @test sol.t == collect(range(0,stop=1,length=100))
end

@testset "Integrator Tstops Tests on the Interval $(["[-1, 0]", "[0, 1]"][i])" for (i, tdir) in enumerate([-1.; 1.])
  prob2 = remake(prob_ode_linear, tspan=(0.0, tdir*1.0))
  integrator = init(prob2,Tsit5())
  tstops = tdir .* [0,1/5,1/4,1/3,1/2,3/4,1]
  for tstop in tstops
    add_tstop!(integrator, tstop)
  end
  @test_throws ErrorException add_tstop!(integrator, -0.1 * tdir)
  solve!(integrator)
  for tstop in tstops
    @test tstop ∈ integrator.sol.t
  end
end

@testset "Tstops Eps" begin
    function de(du, u, p, t) # specific DE does not impact the issue
               a, b = p
               du[1] = a * u[1]
               du[2] = b * u[2]
    end

    saveat = [0.0, 0.0094777, 1.5574]
    tstop = 0.010823

    affect!(integrator) = integrator.u[1] += 1.0
    condition(u, t, integrator) = t == tstop
    callback = DiscreteCallback(condition, affect!)

    prob = ODEProblem(de, zeros(2), (-1, 3.), rand(2))
    sol = solve(prob, Tsit5(), saveat=saveat, tstops=tstop, callback=callback)
    @test sol.t[end] == 1.5574
end

@testset "Tstops Type Conversion" begin
  called = Ref(false)
  tval = rand()
  ff(du,u,p,t) = du .= 0
  cb = DiscreteCallback((u,t,integrator)->t == Float32(tval),integrator->(called[]=true))
  prob = ODEProblem(ff,[0.0],(0f0,1f0))
  sol = solve(prob,Tsit5(),tstops = [tval], callback=cb)
end


text:
# ---------- #
# Grid types #
# ---------- #

abstract type AbstractGrid end

struct Cartesian <: AbstractGrid
    a::Vector{Float64}
    b::Vector{Float64}
    orders::Vector{Int}
end

mutable struct Domain
    states::Vector{Symbol}
    min::Vector{Float64}
    max::Vector{Float64}
end

function Cartesian(stuff::AbstractDict, calib::ModelCalibration)
    kind = get(stuff, :tag, nothing)
    if kind != :Cartesian
        error("Can't build Cartesian from dict with kind $(kind)")
    end

    a = eval_with(calib, stuff[:a])
    b = eval_with(calib, stuff[:b])
    orders = eval_with(calib, stuff[:orders])
    Cartesian(a, b, orders)
end

Cartesian(;a=[], b=[], orders=[]) = Cartesian(a, b, orders)

function _build_grid(data::AbstractDict, calib::ModelCalibration)
    if data[:tag] == :Cartesian
        return Cartesian(data, calib)
    else
        m = "don't know how to handle grid of type $(data[:tag])"
        error(m)
    end
end

# ------------------ #
# Distribution types #
# ------------------ #

abstract type AbstractDistribution end

function _build_dist(data::AbstractDict, calib::ModelCalibration)
    if data[:tag] == :Normal
        n = length(calib[:shocks])
        Sigma = reshape(vcat(data[:Sigma]...), n, n)
        return Distributions.MvNormal(_to_Float64(Sigma))
    else
        m = "don't know how to handle distribution of type $(data[:tag])"
        error(m)
    end
end

# ------------------------- #
# Discrete Transition types #
# ------------------------- #

to_vector(tab::Number) = [tab]
to_matrix(tab::Number) = reshape([tab], 1, 1)
to_matrix(tab::Array) = hcat([Array{Float64}(e) for e in tab]...)
to_matrix(tab::Array{Array{Float64,1},1}) = cat([e' for e in tab]...; dims=1)

function _build_exogenous_entry(data::AbstractDict, calib::ModelCalibration)

    if data[:tag] == :Product
        p1 = _build_exogenous_entry(data[:p1], calib)
        p2 = _build_exogenous_entry(data[:p2], calib)
        return ProductProcess(p1,p2)
    elseif data[:tag] == :MarkovChain
        # need to extract/clean up P and Q
        values = eval_with(calib, data[:values])
        states_values = to_matrix(values)
        transitions = eval_with(calib, data[:transitions])
        Π = to_matrix(transitions)
        return MarkovChain(Π, states_values)
    elseif data[:tag] == :VAR1
        # need to extract rho an dSigma
        rho = eval_with(calib, data[:rho])
        Sigma = eval_with(calib, data[:Sigma])
        N = eval_with(calib, get(data, :N, 10))  # TODO: should default be 10??
        # rho = to_matrix(rho)
        Sigma = to_matrix(Sigma)
        N = to_vector(N)
        return VAR1(rho, Sigma)
    elseif data[:tag] == :Normal
        # need to extract rho an dSigma
        Sigma = eval_with(calib, data[:Sigma])
        Sigma = to_matrix(Sigma)
        if size(Sigma, 1) != size(Sigma, 2)
            sz = size(Sigma)
            msg = "Covariance matrix must be square. Found a $sz matrix"
            throw(DimensionMismatch(msg))
        end
        return Normal(Sigma)
    elseif data[:tag] == :DeathProcess
        # need to extract rho an dSigma
        mu = eval_with(calib, data[:mu])
        return DeathProcess(mu)
    elseif data[:tag] == :PoissonProcess
        # need to extract rho an dSigma
        mu = eval_with(calib, data[:mu])
        K = eval_with(calib, data[:K])
        return PoissonProcess(mu, K)
    elseif data[:tag] == :AgingProcess
        # need to extract rho an dSigma
        mu = eval_with(calib, data[:mu])
        K = eval_with(calib, data[:K])
        return AgingProcess(mu, K)
    end
    m = "don't know how to handle exogenous process of type $(data[:tag])"
    error(m)
end


text:
import TensorToolbox

"""
High-order singular value decomposition (HO-SVD).
"""
function hosvd{T,N}(tensor::StridedArray{T,N}, core_dims::NTuple{N, Int}, eigmethod=trues(N); pad_zeros::Bool=false, compute_error::Bool=false)
	pad_zeros || _check_tensor(tensor, core_dims)

	factors = map(1:N) do i
		X = _col_unfold(tensor, i)
		if eigmethod[i]
			f = eigs(X'X, nev=core_dims[i])[2]
		else
			f = eig(X'X)[2]
		end
		if pad_zeros && size(f, 2) < core_dims[i] # fill missing factors with zeros
			warn("Zero slices ($(core_dims[i]-size(f, 2))) added in dimension $i ")
			f = hcat(f, zeros(T, size(tensor, i), core_dims[i]-size(f, 2)))
		end
		mapslices(_check_sign, f, 1)
	end

	res = Tucker((factors...), tensorcontractmatrices(tensor, factors))
	if compute_error
		_set_rel_residue(res, tensor)
		info("Error: $(res.props[:rel_residue])")
	end
	info("HOSVD core rank: $(TensorToolbox.mrank(res.core))")
	return res
end

hosvd{T,N}(tensor::StridedArray{T,N}, r::Int; compute_error::Bool=false, pad_zeros::Bool=false) =
	hosvd(tensor, (fill(r, N)...); compute_error=compute_error, pad_zeros=pad_zeros)


text:
module MetidaStats

    using MetidaBase

    import MetidaBase: AbstractIdData,
    AbstractIDResult,
    DataSet,
    Tables,
    indsdict!,
    metida_table,
    getid,
    MetidaTable,
    sortbyvec!,
    skipnanormissing, isnanormissing,
    sortbyvec!,
    skipnonpositive, ispositive

    using StatsBase, Statistics

    import Base: ht_keyindex, getindex, names

    export descriptives, dataimport, cvfromsd

    include("types.jl")
    include("descriptive.jl")

end


text:
using ModelingToolkit, DiffEqBase, DiffEqJump, Test, LinearAlgebra
MT = ModelingToolkit

# basic MT SIR model with tweaks
@parameters β γ t
@variables S(t) I(t) R(t)
rate₁ = β * S * I
affect₁ = [S ~ S - 1, I ~ I + 1]
rate₂ = γ * I + t
affect₂ = [I ~ I - 1, R ~ R + 1]
j₁ = ConstantRateJump(rate₁, affect₁)
j₂ = VariableRateJump(rate₂, affect₂)
@named js = JumpSystem([j₁, j₂], t, [S, I, R], [β, γ])
statetoid = Dict(MT.value(state) => i for (i, state) in enumerate(states(js)))
mtjump1 = MT.assemble_crj(js, j₁, statetoid)
mtjump2 = MT.assemble_vrj(js, j₂, statetoid)

# doc version
rate1(u, p, t) = (0.1 / 1000.0) * u[1] * u[2]
function affect1!(integrator)
    integrator.u[1] -= 1
    integrator.u[2] += 1
end
jump1 = ConstantRateJump(rate1, affect1!)
rate2(u, p, t) = 0.01u[2] + t
function affect2!(integrator)
    integrator.u[2] -= 1
    integrator.u[3] += 1
end
jump2 = VariableRateJump(rate2, affect2!)

# test crjs
u = [100, 9, 5]
p = (0.1 / 1000, 0.01)
tf = 1.0
mutable struct TestInt{U, V, T}
    u::U
    p::V
    t::T
end
mtintegrator = TestInt(u, p, tf)
integrator = TestInt(u, p, tf)
@test abs(mtjump1.rate(u, p, tf) - jump1.rate(u, p, tf)) < 10 * eps()
@test abs(mtjump2.rate(u, p, tf) - jump2.rate(u, p, tf)) < 10 * eps()
mtjump1.affect!(mtintegrator)
jump1.affect!(integrator)
@test all(integrator.u .== mtintegrator.u)
mtintegrator.u .= u;
integrator.u .= u;
mtjump2.affect!(mtintegrator)
jump2.affect!(integrator)
@test all(integrator.u .== mtintegrator.u)

# test MT can make and solve a jump problem
rate₃ = γ * I
affect₃ = [I ~ I - 1, R ~ R + 1]
j₃ = ConstantRateJump(rate₃, affect₃)
@named js2 = JumpSystem([j₁, j₃], t, [S, I, R], [β, γ])
u₀ = [999, 1, 0];
p = (0.1 / 1000, 0.01);
tspan = (0.0, 250.0);
u₀map = [S => 999, I => 1, R => 0]
parammap = [β => 0.1 / 1000, γ => 0.01]
dprob = DiscreteProblem(js2, u₀map, tspan, parammap)
jprob = JumpProblem(js2, dprob, Direct(), save_positions = (false, false))
Nsims = 30000
function getmean(jprob, Nsims)
    m = 0.0
    for i in 1:Nsims
        sol = solve(jprob, SSAStepper())
        m += sol[end, end]
    end
    m / Nsims
end
m = getmean(jprob, Nsims)

@variables S2(t)
obs = [S2 ~ 2 * S]
@named js2b = JumpSystem([j₁, j₃], t, [S, I, R], [β, γ], observed = obs)
dprob = DiscreteProblem(js2b, u₀map, tspan, parammap)
jprob = JumpProblem(js2b, dprob, Direct(), save_positions = (false, false))
sol = solve(jprob, SSAStepper(), saveat = tspan[2] / 10)
@test all(2 .* sol[S] .== sol[S2])

# test save_positions is working
jprob = JumpProblem(js2, dprob, Direct(), save_positions = (false, false))
sol = solve(jprob, SSAStepper(), saveat = 1.0)
@test all((sol.t) .== collect(0.0:tspan[2]))

#test the MT JumpProblem rates/affects are correct
rate2(u, p, t) = 0.01u[2]
jump2 = ConstantRateJump(rate2, affect2!)
mtjumps = jprob.discrete_jump_aggregation
@test abs(mtjumps.rates[1](u, p, tf) - jump1.rate(u, p, tf)) < 10 * eps()
@test abs(mtjumps.rates[2](u, p, tf) - jump2.rate(u, p, tf)) < 10 * eps()
mtjumps.affects![1](mtintegrator)
jump1.affect!(integrator)
@test all(integrator.u .== mtintegrator.u)
mtintegrator.u .= u;
integrator.u .= u;
mtjumps.affects![2](mtintegrator)
jump2.affect!(integrator)
@test all(integrator.u .== mtintegrator.u)

# direct vers
p = (0.1 / 1000, 0.01)
prob = DiscreteProblem([999, 1, 0], (0.0, 250.0), p)
r1(u, p, t) = (0.1 / 1000.0) * u[1] * u[2]
function a1!(integrator)
    integrator.u[1] -= 1
    integrator.u[2] += 1
end
j1 = ConstantRateJump(r1, a1!)
r2(u, p, t) = 0.01u[2]
function a2!(integrator)
    integrator.u[2] -= 1
    integrator.u[3] += 1
end
j2 = ConstantRateJump(r2, a2!)
jset = JumpSet((), (j1, j2), nothing, nothing)
jprob = JumpProblem(prob, Direct(), jset, save_positions = (false, false))
m2 = getmean(jprob, Nsims)

# test JumpSystem solution agrees with direct version
@test abs(m - m2) / m < 0.01

# mass action jump tests for SIR model
maj1 = MassActionJump(2 * β / 2, [S => 1, I => 1], [S => -1, I => 1])
maj2 = MassActionJump(γ, [I => 1], [I => -1, R => 1])
@named js3 = JumpSystem([maj1, maj2], t, [S, I, R], [β, γ])
dprob = DiscreteProblem(js3, u₀map, tspan, parammap)
jprob = JumpProblem(js3, dprob, Direct())
m3 = getmean(jprob, Nsims)
@test abs(m - m3) / m < 0.01

# maj jump test with various dep graphs
@named js3b = JumpSystem([maj1, maj2], t, [S, I, R], [β, γ])
jprobb = JumpProblem(js3b, dprob, NRM())
m4 = getmean(jprobb, Nsims)
@test abs(m - m4) / m < 0.01
jprobc = JumpProblem(js3b, dprob, RSSA())
m4 = getmean(jprobc, Nsims)
@test abs(m - m4) / m < 0.01

# mass action jump tests for other reaction types (zero order, decay)
maj1 = MassActionJump(2.0, [0 => 1], [S => 1])
maj2 = MassActionJump(γ, [S => 1], [S => -1])
@named js4 = JumpSystem([maj1, maj2], t, [S], [β, γ])
dprob = DiscreteProblem(js4, [S => 999], (0, 1000.0), [β => 100.0, γ => 0.01])
jprob = JumpProblem(js4, dprob, Direct())
m4 = getmean(jprob, Nsims)
@test abs(m4 - 2.0 / 0.01) * 0.01 / 2.0 < 0.01

# test second order rx runs
maj1 = MassActionJump(2.0, [0 => 1], [S => 1])
maj2 = MassActionJump(γ, [S => 2], [S => -1])
@named js4 = JumpSystem([maj1, maj2], t, [S], [β, γ])
dprob = DiscreteProblem(js4, [S => 999], (0, 1000.0), [β => 100.0, γ => 0.01])
jprob = JumpProblem(js4, dprob, Direct())
sol = solve(jprob, SSAStepper());

# issue #819
@testset "Combined system name collisions" begin
    sys1 = JumpSystem([maj1, maj2], t, [S], [β, γ], name = :sys1)
    sys2 = JumpSystem([maj1, maj2], t, [S], [β, γ], name = :sys1)
    @test_throws ArgumentError JumpSystem([sys1.γ ~ sys2.γ], t, [], [],
                                          systems = [sys1, sys2], name = :foo)
end

# test if param mapper is setup correctly for callbacks
@parameters k1 k2 k3
@variables A(t) B(t)
maj1 = MassActionJump(k1 * k3, [0 => 1], [A => -1, B => 1])
maj2 = MassActionJump(k2, [B => 1], [A => 1, B => -1])
@named js5 = JumpSystem([maj1, maj2], t, [A, B], [k1, k2, k3])
p = [k1 => 2.0, k2 => 0.0, k3 => 0.5]
u₀ = [A => 100, B => 0]
tspan = (0.0, 2000.0)
dprob = DiscreteProblem(js5, u₀, tspan, p)
jprob = JumpProblem(js5, dprob, Direct(), save_positions = (false, false))
@test all(jprob.massaction_jump.scaled_rates .== [1.0, 0.0])

pcondit(u, t, integrator) = t == 1000.0
function paffect!(integrator)
    integrator.p[1] = 0.0
    integrator.p[2] = 1.0
    reset_aggregated_jumps!(integrator)
end
sol = solve(jprob, SSAStepper(), tstops = [1000.0],
            callback = DiscreteCallback(pcondit, paffect!))
@test sol[1, end] == 100

# observed variable handling
@variables OBS(t)
@named js5 = JumpSystem([maj1, maj2], t, [S], [β, γ]; observed = [OBS ~ 2 * S])
OBS2 = OBS
@test isequal(OBS2, @nonamespace js5.OBS)
@unpack OBS = js5
@test isequal(OBS2, OBS)


text:
# -*- coding: utf-8 -*-

using Pkg; Pkg.add("MiniFB")

using LinearAlgebra

using Plots
using Distributions
using MiniFB
using Plots
using FileIO
using ImageTransformations

using GeometryObjects

const WIDTH = 600
const HEIGHT = 600

function populate_buffer!(buffer, img)
    global WIDTH, HEIGHT
    buffer .= zero(UInt32)
    h, w = size(img)
    ratio = min(1, HEIGHT / h, WIDTH / w)
    nimg = imresize(img, ratio=ratio)
    nh, nw = size(nimg)
    oh = (HEIGHT - nh) ÷ 2
    ow = (WIDTH - nw) ÷ 2
    for i in 1:nh
        for j in 1:nw
            buffer[(oh+i-1)*WIDTH+ow+j] = mfb_rgb(nimg[i, j])
        end
    end
end

Base.@kwdef struct CollisionSystem
    xmin::Float64 = -4
    xmax::Float64 = 4
    ymin::Float64 = -4
    ymax::Float64 = 4
    maxT::Int = 500
    num_particles::Int = 20
end

function setup_particle(sys::CollisionSystem)
    particles = Particle2D[]
    num_particles = sys.num_particles
    while length(particles) != num_particles
        θ = 2π * rand()
        v_abs = 0.1
        v = Velocity2D(v_abs * cos(θ), v_abs * sin(θ))
        m = r = rand(Uniform(0.1, 0.3))
        x = rand(Uniform(sys.xmin + r, sys.xmax - r))
        y = rand(Uniform(sys.ymin + r, sys.ymax - r))

        c = Particle2D(
            Circle2D(Point2D(x, y), r),
            v,
            m,
        )

        c_plus_r = Particle2D(
            Circle2D(Point2D(x, y), r + v_abs),
            v,
            m,
        )

        any(has_contact.(Ref(c_plus_r), particles)) && continue
        push!(particles, c)
    end
    return particles
end

function move!(sys::CollisionSystem, pt::Point2D, v::Velocity2D, r::Real=0, Δt=1.0)
    pt .+= Δt .* v # do not use `pt += v`
    if pt.y ≥ sys.ymax - r
        pt.y = 2(sys.ymax - r) - pt.y
        v.y = -v.y
    end
    if pt.y ≤ (sys.ymin + r)
        pt.y = 2(sys.ymin + r) - pt.y
        v.y = -v.y
    end

    if pt.x ≥ (sys.xmax - r)
        pt.x = 2(sys.xmax - r) - pt.x
        v.x = -v.x
    end
    if pt.x ≤ (sys.xmin + r)
        pt.x = 2(sys.xmin + r) - pt.x
        v.x = -v.x
    end
    return pt, v
end

move!(sys::CollisionSystem, p::Particle2D; Δt=1.0) = move!(sys, p.c.pt, p.v, p.c.r, Δt)

dist(pt1, pt2) = sqrt((pt1.x - pt2.x)^2 + (pt1.y - pt2.y)^2)
has_contact(c1::Circle2D, c2::Circle2D) = c1.r + c2.r ≥ dist(c1.pt, c2.pt)
has_contact(p1::Particle2D, p2::Particle2D) = has_contact(p1.c, p2.c)

function minifb(sys::CollisionSystem)
    particles = setup_particle(sys)
    Δt = 0.1
    accumΔt = 0.0
    cnt = 0

    buf = IOBuffer()
    buffer = zeros(UInt32, HEIGHT * WIDTH)
    window = mfb_open_ex(
        "Elastic Collision 2D",
        WIDTH, HEIGHT,
        MiniFB.WF_RESIZABLE,
    )
    while true
        accumΔt += Δt
        move!.(Ref(sys), particles; Δt)
        # naive collision algorithm
        for i in 1:length(particles)
            for j in (i+1):length(particles)
                particle1 = particles[i]
                particle2 = particles[j]
                c1 = particle1.c
                c2 = particle2.c
                # update c1.v and c2.v
                if has_contact(c1, c2)
                    pt1 = c1.pt
                    pt2 = c2.pt

                    v1 = particle1.v
                    v2 = particle2.v
                    m1 = particle1.m
                    m2 = particle2.m

                    # correct position so that c1 and c2 does not contact
                    while has_contact(c1, c2)
                        c1.pt .-= 0.015 .* v1
                        c2.pt .-= 0.015 .* v2
                    end

                    # update velocity
                    v1_next =
                        v1 -
                        2m2 / (m1 + m2) * dot(v1 - v2, pt1 - pt2) * (pt1 - pt2) /
                        dot(pt1 - pt2, pt1 - pt2)
                    v2_next =
                        v2 -
                        2m1 / (m2 + m1) * dot(v2 - v1, pt2 - pt1) * (pt2 - pt1) /
                        dot(pt2 - pt1, pt2 - pt1)
                    particle1.v .= v1_next
                    particle2.v .= v2_next
                end
            end
        end

        if accumΔt ≥ 1.0
            cnt += 1
            s = plot(
                size=(WIDTH, HEIGHT),
                xlim=(sys.xmin, sys.xmax), ylim=(sys.ymin, sys.ymax),
                aspect_ratio=:equal, legend=false,
                grid=false, ticks=false,
                framestyle=:box,
            )
            plot!(s, particles)
            plot!(s, title="$cnt")
            accumΔt = zero(accumΔt)

            show(buf, MIME("image/png"), s)
            img = buf |> load
            take!(buf)
            populate_buffer!(buffer, img)
            state = mfb_update(window, buffer)
            if state != MiniFB.STATE_OK
                break
            end
        end
    end
    mfb_close(window)
end

minifb(CollisionSystem(; xmin=-8, xmax=8, num_particles=20))