text:
"""
    AbstractLagrangeMaster

Abstract type of Lagrangian master methods
"""
abstract type AbstractLagrangeMaster end

"""
    load!

This function loads the Lagrangian dual problem to the master.

# Arguments
- `method`: Lagrangian master method
- `num_coupling_variables`: number of coupling variables
- `num_blocks`: number of blocks (or Lagrangian subproblems)
- `eval_function`: function pointer to evalute Lagrangian dual functions
- `init_sol`: initial solution of the Lagrangian master
"""
function load!(method::AbstractLagrangeMaster, num_coupling_variables::Int, num_blocks::Int, eval_function::Function, init_sol::Vector{Float64}) end

"""
    add_constraints!

This function is to add any constraints to the master problem.

# Arguments
- `LD`: Lagrangian dual
- `method`: Lagrangian master
"""
function add_constraints!(LD::AbstractLagrangeDual, method::AbstractLagrangeMaster) end

"""
    run!

This runs the Lagrangian master method.

# Arguments
- `method`: Lagrangian master
"""
function run!(method::AbstractLagrangeMaster) end

"""
    get_objective

This returns the objective function value.

# Arguments
- `method`: Lagrangian master
"""
function get_objective(method::AbstractLagrangeMaster) end

"""
    get_solution

This returns the Lagrangian master solution.

# Arguments
- `method`: Lagrangian master
"""
function get_solution(method::AbstractLagrangeMaster) end

"""
    get_times

This returns the solution times of Lagrangian master for all iterations.

# Arguments
- `method`: Lagrangian master
"""
get_times(method::AbstractLagrangeMaster)::Vector{Float64} = zeros(1)

"""
set_obj_limit!

This sets the Lagrangian bound limit.

# Arguments
- `method`: Lagrangian master
- `val`: objective limit value
"""
function set_obj_limit!(method::AbstractLagrangeMaster, val::Float64) end


text:
export CUDNNError

struct CUDNNError <: Exception
    code::cudnnStatus_t
    msg::AbstractString
end
Base.show(io::IO, err::CUDNNError) = print(io, "CUDNNError(code $(err.code), $(err.msg))")

function CUDNNError(status::cudnnStatus_t)
    msg = unsafe_string(cudnnGetErrorString(status))
    return CUDNNError(status, msg)
end

macro check(dnn_func)
    quote
        local err::cudnnStatus_t
        err = $(esc(dnn_func))
        if err != CUDNN_STATUS_SUCCESS
            throw(CUDNNError(err))
        end
        err
    end
end

ext: jl


text:
using DataFrames, Dates, CSV
import DataStructures.counter
import DataStructures.Accumulator

include("../../continuous_variable_evolution/src/types.jl")
include("../../continuous_variable_evolution/src/cont_var.jl")
include("../../continuous_variable_evolution/src/run_cont_var.jl")
include("../../continuous_variable_evolution/src/bin_data.jl")
include("../../continuous_variable_evolution/src/propsel.jl")

function test_bins()
  fdc = counter(Int64)
  cutoff = 0.125
  vals= 3.0*rand(20) .- 1.5
  for v in vals
    increment_bins( fdc, v, cutoff )
  end
  println("bins: ",fdc)
  bins_to_vector( fdc, cutoff )
end

function test_run_simulation(mutation_stddev::Float64, N::Int64, ngens::Int64, bins_cutoff::Float64;
    csvfile::String="" )
  #N = 100;
  N_list=[N]
  num_attributes=1; num_attributes_list = [num_attributes]
  #mutation_stddev=0.10;
  mutation_stddev_list = [mutation_stddev]
  N_mut_list = []
  num_trials=1
  num_subpops=1
  num_attributes=1
  #ngens=10
  ideal=1.0
  fit_slope=1.0
  int_burn_in=100
  neutral=:false
  #bins_cutoff = 0.02
  simrecord = cont_var_result_type( N_list, num_attributes_list, mutation_stddev_list, N_mut_list, num_trials,
      N, num_subpops, num_attributes, ngens, int_burn_in,
      mutation_stddev, ideal, fit_slope, neutral, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0,0,0,0 )
  (simrec,bins,lbvec)=cont_var_simulation( simrecord, bins_to_vect=:true, bins_cutoff=bins_cutoff )
  println("bins:  ",bins)
  println("lvvec: ",lbvec)
  df = DataFrame()
  df.bins=bins
  df.lbvec=lbvec
  if length(csvfile) > 0
    open( csvfile, "w" ) do f
      hostname = chomp(open("/etc/hostname") do f read(f,String) end)
      println(f,"# date and time: ",Dates.now())
      println(f,"# host: ",hostname," with ",nprocs()-1,"  processes: " )
      println(f,"# N: ",N)
      println(f,"# mutation_stddev: ",mutation_stddev)
      println(f,"# ngens: ",ngens)
      println(f,"# bins_cutoff: ",bins_cutoff)
      CSV.write( f, df, append=true, writeheader=true )
    end
  end
  (bins, lbvec)
end


text:
using OrdinaryDiffEq
using Trixi

###############################################################################
# semidiscretization of the compressible Euler equations
gamma = 1.4
equations_euler = CompressibleEulerEquations2D(gamma)

"""
    initial_condition_sedov_self_gravity(x, t, equations::CompressibleEulerEquations2D)

Adaptation of the Sedov blast wave with self-gravity taken from
- Michael Schlottke-Lakemper, Andrew R. Winters, Hendrik Ranocha, Gregor J. Gassner (2020)
  A purely hyperbolic discontinuous Galerkin approach for self-gravitating gas dynamics
  [arXiv: 2008.10593](https://arxiv.org/abs/2008.10593)
based on
- http://flash.uchicago.edu/site/flashcode/user_support/flash4_ug_4p62/node184.html#SECTION010114000000000000000
Should be used together with [`boundary_condition_sedov_self_gravity`](@ref).
"""
function initial_condition_sedov_self_gravity(x, t, equations::CompressibleEulerEquations2D)
  # Set up polar coordinates
  r = sqrt(x[1]^2 + x[2]^2)

  # Setup based on http://flash.uchicago.edu/site/flashcode/user_support/flash4_ug_4p62/node184.html#SECTION010114000000000000000
  r0 = 0.125 # = 4.0 * smallest dx (for domain length=8 and max-ref=8)
  E = 1.0
  p_inner   = (equations.gamma - 1) * E / (pi * r0^2)
  p_ambient = 1e-5 # = true Sedov setup

  # Calculate primitive variables
  # use a logistic function to tranfer density value smoothly
  L  = 1.0    # maximum of function
  x0 = 1.0    # center point of function
  k  = -150.0 # sharpness of transfer
  logistic_function_rho = L/(1.0 + exp(-k*(r - x0)))
  rho_ambient = 1e-5
  rho = max(logistic_function_rho, rho_ambient) # clip background density to not be so tiny

  # velocities are zero
  v1 = 0.0
  v2 = 0.0

  # use a logistic function to tranfer pressure value smoothly
  logistic_function_p = p_inner/(1.0 + exp(-k*(r - r0)))
  p = max(logistic_function_p, p_ambient)

  return prim2cons(SVector(rho, v1, v2, p), equations)
end
initial_condition = initial_condition_sedov_self_gravity

"""
    boundary_condition_sedov_self_gravity(u_inner, orientation, direction, x, t,
                                          surface_flux_function,
                                          equations::CompressibleEulerEquations2D)

Adaptation of the Sedov blast wave with self-gravity taken from
- Michael Schlottke-Lakemper, Andrew R. Winters, Hendrik Ranocha, Gregor J. Gassner (2020)
  A purely hyperbolic discontinuous Galerkin approach for self-gravitating gas dynamics
  [arXiv: 2008.10593](https://arxiv.org/abs/2008.10593)
based on
- http://flash.uchicago.edu/site/flashcode/user_support/flash4_ug_4p62/node184.html#SECTION010114000000000000000
Should be used together with [`initial_condition_sedov_self_gravity`](@ref).
"""
function boundary_condition_sedov_self_gravity(u_inner, orientation, direction, x, t,
                                               surface_flux_function,
                                               equations::CompressibleEulerEquations2D)
  # velocities are zero, density/pressure are ambient values according to
  # initial_condition_sedov_self_gravity
  rho = 1e-5
  v1 = 0.0
  v2 = 0.0
  p = 1e-5

  u_boundary = prim2cons(SVector(rho, v1, v2, p), equations)

  # Calculate boundary flux
  if iseven(direction) # u_inner is "left" of boundary, u_boundary is "right" of boundary
    flux = surface_flux_function(u_inner, u_boundary, orientation, equations)
  else # u_boundary is "left" of boundary, u_inner is "right" of boundary
    flux = surface_flux_function(u_boundary, u_inner, orientation, equations)
  end

  return flux
end
boundary_conditions = boundary_condition_sedov_self_gravity

surface_flux = flux_hll
volume_flux  = flux_chandrashekar
polydeg = 3
basis = LobattoLegendreBasis(polydeg)
indicator_sc = IndicatorHennemannGassner(equations_euler, basis,
                                         alpha_max=0.5,
                                         alpha_min=0.001,
                                         alpha_smooth=true,
                                         variable=density_pressure)
volume_integral = VolumeIntegralShockCapturingHG(indicator_sc;
                                                 volume_flux_dg=volume_flux,
                                                 volume_flux_fv=surface_flux)
solver_euler = DGSEM(basis, surface_flux, volume_integral)

coordinates_min = (-4, -4)
coordinates_max = ( 4,  4)
mesh = TreeMesh(coordinates_min, coordinates_max,
                initial_refinement_level=2,
                n_cells_max=100_000,
                periodicity=false)

semi_euler = SemidiscretizationHyperbolic(mesh, equations_euler, initial_condition, solver_euler,
                                          boundary_conditions=boundary_conditions)

###############################################################################
# semidiscretization of the hyperbolic diffusion equations
equations_gravity = HyperbolicDiffusionEquations2D()

"""
    initial_condition_sedov_self_gravity(x, t, equations::HyperbolicDiffusionEquations2D)

Adaptation of the Sedov blast wave with self-gravity taken from
- Michael Schlottke-Lakemper, Andrew R. Winters, Hendrik Ranocha, Gregor J. Gassner (2020)
  A purely hyperbolic discontinuous Galerkin approach for self-gravitating gas dynamics
  [arXiv: 2008.10593](https://arxiv.org/abs/2008.10593)
based on
- http://flash.uchicago.edu/site/flashcode/user_support/flash4_ug_4p62/node184.html#SECTION010114000000000000000
Should be used together with [`boundary_condition_sedov_self_gravity`](@ref).
"""
function initial_condition_sedov_self_gravity(x, t, equations::HyperbolicDiffusionEquations2D)
  # for now just use constant initial condition for sedov blast wave (can likely be improved)
  phi = 0.0
  q1  = 0.0
  q2  = 0.0
  return SVector(phi, q1, q2)
end

"""
    boundary_condition_sedov_self_gravity(u_inner, orientation, direction, x, t,
                                          surface_flux_function,
                                          equations::HyperbolicDiffusionEquations2D)

Adaptation of the Sedov blast wave with self-gravity taken from
- Michael Schlottke-Lakemper, Andrew R. Winters, Hendrik Ranocha, Gregor J. Gassner (2020)
  A purely hyperbolic discontinuous Galerkin approach for self-gravitating gas dynamics
  [arXiv: 2008.10593](https://arxiv.org/abs/2008.10593)
based on
- http://flash.uchicago.edu/site/flashcode/user_support/flash4_ug_4p62/node184.html#SECTION010114000000000000000
Should be used together with [`initial_condition_sedov_self_gravity`](@ref).
"""
function boundary_condition_sedov_self_gravity(u_inner, orientation, direction, x, t,
                                                surface_flux_function,
                                                equations::HyperbolicDiffusionEquations2D)
  u_boundary = initial_condition_sedov_self_gravity(x, t, equations)

  # Calculate boundary flux
  if iseven(direction) # u_inner is "left" of boundary, u_boundary is "right" of boundary
    flux = surface_flux_function(u_inner, u_boundary, orientation, equations)
  else # u_boundary is "left" of boundary, u_inner is "right" of boundary
    flux = surface_flux_function(u_boundary, u_inner, orientation, equations)
  end

  return flux
end

solver_gravity = DGSEM(polydeg, flux_lax_friedrichs)

semi_gravity = SemidiscretizationHyperbolic(mesh, equations_gravity, initial_condition, solver_gravity,
                                            boundary_conditions=boundary_conditions,
                                            source_terms=source_terms_harmonic)

###############################################################################
# combining both semidiscretizations for Euler + self-gravity
parameters = ParametersEulerGravity(background_density=0.0, # aka rho0
                                    gravitational_constant=6.674e-8, # aka G
                                    cfl=2.4,
                                    resid_tol=1.0e-4,
                                    n_iterations_max=100,
                                    timestep_gravity=timestep_gravity_erk52_3Sstar!)

semi = SemidiscretizationEulerGravity(semi_euler, semi_gravity, parameters)

###############################################################################
# ODE solvers, callbacks etc.
tspan = (0.0, 1.0)
ode = semidiscretize(semi, tspan);

summary_callback = SummaryCallback()

amr_indicator = IndicatorHennemannGassner(semi,
                                          alpha_max=1.0,
                                          alpha_min=0.0,
                                          alpha_smooth=false,
                                          variable=density_pressure)
amr_controller = ControllerThreeLevel(semi, amr_indicator,
                                      base_level=2,
                                      max_level =8, max_threshold=0.0003)
amr_callback = AMRCallback(semi, amr_controller,
                           interval=1,
                           adapt_initial_condition=true,
                           adapt_initial_condition_only_refine=true)

stepsize_callback = StepsizeCallback(cfl=1.0)

save_solution = SaveSolutionCallback(interval=100,
                                     save_initial_solution=true,
                                     save_final_solution=true,
                                     solution_variables=cons2prim)

analysis_interval = 100
alive_callback = AliveCallback(analysis_interval=analysis_interval)

analysis_callback = AnalysisCallback(semi_euler, interval=analysis_interval,
                                     save_analysis=true,
                                     extra_analysis_integrals=(energy_total, energy_kinetic, energy_internal))

callbacks = CallbackSet(summary_callback, amr_callback, stepsize_callback,
                        save_solution,
                        analysis_callback, alive_callback)

###############################################################################
# run the simulation
sol = solve(ode, CarpenterKennedy2N54(williamson_condition=false),
            dt=1.0, # solve needs some value here but it will be overwritten by the stepsize_callback
            save_everystep=false, callback=callbacks);
summary_callback() # print the timer summary
println("Number of gravity subcycles: ", semi.gravity_counter.ncalls_since_readout)


