text:
export noteplotter

mpl = PyPlot.matplotlib

const norm_mpl = mpl[:colors][:Normalize](0.0, 1.0)
const c_m = mpl[:cm][:viridis]
const s_m = mpl[:cm][:ScalarMappable](cmap=c_m, norm_mpl=norm_mpl)
s_m[:set_array]([])

# the identity function as dictionary
identimap = Dict{UInt8, UInt8}()
for i = 0:255
    push!(identimap, i => i)
end

# the standard note names for all 255 pitches in a dictionary
standardnames = Dict{UInt8, String}()
for i = 0:255
    push!(standardnames, i => MIDI.pitchname(i))
end

"""
    noteplotter(notes::Notes [, grid]; kwargs...)

Visualize the `notes` by plotting them as lines in a pitch over ticks diagram.
The duration of a `Note` is represented by the length of the corresponding line.

Optionally provide a `grid` (see [`quantize`](@ref))
to additionally plot vertical lines at the positions
specified by it. The beginning of a quarter note is marked by a darker line.

Keyword arguments:
    - `ticknames::Dict{UInt8, String} = standardnames`:
    Change the ticks labels of the pitches to be arbitrary strings. Provide a
    dictionary with the pitches you'd like to assign a label as keys and the
    labels as values. If you leave out pitches, they will be named like standard
    MIDI notes.
    - `reorder::Dict{UInt8, UInt8} = identimap`:
    For better readability you can change the pitches to be mapped to different
    values on the y-axis. Provide a dictionary with the pitches you'd like to
    reorder as keys and the new values they shall be mapped to as values.
    BE CAREFUL, if you map multiple pitches to one value, strange things can happen.
    The reordering does not affect the relabeling via `ticknames` (reordered
    notes get the labels of the original pitch) and if you leave out pitches
    they will just stay where they were.
"""
function noteplotter(notes::MIDI.Notes, grid = nothing; ticknames::Dict{UInt8, String} = standardnames, reorder::Dict{UInt8, UInt8} = identimap)

# positions of the notes separated by velocities (for plotting in different colors)
x1 = Dict{UInt, Vector{UInt}}() #starts of notes
x2 = Dict{UInt, Vector{UInt}}() #ends of notes
y = Dict{UInt, Vector{UInt}}() #pitches

# collect note information
for note in notes
    if !haskey(x1, note.velocity)
        push!(x1, note.velocity => [note.position])
        push!(x2, note.velocity => [note.position + note.duration])
        push!(y, note.velocity => [haskey(reorder, note.pitch) ? reorder[note.pitch] : note.pitch])
    else
        push!(x1[note.velocity], note.position)
        push!(x2[note.velocity], note.position + note.duration)
        push!(y[note.velocity], haskey(reorder, note.pitch) ? reorder[note.pitch] : note.pitch)
    end
end

# reorder pitches with new names
spits = sort(collect(keys(ticknames)))
rpits = copy(spits)
for i in 1:length(rpits)
    if haskey(reorder, rpits[i])
        rpits[i] = reorder[rpits[i]]
    end
end

# find minimum and maximum of >played and reordered< pitches to adjust plot limits
pits = sort(unique(pitches(notes)))
for i in 1:length(pits)
    if haskey(reorder, pits[i])
        pits[i] = reorder[pits[i]]
    end
end
minpit = minimum(pits)
maxpit = maximum(pits)

# find the index of the minimum plotted pitch in the list of all reordered pitches
# permutation array instead of real sorting to be able to reconstruct names
# of reordered pitches
soperm = sortperm(rpits)
minrpit = 1
while minpit > rpits[soperm[minrpit]]
    minrpit += 1
end

# assemble the ticks labels for the graph first look for provided tickname, if
# not given unse standard name.
tickrange = minpit:maxpit
tickstrings = String[]
ind = minrpit
for i = 1:length(tickrange)
    # if pitch was reordered
    if rpits[soperm[ind]] == tickrange[i]
        # revert the reordering with the permutation to get real tick name
        if haskey(ticknames, spits[soperm[ind]])
            push!(tickstrings, ticknames[spits[soperm[ind]]])
        else
            push!(tickstrings, standardnames[spits[soperm[ind]]])
        end
        ind += 1
    # for non reordered pitches
    else
        if haskey(ticknames, i)
            push!(tickstrings, ticknames[i])
        else
            push!(tickstrings, standardnames[i])
        end
    end
end

# get min and max vel to get ready for plotting
vels = keys(y)
maxvel = maximum(vels)

#initialize figure
f = PyPlot.figure()
ax = PyPlot.gca()

#plotting action
for vel in vels
    PyPlot.plot([x1[vel], x2[vel]], [y[vel], y[vel]],
    color = s_m[:to_rgba](vel/maxvel), lw = 3)
end

# function rectangle(x, y, dx, vel)
#     mpl[:patches][:Rectangle]((x, y), 30, 1,
#     facecolor = s_m[:to_rgba](vel/maxvel), edgecolor = "black", lw = 0.5)
# end
#
#
# patches = [rectangle(x1[vel], y[vel], x2[vel], vel) for vel in vels]
# ax[:add_collection](mpl[:collections][:PatchCollection](patches))

# handle grid
if grid != nothing
    gpoints = notes.tpq*grid[2:end-1]
    g = UInt[]
    beat = UInt[]
    for i = 0:ceil(Int,(notes.notes[end].position+notes.notes[end].duration)/notes.tpq)
        append!(g,i*notes.tpq+gpoints)
        push!(beat, i*notes.tpq)
    end
    temp = ones(length(g))
    temp2 = ones(length(beat))
    PyPlot.plot([g,g] ,[temp .* (minpit - 0.5), temp .* (maxpit + 0.5)], lw=0.1, color="grey")
    PyPlot.plot([beat,beat] ,[temp .* (minpit - 0.5), temp .* (maxpit + 0.5)], lw=0.1, color="black")
end

# labels, ranges, ...
ax[:set_ylim](minpit-0.5, maxpit+0.5)
ax[:set_yticks](tickrange)
ax[:set_yticklabels](tickstrings)
ax[:set_xlabel]("Time in ticks")
ax[:set_ylabel]("Pitch")
f[:tight_layout]()

end


text:
function gen_initialstate(cutoffN,x...)

  ketCoeff = zeros(Complex128,length(x),1)
  parameterArray = zeros(Int64,length(x),2)

  for i = 1:length(x)
    str = x[i]
    # Grabbing the coefficient in front of ket
    coeff = str[1:search(str,'|')-1]
    if isempty(coeff)
      ketCoeff[i] = 1
    elseif coeff == '-'
      ketCoeff[i] = -1
    else
      ketCoeff[i] = eval(parse(coeff))
    end

    # Grabbig qubit state
    if str[search(str,'|')+1] == 'g'
      parameterArray[i,1] = 0
    elseif str[search(str,'|')+1] == 'e'
      parameterArray[i,1] = 1
    else
      error("Error, qubit state should be 'g' or 'e'")
    end

    # Grabbing resonator state
    parameterArray[i,2] = eval(parse(str[search(str,',')+1:end-1]))
  end

  # Making sure that cutoffN is large enough
  if cutoffN <=  maximum(parameterArray[:,2])
    error("The size of the resonator space (cutoffN) must be at least the size of the larger number state plus one (due to zero state)!")
  end

  # Assembling the state matrix

  state = zeros(Complex128,2*cutoffN,2*cutoffN)

  for i = 1:length(x)
    for j = 1:length(x)
      # Defining the ket
      qubitPartKet = zeros(Complex128,2,1) # Initialising the qubit ket (2D)
      qubitPartKet[parameterArray[i,1]+1] = 1 # Setting the vector as (1 0) for g or (0 1) for e
      resonatorPartKet = zeros(Complex128,cutoffN,1) # Initialising the resonator ket (cutoffND)
      resonatorPartKet[parameterArray[i,2]+1] = 1 # Setting the vector
      ket = kron(resonatorPartKet, qubitPartKet) # Tensor product
      # Defining the bra
      qubitPartBra =zeros(Complex128,2,1)
      qubitPartBra[parameterArray[j,1]+1] = 1
      resonatorPartBra = zeros(Complex128,cutoffN,1)
      resonatorPartBra[parameterArray[j,2]+1] = 1
      bra = kron(resonatorPartBra, qubitPartBra)'
      # Multiplying the initial ket with its complex conjugate to make
      # the density matrix
      state = state + ketCoeff[i]*ketCoeff[j]'*ket*bra
    end
  end

  norm = 1/sum(abs(ketCoeff[:]).^2)
  state = norm*state # Norming
  info(size(state))
  return state
end

ext: jl


text:
"""
The barrel vault (Scordelis-Lo) roof is one of the benchmarks for linear elastic
analysis of shells.

The candidate element's usefulness in irregular geometries (and most practical
cases involve a high degree of geometric irregularity) is tested. As would be
expected,the irregular mesh results are not as good as those provided by a
regular mesh with the same number of variables.

Problem description

The physical basis of the problem is a deeply arched roof supported only
by diaphragms at its curved edges (an aircraft hanger), deforming under its own
weight. It is interesting to observe that the geometry is such that the
centerpoint of the roof moves upward under the self-weight(downwardly directed)
load. Perhaps this is one reason why the problem is not straightforward
numerically.
"""
module vis_scordelis_lo_examples

using LinearAlgebra
using FinEtools
using FinEtoolsDeforLinear
using FinEtoolsFlexStructures.FESetShellT3Module: FESetShellT3
using FinEtoolsFlexStructures.FESetShellQ4Module: FESetShellQ4
using FinEtoolsFlexStructures.FEMMShellT3FFModule
using FinEtoolsFlexStructures.RotUtilModule: initial_Rfield, update_rotation_field!
using Examples.VisUtilModule: plot_nodes, plot_midline, render, plot_space_box, plot_midsurface, space_aspectratio, save_to_json
using FinEtools.MeshExportModule.VTKWrite: vtkwrite

# analytical solution for the vertical deflection and the midpoint of the
# free edge
analyt_sol=-0.3024;
# Parameters:
E=4.32e8;
nu=0.0;
thickness = 0.25; # geometrical dimensions are in feet
R = 25.0;
L = 50.0;

cylindrical!(csmatout::FFltMat, XYZ::FFltMat, tangents::FFltMat, fe_label::FInt) = begin
    r = vec(XYZ); r[2] = 0.0; r[3] += R
    csmatout[:, 3] .= vec(r)/norm(vec(r))
    csmatout[:, 2] .= (0.0, 1.0, 0.0) #  this is along the axis
    cross3!(view(csmatout, :, 1), view(csmatout, :, 2), view(csmatout, :, 3))
    return csmatout
end

function _execute_dsg_model(formul, n = 8, visualize = true)

    tolerance = R/n/10
    # fens, fes = T3blockrand(40/360*2*pi,L/2,n,n);
    fens, fes = T3block(40/360*2*pi,L/2,n,n,:b);
    fens.xyz = xyz3(fens)
    for i in 1:count(fens)
        a=fens.xyz[i, 1]; y=fens.xyz[i, 2];
        fens.xyz[i, :] .= (R*sin(a), y, R*(cos(a)-1))
    end

    # conn = fill(0, count(fes), 3)
    # for i in 1:count(fes)
    #     conn[i, :] .= fes.conn[i][2], fes.conn[i][3], fes.conn[i][1]
    # end
    # fromarray!(fes, conn)

    mater = MatDeforElastIso(DeforModelRed3D, E, nu)
    ocsys = CSys(3, 3, cylindrical!)

    sfes = FESetShellT3()
    accepttodelegate(fes, sfes)
    femm = formul.make(IntegDomain(fes, TriRule(1), thickness),  mater)
    femm.drilling_stiffness_scale = 1.0e-4
    femm.mult_el_size = 5/12
    stiffness = formul.stiffness
    associategeometry! = formul.associategeometry!

    # Construct the requisite fields, geometry and displacement
    # Initialize configuration variables
    geom0 = NodalField(fens.xyz)
    u0 = NodalField(zeros(size(fens.xyz,1), 3))
    Rfield0 = initial_Rfield(fens)
    dchi = NodalField(zeros(size(fens.xyz,1), 6))

    # Apply EBC's
    # rigid diaphragm
    l1 = selectnode(fens; box = Float64[-Inf Inf 0 0 -Inf Inf], inflate = tolerance)
    for i in [1,3,5]
        setebc!(dchi, l1, true, i)
    end
    # plane of symmetry perpendicular to Y
    l1 = selectnode(fens; box = Float64[-Inf Inf L/2 L/2 -Inf Inf], inflate = tolerance)
    for i in [2,4,6]
        setebc!(dchi, l1, true, i)
    end
    # plane of symmetry perpendicular to X
    l1 = selectnode(fens; box = Float64[0 0 -Inf Inf -Inf Inf], inflate = tolerance)
    for i in [1,5,6]
        setebc!(dchi, l1, true, i)
    end
    applyebc!(dchi)
    numberdofs!(dchi);

    # Assemble the system matrix
    associategeometry!(femm, geom0)
    K = stiffness(femm, geom0, u0, Rfield0, dchi);

    # Midpoint of the free edge
    nl = selectnode(fens; box = Float64[sin(40/360*2*pi)*25 sin(40/360*2*pi)*25 L/2 L/2 -Inf Inf], inflate = tolerance)
    lfemm = FEMMBase(IntegDomain(fes, TriRule(3)))
    fi = ForceIntensity(FFlt[0, 0, -90, 0, 0, 0]);
    F = distribloads(lfemm, geom0, dchi, fi, 3);

    # Solve
    U = K\F
    scattersysvec!(dchi, U[:])
    result =   dchi.values[nl, 3][1]
    @info "Solution: $(result), $(round(result/analyt_sol*100, digits = 4))%"

    # Visualization
    if visualize

        # Generate a graphical display of resultants
        scalars = []
        for nc in 1:3
            fld = fieldfromintegpoints(femm, geom0, dchi, :moment, nc, outputcsys = ocsys)
            push!(scalars, ("m$nc", fld.values))
            fld = elemfieldfromintegpoints(femm, geom0, dchi, :moment, nc, outputcsys = ocsys)
            push!(scalars, ("em$nc", fld.values))
        end
        vtkwrite("vis_scordelis_lo_examples-$(n)-m.vtu", fens, fes; scalars = scalars, vectors = [("u", dchi.values[:, 1:3])])
        scalars = []
        for nc in 1:3
            fld = fieldfromintegpoints(femm, geom0, dchi, :membrane, nc, outputcsys = ocsys)
            push!(scalars, ("n$nc", fld.values))
            fld = elemfieldfromintegpoints(femm, geom0, dchi, :membrane, nc, outputcsys = ocsys)
            push!(scalars, ("en$nc", fld.values))
        end
        vtkwrite("vis_scordelis_lo_examples-$(n)-n.vtu", fens, fes; scalars = scalars, vectors = [("u", dchi.values[:, 1:3])])
        scalars = []
        for nc in 1:2
            fld = fieldfromintegpoints(femm, geom0, dchi, :shear, nc, outputcsys = ocsys)
            push!(scalars, ("q$nc", fld.values))
            fld = elemfieldfromintegpoints(femm, geom0, dchi, :shear, nc, outputcsys = ocsys)
            push!(scalars, ("eq$nc", fld.values))
        end
        vtkwrite("vis_scordelis_lo_examples-o-$(n)-q.vtu", fens, fes; scalars = scalars, vectors = [("u", dchi.values[:, 1:3])])

        # vtkwrite("vis_scordelis_lo_examples-$(n)-uur.vtu", fens, fes; scalars = scalars, vectors = [("u", dchi.values[:, 1:3]), ("ur", dchi.values[:, 4:6])])

        scattersysvec!(dchi, (L/8)/maximum(abs.(U)).*U)
        update_rotation_field!(Rfield0, dchi)
        plots = cat(plot_space_box([[0 0 -L/2]; [L/2 L/2 L/2]]),
            #plot_nodes(fens),
            plot_midsurface(fens, fes; x = geom0.values, facecolor = "rgb(12, 12, 123)"),
            plot_midsurface(fens, fes; x = geom0.values, u = dchi.values[:, 1:3], R = Rfield0.values);
            dims = 1)
        pl = render(plots)
    end

    result
