text:
module GlobalOptimizer

const VecI  = AbstractVector # Input  Vector
const VecO  = AbstractVector # Output Vector
const VecB  = AbstractVector # Buffer Vector
const VecIO = AbstractVector # In/Out Vector

export minimize!, minimizer

############################
# Sorting                  #
############################

function swap!(v::VecIO, i::Int, j::Int) # @code_warntype ✓
    @inbounds temp = v[i]
    @inbounds v[i] = v[j]
    @inbounds v[j] = temp
    return nothing
end

biinsert(arr::VecI, val::T) where T = biinsert(arr, val, 1, length(arr)) # @code_warntype ✓
function biinsert(arr::VecI, val::T, lx::Int, rx::Int) where T           # @code_warntype ✓
    lx ≥ rx && return lx
    ub = rx # upper bound
    while lx < rx
        mx = (lx + rx) >> 1                                    # midpoint (binary search)
        @inbounds isless(val, arr[mx]) ? rx = mx : lx = mx + 1 # arr[mx].f == val in this case
    end
    @inbounds lx == ub && !isless(val, arr[lx]) && (lx += 1)   # lx = upper bound && arr[lx] ≤ val
    return lx
end

binsort!(arr::VecI) = binsort!(arr, 1, length(arr)) # @code_warntype ✓
function binsort!(arr::VecI, lx::Int, rx::Int)      # @code_warntype ✓
    for ix in lx+1:rx
        @inbounds val = arr[ix]
        jx = ix
        lc = biinsert(arr, val, lx, ix) # location
        while jx > lc
            swap!(arr, jx, jx - 1)
            jx -= 1
        end
    end
end

############################
# Agent                    #
############################

mutable struct Agent
    x::Vector{Float64} # parameters passed into models
    f::Float64         # function-value of fn(x)
    v::Bool            # viability / feasibility
    c::Float64         # contravention / violation

    Agent(nd::Int) = new(Vector{Float64}(undef, nd), Inf, false, 0.0) # @code_warntype ✓
end

function Base.isequal(a1::Agent, a2::Agent) # @code_warntype ✓
    # a1, a2 are both feasible
    a1.v && a2.v && return a1.f == a2.f
    # a1, a2 are both infesasible
    !a1.v && !a2.v && return a1.c == a2.c
    return false
end

function Base.isless(a1::Agent, a2::Agent) # @code_warntype ✓
    # a1, a2 are both feasible
    a1.v && a2.v && return a1.f < a2.f
    # a1, a2 are both infesasible
    !a1.v && !a2.v && return a1.c < a2.c
    # if (a1, a2) = (feasible, infeasible), then a1 < a2 is true
    # if (a1, a2) = (infeasible, feasible), then a2 < a1 is false
    return a1.v
end

#### random initialization, @code_warntype ✓
function born!(x::VecIO, lb::NTuple{ND}, ub::NTuple{ND}) where ND
    @simd for i in eachindex(x)
        @inbounds x[i] = lb[i] + rand() * (ub[i] - lb[i])
    end
end

#### groups, @code_warntype ✓
function return_agents(nd::Int, np::Int)
    agents = Vector{Agent}(undef, np)
    @inbounds for i in eachindex(agents)
        agents[i] = Agent(nd)
    end
    return agents
end

#### subgroups, @code_warntype ✓
return_elites(agents::VecI{Agent}, ne::Int)          = view(agents, 1:ne)
return_throng(agents::VecI{Agent}, ne::Int, np::Int) = view(agents, ne+1:np)