text:
module DataDrivenDiffEq

using LinearAlgebra
using ModelingToolkit
using Compat
import Convex, GLPKMathProgInterface;

abstract type abstractBasis end;
abstract type abstractKoopmanOperator end;

include("./basis.jl")
export Basis
export variables, jacobian, dynamics

include("./exact_dmd.jl")
export ExactDMD
export eigen, eigvals, eigvecs
export modes, frequencies, isstable
export dynamics, update!

include("./extended_dmd.jl")
export ExtendedDMD
export dynamics, linear_dynamics
export reduce_basis, update!

include("./dmdc.jl")
export DMDc
export eigen, eigvals, eigvecs
export get_dynamics, get_input_map, dynamics

include("./sindy.jl")
export SInDy

end # module


text:
Base.:(==)(i1::CompositeInterpolation, i2::CompositeInterpolation) = i1.before_first == i2.before_first && i1.inner == i2.inner && i1.after_last == i2.after_last
Base.hash(i::CompositeInterpolation) = 1 + hash(i.before_first) + hash(i.inner) + hash(i.after_last)

# Curve methods implementation

function _zero_rate(method::METHOD, curve::AbstractIRCurve, maturity::T) where {METHOD<:RateInterpolation, T<:Union{Date, YearFraction}}
    _zero_rate(method, curve_get_dtm(curve), curve_get_zero_rates(curve), days_to_maturity(curve, maturity))
end

function _zero_rate(comp::CompositeInterpolation, curve::AbstractIRCurve, maturity::T) where {T<:Union{Date, YearFraction}}
    dtm = days_to_maturity(curve, maturity)

    if dtm < curve_get_dtm(curve)[1]
        return _zero_rate(comp.before_first, curve, maturity)
    elseif dtm > curve_get_dtm(curve)[end]
        return _zero_rate(comp.after_last, curve, maturity)
    else
        return _zero_rate(comp.inner, curve, maturity)
    end
end

# Returns tuple (index_a, index_b) for input vector x
# for interpolands on linear interpolation on point x_out
function _interpolationpoints(x::Vector{T}, x_out::Number) where {T<:Number}
    local index_a::Int
    local index_b::Int

    if x_out <= x[1]
        # Interpolation point is before first vertice
        # Slope will be determined by the 1st and 2nd vertices
        index_a = 1
        index_b = 2
    elseif x_out >= x[end]
        # Interpolation point is after last vertice
        # Slope will be determined by the last and last-1 vertices
        index_b = length(x)
        index_a = index_b - 1
    else
        # Inner point
        index_a =  findlast(a -> a < x_out, x) # last element before x_out on x
        index_b = findfirst(b -> b >=  x_out, x) # first element after x_out on x
    end

    return index_a, index_b
end

# Perform Linear interpolation. Slope is determined by points (Xa, Ya) and (Xb, Yb).
# Interpolation occurs on point (x_out, returnvalue)
_linearinterp(Xa::Number, Ya::Number, Xb::Number, Yb::Number, x_out::Number) = (x_out - Xa) * (Yb - Ya) / (Xb - Xa) + Ya

# Linear interpolation of zero_rates
function _zero_rate(::Linear, x::Vector{N1}, y::Vector{N2}, x_out::Number) where {N1<:Number, N2<:Number}
    # If this curve has only 1 vertice, this will be a flat curve
    if length(x) == 1
        return y[1]
    end

    index_a, index_b = _interpolationpoints(x, x_out)
    return _linearinterp(x[index_a], y[index_a], x[index_b], y[index_b], x_out)
end

# Step Function
function _zero_rate(::StepFunction, x::Vector{N1}, y::Vector{N2}, x_out::Number) where {N1<:Number, N2<:Number}
    # If this curve has only 1 vertice, this will be a flat curve
    if length(x) == 1
        return y[1]
    end

    if x_out <= x[1]
        # Interpolation point is before first vertice
        # The result will be extrapolated using the first vertice zero_rate
        return y[1]
    elseif x_out >= x[end]
        # Interpolation point is after last vertice
        # The result will be extrapolated using the last vertice zero_rate
        return y[end]
    else
        # Inner point
        return y[findlast(a -> a <= x_out, x)] # last element before x_out on x
    end
end

# Flat Forward is linear interpolation on the log of discountfactors
# Maybe not useful for SimpleCompounding curves.
function _zero_rate(::FlatForward, curve::AbstractIRCurve, maturity::T) where {T<:Union{Date, YearFraction}}

    # If this curve has only 1 vertice, this will be a flat curve
    if length(curve_get_zero_rates(curve)) == 1
        return curve_get_zero_rates(curve)[1]
    end

    x_out = days_to_maturity(curve, maturity)
    curve_dtm = curve_get_dtm(curve)
    curve_zero_rates = curve_get_zero_rates(curve)
    index_a, index_b = _interpolationpoints(curve_dtm, x_out)
    Xa = curve_dtm[index_a]
    Ya = curve_zero_rates[index_a]
    Xb = curve_dtm[index_b]
    Yb = curve_zero_rates[index_b]

    year_fraction_a = yearfraction(curve_get_daycount(curve), Xa)
    logPa = log(discountfactor(curve_get_compounding(curve), Ya, year_fraction_a))

    year_fraction_b = yearfraction(curve_get_daycount(curve), Xb)
    logPb = log(discountfactor(curve_get_compounding(curve), Yb, year_fraction_b))

    year_fraction_x = yearfraction(curve_get_daycount(curve), x_out)
    logPx = _linearinterp(value(year_fraction_a), logPa, value(year_fraction_b), logPb, value(year_fraction_x))

    return discountfactor_to_rate(curve_get_compounding(curve), exp(logPx), year_fraction_x)
end

function _zero_rate(::NelsonSiegel, curve::AbstractIRCurve, maturity::YearFraction)

    # beta1 = param[1]
    # beta2 = param[2]
    # beta3 = param[3]
    # lambda = param[4]

    param = curve_get_model_parameters(curve)
    t = value(maturity)
    _exp_lambda_t_ = exp(-param[4]*t)
    F_beta2 = (1.0 - _exp_lambda_t_) / (param[4]*t)

    return param[1] + param[2]*F_beta2 + param[3]*(F_beta2 - _exp_lambda_t_)
end

_zero_rate(method::NelsonSiegel, curve::AbstractIRCurve, maturity::Date) = _zero_rate(method, curve, yearfraction(curve, maturity))

function _zero_rate(::Svensson, curve::AbstractIRCurve, maturity::YearFraction)

    # beta1 = param[1]
    # beta2 = param[2]
    # beta3 = param[3]
    # beta4 = param[4]
    # lambda1 = param[5]
    # lambda2 = param[6]

    param = curve_get_model_parameters(curve)
    t = value(maturity)
    _exp_lambda1_t_ = exp(-param[5]*t)
    _exp_lambda2_t_ = exp(-param[6]*t)
    F_beta2 = (1.0 - _exp_lambda1_t_) / (param[5]*t)

    return param[1] + param[2]*F_beta2 + param[3]*(F_beta2 - _exp_lambda1_t_) +
            param[4]*( (1.0 - _exp_lambda2_t_)/(param[6]*t) - _exp_lambda2_t_)
end

_zero_rate(method::Svensson, curve::AbstractIRCurve, maturity::Date) = _zero_rate(method, curve, yearfraction(curve, maturity))

function _zero_rate(method::CubicSplineOnRates, curve::AbstractIRCurve, maturity::T) where {T<:Union{Date, YearFraction}}
    sp = curve_get_spline_fit_on_rates(curve)
    return splineint(sp, days_to_maturity(curve, maturity))
end

function _zero_rate(::CubicSplineOnRates, curve::AbstractIRCurve, maturity_vec::Vector{Date})
    sp = curve_get_spline_fit_on_rates(curve)

    l = length(maturity_vec)
    rates = Vector{Float64}(undef, l)

    for i in 1:l
        @inbounds dtm = days_to_maturity(curve, maturity_vec[i])
        @inbounds rates[i] = splineint(sp, dtm)
    end
    return rates
end

function _zero_rate(::CubicSplineOnDiscountFactors, curve::AbstractIRCurve, maturity::YearFraction)
    sp = curve_get_spline_fit_on_discount_factors(curve)
    result_discount_factor = splineint(sp, value(maturity))
    return discountfactor_to_rate(curve_get_compounding(curve), result_discount_factor, maturity)
end

function _zero_rate(method::CubicSplineOnDiscountFactors, curve::AbstractIRCurve, maturity::Date)
    return _zero_rate(method, curve, yearfraction(curve, maturity))
end

# to make splineint work in the following method _zero_rate(::CubicSplineOnDiscountFactors, ::AbstractIRCurve, ::Vector{Date})
function splineint(s::Spline{T1}, x_out::YearFraction{T2}) where {T1<:Number, T2<:Number}
    splineint(s, value(x_out))
end

function _zero_rate(::CubicSplineOnDiscountFactors, curve::AbstractIRCurve, maturity_vec::Vector{Date})
    sp = curve_get_spline_fit_on_discount_factors(curve)
    mat_vec_len = length(maturity_vec)

    yf_maturity_vec = Vector{YearFraction}(undef, mat_vec_len)
    for i in 1:mat_vec_len
        @inbounds yf_maturity_vec[i] = yearfraction(curve_get_daycount(curve), curve_get_date(curve), maturity_vec[i])
    end

    return discountfactor_to_rate(curve_get_compounding(curve), splineint(sp, yf_maturity_vec), yf_maturity_vec)
end

# Generate vector functions
for elty in (:FlatForward, :CompositeInterpolation, :StepFunction, :Linear, :NelsonSiegel, :Svensson)
    @eval begin
        function _zero_rate(m::$elty, curve::AbstractIRCurve, maturity_vec::Vector{Date})
            l = length(maturity_vec)
            rates = Vector{Float64}(undef, l)
            for i = 1:l
                @inbounds rates[i] = _zero_rate(m, curve, maturity_vec[i])
            end
        return rates
        end
    end
end


text:
using Documenter
using Coordinates

makedocs(;
    format = Documenter.HTML(prettyurls = get(ENV, "CI", nothing) == "true"),
    modules = [Coordinates],
    sitename = "Coordinates.jl",
    pages=[
        "Home" => "index.md",
    ],
    doctest = true, # :fix
)

deploydocs(
    repo = "github.com/KeitaNakamura/Coordinates.jl.git",
    devbranch = "main",
)


text:
const ALLOWABLE_USELESS_GROWTH = 0.25

"""
    OrderedRobinDict([itr])

`OrderedRobinDict{K,V}()` constructs a ordered dictionary with keys of type `K` and values of type `V`.
It takes advantage of `RobinDict` in maintaining the order of the keys.
Given a single iterable argument, constructs a [`OrderedRobinDict`](@ref) whose key-value pairs
are taken from 2-tuples `(key,value)` generated by the argument.

# Examples
```jldoctest
julia> OrderedRobinDict([("A", 1), ("B", 2)])
OrderedRobinDict{String,Int64} with 2 entries:
  "A" => 1
  "B" => 2
```

Alternatively, a sequence of pair arguments may be passed.

```jldoctest
julia> OrderedRobinDict("A"=>1, "B"=>2)
OrderedRobinDict{String,Int64} with 2 entries:
  "A" => 1
  "B" => 2
```
"""
mutable struct OrderedRobinDict{K,V} <: AbstractDict{K,V}
    dict::RobinDict{K, Int32}
    keys::Vector{K}
    vals::Vector{V}
    count::Int32

    function OrderedRobinDict{K, V}() where {K, V}
        new{K, V}(RobinDict{K, Int32}(), Vector{K}(), Vector{V}(), 0)
    end

    function OrderedRobinDict{K, V}(d::OrderedRobinDict{K, V}) where {K, V}
        new{K, V}(copy(d.dict), copy(d.keys), copy(d.vals), d.count)
    end

    function OrderedRobinDict{K,V}(kv) where {K, V}
        h = OrderedRobinDict{K,V}()
        for (k,v) in kv
            h[k] = v
        end
        return h
    end
    OrderedRobinDict{K,V}(p::Pair) where {K,V} = setindex!(OrderedRobinDict{K,V}(), p.second, p.first)
    function OrderedRobinDict{K,V}(ps::Pair...) where V where K
        h = OrderedRobinDict{K,V}()
        sizehint!(h, length(ps))
        for p in ps
            h[p.first] = p.second
        end
        return h
    end
end

OrderedRobinDict() = OrderedRobinDict{Any,Any}()
OrderedRobinDict(kv::Tuple{}) = OrderedRobinDict()
Base.copy(d::OrderedRobinDict) = OrderedRobinDict(d)
Base.empty(d::OrderedRobinDict, ::Type{K}, ::Type{V}) where {K, V} = OrderedRobinDict{K, V}()

OrderedRobinDict(ps::Pair{K,V}...) where {K,V} = OrderedRobinDict{K,V}(ps)
OrderedRobinDict(ps::Pair...)                  = OrderedRobinDict(ps)

OrderedRobinDict(d::AbstractDict{K, V}) where {K, V} = OrderedRobinDict{K, V}(d)

function OrderedRobinDict(kv)
    try
        return dict_with_eltype((K, V) -> OrderedRobinDict{K, V}, kv, eltype(kv))
    catch e
        if !isiterable(typeof(kv)) || !all(x -> isa(x, Union{Tuple,Pair}), kv)
            !all(x->isa(x,Union{Tuple,Pair}),kv)
            throw(ArgumentError("OrderedRobinDict(kv): kv needs to be an iterator of tuples or pairs"))
        else
            rethrow(e)
        end
    end
end

Base.empty(d::OrderedRobinDict{K,V}) where {K,V} = OrderedRobinDict{K,V}()

Base.length(d::Union{RobinDict, OrderedRobinDict}) = d.count
Base.isempty(d::Union{RobinDict, OrderedRobinDict}) = (length(d) == 0)

"""
    empty!(collection) -> collection

Remove all elements from a `collection`.

# Examples
```jldoctest
julia> A = OrderedRobinDict("a" => 1, "b" => 2)
OrderedRobinDict{String,Int64} with 2 entries:
  "a" => 1
  "b" => 2

julia> empty!(A);

julia> A
OrderedRobinDict{String,Int64} with 0 entries
```
"""
function Base.empty!(h::OrderedRobinDict{K,V}) where {K, V}
    empty!(h.dict)
    empty!(h.keys)
    empty!(h.vals)
    h.count = 0
    return h
end

function _setindex!(h::OrderedRobinDict, v, key)
    hk, hv = h.keys, h.vals
    push!(hk, key)
    push!(hv, v)
    nk = length(hk)
    @inbounds h.dict[key] = Int32(nk)
    h.count += 1