end

function test_convergence(ns = [16, 32], visualize = true)
    formul = FEMMShellT3FFModule
    @info "Scordelis-Lo shell"
    results = []
    for n in ns
        v = _execute_dsg_model(formul, n, visualize)
        push!(results, v/(-0.3024)*100)
    end
    return ns, results
end

function allrun()
    println("#####################################################")
    println("# test_convergence ")
    test_convergence()
    return true
end # function allrun

@info "All examples may be executed with "
println("using .$(@__MODULE__); $(@__MODULE__).allrun()")

end # module
nothing


text:
using Compose
import Distances

include("structs.jl")

function build(contraption::Contraption)
    # Combine components

    # Point masses ##
    x = contraption.position[1, :]
    y = - contraption.position[2, :] # Flip coordinates for graphics

    # Scale point size proportional to cube root of mass
    r = [0.01] .* contraption.mass .^ (1/3)

    points = compose(context(),
                  circle(x, y, r),
                  fill("black"),
                  fillopacity(1.))

    # Simplify drawing if no springs are found
    if sum(contraption.springs) == 0.
        return points
    end

    # Springs ##
    n = length(contraption.mass)
    X_i, Y_i, X_j, Y_j = [], [], [], []
    W = []

    for i in 1:n, j in (i+1):n
        if contraption.springs[i, j] != 0
            x_i, y_i = contraption.position[:, i]
            x_j, y_j = contraption.position[:, j]

            # Flip y-coordinates to match vector graphics convention
            push!(X_i, x_i)
            push!(Y_i, -y_i)
            push!(X_j, x_j)
            push!(Y_j, -y_j)

            # Base width of spring
            w = 1.0mm

            # Scale base spring width proportional to square root of strength
            w *= sqrt(contraption.springs[i,j]) / 100

            # Scale final spring width to conserve total spring volume
            # Base length * base width ^ 2 = current length * current width ^2
            current_length = Distances.norm(contraption.position[:, i] - contraption.position[:, j])
            w *= sqrt(contraption.rest_length[i, j] / current_length)

            push!(W, w)
        end
    end

    spring_array = [[(x_i, y_i), (x_j, y_j)] for (x_i, y_i, x_j, y_j) in zip(X_i, Y_i, X_j, Y_j)]

    lines = compose(context(),
                    line(spring_array),
                    stroke("black"),
                    linewidth(W),
                    strokeopacity(0.5))

    composition = compose(context(), points, lines)

    return composition

end

function build(bounds::Bounds)

    # Create non-square box if not square bounds
    if bounds.width == bounds.height
        scale_x, scale_y = 1., 1.
    elseif bounds.width < bounds.height
        scale_x, scale_y = bounds.width/bounds.height, 1.
    else
        scale_x, scale_y = 1. , bounds.height / bounds.width
    end

    obj = compose(context(), # Defaults to (0, 0, 1, 1)
                rectangle(0, 0, scale_x, scale_y),
                fill("white"), stroke("black"), linewidth(0.4mm))

    return obj

end

function build(contraption::Contraption, bounds::Bounds)
    bounds_sprite = build(bounds)

    contraption_sprite = build(contraption)

    # Rescale contraption based on bound size
    # Center such that (0, 0) is in the middle of the canvas
    scale_x = 1/bounds.width
    scale_y = 1/bounds.height

    scene = compose(bounds_sprite,
                    (context(0.5, 0.5, scale_x, scale_y), contraption_sprite))

    return scene
end


text:
module M2_2
function f(q::Vector{T}, q_dot::Vector{T}, xi::T) where T
    l1 = q[1]
    l2 = q[2]
    l3 = q[3]
    l1_dot = q_dot[1]
    l2_dot = q_dot[2]
    l3_dot = q_dot[3]

    1.38960161444059e-11*l1.^2 + 1.13826433538242e-11*l1.*l2 - 3.92837735785805e-11*l1.*l3 + 1.6885767684885e-14*l1 + 2.33986794159514e-12*l2.^2 - 1.61071218647439e-11*l2.*l3 + 6.89771216933803e-15*l2 + 2.77725822370262e-11*l3.^2 - 2.38497041715743e-14*l3 + 5.13890013369217e-18
end
end


text:
mutable struct HyperplaneType{T}
    v::Vector{T}
    c::T

    HyperplaneType{T}(v,c) where T = new{T}(v,c)
    HyperplaneType{T}() where T = new{T}()
end

mutable struct PartitionDataType{T}
    hp::HyperplaneType{T}
    X::Vector{Vector{T}}
    global_X_indices::Vector{Int}
    index::Int
end

mutable struct BinaryNode{T}
    data::T
    parent::BinaryNode{T}
    left::BinaryNode{T}
    right::BinaryNode{T}

    # Root constructor
    BinaryNode{T}(data) where T = new{T}(data)
    # Child node constructor
    BinaryNode{T}(data, parent::BinaryNode{T}) where T = new{T}(data, parent)
end
BinaryNode(data) = BinaryNode{typeof(data)}(data)

"""
Mutates parent. Taken from AbstractTrees.jl's example code.
"""
function leftchild!(parent::BinaryNode, data)
    !isdefined(parent, :left) || error("left child is already assigned")
    node = typeof(parent)(data, parent)
    parent.left = node
end

"""
Mutates parent. Taken from AbstractTrees.jl's example code.
"""
function rightchild!(parent::BinaryNode, data)
    !isdefined(parent, :right) || error("right child is already assigned")
    node = typeof(parent)(data, parent)
    parent.right = node
end

"""
Taken from AbstractTrees.jl's example code.
"""
function AbstractTrees.children(node::BinaryNode)
    if isdefined(node, :left)
        if isdefined(node, :right)
            return (node.left, node.right)
        end
        return (node.left,)
    end
    isdefined(node, :right) && return (node.right,)
    return ()
end

function splitpoints(u::Vector{T}, X::Vector{Vector{T}}) where T <: Real

    N = length(X)
    indicators = falses(N)

    functional_evals = collect( dot(u, X[n]) for n = 1:N )
    c = Statistics.median(functional_evals)

    for n = 1:N

        if functional_evals[n] < c

            indicators[n] = true
        else
            indicators[n] = false
        end
    end

    return indicators, functional_evals, c
end

function gethyperplane(X::Vector{Vector{T}}) where T

    # center.
    μ = Statistics.mean(X)
    Z = collect( X[n] - μ for n = 1:size(X,2) )

    Z_mat = (array2matrix(Z))'
    U, s, V = svd(Z_mat)
    v = V[:,1]

    indicators, functional_evals, c = splitpoints(v, X)
    hp = HyperplaneType{T}(v, c)

    return hp, indicators
end

"""
current node is p.
"""
function setuppartition(X::Vector{Vector{T}}, level) where T

    # get hyperplane.
    hp, left_indicators = gethyperplane(X)

    X_empty = Vector{Vector{T}}(undef, 0)
    #data = PartitionDataType(hp, X_empty, 0)

    X_empty_inds = Vector{Int}(undef, 0)
    data = PartitionDataType(hp, X_empty, X_empty_inds, 0)

    # add to current node.
    root = BinaryNode(data)

    # might have to use recursion.
    X_inds = collect(1:length(X))
    createchildren(root, left_indicators, "left", X, X_inds, level-1)
    createchildren(root, left_indicators, "right", X, X_inds, level-1)

    # assign consecutive whole numbers as indices for the leaf nodes.
    X_parts, X_parts_inds = labelleafnodes(root, X)

    return root, X_parts, X_parts_inds
end

function labelleafnodes(root::BinaryNode{PartitionDataType{T}},
    X::Vector{Vector{T}}) where T
    #

    X_parts = Vector{Vector{Vector{Float64}}}(undef, 0)
    X_parts_inds = Vector{Vector{Int}}(undef, 0)
    leaves = Vector{BinaryNode{PartitionDataType{T}}}(undef, 0)

    i = 0
    for node in AbstractTrees.Leaves(root)

        # label.
        i += 1
        node.data.index = i

        # get X_parts.
        push!(X_parts, X[node.data.global_X_indices])
        push!(X_parts_inds, node.data.global_X_indices)

        # sanity check.
        if !isempty(node.data.X)
            @assert norm(node.data.X-X_parts[end]) < 1e-10
        end

        #push!(node, leaves) # cache for speed.
    end

    return X_parts, X_parts_inds
end

# might need to include other data, like kernel matrix, etc. at the leaf nodes.
"""
X_p is X associated with parent.
If the input level value is 1, then kid is a leaf node.
"""
function createchildren(parent,
    left_indicators, direction, X_p::Vector{Vector{T}}, X_p_inds::Vector{Int},
    level::Int;
    store_X_flag::Bool = false) where T

    ## prepare children data.
    X_kid = Vector{Vector{T}}(undef, 0)
    X_kid_inds = Vector{Int}(undef, 0)

    if direction == "left"

        X_kid = X_p[left_indicators]
        X_kid_inds = X_p_inds[left_indicators]
        data = PartitionDataType(HyperplaneType{T}(), X_kid, X_kid_inds, 0)

        kid = leftchild!(parent, data)

    else
        right_indicators = .! left_indicators
        X_kid = X_p[right_indicators]
        X_kid_inds = X_p_inds[right_indicators]
        data = PartitionDataType(HyperplaneType{T}(), X_kid, X_kid_inds, 0)

        kid = rightchild!(parent, data)
    end

    if level == 1
        ## kid is a leaf node. Stop propagation.

        if !store_X_flag
            # do not store inputs leaf nodes.
            kid.data.X = Vector{Vector{T}}(undef, 0)
        end

        return nothing
    end

    ## kid is not a leaf node. Propagate.

    # do not store inputs and input info at non-leaf nodes. It is not used during query: i.e., findpartition().
    kid.data.X = Vector{Vector{T}}(undef, 0)
    kid.data.global_X_indices = Vector{Int}(undef, 0)

    # get hyperplane.
    hp_kid, left_indicators_kid = gethyperplane(X_kid)
    kid.data.hp = hp_kid

    createchildren(kid, left_indicators_kid, "left", X_kid, X_kid_inds, level-1)
    createchildren(kid, left_indicators_kid, "right", X_kid, X_kid_inds, level-1)

    return nothing
end

# #### get all leaves. no longer needed because there is a AbstractTrees.leaves() iterator.
# function buildXpart!(X_parts::Vector{Vector{Vector{T}}}, p::BinaryNode{PartitionDataType{T}}) where T

#     #
#     if !isdefined(p, :left) && !isdefined(p, :right)
#         # p is a leaf node. Add its X to X_parts.
#         push!(X_parts, p.data.X)

#         return nothing
#     end

#     # call itself to traverse again.
#     if isdefined(p, :left)
#         buildXpart!(X_parts, p.left)
#     end

#     if isdefined(p, :right)
#         buildXpart!(X_parts, p.right)
#     end

#     return nothing
# end

###### search-related.

# Time critical.
"""
Given a point and the tree, find the leaf node of the tree that corresponds to the region that contains this point.
"""
function findpartition(x::Vector{T},
    root::BinaryNode{PartitionDataType{T}},
    levels::Int) where T

    node = root
    for l = 1:levels-1
        if dot(node.data.hp.v, x) < node.data.hp.c
            node = node.left
        else
            node = node.right
        end
    end

    return node.data.index
end

"""
Given a point the tree setup, find the region it is in, and all regions it is within ε-radius to their boundaries.
Recursive function.
"""
function findεpartitions!(region_list::Vector{Int},
    x::Vector{T},
    node::BinaryNode{PartitionDataType{T}},
    levels::Int,
    ε::T) where T

    ### termination case.
    if !isdefined(node, :left) && !isdefined(node, :right)
        # node is the leaf node of the current region.
        push!(region_list, node.data.index)

        return nothing
    end

    ### call itself to traverse again.

    hyperplane_eval = dot(node.data.hp.v, x)

    if hyperplane_eval < node.data.hp.c + ε

        findεpartitions!(region_list, x, node.left, levels, ε)
    end

    if hyperplane_eval > node.data.hp.c - ε

        findεpartitions!(region_list, x, node.right, levels, ε)
    end

    return nothing
