text:
# MIT License
#
# Copyright (c) 2018 Martin Biel
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

SubWorker{S} = RemoteChannel{Channel{Vector{S}}}
NonSmoothSubWorker{T} = SubWorker{SubProblem{T}}
SmoothSubWorker{T} = SubWorker{SmoothSubProblem{T}}
ScenarioProblemChannel{S} = RemoteChannel{Channel{ScenarioProblems{S}}}

"""
    SynchronousExecution

Functor object for using synchronous execution in an quasi-gradient algorithm (assuming multiple Julia cores are available). Create by supplying a [`Synchronous`](@ref) object through `execution` in `QuasiGradient.Optimizer` or by setting the [`Execution`](@ref) attribute.

"""
struct SynchronousExecution{T <: AbstractFloat, S <: AbstractSubProblem{T}} <: AbstractQuasiGradientExecution
    subworkers::Vector{SubWorker{S}}
    decisions::Vector{DecisionChannel}

    function SynchronousExecution(structure::StageDecompositionStructure{2, 1, <:Tuple{DistributedScenarioProblems}}, x::AbstractVector, subproblems::Unaltered, ::Type{T}) where T <: AbstractFloat
        execution = new{T, SubProblem{T}}(Vector{NonSmoothSubWorker{T}}(undef, nworkers()),
                           scenarioproblems(structure).decisions)
        initialize_subproblems!(execution, scenarioproblems(structure, 2))
        return execution
    end

    function SynchronousExecution(structure::StageDecompositionStructure{2, 1, <:Tuple{DistributedScenarioProblems}}, x::AbstractVector, subproblems::Smoothed, ::Type{T}) where T <: AbstractFloat
        execution = new{T, SmoothSubProblem{T}}(Vector{SmoothSubWorker{T}}(undef, nworkers()),
                                                scenarioproblems(structure).decisions)
        initialize_subproblems!(execution, scenarioproblems(structure, 2), x, subproblems.parameters)
        return execution
    end
end

function initialize_subproblems!(execution::SynchronousExecution{T, SubProblem{T}},
                                 scenarioproblems::DistributedScenarioProblems) where T <: AbstractFloat
    # Create subproblems on worker processes
    @sync begin
        for w in workers()
            execution.subworkers[w-1] = RemoteChannel(() -> Channel{Vector{SubProblem{T}}}(1), w)
            prev = map(2:(w-1)) do p
                scenarioproblems.scenario_distribution[p-1]
            end
            start_id = isempty(prev) ? 0 : sum(prev)
            @async remotecall_fetch(
                w,
                execution.subworkers[w-1],
                scenarioproblems[w-1],
                start_id) do sw, sp_, start_id
                    sp = fetch(sp_)
                    subproblems = Vector{SubProblem{T}}(undef, num_subproblems(sp))
                    for i in 1:num_subproblems(sp)
                        subproblems[i] = SubProblem(
                            subproblem(sp, i),
                            start_id + i,
                            T(probability(scenario(sp, i))))
                    end
                    put!(sw, subproblems)
                    return nothing
                end
        end
    end
    return nothing
end

function initialize_subproblems!(execution::SynchronousExecution{T, SmoothSubProblem{T}},
                                 scenarioproblems::DistributedScenarioProblems,
                                 x::AbstractVector,
                                 params::SmoothingParameters) where T <: AbstractFloat
    # Create subproblems on worker processes
    @sync begin
        for w in workers()
            execution.subworkers[w-1] = RemoteChannel(() -> Channel{Vector{SmoothSubProblem{T}}}(1), w)
            prev = map(2:(w-1)) do p
                scenarioproblems.scenario_distribution[p-1]
            end
            start_id = isempty(prev) ? 0 : sum(prev)
            @async remotecall_fetch(
                w,
                execution.subworkers[w-1],
                execution.decisions[w-1],
                scenarioproblems[w-1],
                x,
                start_id,
                params) do sw, decisions, sp_, x, start_id, params
                    for vi in all_known_decisions(fetch(decisions))
                        # Unfix first-stage decisions
                        fetch(decisions)[vi].state = NotTaken
                    end
                    sp = fetch(sp_)
                    subproblems = Vector{SmoothSubProblem{T}}(undef, num_subproblems(sp))
                    for i in 1:num_subproblems(sp)
                        subproblems[i] = SmoothSubProblem(
                            subproblem(sp, i),
                            start_id + i,
                            T(probability(scenario(sp, i))),
                            x;
                            type2dict(params)...)
                    end
                    put!(sw, subproblems)
                    return nothing
                end
        end
    end
    return nothing
end

function restore_subproblems!(::AbstractQuasiGradient, execution::SynchronousExecution{T,SubProblem{T}}) where T <: AbstractFloat
    return nothing
end

function restore_subproblems!(::AbstractQuasiGradient, execution::SynchronousExecution{T,SmoothSubProblem{T}}) where T <: AbstractFloat
    @sync begin
        for w in workers()
            @async remotecall_fetch(w, execution.subworkers[w-1], execution.decisions[w-1]) do sw, decisions
                for subproblem in fetch(sw)
                    restore_subproblem!(subproblem)
                end
                for vi in all_known_decisions(fetch(decisions))
                    # Remove common projection targets
                    remove_decision!(fetch(decisions), vi)
                end
                for vi in all_decisions(fetch(decisions))
                    # Re-fix first-stage decisions
                    fetch(decisions)[vi].state = Known
                end
            end
        end
    end
    return nothing
end

function resolve_subproblems!(quasigradient::AbstractQuasiGradient, execution::SynchronousExecution{T,SubProblem{T}}) where T <: AbstractFloat
    # Prepare
    partial_gradients = Vector{typeof(quasigradient.gradient)}(undef, nworkers())
    partial_objectives = Vector{T}(undef, nworkers())
    # Initialize subgradient
    quasigradient.gradient .= quasigradient.c
    @sync begin
        for (i,w) in enumerate(workers())
            @async partial_objectives[i], partial_gradients[i] =
                remotecall_fetch(w,
                                 execution.subworkers[w-1],
                                 execution.decisions[w-1],
                                 quasigradient.x) do sw, decisions, x
                                     # Fetch all subproblems stored in worker
                                     subproblems::Vector{SubProblem{T}} = fetch(sw)
                                     # Prepare
                                     partial_subgradient = zero(x)
                                     Q = zero(T)
                                     if length(subproblems) == 0
                                         return Q, partial_subgradient
                                     end
                                     # Update subproblems
                                     update_known_decisions!(fetch(decisions), x)
                                     # Update and solve subproblems
                                     for subproblem in subproblems
                                         update_subproblem!(subproblem)
                                         subgradient::SparseGradient{T} = subproblem(x)
                                         if isinf(subgradient.Q)
                                             return subgradient.Q, partial_subgradient
                                         end
                                         partial_subgradient .-= subgradient.δQ
                                         Q += subgradient.Q
                                     end
                                     return Q, partial_subgradient
                                 end
        end
    end
    # Collect results
    quasigradient.gradient .+= sum(partial_gradients)
    # Return current objective value and cut_added flag
    return current_objective_value(quasigradient, sum(partial_objectives))
end

function resolve_subproblems!(quasigradient::AbstractQuasiGradient, execution::SynchronousExecution{T,SmoothSubProblem{T}}) where T <: AbstractFloat
    # Prepare
    partial_gradients = Vector{typeof(quasigradient.gradient)}(undef, nworkers())
    partial_objectives = Vector{T}(undef, nworkers())
    # Initialize subgradient
    quasigradient.gradient .= quasigradient.c
    @sync begin
        for (i,w) in enumerate(workers())
            @async partial_objectives[i], partial_gradients[i] =
                remotecall_fetch(w,
                                 execution.subworkers[w-1],
                                 execution.decisions[w-1],
                                 quasigradient.x) do sw, decisions, x
                                     # Fetch all subproblems stored in worker
                                     subproblems::Vector{SmoothSubProblem{T}} = fetch(sw)
                                     # Prepare
                                     partial_gradient = zero(x)
                                     Q = zero(T)
                                     if length(subproblems) == 0
                                         return Q, partial_gradient
                                     end
                                     # Update subproblems
                                     update_known_decisions!(fetch(decisions), x)
                                     # Update and solve subproblems
                                     for subproblem in subproblems
                                         update_subproblem!(subproblem)
                                         gradient::DenseGradient{T} = subproblem(x)
                                         if isinf(gradient.Q)
                                             return gradient.Q, partial_gradient
                                         end
                                         partial_gradient .+= gradient.δQ
                                         Q += gradient.Q
                                     end
                                     return Q, partial_gradient
                                 end
        end
    end
    # Collect results
    quasigradient.gradient .+= sum(partial_gradients)
    # Return current objective value and cut_added flag
    return current_objective_value(quasigradient, sum(partial_objectives))
end

# API
# ------------------------------------------------------------
function (execution::Synchronous)(structure::StageDecompositionStructure{2, 1, <:Tuple{DistributedScenarioProblems}},
                                  x::AbstractVector,
                                  subproblems::AbstractSubProblemState,
                                  ::Type{T}) where {T <: AbstractFloat}
    return SynchronousExecution(structure, x, subproblems, T)
end

function str(::Synchronous)
    return "Synchronous "
end


text:
using Vinyl: @primitive, overdub, primitive, @hook
using DataFlow
using JuliaInterpreter

struct Trace
  states::Vector{Any}
end

Trace() = Trace([])

struct Staged{T}
  graph::IVertex{Any}
  val::T
end

struct StagedArray{T,N,R<:AbstractArray{T,N}} <: AbstractArray{T,N}
  s::Staged{R}
end

struct StagedReal{N <: Real} <: Real
  s::Staged{N}
end

const StagedType = Union{StagedArray,StagedReal}

Base.show(io::IO, s::Union{StagedType,Staged}) =
  print(io, "$(typeof(s))")

Base.show(io::IO, ::MIME"text/plain", s::Union{StagedType,Staged}) =
  print(io, "$(typeof(s))($(unwrap(s)), $(val(s)))")

val(x) = x
val(x::Staged) = val(x.val)
val(x::StagedType) = val(x.s)
val(x::Tuple) = val.(x)

graph(x::Staged) = x.graph
graph(x::StagedType) = graph(x.s)
graph(x) = DataFlow.constant(x)
vcall(args...) = DataFlow.vertex(DataFlow.Call(), graph.(args)...)

eval_(f, args) = val(f(val.(args)...))

Staged(f, args...; v=eval_(f, args)) = stage(v, vcall(f, args...))

stage(x::AbstractArray{T,N}, v) where {T,N} = StagedArray(Staged(v, val(x)))
stage(x::StagedType, v) = stage(val(x), v)
stage(x::Tuple, v) = ntuple(n -> stage(x[n], vcall(getindex, v, n - 1)), length(x))
stage(x::Real, v) = StagedReal(Staged(v, val(x)))
stage(x, v) = error("Unsupported type $(typeof(x))")

trace(f, args...; meta = Trace()) = overdub(meta, f, args...)

unwrap(x::Staged) = x.graph
unwrap(xs::Tuple) = vcall(tuple, unwrap.(xs)...)
unwrap(x::Union{Number,AbstractArray{<:Number}}) = DataFlow.constant(x)
unwrap(x::StagedType) = unwrap(x.s)

stagedinputs(xs...) = [stage(x, DataFlow.inputnode(n)) for (n, x) in enumerate(xs)]

function _traceλ(f, args...; meta = Trace())
  out = trace(f, stagedinputs(args...)..., meta = meta)
  v = unwrap(out)
  out, vertex(DataFlow.Lambda(length(args), v))
end

traceλ(f, args...; meta = Trace()) = _traceλ(f, args..., meta = meta)[2]

wrap(args...) = stage(args...)

function tracecall(f, args...; meta = Trace())
  out, v = _traceλ(f, args..., meta = meta)
  v = vcall(v, args...)
  wrap(out, v)
end

iscompiled(f) = parentmodule(typeof(f)) == JuliaInterpreter.CompiledCalls

@primitive ctx::Trace function (f::Any)(args...)
  iscompiled(f) && return Base.invokelatest(f, args...)
  inline = !(all(x -> x isa Union{AbstractArray,Number}, args) &&
             any(x -> x isa StagedType, args))
  inline ?
    trace(f, args..., meta = ctx) :
    tracecall(f, args..., meta = ctx)
end

control(a::IVertex, b::IVertex = DataFlow.inputnode()) = vcall(control, a, b)

@primitive ctx::Trace function (f::Flux.Recur)(args...)
  push!(ctx.states, f.init)
  i = length(ctx.states)-1
  states = control(DataFlow.constant(:states))
  state = stage(f.init, vcall(getindex, states, i))
  h, y = trace(f.cell, state, stagedinputs(args...)...)
  λ = vertex(DataFlow.Lambda(length(args),
                             vertex(DataFlow.Do(),
                                    vcall(setindex!, states, unwrap(h), i),
                                    graph(y))))
  wrap(y, vcall(λ, args...))
end

struct BTrace end
@primitive Trace function (::typeof(Base.Broadcast.broadcasted))(f, args...)
  overdub(BTrace(), () -> f(args...))
end

# don't record type conversions
@primitive Trace function (::typeof(Base.Broadcast.broadcasted))(f::Type{T}, x::StagedType) where {T<:Number}
  stage(f.(val(x)), unwrap(x))
end