end

function Base.setindex!(h::OrderedRobinDict{K, V}, v0, key0) where {K,V}
    key = convert(K, key0)
    v = convert(V, v0)
    index = get(h.dict, key, -2)

    if index < 0
        _setindex!(h, v0, key0)
    else
        @assert haskey(h, key0)
        @inbounds orig_v = h.vals[index]
        (orig_v != v0) && (@inbounds h.vals[index] = v0)
    end

    check_for_rehash(h) && rehash!(h)

    return h
end

# rehash when there are ALLOWABLE_USELESS_GROWTH %
# tombstones, or non-mirrored entries in the dictionary
function check_for_rehash(h::OrderedRobinDict)
    keysl = length(h.keys)
    dictl = length(h)
    return (keysl > (1 + ALLOWABLE_USELESS_GROWTH)*dictl)
end

function rehash!(h::OrderedRobinDict{K, V}) where {K, V}
    keys = h.keys
    vals = h.vals
    hk = Vector{K}()
    hv = Vector{V}()

    for (idx, (k, v)) in enumerate(zip(keys, vals))
        if get(h.dict, k, -1) == idx
            push!(hk, k)
            push!(hv, v)
        end
    end

    h.keys = hk
    h.vals = hv

    for (idx, k) in enumerate(h.keys)
        h.dict[k] = idx
    end
    return h
end

function Base.sizehint!(d::OrderedRobinDict, newsz)
    oldsz = length(d)
    # grow at least 25%
    if newsz < (oldsz*5)>>2
        return d
    end
    sizehint!(d.keys, newsz)
    sizehint!(d.vals, newsz)
    sizehint!(d.dict, newsz)
    return d
end

"""
    get!(collection, key, default)

Return the value stored for the given key, or if no mapping for the key is present, store
`key => default`, and return `default`.

# Examples
```jldoctest
julia> d = OrderedRobinDict("a"=>1, "b"=>2, "c"=>3);

julia> get!(d, "a", 5)
1

julia> get!(d, "d", 4)
4

julia> d
OrderedRobinDict{String,Int64} with 4 entries:
  "a" => 1
  "b" => 2
  "c" => 3
  "d" => 4
```
"""
function Base.get!(h::OrderedRobinDict{K,V}, key0, default) where {K,V}
    index = get(h.dict, key0, -2)
    index > 0 && return h.vals[index]
    v = convert(V, default)
    setindex!(h, v, key0)
    return v
end

"""
    get!(f::Function, collection, key)

Return the value stored for the given key, or if no mapping for the key is present, store
`key => f()`, and return `f()`.

This is intended to be called using `do` block syntax:
```julia
get!(dict, key) do
    # default value calculated here
    time()
end
```
"""
function Base.get!(default::Base.Callable, h::OrderedRobinDict{K,V}, key0) where {K,V}
    index = get(h.dict, key0, -2)
    index > 0 && return @inbounds h.vals[index]
    v = convert(V, default())
    setindex!(h, v, key0)
    return v
end

function Base.getindex(h::OrderedRobinDict{K,V}, key) where {K,V}
    index = get(h.dict, key, -1)
    return (index < 0) ? throw(KeyError(key)) : @inbounds h.vals[index]::V
end

"""
    get(collection, key, default)

Return the value stored for the given key, or the given default value if no mapping for the
key is present.

# Examples
```jldoctest
julia> d = OrderedRobinDict("a"=>1, "b"=>2);

julia> get(d, "a", 3)
1

julia> get(d, "c", 3)
3
```
"""
function Base.get(h::OrderedRobinDict{K,V}, key, default) where {K,V}
    index = get(h.dict, key, -1)
    return (index < 0) ? default : @inbounds h.vals[index]::V
end

"""
    get(f::Function, collection, key)

Return the value stored for the given key, or if no mapping for the key is present, return
`f()`.  Use [`get!`](@ref) to also store the default value in the dictionary.

This is intended to be called using `do` block syntax

```julia
get(dict, key) do
    # default value calculated here
    time()
end
```
"""
function Base.get(default::Base.Callable, h::OrderedRobinDict{K,V}, key) where {K,V}
    index = get(h.dict, key, -1)
    return (index < 0) ? default() : @inbounds h.vals[index]::V
end

"""
    haskey(collection, key) -> Bool

Determine whether a collection has a mapping for a given `key`.

# Examples
```jldoctest
julia> D = OrderedRobinDict('a'=>2, 'b'=>3)
OrderedRobinDict{Char,Int64} with 2 entries:
  'a' => 2
  'b' => 3

julia> haskey(D, 'a')
true

julia> haskey(D, 'c')
false
```
"""
Base.haskey(h::OrderedRobinDict, key) = (get(h.dict, key, -2) > 0)
Base.in(key, v::Base.KeySet{K,T}) where {K,T<:OrderedRobinDict{K}} = (get(v.dict, key, -1) >= 0)

"""
    getkey(collection, key, default)

Return the key matching argument `key` if one exists in `collection`, otherwise return `default`.

# Examples
```jldoctest
julia> D = OrderedRobinDict('a'=>2, 'b'=>3)
OrderedRobinDict{Char,Int64} with 2 entries:
  'a' => 2
  'b' => 3

julia> getkey(D, 'a', 1)
'a': ASCII/Unicode U+0061 (category Ll: Letter, lowercase)

julia> getkey(D, 'd', 'a')
'a': ASCII/Unicode U+0061 (category Ll: Letter, lowercase)
```
"""
function Base.getkey(h::OrderedRobinDict{K,V}, key, default) where {K,V}
    index = get(h.dict, key, -1)
    return (index < 0) ? default : h.keys[index]::K
end

Base.@propagate_inbounds isslotfilled(h::OrderedRobinDict, index) = (h.dict[h.keys[index]] == index)

function _pop!(h::OrderedRobinDict, index)
    @inbounds val = h.vals[index]
    _delete!(h, index)
    return val
end

function Base.pop!(h::OrderedRobinDict)
    check_for_rehash(h) && rehash!(h)
    index = length(h.keys)
    while (index > 0)
        isslotfilled(h, index) && break
        index -= 1
    end
    index == 0 && rehash!(h)
    @inbounds key = h.keys[index]
    return key => _pop!(h, index)
end

function Base.pop!(h::OrderedRobinDict, key)
    index = get(h.dict, key, -1)
    (index > 0) ? _pop!(h, index) : throw(KeyError(key))
end

"""
    pop!(collection, key[, default])

Delete and return the mapping for `key` if it exists in `collection`, otherwise return
`default`, or throw an error if `default` is not specified.

# Examples
```jldoctest
julia> d = OrderedRobinDict("a"=>1, "b"=>2, "c"=>3);

julia> pop!(d, "a")
1

julia> pop!(d, "d")
ERROR: KeyError: key "d" not found
Stacktrace:
[...]

julia> pop!(d, "e", 4)
4
```
"""
function Base.pop!(h::OrderedRobinDict, key, default)
    index = get(h.dict, key, -1)
    (index > 0) ? _pop(h, index) : default
end

"""
    delete!(collection, key)

Delete the mapping for the given key in a collection, and return the collection.

# Examples
```jldoctest
julia> d = OrderedRobinDict("a"=>1, "b"=>2)
OrderedRobinDict{String,Int64} with 2 entries:
  "a" => 1
  "b" => 2

julia> delete!(d, "b")
OrderedRobinDict{String,Int64} with 1 entry:
  "a" => 1
```
"""
function Base.delete!(h::OrderedRobinDict, key)
    pop!(h, key)
    return h
end

function _delete!(h::OrderedRobinDict, index)
    @inbounds h.dict[h.keys[index]] = -1
    h.count -= 1
    check_for_rehash(h) ? rehash!(h) : h
end

function get_first_filled_index(h::OrderedRobinDict)
    index = 1
    while (true)
        isslotfilled(h, index) && return index
        index += 1
    end
end

function get_next_filled_index(h::OrderedRobinDict, index)
    # get the next filled slot, including index and beyond
    while (index <= length(h.keys))
        isslotfilled(h, index) && return index
        index += 1
    end
    return -1
end

Base.@propagate_inbounds function Base.iterate(h::OrderedRobinDict)
    isempty(h) && return nothing
    check_for_rehash(h) && rehash!(h)
    index = get_first_filled_index(h)
    return (Pair(h.keys[index], h.vals[index]), index+1)
end

Base.@propagate_inbounds function Base.iterate(h::OrderedRobinDict, i)
    length(h.keys) < i && return nothing
    index = get_next_filled_index(h, i)
    (index < 0) && return nothing
    return (Pair(h.keys[index], h.vals[index]), index+1)
end

Base.filter!(f, d::Union{RobinDict, OrderedRobinDict}) = Base.filter_in_one_pass!(f, d)

function Base.merge(d::OrderedRobinDict, others::AbstractDict...)
    K,V = _merge_kvtypes(d, others...)
    merge!(OrderedRobinDict{K,V}(), d, others...)
end

function Base.merge(combine::Function, d::OrderedRobinDict, others::AbstractDict...)
    K,V = _merge_kvtypes(d, others...)
    merge!(combine, OrderedRobinDict{K,V}(), d, others...)
end

OrderedCollections.isordered(::Type{T}) where {T <: OrderedRobinDict} = true


text:
using Test: @test, @testset, @test_throws
using Logging: Logging, with_logger
using LoggingExtras: FormatLogger
using LoggingFormats: LoggingFormats, Truncated, JSON, LogFmt
import JSON3

@testset "Truncating" begin
    @test LoggingFormats.shorten_str("αβγαβγ", 3) == "αβ…"
    @test LoggingFormats.shorten_str("αβγαβγ", 4) == "αβγ…"
    @test LoggingFormats.shorten_str("julia", 3) == "ju…"
    @test LoggingFormats.shorten_str("julia", 4) == "jul…"
    @test LoggingFormats.shorten_str("julia", 5) == "julia"

    @test_throws ErrorException Truncated(0)
    @test_throws ErrorException Truncated(-5)

    trunc_fun = Truncated(30)
    io = IOBuffer()
    truncating_logger = FormatLogger(trunc_fun, io)
    with_logger(truncating_logger) do
        @info "a"^50
    end
    str = String(take!(io))

    @test occursin("Info: aaaaaaaaaaaaaaaaaaaaaaaaaaaaa…", str)

    io = IOBuffer()
    truncating_logger = FormatLogger(trunc_fun, io)
    with_logger(truncating_logger) do
        long_var = "a"^50
        @info "a_message" long_var
    end
    str = String(take!(io))

    @test occursin("│   long_var = aaaaaaaaaaaaaa…", str)

    io = IOBuffer()
    truncating_logger = FormatLogger(trunc_fun, io)
    with_logger(truncating_logger) do
        long_var = "a"^50
        short_var = "a"
        @info "a_message" long_var short_var
    end
    str = String(take!(io))
    @test occursin("│   long_var = aaaaaaaaaaaaaa…", str)
    @test occursin("│   short_var = a", str)
end

struct NoStructTypeDefined
    f::Int
end

@testset "JSON" begin
    @test LoggingFormats.lvlstr(Logging.Error + 1) == "error"
    @test LoggingFormats.lvlstr(Logging.Error) == "error"
    @test LoggingFormats.lvlstr(Logging.Error - 1) == "warn"
    @test LoggingFormats.lvlstr(Logging.Warn + 1) == "warn"
    @test LoggingFormats.lvlstr(Logging.Warn) == "warn"
    @test LoggingFormats.lvlstr(Logging.Warn - 1) == "info"
    @test LoggingFormats.lvlstr(Logging.Info + 1) == "info"
    @test LoggingFormats.lvlstr(Logging.Info) == "info"
    @test LoggingFormats.lvlstr(Logging.Info - 1) == "debug"
    @test LoggingFormats.lvlstr(Logging.Debug + 1) == "debug"
    @test LoggingFormats.lvlstr(Logging.Debug) == "debug"

    io = IOBuffer()
    with_logger(FormatLogger(JSON(), io)) do
        @debug "debug msg"
        @info "info msg"
        @warn "warn msg"
        @error "error msg"
    end
    json = [JSON3.read(x) for x in eachline(seekstart(io))]
    @test json[1].level == "debug"
    @test json[1].msg == "debug msg"
    @test json[2].level == "info"
    @test json[2].msg == "info msg"
    @test json[3].level == "warn"
    @test json[3].msg == "warn msg"
    @test json[4].level == "error"
    @test json[4].msg == "error msg"
    for i in 1:4
        @test json[i].line isa Int
        @test json[i].module == "Main"
        @test isempty(json[i].kwargs)
    end

    io = IOBuffer()
    with_logger(FormatLogger(JSON(), io)) do
        y = (1, 2)
        @info "info msg" x = [1, 2, 3] y
    end
    json = JSON3.read(seekstart(io))
    @test json.level == "info"
    @test json.msg == "info msg"
    @test json.module == "Main"
    @test json.line isa Int
    @test json.kwargs.x == "[1, 2, 3]"
    @test json.kwargs.y == "(1, 2)"

    # `recursive=true`
    io = IOBuffer()
    with_logger(FormatLogger(JSON(; recursive=true), io)) do
        y = (1, 2)
        @info "info msg" x = [1, 2, 3] y = Dict("hi" => Dict("hi2" => [1,2]))
    end
    json = JSON3.read(seekstart(io))
    @test json.level == "info"
    @test json.msg == "info msg"
    @test json.module == "Main"
    @test json.line isa Int
    @test json.kwargs.x == [1, 2, 3]
    @test json.kwargs.y == Dict(:hi => Dict(:hi2 => [1,2]))

    # Fallback to strings
    io = IOBuffer()
    with_logger(FormatLogger(JSON(; recursive=true), io)) do
        y = (1, 2)
        @info "info msg" x = [1, 2, 3] y = Dict("hi" => NoStructTypeDefined(1))
    end
    json = JSON3.read(seekstart(io))
    @test json.level == "info"
    @test json.msg == "info msg"
    @test json.module == "Main"
    @test json.line isa Int
    @test json.kwargs.x == "[1, 2, 3]"
    y = json.kwargs.y
    must_have = ("Dict", "\"hi\"", "=>", "NoStructTypeDefined(1)")
    @test all(h -> occursin(h, y), must_have) # avoid issues with printing changing with versions
    @test json.kwargs[Symbol("LoggingFormats.FormatError")] == "ArgumentError: NoStructTypeDefined doesn't have a defined `StructTypes.StructType`"

    # Test logging exceptions
    for recursive in (false, true)
        # no stacktrace
        io = IOBuffer()
        with_logger(FormatLogger(JSON(; recursive=recursive), io)) do
            try
                throw(ArgumentError("no"))
            catch e
                @error "Oh no" exception = e
            end
        end
        logs = JSON3.read(seekstart(io))
        @test logs["msg"] == "Oh no"
        @test logs["kwargs"]["exception"] == "ArgumentError: no"

        # stacktrace
        io = IOBuffer()
        with_logger(FormatLogger(JSON(; recursive=recursive), io)) do
            try
                throw(ArgumentError("no"))
            catch e
                @error "Oh no" exception = (e, catch_backtrace())
            end
        end
        logs = JSON3.read(seekstart(io))
        @test logs["msg"] == "Oh no"

        @test occursin("ArgumentError: no", logs["kwargs"]["exception"])
        # Make sure we get a stacktrace out:
        @test occursin(r"ArgumentError: no\nStacktrace:\s* \[1\]",
                    logs["kwargs"]["exception"])
    end
