text:
# Usage:
#     julia benchmark/run_benchmarks.jl
using BenchmarkTools
using ImageCore
using ImageMorphology
using ImageBinarization
using ImageTransformations
using TestImages

on_CI = haskey(ENV, "GITHUB_ACTIONS")

cameraman = Gray{N0f8}.(testimage("cameraman"))
blobs = binarize(Gray.(testimage("blobs")), Otsu()) .> 0.5

tst_sizes = on_CI ? (64, ) : (256, 512)
tst_types = (Gray{N0f8}, Gray{Float32})

const SUITE = BenchmarkGroup()

SUITE["extreme_filter"] = BenchmarkGroup()
let grp = SUITE["extreme_filter"]
    for T in [tst_types..., Int]
        grp[T] = BenchmarkGroup()
        for sz in tst_sizes
            tst_img = rand(T, sz, sz)

            grp[T]["$sz×$sz"] = BenchmarkGroup()
            se = strel_diamond((3, 3))
            grp[T]["$sz×$sz"]["r1_diamond"] = @benchmarkable extreme_filter(max, $tst_img, $se)
            se = centered(collect(se))
            grp[T]["$sz×$sz"]["r1_generic"] = @benchmarkable extreme_filter(max, $tst_img, $se)
            se = strel_diamond((11, 11))
            grp[T]["$sz×$sz"]["r5_diamond"] = @benchmarkable extreme_filter(max, $tst_img, $se)
            se = centered(collect(se))
            grp[T]["$sz×$sz"]["r5_generic"] = @benchmarkable extreme_filter(max, $tst_img, $se)
        end
    end
end
let grp = SUITE["extreme_filter"]
    T = Bool
    grp[T] = BenchmarkGroup()
    for sz in tst_sizes
        grp[T]["$sz×$sz"] = BenchmarkGroup()
        for (cname, cr) in [("worst", 0.95), ("best", 0.05), ("random", 0.5)]
            tst_img = fill(zero(T), sz, sz)
            tst_img[rand(sz, sz) .>= cr] .= true

            se = strel_diamond((3, 3))
            grp[T]["$sz×$sz"]["r1_diamond_$cname"] = @benchmarkable extreme_filter(max, $tst_img, $se)
            se = centered(collect(se))
            grp[T]["$sz×$sz"]["r1_bool_$cname"] = @benchmarkable extreme_filter(max, $tst_img, $se)

            se = strel_diamond((11, 11))
            grp[T]["$sz×$sz"]["r5_diamond_$cname"] = @benchmarkable extreme_filter(max, $tst_img, $se)
            se = centered(collect(se))
            grp[T]["$sz×$sz"]["r5_bool_$cname"] = @benchmarkable extreme_filter(max, $tst_img, $se)
        end
    end
end

SUITE["dilatation_and_erosion"] = BenchmarkGroup()
let grp = SUITE["dilatation_and_erosion"]
    grp["erode"] = BenchmarkGroup()
    grp["opening"] = BenchmarkGroup()
    for T in tst_types
        grp["erode"][T] = BenchmarkGroup()
        grp["opening"][T] = BenchmarkGroup()
        for sz in tst_sizes
            tst_img = T.(imresize(T.(cameraman), (sz, sz)))
            grp["erode"][T]["$sz×$sz"] = @benchmarkable erode($tst_img)
            grp["opening"][T]["$sz×$sz"] = @benchmarkable opening($tst_img)
        end
    end
end

SUITE["connected"] = BenchmarkGroup()
let grp = SUITE["connected"]
    grp["label_components"] = @benchmarkable label_components($blobs)
end

SUITE["Maxtree"] = BenchmarkGroup()
let grp = SUITE["Maxtree"]
    grp["area_opening"] = BenchmarkGroup()
    for sz in tst_sizes
        tst_img = (imresize((cameraman), (sz, sz)))
        B = similar(tst_img)
        grp["area_opening"]["$sz×$sz"] = @benchmarkable area_opening($B, min_area=50)
    end
end

SUITE["convexhull"] = BenchmarkGroup()
let grp = SUITE["convexhull"]
    grp["convexhull"] = @benchmarkable convexhull($blobs)
end

SUITE["feature_transform"] = BenchmarkGroup()
let grp = SUITE["feature_transform"]
    grp["feature_transform"] = @benchmarkable feature_transform($blobs)
end


text:
#=
    Helper functions
=#
"""
    get_column(model::LinQuadOptimizer, index::MOI.VariableIndex)

Return the column of the variable `index` in `model`.
"""
function get_column(model::LinQuadOptimizer, index::VarInd)
    return model.variable_mapping[index]
end

"""
    get_column(model::LinQuadOptimizer, variable::MOI.SingleVariable)

Return the column of the variable inside a `MOI.SingleVariable` function in
`model`.
"""
function get_column(model::LinQuadOptimizer, variable::SinVar)
    return get_column(model, variable.variable)
end

#=
    Get variable names
=#
MOI.supports(::LinQuadOptimizer, ::MOI.VariableName, ::Type{VarInd}) = true
function MOI.get(model::LinQuadOptimizer, ::MOI.VariableName, index::VarInd)
    return get(model.variable_to_name, index, "")
end

function MOI.set(model::LinQuadOptimizer,
                 ::MOI.VariableName, index::VarInd, name::String)
    if name != ""
        model.variable_to_name[index] = name
    else
        # delete! does nothing if no index exists.
        delete!(model.variable_to_name, index)
    end
    model.name_to_variable = nothing
    return
end

#=
    Get variable by name. Rebuild the lazy reverse-lookup if required.
=#

function _rebuild_name_to_variable(model::LinQuadOptimizer)
    model.name_to_variable = Dict{String, MOI.VariableIndex}()
    for (var, v_name) in model.variable_to_name
        v_name == "" && continue
        if haskey(model.name_to_variable, v_name)
            duplicate_var = model.name_to_variable[v_name]
            model.name_to_variable = nothing
            error("Variable $(var) and $(duplicate_var) have the same ",
                  "name: $(v_name)")
        else
            model.name_to_variable[v_name] = var
        end
    end
    return
end

function MOI.get(model::LinQuadOptimizer, ::Type{MOI.VariableIndex}, name::String)
    if model.name_to_variable === nothing
        _rebuild_name_to_variable(model)
    end
    return get(model.name_to_variable, name, nothing)
end

#=
    Get number of variables
=#

function MOI.get(model::LinQuadOptimizer, ::MOI.NumberOfVariables)
    # TODO(odow): work out why this is not the same as
    # length(model.variable_references)
    return get_number_variables(model)
end

#=
    List of Variable References
=#

function MOI.get(model::LinQuadOptimizer, ::MOI.ListOfVariableIndices)
    return model.variable_references
end

#=
    Add a single variable
=#

function MOI.add_variable(model::LinQuadOptimizer)
    add_variables!(model, 1)
    # assumes we add columns linearly
    model.last_variable_reference += 1
    index = VarInd(model.last_variable_reference)
    model.variable_mapping[index] = MOI.get(model, MOI.NumberOfVariables())
    push!(model.variable_references, index)
    push!(model.variable_primal_solution, NaN)
    push!(model.variable_dual_solution, NaN)
    model.variable_type[index] = CONTINUOUS
    return index
end

#=
    Add multiple variables
=#

function MOI.add_variables(model::LinQuadOptimizer, number_to_add::Int)
    previous_vars = MOI.get(model, MOI.NumberOfVariables())
    add_variables!(model, number_to_add)
    variable_indices = VarInd[]
    sizehint!(variable_indices, number_to_add)
    for i in 1:number_to_add
        # assumes we add columns linearly
        model.last_variable_reference += 1
        index = VarInd(model.last_variable_reference)
        push!(variable_indices, index)
        # m.last_variable_reference might not be the column if we have
        # deleted variables.
        model.variable_mapping[index] = previous_vars + i
        push!(model.variable_references, index)
        push!(model.variable_primal_solution, NaN)
        push!(model.variable_dual_solution, NaN)
        model.variable_type[index] = CONTINUOUS
    end
    return variable_indices
end

#=
    Check if reference is valid
=#

function MOI.is_valid(model::LinQuadOptimizer, index::VarInd)
    if haskey(model.variable_mapping, index)
        column = get_column(model, index)
        if 1 <= column <= MOI.get(model, MOI.NumberOfVariables())
            return true
        end
    end
    return false
end

#=
    Delete a variable
=#

function MOI.delete(model::LinQuadOptimizer, index::VarInd)
    __assert_valid__(model, index)
    column = get_column(model, index)
    delete_variables!(model, column, column)
    # delete from problem
    deleteat!(model.variable_references, column)
    delete!(model.variable_type, index)
    deleteat!(model.variable_primal_solution, column)
    deleteat!(model.variable_dual_solution, column)
    deleteref!(model.variable_mapping, column, index)
    # delete name if not ""
    if haskey(model.variable_to_name, index)
        model.name_to_variable = nothing
        delete!(model.variable_to_name, index)
    end
    # Delete any bounds (deleting from a dict without the key does nothing).
    constraint_map = cmap(model)
    delete_from_dictionary_by_value(constraint_map.upper_bound, index)
    delete_from_dictionary_by_value(constraint_map.lower_bound, index)
    delete_from_dictionary_by_value(constraint_map.fixed_bound, index)
    delete_from_dictionary_by_value(constraint_map.interval_bound, index)
    return
end

"""
    delete_from_dictionary_by_value(dict::Dict{K, V}, index::V)

Delete the entries of a dictionary where the value of a `key=>value` pair is
equal to `index`.

This is used to remove bounds when deleting a variable.
"""
function delete_from_dictionary_by_value(dict::Dict{K, V}, index::V) where {K, V}
    for (key, value) in dict
        if value == index
            delete!(dict, key)
        end
    end
end

#=
    MIP starts
=#
function MOI.supports(
        ::LinQuadOptimizer, ::MOI.VariablePrimalStart, ::Type{VarInd})
    return false
end

function MOI.set(model::LinQuadOptimizer, ::MOI.VariablePrimalStart,
                  indices::Vector{VarInd}, values::Vector{Float64})
    if !MOI.supports(model, MOI.VariablePrimalStart(), MOI.VariableIndex)
        throw(MOI.UnsupportedAttribute(MOI.VariablePrimalStart()))
    end
    add_mip_starts!(model, get_column.(Ref(model), indices), values)
end

function MOI.set(model::LinQuadOptimizer, ::MOI.VariablePrimalStart,
                  index::VarInd, value::Float64)
    MOI.set(model, MOI.VariablePrimalStart(), [index], [value])
end

ext: jl


text:
"""
    core_periphery_deg(g)

Compute the degree-based core-periphery for graph `g`. Return the vertex
assignments (`1` for core and `2` for periphery) for each node in `g`.

References:
    [Lip](http://arxiv.org/abs/1102.5511))
"""
function core_periphery_deg end
@traitfn function core_periphery_deg(g::::(!IsDirected))
    degs = degree(g)
    p = sortperm(degs, rev=true)
    s = sum(degs) / 2.
    sbest = +Inf
    kbest = 0
    for k = 1:nv(g)-1
        s = s + k  - 1 - degree(g, p[k])
        if s < sbest
            sbest = s
            kbest = k
        end
    end
    c = 2 + zeros(Int, nv(g))
    c[p[1:kbest]] = 1
    c
end

max_stars_repo_name: AsileBcd/JohnsonCycles

max_issues_repo_name: AsileBcd/JohnsonCycles

max_forks_repo_name: AsileBcd/JohnsonCycles


text:
macro test_throws_unwrap(terr, args...)
    _test_throws_unwrap(terr, args...)
end
function _test_throws_unwrap(terr, ex; to_match=[])
    @gensym rerr
    match_expr = Expr(:block)
    for m in to_match
        @assert m.head == :(=)
        push!(match_expr.args, :(@test $rerr.$(m.args[1]) == $(m.args[2])))
    end
    quote
        $rerr = try
            $(esc(ex))
        catch err
            Dagger.Sch.unwrap_nested_exception(err)
        end
        @test $rerr isa $terr
        $match_expr
    end
end
function _test_throws_unwrap(terr, args...)
    ex = last(args)
    to_match = args[1:end-1]
    _test_throws_unwrap(terr, ex; to_match=to_match)
end

# NOTE: based on test/pkg.jl::capture_stdout, but doesn't discard exceptions
macro grab_output(ex)
    quote
        mktemp() do fname, fout
            ret = nothing
            open(fname, "w") do fout
                redirect_stderr(fout) do
                    ret = $(esc(ex))
                end
                flush(fout)
            end
            ret, read(fname, String)
        end
    end
end


text:
@with_kw Jet{Tf, Tdf, Td2f, Td3f, Tj, Td}
	δ::Td = 1e-8
	F::Tf
	dF::Tdf = (x,p,dx) -> (F(x .+ δ .* dx, p) - F(x, p)) / δ
	d2F::Td2f = (x,p,dx1,dx2) -> (dF(x .+ δ .* dx2, p, dx1) - dF(x, p, dx1)) / δ
	d3F::Td3f = (x,p,dx1,dx2) -> (d2F(x .+ δ .* dx3, p, dx1, dx2) - d2F(x, p, dx1, dx2)) / δ
	J::Tj = nothing
end

Jet(f) = Jet(F = f)
(j::Jet)(x, p) = j.F(x, p)
(j::Jet)(::Val{Jacobian}, x, p) = j.J(x, p)
(j::Jet)(x, p, dx) = j.dF(x, p, dx)
(j::Jet)(x, p, dx1, dx2) = j.d2F(x, p, dx1, dx2)
(j::Jet)(x, p, dx1, dx2, dx3) = j.d2F(x, p, dx1, dx2, dx3)


text:
using FileIO
using FilePathsBase
using Test
using UUIDs
using CodecZlib

Threads.nthreads() <= 1 && @info "Threads.nthreads() = $(Threads.nthreads()), multithread tests will be disabled"

# Both FileIO and FilePathsBase export filename, but we only want the FileIO definition.
using FileIO: filename

struct MimeSaveTestType
end

