text:
export convnp_1d

"""
    convnp_1d(;
        receptive_field::Float32,
        num_encoder_layers::Integer,
        num_decoder_layers::Integer,
        num_encoder_channels::Integer,
        num_decoder_channels::Integer,
        num_latent_channels::Integer,
        num_global_channels::Integer,
        points_per_unit::Float32,
        margin::Float32=receptive_field,
        noise_type::String="het",
        pooling_type::String="sum",
        σ::Float32=1f-2,
        learn_σ::Bool=true
    )

# Keywords
- `receptive_field::Float32`: Width of the receptive field.
- `num_layers::Integer`: Number of layers of the CNN, excluding an initial
    and final pointwise convolutional layer to change the number of channels
    appropriately.
- `num_encoder_layers::Integer`: Number of layers of the CNN of the encoder.
- `num_decoder_layers::Integer`: Number of layers of the CNN of the decoder.
- `num_encoder_channels::Integer`: Number of channels of the CNN of the encoder.
- `num_decoder_channels::Integer`: Number of channels of the CNN of the decoder.
- `num_latent_channels::Integer`: Number of channels of the latent variable.
- `num_global_channels::Integer`: Number of channels of a global latent variable. Set to
    `0` to not use a global latent variable.
- `points_per_unit::Float32`: Points per unit for the discretisation. See
    `UniformDiscretisation1D`.
- `margin::Float32=receptive_field`: Margin for the discretisation. See
    `UniformDiscretisation1D`.
- `noise_type::String="het"`: Type of noise model. Must be "fixed", "amortised", or "het".
- `pooling_type::String="mean"`: Type of pooling. Must be "mean" or "sum".
- `σ::Float32=1f-2`: Initialisation of the fixed observation noise.
- `learn_σ::Bool=true`: Learn the fixed observation noise.

# Returns
- `Model`: Corresponding model.
"""
function convnp_1d(;
    receptive_field::Float32,
    num_encoder_layers::Integer,
    num_decoder_layers::Integer,
    num_encoder_channels::Integer,
    num_decoder_channels::Integer,
    num_latent_channels::Integer,
    num_global_channels::Integer,
    points_per_unit::Float32,
    margin::Float32=receptive_field,
    noise_type::String="het",
    pooling_type::String="mean",
    σ::Float32=1f-2,
    learn_σ::Bool=true
)
    dim_x = 1
    dim_y = 1
    scale = 2 / points_per_unit

    # Build encoder.
    encoder = Chain(
        set_conv(dim_y, scale; density=true),
        build_conv(
            receptive_field,
            num_encoder_layers,
            num_encoder_channels,
            points_per_unit =points_per_unit,
            dimensionality  =1,
            num_in_channels =dim_y + 1,  # Account for density channel.
            num_out_channels=2num_latent_channels + 2num_global_channels,
        ),
        HeterogeneousGaussianLikelihood()
    )
    if num_global_channels > 0
        # There is a global variable. Split it off and pool instead.
        if pooling_type == "mean"
            pooling = MeanPooling(layer_norm(1, 2num_global_channels, 1))
        elseif pooling_type == "sum"
            pooling = SumPooling(1000)  # Divide by `1000` to help initialisation.
        else
            error("Unknown pooling type \"" * pooling_type * "\".")
        end
        encoder = Chain(
            encoder[1:end - 1]...,  # Remove likelihood.
            Splitter(2num_global_channels),
            Parallel(
                HeterogeneousGaussianLikelihood(),
                Chain(
                    batched_mlp(
                        dim_in    =2num_global_channels,
                        dim_hidden=2num_global_channels,
                        dim_out   =2num_global_channels,
                        num_layers=3
                    ),
                    pooling,
                    batched_mlp(
                        dim_in    =2num_global_channels,
                        dim_hidden=2num_global_channels,
                        dim_out   =2num_global_channels,
                        num_layers=3
                    ),
                    HeterogeneousGaussianLikelihood()
                )
            )
        )
    end

    # Build noise model.
    num_noise_channels, noise = build_noise_model(
        # Need to perform a smoothing before applying the noise model.
        num_channels -> set_conv(num_channels, scale),
        dim_y       =dim_y,
        noise_type  =noise_type,
        pooling_type=pooling_type,
        σ           =σ,
        learn_σ     =learn_σ
    )

    return Model(
        FunctionalCoder(
            UniformDiscretisation1D(points_per_unit, margin),
            encoder
        ),
        Chain(
            Materialise(),
            build_conv(
                receptive_field,
                num_decoder_layers,
                num_decoder_channels,
                points_per_unit =points_per_unit,
                dimensionality  =1,
                num_in_channels =num_latent_channels + num_global_channels,
                num_out_channels=num_noise_channels
            ),
            noise
        )
    )
end


text:
"""

```
printxyz(time,x,v,sys,traj_file_buffer,vel_file_buffer)
```

Subroutine that writes the current coordinates *and* velocities to
output files, previously opened in respective buffers.

"""
function printxyz(time,x,v,sys,traj_file_buffer,vel_file_buffer)
  n = length(x)
  println(traj_file_buffer,n)
  println(traj_file_buffer," time = ", time)
  println(vel_file_buffer,n)
  println(vel_file_buffer," time = ", time)
  for i in 1:n
    xt = image(x[i],sys.sides)
    vt = v[i]
    if length(xt) == 2
      print(traj_file_buffer,@sprintf("He %12.5f %12.5f %12.5f\n", xt[1], xt[2], 0.))
      print(vel_file_buffer,@sprintf("He %12.5f %12.5f %12.5f\n", vt[1], vt[2], 0.))
    else
      print(traj_file_buffer,@sprintf("He %12.5f %12.5f %12.5f\n", xt[1], xt[2], xt[3]))
      print(vel_file_buffer,@sprintf("He %12.5f %12.5f %12.5f\n", vt[1], vt[2], vt[3]))
    end
  end
end

"""

```
printxyz(time,x,sys,traj_file_buffer)
```

Subroutine that writes the current coordinates (or velocities) to
output files, previously opened in respective buffers.

"""
function printxyz(time,x,sys,traj_file_buffer)
  n = length(x)
  println(traj_file_buffer,n)
  println(traj_file_buffer," time = ", time)
  for i in 1:n
    xt = image(x[i],sys.sides)
    if length(xt) == 2
      print(traj_file_buffer,@sprintf("He %12.5f %12.5f %12.5f\n", xt[1], xt[2], 0.))
    else
      print(traj_file_buffer,@sprintf("He %12.5f %12.5f %12.5f\n", xt[1], xt[2], xt[3]))
    end
  end
end

"""

```
printxyz(x,sys,file::String)
```

Prints one set of coordinates (or velocities), given in `x` to file `file`. Overwrites
the file.

"""
function printxyz(x::Vector,sys::System,file::String)
  buffer = open(file,"w")
  printxyz(0,x,sys,buffer)
  close(buffer)
  return nothing
end

ext: jl


text:
raw"""
Computes Cholesky decomposition of a Hermitian positive-definite matrix A, i.e.,

$$
A = L L^*
$$

where $L$ is a lower triangular matrix with real and positive diagonal entries, and $L^*$ denotes
the conjugate transpose of $L$.
"""
module CholeskyDecomposition

using LinearAlgebra

export
    cholesky_outer_product_decompose,
    cholesky_banachiewicz_decompose,
    cholesky_crout_decompose

function get_n(A::Matrix)::Int
    (n_r, n_c) = size(A)
    @assert n_r == n_c "Input matrix is not square"
    n_r
end

"""
The Cholesky (outer-product) algorithm

<https://en.wikipedia.org/wiki/Cholesky_decomposition#The_Cholesky_algorithm>
"""
function cholesky_outer_product_decompose(A::Matrix{T})::LowerTriangular{T} where T<:Number
    n = get_n(A)

    # Some edge case values:
    #
    #     A_1 = A
    #     A_{n + 1} = I(n)
    #     I_i ∈ ℜ_{i ^ 2}
    #     I_0 = ∅
    #     L_0 = I(n)
    #     B_0 = A
    #     B_i ∈ ℂ_{(n - i) ^ 2}
    #     B_n = ∅

    A_i = copy(A)
    A = nothing  # prevents accidental assignment
    L_0 = T.(I(n)) |> Matrix |> LowerTriangular
    L = copy(L_0)  # The product of L1...Ln
    for i in 1:n
        # Makes aliases (no need for `A_ii`)
        b_i = @view A_i[(i + 1):end, i]
        b_i_H = @view A_i[i, (i + 1):end]
        B_i = @view A_i[(i + 1):end, (i + 1):end]

        L_i = copy(L_0)
        sqrtAii = sqrt(A_i[i, i])
        L_i[i, i] = sqrtAii
        L_i[(i + 1):end, i] = b_i / sqrtAii

        L *= L_i

        # Mind the order
        B_i .-= b_i * transpose(b_i_H) / A_i[i, i]  # Ugly, can't do `b_i * b_i_H`
        b_i .= 0.0
        b_i_H .= 0.0
        A_i[i, i] = 1.0
    end
    L
end

"""
The Cholesky-Banachiewicz algorithm

<https://en.wikipedia.org/wiki/Cholesky_decomposition#The_Cholesky–Banachiewicz_and_Cholesky–Crout_algorithms>
"""
function cholesky_banachiewicz_decompose(A::Matrix{T})::LowerTriangular{T} where T<:Number
    n = get_n(A)
    # To write the in-place version, just set `L = A` and do symmetric updating.
    L = zeros(T, n, n) |> LowerTriangular

    # Note that as $A$ is positive-definite, **not positive semi-definite**, $L_{ii}$ will never be
    # zero, so there is no need to check.
    #
    # Also, do NOT use `dot` to compute `L[r, c]`!

    for r in 1:n
        for c in 1:(r - 1)
            L[r, c] = (
                A[r, c] - sum((@view L[r, 1:(c - 1)]) .* conj.(@view L[c, 1:(c - 1)]))
            ) / L[c, c]
        end
        l = @view L[r, 1:(r - 1)]
        L[r, r] = sqrt(A[r, r] - l ⋅ l)
    end
    L
end

"""
The Cholesky-Crout algorithm

It is similar to the Cholesky-Banachiewicz algorithm, but processes column by column instead of row
by row.
"""
function cholesky_crout_decompose(A::Matrix{T})::LowerTriangular{T} where T<:Number
    n = get_n(A)
    L = zeros(T, n, n) |> LowerTriangular
    # Deliberately write the index as `:` rather than `1:(c - 1)`
    for c in 1:n
        l = @view L[c, :]
        L[c, c] = sqrt(A[c, c] - l ⋅ l)
        for r in (c + 1):n
            L[r, c] = (A[r, c] - sum((@view L[r, :]) .* conj.(@view L[c, :]))) / L[c, c]
        end
     end
     L
end

end  # module

max_stars_repo_name: gyk/TrivialSolutions

max_issues_repo_name: gyk/TrivialSolutions

max_forks_repo_name: gyk/TrivialSolutions


