text:
# Optimization and solution query
#
#
# OBS: in xpress all functions with LP in name
#      actually refer to CONTINUOUS optimization
#      LP, QP, QCQP, SOCP
#      as opposed to Integer programming

# use here ismip!
# function optimize(model::Model)
#     @assert model.ptr_model != C_NULL
#     ret = @xprs_ccall(:global, Cint, (Ptr{Void},), model.ptr_model)
#     if ret != 0
#         throw(XpressError(model))
#     end
#     nothing
# end
function optimize(model::Model)
    @assert model.ptr_model != C_NULL
    if is_mip(model)
        ret = @xprs_ccall(mipoptimize, Cint, (Ptr{Void},Ptr{Cchar}),
            model.ptr_model,C_NULL)
        if ret != 0
            throw(XpressError(model))
        end
    else
        ret = @xprs_ccall(lpoptimize, Cint, (Ptr{Void},Ptr{Cchar}),
            model.ptr_model,C_NULL)
        if ret != 0
            throw(XpressError(model))
        end
    end
    nothing
end

function computeIIS(model::Model)
    @assert model.ptr_model != C_NULL
    ret = @xprs_ccall(iisall, Cint, (Ptr{Void},), model.ptr_model)
    if ret != 0
        throw(XpressError(model))
    end
    nothing
end

#################################################
#
#  solution status and optimization info
#
#################################################

const status_symbols_lp = Dict(
    XPRS_LP_UNSTARTED         => :unstarted,
    XPRS_LP_OPTIMAL           => :optimal,
    XPRS_LP_INFEAS            => :infeasible,
    XPRS_LP_CUTOFF            => :cutoff,
    XPRS_LP_UNFINISHED        => :unfinished,
    XPRS_LP_UNBOUNDED         => :unbounded,
    XPRS_LP_CUTOFF_IN_DUAL    => :cutoff_in_dual,
    XPRS_LP_UNSOLVED          => :unsolved,
    XPRS_LP_NONCONVEX         => :nonconvex
)

const status_symbols_mip = Dict(
    XPRS_MIP_NOT_LOADED       => :mip_not_loaded,
    XPRS_MIP_LP_NOT_OPTIMAL   => :mip_lp_not_optimal,
    XPRS_MIP_LP_OPTIMAL       => :mip_lp_optimal,
    XPRS_MIP_NO_SOL_FOUND     => :mip_no_sol_found,
    XPRS_MIP_SOLUTION         => :mip_suboptimal,
    XPRS_MIP_INFEAS           => :mip_infeasible,
    XPRS_MIP_OPTIMAL          => :mip_optimal,
    XPRS_MIP_UNBOUNDED        => :mip_unbounded
)

@xprs_int_attr get_mip_status_code XPRS_MIPSTATUS
@xprs_int_attr get_lp_status_code XPRS_LPSTATUS

get_mip_status(model::Model) = status_symbols_mip[get_mip_status_code(model)]::Symbol
get_lp_status(model::Model) = status_symbols_lp[get_lp_status_code(model)]::Symbol

@xprs_int_attr get_sol_count XPRS_MIPSOLS
@xprs_int_attr get_node_count XPRS_NODES

@xprs_int_attr get_simplex_iter_count XPRS_SIMPLEXITER
@xprs_int_attr get_barrier_iter_count XPRS_BARITER

# objective
@xprs_dbl_attr get_lp_objval XPRS_LPOBJVAL
@xprs_dbl_attr get_mip_objval XPRS_MIPOBJVAL

@xprs_dbl_attr get_bardualobj XPRS_BARDUALOBJ
@xprs_dbl_attr get_barprimalobj XPRS_BARPRIMALOBJ

@xprs_dbl_attr get_bestbound XPRS_BESTBOUND

#iis
@xprs_int_attr num_iis XPRS_NUMIIS

function get_objval(model::Model)
    if num_intents(model)+num_sos(model) > 0
        return get_mip_objval(model)
    else
        return get_lp_objval(model)
    end
end

type OptimInfo
    status_lp::Symbol
    status_mip::Symbol

    #time ??

    sol_count::Int
    node_count::Int

    simplex_iter_count::Int
    barrier_iter_count::Int

end

function get_optiminfo(model::Model)
    OptimInfo(
        get_lp_status(model),
        get_mip_status(model),

        #get_runtime(model),

        get_sol_count(model),
        get_node_count(model),

        get_simplex_iter_count(model),
        get_barrier_iter_count(model)
    )
end

function show(io::IO, s::OptimInfo)
    println(io, "Xpress Optimization Info")
    println(io, "    lp status    = $(s.status_lp)")
    println(io, "    mip status   = $(s.status_mip)")
    #println(io, "    runtime  = $(s.runtime)")
    println(io, "    # solutions     = $(s.sol_count)")
    println(io, "    # nodes         = $(s.node_count)")

    println(io, "    # simplex iters = $(s.simplex_iter_count)")
    println(io, "    # barrier iters = $(s.barrier_iter_count)")
end

#################################################
#
#  solution query
#
#################################################

function get_complete_lp_solution(model::Model)
    cols = num_vars(model)
    rows = num_constrs(model)

    x = Vector{Float64}(cols)
    slack = Vector{Float64}(rows)
    dual = Vector{Float64}(rows)
    red = Vector{Float64}(cols)

    ret = @xprs_ccall(getlpsol, Cint,
        (Ptr{Void},
         Ptr{Float64},
         Ptr{Float64},
         Ptr{Float64},
         Ptr{Float64}
            ), model.ptr_model, x, slack, dual, red)
    if ret != 0
        throw(XpressError(model))
    end
    return x, slack, dual, red
end
function get_lp_solution(model::Model)
    cols = num_vars(model)

    x = Vector{Float64}(cols)

    ret = @xprs_ccall(getlpsol, Cint,
        (Ptr{Void},
         Ptr{Float64},
         Ptr{Float64},
         Ptr{Float64},
         Ptr{Float64}
            ), model.ptr_model, x, C_NULL, C_NULL, C_NULL)
    if ret != 0
        throw(XpressError(model))
    end
    return x
end
function get_lp_slack(model::Model)
    rows = num_constrs(model)

    slack = Vector{Float64}(rows)

    ret = @xprs_ccall(getlpsol, Cint,
        (Ptr{Void},
         Ptr{Float64},
         Ptr{Float64},
         Ptr{Float64},
         Ptr{Float64}
            ), model.ptr_model, C_NULL, slack, C_NULL, C_NULL)
    if ret != 0
        throw(XpressError(model))
    end
    return slack
end
function get_lp_dual(model::Model)
    rows = num_constrs(model)

    dual = Vector{Float64}(rows)

    ret = @xprs_ccall(getlpsol, Cint,
        (Ptr{Void},
         Ptr{Float64},
         Ptr{Float64},
         Ptr{Float64},
         Ptr{Float64}
            ), model.ptr_model, C_NULL, C_NULL, dual, C_NULL)
    if ret != 0
        throw(XpressError(model))
    end
    return dual
end
function get_lp_reducedcost(model::Model)
    cols = num_vars(model)

    red = Vector{Float64}(cols)

    ret = @xprs_ccall(getlpsol, Cint,
        (Ptr{Void},
         Ptr{Float64},
         Ptr{Float64},
         Ptr{Float64},
         Ptr{Float64}
            ), model.ptr_model, C_NULL, C_NULL, C_NULL, red)
    if ret != 0
        throw(XpressError(model))
    end
    return red
end

## mip sol
function get_mip_solution(model::Model)
    cols = num_vars(model)

    x = Vector{Float64}(cols)

    ret = @xprs_ccall(getmipsol, Cint,
        (Ptr{Void},
         Ptr{Float64},
         Ptr{Float64}
            ), model.ptr_model, x, C_NULL)
    if ret != 0
        throw(XpressError(model))
    end
    return x
end
function get_mip_slack(model::Model)
    rows = num_constrs(model)

    slack = Vector{Float64}(rows)

    ret = @xprs_ccall(getmipsol, Cint,
        (Ptr{Void},
         Ptr{Float64},
         Ptr{Float64}
            ), model.ptr_model, C_NULL, slack)
    if ret != 0
        throw(XpressError(model))
    end
    return slack
end

function get_solution(model::Model)
    if is_mip(model)
        return get_mip_solution(model)
    else
        return get_lp_solution(model)
    end
end
function get_slack(model::Model)
    if is_mip(model)
        return get_mip_slack(model)
    else
        return get_lp_slack(model)
    end
end
function get_dual(model::Model)
    if is_mip(model)
        error("Not possible to get MIP duals")
        return Float64[]
    else
        return get_lp_dual(model)
    end
end
function get_reducedcost(model::Model)
    if is_mip(model)
        error("Not possible to get MIP reduced costs")
        return Float64[]
    else
        return get_lp_reducedcost(model)
    end
end

#################################################
#
#  basis query
#
#################################################
const basicmap = Dict(
    0 => :NonbasicAtLower,#_or_SuperbasicAtZero,
    1 => :Basic,
    2 => :NonbasicAtUpper,
    3 => :Superbasic
)

function get_basis(model::Model)
    cval = Array(Cint, num_vars(model))
    cbasis = Array(Symbol, num_vars(model))

    rval = Array(Cint, num_constrs(model))
    rbasis = Array(Symbol, num_constrs(model))

    ret = @xprs_ccall(getbasis, Cint,
        (Ptr{Void},
         Ptr{Cint},
         Ptr{Cint}
            ), model.ptr_model, rval, cval)
    if ret != 0
        throw(XpressError(model))
    end

    for it in 1:length(cval)
        cbasis[it] = basicmap[cval[it]]
    end

    for it in 1:length(rval)
        rbasis[it] = basicmap[rval[it]]
    end
    return cbasis, rbasis
end

#################################################
#
#  IIS query
#
#################################################

function get_iisdata(model::Model, num::Int)
# num is the number of THE IIS to be queried

    rows = Array(Cint,1)
    cols = Array(Cint,1)
    ret = @xprs_ccall(getbasis, Cint,
        (Ptr{Void},
            Cint,# num
         Ptr{Cint},# #row
         Ptr{Cint},# #cols
         Ptr{Cint},# rows set
         Ptr{Cint},# bounds set
         Ptr{Cchar},# cstr type
         Ptr{Cchar},# bnd type
         Ptr{Float64}, # duals
         Ptr{Float64}, # reduced
         Ptr{Cchar}, # isol rows
         Ptr{Cchar} #isol cols
            ), model.ptr_model, Cint(num), rows,
        cols,C_NULL,C_NULL,
        C_NULL,C_NULL,C_NULL,
        C_NULL,C_NULL,C_NULL,)

    if ret != 0
        throw(XpressError(model))
    end

    rows_set = Array(Cint, rows)
    cols_set = Array(Cint, cols)

    ret = @xprs_ccall(getbasis, Cint,
        (Ptr{Void},
            Cint,# num
         Ptr{Cint},# #row
         Ptr{Cint},# #cols
         Ptr{Cint},# rows set
         Ptr{Cint},# bounds set
         Ptr{Cchar},# cstr type
         Ptr{Cchar},# bnd type
         Ptr{Float64}, # duals
         Ptr{Float64}, # reduced
         Ptr{Cchar}, # isol rows
         Ptr{Cchar} #isol cols
            ), model.ptr_model, Cint(num), rows,
        cols,rows_set,cols_set,
        C_NULL,C_NULL,C_NULL,
        C_NULL,C_NULL,C_NULL,)

    if ret != 0
        throw(XpressError(model))
    end

    return rows_set, cols_set
end

#################################################
#
#  infeasibility & unboundedness
#
#################################################

function getdualray(model::Model)

    dray = Array(Float64, num_constrs(model))

    hasray = Array(Cint, 1)

    ret = @xprs_ccall(getdualray, Cint,
        (Ptr{Void},
         Ptr{Float64},
         Ptr{Cint}
            ), model.ptr_model, dray, hasray)
    if ret != 0
        throw(XpressError(model))
    end

    if hasray[1] == 0
        Base.warn("Xpress solver was unable to provide an infeasibility ray")
        return dray
    end

    return dray
end

function getprimalray(model::Model)

    pray = Array(Float64, num_vars(model))

    hasray = Array(Cint, 1)

    ret = @xprs_ccall(getprimalray, Cint,
        (Ptr{Void},
         Ptr{Float64},
         Ptr{Cint}
            ), model.ptr_model, pray, hasray)
    if ret != 0
        throw(XpressError(model))
    end

    if hasray[1] == 0
        Base.warn("Xpress solver was unable to provide an unboundedness ray")
        return pray
    end

    return pray
end


text:
# Force the random number generator to be the same for all runs of the script
using Random
Random.seed!(777)

include("global_curved.jl")

import PGFPlots