@testset "FileIO" begin
    # This threaded test should be put before `CSVFiles` is loaded
    if VERSION >= v"1.3"
        # FIXME: threaded file io is still somehow broken in old Julia versions
        include("threaded_loading.jl")
    end

    include("query.jl")
    include("loadsave.jl")
    include("error_handling.jl")
    include("test_mimesave.jl")
    include("integration.jl")
end


text:
abstract type AbstractLShapedExecution end
# Execution API #
# ------------------------------------------------------------
num_thetas(lshaped::AbstractLShaped) = num_thetas(lshaped, lshaped.execution)
initialize_subproblems!(lshaped::AbstractLShaped, scenarioproblems::AbstractScenarioProblems) = initialize_subproblems!(lshaped.execution, scenarioproblems)
finish_initilization!(lshaped::AbstractLShaped) = finish_initilization!(lshaped, lshaped.execution)
restore_subproblems!(lshaped::AbstractLShaped) = restore_subproblems!(lshaped, lshaped.execution)
solve_master!(lshaped::AbstractLShaped) = solve_master!(lshaped, lshaped.execution)
iterate!(lshaped::AbstractLShaped) = iterate!(lshaped, lshaped.execution)
start_workers!(lshaped::AbstractLShaped) = start_workers!(lshaped, lshaped.execution)
close_workers!(lshaped::AbstractLShaped) = close_workers!(lshaped, lshaped.execution)
resolve_subproblems!(lshaped::AbstractLShaped) = resolve_subproblems!(lshaped, lshaped.execution)
current_decision(lshaped::AbstractLShaped) = current_decision(lshaped, lshaped.execution)
calculate_objective_value(lshaped::AbstractLShaped) = calculate_objective_value(lshaped, lshaped.execution)
timestamp(lshaped::AbstractLShaped) = timestamp(lshaped, lshaped.execution)
incumbent_decision(lshaped::AbstractLShaped, t::Integer, regularizer::AbstractRegularization) = incumbent_decision(lshaped, t, regularizer, lshaped.execution)
incumbent_objective(lshaped::AbstractLShaped, t::Integer, regularizer::AbstractRegularization) = incumbent_objective(lshaped, t, regularizer, lshaped.execution)
incumbent_trustregion(lshaped::AbstractLShaped, t::Integer, regularizer::AbstractRegularization) = incumbent_trustregion(lshaped, t, regularizer, lshaped.execution)
readd_cuts!(lshaped::AbstractLShaped, consolidation::Consolidation) = readd_cuts!(lshaped, consolidation, lshaped.execution)
subobjectives(lshaped::AbstractLShaped) = subobjectives(lshaped, lshaped.execution)
set_subobjectives(lshaped::AbstractLShaped, Qs::AbstractVector) = set_subobjectives(lshaped, Qs, lshaped.execution)
model_objectives(lshaped::AbstractLShaped) = model_objectives(lshaped, lshaped.execution)
set_model_objectives(lshaped::AbstractLShaped, Qs::AbstractVector) = set_model_objectives(lshaped, Qs, lshaped.execution)
fill_submodels!(lshaped::AbstractLShaped, scenarioproblems::AbstractScenarioProblems) = fill_submodels!(lshaped, scenarioproblems, lshaped.execution)

# ------------------------------------------------------------
include("common.jl")
include("serial.jl")
include("distributed.jl")
include("synchronous.jl")
include("asynchronous.jl")


text:
__precompile__(true)
module DFTShims
using Unitful
using AxisArrays

if Pkg.installed("AxisArrays") <= v"0.2.1"
    include("AxisArrays.v0.1.x.jl")
end

include("UnitfulHartree.jl")
using .UnitfulHartree
const UH = UnitfulHartree
export UnitfulHartree, Unitful, @u_str

include("Dispatch.jl")
using .Dispatch
using .Dispatch: DFTArray, DFTAxisArray
export DFTArray, DFTAxisArray, Dispatch

include("Traits.jl")
using .Traits: ColinearSpin, ColinearSpinFirst, ColinearSpinLast, SpinDegenerate,
               SpinCategory, FunctionalCategory, GGA, LDA, components, is_spin_polarized,
               has_axis, concretize_type, SpinAware, standard_name
export ColinearSpin, ColinearSpinFirst, ColinearSpinLast, SpinDegenerate, SpinCategory,
       FunctionalCategory, GGA, LDA, components, is_spin_polarized, has_axis,
       concretize_type, SpinAware, standard_name

include("ConstantArrays.jl")
using .ConstantArrays

include("ArrayInitialization.jl")
using .ArrayInitialization: wrap
export wrap

end # module

max_stars_repo_name: mdavezac/DFTShims

max_issues_repo_name: mdavezac/DFTShims

max_forks_repo_name: mdavezac/DFTShims


text:
#####
##### PrescribedVelocityFields
#####

using Oceananigans.Grids: Center, Face
using Oceananigans.Fields: FunctionField

import Oceananigans.BoundaryConditions: fill_halo_regions!
import Oceananigans.Models.IncompressibleModels: extract_boundary_conditions

using Adapt

struct PrescribedVelocityFields{U, V, W, P}
    u :: U
    v :: V
    w :: W
    parameters :: P
end

@inline Base.getindex(U::PrescribedVelocityFields, i) = getindex((u=U.u, v=U.v, w=U.w), i)

zerofunc(x, y, z) = 0

"""
    PrescribedVelocityFields(; u=zerofunc, v=zerofunc, w=zerofunc, parameters=nothing)

Builds `PrescribedVelocityFields` with prescribed functions `u`, `v`, and `w`.

If `isnothing(parameters)`, then `u, v, w` are called with the signature

    `u(x, y, z, t) = # something interesting`

If `!isnothing(parameters)`, then `u, v, w` are called with the signature

    `u(x, y, z, t, parameters) = # something parameterized and interesting`

In the constructor for `HydrostaticFreeSurfaceModel`, the functions `u, v, w` are wrapped
in `FunctionField` and associated with the model's `grid` and `clock`.
"""
PrescribedVelocityFields(; u=zerofunc, v=zerofunc, w=zerofunc, parameters=nothing) =
    PrescribedVelocityFields(u, v, w, parameters)

function HydrostaticFreeSurfaceVelocityFields(velocities::PrescribedVelocityFields,
                                              arch, grid, clock, bcs)

    u = FunctionField{Face, Center, Center}(velocities.u, grid; clock=clock, parameters=velocities.parameters)
    v = FunctionField{Center, Face, Center}(velocities.v, grid; clock=clock, parameters=velocities.parameters)
    w = FunctionField{Center, Center, Face}(velocities.w, grid; clock=clock, parameters=velocities.parameters)

    return PrescribedVelocityFields(u, v, w, velocities.parameters)
end

@inline fill_halo_regions!(::PrescribedVelocityFields, args...) = nothing
@inline fill_halo_regions!(::FunctionField, args...) = nothing

ab2_step_free_surface!(::Nothing, args...) = nothing
extract_boundary_conditions(::PrescribedVelocityFields) = NamedTuple()
FreeSurface(free_surface::ExplicitFreeSurface{Nothing}, ::PrescribedVelocityFields, arch, grid) = nothing
validate_velocity_boundary_conditions(::PrescribedVelocityFields) = nothing
compute_w_from_continuity!(::PrescribedVelocityFields, args...) = nothing
HorizontalVelocityFields(::PrescribedVelocityFields, arch, grid) = nothing, nothing
hydrostatic_prognostic_fields(::PrescribedVelocityFields, free_surface, tracers) = tracers
calculate_hydrostatic_momentum_tendencies!(tendencies, ::PrescribedVelocityFields, args...) = []

Adapt.adapt_structure(to, velocities::PrescribedVelocityFields) =
    PrescribedVelocityFields(Adapt.adapt(to, velocities.u),
                             Adapt.adapt(to, velocities.v),
                             Adapt.adapt(to, velocities.w),
                             nothing)


text:
module Arcgisscripting

using PyCall
@pyimport arcgisscripting as ags

struct GeoProcessorObject
	o::PyObject
end
create() = ags.create() |> GeoProcessorObject

for f in (:CheckOutExtension, :SingleOutputAlgebra_sa)
	@eval ($f)(g::GeoProcessorObject, args...) = g.o[f](args...)
end

end # module


text:
# Test the scripts

module TestHeatEquationExample
    mktempdir() do dir
        cd(dir) do
            include(joinpath(@__DIR__, "../docs/src/literate/heat_equation.jl"))
        end
    end
end

module TestIncompressibleElasticity
    mktempdir() do dir
        cd(dir) do
            include(joinpath(@__DIR__, "../docs/src/literate/incompressible_elasticity.jl"))
        end
    end
end

module TestHyperElasticity
    mktempdir() do dir
        cd(dir) do
            include(joinpath(@__DIR__, "../docs/src/literate/hyperelasticity.jl"))
        end
    end
end

module TestQuasiIncompressibleHyperElasticity
    mktempdir() do dir
        cd(dir) do
            include(joinpath(@__DIR__, "../docs/src/literate/quasi_incompressible_hyperelasticity.jl"))
        end
    end
end

# module TestNavierStokesDiffeqIntegration
#     mktempdir() do dir
#         cd(dir) do
#             include(joinpath(@__DIR__, "../docs/src/literate/ns_vs_diffeq.jl"))
#         end
#     end
# end

module TestComputationalHomogenization
    # Add this unregistered package here
    import Pkg; Pkg.add(Pkg.PackageSpec(url = "https://github.com/Ferrite-FEM/FerriteGmsh.jl"))
    include(joinpath(@__DIR__, "../docs/download_resources.jl"))
    mktempdir() do dir
        cd(dir) do
            cp(joinpath(@__DIR__, "../docs/src/examples/periodic-rve.msh"),
               joinpath(dir, "periodic-rve.msh")
            )
            include(joinpath(@__DIR__, "../docs/src/literate/computational_homogenization.jl"))
        end
    end
end


text:
#=
:(+) => "f32.add", # these consume 2 args
:(-) => "f32.sub",
:(*) => "f32.mul",
:(/) => "f32.div",
:(==) => "f32.eq",
:(!=) => "f32.ne",
:(<) => "f32.lt",
:(>) => "f32.gt",
:(<=) => "f32.le",
:(>=) => "f32.ge",
:(min) => "f32.min",
:(max) => "f32.max",
:(copysign) => "f32.copysign",

:(abs) => "f32.abs", # these consume 1 arg
:(ceil) => "f32.ceil",
:(floor) => "f32.floor",
:(trunc) => "f32.trunc",
:(round) => "f32.nearest",
:(sqrt) => "f32.sqrt",
:(float) => "f32.convert_i32_s",
:(Int) => "i32.trunc_f32_s",
# :(^) => ["call \$pow",2],
# ""=>["f32.neg",1],
# ""=>["f32.load",1],
# ""=>["f32.store",2],
# ""=>["f32.const",1],
=#

add(a, b) = a + b
sub(a, b) = a - b
mul(a, b) = a * b
_div(a, b) = a / b
eq(a, b) = a == b
ne(a, b) = a != b
lt(a, b) = a < b
gt(a, b) = a > b
le(a, b) = a <= b
ge(a, b) = a >= b
_min(a, b) = min(a, b)
_max(a, b) = max(a, b)
_copysign(a,b) = copysign(a,b)

_abs(a) = abs(a)
_ceil(a) = ceil(a)
_floor(a) = floor(a)
_trunc(a) = trunc(a)
_round(a) = round(a)
_sqrt(a) = sqrt(a)
_float(c) = float(c)
_Int(a) = Int(a)

function ops(a, b)
  add(a, b)
  sub(a, b)
  mul(a, b)
  _div(a, b)
  eq(a, b)
  ne(a, b)
  lt(a, b)
  gt(a, b)
  le(a, b)
  ge(a, b)
  _min(a,b)
  _max(a,b)
  _copysign(a,b)

  _abs(a)
  _ceil(a)
  _floor(a)
  _trunc(a)
  _round(a)
  _sqrt(a)
  _float(3)
  _Int(a)

  return nothing
end

function exports()
  ops(2.0, 3.0) #float ops
  #ops(2, 3) #int ops
end

exports()


text:
# Portions translated from SLICOT-Reference distribution
# Copyright (c) 2002-2020 NICONET e.V.
function run_mb03qd(datfile, io=stdout)
    NIN = 5
    NOUT = 6
    NMAX = 10
    LDA = NMAX
    LDU = NMAX
    LDWORK = 3*NMAX
    A = Array{Float64,2}(undef, LDA,NMAX)
    U = Array{Float64,2}(undef, LDU,NMAX)
    DWORK = Array{Float64,1}(undef, LDWORK)
    WI = Array{Float64,1}(undef, NMAX)
    WR = Array{Float64,1}(undef, NMAX)
    BWORK = Array{Bool,1}(undef, NMAX)
    f = open(datfile,"r")
    readline(f)
    vs = split(readline(f))
    N = parse(BlasInt, vs[1])
    NLOW = parse(BlasInt, vs[2])
    NSUP = parse(BlasInt, vs[3])
    ALPHA = parse(Float64, replace(vs[4],'D'=>'E'))
    DICO =  vs[5][1]
    STDOM =  vs[6][1]
    JOBU =  vs[7][1]
    if ( N<0 || N>NMAX )
        @error "invalid N=$N"
    else

    vs = String[]
    _isz,_jsz = (N,N)
    while length(vs) < _isz*_jsz
        append!(vs, replace.(split(readline(f)),'D'=>'E'))
    end
    for i in 1:_isz
       _i0 = (i-1)*_jsz
       A[i,1:_jsz] .= parsex.(Float64, vs[_i0+1:_i0+_jsz])
    end
# interp call 1

    #FOREIGN.dgees!( 'Vectors', 'Not sorted', SELECT, N, A, LDA, NDIM, WR, WI, U, LDU, DWORK, LDWORK, BWORK, INFO )
    Aschur = schur(A[1:N,1:N])
    WR[1:N] .= real.(Aschur.values)
        WI[1:N] .= imag.(Aschur.values)
        U[1:N,1:N] .= Aschur.Z
        A[1:N,1:N] .= Aschur.T

    println(io, "orig T:")
    _nc = N
    _nr = N
    show(io, "text/plain", A[1:_nr,1:_nc])
    println(io)
