text:
#///////////////////////////////////////
#// File Name: pets.jl
#// Author: Haruki Nishimura (hnishimura@stanford.edu)
#// Date Created: 2020/11/06
#// Description: PETS algorithm
#///////////////////////////////////////

using Random
using Distributions
using LinearAlgebra

using Distributed
import Future.randjump
"""
    CrossEntropyDirectOptimizationSolver(μ_init_array::Vector{Vector{Float64}},
    Σ_init_array::Vector{Matrix{Float64}}; kwargs...)

PETS Solver initialized with `μ_init_array = [μ_0,...,μ_{N-1}]` and
`Σ_init_array = [Σ_0,...,Σ_{N-1}]`, where the initial control distribution at time
`k` is a Gaussian distribution `Distributions.MvNormal(μ_k, Σ_k)`.

# Optional Keyword Arguments
- `num_control_samples::Int64` -- number of Monte Carlo samples for the control
  trajectory. Default: `10`.
- `num_trajectory_samples::Int64` -- number of Monte Carlo samples for the state
  trajectory. Default: `10`.
- `num_elite::Int64` -- number of elite samples. Default: `3`.
- `iter_max::Int64` -- maximum iteration number. Default: `5`.
- `smoothing_factor::Float64` -- smoothing factor in (0, 1), used to update the
  mean and the variance of the Cross Entropy distribution for the next iteration.
  If `smoothing_factor` is `0.0`, the updated distribution is independent of the
  previous iteration. If it is `1.0`, the updated distribution is the same as the
  previous iteration. Default. `0.1`.
"""
mutable struct CrossEntropyDirectOptimizationSolver # a.k.a. "PETS"
    # CE solver parameters
    num_control_samples::Int64
    num_trajectory_samples::Int64
    num_elite::Int64
    iter_max::Int64
    smoothing_factor::Float64

    # action_distributions
    μ_init_array::Vector{Vector{Float64}}
    Σ_init_array::Vector{Matrix{Float64}}
    μ_array::Vector{Vector{Float64}}
    Σ_array::Vector{Matrix{Float64}}
    N::Int64 # control sequence length > 0 (must be the same as N in FiniteHorizonGenerativeOptimalControlProblem)
    iter_current::Int64
end

function CrossEntropyDirectOptimizationSolver(μ_init_array::Vector{Vector{Float64}},
                                              Σ_init_array::Vector{Matrix{Float64}};
                                              num_control_samples=10,
                                              num_trajectory_samples=10,
                                              num_elite=3,
                                              iter_max=5,
                                              smoothing_factor=0.1)

    @assert length(μ_init_array) == length(Σ_init_array)
    μ_array, Σ_array = copy(μ_init_array), copy(Σ_init_array);
    N = length(μ_init_array);
    iter_current = 0;

    return CrossEntropyDirectOptimizationSolver(num_control_samples, num_trajectory_samples, num_elite,
                                                iter_max, smoothing_factor, μ_init_array, Σ_init_array,
                                                μ_array, Σ_array, N, iter_current)
end;

function initialize!(direct_solver::CrossEntropyDirectOptimizationSolver)
    direct_solver.iter_current = 0;
    direct_solver.μ_array = copy(direct_solver.μ_init_array);
    direct_solver.Σ_array = copy(direct_solver.Σ_init_array);
end

function compute_cost_worker(direct_solver::CrossEntropyDirectOptimizationSolver,
                              problem::FiniteHorizonGenerativeOptimalControlProblem,
                              x::Vector{Float64}, #initial state
                              u_array::Vector{Vector{Float64}},
                              rng::AbstractRNG, use_true_model=false)
     # trajectory sampling & cost function integration
     x_history_array = Vector{Vector{Vector{Float64}}}(undef, direct_solver.num_trajectory_samples);
     sampled_cost_array = Vector{Float64}(undef, direct_solver.num_trajectory_samples);
     for kk = 1 : direct_solver.num_trajectory_samples # for-loop over trajectory samples
         x_history_array[kk] = Vector{Vector{Float64}}(undef, direct_solver.N + 1);
         x_history_array[kk][1] = copy(x);
         sampled_cost_array[kk] = 0.0;
         for tt = 1 : direct_solver.N # for-loop over time steps
             # get stage cost and perform stochastic transition
             sampled_cost_array[kk] += problem.c(tt - 1, x_history_array[kk][tt], u_array[tt]);
             x_history_array[kk][tt + 1] = problem.f_stochastic(x_history_array[kk][tt], u_array[tt], rng, use_true_model);
         end
         # add terminal cost
         sampled_cost_array[kk] += problem.h(x_history_array[kk][end]);
     end
     # get average cost for this control sequnece u_array
     cost = mean(sampled_cost_array);
end

function compute_cost(direct_solver::CrossEntropyDirectOptimizationSolver,
                      problem::FiniteHorizonGenerativeOptimalControlProblem,
                      x::Vector{Float64}, #initial state
                      control_sequence_array::Vector{Vector{Vector{Float64}}},
                      rng::AbstractRNG, use_true_model=false)
    @assert length(control_sequence_array) == direct_solver.num_control_samples;
    @assert all([length(sequence) == direct_solver.N for sequence in control_sequence_array]);

    if nprocs() > 1
        proc_id_array = 2 .+ [mod(ii, nprocs() - 1) for ii = 0 : direct_solver.num_control_samples - 1]
    else
        proc_id_array = [1 for ii = 0 : direct_solver.num_control_samples - 1]
    end
    cost_array = Vector{Float64}(undef, direct_solver.num_control_samples);
    rng_array = let m = rng
                    [m; accumulate(randjump, fill(big(10)^20, nprocs()-1), init=m)]
                end;
    @sync begin
        for ii = 1 : direct_solver.num_control_samples
            @inbounds @async cost_array[ii] =
                remotecall_fetch(compute_cost_worker, proc_id_array[ii],
                                 direct_solver, problem, x, control_sequence_array[ii],
                                 rng_array[proc_id_array[ii]], use_true_model)
        end
    end
    return cost_array
end

function compute_cost_serial(direct_solver::CrossEntropyDirectOptimizationSolver,
                             problem::FiniteHorizonGenerativeOptimalControlProblem,
                             x::Vector{Float64}, #initial state
                             control_sequence_array::Vector{Vector{Vector{Float64}}},
                             rng::AbstractRNG, use_true_model=false)
    @assert length(control_sequence_array) == direct_solver.num_control_samples;
    @assert all([length(sequence) == direct_solver.N for sequence in control_sequence_array]);

    cost_array = Vector{Float64}(undef, direct_solver.num_control_samples);
    for ii = 1 : direct_solver.num_control_samples # for-loop over action sequences
        # trajectory sampling & cost function integration
        x_history_array_ii = Vector{Vector{Vector{Float64}}}(undef, direct_solver.num_trajectory_samples);
        sampled_cost_array_ii = Vector{Float64}(undef, direct_solver.num_trajectory_samples);
        for kk = 1 : direct_solver.num_trajectory_samples # for-loop over trajectory samples
            x_history_array_ii[kk] = Vector{Vector{Float64}}(undef, direct_solver.N + 1);
            x_history_array_ii[kk][1] = copy(x);
            sampled_cost_array_ii[kk] = 0.0;
            for tt = 1 : direct_solver.N # for-loop over time steps
                # get stage cost and perform stochastic transition
                sampled_cost_array_ii[kk] += problem.c(tt - 1, x_history_array_ii[kk][tt], control_sequence_array[ii][tt]);
                x_history_array_ii[kk][tt + 1] = problem.f_stochastic(x_history_array_ii[kk][tt], control_sequence_array[ii][tt], rng, use_true_model);
            end
            # add terminal cost
            sampled_cost_array_ii[kk] += problem.h(x_history_array_ii[kk][end]);
        end
        # get average cost for this control sequnece
        cost_array[ii] = mean(sampled_cost_array_ii);
    end
    return cost_array
end

function get_elite_samples(direct_solver::CrossEntropyDirectOptimizationSolver,
                           control_sequence_array::Vector{Vector{Vector{Float64}}},
                           cost_array::Vector{Float64})
    @assert length(cost_array) == direct_solver.num_control_samples;
    @assert length(cost_array) == length(control_sequence_array);
    @assert all([length(sequence) == direct_solver.N for sequence in control_sequence_array]);

    control_sequence_cost_pair_array = collect(zip(control_sequence_array, cost_array));
    control_sequence_cost_pair_sorted_array = sort(control_sequence_cost_pair_array, by=x->x[2]);
    control_sequence_elite_array = [pair[1] for pair in control_sequence_cost_pair_sorted_array[1 : direct_solver.num_elite]];

    return control_sequence_elite_array
end

function compute_new_distribution(direct_solver::CrossEntropyDirectOptimizationSolver,
                                  control_sequence_elite_array::Vector{Vector{Vector{Float64}}})
    @assert length(control_sequence_elite_array) == direct_solver.num_elite
    @assert all([length(sequence) == direct_solver.N for sequence in control_sequence_elite_array]);

    μ_new_array = similar(direct_solver.μ_array);
    Σ_new_array = similar(direct_solver.Σ_array);

    mean_array, cov_array = similar(μ_new_array), similar(Σ_new_array);

    for tt = 1 : direct_solver.N # for-loop over time steps
        mean_array[tt] = mean([elite[tt] for elite in control_sequence_elite_array]);
        cov_array[tt] = Diagonal(var([elite[tt] for elite in control_sequence_elite_array]));

        μ_new_array[tt] = (1.0 - direct_solver.smoothing_factor).*mean_array[tt] + direct_solver.smoothing_factor.*direct_solver.μ_array[tt];
        Σ_new_array[tt] = (1.0 - direct_solver.smoothing_factor).*cov_array[tt] + direct_solver.smoothing_factor.*direct_solver.Σ_array[tt];
    end
    return μ_new_array, Σ_new_array
end

function step!(direct_solver::CrossEntropyDirectOptimizationSolver,
               problem::FiniteHorizonGenerativeOptimalControlProblem,
               x::Vector{Float64},
               rng::AbstractRNG,
               use_true_model=false, verbose=true, serial=false)
    direct_solver.iter_current += 1;
    if verbose
        println("**CE iteration $(direct_solver.iter_current)")
    end
    control_sequence_array = Vector{Vector{Vector{Float64}}}(undef, direct_solver.num_control_samples);
    cost_array = Vector{Float64}(undef, direct_solver.num_control_samples);
    # action sequence sampling
    if verbose
        println("****Drawing $(direct_solver.num_control_samples) samples of control sequences of length $(direct_solver.N)")
    end
    for ii = 1 : direct_solver.num_control_samples # for-loop over action sequences
        # sample and action sequence from current distribution;
        control_sequence_array[ii] = Vector{Vector{Float64}}(undef, direct_solver.N)
        for tt = 1 : direct_solver.N # for-loop over time steps
            d = MvNormal(direct_solver.μ_array[tt], direct_solver.Σ_array[tt]);
            u = rand(rng, d);
            control_sequence_array[ii][tt] = u;
        end
    end

    # cost computation
    if verbose
        println("****Evaluating costs of sampled control sequences")
    end
    if !serial
        # note that the resulting cost_array is different from the serial version, due to randjump.
        cost_array = compute_cost(direct_solver, problem, x, control_sequence_array, rng, use_true_model);
    else
        cost_array = compute_cost_serial(direct_solver, problem, x, control_sequence_array, rng, use_true_model);
    end
    if verbose
        println(cost_array)
    end

    # compute elite samples
    if verbose
        println("****Choosing $(direct_solver.num_elite) elite samples")
    end
    control_sequence_elite_array = get_elite_samples(direct_solver, control_sequence_array, cost_array);

    # update CEM distribution with smoothing
    if verbose
        println("****Updating control sequence distribution")
    end
    μ_new_array, Σ_new_array = compute_new_distribution(direct_solver, control_sequence_elite_array);
    direct_solver.μ_array = μ_new_array;
    direct_solver.Σ_array = Σ_new_array;
end;

"""
    solve!(direct_solver::CrossEntropyDirectOptimizationSolver,
    problem::FiniteHorizonGenerativeOptimalControlProblem, x_0::Vector{Float64},
    rng::AbstractRNG; use_true_model=false, verbose=true, serial=true)

Given `problem` and `direct_solver` (i.e. a PETS Solver), solve stochastic optimal
control with current state `x_0`.

# Return Values (Ordered)
- `μ_array::Vector{Vector{Float64}}` -- array of means `[μ_0,...,μ_{N-1}]` for
  the final Cross Entropy distribution for the control schedule.
- `Σ_array::Vector{Matrix{Float64}}` -- array of covariance matrices `[Σ_0,...,Σ_{N-1}]`
  for the final Cross Entropy distribution for the control schedule.

# Notes
- Returns an open-loop control policy.
- If `use_true_model` is `true`, the solver uses the true stochastic dynamics model
  defined in `problem.f_stochastic`.
- If `serial` is `true`, Monte Carlo sampling of the Cross Entropy method is serialized
  on a single process. If `false` it is distributed on all the available worker processes.
  We recommend to leave this to `true` as distributed processing can be slower for
  this algorithm.
"""
function solve!(direct_solver::CrossEntropyDirectOptimizationSolver,
                problem::FiniteHorizonGenerativeOptimalControlProblem,
                x_0::Vector{Float64},
                rng::AbstractRNG;
                use_true_model=false, verbose=true, serial=true)
    #setting serial=true by default as it turned out this is much faster than serial=false
    initialize!(direct_solver);
    while direct_solver.iter_current < direct_solver.iter_max
        step!(direct_solver, problem, x_0, rng, use_true_model, verbose, serial)
    end
    return copy(direct_solver.μ_array), copy(direct_solver.Σ_array)
end


text:
# This file is a part of BAT.jl, licensed under the MIT License (MIT).