text:
@testset "Boundaries" begin

    # boundary of a point
    ch = boundary(Simplex(1))
    @test iszero(ch)
    @test dim(ch) == -1

    # different order of faces
    @test Simplex(4,1) == Simplex(1,4)
    ch1 = boundary(Simplex(4,1))
    ch2 = boundary(Simplex(1,4))
    @test ch1 != ch2

    cplx = SimplicialComplex(Simplex(1,2,3),
                            Simplex(2,4),
                            Simplex(3,4),
                            Simplex(5,4),
                            Simplex(1,2,3,5))

    # ch = boundary(Simplex(1,2,3))
    # showchain(cplx, ch)

    # simplex boundary
    ch = Chain(1, Int) + (Simplex(1,2),1) + (Simplex(2,3),1) + (Simplex(1,3),-1)
    @testset "Complex boundary chain" for (a,b) in zip(boundary(Simplex(1,2,3)), ch)
        @test a == b
    end

    # ∂∘∂(c) == 0
    @test iszero( boundary(cplx, boundary(Simplex(1,2,3))) )

    # boundary of a path
    cplx = SimplicialComplex(Simplex(1,2), Simplex(2,3), Simplex(3,4))
    path = Chain(1, Dict(hash(c) => 1 for c in cells(cplx,1)))
    ch = Chain(0, Int) + (Simplex(1), -1) + (Simplex(4), 1)
    @testset "Path boundary" for (a,b) in zip(boundary(cplx, path), ch)
        @test a == b
    end

    # boundary of a loop
    push!(cplx, Simplex(4,1))
    loop = Chain(1, Dict(hash(c) => 1 for c in cells(cplx,1)))
    @test iszero(simplify(boundary(cplx, loop)))

    # coboundary
    ch = Chain(2, Int) + (Simplex(1,3,5),1) + (Simplex(1,2,3),-1)
    @testset "Complex coboundary chain" for (a,b) in zip(coboundary(cplx, Simplex(1,3)), ch)
        @test a == b
    end

    # δ∘δ(c) == 0
    @test iszero( coboundary(cplx, coboundary(cplx, Simplex(1,3))) )

    # cw-cell boundary
    a = Cell()
    b = Cell()
    c = Cell(1, a, b)
    ch = Chain(1, Int) + (a,1) + (b,-1)
    @testset "Complex boundary chain" for (a,b) in zip(boundary(c), ch)
        @test a == b
    end

    # cube boundary
    c = Cube([1,1]) # 2-cube
    ch = Chain(1, Int) + (Cube([0,1]), -1) + (Cube([1,0]), 1)
    @testset "Complex boundary chain" for (a,b) in zip(boundary(c), ch)
        @test a == b
    end

end


text:
@testset "linalg" begin
    @testset "symmetric" begin
        adj = [0 1 0 1;
               1 0 1 0;
               0 1 0 1;
               1 0 1 0]
        deg = [2 0 0 0;
               0 2 0 0;
               0 0 2 0;
               0 0 0 2]
        lap = [2 -1 0 -1;
               -1 2 -1 0;
               0 -1 2 -1;
               -1 0 -1 2]
        norm_lap = [1. -.5 0. -.5;
                   -.5 1. -.5 0.;
                   0. -.5 1. -.5;
                   -.5 0. -.5 1.]

        for T in [Int8, Float64]
            @test degree_matrix(adj, T, dir=:out) == T.(deg)
            @test degree_matrix(adj, T, dir=:out) == degree_matrix(adj, T, dir=:in)
            @test degree_matrix(adj, T, dir=:out) == degree_matrix(adj, T, dir=:both)
            @test laplacian_matrix(adj, T) == T.(lap)
        end
        @test normalized_laplacian(adj, Float64) ≈ norm_lap
        @test eltype(normalized_laplacian(adj, Float32)) == Float32
        @test neighbors(adj) == [[2,4], [1,3], [2,4], [1,3]]
    end

    @testset "asymmetric" begin
        adj = [0 2 0 3;
               0 0 4 0;
               2 0 0 1;
               0 0 0 0]
        deg_out = [2 0 0 0;
                   0 2 0 0;
                   0 0 4 0;
                   0 0 0 4]
        deg_in = [5 0 0 0;
                  0 4 0 0;
                  0 0 3 0;
                  0 0 0 0]
        deg_both = [7 0 0 0;
                    0 6 0 0;
                    0 0 7 0;
                    0 0 0 4]

        for T in [Int8, Float64]
            @test degree_matrix(adj, T, dir=:out) == T.(deg_out)
            @test degree_matrix(adj, T, dir=:in) == T.(deg_in)
            @test degree_matrix(adj, T, dir=:both) == T.(deg_both)
            @test_throws DomainError degree_matrix(adj, dir=:other)
            @test laplacian_matrix(adj, T, dir=:out) == T.(deg_out .- adj)
            @test laplacian_matrix(adj, T, dir=:in) == T.(deg_in .- adj)
            @test laplacian_matrix(adj, T, dir=:both) == T.(deg_both .- adj)
        end
        @test neighbors(adj) == [[2,3,4], [1,3], [1,2,4], [1,3]]
    end
end


text:
"""
--------------------------------------------------------------------------------------\n
Author: Huckleberry Febbo, Graduate Student, University of Michigan
Date Create: 6/29/2017, Last Modified: 4/13/2018 \n
-------------------------------------------------------------------------------------\n
"""
function dynamics!(n::NLOpt, dx::Vector{Expr})
    if length(dx)!=n.ocp.state.num
        error("\n The number of differential equations must equal ocp.state.num. \n")
    end
    n.ocp.DXexpr = dx
    return nothing
end

"""
--------------------------------------------------------------------------------------\n
Author: Huckleberry Febbo, Graduate Student, University of Michigan
Date Create: 7/04/2017, Last Modified: 4/13/2018 \n
-------------------------------------------------------------------------------------\n
"""
function constraints!(n::NLOpt, con::Vector{Expr})
    n.ocp.NLcon = con
    return nothing
end

"""
--------------------------------------------------------------------------------------\n
Author: Huckleberry Febbo, Graduate Student, University of Michigan
Date Create: 6/11/2017, Last Modified: 4/13/2018 \n
-------------------------------------------------------------------------------------\n
"""
#function DiffEq(n::NLOpt,x::Matrix{JuMP.JuMPTypes},u::Matrix{JuMP.JuMPTypes},L::Int,st::Int)::Vector{JuMP.NonlinearExpression}
#function DiffEq(n::NLOpt,x,u,L::Int64,st::Int64)
function DiffEq(n::NLOpt,x::Array{Any,2},u::Array{Any,2},L::Int64,st::Int64)
    expr = n.ocp.DXexpr[st]
    return NLExpr(n,expr,x,u,L)
end

"""
--------------------------------------------------------------------------------------\n
Author: Huckleberry Febbo, Graduate Student, University of Michigan
Date Create: 7/04/2017, Last Modified: 4/13/2018 \n
-------------------------------------------------------------------------------------\n
"""
#function addCon(n::NLOpt,x::Matrix{JuMP.JuMPTypes},u::Matrix{JuMP.JuMPTypes},L::Int,num::Int)
function addCon(n::NLOpt,x::Array{Any,2},u::Array{Any,2},L::Int64,num::Int64)
    expr = n.ocp.NLcon[num]
    return NLCon(n,expr,x,u,L)
end

"""
# returns an array of @NLexpression()
--------------------------------------------------------------------------------------\n
Author: Huckleberry Febbo, Graduate Student, University of Michigan
Date Create: 6/29/2017, Last Modified: 12/06/2019 \n
-------------------------------------------------------------------------------------\n
"""
function NLExpr(n::NLOpt, expr::Expr, args...)::Vector{JuMP.NonlinearExpression}
#function NLExpr(n::NLOpt,expr::Expr,args...)

    if length(args) == 3
        xxx = args[1]
        uuu = args[2]
        L = args[3]
    elseif length(args) == 0
        xxx = n.r.ocp.x
        uuu = n.r.ocp.u
        L = n.ocp.state.pts
    else
        error("The length of the args... ($(length(args))) must be either 3 or 0")
    end

    code = quote
        # rename state variables
        state = Array{Expr}(undef, $n.ocp.state.num)
        for st in 1:$n.ocp.state.num
            state[st] = Expr(:(=), $n.ocp.state.name[st], $xxx[:,st])
            eval(state[st])
        end

        # rename control variables
        control = Array{Expr}(undef, $n.ocp.control.num)
        for ctr in 1:$n.ocp.control.num
            control[ctr] = Expr(:(=), $n.ocp.control.name[ctr], $uuu[:,ctr])
            eval(control[ctr])
        end

        @NLexpression($n.ocp.mdl, [j=1:$L], $expr)
    end

    return NLOptControl.eval(code)

end

"""
# returns an array of @NLconstraint()
--------------------------------------------------------------------------------------\n
Author: Huckleberry Febbo, Graduate Student, University of Michigan
Date Create: 7/04/2017, Last Modified: 12/06/2019 \n
-------------------------------------------------------------------------------------\n
"""
function NLCon(n::NLOpt,expr::Expr,args...)::Vector{JuMP.ConstraintRef{JuMP.Model, JuMP.NonlinearConstraint}}
  if length(args) == 3
    xxx = args[1]
    uuu = args[2]
    L = args[3]
  elseif length(args) == 0
    xxx = n.r.ocp.x
    uuu = n.r.ocp.u
    L = n.ocp.state.pts
  else
    error("\n The length of the args... must be either 3 or 0 \n")
  end

  code = quote
    # rename state variables
    state = Array{Expr}(undef, $n.ocp.state.num)
    for st in 1:$n.ocp.state.num
      state[st] = Expr(:(=), $n.ocp.state.name[st], $xxx[:,st])
      eval(state[st])
    end

    # rename control variables
    control = Array{Expr}(undef, $n.ocp.control.num)
    for ctr in 1:$n.ocp.control.num
      control[ctr] = Expr(:(=), $n.ocp.control.name[ctr], $uuu[:,ctr])
      eval(control[ctr])
    end

    @NLconstraint($n.ocp.mdl, [j=1:$L], $expr)
  end
  return NLOptControl.eval(code)
end


text:
# Formulas used for Chapter 2 exercises
yearly_compounded_interest(i, r, y) = (i * (1 + (r / 100)))^y

find_effective_rate(i, n) = (1 + i / n)^n - 1

"internal rate of return"
function irr(func1, func2, init, stop)
    sol = Float64[]
    λₖ = func1(init)
    λₖ₊₁ = λₖ - func1(λₖ) / func2(λₖ)
    push!(sol, float(λₖ))
    if stop == 0
        return [λₖ]
    end
    push!(sol, float(λₖ₊₁))
    if stop == 1
        return [λₖ, λₖ₊₁]
    end
    for k in range(2; stop=stop)
        λₖ = λₖ₊₁
        λₖ₊₁ = λₖ - func1(λₖ) / func2(λₖ)
        push!(sol, float(λₖ₊₁))
    end
    return sol
end

"net present value"
npv(inv, val, rate, year) = -inv + val / (1 + rate)^year

"annual equivalent rate"
aer(r, n) = (1 + r / n)^n - 1

"continuously compounded interest"
cci(r) = exp(r) - 1

"present value"
function pv(p, r, n)
    interest = 0.0
    for i in range(0, n - 1; step=1)
        interest += (1 + r)^-i
    end
    return p * interest
end

# From Lecture Notes

"present value"
function present_value(s, r)
    return mapreduce((k, x) -> x / (1 + r)^k, +, 0:(length(s) - 1), s)
end


text:
# Miscellaneous stuff

