text:
#!/usr/bin/env julia

using Distributed

# this appears to be the point of diminishing returns
max_node_processes = 12

"turns a slurm node list into a list of node names"
function parse_node_names(slurm_node_list::String)
    node_names = String[]
    slurm_node_list_parts = split(slurm_node_list, "[")
    node_name_prefix = slurm_node_list_parts[1]

    #if only one node was found then this length will be 1
    if length(slurm_node_list_parts) == 1
        push!(node_names, node_name_prefix)
    else
        # this assumes all nodes have a common prefix and the nodelist has the form [1-2,4,5-10,12]
        node_ids = split(split(slurm_node_list_parts[2], "]")[1], ",")
        for node_id in node_ids
            node_id_parts = split(node_id, "-")
            start = parse(Int, node_id_parts[1])
            stop = parse(Int, length(node_id_parts) == 1 ? node_id_parts[1] : node_id_parts[2])
            digits = length(node_id_parts[1])
            for i in start:stop
                node_name = "$(node_name_prefix)$(lpad(string(i), digits, "0"))"
                push!(node_names, node_name)
            end
        end
    end
    return node_names
end

function add_local_procs()
    proc_ids = Int[]
    if Distributed.nprocs() <= 2
        node_processes = min(trunc(Int, Sys.CPU_THREADS*0.75), max_node_processes)
        println("local processes: $(node_processes) of $(Sys.CPU_THREADS)")

        proc_ids = Distributed.addprocs(node_processes)
    end
    return proc_ids
end

function add_remote_procs()
    if haskey(ENV, "SLURM_JOB_NODELIST")
        node_list = ENV["SLURM_JOB_NODELIST"]
    elseif haskey(ENV, "SLURM_NODELIST")
        node_list = ENV["SLURM_NODELIST"]
    else
        @info("unable to find slurm node list environment variable")
        return Int[]
    end
    #println(node_list)

    node_names = parse_node_names(node_list)
    #println(node_names)

    @info("host name: $(gethostname())")
    @info("slurm allocation nodes: $(node_names)")

    node_names = [name for name in node_names if name != gethostname()]

    proc_ids = Int[]
    if length(node_names) > 0
        @info("remote slurm nodes: $(node_names)")
        node_processes = min(trunc(Int, Sys.CPU_THREADS*0.75), max_node_processes)
        println("remote processes per node: $(node_processes)/$(Sys.CPU_THREADS)")
        for i in 1:node_processes
            node_proc_ids = Distributed.addprocs(node_names, sshflags="-oStrictHostKeyChecking=no")
            println("process id batch $(i) of $(node_processes): $(node_proc_ids)")
            for npid in node_proc_ids
                push!(proc_ids, npid)
            end
        end
    else
        @info("no remote slurm nodes found")
    end

    #println("process ids: $(proc_ids)")
    return proc_ids
end

function add_procs()
    proc_ids = add_local_procs()
    remote_proc_ids = add_remote_procs()

    for pid in remote_proc_ids
        push!(proc_ids, pid)
    end

    println("process ids: $(proc_ids)")
    return proc_ids
end

println("spin up workers")
add_procs()

println("define hello function on all workers")
@everywhere function hello(task::Int)
    println("I am working on task $(task) on $(gethostname())")
    return "task $(task) complete"
end

println("execute in parallel")
value = pmap(hello, [i for i in 1:Distributed.nprocs()])
sleep(1) # wait for stdout to catchup
println("result: $(value)")
println("execute complete")

max_stars_repo_name: ccoffrin/julia-examples

max_issues_repo_name: ccoffrin/julia-examples

max_forks_repo_name: ccoffrin/julia-examples


text:
module Models

import StatsBase: fit, predict

export Model, Template
export fit, predict, submodels, estimate_type, output_type, predict_input_type
export EstimateTrait, PointEstimate, DistributionEstimate
export OutputTrait, SingleOutput, MultiOutput
export PredictInputTrait, PointPredictInput, PointOrDistributionPredictInput

"""
   Template

A Template is an untrained [`Model`](@ref) that can be [`fit`](@ref) to data.
Defined as well are the traits:
- [`output_type`](@ref): [`SingleOutput`](@ref) or [`MultiOutput`](@ref).
- [`estimate_type`](@ref): [`PointEstimate`](@ref) or [`DistributionEstimate`](@ref).
"""
abstract type Template end

"""
   Model

A Model is a trained [`Template`](@ref) with which one can [`predict`](@ref) on inputs.
Defined as well are the traits:
- [`output_type`](@ref): [`SingleOutput`](@ref) or [`MultiOutput`](@ref).
- [`estimate_type`](@ref): [`PointEstimate`](@ref) or [`DistributionEstimate`](@ref).
"""
abstract type Model end

"""
   fit(::Template, output::AbstractMatrix, input::AbstractMatrix, [weights]) -> Model

Fit the [`Template`](@ref) to the `output` and `input` data and return a trained
[`Model`](@ref).
Convention is that `weights` defaults to `StatsBase.uweights(Float32, size(outputs, 2))`
"""
function fit end

"""
    predict(model::Model, inputs::AbstractMatrix)
    predict(model::Model, inputs::AbstractVector{<:AbstractVector})

Predict targets for the provided the collection of `inputs` and [`Model`](@ref).

A [`Model`](@ref) subtype for which the `predict_input_type(model)` is
[`PointPredictInput`](@ref) will only need to implement a `predict` function that operates
on an `AbstractMatrix` of inputs.

If the `estimate_type(model)` is [`PointEstimate`](@ref) then this function should return
another `AbstractMatrix` in which each column contains the prediction for a single input.

If the `estimate_type(model)` is [`DistributionEstimate`](@ref) then this function should
return a `AbstractVector{<:Distribution}`.
"""
function predict end

function predict(model::Model, inputs::AbstractVector{<:AbstractVector})
   return predict(model, reduce(hcat, inputs))
end

"""
    submodels(::Union{Template, Model})

Return all submodels within a multistage model/template.

Submodels are models within a model that have their own inputs (which may or may not be
combined with outputs of _earlier_ submodels, before actually being passed as input to the submodel).
Such multistage models take a tuple of inputs (which may be nested if the submodel itself
has submodels).
The order of submodels returned by `submodels` is as per the order of the inputs in the
tuple.

For single-stage models, (i.e. ones that simply take a matrix as input), this returns an
empty tuple.
Wrapper models which do not expose their inner models to seperate inputs, including ones
that only wrap a single model, should **not** define `submodels` as they are
(from the outside API perspective) single-stage models.
"""
submodels(::Union{Template, Model}) = ()

include("traits.jl")
include("test_utils.jl")

end # module

ext: jl


text:
N = 6

ug = SimpleWeightedGraph(N)
add_edge!(ug, 1, 2, 2); add_edge!(ug, 1, 3, 2); add_edge!(ug, 2, 3, 1)
add_edge!(ug, 3, 4, 5); add_edge!(ug, 2, 5, 2); add_edge!(ug, 3, 6, 2)

dg = SimpleWeightedDiGraph(N)
add_edge!(dg, 1, 3, 2); add_edge!(dg, 2, 3, 2); add_edge!(dg, 1, 6, 1)
add_edge!(dg, 2, 5, -2); add_edge!(dg, 3, 4, -2); add_edge!(dg, 3, 5, -1)

el_ug = Vector{Int64}[[2, 3], [1, 3, 5], [1, 2, 4, 6], [3], [2], [3]]
el_dg = Vector{Int64}[[3, 6], [3, 5], [4, 5], [], [], []]

@testset "weightedgraph" begin
    fg = FeaturedGraph(ug)
    @test adjacency_list(fg) == el_ug

    fg = FeaturedGraph(dg)
    @test adjacency_list(fg) == el_dg
end


text:
import ChainRulesCore: rrule
import Umlaut: __new__

###############################################################################
#                                 Rules                                       #
###############################################################################

# TODO: figure out failing test_rrules, move to actual tests

# function rrule(::typeof(Broadcast.broadcasted), ::typeof(identity), x)
#     identity_pullback(dy) = (NoTangent(), NoTangent(), dy)
#     return x, identity_pullback
# end

# test_rrule(broadcasted, identity, [1.0, 2.0]) -- fails at the moment

function rrule(::YotaRuleConfig, ::typeof(Core._apply_iterate),
        ::typeof(iterate), f::F, args...) where F
    # flatten nested arguments
    flat = []
    for a in args
        push!(flat, a...)
    end
    # apply rrule of the function on the flat arguments
    y, pb = rrule_via_ad(YotaRuleConfig(), f, flat...)
    sizes = map(length, args)
    function _apply_iterate_pullback(dy)
        if dy isa NoTangent
            return ntuple(_-> NoTangent(), length(args) + 3)
        end
        flat_dargs = pb(dy)
        df = flat_dargs[1]
        # group derivatives to tuples of the same sizes as arguments
        dargs = []
        j = 2
        for i = 1:length(args)
            darg_val = flat_dargs[j:j + sizes[i] - 1]
            if length(darg_val) == 1 && darg_val[1] isa NoTangent
                push!(dargs, darg_val[1])
            else
                darg = Tangent{typeof(darg_val)}(darg_val...)
                push!(dargs, darg)
            end
            j = j + sizes[i]
        end
        return NoTangent(), NoTangent(), df, dargs...
    end
    return y, _apply_iterate_pullback
end

function ChainRulesCore.rrule(::YotaRuleConfig, ::typeof(tuple), args...)
    y = tuple(args...)
    return y, dy -> (NoTangent(), collect(dy)...)
end

# test_rrule(tuple, 1, 2, 3; output_tangent=Tangent{Tuple}((1, 2, 3)), check_inferred=false)

function rrule(::YotaRuleConfig, nt::Type{NamedTuple{names}}, t::Tuple) where {names}
    val = nt(t)
    function namedtuple_pullback(dy)
        return NoTangent(), dy
    end
    return val, namedtuple_pullback
end

# test_rrule(YotaRuleConfig(), NamedTuple{(:dims,)}, (1,))

rrule(::Type{Val{V}}) where V = Val{V}(), dy -> (NoTangent(),)

###############################################################################
#                                 Broadcast                                   #
###############################################################################

# unzip taken from Zygote:
# https://github.com/FluxML/Zygote.jl/blob/d5be4d5ca80e79278d714eaac15ca71904a262e3/src/lib/array.jl#L177-L185
struct StaticGetter{i} end
(::StaticGetter{i})(v) where {i} = v[i]

@generated function _unzip(tuples, ::Val{N}) where {N}
  Expr(:tuple, (:(map($(StaticGetter{i}()), tuples)) for i ∈ 1:N)...)
end

function unzip(tuples)
  N = length(first(tuples))
  _unzip(tuples, Val(N))
end

function rrule(::YotaRuleConfig, ::typeof(Broadcast.broadcasted), f::F, args...) where F
    return rrule_via_ad(YOTA_RULE_CONFIG, broadcasted, f, args...)
end

function rrule(::YotaRuleConfig, ::typeof(Broadcast.materialize), x)
    return Broadcast.materialize(x), dy -> (NoTangent(), dy)
end

# test_rrule(Broadcast.materialize, Broadcast.broadcasted(sin, rand(3)); output_tangent=ones(3))

function rrule(::YotaRuleConfig, ::typeof(Broadcast.broadcasted), ::typeof(+), x, y)
    function plus_bcast_pullback(dy)
        # println("dy- $dy")
        dy = unthunk(dy)
        # println("dy+ $dy")
        return NoTangent(), NoTangent(), unbroadcast(x, dy), unbroadcast(y, dy)
    end
    return x .+ y, plus_bcast_pullback
end

function rrule(::YotaRuleConfig, ::typeof(Broadcast.broadcasted), ::typeof(*), x, y)
    function mul_bcast_pullback(dy)
        dy = unthunk(dy)
        return NoTangent(), NoTangent(), unbroadcast_prod_x(x, y, dy), unbroadcast_prod_y(x, y, dy)
    end
    return x .* y, mul_bcast_pullback
end

function rrule(::YotaRuleConfig, ::typeof(Broadcast.broadcasted), ::typeof(Base.literal_pow),
        ::typeof(^), x, ::Val{p}) where p
    y = Broadcast.broadcasted(Base.literal_pow, ^, x, Val(p))
    function literal_pow_pullback(dy)
        dy = unthunk(dy)
        return NoTangent(), NoTangent(), NoTangent(), (@. p * x ^ (p - 1) * dy), ZeroTangent()
    end
    return y, literal_pow_pullback
end

function rrule(::YotaRuleConfig, ::typeof(Broadcast.broadcasted), ::typeof(^), x, p::Real)
    y = x .^ p
    function pow_pullback(dy)
        dy = unthunk(dy)
        return NoTangent(), NoTangent(), (@. p * x ^ (p - 1) * dy), ZeroTangent()
    end
    return y, pow_pullback
end

###############################################################################
#                        getindex, getfield, __new__                          #
###############################################################################

function rrule(::YotaRuleConfig, ::typeof(getindex), x, I...)
    y = getindex(x, I...)
    function getindex_pullback(dy)
        dy = unthunk(dy)
        return NoTangent(), ungetindex(x, dy, I...), [ZeroTangent() for i in I]...
    end
    return y, getindex_pullback

end

function rrule(::YotaRuleConfig, ::typeof(getproperty), s, f::Symbol)
    y = getproperty(s, f)
    function getproperty_pullback(dy)
        dy = unthunk(dy)
        T = typeof(s)
        nt = NamedTuple{(f,)}((dy,))
        return NoTangent(), Tangent{T}(; nt...), ZeroTangent()
    end
    return y, getproperty_pullback
end

function rrule(::YotaRuleConfig, ::typeof(getfield), s::Tuple, f::Int)
    y = getfield(s, f)
    function tuple_getfield_pullback(dy)
        dy = unthunk(dy)
        T = typeof(s)
        # deriv of a tuple is a Tangent{Tuple}(...) with all elements set to ZeroTangent()
        # except for the one at index f which is set to dy
        return NoTangent(), Tangent{T}([i == f ? dy : ZeroTangent() for i=1:length(s)]...), ZeroTangent()
    end
    return y, tuple_getfield_pullback
end

function rrule(::YotaRuleConfig, ::typeof(__new__), T, args...)
    y = __new__(T, args...)
    function __new__pullback(dy)
        dy = unthunk(dy)
        if dy isa NoTangent || dy isa ZeroTangent
            fld_derivs = [dy for fld in fieldnames(T)]
        else
            fld_derivs = [getproperty(dy, fld) for fld in fieldnames(T)]
        end
        return NoTangent(), NoTangent(), fld_derivs...
    end
    return y, __new__pullback
end

###############################################################################
#                                   iterate                                   #
###############################################################################