# interp call 2

    NDIM, INFO = SLICOT.mb03qd!(DICO, STDOM, JOBU, N, NLOW, NSUP, ALPHA, A, U)
    println(io, "NDIM = $NDIM")
    @test INFO == 0
    INFO == 0 || return

if ( INFO!=0 )
else
# interp output 1
    println(io, "A:")
    _nc = N
    _nr = N
    show(io, "text/plain", A[1:_nr,1:_nc])
    println(io)

# interp output 2
    println(io, "U:")
    _nc = N
    _nr = N
    show(io, "text/plain", U[1:_nr,1:_nc])
    println(io)

end # if
end # if
    close(f)
end # run_mb03qd()


text:
using DataFrames
using CategoricalArrays
using PooledArrays
using BenchmarkTools
using Random

Random.seed!(1)

grouping_benchmarks = BenchmarkGroup()

# `refpool`/`refarray` optimized grouping method
refpool_benchmarks = grouping_benchmarks["refpool"] = BenchmarkGroup()

for k in (10, 10_000), n in (100, 100_000, 10_000_000)
    for x in (PooledArray(rand(1:k, n)),
              CategoricalArray(rand(1:n, 10_000_000)),
              PooledArray(rand([missing; 1:n], 10_000_000)),
              CategoricalArray(rand([missing; 1:n], 10_000_000)))
        df = DataFrame(x=x)

        refpool_benchmarks[k, n, nameof(typeof(x)), "skipmissing=false"] =
            @benchmarkable groupby($df, :x)

        # Skipping missing values
        refpool_benchmarks[k, n, nameof(typeof(x)), "skipmissing=true"] =
            @benchmarkable groupby($df, :x, skipmissing=true)

        # Empty group which requires adjusting group indices
        replace!(df.x, 5 => 6)
        refpool_benchmarks[k, n, nameof(typeof(x)), "empty group"] =
            @benchmarkable groupby($df, :x)
    end
end

# If a cache of tuned parameters already exists, use it, otherwise, tune and cache
# the benchmark parameters. Reusing cached parameters is faster and more reliable
# than re-tuning `suite` every time the file is included.
paramspath = joinpath(dirname(@__FILE__), "params.json")

if isfile(paramspath)
    loadparams!(grouping_benchmarks, BenchmarkTools.load(paramspath)[1], :evals);
else
    tune!(grouping_benchmarks)
    BenchmarkTools.save(paramspath, params(grouping_benchmarks));
end

grouping_results = run(grouping_benchmarks, verbose=true)
# using Serialization
# serialize("grouping_results.jls", grouping_results)
# leaves(judge(median(grouping_results1), median(grouping_results2)))
# leaves(regressions(judge(median(grouping_results1), median(grouping_results2))))
# leaves(improvements(judge(median(grouping_results1), median(grouping_results2))))


text:
using Test, Dates, ExpiringCaches

ExpiringCaches.@cacheable Dates.Second(3) function foo(arg1::Int, arg2::String)::Float64
    sleep(2)
    return arg1 / length(arg2)
end

@testset "ExpiringCaches" begin

cache = ExpiringCaches.Cache{Int, Int}(Dates.Second(5))
@test length(cache) == 0
@test isempty(cache)

@test get(cache, 1, 2) == 2
@test isempty(cache)

@test get!(cache, 1, 2) == 2
@test !isempty(cache)

for (k, v) in cache
    @test v == 2
end

sleep(5)

# test that key isn't used after it expires
@test get(cache, 1, 3) == 3

@test get!(()->4, cache, 1) == 4

@test isempty(delete!(cache, 1))
cache[1] = 5
@test !isempty(cache)
@test isempty(empty!(cache))

@test foo(1, "ffff") == 0.25
tm = @elapsed foo(1, "ffff")
@test tm < 2 # test that normal function body wasn't executed
sleep(3)
tm = @elapsed foo(1, "ffff")
@test tm > 2 # test that normal function body was executed

end


text:
## SegmentedVector
const VectorSegment{T} = SubArray{T,1,Array{T, 1},Tuple{UnitRange{Int}},true} # type of a n:m view into a Vector

"""
$(TYPEDEF)

`SegmentedVector` is an `AbstractVector` backed by another `AbstractVector` (its parent), which additionally stores an [`IndexDict`](@ref)
containing views into the parent. Together, these views cover the parent.

# Examples

```julia-repl
julia> x = [1., 2., 3., 4.]
4-element Array{Float64,1}:
 1.0
 2.0
 3.0
 4.0

julia> viewlength(i) = 2
viewlength (generic function with 1 method)

julia> xseg = SegmentedVector{Int}(x, 1 : 2, viewlength)
4-element RigidBodyDynamics.CustomCollections.SegmentedVector{Int64,Float64,Base.OneTo{Int64},Array{Float64,1}}:
 1.0
 2.0
 3.0
 4.0

julia> segments(xseg)[1]
2-element SubArray{Float64,1,Array{Float64,1},Tuple{UnitRange{Int64}},true}:
 1.0
 2.0

julia> yseg = similar(xseg, Int32); yseg .= 1 : 4 # same view ranges, different element type
4-element RigidBodyDynamics.CustomCollections.SegmentedVector{Int64,Int32,Base.OneTo{Int64},Array{Int32,1}}:
 1
 2
 3
 4

julia> segments(yseg)[2]
2-element SubArray{Int32,1,Array{Int32,1},Tuple{UnitRange{Int64}},true}:
 3
 4
```
"""
struct SegmentedVector{K, T, KeyRange<:AbstractRange{K}, P<:AbstractVector{T}} <: AbstractVector{T}
    parent::P
    segments::IndexDict{K, KeyRange, VectorSegment{T}}

    function SegmentedVector{K, T, KeyRange, P}(p::P, segments::IndexDict{K, KeyRange, VectorSegment{T}}) where {K, T, KeyRange<:AbstractRange{K}, P<:AbstractVector{T}}
        @boundscheck begin
            firstsegment = true
            start = 0
            l = 0
            for segment in values(segments)
                parent(segment) === parent(p) || error()
                indices = first(parentindices(segment))
                if firstsegment
                    start = first(indices)
                    firstsegment = false
                else
                    first(indices) === start || error()
                end
                start = last(indices) + 1
                l += length(indices)
            end
            l == length(p) || error("Segments do not cover input data.")
        end
        new{K, T, KeyRange, P}(p, segments)
    end
end

function SegmentedVector(p::P, segments::IndexDict{K, KeyRange, VectorSegment{T}}) where {K, T, KeyRange<:AbstractRange{K}, P<:AbstractVector{T}}
    SegmentedVector{K, T, KeyRange, P}(p, segments)
end

function SegmentedVector{K, T, KeyRange}(parent::P, keys, viewlengthfun) where {K, T, KeyRange<:AbstractRange{K}, P<:AbstractVector{T}}
    views = Vector{Pair{K, VectorSegment{T}}}()
    start = 1
    for key in keys
        stop = start[] + viewlengthfun(key) - 1
        push!(views, convert(K, key) => view(parent, start : stop))
        start = stop + 1
    end
    SegmentedVector{K, T, KeyRange, P}(parent, IndexDict{K, KeyRange, VectorSegment{T}}(views))
end

function SegmentedVector{K}(parent::P, keys, viewlengthfun) where {K, T, P<:AbstractVector{T}}
    SegmentedVector{K, T, Base.OneTo{K}}(parent, keys, viewlengthfun)
end

function SegmentedVector{K, T, KeyRange}(parent::P, ranges::AbstractDict{K, UnitRange{Int}}) where {K, T, KeyRange, P<:AbstractVector{T}}
    segs = IndexDict{K, KeyRange, VectorSegment{T}}(keys(ranges), [view(parent, range) for range in values(ranges)])
    SegmentedVector{K, T, KeyRange, P}(parent, IndexDict{K, KeyRange, VectorSegment{T}}(segs))
end

Base.size(v::SegmentedVector) = size(v.parent)
@propagate_inbounds Base.getindex(v::SegmentedVector, i::Int) = v.parent[i]
@propagate_inbounds Base.getindex(v::SegmentedVector{K}, key::K) where {K} = v.segments[key] # TODO: deprecate in favor of view?
@propagate_inbounds Base.setindex!(v::SegmentedVector, value, i::Int) = v.parent[i] = value
Base.dataids(x::SegmentedVector) = Base.dataids(x.parent)
@propagate_inbounds Base.view(v::SegmentedVector{K}, key::K) where {K} = v.segments[key]

Base.parent(v::SegmentedVector) = v.parent
segments(v::SegmentedVector) = v.segments
function ranges(v::SegmentedVector{K, <:Any, KeyRange}) where {K, KeyRange}
    segments = v.segments
    IndexDict{K, KeyRange, UnitRange{Int}}(segments.keys, map(segment -> first(parentindices(segment))::UnitRange{Int}, segments.values))
end

function Base.similar(v::SegmentedVector{K, T, KeyRange}, ::Type{S} = T) where {K, T, KeyRange, S}
    SegmentedVector{K, S, KeyRange}(similar(parent(v), S), ranges(v))
end


text:
using CSV

function LoadLowSchwarzschildQnms()
    A = CSV.read(joinpath(@__DIR__, "SwData.csv"),NamedTuple)
    ωlist = Dict{NamedTuple,Complex{Float64}}()
    for i=1:length(A[:l])
        merge!(ωlist,Dict((l = A[:l][i], s = A[:s][i], n = A[:n][i]) => A[:ωreal][i]+im*A[:ωimag][i]))
    end
    ωlist
end

function Low_l_Guess(l,n,s)::Complex{Float64}
    SchwMODESLOWWℓ[(l=l,s=s,n=n)]
end

SchwMODESLOWWℓ = LoadLowSchwarzschildQnms()

function High_l_Guess(l,n,s)
    #Taken Directly from Eq: 17-22 of https://arxiv.org/pdf/0908.0329.pdf
    L = l+0.5
    N = n+0.5
    β = 1-s^2

    ω₋₁ = 1
    ω₀ᵢ = 1
    ω₁ = β/3 - (5/36)*N^2 - (115/432)
    ω₂ᵢ = β/9 + (235/3888)*N^2 - (1415/15552)
    ω₃ = -(1/27)*β^2 + β*(204*N^2 +211)/3888 + (854160*N^4 - 1664760*N^2 - 776939)/(40310784)
    ω₄ᵢ = (1/27)*β^2 + (1100*N^2 - 2719)/(46656) + (11273136*N^4 - 52753800*N^2 + 66480535)/(2902376448)

    ω = (ω₋₁*L + ω₁*(L^(-1)) + ω₃*(L^(-3))) - im*N*(ω₀ᵢ + ω₂ᵢ*(L^(-2)) - ω₄ᵢ*(L^(-4)))
    ω = (1/sqrt(27))*ω
end

function High_n_Guess(l,n,s)::Complex{Float64}
    print("Used higher n guess")
    # From H.-J. Blome, B. Mashhoon, Phys. Lett. A 110, 231 (1984)
    return (0.5 + l - im*(0.5 + n))/(3.0*sqrt(3))
end

function SchwarzschildGuess(l::Int64,n::Int64,s::Int64)::Complex{Float64}
    if l >= 3
        return High_l_Guess(l,n,s)::Complex{Float64}
    else
        return Low_l_Guess(l,n,s)::Complex{Float64}
    end
end

function SolveForParams(l,m,n,s,a,ωguess,Alm,Nmax,lmax)
    let l=l,m=m,n=n,s=s,a=a,Alm=Alm,Nmax=Nmax,lmax=lmax,ωguess=ωguess
        function f!(F,x)
            val = CF(l,m,n,s,a,x[1]+im*x[2],Alm,Nmax,lmax)
            F[1] = real(val)
            F[2] = imag(val)
        end
        sol = nlsolve(f!, [real(ωguess),imag(ωguess)],ftol = 1e-14, method = :newton)
        ωnews = sol.zero
        ωnews[1]+im*ωnews[2]
    end
end

function CF(l,m,n,s,a,ω,Alm,Nmax,lmax)
    let l=l,m=m,n=n,s=s,a=a,ω=ω,Alm=Alm,Nmax=Nmax,lmax=lmax
        _,_,(D₀,D₁,D₂,D₃,D₄) = ParameterTransformations(l,m,n,s,a,ω,Alm,Nmax,lmax)
        N = Nmax;
        αₙ(n) = n^2 + (D₀+1)*n + D₀
        βₙ(n) = -2*n^2 + (D₁+2)*n +D₃
        γₙ(n) = n^2 + (D₂-3)*n  + D₄ - D₂ + 2
        u1,u2,u3,u4 = rNCoeffs(D₀,D₁,D₂,D₃,D₄)
        rN = 1 + u1*N^(-0.5) + u2*N^(-1) + u3*N^(1.5) + u4*N^(-2)
        ContinuedFraction(αₙ,βₙ,γₙ,n,N; rN = rN)
    end
end

function SchwarzschildModes(l,n,s)
    ωguess = SchwarzschildGuess(l,n,s)::Complex{Float64}
    ωguess = Complex(ωguess)
    #println(ωguess)
    Alm = l*(l+1) - s*(s+1)
    lmax = 12; m=0; a = 0

    NmaxTrials = [100,300,500,1000,5_000,10_000,50_000,100_000]
    ω₀ = ωguess; Alm₀ = Alm
    ωnew = ωguess + 1.0;
    iterations = 1;
    while (abs(ωnew-ω₀) > 1e-10) & (iterations < 8)
        Nmax = NmaxTrials[iterations]
        ω₀ = ωnew;
        ωnew = SolveForParams(l,m,n,s,a,ωguess,Alm,Nmax,lmax)
        iterations += 1
        #println("ω diff = ",abs(ωnew-ω₀))
    end
    #println("Number of Iterations = ", iterations)
    ωnew, Alm
end


text:
s = "Hello Julia! Let's start to leanr Julia"

# first 'l'
i = findfirst(isequal('l'), s)
println(i)

# start search at a given offset
r = findnext(isequal('l'), s, 5)
println(r)

# check if a substring is found within a string
println(occursin("Julia", s))

# the number of characters in str
println(length(s))

# regex: Perl-compatible regular expressions
println(match(r"^\s*(?:#|$)", "# a comment"))

