text:
"""
  (S::NLNewton)(integrator) -> (z, η, iter, fail_convergence)

Perform numerically stable modified NLNewton iteration that is specialized for
implicit methods (see [^HS96] and [^HW96]), where `z` is the solution, `η` is
used to measure the iteration error (see [^HW96]), `iter` is the number of
iteration, and `fail_convergence` reports whether the algorithm succeed. It
solves

```math
G(z) = dt⋅f(tmp + γ⋅z, p, t+c⋅h) - z = 0⃗
```

by iterating

```math
W Δᵏ = dt⋅f(tmp + γ⋅zᵏ, p, t+c⋅h) - zᵏ
zᵏ⁺¹ = zᵏ + Δᵏ
```

where `W=M-dt⋅γJ`, `M` is the mass matrix, `dt` is the step size, `γ` is a
constant, and `J` is the Jacobian matrix.

[^HS96]: M.E.Hoseaa and L.F.Shampine, "Analysis and implementation of TR-BDF2",
Applied Numerical Mathematics, Volume 20, Issues 1–2, February 1996, Pages
21-37.
[doi:10.1016/0168-9274(95)00115-8](https://doi.org/10.1016/0168-9274(95)00115-8)

[^HW96]: Ernst Hairer and Gerhard Wanner, "Solving Ordinary Differential
Equations II, Springer Series in Computational Mathematics. ISBN
978-3-642-05221-7. Section IV.8.
[doi:10.1007/978-3-642-05221-7](https://doi.org/10.1007/978-3-642-05221-7)
"""
function (S::NLNewton{false})(integrator)
  nlcache = S.cache
  @unpack t,dt,uprev,u,f,p = integrator
  @unpack z,tmp,W,κ,tol,c,γ,max_iter,min_iter = nlcache
  mass_matrix = integrator.f.mass_matrix
  #alg = unwrap_alg(integrator, true)
  if typeof(integrator.f) <: SplitFunction
    f = integrator.f.f1
  else
    f = integrator.f
  end
  # precalculations
  κtol = κ*tol

  # initial step of NLNewton iteration
  iter = 1
  tstep = t + c*dt
  u = tmp + γ*z
  b = dt*f(u, p, tstep) - z
  dz = W\b
  ndz = integrator.opts.internalnorm(dz)
  z = z + dz

  η = max(nlcache.ηold,eps(eltype(integrator.opts.reltol)))^(0.8)
  do_newton = integrator.success_iter == 0 || η*ndz > κtol

  # NLNewton iteration
  fail_convergence = false
  while (do_newton || iter < min_iter) && iter < max_iter
    iter += 1
    u = tmp + γ*z
    b = dt*f(u, p, tstep) - z
    dz = W\b
    ndzprev = ndz
    ndz = integrator.opts.internalnorm(dz)
    θ = ndz/ndzprev
    if θ > 1 || ndz*(θ^(max_iter - iter)/(1-θ)) > κtol
      fail_convergence = true
      break
    end
    η = θ/(1-θ)
    do_newton = (η*ndz > κtol)
    z = z + dz
  end

  if (iter >= max_iter && do_newton) || fail_convergence
    integrator.force_stepfail = true
    return (z, η, iter, true)
  end

  return (z, η, iter, false)
end

function (S::NLNewton{true})(integrator)
  nlcache = S.cache
  cache = integrator.cache
  @unpack t,dt,uprev,u,f,p = integrator
  @unpack z,dz,tmp,b,W,κ,tol,k,new_W,c,γ,max_iter,min_iter = nlcache
  mass_matrix = integrator.f.mass_matrix
  #alg = unwrap_alg(integrator, true)
  if typeof(integrator.f) <: SplitFunction
    f = integrator.f.f1
  else
    f = integrator.f
  end
  # precalculations
  κtol = κ*tol
  # initial step of NLNewton iteration
  iter = 1
  tstep = t + c*dt
  @. u = tmp + γ*z
  f(k, u, p, tstep)
  if mass_matrix == I
    @. b = dt*k - z
  else
    mul!(vec(b),mass_matrix,vec(z))
    @. b = dt*k - b
  end
  if DiffEqBase.DiffEqBase.has_invW(f)
    mul!(vec(dz),W,vec(b)) # Here W is actually invW
  else
    cache.linsolve(vec(dz),W,vec(b),new_W)
  end
  ndz = integrator.opts.internalnorm(dz)
  z .+= dz

  η = max(nlcache.ηold,eps(eltype(integrator.opts.reltol)))^(0.8)
  do_newton = integrator.success_iter == 0 || η*ndz > κtol

  # NLNewton iteration
  fail_convergence = false
  while (do_newton || iter < min_iter) && iter < max_iter
    iter += 1
    @. u = tmp + γ*z
    f(k, u, p, tstep)
    if mass_matrix == I
      @. b = dt*k - z
    else
      mul!(vec(b),mass_matrix,vec(z))
      @. b = dt*k - b
    end
    if DiffEqBase.DiffEqBase.has_invW(f)
      mul!(vec(dz),W,vec(b)) # Here W is actually invW
    else
      cache.linsolve(vec(dz),W,vec(b),false)
    end
    ndzprev = ndz
    ndz = integrator.opts.internalnorm(dz)
    θ = ndz/ndzprev
    if θ > 1 || ndz*(θ^(max_iter - iter)/(1-θ)) > κtol
      fail_convergence = true
      break
    end
    η = θ/(1-θ)
    do_newton = (η*ndz > κtol)
    z .+= dz
  end

  if (iter >= max_iter && do_newton) || fail_convergence
    integrator.force_stepfail = true
    return (z, η, iter, true)
  end

  return (z, η, iter, false)
end


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

JLLWrappers.@generate_main_file_header("poly2tri")
JLLWrappers.@generate_main_file("poly2tri", UUID("0bbe0808-d418-57ff-bb48-dc0bf000ef4b"))
end  # module poly2tri_jll

ext: jl


text:
export macau_sgd