end

function organizetrainingsets(root::BinaryNode{PartitionDataType{T}},
    levels::Int,
    X0::Vector{Vector{T}},
    ε::T) where T

    N = length(X0)

    # debug.
    problematic_inds = Vector{Vector{Int}}(undef, 0) # for debug.
    regions_list_set = Vector{Vector{Int}}(undef, N)

    # set up training set.
    N_regions = sum( 1 for node in AbstractTrees.Leaves(root) )
    X_set = Vector{Vector{Vector{T}}}(undef, N_regions)
    X_set_inds = Vector{Vector{Int}}(undef, N_regions)

    for r = 1:N_regions
        X_set[r] = Vector{Vector{T}}(undef, 0)
        X_set_inds[r] = Vector{Int}(undef, 0)
    end

    # build training set by inspecting every input x ∈ X0.
    for n = 1:N

        x = X0[n]

        # sanity-check: list of regions must have the region x is in.
        #x_region_ind = findpartition(x, root, levels)

        regions_list = Vector{Int}(undef, 0)
        findεpartitions!(regions_list, x, root, levels, ε)

        # add to training set.
        for m = 1:length(regions_list)
            r::Int = regions_list[m]

            if 0 < r <= length(X_set)
                push!(X_set[r], x)
                push!(X_set_inds[r], n)
            else
                # this shouldn't happen.
                # store for debug.
                push!(problematic_inds, n)
            end
        end

        # if length(regions_list) > 1
        #     @assert 1==234
        # end

        # store for debug.
        regions_list_set[n] = regions_list
    end

    return X_set, X_set_inds, regions_list_set, problematic_inds
end

max_stars_repo_name: RoyCCWang/PatchMixtureKriging

max_issues_repo_name: RoyCCWang/PatchMixtureKriging


text:
include("params.jl")
include("context.jl")
include("info.jl")
include("compress.jl")
include("decompress.jl")


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

using InteractiveUtils
using Printf

# get a temporary directory
dir = ARGS[1]
rm(dir, force=true, recursive=true)
mkdir(dir)

# toggle between unoptimized (CHECK/LOWER) and optimized IR (FINAL).
optimize=false
if length(ARGS) >= 2
    optimize = ARGS[2]=="-O"
end

# Emit LLVM IR to dir
counter = 0
function emit(f, tt...)
    global counter
    name = nameof(f)
    open(joinpath(dir, @sprintf("%05d-%s.ll", counter, name)), "w") do io
        code_llvm(io, f, tt, raw=true, optimize=optimize, dump_module=true, debuginfo=:none)
    end
    counter+=1
end

max_stars_repo_name: TimoLarson/julia

max_issues_repo_name: AsForMe/julia

max_forks_repo_name: AsForMe/julia


text:
using Plots

# Choose backend
gr()
# plotlyjs()

data = readdlm("data.txt")
rows = 1:size(data, 1)
x = data[rows,1]
raw = data[rows,2]
ewma = data[rows,3]

plot(x, raw, background_color=RGB(0.2, 0.2, 0.2), lab="raw")
plot!(x, ewma, lab="EWMA")

title!("Potentiometer reading vs. time")
xaxis!("Time [ms]")
yaxis!("Voltage signal [ADC units]")

savefig("adc.svg")


text:
# This function calculates the largest eigenvalue
# (absolute value wise) by power iteration.
const RKCAlgs = Union{RKC,IRKC,ESERK4,ESERK5,SERK2}
function maxeig!(integrator, cache::OrdinaryDiffEqConstantCache)
  isfirst = integrator.iter == 1 || integrator.u_modified
  @unpack t, dt, uprev, u, f, p, fsalfirst = integrator
  maxiter = (typeof(integrator.alg) <: Union{ESERK4,ESERK5,SERK2}) ? 100 : 50

  safe = (typeof(integrator.alg) <: RKCAlgs) ? 1.0 : 1.2
  # Initial guess for eigenvector `z`
  if isfirst
    if typeof(integrator.alg) <: RKCAlgs
      if integrator.alg isa IRKC
        z = cache.du₂
      else
        z = fsalfirst
      end
    else
      fz = fsalfirst
      z = f(fz, p, t)
      integrator.destats.nf += 1
    end
  else
    z = cache.zprev
  end
  # Perturbation
  u_norm = integrator.opts.internalnorm(uprev,t)
  z_norm = integrator.opts.internalnorm(z,t)
  pert   = eps(u_norm)
  sqrt_pert = sqrt(pert)
  is_u_zero = u_norm == zero(u_norm)
  is_z_zero = z_norm == zero(z_norm)
  # Normalize `z` such that z-u lie in a circle
  if ( !is_u_zero && !is_z_zero )
    dz_u = u_norm * sqrt_pert
    quot = dz_u/z_norm
    z = uprev + quot*z
  elseif !is_u_zero
    dz_u = u_norm * sqrt_pert
    z = uprev + uprev*dz_u
  elseif !is_z_zero
    dz_u = pert
    quot = dz_u/z_norm
    z *= quot
  else
    dz_u = pert
    z = dz_u*ones(z)
  end # endif
  # Start power iteration
  integrator.eigen_est = 0
  for iter in 1:maxiter
    if integrator.alg isa IRKC
      fz = f.f2(z, p, t)
      integrator.destats.nf2 += 1
      tmp = fz - cache.du₂
    else
      fz = f(z, p, t)
      integrator.destats.nf += 1
      tmp = fz - fsalfirst
    end
    Δ  = integrator.opts.internalnorm(tmp,t)
    eig_prev = integrator.eigen_est
    integrator.eigen_est = Δ/dz_u * safe
    # Convergence
    if typeof(integrator.alg) <: RKCAlgs # To match the constants given in the paper
      if iter >= 2 && abs(eig_prev - integrator.eigen_est) < max(integrator.eigen_est,1.0/integrator.opts.dtmax)*0.01
        integrator.eigen_est *= 1.2
        # Store the eigenvector
        cache.zprev = z - uprev
        return true
      end
    else
      if iter >= 2 && abs(eig_prev - integrator.eigen_est) < integrator.eigen_est*0.05
        # Store the eigenvector
        cache.zprev = z - uprev
        return true
      end
    end

    # Next `z`
    if Δ != zero(Δ)
      quot = dz_u/Δ
      z = uprev + quot*tmp
    else
      # An arbitrary change on `z`
      nind = length(z)
      if (nind != 1)
        ind = 1 + iter % nind
        # val = (uprev[ind] - (z[ind] - uprev[ind]))*one(eltype(z))*2
        vec(z) .= vec(z) .* (1 .- 2 .*((1:length(z)) .== ind))
      else
        z = -z
      end
    end
  end
  return false
end

function maxeig!(integrator, cache::OrdinaryDiffEqMutableCache)
  isfirst = integrator.iter == 1 || integrator.u_modified
  @unpack t, dt, uprev, u, f, p, fsalfirst = integrator
  if cache isa IRKCCache
    fz, z, atmp = integrator.fsallast, cache.nlsolver.tmp, cache.atmp
  else
    fz, z, atmp = cache.k, cache.tmp, cache.atmp
  end
  ccache = cache.constantcache
  maxiter = (typeof(integrator.alg) <: Union{ESERK4,ESERK5,SERK2}) ? 100 : 50
  safe = (typeof(integrator.alg) <: RKCAlgs) ? 1.0 : 1.2
  # Initial guess for eigenvector `z`
  if isfirst
    if typeof(integrator.alg) <: RKCAlgs
      if integrator.alg isa IRKC
        @.. z = cache.du₂
      else
        @.. z = fsalfirst
      end
    else
      @.. fz = fsalfirst
      f(z, fz, p, t)
      integrator.destats.nf += 1
    end
  else
    @.. z = ccache.zprev
  end
  # Perturbation
  u_norm = integrator.opts.internalnorm(uprev,t)
  z_norm = integrator.opts.internalnorm(z,t)
  pert   = eps(u_norm)
  sqrt_pert = sqrt(pert)
  is_u_zero = u_norm == zero(u_norm)
  is_z_zero = z_norm == zero(z_norm)
  # Normalize `z` such that z-u lie in a circle
  if ( !is_u_zero && !is_z_zero )
    dz_u = u_norm * sqrt_pert
    quot = dz_u/z_norm
    @.. z = uprev + quot*z
  elseif !is_u_zero
    dz_u = u_norm * sqrt_pert
    @.. z = uprev + uprev*dz_u
  elseif !is_z_zero
    dz_u = pert
    quot = dz_u/z_norm
    @.. z *= quot
  else
    dz_u = pert
    @.. z = dz_u*one(eltype(z))
  end # endif
  # Start power iteration
  integrator.eigen_est = 0
  for iter in 1:maxiter
    if integrator.alg isa IRKC
      f.f2(fz, z, p, t)
      integrator.destats.nf2 += 1
      @.. atmp = fz - cache.du₂
    else
      f(fz, z, p, t)
      integrator.destats.nf += 1
      @.. atmp = fz - fsalfirst
    end
    Δ  = integrator.opts.internalnorm(atmp,t)
    eig_prev = integrator.eigen_est
    integrator.eigen_est = Δ/dz_u * safe
    # Convergence
    if typeof(integrator.alg) <: RKCAlgs # To match the constants given in the paper
      if iter >= 2 && abs(eig_prev - integrator.eigen_est) < max(integrator.eigen_est,1.0/integrator.opts.dtmax)*0.01
        integrator.eigen_est *= 1.2
        # Store the eigenvector
        @.. ccache.zprev = z - uprev
        return true
      end
    else
      if iter >= 2 && abs(eig_prev - integrator.eigen_est) < integrator.eigen_est*0.05
        # Store the eigenvector
        @.. ccache.zprev = z - uprev
        return true
      end
    end
    # Next `z`
    if Δ != zero(Δ)
      quot = dz_u/Δ
      @.. z = uprev + quot*atmp
    else
      # An arbitrary change on `z`
      nind = length(uprev)
      if (nind != 1)
        ind = 1 + iter % nind
        # val = (uprev[ind] - (z[ind] - uprev[ind]))*one(eltype(z))
        vec(z) .= vec(z) .* (1 .- 2 .*((1:length(z)) .== ind))
      else
        z = -z
      end
    end
  end
  return false
end
"""
    choosedeg!(cache) -> nothing

Calculate `mdeg = ms[deg_index]` (the degree of the Chebyshev polynomial)
and `cache.start` (the start index of recurrence parameters for that
degree), where `recf` are the `μ,κ` pairs
for the `mdeg` degree method. The `κ` for `stage-1` for every degree
is 0 therefore it's not included in `recf`
  """
function choosedeg!(cache::T) where T
  isconst = T <: OrdinaryDiffEqConstantCache
  isconst || ( cache = cache.constantcache )
  start = 1
  @inbounds for i in 1:size(cache.ms,1)
    if cache.ms[i] >= cache.mdeg
      cache.deg_index = i
      cache.mdeg = cache.ms[i]
      cache.start = start
      break
    end
    start += cache.ms[i]*2 - 1
  end
  return nothing
end

function choosedeg_SERK!(integrator,cache::T) where T
  isconst = T <: OrdinaryDiffEqConstantCache
  isconst || ( cache = cache.constantcache )
  @unpack ms = cache
  start = 1
  @inbounds for i in 1:size(ms,1)
    if ms[i] < cache.mdeg
      start += ms[i]+1
    else
      cache.start = start
      cache.mdeg = ms[i]
      break
    end
  end
  if integrator.alg isa ESERK5
    if cache.mdeg <= 20
      cache.internal_deg = 2
    elseif cache.mdeg <= 50
      cache.internal_deg = 5
    elseif cache.mdeg <= 100
      cache.internal_deg = 10
    elseif cache.mdeg <= 500
      cache.internal_deg = 50
    elseif cache.mdeg <= 1000
      cache.internal_deg = 100
    elseif cache.mdeg <= 2000
      cache.internal_deg = 200
    end
  end

  if integrator.alg isa ESERK4
    if cache.mdeg <= 20
      cache.internal_deg = 2
    elseif cache.mdeg <= 100
      cache.internal_deg = 10
    elseif cache.mdeg <= 500
      cache.internal_deg = 25
    elseif cache.mdeg <= 1000
      cache.internal_deg = 100
    elseif cache.mdeg <= 4000
      cache.internal_deg = 200
    end
  end

  if integrator.alg isa SERK2
    cache.internal_deg = cache.mdeg/10
  end
  return nothing
end


text:
function check_valid_code(code)
    occursin(r"^[\sa-zA-Z0-9_]+$", code) || throw(ErrorException("`$code` is not a valid series code"))
end

function check_valid_date(date)
    occursin(r"\d{4}-\d{2}-\d{2}", date) && length(date)==10
end

function check_valid_as_at_date(date)
    check_valid_date(date) ||
    throw(ErrorException("Invalid as_at_date `$date`. Date must be in form YYYY-MM-DD."))
end

function check_valid_vintage(vintage)
    (vintage == "current" || vintage == "latest" || vintage == "previous" ||
    check_valid_date(vintage)) ||
    throw(ErrorException("Invalid vintage `$vintage`. Options are a date (in form YYYY-MM-DD) or one of `current` (alias `latest`) or `previous`"))
end


