text:
using ColabINAOE2021

export dqc, dql, dqk
export wijct, wijlt, wijkt

"### `EExchanges.jl` gives the functions and method for the study of energy distribution and exchange within the evolution of initial amplitudes on masses arrays given in `MArrays.jl`"

# 1. ---------------------------------------------------------------------------

"### Circular array energy exchange "

"`dqc` is the derivative of the function `qc` for the evolution on circular masses arrays. It is called as `dqc(k, n, t, N, V::Vector)`"
function dqc(k, n, t, N, V::Vector)
    v = zeros(Float64, (N + 1, N + 1))
	for m in 0:N
		for l in 0:m
			v[l + 1, m + 1] = sin(2 * t * (√k) * sin(π * m / (N + 1))) * sin(π * m / (N + 1)) * cos(2 * π * m * (n - l) / (N + 1)) * V[l + 1]
		end
	end
	return ((-2 * (√k)) / (N + 1)) * sum(v)
end

"`wijc` is the energy exchange function between the position _i_ and _j_ at time _t_ of an initial amplitude vector _V_, such that, but not strictly, _i < j_. It is called `wijc(k, i, j, t, N, V::Vector)`."
wijc(k, i, j, t, N, V) = qc(k, i, t, N, V) * dqc(k, j, t, N, V) - dqc(k, i, t, N, V) * qc(k, j, t, N, V)

"`wijct` is an ad-hoc constructor (comprehension) to evalute the evolution of the energy exchange given a time span _t_ with an initial amplitudes vector _V_. It is called `wijct(k, i, j, t, N, V::Vector)`"
wijct(k, i, j, t, N, V) = [wijc(k, i, j, t[τ], N, V) for τ in 1:length(t)]

"`wijcm` is a function that calculates the energy exchange between the position _i_, and all other positions at time _t_. The item who return is a _N×`length(t)`_ matrix. It is called `wijcm(k, i, t, N, V::Vector)`"
function wijcm(k, i, t, N, V)
	v = zeros(Float64, (N, length(t)))
	for τ in 1:length(t)
			for j in 1:N
				v[j, τ] = wij(k, i, j, t[τ], N, V)
			end
		end
	return v
end

# 2. ---------------------------------------------------------------------------

"""### Lineal array energy exchange """

"`dql` is the derivative of the function `ql` for the evolution on lineal masses arrays. It is called as `dql(k, n, t, N, V::Vector)`"
function dql(k, n, t, N, V::Vector)
	v = zeros(Float64, (N + 1, N + 1))
	for i in 0:N
		for j in 0:N
			v[i + 1, j + 1] = (ColabINAOE2021.twocheb(j, n, i, N) / ColabINAOE2021.nc(j, N)) * sin(2 * t * (√k) * sin(ColabINAOE2021.ϕ(j + 1, N) / 2)) * sin(ColabINAOE2021.ϕ(j + 1, N) / 2) * V[i + 1]
		end
	end
	return (-2 * (√k)) * sum(v)
end

"`wijl` is the energy exchange function between the position _i_ and _j_ at time _t_ of an initial amplitude vector _V_, such that, but not strictly, _i < j_. It is called `wijl(k, i, j, t, N, V::Vector)`."
wijl(k, i, j, t, N, V) = ql(k, i, t, N, V) * dql(k, j, t, N, V) - dql(k, i, t, N, V) * ql(k, j, t, N, V)

"`wijlt` is an ad-hoc constructor (comprehension) to evalute the evolution of the energy exchange given a time span _t_ with an initial amplitudes vector _V_. It is called `wijlt(k, i, j, t, N, V::Vector)`"
wijlt(k, i, j, t, N, V) = [wijl(k, i, j, t[τ], N, V) for τ in 1:length(t)]

# 3. ---------------------------------------------------------------------------

"""### Kravchuk array energy exchange """

"`dqk` is the derivative of the function `ql` for the evolution on Kravchuk masses arrays. It is called as `dqk(n, t, N, V::Vector)`"
function dqk(n, t, N, V::Vector)
	v = zeros(Float64, (N + 1, N + 1))
	for m in 0:N
		for l in 0:N
			v[m + 1, l + 1] = (-√m) *ColabINAOE2021.k2(n, l, m, N) * sin(t * (√m)) * V[l + 1]
		end
	end
	return sum(v)
end

"`wijk` is the energy exchange function between the position _i_ and _j_ at time _t_ of an initial amplitude vector _V_, such that, but not strictly, _i < j_. It is called `wijk(i, j, t, N, V::Vector)`."
wijk(i, j, t, N, V) = qk(i, t, N, V) * dqk(j, t, N, V) - dqk(i, t, N, V) * qk(j, t, N, V)

"`wijkt` is an ad-hoc constructor (comprehension) to evalute the evolution of the energy exchange given a time span _t_ with an initial amplitudes vector _V_. It is called `wijkt(k, i, j, t, N, V::Vector)`"
wijkt(i, j, t, N, V) = [wijk(i, j, t[τ], N, V) for τ in 1:length(t)]


text:
# Direct Search with Probabilistic Descent as described in Gratton2014:
#
#  S. Gratton, C. W. Royer, L. N. Vicente, and Z. Zhang, Direct search
#  based on probabilistic descent, preprint 14-11, Dept. Mathematics, Univ. Coimbra.
#  http://www.mat.uc.pt/~lnv/papers/ds-random.pdf
#

# Generate num random vectors on the n-dimensional, unit (hyper)sphere.
# This is the Muller-Marsaglia method as described on the page:
#   http://mathworld.wolfram.com/HyperspherePointPicking.html
function sample_unit_hypersphere(n, num = 1)
  X = randn(n, num)
  sqrootsums = 1 ./ sqrt(sum( X.^2, 1 ))
  broadcast(*, sqrootsums, X)
end

type RandomDirectionGen <: DirectionGenerator
  numDimensions::Int64
  numDirections::Int64
end

function directions_for_k(rdg::RandomDirectionGen, k)
  sample_unit_hypersphere(rdg.numDimensions, rdg.numDirections)
end

# A MirroredRandomDirectionGen generates half of the directions randomly and then
# mirrors by negating them.
type MirroredRandomDirectionGen <: DirectionGenerator
  numDimensions::Int64
  numDirections::Int64

  MirroredRandomDirectionGen(numDims, numDirections) = begin
    if numDirections % 2 == 1
      throw("the number of directions must be even")
    end
    new(numDims, numDirections)
  end
end

function directions_for_k(rdg::MirroredRandomDirectionGen, k)
  r = sample_unit_hypersphere(rdg.numDimensions, int(rdg.numDirections/2))
  [r -r]
end

DirectSearchProbabilisticDescentDefaultParameters = {
  :NumDirections => 2, # This should be a function of Gamma and Phi for the GSS but 2 is often enough
}

function direct_search_probabilistic_descent(parameters)
  params = Parameters(parameters, DirectSearchProbabilisticDescentDefaultParameters)
  params[:DirectionGenerator] = MirroredRandomDirectionGen(numdims(params[:Evaluator]), params[:NumDirections])
  GeneratingSetSearcher(params)
end

ext: jl


text:
# dimensions
nq = 10
n = 2 * nq
m = 10
T = 20
nz = m * (T - 1) + n * (T - 1)
nd = n * (T - 1)

# problem data
Aa = [SMatrix{nq, nq}(Array(1.0 * Diagonal(0.1 * rand(nq) + ones(nq)))) for t = 1:T-1]
Ab = [SMatrix{nq, nq}(Array(1.0 * Diagonal(0.1 * rand(nq) + ones(nq)))) for t = 1:T-1]
Ac = [SMatrix{nq, nq}(Array(Diagonal(ones(nq)))) for t = 1:T-1]
Ba = [SMatrix{nq,m}(1.0 * rand(nq, m)) for t = 1:T-1]
Qa = [Diagonal(SVector{nq}(0.1 * rand(nq) + ones(nq))) for t = 1:T]
Qb = [Diagonal(SVector{nq}(0.1 * rand(nq) + ones(nq))) for t = 1:T]
Qv = [Diagonal(SVector{nq}(-1.0 * ones(nq) + 0.1 * rand(nq))) for t = 1:T]

# indices
u_idx = [collect((t - 1) * (m + n) .+ (1:m)) for t = 1:T-1]
x_idx = [collect((t - 1) * (m + n) + m .+ (1:n)) for t = 1:T-1]
qa_idx = [collect(x_idx[t][1:nq]) for t = 1:T-1]
qb_idx = [x_idx[t][nq .+ (1:nq)] for t = 1:T-1]
n_idx = [(t - 1) * n .+ (1:n) for t = 1:T-1]