function macau_sgd(data::RelationData;
                   num_latent::Int = 10,
                   verbose::Bool   = true,
                   niter::Int      = 100,
                   lrate::Float64  = 0.001,
                   reset_model     = true,
                   batch_size      = 4000,
                   clamp::Vector{Float64} = Float64[])
  ## initialization
  verbose     && println("Model setup")
  reset_model && reset!(data, num_latent)

  rmse = NaN
  rmse_sample = NaN
  rmse_train  = NaN

  ## data
  df  = data.relations[1].data.df
  mean_value = mean(df[:,end])
  uid = convert(Vector{Int32}, df[:,1])
  vid = convert(Vector{Int32}, df[:,2])
  val = convert(Array, df[:,3]) - mean_value

  Ubatches = create_minibatches(vid, uid, val, batch_size)
  Vbatches = [x' for x in Ubatches]

  Usample  = data.entities[1].model.sample
  Vsample  = data.entities[2].model.sample

  test_uid = convert(Vector{Int}, data.relations[1].test_vec[:,1])
  test_vid = convert(Vector{Int}, data.relations[1].test_vec[:,2])
  test_val = convert(Vector, data.relations[1].test_vec[:,3])
  test_idx = hcat(test_uid, test_vid)

  alpha = data.relations[1].model.alpha
  yhat_post = zeros(length(test_val))

  @time for i in 1:length(Ubatches)
    sgd_update!(Usample, Vsample, Ubatches[i], data.entities[1].model, alpha)
    sgd_update!(Vsample, Usample, Vbatches[i], data.entities[2].model, alpha)
  end
  return Ubatches, Vbatches

end

function create_minibatches{Ti,Tv}(uid::Vector{Ti}, vid::Vector{Ti}, val::Vector{Tv}, bsize::Int)
  perm = randperm(length(uid))
  batches = SparseMatrixCSC{Tv,Ti}[]
  umax = maximum(uid)
  vmax = maximum(vid)

  for i in 1:bsize:length(uid)
    j = min(length(uid), i + bsize - 1)
    push!(batches, sparse(uid[perm[i:j]], vid[perm[i:j]], val[perm[i:j]], umax, vmax) )
  end
  return batches
end

function sgd_update!(sample, Vsample, batch::SparseMatrixCSC, model::EntityModel, alpha)
  for n in 1:size(sample, 2)
    grad!(n, sample, Vsample, batch, model.Lambda, model.mu, alpha)
  end
  nothing
end

function grad!(n,
              sample  ::Matrix{Float64},
              Vsample ::Matrix{Float64},
              Udata   ::SparseMatrixCSC,
              Lambda  ::Matrix{Float64},
              mu      ::Vector{Float64},
              alpha   ::Float64)
  un  = sample[:, n]
  idx = Udata.colptr[n] : Udata.colptr[n+1]-1
  ff  = Udata.rowval[ idx ]
  rr  = Udata.nzval[ idx ]

  MM = Vsample[:, ff]

  sample[:,n] += alpha * (MM * rr - (MM*(MM' * un))) + Lambda * (mu - un)
  nothing
end

function grad_yhat(
              U       ::Matrix{Float64},
              V       ::Matrix{Float64},
              colptr  ::Vector,
              rowval  ::Vector,
              #nzval   ::Vector,
              resid   ::Vector,
              Lambda  ::Matrix{Float64},
              mu      ::Vector{Float64},
              alpha   ::Float64)
  num_latent = size(U, 1)
  grad = Lambda * (mu .- U)
  @inbounds for n in 1:size(U,2)
    @inbounds for i in colptr[n] : colptr[n+1]-1
      err_i = alpha * resid[i]
      v_i   = rowval[i]
      @inbounds @simd for k in 1:num_latent
        grad[k,n] += err_i * V[k, v_i]
      end
    end
  end

  return grad
end

function grad_yhat!(
              grad    ::Matrix{Float64},
              U       ::Matrix{Float64},
              V       ::Matrix{Float64},
              colptr  ::Vector,
              rowval  ::Vector,
              #nzval   ::Vector,
              resid   ::Vector,
              lambda  ::Vector{Float64},
              mu      ::Vector{Float64},
              alpha   ::Float64)
  num_latent = size(U, 1)
  @inbounds for n in 1:size(U,2)
    @inbounds @simd for k in 1:num_latent
      grad[n, k] = (U[n, k] - mu[k]) * lambda[k]
    end
  end
  @inbounds for n in 1:size(U,2)
    @inbounds for i in colptr[n] : colptr[n+1]-1
      err_i = alpha * resid[i]
      v_i   = rowval[i]
      @inbounds @simd for k in 1:num_latent
        grad[k,n] += err_i * V[k, v_i]
      end
    end
  end

  return grad
end


text:
l_exp = []
push!(l_exp, readdlm(ckpt_path * "/../F24_p20bar_phi0.5.txt", ','))
push!(l_exp, readdlm(ckpt_path * "/../F24_p20bar_phi1.0.txt", ','))
push!(l_exp, readdlm(ckpt_path * "/../F24_p20bar_phi1.3.txt", ','))

exp_all = vcat(l_exp...)

n_exp = size(exp_all)[1]
conds = zeros(n_exp, 4)

conds[:, 1] .= 1000.0 ./ exp_all[:, 1]
conds[:, 4] .= exp_all[:, 2] ./ 1000
conds[:, 2] .= 20/ 1.013

conds[1:41, 3] .= 0.5 * 1.08831169
conds[42:68, 3] .= 1.0 * 1.08831169
conds[69:96, 3] .= 1.3 * 1.08831169

n_train = Int64(floor(n_exp * 0.8))

for i_exp in 1:n_exp
    local T0 = conds[i_exp, 1]
    local P = conds[i_exp, 2] * one_atm
    local phi = conds[i_exp, 3]
    @printf("(%d) T %.1f [K] P %.1f [atm] phi %.1f IDT0 = %.2e [s] \n",
            i_exp, conds[i_exp, 1], conds[i_exp, 2],
            conds[i_exp, 3], conds[i_exp, 4])
end
@save string(ckpt_path, "/../conds.bson") conds

Random.seed!(0)
l_all = randperm(n_exp)
l_exp_train = l_all[1:n_train]
l_exp_val = l_all[n_train+1:end]

max_stars_repo_name: DENG-MIT/DeepReduction

max_issues_repo_name: DENG-MIT/DeepReduction

max_forks_repo_name: DENG-MIT/DeepReduction


text:
struct Parameters
    density::CPEGDensityParameters
    gravity::GravityParameters
    aero::AeroParameters
    function Parameters()
        new(CPEGDensityParameters(),GravityParameters(),AeroParameters())
    end
end

struct Planet
    Ω::Float64
    ω::SVector{3,Float64}
    function Planet()
        # mars
        Ω = 7.088218127774194e-5 # rad/s
        ω = SA[0,0,Ω]
        return new(Ω,ω)
    end
end

mutable struct COST
    rf::SVector{3,Float64}
    γ::Float64
    σ_tr::Float64
    function COST()
        new(SA[1.,2.,3.], 1e3 , deg2rad(20))
    end
end

mutable struct CPEGWorkspace
    params::Parameters
    scale::Scaling
    planet::Planet
    qp_solver_opts::SolverSettings
    cost::COST
    U::Vector{SVector{1,Float64}}
    miss_distance::Float64
    dt::Float64
    σ::Vector{Float64}
    max_cpeg_iter::Int64
    miss_distance_tol::Float64
    ndu_tol::Float64
    verbose::Bool
    function CPEGWorkspace()
        U = [SA[0.0] for i = 1:1000]
        new(Parameters(), Scaling(), Planet(), SolverSettings(), COST(), U,0.0,0.0,zeros(2),20,1e3,1e-2,true)
    end
end

# function initialize_CPEG(p::)

# # density methods
# @inline function altitude(ev::CPEGWorkspace,r::SVector{3,T}) where T
#     norm(r) - ev.params.gravity.R
# end
#
# @inline function density(ev::CPEGWorkspace,r::SVector{3,T}) where T
#     return density(ev.params.density,altitude(ev,r))
# end


text:
include("lyapunovs.jl")

# Use this to run: (do it in REPL)
# using PkgBenchmark
#
# benchmarkpkg("ChaosTools")
            # script=defaultscript(pkg),
            # require=defaultrequire(pkg),
            # resultsdir=defaultresultsdir(pkg),
            # saveresults=true,
            # tunefile=defaulttunefile(pkg),
            # retune=false,
            # promptsave=true,
            # promptoverwrite=true)

# This is for copy/paste:
# b = benchmarkpkg("ChaosTools", retune = true, promptsave = false); showall(b)


text:
function update_position!(x, p, fx, fp)
    i_update = (fx .< fp)
    p[i_update, :] = copy(x[i_update, :])
    fp[i_update, :] = fx[i_update, :]
end

function pso(func, lb, ub; swarmsize=100, ω=0.5, ϕp=0.5, ϕg=0.5, maxiter=100,
             verbose=false, swarmIni=[], particle_output=false)

    obj = x -> func(x)

    minstep=1e-8
    minfunc=1e-8

    # Initialize the particle swarm
    vhigh = abs.(ub .- lb) .*0.2
    vlow = -vhigh .*0.2
    S = swarmsize
    D = length(lb)  # the number of dimensions each particle has

    if length(swarmIni)>=1
        x=swarmIni
    else
        x = lb' .+ rand(S, D) .* (ub .- lb)'  # particle positions
    end
    v = vlow' .+ rand(S, D) .* (vhigh .- vlow)'  # particle velocities
    p = zeros(S, D)  # best particle positions

    Vmin = ones(S,1) * vlow'
    Vmax = ones(S,1) * vhigh'

    fx = [obj(x[i, :]) for i = 1:S]  # current particle function values
    fp = ones(S) * Inf  # best particle function values

    g = copy(x[1, :])  # best swarm position
    fg = Inf  # best swarm position starting value

    # Store particle's best position (if constraints are satisfied)
    update_position!(x, p, fx, fp)

    # Update swarm's best position
    i_min = argmin(fp)
    if fp[i_min] < fg
        g = copy(p[i_min, :])
        fg = fp[i_min]
    end

    # Iterate until termination criterion met
    it = 1
    while it <= maxiter
        rp = rand(S, D)
        rg = rand(S, D)

        # Update the particles' velocities and positions
        ω = ω
        v = ω*v .+ ϕp*rp.*(p .- x) .+ ϕg*rg.*(g' .- x)
        v = min.(v,Vmax)
        v = max.(v,Vmin)

        x += v
        # Correct for bound violations
        maskl = x .< lb'
        masku = x .> ub'
        x = x.*(.~(maskl .| masku)) .+ lb'.*maskl .+ ub'.*masku

        # Update objectives and constraints
        for i = 1:S
            fx[i] = obj(x[i, :])
        end

        # Store particle's best position (if constraints are satisfied)
        update_position!(x, p, fx, fp)

        # Compare swarm's best position with global best position
        i_min = argmin(fp)
        if fp[i_min] < fg
            verbose && println("New best for swarm at iteration $(it): $(p[i_min, :]) $(fp[i_min])")

            p_min = copy(p[i_min, :])
            stepsize = √(sum((g .- p_min).^2))

            if abs.(fg .- fp[i_min]) <= minfunc
                verbose && println("Stopping search: Swarm best objective change less than $(minfunc)")
                return (g, fg, p, fp)
            end
            if stepsize <= minstep
                verbose && println("Stopping search: Swarm best position change less than $(minstep)")
                return (g, fg, p, fp)
            end

            g = copy(p_min)
            fg = fp[i_min]
        end

        verbose && println("Best after iteration $(it): $(g) $(fg)")
        it += 1
    end

    verbose && println("Stopping search: maximum iterations reached --> $(maxiter)")
    return particle_output ? (g, fg, p, fp) : (g, fg)
end


text:
using Oceananigans.Buoyancy

"""
    NonDimensionalModel(; N, L, Re, Pr=0.7, Ro=Inf, float_type=Float64, kwargs...)

Construct a "Non-dimensional" `Model` with resolution `N`, domain extent `L`,
precision `float_type`, and the four non-dimensional numbers:

    * `Re = U λ / ν` (Reynolds number)
    * `Pr = U λ / κ` (Prandtl number)
    * `Ro = U / f λ` (Rossby number)

for characteristic velocity scale `U`, length-scale `λ`, viscosity `ν`,
tracer diffusivity `κ`, and Coriolis parameter `f`. Buoyancy is scaled
with `λ U²`, so that the Richardson number is `Ri=B`, where `B` is a
non-dimensional buoyancy scale set by the user via initial conditions or
forcing.

Note that `N`, `L`, and `Re` are required.

Additional `kwargs` are passed to the regular `Model` constructor.
"""
function NonDimensionalModel(; grid, float_type=Float64, Re, Pr=0.7, Ro=Inf,
    buoyancy = BuoyancyTracer(),
    coriolis = FPlane(float_type, f=1/Ro),
     closure = ConstantIsotropicDiffusivity(float_type, ν=1/Re, κ=1/(Pr*Re)),
    kwargs...)

    return Model(; float_type=float_type, grid=grid, closure=closure,
                   coriolis=coriolis, tracers=(:b,), buoyancy=buoyancy, kwargs...)
end


text:
#!/usr/bin/env julia

using Base.Test
using Drs.DrsMathProgSolverInterface

A = [2 3 1 0 0 0; -3 2 0 1 0 0; 0 2 0 0 1 0; 2 1 0 0 0 1]

m = DrsMathProgModel(A)
Drs.DrsMathProgSolverInterface.DrsFindPotentialBasis!(m)

@test m.basis == [3, 4, 5, 6]
@test m.nonbasis == [1, 2]

max_stars_repo_name: abogushov/lplib

max_issues_repo_name: abogushov/lplib

max_forks_repo_name: abogushov/lplib


text:
function WeakRemoverModel(A,b,C,d,XLB,XUB,VariableDef,Y_bar1,YR)

    const EPS = 1e-5;

    m,n = size(A);
    nObj = size(C)[1];
    MyModel = Model();
    @variable(MyModel,XLB[i] <= X[i=1:n] <= XUB[i],category = VariableType(VariableDef,i))#
    @variable(MyModel,Z)

    @objective(MyModel,Min, sum(sum(C[j,i]*X[i] for i = 1:n) + d[j] for j =1:nObj) )
    for i = 1:m
         @constraint(MyModel,sum(A[i,j]*X[j] for j = 1:n) >= b[i] )
    end
    for j = 1:nObj
         @constraint(MyModel,sum(C[j,i]*X[i] for i = 1:n) + d[j]  <= Y_bar1[j] + EPS)
    end

    for j = 1:nObj
         @constraint(MyModel,sum(C[j,i]*X[i] for i = 1:n) + d[j] >= YR[j]  )
    end
    return MyModel;

end


text:
module LatentHeat

# If you want to add a new method here, feel free to do so.
# Remember to also export the function name in GeoParams.jl (in addition to here)

using Parameters, LaTeXStrings, Unitful
using ..Units
using GeoParams: AbstractMaterialParam
using ..MaterialParameters: MaterialParamsInfo
import Base.show, GeoParams.param_info

abstract type AbstractLatentHeat{T} <: AbstractMaterialParam end

export  compute_latent_heat,                  # calculation routines
        compute_latent_heat!,
        param_info,
        ConstantLatentHeat                  # constant

include("../Computations.jl")
include("../Utils.jl")

# Constant  -------------------------------------------------------
"""
    ConstantLatentHeat(Q_L=400kJ/kg)

Set a constant latent heat:
```math
    Q_L  = cst
```
where ``Q_L`` is the latent heat [``kJ/kg``].
"""
@with_kw_noshow struct ConstantLatentHeat{T,U} <: AbstractLatentHeat{T}
    Q_L::GeoUnit{T,U}         =   400kJ/kg                # Latent heat
end
ConstantLatentHeat(args...) = ConstantLatentHeat(convert.(GeoUnit,args)...)

function param_info(s::ConstantLatentHeat) # info about the struct
    return MaterialParamsInfo(Equation = L"Q_L = cst")
end

# Calculation routine
function  (s::ConstantLatentHeat{_T})(;kwargs...) where _T
    @unpack_val Q_L   = s

    return Q_L
end

compute_latent_heat(s::ConstantLatentHeat{_T}; kwargs...) where _T = s()

function (s::ConstantLatentHeat{_T})(I::Integer...) where _T
    @unpack_val Q_L   = s

    return fill(Q_L, I...)
end

function compute_latent_heat!(Hl::AbstractArray{_T,N}, s::ConstantLatentHeat{_T}; kwargs...) where {_T,N}
    @unpack_val Q_L   = s
    for i in eachindex(Hl)
        @inbounds Hl[i] = Q_L
    end
    return nothing
end

# Print info
function show(io::IO, g::ConstantLatentHeat)
    print(io, "Constant latent heat: Q_L=$(Value(g.Q_L))")
end
#-------------------------------------------------------------------------

# Help info for the calculation routines
"""
    Ql = compute_latent_heat(s:<AbstractLatentHeat)

Returns the latent heat `Q_L`

"""
#compute_latent_heat()

# Computational routines needed for computations with the MaterialParams structure
function compute_latent_heat(s::AbstractMaterialParamsStruct, args)
    if isempty(s.LatentHeat)
        return isempty(args) ? 0.0 : zero(typeof(args).types[1])  # return zero if not specified
    else
        return s.LatentHeat[1](args)
    end
end

# add methods programatically
for myType in (:ConstantLatentHeat,)
@eval begin
(s::$(myType))(args)= s(; args...)
compute_latent_heat(s::$(myType), args) = s(args)
compute_latent_heat!(H::AbstractArray{_T,N}, s::$(myType){_T}, args) where {_T,N} = compute_latent_heat!(H, s; args...)
end
end

compute_latent_heat(args...)  = compute_param(compute_latent_heat, args...)
compute_latent_heat!(args...) = compute_param!(compute_latent_heat, args...)

end


text:
#=
Convert an object with R "phylo" structure to a julia Tree object
=#

function ape2Tree(phylo)
    branches = Array{Int64,2}(phylo[:edge])
    branch_lengths = phylo[:edge_length]
    attribute = Array{Float64,1}(ones(length(branch_lengths)...))
    tip_labels = phylo[:tip_label]
    return build_tree(branches, branch_lengths, [attribute], tip_labels, root_attributes=[0.])
end

#=

=#

function Tree2ape(tree::Tree ; id = 1)
    n_nodes = tree.n_nodes
    node_id = (n_nodes + 3)/2

    function aux(subtree, root, next_node, tip)
        if subtree.n_nodes == 0
            return Array{Int64,2}(undef,0,2), Array{Float64,1}(undef,0), Array{Float64,1}(undef,0), next_node, tip, Array{String,1}(undef,0)
        elseif length(subtree.offsprings) == 0
            return [root tip], subtree.branch_length, subtree.attributes[id], next_node, tip+1, subtree.label
        else
            tupple1 = aux(subtree.offsprings[1], next_node, next_node+1, tip)
            tupple2 = aux(subtree.offsprings[2], next_node, tupple1[4], tupple1[5])
            if subtree.branch_length == 0.
                return [tupple1[1]; tupple2[1]],
                    [tupple1[2]; tupple2[2]],
                    [tupple1[3]; tupple2[3]],
                    tupple2[4], tupple2[5],
                    [tupple1[6]; tupple2[6]]
            else
                return [[root next_node]; tupple1[1]; tupple2[1]],      # edges
                    [subtree.branch_length; tupple1[2]; tupple2[2]],    # branch lengths
                    [subtree.attributes[id]; tupple1[3]; tupple2[3]],   # rates
                    tupple2[4], tupple2[5],                             # auxiliary variables
                    [tupple1[6]; tupple2[6]]                            # tip labels
            end
        end
    end

    tupple = aux(tree, -1, node_id, 1)
    return tupple[1], tupple[2], tupple[3], tupple[6]
end

#=
Buil a Tree object from a list of branches relationships and lengths, with branch rates
=#

function build_tree(branches::Array{Array{T1,1}}, branch_lengths::Array{Float64,1}, attributes::Array{Array{T2,1},1} ;
    extinct = [], prune_extinct=true, return_void = false, stem_age=0., root_attributes = Array{T2,1}(undef,0)) where {T1<:Number, T2<:Number}

    n_attributes = length(attributes)
    if length(extinct) == 0
        dead_branches = fill(false, size(branches)[1])
    else
        dead_branches = extinct
    end

    if size(branches)[1] == 0
        if return_void
            return Tree()
        else
            return Tree(Array{Tree,1}(undef,0),stem_age,root_attributes)
        end
    end
    if size(branches)[1] == 1
        if ! dead_branches[1]
            return Tree(Array{Tree,1}(undef,0),branch_lengths[1],root_attributes, true)
        elseif prune_extinct
            return Tree()
        else
            return Tree(Array{Tree,1}(undef,0),branch_lengths[1],root_attributes, false)
        end
    end

    root = -1
    for b1 in branches
        has_no_parent = true
        for b2 in branches
            if b1[1] == b2[2]
                has_no_parent = false
                break
            end
        end
        if has_no_parent
            root = b1[1]
            break
        end
    end

    max_node = -1
    for b in branches
        max_node = max(b[1], b[2], max_node)
    end
    max_node += 1
    offsprings = fill(-1,max_node,2)
    parent_edges = fill(-1,max_node)
    daughter_edges = fill(-1,max_node)

    for i in 1:length(branches)
        individual = branches[i][1]+1
        parent_edges[individual] = i
        daughter_edges[branches[i][2]+1] = i
        if offsprings[individual,1] < 0
            offsprings[individual,1] = branches[i][2]
        else
            offsprings[individual,2] = branches[i][2]
        end

        if dead_branches[i]
            individual = branches[i][2]+1
            offsprings[individual,:]= [-2 -2]
        end
    end

    if prune_extinct
        function aux(node)
            if offsprings[node+1,1]==-1
                return Tree(Array{Tree,1}(undef,0),branch_lengths[daughter_edges[node + 1]], [attributes[p][daughter_edges[node + 1]] for p in 1:n_attributes])
            elseif offsprings[node+1,1]==-2
                return Tree()
            else
                if node != root
                    add_time = branch_lengths[daughter_edges[node + 1]]
                    add_attribute = [attributes[p][daughter_edges[node + 1]] for p in 1:n_attributes]
                else
                    add_time = stem_age
                    add_attribute = root_attributes #repeat([Array{Float64}(undef,0)],n_attributes)
                end

                tree_left = aux(offsprings[node+1,1])
                tree_right = aux(offsprings[node+1,2])

                if tree_left.n_nodes + tree_right.n_nodes == 0
                    return Tree()
                elseif tree_left.n_nodes == 0
                    tree_right = Tree(tree_right.offsprings, tree_right.branch_length + add_time,
                        add_attribute, tree_right.n_nodes)

                    return tree_right
                elseif tree_right.n_nodes == 0
                    tree_left = Tree(tree_left.offsprings, tree_left.branch_length + add_time,
                        add_attribute, tree_left.n_nodes)
                    return tree_left
                else
                    return Tree([tree_left, tree_right], add_time, add_attribute)
                end
            end
        end

        return aux(root)

    else
        function aux_all(node)
            if offsprings[node+1,1]==-1
                return Tree(Array{Tree,1}(undef,0),branch_lengths[daughter_edges[node + 1]], [attributes[p][daughter_edges[node + 1]] for p in 1:n_attributes], true)
            elseif offsprings[node+1,1]==-2
                return Tree(Array{Tree,1}(undef,0),branch_lengths[daughter_edges[node + 1]], [attributes[p][daughter_edges[node + 1]] for p in 1:n_attributes], false)
            else
                if node != root
                    add_time = branch_lengths[daughter_edges[node + 1]]
                    add_attribute = [attributes[p][daughter_edges[node + 1]] for p in 1:n_attributes]
                else
                    add_time = stem_age
                    add_attribute = root_attributes#repeat([Array{Float64}(undef,0)],n_attributes)
                end
                tree_left = aux_all(offsprings[node+1,1])
                tree_right = aux_all(offsprings[node+1,2])
                return Tree([tree_left, tree_right], add_time, add_attribute)
            end
        end

        return aux_all(root)

    end

end

function build_tree(branches::Array{Array{T1,1}}, branch_lengths::Array{Float64,1}, attributes::Array{T2,1} ;
    extinct = [], prune_extinct=true, return_void = false, stem_age=0., root_attributes = Array{T2,1}(undef,0)) where {T1<:Number, T2<:Number}
    build_tree(branches, branch_lengths, [attributes], extinct=extinct, prune_extinct=prune_extinct,
        return_void = return_void, stem_age = stem_age, root_attributes= root_attributes)
end

max_stars_repo_name: hmorlon/JPANDA

max_issues_repo_name: hmorlon/JPANDA

max_forks_repo_name: hmorlon/JPANDA


text:
using PkgBenchmark
using BenchmarkTools
using Test
using Dates
using LibGit2
using Random

const BENCHMARK_DIR = joinpath(@__DIR__, "..", "benchmark")

function temp_pkg_dir(fn::Function; tmp_dir=joinpath(tempdir(), randstring()),
        remove_tmp_dir::Bool=true, initialize::Bool=true)
    # Used in tests below to set up and tear down a sandboxed package directory
    withenv("JULIA_PKGDIR" => tmp_dir) do
        @test !isdir(Pkg.dir())
        try
            if initialize
                Pkg.init()
                @test isdir(Pkg.dir())
                Pkg.resolve()
            else
                mkpath(Pkg.dir())
            end
            fn()
        finally
            remove_tmp_dir && try rm(tmp_dir, recursive=true) end
        end
    end
end

function test_structure(g)
    @test g |> keys |> collect |> Set == ["utf8", "trigonometry"] |> Set
    @test g["utf8"] |> keys |> collect |> Set == ["join", "plots", "replace"] |> Set
    @test g["utf8"]["plots"] |> keys |> collect == ["fnplot"]

    _keys = Set(vec([string((string(f), x)) for x in (0.0, pi), f in (sin, cos, tan)]))
    @test g["trigonometry"]["circular"] |> keys |> collect |> Set == _keys
end

@testset "structure" begin
    results = benchmarkpkg("PkgBenchmark")
    test_structure(PkgBenchmark.benchmarkgroup(results))
    @test PkgBenchmark.name(results) == "PkgBenchmark"
    @test Dates.Year(PkgBenchmark.date(results)) == Dates.Year(now())
    export_markdown(stdout, results)
end

const TEST_PACKAGE_NAME = "Example"

# Set up a test package in a temp folder that we use to test things on
tmp_dir = joinpath(tempdir(), randstring())
old_pkgdir = Pkg.dir()

temp_pkg_dir(;tmp_dir = tmp_dir) do
    test_sig = LibGit2.Signature("TEST", "TEST@TEST.COM", round(time(); digits=0), 0)
    Pkg.add(TEST_PACKAGE_NAME)

    @testset "benchmarkconfig" begin
        PkgBenchmark._withtemp(tempname()) do f
            str = """
            using BenchmarkTools
            using Base.Test
            SUITE = BenchmarkGroup()
            SUITE["foo"] = @benchmarkable 1+1

            @test Base.JLOptions().opt_level == 3
            @test ENV["JL_PKGBENCHMARK_TEST_ENV"] == "10"
            """
            open(f, "w") do file
                print(file, str)
            end

            config = BenchmarkConfig(juliacmd = `$(joinpath(Sys.BINDIR, Base.julia_exename())) -O3`,
                                    env = Dict("JL_PKGBENCHMARK_TEST_ENV" => 10))
            @test typeof(benchmarkpkg(TEST_PACKAGE_NAME, config, script=f; custom_loadpath=old_pkgdir)) == BenchmarkResults
            end
        end

    # Make a commit with a small benchmarks.jl file
    testpkg_path = Pkg.dir(TEST_PACKAGE_NAME)
    repo = LibGit2.GitRepo(testpkg_path)
    LibGit2.branch!(repo, "master")

    mkpath(joinpath(testpkg_path, "benchmark"))

    # Make a small example benchmark file
    open(joinpath(testpkg_path, "benchmark", "benchmarks.jl"), "w") do f
        print(f,
        """
            using BenchmarkTools
            SUITE = BenchmarkGroup()
            SUITE["trig"] = BenchmarkGroup()
            SUITE["trig"]["sin"] = @benchmarkable sin(2.0)
        """)
    end

    LibGit2.add!(repo, "benchmark/benchmarks.jl")
    commit_master = LibGit2.commit(repo, "test"; author=test_sig, committer=test_sig)

    @testset "getting back original commit / branch" begin
        # Test we are on a branch and run benchmark on a commit that we end up back on the branch
        LibGit2.branch!(repo, "PR")
        touch(joinpath(testpkg_path, "foo"))
        LibGit2.add!(repo, "foo")
        commit_PR = LibGit2.commit(repo, "PR commit"; author=test_sig, committer=test_sig)
        LibGit2.branch!(repo, "master")
        PkgBenchmark.benchmarkpkg(TEST_PACKAGE_NAME, "PR"; custom_loadpath=old_pkgdir)
        @test LibGit2.branch(repo) == "master"

        # Test we are on a commit and run benchmark on another commit and end up on the commit
        LibGit2.checkout!(repo, string(commit_master))
        PkgBenchmark.benchmarkpkg(TEST_PACKAGE_NAME, "PR"; custom_loadpath=old_pkgdir)
        @test LibGit2.revparseid(repo, "HEAD") == commit_master
    end

    tmp = tempname() * ".json"

    # Benchmark dirty repo
    cp(joinpath(@__DIR__, "..", "benchmark", "benchmarks.jl"), joinpath(testpkg_path, "benchmark", "benchmarks.jl"); force=true)
    cp(joinpath(@__DIR__, "..", "benchmark", "REQUIRE"), joinpath(testpkg_path, "benchmark", "REQUIRE"))
    LibGit2.add!(repo, "benchmark/benchmarks.jl")
    LibGit2.add!(repo, "benchmark/REQUIRE")
    @test LibGit2.isdirty(repo)
    @test_throws ErrorException PkgBenchmark.benchmarkpkg(TEST_PACKAGE_NAME, "HEAD"; custom_loadpath=old_pkgdir)
    results = PkgBenchmark.benchmarkpkg(TEST_PACKAGE_NAME; custom_loadpath=old_pkgdir, resultfile=tmp)
    test_structure(PkgBenchmark.benchmarkgroup(results))
    @test isfile(tmp)
    rm(tmp)

    # Commit and benchmark non dirty repo
    commitid = LibGit2.commit(repo, "commiting full benchmarks and REQUIRE"; author=test_sig, committer=test_sig)
    @test !LibGit2.isdirty(repo)
    results = PkgBenchmark.benchmarkpkg(TEST_PACKAGE_NAME, "HEAD"; custom_loadpath=old_pkgdir, resultfile=tmp)
    @test PkgBenchmark.commit(results) == string(commitid)
    @test PkgBenchmark.juliacommit(results) == Base.GIT_VERSION_INFO.commit
    test_structure(PkgBenchmark.benchmarkgroup(results))
    @test isfile(tmp)
    r = readresults(tmp)
    @test r.benchmarkgroup == results.benchmarkgroup
    @test r.commit == results.commit
    rm(tmp)

    # Make a dummy commit and test comparing HEAD and HEAD~
    touch(joinpath(testpkg_path, "dummy"))
    LibGit2.add!(repo, "dummy")
    LibGit2.commit(repo, "dummy commit"; author=test_sig, committer=test_sig)

    @testset "judging" begin
        judgement = judge(TEST_PACKAGE_NAME, "HEAD~", "HEAD", custom_loadpath=old_pkgdir)
        test_structure(PkgBenchmark.benchmarkgroup(judgement))
        export_markdown(stdout, judgement)
        judgement = judge(TEST_PACKAGE_NAME, "HEAD", custom_loadpath=old_pkgdir)
        test_structure(PkgBenchmark.benchmarkgroup(judgement))
    end
end


text:
# This file is a part of Julia. License is MIT: https://julialang.org/license

using Test, Random
using SuiteSparse, LinearAlgebra, SparseArrays

if Base.USE_GPL_LIBS
    include("umfpack.jl")
    include("cholmod.jl")
    include("spqr.jl")

    # Test multithreaded execution
    let p, cmd = `$(Base.julia_cmd()) --depwarn=error --startup-file=no threads.jl`
        # test both nthreads==1 and nthreads>1. spawn a process to test whichever
        # case we are not running currently.
        other_nthreads = Threads.nthreads() == 1 ? 4 : 1
        p = run(
                pipeline(
                    setenv(
                        cmd,
                        "JULIA_NUM_THREADS" => other_nthreads,
                        dir=@__DIR__()),
                    stdout = stdout,
                    stderr = stderr),
                wait = false)
        include("threads.jl")
        if !success(p)
            error("SuiteSparse threads test failed with nthreads == $other_nthreads")
        end
    end
end

max_stars_repo_name: TimoLarson/julia

max_issues_repo_name: TimoLarson/julia

max_forks_repo_name: TimoLarson/julia


text:
"""
$(TYPEDEF)

An abstract type the subtypes of which are associated with functions method
overloaded for for new extensions. An instance of the `DefaultExt <:ExtensionType`
structure to the `Optimizer` in the `ext_type` field.
"""
abstract type ExtensionType end
struct DefaultExt <: ExtensionType end
MOIU.map_indices(::Function, x::ExtensionType) = x
MOIU.map_indices(::Function, x::DefaultExt) = x

max_stars_repo_name: MatthewStuber/EAGO

max_issues_repo_name: MatthewStuber/EAGO

max_forks_repo_name: MatthewStuber/EAGO


text:
export FOS,first,last,getindex,setindex!
abstract type AbstractFOS end
"""
    Fermionic Operators String is a String of Fermionic Operator with a series undetermined employed number.
"""
struct FOS <: AbstractFOS
    fos::Array{PFO,1}
end

FOS() = FOS([])

import Base.first,Base.last,Base.getindex,Base.setindex!,Base.length
first(fos::FOS) = first(fos.fos)
last(fos::FOS) = last(fos.fos)
getindex(fos::FOS,key...) = getindex(fos.fos,key...)
setindex!(fos::FOS,value,key...) = setindex!(fos.fos,value,key...)
length(fos::FOS) = length(fos.fos)

import Base.display
export display
function display(fos::FOS)
    fos = fos.fos
    s = ""
    for i=1:length(fos)
        fo = fos[i].fo
        dag = fo.dag ? "^{\\dagger}" : ""
        s=s*"($(fo.symbol)_{$(fo.site)}^{$(fos[i].pow)})"*dag*" "
    end
    display(s)
end


text:
import Base:length, -

"""
    Tiling(ranges::NTuple{N, Tr}) where {N, Tr}

Using a tuple of `ranges` to simulate a tiling.

# Example
```julia
julia> t = Tiling((1:2:5, 10:5:20))
Tiling{2,StepRange{Int64,Int64}}((1:2:5, 10:5:20), [1 3; 2 4])

julia> Ju.encode(t, (2, 12))  # encode into an Int
1

julia> Ju.encode(t, (2, 18))

julia> t2 = t - (1, 3)  # shift a little to get a new Tiling
Tiling{2,StepRange{Int64,Int64}}((0:2:4, 7:5:17), [1 3; 2 4])3
```

See also: [`TilingsV`](@ref), [`TilingsQ`](@ref)
"""
struct Tiling{N, Tr<:AbstractRange}
    ranges::NTuple{N, Tr}
    inds::LinearIndices{N, NTuple{N, Base.OneTo{Int}}}
    Tiling(ranges::NTuple{N, Tr}) where {N, Tr} = new{N, Tr}(ranges, LinearIndices(Tuple(length(r) - 1 for r in ranges)))
end

"""
    (-)(t::Tiling, xs)

Shift `t` along each dimension by each element in `xs`.
"""
function (-)(t::Tiling, xs)
    Tiling(Tuple(r .- x for (r,x) in zip(t.ranges, xs)))
end

length(t::Tiling) = reduce(*, (length(r)-1 for r in t.ranges))

encode(range::AbstractRange, x) = floor(Int, div(x-range[1], step(range)) + 1)
encode(t::Tiling{1}, x::Number) = encode(t.ranges[1], x)
encode(t::Tiling{1}, xs) = encode(t.ranges[1], xs[1])
encode(t::Tiling{2}, xs) = t.inds[CartesianIndex(encode(t.ranges[1], xs[1]), encode(t.ranges[2], xs[2]))]
encode(t::Tiling{3}, xs) = t.inds[CartesianIndex(encode(t.ranges[1], xs[1]), encode(t.ranges[2], xs[2]), encode(t.ranges[3], xs[3]))]

"""
    TilingsV{Tt<:Tiling} <: AbstractVApproximator{Vector{Float64}}
    TilingsV(tilings::Vector{Tt}) where Tt<:Tiling

Using a vector of `tilings` to encode state. Each tiling has an independent weight.

See more details at Section (9.5.4) on Page 217 of the book *Sutton, Richard S., and Andrew G. Barto. Reinforcement learning: An introduction. MIT press, 2018.*

See also: [`Tiling`](@ref), [`TilingsQ`](@ref)
"""
struct TilingsV{Tt<:Tiling} <: AbstractVApproximator{Vector{Float64}}
    tilings::Vector{Tt}
    weights::Vector{Vector{Float64}}
    TilingsV(tilings::Vector{Tt}) where Tt<:Tiling = new{Tt}(tilings, [zeros(Float64, length(t)) for t in tilings])
end

function (ts::TilingsV)(s)
    v = 0.
    for (w, t) in zip(ts.weights, ts.tilings)
        v += w[encode(t, s)]
    end
    v
end

function update!(ts::TilingsV, s, e)
    for i in 1:length(ts.tilings)
        ts.weights[i][encode(ts.tilings[i], s)] += e
    end
end

"""
    TilingsQ{Tt<:Tiling} <: AbstractQApproximator{Vector{Float64}, Int}
    TilingsQ(tilings::Vector{Tt}, nactions) where Tt<:Tiling

The only difference compared to [`TilingsV`](@ref) is that now the weight of each tiling is a matrix.
"""
struct TilingsQ{Tt<:Tiling} <: AbstractQApproximator{Vector{Float64}, Int}
    tilings::Vector{Tt}
    weights::Vector{Array{Float64, 2}}
    TilingsQ(tilings::Vector{Tt}, nactions) where Tt<:Tiling = new{Tt}(tilings, [zeros(Float64, length(t), nactions) for t in tilings])
end

function (ts::TilingsQ)(s, a)
    v = 0.
    for (w, t) in zip(ts.weights, ts.tilings)
        v += w[encode(t, s), a]
    end
    v
end

function (ts::TilingsQ)(s, ::Val{:dist})
    dist = zeros(Float64, size(ts.weights[1],2))
    for (w,t) in zip(ts.weights, ts.tilings)
        dist .+= @view w[encode(t,s), :]
    end
    dist
end

(t::TilingsQ)(s) = t(s, Val(:dist))
(t::TilingsQ)(s, ::Val{:max}) = findmax(t(s))[1]
(t::TilingsQ)(s, ::Val{:argmax}) = sample(findallmax(t(s))[2])

function update!(ts::TilingsQ, s, a, e)
    for (w, t) in zip(ts.weights, ts.tilings)
        w[encode(t, s), a] += e
    end
end


text:
# f(x)=A'*X'
function FAtransposeXtranspose(A,X)
    if size(A)==(1,1)
        return (A[1].*X',nothing)
    elseif size(X)==(1,1)
        return (A'.*X[1],nothing)
    else
        return (A'*X',nothing)
    end
end

function FAtransposeXtranspose_inplace(value,auxvalue,A,X)
    if size(A)==(1,1)
        copy!(value,A[1]*X')
    elseif size(X)==(1,1)
        copy!(value,A'*X[1])
    else
        BLAS.gemm!('T','T',1.0,A,X,0.0,value)
    end
end

function DAtransposeXtranspose(derivativeIDX,f_c,faux_c,grad_c,grad_n,A,X)
    if derivativeIDX==1
        if size(A)==(1,1)
            axpy!(1.0,[sum(X'.*grad_c)],grad_n)
        elseif size(X)==(1,1)
            axpy!(X[1],grad_c',grad_n)
        else
            BLAS.gemm!('T','T',1.0,X,grad_c,1.0,grad_n) # X'*grad_c'
        end
    elseif derivativeIDX==2
        if size(A)==(1,1)
            axpy!(A[1],grad_c,grad_n)
        elseif size(X)==(1,1)
            axpy!(1.0,[sum(A'.*grad_c)],grad_n)
        else
            BLAS.gemm!('T','T',1.0,grad_c,A,1.0,grad_n) # grad_c'*A'
        end
    end
end

if PROC=="GPU"
#    function FAtransposeXtranpose(A::CudaArray,X::CudaArray)
#        if size(A)==(1,1)
#            tmp=CudaArray(Float64,size(X)); gfill!(tmp,A)
#            return (CUBLAS.gemm('N','N',tmp,X),nothing)
#        elseif size(X)==(1,1)
#            tmp=CudaArray(Float64,size(A)); gfill!(tmp,X)
#            return (CUBLAS.gemm('N','N',A,tmp),nothing)
#        else
#            return (CUBLAS.gemm('N','N',A,X),[])
#        end
#    end

    function FAtransposeXtranspose_inplace(value::CudaArray,auxvalue,A::CudaArray,X::CudaArray)
        if size(A)==(1,1)
            CUBLAS.geam!('T','N',1.0,X,0.0,value,value) # nb argument converse of Base.scale!
            scale!(A,value) # nb argument converse of Base.scale!
        elseif size(X)==(1,1)
            CUBLAS.geam!('T','N',1.0,A,0.0,value,value) # nb argument converse of Base.scale!
            scale!(X,value) # nb argument converse of Base.scale!
        else
            CUBLAS.gemm!('T','T',1.0,A,X,0.0,value)
        end
    end

    function DAtransposeXtranspose(derivativeIDX,f_c,faux_c,grad_c,grad_n,A::CudaArray,X::CudaArray)
        if derivativeIDX==1
            if size(A)==(1,1)
                tmp=CudaArray(Float64,size(X))
                CUBLAS.geam!('T','N',1.0,X,0.0,tmp,tmp) # tmp=X'
                vmult!(1.0,tmp,grad_c,tmp) # tmp=X'.*grad_c
                sum_update!(1.0,tmp,1.0,grad_n) # grad_n+=sum(X'.*grad_c)
                free(tmp)
            elseif size(X)==(1,1)
                alphaaxpy!(1.0,X,grad_c,grad_n)
            else
                CUBLAS.gemm!('T','T',1.0,X,grad_c,1.0,grad_n) # X'*grad_c'
            end
        elseif derivativeIDX==2
            if size(A)==(1,1)
                tmp=CudaArray(Float64,size(grad_c))
                CUBLAS.geam!('T','N',1.0,grad_c,0.0,tmp,tmp) # tmp=grad_c'
                alphaaxpy!(1.0,A,tmp,grad_n); free(tmp)
            elseif size(X)==(1,1)
                tmp=CudaArray(Float64,size(A))
                vmult!(1.0,A,grad_c,tmp)
                sum_update!(1.0,tmp,1.0,grad_n);free(tmp)

            else
                CUBLAS.gemm!('T','T',1.0,grad_c,A,1.0,grad_n) # grad_c'*A'
            end
        end
    end
end

Derivative[FAtransposeXtranspose]=DAtransposeXtranspose
export FAtransposeXtranspose

Inplace[FAtransposeXtranspose]=FAtransposeXtranspose_inplace

max_stars_repo_name: davidbarber/AutoDiff

max_issues_repo_name: davidbarber/AutoDiff

max_forks_repo_name: davidbarber/AutoDiff


text:
"""
    eom_chuacircuits_driven(u, p, t) -> SVector{6}

Equations of motion for a bidirectionally coupled system consisting of two
driven Chua circuits [1].

# References
1. Murali, K., and M. Lakshmanan. "Chaotic dynamics of the driven
Chua's circuit." IEEE Transactions on Circuits and Systems I Fundamental
Theory and Applications 40.11 (1993): 836-840.
"""
function eom_chuacircuits_driven(u, p, t)
    α₁, α₂, β₁, β₂, F₁, F₂, ω₁, ω₂, ϵ₁, ϵ₂, m₀, m₁, σ = (p...,)
    x₁, y₁, z₁ = (u[1:3]...,)
    x₂, y₂, z₂ = (u[4:6]...,)

    # Dynamical noise
    if σ == 0
        ξ = 0
    else
        ξ = rand(Normal(0, σ))
    end

    hx₁ = m₁*x₁ + 0.5*(m₀ - m₁)*(abs(x₁+1) - abs(x₁-1))
    hx₂ = m₁*x₂ + 0.5*(m₀ - m₁)*(abs(x₂+1) - abs(x₂-1))

    dx₁ = α₁*(y₁-hx₁) - α₁*ϵ₂*(x₁ - x₂)
    dy₁ = x₁-y₁+z₁
    dz₁ = -β₁*y₁ + F₁*sin(ω₁*t) + ξ

    dx₂ = α₂*(y₂-hx₂) - α₂*ϵ₁*(x₁ - x₂)
    dy₂ = x₂-y₂+z₂
    dz₂ = -β₂*y₂ + F₂*sin(ω₂*t) + ξ
    SVector{6}(dx₁, dy₁, dz₁, dx₂, dy₂, dz₂)
end

"""
    chuacircuits_driven(u₀, α₁, α₂, β₁, β₂, F₁, F₂,
                ω₁, ω₂, ϵ₁, ϵ₂, m₀, m₁, σ) -> ContinuousDynamicalSystem

Initialize a bidirectionally coupled system consisting of two driven Chua
circuits [1], X₁ and X₂. The subsystems are mutually coupled by a linear
resistor, where `ϵ₁` controls the influence of X₁ on X₂, and `ϵ₂` controls the
influence of X₂ on X₁. The parameters for the subsystems are
set equal to each other, as in the original paper, but can be tuned
individually for each subsystem.

# References
1. Murali, K., and M. Lakshmanan. "Chaotic dynamics of the driven Chua's
circuit." IEEE Transactions on Circuits and Systems I Fundamental
Theory and Applications 40.11 (1993): 836-840.
"""
function chuacircuits_driven(u₀, α₁, α₂, β₁, β₂, F₁, F₂,
                                      ω₁, ω₂, ϵ₁, ϵ₂, m₀, m₁, σ)
    p = [α₁, α₂, β₁, β₂, F₁, F₂, ω₁, ω₂, ϵ₁, ϵ₂, m₀, m₁, σ]
    ContinuousDynamicalSystem(eom_chuacircuits_driven, u₀, p)
end

chuacircuits_driven(;u₀ = [0.1, 0.1, 0.2, 0.15, 0.15, 0.22],
                            α₁ = 7.0, α₂ = 7.0,
                            β₁ = 14.286, β₂ = 14.286,
                            F₁ = 1.5, F₂ = 1.5,
                            ω₁ = 3.0, ω₂ = 3.0,
                            σ = 0.1,
                            ϵ₁ = 0.1, ϵ₂ = 0.1,
                            m₀ = -1/7, m₁ = 2/7) =
    chuacircuits_driven(u₀, α₁, α₂, β₁, β₂, F₁, F₂, ω₁, ω₂, ϵ₁, ϵ₂, m₀, m₁, σ)


text:
export
Message,
MessageUpdateRule,
ScheduleEntry,
Schedule

"""Encodes a message, which is a probability distribution with a scaling factor"""
struct Message{family<:FactorFunction, var_type<:VariateType} # Note that parameter order is switched w.r.t. ProbabilityDistribution, for ease of overloading
    dist::ProbabilityDistribution{var_type, family}
    scaling_factor::Any

    Message{F, V}(dist::ProbabilityDistribution{V, F}) where {F, V}= new(dist) # Constructor for unspecified scaling factor
end

Message(dist::ProbabilityDistribution{V, F}) where {F<:FactorFunction, V<:VariateType} = Message{F, V}(dist)

Message(var_type::Type{V}, family::Type{F}; kwargs...) where {F<:FactorFunction, V<:VariateType} = Message{family, var_type}(ProbabilityDistribution(var_type, family; kwargs...))

function Message(family::Type{F}; kwargs...) where F
    dist = ProbabilityDistribution(family; kwargs...)
    var_type = variateType(dist)

    return Message{family, var_type}(dist)
end

family(msg_type::Type{Message{F}}) where F<:FactorFunction = F

function show(io::IO, msg::Message)
    if isdefined(msg, :scaling_factor)
        println(io, "Message: 1/$(format(msg.scaling_factor)) * $(format(msg.dist))")
    else
        println(io, "Message: $(format(msg.dist))")
    end
end

"""Special inheritance rules for parametric Message types"""
matches(::Type{T}, ::Type{T}) where T<:Message = true
matches(::Type{Message{F, V}}, ::Type{Message}) where {F<:FactorFunction, V<:VariateType} = true
matches(::Type{Message{F}}, ::Type{Message}) where {F<:FactorFunction} = true
matches(Ta::Type{Message{Fa, Va}}, Tb::Type{Message{Fb, Vb}}) where {Fa<:FactorFunction, Fb<:FactorFunction, Va<:VariateType, Vb<:VariateType} = (Va==Vb) && (Fa<:Fb)
matches(Ta::Type{Message{Fa, Va}}, Tb::Type{Message{Fb}}) where {Fa<:FactorFunction, Fb<:FactorFunction, Va<:VariateType} = (Fa<:Fb)
matches(Ta::Type{Message{Fa}}, Tb::Type{Message{Fb}}) where {Fa<:FactorFunction, Fb<:FactorFunction} = (Fa<:Fb)
matches(::Type{Nothing}, ::Type{T}) where T<:Message = false
matches(::Type{P}, ::Type{M}) where {P<:ProbabilityDistribution, M<:Message} = false
matches(::Type{M}, ::Type{P}) where {P<:ProbabilityDistribution, M<:Message} = false

function ==(t::Message{fam_t, var_t}, u::Message{fam_u, var_u}) where {fam_t<:FactorFunction, var_t<:VariateType, fam_u<:FactorFunction, var_u<:VariateType}
    (fam_t == fam_u) || return false
    (var_t == var_u) || return false
    (t.dist == u.dist) || return false
    if isdefined(t, :scaling_factor) && isdefined(u, :scaling_factor)
        (t.scaling_factor == u.scaling_factor) || return false
    end
    isdefined(t, :scaling_factor) && !isdefined(u, :scaling_factor) && return false
    !isdefined(t, :scaling_factor) && isdefined(u, :scaling_factor) && return false

    return true
end

"""
A MessageUpdateRule specifies how a Message is calculated from the node function and the incoming messages.
Use `subtypes(MessageUpdateRule)` to list the available rules.
"""
abstract type MessageUpdateRule end

"""
A `ScheduleEntry` defines a message computation.
The `message_update_rule <: MessageUpdateRule` defines the rule that is used
to calculate the message coming out of `interface`.
"""
mutable struct ScheduleEntry
    interface::Interface
    message_update_rule::Type
    internal_schedule::Vector{ScheduleEntry}

    # Fields for algorithm assembly
    schedule_index::Int # Specify the position in the schedule
    initialize::Bool # Indicate whether this mesage is initialized before schedule execution
    family::Type # Specify the outbound family for this entry
    dimensionality::Tuple # Specify the outbound dimensionality
    inbounds::Vector{Any} # Specify the inbounds required for computing the message update

    ScheduleEntry() = new()
    ScheduleEntry(interface::Interface, message_update_rule::Type) = new(interface, message_update_rule)
end

function show(io::IO, entry::ScheduleEntry)
    rule_str = replace(string(entry.message_update_rule), "ForneyLab." => "") # Remove "Forneylab."
    internal_schedule = isdefined(entry, :internal_schedule) ? "(INTERNAL SCHEDULE) " : ""
    print(io, "$(internal_schedule)$(rule_str) on $(entry.interface)")
end

function ==(a::ScheduleEntry, b::ScheduleEntry)
    return (a.interface == b.interface) && (a.message_update_rule == b.message_update_rule)
end

"""
A `Schedule` defines the update order for message computations.
"""
const Schedule = Vector{ScheduleEntry}

function show(io::IO, schedule::Schedule)
    condensed_schedule = condense(schedule)
    idx = 1
    for entry in schedule
        if entry in condensed_schedule
            print(io, "$idx.\t$(entry)")
            idx += 1
        else
            print(io, "\t$(entry)")
        end
    end
end

"""
summaryDependencyGraph(edgeset)

Returns a DependencyGraph (directed graph) that encodes the dependencies among
summary messages (such as sum-product messages) in `edgeset`.
All Interfaces in `edgeset` are vertices in the dependency graph.
The dependency graph can be used for loop detection, scheduling, etc.
"""
function summaryDependencyGraph(edgeset::Set{Edge}; reverse_edges=false)
    # Create dependency graph object
    dg = DependencyGraph{Interface}()

    # Add all Interfaces in edgeset as vertices in dg
    for edge in sort(collect(edgeset))
        isa(edge.a, Interface) && addVertex!(dg, edge.a)
        isa(edge.b, Interface) && addVertex!(dg, edge.b)
    end

    # Add all summary dependencies
    for interface in dg.vertices
        if isa(interface.partner, Interface) # interface is connected to an Edge
            for node_interface in interface.partner.node.interfaces
                (node_interface === interface.partner) && continue
                (node_interface.edge in edgeset) || continue
                if reverse_edges
                    addEdge!(dg, interface, node_interface)
                else
                    addEdge!(dg, node_interface, interface)
                end
            end
        end
    end

    return dg
end

"""
`summaryPropagationSchedule(variables, clusters)` builds a generic
`Schedule` for calculating the marginal distributions of every variable in
`variables` and every cluster in `clusters`. The message update rule in each
schedule entry is set to `Nothing`.
"""
function summaryPropagationSchedule(variables::Vector{Variable}, clusters::Vector{Cluster}=Cluster[]; limit_set=edges(current_graph), target_sites=Interface[], breaker_sites=Interface[])
    # We require the marginal distribution of every variable in variables.
    # If a variable relates to multiple edges, this indicates an equality constraint.
    # Therefore, we only need to consider one arbitrary edge to calculate the marginal.
    for variable in variables
        edge = first(variable.edges) # For the sake of consistency, we always take the first edge.
        (edge.a != nothing) && !isa(edge.a.node, Terminal) && !isa(edge.a.node, Clamp) && push!(target_sites, edge.a)
        (edge.b != nothing) && !isa(edge.b.node, Terminal) && !isa(edge.b.node, Clamp) && push!(target_sites, edge.b)
    end

    # A cluster indicates the computation of a joint marginal belief.
    # In this case, all incoming messages on internal edges are required.
    for cluster in clusters
        for iface in cluster.node.interfaces
            partner = ultimatePartner(iface)
            (iface.edge in limit_set) && (partner != nothing) && !isa(partner.node, Terminal) && !isa(partner.node, Clamp) && push!(target_sites, partner)
        end
    end

    # Determine a feasible ordering of message updates
    dg = summaryDependencyGraph(limit_set)
    iface_list = children(unique(target_sites), dg, breaker_sites=Set(breaker_sites))
    # Build a schedule; Nothing indicates an unspecified message update rule
    schedule = [ScheduleEntry(iface, Nothing) for iface in iface_list]

    return schedule
end

"""
inferUpdateRules!(schedule) infers specific message update rules for all schedule entries.
"""
function inferUpdateRules!(schedule::Schedule; inferred_outbound_types=Dict{Interface, Type}())
    for entry in schedule
        (entry.message_update_rule == Nothing) && error("No msg update rule type specified for $(entry)")
        if !isconcretetype(entry.message_update_rule)
            # In this case entry.message_update_rule is a update rule type, but not a specific rule.
            # Here we infer the specific rule that is applicable, which should be a subtype of entry.message_update_rule.
            inferUpdateRule!(entry, entry.message_update_rule, inferred_outbound_types)
        end
        # Store the rule's outbound type
        inferred_outbound_types[entry.interface] = outboundType(entry.message_update_rule)
    end

    return schedule
end

"""
Flatten a schedule by inlining all internal message passing schedules.
This yields a simple, linear schedule.
"""
function flatten(schedule::Schedule)
    # Check if there are any internal schedules to inline
    if !any([isdefined(entry, :internal_schedule) for entry in schedule])
        return schedule
    end

    flat_schedule = ScheduleEntry[]
    passed_interfaces = Dict{Interface,Bool}()
    for entry in schedule
        if isdefined(entry, :internal_schedule)
            flat_subschedule = flatten(entry.internal_schedule)
            for internal_entry in flat_subschedule
                if !haskey(passed_interfaces, internal_entry.interface)
                    push!(flat_schedule, internal_entry)
                    passed_interfaces[internal_entry.interface] = true
                end
            end
        else
            if !haskey(passed_interfaces, entry.interface)
                push!(flat_schedule, entry)
                passed_interfaces[entry.interface] = true
            end
        end
    end

    return flat_schedule
end

"""
Contruct a condensed schedule.
"""
function condense(schedule::Schedule)
    condensed_schedule = ScheduleEntry[]
    for entry in schedule
        if isdefined(entry, :internal_schedule)
            entry.internal_schedule = condense(entry.internal_schedule)
        elseif !isa(entry.interface.node, Clamp) # Only actual Clamps should be consensed
            push!(condensed_schedule, entry)
        end
    end

    return condensed_schedule
end

"""
Generate a mapping from interface to schedule entry.
Multiple interfaces can map to the same schedule entry if the
graph contains composite nodes.
"""
function interfaceToScheduleEntry(schedule::Schedule)
    mapping = Dict{Interface, ScheduleEntry}()
    for entry in schedule
        interface = entry.interface
        mapping[interface] = entry
        while (interface.partner != nothing) && isa(interface.partner.node, Terminal)
            interface = interface.partner.node.outer_interface
            mapping[interface] = entry
        end
    end

    return mapping
end


text:
using Distributed: @everywhere, @distributed, addprocs, procs
if length(procs()) == 1 addprocs(3) end

@everywhere using Knet: @diff, Param, value, grad
@everywhere using Knet: sigm, tanh, softmax
# using Gadfly: plot
# using Statistics: norm

model_types = ["GRU","GSTM","LSTM"]

input_size  = 2
hiddens     = []
output_size = 1

hm_data    = 1_000
batch_size = 50
hm_test    = 200

seq_len = 200
# 100 # 200 # 500 # 1_000 # 2_000

hm_epochs = 50
lr        = .001

alpha_moments = 0.9
alpha_accugrads = 0.999

layer_test = []#[10], [12], [16], [20], [25], [28], [30], [32], [36], [40], [45], [48], [50], [52], [56], [58], [60], [62], [64], [68], [72], [74], [80], [84], [88], [92], [96], [100]]
hm_trials  = 1

verbose = true

mk_model(in,hiddens,out,type) =
begin
    model = []
    hm_layers = length(hiddens)+1
    if hm_layers != 1
        for i in 1:hm_layers
            if i == 1
                layer = type(in, hiddens[1])
            elseif i == hm_layers
                layer = type(hiddens[end],out)
            else
                layer = type(hiddens[i-1],hiddens[i])
            end
            push!(model, layer)
        end
    else
        push!(model, type(in,out))
    end
model
end

@everywhere prop_model(model, seq) =
begin
    outs = []
    for t in seq
        out = t
        for layer in model
            out = layer(out)
        end
        push!(outs, out)
    end
    for layer in model
        setfield!(layer, :state, zeros(1,size(getfield(layer, :wfi))[end]))
    end
outs
end

@everywhere mutable struct GRU
    wfi::Param
    wfs::Param
    wii::Param
    wki::Param
    wks::Param
    state
end

@everywhere GRU(in_size, layer_size) =
begin
    sq = sqrt(2/(in_size+layer_size))

    wfi = Param(2*sq .* randn(in_size,layer_size)    .-sq)
    wfs = Param(2*sq .* randn(layer_size,layer_size) .-sq)

    wki = Param(2*sq .* randn(in_size,layer_size)    .-sq)
    wks = Param(2*sq .* randn(layer_size,layer_size) .-sq)

    wii = Param(2*sq .* randn(in_size,layer_size)    .-sq)

    state = zeros(1,layer_size)
    layer = GRU(wfi, wfs, wii, wki, wks, state)
layer
end

@everywhere (layer::GRU)(in) =
begin
    focus  = sigm.(in * layer.wfi + layer.state * layer.wfs)
    keep   = sigm.(in * layer.wki + layer.state * layer.wks)
    interm = tanh.(in * layer.wii + layer.state .* focus)

    layer.state = keep .* layer.state + (1 .- keep) .* interm
layer.state
end

@everywhere mutable struct LSTM
    wki::Param
    wks::Param
    wfi::Param
    wfs::Param
    wsi::Param
    wss::Param
    wii::Param
    wis::Param
    state
end

@everywhere LSTM(in_size, layer_size) =
begin
    sq = sqrt(2/(in_size+layer_size))

    wki = Param(2*sq .* randn(in_size,layer_size)    .-sq)
    wks = Param(2*sq .* randn(layer_size,layer_size) .-sq)

    wfi = Param(2*sq .* randn(in_size,layer_size)    .-sq)
    wfs = Param(2*sq .* randn(layer_size,layer_size) .-sq)

    wsi = Param(2*sq .* randn(in_size,layer_size)    .-sq)
    wss = Param(2*sq .* randn(layer_size,layer_size) .-sq)

    wii = Param(2*sq .* randn(in_size,layer_size)    .-sq)
    wis = Param(2*sq .* randn(layer_size,layer_size) .-sq)

    state = zeros(1,layer_size)
    layer = LSTM(wki, wks, wfi, wfs, wsi, wss, wii, wis, state)
layer
end

@everywhere (layer::LSTM)(in) =
begin
    keep   = sigm.(in * layer.wki + layer.state * layer.wfs)
    forget = sigm.(in * layer.wfi + layer.state * layer.wks)
    show   = sigm.(in * layer.wsi + layer.state * layer.wks)
    interm = tanh.(in * layer.wii + layer.state * layer.wis)
    layer.state = forget .* layer.state + keep .* interm
    out = show .* tanh.(layer.state)
out
end

@everywhere mutable struct GSTM
    wfi::Param
    wfs::Param
    wri::Param
    wki::Param
    wks::Param
    wsi::Param
    wss::Param
    state
end

@everywhere GSTM(in_size, layer_size) =
begin
    sq = sqrt(2/(in_size+layer_size))

    wfi = Param(2*sq .* randn(in_size,layer_size)    .-sq)
    wfs = Param(2*sq .* randn(layer_size,layer_size) .-sq)

    wri = Param(2*sq .* randn(in_size,layer_size)    .-sq)

    wki = Param(2*sq .* randn(in_size,layer_size)    .-sq)
    wks = Param(2*sq .* randn(layer_size,layer_size) .-sq)

    wsi = Param(2*sq .* randn(in_size,layer_size)    .-sq)
    wss = Param(2*sq .* randn(layer_size,layer_size) .-sq)

    state = zeros(1,layer_size)
    layer = GSTM(wfi, wfs, wri, wki, wks, wsi, wss, state)
layer
end

@everywhere (layer::GSTM)(in) =
begin
    st = tanh.(layer.state)
    focus    = sigm.(in * layer.wfi + st * layer.wfs)
    reaction = tanh.(in * layer.wri + st .* focus)
    keep     = sigm.(in * layer.wki + st * layer.wks)
    show     = sigm.(in * layer.wsi + st * layer.wss)

    layer.state += reaction .* keep
    out = reaction .* show
out
end

grad_clip(x) =
if x < -1
    -1.0
elseif x > 1
    1.0
else
    x
end

main(model_name, data; hiddens=hiddens) =
begin
    model_type = (@eval $(Symbol(model_name)))

    model = mk_model(input_size,hiddens,output_size,model_type)

    moments = [zeros(size(getfield(layer,param))) for layer in model for param in fieldnames(typeof(layer))]
    accugrads = [zeros(size(getfield(layer,param))) for layer in model for param in fieldnames(typeof(layer))]

    losses = []

    test_sc1 = []
    test_sc2 = []

    for e in 1:hm_epochs

        loss = 0

        for batch in batchify(data, batch_size)

            results = @distributed (vcat) for seq in batch

                d = @diff begin

                    # input  = seq[1:end-1]
                    # label  = seq[2:end]

                    input = seq[1]
                    label = seq[2]

                    output = prop_model(model,input)[end]
                    sum(sum([(lbl-out).^2 for (lbl,out) in zip(label,output)]))

                end

                grads = []

                i = 0
                for layer in model
                    for param in fieldnames(model_type)
                        i +=1
                        g = grad(d,getfield(layer, param))
                        push!(grads, g)
                    end
                end

                grads, value(d)

            end

            for (grads, l) in results

                loss += l

                i = 0
                for layer in model
                    for param in fieldnames(model_type)
                        i +=1
                        g = grads[i]
                        if g != nothing

                            g = grad_clip.(g)

                            moments[i] = alpha_moments * moments[i] + (1 - alpha_moments) * g
                            accugrads[i] = alpha_accugrads * accugrads[i] + (1 - alpha_accugrads) * g .^2

                            moment_hat = moments[i] / (1 - alpha_moments .^ e)
                            accugrad_hat = accugrads[i] / (1 - alpha_accugrads .^ e)

                            setfield!(layer, param, Param(getfield(layer, param)-lr.*moment_hat/sqrt(sum(accugrad_hat) + 1e-8)))

                            # setfield!(layer, param, Param(getfield(layer, param)-lr.*g))
                        end

                    end
                end

            end

        end

    push!(losses, loss)
    verbose ? println("epoch $e loss: $loss") : ()

    # for (inp,lbl) in test_set
    #     println("model prediction: $(sum(prop_model(model,inp)[end]))")
    #     println("actual result: $lbl")
    # end

    test_score = sum([abs(lbl - sum(prop_model(model,inp)[end])) < .01 ? 1 : 0 for (inp,lbl) in test_set])
    test_score2 = sum([abs(lbl - sum(prop_model(model,inp)[end])) < .1 ? 1 : 0 for (inp,lbl) in test_set])

    push!(test_sc1, test_score)
    push!(test_sc2, test_score2)

    if e%1 == 0
        println("\t> test lo_acc: $(test_score2)")
        println("\t> test hi_acc: $(test_score)")
    end

    end

    # verbose ? (begin prev_loss = 999_999_999
    # for (e,loss) in enumerate(losses)
    #     if loss > prev_loss
    #         println("bad loss: ep $e \n \t $prev_loss to $loss")
    #     end
    #     prev_loss = loss
    # end end) : ()

    verbose ? println("\n\t\t $model_name summary:") : ()

    for loss in[losses[1], losses[trunc(Int,length(losses)*1/4)], losses[trunc(Int,length(losses)*2/4)], losses[trunc(Int,length(losses)*3/4)], losses[end]]
        println(loss)
    end ;

    # verbose ? println("progress : $((1-losses[end]/losses[1])*100)") : ()

    println(" ")

    # TODO : graph here.

    # test_sc1 is more accurate
    # test_sc2 is less accurate

    # plot(test_sc1)
    # plot(test_sc2)

(losses[1],losses[end])
end

# d = [[randn(1,input_size) for _ in 1:seq_len] for __ in hm_data]

d = []
for i in 1:hm_data
    y = 1
    x = []
    while ! (y < 1)
        x = []
        for t in 1:seq_len
            arr = [randn(), randn() < 0.2 ? 0 : 1]
            push!(x, reshape(arr, 1, length(arr)))
        end
        y = sum(e1*e2 for (e1,e2) in x)
    end
    push!(d, [x,y])
end

test_set = []
for i in 1:hm_test
    y = 1
    x = []
    while ! (y < 1)
        x = []
        for t in 1:seq_len
            arr = [randn(), randn() < 0.5 ? 0 : 1]
            push!(x, reshape(arr, 1, length(arr)))
        end
        y = sum(e1*e2 for (e1,e2) in x)
    end
    push!(test_set, [x,y])
end

batchify(data, bs) =
begin
    batches = []
    for b in 1:(trunc(Int,length(data)/bs))
        push!(batches, data[(b-1)*bs+1:b*bs])
    end
    hm_left = length(data)%bs
    if hm_left != 0
        push!(batches, data[end-(hm_left-1):end])
    end
batches
end

# runner of main. main of main

for model_type in model_types
    println("\n\t> Running: $model_type \n")
    if length(layer_test) > 0
        progresses = [[0.0,0.0] for _ in 1:length(layer_test)]
        for _ in 1:hm_trials
            data = [[randn(1,input_size) for _ in 1:seq_len] for __ in hm_data]
            for (i,l) in enumerate(layer_test)
                println("layers: $l")
                loss1, lossend = main(model_type, data, hiddens=l)
                progresses[i][1] += loss1
                progresses[i][end] += lossend
            end

            fittest = argmax([(1-progress[end]/progress[1])*100 for progress in progresses])
            println("** Current Optimal layer size: $(layer_test[fittest]) \n")

        end

        fittest = argmax([(1-progress[end]/progress[1])*100 for progress in progresses])
        println(">> General Optimal layer size: $(layer_test[fittest]) \n")

        println("\n Progress list: \n")
        for (i,p) in enumerate(progresses)
            println("$model_type $(hiddens[i]) : $(p[1]) to $(p[end]) : $((1-p[end]/p[1])*100)")
        end

    else main(model_type, d) end
end ; println("\n\ndone.\n")

max_stars_repo_name: developersfox/GSTM

max_issues_repo_name: developersfox/GSTM

max_forks_repo_name: developersfox/GSTM


text:
#Attempt to download the Lalonde data set from http://users.nber.org/~rdehejia/nswdata2.html

function fetchdata(;refetch=false, ds = "lalonde_dw")
    if ds == "lalonde"
        csvname = "lalonde.csv"
        fname = "nsw"
        colnames =["treatment", "age", "education", "black", "hispanic",
                   "married", "nodegree", "RE75", "RE78"]
    elseif ds == "lalonde_dw"
        csvname = "lalonde_dw.csv"
        fname = "nswre74"
        colnames = ["treatment", "age", "education", "black", "hispanic",
                    "married", "nodegree", "RE74", "RE75", "RE78"]
    else
        error("unknown ds")
    end

    datapath = normpath(joinpath(@__FILE__, "..", "data"))
    dataname = joinpath(datapath, csvname)

    if isfile(dataname) && !refetch
        info("$dataname exists. To redownload, call fetchdata(ds=\"$ds\", refetch=true)")
        return
    end

    info("Downloading and parsing http://www.nber.org/~rdehejia/data/$(fname)_treated.txt...")
    treated  = readdlm(download("http://www.nber.org/~rdehejia/data/$(fname)_treated.txt"))
    info("Downloading and parsing http://www.nber.org/~rdehejia/data/$(fname)_control.txt...")
    control = readdlm(download("http://www.nber.org/~rdehejia/data/$(fname)_control.txt"))
    mat = vcat(reshape(colnames, 1, length(colnames)),  treated, control)

    isdir(datapath) || mkdir(datapath)

    info("Writing data set to $dataname")
    writecsv(dataname, mat)
    info("Done!")
end


text:
using Test
using ParametricOptInterface
using MathOptInterface
using GLPK
using Ipopt
using ECOS
using SCS
using LinearAlgebra
using JuMP

const POI = ParametricOptInterface
const MOIU = MOI.Utilities
const MOIT = MOI.Test

const ATOL = 1e-4

include("production_problem_test.jl")
include("MOI_wrapper.jl")
include("basic_tests.jl")
include("dual_tests.jl")
include("quad_tests.jl")
include("sdp_tests.jl")
include("vector_affine_tests.jl")
include("modifications_tests.jl")
include("jump_tests.jl")
include("nlp_test.jl")


text:
mutable struct Lattice
    name::String
    state
    probability

    Lattice(name::String, state, probability) = new(name, state, probability)

function Lattice(identifier::Int64)
    self= new()
    if identifier == 302
        self.name = "Lattice 1x2x2"
        self.state = [[0.0],[-0.856;0.933],[-1.27;1.307]]
        self.probability = [[1.0],[0.52;0.48],[0.75 0.25;0.26 0.74]]
    elseif identifier == 303
        self.name = "Lattice 1x1x4"
        self.state = [[0.0],[0.02],[-1.27;1.307;2.2;1.4]]
        self.probability = [[1.0],[1.0],[0.25 0.25 0.25 0.25]]
    elseif identifier == 304
        self.name = "Tree 1x2x3x4"
        self.state = [[0.0],[-0.856;0.933],[-1.27;1.307;2.0],[-2.5;-0.5;2.5;0.5]]
        self.probability = [[1.0],[0.52;0.48],[0.35 0.35 0.3; 0.1 0.5 0.4],[0.25 0.25 0.25 0.25;0.25 0.25 0.25 0.25;0.25 0.25 0.25 0.25]]
    end
    return self
end
end

######################################################################################
# Include helper functions for lattices here
# similar to Tree structure
"""
    structure(lat::Lattice)

Provides a summary on the branching structure of the lattice `lat`.

# Values
   - `bs`: vector indicating the number of nodes in each stage, i.e. the branching structure.
   - `nodes`: the total number of nodes.
   - `edges`: the total number of edges.
   - `paths`: the total number of unique paths.

# Example
```julia-repl
julia> lat = Lattice(304);
julia> structure(lat)
([1, 2, 3, 4], 10, 20, 24)
```

"""
function structure(lat::Lattice)
    bs = [length(lat.state[i]) for i=1:length(lat.state)];

    edges = 0
    for i=2:length(bs)
        edges += bs[i-1]*bs[i]
    end

    nodes = sum(bs)
    paths = prod(bs)

    return bs, nodes,edges, paths
end

"""
    part_lattice(lat, hist=nothing)

Provides a subtree given an observed history `hist`. If no history is provided the lattice is returned.

# Example
	lat = Lattice(304);
	hist = [1 2]
	part_lattice(lat, hist)
"""
function part_lattice(lat, hist=nothing)
    if (hist === nothing || hist == [])
        return(lat)
    end
    tmp_states = []; tmp_prob = [];
    lat2 = deepcopy(lat)
    for t=1:length(hist)
        push!(tmp_states,lat.state[t][Int(hist[t]),1,1 ])
        push!(tmp_prob,lat.probability[t][1,Int(hist[t]),1])
        if t>1
            deleteat!(lat2.state, 2) # delete everything after initial node up to observed history
            deleteat!(lat2.probability, 2) # stages move up hence 2nd entry is deleted mulitple times
        end
    end

    lat2.state[1] .= tmp_states[end];

    return(lat2)

end


text:
using BAT
using Test

using LinearAlgebra, Distributions, StatsBase, ValueShapes, Random123, DensityInterface
using UnPack, InverseFunctions, ForwardDiff

@testset "mode_estimators" begin
    prior = NamedTupleDist(
        x = Normal(2.0, 1.0),
        c = [4, 5],
        a = MvNormal([1.5, 0.5, 2.5], Matrix{Float32}(I, 3, 3))
    )

    posterior = PosteriorDensity(logfuncdensity(v -> 0), prior)

    true_mode_flat = [2.0, 1.5, 0.5, 2.5]
    true_mode = varshape(prior)(true_mode_flat)

    samples = @inferred(bat_sample(prior, IIDSampling(nsamples = 10^5))).result

    function test_findmode(posterior, algorithm, rtol; inferred::Bool = true)
        res = if inferred
            @inferred(bat_findmode(posterior, algorithm))
        else
            (bat_findmode(posterior, algorithm))
        end
        @test keys(res.result) == keys(true_mode)
        @test isapprox(unshaped(res.result, varshape(posterior)), true_mode_flat, rtol = rtol)

        if hasproperty(res.trace_trafo, :grad_logd)
            @unpack v, logd, grad_logd = res.trace_trafo
            f_logd = logdensityof(posterior) ∘ inverse(res.trafo)
            @test all(f_logd.(v) .≈ logd)
            @test all(grad_logd .≈ ForwardDiff.gradient.(Ref(f_logd), v))
        else
            @test hasproperty(res.trace_trafo, :v)
        end
    end

    function test_findmode_rng(rng, posterior, algorithm, rtol)
        res = (bat_findmode(rng, posterior, algorithm))
        @test keys(res.result) == keys(true_mode)
        @test isapprox(unshaped(res.result, varshape(posterior)), true_mode_flat, rtol = rtol)
    end

    @testset "ModeAsDefined" begin
        @test @inferred(bat_findmode(prior, ModeAsDefined())).result == true_mode
        @test @inferred(bat_findmode(BAT.DistributionDensity(prior), ModeAsDefined())).result == true_mode
        let post_modes = @inferred(bat_findmode(posterior)).result
            for k in keys(post_modes)
                @test isapprox(post_modes[k], true_mode[k], atol=0.001)
            end
        end
    end

    @testset "MaxDensitySampleSearch" begin
        @test @inferred(bat_findmode(samples, MaxDensitySampleSearch())).result isa NamedTuple
        m = bat_findmode(samples, MaxDensitySampleSearch())
        @test samples[m.mode_idx].v == m.result
        @test isapprox(unshaped(m.result, elshape(samples.v)), true_mode_flat, rtol = 0.05)
    end

    @testset "MaxDensityNelderMead" begin
        test_findmode(posterior,  MaxDensityNelderMead(trafo = NoDensityTransform()), 0.01)
    end

    @testset "MaxDensityLBFGS" begin
        # Result Optim.maximize with LBFGS is not type-stable:
        test_findmode(posterior, MaxDensityLBFGS(trafo = NoDensityTransform()), 0.01, inferred = false)

        rng = Philox4x((0, 0))
        test_findmode_rng(rng, posterior, MaxDensityLBFGS(trafo = NoDensityTransform()), 0.01)
    end
end


text:
"""
    mutable struct Atom

This contains information needed to describe an atom in Lammps Data File.
"""
mutable struct Atom <: Unit
    atom::Int64
    mol::Int64
    typ::Int64
    charge::Float64
    coord::Vector{Float64}
end

"""
    Atom(atom::Atom)

Do this will generate a new variable of `Atom` type same as variable `atom`.
"""
function Atom(atom::Atom)
    id = atom.atom
    mol = atom.mol
    typ = atom.typ
    charge = atom.charge
    coord = Vector{Float64}(undef, 3)
    coord[:] = atom.coord[:]
    Atom(id, mol, typ, charge, coord)
end

"""
    mutable struct Bond <: Unit

This contains information needed to describe a bond in Lammps Data File.
"""
mutable struct Bond <: Unit
    id::Int64
    typ::Int64
    atom::Vector{Int64}
end

"""
    Bond(bond::Bond)

Do this will generate a new variable of `Bond` type same as variable `bond`.
"""
function Bond(bond::Bond)
    id = bond.id
    typ = bond.typ
    atom = Vector{Int64}(undef, length(bond.atom))
    atom[:] = bond.atom[:]
    Bond(id, typ, atom)
end

"""
    mutable struct Angle <: Unit

This contains information needed to describe an angle in Lammps Data File.
"""
mutable struct Angle <: Unit
    id::Int64
    typ::Int64
    atom::Vector{Int64}
end

"""
    Angle(angle::Angle)

Do this will generate a new variable of `Angle` type same as variable `angle`.
"""
function Angle(angle::Angle)
    id = angle.id
    typ = angle.typ
    atom = Vector{Int64}(undef, length(angle.atom))
    atom[:] = angle.atom[:]
    Angle(id, typ, atom)
end

"""
    mutable struct Dihedral <: Unit

This contains information needed to describe an dihedral in Lammps Data File.
"""
mutable struct Dihedral <: Unit
    id::Int64
    typ::Int64
    atom::Vector{Int64}
end

"""
    Dihedral(dihedral::Dihedral)

Do this will generate a new variable of `Dihedral` type same as variable `dihedral`.
"""
function Dihedral(dihedral::Dihedral)
    id = dihedral.id
    typ = dihedral.typ
    atom = Vector{Int64}(undef, length(dihedral.atom))
    atom[:] = dihedral.atom[:]
    Dihedral(id, typ, atom)
end

"""
    mutable struct Improper <: Unit

This contains information needed to describe an improper in Lammps Data File.
"""
mutable struct Improper <: Unit
    id::Int64
    typ::Int64
    atom::Vector{Int64}
end

"""
    Improper(improper::Improper)

Do this will generate a new variable of `Improper` type same as variable `improper`.
"""
function Improper(improper::Improper)
    id = improper.id
    typ = improper.typ
    atom = Vector{Int64}(undef, length(improper.atom))
    atom[:] = improper.atom[:]
    Improper(id, typ, atom)
end

"""
    mutable struct Data_Cell

This contains information of cell matrix.

# Arguments
- `cell_mat`: postion of each cell point.
- `cell_vec`: # of cells in 3 direction.
- `num_cells`: # of cells in total.
"""
mutable struct Data_Cell
    cell_mat::Matrix{Int64}
    cell_vec::Vector{Int64}
    num_cells::Int64
end

"""
    mutable struct Data_Basic <:Data

This contains basic information of Lammps Model.

# Arguments
- `num_atoms`: # of atoms in total.
- ... Same formate for bonds, angles, dihedrals, and impropers.
- `num_atom_types`: # of types of atoms.
- ... Same fromate for types of bonds, angles, dihedrals, and impropers.
- `box_size`: 3 x 2 Matrix of lowest and highest postion.
- `box_tilt`: Represent xy xz yz respectively.
"""
mutable struct Data_Basic <:Data
    num_atoms::Int64
    num_bonds::Int64
    num_angles::Int64
    num_dihedrals::Int64
    num_impropers::Int64
    num_atom_types::Int64
    num_bond_types::Int64
    num_angle_types::Int64
    num_dihedral_types::Int64
    num_improper_types::Int64
    box_size::Matrix{Float64}
    box_tilt::Vector{Float64}
end

function Data_Basic(data::Data_Basic)
    box_size = [data.box_size[i, j] for i = 1 : 3, j = 1 :2]
    box_tilt = [data.box_tilt[i] for i = 1 : 3]
    Data_Basic(data.num_atoms, data.num_bonds, data.num_angles, data.num_dihedrals, data.num_impropers, data.num_atom_types, data.num_bond_types, data.num_angle_types, data.num_dihedral_types, data.num_improper_types, box_size, box_tilt)
end

"""
    mutable struct Data_Unit <:Data

This contains all information needed for an unit Lammps Model (only one structure).
"""
mutable struct Data_Unit <: Data
    data_basic::Data_Basic
    data_str::Str
    vec_atom::Union{Vector{Atom}, Int64}
    vec_bond::Union{Vector{Bond}, Int64}
    vec_angle::Union{Vector{Angle}, Int64}
    vec_dihedral::Union{Vector{Dihedral}, Int64}
    vec_improper::Union{Vector{Improper}, Int64}
end

function Data_Unit(data::Data_Unit)
    vec_atom = [Atom(data.vec_atom[n]) for n = 1 : length(data.vec_atom)]
    if typeof(data.vec_bond) != Int64
        vec_bond = [Bond(data.vec_bond[n]) for n = 1 : length(data.vec_bond)]
    else
        vec_bond = 0
    end
    if typeof(data.vec_angle) != Int64
        vec_angle = [Angle(data.vec_angle[n]) for n = 1 : length(data.vec_angle)]
    else
        vec_angle = 0
    end
    if typeof(data.vec_dihedral) != Int64
        vec_dihedral = [Dihedral(data.vec_dihedral[n]) for n = 1 : length(data.vec_dihedral)]
    else
        vec_dihedral = 0
    end
    if typeof(data.vec_improper) != Int64
        vec_improper = [Improper(data.vec_improper[n]) for n = 1 : length(data.vec_improper)]
    else
        vec_improper = 0
    end
    Data_Unit(Data_Basic(data.data_basic), data.data_str,
            vec_atom[:], vec_bond, vec_angle, vec_dihedral, vec_improper)
end

"""
    mutable struct Data_Sum <:Data

This contains all information needed for an unit Lammps Model (could be a combination of different sub-models).
"""
mutable struct Data_Sum <: Data
    data_basic::Data_Basic
    vec_str::Vector{Str}
    vec_atom::Union{Vector{Atom}, Int64}
    vec_bond::Union{Vector{Bond}, Int64}
    vec_angle::Union{Vector{Angle}, Int64}
    vec_dihedral::Union{Vector{Dihedral}, Int64}
    vec_improper::Union{Vector{Improper}, Int64}
end

function Data_Sum(data::Data_Unit)
    Data_Sum(Data_Basic(data.data_basic), [data.data_str], data.vec_atom[:],
        typeof(data.vec_bond)==Int64 ? data.vec_bond : data.vec_bond[:],
        typeof(data.vec_angle)==Int64 ? data.vec_angle : data.vec_angle[:],
        typeof(data.vec_dihedral)==Int64 ? data.vec_dihedral : data.vec_dihedral[:],
        typeof(data.vec_improper)==Int64 ? data.vec_improper : data.vec_improper[:])
end

function Data_Sum(data::Data_Sum)
    Data_Sum(Data_Basic(data.data_basic), [data.vec_str[n] for n = 1:length(data.vec_str)], data.vec_atom[:],
        typeof(data.vec_bond)==Int64 ? data.vec_bond : data.vec_bond[:],
        typeof(data.vec_angle)==Int64 ? data.vec_angle : data.vec_angle[:],
        typeof(data.vec_dihedral)==Int64 ? data.vec_dihedral : data.vec_dihedral[:],
        typeof(data.vec_improper)==Int64 ? data.vec_improper : data.vec_improper[:])
end


text:
using Test

@testset "./src/games.jl" begin

using BellScenario

@testset "Game()" begin
    @testset "int valued game" begin
        G = Game([1 0 0;0 1 0;0 0 1], 2)

        @test G isa Game{Int64}
        @test G == [1 0 0;0 1 0;0 0 1]
        @test G.game == [1 0 0;0 1 0;0 0 1]
        @test G.β == 2
    end

    @testset "float valued game" begin
        G = Game([0.5 0 0;0 0.5 0;0 0 0.5], 1.0)

        @test G isa Game{Float64}
        @test G == 1/2*[1 0 0;0 1 0;0 0 1]
        @test G.game == 1/2*[1 0 0;0 1 0;0 0 1]
        @test G.β == 1
    end

    @testset "errors" begin
        @test_throws MethodError Game([1 0 0;0 1 0;0 0 1], 1im)
    end
end

@testset "BellGame()" begin
    BG = BellGame([1 0 0;0 1 0;0 0 1], 2)

    @test BG.β == 2
    @test BG == [1 0 0;0 1 0;0 0 1]

    @test BG isa AbstractGame{Int64}
    @test BG isa BellGame
end

end


text:
# R Code to generate data
# ```r
# arrow::write_parquet(
#   data.frame(
#     a = c(1, 3, NA, 4, NA, 1, 1, NA, 1:8),
#     b = c(1:9, 1:7)),
#   "c:/scratch/nihao.parquet")
#   ```

using Diban
using Diban: TYPES, read_column

path = "c:/scratch/test.parquet"

#@time read_parquet(path)

path = "c:/scratch/nihao.parquet"

path = "tmp.parquet"

metadata = Diban.metadata(path)

fileio=open(path)

metadata.row_groups[1].columns[1].meta_data.dictionary_page_offset

seek(fileio, metadata.row_groups[1].columns[1].meta_data.dictionary_page_offset)

ph = read_thrift(fileio, PAR2.PageHeader)

position(fileio)

compressed_data = read(fileio, ph.compressed_page_size)

using Snappy

uncompressed_data = Snappy.uncompress(compressed_data)

reinterpret(Int64, uncompressed_data)

ph2 = read_thrift(fileio, PAR2.PageHeader)

compressed_data = read(fileio, ph2.compressed_page_size)
uncompressed_data = Snappy.uncompress(compressed_data)
uncompressed_data

io = IOBuffer(uncompressed_data)

# the definition data is encoded in the `encoded_data_length` bytes
encoded_data_length = Int(read(io, UInt32))

# the next set of data is a LEB128 unsigned int coded by that
# althought it is possible for the number be larger 2^31-1, in practice it is
# not encouraged, and hence not supported by this algorithm

header = _read_varint(io, UInt32)
# if the last binary digit is 1 then the next bits are binary encoded
# otherwise it's rle encoded

bit_pack_encoded = (header & UInt(1)) == 1

if bit_pack_encoded
	bit_packed_run_length = header >> 1
	bit_packed_data = read(io, bit_packed_run_length)
	reduce(vcat, digits.(bit_packed_data, base = 2, pad=8))
else
	rle_run_length = header >> 1
	repeated_value = read(io, UInt8)
end

Int(rle_run_length), Int(repeated_value)

position(io)

encoded_data = read(io)

dataio = IOBuffer(encoded_data)

bitwidth = read(dataio, UInt8)

header = _read_varint(dataio, UInt32)

string(0x10, base=2, pad=8)
string(0x32, base=2, pad=8)

pos = 0

bytes_to_skip, bits_to_shift = divrem(pos, 8)

bp = BitPackedIterator(read(dataio, 2), 3)

iterate(bp)

for value in bp
	println(value)
end

# TODO what happens if the bitwidth is larger than 8?


text:
using Hyperparameters

using FilePathsBase
using JSON
using Memento
using Memento.TestUtils

using Test

const LOGGER = getlogger()

@testset "Hyperparameters.jl" begin
    include("hyperparameters.jl")
end


text:
function check_pbs_tasknum_one()
  pbs_tasknum = @compat Int(parse(get(ENV,"PBS_TASKNUM","1")))
  if (pbs_tasknum!=1) error("Somehow Julia started with a tasknum !=1") end
end

function read_pbs_nodefile(nodefilename = get(ENV,"PBS_NODEFILE",""))
  if length(nodefilename) == 0
    return []
  end
  local nodelist, nodefile
  nodefile = open(nodefilename)
  try
    nodelist = readlines(nodefile)
  finally
    close(nodefile)
  end
  @assert(length(nodelist)>=1)  # ensure read at least one nodename
  # trim CR/LF at end of each node name
  [ nodelist[i] = chomp(nodelist[i]) for i in 1:length(nodelist) ]
  return nodelist
end

function addprocs_pbs( ; nodefilename = get(ENV,"PBS_NODEFILE",""), all_procs_work::Integer = -1 )
  nodelist = read_pbs_nodefile(nodefilename)
  # Automatically decide whether to set the first node as a worker if not specified
  if all_procs_work == -1
     if length(nodelist) <= 16
        all_procs_work = 1
     else
        all_procs_work = 0
     end
  end
  first_proc = all_procs_work > 0 ? 1 : 2
  proclist = try
    addprocs(nodelist[first_proc:end])
  catch
    warn("addprocs threw an exception")
    workers()
  end
end

@everywhere function print_node_status()
  println("# myid= ",myid())
  println("# nprocs= ",nprocs())
  println("# nworkers= ",nworkers())
  #Libc.flush_cstdio()
end

@everywhere function srand_all_procs(seed::Integer; same::Bool=false, debug::Bool = false)
  srand(seed)
  if length(workers()) > 1
     worker_seeds = same ? fill(seed, length(workers())) : rand(1:typemax(Uint32),length(workers()) )
     @sync for (proc,seed) in zip(workers(),worker_seeds)
        if debug println("Running srand($seed) on proc ",proc) end
        remotecall_wait(proc,srand,seed)
     end
     if debug
        @sync for proc in workers()
           println("rand() from proc returns ",@fetchfrom(proc, rand() ) )
        end
     end
  end
end