text:
@testset "fittedeva.jl" begin
    @testset "Base.show(io, obj)" begin
        # Print BayesianEVA does not throw
        fm = Extremes.BayesianEVA(Extremes.BlockMaxima(Variable("y", [100.0])), Mamba.Chains([100.0 log(5.0) .1]))
        buffer = IOBuffer()
        @test_logs Base.show(buffer, fm)

        # Print MaximumLikelihoodEVA does not throw
        fm = Extremes.MaximumLikelihoodEVA(BlockMaxima(Variable("y", [1])), [1.0, 1.0, 0.1])
        @test_logs Base.show(buffer, fm)

        # Print PwmEVA does not throw
        fm = Extremes.pwmEVA{BlockMaxima, GeneralizedExtremeValue}(Extremes.BlockMaxima(Variable("y", [0])), [1.0; 0.0; 0.1])
        @test_logs Base.show(buffer, fm)

    end

    include(joinpath("fittedeva", "bayesianeva_test.jl"))
    include(joinpath("fittedeva", "maximumlikelihoodeva_test.jl"))
    include(joinpath("fittedeva", "pwmeva_test.jl"))

end


text:
# Silo.jl: Wrapper around the Silo library (https://wci.llnl.gov/simulation/computer-codes/silo)
VERSION >= v"0.4.0-dev+6521" && __precompile__()
module Silo
const depsfile = joinpath(dirname(dirname(@__FILE__)), "deps", "deps.jl")
if isfile(depsfile)
  include(depsfile)
else
  error("Silo not properly installed. Please run Pkg.build(\"Silo\")")
end

using Compat

include("libsilo_consts.jl")

include("libsilo_types.jl")

include("pmpio.jl")

include("libsilo.jl")

include("libsilo_h.jl")

include("optlists.jl")

# MAIN Interface:
#	- siloopen(..., "w") mode: creates and initializes. Done
#	- overload write(file, name, values) / silowrite(...) Done (for 1D)
#	- siloclose(...) Done
#	- names(file): values contained in file Done (for 1D)
#	- overload read(file, name) / siloread(...) Done (for 1D)

end # module


text:
using Random, Distributions
using Graphs
using Pkg, Revise
using CSV, DataFrames
using JuMP, CPLEX
using SparseArrays

Pkg.activate(".")
using RRSTExperiments

function generate_graph(seed::UInt32, n::Int; max_weight = 10)
    Random.seed!(seed * n)

    V = collect(1:n) # set of vertices
    E = Tuple{Int64, Int64}[]
    C = Float64[]

    g = Graphs.SimpleGraph(n)
    while length(Graphs.connected_components(g)) > 1
        source = rand(V)
        target = rand(V)
        if source == target || (source, target) ∈ E || (target, source) ∈ E
        continue
        else
        Graphs.add_edge!(g, source, target)
        push!(E, (source, target))
        push!(C, rand(1:max_weight))
        end
    end

    return (E, C)
end

function prepare_graph(seed::UInt32, n::Int, α::Float64 = 1.0, β::Float64 = 3.0)
    # Read vertices, edges and first stage costs
    E, C = generate_graph(seed, n; max_weight = 100)

    # Assume second stage costs are equal to first stage costs
    c = copy(C)

    dist = Beta(α, β)
    d = generate_uncertain_costs(dist, c)

    return (E, C, c, d)
  end

function solve_tree_model(n::Int, E::Vector{InputEdge})
    V = collect(1:n) # set of nodes
    Vminus1 = setdiff(V, [1]) # commodity nodes
    A = [(e.i, e.j) for e ∈ E] ∪ [(e.j, e.i) for e ∈ E]

    # Model
    model = Model(CPLEX.Optimizer)
    set_silent(model)
    set_optimizer_attribute(model, "CPX_PARAM_EPINT", 1e-8)

    # Variables
    @variable(model, y[E] ≥ 0) # yₑ=1 if e∈E belongs to the spanning tree; 0 otherwise
    @variable(model, w[A] ≥ 0)
    @variable(model, f[A, Vminus1] ≥ 0) # flow variables

    # Objective
    @objective(model, Min, sum(E[i].C * y[e] for (i, e) ∈ enumerate(E)))

    #Constraints
    for k ∈ Vminus1 # sources
        @constraint(model, sum(f[(j,1),k] for j ∈ filter(j -> (j,1) ∈ A, V)) -
        sum(f[(1,j),k] for j ∈ filter(j -> (1,j) ∈ A,V)) == -1)
    end

    for k ∈ Vminus1, i ∈ Vminus1 # balances
        if i ≠ k
            @constraint(model, sum(f[(j,i),k] for j ∈ filter(j -> (j,i) ∈ A, V)) -
            sum(f[(i,j),k] for j ∈ filter(j -> (i,j) ∈ A, V)) == 0)
        end
    end

    for k ∈ Vminus1 # sinks
        @constraint(model, sum(f[(j,k),k] for j ∈ filter(j -> (j,k) ∈ A,V)) -
        sum(f[(k,j),k] for j ∈ filter(j -> (k,j) ∈ A,V)) == 1)
    end

    for k ∈ Vminus1, a ∈ A # capacity
        @constraint(model, f[a,k] ≤ w[a])
    end

    # tree
    @constraint(model, sum(w[a] for a ∈ A) == n - 1)

    for e ∈ E
        @constraint(model, y[e] == w[(e.i, e.j)] + w[(e.j, e.i)])
    end

    # Solve

    optimize!(model)

    status=termination_status(model)
    if status == MOI.OPTIMAL
        obj_value = objective_value(model)
        return status, obj_value, Array(value.(y))
    else
        return status, missing, missing
    end
end

function solve_tree_kruskal(n, A)
    weight_map::Dict{Tuple{Int64, Int64}, Float64} = Dict()
    g = SimpleGraph(n)
    for e in A
      add_edge!(g, e.i, e.j)
      push!(weight_map, (e.i, e.j) => e.C)
    end

    w_mat = spzeros(Float64, n, n)
    for e in A
        w_mat[e.i, e.j] = e.C
        w_mat[e.j, e.i] = e.C
    end

    t_x = kruskal_mst(g, w_mat)

    obj_value = get_objective_value(weight_map, t_x)
    x = get_solution(t_x)
    return obj_value, x
end

function get_objective_value(w::Dict{Tuple{Int64, Int64}, Float64}, t_x::Vector{Edge{Int}})
    return sum(map(e -> haskey(w, (e.src, e.dst)) ? w[(e.src, e.dst)] : w[(e.dst, e.src)], t_x))
  end

function get_solution(t_x::Vector{Edge{Int}})
  return map(e -> (e.src, e.dst), t_x)
end

for n in 25:2500, i in 1:3
    print("n=", n, ", k = ", i, ": ")
    seed::UInt32 = n * i
    E, C, c, d = prepare_graph(seed, n)
    A = [InputEdge(a, b, C[i], c[i], d[i]) for (i, (a,b)) in enumerate(E)]
    obj₁, _ = solve_tree_kruskal(n, A)
    print("obj₁=", obj₁)
    @time begin
        _, obj₂, _ = solve_tree_model(n, A)
        println(", obj₂=", obj₂)
    end
    @assert obj₁ == obj₂
end


text:
@doc raw"""
    LinearKernel([a=1 [,c=1]])

The linear kernel is a Mercer kernel given by:

```math
\kappa(\mathbf{x},\mathbf{y}) =
a \mathbf{x}^\intercal \mathbf{y} + c \qquad \alpha > 0, \; c \geq 0
```
"""
struct LinearKernel{T<:AbstractFloat} <: MercerKernel{T}
    a::HyperParameter{T}
    c::HyperParameter{T}
    LinearKernel{T}(a::Real, c::Real) where {T<:AbstractFloat} = new{T}(
        HyperParameter(convert(T,a), interval(OpenBound(zero(T)), nothing)),
        HyperParameter(convert(T,c), interval(ClosedBound(zero(T)), nothing))
    )
end
LinearKernel(a::T1=1.0, c::T2=one(T1)) where {T1<:Real,T2<:Real} = LinearKernel{floattype(T1,T2)}(a,c)

@inline linearkernel(z::T, a::T, c::T) where {T<:AbstractFloat} = a*z + c

@inline basefunction(::LinearKernel) = ScalarProduct()
@inline kappa(κ::LinearKernel{T}, z::T) where {T} = linearkernel(z, getvalue(κ.a), getvalue(κ.c))


text:
module ICR

using SmoothedParticles
using Printf
using SparseArrays

function add_element!(I::Vector{Int64}, J::Vector{Int64}, V::Vector{Float64}, i::Int64, j::Int64, v::Float64)
    if v == 0.0
        return
    end
    push!(I, i)
    push!(J, j)
    push!(V, v)
end

function lhs(sys::ParticleSystem, rDker::Function)::SparseMatrixCSC{Float64}
	N = length(sys.particles)
    I = Int64[]
    J = Int64[]
    V = Float64[]
    for i in 1:N
        p = sys.particles[i]
        S1 = SPHLib.sum(sys, (p,q,r) -> -2*rDker(sys.h,r)*(p.x[1]-q.x[1]), p)
        S2 = SPHLib.sum(sys, (p,q,r) -> -2*rDker(sys.h,r)*(p.x[2]-q.x[2]), p)
        S = RealVector(S1, S2, 0.)
        for j in 1:N
            q = sys.particles[j]
            r = dist(p,q)
            if r > sys.h
                continue
            end
            tmp0 = Float64(p==q)
            tmp1 = rDker(sys.h,r)*(p.x - q.x)
            tmp2 = 0.5*Float64(p==q)*S
            #UL block = I
            add_element!(I,J,V,i+0N,j+0N,tmp0)
            add_element!(I,J,V,i+1N,j+1N,tmp0)
            #UR block = Grad
            add_element!(I,J,V,i+0N,j+2N, tmp1[1]-tmp2[1])
            add_element!(I,J,V,i+1N,j+2N, tmp1[2]-tmp2[2])
            #DL block = Div
            add_element!(I,J,V,i+2N,j+0N, tmp1[1]+tmp2[1])
            add_element!(I,J,V,i+2N,j+1N, tmp1[2]+tmp2[2])
        end
    end
    return sparse(I, J, V, 3N, 3N)
end

function rhs(sys::ParticleSystem, ker::Function, rho0::Float64)
    N = length(sys.particles)
    b = zeros(3N)
    for i in 1:N
        p = sys.particles[i]
        b[i+2N] = SPHLib.sum(sys, (p,q,r) -> ker(sys.h,r), p) - rho0
	end
    return b
end

function renormalize!(sys::ParticleSystem, dr::Float64;
    tol = 1e-6, max_steps = 10,
    ker::Function = wendland2, rDker::Function = rDwendland2)
    @info("renormalizing, please wait...")
    rho0 = 1/dr^2
    for p in sys.particles
        p.x += 0.3*dr*RealVector(rand() - 1.0, rand() - 1.0, 0.)
    end
    for _ in 1:max_steps
        create_cell_list!(sys)
        b = rhs(sys, ker, rho0)
        err = maximum(x -> abs(x), b)
        @show err
        if err < tol
            break
        end
        A = lhs(sys, rDker)
        y = A\b
        N = length(sys.particles)
        for i in 1:N
            sys.particles[i].x += RealVector(y[i], y[i+N], 0.)
        end
	end
end

end


text:
include(joinpath(dirname(@__DIR__), "src", "BestApproximation.jl"))
using Documenter, .BestApproximation

Documenter.makedocs(
    clean = true,
    doctest = true,
    modules = Module[BestApproximation],
    repo = "",
    highlightsig = true,
    sitename = "BestApproximation Documentation",
    expandfirst = [],
    pages = [
        "Home" => "index.md",
        "Usage" => "usage.md"
    ]
)

deploydocs(
    repo  =  "github.com/jakewilliami/BestApproximation.jl.git",
)

max_stars_repo_name: jakewilliami/BestApproximation

max_issues_repo_name: jakewilliami/BestApproximation

max_forks_repo_name: jakewilliami/BestApproximation


text:
coord = Coord()
coord = Coord(1.0,2.0,3.0)
@test coord.x == 1.0
@test coord.y == 2.0
@test coord.z == 3.0
@test coord.h == 0.0
@test coord.p == 0.0
@test coord.r == 0.0

coord = Coord(1.0,2.0,3.0,4.0,5.0,6.0)
@test coord.x == 1.0
@test coord.y == 2.0
@test coord.z == 3.0
@test coord.h == 4.0
@test coord.p == 5.0
@test coord.r == 6.0

@test coord == Coord(1.0,2.0,3.0,4.0,5.0,6.0)

coordim = CoordIm(1.0,2.0,3.0)
@test coordim.x == 1.0
@test coordim.y == 2.0
@test coordim.z == 3.0
@test coordim.h == 0.0
@test coordim.p == 0.0
@test coordim.r == 0.0

coordim = CoordIm(1.0,2.0,3.0,4.0,5.0,6.0)
@test coordim.x == 1.0
@test coordim.y == 2.0
@test coordim.z == 3.0
@test coordim.h == 4.0
@test coordim.p == 5.0
@test coordim.r == 6.0

@test coord == coordim
@test coordim == coord
@test coordim == CoordIm(1.0,2.0,3.0,4.0,5.0,6.0)

@test coord == copy(coord)
@test coord == copy(coordim)
@test coord == deepcopy(coord)
@test coord == deepcopy(coordim)
@test copy!(coord, Coord(2.0,3.0,4.0,5.0,6.0,7.0)) == Coord(2.0,3.0,4.0,5.0,6.0,7.0)
@test copy!(coord, CoordIm(2.0,3.0,4.0,5.0,6.0,7.0)) == Coord(2.0,3.0,4.0,5.0,6.0,7.0)

show(coord); print("\n")
println(coord)
show(coordim); print("\n")
println(coordim)

coordA = Coord(1.0,2.0,3.0,4.0,5.0,6.0)
coordB = Coord(1.5, 2.5, 3.5, 4.5, 5.5, 6.5)
@test isapprox(get_dist(coordA, coordA), 0.0)
@test isapprox(get_dist(coordA, coordB), norm([coordA.x-coordB.x, coordA.y-coordB.y, coordA.z-coordB.z]))
@test isapprox(get_dist(coordB, coordA), norm([coordA.x-coordB.x, coordA.y-coordB.y, coordA.z-coordB.z]))