function rrule(::YotaRuleConfig, ::typeof(iterate), x::AbstractArray)
    y = iterate(x)
    function iterate_pullback(dy)
        dy = unthunk(dy)
        return NoTangent(), ungetindex(x, dy, 1)
    end
    return y, iterate_pullback
end

function rrule(::YotaRuleConfig, ::typeof(iterate), x::AbstractArray, i::Integer)
    y = iterate(x, i)
    function iterate_pullback(dy)
        dy = unthunk(dy)
        return NoTangent(), ungetindex(x, dy, i), ZeroTangent()
    end
    return y, iterate_pullback
end

function rrule(::YotaRuleConfig, ::typeof(iterate), t::Tuple)
    y = iterate(t)
    function iterate_pullback(dy)
        dy = unthunk(dy)
        return NoTangent(), ungetfield(dy[1], t, 1)
    end
    return y, iterate_pullback
end

function rrule(::YotaRuleConfig, ::typeof(iterate), t::Tuple, i::Integer)
    y = iterate(t, i)
    function iterate_pullback(dy)
        dy = unthunk(dy)
        return NoTangent(), ungetfield(dy[1], t, i), ZeroTangent()
    end
    return y, iterate_pullback
end

# here we explicitely stop propagation in iteration
# over ranges (e.g for i=1:3 ... end)
function rrule(::YotaRuleConfig, ::typeof(iterate), x::UnitRange)
    y = iterate(x)
    function iterate_pullback(dy)
        return NoTangent(), ZeroTangent()
    end
    return y, iterate_pullback
end
function rrule(::YotaRuleConfig, ::typeof(iterate), x::UnitRange, i::Integer)
    y = iterate(x, i)
    function iterate_pullback(dy)
        return NoTangent(), ZeroTangent(), ZeroTangent()
    end
    return y, iterate_pullback
end

## tuple unpacking

function rrule(::YotaRuleConfig, ::typeof(Base.indexed_iterate), t::Tuple, i::Int)
    y = Base.indexed_iterate(t, i)
    function indexed_iterate_pullback(dy)
        d_val, d_state = dy
        return NoTangent(), ungetfield(d_val, t, i), ZeroTangent()
    end
    return y, indexed_iterate_pullback
end

function rrule(::YotaRuleConfig, ::typeof(Base.indexed_iterate), t::Tuple, i::Int, state::Int)
    y = Base.indexed_iterate(t, i, state)
    function indexed_iterate_pullback(dy)
        d_val, d_state = dy
        return NoTangent(), ungetfield(d_val, t, i), ZeroTangent(), ZeroTangent()
    end
    return y, indexed_iterate_pullback
end

# ## tuple construction

# ∇tuple(dy, ::typeof(tuple), args...) = (NoTangent(), [dy[i] for i=1:length(args)]...)
# #@drule tuple(args::Vararg) ∇tuple

# ## some no diff functions

# #@drule Core.kwfunc(f::Any) (dy, _, f) -> NoTangent()

# ## cat & co.

# function ∇cat_kw(dy, ::typeof(Core.kwfunc(cat)), kw::Any, ::typeof(cat), arrs...)
#     return (
#         NoTangent(),
#         NoTangent(),
#         NoTangent(),
#         [uncat(dy, i, arrs...; dims=kw.dims) for i=1:length(arrs)]...
#     )
# end
# #@drule Core.kwfunc(cat)(kw::Any, _::typeof(cat), arrs::Vararg) ∇cat_kw

# function ∇vcat(dy, ::typeof(vcat), arrs...)
#     return NoTangent(), [uncat(dy, i, arrs...; dims=1) for i=1:length(arrs)]...
# end
# #@drule vcat(arrs::Vararg) ∇vcat

# function ∇hcat(dy, ::typeof(hcat), arrs...)
#     return NoTangent(), [uncat(dy, i, arrs...; dims=2) for i=1:length(arrs)]...
# end
# #@drule hcat(arrs::Vararg) ∇hcat

## Colon

#@drule Colon()(a::Int, b::Int) NoTangent()

function rrule(::YotaRuleConfig, ::Colon, a::Int, b::Int)
    y = a:b
    function colon_pullback(dy)
        return NoTangent(), NoTangent(), NoTangent()
    end
    return y, colon_pullback
end


text:
"""
A `RateModel` contains the parameters to calculate the contribution of a
growth or stress factor to the overall population growth rate.

A [`rate`](@ref) method corresponding to the model type must be defined specifying
the formulation, and optionally a [`condition`](@ref) method to e.g. ignore data below
thresholds.

## Returns

An `AbstractFloat` with implicit units of `N*N^-1*T^-1`, where `N` is the
number of individuals and `T` is time, usually one day.

It should _not_ return a Unitful.jl `Quantity`.
"""
abstract type RateModel end

"""
The intrinsic rate of population growth is the exponential growth rate of a
population when growth is not limited by density dependent factors.

More formally, if the change in population size ``N`` with time ``t`` is
expressed as ``\\frac{dN}{dt} = rN``, then ``r`` is the intrinsic population
growth rate (individuals per day per individual) which can be decomposed into
per capita reproduction and mortality rate. The intrinsic growth rate parameter
``r`` depends strongly on temperature, with population growth inhibited at low
and high temperatures (Haghani et al. 2006).

This can be described using a variety of non-linear functions,
see [`SchoolfieldIntrinsicGrowth`](@ref) for an implementation.
"""
abstract type GrowthModel <: RateModel end

"""
Extreme stressor mortality can be assumed to occur once an environmental
variable ``s`` exceeds some threshold (e.g. critical thermal maximum), beyond
which the mortality rate scales approximately linearly with the depth of the
stressor (Enriquez and Colinet 2017).

Stressor induced mortality is incorporated in growth rates by quantifying the
`threshold` parameter ``s_c`` beyond which stress associated mortality
commences, and the `mortalityrate` parameter ``m_s`` which reflects the
per-capita mortality per stress unit, per time (e.g. degrees beyond the stress
threshold per day).

When stressors are combined in a model, the mortality rate for each stressor
``s`` is incorporated as ``\\frac{dN}{dt}=(r_p-r_n)N`` where ``r_n = \\sum{s}f(s,
c_s)m_s`` and ``f(s, c_s)`` is a function that provides the positive units by
which s exceeds ``s_c``.

This relies on the simplifying assumption that different stressors contribute
additively to mortality rate, but conveniently allows the mortality response of
stressors to be partitioned and calculated separately to the intrinsic growth
rate, which aids parameterisation and modularity. More complicated stress
responses are more likely to capture reality more completely, but the
availability of data on detailed mortality responses to climatic stressors is
often lacking.

Upper and lower stresses may be in relation to any environmental variable,
specied with the `key` parameter, that will use the data at that key in `GeoStack`
for the current timestep.
"""
abstract type StressModel <: RateModel end

"""
    rate(m::RateModel, x)

Returns a growth rate multiplier for a model given environmental
variable x, for a single cell. Must be defined for all `RateModel`.

## Arguments

- `m`: the object containing the model parameters
- `x`: the value at a particular location for of an environmental variable
  specified by the models `key()` method, using corresponding to its `key` field.

`rate` methods should use the `@inline` macro to force inlining.
"""
function rate end

"""
    condition(m::RateModel, x)

Subset the data before applying the rate method, given the model `m` and
the environmental variable of value x.

Returns a `Bool`. `true` will run the `rate` model, `false` will keep the
current growth rate, which is `1` for a single model.

If not defined for a given `RateModel`, `condition` returns `true` for all cells.

`condition` methods should use the `@inline` macro to force inlining.
"""
function condition end
condition(m, x) = true

abstract type AbstractLowerStress <: StressModel end

@inline rate(m::AbstractLowerStress, x) = (m.threshold - x) * m.mortalityrate
@inline condition(m::AbstractLowerStress, x) = x < m.threshold

"""
    LowerStress(key::Symbol, threshold, mortalityrate)

A [`StressModel`](@ref) where stress occurs below a `threshold` at a
specified `mortalityrate` for some environmental layer `key`.

If units are used, the `Layer` that wraps this models must be in units convertible to the
units used in `threshold`.
"""
struct LowerStress{T,M} <: AbstractLowerStress
    threshold::T
    mortalityrate::M
end

abstract type AbstractUpperStress <: StressModel end

@inline condition(m::AbstractUpperStress, x) = x > m.threshold
@inline rate(m::AbstractUpperStress, x) = (x - m.threshold) * m.mortalityrate

"""
    UpperStress(key::Symbol, threshold, mortalityrate)

A [`StressModel`](@ref) where stress occurs above a `threshold` at the
specified `mortalityrate`, for some environmental layer `key`.
"""
struct UpperStress{T,M} <: AbstractUpperStress
    threshold::T
    mortalityrate::M
end

"""
SchoolfieldIntrinsicGrowth(p, ΔH_A, ΔH_L, ΔH_H, Thalf_L, Thalf_H, T_ref, R)

A [`GrowthModel`](@ref) where the temperature response of positive growth rate is
modelled following Schoolfield et al. 1981, _"Non-linear regression of biological
temperature-dependent rate models base on absolute reaction-rate theory"_.

The value of the specified data layer _must_ be in Kelvin.

## Arguments
- `p::P`: growth rate at reference temperature `T_ref`
- `ΔH_A`: enthalpy of activation of the reaction that is catalyzed by the enzyme, in cal/mol or J/mol.
- `ΔH_L`: change in enthalpy associated with low temperature inactivation of the enzyme.
- `ΔH_H`: change in enthalpy associated with high temperature inactivation of the enzyme.
- `T_halfL`: temperature (Unitful K) at which the enzyme is 1/2 active and 1/2 low temperature inactive.
- `T_halfH`: temperature (Unitful K) at which the enzyme is 1/2 active and 1/2 high temperature inactive.
- `T_ref`: Reference temperature in kelvin.

This can be parameterised from empirical data, (see Zhang et al. 2000; Haghani
et al. 2006; Chien and Chang 2007) and non-linear least squares regression.
"""
struct SchoolfieldIntrinsicGrowth{P,HA,HL,HH,TL,TH,TR} <: GrowthModel
    p::P
    ΔH_A::HA
    ΔH_L::HL
    T_halfL::TL
    ΔH_H::HH
    T_halfH::TH
    T_ref::TR
end

@inline function rate(m::SchoolfieldIntrinsicGrowth, x)
    @fastmath m.p * x/m.T_ref * exp(m.ΔH_A/R * (1/m.T_ref - 1/x)) /
        (1 + exp(m.ΔH_L/R * (1/m.T_halfL - 1/x)) + exp(m.ΔH_H/R * (1/m.T_halfH - 1/x)))
end

# Layers

"""
    Layer{K,U}(model::RateModel)
    Layer(key::Symbol, model::RateModel)
    Layer(key::Symbol, u::Union{Units,Quantity}, model::RateModel)

Layers connect a model to a data source, providing the key to look
up the layer in a GeoData.jl `GeoStack`, and specifying the
scientific units of the layer, if it has units. Using units adds
an extra degree of safety to your calculation, and allows for using
data in different units with the same models.
"""
struct Layer{K,U,M}
    model::M
end
Layer{K,U}(model::RateModel) where {K,U} = Layer{K,U,typeof(model)}(model)
Layer(key::Symbol, model::RateModel) = Layer{key,unit(u"1"),typeof(model)}(model)
Layer(key::Symbol, u::Units, model::RateModel) = Layer{key,u,typeof(model)}(model)
Layer(key::Symbol, q::Quantity, model::RateModel) =
    Layer{key,typeof(q),typeof(model)}(model)

model(l::Layer) = l.model

@inline rate(l::Layer, x) = rate(l.model, x)
@inline condition(l::Layer, x) = condition(l.model, x)

ConstructionBase.constructorof(::Type{<:Layer{K,U}}) where {K,U} = Layer{K,U}

Base.keys(::Layer{K}) where K = K
Base.keys(layers::Tuple{Vararg{<:Layer}}) = map(keys, layers)

Unitful.unit(::Layer{K,U}) where {K,U} = U


text:
using BuildMachines
using Test

@testset "BuildMachines.jl" begin
    @test BuildMachines._hello_world() == "Hello world!"
end


text:
using SparseArrays
using LinearAlgebra
using CUDA, CUDA.CUSPARSE, CUDA.CUBLAS
CUDA.allowscalar(false)

# vec for CuArrays returns a ReshapedArray in some versions, which do not work for mul! with CUDA.CUSPARSE matrices
# so define our own version
myvec(A::CuArray) = unsafe_wrap(CuArray,pointer(A),length(A))
myvec(A) = vec(A)

# L = -i (Heff ρ - ρ Heff^†) + sum_i J_i ρ J_i^†
# L = -i Heff ρ + i ρ Heff^† + sum_i J_i ρ J_i^†
# L = (-i Heff ρ) + (-i Heff ρ^†)^† + sum_i J_i ρ J_i^†
function L!(Lρ,ρ,ps,t)
    miHeff, J, tempρ, tempρ2, hermρ = ps
    LHeff!(Lρ,ρ,miHeff,tempρ,hermρ)
    Lrefill!(Lρ,ρ,J,tempρ,tempρ2)
    Lρ
end

# assume that ρ is Hermitian, so (-i Heff ρ^†)^† = (-i Heff ρ)^†
# miHeff = -i Heff
function LHeff!(Lρ,ρ,miHeff,tempρ,hermρ::Val{true})
    mul!(tempρ,miHeff,ρ)
    adjoint!(Lρ,tempρ)
    Lρ .+= tempρ