# This checks that the local solve is positive definite with random variable
# coefficients
let
  τs = 1

  # Size of the grid
  Nr = Ns = 3 .* (3, 5, 7) .- 1
  SBP_orders = (2, 4, 6)

  # Create the random variable coefficients
  metrics = ntuple(i -> create_metrics(SBP_orders[i], Nr[i], Ns[i]),
                   length(SBP_orders))
  crr = ntuple(i->metrics[i].crr, length(SBP_orders))
  css = ntuple(i->metrics[i].css, length(SBP_orders))
  crs = ntuple(i->metrics[i].crs, length(SBP_orders))

  num_samp = 1000
  min_eig = zeros(length(SBP_orders), 2, num_samp)
  for k = 1:num_samp
    k % 10 == 1 && println("sample $k of $num_samp")

    for (i, SBPp) in enumerate(SBP_orders)
      # Build coefficients which are positive definite
      λ1 = rand(Nr[i]+1, Ns[i]+1)
      λ2 = rand(Nr[i]+1, Ns[i]+1) / 10000

      q  = π * rand(Nr[i]+1, Ns[i]+1)
      @. crr[i] = λ1 * cos(q)^2 + λ2 * sin(q)^2
      @. css[i] = λ1 * sin(q)^2 + λ2 * cos(q)^2
      @. crs[i] = (λ2 - λ1) * cos(q) * sin(q)

      lop = locoperator(SBPp, Nr[i], Ns[i], metrics[i]; τscale = τs)
      eigs = eigen(Matrix(lop.M̃))
      min_eig[i, 1, k] = minimum(eigs.values)
      min_eig[i, 1, k] < 0 && @show (1, k, SBPp, min_eig[i, 1, k])

      lop = locoperator(SBPp, Nr[i], Ns[i], metrics[i],
                        (BC_DIRICHLET, BC_NEUMANN, BC_NEUMANN, BC_NEUMANN);
                        τscale = τs)
      eigs = eigen(Matrix(lop.M̃))
      min_eig[i, 2, k] = minimum(eigs.values)
      min_eig[i, 2, k] < 0 && @show (2, k, SBPp, min_eig[i, 2, k])
    end
  end

  # Dirichlet
  Lx = (1, num_samp)
  Ly = extrema(min_eig[:, 1, :])
  Ly = (floor(10000 * Ly[1]) / 10000, ceil(10000 * Ly[2]) / 10000)
  plt = Plot(BrailleCanvas(40, 20,
                           origin_x = Lx[1], origin_y = Ly[1],
                           width = Lx[2] - Lx[1], height = Ly[2] - Ly[1]))

  annotate!(plt, :l, nrows(plt.graphics), string(Ly[1]), color = :light_black)
  annotate!(plt, :l, 1, string(Ly[2]), color = :light_black)
  annotate!(plt, :bl, string(Lx[1]), color = :light_black)
  annotate!(plt, :br, string(Lx[2]), color = :light_black)
  for (i, SBPp) = enumerate(SBP_orders)
    lineplot!(plt, min_eig[i, 1, :])
  end
  display(plt)

  pgf_axis = PGFPlots.Axis(style="width=5cm, height=5cm", xlabel="realization",
                           ylabel=PGFPlots.L"$\min \lambda$", xmin = 1,
                           xmax = num_samp)
  for (i, SBPp) = enumerate(SBP_orders)
    push!(pgf_axis, PGFPlots.Linear(1:num_samp, min_eig[i, 1, :],
                                    style="no marks, thick",
                                    legendentry = "order = $SBPp"))
  end
  PGFPlots.save("local_eigenvalues_Dirichlet.tikz", pgf_axis)

  # Dirichlet & Neumann
  Lx = (1, num_samp)
  Ly = extrema(min_eig[:, 2, :])
  Ly = (floor(10000 * Ly[1]) / 10000, ceil(10000 * Ly[2]) / 10000)
  plt = Plot(BrailleCanvas(40, 20,
                           origin_x = Lx[1], origin_y = Ly[1],
                           width = Lx[2] - Lx[1], height = Ly[2] - Ly[1]))

  annotate!(plt, :l, nrows(plt.graphics), string(Ly[1]), color = :light_black)
  annotate!(plt, :l, 1, string(Ly[2]), color = :light_black)
  annotate!(plt, :bl, string(Lx[1]), color = :light_black)
  annotate!(plt, :br, string(Lx[2]), color = :light_black)
  for (i, SBPp) = enumerate(SBP_orders)
    lineplot!(plt, min_eig[i, 2, :])
  end
  display(plt)

  pgf_axis = PGFPlots.Axis(style="width=5cm, height=5cm", xlabel="realization",
                           ylabel=PGFPlots.L"$\min \lambda$", xmin = 1,
                           xmax = num_samp)
  for (i, SBPp) = enumerate(SBP_orders)
    push!(pgf_axis, PGFPlots.Linear(1:num_samp, min_eig[i, 2, :],
                                    style="no marks, thick",
                                    legendentry = "order = $SBPp"))
  end
  PGFPlots.save("local_eigenvalues_Neumann.tikz", pgf_axis)
end

# This checks that the local solve is positive definite with random variable
# coefficients
let
  # Size of the grid
  Nr = Ns = 3 .* (3, 5, 7) .- 1
  SBP_orders = (2, 4, 6)

  # Create the random variable coefficients
  metrics = ntuple(i -> create_metrics(SBP_orders[i], Nr[i], Ns[i]),
                   length(SBP_orders))
  crr = ntuple(i->metrics[i].crr, length(SBP_orders))
  css = ntuple(i->metrics[i].css, length(SBP_orders))
  crs = ntuple(i->metrics[i].crs, length(SBP_orders))
  metrics = ntuple(i -> create_metrics(SBP_orders[i], Nr[i], Ns[i]),
                   length(SBP_orders))
  crr = ntuple(i->metrics[i].crr, length(SBP_orders))
  css = ntuple(i->metrics[i].css, length(SBP_orders))
  crs = ntuple(i->metrics[i].crs, length(SBP_orders))

  for (i, SBPp) in enumerate(SBP_orders)
    # Build coefficients which are positive definite
    λ1 = rand(Nr[i]+1, Ns[i]+1)
    λ2 = rand(Nr[i]+1, Ns[i]+1)

    q  = rand(Nr[i]+1, Ns[i]+1)
    @. crr[i] = λ1 * cos(q)^2 + λ2 * sin(q)^2
    @. css[i] = λ1 * sin(q)^2 + λ2 * cos(q)^2
    @. crs[i] = (λ2 - λ1) * cos(q) * sin(q)
  end

  τscales = exp10.(range(-2.;stop=2, length=100))
  max_eig = zeros(length(SBP_orders), length(τscales))
  min_eig = zeros(length(SBP_orders), length(τscales))

  for (k, τs) in enumerate(τscales)
    k % 10 == 1 && println("sample $k of $(length(τscales))")

    for (i, SBPp) in enumerate(SBP_orders)
      lop = locoperator(SBPp, Nr[i], Ns[i], metrics[i]; τscale = τs)
      eigs = eigen(Matrix(lop.M̃))
      max_eig[i,k] = maximum(eigs.values)
      min_eig[i,k] = minimum(eigs.values)
    end
  end

  Lx = (τscales[1], τscales[end])
  Ly = extrema(real.(max_eig))
  Ly = (floor(Ly[1]), ceil(Ly[2]))
  plt_max = Plot(BrailleCanvas(40, 20,
                               origin_x = Lx[1], origin_y = Ly[1],
                               width = Lx[2] - Lx[1], height = Ly[2] - Ly[1]))

  annotate!(plt_max, :l, nrows(plt_max.graphics), string(Ly[1]), color = :light_black)
  annotate!(plt_max, :l, 1, string(Ly[2]), color = :light_black)
  annotate!(plt_max, :bl, string(Lx[1]), color = :light_black)
  annotate!(plt_max, :br, string(Lx[2]), color = :light_black)
  for (i, SBPp) = enumerate(SBP_orders)
    lineplot!(plt_max, τscales, max_eig[i,:])
  end
  display(plt_max)

  pgf_axis = PGFPlots.Axis(style="width=5cm, height=5cm",
                           legendPos="north west",
                           xlabel=PGFPlots.L"$\tau_{s}$",
                           ylabel=PGFPlots.L"$\min \lambda$",
                           xmin = τscales[1],
                           xmax = τscales[end])
  for (i, SBPp) = enumerate(SBP_orders)
    push!(pgf_axis, PGFPlots.Linear(τscales, max_eig[i, :],
                                    style="no marks, thick",
                                    legendentry = "order = $SBPp"))
  end
  PGFPlots.save("tau_scaling_max_eig.tikz", pgf_axis)

  Lx = (τscales[1], τscales[end])
  Ly = extrema(min_eig)
  Ly = (floor(1000*Ly[1])/1000, ceil(1000*Ly[2])/1000)
  plt_max = Plot(BrailleCanvas(40, 20,
                               origin_x = Lx[1], origin_y = Ly[1],
                               width = Lx[2] - Lx[1], height = Ly[2] - Ly[1]))

  annotate!(plt_max, :l, nrows(plt_max.graphics), string(Ly[1]), color = :light_black)
  annotate!(plt_max, :l, 1, string(Ly[2]), color = :light_black)
  annotate!(plt_max, :bl, string(Lx[1]), color = :light_black)
  annotate!(plt_max, :br, string(Lx[2]), color = :light_black)
  for (i, SBPp) = enumerate(SBP_orders)
    lineplot!(plt_max, τscales, min_eig[i,:])
  end
  display(plt_max)

  pgf_axis = PGFPlots.Axis(style="width=5cm, height=5cm",
                           legendPos="north west",
                           xlabel=PGFPlots.L"$\tau_{s}$",
                           ylabel=PGFPlots.L"$\max \lambda$",
                           xmin = τscales[1],
                           xmax = τscales[end])
  for (i, SBPp) = enumerate(SBP_orders)
    push!(pgf_axis, PGFPlots.Linear(τscales, min_eig[i, :],
                                    style="no marks, thick",
                                    legendentry = "order = $SBPp"))
  end
  PGFPlots.save("tau_scaling_min_eig.tikz", pgf_axis)
end

ext: jl


text:
module Filters

using LinearAlgebra, GaussQuadrature, GPUifyLoops
using ..Grids

export AbstractSpectralFilter, AbstractFilter
export ExponentialFilter, CutoffFilter, TMARFilter

abstract type AbstractFilter end
abstract type AbstractSpectralFilter <: AbstractFilter end

include("Filters_kernels.jl")

"""
    spectral_filter_matrix(r, Nc, σ)

Returns the filter matrix that takes function values at the interpolation
`N+1` points, `r`, converts them into Legendre polynomial basis coefficients,
multiplies
```math
σ((n-N_c)/(N-N_c))
```
against coefficients `n=Nc:N` and evaluates the resulting polynomial at the
points `r`.
"""
function spectral_filter_matrix(r, Nc, σ)
  N = length(r)-1
  T = eltype(r)

  @assert N >= 0
  @assert 0 <= Nc <= N

  a, b = GaussQuadrature.legendre_coefs(T, N)
  V = GaussQuadrature.orthonormal_poly(r, a, b)

  Σ = ones(T, N+1)
  Σ[(Nc:N).+1] .= σ.(((Nc:N).-Nc)./(N-Nc))

  V*Diagonal(Σ)/V
end

"""
    ExponentialFilter(grid, Nc=0, s=32, α=-log(eps(eltype(grid))))

Returns the spectral filter with the filter function
```math
σ(η) = \exp(-α η^s)
```
where `s` is the filter order (must be even), the filter starts with
polynomial order `Nc`, and `alpha` is a parameter controlling the smallest
value of the filter function.
"""
struct ExponentialFilter <: AbstractSpectralFilter
  "filter matrix"
  filter

  function ExponentialFilter(grid, Nc=0, s=32,
                             α=-log(eps(eltype(grid))))
    AT = arraytype(grid)
    N = polynomialorder(grid)
    ξ = referencepoints(grid)

    @assert iseven(s)
    @assert 0 <= Nc <= N

    σ(η) = exp(-α*η^s)
    filter = spectral_filter_matrix(ξ, Nc, σ)

    new(AT(filter))
  end
end

"""
    CutoffFilter(grid, Nc=polynomialorder(grid))

Returns the spectral filter that zeros out polynomial modes greater than or
equal to `Nc`.
"""
struct CutoffFilter <: AbstractSpectralFilter
  "filter matrix"
  filter

  function CutoffFilter(grid, Nc=polynomialorder(grid))
    AT = arraytype(grid)
    ξ = referencepoints(grid)

    σ(η) = 0
    filter = spectral_filter_matrix(ξ, Nc, σ)

    new(AT(filter))
  end
end

"""
    TMARFilter()

Returns the truncation and mass aware rescaling nonnegativity preservation
filter.  The details of this filter are described in

    @article{doi:10.1175/MWR-D-16-0220.1,
      author = {Light, Devin and Durran, Dale},
      title = {Preserving Nonnegativity in Discontinuous Galerkin
               Approximations to Scalar Transport via Truncation and Mass
               Aware Rescaling (TMAR)},
      journal = {Monthly Weather Review},
      volume = {144},
      number = {12},
      pages = {4771-4786},
      year = {2016},
      doi = {10.1175/MWR-D-16-0220.1},
    }

Note this needs to be used with a restrictive time step or a flux correction
to ensure that grid integral is conserved.

## Examples

This filter can be applied to the 3rd and 4th fields of an `MPIStateArray` `Q`
with the code

```julia
Filters.apply!(Q, (3, 4), grid, TMARFilter())
```

where `grid` is the associated `DiscontinuousSpectralElementGrid`.
"""
struct TMARFilter <: AbstractFilter
end