#=
    Sine-Cosine Optimizer (https://doi.org/10.1016/j.knosys.2015.12.022)
    params:
    -------
    * Xb := buffer
    * Xn := n-th solution in the pool
    * Xr := referred solution
    * ss := step size
=#
function sco_move!(Xb::VecB, Xn::VecI, Xr::VecI, ss::Real) # @code_warntype ✓
    r = 2.0 * rand()
    s = sinpi(r)
    c = cospi(r)
    @simd for i in eachindex(Xb)
        @inbounds Xb[i] = Xn[i] + ss * abs(Xn[i] - Xr[i]) * ifelse(rand() < 0.5, s, c)
    end
end
#=
    Water-Cycle Algorithm Optimizer (https://doi.org/10.1016/j.compstruc.2012.07.010)
    params:
    -------
    * Xb    := buffer
    * Xbest := the best solution currently
=#
function wca_move!(Xb::VecB, Xbest::VecI) # @code_warntype ✓
    scaled_rand = randn() * 0.31622776601683794 # sqrt(0.1)
    @simd for i in eachindex(Xb)
        @inbounds Xb[i] = Xbest[i] + scaled_rand
    end
end

############################
# Dealing with constraints #
############################

fcall(f::Function, x::VecI) = f(x)

struct BoxBound a::Float64; b::Float64; i::Int end # a * x[i] + b

resolve_lb(lb::Real) = iszero(lb) ? (-1.0, 0.0) : (-abs(inv(lb)),  1.0 * sign(lb)) # @code_warntype ✓
resolve_ub(ub::Real) = iszero(ub) ? ( 1.0, 0.0) : ( abs(inv(ub)), -1.0 * sign(ub)) # @code_warntype ✓

function boxbounds(lb::NTuple{ND}, ub::NTuple{ND}) where ND
    if @generated
        a = Vector{Expr}(undef, 2*ND)
        @inbounds for i in 1:ND
            a[i]    = :(BoxBound(resolve_lb(lb[$i])..., $i))
            a[i+ND] = :(BoxBound(resolve_ub(ub[$i])..., $i))
        end
        return quote
            $(Expr(:meta, :inline))
            @inbounds return $(Expr(:tuple, a...))
        end
    else
        return ntuple(i -> i > ND ? BoxBound(resolve_ub(ub[i - ND])..., i - ND) : BoxBound(resolve_lb(lb[i])..., i), 2*ND)
    end
end

eval_violation(x::VecI, bb::BoxBound) = max(0.0, bb.a * x[bb.i] + bb.b) # @code_warntype ✓
function eval_violation(x::VecI, cons::NTuple{NB,BoxBound}) where NB    # @code_warntype ✓
    if @generated
        a = Vector{Expr}(undef, NB)
        @inbounds for i in eachindex(a)
            a[i] = :(eval_violation(x, cons[$i]))
        end
        return quote
            $(Expr(:meta, :inline))
            @inbounds return $(Expr(:call, :+, a...))
        end
    else
        ret = 0.0
        @inbounds for i in eachindex(cons)
            ret += eval_violation(x, cons[i])
        end
        return ret
    end
end

# check feasibility of agent in throng, @code_warntype ✓
function check!(xnew::VecI, agents::VecIO{Agent}, elites::VecIO{Agent}, throng::VecIO{Agent}, edx::Int, tdx::Int, fn::Function, cons::NTuple)
    violation = eval_violation(xnew, cons)
    violation > 0.0 && return @inbounds check!(xnew, violation, throng[tdx]) # x[new] is infeasible
    return check!(xnew, fcall(fn, xnew), agents, elites, throng, edx, tdx)   # x[new] is feasible
end

# Matchup for a feasible x[new] agent in throng, @code_warntype ✓
function check!(xnew::VecI, fnew::Real, agents::VecIO{Agent}, elites::VecIO{Agent}, throng::VecIO{Agent}, edx::Int, tdx::Int)
    @inbounds xold = throng[tdx]
    # x[old] is infeasible
    if !xold.v
        xold.f = fnew
        xold.v = true
        xold.c = 0.0
        copy!(xold.x, xnew)
        return nothing
    end
    # x[old], x[new] are feasible, x[new] is better than/equals to x[best] (greedy strategy)
    if @inbounds !(elites[edx].f < fnew)
        xold.f = fnew
        copy!(xold.x, xnew)
        swap!(agents, edx, length(elites) + tdx)
        return nothing
    end
    # x[old], x[new] are feasible
    if !(xold.f < fnew)
        xold.f = fnew
        copy!(xold.x, xnew)
        return nothing
    end
end

# check feasibility of agent in elites, @code_warntype ✓
function check!(xnew::VecI, elites::VecIO{Agent}, edx::Int, fn::Function, cons::NTuple)
    violation = eval_violation(xnew, cons)
    violation > 0.0 && return @inbounds check!(xnew, violation, elites[edx]) # x[new] is infeasible
    return check!(xnew, fcall(fn, xnew), elites, edx)                        # x[new] is feasible
end

# Matchup for a feasible x[new] trial in elites, @code_warntype ✓
function check!(xnew::VecI, fnew::Real, elites::VecIO{Agent}, edx::Int)
    @inbounds elite = elites[edx]
    # x[old] is infeasible
    if !elite.v
        elite.f = fnew
        elite.v = true
        elite.c = 0.0
        copy!(elite.x, xnew)
        return nothing
    end
    # x[old], x[new] are feasible, x[new] is better than/equals to x[best] (greedy strategy)
    if @inbounds !(elites[1].f < fnew)
        elite.f = fnew
        copy!(elite.x, xnew)
        swap!(elites, 1, edx)
        return nothing
    end
    # x[old], x[new] are feasible
    if !(elite.f < fnew)
        elite.f = fnew
        copy!(elite.x, xnew)
        return nothing
    end
end

# Matchup for an infeasible x[new] trial, here "fnew = violation", @code_warntype ✓
function check!(xnew::VecI, violation::Real, agent::Agent)
    # x[old], x[new] are infeasible, compare violation
    # There is no `else` condition, if x[old] is feasible, then a matchup is unnecessary.
    if !agent.v && !(agent.c < violation)
        agent.c = violation
        copy!(agent.x, xnew)
    end
    return nothing
end

############################
# Minimization             #
############################

## Generic logistic function
logistic(x::Real, x0::Real, a::Real, k::Real, c::Real) = a / (1.0 + exp(k * (x0 - x))) + c

## Perform a single step of Wolford algorithm
function welford_step(μ::Real, s::Real, v::Real, c::Real)
    isone(c) && return v, zero(v)
    s = s * (c - 1)
    m = μ + (v - μ) / c
    s = s + (v - μ) * (v - m)
    μ = m
    return μ, s / (c - 1)
end

function nrm2(x::VecI{Tx}, y::VecI{Ty}, b::VecB{Tb}) where {Tx<:Real,Ty<:Real,Tb<:Real} # @code_warntype ✓
    @simd for i in eachindex(b)
        @inbounds b[i] = abs2(x[i] - y[i])
    end
    return sqrt(sum(b))
end

struct GlobalMinimizer
    xsol::Vector{Float64}
    xerr::Vector{Float64}
    buff::Vector{Float64}
    fork::Vector{Int64}
    pool::Vector{Agent}
    np::Int
    ne::Int

    function GlobalMinimizer(nd::Int, np::Int, ne::Int)
        xsol = Vector{Float64}(undef, nd)
        xerr = Vector{Float64}(undef, nd)
        buff = Vector{Float64}(undef, nd)
        fork = Vector{Int}(undef, ne)
        pool = return_agents(nd, np)
        return new(xsol, xerr, buff, fork, pool, np, ne)
    end
end

# @code_warntype ✓
function inits!(agents::VecIO{Agent}, lb::NTuple, ub::NTuple)
    for agent in agents
        born!(agent.x, lb, ub)
    end
end

# @code_warntype ✓
function inits!(agents::VecIO{Agent}, f::Function, cons::NTuple)
    fmax = -Inf
    for agent in agents
        violation = eval_violation(agent.x, cons)
        violation > 0.0 && (agent.c = violation; continue) # agent is infeasible
        agent.v = true
        agent.f = fcall(f, agent.x)
        fmax    = max(fmax, agent.f)
    end
    for agent in agents
        !agent.v && (agent.f = agent.c + fmax) # agent is infeasible
    end
end

# @code_warntype ✓
function group!(fork::VecIO{Int}, agents::VecI{Agent}, ne::Int, nc::Int)
    diversity = 0.0
    @inbounds for i in eachindex(fork)
        diversity += agents[ne + 1].f - agents[i].f
    end
    if iszero(diversity) || isnan(diversity)
        fill!(fork, 1)
    else
        @inbounds for i in eachindex(fork)
            fork[i] = max(1, round(Int, nc * (agents[ne + 1].f - agents[i].f) / diversity))
        end
    end
    res = nc - sum(fork) # residue
    idx = 2
    while res > 0
        @inbounds fork[idx] += 1; res -= 1
        idx < ne ? idx += 1 : idx = 2
    end
    while res < 0
        @inbounds fork[idx] = max(1, fork[idx] - 1); res += 1
        idx < ne ? idx += 1 : idx = 2
    end
end

# @code_warntype ✓
function minimize!(o::GlobalMinimizer, fn::Function, lb::NTuple{ND,T}, ub::NTuple{ND,T}, itmax::Int, dmax::Real, avgtimes::Int) where {ND,T<:Real}
    np = o.np
    ne = o.ne
    nc = np - ne

    cons = boxbounds(lb, ub)
    xsol = o.xsol
    xerr = o.xerr
    buff = o.buff
    fork = o.fork

    agents = o.pool
    elites = return_elites(agents, ne)
    throng = return_throng(agents, ne, np)

    generation = 0
    while generation < avgtimes
        generation += 1
        itcount     = 0

        inits!(agents, lb, ub)
        inits!(agents, fn, cons)
        binsort!(agents)

        @inbounds while itcount < itmax
            itcount += 1
            ss = logistic(itcount, 0.5 * itmax, -0.618, 20.0 / itmax, 2.0)
            group!(fork, agents, ne, nc)

            #### Moves: throng → elites, elites → the-best
            rx = 1
            fx = fork[rx]
            # move agents (in throng) → elites
            for ix in eachindex(throng)
                sco_move!(buff, elites[rx].x, throng[ix].x, ss)
                check!(buff, agents, elites, throng, rx, ix, fn, cons)
                fx -= 1
                iszero(fx) && rx < ne && (rx += 1; fx = fork[rx])
            end
            # move agents (in elites) and find the best one
            for rx in 2:ne
                sco_move!(buff, elites[1].x, elites[rx].x, ss)
                check!(buff, elites, rx, fn, cons)
            end

            #### Random searching process
            for ix in 1:fork[1]
                if !(dmax < nrm2(agents[1].x, throng[ix].x, buff))
                    wca_move!(buff, agents[1].x)
                    check!(buff, agents, elites, throng, 1, ix, fn, cons)
                end
            end
            for rx in 2:ne
                if !(dmax < nrm2(agents[1].x, elites[rx].x, buff)) || !(0.1 < rand())
                    born!(buff, lb, ub)
                    check!(buff, elites, rx, fn, cons)
                end
            end

            #### Update the function-value of infeasible candidates
            fmax = -Inf
            for agent in agents
                agent.v && (fmax = max(fmax, agent.f))
            end
            for agent in agents
                !agent.v && (agent.f = agent.c + fmax)
            end

            binsort!(agents)
            dmax -= dmax / itmax
        end

        @inbounds xnew = agents[1].x
        @inbounds for i in eachindex(xsol)
            xsol[i], xerr[i] = welford_step(xsol[i], xerr[i], xnew[i], generation)
        end
    end
end

############################
# Interface                #
############################

"""
    minimizer(nd::Int; np::Int=35*nd, ne::Int=nd+1)
An interface function to create/initialize an object for the minimization.
Arguments:
---
- `nd`:     Dimension of parameters to be minimized.
- `np`:     Desired population size (*optional*).
- `ne`:     Desired size of elites (*optional*).
"""
minimizer(nd::Int; np::Int=35*nd, ne::Int=nd+1) = GlobalMinimizer(nd, np, ne)

"""
    minimize!(o::GlobalMinimizer, fn::Function, lb::NTuple{ND}, ub::NTuple{ND}; itmax::Int=210*ND, dmax::Real=1e-7, avgtimes::Int=1)
An interface function to proceed the "global" minimization.
Arguments:
---
- `o`:        An object for the minimization created by `minimizer(...; args...)`.
- `fn`:       Objective function to be minimized.
              `fn` should be callable with only one argument of `fn(x::Vector)`.
              If you have any additional arguments need to pass into it,
              dispatch the function by `fcall(fn, x; kwargs...) = fn(x; kwargs...)`
- `lb`:       Lower bounds of minimization which are forced to be feasible.
- `ub`:       Upper bounds of minimization which are forced to be feasible.
- `itmax`:    Maximum of minimizing iteration (*optional*).
- `dmax`:     An Euclidean distance acts as a criterion to
              prevent the population falling into local minimal (*optional*).
- `avgtimes`: Number of average times of the whole minimization process (*optional*).
"""
minimize!(o::GlobalMinimizer, fn::Function, lb::NTuple{ND}, ub::NTuple{ND}; itmax::Int=210*ND, dmax::Real=1e-7, avgtimes::Int=1) where ND = minimize!(o, fn, lb, ub, itmax, dmax, avgtimes)

end # module


text:
@define_scenario simple_scenario = begin
    t_xi::Int64
    s_xi::Int64
    @zero begin
        return simple_scenario(0, 0)
    end
    @expectation begin
        t_xi = Int(floor(sum([probability(s)*s.t_xi for s in scenarios])))
        s_xi = Int(floor(sum([probability(s)*s.s_xi for s in scenarios])))
        return simple_scenario(t_xi, s_xi)
    end
end

@sampler simple_sampler = begin
    timesteps::UnitRange{Int64}
    #n::Int64
    simple_sampler(timesteps::UnitRange{Int64}) = new(timesteps)

    @sample simple_scenario begin
        @parameters timesteps
        return simple_scenario(rand(timesteps[1:end-2]), rand([-1, 1]))
    end
end
## Try adding F to the scenario
@define_scenario F_scenario = begin
    t_xi::Int64
    s_xi::Int64
    F_xi::Float64
    @zero begin
        return F_scenario(0, 0, 0.)
    end
    @expectation begin
        t_xi = Int(floor(sum([probability(s)*s.t_xi for s in scenarios])))
        s_xi = Int(floor(sum([probability(s)*s.s_xi for s in scenarios])))
        F_xi = mean([probability(s)*s.F_xi for s in scenarios])
        return F_scenario(t_xi, s_xi, F_xi)
    end
end

@sampler F_sampler = begin
    timesteps::UnitRange{Int64}
    F_range::Tuple{Float64, Float64}
    #n::Int64
    F_sampler(timesteps::UnitRange{Int64}, F_range::Tuple{Float64, Float64}) = new(timesteps, F_range)

    @sample F_scenario begin
        @parameters timesteps F_range
        return F_scenario(rand(timesteps[1:end-2]), rand([-1, 1]), F_range[1]+rand()*F_range[2])
    end
end

ext: jl


text:
include("constants.jl")

"""
Performs bracketing to determine the initial bounds on a minimum.

Arguments:
    fn: objective function
    δ: step size
    limit: maximum number of iterations

Returns:
    3-element array if successful, else nothing
"""
function bracket(fn, δ, limit=100)
    α = 0
    q = 0

    α_values = []
    for i=1:3
        α += G^q*δ
        push!(α_values, α)
        q += 1
    end
    f_0, f_1, f_2 = [fn(α_values[i]) for i in 1:3] # q-2, q-1, q

    for _ = 1:limit
        if f_1 < f_0 && f_2 > f_1
            return (α_values[begin], α_values[2], α_values[end])
        end

        α += G^q*δ
        push!(α_values, α)
        α_values = α_values[2:end]
        q += 1

        f_0 = f_1
        f_1 = f_2
        f_2 = fn(α[end])
    end
end

max_stars_repo_name: camirmas/DesignAutomation

max_issues_repo_name: camirmas/DesignAutomation

max_forks_repo_name: camirmas/DesignAutomation


text:
"""
The TinBasic type contains the state variables (swe), the inputs
(prec, tair) for one time step, the parameters (param) and the time step
length (tstep) for a basic temperature index model.
"""
type TinBasic <: Snow

  swe::Array{Float64,1}
  prec::Array{Float64,1}
  tair::Array{Float64,1}
  q_sim::Array{Float64,1}
  param::Array{Float64,1}
  frac::Array{Float64,1}
  tstep::Float64
  time::DateTime

end

"""
    TinBasic(tstep, time, frac)

Constructor for TinBasic with predefined state variables, parameters and inputs.
The time step (tstep) is given as a fraction of one day. Thus, for hourly input
data tstep should be set to 1/24. The fraction of elevation bands should sum
up to unity.
"""
function TinBasic(tstep::Float64, time::DateTime, frac::Array{Float64,1})

  nreg = length(frac)
  prec   = zeros(Float64, nreg)
  tair   = zeros(Float64, nreg)
  swe    = zeros(Float64, nreg)
  q_sim  = zeros(Float64, nreg)
  param  = [0.0, 3.0, 1.0]

  TinBasic(swe, prec, tair, q_sim, param, frac, tstep, time)

end

"""
    TinBasic(tstep, time, param, frac)

Constructor for TinBasic with predefined state variables and inputs.
The time step (tstep) is given as a fraction of one day. Thus, for hourly input
data tstep should be set to 1/24. The fraction of elevation bands should sum
up to unity.
"""
function TinBasic(tstep::Float64, time::DateTime, param::Array{Float64,1}, frac::Array{Float64,1})

  nreg = length(frac)
  prec   = zeros(Float64, nreg)
  tair   = zeros(Float64, nreg)
  swe    = zeros(Float64, nreg)
  q_sim  = zeros(Float64, nreg)

  TinBasic(swe, prec, tair, q_sim, param, frac, tstep, time)

end

"""
    init_states(mdata::TinBasic)

Initilize the state variables of the model.
"""
function init_states(mdata::TinBasic)

  for ireg in eachindex(mdata.swe)
    mdata.swe[ireg] = 0.0
  end

end

"""
    get_param_range(mdata::TinBasic)

Get allowed parameter ranges for the calibration of the model.
"""
function get_param_range(mdata::TinBasic)

  param_range_snow = [(-3.0, 3.0), (0.1, 10.0), (0.5, 2.0)]

end

"""
    get_states(mdata::TinBasic)

Get state variables for computing penelty during calibration.

"""
function get_states(mdata::TinBasic)

  return [mean(mdata.swe)]

end

"""
    assign_param(mdata::TinBasic, param::Array{Float64,1})

Assign parameter values to the TinBasic type.
"""
function assign_param(mdata::TinBasic, param::Array{Float64,1})

  for ireg in eachindex(mdata.param)
    mdata.param[ireg] = param[ireg]
  end

end

"""
    enkf_snow(mdata::Array{TinBasic}, obs_ens, q_sim)

Implementation of the ensemble Kalman filter for the basic temperature index
snow model.
"""
function enkf_snow(mdata::Array{TinBasic}, obs_ens, q_sim)

  nens = length(mdata)

  # Allocate arrays

  swe = zeros(Float64, length(mdata[1].swe), nens)

  # Add states to arrays

  for iens = 1:nens

    swe[:, iens]    = mdata[iens].swe

  end

  # Run ensemble kalman filter

  swe = enkf(swe, obs_ens, q_sim)

  # Check limits of states

  swe[swe .< 0] = 0.

  # Add arrays to states

  for iens = 1:nens

    mdata[iens].swe = swe[:, iens]

  end

  nothing

end

"""
    run_timestep(mdata::TinBasic)

Propagate the model one time step and compute simulated snowpack dischage.
"""
function run_timestep(mdata::TinBasic)

  # Parameters

  tth   = mdata.param[1]
  ddf   = mdata.param[2]
  pcorr = mdata.param[3]

  for ireg in eachindex(mdata.swe)

    # Compute solid and liquid precipitation

    psolid, pliquid = split_prec(mdata.prec[ireg], mdata.tair[ireg], tth)

    psolid  = psolid * pcorr

    # Compute snow melt

    M = pot_melt(mdata.tair[ireg], ddf, tth)

    M = min(mdata.swe[ireg],M)

    # Update snow water equivalents

    mdata.swe[ireg] += psolid
    mdata.swe[ireg] -= M

    # Compute snowpack runoff

    mdata.q_sim[ireg] = M + pliquid

  end

  mdata.time += Dates.Hour(mdata.tstep)

  return nothing

end

max_stars_repo_name: jmgnve/Vann

max_forks_repo_name: jmgnve/Vann


text:
module QuasiNewtonMethods

using StrideArrays, LoopVectorization, VectorizationBase

using VectorizationBase: gep
using StrideArrays: PtrArray, StrideArray

function logdensity end
function ∂logdensity! end

export optimize!

abstract type AbstractProbabilityModel{D} end# <: LogDensityProblems.AbstractLogDensityProblem end
dimension(::AbstractProbabilityModel{D}) where {D} = StaticInt{D}()
Base.length(::AbstractProbabilityModel{D}) where {D} = D
function Base.show(io::IO, ℓ::AbstractProbabilityModel{D}) where {D}
    print(io, "$D-dimensional Probability Model")
end

function update_state!(s, x_old, α)
    @avx for i ∈ eachindex(s)
        sᵢ = α * s[i]
        s[i] = sᵢ
        x_old[i] += sᵢ
    end
end

"""

Updates B⁻¹ and sets yₖ as the new search direction.
"""
function BFGS_update!(B⁻¹::AbstractMatrix{T}, B⁻¹yₖ, yₖ, sₖ, ∇_new, ∇_old) where {T}
    sₖᵀyₖ = zero(T)
    @avx for i ∈ eachindex(sₖ)
        # yₖᵢ = ∇_new[i] - ∇_old[i]
        yₖᵢ = ∇_old[i] - ∇_new[i]
        sₖᵀyₖ += sₖ[i] * yₖᵢ
        yₖ[i] = yₖᵢ
    end
    sₖᵀyₖ⁻¹ = inv(sₖᵀyₖ)
    yₖᵀB⁻¹yₖ = zero(T)
    @avx for c ∈ axes(B⁻¹,2)
        t = zero(yₖᵀB⁻¹yₖ)
        for r ∈ axes(B⁻¹,1)
            # t += yₖ[r] * B⁻¹[r,c]
            t += yₖ[r] * B⁻¹[c,r]
        end
        B⁻¹yₖ[c] = t * sₖᵀyₖ⁻¹
        yₖᵀB⁻¹yₖ += t * yₖ[c]
    end
    # c₁ = (one(T) + yₖᵀB⁻¹yₖ * sₖᵀyₖ⁻¹) * sₖᵀyₖ⁻¹
    c₁ = muladd(yₖᵀB⁻¹yₖ, sₖᵀyₖ⁻¹, one(T)) * sₖᵀyₖ⁻¹
    m = zero(T)
    @avx for c ∈ axes(B⁻¹,2)
        sₙ = zero(eltype(yₖ))
        for r ∈ axes(B⁻¹,1)
            # B⁻¹ᵣ = B⁻¹[r,c] + c₁ * sₖ[r] * sₖ[c] - B⁻¹yₖ[r]*sₖ[c] - B⁻¹yₖ[c]*sₖ[r]
            # B⁻¹[r,c] = B⁻¹ᵣ
            B⁻¹ᵣ = B⁻¹[c,r] + c₁ * sₖ[r] * sₖ[c] - B⁻¹yₖ[r]*sₖ[c] - B⁻¹yₖ[c]*sₖ[r]
            B⁻¹[c,r] = B⁻¹ᵣ
            sₙ += B⁻¹ᵣ * ∇_new[r]
        end
        yₖ[c] = sₙ
        m += sₙ * ∇_new[c]
    end
    m
end

struct BackTracking{O}
    c₁::Float64
    ρₕ::Float64
    ρₗ::Float64
    iterations::Int
end
# BackTracking{O}(c₁ = 1e-4, ρₕ = 0.5, ρₗ = 0.1, iterations = 1_000) where {O} = BackTracking{O}(c₁, ρₕ, ρₗ, iterations)
BackTracking{O}(c₁ = 1e-4, ρₕ = 0.5, ρₗ = 0.1) where {O} = BackTracking{O}(c₁, ρₕ, ρₗ, 1_000)
BackTracking(c₁ = 1e-4, ρₕ = 0.5, ρₗ = 0.1, iterations = 1_000) = BackTracking{2}(c₁, ρₕ, ρₗ, iterations)

abstract type AbstractBFGSState{P,T,L,LT} end

mutable struct BFGSState{P,T,L,LT} <: AbstractBFGSState{P,T,L,LT}
    x_old::NTuple{L,T}
    ∇_new::NTuple{L,T}
    x_new::NTuple{L,T}
    ∇_old::NTuple{L,T}
    y::NTuple{L,T}
    s::NTuple{L,T}
    B⁻¹y::NTuple{L,T}
    B⁻¹::NTuple{LT,T}
    function BFGSState{P,T,L,LT}(::UndefInitializer) where {P,T,L,LT}
        new{P,T,L,LT}()
    end
    @generated function BFGSState{P}(::UndefInitializer) where {P}
        L =  VectorizationBase.align(P,VectorizationBase.pick_vector_width(Float64))
        :(BFGSState{$P,Float64,$L,$(P*L)}(undef))
    end
    @generated function BFGSState{P,T}(::UndefInitializer) where {P,T}
        L =  VectorizationBase.align(P,VectorizationBase.pick_vector_width(T))
        :(BFGSState{$P,$T,$L,$(P*L)}(undef))
    end
end
BFGSState(::Val{P}, ::Type{T} = Float64) where {P,T} = BFGSState{P,T}(undef)
@inline Base.pointer(s::BFGSState{P,T})  where {P,T} = Base.unsafe_convert(Ptr{T}, pointer_from_objref(s))

function VectorizationBase.lifetime_start!(state::AbstractBFGSState{P,T,L}) where {P,T,L}
    nothing
    # VectorizationBase.lifetime_start!(pointer(ref_x_new(state)), Val(VectorizationBase.staticmul(T, StaticInt{L}()*(StaticInt{5}()+StaticInt{L}()))))
end
function VectorizationBase.lifetime_end!(state::AbstractBFGSState{P,T,L}) where {P,T,L}
    nothing
    # VectorizationBase.lifetime_end!(pointer(ref_x_new(state)), Val(VectorizationBase.staticmul(T, StaticInt{L}()*(StaticInt{5}()+StaticInt{L}()))))
end

struct PtrBFGSState{P,T,L,LT,NI<:Union{Int,Nothing}} <: AbstractBFGSState{P,T,L,LT}
    ptr::Ptr{T}
    offset::NI
end
@inline Base.pointer(state::PtrBFGSState) = state.ptr

@inline function dense_vector(ptr::Ptr{T}, ::StaticInt{P}) where {T,P}
    PtrArray(ptr, (StaticInt{P}(),), (VectorizationBase.static_sizeof(T),), Val{(true,)}())
end
@inline function square_matrix(ptr::Ptr{T}, ::StaticInt{P}, ::StaticInt{X}) where {T,P,X}
    st = VectorizationBase.static_sizeof(T)
    PtrArray(ptr, (StaticInt{P}(),StaticInt{P}()), (st,st*StaticInt{X}()), Val{(true,false)}())
end
@inline function square_matrix(ptr::Ptr{T}, ::StaticInt{P}, ::StaticInt{P}) where {T,P}
    st = VectorizationBase.static_sizeof(T)
    PtrArray(ptr, (StaticInt{P}(),StaticInt{P}()), (st,st*StaticInt{P}()), Val{(true,true)}())
end

@inline ref_x_old(s::AbstractBFGSState{P,T,L,LT}) where {P,T,L,LT} = dense_vector(pointer(s), StaticInt{P}())
@inline ref_∇_old(s::AbstractBFGSState{P,T,L,LT}) where {P,T,L,LT} = dense_vector(gep(pointer(s), fieldoffset(BFGSState{P,T,L,LT}, 4)), StaticInt{P}())
@inline ref_x_new(s::AbstractBFGSState{P,T,L,LT}) where {P,T,L,LT} = dense_vector(gep(pointer(s), fieldoffset(BFGSState{P,T,L,LT}, 3)), StaticInt{P}())
@inline ref_∇_new(s::AbstractBFGSState{P,T,L,LT}) where {P,T,L,LT} = dense_vector(gep(pointer(s), fieldoffset(BFGSState{P,T,L,LT}, 2)), StaticInt{P}())
@inline ref_y(s::AbstractBFGSState{P,T,L,LT}) where {P,T,L,LT} = dense_vector(gep(pointer(s), fieldoffset(BFGSState{P,T,L,LT}, 5)), StaticInt{P}())
@inline ref_s(s::AbstractBFGSState{P,T,L,LT}) where {P,T,L,LT} = dense_vector(gep(pointer(s), fieldoffset(BFGSState{P,T,L,LT}, 6)), StaticInt{P}())
@inline ref_B⁻¹y(s::AbstractBFGSState{P,T,L,LT}) where {P,T,L,LT} = dense_vector(gep(pointer(s), fieldoffset(BFGSState{P,T,L,LT}, 7)), StaticInt{P}())
@inline ref_B⁻¹(s::AbstractBFGSState{P,T,L,LT}) where {P,T,L,LT} = square_matrix(gep(pointer(s), fieldoffset(BFGSState{P,T,L,LT}, 8)), StaticInt{P}(), StaticInt{L}())

function initial_B⁻¹!(B⁻¹::AbstractMatrix)
    @avx for c ∈ axes(B⁻¹,2), r ∈ axes(B⁻¹,1)
        B⁻¹[r,c] = (r == c)
    end
end
@inline optimum(s::AbstractBFGSState) = ref_x_old(s)
@inline gradient(s::AbstractBFGSState) = ref_∇_new(s)

nanmin(a, b) = a < b ? a : (isnan(b) ? a : b)
nanmax(a, b) = a < b ? b : (isnan(a) ? b : a)

sqrttolerance(::Type{T}) where {T} = T(1 / (1 << (Base.Math.significand_bits(T) >> 1)))

function step!(x_new, x_old, s, obj, α)
    @avx for i ∈ eachindex(x_new)
        x_new[i] = x_old[i] + α*s[i]
    end
    logdensity(obj, x_new)
end

function linesearch!(x_new::AbstractVector{T}, x_old, s, obj, ℓ₀, m, ls::BackTracking{order}) where {T, order}
    c₁, ρₕ, ρₗ, iterations = T(ls.c₁), T(ls.ρₕ), T(ls.ρₗ), ls.iterations
    sqrttol = sqrttolerance(T)

    α₀ = one(T)

    # Count the total number of iterations
    ℓx₀, ℓx₁ = ℓ₀, ℓ₀
    α₁, α₂ = α₀, α₀
    ℓx₁ = step!(x_new, x_old, s, obj, α₁)

    # Hard-coded backtrack until we find a finite function value
    # Halve α₂ until function value is finite
    iterfinite = 0
    iterfinitemax = Base.Math.significand_bits(T)
    while !isfinite(ℓx₁) && iterfinite < iterfinitemax
        iterfinite += 1
        α₁, α₂ = (α₂, Base.FastMath.mul_fast(T(0.5), α₂))
        ℓx₁ = step!(x_new, x_old, s, obj, α₂)
    end
    iteration = 0
    # Backtrack until we satisfy sufficient decrease condition
    while !(ℓx₁ ≥ @fastmath(ℓ₀ + α₂*c₁*m))
        # @show α₂, ℓx₁, ℓ₀ + α₂*c₁*m, ℓ₀, α₂, c₁, m
        # Increment the number of steps we've had to perform
        iteration += 1

        # Ensure termination
        iteration > iterations && return zero(T) # linesearch_failure(iterations)

        # Shrink proposed step-size:
        if order == 2 || iteration == 1
            # backtracking via quadratic interpolation:
            # This interpolates the available data
            #    f(0), f'(0), f(α)
            # with a quadractic which is then minimised; this comes with a
            # guaranteed backtracking factor 0.5 * (1-c_1)^{-1} which is < 1
            # provided that c_1 < 1/2; the backtrack_condition at the beginning
            # of the function guarantees at least a backtracking factor ρ.
            # https://github.com/JuliaNLSolvers/LineSearches.jl/blob/0c182ce25c32a385a6156304dc24ff72afbfa308/src/backtracking.jl#L84
            αₜ = @fastmath - (m * α₂*α₂) / ( T(2) * (ℓx₁ - ℓ₀ - m*α₂) )
        else
            denom = @fastmath one(T) / (α₁*α₁ * α₂*α₂ * (α₂ - α₁))
            a = @fastmath (α₁*α₁*(ℓx₁ - ℓ₀ - m*α₂) - α₂*α₂*(ℓx₀ - ℓ₀ - m*α₁))*denom
            b = @fastmath (-α₁*α₁*α₁*(ℓx₁ - ℓ₀ - m*α₂) + α₂*α₂*α₂*(ℓx₀ - ℓ₀ - m*α₁))*denom

            if abs(a) <= eps(T) + sqrttol*abs(a)
                αₜ = @fastmath m / (T(2)*b)
            else
                # discriminant
                d = nanmax(@fastmath(b*b - T(3)*a*m), zero(T))
                # quadratic equation root
                # αₜ = @fastmath (sqrt(d) - b) / (3a)
                αₜ = @fastmath (sqrt(d) + b) / (-3a)
            end
        end
        # @show α₂, ℓx₁, ℓ₀ + α₂*c₁*m, ℓ₀, α₂, c₁, m
        α₁ = α₂

        αₜ = nanmin(αₜ, α₂*ρₕ) # avoid too small reductions
        α₂ = nanmax(αₜ, α₂*ρₗ) # avoid too big reductions

        # Evaluate f(x) at proposed position
        ℓx₀, ℓx₁ = ℓx₁, step!(x_new, x_old, s, obj, α₂)
        # @show α₂, ℓx₁, ℓ₀ + α₂*c₁*m, ℓ₀, α₂, c₁, m
    end
    α₂
end

"""
Optimum value is stored in state.x_old.
"""
function optimize!(state, obj, x::AbstractStrideArray{Tuple{StaticInt{P}},D,T}, ls = BackTracking(), tol = 1e-8) where {P,T,D}
    VectorizationBase.lifetime_start!(state)
    x_old = ref_x_old(state)
    ∇_old = ref_∇_old(state)
    x_new = ref_x_new(state)
    ∇_new = ref_∇_new(state)
    y = ref_y(state)
    s = ref_s(state)
    B⁻¹y = ref_B⁻¹y(state)
    B⁻¹ = ref_B⁻¹(state)
    copyto!(x_old, x)
    # initial_invH!(invH)
    N = 1000
    # f_calls = 0
    # g_calls = 0
    for n ∈ 1:N
        ℓ₀ = ∂logdensity!(∇_new, obj, x_old)#; f_calls +=1; g_calls +=1;
        Base.isfinite(ℓ₀) || break
        # @show (n-1), ℓ₀, maximum(abs, ∇_new), tol ∇_new
        if maximum(abs, ∇_new) < tol
            ∇_new_original = ref_∇_new(state)
            pointer(∇_new) == pointer(∇_new_original) || copyto!(∇_new_original, ∇_new)
            VectorizationBase.lifetime_end!(state)
            return ℓ₀
        end
        if isone(n)
            m = -one(T) # reset search direction
        else
            # @assert false
            m = BFGS_update!(B⁻¹, B⁻¹y, y, s, ∇_new, ∇_old)
            s, y = y, s
            # @show B⁻¹ x_old
            # @show pointer(s) - pointer(y)
        end
        if m <= zero(T) # If bad, reset search direction
            initial_B⁻¹!(B⁻¹)
            m = zero(T)
            @avx for i ∈ eachindex(x)
                sᵢ = ∇_new[i]
                s[i] = sᵢ
                m += sᵢ * sᵢ
            end
        end
        # @show s
        #### Perform line search
        α₂ = linesearch!(x_new, x_old, s, obj, ℓ₀, m, ls)
        iszero(α₂) && break
        update_state!(s, x_old, α₂)
        ∇_old, ∇_new = ∇_new, ∇_old
        # isone(n) && (s = ref_s(state))
    end
    VectorizationBase.lifetime_end!(state)
    T(NaN)
end

# include("precompile.jl")
# _precompile_()

end # module


text:
module Evaluation

include("rouge.jl")
export evaluate, mean_rouge_n, mean_rouge_l

end # module


text:
pg = randbool(SLEN)
pgtail = copy(pg)
while pgtail != mach && pgtail != human
    push!(pg, randbool())
    pgtail = [pgtail[2:end], pg[end]]
end

println(@sprintf("This game lasted %d turns yielding\n    %s",
                 length(pg), pgdecode(pg)))

if human == mach
    println("so you and the computer tied.")
elseif pgtail == mach
    println("so the computer won.")
else
    println("so you won.")
end

max_stars_repo_name: ethansaxenian/RosettaDecode

max_issues_repo_name: ethansaxenian/RosettaDecode

max_forks_repo_name: ethansaxenian/RosettaDecode


text:
#> This code implements the MLJ model interface for models in the
#> DecisionTree.jl package. It is annotated so that it may serve as a
#> template for other regressors of Deterministic type and classifiers
#> of Probabilistic type. The annotations, which begin with "#>",
#> should be removed (but copy this file first!). See also the model
#> interface specification at "doc/adding_new_models.md".

#> Note that all models need to "register" their location by setting
#> `load_path(<:ModelType)` appropriately.

module DecisionTree_

#> export the new models you're going to define (and nothing else):
export DecisionTreeClassifier, DecisionTreeRegressor

import MLJBase

#> needed for classifiers:
using CategoricalArrays

#> import package:
import ..DecisionTree # strange syntax b/s we are lazy-loading

## CLASSIFIER

"""
    DecisionTreeClassifer(; kwargs...)

CART decision tree classifier from
[https://github.com/bensadeghi/DecisionTree.jl/blob/master/README.md](https://github.com/bensadeghi/DecisionTree.jl/blob/master/README.md). Predictions
are Probabilistic.

For post-fit pruning, set `post-prune=true` and set
`min_purity_threshold` appropriately. Other hyperparameters as per
package documentation cited above.

"""
mutable struct DecisionTreeClassifier <: MLJBase.Probabilistic
    pruning_purity::Float64
    max_depth::Int
    min_samples_leaf::Int
    min_samples_split::Int
    min_purity_increase::Float64
    n_subfeatures::Float64
    display_depth::Int
    post_prune::Bool
    merge_purity_threshold::Float64
end

# keywork constructor:
#> all arguments are kwargs with a default value
function DecisionTreeClassifier(
    ; pruning_purity=1.0
    , max_depth=-1
    , min_samples_leaf=1
    , min_samples_split=2
    , min_purity_increase=0.0
    , n_subfeatures=0
    , display_depth=5
    , post_prune=false
    , merge_purity_threshold=0.9)

    model = DecisionTreeClassifier(
        pruning_purity
        , max_depth
        , min_samples_leaf
        , min_samples_split
        , min_purity_increase
        , n_subfeatures
        , display_depth
        , post_prune
        , merge_purity_threshold)

    message = MLJBase.clean!(model)       #> future proof by including these
    isempty(message) || @warn message #> two lines even if no clean! defined below

    return model
end

#> The following optional method (the fallback does nothing, returns
#> empty warning) is called by the constructor above but also by the
#> fit methods below:
function MLJBase.clean!(model::DecisionTreeClassifier)
    warning = ""
    if  model.pruning_purity > 1
        warning *= "Need pruning_purity < 1. Resetting pruning_purity=1.0.\n"
        model.pruning_purity = 1.0
    end
    if model.min_samples_split < 2
        warning *= "Need min_samples_split >= 2. Resetting min_samples_slit=2.\n"
        model.min_samples_split = 2
    end
    return warning
end

#> A required `fit` method returns `fitresult, cache, report`. (Return
#> `cache=nothinobserg` unless you are overloading `update`)
function MLJBase.fit(model::DecisionTreeClassifier
             , verbosity::Int   #> must be here (and typed!!) even if not used (as here)
             , X
             , y)

    Xmatrix = MLJBase.matrix(X)

    yplain = identity.(y) # y as plain not abstact vector
    classes_seen = unique(yplain)

    tree = DecisionTree.build_tree(yplain,
                                   Xmatrix,
                                   model.n_subfeatures,
                                   model.max_depth,
                                   model.min_samples_leaf,
                                   model.min_samples_split,
                                   model.min_purity_increase)
    if model.post_prune
        tree = DecisionTree.prune_tree(tree, model.merge_purity_threshold)
    end

    verbosity < 3 || DecisionTree.print_tree(tree, model.display_depth)

    fitresult = (tree, classes_seen)

    #> return package-specific statistics (eg, feature rankings,
    #> internal estimates of generalization error) in `report`, which
    #> should be a named tuple with the same type every call (can have
    #> empty values):

    cache = nothing
    report = NamedTuple{}()

    return fitresult, cache, report

end

MLJBase.fitted_params(::DecisionTreeClassifier, fitresult) = fitresult[1]

function MLJBase.predict(::DecisionTreeClassifier
                     , fitresult
                     , Xnew)
    Xmatrix = MLJBase.matrix(Xnew)

    tree, classes_seen = fitresult

    y_probabilities = DecisionTree.apply_tree_proba(tree, Xmatrix, classes_seen)
    return [MLJBase.UnivariateFinite(classes_seen, y_probabilities[i,:])
            for i in 1:size(y_probabilities, 1)]
end

## REGRESSOR

"""
    DecisionTreeRegressor(; kwargs...)

CART decision tree classifier from
[https://github.com/bensadeghi/DecisionTree.jl/blob/master/README.md](https://github.com/bensadeghi/DecisionTree.jl/blob/master/README.md). Predictions
are Deterministic.

For post-fit pruning, set `post-prune=true` and set
`pruning_purity_threshold` appropriately. Other hyperparameters as per
package documentation cited above.

"""

mutable struct DecisionTreeRegressor <: MLJBase.Deterministic
    pruning_purity_threshold::Float64
    max_depth::Int
    min_samples_leaf::Int
    min_samples_split::Int
    min_purity_increase::Float64
    n_subfeatures::Int
    post_prune::Bool
end

# constructor:
#> all arguments are kwargs with a default value
function DecisionTreeRegressor(;
    pruning_purity_threshold=0.0
    , max_depth=-1
    , min_samples_leaf=5
    , min_samples_split=2
    , min_purity_increase=0.0
    , n_subfeatures=0
    , post_prune=false)

    model = DecisionTreeRegressor(
       pruning_purity_threshold
       , max_depth
       , min_samples_leaf
       , min_samples_split
       , min_purity_increase
       , n_subfeatures
       , post_prune)

    message = MLJBase.clean!(model)       #> future proof by including these
    isempty(message) || @warn message #> two lines even if no clean! defined below

    return model
end

function MLJBase.clean!(model::DecisionTreeRegressor)
    warning = ""
    if model.min_samples_split < 2
        warning *= "Need min_samples_split >= 2. Resetting min_samples_slit=2.\n"
        model.min_samples_split = 2
    end
    return warning
end

function MLJBase.fit(model::DecisionTreeRegressor
             , verbosity::Int   #> must be here (and typed) even if not used (as here)
             , X
             , y)

    Xmatrix = MLJBase.matrix(X)

    fitresult = DecisionTree.build_tree(float.(y)
				   , Xmatrix
				   , model.n_subfeatures
				   , model.max_depth
				   , model.min_samples_leaf
				   , model.min_samples_split
				   , model.min_purity_increase)

    if model.post_prune
        fitresult = DecisionTree.prune_tree(fitresult,
                                            model.pruning_purity_threshold)
    end
    cache = nothing
    report = nothing

    return fitresult, cache, report
end

MLJBase.fitted_params(::DecisionTreeRegressor, fitresult) = fitresult

function MLJBase.predict(model::DecisionTreeRegressor
                     , fitresult
                     , Xnew)
    Xmatrix = MLJBase.matrix(Xnew)
    return DecisionTree.apply_tree(fitresult,Xmatrix)
end

## METADATA

DTTypes=Union{DecisionTreeClassifier,DecisionTreeRegressor}

MLJBase.package_name(::Type{<:DTTypes}) = "DecisionTree"
MLJBase.package_uuid(::Type{<:DTTypes}) = "7806a523-6efd-50cb-b5f6-3fa6f1930dbb"
MLJBase.package_url(::Type{<:DTTypes}) = "https://github.com/bensadeghi/DecisionTree.jl"
MLJBase.is_pure_julia(::Type{<:DTTypes}) = true

MLJBase.load_path(::Type{<:DecisionTreeClassifier}) = "MLJModels.DecisionTree_.DecisionTreeClassifier"
MLJBase.load_path(::Type{<:DecisionTreeRegressor}) = "MLJModels.DecisionTree_.DecisionTreeRegressor"

MLJBase.input_scitype_union(::Type{<:DecisionTreeClassifier}) = MLJBase.Continuous
MLJBase.input_scitype_union(::Type{<:DecisionTreeRegressor}) = MLJBase.Continuous

MLJBase.target_scitype_union(::Type{<:DecisionTreeClassifier}) = MLJBase.Finite
MLJBase.target_scitype_union(::Type{<:DecisionTreeRegressor}) = MLJBase.Continuous

MLJBase.input_is_multivariate(::Type{<:DecisionTreeClassifier}) = true
MLJBase.input_is_multivariate(::Type{<:DecisionTreeRegressor}) = true

end # module


text:
module q4

import Evolve

steps = 32000                # number of simulation steps
samplesteps = 50            # steps between measurements
L = 2                       # width of the box
xmin = -1                   # leftmost point in the box
N = 2000                    # number of subintervals
ϵ = L/N                     # width of each subinterval
k = 200                     # wavenumbers
σ = 0.05                    # wavepacket width
h = 1000                      # energy is 123.368 for w = 0.2
b = 1
ss = [0.001, 0.01, 0.1]                    # wall thickness
w = 0.2                     # pocket width
x0 = -0.5                   # wavepacket starting position

# Functions for square well states and 1st excited state
Φground = x -> √(2/w)*cos(π*x/w)
Φexcited = x -> √(2/w)*sin(2*π*x/w)
Φs = [:Φground, :Φexcited]

# Run the trials
@sync @parallel for Φ0 in Φs
    for s in ss
        print("Initializing for Φ0 = $(string(Φ0)), s = $s... ")
        V(xs) = map(x -> Evolve.pocket(L, h, b, s, w, xmin, x), xs)      # pocket potential
        AV = Evolve.A(V)            # evolution matrix constructor for square well

        A = AV(L, N, xmin)
        Φ = eval(Φ0)
        ϕ0 = zeros(Complex, N+1)
        xpocket = [(0.5L-0.5w):ϵ:(0.5L+0.5w)]
        ϕ0[int((0.5L-0.5w)N/L):int((0.5L+0.5w)N/L)] = Φ(xpocket)
        x = (L/N .* [0:N] + xmin)   # x positions

        # Evolve the system and record positions every so many steps
        print("evolving... ")
        test = x -> (abs(x[2]) > 1e-2 || abs(x[end-1]) > 1e-2)
        ϕ = Evolve.evolve(A, ϕ0, steps, test)
        phis = [ϕ0 ϕ[:,samplesteps:samplesteps:end]]

        # calculate transmission and reflection probabilities
        print("calculating transmission probabilities...")
        Ts = Float64[]
        Rs = Float64[]
        for i in 1:size(phis)[2]
            phi = phis[:,i]
            T = mean(abs(phi[1:int((0.5L-0.5w-s)N/L),end]).^2)
            T += mean(abs(phi[int((0.5L+0.5w+s)N/L):end,end]).^2)
            R = mean(abs(phi[int((0.5L-0.5w)N/L):int((0.5L+0.5w)N/L),end]).^2)
            I = T + R
            T /= I
            R /= I
            push!(Ts, T)
            push!(Rs, R)
            ### print("T = $T, R = $R...")
        end
        writedlm("trans_pocket_$(string(Φ0))\_$s.dat", [Ts Rs])

        println("done after $(size(ϕ)[2]) steps.")
        writedlm("phis_pocket_im_$(string(Φ0))\_$s.dat", imag(phis))
        writedlm("phis_pocket_re_$(string(Φ0))\_$s.dat", real(phis))
    end
end

### writedlm("decay.dat", [ss Ts Rs])
end


text:
# convert binary into HDF5 data

using HDF5
using Compat

srand(12345678)

datasets = @compat(Dict("train" => ["train-labels-idx1-ubyte","train-images-idx3-ubyte"],
            "test" => ["t10k-labels-idx1-ubyte","t10k-images-idx3-ubyte"]))

for key in keys(datasets)
  label_fn, data_fn = datasets[key]
  label_f = open(label_fn)
  data_f  = open(data_fn)

  label_header = read(label_f, Int32, 2)
  @assert ntoh(label_header[1]) == 2049
  n_label = @compat(Int(ntoh(label_header[2])))
  data_header = read(data_f, Int32, 4)
  @assert ntoh(data_header[1]) == 2051
  n_data = @compat(Int(ntoh(data_header[2])))
  @assert n_label == n_data
  h = @compat(Int(ntoh(data_header[3])))
  w = @compat(Int(ntoh(data_header[4])))

  n_batch = 1
  @assert n_data % n_batch == 0
  batch_size = @compat(Int(n_data / n_batch))

  println("Exporting $n_data digits of size $h x $w")

  h5open("$key.hdf5", "w") do h5
    dset_data = d_create(h5, "data", datatype(Float64), dataspace(w, h, 1, n_data))
    dset_label = d_create(h5, "label", datatype(Float64), dataspace(1, n_data))

    for i = 1:n_batch
      idx = (i-1)*batch_size+1:i*batch_size
      println("  $idx...")

      idx = collect(idx)
      rp = randperm(length(idx))

      img = readbytes(data_f, batch_size * h*w)
      img = convert(Array{Float64},img) / 256 # scale into [0,1)
      class = readbytes(label_f, batch_size)
      class = convert(Array{Float64},class)

      for j = 1:length(idx)
        r_idx = rp[j]
        dset_data[:,:,1,idx[j]] = img[(r_idx-1)*h*w+1:r_idx*h*w]
        dset_label[1,idx[j]] = class[r_idx]
      end
    end
  end

  close(label_f)
  close(data_f)
end


text:
module aoc_08
    export read_input

    const instruction_r = r"([a-z]*) ([+\-]*[0-9]*)"

    const nop = UInt8(0)
    const acc = UInt8(1)
    const jmp = UInt8(2)

    const code_dict = Dict("nop" => nop, "acc" => acc, "jmp" => jmp)

    struct Instruction
        code::UInt8
        value::Int64
    end

    function read_input(filename)
        lines = readlines(filename)
        parse_instruction.(lines)
    end

    function parse_instruction(line::AbstractString)
        m = match(instruction_r , line)
        Instruction(
            code_dict[ m.captures[1] ],
            parse(Int64, m.captures[2] )
        )
    end

    function evaluate(in::Vector{Instruction}, accumulator = 0, i = 1, evaluated = falses( size(in) ))
        while !evaluated[i]
            evaluated[i] = true
            println("i: {$i}")
            if in[i].code == nop
                i += 1
            elseif in[i].code == acc
                accumulator += in[i].value
                i += 1
            elseif in[i].code == jmp
                i += in[i].value
            end
        end
        return accumulator
    end

    # Part 2

    struct State
        accumulator::UInt64
        i::Int
        evaluated::Vector{Bool}
        switched::Bool
    end

    import Base.copy
    copy(s::State) = State(s.accumulator, i, copy(evaluated), switched)
    State(sz::Tuple{Int}) = State(0, 1, falses(sz), false)
    State(in::AbstractArray) = State( size(in) )

    function evaluate_and_fix(in::Vector{Instruction}, accumulator = 0, i = 1, evaluated = falses( size(in) ), switched = false)
        last = length(in)
        while i <= last && !evaluated[i]
            evaluated[i] = true
            # println("i: {$i}")
            if in[i].code == nop || in[i].code == jmp
                # nop
                if in[i].code == nop || !switched
                    nop_eval = evaluate_and_fix(in, accumulator, i + 1, copy(evaluated), switched || in[i].code == jmp)
                    if nop_eval[1]
                        @info "Switched", in[i].code == jmp, i
                        return nop_eval
                    end
                end
                # jmp
                if in[i].code == jmp || !switched
                    jmp_eval = evaluate_and_fix(in, accumulator, i + in[i].value, copy(evaluated), switched || in[i].code == nop)
                    if jmp_eval[1]
                        @info "Switched", in[i].code == nop, i
                        return jmp_eval
                    end
                end
            elseif in[i].code == acc
                accumulator += in[i].value
                i += 1
            end
        end
        if i <= last
            #@info "Failed at ", i
            return (false, accumulator)
        else
            return (true, accumulator)
        end
    end

end


text:
#=
# Moment-matching conditions

The stationarity conditions for ML learning of the RBM imply certain moment-matching
conditions between the data and the RBM distribution.
However, due to biased nature of the PCD approximation, these conditions might not hold
exactly in practice.
=#

import Makie, CairoMakie
import MLDatasets
import Flux
using Statistics: mean, cor
using LinearAlgebra: norm
using RestrictedBoltzmannMachines: RBM, Binary, sample_from_inputs, free_energy, log_pseudolikelihood, training_epochs
using RestrictedBoltzmannMachines: sample_v_from_v, sample_h_from_v, initialize!, pcd!, minibatch_count, moving_average

# load MNIST dataset

train_x = MLDatasets.MNIST(split=:train).features .> 0.5; # binarize
train_y = MLDatasets.MNIST(split=:train).targets;
train_x = train_x[:, :, train_y .== 0]; # only zeros for speed

# floating type we will use

Float = Float32

# initialize RBM

nhidden = 64
batchsize = 64
nupdates = 20000
epochs = training_epochs(; nsamples = size(train_x, 3), nupdates, batchsize)
rbm = RBM(Binary(Float,28,28), Binary(Float,nhidden), randn(Float,28,28,nhidden)/28)
initialize!(rbm, train_x);
lpls = Float64[]
Fm = zeros(batchsize, epochs * minibatch_count(train_x; batchsize))
Fd = zeros(batchsize, epochs * minibatch_count(train_x; batchsize))
t = lastepoch = 0
function mycb(; epoch, vm, vd, kw...)
    global lastepoch
    global t += 1
    Fm[:, t] .= free_energy(rbm, vm)
    Fd[:, t] .= free_energy(rbm, vd)
    if epoch > lastepoch
        push!(lpls, mean(log_pseudolikelihood(rbm, train_x)))
        lastepoch = epoch
    end
end
@time pcd!(
    rbm, train_x;
    epochs, batchsize, callback=mycb, l2_weights=1e-4,
    optim=Flux.ADAM(5e-4, (0.9, 0.999))
);

# Plot log_pseudolikelihood during training

Makie.lines(lpls)

# Average energies of data and model fantasy particles used during training.

fig = Makie.Figure()
ax = Makie.Axis(fig[1,1])
Makie.lines!(ax, vec(mean(Fm; dims=1)), color=(:red, 0.25))
Makie.lines!(ax, vec(mean(Fd; dims=1)), color=(:blue, 0.25))
Makie.lines!(ax, moving_average(vec(mean(Fm; dims=1)), 20), color=:red, label="Fm")
Makie.lines!(ax, moving_average(vec(mean(Fd; dims=1)), 20), color=:blue, label="Fd")
Makie.axislegend(ax)
fig

# Difference in average energies of data and model fantasy particles used during training.

fig = Makie.Figure()
ax = Makie.Axis(fig[1,1])
Makie.lines!(ax, vec(mean(Fm .- Fd; dims=1)), color=(:blue, 0.25))
Makie.lines!(ax, moving_average(vec(mean(Fm .- Fd; dims=1)), 20), color=:red)
fig

# Generate samples, in v-space and h-space.

nsteps = 2000
nsamples = 1000
v_model = falses(28, 28, nsamples, nsteps);
v_model[:,:,:,1] .= sample_from_inputs(rbm.visible, falses(28,28,nsamples));
for t in 2:nsteps
    v_model[:, :, :, t] .= sample_v_from_v(rbm, v_model[:, :, :, t - 1])
end
h_data = sample_h_from_v(rbm, train_x);
h_model = sample_h_from_v(rbm, v_model[:,:,:,end]);

v_data_center = train_x .- mean(train_x; dims=3);
h_data_center = h_data .- mean(h_data; dims=2);
vh_data = reshape(v_data_center, 28*28, :) * h_data_center' / size(v_data_center, 3);

v_model_center = v_model[:,:,:,end] .- mean(v_model[:,:,:,end]; dims=3);
h_model_center = h_model .- mean(h_model; dims=2);
vh_model = reshape(v_model_center, 28*28, :) * h_model_center' / size(v_model_center, 3);

# Average energy of generated and data points

fig = Makie.Figure()
ax = Makie.Axis(fig[1,1])
Makie.hist!(ax, free_energy(rbm, train_x), color=:purple, label="data", normalization=:pdf)
Makie.hist!(ax, free_energy(rbm, v_model[:,:,:,end]), color=:orange, label="model", normalization=:pdf)
Makie.axislegend(ax)
fig

# Plot moment-matching conditions

# <v>_data vs. <v>_model

Makie.scatter(
    vec(mean(train_x; dims=3)),
    vec(mean(v_model[:,:,:,end]; dims=3))
)
Makie.xlims!(0,1)
Makie.ylims!(0,1)
Makie.current_figure()

# <h>_data vs. <h>_model

Makie.scatter(
    vec(mean(h_data; dims=2)),
    vec(mean(h_model; dims=2)),
)
Makie.xlims!(0,1)
Makie.ylims!(0,1)
Makie.current_figure()

# <vh>_data vs. <vh>_model

Makie.scatter(vec(vh_data), vec(vh_model))
Makie.xlims!(0,1)
Makie.ylims!(0,1)
Makie.current_figure()

# Correlation

cor(vec(vh_data), vec(vh_model))

# Verify convergence of sampling

Makie.lines(vec(mean(free_energy(rbm, v_model[:,:,1:100,:]); dims=1)))

# Plot average sampled digit

Makie.heatmap(mean(v_model[:,:,:,end]; dims=3)[:,:,1])

# Plot average data digit

Makie.heatmap(mean(train_x; dims=3)[:,:,1])

# Difference

Makie.heatmap(mean(v_model[:,:,:,end]; dims=3)[:,:,1] - mean(train_x; dims=3)[:,:,1])

# Plot weights  of maximum intensity.

𝒫 = sortperm([norm(rbm.w[:,:,m]) for m in 1:nhidden]; rev=true);

fig = Makie.Figure()
for i = 1:7
    ax = Makie.Axis(fig[1,i], width=50, height=50)
    Makie.heatmap!(ax, rbm.w[:, :, 𝒫[i]])
    Makie.hidedecorations!(ax)
    Makie.hidespines!(ax)
end
Makie.resize_to_layout!(fig)
fig


text:
import Pkg
Pkg.activate(".")
Pkg.instantiate()
using WebIO
WebIO.install_jupyter_nbextension()
using DataFrames
using DataFramesMeta
# using Gadfly # Removed for now because downgrades DataFrames version
using IJulia
using Weave
using CSV
using Revise
using TableView
using Printf
using PrettyTables
using Formatting
using Statistics

# ==================== GENERAL UTILITY FUNCTIONS ====================

broadwrap(f) = function (args...) broadcast(f, args...) end
esc_latex(s) = replace(s, r"(#|%|&|_)" => s"\\\1")
function number2latex(num; enclose = true)
    if ismissing(num)
        "--"
    elseif isa(num, Integer)
        (s -> enclose ? "\\($s\\)" : s)(sprintf1("%'d", num))
    elseif isa(num, AbstractFloat)
        (s -> enclose ? "\\($s\\)" : s)(sprintf1("%'.2f", num))
    elseif isa(num, NTuple{2, Number})
        "$(number2latex(num[1]; enclosed = enclose)) ($(number2latex(num[2]; enclosed = enclose)))"
    else
        error("Unexpected type.")
    end
end

# ==================== METHODS FOR highlight_best_values! ====================
wrap_in_textbf(str) = "\\textbf{$str}"
wrap_in_textit(str) = "\\textit{$str}"
# See https://www.pcre.org/current/doc/html/pcre2syntax.html#SEC2
const MATH_MODE_REGEX = Regex(
    "^(?:\\Q\$\\E|\\Q\\(\\E)(.+)(?:\\Q\$\\E|\\Q\\)\\E)\$"
)
function dirt2number(
  dirt, decimal_separator = '.', failure = NaN
) :: Float64
		gold = filter(c -> isdigit(c) | (c == decimal_separator), dirt)
    number = tryparse(Float64, gold)
    isnothing(number) && return failure
    return number
end
function latex2number(
  str, failure = NaN; separator = nothing
) :: Float64
    str = isnothing(separator) ? str : replace(str, separator => "")
    v = tryparse(Float64, str)
    if isnothing(v)
        m = match(MATH_MODE_REGEX, str)
        if !isnothing(m)
          v = tryparse(Float64, m.captures[1])
        end
    end
    return isnothing(v) ? failure : v
end
find_index_min(a) = isempty(a) ? Int[] : Int[findmin(a)[2]]
find_index_max(a) = isempty(a) ? Int[] : Int[findmax(a)[2]]
find_index_allmin(a) = isempty(a) ? Int[] : findall(isequal(findmin(a)[1]), a)
find_index_allmax(a) = isempty(a) ? Int[] : findall(isequal(findmax(a)[1]), a)

function highlight_best_values!(
    df;
    # Set of columns to be considered.
    columns = 1:ncol(df),
    # Transforms the cell body into another object.
    cleaner = latex2number,
    # Return true if the object returned by cleaner should be ignored.
    ignorer = isnan,
    # Takes a row of clean non-ignored values and returns the
    # indexes of the ones that should be highlighted.
    chooser = find_index_min,
    # Takes the cell body and return the highlighted cell body.
    changer = wrap_in_textbf
)
    for row_id in 1:nrow(df)
        row = df[row_id, :]
        selected_values = []
        selected_columns = empty(columns)
        for col_id in columns
            value = cleaner(row[col_id])
            if !ignorer(value)
                push!(selected_values, value)
                push!(selected_columns, col_id)
            end
        end
        chosen = chooser(selected_values)
        for col_id in selected_columns[chosen]
            df[row_id, col_id] = changer(df[row_id, col_id])
        end
    end
    return
end


text:
module Brownian

using Distributions
using PDMats
using StatsBase
using FFTW

import Base: convert
import Random: rand!, rand
import Distributions: VariateForm, Univariate, Multivariate, ValueSupport, Discrete, Continuous, Sampleable
import StatsBase: IntegerVector, autocov!, autocov

export BrownianMotion, FBM, FGN, rand!, rand, autocov, chol_update

include("BrownianMotion.jl")
include("FBM.jl")

end


text:
using ReachabilityAnalysis, SparseArrays, ExponentialUtilities, JLD2, LazySets
using ReachabilityAnalysis: reach_homog_krylov_LGG09!
using LazySets.Arrays

# HEAT01 5x5x5 grid
function heat01(; δ=0.02) # step size
    @load joinpath(@__DIR__, "HEAT01.jld2") A xind

    n = size(A, 1)
    x0ind = xind.nzind
    # Note: For xc we add +1 because Julia has 1 based indexing in contrast to Python which is 0 based
    xc = 62 + 1

    c = sparsevec(x0ind, fill(1.0, length(x0ind)), n)
    r = sparsevec(x0ind, fill(0.1, length(x0ind)), n)
    Ω₀ = Hyperrectangle(c, r)

    Aᵀδ = transpose(A .* δ)
    ℓ = SingleEntryVector(xc, n, 1.0)

    return A, Aᵀδ, Ω₀, ℓ
end

# HEAT02 10x10x10 grid
function heat02(; δ=0.02) # step size
    @load joinpath(@__DIR__, "HEAT02.jld2") A xind

    n = size(A, 1)
    x0ind = xind.nzind
    xc = 555 + 1

    c = sparsevec(x0ind, fill(1.0, length(x0ind)), n)
    r = sparsevec(x0ind, fill(0.1, length(x0ind)), n)
    Ω₀ = Hyperrectangle(c, r)

    Aᵀδ = transpose(A .* δ)
    ℓ = SingleEntryVector(xc, n, 1.0)

    return A, Aᵀδ, Ω₀, ℓ
end

# HEAT03 20x20x20 grid
function heat03(; δ=0.02) # step size
    @load joinpath(@__DIR__, "HEAT03.jld2") A xind

    n = size(A, 1)
    x0ind = xind.nzind
    xc = 4210 + 1

    c = sparsevec(x0ind, fill(1.0, length(x0ind)), n)
    r = sparsevec(x0ind, fill(0.1, length(x0ind)), n)
    Ω₀ = Hyperrectangle(c, r)

    Aᵀδ = transpose(A .* δ)
    ℓ = SingleEntryVector(xc, n, 1.0)

    return A, Aᵀδ, Ω₀, ℓ
end

# HEAT04 50x50x50 grid
function heat04(; δ=0.02) # step size
    @load joinpath(@__DIR__, "HEAT04.jld2") A xind

    n = size(A, 1)
    x0ind = xind.nzind
    xc = 63775 + 1

    c = sparsevec(x0ind, fill(1.0, length(x0ind)), n)
    r = sparsevec(x0ind, fill(0.1, length(x0ind)), n)
    Ω₀ = Hyperrectangle(c, r)

    Aᵀδ = transpose(A .* δ)
    ℓ = SingleEntryVector(xc, n, 1.0)

    return A, Aᵀδ, Ω₀, ℓ
end


text:
# Example how to use the residual block
# Author: Philipp Witte, pwitte3@gatech.edu
# Date: January 2020

using Test, LinearAlgebra, InvertibleNetworks

# Input
nx = 64
ny = 64
n_in = 10
n_hidden = 20
batchsize = 2

# Residual blocks of various networks
architecture = "Glow"

if architecture == "Glow"
    # Glow
    k1 = 3; p1 = 1; s1 = 1
    k2 = 1; p2 = 0; s2 = 1
elseif architecture == "HINT"
    # HINT
    k1 = 3; p1 = 1; s1 = 1  # standard res-net block
    k2 = 3; p2 = 1; s2 = 1
elseif architecture == "iRIM"
    # i-RIM
    k1 = 4; p1 = 0; s1 = 4  # downsampling/upsampling in first and third layer
    k2 = 3; p2 = 1; s2 = 1  # Standard res-net in center layer
end

# Input image
X = glorot_uniform(nx, ny, n_in, batchsize)

# Residual blocks
RB = ResidualBlock(nx, ny, n_in, n_hidden, batchsize; k1=k1, k2=k2, p1=p1, p2=p2, s1=s1, s2=s2)
RB0 = ResidualBlock(nx, ny, n_in, n_hidden, batchsize; k1=k1, k2=k2, p1=p1, p2=p2, s1=s1, s2=s2)

# Observed data
Y = RB.forward(X)

# Predicted data
Y0 = RB0.forward(X)
@test isnothing(RB0.W1.grad) # after forward pass, gradients are not set

# Residual
ΔY = Y0 - Y

# Backward pass: need to pass data residual and original input X (as layer is not invertible)
ΔX = RB0.backward(ΔY, X)   # returns derivative w.r.t input
@test ~isnothing(RB0.W1.grad)    # after inverse pass, gradients are set
@test ~isnothing(RB0.b1.grad)


text:
"""
    OneHotArray{T,L,N,M,I} <: AbstractArray{Bool,M}

These are constructed by [`onehot`](@ref) and [`onehotbatch`](@ref).
Parameter `I` is the type of the underlying storage, and `T` its eltype.
"""
struct OneHotArray{T<:Integer, L, N, var"N+1", I<:Union{T, AbstractArray{T, N}}} <: AbstractArray{Bool, var"N+1"}
  indices::I
end
OneHotArray{T, L, N, I}(indices) where {T, L, N, I} = OneHotArray{T, L, N, N+1, I}(indices)
OneHotArray(indices::T, L::Integer) where {T<:Integer} = OneHotArray{T, L, 0, 1, T}(indices)
OneHotArray(indices::I, L::Integer) where {T, N, I<:AbstractArray{T, N}} = OneHotArray{T, L, N, N+1, I}(indices)

_indices(x::OneHotArray) = x.indices
_indices(x::Base.ReshapedArray{<: Any, <: Any, <: OneHotArray}) =
  reshape(parent(x).indices, x.dims[2:end])

const OneHotVector{T, L} = OneHotArray{T, L, 0, 1, T}
const OneHotMatrix{T, L, I} = OneHotArray{T, L, 1, 2, I}

@doc @doc(OneHotArray)
OneHotVector(idx, L) = OneHotArray(idx, L)
@doc @doc(OneHotArray)
OneHotMatrix(indices, L) = OneHotArray(indices, L)

# use this type so reshaped arrays hit fast paths
# e.g. argmax
const OneHotLike{T, L, N, var"N+1", I} =
  Union{OneHotArray{T, L, N, var"N+1", I},
        Base.ReshapedArray{Bool, var"N+1", <:OneHotArray{T, L, <:Any, <:Any, I}}}

_isonehot(x::OneHotArray) = true
_isonehot(x::Base.ReshapedArray{<:Any, <:Any, <:OneHotArray{<:Any, L}}) where L = (size(x, 1) == L)

Base.size(x::OneHotArray{<:Any, L}) where L = (Int(L), size(x.indices)...)

function Base.getindex(x::OneHotArray{<:Any, <:Any, N}, i::Integer, I::Vararg{Any, N}) where N
  @boundscheck checkbounds(x, i, I...)
  return x.indices[I...] .== i
end

function Base.getindex(x::OneHotArray{<:Any, L}, ::Colon, I...) where L
  @boundscheck checkbounds(x, :, I...)
  return OneHotArray(x.indices[I...], L)
end

Base.getindex(x::OneHotArray, ::Colon) = BitVector(reshape(x, :))
Base.getindex(x::OneHotArray{<:Any, <:Any, N}, ::Colon, ::Vararg{Colon, N}) where N = x

function Base.showarg(io::IO, x::OneHotArray, toplevel)
  print(io, ndims(x) == 1 ? "OneHotVector(" : ndims(x) == 2 ? "OneHotMatrix(" : "OneHotArray(")
  Base.showarg(io, x.indices, false)
  print(io, ')')
  toplevel && print(io, " with eltype Bool")
  return nothing
end

# this is from /LinearAlgebra/src/diagonal.jl, official way to print the dots:
function Base.replace_in_print_matrix(x::OneHotLike, i::Integer, j::Integer, s::AbstractString)
  x[i,j] ? s : _isonehot(x) ? Base.replace_with_centered_mark(s) : s
end

# copy CuArray versions back before trying to print them:
Base.print_array(io::IO, X::OneHotLike{T, L, N, var"N+1", <:CuArray}) where {T, L, N, var"N+1"} =
  Base.print_array(io, adapt(Array, X))
Base.print_array(io::IO, X::LinearAlgebra.AdjOrTrans{Bool, <:OneHotLike{T, L, N, var"N+1", <:CuArray}}) where {T, L, N, var"N+1"} =
  Base.print_array(io, adapt(Array, X))

_onehot_bool_type(::OneHotLike{<:Any, <:Any, <:Any, N, <:Union{Integer, AbstractArray}}) where N = Array{Bool, N}
_onehot_bool_type(::OneHotLike{<:Any, <:Any, <:Any, N, <:CuArray}) where N = CuArray{Bool, N}

function Base.cat(x::OneHotLike{<:Any, L}, xs::OneHotLike{<:Any, L}...; dims::Int) where L
  if isone(dims) || any(x -> !_isonehot(x), (x, xs...))
    return cat(map(x -> convert(_onehot_bool_type(x), x), (x, xs...))...; dims = dims)
  else
    return OneHotArray(cat(_indices(x), _indices.(xs)...; dims = dims - 1), L)
  end
end

Base.hcat(x::OneHotLike, xs::OneHotLike...) = cat(x, xs...; dims = 2)
Base.vcat(x::OneHotLike, xs::OneHotLike...) = cat(x, xs...; dims = 1)

# optimized concatenation for matrices and vectors of same parameters
Base.hcat(x::T, xs::T...) where {L, T <: OneHotLike{<:Any, L, <:Any, 2}} =
  OneHotMatrix(reduce(vcat, _indices.(xs); init = _indices(x)), L)
Base.hcat(x::T, xs::T...) where {L, T <: OneHotLike{<:Any, L, <:Any, 1}} =
  OneHotMatrix(reduce(vcat, _indices.(xs); init = _indices(x)), L)

MLUtils.batch(xs::AbstractArray{<:OneHotVector{<:Any, L}}) where L = OneHotMatrix(_indices.(xs), L)

Adapt.adapt_structure(T, x::OneHotArray{<:Any, L}) where L = OneHotArray(adapt(T, _indices(x)), L)

Base.BroadcastStyle(::Type{<:OneHotArray{<: Any, <: Any, <: Any, N, <: CuArray}}) where N = CUDA.CuArrayStyle{N}()

Base.map(f, x::OneHotLike) = Base.broadcast(f, x)

Base.argmax(x::OneHotLike; dims = Colon()) =
  (_isonehot(x) && dims == 1) ?
    reshape(CartesianIndex.(_indices(x), CartesianIndices(_indices(x))), 1, size(_indices(x))...) :
    invoke(argmax, Tuple{AbstractArray}, x; dims = dims)


text:
function getSubtypes(T=Any)::Vector{DataType}
    subs = subtypes(T)
    ret = filter(isconcretetype, subs)
    filter!(isabstracttype, subs)

    while !isempty(subs)
        ntype = popfirst!(subs)
        ntype == Any && continue
        nsubs = subtypes(ntype)
        append!(ret, Iterators.filter(isconcretetype, nsubs))
        append!(subs, Iterators.filter(isabstracttype, nsubs))
    end

    ret
end


text:
function main(r::Robot)
    go_to_west_south_corner_and_return_path!(r; go_around_barriers = true)

    max_s = 0
    current_w = 0
    current_h = 0
    border_now = false
    side = East

    while !isborder(r, North)
        while move_if_possible!(r, side)
            current_w, current_h, border_now = find_rectangle_w_h!(r, current_w, current_h, border_now, side)
            if current_h * current_w > max_s
                max_s = current_h * current_w
            end
        end
        current_w, current_h, border_now = find_rectangle_w_h!(r, current_w, current_h, border_now, side)
        if current_h * current_w > max_s
            max_s = current_h * current_w
        end

        side = inverse_side(side)
        move!(r, North)
        current_w, current_h, border_now = find_rectangle_w_h!(r, current_w, current_h, border_now, side)
        if current_h * current_w > max_s
            max_s = current_h * current_w
        end

    end

    return max_s
end

function find_rectangle_w_h!(r::Robot, x::Int, y::Int, border_now::Bool, move_side::HorizonSide)
    if isborder(r, North) && !border_now
        border_now = true
        x = 0
        y = 0
    end

    if !isborder(r, North) && border_now
        border_now = false
        move!(r, North)
        while isborder(r, inverse_side(move_side))
            y += 1
            move!(r, North)
        end
        for i ∈ 1:y
            move!(r, South)
        end
        move!(r, South)

    end

    if border_now
        x += 1
    end

    return x, y, border_now
end

max_stars_repo_name: Savko-fokin/mirea-progs

max_issues_repo_name: Savko-fokin/mirea-progs

max_forks_repo_name: Savko-fokin/mirea-progs


text:
@testset "Carrier replica" begin

    carrier = StructArray(zeros(Complex{Float32}, 2500))

    Tracking.gen_carrier_replica!(
        carrier,
        1500Hz,
        2.5e6Hz,
        0.25, # π / 2
        111,
        2390
    )

    @test carrier[111:2500] ≈ cis.(2π * (0:2389) * 1500Hz / 2.5e6Hz .+ π / 2)

 #   @test sqrt(mean(abs2.(carrier.re[111:2500] ./ 1 << 7 .-
  #                        cos.(2π * (0:2389) * 1500Hz / 2.5e6Hz .+ π / 2)))) < 8e-3
end

@testset "Update carrier phase" begin
    carrier_phase = 0.25
    carrier_frequency = 10Hz
    sampling_frequency = 100Hz
    num_samples = 2000
    phase = @inferred Tracking.update_carrier_phase(
        num_samples,
        carrier_frequency,
        sampling_frequency,
        carrier_phase,
    )
    @test phase ≈ mod(0.25 + 0.1 * 2000 + 0.5, 1) - 0.5

    carrier_phase = 0.25
    carrier_frequency = 10Hz
    sampling_frequency = 2.5e6Hz
    num_samples = 2500
    phase = @inferred Tracking.update_carrier_phase(
        num_samples,
        carrier_frequency,
        sampling_frequency,
        carrier_phase
    )
    @test phase ≈ mod(0.25 + 10 * 2500 / 2.5e6 + 0.5, 1) - 0.5
end


text:
function gen_inv_quote(T, N, L)
    q, qa = initial_quote(L::Int, :U)

    for i ∈ 1:N
        lti = small_triangle(i)
        bti = lti+i
        Ui_i = Symbol(:Ui_, bti)
        U_i = Symbol(:U_, bti)
        push!(qa, :($Ui_i = 1 / $U_i))
        for j ∈ i+1:N
            ltj = small_triangle(j)
            Ui_i = Symbol(:Ui_, bti)
            U_i = Symbol(:U_, ltj+i)
            push!(qa, :(Utemp = $U_i * $Ui_i))
            for k ∈ i+1:j-1
                Ui_i = Symbol(:Ui_, small_triangle(k) + i)
                U_i = Symbol(:U_, ltj+k)
                push!(qa, :(Utemp += $U_i * $Ui_i))
            end
            Ui_i = Symbol(:Ui_, ltj+i)
            U_i = Symbol(:U_, ltj+j)
            push!(qa, :($Ui_i = -Utemp / $U_i))
        end
    end

    push!(q.args, :($T( ( @ntuple $L Ui ) )))
    q
end
function gen_ip_inv_quote!(qa, N, L, Ui = :Ui, U = :U, extract = extract_symbol)

    temp_num = 0

    for i ∈ 1:N
        lti = small_triangle(i)
        bti = lti+i
        Ui_i = extract(Ui, bti)
        U_i = extract(U, bti)
        push!(qa, :($Ui_i = 1 / $U_i))
        for j ∈ i+1:N
            ltj = small_triangle(j)
            Ui_i = extract(Ui, bti)
            U_i = extract(U, ltj+i)
            Utemp = Symbol(:t_, temp_num)
            temp_num += 1 #Comment this line on / off to change.
            push!(qa, :($Utemp = $U_i * $Ui_i)) # I'm curious. Woud renaming Utemps help optimization, or just be wasteful?
            for k ∈ i+1:j-1
                Ui_i = extract(Ui, small_triangle(k) + i)
                U_i = extract(U, ltj+k)
                push!(qa, :($Utemp += $U_i * $Ui_i))
            end
            Ui_i = extract(Ui, ltj+i)
            U_i = extract(U, ltj+j)
            push!(qa, :($Ui_i = -$Utemp / $U_i))
        end
    end

    # push!(q.args, :($T( ( @ntuple $L Ui ) )))
    qa
end

@generated function LinearAlgebra.inv(U::AbstractUpperTriangular{T,N,L}) where {T,N,L}
    gen_inv_quote(UpperTriangularMatrix{T,N,L}, N, L)
end
@generated function LinearAlgebra.inv(U::AbstractLowerTriangular{T,N,L}) where {T,N,L}
    gen_inv_quote(LowerTriangularMatrix{T,N,L}, N, L)
end

function gen_extracted_ip_inv_quote!(qa, N, L, Ui = :Ui, U = :U)
    for i ∈ 1:L
        push!(qa, :($(id_symbol(U, i)) = $(U)[$i]))
    end
    gen_ip_inv_quote!(qa, N, L, Ui, U, id_symbol)
    for i ∈ 1:L
        push!(qa, :($(Ui)[$i] = $(id_symbol(Ui, i))))
    end
end

@generated function LinearAlgebra.inv!(U::Union{UpperTriangularMatrix{T,N,L},LowerTriangularMatrix{T,N,L}}) where {T,N,L}
    q, qa = create_quote()
    gen_extracted_ip_inv_quote!(qa, N, L, :U, :U)
    push!(q.args, :U)
    q
end
@generated function LinearAlgebra.inv!(Ui::UpperTriangularMatrix{T,N,L}, U::AbstractUpperTriangular{T,N,L}) where {T,N,L}
    q, qa = create_quote()
    gen_extracted_ip_inv_quote!(qa, N, L, :Ui, :U)
    push!(q.args, :Ui)
    q
end
@generated function LinearAlgebra.inv!(Ui::LowerTriangularMatrix{T,N,L},U::AbstractLowerTriangular{T,N,L}) where {T,N,L}
    q, qa = create_quote()
    gen_extracted_ip_inv_quote!(qa, N, L, :Ui, :U)
    push!(q.args, :Ui)
    q
end