end
# Take into account that ρ might not be Hermitian
# (because of numerical error or because we propagate a non-physical ρ)
function LHeff!(Lρ,ρ,miHeff,tempρ,hermρ::Val{false})
    mul!(Lρ,miHeff,ρ)
    mulacc!(Lρ,ρ,miHeff',tempρ)
end

function Lrefill!(Lρ,ρ,J::AbstractVector,tempρ,tempρ2)
    for j in J
        mul!(tempρ,j,ρ)
        mulacc!(Lρ,tempρ,j',tempρ2)
    end
end
# if J is a matrix, we assume it is a superoperator
function Lrefill!(Lρ,ρ,J::AbstractMatrix,temp,temp2)
    mul!(myvec(Lρ),J,myvec(ρ),1,1)
end

# calculate C += A*B
function mulacc!(C,A,B,tmp)
    mul!(C,A,B,1,1)
end
# if B is CuSparseMatrixCSR or its adjoint, calculate tmp = B'*A', and then C += tmp'
function mulacc!(C,A,B::U,tmp) where U <: Union{CuSparseMatrixCSR{T},Adjoint{T,CuSparseMatrixCSR{T}}} where T
    mul!(tmp,B',A')
    CUBLAS.geam!('N','C',1,C,1,tmp,C)
end


text:
function branch_time_id(tree::Tree, times::Array{Float64,1})
    nb = node_depths_base(tree)
    ne = node_depths(tree)

    rd= maximum(ne)
    ids = Array{Array{Int64,1},1}(undef,0)

    for t in times
        if t <= rd
            idt = Array{Int64,1}(undef,0)

            for i in 1:length(nb)
                if nb[i] <= t < ne[i]
                    push!(idt,i)
                end
            end
            push!(ids,idt)
        end
    end

    return ids
end

function branch_time_id(tree::Tree, edge_trees::Array{EdgeTreeRates2,1}, times::Array{Float64,1})
    nb = node_depths_base(tree, edge_trees)

    ne = node_depths(tree,edge_trees)

    rd= maximum(ne)
    ids = Array{Array{Int64,1},1}(undef,0)

    for t in times
        if t <= rd
            idt = Array{Int64,1}(undef,0)

            for i in 1:length(nb)
                if nb[i] <= t < ne[i]
                    push!(idt,i+1)
                end
            end
            push!(ids,idt)
        end
    end

    return ids
end

function time_rates(ids::Array{Array{Int64,1},1}, row::Array{Float64,1})
    mean_rates = Array{Float64,1}(undef,0)

    for id in ids
        mr = 0.
        n = 0.
        for i in id
            n += 1.
            mr += row[i]
        end
        push!(mean_rates, mr/n)
    end
    return mean_rates
end

function time_rates(tree::Tree, row::Array{Float64,1}, times::Array{Float64,1})
    ids = branch_time_id(tree, times)

    time_rates(ids, row)
end

function time_rates(tree::Tree, times::Array{Float64,1})
    ids = branch_time_id(tree, times)
    row = extract_rates(tree)
    time_rates(ids, row)
end

function time_rates(tree::Tree, edge_trees::Array{EdgeTreeRates2,1}, times::Array{Float64,1})
    ids = branch_time_id(tree, edge_trees, times)
    row = extract_rates(tree, edge_trees)
    time_rates(ids, row)
end

max_stars_repo_name: hmorlon/JPANDA

max_issues_repo_name: hmorlon/JPANDA

max_forks_repo_name: hmorlon/JPANDA


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

using Cxx

cxxinclude("TTree.h")
cxxinclude("TChain.h")

export create_ttree!
export bind_branch!
export create_branch!
export getname
export gettitle

export TTreePtr
const TTreePtr = pcpp"TTree"

export TChain, TChainPtr, TChainInst
const TChain = cxxt"TChain"
const TChainPtr = pcpp"TChain"
const TChainInst = Union{TChain, TChainPtr}

export ATTreeInst
const ATTreeInst = Union{TTreePtr, TChainInst}

create_ttree!(tdir::ATDirectoryInst, name::AbstractString, title::AbstractString) =
    cd(() -> icxx""" new TTree($name, $title); """, tdir)

Base.size(ttree::ATTreeInst) = (length(ttree),)
Base.length(ttree::ATTreeInst) = Int(@cxx ttree->GetEntries())
Base.firstindex(ttree::ATTreeInst) = 1
Base.lastindex(ttree::ATTreeInst) = length(ttree)

Base.getindex(ttree::ATTreeInst, i::Integer) = begin
    @cxx ttree->GetEntry(i - 1)
    nothing
end

Base.iterate(ttree::ATTreeInst) = iterate(ttree, firstindex(ttree))

function Base.iterate(ttree::ATTreeInst, i::Integer)
    if i <= lastindex(ttree)
        ( getindex(ttree, i), i+1 )
    else
        nothing
    end
end

Base.push!(ttree::TTreePtr) = icxx""" $ttree->Fill(); """

getname(obj::ATTreeInst) = unsafe_string(@cxx obj->GetName())
gettitle(obj::ATTreeInst) = unsafe_string(@cxx obj->GetTitle())

bind_branch!(ttree::ATTreeInst, name::AbstractString, value::Ref{T}) where {T<:Union{Number, CppPtr}} = begin
    ttree_obj = auto_deref(ttree)
    icxx"""
        $ttree_obj.SetBranchStatus($name, true);
        $ttree_obj.AddBranchToCache($name, true);
        $ttree_obj.SetBranchAddress($name, &$value);
    """
end

# Workaround for Int64, else ROOT will give this error:
# The pointer type given "Long_t" does not correspond to the type needed "Long64_t"
bind_branch!(ttree::ATTreeInst, name::AbstractString, value::Ref{Int64}) = begin
    ttree_obj = auto_deref(ttree)
    icxx"""
        $ttree_obj.SetBranchStatus($name, true);
        $ttree_obj.AddBranchToCache($name, true);
        $ttree_obj.SetBranchAddress($name, (Long64_t*)&$value);
    """
end

# Workaround for UInt64, else ROOT will give this error:
# The pointer type given "ULong_t" does not correspond to the type needed "ULong64_t"
bind_branch!(ttree::ATTreeInst, name::AbstractString, value::Ref{UInt64}) = begin
    ttree_obj = auto_deref(ttree)
    icxx"""
        $ttree_obj.SetBranchStatus($name, true);
        $ttree_obj.AddBranchToCache($name, true);
        $ttree_obj.SetBranchAddress($name, (ULong64_t*)&$value);
    """
end

create_branch!(ttree::TTreePtr, name::AbstractString, value::Ref{T}; bufsize::Integer = 32000, splitlevel::Integer = 99) where {T<:Union{Number, CppPtr}} =
    icxx""" $ttree->Branch($name, &$value, $(Int32(bufsize)), $(Int32(splitlevel))); """

create_branch!(ttree::TTreePtr, name::AbstractString, value::CppValue; bufsize::Integer = 32000, splitlevel::Integer = 99) =
    icxx""" $ttree->Branch($name, &$value, $(Int32(bufsize)), $(Int32(splitlevel))); """


text: Dict{String,Any}("a" => Dict{String,Any}("b" => Dict{String,Any}("c" => Dict{String,Any}("answer" => Dict{String,Any}("value" => "42","type" => "integer")))))

max_stars_repo_name: vanillajonathan/julia

max_issues_repo_name: vanillajonathan/julia

max_forks_repo_name: vanillajonathan/julia


text:
apl(f, args...) = f(args...)

function apl(fs::Tuple, args...)
  foreach(f -> f(args...), fs)
end

mutable struct Container
  lmap::NamedTuple
end

const LMAP = Container(NamedTuple())

function apllens(lmap::NamedTuple, nm::Symbol, arg)
  if haskey(lmap, nm)
    fs = getfield(lmap, nm)
    apl(fs, arg)
    arg
  else
    arg
  end
end

function apllens(lmap::NamedTuple, nm::Symbol, args...)
  if haskey(lmap, nm)
    fs = getfield(lmap, nm)
    apl(fs, args...)
    args
  else
    args
  end
end

apllens(nm::Symbol, args...) = apllens(globalmap(), nm, args...)

lens(nm, arg) = apllens(nm, arg)
lens(nm, args...) = apllens(nm, args...)
lens(nm, args::NamedTuple) = apllens(nm, args)

lens(nm; kwargs...) = lens(nm, kwargs)

"Global Map!"
function setgloballmap!(lmap)
  LMAP.lmap::NamedTuple = lmap
end
globalmap()::NamedTuple = LMAP.lmap

function resetglobalmap!()
  setgloballmap!(NamedTuple())
end

"Lens call f(args...)"
function lenscall(lmap::NamedTuple, f, args...; kwargs...)
  setgloballmap!(lmap)
  f(args...; kwargs...)
end

"""
Lensed eval

```julia
function g(x, y)
  lens(:howdy, (x = x, y = y))
  2x + y
end

@leval g(1, 2) (howdy = println ∘ sum ∘ values,)
```
"""
macro leval(expr, lmap)
  quote
    setgloballmap!($(esc(lmap)))
    res = $(esc(expr))
    resetglobalmap!()
    res
  end
end


text:
using DataFrames, Gadfly
df = readtable("/Users/Matthieu/Dropbox/Github/FixedEffectModels/benchmark/benchmark.csv")
mdf = melt(df[[:Command, :Julia, :R, :Stata]], :Command)
mdf = rename(mdf, :variable, :Language)
p = plot(mdf, x = "Command", y = "value", color = "Language", Guide.ylabel("Time (seconds)"), Guide.xlabel("Model"),  Theme(default_point_size=1mm))
draw(SVG("result.svg", 8inch, 5inch), p)


text:
using Revise

using Devito

configuration!("log-level", "DEBUG")
configuration!("language", "openmp")
configuration!("mpi", false)

function ricker(f, _t, t₀)
    t = reshape(_t, length(_t), 1)
    (1.0 .- 2.0 * (pi * f * (t .- t₀)).^2) .* exp.(-(pi * f * (t .- t₀)).^2)
end

grid = Grid(
    shape = (121,121,121),
    origin = (0.0,0.0,0.0),
    extent = (1210.0,1210.0,1210.0),
    dtype = Float32)

b = Devito.Function(name="b", grid=grid, space_order=8)
v = Devito.Function(name="v", grid=grid, space_order=8)
q = Devito.Function(name="woverq", grid=grid, space_order=8)

b_data = data(b)
v_data = data(v)
q_data = data(q)

# note that the folowing is faster if we include the halo
# i.e. use the data_with_halo method rather than the data method.
b_data .= 1
v_data .= 1.5
q_data .= 1/1000

time_range = 0.0f0:1.0f0:750.0f0

p = TimeFunction(name="p", grid=grid, time_order=2, space_order=8)
z,y,x,t = dimensions(p)

src = SparseTimeFunction(name="src", grid=grid, f0=0.01f0, npoint=1, nt=length(time_range), coordinates=[605.0 605.0 10.0])
src_data = data(src)
w = ricker(0.01, collect(time_range), 125)
copy!(src_data, w)
src_term = inject(src; field=forward(p), expr=src*spacing(t)^2*v^2/b)

nz,ny,nx,δz,δy,δx = size(grid)...,spacing(grid)...
rec = SparseTimeFunction(name="rec", grid=grid, npoint=nx, nt=length(time_range))
rec_coords = coordinates(rec)
rec_coords[1,:] .= δx*(0:nx-1)
rec_coords[2,:] .= 605
rec_coords[3,:] .= 20
rec_term = interpolate(rec, expr=p)

g1(field) = dx(field,x0=x+spacing(x)/2)
g2(field) = dy(field,x0=y+spacing(y)/2)
g3(field) = dz(field,x0=z+spacing(z)/2)
g1_tilde(field) = dx(field,x0=x-spacing(x)/2)
g2_tilde(field) = dy(field,x0=y-spacing(y)/2)
g3_tilde(field) = dz(field,x0=z-spacing(z)/2)

update_p = spacing(t)^2 * v^2 / b *
    (g1_tilde(b * g1(p)) + g2_tilde(b * g2(p)) + g3_tilde(b * g3(p))) +
    (2 - spacing(t) * q) * p +
    (spacing(t) * q - 1) * backward(p)

stencil_p = Eq(forward(p), update_p)

dt = step(time_range)
smap = spacing_map(grid)
smap[spacing(t)] = dt

op = Operator([stencil_p, src_term, rec_term], subs=smap, name="OpExampleIso")

summary = apply(op)

using PyPlot

_p = data_with_halo(p)
extrema(_p)

d = data(rec)
extrema(d)

figure();imshow(_p[:,:,60,1])
display(gcf())

figure();imshow(d, aspect="auto",cmap="gray",clim=.1*[-1,1]*maximum(abs,d))
display(gcf())


text:
# hilbert ordering
@mesh(HEXMESH, n)
hilbert_elements([n,n,n]);
ef_nodes();

solve([q1,q2,q3,q4,q5,q6]);

regtime = Base.Libc.time();
for iter=1:times
    solve([q1,q2,q3,q4,q5,q6]);
end
regtime = Base.Libc.time() - regtime;
regtime /= times;
timings[1] = regtime;
println("Finished Hilbet.");


text:
function precompute_lgp(θs::Matrix)
    ks      = Array{KernelFunctions.Kernel}(undef, size(θs, 2))
    for idx = 1:size(θs, 2)
        ks[idx] = construct_kernel(θs[1,idx], θs[2,idx], θs[3:end,idx],)
    end
    ks
end

@inline function gp_predict(x::AbstractVector,
                            X::Matrix,
                            K::PDMats.PDMat,
                            a::Vector,
                            k::KernelFunctions.Kernel)
    k_star = KernelFunctions.kernelmatrix(k, reshape(x, (:,1)), X, obsdim=2)[1,:]
    μ      = dot(k_star, a)
    σ²     = k(x, x) - PDMats.invquad(K, k_star)
    μ, σ²
end

@inline function gp_predict(x::AbstractMatrix,
                            X::Matrix,
                            K::PDMats.PDMat,
                            a::Vector,
                            k::KernelFunctions.Kernel)
    k_star = KernelFunctions.kernelmatrix(k, x, X, obsdim=2)[1,:]
    μ      = k_star * a
    σ²     = @inbounds [k(x[:,i], x[:,i]) - PDMats.invquad(K, k_star[i,:]) for i = 1:size(x,2)]
    μ, σ²
end

@inline function gp_predict(x::AbstractVector,
                            X::Matrix,
                            K::Array{PDMats.PDMat},
                            a::Matrix,
                            ks::Array{KernelFunctions.Kernel})
    N   = length(K)
    μs  = zeros(N)
    σ²s = zeros(N)
    @simd for idx = 1:length(K)
        k        = ks[idx]
        k_star   = KernelFunctions.kernelmatrix(k, reshape(x, (:,1)), X, obsdim=2)[1,:]
        μ        = dot(k_star, a[:,idx])
        σ²       = k(x, x) - PDMats.invquad(K[idx], k_star)
        μs[idx]  = μ
        σ²s[idx] = σ²
    end
    mean(μs), mean(σ²s)
end

@inline function gp_predict(x::AbstractMatrix,
                            X::Matrix,
                            K::Array{PDMats.PDMat},
                            a::Matrix,
                            ks::Array{KernelFunctions.Kernel})
    N   = length(K)
    μs  = zeros(size(x, 2), N)
    σ²s = zeros(size(μs))
    @simd for idx = 1:N
        k      = ks[idx]
        k_star = KernelFunctions.kernelmatrix(k, x, X, obsdim=2)
        μ      = k_star * a[:,idx]
        σ²     = @inbounds [k(x[:,i], x[:,i]) - PDMats.invquad(K[idx], k_star[i,:])
                            for i = 1:size(x,2)]
        μs[:,idx]  = μ
        σ²s[:,idx] = σ²
    end
    mean(μs, dims=2)[:,1], mean(σ²s, dims=2)[:,1]
end


text:
"""
control points of Bézier curve from given function and range
"""
function _bezier(f,a,b)
    f(a),
    3*(f(2*a/3+b/3)-(8*f(a)+f(b))/27)-3*(f(a/3+2*b/3)-(f(a)+8*f(b))/27)/2,
    -3*(f(2*a/3+b/3)-(8*f(a)+f(b))/27)/2+3*(f(a/3+2*b/3)-(f(a)+8*f(b))/27),
    f(b)
end

function _luxor_pt(p::AbstractVector{<:Real}, unitlength)
    Luxor.Point(unitlength*[1,-1].*p...)
end

"""
export svg file
"""
function save_svg(name::AbstractString, M::AbstractBSplineManifold{1}; xlims=(-5,5), ylims=(-5,5), mesh=(10,10), unitlength=100, points=true, thickness=1, backgroundcolor=RGB(1,1,1), maincolor=RGB(1,0,0))
    split(name,'.')[end] ≠ "svg" && error("extension shuould be .svg")
    _save_luxor_1d2d(name, M, xlims=xlims, ylims=ylims, mesh=mesh, unitlength=unitlength, points=points, thickness=thickness, backgroundcolor=backgroundcolor, maincolor=maincolor)
end

"""
export svg file
"""
function save_svg(name::AbstractString, M::AbstractBSplineManifold{2}; xlims=(-5,5), ylims=(-5,5), mesh=(10,10), unitlength=100, points=true, thickness=1, backgroundcolor=RGB(1,1,1), maincolor=RGB(1,0,0))
    split(name,'.')[end] ≠ "svg" && error("extension shuould be .svg")
    _save_luxor_2d2d(name, M, xlims=xlims, ylims=ylims, mesh=mesh, unitlength=unitlength, points=points, thickness=thickness, backgroundcolor=backgroundcolor, maincolor=maincolor)
end

"""
export png file
"""
function save_png(name::AbstractString, M::AbstractBSplineManifold{1}; xlims=(-5,5), ylims=(-5,5), mesh=(10,10), unitlength=100, points=true, thickness=1, backgroundcolor=RGB(1,1,1), maincolor=RGB(1,0,0))
    split(name,'.')[end] ≠ "png" && error("extension shuould be .png")
    _save_luxor_1d2d(name, M, xlims=xlims, ylims=ylims, mesh=mesh, unitlength=unitlength, points=points, thickness=thickness, backgroundcolor=backgroundcolor, maincolor=maincolor)
end

"""
export png file
"""
function save_png(name::AbstractString, M::AbstractBSplineManifold{2}; xlims=(-5,5), ylims=(-5,5), mesh=(10,10), unitlength=100, points=true, thickness=1, backgroundcolor=RGB(1,1,1), maincolor=RGB(1,0,0))
    split(name,'.')[end] ≠ "png" && error("extension shuould be .png")
    _save_luxor_2d2d(name, M, xlims=xlims, ylims=ylims, mesh=mesh, unitlength=unitlength, points=points, thickness=thickness, backgroundcolor=backgroundcolor, maincolor=maincolor)
end

"""
export png file
"""
function save_png(name::AbstractString, M::AbstractBSplineManifold{2}, colors::AbstractArray{<:Colorant,2}; xlims=(-5,5), ylims=(-5,5), unitlength=100)
    split(name,'.')[end] ≠ "png" && error("extension shuould be .png")

    _save_luxor_2d2d_color(name, M, colors, xlims=xlims, ylims=ylims, unitlength=unitlength)
end

"""
export png file
"""
function save_png(name::AbstractString, M::AbstractBSplineManifold{2}, colorfunc::Function; xlims=(-5,5), ylims=(-5,5), unitlength=100)
    split(name,'.')[end] ≠ "png" && error("extension shuould be .png")

    _save_luxor_2d2d_color(name, M, colorfunc, xlims=xlims, ylims=ylims, unitlength=unitlength)
end

function _save_luxor_2d2d(name::AbstractString, M::BSplineManifold{2, Deg, <:StaticVector{2}}; xlims=(-5,5), ylims=(-5,5), mesh=(10,10), unitlength=100, points=true, thickness=1, backgroundcolor=RGB(1,1,1), maincolor=RGB(1,0,0), subcolor=RGB(0.5,0.5,0.5)) where Deg
    left, right = xlims
    down, up = ylims
    linecolor = maincolor
    fillcolor = weighted_color_mean(0.5, maincolor, colorant"white")
    segmentcolor = subcolor
    pointcolor = weighted_color_mean(0.5, subcolor, colorant"black")

    P = bsplinespaces(M)
    p = degree.(P)
    k = knotvector.(P)
    a = controlpoints(M)
    n¹, n² = n = dim.(P)

    K¹, K² = K = [unique(k[i][1+p[i]:end-p[i]]) for i in 1:2]
    N¹, N² = length.(K).-1
    m¹, m² = mesh

    Drawing(unitlength*(right-left),unitlength*(up-down),name)
    Luxor.origin(-unitlength*left,unitlength*up)
    setline(thickness)
    background(backgroundcolor)

    setcolor(fillcolor)
    drawbezierpath(BezierPath(vcat(
        [BezierPathSegment(map(p->_luxor_pt(p,unitlength),_bezier(u¹->M(u¹,K²[1]),K¹[i],K¹[i+1]))...) for i in 1:N¹],
        [BezierPathSegment(map(p->_luxor_pt(p,unitlength),_bezier(u²->M(K¹[end],u²),K²[i],K²[i+1]))...) for i in 1:N²],
        [BezierPathSegment(map(p->_luxor_pt(p,unitlength),_bezier(u¹->M(u¹,K²[end]),K¹[end-i+1],K¹[end-i]))...) for i in 1:N¹],
        [BezierPathSegment(map(p->_luxor_pt(p,unitlength),_bezier(u²->M(K¹[1],u²),K²[end-i+1],K²[end-i]))...) for i in 1:N²]
    )),:fill,close=true)

    setcolor(linecolor)
    for u¹ in range(K¹[1],stop=K¹[end],length=m¹+1)
        drawbezierpath(BezierPath([BezierPathSegment(map(p->_luxor_pt(p,unitlength),_bezier(u²->M(u¹,u²),K²[i],K²[i+1]))...) for i in 1:N²]),:stroke)
    end
    for u² in range(K²[1],stop=K²[end],length=m²+1)
        drawbezierpath(BezierPath([BezierPathSegment(map(p->_luxor_pt(p,unitlength),_bezier(u¹->M(u¹,u²),K¹[i],K¹[i+1]))...) for i in 1:N¹]),:stroke)
    end

    if points
        CtrlPts = [_luxor_pt(a[i,j],unitlength) for i in 1:n¹, j in 1:n²]

        setcolor(segmentcolor)
        setline(thickness)
        for i in 1:n[1]
            poly(CtrlPts[i,:], :stroke)
        end
        for j in 1:n[2]
            poly(CtrlPts[:,j], :stroke)
        end

        setcolor(pointcolor)
        map(p->circle(p,3*thickness,:fill), CtrlPts)
    end
    finish()
    return nothing
end

function _save_luxor_1d2d(name::AbstractString, M::BSplineManifold{1, Deg, <:StaticVector{2}}; xlims=(-5,5), ylims=(-5,5), mesh=10, unitlength=100, points=true, thickness=1, backgroundcolor=RGB(1,1,1), maincolor=RGB(1,0,0), subcolor=RGB(.5,.5,.5)) where Deg
    left, right = xlims
    down, up = ylims
    linecolor = maincolor
    segmentcolor = subcolor
    pointcolor = weighted_color_mean(0.5, subcolor, colorant"black")

    a = controlpoints(M)

    P¹ = bsplinespaces(M)[1]
    p¹ = degree(P¹)
    k¹ = knotvector(P¹)
    n¹ = dim(P¹)

    K¹ = unique(k¹[1+p¹:end-p¹])

    N¹ = length(K¹)-1

    Drawing(unitlength*(right-left),unitlength*(up-down),name)
    Luxor.origin(-unitlength*left,unitlength*up)
    setline(2*thickness)
    background(backgroundcolor)

    setcolor(linecolor)
    drawbezierpath(BezierPath([BezierPathSegment(map(q->_luxor_pt(q,unitlength),_bezier(u¹->M(u¹),K¹[i],K¹[i+1]))...) for i in 1:N¹]),:stroke)

    if points
        CtrlPts = [_luxor_pt(a[i],unitlength) for i in 1:n¹]

        setcolor(segmentcolor)
        setline(thickness)
        poly(CtrlPts[:], :stroke)

        setcolor(pointcolor)
        map(q->circle(q,3*thickness,:fill), CtrlPts)
    end
    finish()
    return nothing
end

function _save_luxor_2d2d_color(name::AbstractString, M::BSplineManifold{2, Deg, <:StaticVector{2}}, colors::AbstractArray{<:Colorant,2}; xlims=(-5,5), ylims=(-5,5), unitlength=100) where Deg
    C = BSplineManifold(colors, bsplinespaces(M))
    _save_luxor_2d2d_color(name, M, C; xlims=xlims, ylims=ylims, unitlength=unitlength)
end

function _save_luxor_2d2d_color(name::AbstractString, M::BSplineManifold{2, Deg, <:StaticVector{2}}, colorfunc; xlims=(-5,5), ylims=(-5,5), unitlength=100) where Deg
    left, right = xlims
    down, up = ylims
    mesh = 10

    P = bsplinespaces(M)
    p¹, p² = degree.(P)
    k¹, k² = knotvector.(P)

    K¹ = unique(vcat([collect(range(k¹[i], k¹[i+1], length=mesh+1)) for i in 1+p¹:length(k¹)-p¹-1]...))
    K² = unique(vcat([collect(range(k²[i], k²[i+1], length=mesh+1)) for i in 1+p²:length(k²)-p²-1]...))

    Drawing(unitlength*(right-left),unitlength*(up-down),name)
    Luxor.origin(-unitlength*left,unitlength*up)
    background(RGBA(0,0,0,0))

    for I₁ in 1:length(K¹)-1, I₂ in 1:length(K²)-1
        BézPth=BezierPath([
                BezierPathSegment(map(q->_luxor_pt(q,unitlength),_bezier(t->M(t,K²[I₂]),K¹[I₁],K¹[I₁+1]))...),
                BezierPathSegment(map(q->_luxor_pt(q,unitlength),_bezier(t->M(K¹[I₁+1],t),K²[I₂],K²[I₂+1]))...),
                BezierPathSegment(map(q->_luxor_pt(q,unitlength),_bezier(t->M(t,K²[I₂+1]),K¹[I₁+1],K¹[I₁]))...),
                BezierPathSegment(map(q->_luxor_pt(q,unitlength),_bezier(t->M(K¹[I₁],t),K²[I₂+1],K²[I₂]))...)])
        mesh1 = Luxor.mesh(BézPth, [
            colorfunc(K¹[I₁],   K²[I₂]  ),
            colorfunc(K¹[I₁+1], K²[I₂]  ),
            colorfunc(K¹[I₁+1], K²[I₂+1]),
            colorfunc(K¹[I₁],   K²[I₂+1]),
            ])
        setmesh(mesh1)
        box(_luxor_pt([right+left,up+down]/2,unitlength), (right-left)*unitlength,(up-down)*unitlength,:fill)
    end
    finish()
    return nothing
end


text:
sources = ["im2col.cpp", "pooling.cpp"]

compiler = get(ENV, "CXX", "g++")
flags    = ["-fPIC", "-Wall", "-O3", "-shared"]
libname  = "libmochaext.so"
openmp   = "-fopenmp"

@static is_apple() ? begin
  if !haskey(ENV, "MOCHA_FORCE_OMP")
    println("OpenMP is currently not officially supported by OS X Clang compiler yet.")
    println("(see http://clang-omp.github.io/ to install OpenMP clang extension, or")
    println("install gcc.")
    println("")
    println("Disabling OpenMP. To force enable OpenMP, set MOCHA_FORCE_OMP environment")
    println("variable and set the CXX environment variable to the appropriate value")
    println("to invoke GCC's g++ frontend, such as g++-5")
    println("")
    openmp = ""
  end
end : nothing

cmd = `$compiler $flags $openmp -o $libname $sources`
println("Running $cmd")
run(cmd)


text:
module LazyJSON

using PropertyDicts: PropertyDict

module JSONjl
    using JSON
end

using OrderedCollections: OrderedDict
using Base: @propagate_inbounds

const JSON = LazyJSON

const LAZY_WARN = false
macro lazywarn(a...)
    LAZY_WARN ? :(@warn $(esc(a))...) :
                :()
end

include("SplicedStrings.jl")            ; using .SplicedStrings: SplicedString
include("IOStrings.jl")                 ; using .IOStrings: IOString, pump

# JSON Value Types

"""
Values are represented by a reference to the JSON text `String`
and the byte index of the value text.

    String: {"foo": 1, "bar": [1, 2, 3, "four"]}
            ▲                 ▲      ▲  ▲
            │                 │      │  │
            ├──────────────┐  │      │  │
            │   JSON.Array(s, i=9)   │  │   == Any[1, 2, 3, "four"]
            │                        │  │
            ├───────────────┐  ┌─────┘  │
            │   JSON.Number(s, i=16)    │   == 3
            │                           │
            ├───────────────┐  ┌────────┘
            │   JSON.String(s, i=19)        == "four"
            │
            └───────────────┬──┐
                JSON.Object(s, i=1)
"""
struct String{T <: AbstractString} <: AbstractString
    s::T
    i::Int
end

struct Number{T <: AbstractString} <: Base.Real
    s::T
    i::Int
end

struct Object{W, T <: AbstractString} <: AbstractDict{AbstractString, Any}
    s::T
    i::Int
end

struct Array{W, T <: AbstractString}  <: AbstractArray{Any, 1}
    s::T
    i::Int
end

const Value = Union{JSON.String, JSON.Number, JSON.Object, JSON.Array}
const Collection = Union{JSON.Object, JSON.Array}

Object(W, s::T, i) where T = Object{W, T}(s, i)
Array(W, s::T, i) where T = Array{W, T}(s, i)

"""
Verbatim JSON text of a `JSON.Value`
"""
Base.string(s::JSON.Value) = SubString(s.s, s.i, lastindex_of_value(s.s, s.i))

"""
JSON text representation of `x`
"""
jsonstring(d::PropertyDict) = jsonstring(PropertyDicts.unwrap(d))
jsonstring(x::JSON.Value) = string(x)
jsonstring(x) = JSONjl.JSON.json(x)

"""
Replace value `v` with `x`.
"""
splice(j::JSON.Value, v::JSON.Value, x) = value(splice(j.s, v.i, x, j.i))

splice(j::Union{JSON.Object{PropertyDict},
                JSON.Array{PropertyDict}}, v::JSON.Value, x) =
    value(splice(j.s, v.i, x, j.i); getproperty=true)

splice(j::JSON.Value, path::Vector, x) =
    value(splice(j.s, getpath(j.s, path, j.i)[1], x, j.i))

splice(s::AbstractString, path::Vector, x) = splice(s, getpath(s, path)[1], x)

splice(s::AbstractString, i::Int, x, start_i = 1) =
    SplicedString(SubString(s, start_i, prevind(s, i)),
                  jsonstring(x),
                  SubString(s, next_i(s, lastindex_of_value(s, i))))

splice(d::PropertyDict, v::PropertyDict, x) = splice(PropertyDicts.unwrap(d),
                                                     PropertyDicts.unwrap(v), x)
splice(d::PropertyDict, v, x) = splice(PropertyDicts.unwrap(d), v, x)

# Get Typed JSON Value

const SupportedString = Union{IOString,
                              SplicedString,
                              Base.String}
"""
    value(jsontext) -> JSON.Value
    JSON.Value <: Union{Number, AbstractString, AbstractVector, AbstractDict}

Create a `JSON.Value` object from a JSON text.
"""
function value(s::SupportedString, path=nothing, i = 1;
               lazy=true, getproperty=false)

    # Check that the string has a C-style termination characer.
    @assert (c = getc(s, ncodeunits(s) + 1); c == 0x00 || c == IOStrings.ASCII_ETB)

    i, c = skip_whitespace(s, i)
    if path != nothing
        i, c = getpath(s, path, i, c)
    end
    W = getproperty ? PropertyDict : Nothing
    v = getvalue(W, s, i, c)
    return lazy ? v : flatten(v)
end

struct ParseError <: Exception
    bytes::Base.String
    index::Int
    c::UInt8
    message::Base.String
end

value(ss::SubString{<:SupportedString}, path=nothing; kw...) =
    value(ss.string, path, 1 + ss.offset; kw...)

value(bytes, path=nothing; kw...) = value(Base.String(bytes), path; kw...)

function value(io::IO, path=nothing; kw...)
    s = IOString(io)
    return pump(() -> value(s, path; kw...), s)
end

IOStrings.recoverable(e::JSON.ParseError) = e.c == IOStrings.ASCII_ETB

Base.string(j::Union{JSON.String{T},
                     JSON.Number{T},
                     JSON.Object{W, T},
                     JSON.Array{W, T}}) where W where T <: IOString =
    pump(() -> SubString(j.s, j.i, lastindex_of_value(j.s, j.i)), j.s)

"""
Get a JSON value object for a value in a JSON text.
 - `s`, the JSON text.
 - `i`, byte index of the value in JSON text.
 - `c`, first byte of the value.
"""
function getvalue(W, s, i, c=getc(s, i))
        if c == '{'                     o = JSON.Object(W, s, i)
                                        W === Nothing ? o : W(o)
    elseif c == '['                     JSON.Array(W, s, i)
    elseif c == '"'                     JSON.String(s, i)
    elseif isnum(c)                     JSON.Number(s, i)
    elseif c == 'f'                     false
    elseif c == 'n'                     nothing
    elseif c == 't'                     true
    else
        throw(JSON.ParseError(s, i, c, "invalid value index"))
    end
end

"""
See `LazyJSON.value(jsontext) -> JSON.Value`.
"""
parse(a...; kw...) = value(a...; kw...)

# Promotion to Base Container types

"""
Promote a JSON.Vaue to an equivalent Base type.
"""
promotejson(x; kw...) = x
promotejson(n::JSON.Number) = convert(Base.Number, n)
promotejson(s::JSON.String) = convert(AbstractString, s)
promotejson(a::JSON.Array)  = convert(Vector, a)
promotejson(o::JSON.Object) = convert(OrderedDict, o)

# Flattening to Base Container types

"""
Get a flattened (non-lazy) Julia object for a value in a JSON text;
and the index of the last character of the value.
"""
function getflat(s, i, c = getc(s, i))
        if c == '{'                     flat_object(s, i)
    elseif c == '['                     flat_array(s, i)
    elseif c == '"'                     parse_string(s, i)
    elseif isnum(c)                     parse_number(s, i)
    elseif c == 'f'                     false, next_i(s, i, 4)
    elseif c == 'n'                     nothing, next_i(s, i, 3)
    elseif c == 't'                     true, next_i(s, i, 3)
    else
        throw(JSON.ParseError(s, i, c, "invalid value index"))
    end
end

flatten(v) = v
flatten(v::PropertyDict) = flatten(PropertyDicts.unwrap(v))
flatten(v::Value) = getflat(v.s, v.i)[1]

function flat_object(s, i)
    o = OrderedDict{SubString{Base.String},Any}()
    i, c = skip_noise(s, i)
    while c != '}'
        k, i = parse_string(s, i)
        i, c = skip_noise(s, i)
        v, i = getflat(s, i, c)
        o[k] = v
        i, c = skip_noise(s, i)
    end
    return o, i
end

function flat_array(s, i)
    a = Any[]
    i, c = skip_noise(s, i)
    while c != ']'
        v, i = getflat(s, i, c)
        push!(a, v)
        i, c = skip_noise(s, i)
    end
    return a, i
end

# JSON Object/Array Iterator

"""
Iterate over the byte-indexes of the values in a JSON Array,
or the byte-indexes of the alternating field names and values in a JSON Object.
"""
struct Indexes{T}
    j::T
end

Base.IteratorSize(::Type{Indexes{T}}) where T = Base.SizeUnknown()
Base.IteratorEltype(::Type{Indexes{T}}) where T = Base.EltypeUnknown()

indexes(j::T) where T <: JSON.Collection = Indexes{T}(j)

function Base.iterate(j::Indexes, (i, c)=(j.j.i, 0x00))
    if c == ']' || c == '}'
        return nothing
    end
    i, c = nextindex(j.j, i, c)
    return (i, c), (i, c)
end

function nextindex(j, i, c)
    if i > j.i
        i = lastindex_of_value(j.s, i, c)
    end
    i, c = skip_noise(j.s, i)
    if c == IOStrings.ASCII_ETB
        throw(JSON.ParseError(j.s, i, c, "input incomplete"))
    end
    return i, c
end

"""
Count the number of values in an Array or Object.
For an Object, the count includes the field names.
"""
function collection_length(c::JSON.Collection)
    count = -1
    for i in indexes(c)
        count += 1
    end
    return count
end

# JSON Array Lookup

"""
Get the index `i` and first byte `c` of the `n`th value in an Array.
"""
function getindex_ic(a::JSON.Array, n::Int)
    for (i, c) in indexes(a)            # Iterate over byte-indexs of array
        if c == ']'                     # values until the end byte ']'.
           throw(BoundsError(a, n))
        end
        if n == 1
            return i, c                 # Return the byte-index of the nth value
        end
        n -= 1
    end
end

# JSON Object Field Name Lookup

memcmp(a, b, l) = ccall(:memcmp, Int32, (Ptr{UInt8}, Ptr{UInt8}, UInt), a, b, l)

"""
Get the index `i` and first byte `c` of the field value for `key` in an Object.
"""
function get_ic(o::JSON.Object, key::AbstractString, default, start::Int=o.i)

    keyl = sizeof(key)                            # Extract length of key and
    key1 = keyl == 0 ? 0x00 : key1 = getc(key, 1) # first byte of key.

    s = o.s                                       # Skip from "begin" byte '{'
    i, c = skip_noise(s, start)                   # to first field name byte.

    while c != '}'
        last_i, has_escape = scan_string(s, i)    # Find end of field name and
        if keyl == 0                              # compare to `key.
            foundkey = last_i == next_i(s, i)
        elseif has_escape || !applicable(pointer, o.s)
            foundkey = key == JSON.String(s, i)
        elseif key1 != next_c(s, i) ||            # {"key": ...}
               last_i != next_i(s, i, 1 + keyl)   #  ^   ^
            foundkey = false                      #  i   last_i
        else
            foundkey = memcmp(pointer(s, i+1),
                              pointer(key), keyl) == 0
        end

        i, c = skip_noise(s, last_i)              # Skip ':' and whitespace.

        if foundkey                               # If the key matched, return
            return i, c                           # index of first value byte.
        end

        i = lastindex_of_value(s, i, c)           # Skip over the value and the
        i, c = skip_noise(s, i)                   # ',' and whitespace.
    end

    if start != o.i
        return get_ic(o, key, default, o.i)
    end

    return default
end

# JSON Key Path Search

"""
Find the value at a specified key `path` in a JSON text.
 - `s`, a JSON text
 - `path`, A vector of Ints and Strings representing JSON Array indexes and
           JSON Object field names.
 - `i`, start index
 - `c`, byte at index `i`
"""
function getpath(s, path, i::Int=1, c::UInt8=getc(s, i))

    i, c = skip_whitespace(s, i, c)

    for key in path
        if c == '[' && key isa Integer
            i, c = getindex_ic(JSON.Array(Nothing, s, i), key)
        elseif c == '{' && key isa AbstractString
            i, c = get_ic(JSON.Object(Nothing, s, i), key, (0, 0x00))
            if i == 0
                throw(KeyError(key))
            end
        else
            throw(KeyError(key))
        end
    end
    return i, c
end

# JSON Token Scanner

"""
Find last index of a value or collection begin/end token end in a JSON text.
 - `s`, a JSON text
 - `i`, token start index
 - `c`, byte at index `i`
"""
function lastindex_of_token(s, i, c)::Int

        if c == '{' ||
           c == '}' ||
           c == '[' ||
           c == ']' ||
           c == 0x00                    i
    elseif c == '"'                     lastindex_of_string(s, i)
    elseif isnum(c)                     lastindex_of_number(s, i)
    elseif c == 'f'                     next_i(s, i, 4)
    elseif c == 'n'                     next_i(s, i, 3)
    elseif c == 't'                     next_i(s, i, 3)
    else
        throw(JSON.ParseError(s, i, c, "invalid input"))
    end
end

"""
Find index of last byte of a JSON value.
 - `s`, a JSON text
 - `i`, value start index
 - `c`, byte at index `i`
"""
function lastindex_of_value(s, i, c=getc(s, i))::Int
    if c == '{' || c == '['
        lastindex_of_collection(s, i, c)
    else
        lastindex_of_token(s, i, c)
    end
end

"""
Find index of last byte of a JSON Collection.
"""
function lastindex_of_collection(s, i, c)::Int

    nest = 1

    while nest > 0
        i, c = skip_noise(s, i)
        if isbegin(c)
            nest += 1
        elseif isend(c)
            nest -= 1
        end
        i = lastindex_of_token(s, i, c)
    end

    return i
end

"""
Find index of last byte of a JSON String.
"""
function lastindex_of_string(s, i)::Int
    i, has_escape = scan_string(s, i)
    return i
end

"""
Find index of last byte of a JSON String
and check for existence of the escape character.

https://tools.ietf.org/html/rfc7159#section-7

    string = quotation-mark *char quotation-mark
    char = unescaped / escape ( ... )
"""
function scan_string(s, i)

    i, c = next_ic(s, i)

    has_escape = false
    while c != '"'
        if isnull(c) || c == IOStrings.ASCII_ETB
            throw(JSON.ParseError(s, i, c, "input incomplete"))
        end
        escape = c == '\\'
        i, c = next_ic(s, i)
        if escape && !(isnull(c) || c == IOStrings.ASCII_ETB)
            has_escape = true
            i, c = next_ic(s, i)
        end
    end

    return i, has_escape
end

"""
Find index of last byte of a JSON Number.
"""
function lastindex_of_number(s, i)::Int

    last = i
    i, c = next_ic(s, i)

    while !isnull(c) && !isnoise(c) && !isend(c)
        if c == IOStrings.ASCII_ETB
            throw(JSON.ParseError(s, i, c, "input incomplete"))
        end
        last = i
        i, c = next_ic(s, i)
    end

    return last
end

"""
Skip over whitespace in String `s` starting at index `i`.
"""
function skip_whitespace(s, i, c = getc(s, i))
    while iswhitespace(c)
        i, c = next_ic(s, i)
    end
    return i, c
end

"""
Skip over non-value characters in String `s` starting at index `i` + 1.
"""
function skip_noise(s, i)
    i, c = next_ic(s, i)
    while isnoise(c)
        i, c = next_ic(s, i)
    end
    return i, c
end

# JSON Character Classes

"""
https://tools.ietf.org/html/rfc7159#section-2

    These are the six structural characters:
        begin-array     = ws %x5B ws  ; [ left square bracket
        begin-object    = ws %x7B ws  ; { left curly bracket
        end-array       = ws %x5D ws  ; ] right square bracket
        end-object      = ws %x7D ws  ; } right curly bracket
        name-separator  = ws %x3A ws  ; : colon
        value-separator = ws %x2C ws  ; , comma
"""
isbegin(c) = c == '{' ||
             c == '['

isend(c) = c == '}' ||
           c == ']'

"""
First character of a Start of a JSON Number
https://tools.ietf.org/html/rfc7159#section-6
"""
isnum(c) = c == '-' ||
           c in UInt8('0'):UInt8('9')

"""
https://tools.ietf.org/html/rfc7159#section-2

    Insignificant whitespace is allowed before or after any of the six
    structural characters.

    ws = *(
            %x20 /              ; Space
            %x09 /              ; Horizontal tab
            %x0A /              ; Line feed or New line
            %x0D )              ; Carriage return
"""
iswhitespace(c) = c == ' '  ||
                  c == '\t' ||
                  c == '\r' ||
                  c == '\n'

"""
The tokeniser ignores whitespace and separators between values.
"""
isnoise(c) = c == ',' ||
             c == ':' ||
             iswhitespace(c)

# Byte Wrangling Utilities

"""
`jl_alloc_string` allocates `n + 1` bytes and sets the last byte to `0x00`
https://github.com/JuliaLang/julia/blob/master/src/array.c#L464
"""
isnull(c) = c == 0x00

"""
Get byte `i` in string `s` without bounds checking.
"""
@generated function getc(s, i)
    if hasmethod(pointer, (s,))
        :(unsafe_load(pointer(s), i))
    else
        :(i > ncodeunits(s) ? 0x00 : @inbounds codeunit(s, i))
    end
end

#getc(s, i) = @inbounds codeunit(s, i)
# FIXME this causes Pkg.test("LazyJSON") to fail becuase the --check-bounds
#       option is passed to the test process.

"""
Increment `i` and get byte `i` in string `s` without bounds checking.
"""
next_ic(s, i) = (i = next_i(s, i); (i, getc(s, i)))
next_c(s, i) = getc(s, next_i(s, i))
next_i(s, i) = i + 1
next_i(s, i, n) = i + n

next_ic(s::SplicedString, i) = SplicedStrings.nextcodeunit(s, i)
next_i(s::SplicedString, i) = SplicedStrings.nextcodeunitindex(s, i)

const Unsafe = Union{SplicedString}
#next_i(s::Unsafe, i) = @inbounds nextind(s, i)
next_i(s::Unsafe, i, n) = @inbounds nextind(s, i, n) # FIXME byte index?

# Allow comparison of UInt8 with Char (e.g. c == '{')
==(a, b) = Base.isequal(a, b)
!=(a, b) = !(a == b)
==(i::T, c::Char) where T <: Integer = Base.isequal(i, T(c))

# Debug Display

function Base.show(io::IO, e::JSON.ParseError)

    s = e.bytes
    l = findprev(isequal('\n'), s, e.index)
    l = l != nothing ? l + 1 : 1
    r = findnext(isequal('\n'), s, max(1, e.index-1))
    r = r != nothing ? r - 1 : length(s)
    l = min(length(s), l)
    line_number = length(split(SubString(s, 1, l), '\n'))
    col_number = min(length(s), e.index) - l + 1
    c = escape_string(string(Char(e.c)))
    print(io, "JSON.ParseError: ", e.message,
              " at line ", line_number, ", col ", col_number, ": '", c, "'\n",
              SubString(s, max(1, l), r), "\n",
              lpad("", col_number - 1, " "), "^")
end

# Interface Protocols

include("AbstractString.jl")
include("Number.jl")
include("AbstractArray.jl")
include("AbstractDict.jl")

end # module


text:
#=
see description in examples/polyenvelope/native.jl
=#

struct PolyEnvelopeJuMP{T <: Real} <: ExampleInstanceJuMP{T}
    n::Int
    rand_halfdeg::Int
    num_polys::Int
    env_halfdeg::Int
end

function build(inst::PolyEnvelopeJuMP{T}) where {T <: Float64}
    n = inst.n
    @assert inst.rand_halfdeg <= inst.env_halfdeg
    domain = PolyUtils.BoxDomain{T}(-ones(T, n), ones(T, n))

    # generate interpolation
    (U, pts, Ps, _, w) = PolyUtils.interpolate(domain, inst.env_halfdeg,
        get_quadr = true)

    # generate random polynomials
    L = binomial(n + inst.rand_halfdeg, n)
    polys = Ps[1][:, 1:L] * rand(-9:9, L, inst.num_polys)

    model = JuMP.Model()
    JuMP.@variable(model, fpv[j in 1:U]) # values at Fekete points
    JuMP.@objective(model, Max, dot(fpv, w)) # integral over domain (quadrature)
    wsosT = Hypatia.WSOSInterpNonnegativeCone{T, T}
    JuMP.@constraint(model, [i in 1:inst.num_polys],
        polys[:, i] .- fpv in wsosT(U, Ps))

    return model
end


text:
using Navigation
using Test

@testset "Navigation.jl" begin
    p1 = Point(0.5, 1.0)
    p2 = Point(1.0, 0.5)
    @test distance(p1, p2) ≈ 3888e3 atol = 1e3
    @test bearing(p1, p2) ≈ 5.81412807071 atol = 0.001
    @test final_bearing(p1, p2) ≈ 5.45864813531 atol = 0.001
    @test midpoint(p1, p2).λ ≈ 0.8106650896663075 atol = 0.0001
    @test midpoint(p1, p2).ϕ ≈ 0.76484691145 atol = 0.0001
    @test intermediate_point(p1, p2, 0.5).λ ≈ 0.8106650896663075 atol = 0.0001
    @test intermediate_point(p1, p2, 0.5).ϕ ≈ 0.76484691145 atol = 0.0001

    dst = 3888000.0
    brg = 5.81412807071
    @test destination_point(p1, dst, brg).ϕ ≈ 1.0 atol = 0.01
    @test destination_point(p1, dst, brg).λ ≈ 0.5 atol = 0.01

    brg13 = 0.5
    brg23 = 1.5
    @test intersection_point(p1, p2, brg13, brg23).ϕ ≈ 0.86121333 atol = 0.001
    @test intersection_point(p1, p2, brg13, brg23).λ ≈ 1.31431535 atol = 0.001

    p1b = deg2rad(Point(30.0, 100.0))
    p2b = deg2rad(Point(50.0, 210.0))
    p3b = deg2rad(Point(40.0, 180.0))
    @test cross_track_distance(p1b, p2b, p3b, 1.0) ≈ 0.297182506587 atol = 0.00001
    @test along_track_distance(p1b, p2b, p3b, 1.0) ≈ 1.09661554384 atol = 0.00001

    @test Vground(110.0, 40.0, deg2rad(315.0), deg2rad(245.0)) ≈ 89.7 atol = 0.1
    @test Vground(120.0, 40.0, deg2rad(315.0), deg2rad(245.0)) ≈ 100.3 atol = 0.1

    @test head_wind(20.0, deg2rad(60.0), deg2rad(30.0)) ≈ 17.32 atol = 0.01
    @test cross_wind(20.0, deg2rad(60.0), deg2rad(30.0)) ≈ 10.00 atol = 0.01

    @test Navigation.normalize(-270.0, -180.0, 180.0) ≈ 90.0 atol = 0.1
    @test Navigation.normalize(181.0, -180.0, 180.0) ≈ -179.0 atol = 0.1
    @test rad2deg(p1).ϕ ≈ 28.6479 atol = 0.0001
    @test rad2deg(p2).λ ≈ 28.6479 atol = 0.0001
    @test deg2rad(p1).ϕ ≈ 0.0087 atol = 0.0001
    @test deg2rad(p2).λ ≈ 0.0087 atol = 0.0001
end


text:
# ---
# title: Multiple color scales
# cover: assets/multiple_color_scales.png
# description: Categorical and continuous color scales in the same plot
# author: "[Pietro Vertechi](https://github.com/piever)"
# ---

using AlgebraOfGraphics, CairoMakie
set_aog_theme!() #src

# Normally, a unique scale is associated to each given attribute. Color is an important
# exception: continuous and discrete color scales can coexist in the same plot.
# This should be used sparingly, as it can make the plot harder to interpret.

x = range(-π, π, length=100)
y = sin.(x)
ŷ = y .+ randn.() .* 0.1
z = cos.(x)
c = rand(["a", "b"], 100)
df = (; x, y, ŷ, z, c)
layers = mapping(:y, color=:z) * visual(Lines) + mapping(:ŷ => "y", color=:c)
plt = data(df) * mapping(:x) * layers
fg = draw(plt)

# save cover image #src
mkpath("assets") #src
save("assets/multiple_color_scales.png", fg) #src


text:
## activate project environment
using Pkg
Pkg.activate(@__DIR__)
Pkg.instantiate()

## load packages
using DelimitedFiles, LinearAlgebra, Statistics
using BenchmarkTools
using MuladdMacro
using StaticArrays
using StrideArrays
using Trixi

import Trixi: flux_ranocha, ln_mean, inv_ln_mean

struct PrimitiveVariablesCompressibleEuler2D{T<:Real} <: FieldVector{4, T}
  rho::T
  v1 ::T
  v2 ::T
  p  ::T
end

@inline function Trixi.cons2prim(u::PrimitiveVariablesCompressibleEuler2D,
                                 equations::CompressibleEulerEquations2D)
  return u
end

struct PrimitiveVariablesCompressibleEuler3D{T<:Real} <: FieldVector{5, T}
  rho::T
  v1 ::T
  v2 ::T
  v3 ::T
  p  ::T
end

@inline function Trixi.cons2prim(u::PrimitiveVariablesCompressibleEuler3D,
                                 equations::CompressibleEulerEquations3D)
  return u
end

@inline function Trixi.get_node_vars(u::AbstractArray{<:StaticVector},
                                     equations, dg::DG, indices...)
  return u[indices...]
end

function cons2prim!(prim::AbstractVector{<:StaticVector},
                    u::AbstractArray{<:Real, 2},
                    equations, solver)
  for i in eachindex(prim)
    u_local = Trixi.get_node_vars(u, equations, solver, i)
    prim_local = cons2prim(u_local, equations)
    prim[i] = prim_local
  end
end

function cons2prim!(_prim::AbstractVector{<:PrimitiveVariablesCompressibleEuler2D},
                    u::AbstractArray{<:Real, 2},
                    equations::CompressibleEulerEquations2D, solver)
  prim = reinterpret(reshape, eltype(eltype(_prim)), _prim)

  Trixi.@turbo for i in Trixi.indices((prim, u), 2)
    rho    = u[1, i]
    rho_v1 = u[2, i]
    rho_v2 = u[3, i]
    rho_e  = u[4, i]

    v1 = rho_v1 / rho
    v2 = rho_v2 / rho
    p = (equations.gamma - 1) * (rho_e - 0.5 * (rho_v1 * v1 + rho_v2 * v2))

    prim[1, i] = rho
    prim[2, i] = v1
    prim[3, i] = v2
    prim[4, i] = p
  end
end

function cons2prim!(_prim::AbstractVector{<:PrimitiveVariablesCompressibleEuler3D},
                    u::AbstractArray{<:Real, 2},
                    equations::CompressibleEulerEquations3D, solver)
  prim = reinterpret(reshape, eltype(eltype(_prim)), _prim)

  Trixi.@turbo for i in Trixi.indices((prim, u), 2)
    rho    = u[1, i]
    rho_v1 = u[2, i]
    rho_v2 = u[3, i]
    rho_v3 = u[4, i]
    rho_e  = u[5, i]

    v1 = rho_v1 / rho
    v2 = rho_v2 / rho
    v3 = rho_v3 / rho
    p = (equations.gamma - 1) * (rho_e - 0.5 * (rho_v1 * v1 + rho_v2 * v2 + rho_v3 * v3))

    prim[1, i] = rho
    prim[2, i] = v1
    prim[3, i] = v2
    prim[4, i] = v3
    prim[5, i] = p
  end
end

struct PrecomputedVariablesCompressibleEulerEquations2D{T<:Real} <: FieldVector{6, T}
    rho::T
    v1::T
    v2::T
    p::T
    log_rho::T
    log_p::T
end

struct PrecomputedVariablesCompressibleEulerEquations3D{T<:Real} <: FieldVector{7, T}
    rho::T
    v1::T
    v2::T
    v3::T
    p::T
    log_rho::T
    log_p::T
end

@muladd begin

"""
    ln_mean(x, y)

Compute the logarithmic mean

    ln_mean(x, y) = (y - x) / (log(y) - log(x)) = (y - x) / log(y / x)

Problem: The formula above has a removable singularity at `x == y`. Thus,
some care must be taken to implement it correctly without problems or loss
of accuracy when `x ≈ y`. Here, we use the approach proposed by
Ismail and Roe (2009).
Set ξ = y / x. Then, we have

    (y - x) / log(y / x) = (x + y) / log(ξ) * (ξ - 1) / (ξ + 1)

Set f = (ξ - 1) / (ξ + 1) = (y - x) / (x + y). Then, we use the expansion

    log(ξ) = 2 * f * (1 + f^2 / 3 + f^4 / 5 + f^6 / 7) + O(ξ^9)

Inserting the first few terms of this expansion yields

    (y - x) / log(ξ) ≈ (x + y) * f / (2 * f * (1 + f^2 / 3 + f^4 / 5 + f^6 / 7))
                        = (x + y) / (2 + 2/3 * f^2 + 2/5 * f^4 + 2/7 * f^6)

Since divisions are usually more expensive on modern hardware than
multiplications (Agner Fog), we try to avoid computing two divisions. Thus,
we use

    f^2 = (y - x)^2 / (x + y)^2
        = (x * (x - 2 * y) + y * y) / (x * (x + 2 * y) + y * y)

Given ε = 1.0e-4, we use the following algorithm.

    if f^2 < ε
        # use the expansion above
    else
        # use the direct formula (y - x) / log(y / x)
    end

# References
- Ismail, Roe (2009).
    Affordable, entropy-consistent Euler flux functions II: Entropy production at shocks.
    [DOI: 10.1016/j.jcp.2009.04.021](https://doi.org/10.1016/j.jcp.2009.04.021)
- Agner Fog.
    Lists of instruction latencies, throughputs and micro-operation breakdowns
    for Intel, AMD, and VIA CPUs.
    https://www.agner.org/optimize/instruction_tables.pdf
"""
@inline function Trixi.ln_mean(x, y, log_x, log_y)
    epsilon_f2 = 1.0e-4
    f2 = (x * (x - 2 * y) + y * y) / (x * (x + 2 * y) + y * y) # f2 = f^2
    if f2 < epsilon_f2
        return (x + y) / @evalpoly(f2, 2, 2/3, 2/5, 2/7)
    else
        return (y - x) / (log_y - log_x) # log(y / x)
    end
end

"""
    inv_ln_mean(x, y)

Compute the inverse `1 / ln_mean(x, y)` of the logarithmic mean
[`ln_mean`](@ref).

This function may be used to increase performance where the inverse of the
logarithmic mean is needed, by replacing a (slow) division by a (fast)
multiplication.
"""
@inline function Trixi.inv_ln_mean(x, y, log_x, log_y)
    epsilon_f2 = 1.0e-4
    f2 = (x * (x - 2 * y) + y * y) / (x * (x + 2 * y) + y * y) # f2 = f^2
    if f2 < epsilon_f2
        return @evalpoly(f2, 2, 2/3, 2/5, 2/7) / (x + y)
    else
        return (log_y - log_x) / (y - x)
    end
end

"""
    flux_ranocha(u_ll, u_rr, orientation_or_normal_direction,
                 equations::CompressibleEulerEquations2D)

Entropy conserving and kinetic energy preserving two-point flux by
- Hendrik Ranocha (2018)
  Generalised Summation-by-Parts Operators and Entropy Stability of Numerical Methods
  for Hyperbolic Balance Laws
  [PhD thesis, TU Braunschweig](https://cuvillier.de/en/shop/publications/7743)
See also
- Hendrik Ranocha (2020)
  Entropy Conserving and Kinetic Energy Preserving Numerical Methods for
  the Euler Equations Using Summation-by-Parts Operators
  [Proceedings of ICOSAHOM 2018](https://doi.org/10.1007/978-3-030-39647-3_42)
"""
@inline function flux_ranocha(u_ll::PrecomputedVariablesCompressibleEulerEquations2D,
                              u_rr::PrecomputedVariablesCompressibleEulerEquations2D,
                              orientation::Integer, equations::CompressibleEulerEquations2D)
  # Unpack left and right state
  rho_ll, v1_ll, v2_ll, p_ll, log_rho_ll, log_p_ll = u_ll
  rho_rr, v1_rr, v2_rr, p_rr, log_rho_rr, log_p_rr = u_rr

  # Compute the necessary mean values
  rho_mean = ln_mean(rho_ll, rho_rr, log_rho_ll, log_rho_rr)
  # Algebraically equivalent to `inv_ln_mean(rho_ll / p_ll, rho_rr / p_rr)`
  # in exact arithmetic since
  #     log((ϱₗ/pₗ) / (ϱᵣ/pᵣ)) / (ϱₗ/pₗ - ϱᵣ/pᵣ)
  #   = pₗ pᵣ log((ϱₗ pᵣ) / (ϱᵣ pₗ)) / (ϱₗ pᵣ - ϱᵣ pₗ)
  inv_rho_p_mean = p_ll * p_rr * inv_ln_mean(rho_ll * p_rr, rho_rr * p_ll, log_rho_ll + log_p_rr, log_rho_rr + log_p_ll)
  v1_avg = 0.5 * (v1_ll + v1_rr)
  v2_avg = 0.5 * (v2_ll + v2_rr)
  p_avg  = 0.5 * (p_ll + p_rr)
  velocity_square_avg = 0.5 * (v1_ll*v1_rr + v2_ll*v2_rr)

  # Calculate fluxes depending on orientation
  if orientation == 1
    f1 = rho_mean * v1_avg
    f2 = f1 * v1_avg + p_avg
    f3 = f1 * v2_avg
    f4 = f1 * ( velocity_square_avg + inv_rho_p_mean * equations.inv_gamma_minus_one ) + 0.5 * (p_ll*v1_rr + p_rr*v1_ll)
  else
    f1 = rho_mean * v2_avg
    f2 = f1 * v1_avg
    f3 = f1 * v2_avg + p_avg
    f4 = f1 * ( velocity_square_avg + inv_rho_p_mean * equations.inv_gamma_minus_one ) + 0.5 * (p_ll*v2_rr + p_rr*v2_ll)
  end

  return SVector(f1, f2, f3, f4)
end

@inline function flux_ranocha(u_ll::PrecomputedVariablesCompressibleEulerEquations3D,
                              u_rr::PrecomputedVariablesCompressibleEulerEquations3D,
                              orientation::Integer, equations::CompressibleEulerEquations3D)
  # Unpack left and right state
  rho_ll, v1_ll, v2_ll, v3_ll, p_ll, log_rho_ll, log_p_ll = u_ll
  rho_rr, v1_rr, v2_rr, v3_rr, p_rr, log_rho_rr, log_p_rr = u_rr

  # Compute the necessary mean values
  rho_mean = ln_mean(rho_ll, rho_rr, log_rho_ll, log_rho_rr)
  # Algebraically equivalent to `inv_ln_mean(rho_ll / p_ll, rho_rr / p_rr)`
  # in exact arithmetic since
  #     log((ϱₗ/pₗ) / (ϱᵣ/pᵣ)) / (ϱₗ/pₗ - ϱᵣ/pᵣ)
  #   = pₗ pᵣ log((ϱₗ pᵣ) / (ϱᵣ pₗ)) / (ϱₗ pᵣ - ϱᵣ pₗ)
  inv_rho_p_mean = p_ll * p_rr * inv_ln_mean(rho_ll * p_rr, rho_rr * p_ll, log_rho_ll + log_p_rr, log_rho_rr + log_p_ll)
  v1_avg = 0.5 * (v1_ll + v1_rr)
  v2_avg = 0.5 * (v2_ll + v2_rr)
  v3_avg = 0.5 * (v3_ll + v3_rr)
  p_avg  = 0.5 * (p_ll + p_rr)
  velocity_square_avg = 0.5 * (v1_ll*v1_rr + v2_ll*v2_rr + v3_ll*v3_rr)

  # Calculate fluxes depending on orientation
  if orientation == 1
    f1 = rho_mean * v1_avg
    f2 = f1 * v1_avg + p_avg
    f3 = f1 * v2_avg
    f4 = f1 * v3_avg
    f5 = f1 * ( velocity_square_avg + inv_rho_p_mean * equations.inv_gamma_minus_one ) + 0.5 * (p_ll*v1_rr + p_rr*v1_ll)
  elseif orientation == 2
    f1 = rho_mean * v2_avg
    f2 = f1 * v1_avg
    f3 = f1 * v2_avg + p_avg
    f4 = f1 * v3_avg
    f5 = f1 * ( velocity_square_avg + inv_rho_p_mean * equations.inv_gamma_minus_one ) + 0.5 * (p_ll*v2_rr + p_rr*v2_ll)
  else # orientation == 3
    f1 = rho_mean * v3_avg
    f2 = f1 * v1_avg
    f3 = f1 * v2_avg
    f4 = f1 * v3_avg + p_avg
    f5 = f1 * ( velocity_square_avg + inv_rho_p_mean * equations.inv_gamma_minus_one ) + 0.5 * (p_ll*v3_rr + p_rr*v3_ll)
  end

  return SVector(f1, f2, f3, f4, f5)
end

end # muladd

@inline function Trixi.get_node_vars(u::AbstractArray{<:StaticVector},
                                     equations, dg::DG, indices...)
  return u[indices...]
end

function cons2prim!(_prim::AbstractVector{<:PrecomputedVariablesCompressibleEulerEquations2D},
                    u::AbstractArray{<:Real, 2},
                    equations::CompressibleEulerEquations2D, solver)
  prim = reinterpret(reshape, eltype(eltype(_prim)), _prim)

  Trixi.@turbo for i in Trixi.indices((prim, u), 2)
    rho    = u[1, i]
    rho_v1 = u[2, i]
    rho_v2 = u[3, i]
    rho_e  = u[4, i]

    v1 = rho_v1 / rho
    v2 = rho_v2 / rho
    p = (equations.gamma - 1) * (rho_e - 0.5 * (rho_v1 * v1 + rho_v2 * v2))

    prim[1, i] = rho
    prim[2, i] = v1
    prim[3, i] = v2
    prim[4, i] = p
    prim[5, i] = log(rho)
    prim[6, i] = log(p)
  end
end

function cons2prim!(_prim::AbstractVector{<:PrecomputedVariablesCompressibleEulerEquations3D},
                    u::AbstractArray{<:Real, 2},
                    equations::CompressibleEulerEquations3D, solver)
  prim = reinterpret(reshape, eltype(eltype(_prim)), _prim)

  Trixi.@turbo for i in Trixi.indices((prim, u), 2)
    rho    = u[1, i]
    rho_v1 = u[2, i]
    rho_v2 = u[3, i]
    rho_v3 = u[4, i]
    rho_e  = u[5, i]

    v1 = rho_v1 / rho
    v2 = rho_v2 / rho
    v3 = rho_v3 / rho
    p = (equations.gamma - 1) * (rho_e - 0.5 * (rho_v1 * v1 + rho_v2 * v2 + rho_v3 * v3))

    prim[1, i] = rho
    prim[2, i] = v1
    prim[3, i] = v2
    prim[4, i] = v3
    prim[5, i] = p
    prim[6, i] = log(rho)
    prim[7, i] = log(p)
  end
end

## set up benchmark code
"""
    initial_condition_isentropic_vortex(x, t, equations)

The classical isentropic vortex test case of
- Chi-Wang Shu (1997)
  Essentially Non-Oscillatory and Weighted Essentially Non-Oscillatory
  Schemes for Hyperbolic Conservation Laws.
  [NASA/CR-97-206253](https://ntrs.nasa.gov/citations/19980007543)
"""
function initial_condition_isentropic_vortex(x, t, equations::CompressibleEulerEquations2D)
  ϱ0 = 1.0               # background density
  v0 = SVector(1.0, 1.0) # background velocity
  p0 = 10.0              # background pressure
  ε = 20.0               # vortex strength
  L = 10.0               # size of the domain per coordinate direction

  T0 = p0 / ϱ0           # background temperature
  γ = equations.gamma    # ideal gas constant

  x0 = v0 * t            # current center of the vortex
  dx = vortex_center.(x - x0, L)
  r2 = sum(abs2, dx)

  # perturbed primitive variables
  T = T0 - (γ - 1) * ε^2 / (8 * γ * π^2) * exp(1 - r2)
  v = v0 + ε / (2 * π) * exp(0.5 * (1 - r2)) * SVector(-dx[2], dx[1])
  ϱ = ϱ0 * (T / T0)^(1 / (γ - 1))
  p = ϱ * T

  return prim2cons(SVector(ϱ, v..., p), equations)
end

function initial_condition_isentropic_vortex(x, t, equations::CompressibleEulerEquations3D)
  ϱ0 = 1.0               # background density
  v0 = SVector(1.0, 1.0, 0.0) # background velocity
  p0 = 10.0              # background pressure
  ε = 20.0               # vortex strength
  L = 10.0               # size of the domain per coordinate direction

  T0 = p0 / ϱ0           # background temperature
  γ = equations.gamma    # ideal gas constant

  x0 = v0 * t            # current center of the vortex
  dx = vortex_center.(x - x0, L)
  r2 = sum(abs2, dx)

  # perturbed primitive variables
  T = T0 - (γ - 1) * ε^2 / (8 * γ * π^2) * exp(1 - r2)
  v = v0 + ε / (2 * π) * exp(0.5 * (1 - r2)) * SVector(-dx[2], dx[1], 0.0)
  ϱ = ϱ0 * (T / T0)^(1 / (γ - 1))
  p = ϱ * T

  return prim2cons(SVector(ϱ, v..., p), equations)
end

vortex_center(x, L) = mod(x + L/2, L) - L/2

function split_form_prim_kernel!(du, u_prim_logs, u_cons,
                                 element, mesh,
                                 nonconservative_terms, equations,
                                 volume_flux::VF, dg, cache) where {VF}

  cons2prim!(vec(u_prim_logs), reshape(u_cons, nvariables(equations), :),
             equations, dg)

  Trixi.split_form_kernel!(du, u_prim_logs, element, mesh,
                           nonconservative_terms, equations,
                           volume_flux, dg, cache)

  return nothing
end

# 2D version
function run_benchmarks_2d(polydeg, numerical_flux, initial_condition)
  println("Benchmarking ", numerical_flux, " in 2D with polydeg=", polydeg)

  # setup basic code
  equations = CompressibleEulerEquations2D(1.4)

  coordinates_min = (-5.0, -5.0)
  coordinates_max = ( 5.0,  5.0)
  mesh = TreeMesh(coordinates_min, coordinates_max,
                  initial_refinement_level=0, n_cells_max=100_000, periodicity=true)

  solver = DGSEM(polydeg=polydeg, surface_flux=numerical_flux,
                 volume_integral=VolumeIntegralFluxDifferencing(numerical_flux))

  semi = SemidiscretizationHyperbolic(mesh, equations,
                                      initial_condition, solver)

  u_ode  = compute_coefficients(0.0, semi)
  du_ode = similar(u_ode)
  GC.@preserve u_ode du_ode begin
    u  = Trixi.wrap_array(u_ode,  semi)
    du = Trixi.wrap_array(du_ode, semi)

    u_prim = StrideArray(undef,
                         PrimitiveVariablesCompressibleEuler2D{eltype(u)},
                         Base.tail(StrideArrays.ArrayInterface.size(u))...)

    u_prim_logs = StrideArray(undef,
                              PrecomputedVariablesCompressibleEulerEquations2D{eltype(u)},
                              Base.tail(StrideArrays.ArrayInterface.size(u))...)
    element = 1

    # benchmark flux differencing volume terms based on conservative variables
    fill!(du, zero(eltype(du)))
    Trixi.split_form_kernel!(
      du, u,
      element, mesh,
      (Trixi.have_nonconservative_terms(equations)), equations,
      numerical_flux, solver, (semi.cache))
    du_cons = copy(du)

    sleep(0.5)
    result_cons = @benchmark Trixi.split_form_kernel!(
      $du, $u,
      $element, $mesh,
      $(Trixi.have_nonconservative_terms(equations)), $equations,
      $numerical_flux, $solver, $(semi.cache)) setup=(fill!($du, zero(eltype($du))))
    display(result_cons)

    # benchmark flux differencing volume terms based on primitive variables
    fill!(du, zero(eltype(du)))
    split_form_prim_kernel!(
      du, u_prim, u,
      element, mesh,
      (Trixi.have_nonconservative_terms(equations)), equations,
      numerical_flux, solver, (semi.cache))
    du_prim = copy(du)
    @assert du_cons ≈ du_prim

    sleep(0.5)
    result_prim = @benchmark split_form_prim_kernel!(
      $du, $u_prim, $u,
      $element, $mesh,
      $(Trixi.have_nonconservative_terms(equations)), $equations,
      $numerical_flux, $solver, $(semi.cache)) setup=(fill!($du, zero(eltype($du))))
    display(result_prim)

    # benchmark both primitive variables and precomputed logarithms
    fill!(du, zero(eltype(du)))
    split_form_prim_kernel!(
      du, u_prim_logs, u,
      element, mesh,
      (Trixi.have_nonconservative_terms(equations)), equations,
      numerical_flux, solver, (semi.cache))
    du_prim_logs = copy(du)
    @assert du_cons ≈ du_prim_logs

    sleep(0.5)
    result_prim_logs = @benchmark split_form_prim_kernel!(
      $du, $u_prim_logs, $u,
      $element, $mesh,
      $(Trixi.have_nonconservative_terms(equations)), $equations,
      $numerical_flux, $solver, $(semi.cache)) setup=(fill!($du, zero(eltype($du))))
    display(result_prim_logs)
  end
  println()

  return result_cons, result_prim, result_prim_logs
end

# 3D version
function run_benchmarks_3d(polydeg, numerical_flux, initial_condition)
  println("Benchmarking ", numerical_flux, " in 3D with polydeg=", polydeg)

  # setup basic code
  equations = CompressibleEulerEquations3D(1.4)

  coordinates_min = (-5.0, -5.0, -5.0)
  coordinates_max = ( 5.0,  5.0,  5.0)
  mesh = TreeMesh(coordinates_min, coordinates_max,
                  initial_refinement_level=0, n_cells_max=100_000, periodicity=true)

  solver = DGSEM(polydeg=polydeg, surface_flux=numerical_flux,
                 volume_integral=VolumeIntegralFluxDifferencing(numerical_flux))

  semi = SemidiscretizationHyperbolic(mesh, equations,
                                      initial_condition, solver)

  u_ode  = compute_coefficients(0.0, semi)
  du_ode = similar(u_ode)
  GC.@preserve u_ode du_ode begin
    u  = Trixi.wrap_array(u_ode,  semi)
    du = Trixi.wrap_array(du_ode, semi)

    u_prim = StrideArray(undef,
                         PrimitiveVariablesCompressibleEuler3D{eltype(u)},
                         Base.tail(StrideArrays.ArrayInterface.size(u))...)

    u_prim_logs = StrideArray(undef,
                         PrecomputedVariablesCompressibleEulerEquations3D{eltype(u)},
                         Base.tail(StrideArrays.ArrayInterface.size(u))...)
    element = 1

    # benchmark flux differencing volume terms based on conservative variables
    fill!(du, zero(eltype(du)))
    Trixi.split_form_kernel!(
      du, u,
      element, mesh,
      (Trixi.have_nonconservative_terms(equations)), equations,
      numerical_flux, solver, (semi.cache))
    du_cons = copy(du)

    sleep(0.5)
    result_cons = @benchmark Trixi.split_form_kernel!(
      $du, $u,
      $element, $mesh,
      $(Trixi.have_nonconservative_terms(equations)), $equations,
      $numerical_flux, $solver, $(semi.cache)) setup=(fill!($du, zero(eltype($du))))
    display(result_cons)

    # benchmark flux differencing volume terms based on primitive variables
    fill!(du, zero(eltype(du)))
    split_form_prim_kernel!(
      du, u_prim, u,
      element, mesh,
      (Trixi.have_nonconservative_terms(equations)), equations,
      numerical_flux, solver, (semi.cache))
    du_prim = copy(du)
    @assert du_cons ≈ du_prim

    sleep(0.5)
    result_prim = @benchmark split_form_prim_kernel!(
      $du, $u_prim, $u,
      $element, $mesh,
      $(Trixi.have_nonconservative_terms(equations)), $equations,
      $numerical_flux, $solver, $(semi.cache)) setup=(fill!($du, zero(eltype($du))))
    display(result_prim)

    # flux differencing volume terms based on primitive variables and precomputed logarithms
    fill!(du, zero(eltype(du)))
    split_form_prim_kernel!(
      du, u_prim_logs, u,
      element, mesh,
      (Trixi.have_nonconservative_terms(equations)), equations,
      numerical_flux, solver, (semi.cache))
    du_prim_logs = copy(du)
    @assert du_cons ≈ du_prim_logs

    sleep(0.5)
    result_prim_logs = @benchmark split_form_prim_kernel!(
      $du, $u_prim_logs, $u,
      $element, $mesh,
      $(Trixi.have_nonconservative_terms(equations)), $equations,
      $numerical_flux, $solver, $(semi.cache)) setup=(fill!($du, zero(eltype($du))))
    display(result_prim_logs)
  end
  println()

  return result_cons, result_prim, result_prim_logs
end

# general code
Base.ndims(::typeof(run_benchmarks_2d)) = 2
Base.ndims(::typeof(run_benchmarks_3d)) = 3

function run_benchmarks(run_benchmarks_inner, polydegs, numerical_flux, initial_condition, icname="")
  results = zeros(length(polydegs), 1)
  results .= polydegs
  header = "# Polydeg"

  # flux differencing
  println("#"^80)
  println("flux differencing\n")
  means_cons = zeros(length(polydegs))
  stds_cons  = zero(means_cons)
  means_prim = zero(means_cons)
  stds_prim  = zero(means_cons)
  means_prim_logs = zero(means_cons)
  stds_prim_logs  = zero(means_cons)
  for (i, polydeg) in enumerate(polydegs)
    result_cons, result_prim, result_prim_logs = run_benchmarks_inner(polydeg, numerical_flux, initial_condition)
    means_cons[i] = time(mean(result_cons)) * 1.0e-9
    stds_cons[i]  = time(std(result_cons))  * 1.0e-9
    means_prim[i] = time(mean(result_prim)) * 1.0e-9
    stds_prim[i]  = time(std(result_prim))  * 1.0e-9
    means_prim_logs[i] = time(mean(result_prim_logs)) * 1.0e-9
    stds_prim_logs[i]  = time(std(result_prim_logs))  * 1.0e-9
  end
  results = hcat(results, means_cons, stds_cons, means_prim, stds_prim, means_prim_logs, stds_prim_logs)
  header = header * "\tConservative-variables-mean\tstd" *
                    "\tPrimitive-variables-mean\tstd" *
                    "\tPrecomputed-log-variables-mean\tstd"

  # write results to a file
  open(joinpath(@__DIR__, "precomputed_$(ndims(run_benchmarks_inner))D_$(numerical_flux)_$icname.dat"), "w") do io
    println(io, header)
    writedlm(io, results)
  end

  return results
end

## run benchmarks
polydegs = 3:15
run_benchmarks(run_benchmarks_2d, polydegs, flux_ranocha, initial_condition_isentropic_vortex, "vortex")
run_benchmarks(run_benchmarks_3d, polydegs, flux_ranocha, initial_condition_isentropic_vortex, "vortex")

initial_condition_random(x, t, equations::CompressibleEulerEquations2D) = prim2cons(SVector{4}(rand(), .1, .2, rand()), equations)
initial_condition_random(x, t, equations::CompressibleEulerEquations3D) = prim2cons(SVector{5}(rand(), .1, .2, .3, rand()), equations)

run_benchmarks(run_benchmarks_2d, polydegs, flux_ranocha, initial_condition_random, "random")
run_benchmarks(run_benchmarks_3d, polydegs, flux_ranocha, initial_condition_random, "random")

function initial_condition_sine(x, t, equations::CompressibleEulerEquations2D)
    @unpack gamma = equations
    rho = 2.0 + sin(pi * x[1] / 5.0) * sin(pi * x[2] / 5.0)
    return prim2cons(SVector{4}(rho, .1, .2, rho^gamma), equations)
end
function initial_condition_sine(x, t, equations::CompressibleEulerEquations3D)
    @unpack gamma = equations
    rho = 2.0 + sin(pi * x[1] / 5.0) * sin(pi * x[2] / 5.0) * sin(pi * x[3] / 5.0)
    return prim2cons(SVector{5}(rho, .1, .2, .3, rho^gamma), equations)
end

run_benchmarks(run_benchmarks_2d, polydegs, flux_ranocha, initial_condition_sine, "sine")
run_benchmarks(run_benchmarks_3d, polydegs, flux_ranocha, initial_condition_sine, "sine")

## compute relative performance
using DelimitedFiles
using Measurements

function compute_relative_performance(filename)
  header = "# Polydeg Precomputed logarithms/primitive-variables-mean std"
  data = readdlm(filename, comments=true)
  cons = data[:, 2] .± data[:, 3]
  prim = data[:, 4] .± data[:, 5]
  prim_logs = data[:, 6] .± data[:, 7]
  ratio = prim_logs ./ prim
  open(filename[1:end-4] * "_relative.dat", "w") do io
    println(io, header)
    writedlm(io, hcat(data[:, 1], # polydeg
                      Measurements.value.(ratio),
                      Measurements.uncertainty.(ratio)))
  end
end

compute_relative_performance(joinpath(@__DIR__, "precomputed_2D_flux_ranocha_vortex.dat"))
compute_relative_performance(joinpath(@__DIR__, "precomputed_3D_flux_ranocha_vortex.dat"))
compute_relative_performance(joinpath(@__DIR__, "precomputed_2D_flux_ranocha_random.dat"))
compute_relative_performance(joinpath(@__DIR__, "precomputed_3D_flux_ranocha_random.dat"))
compute_relative_performance(joinpath(@__DIR__, "precomputed_2D_flux_ranocha_sine.dat"))
compute_relative_performance(joinpath(@__DIR__, "precomputed_3D_flux_ranocha_sine.dat"))


text:
using MAT
using jInv.Mesh
using DivSigGrad
using jInv.LinearSolvers
using jInv.ForwardShare
using jInv.InverseSolve
using jInv.Utils
using ArgParse
using KrylovMethods

include("getDCResistivitySourcesAndRecAll.jl")

function parse_commandline()
	s = ArgParseSettings()

	@add_arg_table s begin

	"--n1"
		help = "number of cells in x1 direction "
		arg_type = Int
		default = 48
	"--n2"
		help = "number of cells in x2 direction "
		arg_type = Int
		default = 48
	"--n3"
		help = "number of cells in x3 direction "
		arg_type = Int
		default = 24
	"--srcSpacing"
		help = "spacing between sources"
		arg_type = Int
		default = 2
	"--nthreads"
		help = "number of threads"
		arg_type = Int
		default = 1
	"--out"
		help ="file to write into"
		default = "times.csv"
	"--solver"
		help = "PDE solver: 1->PCG, 2->BlockPCG, 3->MUMPS"
		arg_type = Int
		default = 1
	end
	return parse_args(s)
end

function main()
	parsed_args = parse_commandline()

	nInv      = [128;128;64]
	m         = 1500 + 3000*rand(tuple(nInv...))
	# random conductivity model can be replaced by SEG model described in
	#
	# Aminzadeh, F., Brac, J., and Kunz, T., 1997. 3D Salt and Overthrust models. SEG/EAGE Modeling Series, No. 1: Distribution CD of Salt and
	# Overthrust models, SEG Book Series Tulsa, Oklahoma
	#
	# matfile   = matread("3Dseg12812864.mat")
	# m         = matfile["VELc"]

	n1 = parsed_args["n1"]
	n2 = parsed_args["n2"]
	n3 = parsed_args["n3"]
	ns = parsed_args["srcSpacing"]
	nthreads = parsed_args["nthreads"]

	out = parsed_args["out"]

	# set number of threads for openblac
	set_num_threads(nthreads)

	solver = parsed_args["solver"]

	domain  = [0;13.5;0;13.5;0;4.2]
	n       = [n1;n2;n3]

	M    = getRegularMesh(domain,n)
	MInv = getRegularMesh(domain,nInv)
	Q,P  = getDCResistivitySourcesAndRecAll(M,srcSpacing=[ns,ns])
	Q    = Q[:,1:10]
	nsrc = size(Q,2)

	@printf "problem size: n1=%d\tn2=%d\tn3=%d\tns=%d\n" n1 n2 n3 nsrc
	@printf "nworkers()=%d\n" nworkers()
	@printf "nthreads=%d \n" nthreads

	if solver==3
		@printf "solver=MUMPS\n\n"
		@printf "warmup\ttotal\tfactorization\tsolve\n"
	elseif solver==2
		@printf "solver=BlockPCG\n\n"
		@printf "warmup\ttotal\titer\tprecond\tcgTime\tmatvec\n"
	elseif solver==1
		@printf "solver=PCG\n\n"
		@printf "warmup\ttotal\titer\tprecond\tcgTime\tmatvec\n"
	end
	Ainv = []
	if solver==1
		Ainv   = getIterativeSolver(KrylovMethods.cg,out=-1,sym=1,PC=:jac);
		Ainv.tol=1e-8
		Ainv.nthreads=nthreads
	elseif solver==2
 		Ainv         = getBlockIterativeSolver(KrylovMethods.blockCG,out=-1,sym=1,PC=:jac);
		Ainv.maxIter = 1000
		Ainv.tol     = 1e-8
		Ainv.nthreads=nthreads
		Ainv.out     =-1
	elseif solver==3
		Ainv = getMUMPSsolver()
	elseif solver 		== 4
		levels      	= 5;
		numCores 		= nthreads;
		maxIter     	= 20;
		relativeTol 	= 1e-8;
		relaxType   	= "SPAI";
		relaxParam  	= 1.0;
		relaxPre 		= 2;
		relaxPost   	= 2;
		cycleType   	='V';
		coarseSolveType = "MUMPS";
		MG 				= getMGparam(levels,numCores,maxIter,relativeTol,relaxType,relaxParam,relaxPre,relaxPost,cycleType,coarseSolveType);
		Ainvt   			= getSA_AMGsolver(MG, "PCG",sym=1,out=-1);
		Ainv   			= getSA_AMGsolver(MG, "PCG",sym=1,out=-1);
	end

	pForp = Array{RemoteRef{Channel{Any}}}(nworkers())
	pFors = Array{RemoteRef{Channel{Any}}}(nworkers())
	Mesh2Mesh  = getInterpolationMatrix(MInv,M)'
	m          = Mesh2Mesh'*vec(m)
	glocp      = Array{RemoteRef{Channel{Any}}}(nworkers())

	for j=1:nworkers()
		pForp[j] = @spawnat workers()[j] DivSigGradParam(M,Q,P,[1.0],Ainv)
		pFors[j] = @spawnat workers()[j] DivSigGradParam(M,Q[:,1],P,[1.0],Ainv)
		glocp[j] = @spawnat workers()[j] sparse(I,M.nc,M.nc)
	end

	# warm up
	wTime = @elapsed begin
	res = getData(vec(m),pFors,glocp)
	end
	@printf "%3.2f\t" wTime
	clear!(res[1])
	clear!(res[2])
	clear!(pFors)

	tTime = @elapsed begin
	getData(vec(m),pForp,glocp)
	end
	@printf "%3.2f\t" tTime

	if solver==3
		fTime = 0.0; sTime =0.0
		for p=1:length(pForp)
			pf = take!(pForp[p])
			fTime += pf.Ainv.facTime
			sTime += pf.Ainv.solveTime
		end
		fTime /= nworkers()
		sTime /= nworkers()
		@printf "%3.4f\t\t%3.4f\n" fTime sTime

		f = open(out,"a")
		str = @sprintf "%d,%d,%d,%d,%d,%d,,%3.5f,%3.5f,%3.5f,%3.5f\n" nworkers() solver n1 n2 n3 nsrc wTime tTime fTime sTime
		write(f,str )
		close(f)
	elseif solver==1 || solver==2
		nIter = 0; pcTime = 0.0; cgTime=0.0; mvTime=0.0
		for p=1:length(pForp)
			pf = take!(pForp[p])
			nIter += pf.Ainv.nIter
			pcTime += pf.Ainv.timePC
			cgTime += pf.Ainv.timeSolve
			mvTime += pf.Ainv.timeMV
			clear!(pf)
		end
		nIter /= nworkers()
		pcTime /=nworkers()
		cgTime /=nworkers()
		mvTime /=nworkers()
		@printf "%5d\t%3.4f\t%3.4f\t%3.4f\n" nIter  pcTime cgTime mvTime

		f = open(out,"a")
		str = @sprintf "%d,%d,%d,%d,%d,%d,%3.5f,%3.5f,%3.5f,%3.5f,%3.5f,%d\n" nworkers() solver n1 n2 n3 nsrc wTime tTime pcTime cgTime mvTime nIter
		write(f,str )
		close(f)
	elseif solver == 4
		nIter = 0; tsetup = 0.0; tSol = 0.0
		for p=1:length(pForp)
			pf = take!(pForp[p])
			nIter  += pf.Ainv.nIter
			tsetup += pf.Ainv.timeSetup
			tSol   += pf.Ainv.timeSolve
			clear!(pf)
		end
		@printf "%5d\t%3.4f\t%3.4f\t\n" nIter tsetup tSol
	end

end
main()


text:
module BonbonRectangle
export GameEnv, GameSpec, Board
include("game.jl")
  module Training
    using AlphaZero
    import ..GameSpec
    include("../../scripts/test_game.jl")
  end
end