@test get_dist2d(coordA, coordA) == 0.0
@test isapprox(get_dist2d(coordA, coordB), norm([coordA.x-coordB.x, coordA.y-coordB.y]))
@test isapprox(get_dist2d(coordB, coordA), norm([coordA.x-coordB.x, coordA.y-coordB.y]))

coordA = CoordIm(1.0,2.0,3.0,4.0,5.0,6.0)
coordB = CoordIm(1.5, 2.5, 3.5, 4.5, 5.5, 6.5)
@test isapprox(get_dist(coordA, coordA), 0.0)
@test isapprox(get_dist(coordA, coordB), norm([coordA.x-coordB.x, coordA.y-coordB.y, coordA.z-coordB.z]))
@test isapprox(get_dist(coordB, coordA), norm([coordA.x-coordB.x, coordA.y-coordB.y, coordA.z-coordB.z]))

coordA = CoordIm(1.0,2.0,3.0,4.0,5.0,6.0)
coordB = Coord(1.5, 2.5, 3.5, 4.5, 5.5, 6.5)
@test isapprox(get_dist(coordA, coordA), 0.0)
@test isapprox(get_dist(coordA, coordB), norm([coordA.x-coordB.x, coordA.y-coordB.y, coordA.z-coordB.z]))
@test isapprox(get_dist(coordB, coordA), norm([coordA.x-coordB.x, coordA.y-coordB.y, coordA.z-coordB.z]))

@test get_dist2d(coordA, coordA) == 0.0
@test isapprox(get_dist2d(coordA, coordB), norm([coordA.x-coordB.x, coordA.y-coordB.y]))
@test isapprox(get_dist2d(coordB, coordA), norm([coordA.x-coordB.x, coordA.y-coordB.y]))

coordA = Coord(1.0,2.0,3.0,4.0,5.0,6.0)
coordB = CoordIm(1.5, 2.5, 3.5, 4.5, 5.5, 6.5)
@test isapprox(get_dist(coordA, coordA), 0.0)
@test isapprox(get_dist(coordA, coordB), norm([coordA.x-coordB.x, coordA.y-coordB.y, coordA.z-coordB.z]))
@test isapprox(get_dist(coordB, coordA), norm([coordA.x-coordB.x, coordA.y-coordB.y, coordA.z-coordB.z]))

@test get_dist2d(coordA, coordA) == 0.0
@test isapprox(get_dist2d(coordA, coordB), norm([coordA.x-coordB.x, coordA.y-coordB.y]))
@test isapprox(get_dist2d(coordB, coordA), norm([coordA.x-coordB.x, coordA.y-coordB.y]))

@test isapprox(get_value(coordA), norm([coordA.x, coordA.y, coordA.z]))
@test isapprox(get_value(coordB), norm([coordB.x, coordB.y, coordB.z]))

@test isa(convert(Coord, coordA), Coord)
@test isa(convert(CoordIm, coordA), CoordIm)
@test isa(convert(Coord, coordB), Coord)
@test isa(convert(CoordIm, coordB), CoordIm)

@test convert(Coord, coordA) == coordA
@test convert(Coord, coordB) == coordB
@test convert(CoordIm, coordA) == coordA
@test convert(CoordIm, coordB) == coordB

@test coordA * 2.0 == Coord(2.0, 4.0, 6.0, 8.0, 10.0, 12.0)
@test coordB * 2.0 == Coord(3.0, 5.0, 7.0, 9.0, 11.0, 13.0)
println(times!(convert(Coord, copy(coordB)), 2.0))
@test times!(convert(Coord, copy(coordB)), 2.0) == Coord(3.0, 5.0, 7.0, 9.0, 11.0, 13.0)
@test times!(Coord(), copy(coordA), 2.0) == Coord(2.0, 4.0, 6.0, 8.0, 10.0, 12.0)
@test times!(Coord(), copy(coordB), 2.0) == Coord(3.0, 5.0, 7.0, 9.0, 11.0, 13.0)

@test coordA * 2.0 == Coord(2.0, 4.0, 6.0, 8.0, 10.0, 12.0)
@test coordB * 2.0 == Coord(3.0, 5.0, 7.0, 9.0, 11.0, 13.0)
@test times!(copy(coordA), 2.0) == Coord(2.0, 4.0, 6.0, 8.0, 10.0, 12.0)
@test times!(convert(Coord, copy(coordB)), 2.0) == Coord(3.0, 5.0, 7.0, 9.0, 11.0, 13.0)
@test times!(Coord(), copy(coordA), 2.0) == Coord(2.0, 4.0, 6.0, 8.0, 10.0, 12.0)
@test times!(Coord(), copy(coordB), 2.0) == Coord(3.0, 5.0, 7.0, 9.0, 11.0, 13.0)

@test coordA + coordB == Coord(2.5, 4.5, 6.5, 8.5, 10.5, 12.5)
@test plus!(copy(coordA), coordB) == Coord(2.5, 4.5, 6.5, 8.5, 10.5, 12.5)
@test plus!(Coord(), coordA, coordB) == Coord(2.5, 4.5, 6.5, 8.5, 10.5, 12.5)
@test coordB - coordA == Coord(0.5, 0.5, 0.5, 0.5, 0.5, 0.5)
@test minus!(convert(Coord, copy(coordB)), coordA) == Coord(0.5, 0.5, 0.5, 0.5, 0.5, 0.5)
@test minus!(Coord(), coordB, coordA) == Coord(0.5, 0.5, 0.5, 0.5, 0.5, 0.5)
@test minus!(convert(Coord, copy(coordB)), convert(CoordIm, coordA)) == Coord(0.5, 0.5, 0.5, 0.5, 0.5, 0.5)
@test minus!(Coord(), coordB,  convert(CoordIm, coordA)) == Coord(0.5, 0.5, 0.5, 0.5, 0.5, 0.5)

@test init!(copy(coordA)) == Coord(0.0, 0.0, 0.0, 0.0, 0.0, 0.0)

max_stars_repo_name: sisl/OpenDriveManagerLite

max_issues_repo_name: sisl/OpenDriveManagerLite

max_forks_repo_name: sisl/OpenDriveManagerLite


text:
"""
connect_mesh(EToV,fv)

Initialize element connectivity matrices, element to element and
element to face connectivity. Works on general element types, so long as

        fv = array of arrays containing (unordered) indices of face vertices

is provided.

# Examples
- connect_mesh(EToV,UniformTriMesh.fv())
- connect_mesh(EToV,[[1,2],[2,3],[3,1]])
```jldoctest
```
"""
function connect_mesh(EToV,fv)
        Nfaces = length(fv)
        K = size(EToV,1)

        # sort and find matches
        fnodes = [[sort(EToV[e,ids]) for ids = fv, e = 1:K]...]
        p = sortperm(fnodes) # sorts by lexicographic ordering by default
        fnodes = fnodes[p,:]

        FToF = reshape(collect(1:Nfaces*K),Nfaces,K)
        for f = 1:size(fnodes,1)-1
                if fnodes[f,:]==fnodes[f+1,:]
                        f1 = FToF[p[f]]
                        f2 = FToF[p[f+1]]
                        FToF[p[f]] = f2
                        FToF[p[f+1]] = f1
                end
        end
        return FToF
end

max_stars_repo_name: jlchan/ESDG-CNS

max_issues_repo_name: jlchan/ESDG-CNS

max_forks_repo_name: jlchan/ESDG-CNS


text:
"""
    TurbulenceConvection

Turbulence convection models, for example
the Eddy-Diffusivity Mass-Flux model
"""
module TurbulenceConvection

using ..BalanceLaws: BalanceLaw, AbstractStateType
using ..VariableTemplates: @vars, Vars, Grad

import ..Mesh.Filters: vars_state_filtered

export TurbulenceConvectionModel, NoTurbConv

export init_state_prognostic!,
    init_aux_turbconv!, turbconv_nodal_update_auxiliary_state!

import ..BalanceLaws:
    vars_state,
    eq_tends,
    prognostic_to_primitive!,
    primitive_to_prognostic!,
    precompute,
    prognostic_vars,
    init_state_prognostic!,
    init_state_auxiliary!,
    update_auxiliary_state!,
    boundary_state!,
    compute_gradient_argument!,
    compute_gradient_flux!,
    integral_load_auxiliary_state!,
    integral_set_auxiliary_state!

using ..MPIStateArrays: MPIStateArray
using ..DGMethods: DGModel, LocalGeometry

abstract type TurbulenceConvectionModel end

"""
    NoTurbConv <: TurbulenceConvectionModel

A "no model" type, which results in kernels that
pass through and do nothing.
"""
struct NoTurbConv <: TurbulenceConvectionModel end

prognostic_vars(::NoTurbConv) = ()

eq_tends(pv, ::NoTurbConv, tt) = ()
precompute(::NoTurbConv, bl, args, ts, tend_type) = NamedTuple()

vars_state(m::TurbulenceConvectionModel, ::AbstractStateType, FT) = @vars()
vars_state_filtered(m::TurbulenceConvectionModel, FT) = @vars()

function init_aux_turbconv!(
    m::TurbulenceConvectionModel,
    bl::BalanceLaw,
    aux::Vars,
    geom::LocalGeometry,
)
    return nothing
end

function update_auxiliary_state!(
    dg::DGModel,
    m::TurbulenceConvectionModel,
    bl::BalanceLaw,
    Q::MPIStateArray,
    t::Real,
    elems::UnitRange,
)
    return nothing
end

function turbconv_nodal_update_auxiliary_state!(
    m::TurbulenceConvectionModel,
    bl::BalanceLaw,
    state::Vars,
    aux::Vars,
    t::Real,
)
    return nothing
end

function compute_gradient_argument!(
    m::TurbulenceConvectionModel,
    bl::BalanceLaw,
    transform::Vars,
    state::Vars,
    aux::Vars,
    t::Real,
)
    return nothing
end

function compute_gradient_flux!(
    m::TurbulenceConvectionModel,
    bl::BalanceLaw,
    diffusive::Vars,
    ∇transform::Grad,
    state::Vars,
    aux::Vars,
    t::Real,
)
    return nothing
end

function integral_load_auxiliary_state!(
    m::TurbulenceConvectionModel,
    bl::BalanceLaw,
    integ::Vars,
    state::Vars,
    aux::Vars,
)
    return nothing
end

function integral_set_auxiliary_state!(
    m::TurbulenceConvectionModel,
    bl::BalanceLaw,
    aux::Vars,
    integ::Vars,
)
    return nothing
end

function init_state_prognostic!(
    m::NoTurbConv,
    bl::BalanceLaw,
    state,
    aux,
    localgeo,
    t,
) end

prognostic_to_primitive!(::NoTurbConv, _...) = nothing
primitive_to_prognostic!(::NoTurbConv, _...) = nothing

include("boundary_conditions.jl")
include("source.jl")

end


text:
"""
    radixsort!(vector_string)

Radixsort on strings

    svec - a vector of strings; sorts it by bits
"""
radixsort_lsd!(x) = radixsort_lsd24!(x)

function radixsort_lsd16!(svec::Vector{String})
    # find the maximum string length
    lens = reduce((x,y) -> max(x,sizeof(y)),0, svec)
    skipbytes = lens
    while lens > 0
       if lens > 8
            skipbytes = max(0, skipbytes - 16)
            sorttwo!(load_bits.(UInt128, svec, skipbytes), svec)
            lens -= 16
        elseif lens > 4
            skipbytes = max(0, skipbytes - 8)
            sorttwo!(load_bits.(UInt64, svec, skipbytes), svec)
            lens -= 8
        else
            skipbytes = max(0, skipbytes - 4)
            sorttwo!(load_bits.(UInt32, svec, skipbytes), svec)
            lens -= 4
        end
    end
    svec
end

function radixsort_lsd24!(svec::Vector{String})
    # find the maximum string length
    lens = reduce((x,y) -> max(x,sizeof(y)),0, svec)
    skipbytes = lens
    while lens > 0
        if lens > 16 && ceil(lens/24) < ceil(lens/16)
            skipbytes = max(0, skipbytes - 24)
            sorttwo!(load_bits.(Bits192, svec, skipbytes), svec)
            lens -= 24
        elseif lens > 8
        # if lens > 8
            skipbytes = max(0, skipbytes - 16)
            sorttwo!(load_bits.(UInt128, svec, skipbytes), svec)
            lens -= 16
        elseif lens > 4
            skipbytes = max(0, skipbytes - 8)
            sorttwo!(load_bits.(UInt64, svec, skipbytes), svec)
            lens -= 8
        else
            skipbytes = max(0, skipbytes - 4)
            sorttwo!(load_bits.(UInt32, svec, skipbytes), svec)
            lens -= 4
        end
    end
    svec
end

function radixsort_lsd32!(svec::Vector{String})
    # find the maximum string length
    lens = reduce((x,y) -> max(x,sizeof(y)),0, svec)
    skipbytes = lens
    while lens > 0
        if lens > 24
            skipbytes = max(0, skipbytes - 32)
            sorttwo!(load_bits.(Bits256, svec, skipbytes), svec)
            lens -= 32
        elseif lens > 16 && ceil(lens/24) < ceil(lens/16)
            skipbytes = max(0, skipbytes - 24)
            sorttwo!(load_bits.(Bits192, svec, skipbytes), svec)
            lens -= 24
        elseif lens > 8
        # if lens > 8
            skipbytes = max(0, skipbytes - 16)
            sorttwo!(load_bits.(UInt128, svec, skipbytes), svec)
            lens -= 16
        elseif lens > 4
            skipbytes = max(0, skipbytes - 8)
            sorttwo!(load_bits.(UInt64, svec, skipbytes), svec)
            lens -= 8
        else
            skipbytes = max(0, skipbytes - 4)
            sorttwo!(load_bits.(UInt32, svec, skipbytes), svec)
            lens -= 4
        end
    end
    svec
end

# radixsort!(svec::Vector{String}) = radixsort!(UInt, svec::Vector{String})
# radixsort!(::Type{T}, svec::Vector{String}) where T = radixsort_lsd!(svec)