struct HistogramAsMvDistribution{T, N} <: Distributions.ContinuousMultivariateDistribution
    h::StatsBase.Histogram{<:Real, N}
    edges::NTuple{N, <:AbstractVector{T}}
    cart_inds::CartesianIndices{N, NTuple{N, Base.OneTo{Int}}}

    probabilty_edges::AbstractVector{T}

    μ::AbstractVector{T}
    var::AbstractVector{T}
    cov::AbstractMatrix{T}
end

Base.length(d::HistogramAsMvDistribution{T, N}) where {T, N} = N
Base.size(d::HistogramAsMvDistribution{T, N}) where {T, N} = (N,)
Base.eltype(d::HistogramAsMvDistribution{T, N}) where {T, N} = T

Statistics.mean(d::HistogramAsMvDistribution{T, N}) where {T, N} = d.μ
Statistics.var(d::HistogramAsMvDistribution{T, N}) where {T, N} = d.var
Statistics.cov(d::HistogramAsMvDistribution{T, N}) where {T, N} = d.cov

_np_valshape(d::HistogramAsMvDistribution{T, N}) where {T, N} = ArrayShape{Real}(size(d)...)

function _np_bounds(d::HistogramAsMvDistribution{T, N}) where {T, N}
    left_bounds  = [ first(d.h.edges[i]) for i in 1:N ]
    right_bounds = [  prevfloat(last(d.h.edges[i])) for i in 1:N ]
    HyperRectBounds(T.(left_bounds), T.(right_bounds), hard_bounds)
end

function HistogramAsMvDistribution(h::StatsBase.Histogram{<:Real, N}, T::DataType = Float64) where {N}
    nh = normalize(h)

    probabilty_widths = nh.weights * inv(sum(nh.weights))
    probabilty_edges::Vector{T} = Vector{Float64}(undef, length(h.weights) + 1)
    probabilty_edges[1] = 0
    for (i, w) in enumerate(probabilty_widths)
        v = probabilty_edges[i] + probabilty_widths[i]
        probabilty_edges[i+1] = v > 1 ? 1 : v
    end

    mean = _mean(h)
    var = _var(h, mean = mean)
    cov = _cov(h, mean = mean)

    return HistogramAsMvDistribution{T, N}(
        nh,
        collect.(nh.edges),
        CartesianIndices(nh.weights),
        probabilty_edges,
        mean,
        var,
        cov
    )
end

function _mean(h::StatsBase.Histogram{<:Real, N}; T::DataType = Float64) where {N}
    s_inv::T = inv(sum(h.weights))
    m::Vector{T} = zeros(T, N)
    mps = StatsBase.midpoints.(h.edges)
    cart_inds = CartesianIndices(h.weights)
    for i in cart_inds
        for idim in 1:N
            m[idim] += s_inv * mps[idim][i[idim]] * h.weights[i]
        end
    end
    return m
end

function _var(h::StatsBase.Histogram{<:Real, N}; T::DataType = Float64, mean = StatsBase.mean(h, T = T), ) where {N}
    s_inv::T = inv(sum(h.weights))
    v::Vector{T} = zeros(T, N)
    mps = StatsBase.midpoints.(h.edges)
    cart_inds = CartesianIndices(h.weights)
    for i in cart_inds
        for idim in 1:N
            v[idim] += s_inv * (mps[idim][i[idim]] - mean[idim])^2 * h.weights[i]
        end
    end
    return v
end

function _cov(h::StatsBase.Histogram{<:Real, N}; T::DataType = Float64, mean = StatsBase.mean(h, T = T)) where {N}
    s_inv::T = inv(sum(h.weights))
    c::Matrix{T} = zeros(T, N, N)
    mps = StatsBase.midpoints.(h.edges)
    cart_inds = CartesianIndices(h.weights)
    for i in cart_inds
        for idim in 1:N
            for jdim in 1:N
                c[idim, jdim] += s_inv * (mps[idim][i[idim]] - mean[idim]) * (mps[jdim][i[jdim]] - mean[jdim]) * h.weights[i]
            end
        end
    end
    return c
end

function Distributions._rand!(r::AbstractRNG, d::HistogramAsMvDistribution{T,N}, A::AbstractVector{<:Real}) where {T, N}
    rand!(r, A)
    next_inds::UnitRange{Int} = searchsorted(d.probabilty_edges::Vector{T}, A[1]::T)
    cell_lin_index::Int = min(next_inds.start, next_inds.stop)
    cell_car_index = d.cart_inds[cell_lin_index]
    for idim in Base.OneTo(N)
        i = cell_car_index[idim]
        sub_int = d.edges[idim][i:i+1]
        sub_int_width::T = sub_int[2] - sub_int[1]
        A[idim] = sub_int[1] + sub_int_width * A[idim]
    end
    return A
end

function Distributions._rand!(r::AbstractRNG, d::HistogramAsMvDistribution{T,N}, A::AbstractMatrix{<:Real}) where {T, N}
    Distributions._rand!.((r,), (d,), nestedview(A))
end

function Distributions.pdf(d::HistogramAsMvDistribution{T, N}, x::AbstractArray{<:Real, 1}) where {T, N}
    return @inbounds d.h.weights[StatsBase.binindex(d.h, Tuple(x))...]
end
function Distributions.logpdf(d::HistogramAsMvDistribution{T, N}, x::AbstractArray{<:Real, 1}) where {T, N}
    return log(pdf(d, x))
end
function Distributions._logpdf(d::HistogramAsMvDistribution{T,N}, x::AbstractArray{<:Real, 1}) where {T, N}
    return logpdf(d, x)
end

ext: jl


text:
immutable Newton <: Optimizer
    linesearch!::Function
end
#= uncomment v0.8.0
Newton(; linesearch::Function = LineSearches.hagerzhang!) =
  Newton(linesearch)
=#
function Newton(; linesearch! = nothing, linesearch::Function = LineSearches.hagerzhang!)
    linesearch = get_linesearch(linesearch!, linesearch)
    Newton(linesearch)
end
type NewtonState{T}
    @add_generic_fields()
    x_previous::Array{T}
    g::Array{T}
    f_x_previous::T
    H
    F
    Hd
    s::Array{T}
    @add_linesearch_fields()
end

function initial_state{T}(method::Newton, options, d, initial_x::Array{T})
        n = length(initial_x)
      # Maintain current gradient in gr
      g = Array(T, n)
      s = Array(T, n)
      x_ls, g_ls = Array(T, n), Array(T, n)
      f_x_previous, f_x = NaN, d.fg!(initial_x, g)
      f_calls, g_calls = 1, 1
      H = Array(T, n, n)
      d.h!(initial_x, H)
      h_calls = 1

      NewtonState("Newton's Method",
                  length(initial_x),
                  copy(initial_x), # Maintain current state in state.x
                  f_x, # Store current f in state.f_x
                  f_calls, # Track f calls in state.f_calls
                  g_calls, # Track g calls in state.g_calls
                  h_calls,
                  copy(initial_x), # Maintain current state in state.x_previous
                  g, # Store current gradient in state.g
                  T(NaN), # Store previous f in state.f_x_previous
                  H,
                  copy(H),
                  copy(H),
                  similar(initial_x), # Maintain current search direction in state.s
                  @initial_linesearch()...) # Maintain a cache for line search results in state.lsr
  end

  function update_state!{T}(d, state::NewtonState{T}, method::Newton)
        lssuccess = true
        # Search direction is always the negative gradient divided by
        # a matrix encoding the absolute values of the curvatures
        # represented by H. It deviates from the usual "add a scaled
        # identity matrix" version of the modified Newton method. More
        # information can be found in the discussion at issue #153.
        state.F, state.Hd = ldltfact!(Positive, state.H)
        state.s[:] = -(state.F\state.g)

        # Refresh the line search cache
        dphi0 = vecdot(state.g, state.s)
        LineSearches.clear!(state.lsr)
        push!(state.lsr, zero(T), state.f_x, dphi0)

        # Determine the distance of movement along the search line
        try
            state.alpha, f_update, g_update =
                method.linesearch!(d, state.x, state.s, state.x_ls, state.g_ls, state.lsr,
                                   state.alpha, state.mayterminate)
            state.f_calls, state.g_calls = state.f_calls + f_update, state.g_calls + g_update
        catch ex
            if isa(ex, LineSearches.LineSearchException)
                lssuccess = false
                state.f_calls, state.g_calls = state.f_calls + ex.f_update, state.g_calls + ex.g_update
                state.alpha = ex.alpha
                Base.warn("Linesearch failed, using alpha = $(state.alpha) and exiting optimization.")
            else
                rethrow(ex)
            end
        end

        # Maintain a record of previous position
        copy!(state.x_previous, state.x)

        # Update current position # x = x + alpha * s
        LinAlg.axpy!(state.alpha, state.s, state.x)
        (lssuccess == false) # break on linesearch error
end


text:
using Oceananigans.Advection:
    _advective_momentum_flux_Uu,
    _advective_momentum_flux_Uv,
    _advective_momentum_flux_Vu,
    _advective_momentum_flux_Vv,
    advective_tracer_flux_x,
    advective_tracer_flux_y,
    vertical_vorticity_U,
    vertical_vorticity_V,
    bernoulli_head_U,
    bernoulli_head_V

using Oceananigans.Grids: AbstractGrid
using Oceananigans.Operators: Ax_qᶠᶜᶜ, Ay_qᶜᶠᶜ

#####
##### Momentum flux operators
#####

@inline momentum_flux_huu(i, j, k, grid, advection, solution) =
    @inbounds _advective_momentum_flux_Uu(i, j, k, grid, advection, solution[1], solution[1]) / solution.h[i, j, k]

@inline momentum_flux_hvu(i, j, k, grid, advection, solution) =
    @inbounds _advective_momentum_flux_Vu(i, j, k, grid, advection, solution[2], solution[1]) / ℑxyᶠᶠᵃ(i, j, k, grid, solution.h)

@inline momentum_flux_huv(i, j, k, grid, advection, solution) =
    @inbounds _advective_momentum_flux_Uv(i, j, k, grid, advection, solution[1], solution[2]) / ℑxyᶠᶠᵃ(i, j, k, grid, solution.h)

@inline momentum_flux_hvv(i, j, k, grid, advection, solution) =
    @inbounds _advective_momentum_flux_Vv(i, j, k, grid, advection, solution[2], solution[2]) / solution.h[i, j, k]

#####
##### Momentum flux divergence operators
#####

@inline div_mom_u(i, j, k, grid, advection, solution, formulation) =
    1 / Azᶠᶜᶜ(i, j, k, grid) * (δxᶠᵃᵃ(i, j, k, grid, momentum_flux_huu, advection, solution) +
                                δyᵃᶜᵃ(i, j, k, grid, momentum_flux_hvu, advection, solution))

@inline div_mom_v(i, j, k, grid, advection, solution, formulation) =
    1 / Azᶜᶠᶜ(i, j, k, grid) * (δxᶜᵃᵃ(i, j, k, grid, momentum_flux_huv, advection, solution) +
                                δyᵃᶠᵃ(i, j, k, grid, momentum_flux_hvv, advection, solution))

@inline div_mom_u(i, j, k, grid, advection, solution, ::VectorInvariantFormulation) = (
    + vertical_vorticity_U(i, j, k, grid, advection, solution[1], solution[2])  # Vertical relative vorticity term
    + bernoulli_head_U(i, j, k, grid, advection, solution[1], solution[2]))     # Bernoulli head term

@inline div_mom_v(i, j, k, grid, advection, solution, ::VectorInvariantFormulation) = (
    + vertical_vorticity_V(i, j, k, grid, advection, solution[1], solution[2])  # Vertical relative vorticity term
    + bernoulli_head_V(i, j, k, grid, advection, solution[1], solution[2]))     # Bernoulli head term

# Support for no advection
@inline div_mom_u(i, j, k, grid::AbstractGrid{FT}, ::Nothing, solution, formulation) where FT = zero(FT)
@inline div_mom_v(i, j, k, grid::AbstractGrid{FT}, ::Nothing, solution, formulation) where FT = zero(FT)
@inline div_mom_u(i, j, k, grid::AbstractGrid{FT}, ::Nothing, solution, ::VectorInvariantFormulation) where FT = zero(FT)
@inline div_mom_v(i, j, k, grid::AbstractGrid{FT}, ::Nothing, solution, ::VectorInvariantFormulation) where FT = zero(FT)

#####
##### Mass transport divergence operator
#####

"""
    div_Uh(i, j, k, grid, advection, solution, formulation)

Calculates the divergence of the mass flux into a cell,

    1/Az * [δxᶜᵃᵃ(Δy * uh) + δyᵃᶜᵃ(Δx * vh)]

which will end up at the location `ccc`.
"""
@inline function div_Uh(i, j, k, grid, advection, solution, formulation)
    return 1/Azᶜᶜᶜ(i, j, k, grid) * (δxᶜᵃᵃ(i, j, k, grid, Δy_qᶠᶜᶜ, solution[1]) +
                                     δyᵃᶜᵃ(i, j, k, grid, Δx_qᶜᶠᶜ, solution[2]))
end

@inline div_Uh(i, j, k, grid, advection, solution, formulation::VectorInvariantFormulation) =
        div_Uc(i, j, k, grid, advection, solution, solution.h, formulation)

#####
##### Tracer advection operator
#####

@inline transport_tracer_flux_x(i, j, k, grid, advection, uh, h, c) =
    @inbounds advective_tracer_flux_x(i, j, k, grid, advection, uh, c) / ℑxᶠᵃᵃ(i, j, k, grid, h)

@inline transport_tracer_flux_y(i, j, k, grid, advection, vh, h, c) =
    @inbounds advective_tracer_flux_y(i, j, k, grid, advection, vh, c) / ℑyᵃᶠᵃ(i, j, k, grid, h)

"""
    div_Uc(i, j, k, grid, advection, U, c)

Calculates the divergence of the flux of a tracer quantity c being advected by
a velocity field U = (u, v), ∇·(Uc),

    1/Az * [δxᶜᵃᵃ(Δy * uh * ℑxᶠᵃᵃ(c) / h) + δyᵃᶜᵃ(Δx * vh * ℑyᵃᶠᵃ(c) / h)]

which will end up at the location `ccc`.
"""

