text:
# extend the CUDA.jl functionality (rand, randn, rand_poisson, etc.)
# to include binomial distributions

const BinomialType = Union{Type{<:Integer}}
const BinomialArray = DenseCuArray{<:Integer}

## exported functions: in-place
rand_binomial!(A::AnyCuArray; kwargs...) =
    error("BinomialGPU.jl does not support generating
    binomially-distributed random numbers of type $(eltype(A))")

## unexported functions: out of place
rand_binomial(T::BinomialType, dims::Dims; kwargs...) = rand_binomial(T, dims; kwargs...)

rand_binomial(T::BinomialType, dim1::Integer, dims::Integer...; kwargs...) =
    rand_binomial(T, Dims((dim1, dims...)); kwargs...)

rand_binomial(T::Type, dims::Dims; kwargs...) =
    rand_binomial!(CuArray{T}(undef, dims...); kwargs...)

rand_binomial(T::Type, dim1::Integer, dims::Integer...; kwargs...) =
    rand_binomial!(CuArray{T}(undef, dim1, dims...); kwargs...)

rand_binomial(dim1::Integer, dims::Integer...; kwargs...) =
    rand_binomial(Int, Dims((dim1, dims...)); kwargs...)

## main internal function
function rand_binomial!(A::BinomialArray; count, prob)
    return rand_binom!(A, count, prob)
end

## dispatching on parameter types

# constant parameters
function rand_binom!(A::BinomialArray, count::Integer, prob::Number)
    # revert to full parameter case (this could be suboptimal,
    # as a table-based method should in principle be faster)
    ns = CUDA.fill(Int(count), size(A))
    ps = CUDA.fill(Float32(prob), size(A))
    return rand_binom!(A, ns, ps)
end

# arrays of parameters
function rand_binom!(A::BinomialArray, count::BinomialArray, prob::Number)
    # revert to full parameter case (this could be suboptimal,
    # as a table-based method should in principle be faster)
    cucount = cu(count)
    ps = CUDA.fill(Float32(prob), size(A))
    return rand_binom!(A, cucount, ps)
end

function rand_binom!(A::BinomialArray, count::Integer, prob::AbstractArray{<:Number})
    # revert to full parameter case (this could be suboptimal,
    # as a table-based method should in principle be faster)
    ns = CUDA.fill(Int(count), size(A))
    cuprob  = cu(prob)
    return rand_binom!(A, ns, cuprob)
end

function rand_binom!(A::BinomialArray, count::BinomialArray, prob::AbstractArray{<:Number})
    cucount = cu(count)
    cuprob  = cu(prob)
    return rand_binom!(A, cucount, cuprob)
end

function rand_binom!(A::BinomialArray, count::BinomialArray, prob::DenseCuArray{Float32})
    if ndims(count) > ndims(A) || ndims(prob) > ndims(A)
        throw(DimensionMismatch("`count` and `prob` need to be scalar
                or have less or equal dimensions than A"))
        return A
    end
    if size(A)[1:ndims(count)] == size(count) && size(A)[1:ndims(prob)] == size(prob)
        count_dim_larger_than_prob_dim = ndims(count) > ndims(prob)
        if count_dim_larger_than_prob_dim
            # indices for count
            R1 = CartesianIndices(prob)
            # indices for prob that are not included in R1
            R2 = CartesianIndices(size(count)[ndims(prob)+1:end])
            # remaining indices in A
            Rr = CartesianIndices(size(A)[ndims(count)+1:end])
        else
            # indices for count
            R1 = CartesianIndices(count)
            # indices for prob that are not included in R1
            R2 = CartesianIndices(size(prob)[ndims(count)+1:end])
            # remaining indices in A
            Rr = CartesianIndices(size(A)[ndims(prob)+1:end])#
        end
        Rp = CartesianIndices((length(R1), length(R2))) # indices for parameters
        Ra = CartesianIndices((length(Rp), length(Rr))) # indices for parameters and A

        seed = rand(UInt32)

        kernel  = @cuda launch=false kernel_BTRS!(
                                        A, count, prob,
                                        R1, R2, Rp, Ra,
                                        count_dim_larger_than_prob_dim,
                                        seed
                                     )
        config  = launch_configuration(kernel.fun)
        threads = Base.min(length(A), config.threads, 256)
        blocks  = cld(length(A), threads)

        kernel(
            A, count, prob,
            R1, R2, Rp, Ra,
            count_dim_larger_than_prob_dim,
            seed;
            threads = threads, blocks=blocks
         )
    else
        throw(DimensionMismatch("`count` and `prob` need have size compatible with A"))
    end
    return A
end


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

# REQUIRES: x86_64
# RUN: julia --startup-file=no %s %t -O && llvm-link -S %t/* -o %t/module.ll
# RUN: cat %t/module.ll | llc - -mtriple=x86_64-- -mattr=fma | FileCheck %s

## Notes:
# This script uses the `emit` function (defined llvmpasses.jl) to emit either
# optimized or unoptimized LLVM IR. Each function is emitted individually and
# `llvm-link` is used to create a single module that can be passed to opt.
# The order in which files are emitted and linked is important since `lit` will
# process the test cases in order.

include(joinpath("..", "testhelpers", "llvmpasses.jl"))

# CHECK-LABEL: julia_dotf_
function dotf(a, b)
    s = 0.0
    @inbounds @simd for i ∈ eachindex(a)
# CHECK:    vfmadd231pd
# CHECK-NEXT:    vfmadd231pd
# CHECK-NEXT:    vfmadd231pd
# CHECK-NEXT:    vfmadd231pd
        s += a[i] * b[i]
    end
    s
end

emit(dotf, Vector{Float64}, Vector{Float64})

ext: jl

max_stars_repo_name: vanillajonathan/julia

max_issues_repo_name: vanillajonathan/julia

max_forks_repo_name: vanillajonathan/julia


text:
using MNISTDataset
using Base.Test

# write your own tests here
@test 1 == 1

labels = parseLabels("test")
@test typeof(labels) == Array{Float64, 2}
@test size(labels)   == (10000, 10)
@test labels[1,1]    == 0.0
@test labels[1,7]    == 1.0

labels = parseLabels("test", multi=true)
@test typeof(labels) == Array{Float64, 2}
@test size(labels)   == (10000, 10)
@test labels[1,1]    == 0.0
@test labels[1,7]    == 1.0

labels = parseLabels("test", multi=false)
@test typeof(labels) == Array{UInt8, 1}
@test size(labels)   == (10000,)
@test labels[1]      == 0x07

labels = parseLabels("test", useBetterValues=true)
@test typeof(labels) == Array{Float64, 2}
@test size(labels)   == (10000, 10)
@test labels[1,1]    == 0.1
@test labels[1,7]    == 0.9

labels = parseLabels("test", useBetterValues=false)
@test typeof(labels) == Array{Float64, 2}
@test size(labels)   == (10000, 10)
@test labels[1,1]    == 0.0
@test labels[1,7]    == 1.0

labels = parseLabels("test", multi=true, useBetterValues=true)
@test typeof(labels) == Array{Float64, 2}
@test size(labels)   == (10000, 10)
@test labels[1,1]    == 0.1
@test labels[1,7]    == 0.9

labels = parseLabels("test", multi=false, useBetterValues=true)
@test typeof(labels) == Array{UInt8, 1}
@test size(labels)   == (10000,)
@test labels[1]      == 0x07

labels = parseLabels("test", multi=false, useBetterValues=false)
@test typeof(labels) == Array{UInt8, 1}
@test size(labels)   == (10000,)
@test labels[1]      == 0x07

labels = parseLabels("train")
@test typeof(labels) == Array{Float64, 2}
@test size(labels)   == (60000, 10)

images = parseImages("test")
@test typeof(images) == Array{Float64, 2}
@test size(images)   == (10000, 784)
@test images[1, 1]   == 0.0
@test images[1, 3]   == 0.9921568627450981

images = parseImages("test"; useDecimal=true)
@test typeof(images) == Array{Float64, 2}
@test size(images)   == (10000, 784)
@test images[1, 1]   == 0.0
@test images[1, 3]   == 0.9921568627450981

images = parseImages("test"; useDecimal=false)
@test typeof(images) == Array{UInt8, 2}
@test size(images)   == (10000, 784)
@test images[1, 1]   == 0x00
@test images[1, 3]   == 0xfd

images = parseImages("test"; addBias=true)
@test typeof(images) == Array{Float64, 2}
@test size(images)   == (10000, 785)
@test images[1, 1]   == 1.0
@test images[1, 2]   == 0.0
@test images[1, 4]   == 0.9921568627450981

images = parseImages("test"; addBias=false)
@test typeof(images) == Array{Float64, 2}
@test size(images)   == (10000, 784)
@test images[1, 1]   == 0.0
@test images[1, 3]   == 0.9921568627450981

images = parseImages("test"; useDecimal=true, addBias=true)
@test typeof(images) == Array{Float64, 2}
@test size(images)   == (10000, 785)
@test images[1, 1]   == 1.0
@test images[1, 2]   == 0.0
@test images[1, 4]   == 0.9921568627450981

images = parseImages("test"; useDecimal=true, addBias=false)
@test typeof(images) == Array{Float64, 2}
@test size(images)   == (10000, 784)
@test images[1, 1]   == 0.0
@test images[1, 3]   == 0.9921568627450981

images = parseImages("test"; useDecimal=false, addBias=true)
@test typeof(images) == Array{UInt8, 2}
@test size(images)   == (10000, 785)
@test images[1, 1]   == 0xff
@test images[1, 2]   == 0x00
@test images[1, 4]   == 0xfd

images = parseImages("test"; useDecimal=false, addBias=false)
@test typeof(images) == Array{UInt8, 2}
@test size(images)   == (10000, 784)
@test images[1, 1]   == 0x00
@test images[1, 3]   == 0xfd

images = parseImages("train")
@test typeof(images) == Array{Float64, 2}
@test size(images)   == (60000, 784)
@test images[1, 1]   == 0.0
@test images[1, 782] == 0.5333333333333333


text:
#n = 5
n = 4
m = 3
# Investment cost
#ic = [16, 5, 32, 2, 0] * 8760
if n == 5
    ic = [16, 5, 32, 2, 0]
else
    ic = [16, 5, 32, 2]
end
# Fuel cost
if n == 5
    C = [25, 80, 6.5, 160, 1000]
else
    C = [25, 80, 6.5, 160]
end
# Duration
#T = [8760, 7000, 1500]
T = [8760, 7000, 1500] / 8760
# Height
D10x = diff([0, 3919, 7329, 10315])
Dref = diff([0, 7086, 9004, 11169])
D21 = D10x
D22 = Dref
D2 = [D21  D22]
D31 = D10x
D32 = Dref
D33 = D10x
D34 = Dref
p2 = [0.9, 0.1]
p3 = [0.81, 0.09, 0.09, 0.01]
nstages = 2


text:
using REPL

abstract type CharStyle end
struct Block <: CharStyle end
struct Braile <: CharStyle end

"""
    REPLOutput <: GraphicOutput

    REPLOutput(init; tspan, kw...)

An output that is displayed directly in the REPL. It can either store or discard
simulation frames.

# Arguments:
- `init`: initialisation `Array` or `NamedTuple` of `Array`s.

# Keywords

- `tspan`: `AbstractRange` timespan for the simulation
- `mask`: `BitArray` for defining cells that will/will not be run.
- `aux`: `NamedTuple` of arbitrary input data. Use `get(data, Aux(:key), I...)`
    to access from a `Rule` in a type-stable way.
- `padval`: padding value for grids with neighborhood rules. The default is `zero(eltype(init))`.
- `fps`: `Real` frames per second to display the simulation
- `store`: `Bool` whether ot store the simulation frames for later use
- `color`: a color from Crayons.jl
- `cutoff`: `Real` cutoff point to display a full or empty cell. Default is `0.5`
- `style`: `CharStyle` `Block()` or `Braile()` printing. `Braile` uses 1/4 the screen space of `Block`.

"""
mutable struct REPLOutput{T,F<:AbstractVector{T},E,GC,Co,St,Cu} <: GraphicOutput{T,F}
    frames::F
    running::Bool
    extent::E
    graphicconfig::GC
    color::Co
    style::St
    cutoff::Cu
end
function REPLOutput(;
    frames, running, extent, graphicconfig,
    color=:white, cutoff=0.5, style=Block(), kw...
)
    REPLOutput(frames, running, extent, graphicconfig, color, style, cutoff)
end

function showframe(frame::AbstractArray, o::REPLOutput, data::AbstractSimData)
    # Print the frame
    _print_to_repl((0, 0), o.color, _replframe(o, frame))
    # Print the timestamp in the top right corner
    _print_to_repl((0, 0), o.color, string("Time $(currenttime(data))"))
end

# Terminal commands
_savepos(io::IO=terminal.out_stream) = print(io, "\x1b[s")
_restorepos(io::IO=terminal.out_stream) = print(io, "\x1b[u")
_movepos(io::IO, c=(0,0)) = print(io, "\x1b[$(c[2]);$(c[1])H")
_cursor_hide(io::IO=terminal.out_stream) = print(io, "\x1b[?25l")
_cursor_show(io::IO=terminal.out_stream) = print(io, "\x1b[?25h")

_print_to_repl(pos, c::Symbol, s::String) = _print_to_repl(pos, Crayon(foreground=c), s)
function _print_to_repl(pos, color::Crayon, str::String)
    io = terminal.out_stream
    _savepos(io)
    _cursor_hide(io)
    _movepos(io, pos)
    print(io, color)
    print(io, str)
    _cursor_show(io)
    _restorepos(io)
end

const YBRAILE = 4
const XBRAILE = 2
const YBLOCK = 2
const XBLOCK = 1

function _replframe(o, frame)
    ystep, xstep, f = _chartype(o)

    # Limit output area to available terminal size.
    dispy, dispx = displaysize(stdout)
    youtput, xoutput = outputsize = size(frame)
    yoffset, xoffset = (0, 0)

    yrange = max(1, ystep * yoffset):min(youtput, ystep * (dispy + yoffset - 1))
    xrange = max(1, xstep * xoffset):min(xoutput, xstep * (dispx + xoffset - 1))
    window = view(adapt(Array, frame), yrange, xrange)
    f(window, o.cutoff)
end

_chartype(o::REPLOutput) = _chartype(o.style)
_chartype(s::Braile) = YBRAILE, XBRAILE, brailize
_chartype(s::Block) = YBLOCK, XBLOCK, blockize


text:
#=
This script generates the "pitfall" data from raw input data
=#

using DrWatson
@quickactivate :FractalDimension # uses DynamicalSystems, PyPlot
include(srcdir("style.jl"))
using DelimitedFiles

# %% nifty 50
name = "nifty50.csv"
file = datadir("pitfall", name)
X = readdlm(file, ','; skipstart = 1)
xx = X[:, 2]
# fix "null" nonsense
for (i, e) in enumerate(xx)
    e isa AbstractFloat && continue
    js = i+1:length(xx)
    j = findfirst(ι -> xx[ι] isa AbstractFloat, js)
    xx[i] = xx[js[j]]
end
x = float.(xx)
file = datadir("pitfall", name)
writedlm(datadir("experimental", "nifty50.txt"), x)

# Run this after `x` is obtained
# Should we de-trend here?
estimate_embedding(x, "nifty50")

# %% Vostok
name = "vostok.csv"
file = datadir("pitfall", name)
X = readdlm(file, '\t')
tim = X[:, 1]
δT = X[:, 2]

using Dierckx
# interpolation object:
spl = Spline1D(tim, δT; k=3, bc="extrapolate")
# equi-spaced time vector:
t = range(minimum(tim), maximum(tim); length = 2000)
T = spl(t)
writedlm(datadir("experimental", "vostok.txt"), T)

# Run this after `x` is obtained
# Should we de-trend here?
estimate_embedding(T, "vostok")

max_stars_repo_name: Datseris/FractalDimension

max_issues_repo_name: Datseris/FractalDimension

max_forks_repo_name: Datseris/FractalDimension


text:
export Intent
export Help
export NothingIntent
export CreateTrial

"""
In general intents are an arbitrary string. But for this specific application
there are a finite set of Intents and it's convenient to represent them as an
enum.
NOTE the names here must literally match the names shown in the Lex dashboard.
"""
abstract type Intent end

"""
The Edgar Lex bot uses the 'Help' intent when replying to questions like
"help", and "what can you do".
"""
struct Help <: Intent end

"""
The Edgar Lex bot uses the 'NothingIntent' to quickly exit a conversation.
The user can trigger this with phrases like "nevermind", and "forget it".
"""
struct NothingIntent <: Intent end

"""
The Edgar Lex bot uses the 'CreateTrial' to schedule a new trial.
The conversation for this intent requires the user to give values for the
date, time, and trial title slots.
"""
struct CreateTrial <: Intent end

const IntentNames = Dict{String, Intent}(
    "CreateTrial" => CreateTrial(),
    "Help" => Help(),
    "Nothing" => NothingIntent()
)


text:
module PGN

using ..Chess

export PGNException, PGNReader

export gamesinfile, gotonextgame!, gamefromstring, gametopgn, readgame

struct PGNException <: Exception
    message::String
end

"""
    PGNReader

A type for reading PGN data from a stream.
"""
mutable struct PGNReader
    io::IO
    unreadchar::Union{Char,Nothing}
end

"""
    PGNReader(io::IO)

Initializes a `PGNReader` from an `IO` object.
"""
PGNReader(io::IO) = PGNReader(io, nothing)

@enum TokenType begin
    str
    integer
    period
    asterisk
    leftbracket
    rightbracket
    leftparen
    rightparen
    leftangle
    rightangle
    nag
    symbol
    comment
    linecomment
    nullmove
    endoffile
end

struct Token
    ttype::TokenType
    value::String
end

function terminatesgame(t::Token)
    t.ttype == asterisk ||
        t.ttype == eof ||
        (
            t.ttype == symbol &&
            (t.value == "1-0" || t.value == "0-1" || t.value == "1/2-1/2")
        )
end

function issymbolstart(c::Char)::Bool
    occursin(c, "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-")
end

function issymbolcontinuation(c::Char)::Bool
    occursin(c, "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_+#=:-/")
end

function readchar(p::PGNReader)::Char
    if !isnothing(p.unreadchar)
        result = p.unreadchar
        p.unreadchar = nothing
        result
    else
        read(p.io, Char)
    end
end

function unread(p::PGNReader, c::Char)
    p.unreadchar = c
end

function peek(p::PGNReader)::Char
    result = readchar(p)
    unread(p, result)
    result
end

function skipwhitespace(p::PGNReader)
    while !eof(p.io)
        c = readchar(p)
        if !isspace(c)
            unread(p, c)
            break
        end
    end
end

function readstring(p::PGNReader)::Token
    c = readchar(p)
    @assert c == '"'
    result = IOBuffer()

    while true
        if eof(p.io)
            throw(PGNException("EOF while reading string"))
        end
        pc = c
        c = readchar(p)
        if c == '"'
            if pc == '\\'
                write(result, c)
            else
                break
            end
        elseif c == '\\'
            if pc == '\\'
                write(result, c)
            end
        else
            write(result, c)
        end
    end
    Token(str, String(take!(result)))
end

function readsymbol(p::PGNReader)::Token
    c = readchar(p)
    @assert issymbolstart(c)

    tt = isdigit(c) ? integer : symbol
    result = IOBuffer()
    write(result, c)

    while !eof(p.io)
        c = readchar(p)
        if !issymbolcontinuation(c)
            break
        end
        write(result, c)
        if !isdigit(c)
            tt = symbol
        end
    end
    unread(p, c)
    Token(tt, String(take!(result)))
end

function readperiod(p::PGNReader)::Token
    c = readchar(p)
    @assert c == '.'
    Token(period, ".")
end

function readasterisk(p::PGNReader)::Token
    c = readchar(p)
    @assert c == '*'
    Token(asterisk, "*")
end

function readleftbracket(p::PGNReader)::Token
    c = readchar(p)
    @assert c == '['
    Token(leftbracket, "[")
end

function readrightbracket(p::PGNReader)::Token
    c = readchar(p)
    @assert c == ']'
    Token(rightbracket, "]")
end

function readleftparen(p::PGNReader)::Token
    c = readchar(p)
    @assert c == '('
    Token(leftparen, "(")
end

function readrightparen(p::PGNReader)::Token
    c = readchar(p)
    @assert c == ')'
    Token(rightparen, ")")
end

function readleftangle(p::PGNReader)::Token
    c = readchar(p)
    @assert c == '<'
    Token(leftangle, "<")
end

function readrightangle(p::PGNReader)::Token
    c = readchar(p)
    @assert c == '>'
    Token(rightangle, ">")
end

function readnag(p::PGNReader)::Token
    c = readchar(p)
    @assert c == '$'
    result = IOBuffer()
    c = readchar(p)
    while isdigit(c)
        write(result, c)
        if eof(p.io)
            break
        end
        c = readchar(p)
    end
    unread(p, c)
    Token(nag, String(take!(result)))
end

function readfakenag(p::PGNReader)::Token
    c = readchar(p)
    @assert c == '!' || c == '?'
    result = IOBuffer()
    write(result, c)
    c = readchar(p)
    while c == '!' || c == '?'
        write(result, c)
        if eof(p.io)
            break
        end
        c = readchar(p)
    end
    unread(p, c)
    s = String(take!(result))
    if startswith(s, "!!")
        Token(nag, "3")
    elseif startswith(s, "??")
        Token(nag, "4")
    elseif startswith(s, "!?")
        Token(nag, "5")
    elseif startswith(s, "?!")
        Token(nag, "6")
    elseif startswith(s, "!")
        Token(nag, "1")
    else
        Token(nag, "2")
    end
end

function readcomment(p::PGNReader)::Token
    c = readchar(p)
    @assert c == '{'
    result = IOBuffer()
    c = readchar(p)
    while c ≠ '}'
        write(result, c)
        if eof(p.io)
            throw(PGNException("EOF while reading comment"))
        end
        c = readchar(p)
    end
    Token(comment, String(take!(result)))
end

function readlinecomment(p::PGNReader)::Token
    c = readchar(p)
    @assert c == ';'
    result = IOBuffer()
    c = readchar(p)
    while c ≠ '\n'
        write(result, c)
        if eof(p.io)
            break
        end
        c = readchar(p)
    end
    Token(linecomment, String(take!(result)))
end

function readnullmove(p::PGNReader)::Token
    c = readchar(p)
    @assert c == '-'
    while c == '-'
        c = readchar(p)
    end
    Token(nullmove, "--")
end

function readtoken(p::PGNReader)::Token
    skipwhitespace(p)

    if eof(p.io)
        return Token(endoffile, "")
    end

    c = peek(p)

    if c == '"'
        readstring(p)
    elseif c == '-'
        readnullmove(p)
    elseif issymbolstart(c)
        readsymbol(p)
    elseif c == '.'
        readperiod(p)
    elseif c == '*'
        readasterisk(p)
    elseif c == '['
        readleftbracket(p)
    elseif c == ']'
        readrightbracket(p)
    elseif c == '('
        readleftparen(p)
    elseif c == ')'
        readrightparen(p)
    elseif c == '<'
        readleftangle(p)
    elseif c == '>'
        readrightangle(p)
    elseif c == '{'
        readcomment(p)
    elseif c == ';'
        readlinecomment(p)
    elseif c == '$'
        readnag(p)
    elseif c == '!' || c == '?'
        readfakenag(p)
    else
        throw(PGNException("Invalid character: " * c))
    end
end

"""
    gotonextgame!(p::PGNReader)::Bool

Tries to go to the next game, returns `true` on success.
"""
function gotonextgame!(p::PGNReader)::Bool
    while true
        t = readtoken(p)
        if t.ttype == endoffile
            return false
        end
        if t.ttype == leftbracket
            unread(p, '[')
            return true
        end
    end
    false
end

function readheader(p::PGNReader)::Tuple{String,String}
    lb = readtoken(p)
    k = readtoken(p)
    v = readtoken(p)
    rb = readtoken(p)
    skipwhitespace(p)
    if lb.ttype == leftbracket &&
       rb.ttype == rightbracket &&
       k.ttype == symbol &&
       v.ttype == str
        (k.value, v.value)
    else
        throw(PGNException("Malformed header pair"))
    end
end

function readheaders(p::PGNReader)::GameHeaders
    result = GameHeaders()
    while peek(p) == '['
        (name, value) = readheader(p)
        setheadervalue!(result, name, value)
    end
    result
end

function skipvariation(p::PGNReader)
    while true
        t = readtoken(p)
        if t.ttype == leftparen
            skipvariation(p)
        elseif t.ttype == rightparen
            break
        elseif t.ttype == endoffile
            throw(PGNException("EOF while reading variation"))
        end
    end
end

"""
    readgame(p::PGNReader; annotations=false)

Attempts to parse a PGN game and use it to create a `Game` or a `SimpleGame`.

If the optional parameter `annotations` is `true`, the return value will be a
`Game` containing all comments, variations and numeric annotation glyphs in the
PGN. Otherwise, it will be a `SimpleGame` with only the game moves.

This function assumes that the `PGNReader` is pointed at the beginning of a
game. If you are not sure this is the case, call `gotonextgame!` on the
`PGNReader` first.

If parsing fails or the notation contains illegal or ambiguous moves, the
function raises a `PGNException`.
"""
function readgame(p::PGNReader; annotations = false)
    headers = readheaders(p)
    if isnothing(headers.fen)
        result = annotations ? Game() : SimpleGame()
    else
        result = annotations ? Game(headers.fen) : SimpleGame(headers.fen)
    end
    result.headers = headers
    precomment = nothing
    nullmoveseen = false
    while true
        t = readtoken(p)
        if terminatesgame(t) || t.ttype == endoffile
            break
        elseif t.ttype == nullmove
            nullmoveseen = true
        elseif t.ttype == leftparen
            if !annotations
                skipvariation(p)
            else
                back!(result)
            end
        elseif t.ttype == rightparen
            @assert annotations
            tobeginningofvariation!(result)
            forward!(result)
        elseif t.ttype == comment
            if annotations
                if isatbeginning(result) || !isatend(result)
                    precomment = t.value
                else
                    addcomment!(result, t.value)
                end
            end
        elseif t.ttype == nag
            if annotations
                addnag!(result, parse(Int, t.value))
            end
        elseif t.ttype == symbol && !nullmoveseen
            # Try to parse the symbol as a move in short algebraic notation,
            # and add it to the game if successful
            m = movefromsan(board(result), t.value)
            if !isnothing(m)
                if annotations
                    addmove!(result, m)
                    if !isnothing(precomment)
                        addprecomment!(result, precomment)
                        precomment = nothing
                    end
                else
                    domove!(result, m)
                end
            end
        end
    end
    tobeginning!(result)
    result
end

function skipgame(p::PGNReader)
    headers = readheaders(p)
    while true
        t = readtoken(p)
        if terminatesgame(t) || t.ttype == endoffile
            break
        elseif t.ttype == leftparen
            skipvariation(p)
        end
    end
end

"""
    gamesinfile(filename::String; annotations=false, skip=0)

Creates a `Channel` of `Game`/`SimpleGame` objects read from the PGN file with
the provided file name.

If the optional parameter `annotations` is `true`, the return value will be a
channel of `Game` objects containing all comments, variations and numeric
annotation glyphs in the PGN. Otherwise, it will consist of `SimpleGame` objects
with only the game moves.

The optional parameter `skip` makes the function skip the first `skip` games of
the file.
"""
function gamesinfile(filename::String; annotations = false, skip = 0)
    function createchannel(ch::Channel)
        open(filename, "r") do io
            pgnr = PGNReader(io)

            if skip > 0
                i = 0
                while !eof(pgnr.io) && i < skip
                    skipgame(pgnr)
                    gotonextgame!(pgnr)
                    i += 1
                end
            end

            while !eof(pgnr.io)
                put!(ch, readgame(pgnr, annotations = annotations))
                gotonextgame!(pgnr)
            end
        end
    end
    Channel(createchannel)
end

"""
    gamefromstring(s::String; annotations=false)

Attempts to create a `Game` or `SimpleGame` object from the provided PGN string.

If the optional parameter `annotations` is `true`, the return value will be a
`Game` containing all comments, variations and numeric annotation glyphs in the
PGN. Otherwise, it will be a `SimpleGame` with only the game moves.

If the string does not parse as valid PGN, or if the notation contains illegal
or ambiguous moves, the function raises a `PGNException`
"""
function gamefromstring(s::String; annotations = false)
    readgame(PGNReader(IOBuffer(s)), annotations = annotations)
end

function formatstring(s::String)::String
    "\"" * replace(replace(s, "\\" => "\\\\"), "\"" => "\\\"") * "\""
end

function formatheader(name::String, value::String)
    "[" * name * " " * formatstring(value) * "]\n"
end

"""
    gametopgn(g)::String

Exports a `Game` or a `SimpleGame` to a PGN string.

# Limitations

- The movetext section is written in a single long line, with no line breaks.
"""
function gametopgn(g)::String
    result = IOBuffer()
    write(result, formatheader("Event", g.headers.event))
    write(result, formatheader("Site", g.headers.site))
    write(result, formatheader("Date", g.headers.date))
    write(result, formatheader("Round", g.headers.round))
    write(result, formatheader("White", g.headers.white))
    write(result, formatheader("Black", g.headers.black))
    write(result, formatheader("Result", g.headers.result))
    if !isnothing(headervalue(g, "FEN"))
        write(result, formatheader("SetUp", "1"))
        write(result, formatheader("FEN", headervalue(g, "FEN")))
    end
    for gh in g.headers.othertags
        write(result, formatheader(gh.name, gh.value))
    end
    write(result, "\n", Chess.formatmoves(g), g.headers.result, "\n")
    String(take!(result))
end

end # module


text:
using FactCheck

import CDLasso

##############################################
#
#  Groups Lasso
#
##############################################

facts("Group Active set") do

  numGroups = 3
  groupToIndex = Array(typeof(1:2), numGroups)
  for i=1:numGroups
    groupToIndex[i] = UnitRange((i-1)*2+1, i*2)
  end

  as = CDLasso.GroupLassoData(zeros(6), groupToIndex)
  @fact as.numActive --> 0

  x = [1., 0., 0., 0., 1., 2.]
  as = CDLasso.GroupLassoData(x, groupToIndex)
  @fact as.numActive --> 2

  @fact CDLasso._group_norm(x, as, 3) --> roughly(sqrt(5))
  @fact CDLasso._group_norm(x, as, 2) --> roughly(0.)

  CDLasso._fill_zero!(x, as, 1)
  @fact x[1:2] --> zeros(2)

  A = eye(6)
  b = [1., 0., 0., 0., 1., 2.]
  x = zeros(6)
  as = CDLasso.GroupLassoData(x, groupToIndex)
  λ = [0.9, 0., 3]

  CDLasso._add_violator!(as, x, A, b, λ)
  @fact as.numActive --> 1
  @fact as.groups[1] --> 1

  A = eye(6)
  b = [1., 0., 0., 0., 1., 2.]
  x = zeros(6)
  as = CDLasso.GroupLassoData(x, groupToIndex)
  λ = [1.9, 0., 1.]

  CDLasso._add_violator!(as, x, A, b, λ)
  @fact as.numActive --> 1
  @fact as.groups[1] --> 3
end

facts("Minimize one group") do
  numGroups = 1
  groupToIndex = Array(typeof(1:2), numGroups)
  for i=1:numGroups
    groupToIndex[i] = UnitRange((i-1)*3+1, i*3)
  end

  b = [1., 1., 1.]
  A = eye(3)
  Asvd = svdfact(A)
  λ = 1.
  x = zeros(3)

  as = CDLasso.GroupLassoData(x, groupToIndex)
  CDLasso._add_violator!(as, x, A, b, fill(λ, 3))
  CDLasso._min_one_group!(x, Asvd, b, as, 1, λ)

  @fact A*x+b+(λ/norm(x))*x --> roughly(zeros(3); atol=1e-5)

  ##

  A = randn(100, 3)
  A = A'A/100
  Asvd = svdfact(A)
  b = [1., 1., 1.]
  λ = 1.
  x = zeros(3)

  as = CDLasso.GroupLassoData(x, groupToIndex)
  CDLasso._add_violator!(as, x, A, b, fill(λ, 3))
  CDLasso._min_one_group!(x, Asvd, b, as, 1, λ)

  @fact A*x+b+(λ/norm(x))*x --> roughly(zeros(3); atol=1e-5)

  ##
  numGroups = 1
  groupToIndex = Array(typeof(1:2), numGroups)
  for i=1:numGroups
    groupToIndex[i] = UnitRange((i-1)*10+1, i*10)
  end

  A = randn(100, 10)
  A = A'A/100
  Asvd = svdfact(A)
  b = randn(10)
  λ = 0.5
  x = zeros(10)

  as = CDLasso.GroupLassoData(x, groupToIndex)
  CDLasso._add_violator!(as, x, A, b, fill(λ, 10))
  CDLasso._min_one_group!(x, Asvd, b, as, 1, λ)

  @fact A*x+b+(λ/norm(x))*x --> roughly(zeros(10); atol=1e-5)

end

facts("Minimize active set") do
  numGroups = 2
  groupToIndex = Array(typeof(1:2), numGroups)
  for i=1:numGroups
    groupToIndex[i] = UnitRange((i-1)*3+1, i*3)
  end

  b = ones(6)
  A = eye(6)
  λ = [1., 0.5]

  x = ones(6)
  as = CDLasso.GroupLassoData(x, groupToIndex)
  CDLasso._add_violator!(as, x, A, b, λ)

  CDLasso._group_lasso!(x, A, b, λ, as)

  for i=1:numGroups
    @fact A[groupToIndex[i], :]*x + b[groupToIndex[i]] + (λ[i]/norm(x[groupToIndex[i]]))*x[groupToIndex[i]]  --> roughly(zeros(length(groupToIndex[i])); atol=1e-5)
  end

  ##

  numGroups = 2
  groupToIndex = Array(typeof(1:2), numGroups)
  for i=1:numGroups
    groupToIndex[i] = UnitRange((i-1)*3+1, i*3)
  end

  b = randn(6)
  A = randn(100, 6)
  A = A'A/100
  λ = [1., 0.5]

  x = ones(6)
  as = CDLasso.GroupLassoData(x, groupToIndex)
  CDLasso._add_violator!(as, x, A, b, λ)

  CDLasso._group_lasso!(x, A, b, λ, as)

  for i=1:numGroups
    @fact A[groupToIndex[i], :]*x + b[groupToIndex[i]] + (λ[i]/norm(x[groupToIndex[i]]))*x[groupToIndex[i]]  --> roughly(zeros(length(groupToIndex[i])); atol=1e-5)
  end

  ##

  numGroups = 2
  groupToIndex = Array(typeof(1:2), numGroups)
  for i=1:numGroups
    groupToIndex[i] = UnitRange((i-1)*3+1, i*3)
  end

  b = randn(6)
  A = randn(100, 6)
  A = A'A/100
  λ = [100., 0.1]

  x = ones(6)
  as = CDLasso.GroupLassoData(x, groupToIndex)
  CDLasso._add_violator!(as, x, A, b, λ)

  CDLasso._group_lasso!(x, A, b, λ, as)

  i = 1
  @fact norm(A[groupToIndex[i], :]*x + b[groupToIndex[i]]) <= λ[i]  --> true
  i = 2
  @fact A[groupToIndex[i], :]*x + b[groupToIndex[i]] + (λ[i]/norm(x[groupToIndex[i]]))*x[groupToIndex[i]]  --> roughly(zeros(length(groupToIndex[i])); atol=1e-4)
end

facts("Group Lasso") do
  numGroups = 200
  groupToIndex = Array(typeof(1:2), numGroups)
  for i=1:numGroups
    groupToIndex[i] = UnitRange((i-1)*5+1, i*5)
  end

  n = 400
  p = 1000
  X = randn(n, p)
  Y = X[:,1:50] * ones(50) + 0.1 * randn(n)
  XX = X' * X / n
  Xyn = -X' * Y / n

  λ = 0.4 .* ones(numGroups)
  x = zeros(p)
  CDLasso.group_lasso!(x, XX, Xyn, λ, groupToIndex)

  for i=1:200
    if norm(x[groupToIndex[i]]) < 1e-6
      @fact norm(XX[groupToIndex[i], :]*x + Xyn[groupToIndex[i]]) <= λ[i] + 1e-4  --> true
    else
      @fact XX[groupToIndex[i], :]*x + Xyn[groupToIndex[i]] + (λ[i]/norm(x[groupToIndex[i]]))*x[groupToIndex[i]] --> roughly(zeros(length(groupToIndex[i])); atol=1e-4)
    end
  end

  λ = 0.01 .* ones(numGroups)
  x = zeros(p)
  CDLasso.group_lasso!(x, XX, Xyn, λ, groupToIndex)

  for i=1:200
    if norm(x[groupToIndex[i]]) < 1e-6
      @fact norm(XX[groupToIndex[i], :]*x + Xyn[groupToIndex[i]]) <= λ[i] + 1e-4  --> true
    else
      @fact XX[groupToIndex[i], :]*x + Xyn[groupToIndex[i]] + (λ[i]/norm(x[groupToIndex[i]]))*x[groupToIndex[i]] --> roughly(zeros(length(groupToIndex[i])); atol=1e-4)
    end
  end

  λ = 40 .* ones(numGroups)
  x = zeros(p)
  CDLasso.group_lasso!(x, XX, Xyn, λ, groupToIndex)

  for i=1:200
    if norm(x[groupToIndex[i]]) < 1e-6
      @fact norm(XX[groupToIndex[i], :]*x + Xyn[groupToIndex[i]]) <= λ[i] + 1e-4  --> true
    else
      @fact XX[groupToIndex[i], :]*x + Xyn[groupToIndex[i]] + (λ[i]/norm(x[groupToIndex[i]]))*x[groupToIndex[i]] --> roughly(zeros(length(groupToIndex[i])); atol=1e-4)
    end
  end

end

##############################################
#
#  Lasso
#
##############################################

facts("Orthogonal design") do
  b = [1.0, 2.0, 3.0]
  A = eye(3)

  λ = 1.5
  x = zeros(3)
  CDLasso.lasso!(x, A, b, λ)
  @fact x --> roughly([0, -0.5, -1.5])

  λ = [0.9, 2.1, 3.1]
  x = zeros(3)
  CDLasso.lasso!(x, A, b, λ)
  @fact x --> roughly([-0.1, 0, 0])
end

facts("Lasso shrink to zero") do
  n = 100
  p = 10

  X = randn(n, p)
  Y = X * ones(p) + 0.1 * randn(n)
  A = X' * X / n
  b = -X' * Y / n

  beta = ones(p)
  λ = maximum(abs(b)) + 0.1
  CDLasso.lasso!(beta, A, b, λ)
  @fact beta --> roughly(zeros(p))
end

##############################################
#
#  Other
#
##############################################

facts("SoftThreshold") do

  @fact CDLasso.shrink(5., 3.) --> 2.
  @fact CDLasso.shrink(-2., 3.) --> 0.
  @fact CDLasso.shrink(-1.,1.) --> 0.
  @fact CDLasso.shrink(-1.,0.9) --> roughly(-0.1, 1e-5)

end

FactCheck.exitstatus()

max_stars_repo_name: mlakolar/CDLasso

max_issues_repo_name: mlakolar/CDLasso

max_forks_repo_name: mlakolar/CDLasso


text:
export AbstractGroup

"""
    AbstractGroup

Abstract type for abstract groups. Currently the only subtype is `FiniteGroup`.
"""
abstract type AbstractGroup end


text:
@testset "GXBeamAssembly" begin

    # number of elements
    nelem = 4

    # point locations
    x = range(0, 60, length=nelem+1)
    y = zero(x)
    z = zero(x)
    points = [[x[i],y[i],z[i]] for i = 1:length(x)]

    # element connectivity
    start = 1:nelem
    stop = 2:nelem+1

    # element orientation
    e1 = [1, 0, 0]
    e2 = [0, 1, 0]
    e3 = [0, 0, 1]
    frame = hcat(e1, e2, e3)

    # element stiffness properties
    stiffness = [
         2.389e9  1.524e6  6.734e6 -3.382e7 -2.627e7 -4.736e8
         1.524e6  4.334e8 -3.741e6 -2.935e5  1.527e7  3.835e5
         6.734e6 -3.741e6  2.743e7 -4.592e5 -6.869e5 -4.742e6
        -3.382e7 -2.935e5 -4.592e5  2.167e7 -6.279e5  1.430e6
        -2.627e7  1.527e7 -6.869e5 -6.279e5  1.970e7  1.209e7
        -4.736e8  3.835e5 -4.742e6  1.430e6  1.209e7  4.406e8
        ]
    compliance = inv(stiffness)
    compliance_entries = [compliance[i,j] for i = 1:6 for j = i:6]

    # element inertial properties
    mass = [
         258.053      0.0        0.0      0.0      7.07839  -71.6871
           0.0      258.053      0.0     -7.07839  0.0        0.0
           0.0        0.0      258.053   71.6871   0.0        0.0
           0.0       -7.07839   71.6871  48.59     0.0        0.0
           7.07839    0.0        0.0      0.0      2.172      0.0
         -71.6871     0.0        0.0      0.0      0.0       46.418
         ]

    # construct assembly
    assembly = Assembly(points, start, stop;
        frames = fill(frame, nelem),
        stiffness = fill(stiffness, nelem),
        mass = fill(mass, nelem))

    # define boundary conditions and applied loads
    prescribed_conditions = Dict(
            # fixed left side
            1 => PrescribedConditions(ux=0, uy=0, uz=0, theta_x=0, theta_y=0, theta_z=0),
            # force on right side
            nelem+1 => PrescribedConditions(Fz=1e5)
            )

    # define submodel
    submodel = Submodel(GXBeamAssembly(assembly, prescribed_conditions))

    # jacobian tests
    submodel_jacobian_tests(submodel)

    # i/o tests
    submodel_io_tests(submodel)
end


text:
struct MNIST <: Name end

# download options
task(::Type{MNIST}) = MultiClass
format(::Type{MNIST}) = GrayImages
source(::Type{MNIST}) = "http://yann.lecun.com/exdb/mnist/"
make_datadep(::Type{MNIST}) = nothing
loadraw(::Type{MNIST}, type) = loadraw_mldatasets(MLDatasets.MNIST, type)

nclasses(::Type{MNIST}) = 10
classes(::Type{MNIST}) = collect(0:9)
classes_orig(::Type{MNIST}) = collect(0:9)
nattributes(::Type{MNIST}) = (28, 28)
ninstances(::Type{MNIST}) = (60000, 0, 10000)

# dataset description
name(::Type{MNIST}) = "The MNIST database of handwritten digits"
function author(::Type{MNIST})
    return ["Yann LeCun", "Corinna Cortes", "Christopher J.C. Burges"]
end

function citation(::Type{MNIST})
    return """
    @article{lecun1998gradient,
        title = {Gradient-based learning applied to document recognition},
        author = {LeCun, Yann and Bottou, L{\\'e}on and Bengio, Yoshua and Haffner, Patrick},
        journal = {Proceedings of the IEEE},
        volume = {86},
        number = {11},
        pages = {2278--2324},
        year = {1998},
        publisher = {Ieee}
    }
    """
end

max_stars_repo_name: VaclavMacha/DatasetProvider

max_issues_repo_name: VaclavMacha/DatasetProvider

max_forks_repo_name: VaclavMacha/DatasetProvider


text:
using StaticArrays

using JuLIP: AbstractCalculator,
             Atoms

import Base: Dict,
             ==

import JuLIP: cutoff,
              energy,
              forces,
              site_energies,
              virial

export NBodyIP,
       fast,
       load_ip,
       save_ip,
       energy,
       forces,
       site_energies,
       virial

# ----------- some auxiliaries we will use throughout -------------

"""
`push_str!(ex::Vector{Expr}, s::String)` : parse the string and
append it to the expressions.
"""
push_str!(ex::Vector{Expr}, s::String) = push!(ex, Meta.parse(s))

"""
`append_str!(ex::Vector{Expr}, s::Vector{String})` : parse the strings and
append them to the expressions.
"""
append_str!(ex::Vector{Expr}, s::Vector{String}) = append!(ex, Meta.parse.(s))

"""
`bo2edges(N)` : bodyorder-to-edges
"""
bo2edges(N::Integer) = (N * (N-1)) ÷ 2
bo2edges(::Val{N}) where {N} = (N * (N-1)) ÷ 2

"""
`edges2bo(M)`: "edges-to-bodyorder", an internal function that translates
the number of edges in a simplex into the body-order
"""
edges2bo(M::Integer) = (M <= 0) ? 1 : round(Int, 0.5 + sqrt(0.25 + 2 * M))

"""
`bo2angles(N)` : bodyorder-to-angles in one corner of a simplex
"""
bo2angles(N::Integer) = ((N-2) * (N-1)) ÷ 2
bo2angles(::Val{N}) where {N} = ((N-2) * (N-1)) ÷ 2

"""
`angles2bo(M)`: "angles-to-bodyorder", an internal function that translates
the number of angles in one corner of a simplex into the body-order
"""
angles2bo(A::Integer) = (M <= 0) ? 2 : round(Int, 3/2 + sqrt((9/4) + (2-2*A)))

# ----------- some generic functions that we like to have available globally
#             to overload as needed

"""
some measure of degree - could be but need not be polymomial degree; this
is basis-specific
"""
function degree end

"""
some measure of degree for bond-angles
"""
function wtdegree end

"""
`basisname(b) -> String`: returns a short string identifying the kind of
basis function.
"""
function basisname end

"""
`combinebasis`:  if `basis::Vector{AbstractCalculator}` with identical
`hash(::BASIS,...)`, then `combinebasis(basis, coeffs)` should return a new
calculator that combines all basis function into a single IP.
"""
function combinebasis end

# prototypes for the invariants
function invariants end
function invariants_d end
function invariants_ed end

# prototypes for obtaining the descriptor
function descriptor end

function evaluate_I end
function evaluate_I_d end
function evaluate_I_ed end

bodyorder(V::NBodyFunction{N}) where {N} = N

==(V1::NBodyIP, V2::NBodyIP) = V1.components == V2.components
cutoff(V::NBodyIP) = maximum( cutoff.(V.components) )
energy(V::NBodyIP, at::Atoms) = sum( energy(Vn, at)  for Vn in V.components )
forces(V::NBodyIP, at::Atoms) = sum( forces(Vn, at)  for Vn in V.components )
virial(V::NBodyIP, at::Atoms) = sum( virial(Vn, at)  for Vn in V.components )

"""
turn a potentially slow representation of an IP into a fast one,
by switching to a different representation.
"""
fast(IP::NBodyIP) = NBodyIP( fast.(IP.components) )

# fall-back
fast(V::Any) = V

function unique_components(basis)
   bods = hash.(Ref(BASIS()), basis)
   un = Int[]
   I = Vector{Int}[]
   for i = 1:length(bods)
      found = false
      for j = 1:length(un)
         if bods[un[j]] == bods[i]
            push!(I[j], i)
            found = true
            break
         end
      end
      if !found
         push!(un, i)
         push!(I, Int[i])
      end
   end
   return I
end

# construct an NBodyIP from a basis
function NBodyIP(basis, coeffs)
   components = AbstractCalculator[]
   newbasis = AbstractCalculator[]
   newcoeffs = Float64[]
   # if the basis contains any NBodyIPs
   # they should be treated differently
   for (i, b) in enumerate(basis)
      if b isa NBodyIP
         append!(components, b.components)
      else
         push!(newbasis, b)
         push!(newcoeffs, coeffs[i])
      end
   end
   basis, coeffs = newbasis, newcoeffs

   I = unique_components(basis)
   for J in I
      push!(components, combinebasis([basis[j] for j in J], coeffs[J]))
   end
   return NBodyIP(components)
end

# functionality for pure NBodyFunctions for computing
#  * site_energies
#  * energy
#  * forces
#  * virial
include("eval_nbody.jl")

# IO of NBodyIPs
include("io.jl")

# space transforms and cutoffs
include("descsupp.jl")

# ======= experimental ============

import JuLIP.Potentials: evaluate, evaluate_d

evaluate(V::NBodyFunction{2, <: BondLengthDesc}, r::AbstractFloat) =
      evaluate(V, SVector(r))
evaluate_d(V::NBodyFunction{2, <: BondLengthDesc}, r::AbstractFloat) =
      evaluate_d(V, SVector(r))[1]

evaluate(V::NBodyFunction{2, <: ClusterBLDesc}, r::AbstractFloat) =
      evaluate(V, SVector(r))
evaluate_d(V::NBodyFunction{2, <: ClusterBLDesc}, r::AbstractFloat) =
      evaluate_d(V, SVector(r))[1]

evaluate(V::NBodyFunction{2, <: BondAngleDesc}, r::T) where {T <: AbstractFloat} =
      evaluate(V, (SVector(r), SVector{0,T}()))
evaluate_d(V::NBodyFunction{2, <: BondAngleDesc}, r::T) where {T <: AbstractFloat} =
      evaluate_d(V, (SVector(r), SVector{0,T}()))[1]

function evaluate(V::NBodyFunction{N}, r::SVector{M}) where {N, M}
   # this assumes that D is a BondLengthDesc
   D = descriptor(V)::Union{BondLengthDesc, ClusterBLDesc}
   return evaluate_I(V, invariants(D, r)) * fcut(D, r)
end

function evaluate_d(V::NBodyFunction{N}, r::SVector{M}) where {N, M}
   D = descriptor(V)::BondLengthDesc
   fc, fc_d = fcut_d(D, r)
   Vn, Vn_d = evaluate_I_ed(V, invariants_ed(D, r))
   return fc * Vn_d + fc_d * Vn
end

function evaluate(V::NBodyFunction{N}, rθ::Tuple{SVector{M1}, SVector{M2}}
                  ) where {N, M1, M2}
   # this assumes that D is a BondAngleDesc
   D = descriptor(V)::BondAngleDesc
   return evaluate_I(V, invariants(D, rθ)) * fcut(D, rθ)
end

function evaluate_d(V::NBodyFunction{N}, rθ::Tuple{SVector{M1}, SVector{M2}}
                  ) where {N, M1, M2}
   D = descriptor(V)::BondAngleDesc
   fc, fc_d = fcut_d(D, rθ)
   Vn, Vn_d = evaluate_I_ed(V, invariants_ed(D, rθ))
   return fc * Vn_d + fc_d * Vn
end

"""
Take a basis and split it into individual basis groups.
"""
function split_basis(basis; splitfun = b -> hash(Val{:basis}(), b))
   # get some basis hashs of the individual basis functions
   tps = splitfun.(basis)
   Iord = Vector{Int}[]
   Bord = Any[]
   for tp in unique(tps)
      # find which elements of basis have type `tp`
      I = findall( [tp == t  for t in tps] )
      push!(Iord, I)
      push!(Bord, [b for b in basis[I]])
   end
   return Bord, Iord
end

max_stars_repo_name: cortner/ManyBodyIPs

max_issues_repo_name: cortner/ManyBodyIPs

max_forks_repo_name: cortner/ManyBodyIPs


text:
@enum Operator crossover mutation none

struct Entry
    id::Int
    index::Int # index in `entry_pool`
    operator::Operator
    tree::GPNode
    id_entry::Tuple{Int, Int, Int} # parents' id
    index_entry::Tuple{Int, Int, Int} #parent's index
    train_semantics::Vector{Float64}
    test_semantics::Vector{Float64}
    train_fitness::Float64
    test_fitness::Float64
end

function Entry(tree, data_train, data_test, cfg; index=get_usable_index())
    train_semantics = Vector{Float64}(undef, length(data_train.target))
    for i in 1:length(train_semantics)
        train_semantics[i] =  evaluate(tree, view(data_train.input, :, i))
    end

    test_semantics = Vector{Float64}(undef, length(data_test.target))
    for i in 1:length(test_semantics)
        test_semantics[i] =  evaluate(tree, view(data_test.input, :, i))
    end

    train_fitness = cfg.fitness(train_semantics, data_train)
    test_fitness = cfg.fitness(test_semantics, data_test)

    id = get_new_id()
    # index = get_usable_index()

    return Entry(id, index, none, tree, (-1, -1, -1), (-1, -1, -1), train_semantics, test_semantics, train_fitness, test_fitness)
end

function Entry(t1, t2, t3, op, data_train, data_test, cfg; index=get_usable_index())
    if op == crossover
        return crossover_entry(t1, t2, t3, data_train, data_test, cfg, index)
    elseif op == mutation
        return mutation_entry(t1, t2, t3, data_train, data_test, cfg, index)
    end
end

sem(idx) = entry_pool[idx].train_semantics
test_sem(idx) = entry_pool[idx].test_semantics

function crossover_entry(t1, t2, r, data_train, data_test, cfg, index)
    train_semantics = Vector{Float64}(undef, length(data_train.target))
    @inbounds @simd for i in 1:length(train_semantics)
        train_semantics[i] = sigmoid(r.train_semantics[i]) * t1.train_semantics[i] + (1 - sigmoid(r.train_semantics[i])) * t2.train_semantics[i]
    end

    test_semantics = Vector{Float64}(undef, length(data_test.target))
    @inbounds @simd for i in 1:length(test_semantics)
        test_semantics[i] = sigmoid(r.test_semantics[i]) * t1.test_semantics[i] + (1 - sigmoid(r.test_semantics[i])) * t2.test_semantics[i]
    end

    train_fitness = cfg.fitness(train_semantics, data_train)
    test_fitness = cfg.fitness(test_semantics, data_test)

    id = get_new_id()
    #index = get_usable_index()

    return Entry(id, index, crossover, dummy_node(), (t1.id, t2.id, r.id), (t1.index, t2.index, r.index), train_semantics, test_semantics, train_fitness, test_fitness)
end

function mutation_entry(t, r1, r2, data_train, data_test, cfg, index)
    train_semantics = Vector{Float64}(undef, length(data_train.target))
    @inbounds @simd for i in 1:length(train_semantics)
        train_semantics[i] = t.train_semantics[i] + cfg.mutation_step * (sigmoid(r1.train_semantics[i]) - sigmoid(r2.train_semantics[i]))
    end

    test_semantics = Vector{Float64}(undef, length(data_test.target))
    @inbounds @simd for i in 1:length(test_semantics)
        test_semantics[i] = t.test_semantics[i] + cfg.mutation_step * (sigmoid(r1.test_semantics[i]) - sigmoid(r2.test_semantics[i]))
    end

    train_fitness = cfg.fitness(train_semantics, data_train)
    test_fitness = cfg.fitness(test_semantics, data_test)

    id = get_new_id()
    #index = get_usable_index()

    return Entry(id, index, mutation, dummy_node(), (t.id, r1.id, r2.id), (t.index, r1.index, r2.index), train_semantics, test_semantics, train_fitness, test_fitness)
end

const id_counter = [0]

function init_id()
    id_counter[1] = 0
end

function get_new_id()
    id_counter[1] += 1
    return id_counter[1]
end

# population like the one in Fig.2 in [2]
const entry_pool = Entry[]

function init_entry_pool()
    while !isempty(entry_pool)
        pop!(entry_pool)
    end
end

function add_entry!(entry)
    if entry.index <= length(entry_pool)
        entry_pool[entry.index] = entry
    elseif entry.index == length(entry_pool)+1
        push!(entry_pool, entry)
    else
        d = entry.index - length(entry_pool)
        append!(entry_pool, Vector{Entry}(undef, d))
        entry_pool[entry.index] = entry
    end
end
get_entry_pool_size() = length(entry_pool)

# reference management

const ref_counter = PriorityQueue{Int, Int}()

function set_ref(entry)
    ref_counter[entry.index] = 0 # new individuals are not referenced
end

function init_ref_counter(cfg)
    while !isempty(ref_counter)
        dequeue!(ref_counter)
    end

    for i in 1:cfg.population_size
        ref_counter[i] = 0
    end
end

function increment_ref(entry)
    ref_counter[entry.index] += 1
end

function avoid_overwrite(idx)
    ref_counter[idx] += 1
end

# return available (unreferenced) index
function get_usable_index()
    if isempty(ref_counter)
        return 1
    elseif peek(ref_counter)[2] == 0
        return peek(ref_counter)[1]
    else
        return length(ref_counter) + 1
    end
end

get_counter_size() = length(ref_counter)

function update_entry_pool!(pop)
    # new individual's ref_counter is 1 to avoid overwritten whether referenced or not, so it has to be decremented
    for i in pop
        decrement_ref!(i)
    end
    _update_entry_pool!(pop)
end

# decrement unreferenced entry's parents' ref_counter recursively
function _update_entry_pool!(pop)
    s = Int[]
    for i in pop
        if ref_counter[i] == 0
            for j in entry_pool[i].index_entry
                if j != -1
                    # decrement_ref!(j)
                    push!(s, j)
                end
            end
        end
    end

    for i in s
        decrement_ref!(i)
    end
    unique!(s)

    isempty(s) && return
    _update_entry_pool!(s)
end

function decrement_ref!(entry)
    ref_counter[entry.index] -= 1
end

function decrement_ref!(index::Int)
    ref_counter[index] -= 1
end

is_rand_tree(index) = entry_pool[index].index_entry[1] == -1

# reconstruction
function reconstruct_tree(idx, ms)
    entry = entry_pool[idx]
    refs = entry.index_entry
    if entry.operator == none
        return entry.tree
    end

    if entry.operator == crossover
        t1 = reconstruct_tree(refs[1], ms)
        t2 = reconstruct_tree(refs[2], ms)
        c =  crossover_tree(t1, t2, entry_pool[refs[3]])
        return c
    else
        t1 = reconstruct_tree(refs[1], ms)
        return mutation_tree(t1, entry_pool[refs[2]], entry_pool[refs[3]], ms)
    end
end

function crossover_tree(t1, t2, r)
    add_node = GPNode(PrimitiveFunction(+, 2))
    sub_node = GPNode(PrimitiveFunction(-, 2))
    mul_node1 = GPNode(PrimitiveFunction(*, 2))
    mul_node2 = GPNode(PrimitiveFunction(*, 2))

    sig_node = GPNode(PrimitiveFunction(sigmoid, 1))
    sig_node.children[1] = r.tree

    if t1 isa GPNode
        mul_node1.children[1] = t1
    else
        mul_node1.children[1] = t1.tree
    end
    mul_node1.children[2] = sig_node

    sub_node.children[1] = GPNode(Const(1.0))
    sub_node.children[2] =  sig_node

    if t2 isa GPNode
        mul_node2.children[1] = t2
    else
        mul_node2.children[1] = t2.tree
    end
    mul_node2.children[2] = sub_node

    add_node.children[1] = mul_node1
    add_node.children[2] = mul_node2

    return add_node
end

function mutation_tree(t, r1, r2, ms)
    add_node = GPNode(PrimitiveFunction(+, 2))
    sub_node = GPNode(PrimitiveFunction(-, 2))
    mul_node = GPNode(PrimitiveFunction(*, 2))

    sig_node1 = GPNode(PrimitiveFunction(sigmoid, 1))
    sig_node2 = GPNode(PrimitiveFunction(sigmoid, 1))
    sig_node1.children[1] = r1.tree
    sig_node2.children[1] = r2.tree

    sub_node.children[1] = sig_node1
    sub_node.children[2] = sig_node2

    mul_node.children[1] = GPNode(Const(ms))
    mul_node.children[2] = sub_node

    if t isa GPNode
        add_node.children[1] = t
    else
        add_node.children[1] =  t.tree
    end
    add_node.children[2] = mul_node

    return add_node
end

max_stars_repo_name: kogad/JuliaGSGP

max_issues_repo_name: kogad/JuliaGSGP

max_forks_repo_name: kogad/JuliaGSGP


text:
# Alphabet
# ========
#
# Alphabet of biological symbols.
#
# Subtypes of `Alphabet` represent a domain of biological characters. For
# example, `DNAAlphabet{2}` has a domain of unambiguous nucleotides (i.e. A, C,
# G, and T). These types are used for parameterizing biological sequences and so
# on. A pair of encoder and decoder is associated with an alphabet, which maps
# values between binary and Julia-level representation.
#
# This file is a part of BioJulia.
# License is MIT: https://github.com/BioJulia/Bio.jl/blob/master/LICENSE.md

"""
Alphabet of biological characters.
"""
abstract Alphabet

"""
DNA nucleotide alphabet.
"""
immutable DNAAlphabet{n} <: Alphabet end

"""
RNA nucleotide alphabet.
"""
immutable RNAAlphabet{n} <: Alphabet end

"""
Amino acid alphabet.
"""
immutable AminoAcidAlphabet <: Alphabet end

"""
General character alphabet.
"""
immutable CharAlphabet <: Alphabet end

"""
Void alphabet (internal use only).
"""
immutable VoidAlphabet <: Alphabet end

typealias NucleotideAlphabet Union{DNAAlphabet,RNAAlphabet}

"""
The number of bits to represent the alphabet.
"""
function bitsof end

for n in (2, 4)
    @eval begin
        bitsof(::Type{DNAAlphabet{$n}}) = $n
        bitsof(::Type{RNAAlphabet{$n}}) = $n
    end
end
bitsof(::Type{AminoAcidAlphabet}) = 8
bitsof(::Type{CharAlphabet}) = 32
bitsof(::Type{VoidAlphabet}) = 0

Base.eltype(::Type{DNAAlphabet}) = DNANucleotide
Base.eltype(::Type{RNAAlphabet}) = RNANucleotide
Base.eltype{n}(::Type{DNAAlphabet{n}}) = DNANucleotide
Base.eltype{n}(::Type{RNAAlphabet{n}}) = RNANucleotide
Base.eltype(::Type{AminoAcidAlphabet}) = AminoAcid
Base.eltype(::Type{CharAlphabet}) = Char
Base.eltype(::Type{VoidAlphabet}) = Void

alphabet(::Type{DNAAlphabet{2}}) = ACGT
alphabet(::Type{RNAAlphabet{2}}) = ACGU
alphabet(::Type{DNAAlphabet{4}}) = alphabet(DNANucleotide)
alphabet(::Type{RNAAlphabet{4}}) = alphabet(RNANucleotide)
alphabet(::Type{AminoAcidAlphabet}) = alphabet(AminoAcid)
# TODO: this alphabet includes invalid Unicode scalar values
alphabet(::Type{CharAlphabet}) = typemin(Char):typemax(Char)
alphabet(::Type{VoidAlphabet}) = nothing

# Encoders & Decoders
# -------------------

"""
Encode biological characters to binary representation.
"""
function encode end

immutable EncodeError{A<:Alphabet,T} <: Exception
    val::T
end

EncodeError{A,T}(::Type{A}, val::T) = EncodeError{A,T}(val)

function Base.showerror{A}(io::IO, err::EncodeError{A})
    print(io, "cannot encode ", err.val, " in ", A)
end

"""
Decode biological characters from binary representation.
"""
function decode end

immutable DecodeError{A<:Alphabet,T} <: Exception
    val::T
end

DecodeError{A,T}(::Type{A}, val::T) = DecodeError{A,T}(val)

function Base.showerror{A}(io::IO, err::DecodeError{A})
    print(io, "cannot decode ", err.val, " in ", A)
end

# DNA and RNA alphabets
# ---------------------

for A in (DNAAlphabet, RNAAlphabet)
    T = eltype(A)
    @eval begin
        # 2-bit encoding
        @inline function encode(::Type{$(A){2}}, nt::$(T))
            if count_ones(nt) != 1 || !isvalid(nt)
                throw(EncodeError($(A){2}, nt))
            end
            return convert(UInt8, trailing_zeros(nt))
        end
        @inline function decode(::Type{$(A){2}}, x::UInt8)
            if x > 0x03
                throw(DecodeError($(A){2}, x))
            end
            return reinterpret($(T), 0x01 << x)
        end
        @inline decode(::Type{$(A){2}}, x::Unsigned) = decode($(A){2}, UInt8(x))

        # 4-bit encoding
        @inline function encode(::Type{$(A){4}}, nt::$(T))
            if !isvalid(nt)
                throw(EncodeError($(A){4}, nt))
            end
            return reinterpret(UInt8, nt)
        end
        @inline function decode(::Type{$(A){4}}, x::UInt8)
            if !isvalid($(T), x)
                throw(DecodeError($(A){4}, x))
            end
            return reinterpret($(T), x)
        end
        @inline decode(::Type{$(A){4}}, x::Unsigned) = decode($(A){4}, UInt8(x))
    end
end

# AminoAcidAlphabet
# -----------------

@inline function encode(::Type{AminoAcidAlphabet}, aa::AminoAcid)
    if aa > AA_Gap
        throw(EncodeError(AminoAcidAlphabet, aa))
    end
    return reinterpret(UInt8, aa)
end

@inline function decode(::Type{AminoAcidAlphabet}, x::UInt8)
    if x > 0x1b
        throw(DecodeError(AminoAcidAlphabet, x))
    end
    return reinterpret(AminoAcid, x)
end

@inline function decode(::Type{AminoAcidAlphabet}, x::Unsigned)
    return decode(AminoAcidAlphabet, UInt8(x))
end

# CharAlphabet
# ------------

@inline function encode(::Type{CharAlphabet}, char::Char)
    if char > '\U10ffff'
        throw(EncodeError(CharAlphabet, char))
    end
    return reinterpret(UInt32, char)
end

@inline function decode(::Type{CharAlphabet}, x::UInt32)
    if x > 0x10ffff
        throw(DecodeError(CharAlphabet, x))
    end
    return reinterpret(Char, x)
end

@inline function decode(::Type{CharAlphabet}, x::Unsigned)
    return decode(CharAlphabet, UInt32(x))
end


text:
tests = [
    "test-types",
    "test-pval-adjustment",
    "test-pval-pi0-adjustment",
    "test-grenander",
    "test-utils",
    "test-model",
    "test-combinations",
    "test-pi0-estimators",
    "test-higher-criticism"
]

for t in tests
    test_file = "$(t).jl"
    include(test_file)
end


text:
"""
    ADI.ADIAlgorithm

This abstract type is used for defining ADI algorithms. Algorithms are stateful objects that define the options for a given algorithm (e.g. the number of components used in PCA decomposition). The most direct usage of an algorithm is to use it to fit HCI data; that is, given a sample of pixels, apply the algorithm using the given options and return an object containing all the necessary information to reconstruct the input data.

See the extended help (`??ADIAlgorithm`) for interface details.

# Extended help
## Interface
To extend `ADIAlgorithm` you may implement the following

    ADI.fit(::Alg, data::AbstractMatrix; kwargs...)

Fit the data (flattened into a matrix). To support RDI, ensure the `ref` keyword argument is usable (`ref` is also a flattened matrix). This is the only method you *need* to implement for a new `ADIAlgorithm`, along with a suitable [`ADIDesign`](@ref).

ADI.jl automatically coaxes the `cube` input into a matrix for use with `fit`, appropriately handling the various geometries. When available, this input is a view, so if the algorithm requires dense arrays, make sure to call `collect` when appropriate. If a given algorithm doesn't support the default operations, all that needs to be done is override the default behavior (for an example, see the [`GreeDS`](@ref) implementation).

---

    reconstruct(::Alg, cube; kwargs...)

Fit the data using the algorithm and return a cube with the estimate of the PSF. By default uses the reconstruction from the [`ADIDesign`](@ref) fit to the data.

---

    subtract(::Alg, cube; kwargs...)

Fit the data using the algorithm and return a cube that has had the PSF estimate subtracted. By default, calls [`reconstruct`](@ref) and subtracts it from `cube`.

---

    process(::ADIAlgorithm, cube; kwargs...)
    (::ADIAlgorithm)(cube; kwargs...)

Fully process the data (estimate, subtract, collapse). By default, derotates and collapses output from [`subtract`](@ref). You only need to define [`process`](@ref), since the functor version is supplied automatically.
"""
abstract type ADIAlgorithm end

"""
    ADI.fit(::ADIAlgorithm, cube; [ref], kwargs...)

Given the description of an algorithm and the appropriate options, take the pixels from `cube` and fit them, returning an ([`ADIDesign`](@ref)) containing the necessary information from the fit (e.g. the principal components from PCA decomposition).

If the algorithm supports reference differential imaging (RDI), the reference cube can be passed by the keyword argument `ref`.
"""
function fit(alg::ADIAlgorithm, cube::AbstractArray{T,3}; kwargs...) where T
    data = flatten(cube)
    if :ref in keys(kwargs)
        kwargs[:ref] isa AbstractArray{<:Any,3} || error("reference data geometry does not match target data")
        ref_data = flatten(kwargs[:ref])
        return fit(alg, data; kwargs..., ref=ref_data)
    end
    return fit(alg, data; kwargs...)
end
function fit(alg::ADIAlgorithm, cube::AnnulusView; kwargs...)
    data = cube(true) # as view
    if :ref in keys(kwargs)
        kwargs[:ref] isa AnnulusView || error("reference data geometry does not match target data")
        ref_data = kwargs[:ref](true)
        return fit(alg, data; kwargs..., ref=ref_data)
    end
    return fit(alg, data; kwargs...)
end

function fit(alg::ADIAlgorithm, cube::MultiAnnulusView; kwargs...)
    if :ref in keys(kwargs)
        kwargs[:ref] isa MultiAnnulusView || error("reference data geometry does not match target data")
        anns = eachannulus(cube, true)
        ref_anns = eachannulus(kwargs[:ref], true)
        return StructArray(fit(alg, ann; kwargs..., ref=ref_ann) for (ann, ref_ann) in zip(anns, ref_anns))
    else
        return StructArray(fit(alg, ann; kwargs...) for ann in eachannulus(cube, true))
    end
end

function fit(algs::AbstractVector{<:ADIAlgorithm}, cube::MultiAnnulusView; kwargs...)
    if :ref in keys(kwargs)
        kwargs[:ref] isa MultiAnnulusView || error("reference data geometry does not match target data")
        anns = eachannulus(cube, true)
        ref_anns = eachannulus(kwargs[:ref], true)
        itr = zip(algs, anns, ref_anns)
        return StructArray(fit(alg, ann; kwargs..., ref=ref_ann) for (alg, ann, ref_ann) in itr)
    else
        itr = zip(algs, eachannulus(cube, true))
        return StructArray(fit(alg, ann; kwargs...) for (alg, ann) in itr)
    end
end

"""
    reconstruct(alg, cube; [ref], kwargs...)

Reconstruct the PSF approximation for the given algorithm, using `ref` as the reference cube if given and supported by the algorithm.

# Examples

```julia
julia> cube, angles = # load data

julia> S = reconstruct(PCA(10), cube);

julia> size(S) == size(cube)
true

julia> flat_res = collapse(cube .- S, angles); # form resid, derotate, and combine
```
"""
function reconstruct(alg, cube; kwargs...)
    design = fit(alg, cube; kwargs...)
    S = reconstruct(design)
    return expand_geometry(cube, S)
end

expand_geometry(::AbstractArray{T,3}, arr) where {T} = expand(arr)
expand_geometry(cube::AnnulusView, arr) = inverse(cube, arr)
expand_geometry(cube::MultiAnnulusView, arrs) = inverse(cube, arrs)

"""
    subtract(alg, cube; [ref], kwargs...)

Reconstruct the PSF approximation for the given algorithm and subtract it from `cube`, using `ref` as the reference cube if given and supported by the algorithm.

# Examples

```julia
julia> cube, angles = # load data

julia> R = subtract(PCA(10), cube);

julia> size(R) == size(cube)
true

julia> flat_res = collapse(R, angles); # derotate, and combine
```
"""
function subtract(alg, cube; kwargs...)
    S = reconstruct(alg, cube; kwargs...)
    return cube .- S
end

"""
    process(alg, cube, angles; [ref], kwargs...)

Fully process an ADI data cube using [`subtract`](@ref) and collapsing the residuals. Keyword arguments will be passed to [`ADI.fit`](@ref).
"""
function process(alg, cube, angles; method=:deweight, kwargs...)
    R = subtract(alg, cube; angles=angles, kwargs...)
    return collapse!(R, angles, method=method)
end

"""
    (::ADIAlgorithm)(cube, angles; [ref], kwargs...)

Fully process an ADI data cube using [`subtract`](@ref) and collapsing the residuals. This is a convenient alias for calling [`process`](@ref) Keyword arguments will be passed to [`ADI.fit`](@ref).
"""
(alg::ADIAlgorithm)(args...; kwargs...) = process(alg, args...; kwargs...)


text:
module BasicUnits

export kg, m, K, mol, s
export W, Pa, J
export cm, km, g, bar, GPa, au

# Use a scale where base SI units are 1
"Kilogram"
const kg = 1
"Metre"
const m = 1
"Kelvin"
const K = 1
"Mole"
const mol = 1
"Second"
const s = 1
# Derived units
"Watt"
const W = kg*m^2/s^3
"Pascal"
const Pa = kg/m/s^2
"Joule"
const J = kg*m^2/s^2
# Multiples of base units
"Centimetre"
const cm = 1e-2m
"Kilometre"
const km = 1000m
"Gram"
const g = 1e-3kg
"Bar"
const bar = 10^5 * Pa
"Gigapascal"
const GPa = 10^9 * Pa
"Astronomical unit"
const au = 1.496e11 * m

end # module


text:
int main() {
   int return;
   return 1;
}

max_stars_repo_name: carlostome/Javalette-Compiler

max_issues_repo_name: carlostome/Javalette-Compiler

max_forks_repo_name: carlostome/Javalette-Compiler


text:
function mean_quantity(f, sim)
    tree = maketree(sim.dir => sim.files)

    quantity = FileTrees.load(tree, lazy=true) do file
        println("Loading $(path(file)) on thread $(Threads.threadid())")
        f(file.value)
    end
    exec(mapvalues(mean, quantity))
end

function phase_space_mean(sim, dir::Symbol; species="electron", rcond=(i,x)->true)
    p1 = sim[1]["p$dir/"*species]
    r1 = getdomain(p1)[dir_to_idx(dir)]
    r_mean = eltype(r1)[]
    p_mean = eltype(p1)[]
    rp_mean = @progress for (idx, file) in enumerate(sim)
        # println("Loading $(file.name)")
        p = file["p$dir/"*species]
        r = getdomain(p)[dir_to_idx(dir)]

        # it't not important to optimize the mean computation
        # It takes 0.07s serial for 82 milion elements and 0.017s parallel
        idx = map(i->rcond(idx, r[i]), eachindex(r))
        @debug "Got $(length(idx)) items that satisfy the condition"
        push!(p_mean, mean(view(p, idx)))
        push!(r_mean, mean(view(r, idx)))
    end

    return r_mean, p_mean
end


text:
using IterTools

cd(dirname(@__FILE__))
lines = open("eight.txt", "r") do f
    readlines(f)
end

parse_op(c::String) =
    let (op, val) = split(c)
        (string(op), parse(Int64, val))
    end

p = map(parse_op, lines)

function check_loop(pc::Int64, acc::Int64, executed::Array{Bool,1})
    if executed[pc]
        throw(ErrorException("Infinite loop at $pc: acc = $acc"))
    end
    executed[pc] = true
end

function accumulate(x::Int64, pc::Int64, acc::Int64, executed::Array{Bool,1})
    check_loop(pc, acc, executed)

    pc+1, acc+x
end

function noop(_::Int64, pc::Int64, acc::Int64, executed::Array{Bool,1})
    check_loop(pc, acc, executed)

    pc+1, acc
end

function jump(x::Int64, pc::Int64, acc::Int64, executed::Array{Bool,1})
    check_loop(pc, acc, executed)

    pc+x, acc
end

cpu = Dict(
    "acc" => accumulate,
    "nop" => noop,
    "jmp" => jump
    )

const Program = Array{Tuple{String,Int64},1}

function make_execute(p::Program)
    executed = fill(false, length(p))
    function execute((pc,acc)::Tuple{Int64,Int64})
        (i,v) = p[pc]
        pc, acc = cpu[i](v, pc, acc, executed)
    end
end

exec = make_execute(p)

try
    for _ in iterated(exec, (1,0)) end
catch e
    println(e)
end


text:
same_length_array(array, value::NativeFont) = Iterators.repeated(value, length(array))

# TODO upgrade GeometryTypes to do these kind of things.
# Problem: I don't really want to introduce a depedency on intervals
function Base.in(point::StaticVector{N}, rectangle::HyperRectangle{N}) where N
    mini, maxi = minimum(rectangle), maximum(rectangle)
    for i = 1:N
        point[i] in (mini[i] .. maxi[i]) || return false
    end
    return true
end


text:
# ------------------------------------------------------------------
# Licensed under the ISC License. See LICENSE in the project root.
# ------------------------------------------------------------------

"""
    AbstractData{T,N}

Spatial data in a `N`-dimensional space with coordinates of type `T`.
"""
abstract type AbstractData{T,N} <: AbstractSpatialObject{T,N} end

"""
    variables(sdata)

Return the variable names in spatial data `sdata` and their types.
"""
variables(sdata::AbstractData) =
  OrderedDict([var => eltype(array) for (var,array) in sdata.data])

#----------------
# DATAFRAME API
#----------------

"""
    sdata[inds,vars]

Return the value of `var` for the `ind`-th point in `sdata`.
"""
Base.getindex(sdata::AbstractData, ind::Int, var::Symbol) =
  sdata.data[var][ind]

Base.getindex(sdata::AbstractData, inds::AbstractVector{Int}, var::Symbol) =
  [getindex(sdata, ind, var) for ind in inds]

Base.getindex(sdata::AbstractData, ind::Int, vars::AbstractVector{Symbol}) =
  [getindex(sdata, ind, var) for var in vars]

Base.getindex(sdata::AbstractData, inds::AbstractVector{Int}, vars::AbstractVector{Symbol}) =
  [getindex(sdata, ind, var) for ind in inds, var in vars]

"""
    getindex(sdata, var)

Return the values of `var` for all points in `sdata` with the shape
of the underlying domain.
"""
Base.getindex(sdata::AbstractData, var::Symbol) =
  [getindex(sdata, ind, var) for ind in 1:npoints(sdata)]

Base.getindex(sdata::AbstractData, vars::AbstractVector{Symbol}) =
  [getindex(sdata, var) for var in vars]

"""
    setindex!(sdata, val, ind, var)

Set the value `val` of variable `var` for point `ind` in spatial
data `sdata`.
"""
Base.setindex!(sdata::AbstractData, val, ind::Int, var::Symbol) =
  sdata.data[var][ind] = val

"""
    setindex!(sdata, vals, var)

Set the values `vals` of variable `var` for all points in spatial
data `sdata`.
"""
Base.setindex!(sdata::AbstractData, vals::AbstractArray, var::Symbol) =
  sdata.data[var] = vals

#-----------
# VIEW API
#-----------

"""
    view(sdata, inds)
    view(sdata, vars)
    view(sdata, inds, vars)

Return a view of `sdata` with all points in `inds` and
all variables in `vars`.
"""
Base.view(sdata::AbstractData, inds::AbstractVector{Int}) =
  DataView(sdata, inds, collect(keys(variables(sdata))))

Base.view(sdata::AbstractData, vars::AbstractVector{Symbol}) =
  DataView(sdata, 1:npoints(sdata), vars)

Base.view(sdata::AbstractData, inds::AbstractVector{Int},
                               vars::AbstractVector{Symbol}) =
  DataView(sdata, inds, vars)

#---------------
# ITERATOR API
#---------------

"""
    iterate(sdata, state=1)

Iterate over samples in `sdata`.
"""
Base.iterate(sdata::AbstractData, state=1) =
  state > npoints(sdata) ? nothing : (sdata[state], state + 1)

"""
    length(sdata)

Return the number of samples in `sdata`.
"""
Base.length(sdata::AbstractData) = npoints(sdata)

#----------------
# INDEXABLE API
#----------------

"""
    getindex(sdata, ind)

Return `ind`-th sample in `sdata`.
"""
function Base.getindex(sdata::AbstractData, ind::Int)
  vars = [var for (var,V) in variables(sdata)]
  vals = [getindex(sdata, ind, var) for var in vars]
  NamedTuple{tuple(vars...)}(vals)
end

Base.getindex(sdata::AbstractData, inds::AbstractVector{Int}) =
  [getindex(sdata, ind) for ind in inds]

"""
    firstindex(sdata)

Return the first index of `sdata`.
"""
Base.firstindex(sdata::AbstractData) = 1

"""
    lastindex(sdata)

Return the last index of `sdata`.
"""
Base.lastindex(sdata::AbstractData) = npoints(sdata)

#-------------
# TABLES API
#-------------
Tables.istable(::Type{<:AbstractData}) = true

Tables.columnaccess(::Type{<:AbstractData}) = true

function Tables.columns(sdata::AbstractData)
  vars = keys(variables(sdata))
  vals = [getindex(sdata, 1:npoints(sdata), var) for var in vars]
  NamedTuple{tuple(vars...)}(vals)
end

#-----------------
# MISSING VALUES
#-----------------
"""
    isvalid(sdata, ind, var)

Return `true` if the `ind`-th point in `sdata` has a valid value for `var`.
"""
function Base.isvalid(sdata::AbstractData, ind::Int, var::Symbol)
  val = getindex(sdata, ind, var)
  !(ismissing(val) || (val isa Number && isnan(val)))
end

"""
    valid(sdata, var)

Return all points in `sdata` with a valid value for `var`. The output is
a tuple with the matrix of coordinates as the first item and the vector
of values as the second item.
"""
function valid(sdata::AbstractData{T,N}, var::Symbol) where {N,T}
  V = variables(sdata)[var]
  npts = npoints(sdata)

  # pre-allocate memory for result
  X = Matrix{T}(undef, N, npts)
  z = Vector{V}(undef, npts)

  nvalid = 0
  for ind in 1:npoints(sdata)
    if isvalid(sdata, ind, var)
      nvalid += 1
      coordinates!(view(X,:,nvalid), sdata, ind)
      z[nvalid] = getindex(sdata, ind, var)
    end
  end

  X[:,1:nvalid], z[1:nvalid]
end

# ------------
# IO methods
# ------------
function Base.show(io::IO, sdata::AbstractData{T,N}) where {N,T}
  npts = npoints(sdata)
  print(io, "$npts SpatialData{$T,$N}")
end

function Base.show(io::IO, ::MIME"text/plain", sdata::AbstractData{T,N}) where {N,T}
  println(io, sdata)
  println(io, "  variables")
  varlines = ["    └─$var ($V)" for (var,V) in variables(sdata)]
  print(io, join(sort(varlines), "\n"))
end

#------------------
# IMPLEMENTATIONS
#------------------
include("data/curve_data.jl")
include("data/geodataframe.jl")
include("data/point_set_data.jl")
include("data/regular_grid_data.jl")
include("data/structured_grid_data.jl")


text:
"""
```
    Canny <: AbstractEdgeDetectionAlgorithm
    Canny(; spatial_scale = 1, high = Percentile(80), low = Percentile(20), thinning_algorithm = NonmaximaSuppression(threshold = low))

    detect_edges([T,] img, f::Canny)
    detect_edges!([out,] img, f::Canny)
    detect_subpixel_edges([T₁, T₂] img, f::Canny)
    detect_subpixel_edges!(out₁, out₂, img, f::Canny)
```

Returns a binary image depicting the edges of the input image.

# Details

TODO

# Options

Various options for the parameters of the `detect_edges` function and `Canny` type are described in more detail below.

# Choices for img

The `detect_edges` function can handle a variety of input types.
By default the type of the returned image matches the type of the
input image.

For colored images, the input is converted to grayscale.

# Choices for `spatial_scale` in `Canny`.

The `spatial_scale` determines the radius (σ) of the Gaussian filter. It must
be a positive real number.

# Choices for `high` and `low` in `Canny`.

The hysteresis thresholds `high` and `low` (`high` > `low`) can be specified as
positive numbers, or as `Percentiles`. If left unspecified, a default value of
`high = Percentile(80)` and `low = Percentile(20)` is assumed.

# Choices for `thinning_algorithm` in `Canny`.
You can specify an [`AbstractEdgeThinningAlgorithm`](@ref). By default, the
[`NonmaximaSuppression`](@ref) algorithm is used which suppresses non-maxima up
to pixel-level accuracy. For subpixel precision specify the
[`SubpixelNonmaximaSuppression`](@ref) algorithm.

# Example

```julia

using TestImages, FileIO, ImageView

img =  testimage("mandril_gray")
img_edges = detect_edges(img, Canny(spatial_scale = 1.4))

imshow(img)
imshow(img_edges)
```

# References
J. Canny, "A Computational Approach to Edge Detection," in IEEE Transactions on Pattern Analysis and Machine Intelligence, vol. PAMI-8, no. 6, pp. 679-698, Nov. 1986, doi: 10.1109/TPAMI.1986.4767851.

"""
@with_kw struct Canny{T₁ <: Union{Real,AbstractGray},
                      T₂ <: Union{Real,AbstractGray, Percentile},
                      T₃ <: Union{Real,AbstractGray, Percentile},
                      T₄ <: AbstractEdgeThinningAlgorithm} <: AbstractEdgeDetectionAlgorithm
    spatial_scale::T₁ = 1
    high::T₂ = Percentile(80)
    low::T₃ = Percentile(20)
    thinning_algorithm::T₄ = NonmaximaSuppression(threshold = low)
end

# Handles pixel-level precision.
function (f::Canny)(out::GenericGrayImage, img::GenericGrayImage)
    @unpack spatial_scale, high, low = f
    @unpack thinning_algorithm = f
    σ = spatial_scale

    # Smooth the image with a Gaussian filter of width σ, which specifies the
    # scale level of the edge detector.
    kernel =  KernelFactors.gaussian((σ,σ))
    imgf = imfilter(img, kernel, NA())

    # Calculate the gradient vector at each position of the filtered image.
    # The derivatives are taken with respect to the first and second dimension
    # (rows first, then columns).
    g₁, g₂ = imgradients(imgf, KernelFactors.scharr)
    # Gradient magnitude
    mag = hypot.(g₁, g₂)

    # In StatsBase quantiles are undefined in the presence of NaNs
    # hence we need to keep only valid magnitudes before we can determine
    # the percentiles.
    valid_indices = map(x-> !isnan(x), mag)
    valid_mag = view(mag, valid_indices)

    low_threshold =  typeof(low) <: Percentile ? StatsBase.percentile(vec(valid_mag), low.p) : low
    high_threshold =  typeof(high) <: Percentile ? StatsBase.percentile(vec(valid_mag), high.p) : high

    thinning_algorithm = @set thinning_algorithm.threshold = low_threshold

    # Isolate local maxima of gradient magnitude by “non-maximum suppression”
    # along the local gradient direction.
    nms = thin_edges(mag, g₁, g₂, thinning_algorithm)

    # Collect sets of connected edge pixels from the local maxima by applying
    # “hysteresis thresholding”. Edge pixels whose gradient magnitude is below
    # the low threshold are removed; edge pixels whose gradient magnitudes are
    # above (or equal to) the high threshold are retained, and edge pixels whose
    # gradient magnituded are between the low and high threshold are only
    # retained if they are connected to edge pixels whose gradient magnitudes
    # are above or equal to the high threshold.
    edges = zeros(Bool, axes(img))
    @inbounds for i in CartesianIndices(nms)
        if nms[i] >= high_threshold && edges[i] == 0
            trace_and_threshold!(edges, nms, i, low_threshold)
        end
    end
    out .= edges
end

# Handles subpixel precision.
function (f::Canny)(out₁::GenericGrayImage, out₂::AbstractArray{<:StaticVector}, img::GenericGrayImage)
    @unpack spatial_scale, high, low = f
    @unpack thinning_algorithm = f
    σ = spatial_scale

    # Smooth the image with a Gaussian filter of width σ, which specifies the
    # scale level of the edge detector.
    kernel = KernelFactors.gaussian((σ,σ))
    imgf = imfilter(img, kernel, NA())

    # Calculate the gradient vector at each position of the filtered image.
    # The derivatives are taken with respect to the first and second dimension
    # (rows first, then columns).
    g₁, g₂ = imgradients(imgf, KernelFactors.scharr)
    # Gradient magnitude
    mag = hypot.(g₁, g₂)

    # In StatsBase quantiles are undefined in the presence of NaNs
    # hence we need to keep only valid magnitudes before we can determine
    # the percentiles.
    valid_indices = map(x-> !isnan(x), mag)
    valid_mag = view(mag, valid_indices)

    low_threshold =  typeof(low) <: Percentile ? StatsBase.percentile(vec(valid_mag), low.p) : low
    high_threshold =  typeof(high) <: Percentile ? StatsBase.percentile(vec(valid_mag), high.p) : high

    thinning_algorithm = @set thinning_algorithm.threshold = low_threshold

    # Isolate local maxima of gradient magnitude by “non-maximum suppression”
    # along the local gradient direction.
    nms, subpixel_offsets = thin_subpixel_edges(mag, g₁, g₂, thinning_algorithm)

    # Collect sets of connected edge pixels from the local maxima by applying
    # “hysteresis thresholding”. Edge pixels whose gradient magnitude is below
    # the low threshold are removed; edge pixels whose gradient magnitudes are
    # above (or equal to) the high threshold are retained, and edge pixels whose
    # gradient magnituded are between the low and high threshold are only
    # retained if they are connected to edge pixels whose gradient magnitudes
    # are above or equal to the high threshold.
    edges = zeros(Bool, axes(img))
    @inbounds for i in CartesianIndices(nms)
        if nms[i] >= high_threshold && edges[i] == 0
            trace_and_threshold!(edges, nms, i, low_threshold)
        end
    end
    out₁ .= edges
    out₂ .= subpixel_offsets

    return out₁, out₂
end

function (f::Canny)(out::AbstractArray{<:Color3}, img::AbstractArray{<:Color3})
    T = eltype(img)
    out_temp = zeros(Gray{eltype(T)}, axes(out))
    # NB This will be refactored in future version once a proper color-based
    # edge detection algorithm is implemented. At the moment we just call the
    # default "grayscale" edge detection algorithm.
    out .= convert.(T, f(out_temp, Gray.(img)))
end

function (f::Canny)(out₁::AbstractArray{<:Color3}, out₂::AbstractArray{<:StaticVector}, img::AbstractArray{<:Color3})
    T = eltype(img)
    out_temp = zeros(Gray{eltype(T)}, axes(out₁))
    # NB This will be refactored in future version once a proper color-based
    # edge detection algorithm is implemented. At the moment we just call the
    # default "grayscale" edge detection algorithm.
    out_temp, offsets = f(out_temp, out₂, Gray.(img))

    out₁ .= convert.(T, out_temp)
    return out₁, offsets
end

(f::Canny)(out::GenericGrayImage, img::AbstractArray{<:Color3}) =
    f(out, of_eltype(Gray, img))

(f::Canny)(out₁::GenericGrayImage, out₂::AbstractArray{<:StaticVector}, img::AbstractArray{<:Color3}) =
    f(out₁, out₂, of_eltype(Gray, img))

"""
```
    trace_and_threshold!(out::AbstractArray,  mag::AbstractArray, i₀::CartesianIndex, low::Number)
```

 Marks all pixels of an edge that are 8-connected to i₀ and exhibit a gradient magnitude above `low`.

"""
function trace_and_threshold!(out::AbstractArray, mag::AbstractArray, i₀::CartesianIndex, low::Number)
    stack = Stack{CartesianIndex{2}}()
    push!(stack, i₀)
    # Trace all the pixels that are reachable from the current edge pixel.
    rows, cols = axes(out)
    while !isempty(stack)
        i = pop!(stack)

        # Mark the pixel as visited.
        out[i] = oneunit(eltype(out))
        rᵢ, cᵢ = i.I

        # Ensure we don't exceed the image bounds
        r₀ = (rᵢ > firstindex(rows)) ? rᵢ - 1 : firstindex(rows)
        r₁ = (rᵢ < lastindex(rows)) ? rᵢ + 1 : lastindex(rows)
        c₀ = (cᵢ > firstindex(cols)) ? cᵢ - 1 : firstindex(cols)
        c₁ = (cᵢ < lastindex(cols)) ? cᵢ + 1 : lastindex(cols)

        # Search the neighbourhood for any connected pixels which exceed
        # the minimum edge magnitude threshold `low`.
        @inbounds for r = r₀:r₁
            for c = c₀:c₁
                j = CartesianIndex(r,c)
                if out[j] == zero(eltype(out)) && mag[j] >= low
                    push!(stack, j)
                end
            end
        end
    end
end


text:
@testset "Database types" begin

@testset "ComponentDatabase" begin
    @test Data.martinez() isa CubicEoSDatabase.AbstractTabularDatabase
    @test_throws ArgumentError ComponentDatabase("foo_db/galaxy.csv")
end

@testset "MixtureDatabase" begin
    @test Data.brusilovsky_mix() isa MixtureDatabase
    @test_throws ArgumentError MixtureDatabase("foo_db/galaxy.csv")
end

end # @testset


text:
mutable struct MMAWorkspace{T, TV1, TV2, TM, TO, TSO, TSubOptions, TTrace<:OptimizationTrace{T}, TModel<:MMAModel{T}, TPD<:PrimalData{T}}
    model::TModel
    optimizer::TO
    suboptimizer::TSO
    suboptions::TSubOptions
    x0::TV1
    x::TV1
	x1::TV1
	x2::TV1
	λ::TV2
	l::TV2
	u::TV2
	∇f_x::TV1
	g::TV2
	ng_approx::TV2
	∇g::TM
	f_x::T
	f_calls::Int
	g_calls::Int
	f_x_previous::T
	primal_data::TPD
	tr::TTrace
	tracing::Bool
	converged::Bool
	x_converged::Bool
	f_converged::Bool
	gr_converged::Bool
	f_increased::Bool
	x_residual::T
	f_residual::T
	gr_residual::T
	asymptotes_updater::AsymptotesUpdater{T, TV1, TModel}
	variable_bounds_updater::VariableBoundsUpdater{T, TPD, TModel}
	cvx_grad_updater::ConvexApproxGradUpdater{T, TPD, TModel}
	lift_updater::LiftUpdater{T, TV2, TPD}
	lift_resetter::LiftResetter{T, TV2}
	x_updater::XUpdater{TPD}
	dual_obj::DualObjVal{T, TPD}
	dual_obj_grad::DualObjGrad{TPD}
    dual_caps::Tuple{T, T}
	outer_iter::Int
    iter::Int
end
const MMA87Workspace{T, TV1, TV2, TM, TSO, TModel, TPD} = MMAWorkspace{T, TV1, TV2, TM, MMA87, TSO, TModel, TPD}
const MMA02Workspace{T, TV1, TV2, TM, TSO, TModel, TPD} = MMAWorkspace{T, TV1, TV2, TM, MMA02, TSO, TModel, TPD}

Base.show(io::IO, w::MMA87Workspace) = print(io, "Workspace for the method of moving asymptotes of 1987.")
Base.show(io::IO, w::MMA02Workspace) = print(io, "Workspace for the method of moving asymptotes of 2002.")

function MMAWorkspace(model::MMAModel{T,TV}, x0::TV, optimizer=MMA02(), suboptimizer=Optim.ConjugateGradient(); suboptions=Optim.Options(x_tol=sqrt(eps(T)), f_tol=sqrt(eps(T)), g_tol=sqrt(eps(T))), s_init=T(0.5), s_incr=T(1.2), s_decr=T(0.7), dual_caps=default_dual_caps(optimizer, T)) where {T, TV}

    n_i = length(constraints(model))
    n_j = dim(model)
    x, x1, x2 = copy(x0), copy(x0), copy(x0)
    TM = MatrixOf(TV)

    # Buffers for bounds and move limits
    α, β, σ = zerosof(TV, n_j), zerosof(TV, n_j), zerosof(TV, n_j)

    # Buffers for p0, pji, q0, qji
    p0, q0 = zerosof(TV, n_j), zerosof(TV, n_j)
    p, q = zerosof(TM, n_j, n_i), zerosof(TM, n_j, n_i)
    if optimizer isa MMA87
        ρ = Vector(zerosof(TV, n_i))
    else
        ρ = Vector(onesof(TV, n_i))
    end
    r = Vector(zerosof(TV, n_i))

    # Initial data and bounds for Optim to solve dual problem
    λ = Vector(onesof(TV, n_i))
    l = Vector(zerosof(TV, n_i))
    u = Vector(infsof(TV, n_i))
    #u .= 1e50
    # Objective gradient buffer

    ∇f_x = nansof(TV, length(x))
    g = Vector(zerosof(TV, n_i))
    ng_approx = Vector(zerosof(TV, n_i))
    ∇g = zerosof(TM, n_j, n_i)

    f_x::T = eval_objective(model, x, ∇f_x)
    f_calls, g_calls = 1, 1
    f_x_previous = T(NaN)

    # Evaluate the constraints and their gradients
    update_constraints!(g, ∇g, model, x)

    # Build a primal data struct storing all primal problem's info
    primal_data = PrimalData(σ, α, β, p0, q0, p, q, ρ, r, Ref(zero(T)), x, x1, Ref(f_x), g, ∇f_x, ∇g)

    tr = OptimizationTrace{T, MMA87}()
    tracing = (model.store_trace || model.extended_trace || model.show_trace)

    converged = false
    # Assess multiple types of convergence
    x_converged, f_converged, gr_converged = false, false, false
    f_increased = false
    x_residual = T(Inf)
    f_residual = T(Inf)
    gr_residual = T(Inf)

    asymptotes_updater = AsymptotesUpdater(model, σ, x, x1, x2, s_init, s_incr, s_decr)
    variable_bounds_updater = VariableBoundsUpdater(primal_data, model, T(μ))
    cvx_grad_updater = ConvexApproxGradUpdater(primal_data, model)
    lift_updater = LiftUpdater(primal_data, ρ, g, ng_approx, n_j)
    lift_resetter = LiftResetter(ρ, T(ρmin))

    x_updater = XUpdater(primal_data)
    dual_obj = DualObjVal(primal_data, x_updater)
    dual_obj_grad = DualObjGrad(primal_data, x_updater)

    # Iteraton counter
    outer_iter = 0
    iter = 0

    TO = typeof(optimizer)
    TSO = typeof(suboptimizer)
    TModel = typeof(model)
    TPD = typeof(primal_data)
    TSubOptions = typeof(suboptions)
    return MMAWorkspace(
        model, optimizer, suboptimizer, suboptions, x0, x, x1, x2, λ, l, u, ∇f_x, g,
        ng_approx, ∇g, f_x, f_calls, g_calls, f_x_previous, primal_data, tr, tracing,
        converged, x_converged, f_converged, gr_converged, f_increased, x_residual,
        f_residual, gr_residual, asymptotes_updater, variable_bounds_updater,
        cvx_grad_updater, lift_updater, lift_resetter, x_updater, dual_obj,
        dual_obj_grad, dual_caps, outer_iter, iter)
end

function update_constraints!(g, ∇g, model, x)
    for i in 1:length(g)
        g[i] = eval_constraint(model, i, x, @view(∇g[:,i]))
    end
    return g
end

function assess_convergence(workspace::MMAWorkspace)
    @unpack model, x, x1, f_x, f_x_previous, ∇f_x = workspace
    return assess_convergence(x, x1, f_x, f_x_previous, ∇f_x, xtol(model), ftol(model), grtol(model))
end


text:
using DifferentialEquations
using Plots
using PyPlot

# Set the required model parameters for the SIRS model with two levels of transmission - Within and between households
N = 10; # Household size - Change to 10 for final analysis
betaHH = 6; # Within household transmission parameter
betaG = 1; # Population wide transmission
gamma = 1; # Rate of recovery from infection
tau = 1; # Rate of loss of protection
params = [betaHH,gamma,tau,betaG,N]; # Put all the parameters together
time = (0.0, 30.0) # Simulation time - note it defined as a float
dim = dim = 0.5*(N+1)*(N+2); # Number of possible configurations - works for three epidemiological classes
y0 = vec(zeros(1,dim)); # Initial condition vector
y0[end-1] = 0.00000001;
y0[end] = 0.99999999;

function hhTransitions(N,dim)
    # Function to generate transition matrices for household model
    # Input: N is the household size

    # Initialize things
    Qinf = zeros(dim,dim);
    Qrec = zeros(dim,dim);
    Qext = zeros(dim,dim);
    Qwane = zeros(dim,dim);
    dataI = Array{Int64}(zeros(dim,3))
    m = 0;
    I = Array{Int64}(zeros(N+1,N+1))

    # To help remember where to store the variables
    for ss = 0:N
        for ii = 0:(N-ss)
            m = m + 1;
            I[ss+1,ii+1] = m
        end
    end

    # Describe the epidemiological transitions

    # Counter for susceptibles
    for ss = 0:N
        # Counter for infecteds
        for ii = 0:(N-ss)
            # If susceptibles and infecteds are more than 1, then infection within the household can occur
            if (ss > 0 && ii > 0)
                Qinf[I[ss+1,ii+1],I[ss,ii+2]] = ii*ss/(N-1);
            end

            # If infecteds are more than 1, recovery can occur
            if ii > 0
                # Rate of recovery
                Qrec[I[ss+1,ii+1],I[ss+1,ii]] = ii;
            end

            # For external infection - just keep track of susceptibles
            if ss > 0
                # Rate of within household infection
                Qext[I[ss+1,ii+1],I[ss,ii+2]] = ss;
            end

            # Loss of protection hence becoming susceptible again. Possible if N-ss-ii = rr > 0
            if (N-ss-ii) > 0
                # Rate of loss of protection
                Qwane[I[ss+1,ii+1],I[ss+2,ii+1]] = N-ss-ii;
            end

            # Store the relevant indices to help identify the household configurations
            dataI[I[ss+1,ii+1],:] = [ss, ii, N-ss-ii];
        end
    end

    Qinf = Qinf - diagm(vec(sum(Qinf,2)),0);
    Qrec = Qrec - diagm(vec(sum(Qrec,2)),0);
    Qext = Qext - diagm(vec(sum(Qext,2)),0);
    Qwane = Qwane - diagm(vec(sum(Qwane,2)),0);

    # Return
    return Qinf, Qrec, Qext, Qwane, dataI
end

Qinf, Qrec, Qext, Qwane, dataI = hhTransitions(N,dim);

function rateSIRS(dy_dt,y0,params,time)

    # Extract the parameters
    betaHH = params[1];
    gamma = params[2];
    tau = params[3];
    betaG = params[4];
    N = params[5];

    # Generate the transition matrices
    Qinf, Qrec, Qext, Qwane, HHconfig = hhTransitions(N,dim);

    # Combine within and external transitions
    Q = betaHH*Qinf + gamma*Qrec + tau*Qwane + (betaG*((HHconfig[:,2]'*y0)/N)*Qext);

    # Generate the rates of change for ODE solver
    dy_dt .= (y0'*Q)';

    # Alternatively this works too
    #=for i=1:length(y0)
        dy_dt[i] = y0'*Q[:,i]
    end=#

end

# Define the ODE problem
prob = ODEProblem(rateSIRS,y0,time,params);

# Solve
sol = solve(prob);

Iconfig = dataI[:,2];
infProp = zeros(length(sol.t),1);

# Prepare the plots
for i = 1:length(sol.t)
    infProp[i,1] = sol[:,i]'*Iconfig/N;
end

# Total infectious in the population
plot(sol.t,infProp,color="blue",xlabel="Time",ylabel="Proportion infectious",label=["Mean infection"],ylims=[0, 1])

# Household profile at endemic prevalence
# Prepare the plots
yprop = zeros(N+1,length(sol.t))
for j = 1:length(sol.t)
    for i = 1:N+1
        index = find(Iconfig.==i-1);
        yprop[i,j] = sum(sol[index,j])
    end
end
step(-0.5:1:10.5,[yprop[:,length(sol.t)];0])

# Household profile at peak prevalence
# Prepare the plots
x = find(infProp.==maximum(infProp))
step(-0.5:1:10.5,[yprop[:,x];0]);

max_stars_repo_name: epimodels/epicookbook

max_issues_repo_name: epimodels/epicookbook

max_forks_repo_name: epimodels/epicookbook