text:
@testset "Fourier Dipole Acceleration" begin

"Tests these field-using functions give the same answer as `ewald_sum_dipole`"
function test_energy_consistency(crystal, latsize)
    μB = Sunny.BOHR_MAGNETON
    μ0 = Sunny.VACUUM_PERM
    sys = SpinSystem(crystal, Sunny.AbstractInteraction[], latsize; μB, μ0)
    rand!(sys)

    dipdip = dipole_dipole(; extent=5, η=0.5)
    dip_real = Sunny.DipoleRealCPU(dipdip, crystal, latsize, sys.sites_info; μB, μ0)
    dip_fourier = Sunny.DipoleFourierCPU(dipdip, crystal, latsize, sys.sites_info; μB, μ0)

    scales = [μB*info.S*info.g for info in sys.sites_info]
    moments = reshape(scales, nbasis(sys), 1, 1, 1) .* sys.sites
    direct_energy = (μ0/4π) * Sunny.ewald_sum_dipole(sys.lattice, moments; extent=5, η=0.5)
    real_energy = Sunny.energy(sys.sites, dip_real)
    fourier_energy = Sunny.energy(sys.sites, dip_fourier)

    @test real_energy ≈ fourier_energy
    @test direct_energy ≈ fourier_energy
end

function test_field_consistency(crystal, latsize)
    sys = SpinSystem(crystal, Sunny.AbstractInteraction[], latsize)
    rand!(sys)

    dipdip = dipole_dipole(; extent=4, η=0.5)
    dip_real = Sunny.DipoleRealCPU(dipdip, crystal, latsize, sys.sites_info)
    dip_fourier = Sunny.DipoleFourierCPU(dipdip, crystal, latsize, sys.sites_info)

    H1 = zero(sys)
    H2 = zero(sys)
    Sunny._accum_neggrad!(H1, sys.sites, dip_real)
    Sunny._accum_neggrad!(H2, sys.sites, dip_fourier)

    @test H1 ≈ H2
end

lat_vecs = lattice_vectors(1.0, 1.0, 2.0, 90., 90., 120.)
basis_vecs = [[0.0, 0.0, 0.0], [0.5, 0.5, 0.5]]
latsize = [5, 5, 5]
crystal = Crystal(lat_vecs, basis_vecs)

test_energy_consistency(crystal, latsize)
test_field_consistency(crystal, latsize)

end


text:
# fbp/z-list.jl

include("image_geom.jl")
include("sino_geom.jl")

include("rotate2d.jl")

include("cuboid_im.jl") # must be after image_geom.jl
include("disk-phantom.jl")
include("ellipse_im.jl")
include("ellipsoid_im.jl")
include("ellipse_sino.jl")
include("rect_im.jl")
include("rect_sino.jl")


text:
@testset "1318.minimum-flips-to-make-a-or-b-equal-to-c.jl" begin
    @test min_flip(2, 6, 5) == 3
    @test min_flip(4, 2, 7) == 1
    @test min_flip(1, 2, 3) == 0
end


text:
"""
    simulateSegment(::S, x0, P, tt)

Simulate path of a target process `P`, started from `x0` on a time-grid `tt`
"""
function simulateSegment(::S, x0, P, tt) where S
    Wnr = Wiener{S}()
    WW = Bridge.samplepath(tt, zero(S))
    sample!(WW, Wnr)
    # forcedSolve will allow sampling of domain-restricted diffusions
    WW, XX = forcedSolve(Euler(), x0, WW, P)
    XX, XX.yy[end]
end

"""
    findCrossings(XX, upLvl, downLvl, upSearch=false)

Find times of first up-crossings of level `upLvl`. The timing starts only after
the diffusion was brought below level `downLvl` first (or alternatively if
`upSearch` flag is set to true)
"""
function findCrossings(XX, upLvl, downLvl, upSearch=false)
    upSearch=upSearch
    upCrossingTimes = Float64[]
    for (x,t) in zip(XX.yy, XX.tt)
        if upSearch && x[1] > upLvl
            upSearch = false
            push!(upCrossingTimes, t)
        elseif !upSearch && x[1] < downLvl
            upSearch = true
        end
    end
    upCrossingTimes, upSearch
end


text:
#
# Copyright (c) 2021 Tobias Thummerer, Lars Mikelsons, Josef Kircher
# Licensed under the MIT license. See LICENSE file in the project root for details.
#

using FMI
using FMIZoo
using Test
import Random

exportingToolsWindows = [("Dymola", "2022x")]
exportingToolsLinux = [("Dymola", "2022x")]
fmuStructs = ["FMU", "FMUCOMPONENT"]

function runtests(exportingTool)
    ENV["EXPORTINGTOOL"] = exportingTool[1]
    ENV["EXPORTINGVERSION"] = exportingTool[2]

    @testset "Testing FMUs exported from $exportingTool" begin

        @testset "Sensitivities" begin
            include("sens.jl")
        end

        for str in fmuStructs
            @testset "Functions for $str" begin
                ENV["FMUSTRUCT"] = str
                @testset "Variable Getters / Setters" begin
                    include("getter_setter.jl")
                end
                @testset "State Manipulation" begin
                    include("state.jl")
                end
                @testset "Directional derivatives" begin
                    include("dir_ders.jl")
                end
                @testset "Automatic Simulation (CS or ME)" begin
                    include("sim_auto.jl")
                end
                @testset "CS Simulation" begin
                    include("sim_CS.jl")
                end
                @testset "ME Simulation" begin
                    include("sim_ME.jl")
                end
                @testset "Support CS and ME simultaneously" begin
                    include("cs_me.jl")
                end
            end
        end

        @testset "Plotting" begin
            include("plots.jl")
        end
    end
end

@testset "FMI.jl" begin
    if Sys.iswindows()
        @info "Automated testing is supported on Windows."
        for exportingTool in exportingToolsWindows
            runtests(exportingTool)
        end
    elseif Sys.islinux()
        @info "Automated testing is supported on Linux."
        for exportingTool in exportingToolsLinux
            runtests(exportingTool)
        end
    elseif Sys.isapple()
        @warn "Test-sets are currrently using Windows- and Linux-FMUs, automated testing for macOS is currently not supported."
    end
end


text:
"""
    simulate(P::PopState,
             F::Fleet,
             M::MovementModel,
             PopDy::PopulationDynamicsModel,
             Ω::AbstractFisheryDomain,
             T::Int)

Simulate the population from P0 forward `T` years, by:

1. Remove individuals through fishing mortality using the `Fleet`,
2. Allow individuals to move according to their preferences,
3. Step the regional Schaefer model forward, distributing new
   individuals proportionally.
"""
function simulate(P0::PopState{Tf},
                  F::Fleet,
                  M::MovementModel,
                  PopDy::PopulationDynamicsModel,
                  Ω::AbstractFisheryDomain,
                  T::Ti) where {Tf<:Real, Ti<:Integer}
    P = copy(P0)
    Pvec = Vector{typeof(P)}()
    Cvec = Vector{Catch{Tf, Ti}}()
    for t in 1:T
        push!(Pvec, copy(P))
        c = fish!(P, F, Ω, t)
        append!(Cvec, c)
        P = M(P)                # TODO: Make these two steps in-place?
        P = PopDy(P)
    end
    Pvec, Cvec
end


text:
module BeaData

using HTTP, JSON, Retry
using DataFrames, DataStructures
using Dates
import Base.show

export
    # Types
    BeaTable,
    # Methods
    bea_datasets,
    bea_parameterlist,
    bea_parametervalues,
    bea_table,
    get_bea_datasets,
    get_bea_parameterlist,
    get_nipa_table

const API_URL = "https://apps.bea.gov/api/data"

function __init__()
    global USER_ID = ""
    if "BEA_USERID" in keys(ENV)
        USER_ID = ENV["BEA_USERID"]
        println("BEA UserID found!")
    elseif isfile(joinpath(homedir(), ".beadatarc"))
        open(joinpath(homedir(),".beadatarc"), "r") do f
            USER_ID = String(strip(read(f, String)))
        end
        println("BEA UserID found!")
    else
        println("No BEA UserID found!")
    end
end

include("metadata_methods.jl")
include("bea_table.jl")

end # module


text:
"""
 function for computation of interpolation matrix from FEM mesh to regular grid
    from displacements in tetraedra fem grid to a regular grid mesh

  Inputs:
    - Nodes: matrix with coordinates of nodes of the FEMmesh
    - Conec: vector of vectors. at i-th entry has the vector with the connectivity indexes
    - numComponentes: number of components of the magnitude to be interpolated
    - order: 0 interpolates from elements centroids
            1 interpolates from nodes
"""

function computeFEM2GridInterpMatrix( Nodes, Conec, intGrid, numCompon, order )

  numVoxelsX   =       intGrid.voxelNums[1]
  numVoxelsXY  = prod( intGrid.voxelNums[1:2] )
  numVoxels    = prod( intGrid.voxelNums )

  numNodes  = size( Nodes, 1)
  numElems  = length( Conec )

  @assert typeof(Conec) == Vector{ Vector{Int64} }

  booleanInterFound = zeros( Bool, numVoxels )

  rowIndexes = Vector{Int}()
  colIndexes = Vector{Int}()
  nonZerVals = Vector{Float64}()

  if order == 1
    totalRows = numCompon * numVoxels
    totalCols = numCompon * numNodes
  end

  # loop in elements for see which grid nodes are in an element.
  for i in (1:numElems)

    nodesElem   = Conec[i]
    coordesElem = Nodes[ nodesElem , :]

    # coordenadas de bounding box de elemento
    mins = vec( minimum( coordesElem, dims=1 ) )
    maxs = vec( maximum( coordesElem, dims=1 ) )

    # indexes of grid within bounding box
    indsIni, indsEnd = ranges( mins, maxs, intGrid )

    if length( indsIni ) > 0

      for kg in (indsIni[3]:indsEnd[3])
        for jg in (indsIni[2]:indsEnd[2])
          for ig in (indsIni[1]:indsEnd[1])

            # number of voxel
            ind = (kg-1) * numVoxelsXY + (jg-1)*numVoxelsX + ig ;

            # if an interpolation of the node was not already obtained
            if !(booleanInterFound[ind])

              pointCand = intGrid.startVox + intGrid.voxelWidths .* ([ ig, jg, kg ] .- 1 )

              # verify if it belongs to the current element
              itIsInTetra, shapeFuncs = checkInTetra( coordesElem, pointCand )

              # if it is in the tetrahedron
              if itIsInTetra
                # then we found it
                booleanInterFound[ind] = true

                dofsGridInd = nodes2dofs( ind, numCompon )

                # if order == 0
                #    # assign interpolation of grid node ind to element i
                #    interpMatrix[ind, i] = 1

                if order == 1

                  for indShapeFunc in (1:length(shapeFuncs))

                    dofsNode = nodes2dofs( nodesElem[indShapeFunc], numCompon )

                    append!( rowIndexes, dofsGridInd )
                    append!( colIndexes, dofsNode    )
                    append!( nonZerVals, shapeFuncs[ indShapeFunc ]*ones(numCompon) )

                  end

                end # if order

              end # if it is in tetra
            end # if node already analyzed
          end # loop x index
        end # loop y index
      end # loop z index
    end # if length inds > 0

  end # loop elements

  return sparse( rowIndexes, colIndexes, nonZerVals, totalRows, totalCols )
end

"""
 funcion para calculo rangos de indices xyz de grid que efecitivamente caen dentro del box del elemento
"""

function ranges( mins, maxs, intGrid )

  # if the maximums of the box are higher than all components of the startVoxel
  # then there could be an intersection
  if ( sum( maxs .>= intGrid.startVox ) == 3 ) && ( sum( mins .<= intGrid.endVox ) == 3 )
    # there is intersection
    indsEnd = Int.( min.( ceil.(  ( maxs - intGrid.startVox ) ./ intGrid.voxelWidths .+ 1 ), intGrid.voxelNums ) )
    indsIni = Int.( max.( floor.( ( mins - intGrid.startVox ) ./ intGrid.voxelWidths .+ 1 ), ones(3)           ) )
  else
    indsIni = []
    indsEnd = []
  end

  return indsIni, indsEnd
end

"""
TO DO
"""

function checkInTetra( nodesTetra, pointCand )

  A          = zeros(4,4)
  A[:, 1]   .= 1.0
  A[:, 2:4] .= nodesTetra
  volTot     = det(A) / 6.0

  volsRel = zeros(4)

  boolInTetra = true
  i=0
  while boolInTetra && i<4
    i = i + 1

    B        = copy(A)
    B[i,2:4] = pointCand
    volsRel[i] = ( det( B ) / 6.0 ) / volTot ;

    boolInTetra = boolInTetra && ( volsRel[i] >= -1.0e-10 )
  end

  if boolInTetra
    shapeFuncs = volsRel
  else
    shapeFuncs = []
  end

  return boolInTetra, shapeFuncs
end

"""
TO DO
"""

function vecGridConvert(ind)

  return 1
end

"""

copied from
https://github.com/jorgepz/FEMAssembler.jl/blob/main/src/mesh/dofsComputations.jl

"""
function nodes2dofs( nodes , degreespernode )
    n    = length(nodes);
    dofs = zeros( Int, n*degreespernode ) ;
    for i in (1:n);
        dofs[ (i-1)*degreespernode .+ Vector(1:degreespernode) ] = (degreespernode*(nodes[i]-1)).+Vector(1:degreespernode)  ;
    end
    return dofs;
end


text:
#----------------------------------------
# aggregate

"""
    aggregate(vsr::VerySimpleRaster, factor::Int, fun::Function [, file])

Aggregates the raster by merging `factor` cells in both directions, using the
aggregation function `fun`. Return value is a new raster. Also supports the
`do` block syntax. Use the optional `file` argument to keep the resulting raster
file in your work directory.
"""
aggregate(fun::Function, vsr, factor, file = "") = aggregate(vsr, factor, fun, file)