@inline function div_Uc(i, j, k, grid, advection, solution, c, formulation)
    return 1/Azᶜᶜᶜ(i, j, k, grid) * (δxᶜᵃᵃ(i, j, k, grid, transport_tracer_flux_x, advection, solution[1], solution.h, c) +
                                     δyᵃᶜᵃ(i, j, k, grid, transport_tracer_flux_y, advection, solution[2], solution.h, c))
end

@inline function div_Uc(i, j, k, grid, advection, solution, c, ::VectorInvariantFormulation)
    return 1/Azᶜᶜᶜ(i, j, k, grid) * (δxᶜᵃᵃ(i, j, k, grid, advective_tracer_flux_x, advection, solution[1], c) +
                                     δyᵃᶜᵃ(i, j, k, grid, advective_tracer_flux_y, advection, solution[2], c))
end

# Support for no advection
@inline div_Uc(i, j, k, grid::AbstractGrid{FT}, ::Nothing, solution, c, formulation) where FT = zero(FT)

@inline u(i, j, k, grid, solution) = @inbounds solution.uh[i, j, k] / ℑxᶠᵃᵃ(i, j, k, grid, solution.h)
@inline v(i, j, k, grid, solution) = @inbounds solution.vh[i, j, k] / ℑyᵃᶠᵃ(i, j, k, grid, solution.h)

"""
    c_div_U(i, j, k, grid, advection, U)

Calculates the product of the tracer concentration c with
the horizontal divergence of the velocity field U = (u, v), c ∇·(U),

    1/Az * [δxᶜᵃᵃ(Δy * uh / h) + δyᵃᶜᵃ(Δx * vh / h]

which will end up at the location `ccc`.
"""
@inline c_div_U(i, j, k, grid, solution, c, formulation) =
    @inbounds c[i, j, k] * 1/Azᶜᶜᶜ(i, j, k, grid) * (δxᶜᵃᵃ(i, j, k, grid, Δy_qᶠᶜᶜ, u, solution) + δyᵃᶜᵃ(i, j, k, grid, Δx_qᶜᶠᶜ, v, solution))

@inline c_div_U(i, j, k, grid, solution, c, ::VectorInvariantFormulation) =
    @inbounds c[i, j, k] * 1/Azᶜᶜᶜ(i, j, k, grid) * (δxᶜᵃᵃ(i, j, k, grid, Δy_qᶠᶜᶜ, solution[1]) + δyᵃᶜᵃ(i, j, k, grid, Δx_qᶜᶠᶜ, solution[2]))

# Support for no advection
@inline c_div_Uc(i, j, k, grid::AbstractGrid{FT}, ::Nothing, solution, c, formulation) where FT = zero(FT)


text:
using CSV
using Plots; pgfplots()
using DataFrames

df = DataFrame(x = 1:10, y1 = 1:10, y2 = (1:10).^2)
plot(df.x, df.y1)
plot!(df.x, df.y2)

savefig("tikz/example_figure.tex")
CSV.write("data/data.csv", df)

max_stars_repo_name: kmundnic/pgfplots-tutorial

max_issues_repo_name: kmundnic/pgfplots-tutorial

max_forks_repo_name: kmundnic/pgfplots-tutorial


text:
module CS2_2
using FaultDetectionTools
using DescriptorSystems
using LinearAlgebra
using Polynomials
using Test
using JLD2

## CS2_2  - Case-study example: Monitoring air data sensor faults
#           Robust least order LPV synthesis
println("Case study CS2_2")

## Part 1 - Model setup
# load matrices of the aircraft multiple-model SYSACM, actuator model ACT,
# output-feedback gain K, percentages of mass variations massi

cd(joinpath(pkgdir(FaultDetectionTools), "Examples"))
@load "cs2data.jld2"

N = size(SYSACSM,1)
p, m = size(SYSACSM[1].D)
T = Float64

# build minimal realizations of AC-models with actuator faults
# set dimensions
nom = 7;             # nominal system
# set sensor indices and set dimensions of inputs
#     [ AoA VCAS ]
sen = [  2,    4  ];
md = 2; mu = m-md; mf = length(sen); n = maximum(order.(SYSACSM));

idm = eye(T,p);

syssen = similar(Vector{DescriptorStateSpace},N)
for i = 1:N
    syssen[i] = [ SYSACSM[i] idm[:,sen]]
end

syssenf = fdimodset(syssen; c = Vector(1:mu), d = mu .+ Vector(1:md), f = (mu+md) .+ (1:mf));

## Part 2 - Setup of the synthesis specifications

# compute the achievable strong specifications for constant faults
S_strong = fdigenspec(syssenf[nom], atol = 1.e-6, atol3 = 1.e-6, FDtol = 1.e-5, FDGainTol = 0.001, FDfreq = 0, sdeg = -0.05);

@test size(S_strong,1) == 2^mf-1

## Part 3 - LPV synthesis using up to 4th order polynomial gain interpolation

# form extended system [ Gu Gd Gf; I 0 0] with outputs extended with u
# and determine constant solutions based on Remark 8.1 of the FDI book
syse = similar(Vector{DescriptorStateSpace},N)
D = zeros(T,mf,p+mu,N)
Dref = [zeros(T,mf,n+mu+md) eye(T,mf)]
for i = 1:N
    syse[i] = [syssenf[i].sys;eye(mu,mu+md+mf)]
    cde = [syse[i].C syse[i].D]
    D[:,:,i] = Dref/cde
end

N4 = similar(Matrix{Float64},N,5)
Qpol = Polynomial.(zeros(T,mf,p+mu))
for deg = 0:4
    for i = 1:mf
        for j = 1:p+mu
            Qpol[i,j] = fit(view(massi,:), view(D,i,j,:), deg)
        end
    end
    for i = 1:N
        rezi = (massi[i] .|> Qpol)*syse[i] - [zeros(mf,mu+md) eye(2)]
        N4[i,deg+1]  = glinfnorm(rezi)[1]
    end
end
@test all(maximum(N4,dims = 1) .> [545.  138.  16.  0.88  0.12])

println(" Table 8.5  Robustness analysis results of interpolation based approximations")
println(" N0        N1         N2         N3         N4")
display(N4)

## Part 4 - LPV synthesis using 3rd order polynomial gain interpolation
#  setup of the LPV synthesis similar to using MATLAB's SYSTUNE (to be implemented)

end  # module

max_stars_repo_name: andreasvarga/FaultDetectionTools

max_issues_repo_name: andreasvarga/FaultDetectionTools

max_forks_repo_name: andreasvarga/FaultDetectionTools


text:
#Quantities base types & core functions
#-------------------------------------------------------------------------------

#==Main data structures
===============================================================================#
abstract type Unit{Symbol} end #Think this might be too much

abstract type QuantityType end
#==TODO: How to limit Units to supported units?
It would be nice to flag an error when people write :M, instead of :m...
==#
const Meter = Unit{:m}
const Kilogram = Unit{:kg}
const Inch = Unit{:in}
const Radians = Unit{:rad}
const Degrees = Unit{:deg}
const RadiansPerSecond = Unit{:rad_s}
const Hertz = Unit{:Hz}

struct Quantity{T<:Number, QuantityType, Unit} #<: Number
	v::T
end
value(q::Quantity) = q.v

#==NOTE:
It would probably be bad to inherit Quantity from Number(Quantity <: Number)
If that happens, how would Julia deal with the recursive structure:
    Quantity{Quantity{Quantity...}} ?
==#

#==Specific quantities
===============================================================================#
abstract type AbstractLengthType <: QuantityType end
abstract type LengthType <: AbstractLengthType end
abstract type DistType <: AbstractLengthType end
#==AbstractLength: Meter is known as a "unit of length"... even though the
actual thing measured might be a length, distance, position, ....
So, "AbstractLength" refers to any quantity measured in "unit of length".
==#

#==Useful aliases
===============================================================================#
const _Length{T<:Number, U<:Unit} = Quantity{T,LengthType,U}
const _Dist{T<:Number, U<:Unit}   = Quantity{T,DistType,U}

#==Problem:
 -"type Length" permits you to create constructor function "Length",
  but "typealias Length" *conflicts* with constructor function "Length"
  (TODO: Can this behaviour be changed in Julia?)
==#

#==Constructor functions
===============================================================================#
Length{T<:Number}(v::T, u::Symbol) = Quantity{T,LengthType, Unit{u}}(v)
Dist{T<:Number}(v::T, u::Symbol) = Quantity{T,DistType, Unit{u}}(v)

#==Support basic math operations
===============================================================================#
#Operators that interoperate with other (same) Quantities
const _operators_self = Symbol[:-, :+]

#Operators that interoperate with Number-s
const _operators_num = Symbol[:/, :*]

#const _operators = vcat(_operators_num, _operators_self)
#_dotop(x)=Symbol(".$x")

for op in _operators_self; @eval begin

#Quantity op Quantity
function Base.$op{T1<:Number,T2<:Number,QT<:QuantityType,U<:Unit}(q1::Quantity{T1,QT,U}, q2::Quantity{T2,QT,U})
	return Quantity{promote_type(T1,T2),QT,U}($op(q1.v, q2.v))
end

end; end

for op in _operators_num; @eval begin

#Quantity op Number
function Base.$op{TN<:Number,T<:Number,QT<:QuantityType,U<:Unit}(q::Quantity{T,QT,U}, n::TN)
	return Quantity{promote_type(T,TN),QT,U}($op(q.v, n))
end

#Number op Quantity
function Base.$op{TN<:Number,T<:Number,QT<:QuantityType,U<:Unit}(n::TN, q::Quantity{T,QT,U})
	return Quantity{promote_type(T,TN),QT,U}($op(n, q.v))
end

end; end

#Last line


text:
struct JuliaNotebookInlineDisplay <: AbstractDisplay end

# supported MIME types for inline display, in descending order
# of preference (descending "richness")
const supported_mime_types = [
    "application/vnd.dataresource+json",
    ["application/vnd.vegalite.v$n+json" for n in 4:-1:1]...,
    ["application/vnd.vega.v$n+json" for n in 5:-1:2]...,
    "application/vnd.plotly.v1+json",
    "text/html",
    "text/latex",
    "image/svg+xml",
    "image/png",
    "image/jpeg",
    "text/plain",
    "text/markdown",
    # "application/javascript"
]

for mime in supported_mime_types
    @eval begin
        function display(::JuliaNotebookInlineDisplay, ::MIME{Symbol($mime)}, x)
            IJuliaCore.flush_all() # so that previous stream output appears in order
            payload = IJuliaCore.limitstringmime(MIME($mime), x)
            JSONRPC.send_notification(conn_endpoint[], "notebook/display", Dict("items" => [Dict{String,Any}("mimetype" => $mime, "data" => payload)]))
        end
        displayable(d::InlineDisplay, ::MIME{Symbol($mime)}) = true
    end
end

function Base.display(::JuliaNotebookInlineDisplay, x)
    things_to_show = IJuliaCore.display_dict(x)

    JSONRPC.send_notification(conn_endpoint[], "notebook/display", Dict("items" => reverse([Dict{String,Any}("mimetype" => k, "data" => v) for (k, v) in things_to_show])))
end

max_stars_repo_name: novitk/julia-vscode

max_issues_repo_name: novitk/julia-vscode

max_forks_repo_name: novitk/julia-vscode


text:
#=
# MNIST

We begin by importing the required packages.
We load MNIST via the MLDatasets.jl package.
=#

import Makie
import CairoMakie
import MLDatasets
import Flux
import RestrictedBoltzmannMachines as RBMs
import ConvolutionalRBMs as ConvRBMs
using Statistics: mean, var, std
using ValueHistories: MVHistory
using Random: bitrand
using RestrictedBoltzmannMachines: visible, hidden, weights, log_pseudolikelihood, transfer_sample
nothing #hide

#=
Useful function to plot MNIST digits.
=#

"""
    imggrid(A)

Given a four dimensional tensor `A` of size `(width, height, ncols, nrows)`
containing `width x height` images in a grid of `nrows x ncols`, this returns
a matrix of size `(width * ncols, height * nrows)`, that can be plotted in a heatmap
to display all images.
"""
function imggrid(A::AbstractArray{<:Any,4})
    (width, height, ncols, nrows) = size(A)
    return reshape(permutedims(A, (1,3,2,4)), width * ncols, height * nrows)
end

imggrid_border(A::AbstractArray{<:Any,4}, borderval = 1) = imggrid(bordered(A, borderval))

function bordered(A::AbstractArray{<:Any,4}, borderval = 1)
    return mapslices(A; dims=(1,2)) do img::AbstractMatrix
        [   fill(borderval, 1, size(img, 2) + 2);
            fill(borderval, size(img, 1)) img fill(borderval, size(img, 1));
            fill(borderval, 1, size(img, 2) + 2)
        ]
    end
end

#=
Load MNIST dataset.
=#

Float = Float32
train_x, train_y = MLDatasets.MNIST.traindata()
digit = 2
train_x = Array{Float}(train_x[:, :, train_y .== digit] .≥ 0.5)
train_y = train_y[train_y .== digit]
println(length(train_y), " training images (for digit = $digit)")
nothing #hide

# Reshape for convolutional input

train_x = reshape(train_x, 1, 28, 28, :) # channel dims, input dims, batch dims
nothing #hide

# Initialize the convolutional RBM.

rbm = ConvRBMs.BinaryConvRBM(Float, 1, 16, (15,15); pad=:same, pool=true)
RBMs.initialize!(rbm, train_x)
nothing #hide

# Pseudolikelihood before training

idx = rand(1:size(train_x)[end], 256)
mean(@time log_pseudolikelihood(rbm, train_x[:,:,:,idx]))

# Initialize training

batchsize = 256
optim = Flux.ADAM()
vm = transfer_sample(visible(rbm), falses(1, 28, 28, batchsize)) # fantasy chains
history = MVHistory()
nothing #hide

# Train!