"""
    apply!(Q, states, grid::DiscontinuousSpectralElementGrid,
           filter::AbstractSpectralFilter; horizontal = true,
           vertical = true)

Applies `filter` to the `states` of `Q`.

The arguments `horizontal` and `vertical` are used to control if the filter
is applied in the horizontal and vertical reference directions, respectively.
Note, it is assumed that the trailing dimension is the vertical dimension and
the rest are horizontal.
"""
function apply!(Q, states, grid::DiscontinuousSpectralElementGrid,
                filter::AbstractSpectralFilter;
                horizontal = true, vertical = true)
  topology = grid.topology

  dim = dimensionality(grid)
  N = polynomialorder(grid)

  nstate = size(Q, 2)

  filtermatrix = filter.filter
  device = typeof(Q.data) <: Array ? CPU() : CUDA()

  nelem = length(topology.elems)
  Nq = N + 1
  Nqk = dim == 2 ? 1 : Nq

  nrealelem = length(topology.realelems)

  @launch(device, threads=(Nq, Nq, Nqk), blocks=nrealelem,
          knl_apply_filter!(Val(dim), Val(N), Val(nstate), Val(horizontal), Val(vertical),
                            Q.data, Val(states), filtermatrix, topology.realelems))
end

"""
    apply!(Q, states, grid::DiscontinuousSpectralElementGrid, ::TMARFilter)

Applies the truncation and mass aware rescaling to `states` of `Q`.  This
rescaling keeps the states nonegative while keeping the element average
the same.
"""
function apply!(Q, states, grid::DiscontinuousSpectralElementGrid,
                ::TMARFilter)
  topology = grid.topology

  device = typeof(Q.data) <: Array ? CPU() : CUDA()

  dim = dimensionality(grid)
  N = polynomialorder(grid)
  Nq = N + 1
  Nqk = dim == 2 ? 1 : Nq

  nrealelem = length(topology.realelems)
  nreduce = 2^ceil(Int, log2(Nq*Nqk))

  @launch(device, threads=(Nq, Nqk, 1), blocks=nrealelem,
          knl_apply_TMAR_filter!(Val(nreduce), Val(dim), Val(N), Q.data,
                                 Val(states), grid.vgeo, topology.realelems))
end

end

max_stars_repo_name: akshaysridhar/CLIMA

max_issues_repo_name: akshaysridhar/CLIMA

max_forks_repo_name: akshaysridhar/CLIMA


text:
__precompile__()

module AstroImages

using FITSIO, FileIO, Images, Interact, Reproject, WCS, MappedArrays

export load, AstroImage, ccd2rgb, set_brightness!, set_contrast!, add_label!, reset!

_load(fits::FITS, ext::Int) = read(fits[ext])
_load(fits::FITS, ext::NTuple{N, Int}) where {N} = ntuple(i-> read(fits[ext[i]]), N)
_load(fits::NTuple{N, FITS}, ext::NTuple{N, Int}) where {N} = ntuple(i -> _load(fits[i], ext[i]), N)

_header(fits::FITS, ext::Int) = WCS.from_header(read_header(fits[ext], String))[1]
_header(fits::FITS, ext::NTuple{N, Int}) where {N} =
    ntuple(i -> WCS.from_header(read_header(fits[ext[i]], String))[1], N)
_header(fits::NTuple{N, FITS}, ext::NTuple{N, Int}) where {N} =
    ntuple(i -> _header(fits[i], ext[i]), N)
"""
    load(fitsfile::String, n=1)

Read and return the data from `n`-th extension of the FITS file.

Second argument can also be a tuple of integers, in which case a
tuple with the data of each corresponding extension is returned.
"""
function FileIO.load(f::File{format"FITS"}, ext::Int=1)
    fits = FITS(f.filename)
    out = _load(fits, ext)
    header = _header(fits,ext)
    close(fits)
    return out, header
end

function FileIO.load(f::File{format"FITS"}, ext::NTuple{N,Int}) where {N}
    fits = FITS(f.filename)
    out = _load(fits, ext)
    header = _header(fits, ext)
    close(fits)
    return out, header
end

function FileIO.load(f::NTuple{N, String}) where {N}
    fits = ntuple(i-> FITS(f[i]), N)
    ext = indexer(fits)
    out = _load(fits, ext)
    header = _header(fits, ext)
    for i in 1:N
        close(fits[i])
    end
    return out, header
end

function indexer(fits::FITS)
    ext = 0
    for (i, hdu) in enumerate(fits)
        if hdu isa ImageHDU && length(size(hdu)) >= 2	# check if Image is atleast 2D
            ext = i
            break
        end
    end
    if ext > 1
       	@info "Image was loaded from HDU $ext"
    elseif ext == 0
        error("There are no ImageHDU extensions in '$(fits.filename)'")
    end
    return ext
end
indexer(fits::NTuple{N, FITS}) where {N} = ntuple(i -> indexer(fits[i]), N)

# Images.jl expects data to be either a float or a fixed-point number.  Here we define some
# utilities to convert all data types supported by FITS format to float or fixed-point:
#
#   * Float numbers are left as they are
#   * Unsigned integers are mapped to [0, 1] with Normed type
#   * Signed integers are mapped to unsigned integers and then to Normed
_float(x::AbstractFloat) = x
for n in (8, 16, 32, 64)
    SIT = Symbol("Int", n) # signed integer type
    UIT = Symbol("UInt", n) # unsigned integer type
    NIT = Symbol("N0f", n) # fixed-point type for unsigned float
    @eval maxint = $UIT(big(2) ^ ($n - 1)) #
    @eval begin
        _float(x::$UIT) = reinterpret($NIT, x)
        _float(x::$SIT) = _float(xor(reinterpret($UIT, x), $maxint))
    end
end

mutable struct Properties{P <: Union{AbstractFloat, FixedPoint}}
    rgb_image::MappedArrays.MultiMappedArray{RGB{P},2,Tuple{Array{P,2},Array{P,2},Array{P,2}},Type{RGB{P}},typeof(ImageCore.extractchannels)}
    contrast::Float64
    brightness::Float64
    label::Array{Tuple{Tuple{Float64,Float64},String},1}
    function Properties{P}(;kvs...) where P
        obj = new{P}()
        obj.contrast = 1.0
        obj.brightness = 0.0
        obj.label = Array{Tuple{Tuple{Float64,Float64},String}}(undef,0)
        for (k,v) in kvs
            setproperty!(obj, k, v)
        end
        return obj
    end
end

struct AstroImage{T<:Real,C<:Color, N, P}
    data::NTuple{N, Matrix{T}}
    minmax::NTuple{N, Tuple{T,T}}
    wcs::NTuple{N, WCSTransform}
    property::Properties{P}
end

"""
    AstroImage([color=Gray,] data::Matrix{Real})
    AstroImage(color::Type{<:Color}, data::NTuple{N, Matrix{T}}) where {T<:Real, N}

Construct an `AstroImage` object of `data`, using `color` as color map, `Gray` by default.
"""
AstroImage(color::Type{<:Color}, data::Matrix{T}, wcs::WCSTransform) where {T<:Real} =
    AstroImage{T,color, 1, Float64}((data,), (extrema(data),), (wcs,), Properties{Float64}())
function AstroImage(color::Type{<:AbstractRGB}, data::NTuple{N, Matrix{T}}, wcs::NTuple{N, WCSTransform}) where {T <: Union{AbstractFloat, FixedPoint}, N}
    if N == 3
        img = ccd2rgb((data[1], wcs[1]),(data[2], wcs[2]),(data[3], wcs[3]))
        return AstroImage{T,color,N, widen(T)}(data, ntuple(i -> extrema(data[i]), N), wcs, Properties{widen(T)}(rgb_image = img))
    end
end
function AstroImage(color::Type{<:AbstractRGB}, data::NTuple{N, Matrix{T}}, wcs::NTuple{N, WCSTransform}) where {T<:Real, N}
    if N == 3
        img = ccd2rgb((data[1], wcs[1]),(data[2], wcs[2]),(data[3], wcs[3]))
        return AstroImage{T,color,N, Float64}(data, ntuple(i -> extrema(data[i]), N), wcs, Properties{Float64}(rgb_image = img))
    end
end
function AstroImage(color::Type{<:Color}, data::NTuple{N, Matrix{T}}, wcs::NTuple{N, WCSTransform}) where {T<:Real, N}
    return AstroImage{T,color, N, Float64}(data, ntuple(i -> extrema(data[i]), N), wcs, Properties{Float64}())
end
AstroImage(data::Matrix{T}) where {T<:Real} = AstroImage{T,Gray,1, Float64}((data,), (extrema(data),), (WCSTransform(2),), Properties{Float64}())
AstroImage(data::NTuple{N, Matrix{T}}) where {T<:Real, N} = AstroImage{T,Gray,N, Float64}(data, ntuple(i -> extrema(data[i]), N), ntuple(i-> WCSTransform(2), N), Properties{Float64}())
AstroImage(data::Matrix{T}, wcs::WCSTransform) where {T<:Real} = AstroImage{T,Gray,1, Float64}((data,), (extrema(data),), (wcs,), Properties{Float64}())
AstroImage(data::NTuple{N, Matrix{T}}, wcs::NTuple{N, WCSTransform}) where {T<:Real, N} = AstroImage{T,Gray,N, Float64}(data, ntuple(i -> extrema(data[i]), N), wcs, Properties{Float64}())

"""
    AstroImage([color=Gray,] filename::String, n::Int=1)
    AstroImage(color::Type{<:Color}, file::String, n::NTuple{N, Int}) where {N}

Create an `AstroImage` object by reading the `n`-th extension from FITS file `filename`.

Use `color` as color map, this is `Gray` by default.
"""
AstroImage(color::Type{<:Color}, file::String, ext::Int) =
    AstroImage(color, file, (ext,))
AstroImage(color::Type{<:Color}, file::String, ext::NTuple{N, Int}) where {N} =
    AstroImage(color, load(file, ext)...)

AstroImage(file::String, ext::Int) = AstroImage(Gray, file, ext)
AstroImage(file::String, ext::NTuple{N, Int}) where {N} = AstroImage(Gray, file, ext)

AstroImage(color::Type{<:Color}, fits::FITS, ext::Int) =
    AstroImage(color, _load(fits, ext), WCS.from_header(read_header(fits[ext],String))[1])
AstroImage(color::Type{<:Color}, fits::FITS, ext::NTuple{N, Int}) where {N} =
    AstroImage(color, _load(fits, ext), ntuple(i -> WCS.from_header(read_header(fits[ext[i]], String))[1], N))
AstroImage(color::Type{<:Color}, fits::NTuple{N, FITS}, ext::NTuple{N, Int}) where {N} =
    AstroImage(color, ntuple(i -> _load(fits[i], ext[i]), N), ntuple(i -> WCS.from_header(read_header(fits[i][ext[i]], String))[1], N))

AstroImage(files::NTuple{N,String}) where {N} =
    AstroImage(Gray, load(files)...)
AstroImage(color::Type{<:Color}, files::NTuple{N,String}) where {N} =
    AstroImage(color, load(files)...)
AstroImage(file::String) = AstroImage((file,))

# Lazily reinterpret the image as a Matrix{Color}, upon request.
function render(img::AstroImage{T,C,N}, header_number = 1) where {T,C,N}
    imgmin, imgmax = extrema(img.minmax[header_number])
    # Add one to maximum to work around this issue:
    # https://github.com/JuliaMath/FixedPointNumbers.jl/issues/102
    f = scaleminmax(_float(imgmin), _float(max(imgmax, imgmax + one(T))))
    return colorview(C, f.(_float.(img.data[header_number])))
end

"""
    set_brightness!(img::AstroImage, value::AbstractFloat)

Sets brightness of `rgb_image` to value.
"""
function set_brightness!(img::AstroImage, value::AbstractFloat)
    if isdefined(img.property, :rgb_image)
        diff = value - img.property.brightness
        img.property.brightness = value
        img.property.rgb_image .+= RGB{typeof(value)}(diff, diff, diff)
    else
        throw(DomainError(value, "Can't apply operation. AstroImage dosen't contain :rgb_image"))
    end
end

"""
    set_contrast!(img::AstroImage, value::AbstractFloat)

Sets contrast of rgb_image to value.
"""
function set_contrast!(img::AstroImage, value::AbstractFloat)
    if isdefined(img.property, :rgb_image)
        diff = (value / img.property.contrast)
        img.property.contrast = value
        img.property.rgb_image = colorview(RGB, red.(img.property.rgb_image) .* diff, green.(img.property.rgb_image) .* diff,
                                                blue.(img.property.rgb_image) .* diff)
    else
        throw(DomainError(value, "Can't apply operation. AstroImage dosen't contain :rgb_image"))
    end
end

"""
    add_label!(img::AstroImage, x::Real, y::Real, label::String)

Stores label to coordinates (x,y) in AstroImage's property label.
"""
function add_label!(img::AstroImage, x::Real, y::Real, label::String)
    push!(img.property.label, ((x,y), label))
end

"""
    reset!(img::AstroImage)

Resets AstroImage property fields.

Sets brightness to 0.0, contrast to 1.0, empties label
and form a fresh rgb_image without any brightness, contrast operations on it.
"""
function reset!(img::AstroImage{T,C,N}) where {T,C,N}
    img.property.contrast = 1.0
    img.property.brightness = 0.0
    img.property.label = []
    if N == 3 && C == RGB
        shape_out = size(img.property.rgb_image)
        img.property.rgb_image = ccd2rgb((img.data[1], img.wcs[1]),(img.data[2], img.wcs[2]),(img.data[3], img.wcs[3]),
                                            shape_out = shape_out)
    end