function aggregate(vsr::VerySimpleRaster, factor::Int, fun::Function, file = "")
   file == "" && (file = tempname())
   file[end-3:end] ∈ (".gri", "grd") && (file = file[1:end-4])
   io = open(file*".gri", "w")
   newm, newn = ceil.(Int,size(vsr.mat)./factor)
   type = typeof(fun([first(vsr.mat)]))
   innerfun(x) = fun(filter(x->x!=vsr.nodata, x))
   for j in 1:newn-1
      for i in 1:newm-1
        write(io, innerfun(vec(@view vsr.mat[(i-1)*factor .+ (1:factor), (j-1)*factor .+ (1:factor)])))
      end
      write(io, innerfun(vec(@view vsr.mat[(newm-1)*factor + 1:size(vsr.mat,1), (j-1)*factor .+ (1:factor)])))
   end
   for i in 1:newm-1
      write(io, innerfun(vec(@view vsr.mat[(i-1)*factor .+ (1:factor), (newn-1)*factor + 1:size(vsr.mat,2)])))
   end
   write(io, innerfun(vec(@view vsr.mat[(newm-1)*factor + 1:size(vsr.mat, 1), (newn-1)*factor + 1:size(vsr.mat,2)])))
   close(io)
   mat = Mmap.mmap(file*".gri", Matrix{type}, (newm, newn))
   vsr_out = VerySimpleRaster(mat, vsr.nodata, minimum(vsr.xs), minimum(vsr.ys), step(vsr.xs)*factor, step(vsr.ys)*factor, vsr.projection, file*".gri")
   write_grd_header(file*".grd", vsr_out)
   vsr_out
end

#--------------------------------------------
# getindex

function getindex(vsr::VerySimpleRaster, x...)
   ret = vsr.mat[x...]
   ret == vsr.nodata && return missing
   ret
end

function getindex(vsr::VerySimpleRaster, x)
   ret = vsr.mat[x]
   ret == vsr.nodata && return missing
   ret
end

setindex!(vsr::VerySimpleRaster, x, y, val) = setindex!(vsr.mat, x, y, val)
setindex!(vsr::VerySimpleRaster, x, val) = setindex!(vsr.mat, x, val)

#-------------------------------------------
# crop
coordinate_to_index(vsr, x) = coordinate_to_index(vsr, x...)
function coordinate_to_index(vsr, x, y)
   bb = bbox(vsr)
   (x < bb.xmin || x > bb.xmax || y < bb.ymin || y > bb.ymax) && (throw(BoundsError()))
   indx, indy = findfirst(i-> i>x, vsr.xs), findfirst(i-> i>y, vsr.ys)
   indx - 1, length(vsr.ys) - indy + 1
end

index_to_coordinate(vsr, x) = index_to_coordinate(vsr, x...)
function index_to_coordinate(vsr, x, y)
   (x < 1 || x > size(vsr.mat, 1) || y < 1 || y > size(vsr.mat, 2)) && throw(BoundsError())
   vsr.xs[x]+step(vsr.xs), vsr.ys[length(vsr.ys)-y+1]
end

"""
    crop(vsr::VerySimpleRaster{T}, xmin, xmax, ymin, ymax [, file])
    crop(vsr::VerySimpleRaster{T}, extent [, file])

Crops the raster to a rectangular window. Returns a new raster, use the
`file` argument to keep this in your working directory.
"""
crop(vsr, inds, file = "") = crop(vsr, inds..., file)
function crop(vsr::VerySimpleRaster{T}, xmin, xmax, ymin, ymax, file = "") where T
   x1,y1 = coordinate_to_index(vsr, xmin, ymax)
   x2,y2 = coordinate_to_index(vsr, xmax, ymin)
   newsize = (x2-x1, y2-y1)

   file == "" && (file = tempname())
   file[end-3:end] ∈ (".gri", "grd") && (file = file[1:end-4])
   open(file*".gri", "w") do IO
      write(IO, vsr.mat[x1:x2-1, y1+1:y2])
   end
   mat = Mmap.mmap(file*".gri", Matrix{T}, newsize)
   vsr_out = VerySimpleRaster(mat, vsr.nodata, vsr.xs[x1], vsr.ys[length(vsr.ys)-y2], step(vsr.xs), step(vsr.ys), vsr.projection, file*".gri")
   write_grd_header(file*".grd", vsr_out)
   vsr_out
end

#------------------------------------------
# extract

"""
   extract(vsr::VerySimpleRaster, x, y)
   extract(vsr::VerySimpleRaster, points)

Extracts the value of the raster at the given points.
"""
extract(vsr::VerySimpleRaster, x, y) = extract(vsr, (x, y))
function extract(vsr::VerySimpleRaster, tup)
   x, y = tup[1], tup[2]
   indx, indy = try
      coordinate_to_index(vsr, x, y)
   catch
      return missing
   end
   vsr[indx, indy]
end

extract(vsr::VerySimpleRaster, tup::Tuple{T, T}) where T <: AbstractVector = extract.(Ref(vsr), zip(tup...))
extract(vsr::VerySimpleRaster, ar::AbstractVector) = extract.(Ref(vsr), ar)

#-------------------------------------------
# mask

midpoints(a::AbstractVector) = [mean(a[i-1:i]) for i in eachindex(a)[2:end]]
const Pt{T<:Real} = Tuple{T, T}

function isinside(r::Pt, poly::AbstractVector)
    # An implementation of Hormann-Agathos (2001) Point in Polygon algorithm
    # See: http://www.sciencedirect.com/science/article/pii/S0925772101000128
    # Code segment adapted from PolygonClipping.jl
    c = false
    detq(q1,q2,r) = (q1[1]-r[1])*(q2[2]-r[2])-(q2[1]-r[1])*(q1[2]-r[2])
    for i in eachindex(poly)[2:end]
        q2 = poly[i]
        q1 = poly[i-1]
        if q1 == r
            @warn("point on polygon vertex - returning false")
            return false
        end
        if q2[2] == r[2]
            if q2[1] == r[1]
                @warn("point on polygon vertex - returning false")
                return false
            elseif (q1[2] == r[2]) && ((q2[1] > x) == (q1[1] < r[1]))
                @warn("point on edge - returning false")
                return false
            end
        end
        if (q1[2] < r[2]) != (q2[2] < r[2]) # crossing
            if q1[1] >= r[1]
                if q2[1] > r[1]
                    c = !c
                elseif ((detq(q1,q2,r) > 0) == (q2[2] > q1[2])) # right crossing
                    c = !c
                end
            elseif q2[1] > r[1]
                if ((detq(q1,q2,r) > 0) == (q2[2] > q1[2])) # right crossing
                    c = !c
                end
            end
        end
    end
    return c
end

"""
    mask(raster, polygon [, filename])`

Masks the raster by a polygon. Creates a new raster where points falling outside
the polygon have been replaced by `missing`
"""
function mask(vsr::VerySimpleRaster{T}, poly::AbstractVector, file = "") where T
   file == "" && (file = tempname())
   file[end-3:end] ∈ (".gri", "grd") && (file = file[1:end-4])
   open(file*".gri", "w") do IO
      for j in axes(vsr.mat, 2)
         for i in axes(vsr.mat, 1)
            if isinside(index_to_coordinate(vsr, i, j), poly)
               write(IO, vsr[i, j])
            else
               write(IO, vsr.nodata)
            end
         end
      end
   end
   mat = Mmap.mmap(file*".gri", Matrix{T}, size(vsr.mat))
   vsr_out = VerySimpleRaster(mat, vsr.nodata, first(vsr.xs), first(vsr.ys), step(vsr.xs), step(vsr.ys), vsr.projection, file*".gri")
   write_grd_header(file*".grd", vsr_out)
   vsr_out
end


text:
# Test 2D modeling
# The receiver positions and the source wavelets are the same for each of the four experiments.
# Author: Philipp Witte, pwitte@eos.ubc.ca
# Date: January 2017
#
# Mathias Louboutin, mlouboutin3@gatech.edu
# Updated July 2020

using JUDI
using ArgParse, Test, Printf
using SegyIO, LinearAlgebra, Distributed, JOLI

const GROUP = get(ENV, "GROUP", "JUDI")

include("utils.jl")

if endswith(GROUP, ".jl")
    @time include(GROUP)
end

base = ["test_abstract_vectors.jl",
        "test_geometry.jl",
        "test_judiVector.jl",
        "test_composite.jl",
        "test_judiWeights.jl",
        "test_judiWavefield.jl",
        "test_linear_operators.jl",
        "test_physicalparam.jl"]

devito = ["test_linearity.jl",
          "test_adjoint.jl",
          "test_all_options.jl",
          "test_jacobian.jl",
          "test_jacobian_extended.jl",
          "test_gradient_fwi.jl",
          "test_gradient_lsrtm.jl",
          "test_parallel.jl"]
        #   "test_gradient_twri.jl"]

extras = ["test_modeling.jl", "test_basics.jl", "test_linear_algebra.jl"]

# Basic JUDI objects tests, no Devito
if GROUP == "JUDI" || GROUP == "All"
    for t=base
        @time include(t)
    end
end

# Generic mdeling tests
if GROUP == "BASICS" || GROUP == "All"
    println("JUDI generic modelling tests")
    VERSION >= v"1.5" && push!(Base.ARGS, "-p 2")
    for t=extras
        @time include(t)
    end
end

# Isotropic Acoustic tests
if GROUP == "ISO_OP" || GROUP == "All"
    println("JUDI iso-acoustic operators tests (parallel)")
    # Basic test of LA/CG/LSQR needs
    VERSION >= v"1.5" && push!(Base.ARGS, "-p 2")
    for t=devito
        @time include(t)
    end
end

# Isotropic Acoustic tests with a free surface
if GROUP == "ISO_OP_FS" || GROUP == "All"
    println("JUDI iso-acoustic operators with free surface tests")
    push!(Base.ARGS, "--fs")
    for t=devito
        @time include(t)
    end
end

# Anisotropic Acoustic tests
if GROUP == "TTI_OP" || GROUP == "All"
    println("JUDI TTI operators tests")
    # TTI tests
    push!(Base.ARGS, "--tti")
    for t=devito
        @time include(t)
    end
end

# Anisotropic Acoustic tests with free surface
if GROUP == "TTI_OP_FS" || GROUP == "All"
    println("JUDI TTI operators with free surface tests")
    push!(Base.ARGS, "--tti")
    push!(Base.ARGS, "--fs")
    for t=devito
        @time include(t)
    end
end


text:
@testset "Test Matrix discrete method for FPDE" begin
    result = solve(0.7, 1.8, 1, 1, 3, 2, FPDEMatrixDiscrete())
    @test result ≈ [0.0  0.0; 0.0  0.0644409; 0.0  0.0]
end

@testset "Test FiniteDiffEx method" begin
    K = 1
    fdorder = 1.9
    dx = pi/2
    dt = 0.5
    n = 2
    xStart = 0
    xEnd = pi
    u0t = 0
    uendt = 0
    u0f(x) = sin(x)

    U=solve(fdorder, dx, dt, xStart, xEnd, n, K, u0t, uendt, u0f, FiniteDiffEx())
    @test isapprox(U, [ 0.0  1.0  1.22465e-16; 0.0   0.793379    0.0; 0.0   0.43766     0.0; 0.0   0.00221205  0.0; 0.0 -0.42797 0.0]; atol=1e-3)
end

@testset "Test FiniteDiffIm method" begin
    K = 1
    α = 0.5
    dx = pi/2
    dt = 0.5
    n = 2
    xStart = 0
    xEnd = pi
    u0t = 0
    uendt = 0
    u0f(x) = sin(x)

    U=solve(α, dx, dt, xStart, xEnd, n, K, u0t, uendt, u0f, FiniteDiffIm())
    @test isapprox(U, [ 0.0  1.0       1.22465e-16
    0.0  0.663152  0.0
    0.0  0.532299  0.0
    0.0  0.459938  0.0
    0.0  0.412321  0.0]; atol=1e-3)
end

@testset "Test Caputo Discrete for Advection Diffusion equation" begin
    fx0(x) = 0
    function fgz(q)
        x = q[1, 2];t=q[1, 1];a=q[1, 3]
        f = exp(x)*t^(6-a)/gamma(7-a)*720
        return f
    end
    function f0t(k)
        x=k[1, 1]
        a=k[1, 2]
        f=x^6
        return f
    end
    function flt(k)
        x = k[1,1];a=k[1,2]
        f = exp(1)*x^6
        return f
    end
    result = solve(3, 0.5, 1, 2, 1, 20, 1, 1, fx0, fgz, f0t, flt, ADV_DIF())
    @test isapprox(result, [0.0  0.0171843  1.06757
    0.0  0.0187449  1.1369
    0.0  0.0203056  1.20805
    0.0  0.0218648  1.2811
    0.0  0.0234204  1.35608
    0.0  0.0249696  1.43307
    0.0  0.0265091  1.51209
    0.0  0.028035   1.59319
    0.0  0.0295425  1.67637
    0.0  0.0310258  1.76166
    0.0  0.0324786  1.84904
    0.0  0.0338932  1.93849
    0.0  0.0352607  2.02997
    0.0  0.0365713  2.12342
    0.0  0.0378133  2.21874
    0.0  0.0389738  2.31582
    0.0  0.0400378  2.41452
    0.0  0.0409888  2.51465
    0.0  0.0418076  2.61599]; atol=1e-3)
end


text:
# SPDX-License-Identifier: MIT

#######################################################################
# Model definition

"""
    struct PLDA{T,D,Q1,Q2}
        hprior_within_class
        hprior_across_class
        trans_within_class
        trans_across_class
        λ
    end

Standard PLDA model. `Q1` is the dimension of the within class subspace
and `Q2` is the dimension across class subspace.
"""
struct PLDA{T,D,Q1,Q2}
    hprior_within_class::Normal{T,Q1}
    hprior_across_class::Normal{T,Q2}
    trans_within_class::AffineTransform{T,D,Q1}
    trans_across_class::AffineTransform{T,D,Q2}
    λ::BayesParam{Gamma{T}}