@time for iter in 1:20
    ConvRBMs.pcd!(rbm, train_x; vm, history, batchsize, optim, epochs=5)
    lpl = log_pseudolikelihood(rbm, train_x[:, :, :, rand(1:size(train_x)[end], 1024)])
    push!(history, :lpl_ave, mean(lpl))
    push!(history, :lpl_std, std(lpl))
end
nothing #hide

# Plot of log-pseudolikelihood of trian data during learning.

fig = Makie.Figure(resolution=(600,300))
ax = Makie.Axis(fig[1,1], xlabel = "train time", ylabel="pseudolikelihood")
Makie.band!(ax, get(history, :lpl_ave)[1],
    get(history, :lpl_ave)[2] - get(history, :lpl_std)[2]/2,
    get(history, :lpl_ave)[2] + get(history, :lpl_std)[2]/2,
    color=:lightblue
)
Makie.lines!(ax, get(history, :lpl_ave)..., color=:blue)
fig

#=
Now let's generate some random RBM samples.
=#

nrows, ncols = 10, 15
nsteps = 1000
fantasy_F = zeros(nrows*ncols, nsteps)
fantasy_x = bitrand(1,28,28,nrows*ncols)
fantasy_F[:,1] .= RBMs.free_energy(rbm, fantasy_x)
@time for t in 2:nsteps
    fantasy_x .= RBMs.sample_v_from_v(rbm, fantasy_x)
    fantasy_F[:,t] .= RBMs.free_energy(rbm, fantasy_x)
end
nothing #hide

# Check equilibration of sampling

fig = Makie.Figure(resolution=(400,300))
ax = Makie.Axis(fig[1,1], xlabel="sampling time", ylabel="free energy")
fantasy_F_μ = vec(mean(fantasy_F; dims=1))
fantasy_F_σ = vec(std(fantasy_F; dims=1))
Makie.band!(ax, 1:nsteps, fantasy_F_μ - fantasy_F_σ/2, fantasy_F_μ + fantasy_F_σ/2)
Makie.lines!(ax, 1:nsteps, fantasy_F_μ)
fig

# Plot the resulting samples.

fig = Makie.Figure(resolution=(40ncols, 40nrows))
ax = Makie.Axis(fig[1,1], yreversed=true)
Makie.image!(ax, imggrid(reshape(fantasy_x, 28, 28, ncols, nrows)), colorrange=(1,0))
Makie.hidedecorations!(ax)
Makie.hidespines!(ax)
fig

# Plot the filters learned

wncols = 4; wnrows = 4
fig = Makie.Figure(resolution=(80wncols, 80wnrows))
ax = Makie.Axis(fig[1,1], yreversed=true)
Makie.image!(ax, imggrid_border(reshape(rbm.w ./ maximum(abs, rbm.w; dims=(2,3)), 15, 15, wncols, wnrows)))
Makie.hidedecorations!(ax)
Makie.hidespines!(ax)
fig


text:
"""
A node in a `Tree`.

# Parameters
* **label** The character or state that the node represents
* **branches** A `Dict` of all this node's children, indexed by what their
    character label is
* **possibleWords** The number of words that can still be formed by traversing
the current branch
"""
mutable struct Node
    label
    branches::Dict
    possibleWords::Int64

    Node(label) = new(label, Dict(), 0)
end

"""
A tree to contain all the possible paths through a language. It is constructed
such that any path that can be drawn by descending the tree will produce an
observed item in the tree.

# Parameters
* **root** `Node` that represents the top of the tree and contains whatever
    branches and leaves it has

# Constructors
```julia
Tree()
Tree(corpus)
```

# Argumentsprintln("help")println("help")println("help")
* **corpus** The corpus of items to add to the tree
"""
mutable struct Tree
    root

    Tree() = new(Node(""))

    function Tree(corpus)
        tree = new(Node(""))
        for item in corpus
            addItem!(tree, item)
        end

        return tree
    end
end

"""
    addItem!(tree, item)

Adds `item` to the specified `Tree` `tree`.

# Parameters
* **tree** The `Tree` to add `item` to
* **item** The item to add to `tree`

# Side-effects
* After running, `tree` will contain `item`
"""
function addItem!(tree::Tree, item)
    curr = tree.root
    curr.possibleWords += 1
    for ch in item
        if ! haskey(curr.branches, ch)
            curr.branches[ch] = Node(ch)
        end
        curr = curr.branches[ch]
        curr.possibleWords += 1
    end
end

"""
    upt(corpus, queries; [inCorpus=true])

Calculates the phonological uniqueness point (upt) the items in `queries` based
on the items in `corpus`. If the items are expected to be in the corpus, this
function will calculate the uniqueness point to be when a branch can be
considered to only represent 1 word. If the items are not expected to be in the
corpus, the uniqueness point will be taken to be the depth at which the tree
can no longer be traversed.

# Parameters
* **corpus** The items comprising the corpus to compare against when calculating
    the uniqueness point of each query
* **queries** The items for which to calculate the uniqueness point
* **inLexicon** Whether the query items are expected to be in the corpus or not

# Returns
* A `DataFrame` with the queries in the first column and the uniqueness points
    in the second
"""
function upt(corpus::Array, queries::Array; inCorpus=true)

    tree = Tree(corpus)

    results = Vector()

    for query in queries
        depth = 1
        curr = tree.root
        while depth <= length(query) && haskey(curr.branches, query[depth])
            curr = curr.branches[query[depth]]
            if inCorpus
                if length(curr.branches) <= 1 && curr.possibleWords == 1
                    break
                end
            end
            depth += 1
        end

        push!(results, depth)
    end

    return DataFrame(Query=queries, UPT=results)
end


text:
module FEA

using ..GPUUtils, ..TopOptProblems, ..Utilities
using JuAFEM, Setfield, TimerOutputs, Preconditioners
using IterativeSolvers, CuArrays, StaticArrays
using LinearAlgebra, GPUArrays
import CUDAdrv
using Parameters: @unpack

const dev = CUDAdrv.device()
const ctx = CUDAdrv.CuContext(dev)

export  AbstractFEASolver,
        AbstractDisplacementSolver,
        DirectDisplacementSolver,
        PCGDisplacementSolver,
        StaticMatrixFreeDisplacementSolver,
        Displacement,
        Direct,
        CG,
        MatrixFree,
        FEASolver,
        Assembly,
        DefaultCriteria,
        EnergyCriteria,
        simulate

const to = TimerOutput()

# FEA solvers
abstract type AbstractFEASolver end

include("matrix_free_operator.jl")
include("convergence_criteria.jl")
include("direct_displacement_solver.jl")
include("assembly_cg_displacement_solvers.jl")
include("matrix_free_cg_displacement_solvers.jl")
include("matrix_free_apply_bcs.jl")
include("simulate.jl")
include("solvers_api.jl")

getcompliance(solver) = solver.u' * solver.globalinfo.K * solver.u

for T in (:DirectDisplacementSolver, :PCGDisplacementSolver)
    @eval @inline CuArrays.cu(p::$T) = error("$T does not support the GPU.")
end

end


text:
function run_lp_model_tests()
    model = MetNets.toy_model()
    M, N = size(model)
    lp_model = MetLP.build_lp_model(model)

    # bounds are correct
    @test all(MetLP._get_con_rhs(lp_model, MetLP._LB_CON_KEY) .== model.lb)
    @test all(MetLP._get_con_rhs(lp_model, MetLP._UB_CON_KEY) .== model.ub)

    # test bound reset
    MetLP.set_lb_con!(lp_model, zeros(N))
    @test all(MetLP._get_con_rhs(lp_model, MetLP._LB_CON_KEY) .== 0.0)
    MetLP.set_ub_con!(lp_model, ones(N))
    @test all(MetLP._get_con_rhs(lp_model, MetLP._UB_CON_KEY) .== 1.0)

end
run_lp_model_tests();


text:
function main(::Robot)
    j = 2
    while !ismarker(r)
      if !ismarker(r)
        move!(r, Nord)
      end
      if !ismarker(r)
        move!(r,West)
      end
      spiral!(r, j)
      j += 2

    end
end

function spiral!(::Robot, j::Int)
    help_moves!(r, j, Ost)
    help_moves!(r, j, Sud)
    help_moves!(r, j, West)
    help_moves!(r, j, Nord)
end

function help_moves!(::Robot, j::Int, side::HorizonSide)
    for i in 1:j
      if !ismarker(r)
        move!(r, side)
      end
    end
end

max_stars_repo_name: VladislavBukin/mirea-progs

max_issues_repo_name: VladislavBukin/mirea-progs

max_forks_repo_name: VladislavBukin/mirea-progs


text:
module Brownian

using Distributions
using PDMats
using StatsBase

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

export
  StochasticProcess,
  UnivariateStochasticProcess,
  MultivariateStochasticProcess,
  DiscreteStochasticProcess,
  ContinuousStochasticProcess,
  #DiscreteUnivariateStochasticProcess,
  ContinuousUnivariateStochasticProcess,
  #DiscreteMultivariateStochasticProcess,
  ContinuousMultivariateStochasticProcess,
  BrownianMotion,
  #MvBrownianMotion,
  FBM,
  FGN,
  convert,
  autocov!,
  autocov,
  rand!,
  rand

include("StochasticProcess.jl")
include(joinpath("univariate", "BrownianMotion.jl"))
include(joinpath("univariate", "FBM.jl"))
#include(joinpath("multivariate", "MvBrownianMotion.jl"))

end # module


text:
function read_xes(path)
    xml_doc = EzXML.readxml(path)
    eventlog = create_eventlog(xml_doc)
    return eventlog
end

function create_eventlog(xml_doc)
    traces = extract_traces(xml_doc)
    event_classifiers = extract_event_classifiers(xml_doc)
    eventlog = EventLog(
        traces,
        event_classifiers
    )
    return eventlog
end

function extract_event_classifiers(xml_doc)
    namespace = get_namespace(xml_doc)
    xpath_expression = "./ns:classifier"
    event_classifier_nodes = EzXML.findall(
        xpath_expression,
        xml_doc.root,
        ["ns"=>namespace]
    )
    event_classifiers = Dict()
    for node in event_classifier_nodes
        key = node["name"]
        value = node["keys"]
        event_classifiers[key] = value
    end
    return event_classifiers
end

function extract_traces(xml_doc)
    namespace = get_namespace(xml_doc)
    xpath_expression = "./ns:trace"
    trace_nodes = EzXML.findall(
        xpath_expression,
        xml_doc.root,
        ["ns"=>namespace]
    )
    traces = [create_trace(t, namespace) for t in trace_nodes]
    return traces
end

function get_namespace(xml_doc)
    try
        return EzXML.namespace(xml_doc.root)
    catch Exception
        return DEFAULT_NAMESPACE
    end
end

function create_trace(trace_node, namespace)
    metadata = extract_metadata(trace_node)
    xpath_expression = "./ns:event"
    event_nodes = EzXML.findall(
        xpath_expression,
        trace_node,
        ["ns"=>namespace]
    )
    name = pop_dict!(metadata, "concept:name")
    id = pop_dict!(metadata, "identity:id")
    events = [create_event(e) for e in event_nodes]
    trace = Trace(name, id, metadata, events)
    return trace
end

function extract_metadata(trace_node)
    metadata_nodes = [
        node
        for node in EzXML.elements(trace_node)
        if node.name != "event"
    ]
    metadata = Dict()
    for node in metadata_nodes
        key = node["key"]
        value = node["value"]
        metadata[key] = value
    end
    return metadata
end

function create_event(event_node)
    event_attribute_nodes = EzXML.elements(event_node)
    event_attributes = Dict()
    for attribute in event_attribute_nodes
        key = attribute["key"]
        value = attribute["value"]
        event_attributes[key] = value
    end
    name = pop_dict!(event_attributes, "concept:name")
    timestamp = pop_dict!(event_attributes, "time:timestamp")
    id = pop_dict!(event_attributes, "identity:id")
    instance = pop_dict!(event_attributes, "concept:instance")
    transition = pop_dict!(event_attributes, "lifecycle:transition")
    resource = pop_dict!(event_attributes, "org:resource")
    role = pop_dict!(event_attributes, "org:role")
    group = pop_dict!(event_attributes, "org:group")
    event = Event(
        name, timestamp, id,
        instance, transition,
        resource, role, group,
        event_attributes
    )
    return event
end

function pop_dict!(dict::Dict, key::String)
    if key in keys(dict)
        id = dict[key]
        delete!(dict, key)
    else
        id = "NA"
    end
    return id
end

# TODO: add streaming option for large files


text:
module NEAT

using Cambrian
using YAML
using JSON
using Random
using LinearAlgebra

include("config.jl")
include("individual.jl")
include("process.jl")
include("mutation.jl")
# include("crossover.jl")

end


text:
using .AMDGPU

struct ROCArrayProc <: Dagger.Processor
    device
end

@gpuproc(ROCArrayProc, ROCArray)

Dagger.execute!(proc::ROCArrayProc, func, args...) = func(args...)

processor(::Val{:ROC}) = ROCArrayProc
cancompute(::Val{:ROC}) = AMDGPU.configured
# FIXME: kernel_backend(::ROCDevice) = ROCArrayProc

if AMDGPU.configured
    Dagger.add_callback!(proc -> begin
        return ROCArrayProc(AMDGPU.get_default_agent())
    end)
end


text:
# Check of handedness
@test rotationmatrix_normalized(qrotation_x(π/2)) * [0, 1, 0] ≈ [0, 0, 1]
@test rotationmatrix_normalized(qrotation_y(π/2)) * [1, 0, 0] ≈ [0, 0, -1]
@test rotationmatrix_normalized(qrotation_z(π/2)) * [1, 0, 0] ≈ [0, 1, 0]

# Simple composition of rotations
@test qrotation_x(π/3) * qrotation_x(π/3) ≈ qrotation_x(2π/3)
@test qrotation_y(π/3) * qrotation_y(π/3) ≈ qrotation_y(2π/3)
@test qrotation_z(π/3) * qrotation_z(π/3) ≈ qrotation_z(2π/3)