end

Images.colorview(img::AstroImage) = render(img)

Base.size(img::AstroImage) = Base.size.(img.data)

Base.length(img::AstroImage{T,C,N}) where {T,C,N} = N

include("showmime.jl")
include("plot-recipes.jl")
include("ccd2rgb.jl")

end # module


text:
function extend_to_dc(inFreqs, inArray)
   deltaf = inFreqs[2]-inFreqs[1]
   numextrapts = round(Int, inFreqs[1]/deltaf)
   phasechange = angle(inArray[1])/numextrapts # amount by which to reverse the phase (in radians)

   # extend the inArray down to DC, rolling down the phase
   extrapts = zeros(Complex64,numextrapts)
   for j = 1:numextrapts
      extrapts[numextrapts-j+1] = inArray[1]*e^(-im*j*phasechange)
   end
   extrapts[1,1] = abs(extrapts[1,1]) # (imag(inArray(1,1)) < 10^-12) -> 0
   outArray = vcat(extrapts, inArray)

   # extend the Freqsuencies
   newFreqs = (inFreqs[1]-numextrapts*deltaf:deltaf:inFreqs[1]-deltaf)
   outFreqs = vcat(newFreqs, inFreqs)
   return (outFreqs, outArray)
end


text:
"""
    count_loops(N::UnipartiteNetwork)

This function is counting the number of loops in a given network. Loops are trophic chains where a species is appearing more than one time.

"""
function count_loops(N::UnipartiteNetwork)

    d = maximum([round(Int, p.weight) for p in bellman_ford(N)])

    paths = Vector{Vector{EcologicalNetworks._species_type(N)}}()
    loops = Vector{Vector{EcologicalNetworks._species_type(N)}}()

    append!(paths, [[sp] for sp in species(N)])
    neighbors = Dict([sp => N[sp,:] for sp in species(N)])

    while !isempty(paths)
        initial_paths = copy(paths)
        done_paths = Vector{Vector{EcologicalNetworks._species_type(N)}}()
        new_paths = Vector{Vector{EcologicalNetworks._species_type(N)}}()
        for path in paths
            for n in neighbors[last(path)]
                if n in path
                    push!(loops, [path..., n])
                    for i in 1:length(path)
                        push!(done_paths, path[1:i])
                    end
                else
                    push!(new_paths, [path..., n])
                end
            end
        end
        unique!(done_paths)
        filter!(p -> length(p) >= d, paths)
        unique!(paths)
        if !isempty(new_paths)
            append!(paths, new_paths)
        end
        if !isempty(done_paths)
            deleteat!(paths, unique(sort(filter(!isnothing, indexin(done_paths, paths)))))
        end
        if paths == initial_paths
            break
        end
    end

    n_loops = length(unique(sort.(unique.(loops))))
    return n_loops
end


text:
function draw_poisson_panel(gui_data::GuiData, app_data::Model.AppData, sock::Comm.SocClient)
    if CImGui.TreeNode("Poisson Patterns")
        CImGui.PushItemWidth(80)

        # Poisson pattern generater properties
        gui_data.buffer = Model.prep_field(Model.firing_rate(app_data.model), 16)
        returned = CImGui.InputText("Firing Rate", gui_data.buffer, 16, CImGui.ImGuiInputTextFlags_CharsDecimal | CImGui.ImGuiInputTextFlags_EnterReturnsTrue)
        if returned
            Model.set_firing_rate!(app_data.model, gui_data.buffer)
        end
        CImGui.SameLine(250)  # Shifted 200px right

        gui_data.buffer = Model.prep_field(Model.poisson_pattern_min(app_data.model), 16)
        returned = CImGui.InputText("Pattern Min", gui_data.buffer, 16, CImGui.ImGuiInputTextFlags_CharsDecimal | CImGui.ImGuiInputTextFlags_EnterReturnsTrue)
        if returned
            Model.set_poisson_pattern_min!(app_data.model, gui_data.buffer)
        end
        CImGui.SameLine(450)

        gui_data.buffer = Model.prep_field(Model.poisson_pattern_max(app_data.model), 16)
        returned = CImGui.InputText("Pattern Max", gui_data.buffer, 16, CImGui.ImGuiInputTextFlags_CharsDecimal | CImGui.ImGuiInputTextFlags_EnterReturnsTrue)
        if returned
            Model.set_poisson_pattern_max!(app_data.model, gui_data.buffer)
        end

        # Row 2 *****************************************************
        gui_data.buffer = Model.prep_field(Model.poisson_pattern_spread(app_data.model), 16)
        returned = CImGui.InputText("Pattern Spread", gui_data.buffer, 16, CImGui.ImGuiInputTextFlags_CharsDecimal | CImGui.ImGuiInputTextFlags_EnterReturnsTrue)
        if returned
            Model.set_poisson_pattern_spread!(app_data.model, gui_data.buffer)
        end

        CImGui.PopItemWidth()

        CImGui.TreePop()
    end
end


text:
"""
This example shows how to use Term.jl to print styled text to the console.

The easiest way to apply style to a String is to use Markup Tags.

Let's say you have this string:

s = "this is my string!"

and you want 'my' to be bold and blue, then simply write:

s = "this is [bold blue]my[/bold blue] string!

The [xx]text[/xx] syntax will result in the style 'xx' being applied to the
text between parentheses.
The style can include things like: `bold`, `underline`, `italic`, `red`, `blue`...
all separated by a space. You can also use hexcodes (like #ff0022) to specify
the color of your text.

Once you're happy with your markup string, you can use `Term.tprintln` to print it!

Note: the second set of parentheses, [/xx], terminates the markup tag and the style
information 'xx' must be identical between the start/end parentheses.
"""

import Term: tprintln

my_string = "[bold underline]This[/bold underline] is [red italic]my[/red italic] [on_green black bold]string[/on_green black bold]"
tprintln(my_string)

"""
Using markup tags is great when you want your styled text to be printed within other `Renderable` object
from Term such as `TextBox` and `Panel` (see examples).

But if you just need to quickly print something with a bit of color, you don't want to be typing all those [,].
Fortunately, Term provides a set of macros to easily print styled strings!!
"""

import Term: @red, @blue, @style

println(@red "This is red")
println(@blue "and this is blue")

# or, if you want some more control
println(@style "this is my beautifully styled text" green italic bold)


text:
using MuZero
using Test

@testset "Connect4" begin
    g = Connect4()
    s = get_state(g)
    m = get_marker(g, 4)
    @test Connect4State() |> typeof == Connect4State
    @test Connect4StateSVG(Connect4State()) |> typeof == Connect4StateSVG
    @test get_coordinate(m) == (6, 4)
    @test get_piece(m) == get_piece(g, :black)
    @test next_state(g, 4) |> typeof == Connect4State
    @test Connect4([4,1,2,4]) |> typeof == Connect4
    @test encode(s) |> typeof == String
    @test step!(g, 4) |> typeof == Float32
end

@testset "Sample game" begin
    g = Connect4()
    s = get_state(g)

    @test to_play(g) == PLAYER_1
    @test get_piece(g, Action(4)) == get_piece(g, :black)

    step!(g, 2)
    @test to_play(g) == PLAYER_2
    @test get_piece(g, Action(4)) == get_piece(g, :red)

    step!(g, 2)
    @test to_play(g) == PLAYER_1
    @test get_piece(g, Action(4)) == get_piece(g, :black)
end

@testset "Coord/ByteIdx inverse test" begin
    g = Connect4()
    s = get_state(g)
    pairs = []
    for row in 1:num_rows(s), col in 1:num_cols(s)
        coord = (row,col)
        byte_idx = coord_to_byte_idx(typeof(s), coord)
        push!(pairs, (byte_idx=byte_idx, coord=coord))
    end
    for pair in pairs
        @test byte_idx_to_coord(typeof(s), pair.byte_idx) == pair.coord
    end
end


text:
export NAC

"""
    NAC(in::Int, out::Int; initW=glorot_uniform, initM=glorot_uniform)

Neural Accumulator as proposed in https://arxiv.org/abs/1808.00508
"""
struct NAC{T<:AbstractMatrix}
    W::T
    M::T
end

function NAC(in::Int, out::Int;
             initW=glorot_uniform, initM=glorot_uniform)
    W = initW(out, in)
    M = initM(out, in)
    NAC(W, M)
end

weights(nac::NAC) = tanh.(nac.W) .* σ.(nac.M)

function (nac::NAC)(x)
    _W = weights(nac)
    _W*x
end

function Base.show(io::IO, l::NAC)
    in = size(l.W, 2)
    out = size(l.W, 1)
    print(io, "NAC(in=$in, out=$out)")
end

Flux.@functor NAC


text:
using QuantileRegressions
using DataFrames, CSV, LinearAlgebra
using Test

include("engel_test.jl")
include("npqreg.jl")


text:
using JuMP, EAGO

                     m = Model()

                     EAGO.register_eago_operators!(m)

                     @variable(m, -1 <= x[i=1:3] <= 1)

                     @variable(m, -13.110901270484293 <= q <= 11.380819522292354)

                     add_NL_constraint(m, :(sigmoid(-0.7773677229963925 + -0.31823111196481735*sigmoid(0.13898320039278467 + -0.4573031975765818*sigmoid(-0.1658975405753922 + 0.8948668185465118*$(x[1]) + 0.5677868700589963*$(x[2]) + -0.004745367697401193*$(x[3])) + -0.027363793963717686*sigmoid(0.8084147176169236 + 0.6887886991760568*$(x[1]) + 0.9976411528938542*$(x[2]) + 0.9554800615635051*$(x[3])) + -0.7118578003946046*sigmoid(-0.23665048660081967 + 0.22941639672835912*$(x[1]) + 0.6694807543091978*$(x[2]) + 0.08612932953468189*$(x[3]))) + -0.05869662098445039*sigmoid(0.7177153834386609 + 0.046197114005355644*sigmoid(-0.1658975405753922 + 0.8948668185465118*$(x[1]) + 0.5677868700589963*$(x[2]) + -0.004745367697401193*$(x[3])) + 0.8341427099559482*sigmoid(0.8084147176169236 + 0.6887886991760568*$(x[1]) + 0.9976411528938542*$(x[2]) + 0.9554800615635051*$(x[3])) + -0.5961814570241306*sigmoid(-0.23665048660081967 + 0.22941639672835912*$(x[1]) + 0.6694807543091978*$(x[2]) + 0.08612932953468189*$(x[3]))) + -0.31920840200501477*sigmoid(-0.05621054225465549 + -0.4685096496776473*sigmoid(-0.1658975405753922 + 0.8948668185465118*$(x[1]) + 0.5677868700589963*$(x[2]) + -0.004745367697401193*$(x[3])) + -0.723697578017108*sigmoid(0.8084147176169236 + 0.6887886991760568*$(x[1]) + 0.9976411528938542*$(x[2]) + 0.9554800615635051*$(x[3])) + 0.643520408356284*sigmoid(-0.23665048660081967 + 0.22941639672835912*$(x[1]) + 0.6694807543091978*$(x[2]) + 0.08612932953468189*$(x[3])))) + sigmoid(-0.2800668423712529 + 0.5241807567635033*sigmoid(0.13898320039278467 + -0.4573031975765818*sigmoid(-0.1658975405753922 + 0.8948668185465118*$(x[1]) + 0.5677868700589963*$(x[2]) + -0.004745367697401193*$(x[3])) + -0.027363793963717686*sigmoid(0.8084147176169236 + 0.6887886991760568*$(x[1]) + 0.9976411528938542*$(x[2]) + 0.9554800615635051*$(x[3])) + -0.7118578003946046*sigmoid(-0.23665048660081967 + 0.22941639672835912*$(x[1]) + 0.6694807543091978*$(x[2]) + 0.08612932953468189*$(x[3]))) + 0.9218858277065758*sigmoid(0.7177153834386609 + 0.046197114005355644*sigmoid(-0.1658975405753922 + 0.8948668185465118*$(x[1]) + 0.5677868700589963*$(x[2]) + -0.004745367697401193*$(x[3])) + 0.8341427099559482*sigmoid(0.8084147176169236 + 0.6887886991760568*$(x[1]) + 0.9976411528938542*$(x[2]) + 0.9554800615635051*$(x[3])) + -0.5961814570241306*sigmoid(-0.23665048660081967 + 0.22941639672835912*$(x[1]) + 0.6694807543091978*$(x[2]) + 0.08612932953468189*$(x[3]))) + 0.0068491484591501894*sigmoid(-0.05621054225465549 + -0.4685096496776473*sigmoid(-0.1658975405753922 + 0.8948668185465118*$(x[1]) + 0.5677868700589963*$(x[2]) + -0.004745367697401193*$(x[3])) + -0.723697578017108*sigmoid(0.8084147176169236 + 0.6887886991760568*$(x[1]) + 0.9976411528938542*$(x[2]) + 0.9554800615635051*$(x[3])) + 0.643520408356284*sigmoid(-0.23665048660081967 + 0.22941639672835912*$(x[1]) + 0.6694807543091978*$(x[2]) + 0.08612932953468189*$(x[3])))) + sigmoid(-0.0923150493931284 + -0.18688964829603716*sigmoid(0.13898320039278467 + -0.4573031975765818*sigmoid(-0.1658975405753922 + 0.8948668185465118*$(x[1]) + 0.5677868700589963*$(x[2]) + -0.004745367697401193*$(x[3])) + -0.027363793963717686*sigmoid(0.8084147176169236 + 0.6887886991760568*$(x[1]) + 0.9976411528938542*$(x[2]) + 0.9554800615635051*$(x[3])) + -0.7118578003946046*sigmoid(-0.23665048660081967 + 0.22941639672835912*$(x[1]) + 0.6694807543091978*$(x[2]) + 0.08612932953468189*$(x[3]))) + 0.4699561510424055*sigmoid(0.7177153834386609 + 0.046197114005355644*sigmoid(-0.1658975405753922 + 0.8948668185465118*$(x[1]) + 0.5677868700589963*$(x[2]) + -0.004745367697401193*$(x[3])) + 0.8341427099559482*sigmoid(0.8084147176169236 + 0.6887886991760568*$(x[1]) + 0.9976411528938542*$(x[2]) + 0.9554800615635051*$(x[3])) + -0.5961814570241306*sigmoid(-0.23665048660081967 + 0.22941639672835912*$(x[1]) + 0.6694807543091978*$(x[2]) + 0.08612932953468189*$(x[3]))) + -0.26284635568255643*sigmoid(-0.05621054225465549 + -0.4685096496776473*sigmoid(-0.1658975405753922 + 0.8948668185465118*$(x[1]) + 0.5677868700589963*$(x[2]) + -0.004745367697401193*$(x[3])) + -0.723697578017108*sigmoid(0.8084147176169236 + 0.6887886991760568*$(x[1]) + 0.9976411528938542*$(x[2]) + 0.9554800615635051*$(x[3])) + 0.643520408356284*sigmoid(-0.23665048660081967 + 0.22941639672835912*$(x[1]) + 0.6694807543091978*$(x[2]) + 0.08612932953468189*$(x[3])))) - $q <= 0.0))

                     @objective(m, Min, q)

                     return m