# replace
println(replace("a", r"." => s"\g<0>1"))

#substring
substr = SubString(s, 1, 4)
println(substr)

# strip
s = strip("hello ")
println(s)

# split
s = split("hello, there,bob", ", ")
println(s)

# join
j = join(collect(1:10), ", ")
println(j)

max_stars_repo_name: startprogress/Hello-Julia

max_issues_repo_name: startprogress/Hello-Julia

max_forks_repo_name: startprogress/Hello-Julia


text:
__precompile__()

module Quats

import Random: AbstractRNG, SamplerType
import LinearAlgebra: norm, normalize
import Base: convert, promote_rule, show, write, read
import Base: real, imag, complex
import Base: float, big, widen, bswap
import Base: isinteger, isfinite, isnan, isinf, iszero, isequal, isone, zero
import Base: vec, conj, abs, abs2, inv, rand, randn, exp, log, round, sqrt
import Base: +, -, *, /, ^, ==

export Quaternion
export QuaternionF16
export QuaternionF32
export QuaternionF64
export QuatF16
export QuatF32
export QuatF64

export quat
export jm
export km
export jmag
export kmag
export cmatrix
export qmatrix
export norm
export normalize

include("quaternion.jl")

end # module


text:
using StaticFFTs, StaticArrays, FFTW
using Test

# write your own tests here
@testset "FFT" begin
    sizes = [6,8,9,12,14,15,16,18,20,22,24,25,26,27,28,30,32]
    for s ∈ sizes
        x = @SVector randn(s)
        @test StaticFFTs.fft(x) ≈ FFTW.fft(Vector(x))
    end
end

max_issues_repo_name: chriselrod/StaticFFTs

max_forks_repo_name: chriselrod/StaticFFTs


text:
ENV["PYTHON"] = ""
ENV["JUPYTER"] = ""
ENV["JULIA_WARN_COLOR"] = :yellow
ENV["JULIA_INFO_COLOR"] = :cyan
ENV["EDITOR"] = "atom"    # "subl"
# ENV["LISP_PROMPT_TEXT"]  = "λ ↦ "

@static if is_windows()
    ENV["VS120COMNTOOLS"] = "C:\\Program Files (x86)\\Microsoft Visual Studio 12.0"
end


text:
include("utils/index.jl")

include("annotate.jl")
include("axes.jl")
include("labels.jl")
include("lines.jl")
include("plots.jl")
include("subplots.jl")
include("ticks.jl")


text:
using
  GridInterpolations

export
  SerialValueIteration,
  ParallelValueIteration,
  discretize_statevariable!

const MaxIter = 1000
const Tol = 1e-4
const Discount = 0.99
const NThreads = CPU_CORES / 2

type LazyDiscrete

  varname::AbstractString
  step::Float64

  LazyDiscrete(varname::AbstractString, step::Float64) = new(varname, step)

end

abstract ValueIteration <: Solver

type SerialValueIteration <: ValueIteration

  verbose::Bool

  maxiter::Int64
  tol::Float64
  discount::Float64

  statemap::Dict{AbstractString, LazyDiscrete}
  actionmap::Dict{AbstractString, LazyDiscrete}
  stategrid::RectangleGrid
  actiongrid::RectangleGrid

  SerialValueIteration(;
      verbose::Bool=true,
      maxiter::Real=MaxIter,
      tol::Float64=Tol,
      discount::Real=Discount) =
    new(
      verbose,
      Int64(maxiter),
      tol,
      Float64(discount),
      Dict{AbstractString, LazyDiscrete}(),
      Dict{AbstractString, LazyDiscrete}(),
      RectangleGrid(),
      RectangleGrid())

end

type ParallelValueIteration <: ValueIteration

  nthreads::Int64
  verbose::Bool

  maxiter::Int64
  tol::Float64
  discount::Float64

  statemap::Dict{AbstractString, LazyDiscrete}
  actionmap::Dict{AbstractString, LazyDiscrete}
  stategrid::RectangleGrid
  actiongrid::RectangleGrid

  ParallelValueIteration(;
      nthreads::Real=NThreads,
      verbose::Bool=true,
      maxiter::Real=MaxIter,
      tol::Float64=Tol,
      discount::Real=Discount) =
    new(
      Int64(nthreads),
      verbose,
      Int64(maxiter),
      tol,
      Float64(discount),
      Dict{AbstractString, LazyDiscrete}(),
      Dict{AbstractString, LazyDiscrete}(),
      RectangleGrid(),
      RectangleGrid())

end

type ValueIterationSolution <: Solution

  qval::Matrix{Float64}  # nactions x nstates Q-value matrix
  stategrid::RectangleGrid
  actiongrid::RectangleGrid

  cputime::Float64
  niter::Int64
  finaltol::Float64

  ValueIterationSolution(
      qval::Matrix{Float64},
      stategrid::RectangleGrid,
      actiongrid::RectangleGrid,
      cputime::Float64,
      niter::Int64,
      finaltol::Float64) =
    new(qval, stategrid, actiongrid, cputime, niter, finaltol)

end

function discretize_statevariable!(vi::ValueIteration, varname::AbstractString, step::Real)
  if haskey(vi.statemap, varname)
    warn(string(
      "state variable ", varname, " already discretized in ValueIteration object, ",
      "replacing existing discretization scheme"))
  end
  vi.statemap[varname] = LazyDiscrete(varname, Float64(step))
end

function discretize_actionvariable!(vi::ValueIteration, varname::AbstractString, step::Real)
  if haskey(vi.actionmap, varname)
    warn(string(
      "action variable ", varname, " already discretized in ValueIteration object, ",
      "replacing existing discretization scheme"))
  end
  vi.actionmap[varname] = LazyDiscrete(varname, Float64(step))
end

function getpolicy(mdp::MDP, solution::ValueIterationSolution)

  statedim = length(mdp.statemap)
  actiondim = length(mdp.actionmap)
  stateargs = mdp.reward.argnames[1:statedim]
  actionargs = mdp.reward.argnames[1 + statedim:end]
  nactions, nstates = size(solution.qval)

  function indexify(statevec::Vector)
    stateidxvec = zeros(statedim)
    for idim in 1:statedim
      statevar = mdp.statemap[stateargs[idim]]
      if isa(statevar, RangeVar)
        stateidxvec[idim] = statevec[idim]
      elseif isa(statevar, ValuesVar)
        stateidxvec[idim] = findfirst(statevar.values, statevec[idim])
      else
        error(string(
          "unknown state variable definition type for ", statevar))
      end
    end
    return stateidxvec
  end

  function policy(state...)

    statevec = [stateelem for stateelem in state]
    stateidxvec = indexify(statevec)
    stateidxs, wts = interpolants(solution.stategrid, stateidxvec)

    iaction_best = 0
    vaction_best = -Inf

    for iaction in 1:nactions
      vaction = 0.0
      for i in 1:length(stateidxs)
        vaction += wts[i] * solution.qval[iaction, stateidxs[i]]
      end

      if vaction > vaction_best
        iaction_best = iaction
        vaction_best = vaction
      end
    end

    rawaction = ind2x(solution.actiongrid, iaction_best)
    action = Array(Any, actiondim)

    for i in 1:actiondim
      actionvar = mdp.actionmap[actionargs[i]]
      if isa(actionvar, RangeVar)
        action[i] = rawaction[i]
      elseif isa(actionvar, ValuesVar)
        action[i] = actionvar.values[Int64(rawaction[i])]
      else
        error(string(
          "unknown action variable definition type for ", actionvar))
      end
    end

    return action

  end

  return policy

end

include("valueiteration_checks.jl")
include("valueiteration_solver.jl")


text:
using Documenter, GPUArrays

makedocs()

deploydocs(
    repo   = "github.com/JuliaGPU/GPUArrays.jl.git",
    julia  = "0.7",
    osname = "linux",
    # no need to build anything here, re-use output of `makedocs`
    deps   = nothing,
    make   = nothing,
)


text:
using DelimitedFiles
using StatsBase, Statistics, StatsPlots
using Random

export estimate_mu, ESS
export compare_efficiencies, view_distances, variance_table, plot_eta_estimates, plot_apost_efficiencies

######### Helper functions

function ESS(weights::Array{Float64,1})
    return sum(weights)^2/sum(weights.^2)
end
function ESS(s::MFABCCloud)
    return ESS([p.w for p in s])
end

function total_cost(s::MFABCCloud)
    return sum([sum(pp.p.cost) for pp in s])
end
function total_cost(s::BenchmarkCloud)
    return sum([sum(p.cost) for p in s])
end

function efficiency(s::MFABCCloud)
    return ESS(s)/total_cost(s)
end

function estimate_mu(s::MFABCCloud, parameterFun::Function)
    F_i = map(parameterFun, [pp.p.k for pp in s])
    w_i = [pp.w for pp in s]
    return sum(w_i .* F_i)/sum(w_i)
end

####################################

using Plots, StatsPlots, KernelDensity, LaTeXStrings, Printf
f(i) = Plots.font("serif",i...)
pyplot(titlefont=f(12), guidefont=f(9), legendfont=f(9), xtickfont=f(8), ytickfont=f(8))

function get_efficiencies(bm::BenchmarkCloud, epsilons::Tuple{Float64, Float64}, size_samples::Int64, eta_vec::Array{Tuple{Float64,Float64},1})
    sample_idx = Iterators.partition(1:length(bm), size_samples)
    Random.seed!(123)
    cloud_set = [MakeMFABCCloud(bm[i], epsilons, etas) for (i,etas) in Iterators.product(sample_idx, eta_vec)]
    Random.seed!()
    return ESS.(cloud_set)./cost.(cloud_set)
end

function compare_efficiencies(bm::BenchmarkCloud, size_samples::Int64, epsilons::Tuple{Float64, Float64};
    output::String="plot")

    eta_mf, phi_mf = get_eta(bm, epsilons, method="mf")
    eta_er, phi_er = get_eta(bm, epsilons, method="er")
    eta_ed, phi_ed = get_eta(bm, epsilons, method="ed")
    eta_abc, phi_abc = get_eta(bm, epsilons, method="abc")
    # Get some other continuation probabilities
    eta_pp = 0.5.*eta_mf .+ 0.5.*(1,1)
    phi_pp = phi(eta_pp,bm,epsilons)
    eta_pm = 0.5.*eta_mf .+ 0.5.*(1,0)
    phi_pm = phi(eta_pm,bm,epsilons)
    eta_mp = 0.5.*eta_mf .+ 0.5.*(0,1)
    phi_mp = phi(eta_mp,bm,epsilons)
    eta_mm = 0.5.*eta_mf .+ 0.5.*(0,0)
    phi_mm = phi(eta_mm,bm,epsilons)

    str_vec = ["Early Accept/Reject", "Early Rejection", "Early Decision", "Rejection Sampling", "+/+", "-/+", "+/-", "-/-"]
    eta_vec = [eta_mf, eta_er, eta_ed, eta_abc, eta_pp, eta_mp, eta_pm, eta_mm]
    phi_vec = [phi_mf, phi_er, phi_ed, phi_abc, phi_pp, phi_mp, phi_pm, phi_mm]
    offset_vec = [(0.03,0.01),(-0.02,0.01),(0.02,0.03),(-0.02,0.00),(0.02,0.02),(0.02,0.02),(0.02,0.02),(0.02,0.02)]
    alpha_vec = [1.0, 1.0, 1.0, 1.0, 0.3, 0.3, 0.3, 0.3]
    position_vec = [:left,:right,:left,:right,:left,:left,:left,:left]
    I = sortperm(phi_vec)

    if output=="plot"
        efficiencies = get_efficiencies(bm, epsilons, size_samples, eta_vec)
        #=
        fig = plot(color_palette=:darkrainbow,
            thickness_scaling=1,
            xlabel="Effective samples per second",
            xtickfontsize=10,
            yticks=[],
            xgrid=false,
            #legend=(0.7,0.25),
            legendfontsize=11,
            title="Distribution of Efficiency Across Multiple Realisations",
            titlefontsize=12)
        for i in I
            if i<=4
                plot!(kde(efficiencies[:,i]),
                label=str_vec[i],
                linewidth=2.5
                )
            else
                plot!(kde(efficiencies[:,i]),
                label=str_vec[i],
                linewidth=1.5,
                linestyle=:dot
                )
            end
        end
        =#
        fig = plot(color_palette=:darkrainbow,
        thickness_scaling=1,
        ylabel="Effective samples per second",
        ytickfontsize=10,
        xticks=[],
        ygrid=false,
        #legend=(0.7,0.25),
        legendfontsize=10,
        title="Observed Distribution of Efficiency",
        titlefontsize=12)
        violin!(efficiencies[:,I], label=hcat(str_vec[I]...), alpha=hcat(alpha_vec[I]...))

        pc_bigger_than(eff1,eff2) = count(F1>F2 for F1 in eff1 for F2 in eff2)/(length(eff1)*length(eff2))
        num_rows = size(efficiencies,2)-1
        tab = Array{Tuple{String, String, Float64},2}(undef,num_rows,num_rows)
        latextab = raw"\mathbb{P}\left(\mathrm{row exceeds column}\right)"*prod(map(s->raw" & \text{"*s*"}", str_vec[I[2:end]]))*raw"\\ \hline "
        for r in 1:num_rows
            latextab *= raw"\\ \text{" * str_vec[I[r]] * "}" * " & "^(r-1)
            for c in r:num_rows
                p=pc_bigger_than(efficiencies[:,I[r]],efficiencies[:,I[c+1]])
                latextab *= @sprintf(" & %.2f",p)
                tab[r,c] = (str_vec[I[r]], str_vec[I[c+1]], p)
            end
        end
        open("./efficiency_table.txt", "w") do f
            write(f, latextab)
        end
        return fig, tab, latextab

    elseif output=="theory"

        fig = plot(color_palette=:darkrainbow, legend=:none, grid=:none,
            xlabel=L"\eta_1",
            ylabel=L"\eta_2",
            labelfontsize=10,
            title="Continuation Probabilities and Efficiency",
            titlefontsize=11)
        for i in I
            scatter!(eta_vec[i], markersize=8, markerstrokewidth=0, alpha=alpha_vec[i], annotations=((eta_vec[i].+offset_vec[i])..., text(str_vec[i],f((8,position_vec[i])))))
        end

        grd = 0:0.01:1
        plot!(collect(grd), x->x, color=[:black], linestyle=:dash, label=[])
        plot!(ones(size(grd)),collect(grd), color=[:black], linestyle=:dash, label=[])

        sp = sample_properties(bm, epsilons)
        phi_plot((x,y)) = phi((x,y), sp...)

        contour!(collect(grd),collect(grd),(x,y)->1/phi_plot((x,y)),
        aspect_ratio=:equal,
        levels = (1/phi_mf).*[0.6, 0.75, 0.8, 0.85, 0.9, 0.95, 0.99],
        c=cgrad(:grays_r),
        linewidth=0.6,
        linestyle=:dot,
        label=[],
        colorbar=:none)

        return fig

    else
        return collect(zip((str_vec[I], eta_vec[I], phi_vec[I], efficiencies[:,I])...))
    end