# More complex compositions
comp = qrotation_y(π/2) * qrotation_x(π/2)
@test rotationmatrix_normalized(comp) * [0, 1, 0] ≈ [1, 0, 0]

comp = qrotation_z(π/2) * qrotation_y(π/2)
@test rotationmatrix_normalized(comp) * [0, 0, 1] ≈ [0, 1, 0]


text:
module BinaryTreeTests

    using PhylogeneticTrees
    @static if VERSION < v"0.7.0-DEV.2005"
        using Base.Test
    else
        using Test
    end

    println("=================================")
    println("Testing BinaryTree data structure")
    println("=================================")

    @testset "Tree construction" begin
        bt4 = PhylogeneticTrees.BinaryTree(4)
        @test bt4.depth == 4
        @test bt4.nnodes == 31
        @test length(bt4.edges) == 30
        @test !PhylogeneticTrees.validnode(bt4,0)
        @test PhylogeneticTrees.validnode(bt4,31)
        @test PhylogeneticTrees.validlayer(bt4,0)
        @test PhylogeneticTrees.validlayer(bt4,4)
        @test PhylogeneticTrees.getlayer(bt4,1) == 0
        @test PhylogeneticTrees.getlayer(bt4,31) == 4
        @test PhylogeneticTrees.getnodes(bt4,0) == 1:1
        @test PhylogeneticTrees.getnodes(bt4,4) == 16:31
        @test all(PhylogeneticTrees.getchildren(bt4,1) .== [2,3])
        @test length(PhylogeneticTrees.getdescendants(bt4,1)) == 30
        @test all(PhylogeneticTrees.getdescendants(bt4,15) .== [30,31])
        @test length(PhylogeneticTrees.getdescendants(bt4,31)) == 0
        @test PhylogeneticTrees.getparent(bt4,31) == 15
        @test all(bt4.codes[1] .== [-1])
        @test all(bt4.codes[2] .== [0])
        @test all(bt4.codes[3] .== [1])
        @test all(bt4.codes[29] .== [1,1,0,1])
        @test all(bt4.codes[30] .== [1,1,1,0])
        @test all(bt4.codes[31] .== [1,1,1,1])
        @test length(unique(bt4.codes)) == bt4.nnodes
    end

    @testset "Path edges" begin
        bt1 = PhylogeneticTrees.BinaryTree(1)
        @test bt1.pathedges[(1,2)] == [(1, 2)]
        @test all(bt1.pathedges[(2,3)] .== [(1, 2), (1,3)])
        @test length(bt1.pathedges[(1,1)]) == 0
    end

end


text:
module TestWithSequential

using Adapt
using CUDA
using Folds.Testing: parse_tests, test_with_sequential
using FoldsCUDA
using Transducers

rawdata = """
count(isodd(x) for x in 1:10)
extrema((x - 5)^2 for x in 1:10)
mapreduce(x -> x^2, +, 1:10)
mapreduce(*, +, 1:10, 11:20)
mapreduce(*, +, 1:10, 11:20, 21:30)
maximum(0:9)
maximum(9:-1:0)
maximum([2, 3, 0, 3, 4, 0, 5, 7, 4, 2])
minimum(0:9)
minimum(9:-1:0)
minimum([2, 3, 0, 3, 4, 0, 5, 7, 4, 2])
prod(1:2:10)
reduce(TeeRF(min, max), 1:10)
reduce(TeeRF(min, max), (2x for x in 1:10))
reduce(TeeRF(min, max), (2x for x in 1:10 if isodd(x)); init = (typemax(Int), typemin(Int)))
sum(1:10)
sum(1:2^30)
sum(1:2^30; init = 0)
sum(first, pairs(1:10))
sum(last, pairs(1:10))
sum(x^2 for x in 1:11)
sum(x^2 for x in 1:11 if isodd(x); init = 0)
sum(y for x in 1:11 if isodd(x) for y in 1:x:x^2; init = 0)
"""

tests = parse_tests(rawdata, @__MODULE__)

upload(x) = adapt(CuArray, x)

test() = test_with_sequential(tests, [CUDAEx()]; upload = upload)

end  # module


text:
"""
    slaterkoster(; kwargs...)

Create a function to be passed to `hamiltonian` or `overlap` for
calculating the matrix elements using the Slater-Koster scheme.

"""

slaterkoster(; kwargs...) = (args...) -> slaterkoster(args...; kwargs...)

function slaterkoster(r, a1; kwargs...)
    return kwargs[Symbol(first(string(last(a1))))]
end

function slaterkoster(r, a1, a2; kwargs...)
    o1, o2 = last(a1), last(a2)
    d = r ./ norm(r)
    s(x) = kwargs[Symbol(first(a1), first(a2), "_", first(string(last(a1))), first(string(last(a2))), x)]

    if (o1, o2) == (:s, :s)
        return s(:σ)
    end

    for (i, x) in [(1, :x), (2, :y), (3, :z)]
        if (o1, o2) == (:s, Symbol(:p, x))
            return d[i]*s(:σ)
        end
        if (o1, o2) == (Symbol(:p, x), :s)
            return d[i]*s(:σ)
        end

        for (j, y) in [(1, :x), (2, :y), (3, :z)]
            if (o1, o2) == (Symbol(:p, x), Symbol(:p, x))
                return d[i]^2*s(:σ) + (1 - d[i]^2)*s(:π)
            end
            if (o1, o2) == (Symbol(:p, x), Symbol(:p, y))
                return d[i]*d[j]*(s(:σ) - s(:π))
            end
        end
    end

    error("orbitals not supported: $a1, $a2")
end


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

"""
    Some{T}

A wrapper type used in `Union{Some{T}, Nothing}` to distinguish between the absence
of a value ([`nothing`](@ref)) and the presence of a `nothing` value (i.e. `Some(nothing)`).

Use [`coalesce`](@ref) to access the value wrapped by a `Some` object.
"""
struct Some{T}
    value::T
end

promote_rule(::Type{Some{S}}, ::Type{Some{T}}) where {S,T} = Some{promote_type(S, T)}
promote_rule(::Type{Some{T}}, ::Type{Nothing}) where {T} = Union{Some{T}, Nothing}

convert(::Type{Some{T}}, x::Some) where {T} = Some{T}(convert(T, x.value))
convert(::Type{Union{Some{T}, Nothing}}, x::Some) where {T} = convert(Some{T}, x)

convert(::Type{Union{T, Nothing}}, x::Any) where {T} = convert(T, x)
convert(::Type{Nothing}, x::Any) = throw(MethodError(convert, (Nothing, x)))
convert(::Type{Nothing}, x::Nothing) = nothing

function show(io::IO, x::Some)
    if get(io, :typeinfo, Any) == typeof(x)
        show(io, x.value)
    else
        print(io, "Some(")
        show(io, x.value)
        print(io, ')')
    end
end

"""
    coalesce(x, y...)

Return the first value in the arguments which is not equal to
either [`nothing`](@ref) or [`missing`](@ref), or the last argument.
Unwrap arguments of type [`Some`](@ref).

# Examples

```jldoctest
julia> coalesce(nothing, 1)
1

julia> coalesce(missing, 1)
1

julia> coalesce(1, nothing)
1

julia> coalesce(nothing, nothing)
nothing

julia> coalesce(Some(1))
1

julia> coalesce(nothing, Some(1))
1
```
"""
function coalesce end

coalesce(x::Any) = x
coalesce(x::Some) = x.value
coalesce(x::Nothing) = nothing
coalesce(x::Missing) = missing
coalesce(x::Any, y...) = x
coalesce(x::Some, y...) = x.value
coalesce(x::Union{Nothing, Missing}, y...) = coalesce(y...)

"""
    notnothing(x)

Throw an error if `x == nothing`, and return `x` if not.
"""
notnothing(x::Any) = x
notnothing(::Nothing) = throw(ArgumentError("nothing passed to notnothing"))

max_stars_repo_name: YingboMa/julia

max_issues_repo_name: YingboMa/julia

max_forks_repo_name: YingboMa/julia


text:
using RoadRunner
using Test

ant_str = """
    const S1
    S1 + S2 -> S3; k1*S1*S2;
    S3 -> 2 S4; k2*S3-k3*S4
    k1 = 0.1; k2 = 0.2; k3 = 0.3
    S1 = 10; S2 = 2.5; S3 = 3.4; S4 = 0
"""

rr = RoadRunner.loada(ant_str)

@testset "compartment" begin
    @test RoadRunner.getNumberOfCompartments(rr) == 1
end

@testset "reaction" begin
    @test RoadRunner.getNumberOfReactions(rr) == 2
    @test RoadRunner.getNumberOfRules(rr) == 0
    @test RoadRunner.getReactionIds(rr) == ["_J0", "_J1"]
end

@testset "parameters" begin
    @test RoadRunner.getNumberOfGlobalParameters(rr) == 3
    @test RoadRunner.getGlobalParameterIds(rr) == ["k1", "k2", "k3"]
    @test RoadRunner.getGlobalParameterValues(rr) == [0.1, 0.2, 0.3]
end

@testset "species" begin
    @test RoadRunner.getNumberOfFloatingSpecies(rr) == 3
    @test RoadRunner.getFloatingSpeciesIds(rr) == ["S2", "S3", "S4"]
    @test RoadRunner.getNumberOfBoundarySpecies(rr) == 1
    @test RoadRunner.getBoundarySpeciesIds(rr) == ["S1"]
    @test RoadRunner.getFloatingSpeciesInitialConcentrationByIndex(rr, 0) == 2.5
    @test RoadRunner.getFloatingSpeciesInitialConcentrationByIndex(rr, 1) == 3.4
end


text:
using GeometryRandomFields
using Base.Test

gf = GaussianCovarianceFunction(10,1)
@test gf(0) == 1
@test_approx_eq gf(10) exp(-1/2)

mf = MatérnCovarianceFunction(1/2,10,1)
@test mf(0) == 1
@test_approx_eq mf(10) exp(-1)

m,n = 1024,512
M = grf(m,n,gf)
@test size(M) == (m,n)

GeometryRandomFields.acf(real(M))

Z1,Z2 = reim(M)

@test Γ(Z1,10) == 0


text:
# module escape_dimer_one_hit
# export escape_exit_num_dimer
using DifferentialEquations
using Roots
using Printf
# using MATLAB
# using ProgressMeter

include("leap_frog_definitions.jl")
max_hit_q1 = 1
barrier = 10

condition_max_hits(u, t, integrator) = u[3] ==max_hit_q1 || u[6]>max_hit_p1 || maximum([abs(u[1]),abs(u[2]),abs(u[4]),abs(u[5])]) > barrier
affect_stop!(integrator) = terminate!(integrator)
function condition_hits_PSS_q1(u, t, integrator) # Event when event_f(u,t) == 0
    u[1]
end

function affect_update_iterator_q1!(integrator)
    integrator.u[3] = integrator.u[3] + 1

callback_max_hits = DiscreteCallback(condition_max_hits, affect_stop!)
callback_hits_PSS_q1 = ContinuousCallback(condition_hits_PSS_q1, nothing, affect_update_iterator_q1!)

cb = CallbackSet(callback_hits_PSS_q1,  callback_max_hits)

function escape_after_one_hit(point, t_end, H)
    Q2 = point[1]; P2 = point[2];
    P1 = P1_find_dimer(Q2, P2, H)
    Q1=0

    if isempty(P1)
        P1 = P1_find_dimer_second(Q2, P2, H)
        t_end=1
    end

    if ~isempty(P1)
        # println(Q1)
        q0, p0 = [zeros(3) for i in 1:2]
        q0[1] = 0
        q0[2] = Q2
        q0[3] = 0
        p0[1] = P1[1]
        p0[2] = P2
        p0[3] = 0
       #constructor for ODE
        prob = HamiltonianProblem{true}(Hamiltonian_Dimer, q0, p0, (0., t_end));
        #solve ode , save_everystep=false is important to prevent sol to include all points, not just event points
        # sol=solve(prob, RK4(),maxiters=1e20,callback=cb,save_start=true,save_end=false,save_everystep=false)
        # sol=solve(prob, RK4(),maxiters=1e20,reltol=1e-8,abstol=1e-10,callback=cb,save_start=true,save_end=true,save_everystep=false)
        sol=solve(prob, Vern9(),maxiters=1e20, reltol=1e-13,abstol=1e-15,callback=cb,save_start=true,save_end=true,save_everystep=true)

        uf=zeros(6)
        uf[1]=sol[1,end]
        uf[2]=sol[2,end]
        uf[3]=sol[3,end]
        uf[4]=sol[4,end]
        uf[5]=sol[5,end]
        uf[6]=sol[6,end]
        q1=uf[1]
        q2=uf[2]
        p1=uf[4]
        p2=uf[5]
        q=[q1, q2]
        p=[p1, p2]
        diff=abs(Hamiltonian_Dimer(q,p,0)-H)

        prob_sign = HamiltonianProblem{true}(Hamiltonian_Dimer, q0, p0, (0., .001));
        sol_sign=solve(prob_sign, RK4(),maxiters=1e20)
        sign_orientation=sign(sol_sign[1,3])

        # if sign_orientation<0 && diff<1e-10 && uf[3]==1 && maximum([abs(q1),abs(q2),abs(p1),abs(p2)]) > barrier
        #     return Q1, Q2, P1[1], P2 , q1,q2, p1,p2
        # end
        return sol.t, sol[1,:], sol[5,:],uf[3]
        # else

        # end
        # if uf[3]==1
        #     return uf[1], uf[2], uf[4], uf[5]
        # end

        # if dH<1e-5
        #     if is_it(uf,H,tol_dist)
        #         return sol.u[end][5]
        #     else
        #         return -1
        #     end
        # end1233
    # else
        # return NaN, NaN, NaN, NaN
    end
end

function add_four!(point,radius, one_hit_exit, t_end, energy)
    X=point[2]
    Y=point[4]
    X_plus=X+radius
    Y_plus=Y+radius
    X_minus=X-radius
    Y_minus=Y-radius

    point_2=[X, Y_plus]

    point_4=[X_minus, Y]

    point_5=[X_plus, Y]

    point_7=[X, Y_minus]

    new_points= [point_2, point_4, point_5, point_7]
    for point in new_points
        new_point=escape_after_one_hit(point,t_end,energy)
        if new_point!=nothing
            push!(one_hit_exit, new_point)
        end
    end