# run-length encoding
"""
    rle(v) -> (vals, lens)

Return the run-length encoding of a vector as a tuple. The
first element of the tuple is a vector of values of the input
and the second is the number of consecutive occurrences of each
element.

# Examples

```jldoctest
julia> using StatsBase

julia> rle([1,1,1,2,2,3,3,3,3,2,2,2])
([1, 2, 3, 2], [3, 2, 4, 3])
```
"""
function rle(v::AbstractVector{T}) where T
    n = length(v)
    vals = T[]
    lens = Int[]

    n>0 || return (vals,lens)

    cv = v[1]
    cl = 1

    i = 2
    @inbounds while i <= n
        vi = v[i]
        if isequal(vi, cv)
            cl += 1
        else
            push!(vals, cv)
            push!(lens, cl)
            cv = vi
            cl = 1
        end
        i += 1
    end

    # the last section
    push!(vals, cv)
    push!(lens, cl)

    return (vals, lens)
end

# inverse run-length encoding
"""
    inverse_rle(vals, lens)

Reconstruct a vector from its run-length encoding (see [`rle`](@ref)).
`vals` is a vector of the values and `lens` is a vector of the corresponding
run lengths.
"""
function inverse_rle(vals::AbstractVector{T}, lens::IntegerVector) where T
    m = length(vals)
    length(lens) == m || raise_dimerror()

    r = Vector{T}(undef, sum(lens))
    p = 0
    @inbounds for i = 1 : m
        j = lens[i]
        v = vals[i]
        while j > 0
            r[p+=1] = v
            j -=1
        end
    end
    return r
end

"""
    indexmap(a)

Construct a dictionary that maps each unique value in `a` to
the index of its first occurrence in `a`.
"""
function indexmap(a::AbstractArray{T}) where T
    d = Dict{T,Int}()
    for i = 1 : length(a)
        @inbounds k = a[i]
        if !haskey(d, k)
            d[k] = i
        end
    end
    return d
end

"""
    levelsmap(a)

Construct a dictionary that maps each of the `n` unique values
in `a` to a number between 1 and `n`.
"""
function levelsmap(a::AbstractArray{T}) where T
    d = Dict{T,Int}()
    index = 1
    for i = 1 : length(a)
        @inbounds k = a[i]
        if !haskey(d, k)
            d[k] = index
            index += 1
        end
    end
    return d
end

"""
    indicatormat(x, k::Integer; sparse=false)

Construct a boolean matrix `I` of size `(k, length(x))` such that
`I[x[i], i] = true` and all other elements are set to `false`.
If `sparse` is `true`, the output will be a sparse matrix, otherwise
it will be dense (default).

# Examples

```jldoctest
julia> using StatsBase

julia> indicatormat([1 2 2], 2)
2×3 Matrix{Bool}:
 1  0  0
 0  1  1
```
"""
function indicatormat(x::IntegerArray, k::Integer; sparse::Bool=false)
    sparse ? _indicatormat_sparse(x, k) : _indicatormat_dense(x, k)
end

"""
    indicatormat(x, c=sort(unique(x)); sparse=false)

Construct a boolean matrix `I` of size `(length(c), length(x))`.
Let `ci` be the index of `x[i]` in `c`. Then `I[ci, i] = true` and
all other elements are `false`.
"""
function indicatormat(x::AbstractArray, c::AbstractArray; sparse::Bool=false)
    sparse ? _indicatormat_sparse(x, c) : _indicatormat_dense(x, c)
end

indicatormat(x::AbstractArray; sparse::Bool=false) =
    indicatormat(x, sort!(unique(x)); sparse=sparse)

function _indicatormat_dense(x::IntegerArray, k::Integer)
    n = length(x)
    r = zeros(Bool, k, n)
    for i = 1 : n
        r[x[i], i] = true
    end
    return r
end

function _indicatormat_dense(x::AbstractArray{T}, c::AbstractArray{T}) where T
    d = indexmap(c)
    m = length(c)
    n = length(x)
    r = zeros(Bool, m, n)
    o = 0
    @inbounds for i = 1 : n
        xi = x[i]
        r[o + d[xi]] = true
        o += m
    end
    return r
end

_indicatormat_sparse(x::IntegerArray, k::Integer) = (n = length(x); sparse(x, 1:n, true, k, n))

function _indicatormat_sparse(x::AbstractArray{T}, c::AbstractArray{T}) where T
    d = indexmap(c)
    m = length(c)
    n = length(x)

    rinds = Vector{Int}(undef, n)
    @inbounds for i = 1 : n
        rinds[i] = d[x[i]]
    end
    return sparse(rinds, 1:n, true, m, n)
end


text:
######################################################################################################
# Copy from https://github.com/xuedong/LinBAI.jl  by Xuedong Shang                                   #
######################################################################################################
using JLD2;
using Distributed;
using Printf;
@everywhere include("runit_threshold.jl");
@everywhere include("../utilities/thresholds.jl");
include("../utilities/experiment_helpers.jl");

# setup
dist = Gaussian();
dim = 6;
μ = zeros(dim); µ[1] = 1.;

arms = Vector{Float64}[]
for k in 1:dim
	v = zeros(dim); v[k] = 1.; push!(arms, v);
end
ω = 0.01; v = zeros(dim); v[1] = cos(ω); v[2] = sin(ω); push!(arms, v)
threshold = 0.9;

pep = LinearThreshold(dist, arms, threshold);
srs = [FWSampling(), ConvexGame(CTracking), LearnerK(CTracking), RoundRobin()];

δs = (0.1,0.01,0.001,0.0001); # confidence
βs = GK16.(δs);
repeats = 1000;
seed = 1234;

# compute
@time data = pmap(
    ((sr,i),) -> runit(seed+i, sr, μ, pep, βs),
    Iterators.product(srs, 1:repeats)
);
dump_stats(pep, μ, δs, βs, srs, data, repeats);
# save
@save isempty(ARGS) ? "Threshold.dat" : ARGS[1]  dist μ pep srs data δs βs repeats seed
# visualise by loading viz.jl


text:
export
    coordarray,
    rmsd,
    displacements,
    sqdistance,
    distance,
    bondangle,
    dihedralangle,
    omegaangle,
    phiangle,
    psiangle,
    omegaangles,
    phiangles,
    psiangles,
    ramachandranangles,
    SpatialMap,
    ContactMap,
    DistanceMap,
    showcontactmap,
    MetaGraph

"""
    coordarray(element, atom_selectors...)

Get the atomic coordinates in Å of a `StructuralElementOrList` as a 2D
`Array` with each column corresponding to one atom.

Additional arguments are atom selector functions - only atoms that return
`true` from all the functions are retained.
"""
function coordarray(el::StructuralElementOrList, atom_selectors::Function...)
    at_list = collectatoms(el, atom_selectors...)
    coords_out = zeros(3, length(at_list))
    for j in eachindex(at_list)
        coords_out[1, j] = x(at_list[j])
        coords_out[2, j] = y(at_list[j])
        coords_out[3, j] = z(at_list[j])
    end
    return coords_out
end

# Selector functions ignored
coordarray(coords_in::Array{Float64}, atom_selectors::Function...) = coords_in

"""
    rmsd(element_one, element_two, atom_selectors...)
    rmsd(coords_one, coords_two)

Get the root-mean-square deviation (RMSD) in Å between two
`StructuralElementOrList`s or two coordinate `Array`s of the same size.

Assumes they are already superimposed.
Additional arguments are atom selector functions - only atoms that return
`true` from the functions are retained.
"""
function rmsd(coords_one::Array{Float64}, coords_two::Array{Float64})
    if size(coords_one) != size(coords_two)
        throw(ArgumentError("Sizes of coordinate arrays are different - cannot calculate RMSD"))
    end
    diff = coords_one - coords_two
    return sqrt.(dot(diff, diff) / size(coords_one, 2))
end

function rmsd(el_one::StructuralElementOrList,
            el_two::StructuralElementOrList,
            atom_selectors::Function...)
    return rmsd(
        coordarray(el_one, atom_selectors...),
        coordarray(el_two, atom_selectors...))
end

"""
    displacements(element_one, element_two, atom_selectors...)
    displacements(coords_one, coords_two)

Get the displacements in Å between atomic coordinates from two
`StructuralElementOrList`s or two coordinate `Array`s of the same size.

Assumes they are already superimposed.
Additional arguments are atom selector functions - only atoms that return
`true` from the functions are retained.
"""
function displacements(coords_one::Array{Float64}, coords_two::Array{Float64})
    if size(coords_one) != size(coords_two)
        throw(ArgumentError("Sizes of coordinate arrays are different - cannot calculate displacements"))
    end
    diff = coords_one - coords_two
    return sqrt.(sum(diff .* diff, dims=1))[:]
end

function displacements(el_one::StructuralElementOrList,
                    el_two::StructuralElementOrList,
                    atom_selectors::Function...)
    return displacements(
        coordarray(el_one, atom_selectors...),
        coordarray(el_two, atom_selectors...))
end

"""
    sqdistance(element_one, element_two, atom_selectors...)

Get the minimum square distance in Å between two
`StructuralElementOrList`s.

Additional arguments are atom selector functions - only atoms that return
`true` from the functions are retained.
"""
function sqdistance(el_one::StructuralElementOrList,
                    el_two::StructuralElementOrList,
                    atom_selectors::Function...)
    coords_one = coordarray(el_one, atom_selectors...)
    coords_two = coordarray(el_two, atom_selectors...)
    min_sq_dist = Inf
    for i in 1:size(coords_one, 2)
        for j in 1:size(coords_two, 2)
            @inbounds sq_dist = (coords_one[1, i] - coords_two[1, j]) ^ 2 + (coords_one[2, i] - coords_two[2, j]) ^ 2 + (coords_one[3, i] - coords_two[3, j]) ^ 2
            if sq_dist < min_sq_dist
                min_sq_dist = sq_dist
            end
        end
    end
    return min_sq_dist
end

function sqdistance(at_one::AbstractAtom, at_two::AbstractAtom)
    return (x(at_one) - x(at_two)) ^ 2 +
           (y(at_one) - y(at_two)) ^ 2 +
           (z(at_one) - z(at_two)) ^ 2
end

"""
    distance(element_one, element_two, atom_selectors...)

Get the minimum distance in Å between two `StructuralElementOrList`s.

Additional arguments are atom selector functions - only atoms that return
`true` from the functions are retained.
"""
function distance(el_one::StructuralElementOrList,
                      el_two::StructuralElementOrList,
                      atom_selectors::Function...)
    return sqrt(sqdistance(el_one, el_two, atom_selectors...))
end

function distance(at_one::AbstractAtom, at_two::AbstractAtom)
    return sqrt(sqdistance(at_one, at_two))
end

"""
    bondangle(atom_a, atom_b, atom_c)
    bondangle(vec_ba, vec_bc)

Calculate the bond or pseudo-bond angle in radians between three
`AbstractAtom`s or two vectors.

The angle between B→A and B→C is returned in the range 0 to π.
"""
function bondangle(at_a::AbstractAtom,
                at_b::AbstractAtom,
                at_c::AbstractAtom)
    return bondangle(
        coords(at_a) - coords(at_b),
        coords(at_c) - coords(at_b)
    )
end

function bondangle(vec_a::Vector{<:Real}, vec_b::Vector{<:Real})
    return acos(dot(vec_a, vec_b) / (norm(vec_a) * norm(vec_b)))