end

function view_distances(s::BenchmarkCloud, epsilons::Tuple{Float64, Float64})

    match = [p.dist for p in s if !xor((p.dist .< epsilons)...)]
    fp = [p.dist for p in s if ((p.dist[1] < epsilons[1]) & (p.dist[2] >= epsilons[2])) ]
    fn = [p.dist for p in s if ((p.dist[1] >= epsilons[1]) & (p.dist[2] < epsilons[2])) ]
    # Compare low and high fidelity
    plot(; title="Distance from data: multifidelity", titlefontsize=12,
    aspect_ratio=:equal, grid=:none,
    xlabel=latexstring("\$ \\tilde{d}(\\tilde{y}, \\tilde{y}_{obs}) \$"), ylabel=latexstring("\$ d(y,y_{obs}) \$"), labelfontsize=10)
    scatter!(match, markersize=1.5, markerstrokewidth=0, alpha=0.6, label="Matching estimator values")
    scatter!(fp, markersize=3, markerstrokewidth=0, label="False positive")
    scatter!(fn, markersize=3, markerstrokewidth=0, label="False negative")
    vline!([epsilons[1]], linestyle=:dash, color=[:black], label="")
    hline!([epsilons[2]], linestyle=:dash, color=[:black], label="")
end
function view_distances(s::BenchmarkCloud, epsilons::Tuple{Float64, Float64}, par_n::Integer, par_name::AbstractString)

    match = [(p.k[par_n], p.dist[2]) for p in s if !xor((p.dist .< epsilons)...)]
    fp = [(p.k[par_n], p.dist[2]) for p in s if ((p.dist[1] < epsilons[1]) & (p.dist[2] >= epsilons[2])) ]
    fn = [(p.k[par_n], p.dist[2]) for p in s if ((p.dist[1] >= epsilons[1]) & (p.dist[2] < epsilons[2])) ]

    # Compare distance by parameter
    plot(; title="Distance from data: by parameter", titlefontsize=12, grid=:none, legend=:none,
    xlabel=par_name, ylabel=L"d(y,y_{obs})", labelfontsize=10)
    scatter!(match, markersize=1.5, markerstrokewidth=0, alpha=0.6, label="Matching estimator values")
    scatter!(fp, markersize=3, markerstrokewidth=0, label="False positive")
    scatter!(fn, markersize=3, markerstrokewidth=0, label="False negative")
    hline!([epsilons[2]], linestyle=:dash, color=[:black], label="")

end
function view_distances(s::BenchmarkCloud, epsilons::Tuple{Float64, Float64}, inset_limits::Tuple{Float64,Float64})

    match = [p.dist for p in s if !xor((p.dist .< epsilons)...)]
    fp = [p.dist for p in s if ((p.dist[1] < epsilons[1]) & (p.dist[2] >= epsilons[2])) ]
    fn = [p.dist for p in s if ((p.dist[1] >= epsilons[1]) & (p.dist[2] < epsilons[2])) ]
    # Compare low and high fidelity
    plot(; title="Distance from data: multifidelity", titlefontsize=12,
    aspect_ratio=:equal, grid=:none,
    xlabel=latexstring("\$ \\tilde{d}(\\tilde{y}, \\tilde{y}_{obs}) \$"), ylabel=latexstring("\$ d(y,y_{obs}) \$"), labelfontsize=10)

    plot!(; inset_subplots=(1,bbox(0.15,0.15,0.3,0.3,:bottom,:right)))

    scatter!(match, markersize=1.5, markerstrokewidth=0, alpha=0.6, label="Matching estimator values", subplot=1)
    scatter!(fp, markersize=3, markerstrokewidth=0, label="False positive", subplot=1)
    scatter!(fn, markersize=3, markerstrokewidth=0, label="False negative", subplot=1)

    scatter!(match, markersize=2.5, markerstrokewidth=0, alpha=0.6, label="Matching estimator values", subplot=2)
    scatter!(fp, markersize=4, markerstrokewidth=0, label="False positive", subplot=2)
    scatter!(fn, markersize=4, markerstrokewidth=0, label="False negative", subplot=2)

    vline!([epsilons[1]], linestyle=:dash, color=[:black], label="", subplot=1)
    hline!([epsilons[2]], linestyle=:dash, color=[:black], label="", subplot=1)

    vline!([epsilons[1]], linestyle=:dash, color=[:black], label="", subplot=2)
    hline!([epsilons[2]], linestyle=:dash, color=[:black], label="", subplot=2)

    plot!(xlim=(0,inset_limits[1]), ylim=(0,inset_limits[2]), legend=:none, aspect_ratio=:equal, grid=:none, subplot=2, xticks=[epsilons[1]], yticks=[epsilons[2]], framestyle=:box)
end

function get_eta(bm::BenchmarkCloud, epsilons::Tuple{Float64,Float64}, F::Array{Function,1})
    method_list = ["abc", "er", "ed", "mf"]
 #   sample_idx = Iterators.partition(1:length(bm), size_samples)
 #   bm_set = [bm[i] for i in sample_idx]

    eta_phi_fun((f,m)) = get_eta(bm, epsilons, method=m, F=f)
    eta_phi = map(eta_phi_fun, Iterators.product(F, method_list))
    ess_eta = get_eta(bm, epsilons, method="mf")[1]
    ess_phi = [phi(ess_eta, bm, epsilons, F=f) for f in F]

    eta_tab = hcat([x[1] for x in eta_phi], repeat([ess_eta], length(F)))
    phi_tab = hcat([x[2] for x in eta_phi], ess_phi)

    return eta_tab, phi_tab
end

function observed_variance(bm_set::Array{BenchmarkCloud,1}, epsilons::Tuple{Float64,Float64}, eta::Tuple{Float64,Float64}, budget::Float64, F::Function)
    mf_set = [MakeMFABCCloud(bm_i, epsilons, eta, budget) for bm_i in bm_set]
    return var(estimate_mu.(mf_set, F))
end

function variance_table(bm::BenchmarkCloud, sample_size::Integer, epsilons::Tuple{Float64, Float64}, eta_tab::Array{Tuple{Float64, Float64}}, Functions::Array{Function,1}, budget::Float64)
    idx = Iterators.partition(1:length(bm), sample_size)
    f_tab = repeat(Functions,1,size(eta_tab,2))
    if size(f_tab)==size(eta_tab)
        shuffle!(bm)
        bm_set = [bm[i] for i in idx]
        return [observed_variance(bm_set, epsilons, eta, budget, f) for (eta, f) in zip(eta_tab, f_tab)]
    else
        error("eta table doesn't match number of functions")
    end
end
function variance_table(bm::BenchmarkCloud, sample_size::Integer, epsilons::Tuple{Float64, Float64}, Functions::Array{Function,1}, budget::Float64)
    eta_tab = get_eta(bm, epsilons, Functions)[1]
    return variance_table(bm, sample_size, epsilons, eta_tab, Functions, budget)
end

function plot_eta_estimates(bm::BenchmarkCloud, epsilons::Tuple{Float64, Float64}, cloud_sets::Tuple{Array{<:Cloud,1},AbstractString}...; method::String="mf", kwargs...)

    fig = plot(grid=:none,
    xlabel=L"\eta_1",
    ylabel=L"\eta_2",
    labelfontsize=10,
    title="Continuation Probabilities and Efficiency",
    titlefontsize=11)

    for (cs,l) in cloud_sets
        eta_estimates = [get_eta(cld, epsilons; method=method, kwargs...)[1] for cld in cs]
        scatter!(eta_estimates, markersize=6, markerstrokewidth=0, xlim=(0,1), ylim=(0,1),label=l)
    end

    eta_real, phi_mf = get_eta(bm, epsilons, method=method)
    scatter!(eta_real, markersize=6, markerstrokewidth=0, label="Benchmark", color=:black)

    grd = 0:0.01:1
    sp = sample_properties(bm, epsilons)
    phi_plot((x,y)) = phi((x,y), sp...)

    contour!(collect(grd),collect(grd),(x,y)->1/phi_plot((x,y)),
    aspect_ratio=:equal,
    levels = (1/phi_mf).*[0.6, 0.75, 0.8, 0.85, 0.9, 0.95, 0.99],
    c=cgrad(:grays_r),
    linewidth=0.6,
    linestyle=:dot,
    label=[],
    colorbar=:none)
end

function plot_apost_efficiencies(cloud_sets::Tuple{Array{<:Cloud,1},AbstractString}...)

    eff(set) = ESS.(set)./cost.(set)

    fig = plot(color_palette=:darkrainbow,
    thickness_scaling=1,
    ylabel="Effective samples per second",
    ytickfontsize=10,
    xticks=[],
    ygrid=false,
    #legend=(0.7,0.25),
    legendfontsize=11,
    title="Observed Distribution of Efficiency",
    titlefontsize=12)

    for (cs, l) in cloud_sets
        violin!(eff(cs), label=l)
    end
    return fig
end

max_stars_repo_name: tpprescott/mf-abc

max_issues_repo_name: tpprescott/mf-abc

max_forks_repo_name: tpprescott/mf-abc


text:
@theme PygmentsTheme Dict(
    :name => "Pygments",
    :description => "Based on the default theme provided by Pygments.",
    :tokens => Dict(
        TEXT               => S"",
        ERROR              => S"fg:ffffff; bg:ff3300",

        WHITESPACE         => S"fg:bbbbbb",
        COMMENT            => S"italic; fg:408080",
        COMMENT_PREPROC    => S"fg:BC7A00",

        KEYWORD            => S"bold; fg:008000",
        KEYWORD_PSEUDO     => S"",
        KEYWORD_TYPE       => S"fg:B00040",

        OPERATOR           => S"fg:666666",
        OPERATOR_WORD      => S"bold; fgAA22FF",

        NAME_BUILTIN       => S"fg:008000",
        NAME_FUNCTION      => S"fg:0000FF",
        NAME_CLASS         => S"bold; fg:0000FF",
        NAME_NAMESPACE     => S"bold; fg:0000FF",
        NAME_EXCEPTION     => S"bold; fg:D2413A",
        NAME_VARIABLE      => S"fg:19177C",
        NAME_CONSTANT      => S"fg:880000",
        NAME_LABEL         => S"fg:A0A000",
        NAME_ENTITY        => S"bold; fg:999999",
        NAME_ATTRIBUTE     => S"fg:7D9029",
        NAME_TAG           => S"bold; fg:008000",
        NAME_DECORATOR     => S"fg:AA22FF",

        STRING             => S"fg:BA2121",
        STRING_DOC         => S"italic",
        STRING_INTERPOL    => S"bold; fg:BB6688",
        STRING_ESCAPE      => S"bold; fg:BB6622",
        STRING_REGEX       => S"fg:BB6688",
        STRING_SYMBOL      => S"fg:19177C",
        STRING_OTHER       => S"fg:008000",

        NUMBER             => S"fg:666666",

        GENERIC            => S"",
        GENERIC_HEADING    => S"bold; fg:000080",
        GENERIC_SUBHEADING => S"bold; fg:800080",
        GENERIC_DELETED    => S"fg:A00000",
        GENERIC_INSERTED   => S"fg:00A000",
        GENERIC_ERROR      => S"fg:FF0000",
        GENERIC_EMPH       => S"italic",
        GENERIC_STRONG     => S"bold",
        GENERIC_PROMPT     => S"bold; fg:000080",
        GENERIC_OUTPUT     => S"fg:888",
        GENERIC_TRACEBACK  => S"fg:04D",
    )
)


text:
#=

  xkcd's knapsack/subset-sum problem in Julia + ConstraintSolver.jl

  http://xkcd.com/287/

  Some amount (or none) of each dish should be ordered to give a total of exact 15.05

  Model created by Hakan Kjellerstrand, hakank@gmail.com
  See also my Julia page: http://www.hakank.org/julia/

=#
using ConstraintSolver, JuMP
using Cbc, GLPK
const CS = ConstraintSolver
using Distributions
include("constraints_utils.jl")

function xkcd(prices=[215, 275, 335, 355, 420, 580],total=1505,max_val=10;all_solutions=false)

    cbc_optimizer = optimizer_with_attributes(Cbc.Optimizer, "logLevel" => 0)
    glpk_optimizer = optimizer_with_attributes(GLPK.Optimizer)

    model = Model(optimizer_with_attributes(CS.Optimizer,   "all_solutions"=> all_solutions,
                                                            # "all_optimal_solutions"=>true,
                                                            "logging"=>[],

                                                            # "traverse_strategy"=>:BFS,
                                                            "traverse_strategy"=>:DFS,
                                                            # "traverse_strategy"=>:DBFS,

                                                            # "branch_split"=>:Smallest,
                                                            # "branch_split"=>:Biggest,
                                                            "branch_split"=>:InHalf,

                                                            # "simplify"=>false,
                                                            "simplify"=>true,

                                                            "time_limit"=>14,
                                                            # "lp_optimizer" => cbc_optimizer,
                                                            # "lp_optimizer" => glpk_optimizer,
                                        ))

    println("prices:$prices")
    println("total:$total   (max_val:$max_val)")
    len = length(prices)

    @variable(model, 0 <= x[1:len] <= max_val, Int)

    # @constraint(model, total == sum([x[i]*prices[i] for i in 1:len]))
    scalar_product(model,x,prices,total)
    # @constraint(model, total == sum(x.*prices))

    # Solve the problem
    println("solve")
    optimize!(model)

    status = JuMP.termination_status(model)
    println("status:$status")
    if status == MOI.OPTIMAL
        num_sols = MOI.get(model, MOI.ResultCount())
        println("\nnum_sols:$num_sols\n")
        # num_sols2 = JuMP.result_count(model) # same result
        # println("num_sols2:$num_sols2\n")

        # The solutions are in reverse order, i.e. the optimal is the first.
        for sol in 1:num_sols
            println("solution #$sol")
            # This works:
            xx = convert.(Integer,JuMP.value.(x,result=sol))
            println("x:$xx")
            println("used: ")
            println([(i,xx[i]) for i in 1:length(xx) if xx[i] > 0  ])
        end
    end