end

function add_eight!(point,radius, one_hit_exit, t_end, energy)
    X=point[2]
    Y=point[4]

    X_plus=X+radius
    Y_plus=Y+radius
    X_minus=X-radius
    Y_minus=Y-radius

    point_1=[X_minus, Y_plus]
    point_2=[X, Y_plus]
    point_3=[X_minus, Y_plus]

    point_4=[X_minus, Y]
    point_5=[X_plus, Y]

    point_6=[X_minus, Y_minus]
    point_7=[X, Y_minus]
    point_8=[X_minus, Y_minus];

    new_points= [point_1, point_2, point_3, point_4, point_5, point_6, point_7, point_8]
    for point in new_points
        new_point=escape_after_one_hit(point,t_end,energy)
        if new_point!=nothing
            push!(one_hit_exit, new_point)
        end
    end
end

function unzipper(array_of_points)
    N=length(array_of_points)
    Q=zeros(N)
    P=zeros(N)
    for (index,point) in enumerate(array_of_points)
        Q[index]=point[1]
        P[index]=point[2]
    end
    return Q,P
end


text:
module Masks

using ..NeuralAttentionlib: @imexport

@imexport import ..NeuralAttentionlib:
    apply_mask, NaiveAttenMaskOp, GenericAttenMaskOp,
    CausalMask, LocalMask, RandomMask, BandPartMask,
    GenericMask, SymLengthMask, BiLengthMask,
    BatchedMask, RepeatMask, getmask

import ..NeuralAttentionlib: AbstractAttenMask, AbstractDatalessMask, AbstractArrayMask

"""
    AbstractDatalessMask <: AbstractAttenMask

Abstract type for mask without array data.
"""
AbstractDatalessMask

"""
    AbstractArrayMask <: AbstractAttenMask

Abstract type for mask with array data
"""
AbstractArrayMask

"""
    CausalMask() <: AbstractDatalessMask

Attention mask that block the future values.

Similar to applying `LinearAlgebra.triu!` on the score matrix
"""
CausalMask

"""
    LocalMask(width::Int) <: AbstractDatalessMask

Attention mask that only allow local (diagonal like) values to pass.

`width` should be ≥ 0 and `A .* LocalMask(1)` is similar to `Diagonal(A)`
"""
LocalMask

"""
    RandomMask(p::Float64) <: AbstractDatalessMask

Attention mask that block value randomly.

`p` specify the percentage of value to block. e.g. `A .* RandomMask(0)` is equivalent to `identity(A)` and
 `A .* RandomMask(1)` is equivalent to `zero(A)`.
"""
RandomMask

"""
    BandPartMask(l::Int, u::Int) <: AbstractDatalessMask

Attention mask that only allow [band_part](https://www.tensorflow.org/api_docs/python/tf/linalg/band_part)
 values to pass.
"""
BandPartMask

"""
    GenericMask <: AbstractArrayMask

Generic attention mask. Just a wrapper over `AbstractArray{Bool}` for dispatch.
"""
GenericMask

"""
    SymLengthMask(len::AbstractArray{Int, N}) <: AbstractArrayMask

Attention mask specified by an array of integer that indicate the length dimension size.
 assuming *Query* length and *Key* length are the same.

# Example

```julia
julia> m = SymLengthMask([2,3])
SymLengthMask{1, Vector{Int32}}(Int32[2, 3])

julia> trues(3,3, 2) .* m
3×3×2 BitArray{3}:
[:, :, 1] =
 1  1  0
 1  1  0
 0  0  0

[:, :, 2] =
 1  1  1
 1  1  1
 1  1  1

```

See also: [`BiLengthMask`](@ref), [`BatchedMask`](@ref), [`RepeatMask`](@ref)
"""
SymLengthMask

"""
    BiLengthMask(q_len::A, k_len::A) where {A <: AbstractArray{Int, N}} <: AbstractArrayMask

Attention mask specified by two arrays of integer that indicate the length dimension size.

# Example

```julia
julia> bm = BiLengthMask([2,3], [3, 5])
BiLengthMask{1, Vector{Int32}}(Int32[2, 3], Int32[3, 5])

julia> trues(5,5, 2) .* bm
5×5×2 BitArray{3}:
[:, :, 1] =
 1  1  0  0  0
 1  1  0  0  0
 1  1  0  0  0
 0  0  0  0  0
 0  0  0  0  0

[:, :, 2] =
 1  1  1  0  0
 1  1  1  0  0
 1  1  1  0  0
 1  1  1  0  0
 1  1  1  0  0

```

See also: [`SymLengthMask`](@ref), [`BatchedMask`](@ref), [`RepeatMask`](@ref)
"""
BiLengthMask

"""
    !m::AbstractAttenMask

Boolean not of an attention mask
"""
Base.:!(m::AbstractAttenMask)

"""
    m1::AbstractAttenMask | m2::AbstractAttenMask

logical or of two attention mask
"""
Base.:|(m1::AbstractAttenMask, m2::AbstractAttenMask)

"""
    m1::AbstractAttenMask & m2::AbstractAttenMask

logical and of two attention mask
"""
Base.:&(m1::AbstractAttenMask, m2::AbstractAttenMask)

"""
    BatchedMask(mask::AbstractArrayMask) <: AbstractWrapperMask

Attention mask wrapper over array mask for applying the same mask within the same batch.

# Example

```julia
julia> m = SymLengthMask([2,3])
SymLengthMask{1, Vector{Int32}}(Int32[2, 3])

julia> trues(3,3, 2) .* m
3×3×2 BitArray{3}:
[:, :, 1] =
 1  1  0
 1  1  0
 0  0  0

[:, :, 2] =
 1  1  1
 1  1  1
 1  1  1

julia> trues(3,3, 2, 2) .* m
ERROR: DimensionMismatch("arrays could not be broadcast to a common size; mask require ndims(A) == 3")
Stacktrace:
[...]

julia> trues(3,3, 2, 2) .* BatchedMask(m) # 4-th dim become batch dim
3×3×2×2 BitArray{4}:
[:, :, 1, 1] =
 1  1  0
 1  1  0
 0  0  0

[:, :, 2, 1] =
 1  1  0
 1  1  0
 0  0  0

[:, :, 1, 2] =
 1  1  1
 1  1  1
 1  1  1

[:, :, 2, 2] =
 1  1  1
 1  1  1
 1  1  1

```
"""
BatchedMask

"""
    RepeatMask(mask::AbstractAttenMask, num::Int) <: AbstractWrapperMask

Attention mask wrapper over array mask for doing inner repeat on the last dimension.

# Example

```julia
julia> m = SymLengthMask([2,3])
SymLengthMask{1, Vector{Int32}}(Int32[2, 3])

julia> trues(3,3, 2) .* m
3×3×2 BitArray{3}:
[:, :, 1] =
 1  1  0
 1  1  0
 0  0  0

[:, :, 2] =
 1  1  1
 1  1  1
 1  1  1

julia> trues(3,3, 4) .* m
ERROR: DimensionMismatch("arrays could not be broadcast to a common size; mask require 3-th dimension to be 2, but get 4")
Stacktrace:
[...]

julia> trues(3,3, 4) .* RepeatMask(m, 2)
3×3×4 BitArray{3}:
[:, :, 1] =
 1  1  0
 1  1  0
 0  0  0

[:, :, 2] =
 1  1  0
 1  1  0
 0  0  0

[:, :, 3] =
 1  1  1
 1  1  1
 1  1  1

[:, :, 4] =
 1  1  1
 1  1  1
 1  1  1

```
"""
RepeatMask

"""
    getmask(m::AbstractAttenMask, score, scale = 1)

Convert `m` into mask array of `AbstractArray` for `score` with `scale`.

# Example
```julia
julia> getmask(CausalMask(), randn(7,7), 2)
7×7 Matrix{Float64}:
 2.0  2.0  2.0  2.0  2.0  2.0  2.0
 0.0  2.0  2.0  2.0  2.0  2.0  2.0
 0.0  0.0  2.0  2.0  2.0  2.0  2.0
 0.0  0.0  0.0  2.0  2.0  2.0  2.0
 0.0  0.0  0.0  0.0  2.0  2.0  2.0
 0.0  0.0  0.0  0.0  0.0  2.0  2.0
 0.0  0.0  0.0  0.0  0.0  0.0  2.0

```
"""
getmask

end


text:
#=
interface for defining an exploration policy
=#

"""
    exploration(exp_policy, policy, env, obs, global_step, rng)
    return an action following an exploration policy
    the use can provide its own exp_policy function
"""
function exploration(f::Function, policy::AbstractNNPolicy, env::AbstractEnv, obs, global_step::Int64, rng::AbstractRNG)
    return f(policy, env, obs, global_step, rng)
end

# Examples

function linear_epsilon_greedy(max_steps::Int64, eps_fraction::Float32, eps_end::Float32)
    function action_epsilon_greedy(policy::AbstractNNPolicy, env::AbstractEnv, obs, global_step::Int64, rng::AbstractRNG)
        eps = update_epsilon(global_step, eps_fraction, eps_end, max_steps)
        if rand(rng) > eps
            return (action(policy, obs), eps)
        else
            return (rand(actions(env)), eps)
        end
    end
    return action_epsilon_greedy
end

function update_epsilon(step::Int64, epsilon_fraction::Float32, epsilon_end::Float32, max_steps::Int64)
    new_eps = 0.
    if step < epsilon_fraction*max_steps
        new_eps = 1 - (1 - epsilon_end)/(epsilon_fraction*max_steps)*step # decay
    else
        new_eps = epsilon_end
    end
    return new_eps
end


text:
@testset "failureBased.jl" begin
    @testset "FailureBasedVariableSelection{TakeObjective=false} " begin
        trailer = SeaPearl.Trailer()
        cpmodel = SeaPearl.CPModel(trailer)
        x1 = SeaPearl.IntVar(1, 4, "x1", trailer)
        x2 = SeaPearl.IntVar(1, 3, "x2", trailer)
        x3 = SeaPearl.IntVar(1, 3, "x3", trailer)
        y = SeaPearl.IntVar(1, 2, "y", trailer)

        SeaPearl.addVariable!(cpmodel, x1)
        SeaPearl.addVariable!(cpmodel, x2)
        SeaPearl.addVariable!(cpmodel, x3)
        SeaPearl.addVariable!(cpmodel, y)
        SeaPearl.addObjective!(cpmodel,y)

        SeaPearl.addConstraint!(cpmodel, SeaPearl.NotEqual(x1, x2, trailer))
        SeaPearl.addConstraint!(cpmodel, SeaPearl.NotEqual(y, x2, trailer))

        variableselection = SeaPearl.FailureBasedVariableSelection{false}()
        @test variableselection(cpmodel) == x2

        SeaPearl.empty!(cpmodel)
        x1 = SeaPearl.IntVar(1, 3, "x1", trailer)
        x2 = SeaPearl.IntVar(1, 3, "x2", trailer)
        x3 = SeaPearl.IntVar(1, 3, "x3", trailer)
        x4 = SeaPearl.IntVar(1, 3, "x4", trailer)
        y = SeaPearl.IntVar(1, 2, "y", trailer)

        SeaPearl.addVariable!(cpmodel, x1)
        SeaPearl.addVariable!(cpmodel, x2)
        SeaPearl.addVariable!(cpmodel, x3)
        SeaPearl.addVariable!(cpmodel, x4)
        SeaPearl.addVariable!(cpmodel, y)
        SeaPearl.addObjective!(cpmodel,y)

        SeaPearl.addConstraint!(cpmodel, SeaPearl.NotEqual(x1, x2, trailer))
        SeaPearl.addConstraint!(cpmodel, SeaPearl.NotEqual(x1, x3, trailer))
        SeaPearl.addConstraint!(cpmodel, SeaPearl.NotEqual(x1, x4, trailer))

        @test variableselection(cpmodel) == x1

        SeaPearl.triggerInfeasible!(SeaPearl.NotEqual(x3, x4, trailer), cpmodel)
        SeaPearl.triggerInfeasible!(SeaPearl.NotEqual(x3, x4, trailer), cpmodel)
        SeaPearl.triggerInfeasible!(SeaPearl.NotEqual(x3, x4, trailer), cpmodel)

        @test variableselection(cpmodel) == x3
    end

    @testset "FailureBasedVariableSelection{TakeObjective=true} " begin
        trailer = SeaPearl.Trailer()
        cpmodel = SeaPearl.CPModel(trailer)
        x1 = SeaPearl.IntVar(1, 4, "x1", trailer)
        x2 = SeaPearl.IntVar(1, 3, "x2", trailer)
        x3 = SeaPearl.IntVar(1, 3, "x3", trailer)
        y = SeaPearl.IntVar(1, 2, "y", trailer)

        SeaPearl.addVariable!(cpmodel, x1)
        SeaPearl.addVariable!(cpmodel, x2)
        SeaPearl.addVariable!(cpmodel, x3)
        SeaPearl.addVariable!(cpmodel, y)
        SeaPearl.addObjective!(cpmodel, y)

        SeaPearl.addConstraint!(cpmodel, SeaPearl.NotEqual(y, x2, trailer))

        variableselection = SeaPearl.FailureBasedVariableSelection{true}()
        @test variableselection(cpmodel) == y

        SeaPearl.empty!(cpmodel)
        x1 = SeaPearl.IntVar(1, 3, "x1", trailer)
        x2 = SeaPearl.IntVar(1, 3, "x2", trailer)
        x3 = SeaPearl.IntVar(1, 3, "x3", trailer)
        x4 = SeaPearl.IntVar(1, 3, "x4", trailer)
        y = SeaPearl.IntVar(1, 2, "y", trailer)

        SeaPearl.addVariable!(cpmodel, x1)
        SeaPearl.addVariable!(cpmodel, x2)
        SeaPearl.addVariable!(cpmodel, x3)
        SeaPearl.addVariable!(cpmodel, x4)
        SeaPearl.addVariable!(cpmodel, y)
        SeaPearl.addObjective!(cpmodel,y)

        SeaPearl.addConstraint!(cpmodel, SeaPearl.NotEqual(x1, x2, trailer))
        SeaPearl.addConstraint!(cpmodel, SeaPearl.NotEqual(x1, x3, trailer))
        SeaPearl.addConstraint!(cpmodel, SeaPearl.NotEqual(x1, x4, trailer))

        @test variableselection(cpmodel) == x1

        SeaPearl.triggerInfeasible!(SeaPearl.NotEqual(x3, x4, trailer), cpmodel)
        SeaPearl.triggerInfeasible!(SeaPearl.NotEqual(x3, x4, trailer), cpmodel)
        SeaPearl.triggerInfeasible!(SeaPearl.NotEqual(x3, x4, trailer), cpmodel)

        @test variableselection(cpmodel) == x3

        @test isnothing(cpmodel.statistics.objectives[1])
        @test isnothing(cpmodel.statistics.objectives[2])
        @test isnothing(cpmodel.statistics.objectives[3])

    end