end

@testset "logfmt" begin
    io = IOBuffer()
    with_logger(FormatLogger(LogFmt(), io)) do
        @debug "debug msg"
        @info "info msg" _file="file with space.jl"
        @warn "msg with \"quotes\""
        @error "error msg with nothings" _module=nothing _file=nothing __line=nothing
        @error :notstring x = [1, 2, 3] y = "hello\" \"world"
    end
    strs = collect(eachline(seekstart(io)))
    @test occursin("level=debug msg=\"debug msg\" module=Main", strs[1])
    @test occursin("file=\"", strs[1])
    @test occursin("group=\"", strs[1])
    @test occursin("level=info msg=\"info msg\" module=Main", strs[2])
    @test occursin("file=\"file with space.jl\"", strs[2])
    @test occursin("group=\"file with space\"", strs[2])
    @test occursin("level=warn msg=\"msg with \\\"quotes\\\"\" module=Main", strs[3])
    @test occursin("file=\"", strs[3])
    @test occursin("group=\"", strs[3])
    @test occursin("level=error msg=\"error msg with nothings\" module=nothing", strs[4])
    @test occursin("file=\"nothing\"", strs[4])
    @test occursin("line=\"nothing\"", strs[4])
    @test occursin("level=error msg=\"notstring\" module=Main", strs[5])
    @test occursin("x=\"[1, 2, 3]\" y=\"hello\\\" \\\"world\"", strs[5])
end


text:
module Sand_02
using LinearAlgebra, Random
#using StaticArrays
"""
    Sand 02: advection of a scalar, with abrupt step in velocity (as a function of position)
"""

mutable struct ring
    N::Integer              # number of cells
    Δt::Float64             # time step
    X_::Array{Float64,1}    # positions (nodes)
    ΔX_::Array{Float64,1}   # cell sizes (cell)
    U_::Array{Float64,1}    # velocity (on nodes)
    ∇U_::Array{Float64,1}   # divergence of velocity (on cells)
    Σ_::Array{Float64,1}    # surface density
    S_::Array{Float64,1}    # integral of the above
    t::Float64              # time
end
function ring(N=100; U = 1.0, xi=0.0, xf=10.0)
    xtemp_ = collect(range(xi,xf,length=N+1))
    xtemp_[2:end-1] += (Random.rand(N-1).-0.5) * (xf-xi)/N * 0.9 # 0.8 is somewhat arbitrary
    @assert prod(diff(xtemp_).>0)  # checks to make sure xtemp_ monotonic
    X_ = xtemp_#[2:end]
    ΔX_ = diff(xtemp_)
    U_ = U * ones(N+1)
    """ XXX putting step-down in speed XXX """
    U_[3*N÷4:end] .*= 0.5
    U_[1:N÷4]     .*= 0.5
    ∇U_ = diff(U_) ./ ΔX_
    # init w/top-hat density:
    Σ_ = ones(N); Σ_[1] = 0.0;
    Σ_[1:(2*N÷5)] .= 0.0; Σ_[(3*N÷5):end] .= 0.0
    S_ = Σ_ .* ΔX_
    #
    Δt = 0.5 / max(maximum(U_[1:end-1]./ΔX_),-minimum(U_[2:end]./ΔX_))
    #
    return ring(N,Δt,X_,ΔX_,U_,∇U_,Σ_,S_,0.0)
end

function step!(r::ring; Δt=r.Δt)
    Σ_=r.Σ_; N = r.N; ΔX_ = r.ΔX_; U_ = r.U_; S_ = r.S_
    #
    ΔΣ_ = cat(Σ_[1]-Σ_[end], diff(Σ_), Σ_[1]-Σ_[end], dims=1)
    S_ = Σ_ .* ΔX_
    U⁺_ = U_[1:end-1] .* (U_[1:end-1].>0) # right-traveling
    U⁻_ = U_[2:end  ] .* (U_[2:end  ].<0) # left-traveling
    S_ -= Δt * U⁺_ .* ΔΣ_[1:end-1]
    S_ -= Δt * U⁻_ .* ΔΣ_[2:end]
    # This part kinda sucks b/c it doesn't conserve very well:
    Σ_ = S_ ./ ΔX_
    #Σ_ += -Δt * r.∇U_ .* Σ_ # <- forward Euler
    Σ_ ./= (1 .+ Δt*r.∇U_)  # <- backward Euler  - conserves a bit better
    S_ = Σ_ .* ΔX_
    #
    r.S_ = S_
    r.Σ_ = S_ ./ ΔX_
    r.t += Δt
    print(sum(S_),"\n")
    return sum(S_)
end

function tstep!(r::ring,Δt)
    M = floor(Δt / r.Δt)
    Δt_last = Δt - M*r.Δt
    @assert Δt_last ≥ 0
    for i in 1:M
        step!(r)
    end
    step!(r, Δt = Δt_last)
end

end#module

max_stars_repo_name: diakosmos/razor

max_issues_repo_name: diakosmos/razor

max_forks_repo_name: diakosmos/razor


text:
"""
    abstract type GNNLayer end

An abstract type from which graph neural network layers are derived.

See also [`GNNChain`](@ref).
"""
abstract type GNNLayer end

# Forward pass with graph-only input.
# To be specialized by layers also needing edge features as input (e.g. NNConv).
(l::GNNLayer)(g::GNNGraph) = GNNGraph(g, ndata=l(g, node_features(g)))

"""
    GNNChain(layers...)
    GNNChain(name = layer, ...)

Collects multiple layers / functions to be called in sequence
on given input graph and input node features.

It allows to compose layers in a sequential fashion as `Flux.Chain`
does, propagating the output of each layer to the next one.
In addition, `GNNChain` handles the input graph as well, providing it
as a first argument only to layers subtyping the [`GNNLayer`](@ref) abstract type.

`GNNChain` supports indexing and slicing, `m[2]` or `m[1:end-1]`,
and if names are given, `m[:name] == m[1]` etc.

# Examples

```juliarepl
julia> m = GNNChain(GCNConv(2=>5), BatchNorm(5), x -> relu.(x), Dense(5, 4));

julia> x = randn(Float32, 2, 3);

julia> g = GNNGraph([1,1,2,3], [2,3,1,1]);

julia> m(g, x)
4×3 Matrix{Float32}:
  0.157941    0.15443     0.193471
  0.0819516   0.0503105   0.122523
  0.225933    0.267901    0.241878
 -0.0134364  -0.0120716  -0.0172505
```
"""
struct GNNChain{T} <: GNNLayer
  layers::T

  GNNChain(xs...) = new{typeof(xs)}(xs)

  function GNNChain(; kw...)
    :layers in Base.keys(kw) && throw(ArgumentError("a GNNChain cannot have a named layer called `layers`"))
    isempty(kw) && return new{Tuple{}}(())
    new{typeof(values(kw))}(values(kw))
  end
end

@forward GNNChain.layers Base.getindex, Base.length, Base.first, Base.last,
    Base.iterate, Base.lastindex, Base.keys

Flux.functor(::Type{<:GNNChain}, c) = c.layers, ls -> GNNChain(ls...)

applylayer(l, g::GNNGraph, x) = l(x)
applylayer(l::GNNLayer, g::GNNGraph, x) = l(g, x)

# Handle Flux.Parallel
applylayer(l::Parallel, g::GNNGraph, x::AbstractArray) = mapreduce(f -> applylayer(f, g, x), l.connection, l.layers)
applylayer(l::Parallel, g::GNNGraph, xs::Vararg{<:AbstractArray}) = mapreduce((f, x) -> applylayer(f, g, x), l.connection, l.layers, xs)
applylayer(l::Parallel, g::GNNGraph, xs::Tuple) = applylayer(l, g, xs...)

applychain(::Tuple{}, g::GNNGraph, x) = x
applychain(fs::Tuple, g::GNNGraph, x) = applychain(tail(fs), g, applylayer(first(fs), g, x))

(c::GNNChain)(g::GNNGraph, x) = applychain(Tuple(c.layers), g, x)

Base.getindex(c::GNNChain, i::AbstractArray) = GNNChain(c.layers[i]...)
Base.getindex(c::GNNChain{<:NamedTuple}, i::AbstractArray) =
    GNNChain(; NamedTuple{Base.keys(c)[i]}(Tuple(c.layers)[i])...)

function Base.show(io::IO, c::GNNChain)
    print(io, "GNNChain(")
    _show_layers(io, c.layers)
    print(io, ")")
end
_show_layers(io, layers::Tuple) = join(io, layers, ", ")
_show_layers(io, layers::NamedTuple) = join(io, ["$k = $v" for (k, v) in pairs(layers)], ", ")


text:
# See https://udel.edu/~arthij/MD.pdf for information on forces
# See https://arxiv.org/pdf/1401.1181.pdf for applying forces to atoms
# See Gromacs manual for other aspects of forces

export
    LennardJones,
    force!,
    SoftSphere,
    Coulomb,
    Gravity,
    HarmonicBond,
    HarmonicAngle,
    Torsion

"Square of the non-bonded interaction distance cutoff in nm^2."
const sqdist_cutoff_nb = 1.0 ^ 2

"The constant for Coulomb interaction, 1/(4*π*ϵ0*ϵr)."
const coulomb_const = 138.935458 / 70.0 # Treat ϵr as 70 for now

"""
    LennardJones(nl_only)

The Lennard-Jones 6-12 interaction.
"""
struct LennardJones <: GeneralInteraction
    nl_only::Bool
end

LennardJones() = LennardJones(false)

"""
    force!(forces, interaction, simulation, atom_i, atom_j)

Update the force for an atom pair in response to a given interation type.
Custom interaction types should implement this function.
"""
function force! end

@fastmath @inbounds function force!(forces,
                                    inter::LennardJones,
                                    s::Simulation,
                                    i::Integer,
                                    j::Integer)
    if iszero(s.atoms[i].σ) || iszero(s.atoms[j].σ) || i == j
        return
    end
    σ = sqrt(s.atoms[i].σ * s.atoms[j].σ)
    ϵ = sqrt(s.atoms[i].ϵ * s.atoms[j].ϵ)
    dr = vector(s.coords[i], s.coords[j], s.box_size)
    r2 = sum(abs2, dr)
    r2 > sqdist_cutoff_nb && return
    invr2 = inv(r2)
    six_term = (σ ^ 2 * invr2) ^ 3
    # Limit this to 100 as a fudge to stop it exploding
    f = min((24ϵ * invr2) * (2 * six_term ^ 2 - six_term), 100)
    fdr = f * dr
    forces[i] -= fdr
    forces[j] += fdr
    return nothing
end

"""
    SoftSphere(nl_only)

The soft-sphere potential.
"""
struct SoftSphere <: GeneralInteraction
    nl_only::Bool
end

@fastmath @inbounds function force!(forces,
                                    inter::SoftSphere,
                                    s::Simulation,
                                    i::Integer,
                                    j::Integer)
    if iszero(s.atoms[i].σ) || iszero(s.atoms[j].σ) || i == j
        return
    end
    σ = sqrt(s.atoms[i].σ * s.atoms[j].σ)
    ϵ = sqrt(s.atoms[i].ϵ * s.atoms[j].ϵ)
    dr = vector(s.coords[i], s.coords[j], s.box_size)
    r2 = sum(abs2, dr)
    r2 > sqdist_cutoff_nb && return
    invr2 = inv(r2)
    six_term = (σ ^ 2 * invr2) ^ 3
    # Limit this to 100 as a fudge to stop it exploding
    f = min((24ϵ * invr2) * 2 * six_term ^ 2, 100)
    fdr = f * dr
    forces[i] -= fdr
    forces[j] += fdr
    return nothing
end

"""
    Coulomb(nl_only)

The Coulomb electrostatic interaction.
"""
struct Coulomb <: GeneralInteraction
    nl_only::Bool
end

@fastmath @inbounds function force!(forces,
                                    inter::Coulomb,
                                    s::Simulation,
                                    i::Integer,
                                    j::Integer)
    i == j && return
    dr = vector(s.coords[i], s.coords[j], s.box_size)
    r2 = sum(abs2, dr)
    r2 > sqdist_cutoff_nb && return
    T = typeof(r2)
    f = (T(coulomb_const) * s.atoms[i].charge * s.atoms[j].charge) / sqrt(r2 ^ 3)
    fdr = f * dr
    forces[i] -= fdr
    forces[j] += fdr
    return nothing
end

"""
    Gravity(nl_only, G)

The gravitational interaction.
"""
struct Gravity{T} <: GeneralInteraction
    nl_only::Bool
    G::T
end

function force!(forces,
                inter::Gravity,
                s::Simulation,
                i::Integer,
                j::Integer)
    i == j && return
    dr = vector(s.coords[i], s.coords[j], s.box_size)
    f = -inter.G * s.atoms[i].mass * s.atoms[j].mass * inv(sum(abs2, dr))
    fdr = f * normalize(dr)
    forces[i] -= fdr
    forces[j] += fdr
    return nothing
end

"""
    HarmonicBond(i, j, b0, kb)

A harmonic bond between two atoms.
"""
struct HarmonicBond{T} <: SpecificInteraction
    i::Int
    j::Int
    b0::T
    kb::T
end

function force!(forces,
                b::HarmonicBond,
                s::Simulation)
    ab = vector(s.coords[b.i], s.coords[b.j], s.box_size)
    c = b.kb * (norm(ab) - b.b0)
    f = c * normalize(ab)
    forces[b.i] += f
    forces[b.j] -= f
    return nothing
end

"""
    HarmonicAngle(i, j, k, th0, cth)

A bond angle between three atoms.
"""
struct HarmonicAngle{T} <: SpecificInteraction
    i::Int
    j::Int
    k::Int
    th0::T
    cth::T