end

"""
    dihedralangle(atom_a, atom_b, atom_c, atom_d)
    dihedralangle(vec_ab, vec_bc, vec_cd)

Calculate the dihedral angle in radians defined by four `AbstractAtom`s or
three vectors.

The angle between the planes defined by atoms (A, B, C) and (B, C, D) is
returned in the range -π to π.
"""
function dihedralangle(at_a::AbstractAtom,
            at_b::AbstractAtom,
            at_c::AbstractAtom,
            at_d::AbstractAtom)
    return dihedralangle(
        coords(at_b) - coords(at_a),
        coords(at_c) - coords(at_b),
        coords(at_d) - coords(at_c))
end

function dihedralangle(vec_a::Vector{<:Real},
                    vec_b::Vector{<:Real},
                    vec_c::Vector{<:Real})
    return atan(
        dot(cross(cross(vec_a, vec_b), cross(vec_b, vec_c)), vec_b / norm(vec_b)),
        dot(cross(vec_a, vec_b), cross(vec_b, vec_c)))
end

"""
    omegaangle(res, res_previous)
    omegaangle(chain, res_id)

Calculate the omega angle in radians for an `AbstractResidue`.

Arguments can either be a residue and the previous residue or a chain and
the position as a residue ID.
The first residue (or one at the given index) requires the atoms "N" and
"CA" and the previous residue requires the atoms "CA" and "C".
The angle is in the range -π to π.
"""
function omegaangle(res::AbstractResidue, res_prev::AbstractResidue)
    if !("CA" in atomnames(res_prev))
        throw(ArgumentError("Atom with atom name \"CA\" not found in previous residue"))
    elseif !("C"  in atomnames(res_prev))
        throw(ArgumentError("Atom with atom name \"C\" not found in previous residue"))
    elseif !("N"  in atomnames(res))
        throw(ArgumentError("Atom with atom name \"N\" not found in residue"))
    elseif !("CA" in atomnames(res))
        throw(ArgumentError("Atom with atom name \"CA\" not found in residue"))
    end
    return dihedralangle(res_prev["CA"], res_prev["C"], res["N"], res["CA"])
end

function omegaangle(chain::Chain, res_id::Union{Integer, AbstractString})
    inds = findall(r -> r == string(res_id), resids(chain))
    if length(inds) != 1
        throw(ArgumentError("\"$res_id\" is an invalid residue ID"))
    end
    i = inds[1]
    if i == 1 || !sequentialresidues(chain[resids(chain)[i - 1]], chain[resids(chain)[i]])
        throw(ArgumentError("Cannot calculate omega angle for residue \"$res_id\" due to a lack of connected residues"))
    end
    return omegaangle(chain[resids(chain)[i]], chain[resids(chain)[i - 1]])
end

"""
    phiangle(res, res_previous)
    phiangle(chain, res_id)

Calculate the phi angle in radians for an `AbstractResidue`.

Arguments can either be a residue and the previous residue or a chain and the
position as a residue ID.
The first residue (or one at the given index) requires the atoms "N", "CA" and
"C" and the previous residue requires the atom "C".
The angle is in the range -π to π.
"""
function phiangle(res::AbstractResidue, res_prev::AbstractResidue)
    if !("C"  in atomnames(res_prev))
        throw(ArgumentError("Atom with atom name \"C\" not found in previous residue"))
    elseif !("N"  in atomnames(res))
        throw(ArgumentError("Atom with atom name \"N\" not found in residue"))
    elseif !("CA" in atomnames(res))
        throw(ArgumentError("Atom with atom name \"CA\" not found in residue"))
    elseif !("C"  in atomnames(res))
        throw(ArgumentError("Atom with atom name \"C\" not found in residue"))
    end
    return dihedralangle(res_prev["C"], res["N"], res["CA"], res["C"])
end

function phiangle(chain::Chain, res_id::Union{Integer, AbstractString})
    inds = findall(r -> r == string(res_id), resids(chain))
    if length(inds) != 1
        throw(ArgumentError("\"$res_id\" is an invalid residue ID"))
    end
    i = inds[1]
    if i == 1 || !sequentialresidues(chain[resids(chain)[i - 1]], chain[resids(chain)[i]])
        throw(ArgumentError("Cannot calculate phi angle for residue \"$res_id\" due to a lack of connected residues"))
    end
    return phiangle(chain[resids(chain)[i]], chain[resids(chain)[i - 1]])
end

"""
    psiangle(res, res_next)
    psiangle(chain, res_id)

Calculate the psi angle in radians for an `AbstractResidue`.

Arguments can either be a residue and the next residue or a chain and the
position as a residue ID.
The first residue (or one at the given index) requires the atoms "N", "CA" and
"C" and the next residue requires the atom "N".
The angle is in the range -π to π.
"""
function psiangle(res::AbstractResidue, res_next::AbstractResidue)
    if !("N"  in atomnames(res))
        throw(ArgumentError("Atom with atom name \"N\" not found in residue"))
    elseif !("CA" in atomnames(res))
        throw(ArgumentError("Atom with atom name \"CA\" not found in residue"))
    elseif !("C"  in atomnames(res))
        throw(ArgumentError("Atom with atom name \"C\" not found in residue"))
    elseif !("N"  in atomnames(res_next))
        throw(ArgumentError("Atom with atom name \"N\" not found in next residue"))
    end
    return dihedralangle(res["N"], res["CA"], res["C"], res_next["N"])
end

function psiangle(chain::Chain, res_id::Union{Integer, AbstractString})
    inds = findall(r -> r == string(res_id), resids(chain))
    if length(inds) != 1
        throw(ArgumentError("\"$res_id\" is an invalid residue ID"))
    end
    i = inds[1]
    if i == length(chain) || !sequentialresidues(chain[resids(chain)[i]], chain[resids(chain)[i + 1]])
        throw(ArgumentError("Cannot calculate psi angle for residue \"$res_id\" due to a lack of connected residues"))
    end
    return psiangle(chain[resids(chain)[i]], chain[resids(chain)[i + 1]])
end

"""
    omegaangles(element, residue_selectors...)

Calculate the `Vector` of omega angles of a `StructuralElementOrList`.

The vectors have `NaN` for residues where an angle cannot be calculated, e.g.
due to missing atoms or lack of an adjacent residue.
The angle is in the range -π to π.
Additional arguments are residue selector functions - only residues that return
`true` from the functions are retained.
"""
function omegaangles(el::StructuralElementOrList,
                    residue_selectors::Function...)
    res_list = collectresidues(el, residue_selectors...)
    if length(res_list) < 2
        throw(ArgumentError("At least 2 residues required to calculate dihedral angles"))
    end

    omega_angles = Float64[NaN]

    for i in 2:length(res_list)
        res = res_list[i]
        res_prev = res_list[i - 1]
        if sequentialresidues(res_prev, res)
            try
                omega_angle = omegaangle(res, res_prev)
                push!(omega_angles, omega_angle)
            catch ex
                isa(ex, ArgumentError) || rethrow()
                push!(omega_angles, NaN)
            end
        else
            push!(omega_angles, NaN)
        end
    end

    return omega_angles
end

"""
    phiangles(element, residue_selectors...)

Calculate the `Vector` of phi angles of a `StructuralElementOrList`.

The vectors have `NaN` for residues where an angle cannot be calculated, e.g.
due to missing atoms or lack of an adjacent residue.
The angle is in the range -π to π.
Additional arguments are residue selector functions - only residues that return
`true` from the functions are retained.
"""
function phiangles(el::StructuralElementOrList,
                    residue_selectors::Function...)
    res_list = collectresidues(el, residue_selectors...)
    if length(res_list) < 2
        throw(ArgumentError("At least 2 residues required to calculate dihedral angles"))
    end

    phi_angles = Float64[NaN]

    for i in 2:length(res_list)
        res = res_list[i]
        res_prev = res_list[i - 1]
        if sequentialresidues(res_prev, res)
            try
                phi_angle = phiangle(res, res_prev)
                push!(phi_angles, phi_angle)
            catch ex
                isa(ex, ArgumentError) || rethrow()
                push!(phi_angles, NaN)
            end
        else
            push!(phi_angles, NaN)
        end
    end

    return phi_angles
end

"""
    psiangles(element, residue_selectors...)

Calculate the `Vector` of psi angles of a `StructuralElementOrList`.

The vectors have `NaN` for residues where an angle cannot be calculated, e.g.
due to missing atoms or lack of an adjacent residue.
The angle is in the range -π to π.
Additional arguments are residue selector functions - only residues that return
`true` from the functions are retained.
"""
function psiangles(el::StructuralElementOrList,
                    residue_selectors::Function...)
    res_list = collectresidues(el, residue_selectors...)

    if length(res_list) < 2
        throw(ArgumentError("At least 2 residues required to calculate dihedral angles"))
    end

    psi_angles = Float64[]

    for i in 1:(length(res_list) - 1)
        res = res_list[i]
        res_next = res_list[i + 1]
        if sequentialresidues(res, res_next)
            try
                psi_angle = psiangle(res, res_next)
                push!(psi_angles, psi_angle)
            catch ex
                isa(ex, ArgumentError) || rethrow()
                push!(psi_angles, NaN)
            end
        else
            push!(psi_angles, NaN)
        end
    end

    push!(psi_angles, NaN)

    return psi_angles
end

"""
    ramachandranangles(element, residue_selectors...)

Calculate the `Vector`s of phi and psi angles of a `StructuralElementOrList`.

The vectors have `NaN` for residues where an angle cannot be calculated, e.g.
due to missing atoms or lack of an adjacent residue.
The angles are in the range -π to π.
Additional arguments are residue selector functions - only residues that return
`true` from the functions are retained.
"""
function ramachandranangles(el::StructuralElementOrList,
                    residue_selectors::Function...)
    return phiangles(el, residue_selectors...), psiangles(el, residue_selectors...)
end

"A map of a structural property, e.g. a `ContactMap` or a `DistanceMap`."
abstract type SpatialMap end

"""
    ContactMap(element, contact_distance)
    ContactMap(element_one, element_two, contact_distance)
    ContactMap(bit_array_2D)

Calculate the contact map for a `StructuralElementOrList`, or between two
`StructuralElementOrList`s.

This returns a `ContactMap` type containing a `BitArray{2}` with `true` where
the sub-elements are no further than the contact distance and `false` otherwise.
When one element is given as input this returns a symmetric square matrix.

# Examples
```julia
cbetas_A = collectatoms(struc["A"], cbetaselector)
cbetas_B = collectatoms(struc["B"], cbetaselector)

# Contact map of chain A using standard C-beta and 8.0 Å definitions
ContactMap(cbetas_A, 8.0)

# Rectangular contact map of chains A and B
ContactMap(cbetas_A, cbetas_B, 8.0)
```
"""
struct ContactMap <: SpatialMap
    data::BitArray{2}
end

"""
    DistanceMap(element)
    DistanceMap(element_one, element_two)
    DistanceMap(float_array_2D)

Calculate the distance map for a `StructuralElementOrList`, or between two
`StructuralElementOrList`s.

This returns a `DistanceMap` type containing a `Array{Float64, 2}` with minimum
distances between the sub-elements.
When one element is given as input this returns a symmetric square matrix.

# Examples
```julia
cbetas_A = collectatoms(struc["A"], cbetaselector)
cbetas_B = collectatoms(struc["B"], cbetaselector)

# Distance map of chain A showing how far each residue is from the others
DistanceMap(cbetas_A)

# Rectangular distance map of chains A and B
DistanceMap(cbetas_A, cbetas_B)
```
"""
struct DistanceMap <: SpatialMap
    data::Array{Float64, 2}