end


text:
# calculate eigen values/vectors
# plot eigen vectors on the contour

module CalcPlotEigen
    using Plots, Base.Iterators, LinearAlgebra
    pyplot()

    function pdf(x, mu, cov_mat)
        # determinant
        det_mat = det(cov_mat)

        diff_x = x - mu
        return exp(-((diff_x'/cov_mat)*diff_x)/2)/(2*pi*sqrt(det_mat))
    end

    function main(is_test=false)
        # x-y grid
        vx = 0:200
        vy = 0:100

        # mu, covariance matrix
        mu = [150; 50]
        cov = [100 -25*sqrt(3); -25*sqrt(3) 50]

        # calculate eigen values/vectors
        eig_vals, eig_vecs = eigen(cov)
        println("Eigen Values: $(eig_vals)")
        println("Eigen Vectors: $(eig_vecs)")
        println("Eigen Vector1: $(eig_vecs[:, 1])")
        println("Eigen Vector2: $(eig_vecs[:, 2])")

        # long or short
        if eig_vals[1] > eig_vals[2]
            long_val = eig_vals[1]
            long_vec = eig_vecs[:, 1]
            short_val = eig_vals[2]
            short_vec = eig_vecs[:, 2]
        else
            long_val = eig_vals[2]
            long_vec = eig_vecs[:, 2]
            short_val = eig_vals[1]
            short_vec = eig_vecs[:, 1]
        end
        println("Long Value: $(long_val)")
        println("Long Vector: $(long_vec)")
        println("Short Value: $(short_val)")
        println("Short Vector: $(short_vec)")

        # plot
        z = [pdf([x; y], mu, cov) for x in vx, y in vy]
        contour(vx, vy, z', c=:haline, aspect_ratio=:equal)
        long = 2 * sqrt(long_val) * long_vec
        short = 2 * sqrt(short_val) * short_vec
        quiver!([mu[1]], [mu[2]], quiver=([long[1]], [long[2]]),
                c=:red, aspect_ratio=:equal)
        quiver!([mu[1]], [mu[2]], quiver=([short[1]], [short[2]]),
                c=:blue, aspect_ratio=:equal)

        if is_test == false
            save_path = joinpath(split(@__FILE__, "src")[1], "src/prob_stats/multi_dim_gauss_dist/eigen/contour_eigen.png")
            savefig(save_path)
        end
    end
end

max_stars_repo_name: ShisatoYano/JuliaAutonomy

max_issues_repo_name: ShisatoYano/JuliaAutonomy

max_forks_repo_name: ShisatoYano/JuliaAutonomy


text:
struct HandlerData
	user::Any
	func::Function
end

EventHandlers = Dict{Symbol, Vector{HandlerData}}

get_events(owner) = owner.events

function add_event(handler::Function, owner, event::Symbol, user = nothing)
	handlers = get!(get_events(owner), event) do; HandlerData[] end
	push!(handlers, HandlerData(user, handler))
	handler
end

function remove_event(handler::Function, owner, event::Symbol)
	handlers = get_events(owner)[event]
	deleteat!(handlers, findfirst(h->h.func == handler, handlers))
	nothing
end

function remove_event(user, owner, event::Symbol)
	handlers = get_events(owner)[event]
	deleteat!(handlers, findfirst(h->h.user == user, handlers))
	nothing
end

function remove_events(user, owner)
	for (event, handlers) in get_events(owner)
		filter!(h->h.user != user, handlers)
	end
end

function call_event(owner, event::Symbol, args...)
	#@info("call_event $(get_id(owner)) $event")
	handlers = get!(get_events(owner), event) do; Function[] end
	for h in handlers
		h.func(owner, event, args...)
	end
	#@info("call_event $(get_id(owner)) $event done")
	nothing
end


text:
@doc raw"""
    Point3

An object of the type [`Point3`](@ref) is a point in the three-dimensional
Euclidean space ``Ε^3``.
"""
Point3

"""
    Point3(::Origin = ORIGIN)

Introduces a variable `p` with Cartesian coordinates ``(0, 0, 0)``.
"""
Point3(::Origin = ORIGIN)

"""
    Point3(x::Real, y::Real, z::Real)

Introduces a point `p` initialized to ``(x, y, z)``.
"""
Point3(x::Real, y::Real, z::Real) =
    Point3(convert(FT, x), convert(FT, y), convert(FT, z))

@doc raw"""
    Point3(hx::Real, hy::Real, hz::Real, hw::Real = 1)

Introduces a point `p` initialized to `(hx/hw,hy/hw,hz/hw)`.

!!! info "Precondition"

    `hw` ``≠`` `0`
"""
Point3(x::Real, y::Real, z::Real, hw::Real) =
    Point3(convert(RT, x), convert(RT, y), convert(RT, z), convert(RT, hw))

"""
    Point3(wp::WeightedPoint3)

Introduces a point from a weighted point.
"""
Point3(wp::WeightedPoint3)

"""
    ==(p::Point3, q::Point3)

Test for equality.

Two points are equal, iff their ``x``, ``y`` and ``z`` coordinates are equal.
The point can be compared with [`ORIGIN`](@ref).
"""
==(p::Point3, q::Point3)

"""
    hx(p::Point3)

Returns the homogeneous ``x`` coordinate.
"""
hx(p::Point3)

"""
    hy(p::Point3)

Returns the homogeneous ``y`` coordinate.
"""
hy(p::Point3)

"""
    hz(p::Point3)

Returns the homogeneous ``z`` coordinate.
"""
hz(p::Point3)

"""
    hw(p::Point3)

Returns the homogenizing coordinate.
"""
hw(p::Point3)

"""
    x(p::Point3)

Returns the Cartesian ``x`` coordinate, that is [`hx()`](@ref)/[`hw()`](@ref).
"""
x(p::Point3)

"""
    y(p::Point3)

Returns the Cartesian ``y`` coordinate, that is [`hy()`](@ref)/[`hw()`](@ref).
"""
y(p::Point3)

"""
    z(p::Point3)

Returns the Cartesian ``y`` coordinate, that is [`hz()`](@ref)/[`hw()`](@ref).
"""
z(p::Point3)

@doc raw"""
    homogeneous(p::Point3, i::Integer)

Returns the `i`ᵗʰ homogeneous coordinate of `p`.

!!! info "Precondition"

    ``0 ≤ i ≤ 3``.
"""
homogeneous(p::Point3, i::Integer)

@doc raw"""
    cartesian(p::Point3, i::Integer)

Returns the `i`ᵗʰ Cartesian coordinate of `p`.

!!! info "Precondition"

    ``0 ≤ i ≤ 2``.
"""
cartesian(p::Point3, i::Integer)

"""
    dimension(p::Point3)

Returns the dimension (the constant 3).
"""
dimension(p::Point3)

"""
    bbox(p::Point3)

Returns a bounding box containing `p`.
"""
bbox(p::Point3)

"""
    transform(p::Point3, t::AffTransformation3)

Returns the point obtained by applying `t` on `p`.
"""
transform(p::Point3, t::AffTransformation3)

"""
    <(p::Point3, q::Point3)

Return `true` iff `p` is lexicographically smaller than `q` (the lexicographical
order being defined on the Cartesian coordinates).
"""
<(p::Point3, q::Point3)

"""
    >(p::Point3, q::Point3)

Return `true` iff `p` is lexicographically greater than `q`.
"""
>(p::Point3, q::Point3)

"""
    <=(p::Point3, q::Point3)

Return `true` iff `p` is lexicographically smaller or equal to `q`.
"""
<=(p::Point3, q::Point3)

"""
    >=(p::Point3, q::Point3)

Return `true` iff `p` is lexicographically greater or equal to `q`.
"""
>=(p::Point3, q::Point3)

"""
    -(p::Point3, q::Point3)

Returns the difference vector between `p` and `q`.

You can substitute [`ORIGIN`](@ref) for either `p` or `q`, but not for both.
"""
-(p::Point3, q::Point3)


text:
module RationalSimplex

import LinearAlgebra: I

export simplex

"""
    simplex(c, obj, A, b, senses)

Solve the linear program

     {min/max}  dot(c, x)
    subject to  A x {>=|==|<=} b
                x >= 0,

where `x` is in the set of Rationals (so algorithm is exact).
`obj` should be either `:Min` or `:Max`.
`senses` should be a `Vector` of `Char`s: `'<','=','>'`.
"""
function simplex(c::Vector{T}, obj::Symbol, A::Matrix{T}, b::Vector{T},
                 senses::Vector{Char}) where {T <: Rational}
    # Any constraint that isn't an equality needs an auxiliary variable.
    num_auxiliary = count(sense -> sense != '=', senses)
    num_constraints, num_variables = size(A)
    num_variables_aux = num_variables + num_auxiliary
    # Extend constraint matrix and objective function, normalize objective.
    A_aux = zeros(T, num_constraints, num_variables_aux)
    A_aux[:, 1:num_variables] = A
    c_aux = zeros(T, num_variables_aux)
    c_aux[1:num_variables] = (obj == :Max) ? -c : c
    b_aux = copy(b)  # In case we modify signs.
    # Add the auxiliaries to the constraint matrix, and make sure b ≥ 0.
    offset = 1
    for (row, sense) in enumerate(senses)
        if sense == '<'
            A_aux[row, num_variables + offset] = one(T)
            offset += 1
        elseif sense == '>'
            A_aux[row, num_variables + offset] = -one(T)
            offset += 1
        end
        if b[row] < zero(T)
            A_aux[row, :] *= -one(T)
            b_aux[row] *= -one(T)
        end
    end
    # Solve problem with auxiliaries, but remove auxiliaries before returning.
    sense, x = simplex(c_aux, A_aux, b_aux)
    return sense, x[1:num_variables]
end

"""
    simplex(c, A, b)

Solve the linear program in standard computational form

           min  dot(c, x)
    subject to  A x == b [≥ 0]
                x >= 0,

where `x` is in the set of Rationals (so algorithm is exact).

The algorithm is the "two-phase primal revised simplex method".
In the first phase auxiliaries are created to get an initial basis.
We then eliminate the auxiliaries from the basis until it consists only of
actual variables. This is the "textbook" algorithm, and shouldn't be used for
anything that matters. It doesn't exploit sparsity at all. You could use it
with floating points but it won't work for anything except the most simple
problem due to accumulated errors and comparisons with zero.
"""
function simplex(c::Vector{T}, A::Matrix{T}, b::Vector{T}) where {T<:Rational}
    @assert all(b .> zero(T))

    # Set up data structures.
    num_constraints, num_variables = size(A)
    is_basic = zeros(Bool, num_variables + num_constraints)
    basic = zeros(Int, num_constraints)  # Indices of current basis.
    Binv = Matrix{T}(I, num_constraints, num_constraints)  # Basis inverse.
    cB = ones(T, num_constraints)  # Costs of basic variables.
    x = zeros(T, num_variables + num_constraints)  # Current solution.

    # Intialize phase one of RSM by setting basis = auxiliaries.
    for aux_index in 1:num_constraints
        basic[aux_index] = num_variables + aux_index
        is_basic[num_variables + aux_index] = true
        x[num_variables + aux_index] = b[aux_index]
    end
    phase_one = true

    # Begin simplex iterations.
    status = :Unknown
    while true
        # Calculate dual solution.
        π = vec(cB' * Binv)
        # Use it to calculate the reduced costs of the variables. Don't
        # calculate for auxiliaries - they can't re-enter the basis.
        rc = (phase_one ? zero(T) : c) .- vec(π' * A)
        @. rc *= !is_basic[1:num_variables]  # In basis, so don't consider.
        # Find variable to enter basis.
        min_rc, entering = findmin(rc)
        # If none have negative reduced cost, there are no variables that can
        # enter, and we are at optimality for this phase. This may or may not
        # be optimal for the actual problem.
        if min_rc >= zero(T)
            if phase_one
                phase_one = false
                # If any auxiliary still nonzero, we couldn't find a feasible
                # basis without auxiliaries.
                if any(x[num_variables+1:end] .> zero(T))
                    return :Infeasible, x[1:num_variables]
                end
                # Otherwise, start phase two with nice feasible basis.
                for i in 1:num_constraints
                    cB[i] = basic[i] > num_variables ? zero(T) : c[basic[i]]
                end
                continue
            else  # Phase two. We can't improve further, so we are optimal.
                return :Optimal, x[1:num_variables]
            end
        end
        # Calculate how the solution will change when our new variable enters
        # the basis and increases from zero.
        BinvAs = Binv * A[:, entering]
        # Perform a "ratio test" on each variable to determine which will
        # reach zero first.
        leaving = 0
        min_ratio = zero(T)
        for j in 1:num_constraints
            if BinvAs[j] > zero(T)
                ratio = x[basic[j]] / BinvAs[j]
                if ratio < min_ratio || leaving == 0
                    min_ratio = ratio
                    leaving = j
                end
            end
        end
        # If no variable will leave basis, then we have an unbounded problem.
        if leaving == 0
            return :Unbounded, x[1:num_variables]
        end
        # Update solution.
        for j in 1:num_constraints
            x[basic[j]] -= min_ratio * BinvAs[j]
        end
        x[entering] = min_ratio
        # Update basis inverse.
        # Our tableau is [ Binv b | Binv | BinvAs ] and we doing a pivot on the
        # `leaving` row of BinvAs.
        pivot_value = BinvAs[leaving]
        for basis_row in 1:num_constraints
            basis_row == leaving && continue  # All rows except `leaving` row.
            factor = BinvAs[basis_row] / pivot_value
            for basis_col in 1:num_constraints
                Binv[basis_row, basis_col] -= factor * Binv[leaving, basis_col]
            end
        end
        # Finally, the `leaving` row.
        for basis_col in 1:num_constraints
            Binv[leaving, basis_col] /= pivot_value
        end
        # Update variable status flags.
        is_basic[basic[leaving]] = false
        is_basic[entering] = true
        cB[leaving] = phase_one ? zero(T) : c[entering]
        basic[leaving] = entering
    end
end

end  # module RationalSimplex.


text:
"""
    TV(p, q)

Returns the total variation distance between `p` and `q` (half the sum of the absolute value of the differences).
"""
TV(p::AbstractVector, q::AbstractVector) = (1 / 2) * norm(p - q, 1)

function TV(p::AbstractVector{T1}, q::AbstractVector{T2}) where {T1<:Rational,T2<:Rational}
    length(p) == length(q) ||
    throw(DimensionMismatch("Input vectors must be of the same length"))
    return (1 // 2) * sum(abs, p - q)
end

"""
    ≺(p::AbstractVector{T1}, q::AbstractVector{T2}; tol = (T1 <: AbstractFloat || T2 <: AbstractFloat) ? T2(1e-8) : zero(T2) ) where {T1, T2}

Returns true if `q` majorizes `p` and false otherwise. The keyword argument `tol` specifies a tolerance for the comparisons. Can be used in infix form, i.e. `p ≺ q`.
"""
function ≺(
    p::AbstractVector{T1},
    q::AbstractVector{T2};
    tol = (T1 <: AbstractFloat || T2 <: AbstractFloat) ? T2(1e-8) : zero(T2),
) where {T1,T2}
    length(p) == length(q) ||
    throw(DimensionMismatch("Input vectors must be of the same length"))
    pv = sort(p, rev = true)
    qv = sort(q, rev = true)
    isapprox(sum(p), sum(q), atol = tol, rtol = 0) || return false
    return all(cumsum(pv) .<= (cumsum(qv) .+ tol))
end

function ≻(
    p::AbstractVector{T1},
    q::AbstractVector{T2};
    tol = (T1 <: AbstractFloat || T2 <: AbstractFloat) ? T2(1e-8) : zero(T2),
) where {T1,T2}
    ≺(q, p; tol = tol)
end

"""
    majmax(q::AbstractVector, ϵ)

Returns the maximum in majorization order over the total variation ball (of probability vectors) of radius `ϵ` around a probability vector `q`.
"""
function majmax(q::AbstractVector, ϵ)
    d = length(q)
    T = eltype(q[1] + ϵ)
    inds = sortperm(q)
    q = q[inds]
    p = zeros(T, d)
    budget = ϵ

    p[1] = max(zero(T), q[1] - budget)
    budget = budget - abs(q[1] - p[1])

    for j = 2:d-1
        p[j] = max(zero(T), q[j] - budget)
        budget = budget - abs(q[j] - p[j])
    end
    p[d] = q[d] + ϵ - budget
    invpermute!(p, inds)
    return p
end

@inline function get_alpha1(qup::AbstractVector, ϵ)
    d = length(qup)
    cs = qup[1] + ϵ
    for j = 1:d-1
        if cs <= j * qup[j+1]
            return j, cs / j
        end
        cs += qup[j+1]
    end
    return nothing
end

@inline function get_alpha2(qup::AbstractVector, ϵ)
    d = length(qup)
    cs = qup[d] - ϵ
    for (idx, j) in enumerate(d:-1:2)
        if cs >= idx * qup[j-1]
            return idx, cs / idx
        end
        cs += qup[j-1]
    end
    return nothing
end

uniform(T, d) = inv(T(d)) * ones(T, d)

"""
    majmin(q::AbstractVector, ϵ)

Returns the minimum in majorization order over the total variation ball (of probability vectors) of radius `ϵ` around a probability vector `q`.
"""
function majmin(q::AbstractVector, ϵ)
    T = typeof((q[1] + ϵ) / 2)
    d = length(q)
    inds = sortperm(q, rev = false)

    qup = convert.(T, q[inds])
    VT = typeof(qup)

    n1a1 = get_alpha1(qup, ϵ)
    n2a2 = get_alpha2(qup, ϵ)

    if n1a1 === nothing || n2a2 === nothing
        return VT(uniform(T, d))
    end

    n1, alpha1 = n1a1
    n2, alpha2 = n2a2

    if alpha1 > 1 / d || alpha2 < 1 / d
        return VT(uniform(T, d))
    end

    out = qup
    out[1:n1] .= alpha1
    out[d-n2+1:d] .= alpha2
    invpermute!(out, inds)
    return VT(out)
end


text:
cache(args...) = joinpath(@__DIR__, "..", "cache", args...)

"""
    get_id(x)

Get unique id of arbitrary struct `x`. Works by serializing the struct and computing the
SHA256 checksum of the serialization.
"""
function get_id(x)
    buff = IOBuffer()
    serialize(buff, x)
    id = bytes2hex(sha256(buff.data))
end

"""
    reference(H; force_recompute=false) -> (E0, reference_vector)

Get the exact ground-state energy a wave function of Hamiltonian `H`. Values are cached to
`Sekira/cache`.
"""
function reference(H; force_recompute=false)
    !isdir(cache()) && mkpath(cache())

    Hs = Rimu.Hamiltonians
    id = get_id(H)
    if !force_recompute && isfile(cache("$id.bson"))
        @mpi_root @info "Reference found."
        ref = RimuIO.load_dvec(cache("$id.bson"))
        E0 = read(cache("$id.energy"), Float64)
    else
        @mpi_root @info "Computing reference."
        dv = DVec(starting_address(H) => 1.0)
        ishermitian = Hs.LOStructure(H) == Hs.Hermitian()
        issymmetric = ishermitian && eltype(H) <: Real
        vals, vecs, info = eigsolve(H, dv, 1, :SR; issymmetric, ishermitian)
        ref = vecs[1]
        E0 = vals[1]
        @assert info.converged ≥ 1

        RimuIO.save_dvec(cache("$id.bson"), ref)
        write(cache("$id.energy"), E0)
    end
    return E0, ref
end

function deterministic_reference(H)
    ref = MPIData(DVec(starting_address(H) => 1.0, style=IsDeterministic()))
    r_strat = ReportToFile(
        chunk_size=100, save_if=is_mpi_root(), io=stderr, filename="ref.arrow"
    )
    lomc!(H, ref; laststep=10_000, r_strat)

    RimuIO.save_dvec("ref.bson", ref)

end


text:
using KernelAbstractions
using Test

@kernel function localmem(A)
    N = @uniform prod(@groupsize())
    @uniform begin
        N2 = prod(@groupsize())
    end
    I = @index(Global, Linear)
    i = @index(Local, Linear)
    lmem = @localmem Int (N,) # Ok iff groupsize is static
    @inbounds begin
        lmem[i] = i
        @synchronize
        A[I] = lmem[N2 - i + 1]
    end
end

@kernel function localmem2(A)
    N = @uniform prod(@groupsize())
    @uniform begin
        N2 = prod(@groupsize())
    end
    I = @index(Global, Linear)
    i = @index(Local, Linear)
    lmem = @localmem Int (N,) # Ok iff groupsize is static
    @inbounds begin
        lmem[i] = i + 3
        for j in 1:2
            lmem[i] -= j
            @synchronize
        end
        A[I] = lmem[N2 - i + 1]
    end
end

function localmem_testsuite(backend, ArrayT)
    @testset "kernels" begin
        @testset for kernel! in (localmem(backend(), 16), localmem2(backend(), 16))
            A = ArrayT{Int}(undef, 64)
            wait(kernel!(A, ndrange=size(A)))
            @test all(A[1:16] .== 16:-1:1)
            @test all(A[17:32] .== 16:-1:1)
            @test all(A[33:48] .== 16:-1:1)
            @test all(A[49:64] .== 16:-1:1)
        end
    end
end


text:
"""
    ETA(steps, header, t₀, fmt)

This trace column can be used to indicate the estimated time of
completion of a calculation.
"""
struct ETA <: TraceColumn
    steps::Int
    header::String
    t₀::Float64
    fmt::FormatExpr
end

"""
    ETA(steps[, name="ETA"])

Constructe a [`ETA`](@ref) trace column, for a specific number of
`steps`.
"""
function ETA(steps::Int, name::String="ETA")
    # It would possibly be good to allow for more than two-digit hour
    # counts, but still keep two-digit display for the standard case,
    # i.e. less than 100 hours. Is there a way to make fmt always
    # occupy e.g. 12 characters (and right-shifted), but keep, the
    # standard format 02d?
    ETA(steps, fmt(">8s", name), time(), FormatExpr("{1:02d}:{2:02d}:{3:02d}"))
end

function (p::ETA)(i::Integer)
    elapsed = time() - p.t₀
    seconds = (elapsed > 0 ? (p.steps-i)*elapsed/i : Inf)

    hours = div(seconds, 3600)
    seconds -= 3600hours

    minutes = div(seconds, 60)
    seconds -= 60minutes

    hours,minutes,trunc(Int, seconds)
end

export ETA


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

function transform!(context::T) where {T<:SHA3_CTX}
    # First, update state with buffer
    pbuf = Ptr{eltype(context.state)}(pointer(context.buffer))
    for idx in 1:div(blocklen(T),8)
        context.state[idx] = context.state[idx] ⊻ unsafe_load(pbuf, idx)
    end
    bc = context.bc
    state = context.state

    # We always assume 24 rounds
    @inbounds for round in 0:23
        # Theta function
        for i in 1:5
            bc[i] = state[i] ⊻ state[i + 5] ⊻ state[i + 10] ⊻ state[i + 15] ⊻ state[i + 20]
        end

        for i in 0:4
            temp = bc[rem(i + 4, 5) + 1] ⊻ L64(1, bc[rem(i + 1, 5) + 1])
            j = 0
            while j <= 20
                state[Int(i + j + 1)] = state[i + j + 1] ⊻ temp
                j += 5
            end
        end

        # Rho Pi
        temp = state[2]
        for i in 1:24
            j = SHA3_PILN[i]
            bc[1] = state[j]
            state[j] = L64(SHA3_ROTC[i], temp)
            temp = bc[1]
        end

        # Chi
        j = 0
        while j <= 20
            for i in 1:5
                bc[i] = state[i + j]
            end
            for i in 0:4
                state[j + i + 1] = state[j + i + 1] ⊻ (~bc[rem(i + 1, 5) + 1] & bc[rem(i + 2, 5) + 1])
            end
            j += 5
        end

        # Iota
        state[1] = state[1] ⊻ SHA3_ROUND_CONSTS[round+1]
    end

    return context.state
end

# Finalize data in the buffer, append total bitlength, and return our precious hash!
function digest!(context::T) where {T<:SHA3_CTX}
    usedspace = context.bytecount % blocklen(T)
    # If we have anything in the buffer still, pad and transform that data
    if usedspace < blocklen(T) - 1
        # Begin padding with a 0x06
        context.buffer[usedspace+1] = 0x06
        # Fill with zeros up until the last byte
        context.buffer[usedspace+2:end-1] .= 0x00
        # Finish it off with a 0x80
        context.buffer[end] = 0x80
    else
        # Otherwise, we have to add on a whole new buffer just for the zeros and 0x80
        context.buffer[end] = 0x06
        transform!(context)

        context.buffer[1:end-1] = 0x0
        context.buffer[end] = 0x80
    end

    # Final transform:
    transform!(context)

    # Return the digest
    return reinterpret(UInt8, context.state)[1:digestlen(T)]
end

max_stars_repo_name: greimel/julia

max_issues_repo_name: greimel/julia

max_forks_repo_name: greimel/julia


text:
# Autogenerated wrapper script for MariaDB_Connector_C_jll for i686-w64-mingw32
export libmariadb

using LibCURL_jll
using Libiconv_jll
using OpenSSL_jll
using Zlib_jll
JLLWrappers.@generate_wrapper_header("MariaDB_Connector_C")
JLLWrappers.@declare_library_product(libmariadb, "libmariadb.dll")
function __init__()
    JLLWrappers.@generate_init_header(LibCURL_jll, Libiconv_jll, OpenSSL_jll, Zlib_jll)
    JLLWrappers.@init_library_product(
        libmariadb,
        "bin\\libmariadb.dll",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@generate_init_footer()
end  # __init__()


text:
__precompile__()
module Metalhead

using Flux, Images, ImageFiltering, BSON, REPL, Requires, Statistics
using Flux: @treelike

# Models
export VGG19, SqueezeNet, DenseNet, ResNet, GoogleNet

# Useful re-export from Images
export load

# High-level classification APIs
export predict, classify

# Data Sets
export ImageNet, CIFAR10

# Data set utilities
export trainimgs, testimgs, valimgs, dataset, datasets

function __init__()
    @require TerminalExtensions="d3a6a179-465e-5219-bd3e-0137f7fd17c7" include("display/terminal_extensions.jl")
end

include("datasets/utils.jl")
include("model.jl")
include("utils.jl")
include("display/terminal.jl")
include("datasets/imagenet.jl")
include("datasets/cifar10.jl")
include("datasets/autodetect.jl")
include("vgg19.jl")
include("squeezenet.jl")
include("densenet.jl")
include("resnet.jl")
include("googlenet.jl")
end # module