max_stars_repo_name: PSORLab/RSActivationFunctions

max_issues_repo_name: PSORLab/RSActivationFunctions

max_forks_repo_name: PSORLab/RSActivationFunctions


text:
@compat abstract type ImplicitStack end
immutable ImplicitIndexStack{S} <: ImplicitStack
    stack::Vector{S}
end
Base.copy(s::ImplicitIndexStack) = typeof(s)(copy(s.stack))
"""
Keeps a stack of nodes and their corresponding indices. Note that the last node
is not explicitly stored in the node_stack, such that length(node_stack) ==
length(idx_stack)-1 (unless we're at the root in which case both are empty)
"""
immutable ImplicitNodeStack{T, S} <: ImplicitStack
    node_stack::Vector{T}
    idx_stack::ImplicitIndexStack{S}
end
Base.copy(s::ImplicitNodeStack) = typeof(s)(copy(s.node_stack), copy(s.idx_stack))
Base.isempty(s::ImplicitNodeStack) = isempty(s.idx_stack)
Base.isempty(s::ImplicitIndexStack) = isempty(s.stack)
getnode(tree, ns::ImplicitIndexStack) = isempty(ns.stack) ? tree[rootstate(tree)] :
    (@assert isa(treekind(tree), IndexedTree); tree[ns.stack[end]])
function getnode(tree, stack::ImplicitNodeStack)
    isempty(stack.node_stack) ?
    (isempty(stack.idx_stack) ? tree : children(tree)[stack.idx_stack.stack[end]]) :
    children(stack.node_stack[end])[stack.idx_stack.stack[end]]
end
immutable ImplicitChildStates{T, S}
    tree::T
    stack::S
end
Base.iteratorsize{T<:ImplicitChildStates}(::Type{T}) = Base.SizeUnknown()
children(states::ImplicitChildStates) = children(states.tree, states.stack)

parentstate(tree, state::ImplicitNodeStack) =
    ImplicitNodeStack(state.node_stack[1:end-1], parentstate(tree, state.idx_stack))
parentstate(tree, state::ImplicitIndexStack) = ImplicitIndexStack(state.stack[1:end-1])
childstates(tree, state) = childstates(tree, state, treekind(tree))
childstates(tree, state::ImplicitStack) = ImplicitChildStates(tree, state)
children(tree, state::ImplicitNodeStack) = children(tree, getnode(tree, state))
children(tree, state::ImplicitIndexStack) = isempty(state) ?
    children(tree) : children(tree, tree[state.stack[end]])

childstates(s::ImplicitChildStates) = isempty(s.stack) ?
  childstates(s.tree, s.tree) : childstates(s.tree,
    isa(s.stack, ImplicitNodeStack) ? getnode(s.tree, s.stack) : s.stack.stack[end])
nextind(s::ImplicitChildStates, ind) = nextind(childstates(s), ind)
start(s::ImplicitChildStates) = start(childstates(s))
function next(s::ImplicitChildStates, ind)
    ni = next(childstates(s), ind)[2]
    (update_state!(s.tree, copy(s.stack), childstates(s), ind, treekind(s.tree)), ni)
end
done(s::ImplicitChildStates, ind) = done(childstates(s), ind)

update_state!(tree, ns, cs, idx, _) = update_state!(tree, ns, cs, idx)
update_state!(tree, ns::ImplicitIndexStack, cs, idx) = (push!(ns.stack, idx); ns)
update_state!(tree, ns::ImplicitIndexStack, cs, idx, ::IndexedTree) = (push!(ns.stack, next(cs, idx)[1]); ns)
function update_state!(tree, ns::ImplicitNodeStack, cs, idx)
    !isempty(ns) && push!(ns.node_stack, getnode(tree, ns))
    update_state!(tree, ns.idx_stack, cs, idx)
    ns
end

function joinstate(tree, state::ImplicitNodeStack, new_state::ImplicitNodeStack)
    isempty(new_state) && return state
    ImplicitNodeStack([push!(copy(state.node_stack), getnode(tree, state)); new_state.node_stack],
        joinstate(tree, state.idx_stack, new_state.idx_stack))
end
joinstate(tree, state::ImplicitIndexStack, new_state::ImplicitIndexStack) =
    ImplicitIndexStack([state.stack; new_state.stack])

isroot(tree, state::ImplicitStack) = isempty(state)

max_stars_repo_name: femtocleaner/DeprecationBotTest

max_issues_repo_name: femtocleaner/DeprecationBotTest

max_forks_repo_name: femtocleaner/DeprecationBotTest


text:
# ---
# title: 1335. Minimum Difficulty of a Job Schedule
# id: problem1335
# author: Tian Jun
# date: 2020-10-31
# difficulty: Hard
# categories: Dynamic Programming
# link: <https://leetcode.com/problems/minimum-difficulty-of-a-job-schedule/description/>
# hidden: true
# ---
#
# You want to schedule a list of jobs in `d` days. Jobs are dependent (i.e To
# work on the `i-th` job, you have to finish all the jobs `j` where `0 <= j <
# i`).
#
# You have to finish **at least** one task every day. The difficulty of a job
# schedule is the sum of difficulties of each day of the `d` days. The
# difficulty of a day is the maximum difficulty of a job done in that day.
#
# Given an array of integers `jobDifficulty` and an integer `d`. The difficulty
# of the `i-th` job is `jobDifficulty[i]`.
#
# Return _the minimum difficulty_ of a job schedule. If you cannot find a
# schedule for the jobs return **-1**.
#
#
#
# **Example 1:**
#
# ![](https://assets.leetcode.com/uploads/2020/01/16/untitled.png)
#
#
#
#     Input: jobDifficulty = [6,5,4,3,2,1], d = 2
#     Output: 7
#     Explanation: First day you can finish the first 5 jobs, total difficulty = 6.
#     Second day you can finish the last job, total difficulty = 1.
#     The difficulty of the schedule = 6 + 1 = 7
#
#
# **Example 2:**
#
#
#
#     Input: jobDifficulty = [9,9,9], d = 4
#     Output: -1
#     Explanation: If you finish a job per day you will still have a free day. you cannot find a schedule for the given jobs.
#
#
# **Example 3:**
#
#
#
#     Input: jobDifficulty = [1,1,1], d = 3
#     Output: 3
#     Explanation: The schedule is one job per day. total difficulty will be 3.
#
#
# **Example 4:**
#
#
#
#     Input: jobDifficulty = [7,1,7,1,7,1], d = 3
#     Output: 15
#
#
# **Example 5:**
#
#
#
#     Input: jobDifficulty = [11,111,22,222,33,333,44,444], d = 6
#     Output: 843
#
#
#
#
# **Constraints:**
#
#   * `1 <= jobDifficulty.length <= 300`
#   * `0 <= jobDifficulty[i] <= 1000`
#   * `1 <= d <= 10`
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


text:
# Goal: Make sure that our solvers are "correct" for small steps or tolerances
# Verify this for many (ideally all) combinations
# Compare with an algorithm from OrdinaryDiffEq.jl with high precision
using ODEFilters
using Test
using OrdinaryDiffEq
using LinearAlgebra
using Statistics: mean
using DiffEqProblemLibrary.ODEProblemLibrary: importodeproblems; importodeproblems()
import DiffEqProblemLibrary.ODEProblemLibrary: prob_ode_lotkavoltera, prob_ode_fitzhughnagumo

import ODEFilters: remake_prob_with_jac

for (prob, probname) in [
    (remake_prob_with_jac(prob_ode_lotkavoltera), "lotkavolterra"),
    (remake_prob_with_jac(prob_ode_fitzhughnagumo), "fitzhughnagumo"),
]
    @testset "Constant steps: $probname" begin

        true_sol = solve(remake(prob, u0=big.(prob.u0)), Tsit5(), abstol=1e-20, reltol=1e-20)

        for Alg in (EK0, EK1),
            diffusion in [:fixed, :dynamic, :fixedMAP, :fixedMV, :dynamicMV],
            q in [1, 3, 5]

            if Alg == EK1 && diffusion in (:fixedMV, :dynamicMV) continue end

            @testset "Constant steps: $probname; q=$q, diffusion=$diffusion, alg=$Alg" begin

            @debug "Testing for correctness: Constant steps" probname alg diffusion q dt

            sol = solve(prob, Alg(order=q, diffusionmodel=diffusion),
                        adaptive=false, dt=5e-3)
            @test sol.u ≈ true_sol.(sol.t) rtol=1e-5
            end
        end
    end
end

for (prob, probname) in [
    (remake_prob_with_jac(prob_ode_lotkavoltera), "lotkavolterra"),
    (remake_prob_with_jac(prob_ode_fitzhughnagumo), "fitzhughnagumo"),
]
    @testset "Adaptive steps: $probname" begin

        t_eval = prob.tspan[1]:0.01:prob.tspan[end]
        true_sol = solve(remake(prob, u0=big.(prob.u0)), Tsit5(), abstol=1e-20, reltol=1e-20)
        true_dense_vals = true_sol.(t_eval)

        for Alg in (EK0, EK1),
            diffusion in [:fixed, :dynamic, :fixedMAP, :fixedMV, :dynamicMV],
            q in [2, 4, 6]

            if Alg == EK1 && diffusion in (:fixedMV, :dynamicMV) continue end

            @testset "Adaptive steps: $probname; q=$q, diffusion=$diffusion, Alg=$Alg" begin

            @debug "Testing for correctness: Adaptive steps" probname Alg diffusion q

            sol = solve(prob, Alg(order=q, diffusionmodel=diffusion),
                        adaptive=true)

            @test sol.u ≈ true_sol.(sol.t) rtol=1e-4
            @test mean.(sol.(t_eval)) ≈ true_dense_vals rtol=1e-4

            end
        end
    end
end


text:
using Optim
using FileIO
using DataFrames
using StatsBase
using StatsFuns
using StructArrays
using Plots
import Flipping.nextcount
##
function evidencepertrial(p, leave::Bool, failure_index)
    β,λ,c = p
    param = β*(1 - c*exp(-(λ*failure_index)))
    # - param only if event doesn't occur
    return - log(exp(-param)+1) - !leave * param
end

function evidenceofdatasmart(p, leaves, failure_indices)
    soa = StructArray((leaves, failure_indices))
    cm = countmap(soa)
    evidenceofdatasmart(p, cm)
end

function evidenceofdatasmart(p, cm)
    sum(n*evidencepertrial(p, value...) for (value, n) in cm)
end

probability(p, failure_index) = exp(evidencepertrial(p, true, failure_index))

#map(t -> probability(p, t), 1:10)
########
function D_evidencepertrial(p, leave::Bool, failure_index)
    β,λ,c = p
    param = β*(exp(-(λ*failure_index))-c)
    # - param only if event doesn't occur
    return - log(exp(-param)+1) - !leave * param
end

function D_evidenceofdatasmart(p, leaves, failure_indices)
    soa = StructArray((leaves, failure_indices))
    cm = countmap(soa)
    D_evidenceofdatasmart(p, cm)
end

function D_evidenceofdatasmart(p, cm)
    sum(n*D_evidencepertrial(p, value...) for (value, n) in cm)
end

########
function failure_idx(count::T, rewarded) where {T <: Number}
    if rewarded
        count = 0.0
    else
        count + 1.0
    end