end

Base.getindex(m::SpatialMap, args::Integer...) = m.data[args...]

function Base.setindex!(m::SpatialMap, v, args::Integer...)
    m.data[args...] = v
    return m
end

Base.size(m::SpatialMap) = size(m.data)
Base.size(m::SpatialMap, dim::Integer) = size(m.data, dim)

function Base.show(io::IO, cm::ContactMap)
    print(io, "Contact map of size $(size(cm))")
end

function Base.show(io::IO, dm::DistanceMap)
    print(io, "Distance map of size $(size(dm))")
end

function ContactMap(el_one::StructuralElementOrList,
                el_two::StructuralElementOrList,
                contact_dist::Real)
    sq_contact_dist = contact_dist ^ 2
    contacts = falses(length(el_one), length(el_two))
    for (i, subel_one) in enumerate(el_one)
        for (j, subel_two) in enumerate(el_two)
            if sqdistance(subel_one, subel_two) <= sq_contact_dist
                contacts[i, j] = true
            end
        end
    end
    return ContactMap(contacts)
end

function ContactMap(el::StructuralElementOrList, contact_dist::Real)
    sq_contact_dist = contact_dist ^ 2
    contacts = falses(length(el), length(el))
    el_list = collect(el)
    for i in 1:length(el)
        contacts[i, i] = true
        for j in 1:(i - 1)
            if sqdistance(el_list[i], el_list[j]) <= sq_contact_dist
                contacts[i, j] = true
                contacts[j, i] = true
            end
        end
    end
    return ContactMap(contacts)
end

function DistanceMap(el_one::StructuralElementOrList,
                el_two::StructuralElementOrList)
    dists = zeros(length(el_one), length(el_two))
    for (i, subel_one) in enumerate(el_one)
        for (j, subel_two) in enumerate(el_two)
            dists[i, j] = distance(subel_one, subel_two)
        end
    end
    return DistanceMap(dists)
end

function DistanceMap(el::StructuralElementOrList)
    dists = zeros(length(el), length(el))
    el_list = collect(el)
    for i in 1:length(el)
        for j in 1:(i - 1)
            dist = distance(el_list[i], el_list[j])
            dists[i, j] = dist
            dists[j, i] = dist
        end
    end
    return DistanceMap(dists)
end

# Plot recipe to show a ContactMap
@recipe function plot(cm::ContactMap)
    seriestype := :heatmap
    fillcolor --> :dense
    aspectratio --> 1
    xmirror --> true
    colorbar --> false
    xs = string.(1:size(cm, 2))
    ys = string.(size(cm, 1):-1:1)
    xs, ys, reverse(cm.data, dims=1)
end

# Plot recipe to show a DistanceMap
@recipe function plot(dm::DistanceMap)
    seriestype := :heatmap
    fillcolor --> :inferno
    aspectratio --> 1
    xmirror --> true
    colorbar --> true
    xs = string.(1:size(dm, 2))
    ys = string.(size(dm, 1):-1:1)
    xs, ys, reverse(dm.data, dims=1)
end

"""
    showcontactmap(contact_map)
    showcontactmap(io, contact_map)

Print a representation of a `ContactMap` to `stdout`, or a specified `IO`
instance.
A fully plotted version can be obtained with `plot(contact_map)` but that
requires Plots.jl; `showcontactmap` works without that dependency.
"""
function showcontactmap(io::IO, cm::ContactMap)
    size1 = size(cm, 1)
    # Print two y values to each line for a nicer output
    for i in 1:2:size1
        for j in 1:size(cm, 2)
            cont_one = cm[i, j]
            # Check we aren't going over the end of the map
            cont_two = i + 1 <= size1 && cm[i + 1, j]
            if cont_one && cont_two
                char = "█"
            elseif cont_one
                char = "▀"
            elseif cont_two
                char = "▄"
            else
                char = " "
            end
            print(io, char)
        end
        println(io)
    end
end

showcontactmap(cm::ContactMap) = showcontactmap(stdout, cm)

# Construct a graph of atoms where edges are contacts
function MetaGraphs.MetaGraph(el::StructuralElementOrList, contact_dist::Real)
    sq_contact_dist = contact_dist ^ 2
    el_list = collect(el)
    mg = MetaGraph(length(el_list))
    set_prop!(mg, :contactdist, Float64(contact_dist))
    for (i, el) in enumerate(el_list)
        set_prop!(mg, i, :element, el)
        for j in 1:(i - 1)
            if sqdistance(el, el_list[j]) <= sq_contact_dist
                add_edge!(mg, j, i)
            end
        end
    end
    set_indexing_prop!(mg, :element)
    return mg
end


text:
x0 = [1., 0.]
P0 = eye(2)

Phi = [0.8 0.5; 0.0 0.8]
b = zeros(2)
Q = [0.1 0.0; 0.0 1.0]

yshadow = [1.0]
H = [1.0 0.0]
R = eye(1)
M = LinearHomogSystem(x0, P0, Phi, b, Q, yshadow, H, R)


text:
export rule

@rule BIFMHelper(:out, Marginalisation) (q_in::Any, ) = ProdFinal(q_in)


text:
export DepParser
export execute_transition, form_batch, predict_transition

using ..DependencyParser.Units
import .ArcEager: execute_transition

using TextAnalysis

struct DepParser <: AbstractDepParser
  settings::Settings
  parsing_system::ParsingSystem
  model::Model

  function DepParser(settings::Settings, model_file::String)
    model = Model(model_file)
    system =  ArcEager.ArcEagerSystem()
    sort(collect(model.label_ids), by=pair->pair[end]) |>
      pairs -> map(pair -> pair[begin], pairs) |>
      labels -> set_labels!(system, labels)

    new(settings, system, model)
  end

  function DepParser(settings::Settings, model::Model, system::ParsingSystem)
    system =  ArcEager.ArcEagerSystem()
    sort(collect(model.label_ids), by=pair->pair[end]) |>
      pairs -> map(pair -> pair[begin], pairs) |>
      labels -> set_labels!(system, labels)

    new(settings, system, model)
  end
end

function (parser::DepParser)(tokens_with_tags)
  sentence = Sentence(tokens_with_tags)
  config = Configuration(sentence)
  transitions_number = 0

  while !is_terminal(config)
    transition = predict_transition(parser, config)
    transition === nothing && break
    execute_transition(parser, config, transition)

    transitions_number += 1
    transitions_number > LIMIT_TRANSITIONS_NUMBER && break
  end

  config.tree
end

function execute_transition(parser::DepParser, config::Configuration, transition::Transition)
  execute_transition(config, transition, parser.parsing_system)
end

function train!(system::ParsingSystem, train_file::String, embeddings_file::String, model_file::String)
  println("Parse train file...")
  connlu_sentences = load_connlu_file(train_file)
  settings = Settings()

  println("Build initial model...")
  model = Model(settings, system, embeddings_file, connlu_sentences)

  train!(model, settings, system, connlu_sentences)
end

function train!(train_file::String, test_file::String, results_file::String, model::Model, system::ParsingSystem, model_file::String)
  println("Parse train file...")
  connlu_sentences = load_connlu_file(train_file)
  test_sentences = load_connlu_file(test_file)
  settings = Settings(embeddings_size=100)

  training_context = TrainingContext(
    system,
    settings,
    train_file,
    connlu_sentences,
    test_sentences,
    test_file,
    results_file,
    model_file,
    beam_coef = 0.05
  )

  train!(model, training_context)
end

# Temp code for test

function default_train()
  system = ArcEager.ArcEagerSystem()
  train_file = "F:\\ed_soft\\parser_materials\\UD_English-EWT-master\\en_ewt-ud-train.conllu"
  test_file = "F:\\ed_soft\\parser_materials\\UD_English-EWT-master\\en_ewt-ud-dev.conllu"
  embeddings_file = "F:\\ed_soft\\parser_materials\\model.txt"
  model_file = "tmp/model_b3000_adam_fl2_e100_beam_nt_2_new_batch_48"
  results_file = "tmp/results_b3000_adam_fl2_e100_beam_nt_2_new_batch_48"

  connlu_sentences = load_connlu_file(train_file)
  settings = Settings(embeddings_size=100)
  model = cache_data((args...) -> Model(args...), "tmp/cache-new-target-2", "model_cache_e100-ewt", settings, system, embeddings_file, connlu_sentences)
  enable_cuda(model)

  model = Model(model_file * "_last.txt")

  sort(collect(model.label_ids), by=pair->pair[end]) |>
        pairs -> map(pair -> pair[begin], pairs) |>
        labels -> set_labels!(system, labels)

  train!(train_file, test_file, results_file, model, system, model_file)
end

function predict_transition(parser::DepParser, config::Configuration)
  predict_transition(parser.model, parser.settings, parser.parsing_system, config)
end

max_stars_repo_name: LarsWl/DependencyParser

max_issues_repo_name: LarsWl/DependencyParser

max_forks_repo_name: LarsWl/DependencyParser


text:
# Test variable reference makers
@testset "Variable Reference Makers" begin
    # initialize model and references
    m = InfiniteModel()
    @infinite_parameter(m, 1 <= par1 <= 2)
    @infinite_parameter(m, 1 <= par2 <= 2)
    @infinite_variable(m, inf1(par1))
    @infinite_variable(m, inf2(par1, par2))
    # test make_point_variable_ref
    @testset "make_point_variable_ref" begin
        idx = m.next_var_index + 1
        @test make_point_variable_ref(inf1) == PointVariableRef(m, idx)
    end
    # test make_reduced_variable_ref (from ivref)
    @testset "make_reduced_variable_ref (from ivref)" begin
        # test first addition
        idx = m.next_var_index + 1
        rvref = ReducedInfiniteVariableRef(m, idx)
        @test make_reduced_variable_ref(inf2, 1, 1) == rvref
        @test m.reduced_info[idx].infinite_variable_ref == inf2
        @test m.reduced_info[idx].eval_supports == Dict(1 => 1)
        @test m.infinite_to_reduced[JuMP.index(inf2)] == [idx]
        # test second addition
        idx = m.next_var_index + 1
        rvref = ReducedInfiniteVariableRef(m, idx)
        @test make_reduced_variable_ref(inf2, 1, 1) == rvref
        @test m.reduced_info[idx].infinite_variable_ref == inf2
        @test m.reduced_info[idx].eval_supports == Dict(1 => 1)
        @test m.infinite_to_reduced[JuMP.index(inf2)] == [idx - 1, idx]
        # undo changes
        delete!(m.infinite_to_reduced, JuMP.index(inf2))
    end
    # test make_reduced_variable_ref (from rvref)
    @testset "make_reduced_variable_ref (from rvref)" begin
        # test first addition
        idx = m.next_var_index + 1
        rvref = ReducedInfiniteVariableRef(m, idx)
        @test make_reduced_variable_ref(inf2, Dict(1 => 1)) == rvref
        @test m.reduced_info[idx].infinite_variable_ref == inf2
        @test m.reduced_info[idx].eval_supports == Dict(1 => 1)
        @test m.infinite_to_reduced[JuMP.index(inf2)] == [idx]
        # test second addition
        idx = m.next_var_index + 1
        rvref = ReducedInfiniteVariableRef(m, idx)
        @test make_reduced_variable_ref(inf2, Dict(1 => 1)) == rvref
        @test m.reduced_info[idx].infinite_variable_ref == inf2
        @test m.reduced_info[idx].eval_supports == Dict(1 => 1)
        @test m.infinite_to_reduced[JuMP.index(inf2)] == [idx - 1, idx]
    end