end

# Sometimes domain error occurs for acos if the value is > 1.0 or < -1.0
acosbound(x::Real) = acos(clamp(x, -1, 1))

function force!(forces,
                a::HarmonicAngle,
                s::Simulation)
    ba = vector(s.coords[a.j], s.coords[a.i], s.box_size)
    bc = vector(s.coords[a.j], s.coords[a.k], s.box_size)
    pa = normalize(ba × (ba × bc))
    pc = normalize(-bc × (ba × bc))
    angle_term = -a.cth * (acosbound(dot(ba, bc) / (norm(ba) * norm(bc))) - a.th0)
    fa = (angle_term / norm(ba)) * pa
    fc = (angle_term / norm(bc)) * pc
    fb = -fa - fc
    forces[a.i] += fa
    forces[a.j] += fb
    forces[a.k] += fc
    return nothing
end

"""
    Torsion(i, j, k, l, f1, f2, f3, f4)

A dihedral torsion angle between four atoms.
"""
struct Torsion{T} <: SpecificInteraction
    i::Int
    j::Int
    k::Int
    l::Int
    f1::T
    f2::T
    f3::T
    f4::T
end

function force!(forces,
                d::Torsion,
                s::Simulation)
    ba = vector(s.coords[d.j], s.coords[d.i], s.box_size)
    bc = vector(s.coords[d.j], s.coords[d.k], s.box_size)
    dc = vector(s.coords[d.l], s.coords[d.k], s.box_size)
    p1 = normalize(ba × bc)
    p2 = normalize(-dc × -bc)
    θ = atan(dot((-ba × bc) × (bc × -dc), normalize(bc)), dot(-ba × bc, bc × -dc))
    angle_term = (d.f1*sin(θ) - 2*d.f2*sin(2*θ) + 3*d.f3*sin(3*θ)) / 2
    fa = (angle_term / (norm(ba) * sin(acosbound(dot(ba, bc) / (norm(ba) * norm(bc)))))) * p1
    # fd clashes with a function name
    f_d = (angle_term / (norm(dc) * sin(acosbound(dot(bc, dc) / (norm(bc) * norm(dc)))))) * p2
    oc = bc / 2
    tc = -(oc × f_d + (-dc × f_d) / 2 + (ba × fa) / 2)
    fc = (1 / dot(oc, oc)) * (tc × oc)
    fb = -fa - fc - f_d
    forces[d.i] += fa
    forces[d.j] += fb
    forces[d.k] += fc
    forces[d.l] += f_d
    return nothing
end

"""
    Mie(m,n,nl_only)

The Mie generalized interaction.
Mie force with m = 6, n = 12 is equivalent to Lennard Jones potential
"""
struct Mie{T<:Real} <: GeneralInteraction
    m::T
    n::T
    mn::T #this is to save on computation
    nl_only::Bool
end

Mie(m,n) = Mie(m,n,(n/(n-m))*(n/m)^(m/(n-m)),false)

@fastmath @inbounds function force!(forces,
    inter::Mie{T},
    s::Simulation,
    i::Integer,
    j::Integer) where T
    if iszero(s.atoms[i].σ) || iszero(s.atoms[j].σ) || i == j
    return
    end
    m = inter.m
    n = inter.n
    σ = sqrt(s.atoms[i].σ * s.atoms[j].σ)
    ϵ = sqrt(s.atoms[i].ϵ * s.atoms[j].ϵ)
    dr = vector(s.coords[i], s.coords[j], s.box_size)
    rr = norm(dr)
    abs2(rr) > sqdist_cutoff_nb && return
    #derivative obtained via wolfram
    const_mn = inter.mn*ϵ/rr
    σ_r = σ/rr
    f= m*(σ_r)^m - n*(σ_r)^n
    f = f*const_mn
    # Limit this to 100 as a fudge to stop it exploding
    f = min(f, 100)
    fdr = f * dr
    forces[i] -= fdr
    forces[j] += fdr
    return nothing
end


text:
@testset "Updates" begin

    for i in testplatforms

        #Update Local
        m = @agent(
            3,

            [w]::Local,

            UpdateLocal =
            begin
                x += 1
                z += 2
                w = Uniform(1.,2.)
            end
        )
        mo = compile(m, platform = i, debug=false)
        # println(mo.program)

        @test_nowarn begin

            com = Community(mo,N=1000)
            com.x .= 0.
            com.y .= 0.
            com.z .= 0.
            com.w = rand(Uniform(1,2),com.N)
            comt = mo.evolve(com,dt=1.,tMax=100)#,t=0.,N=com.N,nMax=com.N)

            for i in 1:10
                if comt[i].x != (i-1)*ones(com.N) println("x ", i, " ", comt[i].x);error() end
                if comt[i].y != 7*zeros(com.N) println("y ", i, " ", comt[i].y);error() end
                if comt[i].z != 2*(i-1)*ones(com.N) println("z ", i, " ", comt[i].z);error() end
                if !prod(Bool.((comt[i].w .>= 1) .* (comt[i].w .<= 2))) println("w ", i, " ", comt[i].w);error() end
            end
        end

        #Update Identity
        m = @agent(
            3,

            [w1,w2]::Identity,

            UpdateLocal =
            begin
                nothing
                x += 1
                z += 2
                w1 = floor(Uniform(1.,2.))
                w2 = ceil(Uniform(1.,2.))
            end
        )
        mo = compile(m, platform = i, debug=false)
        #println(mo.program)

        @test_nowarn begin

            com = Community(mo,N=10)
            com.x .= 0
            com.y .= 0
            com.z .= 0
            com.w1 .= 1
            com.w2 .= 2
            comt = mo.evolve(com,dt=1.,tMax=10)#,t=0.,N=com.N,nMax=com.N)

            for i in 1:10
                if comt[i].x != (i-1)*ones(Int,10) println("x ", i, " ", comt[i].x);error() end
                if comt[i].y != 7*zeros(Int,10) println("y ", i, " ", comt[i].y);error() end
                if comt[i].z != 2*(i-1)*ones(Int,10) println("z ", i, " ", comt[i].z);error() end
                if !prod((comt[i].w1 .== 1)) println("w ", i, " ", comt[i].w1);error() end
                if !prod((comt[i].w2 .== 2)) println("w ", i, " ", comt[i].w2);error() end
            end
        end

        #Update Global
        m = @agent(
            3,

            [x1,y1,z1,w]::Global,

            UpdateGlobal =
            begin
                x1 += 1
                z1 += 2
                w = Uniform(1.,2.)
            end
        )
        mo = compile(m, platform = i, debug=false)
        #println(mo.program)

        @test_nowarn begin

            com = Community(mo,N=10)
            com.x1 = 0.
            com.y1 = 0.
            com.z1 = 0.
            com.w = 1.2
            comt = mo.evolve(com,dt=1.,tMax=10)#,t=0.,N=com.N,nMax=com.N)

            for i in 1:10
                if comt[i].x1 != 1. *(i-1) println("x ", i, " ", comt[i].x);error() end
                if comt[i].y1 != 0. println("y ", i, " ", comt[i].y);error() end
                if comt[i].z1 != 2. *(i-1) println("z ", i, " ", comt[i].z);error() end
                if !prod(Bool((comt[i].w .>= 1) * (comt[i].w .<= 2))) println("w ", i, " ", comt[i].w);error() end
            end
        end

        #Update Local Interaction
        m = @agent(
            3,

            UpdateLocalInteraction =
            begin
                x += 1
                z += 2
            end
        )
        mo = compile(m, platform = i, debug=false)
        #println(mo.program)

        @test_nowarn begin

            com = Community(mo,N=10)
            com.x .= 0.
            com.y .= 0.
            com.z .= 0.
            comt = mo.evolve(com,dt=1.,tMax=10)#,t=0.,N=com.N,nMax=com.N)

            for i in 1:10
                if comt[i].x != 10*ones(10) println("x ", i, " ", comt[i].x);error() end
                if comt[i].y != zeros(10) println("y ", i, " ", comt[i].y);error() end
                if comt[i].z != 20*ones(10) println("z ", i, " ", comt[i].z);error() end
            end
        end

    end

end

max_stars_repo_name: gatocor/embryogenesisJulia

max_issues_repo_name: gatocor/embryogenesisJulia

max_forks_repo_name: gatocor/embryogenesisJulia


text:
#### IMPLEMENTATION OF JEFF SKILLINGS' NESTED SAMPLING ALGORITHM ####
function nested_step!(e::IPM_Ensemble, instruction::Permute_Instruct)
    N = length(e.models) #number of sample models/particles on the posterior surface
    i = length(e.log_Li) #iterate number, index for last values
    j = i+1 #index for newly pushed values

    e.contour, least_likely_idx = findmin([model.log_Li for model in e.models])
    Li_model = e.models[least_likely_idx]

    #SELECT NEW MODEL, SAVE TO ENSEMBLE DIRECTORY, CREATE RECORD AND PUSH TO ENSEMBLE
    model_selected=false; step_report=0
    while !model_selected
        candidate,step_report=permute_IPM(e, instruction)
        if !(candidate===nothing)
            if !(candidate.log_Li in [m.log_Li for m in e.models])
                model_selected=true

                new_model_record = Model_Record(string(e.path,'/',e.model_counter), candidate.log_Li);
                e.sample_posterior && push!(e.posterior_samples, Li_model)#if sampling posterior, push the model record to the ensemble's posterior samples vector
                deleteat!(e.models, least_likely_idx)
                push!(e.models, new_model_record);

                final_model=ICA_PWM_Model(string(e.model_counter), candidate.origin, candidate.sources, candidate.source_length_limits, candidate.mix_matrix, candidate.log_Li, candidate.permute_blacklist)
                serialize(new_model_record.path, final_model)

                e.model_counter +=1
            end
        else
            push!(e.models, Li_model)
            return 1, step_report
        end
    end

    #UPDATE ENSEMBLE QUANTITIES
    push!(e.log_Li, e.contour) #log likelihood of the least likely model - the current ensemble ll contour at Xi
    push!(e.log_Xi, -i/N) #log Xi - crude estimate of the iterate's enclosed prior mass
    push!(e.log_wi, logsubexp(e.log_Xi[i], -j/N′) - log(2)) #log width of prior mass spanned by the last step-trapezoidal approx
    push!(e.log_Liwi, lps(e.log_Li[j],e.log_wi[j])) #log likelihood + log width = increment of evidence spanned by iterate
    push!(e.log_Zi, logaddexp(e.log_Zi[i],e.log_Liwi[j]))    #log evidence
    #information- dimensionless quantity. cf. MultiNest @ https://github.com/farhanferoz/MultiNest/blob/master/MultiNest_v3.12/nested.F90- MultiNest now gives info only at the final step
    Hj=lps(
        (exp(lps(e.log_Liwi[j],-e.log_Zi[j])) * e.log_Li[j]), #information contribution of this step
        (exp(lps(e.log_Zi[i],-e.log_Zi[j])) * lps(e.Hi[i],e.log_Zi[i])), #rescale last information by evidence
        -e.log_Zi[j])
    Hj === -Inf ? push!(e.Hi,0.) : push!(e.Hi, Hj) #prevent problems from early strings of models with -Inf log likelihoods

    return 0, step_report
end

function nested_step!(e::IPM_Ensemble, model_chan::RemoteChannel, wk_mon::Worker_Monitor)
    N = length(e.models)+1 #number of sample models/particles on the posterior surface- +1 as one has been removed in the distributed dispatch for converge_ensemble
    i = length(e.log_Li) #iterate number, index for last values
    j = i+1 #index for newly pushed values

    e.contour, least_likely_idx = findmin([model.log_Li for model in e.models])
    Li_model = e.models[least_likely_idx]

    #SELECT NEW MODEL, SAVE TO ENSEMBLE DIRECTORY, CREATE RECORD AND REPLACE LEAST LIKELY MODEL
    model_selected=false;wk=0;step_report=0
    while !model_selected
        @async wait(model_chan)
        candidate,wk,step_report = take!(model_chan)
        if !(candidate=="quit")
            if (candidate.log_Li > e.contour) && !(candidate.log_Li in [m.log_Li for m in e.models])
                model_selected=true

                new_model_record = Model_Record(string(e.path,'/',e.model_counter), candidate.log_Li);
                e.sample_posterior && push!(e.posterior_samples, Li_model)#if sampling posterior, push the model record to the ensemble's posterior samples vector
                deleteat!(e.models, least_likely_idx)
                push!(e.models, new_model_record);

                final_model=ICA_PWM_Model(string(e.model_counter), candidate.origin, candidate.sources, candidate.source_length_limits, candidate.mix_matrix, candidate.log_Li, candidate.permute_blacklist)
                serialize(new_model_record.path, final_model)

                e.model_counter +=1
            end
            update_worker_monitor!(wk_mon,wk,true)
        else
            update_worker_monitor!(wk_mon,wk,false)
            !any(wk_mon.persist) && (return 1,step_report)
        end
    end

    #UPDATE ENSEMBLE QUANTITIES
    push!(e.log_Li, minimum([model.log_Li for model in e.models])) #log likelihood of the least likely model - the current ensemble ll contour at Xi
    push!(e.log_Xi, -i/N) #log Xi - crude estimate of the iterate's enclosed prior mass
    push!(e.log_wi, lps(e.log_Xi[i], -((j+1)/N)-log(2))) #log width of prior mass spanned by the last step-trapezoidal approx
    push!(e.log_Liwi, lps(e.log_Li[j],e.log_wi[j])) #log likelihood + log width = increment of evidence spanned by iterate
    push!(e.log_Zi, logaddexp(e.log_Zi[i],e.log_Liwi[j]))    #log evidence

    #information- dimensionless quantity
    push!(e.Hi, lps(
            (exp(lps(e.log_Liwi[j],-e.log_Zi[j])) * e.log_Li[j]), #term1
            (exp(lps(e.log_Zi[i],-e.log_Zi[j])) * lps(e.Hi[i],e.log_Zi[i])), #term2
            -e.log_Zi[j])) #term3

    return 0, step_report
end


text:
using DiffResults, ForwardDiff, Distributions, StatsBase

export BareParticleContainer

struct BareParticleContainer <: ParticleContainer
    X # Particles
    w # weights
    ℓ # Estimate of the log likelihood
end

Base.size(ps::BareParticleContainer) = (size(ps.X,2),)
Base.length(ps::BareParticleContainer) = prod(size(ps))
Base.getindex(ps::BareParticleContainer,i::Int) = ps.X[:,i]
Base.setindex!(ps::BareParticleContainer,v,i::Int) = ps.X[:,i] = v
Base.IndexStyle(ps::BareParticleContainer) = IndexLinear()