@primitive Trace (::typeof(Base.Broadcast.materialize))(bc) = bc
@primitive Trace (::typeof(broadcast))(f, args...) = overdub(BTrace(), () -> f(args...))

bcast_ndims(args...) = maximum(arg isa AbstractArray ? ndims(arg) : 0 for arg in args)

function bcastable(op)
  @eval @primitive BTrace (::typeof($op))(args...) = Staged(broadcast, $op, args...)
end

bcastable(op, ops...) = (bcastable(op); bcastable(ops...))

ext: jl


text:
# These functions are excluded until the register map can be parsed automatically,
# and package-specific types are defined to avoid potential type piracy with Base.
# write functions
function Base.write(loc::String, val, dev::LJDevice=DEFAULT_LJDEVICE)
    # look-up datatype for register, convert `val` then call:
    LJM.eWriteName(dev.handle, loc, val) |> errorcheck
    return
end

# read functions
function Base.read(loc::String, dev::LJDevice=DEFAULT_LJDEVICE)
    val = Ref{Float64}() # all read values are written into a 64-bit word
    LJM.eReadName(dev.handle, loc, val) |> errorcheck
    # convert return value to appropriate type based on register data type
    # instead of returning the unloaded Ref value.
    return val[]
end


text:
using PyCall
using DataFrames
using JSONTables
using Plots
using DFTK
using JSON3

const DEBUG = false

debugsuffix() = DEBUG ? "_DEBUG" : ""

function all_elements()
    kgrids = open(JSON3.read, joinpath(@__DIR__, "kgrids.json"))
    string.(keys(kgrids))
end

function loadjson(file::AbstractString)
    open(file) do fp
        DataFrame(jsontable(read(fp, String)))
    end
end

function savejson(file::AbstractString, df::DataFrame)
    open(file, "w") do fp
        write(fp, arraytable(df))
    end
end

function run_dftk(ase_atoms)
    # Pretty much follows the values from SupplMat-Abinit-HGHsc.pdf

    # Try also core=:semicore
    getpsp(el) = load_psp(el.symbol, functional="pbe", core=:fullcore)

    magnetic_moments = load_magnetic_moments(ase_atoms)
    if iszero(sum(pair -> sum(pair[2]), magnetic_moments))
        empty!(magnetic_moments)
    end

    lattice = load_lattice(ase_atoms)
    atoms   = [ElementPsp(el.symbol, psp=getpsp(el)) => position
               for (el, position) in load_atoms(ase_atoms)]
    model   = model_PBE(lattice, atoms, smearing=Smearing.FermiDirac(),
                        temperature=0.01 * DFTK.units.eV,
                        magnetic_moments=magnetic_moments)
    model.spin_polarization != :collinear && empty!(magnetic_moments)

    if DEBUG
        basis = PlaneWaveBasis(model, 10, kgrid=[3, 3, 3])
    else
        symbol = atoms[1][1].symbol
        kgrid  = open(JSON3.read, joinpath(@__DIR__, "kgrids.json"))[symbol]
        basis  = PlaneWaveBasis(model, 250 * DFTK.units.Ry, kgrid=kgrid)
    end
    ρspin  = guess_spin_density(basis, magnetic_moments)
    scfres = self_consistent_field(basis, ρspin=ρspin, tol=1e-8, mixing=HybridMixing())

    scfres.energies.total / DFTK.units.eV
end

function compute_volumes(symbol)
    # Follows https://wiki.fysik.dtu.dk/ase/tutorials/deltacodesdft/deltacodesdft.html
    dcdft = pyimport("ase.collections").dcdft

    # Table of volume and energy per atom
    df = DataFrame(code=String[], symbol=String[], volume=Float64[], energy=Float64[])
    for cellfactor in (94:2:108) ./ 100
        println("#\n# Working on $symbol $(100cellfactor)%\n#")

        atoms = dcdft.__getitem__(symbol)
        atoms.set_cell(atoms.cell * cbrt(cellfactor), scale_atoms=true)
        push!(df, ("DFTK", symbol,
                   atoms.get_volume() / length(atoms),
                   run_dftk(atoms) / length(atoms)))

        println()
        println()
    end
    df
end

function fit_eos(df)
    # For each Symbol fit an eos and return a new DataFrame
    EquationOfState = pyimport("ase.eos").EquationOfState

    # See https://juliadata.github.io/DataFrames.jl/stable/man/split_apply_combine/
    combine(groupby(df, [:symbol, :code])) do gdf
        eos = EquationOfState(gdf.volume, gdf.energy, "birchmurnaghan")
        eos.fit(warn=true)
        e0, B, Bp, v0 = eos.eos_parameters
        (energy=e0, B=B, Bp=Bp, volume=v0)
    end
end

function reference_eos()
    dcdft = pyimport("ase.collections").dcdft
    kJ = pyimport("ase.units").kJ
    df = DataFrame(symbol=String[], code=String[], energy=Union{Float64,Missing}[],
                   B=Float64[], Bp=Float64[], volume=Float64[])
    # df uses the units of ASE:
    # energy  Total energy in eV
    # B       Bulk modulus in eV/Å^3
    # Bp      Pressure derivative of bulk modulus
    # volume  Volume per atom in Å^3

    for (symbol, data) in dcdft.data
        for code in ("wien2k", "exp")
            if "$(code)_B" in keys(data)
                # Note unit conversion: Bulk modulus in dcdft.data stored in GPa,
                # but we need eV / Å^3 for consistency with ASE
                push!(df, (symbol, "$(code)",
                           get(data, "$(code)_energy", missing),
                           data["$(code)_B"] * 1e-24 * kJ,
                           get(data, "$(code)_Bp",     missing),
                           get(data, "$(code)_volume", missing)))
            end
        end
    end

    df
end

function amend_deltas(df)
    function calc_delta(compare, symbol, volume, B, Bp)
        delta = pyimport("ase.utils.deltacodesdft").delta
        dfcompare = filter(:code => isequal(compare), df)
        compare = filter(:symbol => isequal(symbol), dfcompare)
        isempty(compare) && return missing
        ctuple = (only(compare.volume), only(compare.B), only(compare.Bp))
        1000 * delta(ctuple..., volume, B, Bp)
    end

    transform(df,
        [:symbol, :volume, :B, :Bp] => (
            (symbol, volume, B, Bp) -> calc_delta.("wien2k", symbol, volume, B, Bp)
        ) => "Δ wien2k",
        [:symbol, :volume, :B, :Bp] => (
            (symbol, volume, B, Bp) -> calc_delta.("exp", symbol, volume, B, Bp)
        ) => "Δ exp"
    )
end

function volumedata(elements=all_elements())
    datafile = "volumedata$(debugsuffix()).json"
    !isfile(datafile) && savejson(datafile, DataFrame())
    dfvol = loadjson(datafile)

    for symbol in elements
        !isempty(dfvol) && symbol in dfvol.symbol && continue
        res = compute_volumes(symbol)

        # Update dfvol and make a checkpoint save
        dfvol = loadjson(datafile)
        append!(dfvol, res)
        savejson(datafile, dfvol)
    end
    dfvol
end

function compute_deltas(elements=all_elements())
    dfeos = append!(reference_eos(), fit_eos(volumedata(elements)))
    dfdelta = amend_deltas(dfeos)
    savejson("delta$(debugsuffix()).json", dfdelta)
    dfdelta
end

function plot_case(df, symbol)
    birchmurnaghan = pyimport("ase.eos").birchmurnaghan

    df = filter(:symbol => isequal(symbol), df)
    vmin = minimum(df.volume)
    vmax = maximum(df.volume)
    v = collect((0.94 * vmin):0.05:(1.06 * vmax))

    p = plot()
    for (i, code) in enumerate(("wien2k", "DFTK", "exp"))
        dfcode = filter(:code => isequal(code), df)
        isempty(dfcode) && continue
        ctuple = (only(dfcode.volume), only(dfcode.B), only(dfcode.Bp))
        println(ctuple)
        plot!(p, v, birchmurnaghan(v, 0.0, ctuple...), label=code, color=i)
    end

    dfdftk = filter(:code => isequal("DFTK"), df)
    if !isempty(dfdftk)
        eDFTK = only(dfdftk.energy)
        computed = filter(:symbol => isequal(symbol), volumedata())
        println(eDFTK, " ", computed.energy .- eDFTK)
        scatter!(p, computed.volume, computed.energy .- eDFTK)
    end

    xlabel!(p, "volume [Å^3]")
    ylabel!(p, "energy [eV/atom]")
    p
end

function main()
    DEBUG && @warn "Running in debug mode. Results not converged."
    elements = isempty(ARGS) ? all_elements() : ARGS
    compute_deltas(elements)
end
(abspath(PROGRAM_FILE) == @__FILE__) && main()

max_stars_repo_name: mfherbst/DFTK-deltatest

max_issues_repo_name: mfherbst/DFTK-deltatest

max_forks_repo_name: mfherbst/DFTK-deltatest


text:
#
# In the following two examples V is the factored covariance matrix,
# i.e. Q = VV', μ is the vector of expected returns for a set of
# stocks, δ is the target expected return and S is a measure for the
# risk (specifically, the maximum variance).
#
# The interesting constraint here is the constraint t > x'Q x. We notice that
#   x' Q x = x'VVx = (x'V)^2 = || x' V ||^2
# And write
#   t > || x' V ||^2
# Or equivalently
#   (0.5, t, x' V) in RotatedSecondOrderCone
#
# Our initial total wealth is 1, thus sum(x)=1.
#

using JuMP
using MosekTools