end
##
filename = "/home/beatriz/mainen.flipping.5ht@gmail.com/Flipping/Datasets/Stimulations/DRN_Opto_Flipping/pokesDRN_Opto_Flipping.csv"
data = FileIO.load(filename) |> DataFrame
##
stim = launch();
t = stim[][:data][]
data = DataFrame(t)
data[!,:LastPoke] .= false
data[!,:FailuresIdx] .= 0.0
by(data,[:Session,:Streak]) do dd
    dd.LastPoke[end] = true
    dd[:,:FailuresIdx] = accumulate(failure_idx, occursin.(dd.Reward,"true");init=0.0)
end
checkpoint = by(data,:Session) do d
    d.PokeIn[1] == d.PokeIn[2]
end
any(checkpoint[:,2])
println(describe(data))
# coll = DataFrame(Temp = Float64[], Integration = Float64[], Cost = Float64[])
##
using LinearAlgebra
#filtered = filter(t -> t[:Protocol].!="45/15", data)
filtered = data
##
coll =by(filtered,[:Protocol,:Wall,:MouseID]) do dd
    leaves = dd[:,:LastPoke]
    failures_indices = dd[:,:FailuresIdx]
    ϵ = 10^(-2)
    res = optimize(p -> -evidenceofdatasmart(p, leaves, failures_indices) + ϵ*norm(p)^2, [1.0,1.0,1.0])
    mins = Optim.minimizer(res)
    # res = optimize(p -> -evidenceofdatasmart(p, leaves, failures_indices) + ϵ*norm(p)^2, [1.0,1.0,1.0])
    # mins = Optim.minimizer(res)
    DataFrame(Inverse_Temp = mins[1], Integration = mins[2], Cost = mins[3])
    # D_Inverse_Temp = D_mins[1], D_Integration = D_mins[2], D_Cost = D_mins[3])
end
###
mm1 = fit(LinearMixedModel, @formula( Integration ~ 1 + Protocol + Wall + (1|MouseID)), coll)
mm2 =  lm( @formula( Integration ~ 1 + Protocol + Wall),coll)
###
describe(coll)
f = plot(; legend = :bottomright)

for r in eachrow(coll)
    p = [r[:Inverse_Temp],r[:Integration],r[:Cost]]
    plot!(0:10,map(t -> probability(p, t), 0:10),label = r[:Protocol]*string(r[:Wall]))
end
f

####################
function check_double(that)
    checkpoint = by(that,:Session) do d
        d.PokeIn[1] == d.PokeIn[2]
    end
    return checkpoint[findall(checkpoint[:,2]),:]
end
checkpoint[findall(checkpoint[:,2]),:]
###################

sort!(coll,[:MouseID,:Protocol,:Wall])
z = []
for par in [:Inverse_Temp, :Integration, :Cost,:D_Inverse_Temp, :D_Integration, :D_Cost]
    idxs = occursin.("true",coll[:,:Wall])
    o = scatter(coll[.!idxs,par],coll[idxs,par],
    group = coll[idxs,:MouseID],
    legend = false,
    xlabel = "no barrier",
    ylabel = "barrier",
    title = string(par),
    markersize = 1)
    axis_list = (xlims(o)...,ylims(o)...)
    Plots.abline!(1,0)
    push!(z,o)
end

plot(z...)
Plots.abline!(1,0)
z
#####
function param_plot(df::AbstractDataFrame,split_var::Symbol,params::Union{Tuple,AbstractVector};groups = :MouseID)
    z = []
    sort!(df, [groups,split_var])
    dof = union(df[:,split_var])
    idxs_x = df[:,split_var].== dof[1]
    idxs_y = df[:,split_var].== dof[end]
    for par in params
        o = scatter(df[idxs_x,par],df[idxs_y,par],
        group = df[idxs_x,groups],
        legend = false,
        xlabel = dof[1],
        ylabel = dof[end],
        title = string(par),
        markersize = 1)
        Plots.abline!(1,0)
        push!(z,o)
    end
    plot(z...)
end

#####
param_plot(coll,:Protocol,(:Inverse_Temp, :Integration, :Cost))
#####
describe(coll)
using StatPlots
##
P_int = @df coll boxplot(:Protocol,:Integration,xlabel = "Protocol", ylabel = "Integration")
P_inv = @df coll boxplot(:Protocol,:Inverse_Temp, xlabel = "Protocol", ylabel = "Inverse_Temp")
P_costo = @df coll boxplot(:Protocol,:Cost, xlabel = "Protocol", ylabel = "Cost")
W_int = @df coll boxplot(:Wall,:Integration, xlabel = "Wall", ylabel = "Integration")
W_inv = @df coll boxplot(:Wall,:Inverse_Temp, xlabel = "Wall", ylabel = "Inverse_Temp")
W_costo = @df coll boxplot(:Wall,:Cost, xlabel = "Wall", ylabel = "Cost");
##
plot(P_int,
    P_inv,
    P_costo,
    W_int,
    W_inv,
    W_costo,legend = false)


text:
using Revise
	using BifurcationKit, LinearAlgebra, Plots, Setfield, Parameters

N(x; a = 0.5, b = 0.01) = 1 + (x + a*x^2)/(1 + b*x^2)
dN(x; a = 0.5, b = 0.01) = (1-b*x^2+2*a*x)/(1+b*x^2)^2
d2N(x; a = 0.5, b = 0.01) = -(2*(-a+3*a*b*x^2+3*b*x-b^2*x^3))/(1+b*x^2)^3

function F_chan(x, p)
	@unpack α, β = p
	f = similar(x)
	n = length(x)
	f[1] = x[1] - β
	f[n] = x[n] - β
	for i=2:n-1
		f[i] = (x[i-1] - 2 * x[i] + x[i+1]) * (n-1)^2 + α * N(x[i], b = β)
	end
	return f
end

function Jac_mat(u, p)
	@unpack α, β = p
	n = length(u)
	J = zeros(n, n)
	J[1, 1] = 1.0
	J[n, n] = 1.0
	for i = 2:n-1
		J[i, i-1] = (n-1)^2
		J[i, i+1] = (n-1)^2
		J[i, i] = -2 * (n-1)^2 + α * dN(u[i], b = β)
	end
	return J
end

n = 101
	par = (α = 3.3, β = 0.01)
	sol = [(i-1)*(n-i)/n^2+0.1 for i=1:n]
	optnewton = NewtonPar(tol = 1e-8, verbose = true)
	# ca fait dans les 63.59k Allocations
	out, hist, flag = @time newton( F_chan,	Jac_mat, sol, par, optnewton)

optscont = ContinuationPar(dsmin = 0.01, dsmax = 0.2, ds= 0.1, pMax = 4.1, nev = 5, detectFold = true, plotEveryStep = 40, newtonOptions = NewtonPar(maxIter = 10, tol = 1e-9), maxSteps = 150)
	br, = @time continuation(
		F_chan, Jac_mat,
		out, par, (@lens _.α),
		optscont; plot = true, verbosity = 0,
		plotSolution = (x, p; kwargs...) -> (plot!(x;ylabel="solution",label="", kwargs...))
		)
###################################################################################################
# Example with deflation technique
deflationOp = DeflationOperator(2.0, dot, 1.0, [out])

optdef = setproperties(optnewton; tol = 1e-10, maxIter = 1000)

outdef1, = @time newton(
	F_chan, Jac_mat,
	out .* (1 .+ 0.01*rand(n)), par,
	optdef, deflationOp)

plot(out, label="newton")
	plot!(sol, label="init guess")
	plot!(outdef1, label="deflation-1")

#save newly found point to look for new ones
push!(deflationOp, outdef1)
outdef2, = @time newton(
	F_chan, Jac_mat,
	outdef1.*(1 .+ 0.01*rand(n)), par,
	optdef, deflationOp)
plot!(outdef2, label="deflation-2")
#################################################################################################### Continuation of the Fold Point using minimally augmented formulation
optscont = (@set optscont.newtonOptions = setproperties(optscont.newtonOptions; verbose = true, tol = 1e-10))

indfold = 2

outfold, _, flag = @time newton(
	F_chan, Jac_mat,
	#index of the fold point
	br, indfold, (@lens _.α))
	flag && printstyled(color=:red, "--> We found a Fold Point at α = ", outfold.p, ", β = 0.01, from ", br.foldpoint[indfold].param,"\n")

optcontfold = ContinuationPar(dsmin = 0.001, dsmax = 0.05, ds= 0.05, pMax = 4.1, pMin = 0., newtonOptions = NewtonPar(verbose=false, tol = 1e-8), maxSteps = 1300)
	foldbranch, = @time continuation(
		F_chan, Jac_mat,
		br, indfold, (@lens _.α), (@lens _.β),
		plot = true, verbosity = 2,
		optcontfold)
plot(foldbranch, label = "");title!("")
################################################################################################### Fold Newton / Continuation when Hessian is known. Does not require state to be AbstractVector
d2F(x, p, u, v; b = 0.01) = p.α .* d2N.(x; b = b) .* u .* v

outfold, _, flag = @time newton(
		F_chan, Jac_mat,
		br, indfold, (@lens _.α);
		d2F = d2F
		)
	flag && printstyled(color=:red, "--> We found a Fold Point at α = ", outfold.p, ", β = 0.01, from ", br.foldpoint[indfold].param,"\n")

optcontfold = ContinuationPar(dsmin = 0.001, dsmax = 0.05, ds= 0.01, pMax = 4.1, pMin = 0., newtonOptions = NewtonPar(verbose=true), maxSteps = 1300)

outfoldco, = @time continuation(
		F_chan, Jac_mat,
		br, indfold, (@lens _.α), (@lens _.β), optcontfold;
		# Jt = (x, p) -> transpose(Jac_mat(x, p)),
		d2F = d2F,
		plot = true)
###################################################################################################
# Matrix Free example
function dF_chan(x, dx, p)
	@unpack α, β = p
	out = similar(x)
	n = length(x)
	out[1] = dx[1]
	out[n] = dx[n]
	for i=2:n-1
		out[i] = (dx[i-1] - 2 * dx[i] + dx[i+1]) * (n-1)^2 + α * dN(x[i], b = β) * dx[i]
	end
	return out
end

ls = GMRESKrylovKit(dim = 100)
	optnewton_mf = NewtonPar(tol = 1e-11, verbose = true, linsolver = ls, eigsolver = DefaultEig())
	out_mf, _, flag = @time newton(
		F_chan,
		(x, p) -> (dx -> dF_chan(x, dx, p)),
		sol, par, optnewton_mf)

opts_cont_mf  = ContinuationPar(dsmin = 0.01, dsmax = 0.1, ds= 0.01, pMax = 4.1, nev = 5, plotEveryStep = 40, newtonOptions = setproperties(optnewton_mf; maxIter = 70, tol = 1e-8), maxSteps = 150)
	brmf, _ = @time continuation(
		F_chan,
		(x, p) -> (dx -> dF_chan(x, dx, p)),
		out, par, (@lens _.α), opts_cont_mf,
		linearAlgo = MatrixFreeBLS(),
		)

plot(brmf)

using SparseArrays
P = spdiagm(0 => -2 * (n-1)^2 * ones(n), -1 => (n-1)^2 * ones(n-1), 1 => (n-1)^2 * ones(n-1))
P[1,1:2] .= [1, 0.];P[end,end-1:end] .= [0, 1.]

ls = GMRESIterativeSolvers(reltol = 1e-5, N = length(sol), restart = 20, maxiter=10, Pl = lu(P))
	optnewton_mf = NewtonPar(tol = 1e-9, verbose = true, linsolver = ls)
	out_mf, _, flag = @time newton(
		F_chan,
		(x, p) -> (dx -> dF_chan(x, dx, p)),
		sol, par, optnewton_mf)

plot(brmf,color=:red)

# matrix free with different tangent predictor
brmf, = @time continuation(
	F_chan,
	(x, p) -> (dx -> dF_chan(x, dx, p)),
	out, par, (@lens _.α), (@set opts_cont_mf.newtonOptions = optnewton_mf),
	tangentAlgo = BorderedPred(),
	)

plot(brmf,color=:blue)

brmf, = @time continuation(
	F_chan,
	(x, p) -> (dx -> dF_chan(x, dx, p)),
	out, par, (@lens _.α), opts_cont_mf,
	tangentAlgo = SecantPred(),
	linearAlgo = MatrixFreeBLS()
	)

plot(brmf,color=:green)

brmf, = @time continuation(
	F_chan,
	(x, p) -> (dx -> dF_chan(x, dx, p)),
	out, par, (@lens _.α), opts_cont_mf,
	tangentAlgo = BorderedPred(),
	linearAlgo = MatrixFreeBLS())

plot(brmf,color=:orange)


text:
abstract type Snapshot end

mutable struct HistoryTimeline
    snapshots::Array{Snapshot, 1}
    index::Integer
    skip::Bool

    HistoryTimeline(snapshots::Array{Snapshot, 1}=Snapshot[], index::Integer=length(snapshots)) = new(snapshots, index, false)
end

function Base.push!(history::HistoryTimeline, snapshot::Snapshot)
    if history.index < length(history.snapshots)
        history.snapshots = history.snapshots[1:history.index]
    end

    push!(history.snapshots, snapshot)
    history.index += 1
end