end

function PLDA(T::Type{<:AbstractFloat}; datadim, latentdim_within_class,
              latentdim_across_class, pstrength = 1e-3, W_MAP = false)

    Q1, Q2 = latentdim_within_class, latentdim_across_class

    hprior_w = Normal(zeros(T, Q1), Symmetric(Matrix{T}(I, Q1, Q1)))
    hprior_a = Normal(zeros(T, Q2), Symmetric(Matrix{T}(I, Q2, Q2)))

    trans_within = AffineTransform(T, outputdim = datadim, inputdim = Q1,
                                   pstrength = pstrength, W_MAP = W_MAP)
    trans_across = AffineTransform(T, outputdim = datadim, inputdim = Q2,
                                   pstrength = pstrength, W_MAP = W_MAP)
    λprior = Gamma{T}(pstrength, pstrength)
    λposterior = Gamma{T}(pstrength, pstrength)
    λ = BayesParam(λprior, λposterior)

    PLDA{T,datadim,Q1,Q2}(hprior_w, hprior_a, trans_within, trans_across, λ)
end

function PLDA(;datadim, latentdim_within_class, latentdim_across_class,
              pstrength = 1e-3, W_MAP = false)
    PLDA(Float64; datadim = datadim, latentdim = latentdim,
         pstrength = pstrength, W_MAP = W_MAP)
end

#######################################################################
# Estimate the latent variables

function (m::PLDA{T,D,Q})(X::AbstractVector, z, uposts::Dict) where {T,D,Q}
    # Project the means in the data space
    μ = mean.([uposts[zᵢ] for zᵢ in z]) |> m.trans_across_class

    S₁, S₂ = hstats(m.trans_within_class, X - μ)
    λ̄ = mean(m.λ.posterior)
    S₁, S₂ = λ̄*S₁, λ̄*S₂

    Λ₀ = inv(m.hprior_within_class.Σ)
    Λ₀μ₀ = Λ₀ * m.hprior_within_class.μ
    Σ = Symmetric(inv(Λ₀ + S₂))
    [Normal(Σ * (Λ₀μ₀ + mᵢ), Σ) for mᵢ in S₁]
end

function (m::PLDA{T,D,Q})(X::AbstractVector, hposts) where {T,D,Q}

    # Project the embeddings in the data space
    μ = mean.(hposts) |> m.trans_within_class

    S₁, S₂ = hstats(m.trans_across_class, X - μ)
    λ̄ = mean(m.λ.posterior)
    S₁, S₂ = λ̄*S₁, λ̄*S₂

    Λ₀ = inv(m.hprior_across_class.Σ)
    Λ₀μ₀ = Λ₀ * m..hprior_across_class.μ
    Σ = Symmetric(inv(Λ₀ + S₂))
    [Normal(Σ * (Λ₀μ₀ + mᵢ), Σ) for mᵢ in S₁]
end

(m::PLDA)(t::Tuple) = m(t...)

#######################################################################
# Pretty print

function Base.show(io::IO, ::MIME"text/plain", model::PLDA)
    println(io, typeof(model), ":")
    println(io, "  trans_within_class: $(typeof(model.trans_within_class))")
    println(io, "  trans_across_class: $(typeof(model.trans_across_class))")
    println(io, "  λ: $(typeof(model.λ))")
end

#######################################################################
# Log-likelihood

function _llh_d(::PLDA, x, Tŵ, Tv̂, Tλ, Th, Tu)
    λ, lnλ = Tλ
    h, hhᵀ = Th
    u, uuᵀ = Tu

    # Extract the bias parameter
    wwᵀ = Tŵ[2][1:end-1, 1:end-1]
    w = Tŵ[1][1:end-1]
    μ = Tŵ[1][end]
    μ² = Tŵ[2][end, end]

    vvᵀ = Tv̂[2][1:end-1, 1:end-1]
    v = Tv̂[1][1:end-1]
    a = Tv̂[1][end]
    a² = Tv̂[2][end, end]

    x̄ = dot(w, h) + μ + dot(v, u) + a
    lognorm = (-log(2π) + lnλ - λ*(x^2))/2
    K = λ * (x̄*x - dot(w, h)*μ - dot(v, u)*a)
    K += -λ * (dot(vec(hhᵀ), vec(wwᵀ)) + dot(vec(uuᵀ), vec(vvᵀ)) + μ² + a²)/2
    K += -λ * ((dot(w, h) + μ) * (dot(v, u) + a))

    lognorm + K
end

function _llh(m::PLDA, x, Tŵs, Tv̂s, Tλ, Th, Tu)
    f = (a,b) -> begin
        xᵢ, Tŵᵢ, Tv̂ᵢ = b
        a + _llh_d(m, xᵢ, Tŵᵢ, Tv̂ᵢ, Tλ, Th, Tu)
    end
    foldl(f, zip(x, Tŵs, Tv̂s), init = 0)
end

function loglikelihood(m::PLDA, X, z, uposts)
    hposts = (X, z, uposts) |> m
    _llh.(
        [m],
        X,
        [[gradlognorm(w.posterior, vectorize = false) for w in m.trans_within_class.W]],
        [[gradlognorm(w.posterior, vectorize = false) for w in m.trans_across_class.W]],
        [gradlognorm(m.λ.posterior, vectorize = false)],
        [gradlognorm(p, vectorize = false) for p in hposts],
        [gradlognorm(uposts[zᵢ], vectorize = false) for zᵢ in z]
   ) - kldiv.(hposts, [m.hprior_within_class])
end

loglikelihood(m::PLDA, t::Tuple, uposts) = loglikelihood(m, t..., uposts)

#######################################################################
# Update of the precision parameter λ

function _λstats_d(::PLDA, x::Real, Tŵ, Tv̂, Th, Tu)
    h, hhᵀ = Th
    u, uuᵀ = Tu

    # Extract the bias parameter
    w = Tŵ[1][1:end-1]
    wwᵀ = Tŵ[2][1:end-1, 1:end-1]
    μ = Tŵ[1][end]
    μ² = Tŵ[2][end, end]

    v = Tv̂[1][1:end-1]
    vvᵀ = Tv̂[2][1:end-1, 1:end-1]
    a = Tv̂[1][end]
    a² = Tv̂[2][end, end]

    x̄ = dot(w, h) + μ + dot(v, u) + a
    Tλ₁ = -.5 * x.^2 + (x̄*x - dot(w, h)*μ - dot(v, u)*a)
    Tλ₁ += -(dot(vec(hhᵀ), vec(wwᵀ)) + dot(vec(uuᵀ), vec(vvᵀ)) + μ² + a²)/2
    Tλ₁ += -((dot(w, h) + μ) * (dot(v, u) + a))
    Tλ₂ = 1/2

    vcat(Tλ₁, Tλ₂)
end

function _λstats(m::PLDA, x::AbstractVector, Tŵ, Tv̂, Th, Tu)
    sum(_λstats_d.([m], x, Tŵ, Tv̂, [Th], [Tu]))
end

function λstats(m::PLDA, X, z, uposts, hposts)
    Tŵ = [gradlognorm(w.posterior, vectorize = false) for w in m.trans_within_class.W]
    Tv̂ = [gradlognorm(w.posterior, vectorize = false) for w in m.trans_across_class.W]
    Th = [gradlognorm(p, vectorize = false) for p in hposts]
    Tu = [gradlognorm(uposts[zᵢ], vectorize = false) for zᵢ in z]
    sum(_λstats.([m], X, [Tŵ], [Tv̂], Th, Tu))
end

function update_λ!(m::PLDA, accstats)::Nothing
    η₀ = naturalparam(m.λ.prior)
    update!(m.λ.posterior, η₀ + accstats)
    nothing
end

#######################################################################
# Update of the bases W (within class)

function wstats_within_class(m::PLDA, X, z, uposts, hposts)
    λ, _ = gradlognorm(m.λ.posterior, vectorize = false)
    μ = mean.([uposts[zᵢ] for zᵢ in z]) |> m.trans_across_class
    S₁, S₂ = wstats(m.trans_within_class, X-μ, hposts)
    [λ*S₁, λ*S₂]
end

update_W_within_class!(m::PLDA, accstats) = update_W!(m.trans_within_class, accstats)

#######################################################################
# Update of the bases W (across class)

function wstats_across_class(m::PLDA, X, z, uposts, hposts)
    λ, _ = gradlognorm(m.λ.posterior, vectorize = false)
    μ = mean.(hposts) |> m.trans_within_class
    posts = [uposts[zᵢ] for zᵢ in z]
    S₁, S₂ = wstats(m.trans_across_class, X-μ, posts)
    [λ*S₁, λ*S₂]
end

update_W_across_class!(m::PLDA, accstats) = update_W!(m.trans_across_class, accstats)

#######################################################################
# Update of the mean embeddings

function ustats(m::PLDA, X, z, uposts, hposts) where {T,D,Q}
    # Project the embeddings in the data space
    μ = mean.(hposts) |> m.trans_within_class

    S₁, S₂ = hstats(m.trans_across_class, X - μ)
    λ̄ = mean(m.λ.posterior)
    S₁, S₂ = λ̄*S₁, λ̄*S₂
    ustats = Dict()
    for key in keys(uposts)
        idxs = z .== key
        N = sum(idxs) # Number of points to the cluster `key`
        if N == 0
            [zero(S₁[1]), 0*S₂]
        else
            ustats[key] = [sum(S₁[idxs]), N*S₂]
        end
    end
    ustats
end

function update_u!(m::PLDA, uposts, accstats)::Nothing where {T,D,Q}
    Λ₀ = inv(m.hprior_across_class.Σ)
    Λ₀μ₀ = Λ₀ * m.hprior_across_class.μ
    for key in keys(accstats)
        S₁, S₂ = accstats[key]
        Σ = Symmetric(inv(Λ₀ + S₂))
        uposts[key] = Normal(Σ * (Λ₀μ₀ + S₁), Σ)
    end
    nothing
end

#######################################################################
# Training

"""
    fit!(model, dataloader, uposts[, epochs = 1, callback = x -> x])

Fit a PPCA model to a data set by estimating the variational posteriors
over the parameters.
"""
function fit!(model::PLDA, dataloader, uposts; epochs = 1, callback = x -> x)

    @everywhere dataloader = $dataloader

    for e in 1:epochs
        # Propagate the model to all the workers
        @everywhere model = $model
        @everywhere uposts = $uposts

        ###############################################################
        # Step 1: update the posterior of the bases (within class)
        waccstats_within = @distributed (+) for batch in dataloader
            X, z = batch

            # E-step: estimate the posterior of the embeddings
            hposts = (X, z, uposts) |> model

            # Accumulate statistics for the bases w
            wstats_within_class(model, X, z, uposts, hposts)
        end
        update_W_within_class!(model, waccstats_within)

        # Propagate the model to all the workers
        @everywhere model = $model
        @everywhere uposts = $uposts

        ###############################################################
        # Step 2: update the class mean embeddings
        reducer = (a,b) -> merge(+,a,b)
        uaccstats = @distributed reducer for batch in dataloader
            X, z = batch

            # E-step: estimate the posterior of the embeddings
            hposts = (X, z, uposts) |> model

            # Accumulate statistics for the bases w
            ustats(model, X, z, uposts, hposts)
        end
        update_u!(model, uposts, uaccstats)

        # Propagate the model to all the workers
        @everywhere model = $model
        @everywhere uposts = $uposts

        ###############################################################
        # Step 3: update the posterior of the bases (across class)
        waccstats_across = @distributed (+) for batch in dataloader
            X, z = batch

            # E-step: estimate the posterior of the embeddings
            hposts = (X, z, uposts) |> model

            # Accumulate statistics for the bases w
            wstats_across_class(model, X, z, uposts, hposts)
        end
        update_W_across_class!(model, waccstats_across)

        # Propagate the model to all the workers
        @everywhere model = $model
        @everywhere uposts = $uposts

        ###############################################################
        # Step 4: update the posterior of the precision λ
        λaccstats = @distributed (+) for batch in dataloader
            X, z = batch

            # E-step: estimate the posterior of the embeddings
            hposts = (X, z, uposts) |> model

            # Accumulate statistics for λ
            λstats(model, X, z, uposts, hposts)
        end
        update_λ!(model, λaccstats)

        # Notify the caller
        callback(e)
    end
end

max_stars_repo_name: BUTSpeechFIT/PPCA

max_issues_repo_name: BUTSpeechFIT/PPCA

max_forks_repo_name: BUTSpeechFIT/PPCA


text:
include("INCLUDEME.jl")

using Yao, Circuit, UnicodePlots, GradOptim, Utils
import Kernels

function train!(qcbm::QCBM, ptrain, optim; learning_rate=0.1, maxiter=100)
    initialize!(qcbm)
    kernel = Kernels.RBFKernel(nqubits(qcbm), [0.25], false)
    history = Float64[]

    for i = 1:maxiter
        grad = gradient(qcbm, kernel, ptrain)
        curr_loss = loss(qcbm, kernel, ptrain)
        push!(history, curr_loss)
        println(i, " step, loss = ", curr_loss)

        # Warn: we need a primitive block to enable
        # BLAS here.
        params = parameters(qcbm)
        update!(params, grad, optim)
        dispatch!(qcbm, params)
    end
    history
end

# parameters
n = 6
maxiter = 20

#data to learn
pg = gaussian_pdf(n, 2^5-0.5, 2^4)
fig = lineplot(0:1<<n - 1, pg)
display(fig)

# solver setup
qcbm = QCBM{n, 10}(get_nn_pairs(n))
optim = Adam(lr=0.1);

his = train!(qcbm, pg, optim, maxiter=maxiter);

# analyze result
display(lineplot(his, title = "loss"))
psi = qcbm()
p = statevec(psi) .|> abs2
lineplot!(fig, p, color=:yellow, name="trained")
display(fig)

max_stars_repo_name: Roger-luo/QCBM

max_issues_repo_name: Roger-luo/QCBM

max_forks_repo_name: Roger-luo/QCBM