end

xkcd(;all_solutions=true)

println("\nA random instance")

n = 120
prices = rand(DiscreteUniform(100,1000),n)
s = sum(prices)
total = rand(DiscreteUniform(0,s))
max_val = rand(DiscreteUniform(1,100))
xkcd(prices,total,max_val)

max_stars_repo_name: Wikunia/hakank

max_issues_repo_name: Wikunia/hakank

max_forks_repo_name: Wikunia/hakank


text:
println(pwd())

println(readdir("../doc"))

for (root, dirs, files) in walkdir(".")
    println("Directories in $root")
    for dir in dirs
        println(joinpath(root, dir)) # path to directories
    end
    println("Files in $root")
    for file in files
        println(joinpath(root, file)) # path to files
    end
end

# mkpath("../mkpath")

# mkdir("testingdir")

write("hello.txt", "world")

mv("hello.txt", "goodbye.txt")

readline("goodbye.txt")

rm("goodbye.txt")

println(isdir("../"))
println(isfile("../"))


text:
using MPI, Test
include("../testhelpers.jl")

@testset "Ocean" begin
  tests = [
    # (1,"Hydrostatic_Boussinesq/test_divergence_free.jl"),
    # (1,"shallow_water/GyreDriver.jl"),
   ]

  runmpi(tests, @__FILE__)
end

max_stars_repo_name: emassoud/CLIMA

max_issues_repo_name: emassoud/CLIMA

max_forks_repo_name: emassoud/CLIMA


text:
# Reactive message-passing in linear Gaussian dynamical system
# Experiment to infer states and transition coefficient simultaneously
#
# Wouter Kouw, BIASlab
# 02-01-2020

using Random
using Revise
using Distributions
using DataStructures
using LightGraphs
using MetaGraphs
using CPUTime
using Plots
pyplot()

# Factor graph components
include(joinpath(@__DIR__, "../factor_nodes/factor_gaussian.jl"))
include(joinpath(@__DIR__, "../factor_nodes/factor_gamma.jl"))
include(joinpath(@__DIR__, "../variables/var_gaussian.jl"))
include(joinpath(@__DIR__, "../variables/var_gamma.jl"))
include(joinpath(@__DIR__, "../variables/var_delta.jl"))
include(joinpath(@__DIR__, "../prob_operations.jl"))
include(joinpath(@__DIR__, "../graph_operations.jl"))
include(joinpath(@__DIR__, "../util.jl"))

# Data
include(joinpath(@__DIR__, "../gen_data.jl"))

"""
Experiment parameters
"""

# Signal time horizon
T = 50

# Reaction-time clock
TT = 10

# Transition and observation matrices
gain = 0.8
emission = 1.0

# Noises (variance form)
process_noise = 2.0
measurement_noise = 0.5

# Prior parameters
px0 = Normal(0.0, 1.0)
pA = Normal(0.0, 1.0)

# Generate data
Random.seed!(256)
observed, hidden = gendata_LGDS(gain,
                                emission,
                                process_noise,
                                measurement_noise,
                                mean(px0),
                                inv(var(px0)),
                                time_horizon=T)

"""
Model/graph specification

This assumes the following type of models
p(y_{1:T}, x_{0:T} | u_{1:T}) = p(x_0) Π_t p(y_t | x_t) p(x_t | x_{t-1})

In other words, Markov chains of time-slices of a state-space models.
Below, we specify the following model through the time-slice subgraph
				 _
				[A]
                 |
                (α)
         ________|________________
         |      _|_       ___    |
(x_0)--  | --->[g_t]---->(x_t)---|
         |                 |     |
         |                _|_    |
         |               [f_t]   |
         |                 |     |
         | T            [[y_t]]  |
		 -------------------------

x_t-1 = previous state variable
g_t = state transition factor node
α = transition coefficient
A = transition coefficient prior factor node
x_t = current state variable
f_t = likelihood factor node
y_t = observed variable

---------------------------------
#TODO: Graph specification using a graphical user interface
"""

CPUtic()

# Start graph
global graph = MetaGraph(SimpleGraph(4*T+3))

let

# Set initial state variable
set_props!(graph, 1, Dict(:id => :x_0,
						  :time => 0,
						  :node => VarGaussian(:x_0, marginal=px0, time=0)))

# Connect initial state prior to first state transition
add_edge!(graph, (1, 4))

# Transition coefficient variable
set_props!(graph, 2, Dict(:id => :α,
						  :time => nothing, #TODO: list of time-points
						  :node => VarGaussian(:α, time=nothing)))

# Transition coefficient prior
set_props!(graph, 3, Dict(:id => :A,
						  :time => nothing,
						  :node => FactorGaussian(:A, out=:α, mean=0.0, precision=0.1, time=nothing)))

# Connect transition coefficient prior to noise variable
add_edge!(graph, (2, 3))

# Start node numbering
node_num = 3

for t = 1:T

	# Cast current time to string
	t_ = string(t)

	# State transition node
	props_gt = Dict(:id => Symbol("g_"*t_),
				    :time => t,
				    :node => FactorGaussian(Symbol("g_"*t_),
								            out=Symbol("x_"*t_),
	                                        mean=Symbol("x_"*string(t-1)),
	                                        precision=inv(process_noise),
	                                        transition=:α))

	# Current state
	props_xt = Dict(:id => Symbol("x_"*t_),
				    :time => t,
				    :node => VarGaussian(Symbol("x_"*t_), time=t))

	# Observation likelihood node
	props_ft = Dict(:id => Symbol("f_"*t_),
				    :time => t,
				    :node => FactorGaussian(Symbol("f_"*t_),
		 						            out=Symbol("y_"*t_),
		 							        mean=Symbol("x_"*t_),
		 							        precision=inv(measurement_noise),
		 							        transition=emission,
		 							        time=t))

	# Observation
	props_yt = Dict(:id => Symbol("y_"*t_),
				    :time => t,
				    :node => VarDelta(Symbol("y_"*t_), observed[t], time=t))

	# Add properties to nodes in current time-slice
	set_props!(graph, node_num+1, props_gt)
	set_props!(graph, node_num+2, props_xt)
	set_props!(graph, node_num+3, props_ft)
	set_props!(graph, node_num+4, props_yt)

	# Add edges between factors and variables
	add_edge!(graph, (node_num-2, node_num+1)) # x_t-1 -- g_t
	add_edge!(graph, (node_num+1, node_num+2)) # g_t -- x_t
	add_edge!(graph, (node_num+2, node_num+3)) # x_t -- f_t
	add_edge!(graph, (node_num+3, node_num+4)) # f_t -- y_t
	add_edge!(graph, (node_num+1, 2)) # g_t -- γ

	# Increment node number
	node_num += 4
end
end

# Ensure vertices can be recalled from given id
set_indexing_prop!(graph, :id)

"""
Inference: filtering
"""

function infer!(graph::MetaGraph; start_nodes=[:x_0, :α], time=T, num_iterations=TT)

	# Start message routine
	for node in start_nodes
		act!(graph, node)
	end

	for t = 1:time

		# Report progress
		if mod(t, time/10) == 1
		  println("At iteration "*string(t)*"/"*string(time))
		end

		# Start clock for reactions
		for tt = 1:num_iterations

			# Iterate over all nodes to react
			for node in nodes_t(graph, t, include_notime=true) #TODO: parallelize
				react!(graph, node)
			end
		end
	end
end

infer!(graph)

CPUtoc()

"""
Visualize experimental results
"""

# Preallocate state vectors
estimated_states = zeros(T,2)
energies = zeros(T,)

# Loop over time
for t = 1:T

	# Loop over nodes in timeslice
	for node in nodes_t(graph, t)

		# Check for state node
		if string(node)[1] == 'x'

			# Retrieve object from node
			node_x = graph[graph[node, :id], :node]

			# Extract moments of marginals
			estimated_states[t,1] = mean(node_x.marginal)
			estimated_states[t,2] = sqrt(var(node_x.marginal))

			# Free energy
			energies[t] = node_x.free_energy

		end
	end
end

# Plot estimated states
scatter(1:T, observed, color="black", label="observations")
plot!(hidden[2:end], color="red", label="true states")
plot!(estimated_states[:,1], color="blue", label="inferred")
plot!(estimated_states[:,1],
	  ribbon=[estimated_states[:,2], estimated_states[:,2]],
	  linewidth=2,
      color="blue",
      fillalpha=0.2,
      fillcolor="blue",
	  label="")
xlabel!("time (t)")
savefig(joinpath(@__DIR__, "viz/state_estimates_experiment.png"))

# Plot estimated transition coefficient
x_range = 0.05:0.02:4
plot(x_range, pdf.(graph[graph[:α, :id], :node].marginal, x_range), color="blue", label="posterior")
vline!([gain], color="red", label="true")
xlabel!("transition coefficient (α)")
ylabel!("p(α)")
savefig(joinpath(@__DIR__, "viz/transition-coefficient_estimates_experiment.png"))

# Plot free energy
plot(1:T, energies, color="green", label="")
xlabel!("time (t)")
ylabel!("free energy (F)")
savefig(joinpath(@__DIR__, "viz/free_energy_experiment.png"))

max_stars_repo_name: wmkouw/schedulefreeMP

max_issues_repo_name: wmkouw/schedulefreeMP

max_forks_repo_name: wmkouw/schedulefreeMP


text:
export DefaultServer

struct DefaultServer{K, ELT} <: AbstractCacheServer{K, ELT}
    storage::Dict{UInt, ELT} # objectid => element

    DefaultServer{K, ELT}() where K where ELT = new{K, ELT}(Dict{UInt, ELT}())
end

function alloc!(server::DefaultServer{K}, object::K, fragment) where K
    server.storage[objectid(object)] = fragment
    server
end

function iscacheable(server::DefaultServer{K}, object::K) where K
    objectid(object) in keys(server.storage)
end

function iscached(server::DefaultServer{K}, object::K, params...) where K
    iscacheable(server, object) && iscached(server.storage[objectid(object)], params...)
end

function push!(server::DefaultServer{K}, val, index::K) where K
    update!(server.storage[objectid(index)], val)
    server
end

function pull(server::DefaultServer{K}, object::K, params...) where K
    pull(server.storage[objectid(object)], params...)
end

function delete!(server::DefaultServer{K}, object::K) where K
    delete!(server.storage, objectid(object))
    server
end

function clear!(server::DefaultServer{K}, object::K) where K
    clear!(server.storage[objectid(object)])
    server
end

function show(io::IO, server::DefaultServer)
    summary(io, server)
    for (key, elem) in server.storage
        print(io, key, " => ", elem)
    end
end


text:
# %%
using Pkg; Pkg.activate("GraphFusedElasticNet.jl")
# using Revise
using Distributions
import StatsBase.weights
using Random
# using RCall
using ProgressMeter
using FileIO
using DataFrames, CSV
using Printf
using Base.Threads
using NPZ
using GraphFusedElasticNet

Random.seed!(138590)

println("Running with $(Threads.nthreads())")

debug = true
if debug
    runid = 0
else
    runid =  parse(Int, ENV["SLURM_PROCID"])
end

# %%
function generate_trace(task, N)
    cuts = [N ÷ 3, 2(N ÷ 3)]
#     cuts = sort(sample(2:N-1, 2, replace=false))
    x1 = 1:cuts[1]
    x2 = (cuts[1] + 1):cuts[2]
    x3 = (cuts[2] + 1):N
    x = [x1; x2; x3]

    values1 = 2rand(Uniform(), 4) .- 1.0

    if task == "smooth"
        y1 = values1[1] .+ (values1[2] - values1[1]) .* (x1 ./ cuts[1])
        y2 = values1[2] .+ (values1[3] - values1[2]) .* (x2 .- cuts[1]) ./ (cuts[2] - cuts[1])
        y3 = values1[3] .+ (values1[4] - values1[3]) .* (x3 .- cuts[2]) ./ (N - cuts[2])
    elseif task == "constant"
        y1 = fill(values1[1], cuts[1])
        y2 = fill(values1[2], cuts[2] - cuts[1])
        y3 = fill(values1[3], N - cuts[2])
    elseif task == "mixed"
        y1 = values1[1] .+ Int(rand() < 0.5) .* (values1[2] - values1[1]) .* (x1 ./ cuts[1])
        y2 = values1[2] .+ Int(rand() < 0.5) .* (values1[3] - values1[2]) .* (x2 .- cuts[1]) ./ (cuts[2] - cuts[1])
        y3 = values1[3] .+ Int(rand() < 0.5) .* (values1[4] - values1[3]) .* (x3 .- cuts[2]) ./ (N - cuts[2])
    else
        throw(ArgumentError)
    end
    μ = [y1; y2; y3]
    return μ
end