function Base.pop!(history::HistoryTimeline)
    res = history.snapshots[history.index]
    history.index -= 1

    return res
end


text:
using TypeSerializers
using Test

struct Foo
    i::Int
    s::String
end

@testset "TypeSerializers" begin
    integer_value = 42
    stream = serialize(SerializationTypeSerializer{Int}, integer_value)
    seekstart(stream)
    @test deserialize(SerializationTypeSerializer{Int}, stream) == integer_value

    string_value = "Hello World!"
    stream = serialize(SerializationTypeSerializer{String}, string_value)
    seekstart(stream)
    @test deserialize(SerializationTypeSerializer{String}, stream) == string_value

    foo_value = Foo(integer_value, string_value)
    stream = serialize(SerializationTypeSerializer{Foo}, foo_value)
    seekstart(stream)
    @test deserialize(SerializationTypeSerializer{Foo}, stream) == foo_value
end


text:
function centroid!{T <: Real}(o::Vector{T},
	                          a::Vector{T},
	                          b::Vector{T})
	for i in 1:length(o)
		o[i] = (a[i] + b[i]) / 2
	end
	return
end


text:
"""
    module ElectronicStructure

Provides an interface to electronic structure calculation packages.

PySCF is supported by the package ElectronicStructurePySCF
"""
module ElectronicStructure

import PeriodicTable
import ZChop
using Tullio: @tullio

export Atom, Geometry, MolecularSpec, InteractionOperator, InteractionOperator!
export non_zero_elements, non_zero_elements_python, phys_to_chem, chem_to_phys,
    check_two_body_symmetries, find_index_order,
    to_index_order

export to_openfermion

include("utils.jl")
include("molecular_spec.jl")
include("molecular_data.jl")
include("interaction_operator.jl")
include("openfermion.jl")

end # module ElectronicStructure


text:
# By default, Julia/LLVM does not use fused multiply-add operations (FMAs).
# Since these FMAs can increase the performance of many numerical algorithms,
# we need to opt-in explicitly.
# See https://ranocha.de/blog/Optimizing_EC_Trixi for further details.
@muladd begin

# The methods below are specialized on the mortar type
# and called from the basic `create_cache` method at the top.
function create_cache(mesh::P4estMesh{2}, equations, mortar_l2::LobattoLegendreMortarL2, uEltype)
  # TODO: Taal performance using different types
  MA2d = MArray{Tuple{nvariables(equations), nnodes(mortar_l2)},
                uEltype, 2,
                nvariables(equations) * nnodes(mortar_l2)}
  fstar_upper_threaded = MA2d[MA2d(undef) for _ in 1:Threads.nthreads()]
  fstar_lower_threaded = MA2d[MA2d(undef) for _ in 1:Threads.nthreads()]
  u_threaded =           MA2d[MA2d(undef) for _ in 1:Threads.nthreads()]

  (; fstar_upper_threaded, fstar_lower_threaded, u_threaded)
end

#     index_to_start_step_2d(index::Symbol, index_range)
#
# Given a symbolic `index` and an `indexrange` (usually `eachnode(dg)`),
# return `index_start, index_step`, i.e., a tuple containing
# - `index_start`, an index value to begin a loop
# - `index_step`,  an index step to update during a loop
# The resulting indices translate surface indices to volume indices.
#
# !!! warning
#     This assumes that loops using the return values are written as
#
#     i_volume_start, i_volume_step = index_to_start_step_2d(symbolic_index_i, index_range)
#     j_volume_start, j_volume_step = index_to_start_step_2d(symbolic_index_j, index_range)
#
#     i_volume, j_volume = i_volume_start, j_volume_start
#     for i_surface in index_range
#       # do stuff with `i_surface` and `(i_volume, j_volume)`
#
#       i_volume += i_volume_step
#       j_volume += j_volume_step
#     end
@inline function index_to_start_step_2d(index::Symbol, index_range)
  index_begin = first(index_range)
  index_end   = last(index_range)

  if index === :begin
    return index_begin, 0
  elseif index === :end
    return index_end, 0
  elseif index === :i_forward
    return index_begin, 1
  else # if index === :i_backward
    return index_end, -1
  end
end

function prolong2interfaces!(cache, u,
                             mesh::P4estMesh{2},
                             equations, surface_integral, dg::DG)
  @unpack interfaces = cache
  index_range = eachnode(dg)

  @threaded for interface in eachinterface(dg, cache)
    # Copy solution data from the primary element using "delayed indexing" with
    # a start value and a step size to get the correct face and orientation.
    # Note that in the current implementation, the interface will be
    # "aligned at the primary element", i.e., the index of the primary side
    # will always run forwards.
    primary_element = interfaces.neighbor_ids[1, interface]
    primary_indices = interfaces.node_indices[1, interface]

    i_primary_start, i_primary_step = index_to_start_step_2d(primary_indices[1], index_range)
    j_primary_start, j_primary_step = index_to_start_step_2d(primary_indices[2], index_range)

    i_primary = i_primary_start
    j_primary = j_primary_start
    for i in eachnode(dg)
      for v in eachvariable(equations)
        interfaces.u[1, v, i, interface] = u[v, i_primary, j_primary, primary_element]
      end
      i_primary += i_primary_step
      j_primary += j_primary_step
    end

    # Copy solution data from the secondary element using "delayed indexing" with
    # a start value and a step size to get the correct face and orientation.
    secondary_element = interfaces.neighbor_ids[2, interface]
    secondary_indices = interfaces.node_indices[2, interface]

    i_secondary_start, i_secondary_step = index_to_start_step_2d(secondary_indices[1], index_range)
    j_secondary_start, j_secondary_step = index_to_start_step_2d(secondary_indices[2], index_range)

    i_secondary = i_secondary_start
    j_secondary = j_secondary_start
    for i in eachnode(dg)
      for v in eachvariable(equations)
        interfaces.u[2, v, i, interface] = u[v, i_secondary, j_secondary, secondary_element]
      end
      i_secondary += i_secondary_step
      j_secondary += j_secondary_step
    end
  end

  return nothing
end

function calc_interface_flux!(surface_flux_values,
                              mesh::P4estMesh{2},
                              nonconservative_terms,
                              equations, surface_integral, dg::DG, cache)
  @unpack neighbor_ids, node_indices = cache.interfaces
  @unpack contravariant_vectors = cache.elements
  index_range = eachnode(dg)
  index_end = last(index_range)

  @threaded for interface in eachinterface(dg, cache)
    # Get element and side index information on the primary element
    primary_element = neighbor_ids[1, interface]
    primary_indices = node_indices[1, interface]
    primary_direction = indices2direction(primary_indices)

    # Create the local i,j indexing on the primary element used to pull normal direction information
    i_primary_start, i_primary_step = index_to_start_step_2d(primary_indices[1], index_range)
    j_primary_start, j_primary_step = index_to_start_step_2d(primary_indices[2], index_range)

    i_primary = i_primary_start
    j_primary = j_primary_start

    # Get element and side index information on the secondary element
    secondary_element = neighbor_ids[2, interface]
    secondary_indices = node_indices[2, interface]
    secondary_direction = indices2direction(secondary_indices)

    # Initiate the secondary index to be used in the surface for loop.
    # This index on the primary side will always run forward but
    # the secondary index might need to run backwards for flipped sides.
    if :i_backward in secondary_indices
      node_secondary = index_end
      node_secondary_step = -1
    else
      node_secondary = 1
      node_secondary_step = 1
    end

    for node in eachnode(dg)
      # Get the normal direction on the primary element.
      # Contravariant vectors at interfaces in negative coordinate direction
      # are pointing inwards. This is handled by `get_normal_direction`.
      normal_direction = get_normal_direction(primary_direction, contravariant_vectors,
                                              i_primary, j_primary, primary_element)

      calc_interface_flux!(surface_flux_values, mesh, nonconservative_terms, equations,
                           surface_integral, dg, cache,
                           interface, normal_direction,
                           node, primary_direction, primary_element,
                           node_secondary, secondary_direction, secondary_element)

      # Increment primary element indices to pull the normal direction
      i_primary += i_primary_step
      j_primary += j_primary_step
      # Increment the surface node index along the secondary element
      node_secondary += node_secondary_step
    end
  end

  return nothing
end

# Inlined version of the interface flux computation for conservation laws
@inline function calc_interface_flux!(surface_flux_values,
                                      mesh::P4estMesh{2},
                                      nonconservative_terms::Val{false}, equations,
                                      surface_integral, dg::DG, cache,
                                      interface_index, normal_direction,
                                      primary_node_index, primary_direction_index, primary_element_index,
                                      secondary_node_index, secondary_direction_index, secondary_element_index)
  @unpack u = cache.interfaces
  @unpack surface_flux = surface_integral

  u_ll, u_rr = get_surface_node_vars(u, equations, dg, primary_node_index, interface_index)

  flux_ = surface_flux(u_ll, u_rr, normal_direction, equations)

  for v in eachvariable(equations)
    surface_flux_values[v, primary_node_index, primary_direction_index, primary_element_index] = flux_[v]
    surface_flux_values[v, secondary_node_index, secondary_direction_index, secondary_element_index] = -flux_[v]
  end
end

# Inlined version of the interface flux computation for equations with conservative and nonconservative terms
@inline function calc_interface_flux!(surface_flux_values,
                                      mesh::P4estMesh{2},
                                      nonconservative_terms::Val{true}, equations,
                                      surface_integral, dg::DG, cache,
                                      interface_index, normal_direction,
                                      primary_node_index, primary_direction_index, primary_element_index,
                                      secondary_node_index, secondary_direction_index, secondary_element_index)
  @unpack u = cache.interfaces
  surface_flux, nonconservative_flux = surface_integral.surface_flux

  u_ll, u_rr = get_surface_node_vars(u, equations, dg, primary_node_index, interface_index)

  flux_ = surface_flux(u_ll, u_rr, normal_direction, equations)

  # Compute both nonconservative fluxes
  # In general, nonconservative fluxes can depend on both the contravariant
  # vectors (normal direction) at the current node and the averaged ones.
  # However, both are the same at watertight interfaces, so we pass the
  # `normal_direction` twice.
  noncons_primary   = nonconservative_flux(u_ll, u_rr, normal_direction, normal_direction, equations)
  noncons_secondary = nonconservative_flux(u_rr, u_ll, normal_direction, normal_direction, equations)

  # Store the flux with nonconservative terms on the primary and secondary elements
  for v in eachvariable(equations)
    # Note the factor 0.5 necessary for the nonconservative fluxes based on
    # the interpretation of global SBP operators coupled discontinuously via
    # central fluxes/SATs
    surface_flux_values[v, primary_node_index, primary_direction_index, primary_element_index] = (
      flux_[v] + 0.5 * noncons_primary[v])
    surface_flux_values[v, secondary_node_index, secondary_direction_index, secondary_element_index] = -(
      flux_[v] + 0.5 * noncons_secondary[v])
  end
end

function prolong2boundaries!(cache, u,
                             mesh::P4estMesh{2},
                             equations, surface_integral, dg::DG)
  @unpack boundaries = cache
  index_range = eachnode(dg)

  @threaded for boundary in eachboundary(dg, cache)
    # Copy solution data from the element using "delayed indexing" with
    # a start value and a step size to get the correct face and orientation.
    element       = boundaries.neighbor_ids[boundary]
    node_indices  = boundaries.node_indices[boundary]

    i_node_start, i_node_step = index_to_start_step_2d(node_indices[1], index_range)
    j_node_start, j_node_step = index_to_start_step_2d(node_indices[2], index_range)

    i_node = i_node_start
    j_node = j_node_start
    for i in eachnode(dg)
      for v in eachvariable(equations)
        boundaries.u[v, i, boundary] = u[v, i_node, j_node, element]
      end
      i_node += i_node_step
      j_node += j_node_step
    end
  end

  return nothing
end

function calc_boundary_flux!(cache, t, boundary_condition, boundary_indexing,
                             mesh::P4estMesh{2},
                             equations, surface_integral, dg::DG)
  @unpack boundaries = cache
  @unpack surface_flux_values = cache.elements
  index_range = eachnode(dg)

  @threaded for local_index in eachindex(boundary_indexing)
    # Use the local index to get the global boundary index from the pre-sorted list
    boundary = boundary_indexing[local_index]

    # Get information on the adjacent element, compute the surface fluxes,
    # and store them
    element       = boundaries.neighbor_ids[boundary]
    node_indices  = boundaries.node_indices[boundary]
    direction     = indices2direction(node_indices)

    i_node_start, i_node_step = index_to_start_step_2d(node_indices[1], index_range)
    j_node_start, j_node_step = index_to_start_step_2d(node_indices[2], index_range)

    i_node = i_node_start
    j_node = j_node_start
    for node in eachnode(dg)
      calc_boundary_flux!(surface_flux_values, t, boundary_condition,
                          mesh, have_nonconservative_terms(equations),
                          equations, surface_integral, dg, cache,
                          i_node, j_node,
                          node, direction, element, boundary)

      i_node += i_node_step
      j_node += j_node_step
    end
  end
end