end

# Test _get_param_value_list
@testset "_get_param_value_list" begin
    # initialize model and data
    m = InfiniteModel()
    @infinite_parameter(m, 1 <= par1 <= 2)
    @infinite_parameter(m, 0 <= pars1[1:2] <= 2)
    data1 = DiscreteMeasureData(par1, [0.5, 0.5], [1, 2])
    data2 = DiscreteMeasureData(pars1, [1, 1], [[1, 0.5], [2, 1.5]])
    # test with DiscreteMeasureData
    @testset "DiscreteMeasureData" begin
        @test InfiniteOpt._get_param_value_list(par1, data1) == [1, 2]
    end
    # test with MultiDiscreteMeasureData
    @testset "MultiDiscreteMeasureData" begin
        @test InfiniteOpt._get_param_value_list(pars1[1], data2) == [1, 2]
        @test InfiniteOpt._get_param_value_list(pars1[2], data2) == [0.5, 1.5]
    end
end

# Test measure expansion methods
@testset "expand_measure" begin
    # initialize model and references
    m = InfiniteModel()
    @infinite_parameter(m, 1 <= par1 <= 2)
    @infinite_parameter(m, 1 <= par2 <= 2)
    @infinite_parameter(m, 1 <= pars1[1:2] <= 2)
    @infinite_parameter(m, 1 <= pars2[1:2] <= 2)
    @infinite_variable(m, inf1(par1))
    @infinite_variable(m, inf2(par1, par2))
    @infinite_variable(m, inf3(par2))
    @infinite_variable(m, inf4(pars1))
    @infinite_variable(m, inf5(pars1, pars2))
    @infinite_variable(m, inf6(pars2))
    @infinite_variable(m, inf7(par1, par2, pars1))
    @hold_variable(m, x)
    # prepare measures
    w(t) = 3
    data1 = DiscreteMeasureData(par1, [0.5, 0.5], [1, 2])
    data2 = DiscreteMeasureData(par2, [0.5, 0.5], [1, 1], weight_function = w)
    data3 = DiscreteMeasureData(pars1, [1, 1], [[1, 1], [2, 2]])
    data4 = DiscreteMeasureData(pars2, [2, 2], [[1, 1], [1, 1]])
    meas1 = measure(inf1, data1)
    meas2 = measure(2 * inf3 * x - par1, data2)
    # prepare point_mapper function and info
    function mapper(m2, pvref, ivref, support)
        set_name(pvref, string(InfiniteOpt._root_name(ivref), support))
        return
    end
    trans_model = Model()
    map_args = (trans_model, mapper)
    # test expand_measure (infinite variable)
    @testset "Infinite Variable" begin
        # test single param infinite var with measure param
        idx = m.next_var_index + 1
        pts = [PointVariableRef(m, idx), PointVariableRef(m, idx + 1)]
        expected = 0.5 * (pts[1] + pts[2])
        @test expand_measure(inf1, data1, map_args...) == expected
        # test single param infinite var without measure param
        expected = 3inf1
        @test expand_measure(inf1, data2, map_args...) == expected
        # test single param infinite var with measure param and others
        idx = m.next_var_index + 1
        rv1 = ReducedInfiniteVariableRef(m, idx)
        m.reduced_info[idx] = ReducedInfiniteInfo(inf2, Dict(1 => 1))
        rv2 = ReducedInfiniteVariableRef(m, idx + 1)
        m.reduced_info[idx + 1] = ReducedInfiniteInfo(inf2, Dict(1 => 2))
        expected = 0.5 * (rv1 + rv2)
        @test expand_measure(inf2, data1, map_args...) == expected
        # test array param infinite var with measure param
        idx = m.next_var_index + 1
        pts = [PointVariableRef(m, idx), PointVariableRef(m, idx + 1)]
        expected = pts[1] + pts[2]
        @test expand_measure(inf4, data3, map_args...) == expected
        # test array param infinite var without measure param
        expected = 4inf4
        @test expand_measure(inf4, data4, map_args...) == expected
        # test array param infinite var with measure param and others
        idx = m.next_var_index + 1
        supp1 = convert(JuMP.Containers.SparseAxisArray, [1, 1])
        supp2 = convert(JuMP.Containers.SparseAxisArray, [2, 2])
        rv1 = ReducedInfiniteVariableRef(m, idx)
        m.reduced_info[idx] = ReducedInfiniteInfo(inf5, Dict(1 => supp1))
        rv2 = ReducedInfiniteVariableRef(m, idx + 1)
        m.reduced_info[idx + 1] = ReducedInfiniteInfo(inf5, Dict(1 => supp2))
        expected = rv1 + rv2
        @test expand_measure(inf5, data3, map_args...) == expected
    end
    # test expand_measure (reduced infinite variable)
    @testset "Reduced Variable" begin
        # test single param reduced var without measure param
        rv = ReducedInfiniteVariableRef(m, -1)
        m.reduced_info[-1] = ReducedInfiniteInfo(inf2, Dict(1 => 1))
        expected = 0.5 * (rv + rv)
        @test expand_measure(rv, data1, map_args...) == expected
        # test single param reduced var with measure param
        idx = m.next_var_index + 1
        pts = [PointVariableRef(m, idx), PointVariableRef(m, idx + 1)]
        expected = 1.5 * (pts[1] + pts[2])
        @test expand_measure(rv, data2, map_args...) == expected
        # test single param reduced var with measure param and others
        rv = ReducedInfiniteVariableRef(m, -2)
        m.reduced_info[-2] = ReducedInfiniteInfo(inf7, Dict(1 => 1))
        idx = m.next_var_index + 1
        rv1 = ReducedInfiniteVariableRef(m, idx)
        m.reduced_info[idx] = ReducedInfiniteInfo(inf7, Dict(1 => 1, 2 => 1))
        rv2 = ReducedInfiniteVariableRef(m, idx + 1)
        m.reduced_info[idx + 1] = ReducedInfiniteInfo(inf7,
                                                         Dict(1 => 2, 2 => 1))
        expected = 1.5 * (rv1 + rv2)
        @test expand_measure(rv, data2, map_args...) == expected
        # test array param reduced var without measure param
        supp1 = convert(JuMP.Containers.SparseAxisArray, [1, 1])
        rv = ReducedInfiniteVariableRef(m, -1)
        m.reduced_info[-1] = ReducedInfiniteInfo(inf5, Dict(1 => supp1))
        expected = rv + rv
        @test expand_measure(rv, data3, map_args...) == expected
        # test array param reduced var with measure param
        idx = m.next_var_index + 1
        pts = [PointVariableRef(m, idx), PointVariableRef(m, idx + 1)]
        expected = 2 * (pts[1] + pts[2])
        @test expand_measure(rv, data4, map_args...) == expected
        # test array param reduced var with measure param and others
        rv = ReducedInfiniteVariableRef(m, -2)
        m.reduced_info[-2] = ReducedInfiniteInfo(inf7, Dict(1 => 1))
        idx = m.next_var_index + 1
        supp1 = convert(JuMP.Containers.SparseAxisArray, [1, 1])
        supp2 = convert(JuMP.Containers.SparseAxisArray, [2, 2])
        rv1 = ReducedInfiniteVariableRef(m, idx)
        m.reduced_info[idx] = ReducedInfiniteInfo(inf7,
                                                     Dict(1 => 1, 3 => supp1))
        rv2 = ReducedInfiniteVariableRef(m, idx + 1)
        m.reduced_info[idx + 1] = ReducedInfiniteInfo(inf7,
                                                       Dict(1 => 2, 3 => supp2))
        expected = rv1 + rv2
        @test expand_measure(rv, data3, map_args...) == expected
    end
    # test expand_measure (finite variable)
    @testset "Finite Variable" begin
        # test with single parameter measure
        expected = 0.5 * (x + x)
        @test expand_measure(x, data1, map_args...) == expected
        # test with multi parameter measure
        expected = x + x
        @test expand_measure(x, data3, map_args...) == expected
    end
    # test expand_measure (parameter)
    @testset "Parameter" begin
        # test with different parameter
        expected = 1.5 * (par1 + par1)
        @test expand_measure(par1, data2, map_args...) == expected
        # test with same parameter
        expected = zero(JuMP.GenericAffExpr{Float64, GeneralVariableRef})
        expected.constant = 0.5 * (1 + 2)
        @test expand_measure(par1, data1, map_args...) == expected
    end
    # test expand_measure (multi measure data with parameter)
    @testset "MultiMeasureData Parameter" begin
        # test with different parameter
        expected = par1 + par1
        @test expand_measure(par1, data3, map_args...) == expected
        # test with same parameter
        expected = zero(JuMP.GenericAffExpr{Float64, GeneralVariableRef})
        expected.constant = 1 + 2
        @test expand_measure(pars1[2], data3, map_args...) == expected
    end
    # test expand_measure (AffExpr)
    @testset "AffExpr" begin
        # test single param AffExpr, no measures
        expr = 2inf1 + par1 - x + 3par2 - 3
        idx = m.next_var_index + 1
        pts = [PointVariableRef(m, idx), PointVariableRef(m, idx + 1)]
        expected = 0.5 * (2pts[1] + 2pts[2] + 1 + 2 - x - x + 3par2 + 3par2 - 6)
        @test expand_measure(expr, data1, map_args...) == expected
        # test single param AffExpr, with measures
        expr = meas2 - x + par1
        idx = m.next_var_index + 1
        pts = [PointVariableRef(m, idx), PointVariableRef(m, idx + 1)]
        expected = 0.5 * (6 * (pts[1] * x + pts[2] * x) - 2x - 6)
        @test expand_measure(expr, data1, map_args...) == expected
        # test array param AffExpr, no measures
        expr = inf4 + par1 - x + 3pars1[2] - 1
        idx = m.next_var_index + 1
        pts = [PointVariableRef(m, idx), PointVariableRef(m, idx + 1)]
        expected = pts[1] + pts[2] + 2par1 - 2x + 9 - 2
        @test expand_measure(expr, data3, map_args...) == expected
        # test array param AffExpr, with measures
        expr = meas2 - x + par1 - inf4
        idx = m.next_var_index + 1
        pts = [PointVariableRef(m, idx), PointVariableRef(m, idx + 1),
               PointVariableRef(m, idx + 2), PointVariableRef(m, idx + 3)]
        expected = 6 * (pts[1] * x + pts[2] * x) - 4par1 - 2x - pts[3] - pts[4]
        @test expand_measure(expr, data3, map_args...) == expected
    end
    # test expand_measure (QuadExpr)
    @testset "QuadExpr" begin
        # test single param QuadExpr with both variables integrated or not
        expr = 2 * inf1 * inf2 - inf3 * inf4 + x + 2
        idx = m.next_var_index + 1
        pts = [PointVariableRef(m, idx), PointVariableRef(m, idx + 1)]
        rv1 = ReducedInfiniteVariableRef(m, idx + 2)
        m.reduced_info[idx + 2] = ReducedInfiniteInfo(inf2, Dict(1 => 1))
        rv2 = ReducedInfiniteVariableRef(m, idx + 3)
        m.reduced_info[idx + 3] = ReducedInfiniteInfo(inf2, Dict(1 => 2))
        expected = 0.5 * (2 * pts[1] * rv1 + 2 * pts[2] * rv2 - 2 * inf3 *
                          inf4 + 2x + 4)
        @test expand_measure(expr, data1, map_args...) == expected
        # test single param QuadExpr with first variable not integrated
        expr = 3 * inf3 * inf1 + pars1[1] - 1
        idx = m.next_var_index + 1
        pts = [PointVariableRef(m, idx), PointVariableRef(m, idx + 1)]
        expected = 0.5 * (3 * inf3 * pts[1] + 3 * inf3 * pts[2] + 2pars1[1] - 2)
        @test expand_measure(expr, data1, map_args...) == expected
        # test single param QuadExpr with first variable integrated
        expr = 3 * inf2 * inf1 + pars1[1] + 1
        idx = m.next_var_index + 1
        rv1 = ReducedInfiniteVariableRef(m, idx)
        m.reduced_info[idx] = ReducedInfiniteInfo(inf2, Dict(2 => 1))
        rv2 = ReducedInfiniteVariableRef(m, idx + 1)
        m.reduced_info[idx + 1] = ReducedInfiniteInfo(inf2, Dict(2 => 1))
        expected = 1.5 * (3 * rv1 * inf1 + 3 * rv2 * inf1 + 2pars1[1] + 2)
        @test expand_measure(expr, data2, map_args...) == expected
        # test array param QuadExpr with both variables integrated
        expr = 2 * inf4 * inf5 - inf1 * inf2 + x + 2
        idx = m.next_var_index + 1
        pts = [PointVariableRef(m, idx), PointVariableRef(m, idx + 1)]
        supp1 = convert(JuMP.Containers.SparseAxisArray, [1, 1])
        supp2 = convert(JuMP.Containers.SparseAxisArray, [2, 2])
        rv1 = ReducedInfiniteVariableRef(m, idx + 2)
        m.reduced_info[idx + 2] = ReducedInfiniteInfo(inf5, Dict(1 => supp1))
        rv2 = ReducedInfiniteVariableRef(m, idx + 3)
        m.reduced_info[idx + 3] = ReducedInfiniteInfo(inf5, Dict(1 => supp2))
        expected = 2 * pts[1] * rv1 + 2 * pts[2] * rv2 - 2 * inf1 * inf2 + 2x + 4
        @test expand_measure(expr, data3, map_args...) == expected
        # test array param QuadExpr with first variable not integrated
        expr = 3 * inf1 * inf4 + par1 - 1
        idx = m.next_var_index + 1
        pts = [PointVariableRef(m, idx), PointVariableRef(m, idx + 1)]
        expected = 3 * inf1 * pts[1] + 3 * inf1 * pts[2] + 2par1 - 2
        @test expand_measure(expr, data3, map_args...) == expected
        # test array param QuadExpr with first variable integrated
        expr = 3 * inf5 * inf1 + pars1[1] + 1
        idx = m.next_var_index + 1
        supp = convert(JuMP.Containers.SparseAxisArray, [1, 1])
        rv1 = ReducedInfiniteVariableRef(m, idx)
        m.reduced_info[idx] = ReducedInfiniteInfo(inf5, Dict(2 => supp))
        rv2 = ReducedInfiniteVariableRef(m, idx + 1)
        m.reduced_info[idx + 1] = ReducedInfiniteInfo(inf5, Dict(2 => supp))
        expected = 2 * (3 * rv1 * inf1 + 3 * rv2 * inf1 + 2pars1[1] + 2)
        @test expand_measure(expr, data4, map_args...) == expected
    end
    # test expand_measure (QuadExpr) with parameter quadratics that become numbers
    @testset "QuadExpr with Number Quadratics" begin
        # test single parameter with first quadratic term becomes number
        expr = par1 * inf1 + 2
        idx = m.next_var_index + 1
        pts = [PointVariableRef(m, idx), PointVariableRef(m, idx + 1)]
        expected = zero(JuMP.GenericQuadExpr{Float64, GeneralVariableRef})
        expected.aff = 0.5 * (pts[1] + 2 * pts[2] + 4)
        @test expand_measure(expr, data1, map_args...)== expected
        # test single parameter with first quadratic term becomes number, 2nd constant
        expr = par1 * x + 2
        expected = zero(JuMP.GenericQuadExpr{Float64, GeneralVariableRef})
        expected.aff = 0.5 * (x + 2 * x + 4)
        @test expand_measure(expr, data1, map_args...) == expected
        # test single parameter with both quadratic terms become numbers
        expr = par1 * par1 + 2
        expected = zero(JuMP.GenericQuadExpr{Float64, GeneralVariableRef})
        expected.aff = 0.5 * (1 + 2 * 2 + 4)
        @test expand_measure(expr, data1, map_args...) == expected
        # test single parameter with 2nd quadratic term becomes number
        expr = inf1 * par1 + 2
        idx = m.next_var_index + 1
        pts = [PointVariableRef(m, idx), PointVariableRef(m, idx + 1)]
        expected = zero(JuMP.GenericQuadExpr{Float64, GeneralVariableRef})
        expected.aff = 0.5 * (pts[1] + 2 * pts[2] + 4)
        @test expand_measure(expr, data1, map_args...) == expected
        # test single parameter with 2nd quadratic term becomes number, 1st constant
        expr = x * par1 + 2
        expected = zero(JuMP.GenericQuadExpr{Float64, GeneralVariableRef})
        expected.aff = 0.5 * (x + 2 * x + 4)
        @test expand_measure(expr, data1, map_args...) == expected
        # test array parameter with first quadratic term becomes number
        expr = pars1[1] * inf4 + 2
        idx = m.next_var_index + 1
        pts = [PointVariableRef(m, idx), PointVariableRef(m, idx + 1)]
        expected = zero(JuMP.GenericQuadExpr{Float64, GeneralVariableRef})
        expected.aff = pts[1] + 2 * pts[2] + 4
        @test expand_measure(expr, data3, map_args...) == expected
        # test array parameter with first quadratic term becomes number, 2nd constant
        expr = pars1[1] * x + 2
        expected = zero(JuMP.GenericQuadExpr{Float64, GeneralVariableRef})
        expected.aff = x + 2 * x + 4
        @test expand_measure(expr, data3, map_args...) == expected
        # test array parameter with both quadratic terms become numbers
        expr = pars1[1] * pars1[2] + 2
        expected = zero(JuMP.GenericQuadExpr{Float64, GeneralVariableRef})
        expected.aff = 1 + 2 * 2 + 4
        @test expand_measure(expr, data3, map_args...) == expected
        # test array parameter with 2nd quadratic term becomes number
        expr = inf4 * pars1[1] + 2
        idx = m.next_var_index + 1
        pts = [PointVariableRef(m, idx), PointVariableRef(m, idx + 1)]
        expected = zero(JuMP.GenericQuadExpr{Float64, GeneralVariableRef})
        expected.aff = pts[1] + 2 * pts[2] + 4
        @test expand_measure(expr, data3, map_args...) == expected
        # test array parameter with 2nd quadratic term becomes number, 1st constant
        expr = x * pars1[1] + 2
        expected = zero(JuMP.GenericQuadExpr{Float64, GeneralVariableRef})
        expected.aff = x + 2 * x + 4
        @test expand_measure(expr, data3, map_args...) == expected
    end
    # test expand_measure (measure)
    @testset "Measure" begin
        # test single param measure
        idx = m.next_var_index + 1
        pts = [PointVariableRef(m, idx), PointVariableRef(m, idx + 1)]
        expected = 0.5 * (pts[1] + pts[2])
        @test expand_measure(meas1, data1, map_args...) == expected
        # test another single param
        idx = m.next_var_index + 1
        pts = [PointVariableRef(m, idx), PointVariableRef(m, idx + 1)]
        expected = 0.5 * (6 * pts[1] * x + 6 * pts[2] * x - 3 * (1 + 2))
        @test expand_measure(meas2, data1, map_args...) == expected
        # test array param measure
        idx = m.next_var_index + 1
        pts = [PointVariableRef(m, idx), PointVariableRef(m, idx + 1)]
        expected = pts[1] + pts[2]
        @test expand_measure(meas1, data3, map_args...) == expected
    end
    # test expand_measure (other)
    @testset "Other" begin
        # prepare test
        @variable(trans_model, y)
        # test it
        @test_throws ErrorException expand_measure(y, BadData(), map_args...)
    end