resample!(ps::BareParticleContainer,m,y,t,θ,sample_function=resample_stratified) = resample!(BareParticleContainer,ps.X,ps.w,ps.w[:],m.ξ,y,t,θ,length(ps.w),sample_function)

function resample!(::Type{BareParticleContainer},X,w,ξs,ξ,y,t,θ,N,sample_function)
    for i in 1:N
        ξs[i] += ξ(X[:,i,t],y,t,θ)
    end
    Σ = sum(exp,ξs)
    idx = sample_function(exp.(ξs)./Σ,N)

    X[:,:,1:t] .= X[:,idx,1:t]
    w[:] .= w[idx]
    ξs[idx],Σ
end

propagate!(ps::BareParticleContainer,ξ,Σ,m,y,t,θ,λ) = propagate!(BareParticleContainer,ps.X,ps.w,ps.ℓ,ξ,Σ,m.f,m.g,m.q,m.dp,y,t,θ,length(ps.w),λ)

function propagate!(::Type{BareParticleContainer},X,w,ℓ,ξ,Σ,f,g,q,dp,y,t,θ,N,λ)
    for i in 1:N
        dq = q(X[:,i,t],y,t,θ)
        X[:,i,t+1] = rand(dq)

        w[i] += logpdf(f(X[:,i,t],t,θ),X[:,i,t+1])
        w[i] += logpdf(g(X[:,i,t+1],t,θ),y)
        #w[i] *= pdf(f(X[:,i],t,θ),Xn)
        #w[i] *= pdf(g(Xn,t,θ),y)
        w[i] -= ξ[i]
        w[i] -= logpdf(dq,X[:,i,t+1])

    end
    ℓn = log(sum(exp,w))
    ℓ[t+1] = ℓ[t] + ℓn - log(N) + log(Σ)
    w .-= ℓn
end

function step(ps::BareParticleContainer,m::ProposalStateSpaceModel,y,t,θ,λ=0.95;Nthreshold=0.5,sample_function=resample_stratified)

    # Resample

    #rs = Resampler(ps,m,y,t,θ,sample_function)

    ξs,Σ = resample!(ps,m,y,t,θ,sample_function)

    # Propagate

    #ps2 = ScoreParticleContainer([propose(t,x,w,ξ,α,ps.S,y,θ,m.f,m.g,m.q,m.dp,ps.dg,λ) for (x,w,ξ,α) in rs],ps.ℓ,ps.dg,sum(rs.ξ),λ)

    propagate!(ps,ξs,Σ,m,y,t,θ,λ)

    #normalize_weights!(ps2)

    ps
end

function bare_filter(m::ProposalStateSpaceModel,Y,θ,N;λ=0.95,Nthreshold=0.5,sample_function=resample_stratified,verbose=false)

    T = length(Y)

    μ0 = m.μ0
    f = m.f
    g = m.g

    ps = initialize_particles(BareParticleContainer,m,θ,Y,N,λ)

    for t in 1:T
        if verbose
            println("Step: ",t)
            println("ESS: ",1 ./ sum(abs2,exp.(ps.w)))
        end
        ps = step(ps,m,Y[t],t,θ,λ,Nthreshold=Nthreshold,sample_function=sample_function)
    end

    ps
end

function initialize_particles(::Type{BareParticleContainer},μ0,f,g,θ,Y,N,T,λ;save_filter=false)

    # We should probably think of μ0 as a distribution for p(x_0)
    # not one for p(x_1)
    # Initialization gives X0, w = 1/N
    #x0 = rand(μ0(θ),N)

    x0 = rand(μ0(θ),N)
    X0 = zeros(size(x0,1),N,T+1)
    X0[:,:,1] = x0[:,:]

    w = log.(ones(N)/N)
    ℓ = zeros(T+1)
    #w = [pdf(g(X0[:,end,i],1,θ),y)/N for i in 1:N]
    #ℓ = sum(w)
    #w ./= ℓ

    BareParticleContainer(X0,w,ℓ)
end
initialize_particles(::Type{BareParticleContainer},m::ProposalStateSpaceModel,θ,Y,N,λ) = initialize_particles(BareParticleContainer,m.μ0,m.f,m.g,θ,Y,N,length(Y),λ)

#######################
# Filtering distribution
# This does not save trajectories, but saves weights and particles at each time step
# Thus each corresponds to the filtering distribution.

struct FilteringParticleContainer <: ParticleContainer
    X # Particles
    w # weights
    ℓ # Estimate of the log likelihood
end

Base.size(ps::FilteringParticleContainer) = (size(ps.X,2),)
Base.length(ps::FilteringParticleContainer) = prod(size(ps))
Base.getindex(ps::FilteringParticleContainer,i::Int) = ps.X[:,i]
Base.setindex!(ps::FilteringParticleContainer,v,i::Int) = ps.X[:,i] = v
Base.IndexStyle(ps::FilteringParticleContainer) = IndexLinear()

resample!(ps::FilteringParticleContainer,m,y,t,θ,sample_function=resample_stratified) = resample!(FilteringParticleContainer,ps.X,ps.w,ps.w[:,t],m.ξ,y,t,θ,size(ps.w,1),sample_function)

function resample!(::Type{FilteringParticleContainer},X,w,ξs,ξ,y,t,θ,N,sample_function)
    for i in 1:N
        ξs[i] += ξ(X[:,i,t],y,t,θ)
    end
    Σ = sum(exp,ξs)
    idx = sample_function(exp.(ξs)./Σ,N)

    X[:,:,t] .= X[:,idx,t]
    w[:,t] .= w[idx,t]
    ξs[idx],Σ
end

propagate!(ps::FilteringParticleContainer,ξ,Σ,m,y,t,θ,λ) = propagate!(FilteringParticleContainer,ps.X,ps.w,ps.ℓ,ξ,Σ,m.f,m.g,m.q,m.dp,y,t,θ,size(ps.w,1),λ)

function propagate!(::Type{FilteringParticleContainer},X,w,ℓ,ξ,Σ,f,g,q,dp,y,t,θ,N,λ)
    for i in 1:N
        dq = q(X[:,i,t],y,t,θ)
        X[:,i,t+1] = rand(dq)

        w[i,t+1] = w[i,t]
        w[i,t+1] += logpdf(f(X[:,i,t],t,θ),X[:,i,t+1])
        w[i,t+1] += logpdf(g(X[:,i,t+1],t,θ),y)
        #w[i] *= pdf(f(X[:,i],t,θ),Xn)
        #w[i] *= pdf(g(Xn,t,θ),y)
        w[i,t+1] -= ξ[i]
        w[i,t+1] -= logpdf(dq,X[:,i,t+1])

    end
    ℓn = log(sum(exp,w[:,t+1]))
    ℓ[t+1] = ℓ[t] + ℓn - log(N) + log(Σ)
    w[:,t+1] .-= ℓn
end

function step(ps::FilteringParticleContainer,m::ProposalStateSpaceModel,y,t,θ,λ=0.95;Nthreshold=0.5,sample_function=resample_stratified)

    # Resample

    #rs = Resampler(ps,m,y,t,θ,sample_function)

    ξs,Σ = resample!(ps,m,y,t,θ,sample_function)

    # Propagate

    #ps2 = ScoreParticleContainer([propose(t,x,w,ξ,α,ps.S,y,θ,m.f,m.g,m.q,m.dp,ps.dg,λ) for (x,w,ξ,α) in rs],ps.ℓ,ps.dg,sum(rs.ξ),λ)

    propagate!(ps,ξs,Σ,m,y,t,θ,λ)

    #normalize_weights!(ps2)

    ps
end

function filtering_filter(m::ProposalStateSpaceModel,Y,θ,N;λ=0.95,Nthreshold=0.5,sample_function=resample_stratified,verbose=false)

    T = length(Y)

    μ0 = m.μ0
    f = m.f
    g = m.g

    if verbose
        p = Progress(T,1)
    end

    ps = initialize_particles(FilteringParticleContainer,m,θ,Y,N,λ)

    for t in 1:T
        ps = step(ps,m,Y[t],t,θ,λ,Nthreshold=Nthreshold,sample_function=sample_function)
        if verbose
            next!(p)
        end
    end

    ps
end

function initialize_particles(::Type{FilteringParticleContainer},μ0,f,g,θ,Y,N,T,λ;save_filter=false)

    # We should probably think of μ0 as a distribution for p(x_0)
    # not one for p(x_1)
    # Initialization gives X0, w = 1/N
    #x0 = rand(μ0(θ),N)

    x0 = rand(μ0(θ),N)
    X0 = zeros(size(x0,1),N,T+1)
    X0[:,:,1] = x0[:,:]

    w = log.(ones(N,T+1)/N)
    ℓ = zeros(T+1)
    #w = [pdf(g(X0[:,end,i],1,θ),y)/N for i in 1:N]
    #ℓ = sum(w)
    #w ./= ℓ

    FilteringParticleContainer(X0,w,ℓ)
end
initialize_particles(::Type{FilteringParticleContainer},m::ProposalStateSpaceModel,θ,Y,N,λ) = initialize_particles(FilteringParticleContainer,m.μ0,m.f,m.g,θ,Y,N,length(Y),λ)


text:
#=
function init!(ev :: Env)

    # Create a master ModelBlock
    ev.id = 0
    master_mb = ModelBlock(ev)

    # Setup DataExchanger

    # Create ModelBlocks
    wkrs  =  workers()
    @sync for (i, p) in enumerate(wkrs)
        # We have P processors, N workers, N blocks
        # Block ids are numbered from 1 to N
        @spawnat p let
            global subocn = makeSubOcean(
                master_ocn  = ocn,
                block_id    = i,
                pull_fr_rng      = pull_fr_rngs[i],
                push_fr_rng      = push_fr_rngs[i],
                push_to_rng      = push_to_rngs[i],
                pull_fr_rng_bnd  = pull_fr_rngs_bnd[i, :],
                pull_to_rng_bnd  = pull_to_rngs_bnd[i, :],
                push_fr_rng_bnd  = push_fr_rngs_bnd[i, :],
                push_to_rng_bnd  = push_to_rngs_bnd[i, :],
            )
        end

end

=#

max_stars_repo_name: meteorologytoday/IOM

max_issues_repo_name: meteorologytoday/IOM

max_forks_repo_name: meteorologytoday/IOM


text:
using Test, Ingolstadt

out = cheers(3)

@test out == 5               # Test condition


text:
#implement recycle index
# bidm[i] is bond dimension between tensors[i] and tensors[i+1]
struct Bdims
    L::Int
    bdim::Array{Int,1}
end
getindex(bdim::Bdims,target::Int) = getindex(bdim.bdim,(target+bdim.L-1) % bdim.L+1)
setindex!(bdim::Bdims,value::Int,target::Int) = setindex!(bdim.bdim,value,(target+bdim.L-1) % bdim.L + 1)

max_stars_repo_name: qiyang-ustc/JMPS

max_issues_repo_name: qiyang-ustc/JMPS

max_forks_repo_name: qiyang-ustc/JMPS


text:
module AAABK2018

using NLsolve
using LinearAlgebra
using OrdinaryDiffEq
using DelayDiffEq
using FastGaussQuadrature
using Statistics
using NLopt
using Printf

include("functionsSolveModel.jl")
include("structures.jl")
include("optimalPolicy.jl")
include("socialPlanner.jl")

export Params, EqObj, solveBGP,eqfunc!, policy_opt, solveSocialPlanner

end # module


text:
using FANN

typealias _ANN Ptr{FANN.fann}

type ANN
	ann::_ANN
	nin::Int
	nout::Int
	function ANN(ann::_ANN)
		# additional constructor for when we load a net from file
		nin = ccall((:fann_get_num_input, libfann), Uint32, (Ptr{fann},),ann)
		nout = ccall((:fann_get_num_output, libfann), Uint32, (Ptr{fann},),ann)
		ann = new(ann, nin, nout)
		finalizer(ann, destroy)
		ann
	end
	function ANN(layers::Vector{Int}, activation::Vector{Symbol}; b::Float64=0.1, errorfunc::Symbol=:tanh, steepness::Float64=1.0)
		# Artificial Neural Network type
		#
		# Parameters
		# ----------
		# layers : number of nodes for each layer
		# activation : array of symbols with activation function for each layer excluded the input layer
		# b : [-b, b] defines the range for random initialisation of the weights
		# errorfunc : the error function used for training
		# steepness : The value of the steepness for all layers

		# activation function for hidden and output layers
		length(activation) == length(layers) - 1 || error("wrong dimension of activation function vector ")
		# allocate object
		ann = ccall((:fann_create_standard_array, libfann),
			        Ptr{fann},
			        (Uint32, Ptr{Uint32}),
			        length(layers),
			        pointer(uint32(layers)))
		if ann == C_NULL
            error("Error in fann_create_standard_array")
        end
        # set activation function and steepness for each layer
        for layer = 1:length(layers)-1
            ccall((:fann_set_activation_function_layer, libfann),
            	  Void,
            	  (Ptr{fann}, fann_activationfunc_enum, Cint),
            	  ann, act2uint(activation[layer]), layer)
            ccall((:fann_set_activation_steepness_layer, libfann),
       			  Void,
       			  (Ptr{fann}, fann_type, Cint),
       			  ann, steepness, layer)
        end
        # randomize weights
       	ccall((:fann_randomize_weights, libfann),
       		  Void,
       		  (Ptr{fann}, fann_type, fann_type),
       		  ann, -b, b)
       	# set train error function
       	ccall((:fann_set_train_error_function, libfann),
       		  Void,
       		  (Ptr{fann}, fann_errorfunc_enum),
       		  ann, errorfunc2uint(errorfunc))
		ANN(ann)
	end
end

Base.convert(::Type{_ANN}, ann::ANN) = ann.ann
destroy(ann::ANN) = ccall((:fann_destroy, libfann), Void, (_ANN,), ann)
Base.show(ann::ANN) = ccall((:fann_print_parameters, libfann), Void, (Ptr{fann},), ann)

# ~~~~~~ IO functions ~~~~~~~~
savenet(ann::ANN, file::ASCIIString) = ccall((:fann_save, libfann), Cint, (Ptr{fann}, Ptr{Uint8}), ann, file)
loadnet(file::ASCIIString) = ANN(ccall((:fann_create_from_file, libfann), Ptr{fann}, (Ptr{Uint8},), file))

# ~~~~~~~~~ Get/Set weights ~~~~~~~~~~~~~
nweights(ann::ANN) = int(ccall((:fann_get_total_connections, libfann), Uint32, (Ptr{fann},), ann))