# inlined version of the boundary flux calculation along a physical interface
@inline function calc_boundary_flux!(surface_flux_values, t, boundary_condition,
                                     mesh::P4estMesh{2},
                                     nonconservative_terms::Val{false}, equations,
                                     surface_integral, dg::DG, cache,
                                     i_index, j_index,
                                     node_index, direction_index, element_index, boundary_index)
  @unpack boundaries = cache
  @unpack node_coordinates, contravariant_vectors = cache.elements
  @unpack surface_flux = surface_integral

  # Extract solution data from boundary container
  u_inner = get_node_vars(boundaries.u, equations, dg, node_index, boundary_index)

  # Outward-pointing normal direction (not normalized)
  normal_direction = get_normal_direction(direction_index, contravariant_vectors,
                                          i_index, j_index, element_index)

  # Coordinates at boundary node
  x = get_node_coords(node_coordinates, equations, dg, i_index, j_index, element_index)

  flux_ = boundary_condition(u_inner, normal_direction, x, t, surface_flux, equations)

  # Copy flux to element storage in the correct orientation
  for v in eachvariable(equations)
    surface_flux_values[v, node_index, direction_index, element_index] = flux_[v]
  end
end

# inlined version of the boundary flux with nonconservative terms calculation along a physical interface
@inline function calc_boundary_flux!(surface_flux_values, t, boundary_condition,
                                     mesh::P4estMesh{2},
                                     nonconservative_terms::Val{true}, equations,
                                     surface_integral, dg::DG, cache,
                                     i_index, j_index,
                                     node_index, direction_index, element_index, boundary_index)
  @unpack boundaries = cache
  @unpack node_coordinates, contravariant_vectors = cache.elements
  surface_flux, nonconservative_flux = surface_integral.surface_flux

  # Extract solution data from boundary container
  u_inner = get_node_vars(boundaries.u, equations, dg, node_index, boundary_index)

  # Outward-pointing normal direction (not normalized)
  normal_direction = get_normal_direction(direction_index, contravariant_vectors,
                                          i_index, j_index, element_index)

  # Coordinates at boundary node
  x = get_node_coords(node_coordinates, equations, dg, i_index, j_index, element_index)

  # Call pointwise numerical flux function for the conservative part
  # in the normal direction on the boundary
  flux_ = boundary_condition(u_inner, normal_direction, x, t, surface_flux, equations)

  # Compute pointwise nonconservative numerical flux at the boundary.
  # Note: This does not set any type of boundary condition for the nonconservative term
  noncons_ = nonconservative_flux(u_inner, u_inner, normal_direction, normal_direction, equations)

  # Copy flux to element storage in the correct orientation
  for v in eachvariable(equations)
    # Note the factor 0.5 necessary for the nonconservative fluxes based on
    # the interpretation of global SBP operators coupled discontinuously via
    # central fluxes/SATs
    surface_flux_values[v, node_index, direction_index, element_index] = flux_[v] + 0.5 * noncons_[v]
  end
end

function prolong2mortars!(cache, u,
                          mesh::P4estMesh{2}, equations,
                          mortar_l2::LobattoLegendreMortarL2,
                          surface_integral, dg::DGSEM)
  @unpack neighbor_ids, node_indices = cache.mortars
  index_range = eachnode(dg)

  @threaded for mortar in eachmortar(dg, cache)
    # Copy solution data from the small elements using "delayed indexing" with
    # a start value and a step size to get the correct face and orientation.
    small_indices = node_indices[1, mortar]

    i_small_start, i_small_step = index_to_start_step_2d(small_indices[1], index_range)
    j_small_start, j_small_step = index_to_start_step_2d(small_indices[2], index_range)

    for position in 1:2
      i_small = i_small_start
      j_small = j_small_start
      element = neighbor_ids[position, mortar]
      for i in eachnode(dg)
        for v in eachvariable(equations)
          cache.mortars.u[1, v, position, i, mortar] = u[v, i_small, j_small, element]
        end
        i_small += i_small_step
        j_small += j_small_step
      end
    end

    # Buffer to copy solution values of the large element in the correct orientation
    # before interpolating
    u_buffer = cache.u_threaded[Threads.threadid()]

    # Copy solution of large element face to buffer in the
    # correct orientation
    large_indices = node_indices[2, mortar]

    i_large_start, i_large_step = index_to_start_step_2d(large_indices[1], index_range)
    j_large_start, j_large_step = index_to_start_step_2d(large_indices[2], index_range)

    i_large = i_large_start
    j_large = j_large_start
    element = neighbor_ids[3, mortar]
    for i in eachnode(dg)
      for v in eachvariable(equations)
        u_buffer[v, i] = u[v, i_large, j_large, element]
      end
      i_large += i_large_step
      j_large += j_large_step
    end

    # Interpolate large element face data from buffer to small face locations
    multiply_dimensionwise!(view(cache.mortars.u, 2, :, 1, :, mortar),
                            mortar_l2.forward_lower,
                            u_buffer)
    multiply_dimensionwise!(view(cache.mortars.u, 2, :, 2, :, mortar),
                            mortar_l2.forward_upper,
                            u_buffer)
  end

  return nothing
end

function calc_mortar_flux!(surface_flux_values,
                           mesh::P4estMesh{2},
                           nonconservative_terms, equations,
                           mortar_l2::LobattoLegendreMortarL2,
                           surface_integral, dg::DG, cache)
  @unpack neighbor_ids, node_indices = cache.mortars
  @unpack contravariant_vectors = cache.elements
  @unpack fstar_upper_threaded, fstar_lower_threaded = cache
  index_range = eachnode(dg)

  @threaded for mortar in eachmortar(dg, cache)
    # Choose thread-specific pre-allocated container
    fstar = (fstar_lower_threaded[Threads.threadid()],
             fstar_upper_threaded[Threads.threadid()])

    # Get index information on the small elements
    small_indices = node_indices[1, mortar]
    small_direction = indices2direction(small_indices)

    i_small_start, i_small_step = index_to_start_step_2d(small_indices[1], index_range)
    j_small_start, j_small_step = index_to_start_step_2d(small_indices[2], index_range)

    for position in 1:2
      i_small = i_small_start
      j_small = j_small_start
      element = neighbor_ids[position, mortar]
      for node in eachnode(dg)
        # Get the normal direction on the small element.
        # Note, contravariant vectors at interfaces in negative coordinate direction
        # are pointing inwards. This is handled by `get_normal_direction`.
        normal_direction = get_normal_direction(small_direction, contravariant_vectors,
                                                i_small, j_small, element)

        calc_mortar_flux!(fstar, mesh, nonconservative_terms, equations,
                          surface_integral, dg, cache,
                          mortar, position, normal_direction,
                          node)

        i_small += i_small_step
        j_small += j_small_step
      end
    end

    # Buffer to interpolate flux values of the large element to before
    # copying in the correct orientation
    u_buffer = cache.u_threaded[Threads.threadid()]

    mortar_fluxes_to_elements!(surface_flux_values,
                               mesh, equations, mortar_l2, dg, cache,
                               mortar, fstar, u_buffer)
  end

  return nothing
end

# Inlined version of the mortar flux computation on small elements for conservation laws
@inline function calc_mortar_flux!(fstar,
                                   mesh::P4estMesh{2},
                                   nonconservative_terms::Val{false}, equations,
                                   surface_integral, dg::DG, cache,
                                   mortar_index, position_index, normal_direction,
                                   node_index)
  @unpack u = cache.mortars
  @unpack surface_flux = surface_integral

  u_ll, u_rr = get_surface_node_vars(u, equations, dg, position_index, node_index, mortar_index)

  flux = surface_flux(u_ll, u_rr, normal_direction, equations)

  # Copy flux to buffer
  set_node_vars!(fstar[position_index], flux, equations, dg, node_index)
end

# Inlined version of the mortar flux computation on small elements for equations with conservative and
# nonconservative terms
@inline function calc_mortar_flux!(fstar,
                                   mesh::P4estMesh{2},
                                   nonconservative_terms::Val{true}, equations,
                                   surface_integral, dg::DG, cache,
                                   mortar_index, position_index, normal_direction,
                                   node_index)
  @unpack u = cache.mortars
  surface_flux, nonconservative_flux = surface_integral.surface_flux

  u_ll, u_rr = get_surface_node_vars(u, equations, dg, position_index, node_index, mortar_index)

  # Compute conservative flux
  flux = surface_flux(u_ll, u_rr, normal_direction, equations)

  # Compute nonconservative flux and add it to the conservative flux.
  # The nonconservative flux is scaled by a factor of 0.5 based on
  # the interpretation of global SBP operators coupled discontinuously via
  # central fluxes/SATs
  noncons = nonconservative_flux(u_ll, u_rr, normal_direction, normal_direction, equations)

  flux_plus_noncons = flux + 0.5 * noncons

  # Copy to buffer
  set_node_vars!(fstar[position_index], flux_plus_noncons, equations, dg, node_index)
end

@inline function mortar_fluxes_to_elements!(surface_flux_values,
                                            mesh::P4estMesh{2}, equations,
                                            mortar_l2::LobattoLegendreMortarL2,
                                            dg::DGSEM, cache, mortar, fstar, u_buffer)
  @unpack neighbor_ids, node_indices = cache.mortars

  # Copy solution small to small
  small_indices   = node_indices[1, mortar]
  small_direction = indices2direction(small_indices)

  for position in 1:2
    element = neighbor_ids[position, mortar]
    for i in eachnode(dg)
      for v in eachvariable(equations)
        surface_flux_values[v, i, small_direction, element] = fstar[position][v, i]
      end
    end
  end

  # Project small fluxes to large element.
  multiply_dimensionwise!(u_buffer,
                          mortar_l2.reverse_upper, fstar[2],
                          mortar_l2.reverse_lower, fstar[1])

  # The flux is calculated in the outward direction of the small elements,
  # so the sign must be switched to get the flux in outward direction
  # of the large element.
  # The contravariant vectors of the large element (and therefore the normal
  # vectors of the large element as well) are twice as large as the
  # contravariant vectors of the small elements. Therefore, the flux needs
  # to be scaled by a factor of 2 to obtain the flux of the large element.
  u_buffer .*= -2

  # Copy interpolated flux values from buffer to large element face in the
  # correct orientation.
  # Note that the index of the small sides will always run forward but
  # the index of the large side might need to run backwards for flipped sides.
  large_element  = neighbor_ids[3, mortar]
  large_indices  = node_indices[2, mortar]
  large_direction = indices2direction(large_indices)

  if :i_backward in large_indices
    for i in eachnode(dg)
      for v in eachvariable(equations)
        surface_flux_values[v, end + 1 - i, large_direction, large_element] = u_buffer[v, i]
      end
    end
  else
    for i in eachnode(dg)
      for v in eachvariable(equations)
        surface_flux_values[v, i, large_direction, large_element] = u_buffer[v, i]
      end
    end
  end

  return nothing
end

function calc_surface_integral!(du, u,
                                mesh::P4estMesh{2},
                                equations,
                                surface_integral::SurfaceIntegralWeakForm,
                                dg::DGSEM, cache)
  @unpack boundary_interpolation = dg.basis
  @unpack surface_flux_values = cache.elements

  # Note that all fluxes have been computed with outward-pointing normal vectors.
  # Access the factors only once before beginning the loop to increase performance.
  # We also use explicit assignments instead of `+=` to let `@muladd` turn these
  # into FMAs (see comment at the top of the file).
  factor_1 = boundary_interpolation[1,          1]
  factor_2 = boundary_interpolation[nnodes(dg), 2]
  @threaded for element in eachelement(dg, cache)
    for l in eachnode(dg)
      for v in eachvariable(equations)
        # surface at -x
        du[v, 1,          l, element] = (
          du[v, 1,          l, element] + surface_flux_values[v, l, 1, element] * factor_1)

        # surface at +x
        du[v, nnodes(dg), l, element] = (
          du[v, nnodes(dg), l, element] + surface_flux_values[v, l, 2, element] * factor_2)

        # surface at -y
        du[v, l, 1,          element] = (
          du[v, l, 1,          element] + surface_flux_values[v, l, 3, element] * factor_1)

        # surface at +y
        du[v, l, nnodes(dg), element] = (
          du[v, l, nnodes(dg), element] + surface_flux_values[v, l, 4, element] * factor_2)
      end
    end
  end

  return nothing
end

end # @muladd


text:
using QuNet

function make_user_pairs_wbuffer(squarenet::QNetwork, num_pairs::Int; buffer::Int=1)::Vector{Tuple{Int64, Int64}}
    num_nodes = length(squarenet.nodes)
    d = sqrt(num_nodes)
    @assert isinteger(d) "Grid network is not square"
    @assert buffer > 0
    @assert d - 2 * buffer > 1 "Gridsize too small for given buffer"
    rand_space = Array(collect(0:num_nodes-1))
    # Filter out nodes in the margin
    # NOTE Cool trick: n % d == x coord, n ÷ d == y coord
    filter!(i -> (buffer <= (i%d)) && (buffer <= (i÷d)) && (i%d < d-buffer) && (i÷d < d-buffer), rand_space)
    # Set first index to 1
    rand_space .+= 1
    pairs = make_user_pairs(squarenet, num_pairs, node_list=rand_space)
    return pairs
end

G = GridNetwork(5, 5)
p = make_user_pairs_wbuffer(G, 3, buffer = 1)
println(p)

max_stars_repo_name: FalafelGood/QuNet

max_issues_repo_name: FalafelGood/QuNet

max_forks_repo_name: FalafelGood/QuNet