function generate_spt_task(task_space, task_time, N, pmiss; σ=0.3, outliers=false, type="mixed")
    μ1s = 2generate_trace(task_space, N)
    μ2s = 2generate_trace(task_space, N)
    ts1 = 1.5generate_trace(task_time, N)
    ts2 = 1.5generate_trace(task_time, N)

    min_x = -2.5
    max_x = 2.5
    evalpts = collect(range(min_x, max_x, length=100))

    if type == "add"
        μs = [(t1 + μ1, t2 + μ2) for (μ1, μ2) in zip(μ1s, μ2s), (t1, t2) in zip(ts1, ts2)]
    elseif type == "mixed"
        μs = [(t1 * μ1, t2 * μ2) for (μ1, μ2) in zip(μ1s, μ2s), (t1, t2) in zip(ts1, ts2)]
    else
        throw(Exception)
    end

    dmodels = [
        begin
            dist1 = Truncated(Normal(μ1, σ), min_x, max_x)
            dist2 = Truncated(Normal(μ2, σ), min_x, max_x)
            MixtureModel([dist1, dist2], [0.5, 0.5])
        end
        for (μ1, μ2) in μs
    ]

    devals = [pdf.(d, evalpts) for d in dmodels];
    ndata = [sample([0, 10], weights([pmiss, 1.0 - pmiss])) for d in dmodels]
    y = [rand(d, n) for (d, n) in zip(dmodels, ndata)]

    if outliers
        Nobs = sum([1 for n in ndata if n > 0])
        K = Int(floor(Nobs * 0.5))
        idx = sample([i for (i, n) in enumerate(ndata) if n > 0], K, replace=false)
        for i in idx
            j = rand(1:length(y[i]))
            y[i][j] += rand([-1, 1]) * 5.0
        end
    end

    # make matrix pts
    xrange = collect(1:N)
    # temporal
    ptr = Int[]
    brks = Int[1]
    for i in 1:N
        append!(ptr, xrange .+ (i - 1) * N)
        push!(brks, brks[end] + N)
    end
    istemporal = fill(true, N^2)
    # spatial
    xrange = [(i - 1) * N + 1 for i in 1:N]
    for i in 1:N
        append!(ptr, xrange .+ (i - 1))
        push!(brks, brks[end] + N)
    end
    append!(istemporal, fill(false, N^2))

    return Dict("evalpts" => evalpts,
                "dmodels" => dmodels,
                "devals" => devals,
                "y" => y,
                "ndata" => ndata,
                "mean1" => μ1s,
                "mean2" => μ2s,
                "t" => ts1,
                "t" => ts2,
                "means" => μs,
                "ptr" => ptr,
                "brks" => brks,
                "istemporal" => istemporal)
end

# function for cross-validation fit

function generate_cvsets(y, nsplits)
    # make the cv splits
    N = length(y)
    cvsets = [Set{Int}() for i in 1:nsplits]
    iobs = shuffle([i for (i, yi) in enumerate(y) if !isempty(yi)])
    Nobs = length(iobs)
    splitsize = Nobs ÷ nsplits
    for k in 1:nsplits
        for i in ((k - 1) * splitsize + 1):(k * splitsize)
            push!(cvsets[k], iobs[i])
        end
    end
    return cvsets
end

function fit2(ytrain, ptr, brks, λ1, λ2, η1, η2, istemporal; parallel=false)
    # create the tree
    depth = 5
    nsplits = 2^depth - 1
    splits = uniform_binary_splits(-2.5, 2.5, depth)

    modelopts = Dict{Symbol, Any}(
        :admm_balance_every => 20,
        :admm_init_penalty => 10.0,
        :admm_residual_balancing => true,
        :admm_adaptive_inflation => false,
        :reltol => 1e-2,
        :admm_min_penalty => 0.1,
        :admm_max_penalty => 32.0,
        :abstol => 1e-6,
        :save_norms => false,
        :save_loss => false
    )
    fitopts = Dict{Symbol, Any}(
        :walltime => 60.0,
        :parallel => parallel
    )

    # fit binomial model in each split
    num_nodes = length(ytrain)
    lambdasl1 = Float64[temp ? η1 : λ1 for temp in istemporal]
    lambdasl2 = Float64[temp ? η2 : λ2 for temp in istemporal]
    betas = zeros(nsplits, num_nodes)
    for j in 1:nsplits
        l, m, u = splits.lows[j], splits.mids[j], splits.ups[j]
        successes = [sum(l .< yᵢ .<= m) for yᵢ in ytrain] .+ 1e-8
        attempts = [sum(l .< yᵢ .<= u) for yᵢ in ytrain] .+ 2e-8
        model = BinomialGFEN(ptr, brks, lambdasl1, lambdasl2; modelopts...)
        fit!(model, successes, attempts; fitopts...)
        betas[j, :] = model.beta
    end
    splits, betas
end

function cv_fit2(y, evalpts, ptr, brks, istemporal, lambdas, cvsets, models)
    # override only one cvsplit
    nsplits = 1

    # for each cv split get the mse error
    num_nodes = length(y)
    nsplits = length(cvsets)
    nlambdas = length(lambdas)
    test_nloglikelihood = zeros(nlambdas, nsplits)

    # prepare the tree structure and the bint
    for k in 1:nlambdas
        λ1, λ2, η1, η2 = lambdas[k]
        for i in 1:nsplits
            # get the cv vector with missing data
            ytrain = [(j ∈ cvsets[i] ? Float64[] : yi) for (j, yi) in enumerate(y)]

            splits, betas = fit2(ytrain, ptr, brks, λ1, λ2, η1, η2, istemporal)
            lows, mids, ups = splits.lows, splits.mids, splits.ups

            # compute the out-of-sample likelihood
            Ntest = 0
            nloglikelihood = 0.0
            for j in cvsets[i]
                root = make_tree_from_bfs(lows, mids, ups, betas[:, j])
                test_eval = y[j]
                Ntest += length(test_eval)
                for yᵢ in test_eval
                    nloglikelihood += - eval_logdens(root, yᵢ)
                end
            end
            nloglikelihood /= Ntest
            test_nloglikelihood[k, i] = nloglikelihood
        end
    end
    # println(".")

    # now choose the best lambdas
    test_nloglikelihood = dropdims(mean(test_nloglikelihood, dims=2), dims=2)
    best_lambdas = lambdas[argmin(test_nloglikelihood)]
    best_nloglikelihood = minimum(test_nloglikelihood)

    # compute overall likelihood
    nsims = 100
    samples = [rand(model, nsims) for model in models]
    λ1, λ2, η1, η2 = best_lambdas
    splits, betas = fit2(y, ptr, brks, λ1, λ2, η1, η2, istemporal, parallel=true)
    lows, mids, ups = splits.lows, splits.mids, splits.ups

    lls = zeros(num_nodes)
    for i in 1:num_nodes
        beta = betas[:, i]
        root = make_tree_from_bfs(lows, mids, ups, beta)
        ll = mean([-eval_logdens(root, yᵢ) for yᵢ in samples[i]])
        lls[i] = ll
    end
    validation_nloglikelihood = mean(lls)

    return Dict("best_lambdas" => best_lambdas,
                "cv_nloglikelihood" => best_nloglikelihood,
                "val_nloglikelihood" => validation_nloglikelihood)
end

function get_hypers(method)
    lambdas_dict = Dict(
        "fl" => [(10^l, 1e-12) for l in range(-2, log10(5.0), length=20)],
        "kal" => [(1e-12, 10^x) for x in range(-2, log10(5.0), length=20)],
        "enet" => [
            (10^l1, 10.0^x) for l1 in range(-2, log10(5.0), length=10) for x in range(-2, log10(5.0), length=10)
        ]
    )
    ls = lambdas_dict[method]
    hypers = [(λ1, λ2, η1, η2) for (λ1, λ2) in ls for (η1, η2) in ls]
    # if method == "enet"
    #     M = 400
    # else
    #     M = 100
    # end
    M = 10
    return rand(hypers, M)
end

##

function run_benchmarks(N, pmiss;
                        nsims=100,
                        nsplits=5,
                        tasks=("constant", "smooth", "mixed"),
                        outliers=false)
    experiment_results = []
    for task_space in tasks
        for task_time in tasks
            for method in ("fl", "kal", "enet")
                println("Running task_space $task_space task_time $task_time for method $method, outliers $outliers")

                for l in 1:nsims
                    D = generate_spt_task(task_space, task_time, N, pmiss, outliers=outliers)
                    y = vec(D["y"])
                    models = vec(D["dmodels"])
                    ndata = vec(D["ndata"])
                    devals = vec(D["devals"])
                    ptr = D["ptr"]
                    brks = D["brks"]
                    evalpts = D["evalpts"]
                    istemporal = D["istemporal"]

                    cvsets = generate_cvsets(y, nsplits)

                    lambdas = get_hypers(method)
                    results = cv_fit2(y, evalpts, ptr, brks, istemporal, lambdas, cvsets, models)

                    new_result = Dict(
                        :experiment => l,
                        :task_space => task_space,
                        :task_time => task_time,
                        :method => method,
                        :cv_nloglikelihood => results["cv_nloglikelihood"],
                        :val_nloglikelihood => results["val_nloglikelihood"])
                    push!(experiment_results, new_result)
                end
            end
        end
    end
    return experiment_results
end

# %%
# see examples
task = generate_spt_task("mixed", "mixed", 30, 0.8, outliers=false)
evalpts = task["evalpts"]
neval = length(evalpts)
devals = zeros(30, 30, neval)
for j in 1:30, i in 1:30
    devals[i, j, :] = task["devals"][i, j]
end
# fit on sample data tor example

λ1, λ2, η1, η2 = 0.25, 0.05, 0.25, 0.05
ptr, brks, istemporal = task["ptr"], task["brks"], task["istemporal"]
y = reshape(task["y"], :)
splits, betas = fit2(y, ptr, brks, λ1, λ2, η1, η2, istemporal)
betas = reshape(betas, :, 30, 30)
splitvals = splits.splitvals
fitprobs = zeros(30, 30, length(splitvals))
fitdens = zeros(30, 30, length(splitvals))
nsplits = size(betas, 1)
δx = diff(splitvals)
δx = [δx[1]; δx]
for i in 1:30, j in 1:30
    root = make_tree_from_bfs(splits.lows, splits.mids, splits.ups, betas[:, i, j])
    fitprobs[i, j, :] = [exp(eval_logprob(root, x)) for x in splitvals]
    fitdens[i, j, :] = [exp(eval_logdens(root, x)) for x in splitvals]
end

tosave = Dict(
    "evalpts" => evalpts,
    "devals" => devals,
    "ndata" => task["ndata"],
    "splitvals" => splitvals,
    "fitprobs" => fitprobs,
    "fitdens" => fitdens,
)
# npzwrite("simulations/example_simulation_task.npz", tosave)

##
# %%
N = 30
pmiss = 0.8
nsims = 1
tasks = ("smooth", "constant", "mixed")

experiment_results = run_benchmarks(N, pmiss, nsims=nsims, tasks=tasks)
experiment_results_out = run_benchmarks(N, pmiss, nsims=nsims, tasks=("mixed", ), outliers=true)

# %%
df = DataFrame(experiment = Int[],
               task_space=String[],
               task_time=String[],
               method=String[],
               cv_nloglikelihood=Float64[],
               cv_nloglikelihood=Float64[])
for record in [experiment_results; experiment_results_out]
    push!(df, record)
end
df.runid = fill(runid, size(df, 1))

fname = @sprintf("simulations/benchmarks/results-spt-%03d.csv", runid)
CSV.write(fname, df)

max_stars_repo_name: mauriciogtec/gfen-reproduce

max_issues_repo_name: mauriciogtec/gfen-reproduce

max_forks_repo_name: mauriciogtec/gfen-reproduce


text:
using LinearAlgebra

struct ValueOne; end
ValueOne()

# Compute X <- a X + b I.
function matfun_axpby!(X,a,b,Y::UniformScaling)
    m,n=size(X)
    if ~(a isa ValueOne)
        rmul!(X,a)
    end
    @inbounds for i=1:n
        X[i,i]+=(b isa ValueOne) ? 1 : b
    end
end

@inline function exp_mono_m5_opt_rho1_9(A)
    T=promote_type(eltype(A),Float64)
    A_copy=similar(A,T); A_copy .= A;
    return exp_mono_m5_opt_rho1_9!(A_copy)
end