end

# Test _expand_measures
@testset "_expand_measures" begin
    # initialize model and references
    m = InfiniteModel()
    @infinite_parameter(m, 1 <= par1 <= 2)
    @infinite_parameter(m, 1 <= par2 <= 2)
    @infinite_parameter(m, 1 <= pars1[1:2] <= 2)
    @infinite_variable(m, inf1(par1) >= 1)
    @infinite_variable(m, inf2(par1, par2))
    @infinite_variable(m, inf3(par2))
    @infinite_variable(m, inf4(pars1))
    @hold_variable(m, x)
    # prepare measures
    data1 = DiscreteMeasureData(par1, [0.5, 0.5], [1, 2])
    data2 = DiscreteMeasureData(pars1, [1, 1], [[1, 1], [2, 2]])
    meas1 = measure(inf1 + 3x - inf2 + inf3 - 2, data1)
    meas2 = measure(2 * inf4 * x - pars1[2] + inf2, data2)
    # prepare point_mapper function and info
    function mapper(m2, pvref, ivref, support)
        set_name(pvref, string(InfiniteOpt._root_name(ivref), support))
        return
    end
    trans_model = Model()
    map_args = (trans_model, mapper)
    # test _expand_measures (measure)
    @testset "MeasureRef" begin
        # test the first measure
        idx = m.next_var_index + 1
        pts = [PointVariableRef(m, idx), PointVariableRef(m, idx + 1)]
        rv1 = ReducedInfiniteVariableRef(m, idx + 2)
        m.reduced_info[idx + 2] = ReducedInfiniteInfo(inf2, Dict(1 => 1))
        rv2 = ReducedInfiniteVariableRef(m, idx + 3)
        m.reduced_info[idx + 3] = ReducedInfiniteInfo(inf2, Dict(1 => 2))
        expected = 0.5 * (pts[1] + pts[2] + 6x - rv1 - rv2 + 2inf3 - 4)
        @test InfiniteOpt._expand_measures(meas1, map_args...) == expected
        # test the second measure
        idx = m.next_var_index + 1
        pts = [PointVariableRef(m, idx), PointVariableRef(m, idx + 1)]
        expected = 2 * pts[1] * x + 2 * pts[2] * x - 3 + 2inf2
        @test InfiniteOpt._expand_measures(meas2, map_args...) == expected
        # test error
        m2 = Model()
        @variable(m2, y)
        m.measures[1] = Measure(y, BadData())
        @test_throws ErrorException InfiniteOpt._expand_measures(meas1,
                                                                 map_args...)
        # Undo changes
        m.measures[1] = Measure(inf1 + 3x - inf2 + inf3 - 2, data1)
    end
    # test _expand_measures (GenericAffExpr)
    @testset "AffExpr" begin
        # test returning AffExpr
        expr = 2inf4 + x + 2meas1
        idx = m.next_var_index + 1
        pts = [PointVariableRef(m, idx), PointVariableRef(m, idx + 1)]
        rv1 = ReducedInfiniteVariableRef(m, idx + 2)
        m.reduced_info[idx + 2] = ReducedInfiniteInfo(inf2, Dict(1 => 1))
        rv2 = ReducedInfiniteVariableRef(m, idx + 3)
        m.reduced_info[idx + 3] = ReducedInfiniteInfo(inf2, Dict(1 => 2))
        expected = 2inf4 + 7x + pts[1] + pts[2] - rv1 - rv2 + 2inf3 - 4
        @test InfiniteOpt._expand_measures(expr, map_args...) == expected
        # test returning QuadExpr
        expr = x + 3par1 + meas2
        idx = m.next_var_index + 1
        pts = [PointVariableRef(m, idx), PointVariableRef(m, idx + 1)]
        expected = x + 3par1 + 2 * pts[1] * x + 2 * pts[2] * x - 3 + 2inf2
        @test InfiniteOpt._expand_measures(expr, map_args...) == expected
    end
    # test _expand_measures (GenericQuadExpr)
    @testset "QuadExpr" begin
        # prepare first measure expansion
        expr = 2 * meas1 * x + inf1 * meas1 - x + meas1 - 3
        idx = m.next_var_index + 1
        pts = [PointVariableRef(m, idx), PointVariableRef(m, idx + 1)]
        rv1 = ReducedInfiniteVariableRef(m, idx + 2)
        m.reduced_info[idx + 2] = ReducedInfiniteInfo(inf2, Dict(1 => 1))
        rv2 = ReducedInfiniteVariableRef(m, idx + 3)
        m.reduced_info[idx + 3] = ReducedInfiniteInfo(inf2, Dict(1 => 2))
        m1 = 0.5 * (pts[1] + pts[2] + 6x - rv1 - rv2 + 2inf3 - 4)
        # prepare second measure expansion
        pts = [PointVariableRef(m, idx + 4), PointVariableRef(m, idx + 5)]
        rv1 = ReducedInfiniteVariableRef(m, idx + 6)
        m.reduced_info[idx + 6] = ReducedInfiniteInfo(inf2, Dict(1 => 1))
        rv2 = ReducedInfiniteVariableRef(m, idx + 7)
        m.reduced_info[idx + 7] = ReducedInfiniteInfo(inf2, Dict(1 => 2))
        m2 = 0.5 * (pts[1] + pts[2] + 6x - rv1 - rv2 + 2inf3 - 4)
        # prepare third measure expansion
        pts = [PointVariableRef(m, idx + 8), PointVariableRef(m, idx + 9)]
        rv1 = ReducedInfiniteVariableRef(m, idx + 10)
        m.reduced_info[idx + 10] = ReducedInfiniteInfo(inf2, Dict(1 => 1))
        rv2 = ReducedInfiniteVariableRef(m, idx + 11)
        m.reduced_info[idx + 11] = ReducedInfiniteInfo(inf2, Dict(1 => 2))
        m3 = 0.5 * (pts[1] + pts[2] + 6x - rv1 - rv2 + 2inf3 - 4)
        # test with this thorough expression
        quad = zero(GenericQuadExpr{Float64, GeneralVariableRef})
        quad.aff = - x + m1 - 3
        expected = quad + 2 * m2 * x + inf1 * m3
        @test InfiniteOpt._expand_measures(expr, map_args...) == expected
    end