Q = [[Qa[t] Qv[t]; Qv[t]' Qb[t]] for t = 1:T]
R = [Diagonal(SVector{m}(0.1 * rand(m) + ones(m))) for t = 1:T-1]
A = [SMatrix{n,n}([zeros(nq, nq) I; Aa[t] Ab[t]]) for t = 1:T-1]
B = [SMatrix{n,m}([zeros(nq, m); Ba[t];]) for t = 1:T-1]
P = [SMatrix{n,n}([I zeros(nq, nq); zeros(nq, nq) Ac[t]]) for t = 1:T-1]

# direct problem
S = zeros(nz, nz)
C = zeros(nd, nz)

for t = 1:T-1
	S[u_idx[t], u_idx[t]] = R[t]
	S[x_idx[t], x_idx[t]] = Q[t+1]

	C[n_idx[t], u_idx[t]] = -B[t]
	C[n_idx[t], x_idx[t]] = P[t]
	t == 1 && continue
	C[n_idx[t], x_idx[t-1]] = -A[t]
end

Q̃ = [inv(Array(Q[t])) for t = 1:T]

Q̃a = [Q̃[t][1:nq, 1:nq] for t = 1:T]
Q̃b = [Q̃[t][nq .+ (1:nq), nq .+ (1:nq)] for t = 1:T]
Q̃v = [Q̃[t][1:nq, nq .+ (1:nq)] for t = 1:T]

R̃ = [inv(R[t]) for t = 1:T-1]
S̃ = zeros(nz, nz)

for t = 1:T-1
	S̃[u_idx[t], u_idx[t]] = R̃[t]
	S̃[x_idx[t], x_idx[t]] = Q̃[t+1]
end

norm(S̃ - inv(S))

# build kkt system
J = [S C'; C zeros(nd, nd)]
Js = sparse(J)
r = rand(nz + nd)
Δ = J \ r
Y = C * S̃ * C'
rlag = view(r, 1:nz)
rdyn = view(r, nz .+ (1:nd))
β = -rdyn + C * S̃ * rlag
Δν = Y \ β
Δz = S̃ * (rlag - C' * Δν)

# LU
solver_lu = lu_solver(Js)
Δlu = zeros(nz + nd)
@benchmark linear_solve!($solver_lu, $Δlu, $Js, $r)

solver_lu = lu_solver(J)
Δlu = zeros(nz + nd)
@benchmark linear_solve!($solver_lu, $Δlu, $J, $r)

# # QDLDL
# solver_ldl = ldl_solver(J)
# Δldl = zeros(nz + nd)
# @benchmark linear_solve!($solver_ldl, $Δldl, $J, $r)
# Δldl

norm(Δ - Δlu)
norm(Δ - [Δz; Δν])


text:
# This file is part of the TaylorSeries.jl Julia package, MIT license
#
# Luis Benet & David P. Sanders
# UNAM
#
# MIT Expat license
#

for T in (:Taylor1, :HomogeneousPolynomial, :TaylorN)

    ## real, imag, conj and ctranspose ##
    for f in (:real, :imag, :conj)
        @eval ($f)(a::$T) = $T(($f).(a.coeffs), a.order)
    end

    @eval adjoint(a::$T) = conj(a)

    ## isinf and isnan ##
    @eval isinf(a::$T) = any( isinf.(a.coeffs) )

    @eval isnan(a::$T) = any( isnan.(a.coeffs) )
end

## Division functions: rem and mod ##
for op in (:mod, :rem)
    for T in (:Taylor1, :TaylorN)
        @eval begin
            function ($op)(a::$T{T}, x::T) where {T<:Real}
                coeffs = copy(a.coeffs)
                @inbounds coeffs[1] = ($op)(constant_term(a), x)
                return $T(coeffs, a.order)
            end

            function ($op)(a::$T{T}, x::S) where {T<:Real, S<:Real}
                R = promote_type(T, S)
                a = convert($T{R}, a)
                return ($op)(a, convert(R,x))
            end
        end
    end

    @eval begin
        function ($op)(a::TaylorN{Taylor1{T}}, x::T) where {T<:Real}
            coeffs = copy(a.coeffs)
            @inbounds coeffs[1] = ($op)(constant_term(a), x)
            return TaylorN( coeffs, a.order )
        end

        function ($op)(a::TaylorN{Taylor1{T}}, x::S) where {T<:Real, S<:Real}
            R = promote_type(T,S)
            a = convert(TaylorN{Taylor1{R}}, a)
            return ($op)(a, convert(R,x))
        end

        function ($op)(a::Taylor1{TaylorN{T}}, x::T) where {T<:Real}
            coeffs = copy(a.coeffs)
            @inbounds coeffs[1] = ($op)(constant_term(a), x)
            return Taylor1( coeffs, a.order )
        end

        @inbounds function ($op)(a::Taylor1{TaylorN{T}}, x::S) where
                {T<:Real, S<:Real}

            R = promote_type(T,S)
            a = convert(Taylor1{TaylorN{R}}, a)
            return ($op)(a, convert(R,x))
        end
    end
end

## mod2pi and abs ##
for T in (:Taylor1, :TaylorN)
    @eval begin
        function mod2pi(a::$T{T}) where {T<:Real}
            coeffs = copy(a.coeffs)
            @inbounds coeffs[1] = mod2pi( constant_term(a) )
            return $T( coeffs, a.order)
        end

        function abs(a::$T{T}) where {T<:Real}
            if constant_term(a) > 0
                return a
            elseif constant_term(a) < 0
                return -a
            else
                throw(ArgumentError(
                """The 0th order Taylor1 coefficient must be non-zero
                (abs(x) is not differentiable at x=0)."""))
            end
        end

        abs2(a::$T) = a^2
    end
end

function mod2pi(a::TaylorN{Taylor1{T}}) where {T<:Real}
    coeffs = copy(a.coeffs)
    @inbounds coeffs[1] = mod2pi( constant_term(a) )
    return TaylorN( coeffs, a.order )
end

function mod2pi(a::Taylor1{TaylorN{T}}) where {T<:Real}
    coeffs = copy(a.coeffs)
    @inbounds coeffs[1] = mod2pi( constant_term(a) )
    return Taylor1( coeffs, a.order )
end

function abs(a::TaylorN{Taylor1{T}}) where {T<:Real}
    if constant_term(a)[0] > 0
        return a
    elseif constant_term(a)[0] < 0
        return -a
    else
        throw(ArgumentError(
        """The 0th order TaylorN{Taylor1{T}} coefficient must be non-zero
        (abs(x) is not differentiable at x=0)."""))
    end
end

function abs(a::Taylor1{TaylorN{T}}) where {T<:Real}
    if constant_term(a[0]) > 0
        return a
    elseif constant_term(a[0]) < 0
        return -a
    else
        throw(ArgumentError(
        """The 0th order Taylor1{TaylorN{T}} coefficient must be non-zero
        (abs(x) is not differentiable at x=0)."""))
    end
end

@doc doc"""
    abs(a)

Returns `a` if `constant_term(a) > 0` and `-a` if `constant_term(a) < 0` for
`a <:Union{Taylor1,TaylorN}`.
Notice that `typeof(abs(a)) <: AbstractSeries`.

""" abs

#norm
@doc doc"""
    norm(x::AbstractSeries, p::Real)

Returns the p-norm of an `x::AbstractSeries`, defined by

```math
\begin{equation*}
\left\Vert x \right\Vert_p =  \left( \sum_k | x_k |^p \right)^{\frac{1}{p}},
\end{equation*}
```
which returns a non-negative number.

""" norm

norm(x::AbstractSeries, p::Real=2) = norm( norm.(x.coeffs, p), p)
norm(x::Union{Taylor1{T},HomogeneousPolynomial{T}}, p::Real=2) where
    {T<:NumberNotSeries} = norm(x.coeffs, p)
#norm for Taylor vectors
norm(v::Vector{T}, p::Real=2) where T<:AbstractSeries = norm( norm.(v, p), p)

# rtoldefault
for T in (:Taylor1, :HomogeneousPolynomial, :TaylorN)
    @eval rtoldefault(::Type{$T{T}}) where {T<:Number} = rtoldefault(T)
    @eval rtoldefault(::$T{T}) where {T<:Number} = rtoldefault(T)
end

# isfinite
"""
    isfinite(x::AbstractSeries) -> Bool

Test whether the coefficients of the polynomial `x` are finite.
"""
isfinite(x::AbstractSeries) = !isnan(x) && !isinf(x)

# isapprox; modified from Julia's Base.isapprox
"""
    isapprox(x::AbstractSeries, y::AbstractSeries; rtol::Real=sqrt(eps), atol::Real=0, nans::Bool=false)

Inexact equality comparison between polynomials: returns `true` if
`norm(x-y,1) <= atol + rtol*max(norm(x,1), norm(y,1))`, where `x` and `y` are
polynomials. For more details, see [`Base.isapprox`](@ref).
"""
function isapprox(x::T, y::S; rtol::Real=rtoldefault(x,y,0), atol::Real=0.0,
        nans::Bool=false) where {T<:AbstractSeries, S<:AbstractSeries}

    x == y || (isfinite(x) && isfinite(y) &&
        norm(x-y,1) <= atol + rtol*max(norm(x,1), norm(y,1))) ||
        (nans && isnan(x) && isnan(y))
end
#isapprox for vectors of Taylors
function isapprox(x::Vector{T}, y::Vector{S}; rtol::Real=rtoldefault(T,S,0), atol::Real=0.0,
        nans::Bool=false) where {T<:AbstractSeries, S<:AbstractSeries}

    x == y || norm(x-y,1) <= atol + rtol*max(norm(x,1), norm(y,1)) ||
        (nans && isnan(x) && isnan(y))
end

#taylor_expand function for Taylor1
"""
    taylor_expand(f, x0; order)

Computes the Taylor expansion of the function `f` around the point `x0`.

If `x0` is a scalar, a `Taylor1` expansion will be returned. If `x0` is a vector,
a `TaylorN` expansion will be computed. If the dimension of x0 (`length(x0)`)
is different from the variables set for `TaylorN` (`get_numvars()`), an
`AssertionError` will be thrown.
"""
function taylor_expand(f::Function; order::Int=15)
   a = Taylor1(order)
   return f(a)
end

function taylor_expand(f::Function, x0::T; order::Int=15) where {T<:Number}
   a = Taylor1([x0, one(T)], order)
   return f(a)
end

#taylor_expand function for TaylorN
function taylor_expand(f::Function, x0::Vector{T};
        order::Int=get_order()) where {T<:Number}

    ll = length(x0)
    @assert ll == get_numvars() && order <= get_order()
    X = Array{TaylorN{T}}(undef, ll)

    for i in eachindex(X)
        X[i] = x0[i] + TaylorN(T, i, order=order)
    end

    return f( X )
end

function taylor_expand(f::Function, x0...; order::Int=get_order())
    x0 = promote(x0...)
    T = eltype(x0[1])
    ll = length(x0)
    @assert ll == get_numvars() && order <= get_order()
    X = Array{TaylorN{T}}(undef, ll)

    for i in eachindex(X)
        X[i] = x0[i] + TaylorN(T, i, order=order)
    end

    return f( X... )
end

#update! function for Taylor1
"""
    update!(a, x0)

Takes `a <: Union{Taylo1,TaylorN}` and expands it around the coordinate `x0`.
"""
function update!(a::Taylor1, x0::T) where {T<:Number}
    a.coeffs .= evaluate(a, Taylor1([x0, one(x0)], a.order) ).coeffs
    nothing
end

#update! function for TaylorN
function update!(a::TaylorN, vals::Vector{T}) where {T<:Number}
    a.coeffs .= evaluate(a, get_variables() .+ vals).coeffs
    nothing
end

function update!(a::Union{Taylor1,TaylorN})
    #shifting around zero shouldn't change anything...
    nothing
end

for T in (:Taylor1, :TaylorN)
    @eval deg2rad(z::$T{T}) where {T<:AbstractFloat} = z * (convert(T, pi) / 180)
    @eval deg2rad(z::$T{T}) where {T<:Real} = z * (convert(float(T), pi) / 180)

    @eval rad2deg(z::$T{T}) where {T<:AbstractFloat} = z * (180 / convert(T, pi))
    @eval rad2deg(z::$T{T}) where {T<:Real} = z * (180 / convert(float(T), pi))
end

# Internal mutating deg2rad!, rad2deg! functions
for T in (:Taylor1, :TaylorN)
    @eval @inline function deg2rad!(v::$T{T}, a::$T{T}, k::Int) where {T<:AbstractFloat}
        @inbounds v[k] = a[k] * (convert(T, pi) / 180)
        return nothing
    end
    @eval @inline function deg2rad!(v::$T{S}, a::$T{T}, k::Int) where {S<:AbstractFloat, T<:Real}
        @inbounds v[k] = a[k] * (convert(float(T), pi) / 180)
        return nothing
    end
    @eval @inline function rad2deg!(v::$T{T}, a::$T{T}, k::Int) where {T<:AbstractFloat}
        @inbounds v[k] = a[k] * (180 / convert(T, pi))
        return nothing
    end
    @eval @inline function rad2deg!(v::$T{S}, a::$T{T}, k::Int) where {S<:AbstractFloat, T<:Real}
        @inbounds v[k] = a[k] * (180 / convert(float(T), pi))
        return nothing
    end
end


text:
RMSEA(sem_fit::SemFit{Mi, So, St, Mo, O} where {Mi, So, St, Mo <: AbstractSemSingle, O}) = RMSEA(df(sem_fit), χ²(sem_fit), n_obs(sem_fit))

RMSEA(sem_fit::SemFit{Mi, So, St, Mo, O} where {Mi, So, St, Mo <: SemEnsemble, O}) = sqrt(length(sem_fit.model.sems))*RMSEA(df(sem_fit), χ²(sem_fit), n_obs(sem_fit))

function RMSEA(df, chi2, n_obs)
    rmsea = (chi2 - df) / (n_obs*df)
    rmsea > 0 ? nothing : rmsea = 0
    return sqrt(rmsea)
end

#= # SemFit splices loss functions ---------------------------------------------------------------------
RMSEA(sem_fit::SemFit{Mi, So, St, Mo, O} where {Mi, So, St, Mo <: AbstractSemSingle, O}) =
    RMSEA(
        sem_fit,
        sem_fit.model.observed,
        sem_fit.model.imply,
        sem_fit.model.diff,
        sem_fit.model.loss.functions...
        ) =#

#=
# RAM + SemML
function RMSEA(sem_fit::SemFit, obs, imp::Union{RAM, RAMSymbolic}, diff, loss_ml::Union{SemML, SemFIML, SemWLS})
    df_ = df(sem_fit)
    rmsea = (χ²(sem_fit) - df_) / ((n_obs(sem_fit))*df_)
    rmsea > 0 ? nothing : rmsea = 0
    return sqrt(rmsea)
end =#


text:
using PyPlot
using JLD

function main(code::AbstractString)
    log_path_record = "$(homedir())/open_loop/output-record-$(code).jld"
    log_path_record = "$(homedir())/simulations/output-record-$(code).jld"
    d_rec = load(log_path_record)

    imu_meas    = d_rec["imu_meas"]
    t0 = imu_meas.t[1]
    sz = size(imu_meas.t,1)
    a_norm = zeros(sz,3)
    for i=1:sz
        a_norm[i,:] = (rotMatrix('y',-imu_meas.z[i,5])*rotMatrix('x',-imu_meas.z[i,4])*imu_meas.z[i,7:9]')'
    end
    figure()
    plot(imu_meas.t-t0,imu_meas.z[:,7:8])
    grid("on")
    figure()
    plot(imu_meas.t-t0,a_norm[:,1:2])
    grid("on")
    println("Mean raw: $(mean(imu_meas.z[:,8])), std raw: $(std(imu_meas.z[:,8]))")
    println("Mean filtered: $(mean(a_norm[:,2])), std filtered: $(std(a_norm[:,2]))")
end

function rotMatrix(s::Char,deg::Float64)
    A = zeros(3,3)
    if s=='x'
        A = [1 0 0;
             0 cos(deg) sin(deg);
             0 -sin(deg) cos(deg)]
    elseif s=='y'
        A = [cos(deg) 0 -sin(deg);
             0 1 0;
             sin(deg) 0 cos(deg)]
    elseif s=='z'
        A = [cos(deg) sin(deg) 0
             -sin(deg) cos(deg) 0
             0 0 1]
    else
        warn("Wrong angle for rotation matrix")
    end
    return A
end

max_stars_repo_path: eval_sim/Other tests/NormalizeImu.jl

max_stars_repo_name: RoyaFiroozi/barc

max_issues_repo_path: eval_sim/Other tests/NormalizeImu.jl

max_issues_repo_name: RoyaFiroozi/barc

max_forks_repo_path: eval_sim/Other tests/NormalizeImu.jl

max_forks_repo_name: RoyaFiroozi/barc


text:
using SimpleSATSolver
using SatisfiabilityInterface
using Test

function is_sat(p::SATProblem)
    status, results = solve(p, SimpleSAT())

    return status == :sat && satisfies(p, results)
end

function is_unsat(p::SATProblem)
    status, results = solve(p, SimpleSAT())

    return status == :unsat
end

is_sat(v::Vector{Vector{Int}}) = is_sat(SATProblem(v))
is_unsat(v::Vector{Vector{Int}}) = is_unsat(SATProblem(v))

@testset "SimpleSATSolver.jl" begin
    @test is_sat([ [2] ])
    @test is_sat([ [1, -5, 4], [-1, 5, 3, 4], [-3, -4] ])
    @test is_sat([[1, 2], [-1, 3], [-3, -2], [-1, 3], [-1, -3]])

    @test is_unsat([ [1], [-1] ])

    @test is_sat([ [1, 2], [-1, 3], [-3, -2], [-1, 3], [-1, -3] ])  # forces backtracking

    # 2-colouring of linear graph with 3 vertices:
    @test is_sat( [[1, 2], [-1, -2], [3, 4], [-3, -4], [5, 6], [-5, -6], [-1, -3], [-2, -4], [-3, -5], [-4, -6]] )

    # 3-colouring of ring with 11 vertices:
    @test is_sat( [[1, 2, 3], [-1, -2], [-1, -3], [-2, -3], [4, 5, 6], [-4, -5], [-4, -6], [-5, -6], [7, 8, 9], [-7, -8], [-7, -9], [-8, -9], [10, 11, 12], [-10, -11], [-10, -12], [-11, -12], [13, 14, 15], [-13, -14], [-13, -15], [-14, -15], [16, 17, 18], [-16, -17], [-16, -18], [-17, -18], [19, 20, 21], [-19, -20], [-19, -21], [-20, -21], [22, 23, 24], [-22, -23], [-22, -24], [-23, -24], [25, 26, 27], [-25, -26], [-25, -27], [-26, -27], [28, 29, 30], [-28, -29], [-28, -30], [-29, -30], [31, 32, 33], [-31, -32], [-31, -33], [-32, -33], [-1, -4], [-2, -5], [-3, -6], [-4, -7], [-5, -8], [-6, -9], [-7, -10], [-8, -11], [-9, -12], [-10, -13], [-11, -14], [-12, -15], [-13, -16], [-14, -17], [-15, -18], [-16, -19], [-17, -20], [-18, -21], [-19, -22], [-20, -23], [-21, -24], [-22, -25], [-23, -26], [-24, -27], [-25, -28], [-26, -29], [-27, -30], [-28, -31], [-29, -32], [-30, -33], [-31, -1], [-32, -2], [-33, -3]])

    @testset "Load CNF file" begin
        p = read_cnf("factorisation.cnf")
        @test is_sat(p)
    end

end

max_stars_repo_name: dpsanders/SimpleSATSolver

max_issues_repo_name: dpsanders/SimpleSATSolver

max_forks_repo_name: dpsanders/SimpleSATSolver


text:
#Julia test

#==Constants
===============================================================================#
module PSFResult
	const Success = 0

	#Exceptions:
	const NotImplemented = 1
	const FileOpenError = 2
	const InvalidFileError = 3
	const FileCloseError = 4
	const NotFound = 5
	const DataSetNotOpen = 6
	const PropertyNotFound = 7
	const Unknown = 8
end #module PSFResult

module PSFElemType
	const E_UNKNOWN = 0
	const E_INT8 = 1
	const E_INT32 = 2
	const E_DOUBLE = 3
	const E_COMPLEXDOUBLE = 4
	const E_STRING = 5
	const E_STRUCT = 6
end #module PSFElemType

#==Main types
===============================================================================#
const PropDict = Dict{String, Any}

mutable struct PSFDataSetRef
	ref::Ptr{Nothing}
end
PSFDataSetRef() = PSFDataSetRef(0)

mutable struct DataReader
	ds::PSFDataSetRef
	filepath::String
	properties::PropDict
	nsig::Int #Not currently used
	npts::Int #Not currently used
	strbuf1::Vector{UInt8} #For reading
	strbuf2::Vector{UInt8} #For reading
end
function DataReader(filepath::String)
	result = DataReader(PSFDataSetRef(), filepath, PropDict(), 0, 0, Array{UInt8}(undef, 1000), Array{UInt8}(undef, 1000))
	result.strbuf1[1] = 0
	result.strbuf2[1] = 0
	return result
end

#==Helper functions
===============================================================================#
function raiseonerror(returnval::Cint)
	if PSFResult.Success == returnval
		return
	elseif PSFResult.NotImplemented == returnval
		error("Not implemented")
	elseif PSFResult.FileOpenError == returnval
		error("File open error")
	elseif PSFResult.InvalidFileError == returnval
		error("Invalid file error")
	elseif PSFResult.FileCloseError == returnval
		error("File close error")
	elseif PSFResult.NotFound == returnval
		error("Not found")
	elseif PSFResult.DataSetNotOpen == returnval
		error("Dataset not open")
	elseif PSFResult.PropertyNotFound == returnval
		error("Property not found")
	else
		error("Unknown error: $returnval")
	end
end

function getpsfelemtype(elemtypeid::Cint)
	if PSFElemType.E_INT8 == elemtypeid
		return Int8;
	elseif PSFElemType.E_INT32 == elemtypeid
		return Int32;
	elseif PSFElemType.E_DOUBLE == elemtypeid
		return Cdouble;
	elseif PSFElemType.E_COMPLEXDOUBLE == elemtypeid
		return Complex{Cdouble};
#	elseif PSFElemType.E_STRING == elemtypeid
#		return Cstring;
#	elseif PSFElemType.E_STRUCT == elemtypeid
#		return ??;
	else
		error("PSF element type not recognized: $elemtypeid.")
	end
end

#readpsfscalar: Read in buffered scalar value.
readpsfscalar(reader::DataReader, ::Type{T}) where T = error("PSF element type not supported $T.")
readpsfscalar(reader::DataReader, ::Type{Int8}) =
	ccall((:PSF_GetScalarInt8, objfile), Int8, (Ptr{Nothing},), reader.ds.ref)
readpsfscalar(reader::DataReader, ::Type{Int32}) =
	ccall((:PSF_GetScalarInt32, objfile), Int32, (Ptr{Nothing},), reader.ds.ref)
readpsfscalar(reader::DataReader, ::Type{Cdouble}) =
	ccall((:PSF_GetScalarDouble, objfile), Cdouble, (Ptr{Nothing},), reader.ds.ref)
readpsfscalar(reader::DataReader, ::Type{Complex{Cdouble}}) =
	ccall((:PSF_GetScalarCplxDouble, objfile), Complex{Cdouble}, (Ptr{Nothing},), reader.ds.ref)

#readpsfvec assuption: v is sized large enough to read in entire vector
readpsfvec(reader::DataReader, v::Vector{T}) where T = error("PSF element type not supported $T.")
readpsfvec(reader::DataReader, v::Vector{Int8}) =
	raiseonerror(ccall((:PSF_CopySigInt8, objfile), Cint, (Ptr{Nothing}, Ptr{Int8}),
		reader.ds.ref, v))
readpsfvec(reader::DataReader, v::Vector{Int32}) =
	raiseonerror(ccall((:PSF_CopySigInt32, objfile), Cint, (Ptr{Nothing}, Ptr{Int32}),
		reader.ds.ref, v))
readpsfvec(reader::DataReader, v::Vector{Cdouble}) =
	raiseonerror(ccall((:PSF_CopySigDouble, objfile), Cint, (Ptr{Nothing}, Ptr{Cdouble}),
		reader.ds.ref, v))
#==Question:
Can we guarantee that all c++ implementations of std::complex<double> uses the same
real/imag ordering as Julia's Complex{Cdouble}?==#
readpsfvec(reader::DataReader, v::Vector{Complex{Cdouble}}) =
	raiseonerror(ccall((:PSF_CopySigComplexDouble, objfile), Cint, (Ptr{Nothing}, Ptr{Complex{Cdouble}}),
		reader.ds.ref, v))

#==Misc readers
===============================================================================#

function readproperties(reader::DataReader)
	result = PropDict()

	if ccall((:PSFHeader_GetPropFirst, objfile), Cint,
		(Ptr{Nothing}, Ptr{UInt8}, Cint, Ptr{UInt8}, Cint),
		reader.ds.ref, reader.strbuf1, length(reader.strbuf1), reader.strbuf2, length(reader.strbuf2)
	) != PSFResult.Success; return result; end
	push!(result, unsafe_string(pointer(reader.strbuf1)) => unsafe_string(pointer(reader.strbuf2)))

	while PSFResult.Success == ccall((:PSFHeader_GetPropNext, objfile), Cint,
		(Ptr{Nothing}, Ptr{UInt8}, Cint, Ptr{UInt8}, Cint),
		reader.ds.ref, reader.strbuf1, length(reader.strbuf1), reader.strbuf2, length(reader.strbuf2))

		push!(result, unsafe_string(pointer(reader.strbuf1)) => unsafe_string(pointer(reader.strbuf2)))
	end

	return result
end

function Base.names(reader::DataReader)
	result = String[]

	if ccall((:PSF_GetSigNamesFirst, objfile), Cint, (Ptr{Nothing}, Ptr{UInt8}, Cint),
		reader.ds.ref, reader.strbuf1, length(reader.strbuf1)
	) != PSFResult.Success; return result; end
	push!(result, unsafe_string(pointer(reader.strbuf1)))

	while PSFResult.Success == ccall((:PSF_GetSigNamesNext, objfile), Cint,
		(Ptr{Nothing}, Ptr{UInt8}, Cint), reader.ds.ref, reader.strbuf1, length(reader.strbuf1))

		push!(result, unsafe_string(pointer(reader.strbuf1)))
	end

	return result
end

#==Open/close/read functions
===============================================================================#
function Base.open(::Type{PSFDataSetRef}, filepath::String)
	result = PSFDataSetRef()
	raiseonerror(ccall((:PSFDataSet_Open, objfile), Cint, (Cstring, Ptr{PSFDataSetRef}),
		filepath, pointer_from_objref(result)
	))
	finalizer(close, result)
	return result
end
function Base.open(::Type{DataReader}, filepath::String)
	result = DataReader(filepath)
	result.ds = open(PSFDataSetRef, result.filepath)
	result.properties = readproperties(result)
	return result
end
_open(filepath::String) = open(DataReader, filepath)

function Base.close(ds::PSFDataSetRef)
	raiseonerror(ccall((:PSFDataSet_Close, objfile), Cint, (Ptr{Nothing},),
		ds.ref
	))
	ds.ref = 0
end

Base.close(reader::DataReader) = close(reader.ds)

#Read in vector buffered by C++ reader:
function _readbufferedvec(reader::DataReader)
	siglen = ccall((:PSF_GetSigLen, objfile), Cint, (Ptr{Nothing},),
		reader.ds.ref
	)
	elemtypeid = ccall((:PSF_GetSigType, objfile), Cint, (Ptr{Nothing},),
		reader.ds.ref
	)
	v = Array{getpsfelemtype(elemtypeid)}(undef, siglen)
	readpsfvec(reader, v)
	return v
end

#Read in sweep vector:
function readsweep(reader::DataReader)
	raiseonerror(ccall((:PSF_ReadSigSweep, objfile), Cint, (Ptr{Nothing},),
		reader.ds.ref
	))
	return _readbufferedvec(reader)
end

#Read in a signal by name:
function Base.read(reader::DataReader, signame::String)
	raiseonerror(ccall((:PSF_ReadSig, objfile), Cint, (Ptr{Nothing}, Cstring),
		reader.ds.ref, signame
	))
	return _readbufferedvec(reader)
end

#Read in a scalar value by name:
function readscalar(reader::DataReader, signame::String)
	raiseonerror(ccall((:PSF_ReadScalar, objfile), Cint, (Ptr{Nothing}, Cstring),
		reader.ds.ref, signame
	))
	elemtypeid = ccall((:PSF_GetScalarType, objfile), Cint, (Ptr{Nothing},),
		reader.ds.ref
	)
	return readpsfscalar(reader, getpsfelemtype(elemtypeid))
end

#Last line


text:
using LinearAlgebra

mutable struct AndersonAcceleration{
    R<:Real,
    C<:Union{R,Complex{R}},
    I<:Integer,
    T<:AbstractArray{C},
    M,
}
    currmem::I
    curridx::I
    s::T
    y::T
    s_M::Vector{T}
    y_M::Vector{T}
end

function AndersonAcceleration(
    x::T,
    M::I,
) where {R<:Real,C<:Union{R,Complex{R}},I<:Integer,T<:AbstractArray{C}}
    s_M = [zero(x) for i = 1:M]
    y_M = [zero(x) for i = 1:M]
    s = zero(x)
    y = zero(x)
    AndersonAcceleration{R,C,I,T,M}(0, 0, s, y, s_M, y_M)
end

function update!(L::AndersonAcceleration{R,C,I,T,M}, s, y) where {R,C,I,T,M}
    L.s .= s
    L.y .= y
    L.curridx += 1
    if L.curridx > M
        L.curridx = 1
    end
    L.currmem += 1
    if L.currmem > M
        L.currmem = M
    end
    copyto!(L.s_M[L.curridx], L.s)
    copyto!(L.y_M[L.curridx], L.y)
    return L
end

function reset!(L::AndersonAcceleration{R,C,I,T,M}) where {R,C,I,T,M}
    L.currmem, L.curridx = zero(I), zero(I)
end

import Base: *

function (*)(L::AndersonAcceleration, v)
    w = similar(v)
    mul!(w, L, v)
end

import LinearAlgebra: mul!

function mul!(d::T, L::AndersonAcceleration{R,C,I,T,M}, v::T) where {R,C,I,T,M}
    if L.currmem == 0
        d .= v
    else
        # TODO: optimize this
        S = hcat(L.s_M[1:L.currmem]...)
        Y = hcat(L.y_M[1:L.currmem]...)
        # H = I + (S - Y)inv(Y'Y)Y'
        d .= v .+ (S - Y) * (pinv(Y' * Y) * (Y' * v))
    end
end


text:
module LinearElasticity

using ConstructionBase: constructorof
using Tensorial: SymmetricSecondOrderTensor, SymmetricFourthOrderTensor, Vec

export TensorStress,
    TensorStrain,
    StiffnessTensor,
    ComplianceTensor,
    EngineeringStress,
    EngineeringStrain,
    ComplianceMatrix,
    StiffnessMatrix

abstract type Stress{T,N} <: AbstractArray{T,N} end
abstract type Strain{T,N} <: AbstractArray{T,N} end
abstract type Stiffness{T,N} <: AbstractArray{T,N} end
abstract type Compliance{T,N} <: AbstractArray{T,N} end
struct TensorStress{T} <: Stress{T,2}
    data::SymmetricSecondOrderTensor{3,T,6}
end
struct TensorStrain{T} <: Strain{T,2}
    data::SymmetricSecondOrderTensor{3,T,6}
end
struct StiffnessTensor{T} <: Stiffness{T,4}
    data::SymmetricFourthOrderTensor{3,T}
end
struct ComplianceTensor{T} <: Compliance{T,4}
    data::SymmetricFourthOrderTensor{3,T}
end
struct EngineeringStress{T} <: Stress{T,1}
    data::Vec{6,T}
end
struct EngineeringStrain{T} <: Strain{T,1}
    data::Vec{6,T}
end
struct StiffnessMatrix{T} <: Stiffness{T,2}
    data::SymmetricSecondOrderTensor{6,T,21}
end
struct ComplianceMatrix{T} <: Compliance{T,2}
    data::SymmetricSecondOrderTensor{6,T,21}
end

TensorStress(x::EngineeringStress) = convert(TensorStress{eltype(x)}, x)

TensorStrain(x::EngineeringStrain) = convert(TensorStrain{eltype(x)}, x)

StiffnessTensor(x::StiffnessMatrix) = convert(StiffnessTensor{eltype(x)}, x)

ComplianceTensor(x::ComplianceMatrix) = convert(ComplianceTensor{eltype(x)}, x)

EngineeringStress(x::TensorStress) = convert(EngineeringStress{eltype(x)}, x)

EngineeringStrain(x::TensorStrain) = convert(EngineeringStrain{eltype(x)}, x)

StiffnessMatrix(x::StiffnessTensor) = convert(StiffnessMatrix{eltype(x)}, x)

ComplianceMatrix(x::ComplianceTensor) = convert(ComplianceMatrix{eltype(x)}, x)

Base.size(::Union{TensorStress,TensorStrain}) = (3, 3)
Base.size(::Union{StiffnessTensor,ComplianceTensor}) = (3, 3, 3, 3)
Base.size(::Union{EngineeringStress,EngineeringStrain}) = (6,)
Base.size(::Union{StiffnessMatrix,ComplianceMatrix}) = (6, 6)

Base.getindex(A::Union{Stress,Strain,Stiffness,Compliance}, i...) = getindex(A.data, i...)

Base.parent(A::Union{Stress,Strain,Stiffness,Compliance}) = A.data

Base.IndexStyle(::Type{<:Union{Stress,Strain,Stiffness,Compliance}}) = IndexLinear()

for T in (:TensorStress, :TensorStrain)
    @eval begin
        $T(m::AbstractMatrix) = $T(SymmetricSecondOrderTensor{3}(m))
        $T(data...) = $T(SymmetricSecondOrderTensor{3}(data...))
    end
end
for T in (:EngineeringStress, :EngineeringStrain)
    @eval begin
        $T(v::AbstractVector) = $T(Vec{6}(v))
        $T(data...) = $T(Vec{6}(data...))
    end
end
for T in (:StiffnessMatrix, :ComplianceMatrix)
    @eval begin
        $T(m::AbstractMatrix) = $T(SymmetricSecondOrderTensor{6}(m))
        $T(data...) = $T(SymmetricSecondOrderTensor{6}(data...))
    end
end

# See https://github.com/JuliaLang/julia/blob/cb9acf5/base/arraymath.jl#L19-L26
for op in (:*, :/)
    @eval Base.$op(A::Union{Stress,Strain,Stiffness,Compliance}, B::Number) =
        constructorof(typeof(A))(Base.broadcast_preserving_zero_d($op, A, B))
end
Base.:*(B::Number, A::Union{Stress,Strain,Stiffness,Compliance}) = A * B

Base.:-(A::Union{Stress,Strain,Stiffness,Compliance}) =
    constructorof(typeof(A))(Base.broadcast_preserving_zero_d(-, A))

for op in (:+, :-)
    for T in (
        :TensorStress,
        :TensorStrain,
        :EngineeringStress,
        :EngineeringStrain,
        :StiffnessMatrix,
        :ComplianceMatrix,
        :StiffnessTensor,
        :ComplianceTensor,
    )
        @eval Base.$op(A::$T, B::$T) = $T(Base.broadcast_preserving_zero_d($op, A, B))
    end
end

include("conversion.jl")
include("invariants.jl")
include("SymmetryCriteria.jl")
# include("StabilityCriteria.jl")
include("Isotropic.jl")
include("distort.jl")
include("fitting.jl")
include("calc.jl")

end


text:
# # SNAP example for hydrogen diffusing in palladium.

using LAMMPS
using Printf

const DATA = joinpath(dirname(pathof(LAMMPS)), "..", "examples", "simple-diffusion", "2_run_md", "data")
const ntypes = 2
const rcut = 2.0
const twojmax = 6
const M  = 100 # number of input files
const N1 = 1 # number of atoms of the first type
const N2 = 64 # number of atoms of the second type
# SNAP radii and weights (see documentation on compute sna/atom).
const r1 = 0.5
const r2 = 0.5
const w1 = 0.5
const w2 = 1.0

if (twojmax%2==0)
    mm = (twojmax/2) + 1
    K = mm*(mm+1)*(2*mm+1)/6
end

@printf("Number of bispectrum components per atom: %d\n", K)

J = twojmax / 2
ncoeff = round(Int, (J+1)*(J+2)*((J+(1.5))/3) + 1)
@printf("Number of bispectrum coefficients: %d\n", ncoeff)

function calc_bispectrum(lmp, path, rcut, twojmax)
    read_data_str = "read_data " * path

    #print(read_data_str)

    command(lmp, "log none")
    command(lmp, "units metal")
    command(lmp, "boundary p p p")
    command(lmp, "atom_style atomic")
    command(lmp, "atom_modify map array")
    command(lmp, read_data_str)
    command(lmp, "pair_style zero $rcut")
    command(lmp, "pair_coeff * *")
    command(lmp, "compute PE all pe")
    command(lmp, "compute S all pressure thermo_temp")
    #command(lmp, "compute SNAP all snap $rcut 0.99363 $twojmax 0.5 0.5 1.0 0.5 rmin0 0.0 bzeroflag 0 quadraticflag 0 switchflag 1")
    command(lmp, "compute SNA all sna/atom $rcut 0.99363 $twojmax $r1 $r2 $w1 $w2 rmin0 0.0 bzeroflag 0 quadraticflag 0 switchflag 1")
    command(lmp, "compute SNAD all snad/atom $rcut 0.99363 $twojmax $r1 $r2 $w1 $w2 rmin0 0.0 bzeroflag 0 quadraticflag 0 switchflag 1")
    command(lmp, "compute SNAV all snav/atom $rcut 0.99363 $twojmax $r1 $r2 $w1 $w2 rmin0 0.0 bzeroflag 0 quadraticflag 0 switchflag 1")
    command(lmp, "thermo_style custom pe")
    command(lmp, "run 0")

    ## Extract bispectrum
    bs = extract_compute(lmp, "SNA", LAMMPS.API.LMP_STYLE_ATOM,
                                     LAMMPS.API.LMP_TYPE_ARRAY)

    bsd = extract_compute(lmp, "SNAD", LAMMPS.API.LMP_STYLE_ATOM, LAMMPS.API.LMP_TYPE_ARRAY)

    #btest = extract_compute(lmp, "SNAP", LAMMPS.API.LMP_STYLE_GLOBAL, LAMMPS.API.LMP_TYPE_ARRAY)
    #sz_btest = size(btest)
    #print(sz_btest)
    #print("\n")

    types = extract_atom(lmp, "type")
    #print(types)

    sz_bs = size(bs)
    #print(sz_bs)
    #print("\n")
    sz_bsd = size(bsd)
    #print(sz_bsd)
    #print("\n")

    #display("text/plain", bs)

    natoms = get_natoms(lmp)
    #@printf("%d atoms\n", natoms)
    return bs, natoms
end

function generate_bispectrum_matrix(path, rcut, twojmax, M, N1, N2)

    A = LMP(["-screen","none"]) do lmp
        A = Array{Float64}(undef, M, 2*ncoeff)
        for m in 1:M

            if (m%100==0)
                @printf "Config %d\n" m
            end
            data = joinpath(path, string(m), "DATA")
            bs, natoms = calc_bispectrum(lmp, data, rcut, twojmax)

            ## Make bispectrum sum vector
            row = Float64[]

            push!(row, N1)

            for k in 1:(ncoeff-1)
                acc = 0.0
                for n in 1:N1
                    acc += bs[k, n]
                end
                push!(row, acc)
            end

            push!(row, N2)
            for k in 1:(ncoeff-1)
                acc = 0.0
                for n in N1 .+ (1:N2)
                    acc += bs[k, n]
                end
                push!(row, acc)
            end

            A[m, :] = row

            command(lmp, "clear")
        end
        return A
    end
    return A, ncoeff
end

function generate_target_array(path,M,N1,N2)

    y = Array{Float64}(undef, M, 1)
    #sz_y = size(y)
    #print(sz_y)
    #print("\n")

    for m in 1:M
        if (m%100==0)
            @printf "Config %d\n" m
        end
        open(joinpath(DATA, string(m), "PE")) do f
            s = split(readline(f))
            #print(s)
            r = parse(Float64, s[1])
            #print(r)
            #print("\n")
            #push!(y, r)
            y[m]=r
            #pe = parse(Float64, s)
            #print(pe)
        end
    end

    return y
end

print("Generating bispectrum matrix (A).\n")
A, ncoeff = generate_bispectrum_matrix(DATA, rcut, twojmax, M, N1, N2)
print("Generating target array (y).\n")
y = generate_target_array(DATA, M, N1, N2)
#sz_A = size(A)
#print(sz_A)
#print("\n")

# Solve for the coefficients.
print("Solving for coefficients (β).\n")
β = A \ y

# Write a LAMMPS snapcoeff file
open("fit.snapcoeff", "w") do io
    write(io, "#LAMMPS SNAP coeffs for hydrogen diffusing in palladium.\n")
    write(io, "\n")
    println(io, ntypes, " ", ncoeff)
    println(io, "Pd ", r1, " ", w1)
    for b in 1:ncoeff
        println(io, β[b])
    end
    println(io, "H ", r2, " ", w2)
    for b in (ncoeff+1):(2*ncoeff)
        println(io, β[b])
    end
end
# Write a LAMMPS snapparam file
open("fit.snapparam", "w") do io
    write(io, "#LAMMPS SNAP coeffs for hydrogen diffusing in palladium.\n")
    write(io, "\n")
    write(io, "#required\n")
    println(io, "rcutfac ", rcut)
    println(io, "twojmax ", twojmax)
    write(io, "\n")
    write(io, "#optional\n")
    println(io, "rfac0 ", 0.99363)
    println(io, "rmin0 ",  0)
    println(io, "bzeroflag ", 0)
    println(io, "quadraticflag ", 0)
end


text:
# Generic concepts and functions

# a premetric is a function d that satisfies:
#
#   d(x, y) >= 0
#   d(x, x) = 0
#
abstract type PreMetric end

# a semimetric is a function d that satisfies:
#
#   d(x, y) >= 0
#   d(x, x) = 0
#   d(x, y) = d(y, x)
#
abstract type SemiMetric <: PreMetric end

# a metric is a semimetric that satisfies triangle inequality:
#
#   d(x, y) + d(y, z) >= d(x, z)
#
abstract type Metric <: SemiMetric end

evaluate(dist::PreMetric, a, b) = dist(a, b)

# Generic functions

"""
    result_type(dist::PreMetric, Ta::Type, Tb::Type) -> T
    result_type(dist::PreMetric, a::AbstractArray, b::AbstractArray) -> T

Infer the result type of metric `dist` with input type `Ta` and `Tb`, or input
data `a` and `b`.
"""
result_type(::PreMetric, ::Type, ::Type) = Float64 # fallback
result_type(dist::PreMetric, a::AbstractArray, b::AbstractArray) = result_type(dist, eltype(a), eltype(b))

# Generic column-wise evaluation

function colwise!(r::AbstractArray, metric::PreMetric, a::AbstractVector, b::AbstractMatrix)
    n = size(b, 2)
    length(r) == n || throw(DimensionMismatch("Incorrect size of r."))
    @inbounds for j = 1:n
        r[j] = metric(a, view(b, :, j))
    end
    r
end

function colwise!(r::AbstractArray, metric::PreMetric, a::AbstractMatrix, b::AbstractVector)
    n = size(a, 2)
    length(r) == n || throw(DimensionMismatch("Incorrect size of r."))
    @inbounds for j = 1:n
        r[j] = metric(view(a, :, j), b)
    end
    r
end

function colwise!(r::AbstractArray, metric::PreMetric, a::AbstractMatrix, b::AbstractMatrix)
    n = get_common_ncols(a, b)
    length(r) == n || throw(DimensionMismatch("Incorrect size of r."))
    @inbounds for j = 1:n
        r[j] = metric(view(a, :, j), view(b, :, j))
    end
    r
end

function colwise!(r::AbstractArray, metric::SemiMetric, a::AbstractMatrix, b::AbstractVector)
    colwise!(r, metric, b, a)
end

function colwise(metric::PreMetric, a::AbstractMatrix, b::AbstractMatrix)
    n = get_common_ncols(a, b)
    r = Vector{result_type(metric, a, b)}(undef, n)
    colwise!(r, metric, a, b)
end

function colwise(metric::PreMetric, a::AbstractVector, b::AbstractMatrix)
    n = size(b, 2)
    r = Vector{result_type(metric, a, b)}(undef, n)
    colwise!(r, metric, a, b)
end

function colwise(metric::PreMetric, a::AbstractMatrix, b::AbstractVector)
    n = size(a, 2)
    r = Vector{result_type(metric, a, b)}(undef, n)
    colwise!(r, metric, a, b)
end

# Generic pairwise evaluation

function _pairwise!(r::AbstractMatrix, metric::PreMetric,
                    a::AbstractMatrix, b::AbstractMatrix=a)
    na = size(a, 2)
    nb = size(b, 2)
    size(r) == (na, nb) || throw(DimensionMismatch("Incorrect size of r."))
    @inbounds for j = 1:size(b, 2)
        bj = view(b, :, j)
        for i = 1:size(a, 2)
            r[i, j] = metric(view(a, :, i), bj)
        end
    end
    r
end

function _pairwise!(r::AbstractMatrix, metric::SemiMetric, a::AbstractMatrix)
    n = size(a, 2)
    size(r) == (n, n) || throw(DimensionMismatch("Incorrect size of r."))
    @inbounds for j = 1:n
        aj = view(a, :, j)
        for i = (j + 1):n
            r[i, j] = metric(view(a, :, i), aj)
        end
        r[j, j] = 0
        for i = 1:(j - 1)
            r[i, j] = r[j, i]   # leveraging the symmetry of SemiMetric
        end
    end
    r
end

function deprecated_dims(dims::Union{Nothing,Integer})
    if dims === nothing
        Base.depwarn("implicit `dims=2` argument now has to be passed explicitly " *
                     "to specify that distances between columns should be computed",
                     :pairwise!)
        return 2
    else
        return dims
    end
end

"""
    pairwise!(r::AbstractMatrix, metric::PreMetric,
              a::AbstractMatrix, b::AbstractMatrix=a; dims)

Compute distances between each pair of rows (if `dims=1`) or columns (if `dims=2`)
in `a` and `b` according to distance `metric`, and store the result in `r`.
If a single matrix `a` is provided, compute distances between its rows or columns.

`a` and `b` must have the same numbers of columns if `dims=1`, or of rows if `dims=2`.
`r` must be a square matrix with size `size(a, dims) == size(b, dims)`.
"""
function pairwise!(r::AbstractMatrix, metric::PreMetric,
                   a::AbstractMatrix, b::AbstractMatrix;
                   dims::Union{Nothing,Integer}=nothing)
    dims = deprecated_dims(dims)
    dims in (1, 2) || throw(ArgumentError("dims should be 1 or 2 (got $dims)"))
    if dims == 1
        na, ma = size(a)
        nb, mb = size(b)
        ma == mb || throw(DimensionMismatch("The numbers of columns in a and b " *
                                            "must match (got $ma and $mb)."))
    else
        ma, na = size(a)
        mb, nb = size(b)
        ma == mb || throw(DimensionMismatch("The numbers of rows in a and b " *
                                            "must match (got $ma and $mb)."))
    end
    size(r) == (na, nb) ||
        throw(DimensionMismatch("Incorrect size of r (got $(size(r)), expected $((na, nb)))."))
    if dims == 1
        _pairwise!(r, metric, transpose(a), transpose(b))
    else
        _pairwise!(r, metric, a, b)
    end
end

function pairwise!(r::AbstractMatrix, metric::PreMetric, a::AbstractMatrix;
                   dims::Union{Nothing,Integer}=nothing)
    dims = deprecated_dims(dims)
    dims in (1, 2) || throw(ArgumentError("dims should be 1 or 2 (got $dims)"))
    if dims == 1
        n, m = size(a)
    else
        m, n = size(a)
    end
    size(r) == (n, n) ||
        throw(DimensionMismatch("Incorrect size of r (got $(size(r)), expected $((n, n)))."))
    if dims == 1
        _pairwise!(r, metric, transpose(a))
    else
        _pairwise!(r, metric, a)
    end
end

"""
    pairwise(metric::PreMetric, a::AbstractMatrix, b::AbstractMatrix=a; dims)

Compute distances between each pair of rows (if `dims=1`) or columns (if `dims=2`)
in `a` and `b` according to distance `metric`. If a single matrix `a` is provided,
compute distances between its rows or columns.

`a` and `b` must have the same numbers of columns if `dims=1`, or of rows if `dims=2`.
"""
function pairwise(metric::PreMetric, a::AbstractMatrix, b::AbstractMatrix;
                  dims::Union{Nothing,Integer}=nothing)
    dims = deprecated_dims(dims)
    dims in (1, 2) || throw(ArgumentError("dims should be 1 or 2 (got $dims)"))
    m = size(a, dims)
    n = size(b, dims)
    r = Matrix{result_type(metric, a, b)}(undef, m, n)
    pairwise!(r, metric, a, b, dims=dims)
end

function pairwise(metric::PreMetric, a::AbstractMatrix;
                  dims::Union{Nothing,Integer}=nothing)
    dims = deprecated_dims(dims)
    dims in (1, 2) || throw(ArgumentError("dims should be 1 or 2 (got $dims)"))
    n = size(a, dims)
    r = Matrix{result_type(metric, a, a)}(undef, n, n)
    pairwise!(r, metric, a, dims=dims)
end


text:
# Various functions regarding graph connectivity

export components, num_components, is_connected, spanning_forest, random_spanning_forest
export find_path, dist, diam, is_cut_edge, is_acyclic, wiener_index
export eccentricity, radius, graph_center

"""
`components(G)` returns the vertex sets of the connected components of
`G` (as a `Partition`).
"""
function components(G::SimpleGraph{T}) where {T}
  if cache_check(G,:components)
    return cache_recall(G,:components)
  end

  P = Partition(G.V)
  for e in G.E
    (u,v) = e
    merge_parts!(P,u,v)
  end
  cache_save(G,:components,P)
  return P
end

"""
`num_components(G)` returns the number of connected components in `G`.
"""
function num_components(G::SimpleGraph{T})::Int where {T}
  if cache_check(G,:num_components)
    return cache_recall(G,:num_components)
  end
  result = num_parts(components(G))
  cache_save(G,:num_components,result)
  return result
end

# determine if the graph is connected
"""
`is_connected(G)` determines if `G` is connected.
"""
function is_connected(G::SimpleGraph{T}) where {T}
    return num_components(G) <= 1
end

# create a spanning forest of G, i.e., a maximal acyclic spanning
# subgraph. If G is connected, this is a tree.

"""
`spanning_forest(G)` creates a maximal acyclic subgraph of `G`.
"""
function spanning_forest(G::SimpleGraph{T}) where {T}
    if cache_check(G,:spanning_forest)
      return cache_recall(G,:spanning_forest)
    end
    H = SimpleGraph{T}()
    if NV(G) == 0
        return H
    end

    for v in G.V
        add!(H,v)
    end

    P = Partition(G.V)

    for e in G.E
        (u,v) = e
        if in_same_part(P,u,v)
            continue
        end
        add!(H,u,v)
        merge_parts!(P,u,v)
        if num_parts(P)==1
            break
        end
    end
    cache_save(G,:spanning_forest,H)
    return H
end

"""
`random_spanning_forest(G)` returns a random spanning forest of the graph `G`.
Each component of the result spans a component of `G`.

This differs from `spanning_forest` in that repeated invocations of this function
can return different results.
"""
function random_spanning_forest(G::SimpleGraph)
    VT = vertex_type(G)
    ET = Tuple{VT,VT}

    T = SimpleGraph{VT}()  # output

    VV = vlist(G)
    EE = elist(G)

    for v in VV
        add!(T,v)
    end

    wt = Dict{ET,Float64}()
    for e in EE
        wt[e] = rand()
    end

    wt_list = [ wt[e] for e in EE ]
    p = sortperm(wt_list)
    E_list = [ EE[j] for j in p ]

    P = Partition(G.V)

    for e in E_list
        a,b = e
        if !in_same_part(P,a,b)
            add!(T,a,b)
            merge_parts!(P,a,b)
        end
    end
    return T
end

"""
`find_path(G,s,t)` finds a shortest path from `s` to `t`. If no
such path exists, an empty list is returned.
"""
function find_path(G::AbstractSimpleGraph,s,t)
    T = vertex_type(G)
    if ~has(G,s) || ~has(G,t)
        error("Source and/or target vertex is not in this graph")
    end
    if s==t
        #   result = Array{T}(undef,1)
        #   result[1] = s
        return [s]
    end

    # set up a queue for vertex exploration
    Q = Queue{T}()
    enqueue!(Q,s)

    # set up trace-back dictionary
    tracer = Dict{T,T}()
    tracer[s] = s

    while length(Q) > 0
        v = dequeue!(Q)
        Nv = G.N[v]
        for w in Nv
            if haskey(tracer,w)
                continue
            end
            tracer[w] = v
            enqueue!(Q,w)

            if w==t  # success!
                path = Array{T}(undef,1)
                path[1] = t
                while path[1] != s
                    v = tracer[path[1]]
                    pushfirst!(path,v)
                end
                return path

            end
        end
    end
    return T[]   # return empty array if no path found
end

# functions for finding distances between vertices in a graph. the
# distance between two vertices is the number of edges in a shortest
# path between them. if there is no such path, the distance is
# undefined (or infinite) but since we want the return value to be an
# Int, we use -1 to signal this.

# find the distance between specified vertices

"""
`dist(G,u,v)` finds the length of a shortest path from `u` to `v` in
`G`. Returns `-1` if no such path exists.

`dist(G,u)` finds the distance from `u` to all other vertices in
`G`. Result is returned as a `Dict`.

`dist(G)` finds all pairs of distances in `G`. Result is a `Dict`
whose `[u,v]` entry is the distance from `u` to `v`.
"""
function dist(G::AbstractSimpleGraph,u,v)
    if !has(G,u) || !has(G,v)
        error("One or both of $u and $v are not vertices of this graph")
    end
    if typeof(G) <: SimpleGraph && cache_check(G,:dist)
        d = cache_recall(G,:dist)
        return d[u,v]
    end
    if u==v
        return 0
    end
    return length(find_path(G,u,v))-1
end

# find all distances from a given vertex
function dist(G::AbstractSimpleGraph, v)
    T = vertex_type(G)
    d = Dict{T,Int}()
    if !has(G,v)
        error("Given vertex is not in this graph")
    end

    d[v] = 0
    Q = Queue{T}()
    enqueue!(Q,v)

    while length(Q)>0
        w = dequeue!(Q)  # get 1st vertex in the queue
        Nw = G.N[w]
        for x in Nw
            if !haskey(d,x)
                d[x] = d[w]+1
                enqueue!(Q,x)
            end
        end
    end

    # record -1 for any vertices we missed
    for x in G.V
        if !haskey(d,x)
            d[x] = -1
        end
    end

    return d
end

# find all distances between all vertices
function dist(G::AbstractSimpleGraph)
    if cache_check(G,:dist)
      return cache_recall(G,:dist)
    end

    T = vertex_type(G)
    dd = Dict{Tuple{T,T},Int}()
    vtcs = vlist(G)

    for v in vtcs
        d = dist(G,v)
        for w in vtcs
            dd[(v,w)] = d[w]
        end
    end
    cache_save(G,:dist,dd)
    return dd
end

"""
`eccentricity(G,v)` returns the eccentricty of vertex `v`
in the graph `G`. This is the maximum distance from `v`
to another vertex (or -1 if the graph is not connected).
"""
function eccentricity(G::SimpleGraph, v)
  if !has(G,v)
    error("$v is not a vertex of this graph")
  end
  d = collect(values(dist(G,v)))
  if minimum(d)<0
    return -1
  end
  return maximum(d)
end

"""
`graph_center(G)` returns the set of vertices of a `SimpleGraph` with minimum
eccentricities.
"""
function graph_center(G::SimpleGraph)::Set
    if cache_check(G,:graph_center)
        return cache_recall(G,:graph_center)
    end
    if G.cache_flag
        dist(G) # force all pairs distance computation
    end

    xtable = Dict( (v, eccentricity(G,v)) for v in G.V )
    min_r = minimum(values(xtable))

    if min_r < 0
        return deepcopy(G.V)
    end

    A = Set( v for v in keys(xtable) if xtable[v]==min_r)
    cache_save(G,:graph_center, A)
    return A
end

"""
`radius(G)` returns the radius of the graph `G`. This is the
minimum `eccentricity` of a vertex of `G` (or -1 if the graph
is not connected).
"""
function radius(G::SimpleGraph)
  if cache_check(G,:radius)
    return cache_recall(G,:radius)
  end
  D = dist_matrix(G)
  if minimum(D)<0
    return -1
  end
  r = minimum(maximum(D,dims=1))
  cache_save(G,:radius,r)
  return r
end

"""
`wiener_index(G)` is the sum of the distances between vertices in `G`.
Returns -1 if `G` is not connected.
"""
function wiener_index(G::SimpleGraph)::Int
  if is_connected(G)
    return div(sum(values(dist(G))),2)
  end
  return -1
end

# Calculate the diameter of a graph, but return -1 if not connected.

"""
`diam(G)` returns the diameter of `G` or `-1` if `G` is not connected.
"""
function diam(G::SimpleGraph)
    if is_connected(G)
        return maximum(values(dist(G)))
    end
    return -1
end

# Determine if a given edge in a graph is a cut edge. If there is no
# such edge in the graph, an error is raised.

"""
`is_cut_edge(G,u,v)` [or `is_cut_edge(G,e)`] determins if `(u,v)` [or
`e`] is a cut edge of `G`.
"""
function is_cut_edge(G::SimpleGraph, u, v)
    if !has(G,u,v)
        error("No such edge in this graph")
    end

    SimpleGraphs.delete!(G,u,v)
    P = find_path(G,u,v)
    result = false
    if length(P)==0
        result = true
    end
    add!(G,u,v)
    return result
end

# When called as is_cut_edge(G,e), we assume e is a tuple or list
# whose first two entries are the end points of the edge
function is_cut_edge(G::SimpleGraph, e)
    return is_cut_edge(G,e[1],e[2])
end

"""
`is_acyclic(G)` returns `true` if `G` has no cycles and `false`
otherwise.
"""
function is_acyclic(G::SimpleGraph)
  n = NV(G)
  m = NE(G)
  c = num_components(G)
  return m == n-c
end


text:
#=
    Figure 5
=#

using Inference
using Model

using ColorSchemes
using JLD2
using Latexify
using Plots
using StatsPlots
using Random

include("../figure_defaults.jl")

##############################################################
## Load model and results
##############################################################

include("../../Results/MainResult_Setup.jl")
@load "Results/MainResult.jld2"

##############################################################
## Table S1
##############################################################

tabS1 = DataFrame(summarystats(C))[:,[1,2,3,6,7]]
tabS1.bestfit = θ

# Reorder
tabS1 = tabS1[!,[1,6,2,3,4,5]]

# Rounding
for i in 1:ncol(tabS1)
    if typeof(tabS1[1,i]) <: Number
        tabS1[:,i] = round.(tabS1[:,i],sigdigits=4)
    end
end

res = latexify(tabS1,env=:tabular)
write("$(@__DIR__)/tabS1.tex",res)

max_stars_repo_name: ap-browning/internalisation

max_issues_repo_name: ap-browning/internalisation

max_forks_repo_name: ap-browning/internalisation


text:
function dotAfterFoldAlongX(lineX, dot) # Vertical Line
    if (dot[1] > lineX)
        return [lineX - abs(lineX-dot[1]), dot[2]]
    end

    return dot
end

function dotAfterFoldAlongY(lineY, dot) # Horizontal Line
    if (dot[2] > lineY)
        return [dot[1], lineY - abs(lineY-dot[2])]
    end

    return dot
end

content = readlines("input.txt")
splitIndex = findfirst(isempty, content)

atEnd(vec) = vec[end]

myParse(x) = parse.(Int, x)

dots, instructions = myParse.(split.(content[1:splitIndex-1], ",")), split.(atEnd.(split.(content[splitIndex+1:end], " ")), "=")

function doFold(instruction)
    if instruction[1] == "x"
        return Set([dotAfterFoldAlongX(parse(Int, instruction[2]), dot) for dot in dots])
    else
        return Set([dotAfterFoldAlongY(parse(Int, instruction[2]), dot) for dot in dots])
    end
end

#! Part 1

length(doFold(instructions[1]))

#! Part 2

for instruction in instructions
    dots = doFold(instruction)
end

using Plots

plot([(dot[1], dot[2]) for dot in dots], seriestype = :scatter, yflip = true, aspect_ratio=:equal, c=:black)

max_stars_repo_path: Day 13/main.jl

max_issues_repo_path: Day 13/main.jl

max_forks_repo_path: Day 13/main.jl


text:
import Base.Test: @test,@test_approx_eq,@test_approx_eq_eps,@test_throws

@everywhere import GeneralizedMetropolisHastings:
    ###Types
    MHRunnerType,InitializeFrom,ProposeFrom,IndicatorType,JobSegments,ChainType,StoreDuring,
    MHRuntimePolicy,
    ###Types
    AbstractParameter,ParameterDefault,ParameterUnivariate,
    DataArray,DataFunction,
    AbstractBactrian,Bactrian,
    AbstractNoiseModel,NoiseModelGaussian,
    BaseSample,GradientSample,TensorSample,TangentTensorSample,
    SymmetricDensity,ASymmetricDensity,DistributionWrapper,NormalDensity,LogNormalDensity,
    MonitorTuner,ScaleTuner,
    AbstractModel,TargetModel,ODEModel,
    ChainStandard,ChainGradient,
    IndicatorStationary,
    GMHSegment,MHRemoteSegments,
    AbstractMHRunner,SMHRunner,GMHRunner,
    ###Functions
    trait,traitvalue,traittype,policy, #form policies.jl
    parameter,parameters,initvalues!,initvalues,logprior!,logprior, #from parameters.jl
    data,numvalues,numvars,generate!,dataindex,datavalues, #from data.jl
    noise,loglikelihood,applynoise!, #from noise.jl
    samples,numparas,numsamples,numtangents,sampletype,calculationtype,similar,offset!, #from samples.jl
    distribution,distributions,recenter,rescale,update, #from distributions.jl
    density,condition!,scale!,update!,propose!,logprobability,logprobability!,issymmetric, #from densities.jl
    sampler,samplerstate,setfrom!,propose!,acceptance!,acceptance,tune!,from,proposals,prepare!,prepareindicator!,prepareauxiliary!, #from samplers
    tuner,tunerstate,rate,accepted,proposed,total,index,numtunesteps,current,tune,nextindex!,period,verbose,needstuning,accepted!,showstep, #from tuners.jl
    model,geometry!,evaluate!,loglikelihood,measurements,noisemodel,#from models.jl
    chain,store!,accepted!,logposterior, #from chains.jl
    indicator,numproposals,transitionprobability!,sampleindicator!,indicatorsamples,accepted, #from indicators.jl
    segment,numproposals,iterate!,getsamples, #from jobsegments.jl
    remotesegments,numsegments,numproposalspersegment,numtotalproposals,iterate!,retrievesamples!,
    runner,run!,initialize!,preparenext!,auxiliary!


text:
"""
    check_existing_dir(dir_path::String)

Checks if the directory under `dir_path` exists. If not, throws IOError
"""
function check_existing_dir(dir_path::String)
    if !isdir(dir_path)
        @warn "Folder $(data_path) does not exists in current directory."
        @warn "Terminating execution."
        throw(ErrorException("Can nor find folder \""*dir_path*"\"."))
    end
end


text:
# This file contains all the functions for simulating the physics of the game.

"""Collisions between ball and the walls"""
function wall_collision!(new_position::Vector_, ball::Ball, arena::Arena)
    # If the height of the ball, in this new position is out of range,
    # then we perform a correction and flip the sign of the vertical velocity.
    if new_position.y >= (arena.dims.y - ball.radius)
        ball.velocity = rebound(ball, :wall)
        return Vector_(new_position.x, arena.dims.y - ball.radius)
    elseif new_position.y <= ball.radius
        ball.velocity = rebound(ball, :wall)
        return Vector_(new_position.x, ball.radius)
    end
    return new_position
end

"""Collisions between ball and the paddle"""
function paddle_collision!(new_position::Vector_, player_a_pos::Vector_, player_b_pos::Vector_, paddle_length::AbstractFloat,
                            ball::Ball, arena::Arena)
    # For ball - paddle collision, we find the whether the ball is in the
    # horizontal and vertical "sweet spot" of the paddle. If it is, we apply a rebound.
    # horizontal sweet spot lies in arena_width - ball.radius to arena_width
    # vertical sweet spot is defined by the lower and higher bounds of the paddles.
    if new_position.x >= arena.dims.x - ball.radius
        lower_bound_b  = max(player_b_pos.y - 3ball.radius/4f0, 0)
        higher_bound_b = min(player_b_pos.y + paddle_length + 3ball.radius/4f0, arena.dims.y)

        if lower_bound_b < new_position.y < higher_bound_b
            ball.velocity = rebound(ball, :paddle)
            return Vector_(arena.dims.x - ball.radius, new_position.y)
        end
    end
    if new_position.x <= ball.radius
        lower_bound_a  = max(player_a_pos.y - 3ball.radius/4f0, 0)
        higher_bound_a = min(player_a_pos.y + paddle_length + 3ball.radius/4f0, arena.dims.y)

        if lower_bound_a < new_position.y < higher_bound_a
            ball.velocity = rebound(ball, :paddle)
            return Vector_(ball.radius, new_position.y)
        end
    end
    return new_position
end

function rebound(ball::Ball, sym::Symbol)
    ball_vx = ball.velocity.x
    ball_vy = ball.velocity.y
    sym == :wall   && (ball_vy *= -1)  # Ball hits the walls
    sym == :paddle && (ball_vx *= -1)  # Ball hits the paddle
    Vector_(ball_vx, ball_vy)
end


text:
@testset "fsm_active_close.jl" begin
    base_seq = WrappingInt32(1 << 31)
    DEFAULT_CAPACITY = 64000
    TIMEOUT_DFLT = 1000
    @testset "start in TIME_WAIT, timeout" begin
        conn = fsm_in_time_wait()
        tick!(conn, 10TIMEOUT_DFLT - 1)
        expect_state(conn, JLSponge.TIME_WAIT)
        tick!(conn, 1)
        expect_notin_state(conn, JLSponge.TIME_WAIT)
        tick!(conn, 10TIMEOUT_DFLT)
        expect_state(conn, JLSponge.CLOSED)
    end

    @testset "start in CLOSING, send ack, time out" begin
        conn = fsm_in_closing()
        tick!(conn, 4TIMEOUT_DFLT)
        expect_one_seg(conn, fin=true)
        expect_state(conn, JLSponge.CLOSING)
        send_ack!(conn, WrappingInt32(2), WrappingInt32(2))
        expect_no_seg(conn)
        expect_state(conn, JLSponge.TIME_WAIT)

        tick!(conn, 10TIMEOUT_DFLT - 1)
        expect_state(conn, JLSponge.TIME_WAIT)
        tick!(conn, 2)
        expect_state(conn, JLSponge.CLOSED)
    end

    @testset "start in FIN_WAIT_2, send FIN, time out" begin
        conn = fsm_in_fin_wait_2()
        tick!(conn, 4TIMEOUT_DFLT)
        expect_state(conn, JLSponge.FIN_WAIT_2)

        rx_seqno = WrappingInt32(1)
        send_fin!(conn, rx_seqno, WrappingInt32(2))

        tick!(conn, 1)
        expect_one_seg(conn; ack=true, ackno=rx_seqno + 1)
        expect_state(conn, JLSponge.TIME_WAIT)

        tick!(conn, 10TIMEOUT_DFLT)
        expect_state(conn, JLSponge.CLOSED)
    end

    @testset "start in FIN_WAIT_1, ack, FIN, time out" begin
        conn = fsm_in_fin_wait_1()
        tick!(conn, 4TIMEOUT_DFLT)
        expect_one_seg(conn; fin=true)
        rx_seqno = WrappingInt32(1)
        send_ack!(conn, rx_seqno, WrappingInt32(2))
        tick!(conn, 5)
        send_fin!(conn, rx_seqno, WrappingInt32(2))
        ack_expect = rx_seqno + 1
        tick!(conn, 1)
        expect_one_seg(conn; no_flag=true, ack=true, ackno=ack_expect)
        tick!(conn, 10TIMEOUT_DFLT)
        expect_state(conn, JLSponge.CLOSED)
    end

    @testset "start in FIN_WAIT_1, ack, FIN, FIN again, time out" begin
        conn = fsm_in_fin_wait_1()

        rx_seqno = WrappingInt32(1)
        send_ack!(conn, rx_seqno, WrappingInt32(2))
        expect_state(conn, JLSponge.FIN_WAIT_2)
        tick!(conn, 5)
        send_fin!(conn, rx_seqno, WrappingInt32(2))
        expect_state(conn, JLSponge.TIME_WAIT)
        @test conn |> time_since_last_segment_received == 0
        ack_expect = rx_seqno + 1
        tick!(conn, 1)
        @test conn |> time_since_last_segment_received == 1
        expect_one_seg(conn, no_flag=true, ack=true, ackno=ack_expect)
        tick!(conn, 10TIMEOUT_DFLT - 10)
        @test conn |> time_since_last_segment_received == 10TIMEOUT_DFLT - 9
        send_fin!(conn, rx_seqno, WrappingInt32(2))
        @test conn |> time_since_last_segment_received == 0
        tick!(conn, 1)
        expect_one_seg(conn, ack=true, ackno=ack_expect)
        expect_state(conn, JLSponge.TIME_WAIT)

        tick!(conn, 10TIMEOUT_DFLT - 10)
        @test conn |> time_since_last_segment_received == 10TIMEOUT_DFLT - 9
        expect_no_seg(conn)
        tick!(conn, 10)
        expect_state(conn, JLSponge.CLOSED)
    end

    @testset "start in FIN_WAIT_1, ack, FIN, FIN again, time out" begin
        conn = fsm_in_established()
        end_input_stream!(conn)

        tick!(conn, 1)
        seg1 = expect_one_seg(conn; fin=true)
        seg1_hdr = seg1.header
        tick!(conn, TIMEOUT_DFLT - 2)
        expect_no_seg(conn)
        tick!(conn, 2)
        seg2 = expect_one_seg(conn; fin=true, seqno=seg1_hdr.seqno)
        seg2_hdr = seg2.header
        rx_seqno = WrappingInt32(1)
        send_fin!(conn, rx_seqno, WrappingInt32(0))
        ack_expect = rx_seqno + 1
        tick!(conn, 1)

        expect_state(conn, JLSponge.CLOSING)
        expect_one_seg(conn, ack=true, ackno=ack_expect)

        send_ack!(conn, ack_expect, seg2_hdr.seqno + 1)
        tick!(conn, 1)
        expect_state(conn, JLSponge.TIME_WAIT)

        tick!(conn, 10TIMEOUT_DFLT)
        expect_state(conn, JLSponge.CLOSED)
    end
end

max_stars_repo_name: AquaIndigo/JLSponge

max_issues_repo_name: AquaIndigo/JLSponge

max_forks_repo_name: AquaIndigo/JLSponge


text:
#========== use TensorOperations when you can ==========#
# This seems to always be faster, when applicable.
# When not, it will return nothing, and we go back the the loops.

function try_tensor(expr, ranges, store)

    fail = nothing
    if isexpr(expr, [:(:=), :(=), :(+=)])
    else
        fail = "TensorOperations not used, expected left := right etc"
    end
    if @capture_(expr.args[1], Z_[leftind__]) && all(a -> a isa Symbol, leftind)
        if expr.head == :(:=)
            store.newarray = true
        end
    elseif expr.args[1] isa Symbol # scalar output
        store.scalar = true # not used?
        Z, leftind = expr.args[1], nothing
        if expr.head == :(:=)
            store.newarray = true
            expr.head = :(=) # mutate it as @tensor doesn't accept scalar :=
        elseif expr.head == :(=)
            store.newarray = true
        end # for scalars, only += case isn't :newarray
    else
        fail = "TensorOperations not used, expected A[i,j,k] := ..."
    end
    MacroTools_postwalk(expr.args[2]) do ex
        ex isa Expr || return ex
        if ex.head == :call && ex.args[1] == :* && all(a -> a isa Expr || a isa Number, ex.args[2:end])
            # Todo: allow A[i] * $c
        elseif ex.head == :ref && all(a -> a isa Symbol, ex.args)
        elseif ex.head == :call && ex.args[1] in [:+, :-] && length(ex.args)==2
            # Allows -A[i]. Could allow conj() too, but gradient would be wrong.
        else
            # Disallows anything containing +, since A[i] + B[i,k,k] has differing meanings.
            fail = "TensorOperations not used, can't handle $(ex)"
        end
        ex
    end
    if fail != nothing
        store.verbose>0 && @warn fail
        return nothing
    end

    outex = [] # you could simplify, only one expression really
    try
        tex = macroexpand(store.mod, :(TensorOperations.@tensor $expr))

        if store.newarray
            left, right = expr.args
            #===== new array =====#

            MacroTools_postwalk(right) do ex
                ex isa Expr || return ex
                # Save array and scalar arguments
                if @capture_(ex, A_[ijk__])
                    A1 = arrayonly(A)
                    push!(store.arrays, A1)
                    push!(store.indices, ijk)
                    n = length(ijk)
                    str = "expected a $n-array $A1"
                    push!(outex, :( $ndims($A1) == $n || $error($str) ))
                elseif ex.head == :call && ex.args[1] == :*
                    foreach(ex.args[2:end]) do a
                        a isa Symbol && push!(store.scalars, a)
                    end
                end
                ex
            end

            if store.grad == false
                push!(outex, tex)
            else
                args = unique(vcat(store.arrays, store.scalars))
                push!(outex, quote
                    local function $MAKE($(args...),)
                        local $Z
                        $tex
                    end
                end)

                ∇make, backdefs = tensor_grad(right, leftind, store)
                append!(outex, backdefs)
                outex = [:($Z = let
                    $(outex...)
                    $Eval($MAKE, $∇make)($(args...))
                end)]
            end
        else
            #===== in-place =====#
            push!(outex, tex)
        end

        # @tensor may return "throw(TensorOperations.IndexError("non-matching indices ..."
        for line in outex
            MacroTools_postwalk(line) do ex
                isexpr(ex, :call) && ex.args[1] == :throw && error(string(ex.args[2]))
                ex
            end
        end
        store.verbose>1 && verbose_tensor(outex, store)
        return outex

    catch err
        store.verbose>0 && @warn "TensorOperations failed" err
        return nothing
    end
end

verbose_tensor(outex, store) = begin
    verboseprint(store)
    printstyled("TensorOperations outex =\n", color=:blue)
    foreach(ex -> printstyled(Base.remove_linenums!(ex) , "\n", color=:green), outex)
end

#========== symbolic gradient ==========#
# Originally TensorGrad.jl (an unregistered package),
# all terms are again @tensor expressions.

function tensor_grad(right, leftind, store)
    dZ = Symbol(DEL, ZED)
    ∇make = Symbol(:∇, MAKE)
    backsteps, backseen = [], []

    for (B, Binds) in zip(store.arrays, store.indices)
        deltaB = Symbol(DEL, B)

        newright, extra, ijk = replace_B_with_Δ(B, Binds, right, leftind)

        append!(backsteps, extra)

        if B in backseen
            addon = macroexpand(store.mod, :( @tensor $deltaB[$(ijk...)] = $deltaB[$(ijk...)] + $newright ))
            push!(backsteps, addon)
            store.verbose>0 && @info "gradient @tensor $deltaB[$(join(ijk,','))] += $newright"
        else
            push!(backseen, B)
            symB = Symbol(DEL, B, '_', join(ijk))
            create = macroexpand(store.mod, :( @tensor( $deltaB[$(ijk...)] := $newright ) ))
            push!(backsteps, create)
            store.verbose>0 && @info "gradient @tensor $deltaB[$(join(ijk,','))] := $newright"
        end
    end

    args = unique(vcat(store.arrays, store.scalars))
    backtuple = vcat(
        map(B -> Symbol(DEL, B), unique(store.arrays)),
        map(_ -> nothing, unique(store.scalars)),
        )

    outex = [:(
        local function $∇make($dZ, $ZED, $(args...))
            $(backsteps...)
            return ($(backtuple...),)
        end
    )]

    if !isnothing(leftind) && isdefined(store.mod, :Zygote) # special case for FillArrays
        # backsteps_fill = fillarrayreplace(backsteps, dZ)
        # ex_value = :($(Symbol(dZ, :_value)) = $dZ.value)
        push!(outex, :(
            local $∇make($dZ::Zygote.Fill, $ZED, $(args...)) = $∇make(collect($dZ), $ZED, $(args...))
            # Todo: make this work without collect! Ideally it would write simpler @tensor expr.
            # local function $∇make($dZ::Zygote.Fill, $ZED, $(args...))
            #     $ex_value
            #     $(backsteps_fill...)
            #     return ($(backtuple...),)
            # end
        ))
    end

    ∇make, outex
end

using LinearAlgebra

function replace_B_with_Δ(B, Bijk, right, leftind)
    dZ = Symbol(DEL, ZED)

    # If B[ijk] occurs twice this will be wrong:
    countB = 0

    # Construct the new RHS
    out = MacroTools_postwalk(right) do x
        if @capture_(x, A_[ijk__]) && A==B && ijk == Bijk
            countB += 1
            if isnothing(leftind)
                return :(conj($dZ)) # scalar case
            else
                return :( conj($dZ[$(leftind...)]) )
            end
        else
            return x
        end
    end
    out = :(conj($out))

    # Deal with partial traces -- repeated indices on same array
    extra, deltas = [], []
    newijk = copy(Bijk)
    if !allunique(Bijk)
        for n in 1:length(Bijk)
            i = newijk[n]
            m = findfirst(isequal(i), newijk[n+1:end])
            if m != nothing
                j = Symbol('_',i,'′')
                newijk[n] = j
                delta = Symbol("_δ_",i,j)

                # This definition is added up front:
                push!(extra, quote
                    local $delta = $Diagonal(fill!(similar($B, real(eltype($B)), size($B,$n)),true))
                end)
                # This factor is included in the new RHS:
                push!(deltas, :( $delta[$i,$j] ))
            end
        end
    end
    if length(extra) > 0
        out = :( *($out, $(deltas...)) )
    end

    # I said:
    # Gradient has indices appearing only on LHS... so you need * ones()[i,j]?

    countB > 1 && error("can't handle case of $B appearing twice with same indices")
    # Could also multiply by countB, and replace just once, would that be safe?

    return out, extra, newijk
end

#========== the end ==========#


text:
module test_julia_pointer1

using Test

struct D
end

struct G
    ptr_d::Ptr{D}
    val
end

d = D()
refd = Ref(d)
ptr_d = pointer_from_objref(refd)
g = G(ptr_d, 3)
refg = Ref(g)
p = Ptr{Base.RefValue{G}}(pointer_from_objref(refg))
refg2 = unsafe_load(p)

@test refg2[].ptr_d == ptr_d
@test refg2[].val == g.val

struct E
end

struct H
    ptr_d::Ptr{D}
    val
    ptr_e::Ptr{E}
end

e = E()
refe = Ref(e)
ptr_e = pointer_from_objref(refe)
h = H(ptr_d, 3, ptr_e)
refh = Ref(h)
p2 = Ptr{Base.RefValue{G}}(pointer_from_objref(refh))
refh2 = unsafe_load(p2)
@test refh2[].ptr_d == ptr_d
@test refh2[].val == h.val
# @test refh2[].ptr_e == ptr_e

end # module test_julia_pointer1

module test_julia_pointer2

using Test

@testset "Ptr{T} Ref{T}" begin
@test C_NULL isa Ptr{Cvoid}

ref = Ref(3)
@test ref[] == 3
end # @testset "Ptr{T} Ref{T}"

@testset "pointer" begin
# pointer(array [, index])
@test pointer([1,2,3]) isa Ptr{Int}
@test pointer([0xff, 0x00]) isa Ptr{UInt8}
end # @testset "pointer"

@testset "pointer_from_objref unsafe_pointer_to_objref" begin
x = [1,2,3]
p = pointer_from_objref(x)
@test unsafe_pointer_to_objref(p) == x
end # @testset "pointer_from_objref unsafe_pointer_to_objref"

@testset "unsafe_load unsafe_store!" begin
# unsafe_load(p::Ptr{T}, i::Integer=1)
# unsafe_store!(p::Ptr{T}, x, i::Integer=1)
# julia test/core.jl - pull request 1270
a = [11,12,13]
p = pointer(a)
@test unsafe_load(p, 1) == 11
unsafe_store!(p, 99, 2)
@test a == [11,99,13]
end # @testset "unsafe_load unsafe_store!"

@testset "unsafe_wrap" begin
# unsafe_wrap(Array, pointer::Ptr{T}, dims; own = false)
a = [1,2,3]
@test unsafe_wrap(Array, pointer(a), length(a)) == a
end # @testset "unsafe_wrap"

foo(x::Int, y::Int) = x + y
@testset "@cfunction Base.unsafe_convert" begin
# @cfunction(callable, ReturnType, (ArgumentTypes...,)) -> Ptr{Cvoid}
fa = @cfunction(foo, Int, (Int, Int))
@test ccall(fa, Int, (Int,Int), 1, 2) == 3

# @cfunction(\$callable, ReturnType, (ArgumentTypes...,)) -> CFunction
fb = @cfunction($foo, Int, (Int, Int))
@test fb isa Base.CFunction
@test ccall(Base.unsafe_convert(Ptr{Cvoid}, fb), Int, (Int,Int), 1, 2) == 3
end # @testset "@cfunction"

@testset "GC.@preserve" begin
f() = 42
cf = @cfunction $f Int ()
GC.@preserve cf begin
    fptr = Base.unsafe_convert(Ptr{Cvoid}, cf)
    b = ccall(fptr, Int, ())
end
@test b == 42
end # @testset "GC.@preserve"

end # module test_julia_pointer2

module test_julia_pointer3

using Test

p = Ptr{Cvoid}(UInt(0x03))
@test UInt32(UInt(p)) == 0x03

end # module test_julia_pointer3


text:
function normalise(logw::Vector)
    nrml = logsumexp(logw)
    return logw - repeat([nrml], length(logw))
end


text:
"""
Bridges
`GlobalCardinality{CP.VARIABLE_COUNTED_VALUES, CP.CLOSED_COUNTED_VALUES, T}`
to `GlobalCardinality{CP.VARIABLE_COUNTED_VALUES, CP.OPEN_COUNTED_VALUES, T}`.
"""
struct GlobalCardinalityVariableClosed2GlobalCardinalityVariableOpenBridge{T} <: MOIBC.AbstractBridge
    cons_domain::Vector{MOI.ConstraintIndex{MOI.VectorAffineFunction{T}, CP.Membership}}
    con_gcv::MOI.ConstraintIndex{MOI.VectorAffineFunction{T}, CP.GlobalCardinality{CP.VARIABLE_COUNTED_VALUES, CP.OPEN_COUNTED_VALUES, T}}
end

function MOIBC.bridge_constraint(
    ::Type{GlobalCardinalityVariableClosed2GlobalCardinalityVariableOpenBridge{T}},
    model,
    f::MOI.VectorOfVariables,
    s::CP.GlobalCardinality{CP.VARIABLE_COUNTED_VALUES, CP.CLOSED_COUNTED_VALUES, T},
) where {T}
    return MOIBC.bridge_constraint(
        GlobalCardinalityVariableClosed2GlobalCardinalityVariableOpenBridge{T},
        model,
        MOI.VectorAffineFunction{T}(f),
        s,
    )
end

function MOIBC.bridge_constraint(
    ::Type{GlobalCardinalityVariableClosed2GlobalCardinalityVariableOpenBridge{T}},
    model,
    f::MOI.VectorAffineFunction{T},
    s::CP.GlobalCardinality{CP.VARIABLE_COUNTED_VALUES, CP.CLOSED_COUNTED_VALUES, T},
) where {T}
    f_scalars = MOIU.scalarize(f)
    f_sought = f_scalars[(s.dimension + s.n_values + 1):(s.dimension + 2 * s.n_values)]

    cons_domain = MOI.ConstraintIndex{MOI.VectorAffineFunction{T}, CP.Membership}[
        MOI.add_constraint(
            model,
            MOIU.vectorize(
                [
                    f_scalars[i],
                    f_sought...
                ]
            ),
            CP.Membership(s.n_values),
        )
        for i in 1:s.dimension
    ]

    con_gcv = MOI.add_constraint(
        model,
        f,
        CP.GlobalCardinality{CP.VARIABLE_COUNTED_VALUES, CP.OPEN_COUNTED_VALUES, T}(s.dimension, s.n_values),
    )

    return GlobalCardinalityVariableClosed2GlobalCardinalityVariableOpenBridge(cons_domain, con_gcv)
end

function MOI.supports_constraint(
    ::Type{GlobalCardinalityVariableClosed2GlobalCardinalityVariableOpenBridge{T}},
    ::Union{Type{MOI.VectorOfVariables}, Type{MOI.VectorAffineFunction{T}}},
    ::Type{CP.GlobalCardinality{CP.VARIABLE_COUNTED_VALUES, CP.CLOSED_COUNTED_VALUES, T}},
) where {T}
    return true
end

function MOIB.added_constrained_variable_types(::Type{GlobalCardinalityVariableClosed2GlobalCardinalityVariableOpenBridge{T}}) where {T}
    return Tuple{Type}[]
end

function MOIB.added_constraint_types(::Type{GlobalCardinalityVariableClosed2GlobalCardinalityVariableOpenBridge{T}}) where {T}
    return [
        (MOI.VectorAffineFunction{T}, CP.Membership),
        (MOI.VectorAffineFunction{T}, CP.GlobalCardinality{CP.VARIABLE_COUNTED_VALUES, CP.OPEN_COUNTED_VALUES, T}),
    ]
end

function MOI.get(
    b::GlobalCardinalityVariableClosed2GlobalCardinalityVariableOpenBridge{T},
    ::MOI.NumberOfConstraints{
        MOI.VectorAffineFunction{T}, CP.Membership,
    },
) where {T}
    return length(b.cons_domain)
end

function MOI.get(
    ::GlobalCardinalityVariableClosed2GlobalCardinalityVariableOpenBridge{T},
    ::MOI.NumberOfConstraints{
        MOI.VectorAffineFunction{T}, CP.GlobalCardinality{CP.VARIABLE_COUNTED_VALUES, CP.OPEN_COUNTED_VALUES, T},
    },
) where {T}
    return 1
end

function MOI.get(
    b::GlobalCardinalityVariableClosed2GlobalCardinalityVariableOpenBridge{T},
    ::MOI.ListOfConstraintIndices{
        MOI.VectorAffineFunction{T}, CP.Membership,
    },
) where {T}
    return copy(b.cons_domain)
end

function MOI.get(
    b::GlobalCardinalityVariableClosed2GlobalCardinalityVariableOpenBridge{T},
    ::MOI.ListOfConstraintIndices{
        MOI.VectorAffineFunction{T}, CP.GlobalCardinality{CP.VARIABLE_COUNTED_VALUES, CP.OPEN_COUNTED_VALUES, T},
    },
) where {T}
    return [b.con_gcv]
end


text:
using CovarianceEstimation
using Statistics
using LinearAlgebra
using Random:seed!
using DelimitedFiles
using ProgressMeter
using PyPlot

const CE = CovarianceEstimation

seed!(6312)

# dimensions

const np = (
    # fat matrices
    (15, 20), (20, 30), (20, 50), (100, 200),
    # tall matrices
    (20, 15), (30, 20), (50, 20), (200, 100),
    )

const LSE = LinearShrinkageEstimator

const estimators = Dict(
    # BASELINE ESTIMATORS
    "s_uncorr" => Simple(),
    "s_corr"   => Simple(corrected=true),
    # LINEAR SHRINKAGE, BASED ON UNCORRECTED SAMPLE COV
    "d1v_lw_uncorr"   => LSE(target=DiagonalUnitVariance(),
                              shrinkage=:lw),
    "d1v_ss_uncorr"   => LSE(target=DiagonalUnitVariance(),
                              shrinkage=:ss),
    "dcv_lw_uncorr"   => LSE(target=DiagonalCommonVariance(),
                              shrinkage=:lw),
    "dcv_ss_uncorr"   => LSE(target=DiagonalCommonVariance(),
                              shrinkage=:ss),
    "dcv_rblw_uncorr" => LSE(target=DiagonalCommonVariance(),
                              shrinkage=:rblw),
    "dcv_oas_uncorr"  => LSE(target=DiagonalCommonVariance(),
                              shrinkage=:oas),
    "duv_lw_uncorr"   => LSE(target=DiagonalUnequalVariance(),
                              shrinkage=:lw),
    "duv_ss_uncorr"   => LSE(target=DiagonalUnequalVariance(),
                              shrinkage=:ss),
    "ccov_lw_uncorr"  => LSE(target=CommonCovariance(),
                              shrinkage=:lw),
    "ccov_ss_uncorr"  => LSE(target=CommonCovariance(),
                              shrinkage=:ss),
    "ppc_lw_uncorr"   => LSE(target=PerfectPositiveCorrelation(),
                              shrinkage=:lw),
    "ppc_ss_uncorr"   => LSE(target=PerfectPositiveCorrelation(),
                              shrinkage=:ss),
    "ccor_lw_uncorr"  => LSE(target=ConstantCorrelation(),
                              shrinkage=:lw),
    "ccor_ss_uncorr"  => LSE(target=ConstantCorrelation(),
                              shrinkage=:ss),
    # LINEAR SHRINKAGE, BASED ON CORRECTED SAMPLE COV
    "d1v_lw_corr"   => LSE(target=DiagonalUnitVariance(),
                            shrinkage=:lw, corrected=true),
    "d1v_ss_corr"   => LSE(target=DiagonalUnitVariance(),
                            shrinkage=:ss, corrected=true),
    "dcv_lw_corr"   => LSE(target=DiagonalCommonVariance(),
                            shrinkage=:lw, corrected=true),
    "dcv_ss_corr"   => LSE(target=DiagonalCommonVariance(),
                            shrinkage=:ss, corrected=true),
    "dcv_rblw_corr" => LSE(target=DiagonalCommonVariance(),
                            shrinkage=:rblw, corrected=true),
    "dcv_oas_corr"  => LSE(target=DiagonalCommonVariance(),
                            shrinkage=:oas, corrected=true),
    "duv_lw_corr"   => LSE(target=DiagonalUnequalVariance(),
                            shrinkage=:lw, corrected=true),
    "duv_ss_corr"   => LSE(target=DiagonalUnequalVariance(),
                            shrinkage=:ss, corrected=true),
    "ccov_lw_corr"  => LSE(target=CommonCovariance(),
                            shrinkage=:lw, corrected=true),
    "ccov_ss_corr"  => LSE(target=CommonCovariance(),
                            shrinkage=:ss, corrected=true),
    "ppc_lw_corr"   => LSE(target=PerfectPositiveCorrelation(),
                            shrinkage=:lw, corrected=true),
    "ppc_ss_corr"   => LSE(target=PerfectPositiveCorrelation(),
                            shrinkage=:ss, corrected=true),
    "ccor_lw_corr"  => LSE(target=ConstantCorrelation(),
                            shrinkage=:lw, corrected=true),
    "ccor_ss_corr"  => LSE(target=ConstantCorrelation(),
                            shrinkage=:ss, corrected=true),
    # NONLINEAR SHRINKAGE, BASED ON UNCORRECTED SAMPLE COV
    "anshrink_uncorr" => AnalyticalNonlinearShrinkage(),
    # NONLINEAR SHRINKAGE, BASED ON CORRECTED SAMPLE COV
    "anshrink_uncorr" => AnalyticalNonlinearShrinkage()
    )

# warmup case
begin
    println("Warmup...")
    X = randn(15, 5)
    cov(Simple(), X)
    cov(Simple(corrected=true), X)
    @showprogress for (_, estimator) ∈ estimators
        cov(estimator, X)
    end
end

# number of time each experiment is ran
nrounds = 50

results = Dict{String, Float64}()

for (n, p) ∈ np
    # Generate a colouring matrix (cholesky of true cov)
    # the first randn allows for different amplitudes, the last one allows
    # for the various dimensions to have their variance on different scales
    L = randn() .* randn(p, p) .* randn(p)
    Σ = L * L'
    p2 = p^2
    println("Dims: $n x $p...")
    @showprogress for round ∈ 1:nrounds
        # generate a sample matrix
        X = randn(n, p) * L
        # estimate the covariance
        for (name, estimator) ∈ estimators
            results[name * "_$(n)x$(p)_$round"] =
                norm(cov(estimator, X) - Σ) / p2
        end
    end
end

reskeys = collect(keys(results))

results2 = Dict{String, NTuple{2, Float64}}()

# print results for each size and every estimator (cumbersome)
for (name, estimator) ∈ estimators
    for (n, p) ∈ np
        # find all the results that match
        matchkeys = filter(k -> occursin("$(name)_$(n)x$(p)_", k), reskeys)
        r = [results[key] for key ∈ matchkeys]
        med, iqr = mean(r), -(quantile(r, [0.75, 0.25])...)
        results2["$(name)_$(n)x$(p)_"] = (med, iqr)
    end
end

res2keys = collect(keys(results2))

# print best estimator per size
for (n, p) ∈ np
    matchkeys = filter(k -> occursin("_$(n)x$(p)_", k), res2keys)
    # dirty way to find the min
    name_best = ""
    val_best = Inf
    iqr_best = Inf
    for key ∈ matchkeys
        med, iqr = results2[key]
        if med < val_best
            name_best = key
            val_best = med
            iqr_best = iqr
        end
    end
    println("Size $(n)x$(p), best: $name_best ($val_best; $iqr_best)")
end

for (n, p) ∈ np
    figure()
    rkeys = filter(κ->occursin("$(n)x$(p)_", κ), res2keys)
    data = Vector{Vector{Float64}}(undef, length(rkeys))
    for (i, k) ∈ enumerate(rkeys)
        # find the results
        matchkeys = filter(κ -> occursin(k, κ), reskeys)
        data[i] = [results[κ] for κ ∈ matchkeys]
    end
    perm = sortperm(rkeys)
    boxplot(data[perm])
    title("Size: $(n)x$(p)")
    xticks(1:length(rkeys), rkeys[perm], rotation=90)
    tight_layout()
    savefig("docs/src/assets/mse_comp/bm_$(n)x$(p).png")
end


text:
#Permuation Recombinations

function pmx{T <: Integer}(v1::Vector{T}, v2::Vector{T})
    s = length(v1)
    from, to = rand(1:s, 2)
    from, to = from > to ? (to, from)  : (from, to)
    c1 = similar(v1)
    c2 = similar(v2)

    # Swap
    c1[from:to] = v2[from:to]
    c2[from:to] = v1[from:to]

    # Fill in from parents
    for i in vcat(1:from-1, to+1:s)
        # Check conflicting offspring
        in1 = inmap(v1[i], c1, from, to)
        if in1 == 0
            c1[i] = v1[i]
        else
            tmpin = in1
            while tmpin > 0
                tmpin = inmap(c2[in1], c1, from, to)
                in1 = tmpin > 0 ? tmpin : in1
            end
            c1[i] = v1[in1]
        end

        in2 = inmap(v2[i], c2, from, to)
        if in2 == 0
            c2[i] = v2[i]
        else
            tmpin = in2
            while tmpin > 0
                tmpin = inmap(c1[in2], c2, from, to)
                in2 = tmpin > 0 ? tmpin : in2
            end
            c2[i] = v2[in2]
        end
    end
    return c1, c2
end

# Order crossover
function ox1{T <: Integer}(v1::Vector{T}, v2::Vector{T})
  s = length(v1)
  from, to = rand(1:s, 2)
  from, to = from > to ? (to, from)  : (from, to)
  c1 = zeros(v1)
  c2 = zeros(v2)
  # Swap
  c1[from:to] = v2[from:to]
  c2[from:to] = v1[from:to]
  # Fill in from parents
  k = to+1 > s ? 1 : to+1 #child1 index
  j = to+1 > s ? 1 : to+1 #child2 index
  for i in vcat(to+1:s,1:from-1)
    while in(v1[k],c1)
      k = k+1 > s? 1 : k+1
    end
    c1[i] = v1[k]
    while in(v2[j],c2)
      j = j+1 > s? 1 : j+1
    end
    c2[i] = v2[j]
  end
  return c1, c2
end

# Cycle crossover
function cx{T <: Integer}(v1::Vector{T}, v2::Vector{T})
  s = length(v1)
  c1 = zeros(v1)
  c2 = zeros(v2)

  f1 = true #switch
  k = 1
  while k > 0
    idx = k
    if f1
      #cycle from v1
      while c1[idx] == zero(T)
        c1[idx] = v1[idx]
        c2[idx] = v2[idx]
        idx = inmap(v2[idx],v1)
      end
    else
      #cycle from v2
      while c2[idx] == zero(T)
        c1[idx] = v2[idx]
        c2[idx] = v1[idx]
        idx = inmap(v1[idx],v2)
      end
    end
    f1 $= true
    k = inmap(zero(T),c2)
  end
  return c1,c2
end

# Order-based crossover
function ox2{T <: Integer}(v1::Vector{T}, v2::Vector{T})
  s = length(v1)
  c1 = copy(v1)
  c2 = copy(v2)

  for i in 1:s
    if rand(Bool)
      idx1 = inmap(v2[i],v1)
      idx2 = inmap(v1[i],v2)
      c1[idx1] = zero(T)
      c2[idx2] = zero(T)
    end
  end

  for i in 1:s
    if !in(v2[i],c1)
      tmpin = inmap(zero(T),c1)
      c1[tmpin] = v2[i]
    end
    if !in(v1[i],c2)
      tmpin = inmap(zero(T),c2)
      c2[tmpin] = v1[i]
    end
  end
  return c1,c2
end

# Position-based crossover
function pos{T <: Integer}(v1::Vector{T}, v2::Vector{T})
  s = length(v1)
  c1 = zeros(v1)
  c2 = zeros(v2)

  for i in 1:s
    if rand(Bool)
      c1[i] = v2[i]
      c2[i] = v1[i]
    end
  end

  for i in 1:s
    if !in(v1[i],c1)
      tmpin = inmap(zero(T),c1)
      c1[tmpin] = v1[i]
    end
    if !in(v2[i],c2)
      tmpin = inmap(zero(T),c2)
      c2[tmpin] = v2[i]
    end
  end
  return c1,c2
end


text:
import Evolutionary: value, population_size, default_options, minimizer,
                     initial_state, update_state!

# mock state
mutable struct TestOptimizerState <: Evolutionary.AbstractOptimizerState
    individual
    fitness
end
Evolutionary.value(state::TestOptimizerState) = state.fitness
Evolutionary.minimizer(state::TestOptimizerState) = state.individual

# mock optimizer
Base.@kwdef struct TestOptimizer <: Evolutionary.AbstractOptimizer
    metrics::Vector{ConvergenceMetric}=[Evolutionary.AbsDiff(1e-10)]
end
Evolutionary.population_size(method::TestOptimizer) = 5
Evolutionary.default_options(method::TestOptimizer) = Dict(:iterations =>10,)
Evolutionary.initial_state(method, options, d, population) =
    TestOptimizerState(population[end], value(d, population[end]))
function Evolutionary.update_state!(d, constraints, state::TestOptimizerState, population::AbstractVector, method, opts, itr)
    i = rand(1:population_size(method))
    state.individual = population[i]
    state.fitness = value!(d, state.individual)
    return false
end


text:
module TokamakIterCylindrical

    import ElectromagneticFields.AxisymmetricTokamakCylindrical

    export pauli_particle_3d_pode, hamiltonian, toroidal_momentum

    AxisymmetricTokamakCylindrical.@code_iter() # inject magnetic field code

    const qᵢ = [7.0-1.4, 0.0, 0.0]
    const vᵢ = [3.43E-3, 6.75, -3.41E-1]

    include("pauli_particle_3d.jl")

    function initial_conditions(x₀, u₀, μ)
        v₀ = u₀ * [b¹(0, x₀), b²(0, x₀), b³(0, x₀)]
        (x₀, v₀, (μ = μ,))
    end

    initial_conditions_barely_passing() = initial_conditions([2.5, 0., 0.], 3.425E-1, 1E-2)
    initial_conditions_barely_trapped() = initial_conditions([2.5, 0., 0.], 3.375E-1, 1E-2)
    initial_conditions_deeply_passing() = initial_conditions([2.5, 0., 0.],  5E-1,    1E-2)
    initial_conditions_deeply_trapped() = initial_conditions([2.5, 0., 0.],  1E-1,    1E-2)
    initial_conditions_trapped()        = initial_conditions([7.0, 0., 0.], -2E-3,    1.88E-7)

end


text:
using Pkg
Pkg.add("DataFrames")
Pkg.add("BioCore")
Pkg.add("BioSymbols")
Pkg.add("GenomeGraphs")
Pkg.add("BioGenerics")
Pkg.add("BigBed")
Pkg.add("BioSequences")
Pkg.add("XAM")

max_stars_repo_name: steinbaugh/koopa

max_issues_repo_name: steinbaugh/koopa

max_forks_repo_name: steinbaugh/koopa


text:
# LuGreModels.jl

# The Lu-Gre state equation (bristle velocity)
function heuristic_state_equation(
    mdl::LuGreModel,
    t::T,
    nrm::T,
    pos::T,
    vel::T,
    z::Array{T}
) where T <: Number
    z0 = z[1]
    Fc = mdl.coulomb_coefficient * nrm
    Fs = mdl.static_coefficient * nrm
    g = Fc + (Fs - Fc) * exp(-(vel / mdl.stribeck_velocity)^2)
    return [vel - mdl.bristle_stiffness * abs(vel) * z0 / g]
end

# The Lu-Gre friction force equation.
function heuristic_force_equation(
    mdl::LuGreModel,
    t::T,
    nrm::T,
    pos::T,
    vel::T,
    z::Array{T},
    dzdt::Array{T}
) where T <: Number
    return mdl.bristle_stiffness * z[1] + mdl.bristle_damping * dzdt[1] +
        mdl.viscous_damping * vel
end

# ------------------------------------------------------------------------------
# Required routines to support friction model fitting.

function model_from_array(mdl::LuGreModel, x::Array{T}) where T <: Number
    LuGreModel(x[1], x[2], x[3], x[4], x[5], x[6])
end

function model_to_array(x::LuGreModel)
    [
        x.static_coefficient,
        x.coulomb_coefficient,
        x.stribeck_velocity,
        x.bristle_stiffness,
        x.bristle_damping,
        x.viscous_damping
    ]
end


text:
#
# This file is part of the YAActL.jl Julia package, MIT license
#
# Paul Bayer, 2020
#

using YAActL, Distributed, Test

length(procs()) == 1 && addprocs(1)

@everywhere using YAActL
@everywhere function ident(id, from)
    id == from ?
        ("local actor",  id, from) :
        ("remote actor", id, from)
end

#
# test registry with one local actor
#
@test register(:act1, Actor(ident, 1))
a1 = YAActL.diag!(:act1)
@test a1.name == :act1  # does it have its name ?
@test call!(:act1, myid()) == ("local actor", 1, 1)
act1 = whereis(:act1)
unregister(:act1)
sleep(0.1)
@test isnothing(a1.name) # did it delete the name ?
@test isempty(registered())
@test register(:act1, act1)
sleep(0.1)
@test length(registered()) == 1
exit!(:act1)
sleep(0.1)
@test act1.chn.state == :closed
@test isempty(registered()) # did it unregister at exit ?
#
# test registry across workers
#
@test register(:act1, Actor(ident, 1))
@test register(:act2, Actor(2, ident, 2))
@test call!(:act2, myid()) == ("remote actor", 2, 1)
@test fetch(@spawnat 2 call!(:act1, myid())) == ("remote actor", 1, 2)
@test fetch(@spawnat 2 call!(:act2, myid())) == ("local actor", 2, 2)
@test whereis(:act1).type == :local
@test whereis(:act2).type == :remote
@test fetch(@spawnat 2 whereis(:act1)).type == :remote
@test fetch(@spawnat 2 whereis(:act2)).type == :remote
let
    r = registered()
    l = [i[1] for i in r]
    @test length(r) == 2
    @test :act1 in l
    @test :act2 in l
    r = fetch(@spawnat 2 registered())
    @test all([i[2].type for i in r]) do x
        x == :remote
    end
end
unregister(:act1)
unregister(:act2)
@test isempty(registered())
#
# test API functions with registered actors
#
f(a, b) = a + b
@test register(:act1, Actor(f, 1))
a1 = YAActL.diag!(:act1)
send!(:act1, YAActL.Cast((1,)))
sleep(0.1)
@test a1.res == 2
@test request!(:act1, YAActL.Call, 2) == 3
become!(:act1, f, 0)
@test call!(:act1, 1) == 1
cast!(:act1, 2)
sleep(0.1)
@test a1.res == 2
@test exec!(:act1, Func(f, 5, 5)) == 10
@test query!(:act1, :res) == 2
set!(:act1, state)
@test query!(:act1, :dsp) == state
update!(:act1, 5)
@test call!(:act1, 5) == 10
unregister(:act1)


text:
# sample from a double well in 1D using HMC
using Plots
using Printf
using Random
using LinearAlgebra
using ForwardDiff
using QuadGK
using BasicMD

include("potentials.jl")

β = 5.0;
x₀ = [-1.0];
seed = 100;
Δt = 1e-2;
n_iters = 10^5; # number of samples

V = x -> DoubleWell(x);
cfg = ForwardDiff.GradientConfig(V, x₀);
gradV! = (gradV, x) -> ForwardDiff.gradient!(gradV, V, x, cfg);

sampler = EM(gradV!, β, Δt);

# define the constraint to to keep the trajectory in the set A = [a,b]
# this occurs after an update is taken
a = -1.2;
b = 0.5;
function constraintA!(state::BasicMD.EMState)
    if state.x[1] < a
        @. state.x = a
        gradV!(state.∇V, state.x)
    end
    if state.x[1] > b
        @. state.x = b
        gradV!(state.∇V, state.x)
    end
    state
end
constraints = Constraints(trivial_constraint!, constraintA!, 1, 1)

Random.seed!(100);
X₀ = copy(x₀);
sample_trajectory!(X₀, sampler, constraints, options = MDOptions(n_iters = n_iters));
@printf("In Place X after %d iterations: %g\n", n_iters, X₀[1])

Random.seed!(100);
Xvals = sample_trajectory(x₀, sampler, constraints, options = MDOptions(n_iters = n_iters, n_save_iters = n_iters));
X = Xvals[end];
@printf("X after %d iterations: %g\n", n_iters, X[1])

Random.seed!(100);
X_vals = sample_trajectory(x₀, sampler, constraints, options = MDOptions(n_iters = n_iters));
histogram([X[1] for X in X_vals], label = "Samples", normalize = true)
xlabel!("x")
ylabel!("Frequency")