function weights(ann::ANN)
	w = Array(Float64, nweights(ann))
	ccall((:fann_get_weights, libfann), Void, (Ptr{fann}, Ptr{fann_type}), ann, w)
	return w
end

function weights!(ann::ANN, weights::Vector{Float64})
	ccall((:fann_set_weights, libfann), Void, (Ptr{fann}, Ptr{fann_type}), ann, weights)
end

# ~~~~~~~~~ Tranining algorithms ~~~~~~~~~~~~~
# QuickProp training algorithm
function setup_qprop!(ann::ANN; mu::Float64=1.75, decay::Float64=-0.0001)
    ccall((:fann_set_training_algorithm, libfann), Void, (Ptr{fann}, fann_train_enum), ann, FANN_TRAIN_QUICKPROP)
    ccall((:fann_set_quickprop_mu,       libfann), Void, (Ptr{fann}, Cfloat), ann, mu)
    ccall((:fann_set_quickprop_decay,    libfann), Void, (Ptr{fann}, Cfloat), ann, decay)
end

# Rprop training algorithm, this is the default one
function setup_rprop!(ann::ANN; increase_factor::Float64=1.2, decrease_factor::Float64=0.5, delta_min::Float64=0.0, delta_max::Float64=50.0, delta_zero::Float64=0.1)
	ccall((:fann_set_training_algorithm, libfann), Void, (Ptr{fann}, fann_train_enum), ann, FANN_TRAIN_RPROP)
	ccall((:fann_set_rprop_increase_factor, libfann), Void, (Ptr{fann}, Cfloat), ann, increase_factor)
	ccall((:fann_set_rprop_decrease_factor, libfann), Void, (Ptr{fann}, Cfloat), ann, decrease_factor)
	ccall((:fann_set_rprop_delta_min,       libfann), Void, (Ptr{fann}, Cfloat), ann, delta_min)
	ccall((:fann_set_rprop_delta_max,       libfann), Void, (Ptr{fann}, Cfloat), ann, delta_max)
	ccall((:fann_set_rprop_delta_zero,      libfann), Void, (Ptr{fann}, Cfloat), ann, delta_zero)
end

# Standard backpropagation - batch or incremental
function setup_bprop!(ann::ANN; batch::Bool=true, learning_rate::Float64=0.7, learning_momentum::Float64=0.0)
	ccall((:fann_set_training_algorithm, libfann), Void, (Ptr{fann}, fann_train_enum), ann, batch == true ? FANN_TRAIN_BATCH : FANN_TRAIN_INCREMENTAL)
	ccall((:fann_set_learning_rate,      libfann), Void, (Ptr{fann}, Cfloat), ann, learning_rate)
	ccall((:fann_set_learning_momentum,  libfann), Void, (Ptr{fann}, Cfloat), ann, learning_momentum)
end

# Standard training function
function train!(ann::ANN, dset::DataSet; max_epochs::Int=100, desired_error::Float64=1e-5, epochs_between_reports::Int=10)
	# first check size
	checksizes(ann, dset)
	# verbosity
	isverbose = epochs_between_reports > 0
	# specify number of epoch in each iteration between reports
	epochs_between_reports < 0 && error("epochs_between_reports cannot be less than 0")
	nepochs = epochs_between_reports == 0 ? max_epochs : epochs_between_reports
	# debug header only if verbose
	isverbose && printprog(0 , mse(ann, dset))
	# init training epochs
	epoch = 0
	while true
		# train for nepochs
		ccall((:fann_train_on_data, libfann),
		  	  Void,
		  	  (Ptr{fann}, Ptr{fann_train_data}, Uint32, Uint32,  Cfloat),
		  	  ann, dset, nepochs, 0, desired_error)
		# update nepochs
		epoch += nepochs
		# get error
		tmse = mse(ann, dset)
		# print progress
		isverbose && printprog(epoch, tmse)
		# exit condition
		tmse < desired_error && break
		epoch >= max_epochs  && break
	end
end

# Early-stop training with validation test
function train!(ann::ANN, tset::DataSet, vset::DataSet; max_epochs::Int=100, desired_error::Float64=1e-5, epochs_between_checks::Int=10, minratio::Float64=0.95)
	# check this. Otherwise what is the point of checking?
	epochs_between_checks > 0 || error("error_between_checks must be greater than one")
	# check sizes
	checksizes(ann, tset); checksizes(ann, vset)
	# run ann on validation set to have first value
	vmse_curr, vmse_last= mse(ann, vset), 0.0f0
	# debug header
	printprog(0, mse(ann, tset), vmse_curr)
	for i = epochs_between_checks:epochs_between_checks:max_epochs
		train!(ann, tset;
			   max_epochs=epochs_between_checks,
			   desired_error=desired_error,
			   epochs_between_reports=0)
		# update validation error
		vmse_curr, vmse_last = mse(ann, vset), vmse_curr
		# update training error
		tmse_curr = mse(ann, tset)
		# print progress
		printprog(i, tmse_curr, vmse_curr)
		# check converged
		tmse_curr/vmse_curr < minratio && return nothing
	end
end

# Print output (can be done better though)
function printprog(epoch, tmse)
	epoch == 0 && println("Epoch  Train MSE\n~~~~~~~~~~~~~~~~")
	a = @sprintf "%5d" epoch
	b = @sprintf "%.3e" tmse
	println(a, "  ", b)
end

function printprog(epoch, tmse, vmse)
	epoch == 0 && println("Epoch  Train MSE  Valid MSE  Ratio\n~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~")
	a = @sprintf "%5d" epoch
	b = @sprintf "%.3e" tmse
	c = @sprintf "%.3e" vmse
	r = @sprintf "%.3f" tmse/vmse
	println(a, "  ", b, "  ", c, "  ", r)
end

# Compute mean square error on dataset
mse(ann::ANN, dset::DataSet) = ccall((:fann_test_data, libfann), Cfloat, (Ptr{fann}, Ptr{fann_train_data}), ann, dset)

# Check that sizes of DataSet and ANN match
function checksizes(ann::ANN, dset::DataSet)
	ret = ccall((:fann_check_input_output_sizes, libfann),
		        Cint,
		        (Ptr{fann}, Ptr{fann_train_data}),
		        ann, dset)
	ret == 0 || error("wrong DataSet or ANN sizes")
end

# ~~~~~~ Prediction function ~~~~~~
# Predict output on input X
function predict{T<:FloatingPoint}(ann::ANN, X::Matrix{T})
	nin, nobs = size(X)
	out = zeros(T, ann.nout, nobs)
	for i = 1:nobs
		out_ptr = ccall((:fann_run, libfann),
						Ptr{fann_type},
						(Ptr{fann}, Ptr{fann_type}),
						ann, pointer(X, (i-1)*nin + 1))
		out[:, i] = pointer_to_array(out_ptr, ann.nout)
	end
	return out
end

predict{T<:FloatingPoint}(ann::ANN, X::Vector{T}) =
	pointer_to_array(ccall((:fann_run, libfann),
					 Ptr{fann_type},
					 (Ptr{fann}, Ptr{fann_type}),
					 ann, X), ann.nout)


text:
module MyApp

using Example
using HelloWorldC_jll
using Pkg.Artifacts

fooifier_path() = joinpath(artifact"fooifier", "bin", "fooifier" * (Sys.iswindows() ? ".exe" : ""))

function julia_main()
    try
        real_main()
    catch
        Base.invokelatest(Base.display_error, Base.catch_stack())
        return 1
    end
    return 0
end

function real_main()
    @show ARGS
    @show Base.PROGRAM_FILE
    @show DEPOT_PATH
    @show LOAD_PATH
    @show pwd()
    @show Base.active_project()
    @show Sys.BINDIR

    @show Base.JLOptions().opt_level
    @show Base.JLOptions().nthreads
    @show Base.JLOptions().check_bounds

    display(Base.loaded_modules)
    println()

    println("Running a jll package:")
    HelloWorldC_jll.hello_world() do x
        println("HelloWorld artifact at $(realpath(x))")
        run(`$x`)
    end
    println()

    println("Running the artifact")
    res = read(`$(fooifier_path()) 5 10`, String)
    println("The result of 2*5^2 - 10 == $res")

    @show unsafe_string(Base.JLOptions().image_file)
    @show Example.domath(5)
    @show sin(0.0)
    return
end

if abspath(PROGRAM_FILE) == @__FILE__
    real_main()
end

end # module


text:
using MTHE493
using Test

@testset "MTHE493.jl" begin
    # Small Dataset
    image_dir = "./data/small_dataset/"
    dims = (112,92)
    images = load_images(image_dir; dims=dims)
    @test length(images) == 400

    image_matrix = images_to_image_matrix(map(x->x[1], images))
    @test size(image_matrix) == (10304, 400)

    labels = map(x->x[2], images)
    @test typeof(labels) == Array{SubString{String},1}

    label_map = Dict(String.(unique(labels)) .=> collect(1:length(unique(labels))))
    @test label_map.count == 40

    proj, _ = get_projection(image_matrix, 400)
    @test size(proj) == (10304, 400)

    eigenfaces = get_eigenfaces(proj, 25, dims)
    @test size(eigenfaces[1]) == (112,92)

end

#%%
#= Large Dataset
image_dir = "./data/large_dataset/"
dims = (196,180)
images = load_images(image_dir, 1:100; dims=dims)
image_matrix = images_to_image_matrix(map(x->x[1], images))
labels = map(x->x[2], images)
label_map = Dict(String.(unique(labels)) .=> collect(1:length(unique(labels))))
tile_images(image_matrix_to_images(image_matrix, dims), 10)
=#


text:
#
# A Lee Group as a trait to manifolds
#
#
# Manopt.jl, R. Bergmann, 2018-07-07
export IsLieGroupM, IsLieGroupP, IsLieGroupV, ⊗

"""
    IsLieGroupM{X}
Indicates that `X` is a [`Manifold`](@ref) with a Lie group structure. This also
introdcues a group operation [`⊗`](@ref) of two [`MPoint`](@ref)s of `X`.
"""
@traitdef IsLieGroupM{X}
"""
    IsLieGroupP{X}
An abstract [`MPoint`](@ref) belonging to a Lie group manifold.
"""
@traitdef IsLieGroupP{X}
"""
    IsLieGroupV{X}
An abstract [`TVector`](@ref) belonging to a Lie group manifold.
"""
@traitdef IsLieGroupV{X}
@doc doc"""
    ⊗(x,y)

the binary operator `x ⊗ y` represents the Lie group action for a Lie group
[`Manifold`](@ref), see [`IsLieGroupM`](@ref).
"""
⊗(x::P,y::P) where {P <: MPoint}= LieGroupOp(x,y)
@traitfn function LieGroupOp(x::P,y::P) where {P <: MPoint; !IsLieGroupP{P}}
    sig1 = string( typeof(x) )
    throw( ErrorException("The manifold point of type $sig1 does not belong to a Lie Group") );
end


text:
Lar = LinearAlgebraicRepresentation

"""
	extrudeSimplicial(model::LAR, pattern::Array)::LAR

Algorithm for multimensional extrusion of a simplicial complex.
Can be applied to 0-, 1-, 2-, ... simplicial models, to get a 1-, 2-, 3-, .... model.
The pattern `Array` is used to specify how to decompose the added dimension.

A `model` is a LAR model, i.e. a pair (vertices,cells) to be extruded, whereas pattern is an array of `Int64`, to be used as lateral measures of the *extruded* model. `pattern` elements are assumed as either *solid* or *empty* measures, according to their (+/-) sign.

# Example
```julia
julia> V = [[0,0] [1,0] [2,0] [0,1] [1,1] [2,1] [0,2] [1,2] [2,2]];

julia> FV = [[1,2,4],[2,3,5],[3,5,6],[4,5,7],[5,7,8],[6,8,9]];

julia> pattern = repeat([1,2,-3],outer=4);

julia> model = (V,FV);

julia> W,FW = extrudeSimplicial(model, pattern);

julia> Plasm.view(W,FW)
```
"""
function extrudeSimplicial(model::Lar.LAR, pattern)
	V = [model[1][:,k] for k=1:size(model[1],2)]
    FV = model[2]
    d, m = length(FV[1]), length(pattern)
    coords = collect(cumsum(append!([0], abs.(pattern))))
    offset, outcells, rangelimit, i = length(V), [], d*m, 0
    for cell in FV
    	i += 1
        tube = [v+k*offset for k in range(0, length=m+1) for v in cell]
        cellTube = [tube[k:k+d] for k in range(1, length=rangelimit)]
        if i==1 outcells = reshape(cellTube, d, m)
        else outcells = vcat(outcells, reshape(cellTube, d, m)) end
    end
    cellGroups = []
    for i in 1:size(outcells, 2)
        if pattern[i]>0
            cellGroups = vcat(cellGroups, outcells[:, i])
        end
    end
    outVertices = [vcat(v, [z]) for z in coords for v in V]
    cellGroups = convert(Array{Array{Int, 1}, 1}, cellGroups)
    outModel = outVertices, cellGroups
    hcat(outVertices...), cellGroups
end
function extrudeSimplicial(model::Union{Any,Lar.Cells}, pattern)
	V,FV = model
    d, m = length(FV[1]), length(pattern)
    coords = collect(cumsum(append!([0], abs.(pattern))))
    offset, outcells, rangelimit, i = length(V), [], d*m, 0
    for cell in FV
    	i += 1
        tube = [v+k*offset for k in range(0, length=m+1) for v in cell]
        cellTube = [tube[k:k+d] for k in range(1, length=rangelimit)]
        if i==1 outcells = reshape(cellTube, d, m)
        else outcells = vcat(outcells, reshape(cellTube, d, m)) end
    end
    cellGroups = []
    for i in 1:size(outcells, 2)
        if pattern[i]>0
            cellGroups = vcat(cellGroups, outcells[:, i])
        end
    end
    outVertices = [vcat(v, [z]) for z in coords for v in V]
    cellGroups = convert(Array{Array{Int, 1}, 1}, cellGroups)
    outModel = outVertices, cellGroups
    hcat(outVertices...), cellGroups
end