@inline function exp_mono_m5_opt_rho1_9!(A)
    T=promote_type(eltype(A),Float64) # Make it work for many 'bigger' types (matrices and scalars)
    # max_memslots=8
    n=size(A,1)
    # The first slots are precomputed nodes [:A]
    memslots2 = similar(A,T)
    memslots3 = similar(A,T)
    memslots4 = similar(A,T)
    memslots5 = similar(A,T)
    memslots6 = similar(A,T)
    memslots7 = similar(A,T)
    memslots8 = similar(A,T)
    # Assign precomputed nodes memslots
    memslots1=A # overwrite A
    # Uniform scaling is exploited.
    # No matrix I explicitly allocated.
    # Computation order: Bb2 Ba2 B2 Ba3 Bb3 B3 Ba4 Bb4 B4 Ba5 Bb5 B5 Ba6 Bb6 B6 T2k8
    # Computing Bb2 = x*I+x*A
    coeff1=2.925822675582923
    coeff2=0.2429580790486332
    memslots2 .= coeff2.*memslots1
    mul!(memslots2,true,I*coeff1,true,true)
    # Computing Ba2 = x*I+x*A
    coeff1=2.925822675582923
    coeff2=0.2429580790486332
    memslots3 .= coeff2.*memslots1
    mul!(memslots3,true,I*coeff1,true,true)
    # Computing B2 with operation: mult
    mul!(memslots4,memslots3,memslots2)
    # Deallocating Ba2 in slot 3
    # Deallocating Bb2 in slot 2
    # Computing Ba3 = x*I+x*A+x*B2
    coeff1=1.6001548380533195
    coeff2=-0.9896362789001155
    coeff3=2.535123243045754
    memslots2 .= coeff2.*memslots1 .+ coeff3.*memslots4
    mul!(memslots2,true,I*coeff1,true,true)
    # Computing Bb3 = x*I+x*A+x*B2
    coeff1=2.477932656299724
    coeff2=0.12561938455295804
    coeff3=0.39676277541410987
    memslots3 .= coeff2.*memslots1 .+ coeff3.*memslots4
    mul!(memslots3,true,I*coeff1,true,true)
    # Computing B3 with operation: mult
    mul!(memslots5,memslots2,memslots3)
    # Deallocating Ba3 in slot 2
    # Deallocating Bb3 in slot 3
    # Computing Ba4 = x*I+x*A+x*B2+x*B3
    coeff1=0.11062999124786087
    coeff2=0.5847339177180232
    coeff3=1.5630308051212967
    coeff4=1.4346195499890089
    memslots2 .= coeff2.*memslots1 .+ coeff3.*memslots4 .+ coeff4.*memslots5
    mul!(memslots2,true,I*coeff1,true,true)
    # Computing Bb4 = x*I+x*A+x*B2+x*B3
    coeff1=3.2853771337963456
    coeff2=0.3808592456926523
    coeff3=-0.5832562723145828
    coeff4=0.01314523058371281
    memslots3 .= coeff2.*memslots1 .+ coeff3.*memslots4 .+ coeff4.*memslots5
    mul!(memslots3,true,I*coeff1,true,true)
    # Computing B4 with operation: mult
    mul!(memslots6,memslots2,memslots3)
    # Deallocating Ba4 in slot 2
    # Deallocating Bb4 in slot 3
    # Computing Ba5 = x*I+x*A+x*B2+x*B3+x*B4
    coeff1=0.019638592876749326
    coeff2=-0.06879708908316624
    coeff3=0.38255885470355305
    coeff4=1.195010535882028
    coeff5=1.3539321663418642
    memslots2 .= coeff2.*memslots1 .+ coeff3.*memslots4 .+ coeff4.*memslots5 .+ coeff5.*memslots6
    mul!(memslots2,true,I*coeff1,true,true)
    # Computing Bb5 = x*I+x*A+x*B2+x*B3+x*B4
    coeff1=1.7177215564670094
    coeff2=1.356661682335682
    coeff3=-1.283038970541202
    coeff4=0.04857205018560095
    coeff5=0.0010901570728226485
    memslots3 .= coeff2.*memslots1 .+ coeff3.*memslots4 .+ coeff4.*memslots5 .+ coeff5.*memslots6
    mul!(memslots3,true,I*coeff1,true,true)
    # Computing B5 with operation: mult
    mul!(memslots7,memslots2,memslots3)
    # Deallocating Ba5 in slot 2
    # Deallocating Bb5 in slot 3
    # Computing Ba6 = x*I+x*A+x*B2+x*B3+x*B4+x*B5
    coeff1=-0.004633918773394507
    coeff2=0.03077253772413558
    coeff3=0.0048951603098945345
    coeff4=0.39705603974476195
    coeff5=0.6636659134668268
    coeff6=1.283915906769389
    memslots2 .= coeff2.*memslots1 .+ coeff3.*memslots4 .+ coeff4.*memslots5 .+ coeff5.*memslots6 .+ coeff6.*memslots7
    mul!(memslots2,true,I*coeff1,true,true)
    # Computing Bb6 = x*I+x*A+x*B2+x*B3+x*B4+x*B5
    coeff1=0.00017310181177139793
    coeff2=0.7091796380446608
    coeff3=-0.8509852584224694
    coeff4=0.029768444198531262
    coeff5=0.002161503231210118
    coeff6=2.4473617044179496e-5
    memslots3 .= coeff2.*memslots1 .+ coeff3.*memslots4 .+ coeff4.*memslots5 .+ coeff5.*memslots6 .+ coeff6.*memslots7
    mul!(memslots3,true,I*coeff1,true,true)
    # Computing B6 with operation: mult
    mul!(memslots8,memslots2,memslots3)
    # Deallocating Ba6 in slot 2
    # Deallocating Bb6 in slot 3
    # Computing T2k8 = x*I+x*A+x*B2+x*B3+x*B4+x*B5+x*B6
    coeff1=1.089645716755357
    coeff2=-0.05894345454785517
    coeff3=1.9565192611905264
    coeff4=-0.07941140929885598
    coeff5=0.06299059278775863
    coeff6=0.022010502754629162
    coeff7=0.0020981712432435614
    # Smart lincomb recycle A
    memslots1 .= coeff2.*memslots1 .+ coeff3.*memslots4 .+ coeff4.*memslots5 .+ coeff5.*memslots6 .+ coeff6.*memslots7 .+ coeff7.*memslots8
    mul!(memslots1,true,I*coeff1,true,true)
    # Deallocating B2 in slot 4
    # Deallocating B3 in slot 5
    # Deallocating B4 in slot 6
    # Deallocating B5 in slot 7
    # Deallocating B6 in slot 8
    return memslots1 # Returning T2k8
end

max_stars_repo_name: matrixfunctions/GraphMatFunData

max_issues_repo_name: matrixfunctions/GraphMatFunData

max_forks_repo_name: matrixfunctions/GraphMatFunData


text:
module SpatialJackknife

using Statistics, LinearAlgebra
using NearestNeighbors, Distances

export jackknife, get_subvols

"""
This function gets dimensions of a given data array and returns an error if
they are not appropriate or consistent with the dimension of optional randoms.
"""
function get_dims(data::Array{Float64, 2},
                  randmask::Array{Float64, 2} = ones(0, 0))::Tuple{Int, Int}

    ndat, ndims = size(data)

    if ndat < ndims
        throw(ArgumentError("data array must have shape (ndat, ndims)"))
    end

    if !isempty(randmask)
        if ndims != size(randmask)[2]
            throw(ArgumentError("data array and mask array must have same number of dimensions"))
        end
    end

    ndat, ndims
end

"""
This function takes a dataset and returns an array of subvolume indices for
each datapoint. In this case, it takes a set of random mask points for
dividing arbitrary volumes. It computes the quantiles of the mask and then
finds the closest mask point to each data point to determine the subvolume
the data point belongs in.

The function expects a data array with shape (ndat, ndims) for Ndat samples in
ndims dimensions. The mask array should have shape (nmask, ndims) and then
there is an integere side_divs which specifies the number of times each
dimension is divided up to make side_divs^ndims subvolumes.
"""
function get_subvols(data::Array{Float64, 2},
                     randmask::Array{Float64, 2},
                     side_divs::Int,
                     metric::Metric = Euclidean())::Array{Int, 1}

    ndat, ndims = get_dims(data, randmask)
    nmask = size(randmask)[1]

    # now to figure out which subvolume each of the mask points is in by
    # getting the quantiles in each dimension
    maskvols = ones(Int, nmask)
    volnums = [1]  # start off with all points in the same volume
    for i in 1:ndims

        # creates a set of numbers to assign volumes for the ith dimension
        nuvolnums = [(a - 1) * side_divs^(ndims - i) for a in 1:(side_divs^i)] .+ 1

        for (j, volnum) in enumerate(volnums)
            dimmask = findall(maskvols .== volnum)
            dimquants = quantile(randmask[dimmask, i],
                                 range(0, stop = 1, length = side_divs + 1))
            dimquants[end] += 1

            for k in 2:side_divs
                submask = findall(dimquants[k] .<= randmask[dimmask, i] .< dimquants[k + 1])
                maskvols[dimmask[submask]] .= nuvolnums[side_divs * (j - 1) + k]
            end
        end

        volnums = nuvolnums
    end

    # and find the closest random to each data point to find the subvolume
    balltree = BallTree(collect(randmask'), metric)
    nninds, dists = knn(balltree, collect(data'), 1)

    subvols = maskvols[[n[1] for n in nninds]]
    subvols
end

"""
As with the other method for get_subvols, this function computes the subvolume
indices for given samples. It assumes a constant density throughout the
volume. The method computes the subvolumes based on a set of extrema for
sample values in each dimension and a number of volumes to divide the sample
on per side. If a single set of extrema are given, the assumption is that they are the same in each dimension. Alternatively, it can take a set of volume
edges in each dimension.
"""
function get_subvols(data::Array{Float64, 2};
                     side_divs::Int = 3,
                     edges::Array{Array{Float64, 1}, 1} = Array{Array{Float64, 1}, 1}(undef, 0))::Array{Int, 1}

    ndat, ndims = get_dims(data)

    if length(edges) == 1
        edges = [edges[1] for a in 1:ndims]
    elseif length(edges) != ndims
        throw(ArgumentError("edges argument must have either one array or as many arrays as dimensions in the data"))
    end

    if !all(length.(edges) .== length(edges[1]))
        throw(ArgumentError("expect the edges arrays to all have the same length since volume is a cube"))
    end

    if length(edges[1]) == 2
        edges = [collect(range(ed[1], stop = ed[2],
                               length = side_divs + 1)) for ed in edges]
    end

    # now to use the edges to assign subvolume indices
    subvols = ones(Int, ndat)
    diminds = ones(Int, ndims)
    dimfacs = [side_divs^(ndims - a) for a in 1:ndims]
    for i in 1:ndat
        for j in 1:ndims
            diminds[j] = maximum(findall(data[i, j] .>= edges[j]))
        end
        subvols[i] = sum(dimfacs .* (diminds .- 1)) + 1
    end

    subvols
end

"""
This function computes the jackknife mean and variance over the observations
computed from the data with obsfunc. They are computed by repeatedly applying
obsfunc to the data with one of the subvolumes removed. It therefore assumes
that obsfunc has normalised the observable for that remaining volume. The
function obsfunc must be callable with the form

obsfunc(data::Array{Float64, 2}, args...)

for optional args tuple and the value returned by obsfunc must be in the form
of an array of floats. If any of the returned values are NaNs, an error will be
raised. By default, the covariance matrix is computed for observables in more
than one dimension, but setting covar to false will result in only the diagonal
variances being returned.
"""
function jackknife(obsfunc::Function,
                   data::Array{Float64, 2},
                   subvolinds::Array{Int, 1},
                   args::Tuple = ();
                   covar::Bool = true,
                   pass_vol = false)::Tuple

    volset = BitSet(subvolinds)
    nvols = length(volset)
    meanfac = 1 / nvols
    varfac = (nvols - 1) / nvols

    # get the observables for the data with each volume left out
    local ivals = Array{Array{Float64, 1}, 1}(undef, nvols)
    for (i, vol) in enumerate(volset)
        if pass_vol
            argsvol = (args..., vol)
            ivals[i] = obsfunc(data[findall(subvolinds .!= vol), :],
                               argsvol...)
        else
            ivals[i] = obsfunc(data[findall(subvolinds .!= vol), :], args...)
        end
    end
    meanvals = sum(ivals) .* meanfac

    if any(isnan.(meanvals))
        throw(ErrorException("nans returned by obsfunc, perhaps the data is too sparse for the selected size/number of subvolumes"))
    end

    ndims = length(meanvals)

    # now to get the (co)variance
    covarmat = zeros(ndims, ndims)
    for i in 1:ndims
        for j in i:ndims
            if !covar && i != j
                continue
            end

            covarmat[i, j] = varfac * sum([(ivals[k][i] - meanvals[i]) *
                                           (ivals[k][j] - meanvals[j]) for k in 1:nvols])
            covarmat[j, i] = covarmat[i, j]
        end
    end

    if ndims == 1
        return meanvals, covarmat[1, 1]
    elseif !covar
        return meanvals, diag(covarmat)
    else
        return meanvals, covarmat
    end
end

end


text:
#=
Buffers are just normal plots, with the benefit of being setup correctly to
efficiently append + push new values to them
=#

function LinesegmentBuffer(
        scene::SceneLike, ::Type{Point{N}} = Point{2};
        color = RGBAf0[], linewidth = Float32[], raw = true,
        kw_args...
    ) where N
    linesegments!(
        scene, Point{N, Float32}[]; color = color,
        linewidth = linewidth, raw = raw, kw_args...
    ).plots[end]
end

function append!(lsb::LineSegments, positions::Vector{Point{N, Float32}}; color = :black, linewidth = 1.0) where N
    thickv = same_length_array(positions, linewidth, key"linewidth"())
    colorv = same_length_array(positions, color, key"color"())
    append!(lsb[1][], positions)
    append!(lsb[:color][], colorv)
    append!(lsb[:linewidth][], thickv)
    return
end

function push!(tb::LineSegments, positions::Point{N, Float32}; kw_args...) where N
    append!(tb, [positions]; kw_args...)
end

function start!(lsb::LineSegments)
    resize!(lsb[1][], 0)
    resize!(lsb[:color][], 0)
    resize!(lsb[:linewidth][], 0)
    return
end

function finish!(lsb::LineSegments)
    # update the signal!
    lsb[1][] = lsb[1][]
    lsb[:color][] = lsb[:color][]
    lsb[:linewidth][] = lsb[:linewidth][]
    return
end

function TextBuffer(
        scene::SceneLike, ::Type{Point{N}} = Point{2};
        rotation = [Quaternionf0(0,0,0,1)],
        color = RGBAf0[RGBAf0(0,0,0,0)],
        textsize = Float32[0],
        font = [defaultfont()],
        align = [Vec2f0(0)],
        raw = true,
        kw_args...
    ) where N
    annotations!(
        scene, String[" "], [Point{N, Float32}(0)];
        rotation = rotation,
        color = color,
        textsize = textsize,
        font = font,
        align = align,
        raw = raw,
        kw_args...
    ).plots[end]
end

function start!(tb::Annotations)
    for key in (1, :color, :rotation, :textsize, :font, :align)
        empty!(tb[key][])
    end
    return
end

function finish!(tb::Annotations)
    # update the signal!
    # now update all callbacks
    # TODO this is a bit shaky, buuuuhut, in theory the whole lift(color, ...)
    # in basic_recipes annotations should depend on all signals here, so updating one should be enough
    if length(tb[1][]) != length(tb.textsize[])
        error("Inconsistent buffer state for $(tb[1][])")
    end
    notify!(tb[1])
    return
end

function push!(tb::Annotations, text::String, position::VecTypes{N}; kw_args...) where N
    append!(tb, [(String(text), Point{N, Float32}(position))]; kw_args...)
end

function append!(tb::Annotations, text::Vector{String}, positions::Vector{Point{N, Float32}}; kw_args...) where N
    text_positions = convert_argument(Annotations, text, positions)[1]
    append!(tb, text_positions; kw_args...)
    return
end

function append!(tb::Annotations, text_positions::Vector{Tuple{String, Point{N, Float32}}}; kw_args...) where N
    append!(tb[1][], text_positions)
    kw = Dict(kw_args)
    for key in (:color, :rotation, :textsize, :font, :align)
        val = get(kw, key) do
            isempty(tb[key][]) && error("please provide default for $key")
            return last(tb[key][])
        end
        val_vec = same_length_array(text_positions, val, Key{key}())
        append!(tb[key][], val_vec)
    end
    return
end