# Find the position x with minimal risk such that we get at least thge
# expected return δ.
#
#
function portfolio_min_risk(solver, V :: Array{Float64,2}, μ :: Vector{Float64}, δ :: Float64)
    model = Model(with_optimizer(solver))
    n = size(μ,1)
    m = size(V,2)

    @variable(model, t >= 0)
    @variable(model, x[1:n] >= 0)

    @constraint(model, μ'x .>= δ)

    @constraint(model, sum(x) == 1.0)
    @constraint(model, [0.5; t; V' * x] in MOI.RotatedSecondOrderCone(m+2))

    @objective(model, Min, t)

    optimize!(model)

    value(t), [ value(item) for item in x ]
end

# Find the position x with maximum return such that the variance is at most σ.
function portfolio_max_return(solver, V :: Array{Float64,2}, μ :: Vector{Float64}, R :: Float64)
    model = Model(with_optimizer(solver))
    n = length(μ)
    m = size(V,2)

    @variable(model, t >= 0)
    @variable(model, x[1:n] >= 0)

    @constraint(model, sum(x) == 1.0)
    @constraint(model, [0.5; t; V' * x] in MOI.RotatedSecondOrderCone(m+2))
    @constraint(model, t <= R)

    @objective(model, Max, μ'x)

    optimize!(model)

    value(t), [ value(item) for item in x ]
end

# Find the position x that maximizes the Sharpe ratio S(x) =
# (μ'x-rf)/||Vx||. See MOSEK Cookbook 2018 section 3.3.4 for details.
#
# Here rf is the return of the risk-free asset.
#
# Note that in the implementation we minimize 1/S(x).
function portfolio_sharpe_ratio(solver, V :: Array{Float64,2}, μ :: Vector{Float64}, rf :: Float64)
    model = Model(with_optimizer(solver))
    n = length(μ)
    m = size(V,2)

    @variable(model, t >= 0)
    @variable(model, y[1:n] >= 0)
    @variable(model, z >= 0)

    @constraint(model, [t; V' * y] in MOI.SecondOrderCone(m+1))
    @constraint(model, sum(y) - z == 0.0)
    @constraint(model, (μ .- rf)'y == 1)

    @objective(model, Min, t)

    optimize!(model)

    zres = value(z)

    1.0/value(t), [ value(item)/zres for item in y ]
end

#
# Find the position x with minimal risk, including a risk parity term,
# such that we get at least the expected return δ. See MOSEK Cookbook
# 2018 section 5.4.1 for details.
#
# Here c is the weight given to the parity term.
function portfolio_risk_parity(solver, V :: Array{Float64,2}, μ :: Vector{Float64}, δ :: Float64, c :: Float64)
    model = direct_model(solver())
    n = size(μ,1)
    m = size(V,2)

    @variable(model, t >= 0)
    @variable(model, x[1:n] >= 0)
    @variable(model, s[1:n] >= 0)

    @constraint(model, μ'x .>= δ)

    @constraint(model, sum(x) == 1.0)
    @constraint(model, [0.5; t; V' * x] in MOI.RotatedSecondOrderCone(m+2))
    for i in 1:n
        @constraint(model, [ s[i], 1.0, x[i] ] in MOI.ExponentialCone())
    end

    @objective(model, Min, t + c * sum(s))

    optimize!(model)

    value(t), [ value(item) for item in x ]
end

solver = Mosek.Optimizer
## Example data:

# Vector of expected returns
μ = [1.05, 1.3, 0.9, 1.0, 0.9, 1.5]

# lower bound on return of investment
δ = 1.1

# upper bound for standard deviation
σ = 0.018

# factor of covariance (we use the data-matrix)
V = [-0.0644  -0.0352  -0.0019   0.1015 ;
      0.1183  -0.0428  -0.0013  -0.0741 ;
     -0.0798  -0.1201   0.1040   0.0958 ;
     -0.0855  -0.1174   0.1042   0.0987 ;
     -0.0158  -0.0074   0.0110   0.0123 ;
      0.0636  -0.1341   0.0206   0.0497 ]

t1,x1 = portfolio_min_risk(solver,V,μ,δ)
t2,x2 = portfolio_max_return(solver,V,μ,σ)
t3,x3 = portfolio_sharpe_ratio(solver,V,μ,0.0)
t4,x4 = portfolio_risk_parity(solver,V,μ,δ,0.05)

println("Minimize risk with δ = $δ:")
println("  variance = $t1")
println("  x = $x1")
println("  expected return = $(μ'x1)")
println("")
println("Maximize expected return with variance <= $σ:")
println("  variance = $t2")
println("  x = $x2")
println("  expected return = $(μ'x2)")
println("")
println("Maximize Sharpe ratio:")
println("  Sharpe ratio = $t3")
println("  x = $x3")
println("  expected return = $(μ'x3)")
println("")
println("Minimize risk with parity:")
println("  risk = $t4")
println("  x = $x4")
println("  expected return = $(μ'x4)")


text:
using PlanktonIndividuals.Output
using JLD2

function test_output()
    grid = RectilinearGrid(size = (16, 16, 16), x = (0,32), y = (0,32), z = (0,-32))
    model = PlanktonModel(CPU(), grid; mode = QuotaMode())
    diags = PlanktonDiagnostics(model; tracer=(:PAR,),
                                   plankton = (:num, :graz, :mort, :dvid, :PS),
                                   time_interval = 60seconds)

    sim = PlanktonSimulation(model, ΔT = 60, iterations = 4, diags = diags)

    sim.output_writer = PlanktonOutputWriter(dir = "./result",
                                             write_log = true,
                                             save_diags = true,
                                             save_plankton = true,
                                             plankton_time_interval = 60,
                                             max_filesize = 256KiB)

    @test sim.output_writer isa PlanktonOutputWriter

    update!(sim)

    @test isfile("result/dynamic_species001.txt")
    @test isfile("result/diags_part1.jld2")
    @test isfile("result/plankton_part1.jld2")

    ds = jldopen("result/diags_part1.jld2")
    time_length = length(keys(ds["timeseries/t"]))
    @test time_length == 2
    close(ds)

    ds1 = jldopen("result/plankton_part1.jld2")
    time_length1 = length(keys(ds1["timeseries/t"]))
    @test time_length1 == 1
    close(ds1)

    rm("result", recursive=true)
    return nothing
end

@testset "Output" begin
    test_output()
end


text:
using CuNESSie: _clamp, _cos, _cross, _dot, _norm, _sign, _smul, _sub
using LinearAlgebra: ×

@test_skip "_cathethus"

@testset "_clamp" begin
    function _clamp_kernel!(dst::CuDeviceVector{T}, vals::CuDeviceVector{T}) where T
        for i in 1:9
            dst[i] = _clamp(vals[i+2], vals[1], vals[2])
        end
        nothing
    end

    for T in [Float32, Float64]
        vals = CuArray(T[-1, 1, -2, -1.0001, -1, -.999, 0, .999, 1, 1.0001, 2])
        dst  = CuArray{T}(undef, 9)

        @cuda _clamp_kernel!(dst, vals)
        res = Array(dst)
        @test res ≈ T[-1, -1, -1, -.999, 0, .999, 1, 1, 1]
    end
end

@testset "_cos" begin
    function _cos_kernel!(dst::CuDeviceVector{T}, vals::CuPositionVector{T}) where T
        u = vals[1]
        un = _norm(u)
        dst[1] = _cos(u, un, u, un)
        for i in 0:5
            v = vals[i + 2]
            vn = _norm(v)
            dst[2i + 2] = _cos(u, un, v, vn)
            dst[2i + 3] = _cos(v, vn, u, un)
        end
        nothing
    end

    for T in [Float32, Float64]
        vals = PositionVector([
            T[1, 0, 0],
            T[0, 1, 0], T[-1, 0, 0], T[0, -1, 0],
            T[1, 1, 0], T[-1, 1, 0], T[1, 0, 1]
        ])
        dst  = CuArray{T}(undef, 13)

        @cuda _cos_kernel!(dst, vals)
        res = Array(dst)
        @test res[1] ≈ one(T)
        @test res[2] ≈ zero(T)
        @test res[3] ≈ zero(T)
        @test res[4] ≈ -one(T)
        @test res[5] ≈ -one(T)
        @test res[6] ≈ zero(T)
        @test res[7] ≈ zero(T)
        @test res[8] ≈ T(cos(π/4))
        @test res[9] ≈ T(cos(π/4))
        @test res[10] ≈ T(cos(3π/4))
        @test res[11] ≈ T(cos(3π/4))
        @test res[12] ≈ T(cos(π/4))
        @test res[13] ≈ T(cos(π/4))
    end
end

@testset "_cross" begin
    function _cross_kernel!(dst::CuDeviceVector{T}, vals::CuPositionVector{T}) where T
        v1 = vals[1]
        v2 = vals[2]
        v3 = vals[3]
        v4 = vals[4]
        dst[1], dst[2], dst[3] = _cross(v1, v2)
        dst[4], dst[5], dst[6] = _cross(v2, v1)
        dst[7], dst[8], dst[9] = _cross(v1, v3)
        dst[10], dst[11], dst[12] = _cross(v3, v1)
        dst[13], dst[14], dst[15] = _cross(v1, v4)
        dst[16], dst[17], dst[18] = _cross(v4, v1)
        nothing
    end

    for T in [Float32, Float64]
        vals = PositionVector([T[1, 2, 3], T[4, 6, 8], T[-1, -1, -1], T[0, 0, 0]])
        dst  = CuArray{T}(undef, 3 * 6)

        @cuda _cross_kernel!(dst, vals)
        res = Array(dst)
        @test res[1:3] ≈ T[1, 2, 3] × T[4, 6, 8]
        @test res[4:6] ≈ T[4, 6, 8] × T[1, 2, 3]
        @test res[7:9] ≈ T[1, 2, 3] × T[-1, -1, -1]
        @test res[10:12] ≈ T[-1, -1, -1] × T[1, 2, 3]
        @test res[13:15] ≈ T[1, 2, 3] × T[0, 0, 0]
        @test res[16:18] ≈ T[0, 0, 0] × T[1, 2, 3]
    end
end

@testset "_dot" begin
    function _dot_kernel!(dst::CuDeviceVector{T}, vals::CuPositionVector{T}) where T
        v1 = vals[1]
        v2 = vals[2]
        v3 = vals[3]
        v4 = vals[4]
        dst[1] = _dot(v1, v2)
        dst[2] = _dot(v2, v1)
        dst[3] = _dot(v1, v3)
        dst[4] = _dot(v3, v1)
        dst[5] = _dot(v1, v4)
        dst[6] = _dot(v4, v1)
        nothing
    end

    for T in [Float32, Float64]
        vals = PositionVector([T[1, 2, 3], T[4, 6, 8], T[-1, -1, -1], T[0, 0, 0]])
        dst  = CuArray{T}(undef, 6)

        @cuda _dot_kernel!(dst, vals)
        res = Array(dst)
        @test res ≈ T[40, 40, -6, -6, 0, 0]
    end
end

@testset "_norm" begin
    function _norm_kernel!(dst::CuDeviceVector{T}, vals::CuPositionVector{T}) where T
        dst[1] = _norm(vals[1])
        dst[2] = _norm(vals[2])
        dst[3] = _norm(vals[3])
        dst[4] = _norm(vals[4])
        nothing
    end

    for T in [Float32, Float64]
        vals = PositionVector([T[1, 2, 3], T[4, 6, 8], T[-1, -1, -1], T[0, 0, 0]])
        dst  = CuArray{T}(undef, 4)

        @cuda _norm_kernel!(dst, vals)
        res = Array(dst)
        @test res ≈ T[√14, √116, √3, 0]
    end
end

@testset "_sign" begin
    function _sign_kernel!(dst::CuDeviceVector{T}, vals::CuDeviceVector{T}) where T
        for i in 1:6
            dst[i] = _sign(vals[i])
        end
        nothing
    end

    for T in [Float32, Float64]
        vals = CuArray(T[0, -0, 1, 2, -1, -2])
        dst  = CuArray{T}(undef, 6)

        @cuda _sign_kernel!(dst, vals)
        res = Array(dst)
        @test res ≈ T[0, 0, 1, 1, -1, -1]
    end
end

@testset "_smul" begin
    function _smul_kernel!(
        dst::CuDeviceVector{T},
        vals::CuPositionVector{T},
        s::T
    ) where T
        dst[1], dst[2], dst[3] = _smul(vals[1], s)
        dst[4], dst[5], dst[6] = _smul(vals[2], s)
        dst[7], dst[8], dst[9] = _smul(vals[3], s)
        dst[10], dst[11], dst[12] = _smul(vals[4], s)
        nothing
    end

    for T in [Float32, Float64]
        vals = PositionVector([T[1, 2, 3], T[4, 6, 8], T[-1, -1, -1], T[0, 0, 0]])
        dst  = CuArray{T}(undef, 12)

        @cuda _smul_kernel!(dst, vals, T(2))
        res = Array(dst)
        @test res[1:3] ≈ T[2, 4, 6]
        @test res[4:6] ≈ T[8, 12, 16]
        @test res[7:9] ≈ -T[2, 2, 2]
        @test res[10:12] ≈ T[0, 0, 0]

        @cuda _smul_kernel!(dst, vals, T(-2))
        res = Array(dst)
        @test res[1:3] ≈ -T[2, 4, 6]
        @test res[4:6] ≈ -T[8, 12, 16]
        @test res[7:9] ≈ T[2, 2, 2]
        @test res[10:12] ≈ -T[0, 0, 0]
    end
end

@testset "_sub" begin
    function _sub_kernel!(dst::CuDeviceVector{T}, vals::CuPositionVector{T}) where T
        v1 = vals[1]
        v2 = vals[2]
        v3 = vals[3]
        v4 = vals[4]
        dst[1], dst[2], dst[3] = _sub(v1, v2)
        dst[4], dst[5], dst[6] = _sub(v2, v1)
        dst[7], dst[8], dst[9] = _sub(v1, v3)
        dst[10], dst[11], dst[12] = _sub(v3, v1)
        dst[13], dst[14], dst[15] = _sub(v1, v4)
        dst[16], dst[17], dst[18] = _sub(v4, v1)
        nothing
    end

    for T in [Float32, Float64]
        vals = PositionVector([T[1, 2, 3], T[4, 6, 8], T[-1, -1, -1], T[0, 0, 0]])
        dst  = CuArray{T}(undef, 3 * 6)

        @cuda _sub_kernel!(dst, vals)
        res = Array(dst)
        @test res[1:3] ≈ T[-3, -4, -5]
        @test res[4:6] ≈ T[3, 4, 5]
        @test res[7:9] ≈ T[2, 3, 4]
        @test res[10:12] ≈ T[-2, -3, -4]
        @test res[13:15] ≈ T[1, 2, 3]
        @test res[16:18] ≈ T[-1, -2, -3]
    end
end


text:
@with_kw struct ArmijoGoldstein <: ManifoldLinesearch
    ω₁::Float64 = 1e-4
    ω₂::Float64 = 0.99
    τₑ::Float64 = 1.5        # expansion factor when interval is unbounded
    maxit::Int = 50
end

function linesearch(ls::ArmijoGoldstein, state, linesearchstate, pb::CompositeProblem, M::Manifold, x, ∇fₘ, d; hist=Dict())
    @unpack ω₁, ω₂, τₑ = ls

    # TODO: replace 0 by -ϵ, ϵ>0 for descent dirction criterion.
    # TODO: extrapolation step for gradient...
    if inner(M, x, ∇fₘ, d) / (norm(M, x, ∇fₘ)*norm(M, x, d)) > 0
        @warn "ArmijoGoldstein: non negative direction provided, taking opposite direction." inner(M, x, ∇fₘ, d), norm(M, x, ∇fₘ), norm(M, x, d)
        d = -d
    end

    α = 1
    α_low, α_up = 0, Inf

    F_x = F(pb, x)
    F_cand = Inf
    x_cand = deepcopy(x)
    dh_0 = inner(M, x, ∇fₘ, d)

    it_ls = 0
    validpoint = false
    while !validpoint
        x_cand = retract(M, x, α*d); state.ncalls_retr += 1
        F_cand = F(pb, x_cand); state.ncalls_f += 1; state.ncalls_g += 1

        if F_x > F_cand > F_x - 3*eps(F_cand)
            @warn "Linesearch: reached function conditionning of funtion here, #it ls: $it_ls"
            # @printf "F_x        : %.16e\n" F_x
            # @printf "F_cand     : %.16e\n" F_cand
            # @printf "eps(F_cand): %.16e\n" eps(F_cand)
            break
        end

        if (F_x + ω₁*α*dh_0 ≥ F_cand) && (F_cand ≥ F_x + ω₂*α*dh_0)
            # α = 1 should be accepted for superlinear convergence.
            validpoint = true
            break
        end

        if F_x + ω₁*α*dh_0 < F_cand
            α_up = α

            α = α_up*0.1 + α_low*0.9    # NOTE : Might do interpolation here?
        elseif F_cand < F_x + ω₂*α*dh_0
            α_low = α

            if isinf(α_up)
                α = τₑ * α_low
            else
                α = α_up*0.9 + α_low*0.1
            end
        else
            validpoint = true
        end

        it_ls += 1
        (it_ls > ls.maxit) && (break)
    end

    if F_cand > F_x + ω₁*α*dh_0
        @debug "Linesearch: no suficient decrease" F_cand F_x + ω₁*α*dh_0
    end
    if F_x + ω₂*α*dh_0 > F_cand
        @debug "Linesearch: too small step" F_x + ω₂*α*dh_0 F_cand
    end

    state.niter_manls += it_ls
    hist[:niter] = it_ls
    hist[:end_fval] = F_cand

    return x_cand
end


text:
using Redpitaya
using PyPlot

#rp = RedPitaya("192.168.1.26")
rp = RedPitaya("10.167.6.99")

freqs  = linspace( 10e3, 40e3, 50)
tf = measureTransferFunction(rp, freqs)

figure(1)
clf()
subplot(2,1,1)
plot(freqs,angle(tf),"o-r",lw=2)
subplot(2,1,2)
semilogy(freqs,abs(tf),"o-b",lw=2)


text:
#=
We want to simulate the life of a puppy. It's life consists of three activities:
eating, sleeping and playing. This simple life can be disturbed (i.e. interrupted)
when it gets picked up by a human. If a puppy likes you, it might lick your face.
After being picked up, a puppy continues its life as before.

We create our own type `Puppy` that has an associated process that models its life.
We also create a type `Human` that has an associated process to pick up a random
puppy from time to time.

This application illustrates how you can interrupt an ongoing process and even do
something with the cause of the interruption (in this case keeping track of who
got liked by a puppy).
=#
using SimJulia
using ResumableFunctions
using Logging

# known puppy properties
const puppyparams = Dict(:eat => 1:5, :sleep => 5:30, :play => 5:10)

mutable struct Puppy
    name::String
    state::Symbol
    proc::Process
    pickedup::Int
    licked::Int
    function Puppy(env::Environment, n::String)
        p = new()
        p.name = n
        p.pickedup=0
        p.licked=0
        p.state = rand(keys(puppyparams)) # start in random state
        p.proc = @process puppylife(env,p)
        return p
    end
end

mutable struct Human
    name::String
    picked::Int
    licked::Int
    proc::Process
    function Human(env::Environment, litter::Array{Puppy,1}, name::String)
        h = new()
        h.name = name
        h.picked = 0
        h.licked = 0
        h.proc = @process humanlife(env, litter, h)
        return h
    end
end

@resumable function puppylife(env::Environment, p::Puppy)
    while true
        try
            # do current activity for an amount of time
            duration = rand(puppyparams[p.state])
            @debug """$(now(env)) - Puppy '$(p.name)' current state: $(p.state) for $(duration)"""
            @yield timeout(env, duration)
            # change state (avoiding the current one)
            newstate = rand(filter(x -> x ≠ p.state, keys(puppyparams)))
            @debug """$(now(env)) - Puppy '$(p.name)' will go from $(p.state) to $(newstate)"""
            p.state = newstate
        catch err
            # if interuppted this will happen
            @debug "$(now(env)) - '$(p.name)' is picked up by $(err.cause.name)"
            p.pickedup += 1
            if rand() < 0.5
                @debug "$(now(env)) - '$(p.name)' likes $(err.cause.name)"
                # track how many humans a puppy has liked
                p.licked += 1
                # track how many this specific human was liked by a puppy
                err.cause.licked += 1
            end
        end
    end
end

@resumable function humanlife(env::Environment, litter::Array{Puppy,1}, human::Human)
    while true
        @yield timeout(env, rand(5:20))
        choice = rand(litter)
        @debug "$(now(env)) - $(human.name) chooses to pick up $(choice.name)"
        human.picked += 1
        @yield interrupt(choice.proc, human)
    end
end

function mysim(humannames=["Anaïs", "Pieter", "Bart"],
               dognames=["Django", "Zappa", "Squeez", "Fluffy"])
    @info "\n$("-"^70)\nPuppy life\n$("-"^70)\n"
    sim = Simulation()
    # add puppies
    litter = [Puppy(sim, name) for name in dognames]
    # add humans
    humans = [Human(sim, litter, name) for name in humannames]
    # run simulation
    run(sim,100)
    # show results
    overview_h = ["\t$(human.name) was liked $(human.licked)/$(human.picked) times\n" for human in humans]
    overview_p = ["\t$(p.name) liked getting picked up $(p.licked)/$(p.pickedup) times\n" for p in litter]
    msg = "Results\n$("-"^15)\n$(string("\tHumans:\n",overview_h...))\n$(string("\tPuppies:\n",overview_p...))"
    @info msg
end

# logging settings
Logging.disable_logging(LogLevel(-1000)) # (de-)activate debug messages
logger = Logging.ConsoleLogger(stdout, Logging.Debug)

# run simulation
with_logger(logger) do
    mysim()
end


text:
include("../../fancy_toys.jl")

version = v"1.2.0"

name = "CUTENSOR_CUDA$(cuda_version.major)$(cuda_version.minor)"

sources_linux_x64 = [
    ArchiveSource("https://developer.nvidia.com/compute/cutensor/secure/1.2.0/local_installers/libcutensor-linux-x86_64-1.2.0.tar.gz",
                  "0b33694d391bca537cad0f349b77b31fe45f668abdaee7de9133ca30d3bded6e")
]
sources_linux_ppc64le = [
    ArchiveSource("https://developer.nvidia.com/compute/cutensor/secure/1.2.0/local_installers/libcutensor-linux-ppc64le-1.2.0.tar.gz",
                  "892fa6b3f48bd9f46dcc3001fe702e88605c2b6b65709815aa558e4823539617")
]
sources_windows = [
    FileSource("https://developer.nvidia.com/compute/cutensor/secure/1.2.0/local_installers/libcutensor_1.2.0.exe",
                  "1e1dcffc9a91d88cbe89d039838bf0d0be0e4928751cd03a8274a262ec1ebba3")
]

products = [
    LibraryProduct(["libcutensor", "cutensor"], :libcutensor),
]

dependencies = [
    Dependency(PackageSpec(name="CUDA_jll", version=cuda_version)),
    Dependency(PackageSpec(name="CompilerSupportLibraries_jll", uuid="e66e0078-7015-5450-92f7-15fbd957f2ae"))
]

non_reg_ARGS = filter(arg -> arg != "--register", ARGS)

if @isdefined(sources_linux_x64) && should_build_platform("x86_64-linux-gnu")
    build_tarballs(non_reg_ARGS, name, version, sources_linux_x64, script,
                   [Platform("x86_64", "linux")], products, dependencies)
end

if @isdefined(sources_linux_ppc64le) && should_build_platform("powerpc64le-linux-gnu")
    build_tarballs(non_reg_ARGS, name, version, sources_linux_ppc64le, script,
                   [Platform("powerpc64le", "linux")], products, dependencies)
end

if @isdefined(sources_windows) && should_build_platform("x86_64-w64-mingw32")
    build_tarballs(ARGS, name, version, sources_windows, script,
                   [Platform("x86_64", "windows")], products, dependencies)
end

max_stars_repo_name: mortenpi/Yggdrasil

max_issues_repo_name: mortenpi/Yggdrasil

max_forks_repo_name: mortenpi/Yggdrasil


text:
using Zygote
using Zygote:@adjoint
using Random

export gradient_check

@adjoint function qr(A)
		Q, R = qr(A)
    (Q, R), dy -> (qr_back(A, Q, R, dy...),)
end

@adjoint function qr(A::AbstractMatrix, pivot::Val{true})
		Q, R, P = qr(A, pivot)
    (Q, R, P), dy -> (qr_back(Q*R, Q, R, dy[1], dy[2])*P',nothing)
end

@adjoint function lq(A)
    L, Q = lq(A)
    (L, Q), dy -> (lq_back(A, L, Q, dy...),)
end

@adjoint function svd(A)
    U, S, V = svd(A)
    (U, S, V), dy -> (svd_back(U, S, V, dy...),)
end

@adjoint function rsvd(A, args...; kwargs...)
    U, S, V = rsvd(A, args...; kwargs...)
    (U, S, V), dy -> (svd_back(U, S, V, dy...),)
end

@adjoint function symeigen(A)
    E, U = symeigen(A)
    (E, U), adjy -> (symeigen_back(E, U, adjy...),)
end

function gradient_check(f, args...; η = 1e-5)
    g = gradient(f, args...)
    dy_expect = η*sum(abs2.(g[1]))
    dy = f(args...)-f([gi == nothing ? arg : arg.-η.*gi for (arg, gi) in zip(args, g)]...)
    @show dy
    @show dy_expect
    isapprox(dy, dy_expect, rtol=1e-2, atol=1e-8)
end

@adjoint Random.seed!(n) = Random.seed!(n), _ -> nothing


text:
export
ruleSPSampleListOutNPP,
ruleVBSampleListOut

function ruleSPSampleListOutNPP(msg_out::Nothing,
	                            msg_s::Message{PointMass, Multivariate}, # Multivariate, because vectors of samples and weight are passed as parameters
	                            msg_w::Message{PointMass, Multivariate})

    # Extract the variate type from the first sample
    s = msg_s.dist.params[:m][1]
    if isa(s, Number)
	    V = Univariate
	elseif isa(s, AbstractVector)
        V = Multivariate
    elseif isa(s, AbstractMatrix)
       	V = MatrixVariate
    else
    	error("Unexpected sample type: $(typeof(s))")
    end

    return Message(V, SampleList, s=deepcopy(msg_s.dist.params[:m]), w=deepcopy(msg_w.dist.params[:m]))
end

function ruleVBSampleListOut(dist_out::Any,
                             dist_s::ProbabilityDistribution{Multivariate, PointMass}, # Multivariate, because vectors of samples and weight are passed as parameters
                             dist_w::ProbabilityDistribution{Multivariate, PointMass})

    # Extract the variate type from the first sample
    s = dist_s.params[:m][1]
    if isa(s, Number)
        V = Univariate
    elseif isa(s, AbstractVector)
        V = Multivariate
    elseif isa(s, AbstractMatrix)
        V = MatrixVariate
    else
        error("Unexpected sample type: $(typeof(s))")
    end

    return Message(V, SampleList, s=deepcopy(dist_s.params[:m]), w=deepcopy(dist_w.params[:m]))
end


text:
function my_custom_train!(m, loss, ps, data, opt; data_range=nothing, lower=nothing, upper=nothing, cb=()->nothing)
  # training_loss is declared local so it will be available for logging outside the gradient calculation.
  local training_loss
  ps = Zygote.Params(ps)
  for d in data
    gs = Zygote.gradient(ps) do
      training_loss = loss(d...)
      # Code inserted here will be differentiated, unless you need that gradient information
      # it is better to do the work outside this block.
      return training_loss
    end
    cb(d..., training_loss, m)
    # Insert whatever code you want here that needs training_loss, e.g. logging.
    # logging_callback(training_loss)
    # Insert what ever code you want here that needs gradient.
    # E.g. logging with TensorBoardLogger.jl as histogram so you can see if it is becoming huge.
    Flux.Optimise.update!(opt, ps, gs)
    # Here you might like to check validation set accuracy, and break out to do early stopping.
    for i = 1:length(lower)
      # clamp!(ps[1][i], lower[i], upper[i])
      ps[1][i] = max(lower[i],ps[1][i])
      ps[1][i] = min(upper[i],ps[1][i])
    end
  end
end

function my_custom_train2!(m, loss, ps, data, opt; data_range=nothing, lower=nothing, upper=nothing, cb=()->nothing)
  ps = Flux.Params(ps)
  #losses = []
  for d in data

    #Flux.reset!(m)

    x, y = d[1], d[2]

    # if data_range !== nothing
    #   seq_start = data_range[1]
    #   seq_end = length(data_range) == 2 ? data_range[2] : length(x)
    #
    #   @views x = x[1:seq_end]
    #   @views y = y[1:seq_end]
    #
    #   if seq_start != 1
    #     skipx = x[1:seq_start-1]
    #     m.(skipx)
    #     @views x = x[seq_start:end]
    #     @views y = y[seq_start:end]
    #   end
    # end

    # back is a method that computes the product of the gradient so far with its argument.
    train_loss, back = Zygote.pullback(() -> loss(x,y), ps)
    #push!(losses, train_loss)
    cb(x,y,train_loss,m)
    # Insert whatever code you want here that needs training_loss, e.g. logging.
    # logging_callback(training_loss)
    # Apply back() to the correct type of 1.0 to get the gradient of loss.
    gs = back(one(train_loss))
    # Insert what ever code you want here that needs gradient.
    # E.g. logging with TensorBoardLogger.jl as histogram so you can see if it is becoming huge.

    Flux.Optimise.update!(opt, ps, gs)

    # Here you might like to check validation set accuracy, and break out to do early stopping.

    # lower == nothing && continue
    # upper == nothing && continue
    #
    #
    for (i,p) in enumerate(ps[1])
      # clamp!(ps[i], lower[i], upper[i])
      ps[1][i] = max(lower[i],p)
      ps[1][i] = min(upper[i],p)
    end
  end
  #x,y = first(data)
  #cb(x,y,sum(losses)/length(losses),m)
end


text:
if lowercase(get(ENV, "CI", "false")) == "true"
    @info "Executing in CI. Instantiating benchmark environment..."
    using Pkg
    Pkg.activate(@__DIR__)
    Pkg.instantiate()
end

using BenchmarkTools
SUITE = BenchmarkGroup()
for file in readdir(@__DIR__)
    file == "bench_words.jl" && continue
    if startswith(file, "bench_") && endswith(file, ".jl")
        SUITE[file[length("bench_") + 1:end - length(".jl")]] =
            include(file)
    end
end


text:
n = 256
w = exp(-2im*pi/n)
DFT = zeros(Complex, n, n)

for i = 1:n
    for j = 1:n
        DFT[i,j] = w^((i-1)*(j-1))
    end
end


text:
# These are some macros (and supporting functions) to make it easier to define rules.

# Note: must be declared before it is used, which is later in this file.
macro strip_linenos(expr)
    return esc(Base.remove_linenums!(expr))
end

############################################################################################
### @scalar_rule

"""
    @scalar_rule(f(x₁, x₂, ...),
                 @setup(statement₁, statement₂, ...),
                 (∂f₁_∂x₁, ∂f₁_∂x₂, ...),
                 (∂f₂_∂x₁, ∂f₂_∂x₂, ...),
                 ...)

A convenience macro that generates simple scalar forward or reverse rules using
the provided partial derivatives. Specifically, generates the corresponding
methods for `frule` and `rrule`:

    function ChainRulesCore.frule((NoTangent(), Δx₁, Δx₂, ...), ::typeof(f), x₁::Number, x₂::Number, ...)
        Ω = f(x₁, x₂, ...)
        \$(statement₁, statement₂, ...)
        return Ω, (
                (∂f₁_∂x₁ * Δx₁ + ∂f₁_∂x₂ * Δx₂ + ...),
                (∂f₂_∂x₁ * Δx₁ + ∂f₂_∂x₂ * Δx₂ + ...),
                ...
            )
    end

    function ChainRulesCore.rrule(::typeof(f), x₁::Number, x₂::Number, ...)
        Ω = f(x₁, x₂, ...)
        \$(statement₁, statement₂, ...)
        return Ω, ((ΔΩ₁, ΔΩ₂, ...)) -> (
                NoTangent(),
                ∂f₁_∂x₁ * ΔΩ₁ + ∂f₂_∂x₁ * ΔΩ₂ + ...),
                ∂f₁_∂x₂ * ΔΩ₁ + ∂f₂_∂x₂ * ΔΩ₂ + ...),
                ...
            )
    end

If no type constraints in `f(x₁, x₂, ...)` within the call to `@scalar_rule` are
provided, each parameter in the resulting `frule`/`rrule` definition is given a
type constraint of `Number`.
Constraints may also be explicitly be provided to override the `Number` constraint,
e.g. `f(x₁::Complex, x₂)`, which will constrain `x₁` to `Complex` and `x₂` to
`Number`.

At present this does not support defining for closures/functors.
Thus in reverse-mode, the first returned partial,
representing the derivative with respect to the function itself, is always `NoTangent()`.
And in forward-mode, the first input to the returned propagator is always ignored.

The result of `f(x₁, x₂, ...)` is automatically bound to `Ω`. This
allows the primal result to be conveniently referenced (as `Ω`) within the
derivative/setup expressions.

This macro assumes complex functions are holomorphic. In general, for non-holomorphic
functions, the `frule` and `rrule` must be defined manually.

If the derivative is one, (e.g. for identity functions) `true` can be used as the most
general multiplicative identity.

The `@setup` argument can be elided if no setup code is need. In other
words:

    @scalar_rule(f(x₁, x₂, ...),
                 (∂f₁_∂x₁, ∂f₁_∂x₂, ...),
                 (∂f₂_∂x₁, ∂f₂_∂x₂, ...),
                 ...)

is equivalent to:

    @scalar_rule(f(x₁, x₂, ...),
                 @setup(nothing),
                 (∂f₁_∂x₁, ∂f₁_∂x₂, ...),
                 (∂f₂_∂x₁, ∂f₂_∂x₂, ...),
                 ...)

For examples, see ChainRules' `rulesets` directory.

See also: [`frule`](@ref), [`rrule`](@ref).
"""
macro scalar_rule(call, maybe_setup, partials...)
    call, setup_stmts, inputs, partials = _normalize_scalarrules_macro_input(
        call, maybe_setup, partials
    )
    f = call.args[1]

    # Generate variables to store derivatives named dfi/dxj
    derivatives = map(keys(partials)) do i
        syms = map(j -> Symbol("∂f", i, "/∂x", j), keys(inputs))
        return Expr(:tuple, syms...)
    end

    derivative_expr = scalar_derivative_expr(__source__, f, setup_stmts, inputs, partials)
    frule_expr = scalar_frule_expr(__source__, f, call, [], inputs, derivatives)
    rrule_expr = scalar_rrule_expr(__source__, f, call, [], inputs, derivatives)

    # Final return: building the expression to insert in the place of this macro
    return quote
        if !($f isa Type) && fieldcount(typeof($f)) > 0
            throw(
                ArgumentError(
                    "@scalar_rule cannot be used on closures/functors (such as $($f))"
                ),
            )
        end

        $(derivative_expr)
        $(frule_expr)
        $(rrule_expr)
    end
end

"""
    _normalize_scalarrules_macro_input(call, maybe_setup, partials)

returns (in order) the correctly escaped:
    - `call` with out any type constraints
    - `setup_stmts`: the content of `@setup` or `[]` if that is not provided,
    -  `inputs`: with all args having the constraints removed from call, or
        defaulting to `Number`
    - `partials`: which are all `Expr{:tuple,...}`
"""
function _normalize_scalarrules_macro_input(call, maybe_setup, partials)
    # Setup: normalizing input form etc

    if Meta.isexpr(maybe_setup, :macrocall) && maybe_setup.args[1] == Symbol("@setup")
        setup_stmts = Any[esc(ex) for ex in maybe_setup.args[3:end]]
    else
        setup_stmts = []
        partials = (maybe_setup, partials...)
    end
    @assert Meta.isexpr(call, :call)

    # Annotate all arguments in the signature as scalars
    inputs = esc.(_constrain_and_name.(call.args[2:end], :Number))
    # Remove annotations and escape names for the call
    call.args[2:end] .= _unconstrain.(call.args[2:end])
    call.args = esc.(call.args)

    # For consistency in code that follows we make all partials tuple expressions
    partials = map(partials) do partial
        if Meta.isexpr(partial, :tuple)
            Expr(:tuple, map(esc, partial.args)...)
        else
            length(inputs) == 1 || error("Invalid use of `@scalar_rule`")
            Expr(:tuple, esc(partial))
        end
    end

    return call, setup_stmts, inputs, partials
end

"""
    derivatives_given_output(Ω, f, xs...)

Compute the derivative of scalar function `f` at primal input point `xs...`,
given that it had primal output `Ω`.
Return a tuple of tuples with the partial derivatives of `f` with respect to the `xs...`.
The derivative of the `i`-th component of `f` with respect to the `j`-th input can be
accessed as `Df[i][j]`, where `Df = derivatives_given_output(Ω, f, xs...)`.

!!! warning "Experimental"
    This function is experimental and not part of the stable API.
    At the moment, it can be considered an implementation detail of the macro
    [`@scalar_rule`](@ref), in which it is used.
    In the future, the exact semantics of this function will stabilize, and it
    will be added to the stable API.
    When that happens, this warning will be removed.

"""
function derivatives_given_output end

function scalar_derivative_expr(__source__, f, setup_stmts, inputs, partials)
    return @strip_linenos quote
        function ChainRulesCore.derivatives_given_output(
            $(esc(:Ω)), ::Core.Typeof($f), $(inputs...)
        )
            $(__source__)
            $(setup_stmts...)
            return $(Expr(:tuple, partials...))
        end
    end
end

function scalar_frule_expr(__source__, f, call, setup_stmts, inputs, partials)
    n_outputs = length(partials)
    n_inputs = length(inputs)

    # Δs is the input to the propagator rule
    # because this is push-forward there is one per input to the function
    Δs = _propagator_inputs(n_inputs)
    pushforward_returns = map(1:n_outputs) do output_i
        ∂s = partials[output_i].args
        propagation_expr(Δs, ∂s)
    end
    if n_outputs > 1
        # For forward-mode we return a Tangent if output actually a tuple.
        pushforward_returns = Expr(
            :call, :(Tangent{typeof($(esc(:Ω)))}), pushforward_returns...
        )
    else
        pushforward_returns = first(pushforward_returns)
    end

    return @strip_linenos quote
        # _ is the input derivative w.r.t. function internals. since we do not
        # allow closures/functors with @scalar_rule, it is always ignored
        function ChainRulesCore.frule((_, $(Δs...)), ::Core.Typeof($f), $(inputs...))
            $(__source__)
            $(esc(:Ω)) = $call
            $(setup_stmts...)
            $(Expr(:tuple, partials...)) = ChainRulesCore.derivatives_given_output(
                $(esc(:Ω)), $f, $(inputs...)
            )
            return $(esc(:Ω)), $pushforward_returns
        end
    end
end

function scalar_rrule_expr(__source__, f, call, setup_stmts, inputs, partials)
    n_outputs = length(partials)
    n_inputs = length(inputs)

    # Δs is the input to the propagator rule
    # because this is a pull-back there is one per output of function
    Δs = _propagator_inputs(n_outputs)

    # Make a projector for each argument
    projs, psetup = _make_projectors(call.args[2:end])
    append!(setup_stmts, psetup)

    # 1 partial derivative per input
    pullback_returns = map(1:n_inputs) do input_i
        ∂s = [partial.args[input_i] for partial in partials]
        propagation_expr(Δs, ∂s, true, projs[input_i])
    end

    # Multi-output functions have pullbacks with a tuple input that will be destructured
    pullback_input = n_outputs == 1 ? first(Δs) : Expr(:tuple, Δs...)
    pullback = @strip_linenos quote
        @inline function $(esc(propagator_name(f, :pullback)))($pullback_input)
            $(__source__)
            return (NoTangent(), $(pullback_returns...))
        end
    end

    return @strip_linenos quote
        function ChainRulesCore.rrule(::Core.Typeof($f), $(inputs...))
            $(__source__)
            $(esc(:Ω)) = $call
            $(setup_stmts...)
            $(Expr(:tuple, partials...)) = ChainRulesCore.derivatives_given_output(
                $(esc(:Ω)), $f, $(inputs...)
            )
            return $(esc(:Ω)), $pullback
        end
    end
end

# For context on why this is important, see
# https://github.com/JuliaDiff/ChainRulesCore.jl/pull/276
"Declares properly hygenic inputs for propagation expressions"
_propagator_inputs(n) = [esc(gensym(Symbol(:Δ, i))) for i in 1:n]

"given the variable names, escaped but without types, makes setup expressions for projection operators"
function _make_projectors(xs)
    projs = map(x -> Symbol(:proj_, x.args[1]), xs)
    setups = map((x, p) -> :($p = ProjectTo($x)), xs, projs)
    return projs, setups
end

"""
    propagation_expr(Δs, ∂s, [_conj=false, proj=identity])

Returns the expression for the propagation of
the input gradient `Δs` though the partials `∂s`.
Specify `_conj = true` to conjugate the partials.
Projector `proj` is a function that will be applied at the end;
    for `rrules` it is usually a `ProjectTo(x)`, for `frules` it is `identity`
"""
function propagation_expr(Δs, ∂s, _conj=false, proj=identity)
    # This is basically Δs ⋅ ∂s
    _∂s = map(∂s) do ∂s_i
        if _conj
            :(conj($∂s_i))
        else
            ∂s_i
        end
    end

    # Apply `muladd` iteratively.
    # Explicit multiplication is only performed for the first pair of partial and gradient.
    (∂s_1, Δs_1), _∂s_Δs_tail = Iterators.peel(zip(_∂s, Δs))
    init_expr = :($∂s_1 * $Δs_1)
    summed_∂_mul_Δs = foldl(_∂s_Δs_tail; init=init_expr) do ex, (∂s_i, Δs_i)
        :(muladd($∂s_i, $Δs_i, $ex))
    end
    return :($proj($summed_∂_mul_Δs))
end

"""
    propagator_name(f, propname)

Determines a reasonable name for the propagator function.
The name doesn't really matter too much as it is a local function to be returned
by `frule` or `rrule`, but a good name make debugging easier.
`f` should be some form of AST representation of the actual function,
`propname` should be either `:pullback` or `:pushforward`

This is able to deal with fairly complex expressions for `f`:

    julia> propagator_name(:bar, :pushforward)
    :bar_pushforward

    julia> propagator_name(esc(:(Base.Random.foo)), :pullback)
    :foo_pullback
"""
propagator_name(f::Expr, propname::Symbol) = propagator_name(f.args[end], propname)
propagator_name(fname::Symbol, propname::Symbol) = Symbol(fname, :_, propname)
propagator_name(fname::QuoteNode, propname::Symbol) = propagator_name(fname.value, propname)

############################################################################################
### @non_differentiable

"""
    @non_differentiable(signature_expression)

A helper to make it easier to declare that a method is not differentiable.
This is a short-hand for defining an [`frule`](@ref) and [`rrule`](@ref) that
return [`NoTangent()`](@ref) for all partials (even for the function `s̄elf`-partial
itself)

Keyword arguments should not be included.

```jldoctest
julia> @non_differentiable Base.:(==)(a, b)

julia> _, pullback = rrule(==, 2.0, 3.0);

julia> pullback(1.0)
(NoTangent(), NoTangent(), NoTangent())
```

You can place type-constraints in the signature:
```jldoctest
julia> @non_differentiable Base.length(xs::Union{Number, Array})

julia> frule((ZeroTangent(), 1), length, [2.0, 3.0])
(2, NoTangent())
```

!!! warning
    This helper macro covers only the simple common cases.
    It does not support `where`-clauses.
    For these you can declare the `rrule` and `frule` directly

"""
macro non_differentiable(sig_expr)
    Meta.isexpr(sig_expr, :call) || error("Invalid use of `@non_differentiable`")
    has_vararg = _isvararg(sig_expr.args[end])

    primal_name, orig_args = Iterators.peel(sig_expr.args)

    primal_name_sig, primal_name = _split_primal_name(primal_name)
    constrained_args = _constrain_and_name.(orig_args, :Any)
    primal_sig_parts = [primal_name_sig, constrained_args...]

    unconstrained_args = _unconstrain.(constrained_args)

    primal_invoke = if !has_vararg
        :($(primal_name)($(unconstrained_args...)))
    else
        normal_args = unconstrained_args[1:(end - 1)]
        var_arg = unconstrained_args[end]
        :($(primal_name)($(normal_args...), $(var_arg)...))
    end

    quote
        $(_nondiff_frule_expr(__source__, primal_sig_parts, primal_invoke))
        $(_nondiff_rrule_expr(__source__, primal_sig_parts, primal_invoke))
    end
end

"changes `f(x,y)` into `f(x,y; kwargs....)`"
function _with_kwargs_expr(call_expr::Expr, kwargs)
    @assert isexpr(call_expr, :call)
    return Expr(
        :call,
        call_expr.args[1],
        Expr(:parameters, :($(kwargs)...)),
        call_expr.args[2:end]...,
    )
end

function _nondiff_frule_expr(__source__, primal_sig_parts, primal_invoke)
    @gensym kwargs
    return @strip_linenos quote
        # Manually defined kw version to save compiler work. See explanation in rules.jl
        function (::Core.kwftype(typeof(ChainRulesCore.frule)))(
            @nospecialize($kwargs::Any),
            frule::typeof(ChainRulesCore.frule),
            @nospecialize(::Any),
            $(map(esc, primal_sig_parts)...),
        )
            return ($(esc(_with_kwargs_expr(primal_invoke, kwargs))), NoTangent())
        end
        function ChainRulesCore.frule(
            @nospecialize(::Any), $(map(esc, primal_sig_parts)...)
        )
            $(__source__)
            # Julia functions always only have 1 output, so return a single NoTangent()
            return ($(esc(primal_invoke)), NoTangent())
        end
    end
end

function tuple_expression(primal_sig_parts)
    has_vararg = _isvararg(primal_sig_parts[end])
    return if !has_vararg
        num_primal_inputs = length(primal_sig_parts)
        Expr(:tuple, ntuple(_ -> NoTangent(), num_primal_inputs)...)
    else
        num_primal_inputs = length(primal_sig_parts) - 1 # - vararg
        length_expr =
            :($num_primal_inputs + length($(esc(_unconstrain(primal_sig_parts[end])))))
        @strip_linenos :(ntuple(i -> NoTangent(), $length_expr))
    end
end

function _nondiff_rrule_expr(__source__, primal_sig_parts, primal_invoke)
    esc_primal_sig_parts = map(esc, primal_sig_parts)
    tup_expr = tuple_expression(primal_sig_parts)
    primal_name = first(primal_invoke.args)
    pullback_expr = @strip_linenos quote
        function $(esc(propagator_name(primal_name, :pullback)))(@nospecialize(_))
            return $(tup_expr)
        end
    end

    @gensym kwargs
    return @strip_linenos quote
        # Manually defined kw version to save compiler work. See explanation in rules.jl
        function (::Core.kwftype(typeof(rrule)))(
            $(esc(kwargs))::Any, ::typeof(rrule), $(esc_primal_sig_parts...)
        )
            return ($(esc(_with_kwargs_expr(primal_invoke, kwargs))), $pullback_expr)
        end
        function ChainRulesCore.rrule($(esc_primal_sig_parts...))
            $(__source__)
            return ($(esc(primal_invoke)), $pullback_expr)
        end
    end
end

############################################################################################
# @opt_out

"""
    @opt_out frule([config], _, f, args...)
    @opt_out rrule([config], f, args...)

This allows you to opt-out of an `frule` or an `rrule` by providing a more specific method,
that says to use the AD system to differentiate it.

For example, consider some function `foo(x::AbtractArray)`.
In general, you know an efficient and generic way to implement its `rrule`.
You do so, (likely making use of [`ProjectTo`](@ref)).
But it actually turns out that for some `FancyArray` type it is better to let the AD do its
thing.

Then you would write something like:
```julia
function rrule(::typeof(foo), x::AbstractArray)
    foo_pullback(ȳ) = ...
    return foo(x), foo_pullback
end

@opt_out rrule(::typeof(foo), ::FancyArray)
```

This will generate an [`rrule`](@ref) that returns `nothing`,
and will also add a similar entry to [`ChainRulesCore.no_rrule`](@ref).

Similar applies for [`frule`](@ref) and [`ChainRulesCore.no_frule`](@ref)

For more information see the [documentation on opting out of rules](@ref opt_out).
"""
macro opt_out(expr)
    no_rule_target = _target_rewrite!(deepcopy(expr), true)
    rule_target = _target_rewrite!(deepcopy(expr), false)

    return @strip_linenos quote
        $(esc(no_rule_target)) = nothing
        $(esc(rule_target)) = nothing
    end
end

"""
    _target_rewrite!(expr::Expr, no_rule)

Rewrite method sig `expr` for `rrule` to be for `no_rrule` or `ChainRulesCore.rrule`
(with the CRC namespace qualification), depending on the `no_rule` argument.
Does the equivalent for `frule`.
"""
function _target_rewrite!(expr::Expr, no_rule)
    length(expr.args) === 0 && error("Malformed method expression. $expr")
    if expr.head === :call || expr.head === :where
        expr.args[1] = _target_rewrite!(expr.args[1], no_rule)
    elseif expr.head == :(.) && expr.args[1] == :ChainRulesCore
        expr = _target_rewrite!(expr.args[end], no_rule)
    else
        error("Malformed method expression. $(expr)")
    end
    return expr
end
_target_rewrite!(qt::QuoteNode, no_rule) = _target_rewrite!(qt.value, no_rule)
function _target_rewrite!(call_target::Symbol, no_rule)
    return if call_target == :rrule && no_rule
        :($ChainRulesCore.no_rrule)
    elseif call_target == :rrule && !no_rule
        :($ChainRulesCore.rrule)
    elseif call_target == :frule && no_rule
        :($ChainRulesCore.no_frule)
    elseif call_target == :frule && !no_rule
        :($ChainRulesCore.frule)
    else
        error("Unexpected opt-out target. Expected frule or rrule, got: $call_target")
    end
end

############################################################################################
# Helpers

"""
    _isvararg(expr)

returns true if the expression could represent a vararg

```jldoctest
julia> ChainRulesCore._isvararg(:(x...))
true

julia> ChainRulesCore._isvararg(:(x::Int...))
true

julia> ChainRulesCore._isvararg(:(::Int...))
true

julia> ChainRulesCore._isvararg(:(x::Vararg))
true

julia> ChainRulesCore._isvararg(:(x::Vararg{Int}))
true

julia> ChainRulesCore._isvararg(:(::Vararg))
true

julia> ChainRulesCore._isvararg(:(::Vararg{Int}))
true

julia> ChainRulesCore._isvararg(:(x))
false
````
"""
_isvararg(expr) = false
function _isvararg(expr::Expr)
    Meta.isexpr(expr, :...) && return true
    if Meta.isexpr(expr, :(::))
        constraint = last(expr.args)
        constraint == :Vararg && return true
        Meta.isexpr(constraint, :curly) && first(constraint.args) == :Vararg && return true
    end
    return false
end

"""
splits the first arg of the `call` expression into an expression to use in the signature
and one to use for calling that function
"""
function _split_primal_name(primal_name)
    # e.g. f(x, y)
    is_plain = primal_name isa Symbol
    is_qualified = Meta.isexpr(primal_name, :(.))
    is_parameterized = Meta.isexpr(primal_name, :curly)
    if is_plain || is_qualified || is_parameterized
        primal_name_sig = :(::$Core.Typeof($primal_name))
        return primal_name_sig, primal_name
    elseif Meta.isexpr(primal_name, :(::))  # e.g. (::T)(x, y)
        _primal_name = gensym(Symbol(:instance_, primal_name.args[end]))
        primal_name_sig = Expr(:(::), _primal_name, primal_name.args[end])
        return primal_name_sig, _primal_name
    else
        error("invalid primal name: `$primal_name`")
    end
end

"turn both `a` and `a::S` into `a`"
_unconstrain(arg::Symbol) = arg
function _unconstrain(arg::Expr)
    Meta.isexpr(arg, :(::), 2) && return arg.args[1]  # drop constraint.
    Meta.isexpr(arg, :(...), 1) && return _unconstrain(arg.args[1])
    return error("malformed arguments: $arg")
end

"turn both `a` and `::constraint` into `a::constraint` etc"
function _constrain_and_name(arg::Expr, _)
    Meta.isexpr(arg, :(::), 2) && return arg  # it is already fine.
    Meta.isexpr(arg, :(::), 1) && return Expr(:(::), gensym(), arg.args[1]) # add name
    Meta.isexpr(arg, :(...), 1) &&
        return Expr(:(...), _constrain_and_name(arg.args[1], :Any))
    return error("malformed arguments: $arg")
end
_constrain_and_name(name::Symbol, constraint) = Expr(:(::), name, constraint)  # add type


text:
using Dates

"""
    optimize!

Optimizes a local search index for an specific algorithm to get the desired performance.
Note that optimizing for low-recall will yield to faster searches; the train queries
are specified as part of the `perf` struct.
"""
function optimize!(algosearch::LocalSearchAlgorithm,
                   index::SearchGraph{T},
                   dist::Function,
                   recall::Float64,
                   perf::Performance;
                   bsize::Int=4,
                   tol::Float64=0.01) where T
    n = length(index.db)
    score_function(p) = p.recall < recall ? p.recall : 1.0 + n / p.evaluations
    #score_function(p) = p.recall < recall ? p.recall : 1.0 + 1.0 / (1.0 + sum(p.distances))
    #score_function(p) = 1.0 / (1.0 + sum(p.distances))
    optimize_neighborhood!(index.neighborhood_algo, index, dist, perf, recall)
    p = probe(perf, index, dist)
    best_list = [(score=score_function(p), state=algosearch, perf=p)]
    exploration = Dict(algosearch => 0)  ## -1 unexplored; 0 visited; 1 visited & expanded

    @info "==== BEGIN $(typeof(algosearch)). Starting parameter optimization; expected recall: $recall, n: $n"
    prev_score = -1.0
    iter = 0
    while abs(best_list[1].score - prev_score) > tol
        iter += 1
        prev_score = best_list[1].score
        @info "  == begin iteration $(typeof(algosearch)). Iteration: $iter, expected recall: $recall, n: $n"

        for prev in @view best_list[1:end]  ## the view also fixes the size of best_list even after push!
            S = get(exploration, prev.state, -1)
            if S == 1
                continue  # visited and explored
            elseif S == 0
                exploration[prev.state] = 1
            end

            opt_expand_neighborhood(prev.state, n, iter) do state
                S = get(exploration, state, -1)
                if S == -1
                    exploration[state] = 0
                    index.search_algo = state
                    # p = probe(perf, index, repeat=3, aggregation=:median, field=:seconds)
                    p = probe(perf, index, dist, repeat=1, field=:seconds)
                    score = score_function(p)
                    push!(best_list, (score=score, state=state, perf=p))
                    if score > best_list[1].score
                        @info "  ** $(typeof(algosearch)). A new best conf was found> score: $score, conf: $(JSON.json(state)), perf: $(JSON.json(p)), best_list's length: $(length(best_list)), n: $(n)"
                    end
                end
            end
        end

        sort!(best_list, by=(x) -> -x.score)
        if length(best_list) > bsize
            best_list = best_list[1:bsize]
        end
        @info "  == end $(typeof(algosearch)). Iteration finished; $(JSON.json(best_list[1])), beam: $(length(best_list)), n: $(n)"
    end

    index.search_algo = best_list[1].state
    @info "==== END $(typeof(algosearch)). Finished, best: $(JSON.json(best_list[1])), n: $(n)"
    return index
end


text:
using Random
#####
##### Interface for players
#####
using Random: AbstractRNG, GLOBAL_RNG
"""
    AbstractPlayer

Abstract type for a game player.
"""
abstract type AbstractPlayer end

"""
    think(::AbstractPlayer, game)

Return a probability distribution over available actions as a `(actions, π)` pair.
"""
function think end

"""
    reset_player!(::AbstractPlayer)

Reset the internal memory of a player (e.g. the MCTS tree).
The default implementation does nothing.
"""
function reset_player!(::AbstractPlayer)
  return
end

"""
    player_temperature(::AbstractPlayer, game, turn_number)

Return the player temperature, given the number of actions that have
been played before by both players in the current game.

A default implementation is provided that always returns 1.
"""
function player_temperature(::AbstractPlayer, game, turn)
  return 1.0
end

"""
    select_move(player::AbstractPlayer, game, turn_number)

Return a single action. A default implementation is provided that samples
an action according to the distribution computed by [`think`](@ref), with a
temperature given by [`player_temperature`](@ref).
"""
function select_move(player::AbstractPlayer, game, turn_number)
  actions, π = think(player, game)
  τ = player_temperature(player, game, turn_number)
  π = apply_temperature(π, τ)
  return actions[Util.rand_categorical(π)]
end

#####
##### Random Player
#####

"""
    RandomPlayer <: AbstractPlayer

A player that picks actions uniformly at random.
"""
struct RandomPlayer <: AbstractPlayer end

function think(::RandomPlayer, game)
  actions = GI.available_actions(game)
  n = length(actions)
  π = ones(n) ./ length(actions)
  return actions, π
end

#####
##### Epsilon-greedy player
#####

"""
    EpsilonGreedyPlayer{Player} <: AbstractPlayer

A wrapper on a player that makes it choose a random move
with a fixed ``ϵ`` probability.
"""
struct EpsilonGreedyPlayer{P} <: AbstractPlayer
  player :: P
  ϵ :: Float64
end

function think(p::EpsilonGreedyPlayer, game)
  actions, π = think(p.player, game)
  n = length(actions)
  η = ones(n) ./ n
  return actions, (1 - p.ϵ) * π + p.ϵ * η
end

function reset!(p::EpsilonGreedyPlayer)
  reset!(p.player)
end

function player_temperature(p::EpsilonGreedyPlayer, game, turn)
  return player_temperature(p.player, game, turn)
end

#####
##### Player with a custom temperature
#####

"""
    PlayerWithTemperature{Player} <: AbstractPlayer

A wrapper on a player that enables overwriting the temperature schedule.
"""
struct PlayerWithTemperature{P} <: AbstractPlayer
  player :: P
  temperature :: AbstractSchedule{Float64}
end

function think(p::PlayerWithTemperature, game)
  return think(p.player, game)
end

function reset!(p::PlayerWithTemperature)
  reset!(p.player)
end

function player_temperature(p::PlayerWithTemperature, game, turn)
  return p.temperature[turn]
end

#####
##### An MCTS-based player
#####

"""
    MctsPlayer{MctsEnv} <: AbstractPlayer

A player that selects actions using MCTS.

# Constructors

    MctsPlayer(mcts::MCTS.Env; τ, niters, timeout=nothing)

Construct a player from an MCTS environment. When computing each move:

- if `timeout` is provided,
  MCTS simulations are executed for `timeout` seconds by groups of `niters`
- otherwise, `niters` MCTS simulations are run

The temperature parameter `τ` can be either a real number or a
[`AbstractSchedule`](@ref).

    MctsPlayer(game_spec::AbstractGameSpec, oracle,
               params::MctsParams; timeout=nothing)

Construct an MCTS player from an oracle and an [`MctsParams`](@ref) structure.
"""
struct MctsPlayer{M} <: AbstractPlayer
  mcts :: M
  niters :: Int
  timeout :: Union{Float64, Nothing}
  τ :: AbstractSchedule{Float64} # Temperature
  function MctsPlayer(mcts::MCTS.Env; τ, niters, timeout=nothing)
    @assert niters > 0
    @assert isnothing(timeout) || timeout > 0
    new{typeof(mcts)}(mcts, niters, timeout, τ)
  end
end

# Alternative constructor
function MctsPlayer(
    game_spec::AbstractGameSpec, oracle, params::MctsParams; timeout=nothing)
  mcts = MCTS.Env(game_spec, oracle,
    gamma=params.gamma,
    cpuct=params.cpuct,
    noise_ϵ=params.dirichlet_noise_ϵ,
    noise_α=params.dirichlet_noise_α,
    prior_temperature=params.prior_temperature,
    adaptive_cpuct=params.adaptive_cpuct)
  return MctsPlayer(mcts,
    niters=params.num_iters_per_turn,
    τ=params.temperature,
    timeout=timeout)
end

# MCTS with random oracle
function RandomMctsPlayer(game_spec::AbstractGameSpec, params::MctsParams)
  oracle = MCTS.RandomOracle()
  mcts = MCTS.Env(game_spec, oracle,
    cpuct=params.cpuct,
    gamma=params.gamma,
    noise_ϵ=params.dirichlet_noise_ϵ,
    noise_α=params.dirichlet_noise_α,
    adaptive_cpuct=params.adaptive_cpuct)
  return MctsPlayer(mcts,
    niters=params.num_iters_per_turn,
    τ=params.temperature)
end

function think(p::MctsPlayer, game)
  if isnothing(p.timeout) # Fixed number of MCTS simulations
    MCTS.explore!(p.mcts, game, p.niters)
  else # Run simulations until timeout
    start = time()
    while time() - start < p.timeout
      MCTS.explore!(p.mcts, game, p.niters)
    end
  end
  return MCTS.policy(p.mcts, game)
end

function player_temperature(p::MctsPlayer, game, turn)
  return p.τ[turn]
end

function reset_player!(player::MctsPlayer)
  MCTS.reset!(player.mcts)
end

#####
##### Network player
#####

"""
    NetworkPlayer{Net} <: AbstractPlayer

A player that uses the policy output by a neural network directly,
instead of relying on MCTS. The given neural network must be in test mode.
"""
struct NetworkPlayer{N} <: AbstractPlayer
  network :: N
end

function think(p::NetworkPlayer, game)
  actions = GI.available_actions(game)
  state = GI.current_state(game)
  π, _ = p.network(state)
  return actions, π
end

#####
##### Merging two players into one
#####

"""
    TwoPlayers <: AbstractPlayer

If `white` and `black` are two [`AbstractPlayer`](@ref), then
`TwoPlayers(white, black)` is a player that behaves as `white` when `white`
is to play and as `black` when `black` is to play.
"""
struct TwoPlayers{W, B} <: AbstractPlayer
  white :: W
  black :: B
end

flipped_colors(p::TwoPlayers) = TwoPlayers(p.black, p.white)

function think(p::TwoPlayers, game)
  if GI.white_playing(game)
    return think(p.white, game)
  else
    return think(p.black, game)
  end
end

function select_move(p::TwoPlayers, game, turn)
  if GI.white_playing(game)
    return select_move(p.white, game, turn)
  else
    return select_move(p.black, game, turn)
  end
end

function reset!(p::TwoPlayers)
  reset!(p.white)
  reset!(p.black)
end

function player_temperature(p::TwoPlayers, game, turn)
  if GI.white_playing(game)
    return player_temperature(p.white, game, turn)
  else
    return player_temperature(p.black, game, turn)
  end
end

#####
##### Players can play against each other
#####

"""
    play_game(gspec::AbstractGameSpec, player; flip_probability=0.) :: Trace

Simulate a game by an [`AbstractPlayer`](@ref).

- For two-player games, please use [`TwoPlayers`](@ref).
- If the `flip_probability` argument is set to ``p``, the board
  is _flipped_ randomly at every turn with probability ``p``,
  using [`GI.apply_random_symmetry!`](@ref).
"""
function play_game(gspec, player; flip_probability=0., itc=1, rng::AbstractRNG=Random.GLOBAL_RNG)
  game = GI.init(gspec, itc, rng)
  trace = Trace(GI.current_state(game))
  while true
    if GI.game_terminated(game)
      return trace
    end
    if !iszero(flip_probability) && rand() < flip_probability
      GI.apply_random_symmetry!(game)
    end
    actions, π_target = think(player, game)
    τ = player_temperature(player, game, length(trace))
    π_sample = apply_temperature(π_target, τ)
    a = actions[Util.rand_categorical(π_sample)]
    GI.play!(game, a)
    push!(trace, π_target, GI.white_reward(game), GI.current_state(game))
  end
end

#####
##### Simple utilities to play an interactive game
#####

"""
    Human <: AbstractPlayer

Human player that queries the standard input for actions.

Does not implement [`think`](@ref) but instead implements
[`select_move`](@ref) directly.
"""
struct Human <: AbstractPlayer end

struct Quit <: Exception end

function select_move(::Human, game, turn)
  a = nothing
  while isnothing(a) || a ∉ GI.available_actions(game)
    print("> ")
    str = readline()
    print("\n")
    isempty(str) && throw(Quit())
    a = GI.parse_action(GI.spec(game), str)
  end
  return a
end

function play_out(gspec, player, state)
  game = GI.init(gspec, state)
  i = 1
  while true
    if GI.game_terminated(game)
      GI.render(game)
      return
    end
    actions, π_target = think(player, game)
    τ = player_temperature(player, game, i)
    π_sample = apply_temperature(π_target, τ)
    a = actions[Util.rand_categorical(π_sample)]
    GI.play!(game, a)
    i += 1
    print(i, "\r")
  end

end

"""
    interactive!(game)
    interactive!(gspec)
    interactive!(game, player)
    interactive!(gspec, player)
    interactive!(game, white, black)
    interactive!(gspec, white, black)

Launch a possibly interactive game session.

This function takes either an `AbstractGameSpec` or `AbstractGameEnv` as an argument.
"""
function interactive! end

function interactive!(game::AbstractGameEnv, player)
  try
  GI.render(game)
  turn = 0
  while !GI.game_terminated(game)
    action = select_move(player, game, turn)
    GI.play!(game, action)
    GI.render(game)
    turn += 1
  end
  catch e
    isa(e, Quit) || rethrow(e)
    return
  end
end

interactive!(gspec::AbstractGameSpec, player) = interactive!(GI.init(gspec), player)

interactive!(game, white, black) = interactive!(game, TwoPlayers(white, black))

interactive!(game) = interactive!(game, Human(), Human())


text:
module Bezier

export evalbezier, subdivide

include("evalbezier.jl")
include("subdivide.jl")

end # module


text:
using AccelerationBenchmark
using JLD, DataFrames, StatPlots, NaNMath, LaTeXStrings

searchdir(path,key) = filter(x->contains(x,key), readdir(path))
mdf = reduce(append!, [load(fname, "mdf") for fname in ["data/"*fname for fname in searchdir("data", "jld")]])
rdf = AccelerationBenchmark.createratiodataframe(mdf)

# TODO: GroupedErrors is using a uniform spread in the x-values (\tau-values),
# so there is quite low resolution for small x-values.
# I need to deal with this somehow, and combine that with a downsampling approach.
# Alternatively, we can create fevals_AG.csv and run it through "mergealldata.jl" with BenchmarkProfiles.

rdf[:fcalls] .=  NaNMath.log2.(rdf[:fcalls])
plt = AccelerationBenchmark.createperfprofiles(rdf,:fcalls; xscale=:linear,
                                              xlabel=L"$\tau$ ($\log_2$ scale)")


text:
# Add Exasim to Julia search path
versiondir = cdir[1:ii[end]] * "/" * version * "/Julia";
push!(LOAD_PATH, versiondir * "/Gencode");
push!(LOAD_PATH, versiondir * "/Mesh");
push!(LOAD_PATH, versiondir * "/Preprocessing");
push!(LOAD_PATH, versiondir * "/Postprocessing");
push!(LOAD_PATH, versiondir * "/Utilities");

# Set Julia's PATH enviroment variable so that Exasim can call external programs
ENV["PATH"] = "/usr/local/bin:/usr/bin:/opt/local/bin:/bin:/usr/sbin:/sbin";
# Add more paths if neccesary
ENV["PATH"] =  ENV["PATH"] * ":/Applications/ParaView-5.8.1.app/Contents/MacOS";

print("==> Exasim " * version * " ...\n");

max_stars_repo_name: glwagner/Exasim

max_issues_repo_name: glwagner/Exasim

max_forks_repo_name: glwagner/Exasim


text:
module StdLib

export include_markdown, include_text

using ..RemarkStates
using SExpressions.Lists
using Documenter.Writers.HTMLWriter: mdconvert
using Documenter.Utilities.DOM: Node, TEXT
using Markdown

undomify(d::Node) = if d.name === TEXT
    d.text
else
    attributes = convert(list, [list(k, v) for (k, v) in d.attributes])
    undomified = map(undomify, d.nodes)
    if isempty(attributes)
        list(d.name, undomified...)
    else
        list(d.name, attributes, undomified...)
    end
end
rendermd(x) = undomify.(mdconvert(Markdown.parse(x)))

let state = nothing
    global function setstate!(st)
        state = st
    end
    global function include_markdown(filename)
        file = relativeto(state, filename)
        StdLib.rendermd(read(file, String))
    end
    global function include_text(filename)
        read(relativeto(state, filename), String)
    end
end

end


text:
module NordicEcon

include("dataloader.jl")
include("plotter.jl")

end

max_stars_repo_name: ordovician/NordicEcon

max_issues_repo_name: ordovician/NordicEcon

max_forks_repo_name: ordovician/NordicEcon


text:
using DiscreteChoice
using ForwardDiff
using Test

# Rosenbrock.
f(x) = (1.0 - x[1])^2 + 100.0 * (x[2] - x[1]^2)^2;
g(x) = ForwardDiff.gradient(f, x);
H(x) = ForwardDiff.hessian(f, x);

# Gradient descent.
@test optimize(f, g, zeros(2), 100000)[1] ≈ ones(2)

# Newton.
@test optimize(f, g, H, zeros(2))[1] ≈ ones(2)

# Quasi-Newton using inverse BFGS.
#@test optimize(f, g, inv_BFGS, zeros(2), true)[1] ≈ ones(2)

# Trust region using Steihaug-Toint's method and exact Hessian.
@test optimize(f, g, H, tcg, zeros(2))[1] ≈ ones(2)

# Trust region using Steihaug-Toint's method and approximative Hessian using BFGS.
@test optimize(f, g, BFGS!, tcg, zeros(2), true)[1] ≈ ones(2)

# Trust region using Steihaug-Toint's method and approximative Hessian using SR1.
@test optimize(f, g, SR1!, tcg, zeros(2), true)[1] ≈ ones(2)


text:
using BenchmarkTools
BenchmarkTools.DEFAULT_PARAMETERS.samples = 100

compute(n::Int)::Int = binomial(2n, n)

compute(2)

compute(20)

@benchmark compute(20)

max_stars_repo_name: Dynortice/Project-Euler

max_issues_repo_name: Dynortice/Project-Euler

max_forks_repo_name: Dynortice/Project-Euler


text:
using PyCall
using DFTK
import JSON
using LinearAlgebra

function ase_bulk(args...; kwargs...)
    ase_build = pyimport("ase.build")
    if first(args) == "SiO2"
        # Lattice parameters for α-quartz from
        # https://www.materialsproject.org/materials/mp-549166/
        # All lengths in Ångström (used by ASE)
        a = b = 4.98530291
        c = 5.47852971
        α = β = 90
        γ = 120

        silicon = [  # Fractional coordinates
            [0.000000  0.471138  0.833333],
            [0.471138  0.000000  0.166666],
            [0.528862  0.528862  0.500000],
        ]
        oxygen = [  # Fractional coordinates
            [0.148541  0.734505  0.046133],
            [0.265495  0.414036  0.712800],
            [0.585964  0.851459  0.379466],
            [0.734505  0.148541  0.953867],
            [0.414036  0.265495  0.287200],
            [0.851459  0.585964  0.620534],
        ]
        symbols = append!(fill("Si", 3), fill("O", 6))
        pyimport("ase").Atoms(symbols=symbols, cell=[a, b, c, α, β, γ], pbc=true,
                              scaled_positions=vcat(vcat(silicon, oxygen)...))
    elseif first(args) == "GaAs"
        # Lattice parameters for ZnS-type GaAs from
        # Haynes, William M., ed. (2011). CRC Handbook of Chemistry and Physics
        # (92nd ed.). CRC Press. ISBN 978-1439855119.
        a = 5.65315  # for zinc-blende cell
        ase_build.bulk("GaAs", "zincblende"; a=a, kwargs...)
    elseif first(args) == "AlN"
        a = b = 3.12858814
        c = 5.01695500
        α = β = 90
        γ = 120
        aluminium = [
            [1/3  2/3  0.99928700],
            [2/3  1/3  0.49928700],
        ]
        nitrogen = [
            [1/3  2/3  0.38071300],
            [2/3  1/3  0.88071300],
        ]
        symbols = ["Al", "Al", "N", "N"]
        pyimport("ase").Atoms(symbols=symbols, cell=[a, b, c, α, β, γ], pbc=true,
                              scaled_positions=vcat(vcat(aluminium, nitrogen)...))
    else
        ase_build.bulk(args...; kwargs...)
    end
end

function bulk(args...; kwargs...)
    atoms_ase = ase_bulk(args...; kwargs...)
    atoms = load_atoms(atoms_ase)
    atoms = [ElementPsp(el.symbol, psp=load_psp(el.symbol, functional="pbe")) => position
             for (el, position) in atoms]
    load_lattice(atoms_ase), atoms
end

function randomised_supercell(name; n_repeat=5)
    kwargs = ()
    if name == "Al"
        factor = 1/4  # 4 Al per repeat
        kwargs = (cubic=true, )
        amplitude = 1e-2
    elseif name == "GaAs"
        factor = 1    # One Ga and one As per repeat
        amplitude = 1e-2
    elseif name == "SiO2"
        factor = 1/3  # One Si and two O per repeat
        amplitude = 1e-2
    else
        error("Not implemented: $system")
    end
    @assert isinteger(n_repeat * factor)
    n_repeat = Int(n_repeat * factor)

    # Make supercell
    prim = ase_bulk(name; kwargs...)
    super = pyimport("ase.build").make_supercell(prim,
                                                 Array(Diagonal([1, 1, n_repeat])))

    # Introduce some random noise
    cartesian_pos = super.get_positions()
    ampl_Angstr = amplitude / DFTK.units.Å
    # noise = ampl_Angstr * randn(size(cartesian_pos))
    noise = 4ampl_Angstr .* (-0.5 .+ 0.5*rand(size(cartesian_pos)...))
    super.set_positions(cartesian_pos .+ noise)

    # Return DFTK datastructures
    atoms = [ElementPsp(el.symbol, psp=load_psp(el.symbol, functional="pbe")) => position
             for (el, position) in load_atoms(super)]
    load_lattice(super), atoms
end

function ase_surface(identifier; n_repeat=5, distance=nothing)
    # n_repeat gives the number of repeats in the surfaces
    # (i.e. n_repeat == 5 means 5 cells one material and 5 cells the other)
    ase_build = pyimport("ase.build")
    repeat = Array(Diagonal([1, 1, n_repeat]))  # repeat along z-dirn

    function get_distance(surface, user_set)  # in Å
        if isnothing(user_set)
            return maximum(abs, hcat(surface.cell...))
        else
            return user_set / DFTK.units.Å
        end
    end

    if identifier == "AlGaAs"
        gaas_110 = ase_lego("GaAs", (1, 1, 0), n_repeat=n_repeat)
        al_100 = ase_lego("Al", (1, 0, 0), n_repeat=n_repeat)
        return ase_build.stack(gaas_110, al_100 * (2, 1, 1), fix=0, maxstrain=1.7, distance=3.0)
    elseif identifier == "AlSiO2H"
        sio2h_110 = ase_lego("SiO2H", (1, 1, 0), cubic=false, n_repeat=n_repeat)
        ase_build.add_vacuum(sio2h_110, 1.0)
        sio2h_110.center()

        al_100 = ase_lego("Al", (1, 0, 0), n_repeat=n_repeat) * (2, 1, 1)
        ase_build.add_vacuum(al_100, 1.0)
        al_100.center()

        return ase_build.stack(sio2h_110, al_100, fix=0, maxstrain=1.6, distance=2.5)
    elseif identifier == "AlVac"
        al_100 = ase_lego("Al", (1, 0, 0), n_repeat=n_repeat)
        ase_build.add_vacuum(al_100, get_distance(al_100, distance))
        return al_100
    elseif identifier == "GaAsVac"
        gaas_110 = ase_lego("GaAs", (1, 1, 0), n_repeat=n_repeat)
        ase_build.add_vacuum(gaas_110, get_distance(gaas_110, distance))
        return gaas_110
    elseif identifier == "GaAsSiO2H"
        gaas_110 = ase_lego("GaAs", (1, 1, 0), n_repeat=n_repeat)
        sio2h_110 = ase_lego("SiO2H", (1, 1, 0), cubic=false, n_repeat=n_repeat)
        return ase_build.stack(gaas_110, sio2h_110, fix=0, maxstrain=0.7, distance=2.5)
    elseif identifier == "SiO2HVac"
        sio2h_110 = ase_lego("SiO2H", (1, 1, 0), cubic=false, n_repeat=n_repeat)
        ase_build.add_vacuum(sio2h_110, get_distance(sio2h_110, distance))
        return sio2h_110
    elseif identifier == "AlGaAsSiO2H"
        gaas_110  = ase_lego("GaAs", (1, 1, 0), n_repeat=n_repeat)
        sio2h_110 = ase_lego("SiO2H", (1, 1, 0), cubic=false, n_repeat=n_repeat)
        interm = ase_build.stack(gaas_110, sio2h_110, fix=0, maxstrain=0.7, distance=2.5)

        al_100 = ase_lego("Al", (1, 0, 0), n_repeat=n_repeat) * (2, 1, 1)
        return ase_build.stack(interm, al_100, fix=0, maxstrain=1.7, distance=2.5)
    else
        error("Unknown surface identifier: $identfier")
    end
end

function surface(args...; kwargs...)
    atoms_ase = ase_surface(args...; kwargs...)
    atoms = load_atoms(atoms_ase)
    atoms = [ElementPsp(el.symbol, psp=load_psp(el.symbol, functional="pbe")) => position
             for (el, position) in atoms]
    load_lattice(atoms_ase), atoms
end

function store_system(lattice, atoms, name, repeat)
    atnums = Int[]
    pspfiles = String[]
    positions = Vector{Float64}[]
    for (elem, pos) in atoms
        append!(atnums, fill(elem.Z, length(pos)))
        append!(pspfiles, fill(elem.psp.identifier, length(pos)))
        append!(positions, pos)
    end

    if !isdir(joinpath(@__DIR__, "systems"))
        mkdir(joinpath(@__DIR__, "systems"))
    end
    open(joinpath(@__DIR__, "systems/$(name)_$repeat.json"), "w") do fp
        save_dict = Dict(
            "lattice" => Array(lattice),
            "atnums" => atnums,
            "pspfiles" => pspfiles,
            "positions" => positions,
        )
        JSON.print(fp, save_dict)
    end
    nothing
end

function load_system(name, repeat)
    store = joinpath(@__DIR__, "systems/$(name)_$repeat.json")
    if !isfile(store)
        if name in ("Al", "SiO2", "GaAs")
            lattice, atoms = randomised_supercell(name, n_repeat=repeat)
        elseif name == "SiO2H"
            error("Sorry cannot build passivated SiO2")
        else
            lattice, atoms = surface(name, n_repeat=repeat)
        end
        store_system(lattice, atoms, name, repeat)
    end

    saved_dict = open(JSON.parse, store, "r")
    lattice = Mat3(hcat(saved_dict["lattice"]...))
    atoms = map(unique(saved_dict["atnums"])) do atnum
        pspid = saved_dict["pspfiles"][findfirst(atnum .== saved_dict["atnums"])]
        elem = ElementPsp(atnum, psp=load_psp(pspid))
        coords = Vec3.(saved_dict["positions"][atnum .== saved_dict["atnums"]])
        elem => coords
    end
    lattice, atoms
end

function ase_lego(system, miller; cubic=true, n_repeat=1)
    ase_build = pyimport("ase.build")
    ase_io = pyimport("ase.io")

    extra = ""
    if cubic
        extra = "_cubic"
    end
    mstr = prod(string, miller)

    store = joinpath(@__DIR__, "systems/ase_lego_$(system)$(extra)_$(mstr)_$(n_repeat).xyz")
    if !isfile(store)
        atoms = ase_bulk(system, cubic=cubic)

        factor = 1.0
        if system == "GaAs"
            factor = 0.5  # Two Ga and two As per layer
        elseif system == "SiO2"
            factor = 0.5  # Three Si and 6 O per layer
        elseif system == "Al"
            factor = 0.5  # 2 Al per layer
        else
            error("Not implemented: $system")
        end
        @assert isinteger(n_repeat * factor)
        n_repeat = Int(n_repeat * factor)
        atoms = ase_build.surface(atoms, miller, n_repeat, 0; periodic=true)
        ase_io.write(store, atoms)
    end

    ase_io.read(store)
end


text:
function FacesInNodes(Grid)
NumFacesPerNode=zeros(Int,Grid.NumNodes,1);
for iF=1:Grid.NumFaces
  Face=Grid.Faces[iF];
  for iN=1:size(Face.N,2)
    NumFacesPerNode[Face.N[iN]]=NumFacesPerNode[Face.N[iN]]+1;
  end
end
FacesPerNode=zeros(Int,Grid.NumNodes,maximum(NumFacesPerNode));
NumFacesPerNode=zeros(Int,Grid.NumNodes,1);
for iF=1:Grid.NumFaces
  Face=Grid.Faces[iF];
  for iN=1:size(Face.N,2)
    NumFacesPerNode[Face.N[iN]]=NumFacesPerNode[Face.N[iN]]+1;
    FacesPerNode[Face.N[iN],NumFacesPerNode[Face.N[iN]]]=iF;
  end
end
for iN=1:Grid.NumNodes
  Grid.Nodes[iN].F=FacesPerNode[iN,1:NumFacesPerNode[iN]];
end
return Grid
end


text:
import Random

function collect_distribution(model_creator, step_function, print_function, num_steps, num_trials, slice_size)
    for t in 1:num_trials
        trial_id = time_ns()

        Random.seed!(trial_id)

        logger = RunLogger(trial_id)
        model = model_creator(logger)

        for s in 1:convert(UInt64, floor(num_steps / slice_size))
            step_function(model)
            print_function(t, s*slice_size, model)
        end

        save_log(logger)
    end
end

max_stars_repo_name: tochsner/ALifeBenchmark

max_issues_repo_name: tochsner/ALifeBenchmark

max_forks_repo_name: tochsner/ALifeBenchmark


text:
include(joinpath(ROOT,"..","src_PowSysMod", "PowSysMod_body.jl"))

if length(ARGS) != 2
  error("Expected input:\n  julia compare_dats.jl file1.dat file2.dat")
end

file1, file2 = ARGS

maxerror, errors = compare_dat(file1, file2, 1e-10)

count = 0
for key in keys(errors)
  count += length(errors[key])
end

println("Max error on coeff: $maxerror")
println("Nb of coeffs with error > 1e-10: $(count)")

max_stars_repo_name: Gimaju/GOC

max_issues_repo_name: Gimaju/GOC

max_forks_repo_name: Gimaju/GOC


text:
using Test, Workbooks

out = plusTwo(3)

@test out == 5


text:
using ExamplePkg
using Test

@testset "ExamplePkg.jl" begin
    # Write your own tests here.
end