"""
	simplexGrid(shape::Array)::LAR

Generate a simplicial complex decomposition of a cubical grid of ``d``-cuboids, where ``d`` is the length of `shape` array. Vertices (0-cells) of the grid have `Int64` coordinates.

# Examples
```julia
julia> simplexGrid([0]) # 0-dimensional complex
# output
([0], Array{Int64,1}[])

julia> V,EV = simplexGrid([1]) # 1-dimensional complex
# output
([0 1], Array{Int64,1}[[1, 2]])

julia> V,FV = simplexGrid([1,1]) # 2-dimensional complex
# output
([0 1 0 1; 0 0 1 1], Array{Int64,1}[[1, 2, 3], [2, 3, 4]])

julia> V,CV = simplexGrid([10,10,1]) # 3-dimensional complex
# output
([0 1 … 9 10; 0 0 … 10 10; 0 0 … 1 1], Array{Int64,1}[[1, 2, 12, 122], [2, 12, 122, 123], [12, 122, 123, 133], [2, 12, 13, 123], [12, 13, 123, 133], [13, 123, 133, 134], [2, 3, 13, 123], [3, 13, 123, 124], [13, 123, 124, 134], [3, 13, 14, 124]  …  [119, 229, 230, 240], [109, 119, 120, 230], [119, 120, 230, 240], [120, 230, 240, 241], [109, 110, 120, 230], [110, 120, 230, 231], [120, 230, 231, 241], [110, 120, 121, 231], [120, 121, 231, 241], [121, 231, 241, 242]])

julia> V
# output
3×242 Array{Int64,2}:
 0  1  2  3  4  5  6  7  8  9  10  0  1  2  3  …   1   2   3   4   5   6   7   8   9  10
 0  0  0  0  0  0  0  0  0  0   0  1  1  1  1     10  10  10  10  10  10  10  10  10  10
 0  0  0  0  0  0  0  0  0  0   0  0  0  0  0      1   1   1   1   1   1   1   1   1   1

julia> using Plasm

julia> hpc = Plasm.lar2exploded_hpc(V,CV) # exploded visualization of the grid

julia> Plasm.view(hpc)

julia> V,HV = simplexGrid([1,1,1,1]) # 4-dim cellular complex from the 4D simplex
# output
([0 1 … 0 1; 0 0 … 1 1; 0 0 … 1 1; 0 0 … 1 1], Array{Int64,1}[[1, 2, 3, 5, 9], [2, 3, 5, 9, 10], [3, 5, 9, 10, 11], [5, 9, 10, 11, 13], [2, 3, 5, 6, 10], [3, 5, 6, 10, 11], [5, 6, 10, 11, 13], [6, 10, 11, 13, 14], [3, 5, 6, 7, 11], [5, 6, 7, 11, 13]  …  [4, 6, 10, 11, 12], [6, 10, 11, 12, 14], [3, 4, 6, 7, 11], [4, 6, 7, 11, 12], [6, 7, 11, 12, 14], [7, 11, 12, 14, 15], [4, 6, 7, 8, 12], [6, 7, 8, 12, 14], [7, 8, 12, 14, 15], [8, 12, 14, 15, 16]])
```
"""
function simplexGrid(shape)
    model = [[]], [[1]]
    for item in shape
        model = extrudeSimplicial(model, fill(1, item))
    end
    V, CV = model
    V = convert(Array{Float64,2}, V)
    return V, CV
end

"""
	simplexFacets(simplices::Cells)::Cells

Compute the `(d-1)`-skeleton (set of `facets`) of a simplicial `d`-complex.

# Example
```julia
julia> V,FV = Lar.simplexGrid([1,1]) # 2-dimensional complex
# output
([0 1 0 1; 0 0 1 1], Array{Int64,1}[[1, 2, 3], [2, 3, 4]])

julia> Plasm.view(V,FV)

julia> W,CW = Lar.extrudeSimplicial((V,FV), [1])
([0.0 1.0 … 0.0 1.0; 0.0 0.0 … 1.0 1.0; 0.0 0.0 … 1.0 1.0],
Array{Int64,1}[[1,2,3,5],[2,3,5,6],[3,5,6,7],[2,3,4,6],[3,4,6,7],[4,6,7,8]])

julia> FW = Lar.simplexFacets(CW)
18-element Array{Any,1}:
[[1,3,5],[5,6,7],[3,5,7],[3,6,7],[4,6,7],[4,7,8],[4,6,8],
[6,7,8],[3,5,6],[2,3,5],[2,3,4],[3,4,7],[1,2,3],[2,4,6],[2,5,6],
[1,2,5],[2,3,6],[3,4,6]]

julia> Plasm.view(W,FW)
```

# Example

```julia
julia> V,(VV,EV,FV,CV) = Lar.cuboidGrid([3,3,3],true)

julia> TV = Lar.simplexFacets(CV)

julia> Plasm.view(V,TV)

```
"""
function simplexFacets(simplices) # TODO: solve bug
    out = []
    d = length(simplices[1])
    for simplex in simplices
    	non_oriented_facets = [ append!(simplex[1:k-1],simplex[k+1:d]) for k=1:d ]
    	if (-1)^d == 1
    		oriented_facets = non_oriented_facets
    	else
    		oriented_facets = [ length(f)==2 ? [f[2],f[1]] : [f[2],f[1],f[3:end]]
    			for f in non_oriented_facets ]
    	end
        append!(out, oriented_facets)
    end
    out = collect(Set(out))
    return convert(Lar.Cells, out)
end

"""
	quads2triangles(quads::Cells)::Cells

Convert an array of *quads* with type `::Lar.Cells` into an array of *triangles*
with the same type.

# Examples

The transformation from quads to triangles works for any 2-complex, embedded in any dimensional space

## 2D example
```
V,FV = Lar.cuboidGrid([4,5])
triangles = Lar.quads2triangles(FV::Lar.Cells)::Lar.Cells
using Plasm
Plasm.view((V,[triangles]))
```
## 3D example
```
V,(VV,EV,FV,CV) = Lar.cuboidGrid([4,5,3],true)
triangles = Lar.quads2triangles(FV::Lar.Cells)::Lar.Cells
using Plasm
Plasm.view((V,[triangles]))
```
"""
function quads2triangles(quads::Lar.Cells)::Lar.Cells
	pairs = [[ Int[v1,v2,v3], Int[v3,v2,v4]] for (v1,v2,v3,v4) in quads ]
	return cat(pairs)
end


text:
# # Nonlinear constrained optimization
#
#-
#md # !!! tip
#md #     This example is also available as a Jupyter notebook:
#md #     [`ipnewton_basics.ipynb`](@__NBVIEWER_ROOT_URL__examples/generated/ipnewton_basics.ipynb)
#-
#
# The nonlinear constrained optimization interface in
# `Optim` assumes that the user can write the optimization
# problem in the following way.

# ```math
# \min_{x\in\mathbb{R}^n} f(x) \quad \text{such that}\\
# l_x \leq \phantom{c(}x\phantom{)} \leq u_x \\
# l_c \leq c(x) \leq u_c.
# ```

# For equality constraints on ``x_j`` or ``c(x)_j`` you set those
# particular entries of bounds to be equal, ``l_j=u_j``.
# Likewise, setting ``l_j=-\infty`` or ``u_j=\infty`` means that the
# constraint is unbounded from below or above respectively.

using Optim, NLSolversBase #hide
import NLSolversBase: clear! #hide

# # Constrained optimization with `IPNewton`

# We will go through examples on how to use the constraints interface
# with the interior-point Newton optimization algorithm [IPNewton](../../algo/ipnewton.md).

# Throughout these examples we work with the standard Rosenbrock function.
# The objective and its derivatives are given by

fun(x) =  (1.0 - x[1])^2 + 100.0 * (x[2] - x[1]^2)^2

function fun_grad!(g, x)
g[1] = -2.0 * (1.0 - x[1]) - 400.0 * (x[2] - x[1]^2) * x[1]
g[2] = 200.0 * (x[2] - x[1]^2)
end

function fun_hess!(h, x)
h[1, 1] = 2.0 - 400.0 * x[2] + 1200.0 * x[1]^2
h[1, 2] = -400.0 * x[1]
h[2, 1] = -400.0 * x[1]
h[2, 2] = 200.0
end;

# ## Optimization interface
# To solve a constrained optimization problem we call the `optimize`
# method
# ``` julia
# optimize(d::AbstractObjective, constraints::AbstractConstraints, initial_x::Tx, method::ConstrainedOptimizer, options::Options)
# ```

# We can create instances of `AbstractObjective` and
# `AbstractConstraints` using the types `TwiceDifferentiable` and
# `TwiceDifferentiableConstraints` from the package `NLSolversBase.jl`.

# ## Box minimization
# We want to optimize the Rosenbrock function in the box
# ``-0.5 \leq x \leq 0.5``, starting from the point ``x_0=(0,0)``.
# Box constraints are defined using, for example,
# `TwiceDifferentiableConstraints(lx, ux)`.

x0 = [0.0, 0.0]
df = TwiceDifferentiable(fun, fun_grad!, fun_hess!, x0)

lx = [-0.5, -0.5]; ux = [0.5, 0.5]
dfc = TwiceDifferentiableConstraints(lx, ux)

res = optimize(df, dfc, x0, IPNewton())
## Test the results             #src
using Test                 #src
@test Optim.converged(res)      #src
@test Optim.minimum(res) ≈ 0.25 #src

# Like the rest of Optim, you can also use `autodiff=:forward` and just pass in
# `fun`.

# If we only want to set lower bounds, use `ux = fill(Inf, 2)`

ux = fill(Inf, 2)
dfc = TwiceDifferentiableConstraints(lx, ux)

clear!(df)
res = optimize(df, dfc, x0, IPNewton())
@test Optim.converged(res)                   #src
@test Optim.minimum(res) < 0.0 + sqrt(eps()) #src

# ## Defining "unconstrained" problems

# An unconstrained problem can be defined either by passing
# `Inf` bounds or empty arrays.
# **Note that we must pass the correct type information to the empty `lx` and `ux`**

lx = fill(-Inf, 2); ux = fill(Inf, 2)
dfc = TwiceDifferentiableConstraints(lx, ux)

clear!(df)
res = optimize(df, dfc, x0, IPNewton())
@test Optim.converged(res)                   #src
@test Optim.minimum(res) < 0.0 + sqrt(eps()) #src

lx = Float64[]; ux = Float64[]
dfc = TwiceDifferentiableConstraints(lx, ux)

clear!(df)
res = optimize(df, dfc, x0, IPNewton())
@test Optim.converged(res)                   #src
@test Optim.minimum(res) < 0.0 + sqrt(eps()) #src

# ## Generic nonlinear constraints

# We now consider the Rosenbrock problem with a constraint on
# ```math
#    c(x)_1 = x_1^2 + x_2^2.
# ```

# We pass the information about the constraints to `optimize`
# by defining a vector function `c(x)` and its Jacobian `J(x)`.

# The Hessian information is treated differently, by considering the
# Lagrangian of the corresponding slack-variable transformed
# optimization problem. This is similar to how the [CUTEst
# library](https://github.com/JuliaSmoothOptimizers/CUTEst.jl) works.
# Let ``H_j(x)`` represent the Hessian of the ``j``th component
# ``c(x)_j`` of the generic constraints.
# and ``\lambda_j`` the corresponding dual variable in the
# Lagrangian. Then we want the `constraint` object to
# add the values of ``H_j(x)`` to the Hessian of the objective,
# weighted by ``\lambda_j``.

# The Julian form for the supplied function ``c(x)`` and the derivative
# information is then added in the following way.

con_c!(c, x) = (c[1] = x[1]^2 + x[2]^2; c)
function con_jacobian!(J, x)
    J[1,1] = 2*x[1]
    J[1,2] = 2*x[2]
    J
end
function con_h!(h, x, λ)
    h[1,1] += λ[1]*2
    h[2,2] += λ[1]*2
end;

# **Note that `con_h!` adds the `λ`-weighted Hessian value of each
# element of `c(x)` to the Hessian of `fun`.**

# We can then optimize the Rosenbrock function inside the ball of radius
# ``0.5``.

lx = Float64[]; ux = Float64[]
lc = [-Inf]; uc = [0.5^2]
dfc = TwiceDifferentiableConstraints(con_c!, con_jacobian!, con_h!,
                                     lx, ux, lc, uc)
res = optimize(df, dfc, x0, IPNewton())
@test Optim.converged(res)                    #src
@test Optim.minimum(res) ≈ 0.2966215688829263 #src

# We can add a lower bound on the constraint, and thus
# optimize the objective on the annulus with
# inner and outer radii ``0.1`` and ``0.5`` respectively.

lc = [0.1^2]
dfc = TwiceDifferentiableConstraints(con_c!, con_jacobian!, con_h!,
                                     lx, ux, lc, uc)
res = optimize(df, dfc, x0, IPNewton())
@test Optim.converged(res)                    #src
@test Optim.minimum(res) ≈ 0.2966215688829255 #src

# **Note that the algorithm warns that the Initial guess is not an
# interior point.** `IPNewton` can often handle this, however, if the
# initial guess is such that `c(x) = u_c`, then the algorithm currently
# fails. We may fix this in the future.

# ## Multiple constraints
# The following example illustrates how to add an additional constraint.
# In particular, we add a constraint function
# ```math
#    c(x)_2 = x_2\sin(x_1)-x_1
# ```

function con2_c!(c, x)
    c[1] = x[1]^2 + x[2]^2     ## First constraint
    c[2] = x[2]*sin(x[1])-x[1] ## Second constraint
    c
end
function con2_jacobian!(J, x)
    ## First constraint
    J[1,1] = 2*x[1]
    J[1,2] = 2*x[2]
    ## Second constraint
    J[2,1] = x[2]*cos(x[1])-1.0
    J[2,2] = sin(x[1])
    J
end
function con2_h!(h, x, λ)
    ## First constraint
    h[1,1] += λ[1]*2
    h[2,2] += λ[1]*2
    ## Second constraint
    h[1,1] += λ[2]*x[2]*-sin(x[1])
    h[1,2] += λ[2]*cos(x[1])
    ## Symmetrize h
    h[2,1]  = h[1,2]
    h
end;

# We generate the constraint objects and call `IPNewton` with
# initial guess ``x_0 = (0.25,0.25)``.

x0 = [0.25, 0.25]
lc = [-Inf, 0.0]; uc = [0.5^2, 0.0]
dfc = TwiceDifferentiableConstraints(con2_c!, con2_jacobian!, con2_h!,
                                     lx, ux, lc, uc)
res = optimize(df, dfc, x0, IPNewton())
@test Optim.converged(res)                                       #src
@test Optim.minimum(res) ≈ 1.0                                   #src
@test isapprox(Optim.minimizer(res), zeros(2), atol=sqrt(eps())) #src

#md # ## [Plain Program](@id ipnewton_basics-plain-program)
#md #
#md # Below follows a version of the program without any comments.
#md # The file is also available here: [ipnewton_basics.jl](ipnewton_basics.jl)
#md #
#md # ```julia
#md # @__CODE__
#md # ```

max_stars_repo_name: ethansaxenian/RosettaDecode

max_issues_repo_name: ethansaxenian/RosettaDecode

max_forks_repo_name: ethansaxenian/RosettaDecode