end

# Test user expansion methods
@testset "User Methods" begin
    # initialize model and references
    m = InfiniteModel()
    @infinite_parameter(m, 1 <= par1 <= 2)
    @infinite_parameter(m, 1 <= par2 <= 2)
    @infinite_parameter(m, 1 <= pars1[1:2] <= 2)
    @infinite_variable(m, inf1(par1) >= 1)
    @infinite_variable(m, inf2(par1, par2) >= 0)
    @infinite_variable(m, inf3(par2))
    @infinite_variable(m, inf4(pars1))
    @hold_variable(m, x >= 0)
    # prepare measures
    data1 = DiscreteMeasureData(par1, [0.5, 0.5], [1, 2])
    data2 = DiscreteMeasureData(pars1, [1, 1], [[1, 1], [2, 2]])
    data3 = DiscreteMeasureData(par2, [2, 2], [1, 2])
    meas1 = measure(inf1 + 3x - inf2 + inf3 - 2, data1)
    meas2 = measure(2 * inf4 * x - pars1[2] + inf2, data2)
    # test _add_mapped_point_variable
    @testset "_add_mapped_point_variable" begin
        # prepare the partially made point variable reference
        pref = InfiniteOpt.make_point_variable_ref(inf1)
        # try mapping/adding it
        @test isa(InfiniteOpt._add_mapped_point_variable(m, pref, inf1, (1,)),
                  Nothing)
        @test name(pref) == "inf1(1)"
        @test haskey(m.vars, JuMP.index(pref))
        @test lower_bound(pref) == 1
    end
    # test expand
    @testset "expand" begin
        # test the first measure
        idx = m.next_var_index + 1
        pts = [PointVariableRef(m, idx), PointVariableRef(m, idx + 1)]
        rv1 = ReducedInfiniteVariableRef(m, idx + 2)
        m.reduced_info[idx + 2] = ReducedInfiniteInfo(inf2, Dict(1 => 1))
        rv2 = ReducedInfiniteVariableRef(m, idx + 3)
        m.reduced_info[idx + 3] = ReducedInfiniteInfo(inf2, Dict(1 => 2))
        expected = 0.5 * (pts[1] + pts[2] + 6x - rv1 - rv2 + 2inf3 - 4)
        @test expand(meas1) == expected
        # test the second measure
        idx = m.next_var_index + 1
        pts = [PointVariableRef(m, idx), PointVariableRef(m, idx + 1)]
        expected = 2 * pts[1] * x + 2 * pts[2] * x - 3 + 2inf2
        @test expand(meas2) == expected
        # test error
        m2 = Model()
        @variable(m2, y)
        m.measures[1] = Measure(y, BadData())
        @test_throws ErrorException expand(meas1)
        # Undo changes
        m.measures[1] = Measure(inf1 + 3x - inf2 + inf3 - 2, data1)
    end
    # test expand_all_measures!
    @testset "expand_all_measures!" begin
        # test error
        @objective(m, Min, x + meas1)
        @test_throws ErrorException expand_all_measures!(m)
        # prepare the model for testing
        set_objective_function(m, x + measure(inf1 * par1 + 3, data1))
        @constraint(m, c1, inf1 + x >= 42.)
        @constraint(m, c2, 2x - measure(measure(inf1 * x + par1 + inf2, data1),
                                        data3) == 0.)
        @BDconstraint(m, c3(par2 in [1, 1.5]), measure(inf2, data1) + inf1 >= 0.)
        @constraint(m, c4, measure(inf4, data2) <= 3.)
        # prepare comparison expressions
        idx = m.next_var_index + 1
        pts = [PointVariableRef(m, idx), PointVariableRef(m, idx + 1)]
        obj = x + 0.5pts[1] + pts[2] + 3
        pts = [PointVariableRef(m, idx + 4), PointVariableRef(m, idx + 5),
               PointVariableRef(m, idx + 6), PointVariableRef(m, idx + 7),
               PointVariableRef(m, idx + 8), PointVariableRef(m, idx + 9)]
        c2_expected = 2x - (2 * pts[1] * x + 2 * pts[2] * x + pts[3] + pts[4] +
                      pts[5] + pts[6])
        rv1 = ReducedInfiniteVariableRef(m, idx + 10)
        m.reduced_info[idx + 10] = ReducedInfiniteInfo(inf2, Dict(1 => 1))
        rv2 = ReducedInfiniteVariableRef(m, idx + 11)
        m.reduced_info[idx + 11] = ReducedInfiniteInfo(inf2, Dict(1 => 2))
        c3_expected = 0.5rv1 + 0.5rv2 + inf1
        pts = [PointVariableRef(m, idx + 12), PointVariableRef(m, idx + 13)]
        c4_expected = pts[1] + pts[2]
        # test the expansion
        @test isa(expand_all_measures!(m), Nothing)
        @test objective_function(m) == obj
        @test is_valid(m, c2)
        @test is_valid(m, c3)
        @test is_valid(m, c4)
        @test name(c1) == "c1"
        @test name(c2) == "c2"
        @test name(c3) == "c3"
        @test name(c4) == "c4"
        @test m.constrs[JuMP.index(c1)].func == inf1 + x
        @test m.constrs[JuMP.index(c2)].func == c2_expected
        @test m.constrs[JuMP.index(c3)].func == c3_expected
        @test m.constrs[JuMP.index(c4)].func == c4_expected
        @test m.constrs[JuMP.index(c1)].set == MOI.GreaterThan(42.)
        @test m.constrs[JuMP.index(c2)].set == MOI.EqualTo(6.)
        @test m.constrs[JuMP.index(c3)].set == MOI.GreaterThan(0.)
        @test m.constrs[JuMP.index(c4)].set == MOI.LessThan(3.)
        @test isa(m.constrs[JuMP.index(c1)], ScalarConstraint)
        @test isa(m.constrs[JuMP.index(c2)], ScalarConstraint)
        @test isa(m.constrs[JuMP.index(c3)], BoundedScalarConstraint)
        @test isa(m.constrs[JuMP.index(c4)], ScalarConstraint)
        @test m.constrs[JuMP.index(c3)].bounds.intervals == Dict(par2 => IntervalSet(1, 1.5))
    end
end


text:
using SpecialFunctions: gamma
export c2, c3

function c2(psi)
    eps = 1.0
    if psi > eps
        res = (1 - cos(sqrt(psi))) / psi
    elseif psi < -eps
        res = (cosh(sqrt(-psi)) - 1) / (-psi)
    else
        res = 1.0 / 2.0
        delta = (-psi) / gamma(2 + 2 + 1)
        k = 1
        while res + delta != res
            res += delta
            k += 1
            delta = (-psi)^k / gamma(2*k + 2 + 1)
        end
    end
    return res
end

function c3(psi)
    eps = 1.0
    if psi > eps
        res = (sqrt(psi) - sin(sqrt(psi))) / (psi * sqrt(psi))
    elseif psi < -eps
        res = (sinh(sqrt(-psi)) - sqrt(-psi)) / (-psi * sqrt(-psi))
    else
        res = 1.0 / 6.0
        delta = (-psi) / gamma(2 + 3 + 1)
        k = 1
        while res + delta != res
            res += delta
            k += 1
            delta = (-psi)^k / gamma(2*k + 3 + 1)
        end
    end
    return res
end


text:
function calculate_pi(NBIN::Int= 100000)::Float64
    step = 1/NBIN
    sum = 0
    for i=1:NBIN
        x = (i+0.5)*step;
        sum += 4 / (1 + x^2)
    end
    pi = sum*step
    return pi
end

@time @show calculate_pi()

max_stars_repo_name: naezzell/accelqat

max_issues_repo_name: naezzell/accelqat

max_forks_repo_name: naezzell/accelqat