text:
# include("simulate.jl")

using ForwardDiff
using FiniteDifferences
using LazyArrays

## helper functions for matrix derivatives
function kronSum(A,B)
  # kronecker sum of 2 arrays
  out = kronProd(A, Matrix(I,(size(B)))) + kronProd(Matrix(I,(size(A))),B)
  return out
end

function kronProd(A,B)
    K = ApplyArray(kron,A[:,:],B[:,:]);
    out = Matrix{Float64}(undef,size(K)...); copyto!(out,K);
    return out
end

function createCommMat(M)
  # commutation matrix
  r,m = size(M);
  K = zeros(m*r,m*r);
  for i=1:r
      for j=1:m
          ei = 1.0I(r)[:,i] # ith-canonical unit vector of dimension r
          ej = 1.0I(m)[:,j] # jth-canonical unit vector of dimension m
          K += kronProd(ei*permutedims(ej), ej*permutedims(ei))
      end
  end
  return K
end

## main function for linearize
function linearize(xv,uv,j,GravityInInertial)
    function Mdiff1(x)
        # First derivatives of mass matrix and powers (1/2, -1/2, -1)
        M = assembleM(x,j)
        Ms = real(sqrt(M))
        M_inv = inv(M)
        Ms_inv = real(sqrt(M_inv))

        dM = ForwardDiff.jacobian(z -> assembleM(z,j),x)
        dM_inv = -(kronProd(permutedims(M_inv),M_inv))*dM;
        dMs = inv(kronSum(permutedims(Ms), Ms))*dM
        dMs_inv = -kronProd(permutedims(Ms_inv),Ms_inv)*(dMs);
        return M, Ms, Ms_inv, dM, dMs, dMs_inv, M_inv, dM_inv
    end
    M, Ms, Ms_inv, dM, dMs, dMs_inv, M_inv, dM_inv = Mdiff1(xv);

    function Adiff1(x)
        A, _ = Ab_VecOfMat(x,j);
        dA = ForwardDiff.jacobian(z->Ab_VecOfMat(z,j)[1],x);
        return A, dA
    end
    A, dA = Adiff1(xv);

    function bdiff1(x)
        _, b = Ab_VecOfMat(x,j);
        db = ForwardDiff.jacobian(z->Ab_VecOfMat(z,j)[2],x);
        return b, db
    end
    b, db = bdiff1(xv);

    function Gpdiff1(x)
        function t1diff1(x)
            # t1 = A'*(AM)^{-1}*A';
            sz_A1 = size(A,1)
            t1 = permutedims(A)*inv(A*inv(M)*permutedims(A))

            dA_tr = createCommMat(A)*dA;

            Y = A*M_inv*permutedims(A); invY = inv(Y);
            dY1 = kronProd(A*M_inv,1.0I(sz_A1))*dA;
            dY2 = kronProd(A,A)*dM_inv;
            dY3 = kronProd(1.0I(sz_A1), A*M_inv)*dA_tr;
            dY = dY1 + dY2 + dY3;
            dY_inv = -kronProd(permutedims(invY),invY)*dY;

            dt1_1 = kronProd(permutedims(invY),1.0I(size(A,2)))*dA_tr;
            dt1_2 = kronProd(1.0I(sz_A1),permutedims(A))*dY_inv;
            dt1 = dt1_1 + dt1_2;
            return t1, dt1
        end
        G = A*Ms_inv;
        Gp = permutedims(G)*((G*permutedims(G))\I(size(G,1)))

        t1, dt1 = t1diff1(x)
        dGp1 = kronProd(permutedims(t1),1.0I(size(M,1)))*dMs_inv

        sz_t2 =size(t1,2)
        t2 = kronProd(1.0I(sz_t2), Ms_inv)
        dGp2 = t2*dt1;

        dGp = dGp1 + dGp2
    return Gp, dGp
    end
    Gp, dGp = Gpdiff1(xv);

    function Fudiff1(x,u)
        Fu = assembleF(x,u,j,g);
        dFu = ForwardDiff.jacobian(y->assembleF(y,u,j,g),x);
        dFu_u = ForwardDiff.jacobian(y->assembleF(x,y,j,g),u);
        return Fu, dFu, dFu_u
    end
    Fu, dFu, dFu_u = Fudiff1(xv,uv);

    function hdiff1(x,u)
        # has derivatives wrt both x,u
        h = b - A*inv(M)*Fu;

        # dh
        term1_1 = kronProd(permutedims(M_inv*Fu), 1.0I(size(A,1)));
        term1 = term1_1*dA;
        term2 = kronProd(permutedims(Fu), A)*dM_inv;
        term3 = A*M_inv*dFu;
        dh_x = db - (term1 + term2 + term3);

        # dh_u
        dh_u = -A*M_inv*dFu_u;

        return h, dh_x, dh_u
    end
    h, dh_x, dh_u = hdiff1(xv,uv);

    function Fcdiff1(x,u)
        Fc = Ms*Gp*h

        dFc1 = kronProd(permutedims(Gp*h),I(size(M,1)))*dMs
        dFc2 = kronProd(permutedims(h),Ms)*dGp
        dFc3 = Ms*Gp*dh_x;

        dFc = dFc1 + dFc2 + dFc3;

        dFc_u = Ms*Gp*dh_u;
        return Fc, dFc, dFc_u
    end
    Fc, dFc, dFc_u = Fcdiff1(xv,uv);

    function accdiff1(x,u)
        # derivative of qdd w.r.t x
        dacc1 = kronProd(permutedims(Fc+Fu),1.0I(size(M_inv,1)))*dM_inv
        dacc2 = M_inv*(dFc + dFu)
        dqdd_x = dacc1 + dacc2

        return dqdd_x
    end
    dqdd_x = accdiff1(xv,uv);
    dqdd_u = M_inv*(dFc_u + dFu_u);

    function fxdotDiff(x,u)
        len_x = length(x); len_u = length(u);
        df_x = zeros(len_x,len_x)

        df_u = zeros(len_x, len_u);

        for k=1:length(j)
            df_x[14*k+1:14*k+7,14*k+8:14*k+14] = 1.0*I(7)
            df_x[14*k+8:14*k+14,:] = dqdd_x[7*(k-1)+1:7*k,:]

            df_u[14*k+8:14*k+14,:] = dqdd_u[7*(k-1)+1:7*k,:];
        end
        df_x_woIn = df_x[15:end,15:end]; # without the inertial frame
        df_u_woIn = df_u[15:end,7:end];  # without the inertial frame
        return df_x_woIn, df_u_woIn
    end

    dfx_x, dfx_u = fxdotDiff(xv,uv);
    return dfx_x, dfx_u
end
## Test for different robots
# x0Orig, u0Orig = getXU_0(j);
# fdJ_A, fdJ_B = linearize(x0Orig, u0Orig, j, g);
#
# m = central_fdm(5,1); # first order differential using fifth order central difference method
# finJ_A = FiniteDifferences.jacobian(m, z->fxdot(z,u0Orig,j,g),x0Orig)[1];
# finJ_B = FiniteDifferences.jacobian(m, z->fxdot(x0Orig,z,j,g),u0Orig)[1];
#
# println("err_A = ", norm(fdJ_A - finJ_A))
# println("err_B = ", norm(fdJ_B - finJ_B))

max_stars_repo_name: isrlab/Multibody-Dynamics

max_issues_repo_name: isrlab/Multibody-Dynamics

max_forks_repo_name: isrlab/Multibody-Dynamics


text:
# Run using: mpiexec -n 2 julia parallel_solve.jl

using MiniPETSc

vec = PetscVec()

setSize!(vec, n_local=(Int32)(5), n_global=(Int32)(10))

if MPI.Comm_rank(MPI.COMM_WORLD) == 0
    vec[1:5] = (Float64)[1, 2, 3, 4, 5]
else
    vec[6:10] = (Float64)[6, 7, 8, 9, 10]
end

assemble!(vec)
viewVec(vec)

sol = PetscVec()
setSize!(sol, n_local=(Int32)(5), n_global=(Int32)(10))
assemble!(sol)
viewVec(sol)

mat = PetscMat()
setSize!(mat, m_local=(Int32)(5), n_local=(Int32)(5))
setPreallocation!(mat, (Int32)[1,1,1,1,1], (Int32)[0,0,0,0,0])

if MPI.Comm_rank(MPI.COMM_WORLD) == 0
    mat[1,1] = (Float64)(1)
    mat[2,2] = (Float64)(2)
    mat[3,3] = (Float64)(3)
    mat[4,4] = (Float64)(4)
    mat[5,5] = (Float64)(5)
else
    mat[6,6] = (Float64)(6)
    mat[7,7] = (Float64)(7)
    mat[8,8] = (Float64)(8)
    mat[9,9] = (Float64)(9)
    mat[10,10] = (Float64)(10)
end

assemble!(mat)
viewMat(mat)

ksp = PetscKSP()

setOperators(ksp, mat)

solve!(ksp, vec, sol)

viewVec(sol)

ext: jl


text:
using CrossEntropyVariants
using Distributions

## Gaussian process surrogate model example

f = sierra # objective function
𝛍 = [0, 0]
𝚺 = [200 0; 0 200]
𝐌 = MvNormal(𝛍, 𝚺) # proposal distribution

(𝐌, bestₓ, bestᵥ) = ce_surrogate(f, 𝐌; plot=true)

## Radial basis function example

f = paraboloid # objective function
𝐌 = MvNormal([0, 0], [200 0; 0 200]) # proposal distribution

(𝐌, bestₓ, bestᵥ) = ce_surrogate(f, 𝐌; plot=true, basis=:squared)


text:
# put units tests.

# something like
#using DCEMRI
#
#@test_approx_eq validate6() == mean error
#@test_approx_eq validate4() == mean error


text:
"""
Data structure for holding a routing in a path-based manner.
"""
struct Routing
    data::RoutingData
    path_flows::Dict{Edge{Int}, Dict{Int, Float64}} # Demand -> path -> flow.
    edge_flows::Dict{Edge{Int}, Dict{Edge{Int}, Float64}} # Demand -> edge -> flow.
    demands::Vector{Edge} # List of demand; the only goal of this vector is to fix
    # the order of the demands (i.e. have demand IDs).
    paths::Vector{Vector{Edge}} # A path is a list of edges. This vector stores
    # a mapping from integers (indices in the vector) to paths.

    function Routing(data::RoutingData, routing::AbstractMatrix{Float64})
        if data.model_type.type == FlowFormulation()
            # If the model is not path-based, translate it.
            path_flows, edge_flows, paths = flow_routing_to_path(data, routing)
            return new(data, path_flows, edge_flows, collect(demands(data)), paths)
        else
            @assert data.model_type.type == PathFormulation()
            # Otherwise, much simpler transformation: the solution is already written in terms of paths.
            path_flows, edge_flows = path_routing_to_flow(data, routing)
            return new(data, path_flows, edge_flows, collect(demands(data)), data.paths_edges)
        end
    end
end

function routing_to_matrix(r::Routing)
    routing_matrix = zeros(length(r.demands), length(r.paths))
    for demand_id in 1:length(r.demands)
        demand = r.demands[demand_id]
        for (path, weight) in r.path_flows[demand]
            routing_matrix[demand_id, path] = weight
        end
    end
    return routing_matrix
end

function total_flow_in_edge(r::Routing, e::Edge, dm)
    # dm is an optimisation variable, as the routing is fixed.
    return total_flow_in_edge(r, e, dm, r.data.model_type.type)
end

function total_flow_in_edge(r::Routing, e::Edge, dm, ::FormulationType; ε::Float64=1.e-3)
    # dm is an optimisation variable, as the routing is fixed.
    # Works for both flow-based and path-based formulations.
    # Cannot use sum(), as the collection is sometimes empty.
    flow = 0.0
    for d in demands(r.data)
        if e in keys(r.edge_flows[d]) && abs(r.edge_flows[d][e]) > ε
            flow += r.edge_flows[d][e] * dm[d]
        end
    end
    return flow
end

"""
Data structure holding results of the computations for a routing.
All information from intermediate iterations are kept within this object.

  * `data`: a pointer to the `RoutingData` object that was used for the
    computations.
  * `result`: an `MOI.TerminationStatusCode` indicating the result of the
    optimiser. This is usually `MOI.OPTIMAL`, but other codes indicate why
    the solver stopped (`MOI.SLOW_PROGRESS`, `MOI.INFEASIBLE`, etc.).
    This code is not supposed to be the result of the last iteration, but
    a summary of the termination of the process: for instance, if the last
    iteration gives an `MOI.INFEASIBLE` code at the last iteration due to
    numerical errors, the whole process may show `MOI.SLOW_PROGRESS`.
  * `n_iter`: the number of iterations (the exact meaning of this field depends
    on the algorithm). Iterations are supposed to be spent in this package,
    improving/finding a network routing; this should not be a copy of a field
    from the underlying solver (e.g., simplex iterations).
    (This field is computed and not explicitly stored in the object.)
  * `n_matrices`: the number of traffic matrices added into the formulation. It
    may be equal to the number of iterations, depending on the algorithm.
    (This field is computed and not explicitly stored in the object.)
  * `n_cuts`: the number of added constraints.
  * `n_columns`: the number of added columns, for column-generation based
    algorithms. Its value should be 0 for other algorithms.

Three different timings are measured for the whole execution (always in
milliseconds):

  * `time_precompute_ms`: time to prepare the data structures before the actual
    computations. For instance, it may account for detection of loops or
    unroutable demands; for column-generation algorithms, it includes the time
    to generate the first few paths (i.e. before any pricing process can take
    place). This parameter takes precedence on the one from the given
	`RoutingData` object (i.e. it may include more operations).
  * `time_create_master_model_ms`: time to create the mathematical formulation
    of the master problem (or the only problem, if the algorithm is not
	iterative). In particular, this does not include instanciation of
	subproblems.
  * `total_time_solve_ms`: total time spent in solving.
    (This field is computed and not explicitly stored in the object.)
  * `time_solve_ms`: time to compute the routing, one value per iteration.
  * `time_intermediate_export_ms`: time to export the intermediate results,
    when requested, one value per iteration where there is export.
  * `time_final_export_ms`: time to export the final solution, when requested.
  * `total_time_export_ms`: total time spent in exporting.
    (This field is computed and not explicitly stored in the object.)
  * `total_time_ms`: total time spent in the whole process (precomputing,
    solving, exporting).
    (This field is computed and not explicitly stored in the object.)

The following vectors contain information about the execution of the algorithm.
Not all vectors have as many entries as iterations, even though it is expected
to be the most likely case.

  * `objectives`: the value of the objective function that is being optimised,
    at most one per iteration.
  * `matrices`: the demand matrices generated during the execution, when their
    output is requested. It may contain a single matrix if the algorithm does
    not generate new matrices during its execution. There may be no such matrix
    at the last iteration. There may be several matrices for some iterations.
  * `routings`: the various routings generated during the execution. There must
    be at most one routing per iteration, except when requested.
  * `master_model`: the final optimisation model, with all generated cuts and
    columns. Solving it should give the same solution as `routings[end]`.
"""
struct RoutingSolution
    data::RoutingData
    result::MOI.TerminationStatusCode
    n_cuts::Int
    n_columns::Int

    time_precompute_ms::Float64
	time_create_master_model_ms::Float64
    time_solve_ms::Dict{Int, Float64}
    time_intermediate_export_ms::Dict{Int, Float64}
    time_final_export_ms::Float64

    objectives::Dict{Int, Float64}
    matrices::Dict{Int, Vector{Dict{Edge{Int}, Float64}}}
    routings::Dict{Int, Routing}
    master_model::RoutingModel
end

function Base.getproperty(obj::RoutingSolution, sym::Symbol)
    if sym === :n_matrices
        return length(obj.matrices)
    elseif sym === :n_iter
        return length(obj.time_solve_ms)
    elseif sym === :total_time_solve_ms
        return sum(values(obj.time_solve_ms))
    elseif sym === :total_time_export_ms
        return sum(values(obj.time_intermediate_export_ms)) + obj.time_final_export_ms
    elseif sym === :total_time_ms
        return obj.time_precompute_ms + obj.time_create_master_model_ms + sum(values(obj.time_solve_ms)) + sum(values(obj.time_intermediate_export_ms)) + obj.time_final_export_ms
    end
    return getfield(obj, sym)
end

_parse_routingsolution_input(x::T) where T = Dict{Int, T}(1 => x)
_parse_routingsolution_input(x::Vector{T}) where T = Dict{Int, T}(i => x[i] for i in 1:length(x))
_parse_routingsolution_input(x::Dict{Int, T}) where T = x

_parse_routingsolution_matrices(x::Dict{Edge{Int}, Float64}) = length(x) == 0 ? Dict{Int, Vector{Dict{Edge{Int}, Float64}}}() : Dict(1 => [x])
_parse_routingsolution_matrices(x::Vector{Dict{Edge{Int}, Float64}}) = length(x) == 0 ? Dict{Int, Vector{Dict{Edge{Int}, Float64}}}() : Dict(1 => x)
_parse_routingsolution_matrices(x::Dict{Int, Vector{Dict{Edge{Int}, Float64}}}) = x

function RoutingSolution(data::RoutingData;
                         result::MOI.TerminationStatusCode=MOI.OPTIMAL,
                         n_cuts::Int=0,
                         n_columns::Int=0,
                         time_precompute_ms::Float64=0.0,
						 time_create_master_model_ms::Float64=0.0,
						 time_solve_ms::Union{Float64, Vector{Float64}, Dict{Int, Float64}}=0.0,
						 time_intermediate_export_ms::Union{Float64, Vector{Float64}, Dict{Int, Float64}}=Dict{Int, Float64}(),
						 time_final_export_ms::Float64=0.0,
						 objectives::Union{Float64, Vector{Float64}, Dict{Int, Float64}}=error("Missing parameter `objectives` when building a `RoutingSolution` object"),
						 matrices::Union{Dict{Edge{Int}, Float64}, Vector{Dict{Edge{Int}, Float64}}, Dict{Int, Vector{Dict{Edge{Int}, Float64}}}}=Dict{Edge{Int}, Float64}(),
						 routings::Union{Routing, Vector{Routing}, Dict{Int, Routing}}=error("Missing parameter `routings` when building a `RoutingSolution` object"),
						 master_model::RoutingModel=error("Missing parameter `master_model` when building a `RoutingSolution` object")
						)
    return RoutingSolution(data,
	                       result,
						   n_cuts,
						   n_columns,
	                       time_precompute_ms,
						   time_create_master_model_ms,
						   _parse_routingsolution_input(time_solve_ms),
						   _parse_routingsolution_input(time_intermediate_export_ms),
						   time_final_export_ms,
						   _parse_routingsolution_input(objectives),
						   _parse_routingsolution_matrices(matrices),
						   _parse_routingsolution_input(routings),
						   master_model)
end

function flow_routing_to_path(data::RoutingData, routing::AbstractMatrix{Float64}; demand=nothing, ε::Float64=1.e-5)
    if data.model_type.type != FlowFormulation()
        error("This function can only be used for flow-based formulations; it makes no sense for formulation type $(data.model_type)")
    end

    path_flows = Dict{Edge{Int}, Dict{Int, Float64}}() # demand -> path index -> flow
    edge_flows = Dict{Edge{Int}, Dict{Edge{Int}, Float64}}() # demand -> edge -> flow
    paths = Vector{Edge}[] # list of paths (indexes match those of path_flows)

    for d in demands(data)
        # Prepare to store the results for this demand.
        path_flows[d] = Dict{Int, Float64}()
        edge_flows[d] = Dict{Edge{Int}, Float64}()

        # Copy the flows into a new graph (with only those edges).
        # Don't use a weighted graph: removing edges does not always work
        # as expected.
        ug = MetaGraph(nv(graph(data)))
        for e in edges(data)
            if has_edge(ug, src(e), dst(e))
                continue
            end

            flow_a = try; routing[d, e]; catch; 0.0; end
            flow_b = try; routing[d, reverse(e)]; catch; 0.0; end
            flow = abs(flow_a - flow_b)
            if flow > ε
                add_edge!(ug, src(e), dst(e), :weight, flow)
            end
        end

        # Compute paths and remove the corresponding flow as long as there is
        # something in there.
        weight_e(e) = get_prop(ug, e, :weight)
        while ne(ug) > 0
            # Get a path arbitrarily (easiest to compute: shortest one)
            sp = desopo_pape_shortest_paths(ug, src(d))
            path_vertices = enumerate_paths(sp, dst(d))
            path = map(Edge, zip(path_vertices[1:end-1], path_vertices[2:end]))
            if length(path) == 0
                break
            end
            path_weight = minimum(weight_e(e) for e in path)

            # Update the data structures with this new path.
            push!(paths, path)
            path_idx = length(paths)
            path_flows[d][path_idx] = path_weight
            for e in path # TODO: DataStructures.jl's DefaultDict? Potential problems when returning this instead of a Dict?
                if ! haskey(path_flows[d], e)
                    edge_flows[d][e] = 0.0
                end
                edge_flows[d][e] += path_weight
            end

            # Remove the path from the graph.
            for e in path
                new_weight = weight_e(e) - path_weight
                if new_weight > ε
                    set_prop!(ug, e, :weight, new_weight)
                else
                    rem_edge!(ug, e)
                end
            end
        end
    end

    return path_flows, edge_flows, paths
end

function path_routing_to_flow(data::RoutingData, routing::AbstractMatrix{Float64}; ε::Float64=1.e-5)
    if data.model_type.type != PathFormulation()
        error("This function can only be used for path-based formulations; it makes no sense for formulation type $(data.model_type)")
    end

    # Sum over all paths that cross a given edge to get the total weight of each edge, demand per demand.
    path_flows = Dict{Edge{Int}, Dict{Int, Float64}}()

    for k in eachindex(routing)
        demand = k[1]
        path = k[2]

        if ! (demand in keys(path_flows))
            path_flows[demand] = Dict{Int, Float64}()
        end

        if abs(routing[k]) > ε
            path_flows[demand][path] = routing[k]
        end
    end

    # Sum the flows over the demands.
    edge_flows = Dict{Edge{Int}, Dict{Edge{Int}, Float64}}()
    for demand in demands(data)
        edge_flows[demand] = Dict{Edge{Int}, Float64}()
        for edge in edges(data)
            total_weight = zero(Float64)
            for p_id in find_path_ids_with_edge(data, edge)
                if data.path_id_to_demand[p_id] == demand
                    total_weight += routing[demand, p_id]
                end
            end

            if abs(total_weight) > ε
                edge_flows[demand][edge] = total_weight
            end
        end
    end

    return path_flows, edge_flows
end


text:
using DataFrames, Plots, FileIO
# Dimensions of the logo .png
L = [480, 1090]
# The logo dewetting data
logo_data = load("E:\\JuliaStuff\\Data_PhD\\Tutorials\\Logo_dewetting.jld2") |> DataFrame

heatmap(reshape(logo_data[!, :h_10000],480, 1090),
        c=:viridis,
        aspect_ratio=1,
        grid=:none,
        tickfont = (16, "Arial"),	# tick font and size
        guidefont = (18, "Arial"),	# label font and size
        ylim=(0, 480),
        xlim=(0,1090),
		label="t = 10000Δt", 					# labels
		xlabel="x/Δx", 						# x-axis label
		ylabel="y/Δx",					# y-axis label
        colorbar_title="thickness",
        framestyle = :none,
        )

max_stars_repo_name: Zitzeronion/NotebooksnStuff

max_issues_repo_name: Zitzeronion/NotebooksnStuff

max_forks_repo_name: Zitzeronion/NotebooksnStuff


text:
using SequentialMonteCarlo
using RNGPool
import SMCExamples.Lorenz96: defaultLorenzModel, LorenzTheta, makeLorenzModel
using Test

setRNGs(0)

N = 2^10
n = 10
nt = Threads.nthreads()
model, theta, ys = defaultLorenzModel(8, 100)

smcio = SMCIO{model.particle, model.pScratch}(N, n, nt, false)
smc!(model, smcio)
lzh1 = copy(smcio.logZhats)

theta2 = LorenzTheta(theta.σ, theta.F, theta.δ, theta.Δ, theta.steps*2)
model = makeLorenzModel(theta2, ys)
smc!(model, smcio)

@test smcio.logZhats ≈ lzh1 atol=1.0


text:
# A BallTree (also called Metric tree) is a tree that is created
# from successively splitting points into surrounding hyper spheres
# which radius are determined from the given metric.
# The tree uses the triangle inequality to prune the search space
# when finding the neighbors to a point,
struct BallTree{V <: AbstractVector,N,T,M <: Metric} <: NNTree{V,M}
    data::Vector{V}
    hyper_spheres::Vector{HyperSphere{N,T}} # Each hyper sphere bounds its children
    indices::Vector{Int}                  # Translates from tree index -> point index
    metric::M                             # Metric used for tree
    tree_data::TreeData                   # Some constants needed
    reordered::Bool                       # If the data has been reordered
end

# When we create the bounding spheres we need some temporary arrays.
# We create a type to hold them to not allocate these arrays at every
# function call and to reduce the number of parameters in the tree builder.
struct ArrayBuffers{N,T <: AbstractFloat}
    center::MVector{N,T}
end

function ArrayBuffers(::Type{Val{N}}, ::Type{T}) where {N, T}
    ArrayBuffers(zeros(MVector{N,T}))
end

"""
    BallTree(data [, metric = Euclidean(), leafsize = 10]) -> balltree

Creates a `BallTree` from the data using the given `metric` and `leafsize`.
"""
function BallTree(data::Vector{V},
                  metric::M = Euclidean();
                  leafsize::Int = 10,
                  reorder::Bool = true,
                  storedata::Bool = true,
                  reorderbuffer::Vector{V} = Vector{V}()) where {V <: AbstractArray, M <: Metric}
    reorder = !isempty(reorderbuffer) || (storedata ? reorder : false)

    tree_data = TreeData(data, leafsize)
    n_d = length(V)
    n_p = length(data)

    array_buffs = ArrayBuffers(Val{length(V)}, get_T(eltype(V)))
    indices = collect(1:n_p)

    # Bottom up creation of hyper spheres so need spheres even for leafs)
    hyper_spheres = Vector{HyperSphere{length(V),eltype(V)}}(undef, tree_data.n_internal_nodes + tree_data.n_leafs)

    if reorder
        indices_reordered = Vector{Int}(undef, n_p)
        if isempty(reorderbuffer)
            data_reordered = Vector{V}(undef, n_p)
        else
            data_reordered = reorderbuffer
        end
    else
        # Dummy variables
        indices_reordered = Vector{Int}()
        data_reordered = Vector{V}()
    end

    if n_p > 0
        # Call the recursive BallTree builder
        build_BallTree(1, data, data_reordered, hyper_spheres, metric, indices, indices_reordered,
                       1,  length(data), tree_data, array_buffs, reorder)
    end

    if reorder
       data = data_reordered
       indices = indices_reordered
    end

    BallTree(storedata ? data : similar(data, 0), hyper_spheres, indices, metric, tree_data, reorder)
end

 function BallTree(data::Matrix{T},
                  metric::M = Euclidean();
                  leafsize::Int = 10,
                  storedata::Bool = true,
                  reorder::Bool = true,
                  reorderbuffer::Matrix{T} = Matrix{T}(undef, 0, 0)) where {T <: AbstractFloat, M <: Metric}
    dim = size(data, 1)
    npoints = size(data, 2)
    points = copy_svec(T, data, Val(dim))
    if isempty(reorderbuffer)
        reorderbuffer_points = Vector{SVector{dim,T}}()
    else
        reorderbuffer_points = copy_svec(T, reorderbuffer, Val(dim))
    end
    BallTree(points, metric, leafsize = leafsize, storedata = storedata, reorder = reorder,
            reorderbuffer = reorderbuffer_points)
end

# Recursive function to build the tree.
function build_BallTree(index::Int,
                        data::Vector{V},
                        data_reordered::Vector{V},
                        hyper_spheres::Vector{HyperSphere{N,T}},
                        metric::Metric,
                        indices::Vector{Int},
                        indices_reordered::Vector{Int},
                        low::Int,
                        high::Int,
                        tree_data::TreeData,
                        array_buffs::ArrayBuffers{N,T},
                        reorder::Bool) where {V <: AbstractVector, N, T}

    n_points = high - low + 1 # Points left
    if n_points <= tree_data.leafsize
        if reorder
            reorder_data!(data_reordered, data, index, indices, indices_reordered, tree_data)
        end
        # Create bounding sphere of points in leaf node by brute force
        hyper_spheres[index] = create_bsphere(data, metric, indices, low, high, array_buffs)
        return
    end

    # Find split such that one of the sub trees has 2^p points
    # and the left sub tree has more points
    mid_idx = find_split(low, tree_data.leafsize, n_points)

    # Brute force to find the dimension with the largest spread
    split_dim = find_largest_spread(data, indices, low, high)

    # Sort the data at the mid_idx boundary using the split_dim
    # to compare
    select_spec!(indices, mid_idx, low, high, data, split_dim)

    build_BallTree(getleft(index), data, data_reordered, hyper_spheres, metric,
                   indices, indices_reordered, low, mid_idx - 1,
                   tree_data, array_buffs, reorder)

    build_BallTree(getright(index), data, data_reordered, hyper_spheres, metric,
                  indices, indices_reordered, mid_idx, high,
                  tree_data, array_buffs, reorder)

    # Finally create bounding hyper sphere from the two children's hyper spheres
    hyper_spheres[index]  =  create_bsphere(metric, hyper_spheres[getleft(index)],
                                            hyper_spheres[getright(index)],
                                            array_buffs)
end

function _knn(tree::BallTree,
              point::AbstractVector,
              best_idxs::Vector{Int},
              best_dists::Vector,
              skip::Function)
    knn_kernel!(tree, 1, point, best_idxs, best_dists, skip)
    return
end

function knn_kernel!(tree::BallTree{V},
                           index::Int,
                           point::AbstractArray,
                           best_idxs::Vector{Int},
                           best_dists::Vector,
                           skip::F) where {V, F}
    @NODE 1
    if isleaf(tree.tree_data.n_internal_nodes, index)
        add_points_knn!(best_dists, best_idxs, tree, index, point, true, skip)
        return
    end

    left_sphere = tree.hyper_spheres[getleft(index)]
    right_sphere = tree.hyper_spheres[getright(index)]

    left_dist = max(zero(eltype(V)), evaluate(tree.metric, point, left_sphere.center) - left_sphere.r)
    right_dist = max(zero(eltype(V)), evaluate(tree.metric, point, right_sphere.center) - right_sphere.r)

    if left_dist <= best_dists[1] || right_dist <= best_dists[1]
        if left_dist < right_dist
            knn_kernel!(tree, getleft(index), point, best_idxs, best_dists, skip)
            if right_dist <=  best_dists[1]
                 knn_kernel!(tree, getright(index), point, best_idxs, best_dists, skip)
             end
        else
            knn_kernel!(tree, getright(index), point, best_idxs, best_dists, skip)
            if left_dist <=  best_dists[1]
                 knn_kernel!(tree, getleft(index), point, best_idxs, best_dists, skip)
            end
        end
    end
    return
end

function _inrange(tree::BallTree{V},
                  point::AbstractVector,
                  radius::Number,
                  idx_in_ball::Vector{Int}) where {V}
    ball = HyperSphere(convert(V, point), convert(eltype(V), radius))  # The "query ball"
    inrange_kernel!(tree, 1, point, ball, idx_in_ball)  # Call the recursive range finder
    return
end

function inrange_kernel!(tree::BallTree,
                         index::Int,
                         point::AbstractVector,
                         query_ball::HyperSphere,
                         idx_in_ball::Vector{Int})
    @NODE 1

    if index > length(tree.hyper_spheres)
        return
    end

    sphere = tree.hyper_spheres[index]

    # If the query ball in the bounding sphere for the current sub tree
    # do not intersect we can disrecard the whole subtree
    if !intersects(tree.metric, sphere, query_ball)
        return
    end

    # At a leaf node, check all points in the leaf node
    if isleaf(tree.tree_data.n_internal_nodes, index)
        add_points_inrange!(idx_in_ball, tree, index, point, query_ball.r, true)
        return
    end

    # The query ball encloses the sub tree bounding sphere. Add all points in the
    # sub tree without checking the distance function.
    if encloses(tree.metric, sphere, query_ball)
         addall(tree, index, idx_in_ball)
    else
        # Recursively call the left and right sub tree.
        inrange_kernel!(tree,  getleft(index), point, query_ball, idx_in_ball)
        inrange_kernel!(tree, getright(index), point, query_ball, idx_in_ball)
    end
end


text:
function initial_ub(X::Array{Float64,2}, Y::YVector, sparsity::Int, γ::Float64)
    αstar = -SubsetSelectionCIO.sparse_inverse(OLS(), Y, X[:, 1:sparsity], γ)
    X_sparse = @view(X[:, 1:sparsity])
    ax_sparse = X_sparse' * αstar
    getlb(Y, αstar, ax_sparse, γ)
end

function getlb(Y::YVector, α::Vector{Float64}, ax_sparse::Vector{Float64}, γ::Float64)
    -0.5 * dot(α, α) + dot(Y, α) - γ * 0.5 * sum(ax_sparse.^2)
end

function getslope!(∇::Vector{Float64}, X::Array{Float64,2},
            Y::YVector, α::Vector{Float64},
            indices::Vector{Int}, γ::Float64, sparsity::Int,
            ax_sparse::Vector{Float64})
    ∇ .= -α + Y
    g = zeros(length(α))
    @inbounds for j = 1:sparsity
        x = @view(X[:, indices[j]])
        g .+= ax_sparse[j] * x
    end
    ∇ .-= γ * g
    nothing
end

"""
    solve_dualcutting(X::Array{Float64,2},
            Y::YVector,
            sparsity::Int,
            γ::Float64,
            solver::MathProgBase.AbstractMathProgSolver)
"""
function solve_dualcutting(X::Array{Float64,2},
            Y::YVector,
            sparsity::Int,
            γ::Float64,
            solver::MathProgBase.AbstractMathProgSolver)

    if solver == UnsetSolver()
        error("You need to set a solver for the cutting plane method. E.g.:
                `RelaxDualCutting(0.1, GurobiSolver())`")
    end

    maxiter = 10_0000

    n, p = size(X)
    @assert sparsity <= p
    indices = collect(1:sparsity)
    initial_bound = initial_ub(X, Y, sparsity, γ)

    ax = Vector{Float64}(p)
    ∇  = Vector{Float64}(n)
    ax_sparse = Vector{Float64}(sparsity)
    α_hat = Vector{Float64}(n)

    ub = initial_bound
    lb = -Inf

    m = Model(solver = solver)
    @variables(m, begin
        α[1:n]
        0 <= η <= initial_bound
    end)
    @objective(m, Max, η)

    iter = 0

    while (ub - lb) / (1 + abs(ub)) > 1e-4
        # Solve cutting plane model
        @assert solve(m) == :Optimal
        # Upper bound is current optimum
        ub = getobjectivevalue(m)
        # Get the α that attains the optimum
        α_hat .= getvalue(m[:α])
        # Recover a primal solution
        ax .= X' * α_hat # TODO don't need to store all at the same time
        indices .= sortperm(-abs.(ax))[1:sparsity]
        X_sparse = @view(X[:, indices])
        ax_sparse .= X_sparse' * α_hat        # Use it to update gradient
        getslope!(∇, X, Y, α_hat, indices, γ, sparsity, ax_sparse)
        # Lower bound is the actual function value at α_hat
        lb = getlb(Y, α_hat, ax_sparse, γ)
        # Add a cut
        @constraint(m, η <= lb + dot(∇, α - α_hat))
        if iter > maxiter
            break
        end
    end
    # Recover weights
    w = SubsetSelection.recover_primal(OLS(), Y, X[:,indices], γ)
    perm = sortperm(indices)
    indices[perm], w[perm]
end


text:
# change to course environment
using Pkg
Pkg.activate(dirname(@__FILE__))

# start Pluto server
using Pluto
Pluto.run()


text:
module NonLinearSolvers

using Gridap.LinearSolvers
using Gridap.Helpers

export NonLinearOperator
export residual
export jacobian
export residual!
export jacobian!
export create_in_domain

export NonLinearSolver
export NextNonLinearSolver
import Gridap: solve
import Gridap: solve!

export NewtonRaphsonSolver

abstract type NonLinearOperator end

function residual!(b::AbstractVector,::NonLinearOperator,x::AbstractVector)
  @abstractmethod
end

function jacobian!(A::AbstractMatrix,::NonLinearOperator,x::AbstractVector)
  @abstractmethod
end

function jacobian(::NonLinearOperator,x::AbstractVector)::AbstractMatrix
  @abstractmethod
end

function create_in_domain(::NonLinearOperator)::AbstractVector
  @abstractmethod
end

function residual(op::NonLinearOperator,x::AbstractVector)
  b = similar(x)
  residual!(b,op,x)
  b
end

abstract type NonLinearSolver end

function solve!(x::AbstractVector,::NonLinearSolver,::NonLinearOperator)::Any
  @abstractmethod
end

function solve!(x::AbstractVector,::NonLinearSolver,::NonLinearOperator,cache::Any)
  @abstractmethod
end

function solve(nls::NonLinearSolver,op::NonLinearOperator)
  x = create_in_domain(op)
  E = eltype(x)
  x .= zero(E)
  solve!(x,nls,op)
  x
end

"""
Newton-Raphson method
"""
struct NewtonRaphsonSolver <:NonLinearSolver
  ls::LinearSolver
  tol::Float64
  max_nliters::Int
end

struct NewtonRaphsonCache
  A::AbstractMatrix
  b::AbstractVector
  dx::AbstractVector
  ns::NumericalSetup
end

function solve!(x::AbstractVector,nls::NewtonRaphsonSolver,op::NonLinearOperator)

  b = residual(op, x)
  A = jacobian(op, x)
  dx = similar(b)
  ss = symbolic_setup(nls.ls, A)
  ns = numerical_setup(ss,A)

  _solve_nr!(x,A,b,dx,ns,nls,op)

  NewtonRaphsonCache(A,b,dx,ns)

end

function solve!(
  x::AbstractVector,nls::NewtonRaphsonSolver,op::NonLinearOperator,cache::NewtonRaphsonCache)

  b = cache.b
  A = cache.A
  dx = cache.dx
  ns = cache.ns

  residual!(b, op, x)
  jacobian!(A, op, x)
  numerical_setup!(ns,A)

  _solve_nr!(x,A,b,dx,ns,nls,op)

end

function _solve_nr!(x,A,b,dx,ns,nls,op)

  # Check for convergence on the initial residual
  isconv, conv0 = _check_convergence(nls,b)
  if isconv; return; end

  # Newton-like iterations
  for nliter in 1:nls.max_nliters

    # Solve linearized problem
    broadcast!(*,b,b,-1)
    solve!(dx,ns,b)
    broadcast!(+,x,x,dx)

    # Check convergence for the current residual
    residual!(b, op, x)
    isconv = _check_convergence(nls, b, conv0)
    if isconv; return; end

    if nliter == nls.max_nliters
      @unreachable
    end

    # Assemble jacobian (fast in-place version)
    # and prepare solver
    jacobian!(A, op, x)
    numerical_setup!(ns,A)

  end

end

function _check_convergence(nls,b)
  m0 = _inf_norm(b)
  (false, m0)
end

function _check_convergence(nls,b,m0)
  m = _inf_norm(b)
  m < nls.tol * m0
end

function _inf_norm(b)
  m = 0
  for bi in b
    m = max(m,abs(bi))
  end
  m
end

end # module NonLinearSolvers


text:
## ExoplanetsSysSim/src/limb_darkening.jl
## (c) 2018 Eric B. Ford

abstract type LimbDarkeningParamAbstract end
# Note this file is currently not used by SysSim.

struct LimbDarkeningParamLinear <: LimbDarkeningParamAbstract
  coeff::Tuple{Float64}
  # TODO SCI DETAIL: Repalce with sensible limits on LD params
  LimbDarkeningParamLinear(_u1::Real ) = (!( (-2.0<=_u1<=2.0) )) ? error(string("Invalid quadratic limb darkening parameters: ",_u1," & ", _u2)) : new((_u1,))
end

function depth_at_midpoint(radius_ratio::Float64, ld::LimbDarkeningParamLinear)
    c0 = 1-sum(ld.coeff)
    omega = c0/4+ld.coeff[1]/6
    ksq = 1-radius_ratio^2
    tmp0 = c0/4*ksq
    tmp2 = ld.coeff[1]/6*ksq^(3//2)
    return  (tmp0+tmp2)/omega
end

struct LimbDarkeningParamQuadratic <: LimbDarkeningParamAbstract
  coeff::Tuple{Float64,Float64}
  # TODO SCI DETAIL: Repalce with sensible limits on LD params
  LimbDarkeningParamQuadratic(_u1::Real, _u2::Real ) = (!( (-2.0<=_u1<=2.0) && (-2.0<=_u2<=2.0) )) ? error(string("Invalid quadratic limb darkening parameters: ",_u1," & ", _u2)) : new((_u1,_u2))
end

function depth_at_midpoint(radius_ratio::Float64, ld::LimbDarkeningParamQuadratic)
    c0 = 1-sum(ld.coeff)
    omega = c0/4+(ld.coeff[1]+2*ld.coeff[2])/6-ld.coeff[2]/8
    ksq = 1-radius_ratio^2
    tmp0 = c0/4*ksq
    tmp2 = (ld.coeff[1]+2*ld.coeff[2])/6*ksq^(3//2)
    tmp4 = -ld.coeff[2]/8*ksq^2
    return  (tmp0+tmp2+tmp4)/omega
end

struct LimbDarkeningParam4thOrder <: LimbDarkeningParamAbstract
  coeff::Tuple{Float64,Float64,Float64,Float64}
  # TODO SCI DETAIL: Repalce with sensible limits on LD params
  LimbDarkeningParam4thOrder(_c1::Real, _c2::Real, _c3::Real, _c4::Real ) = (!( (-2.0<=_c1<=2.0) && (-2.0<=_c2<=2.0)  && (-2.0<=_c3<=2.0) && (-2.0<=_c4<=2.0) )) ? error(string("Invalid limb darkening parameters: ",_c1," , ", _c2, ", ",_c3," , ",_c4)) : new((_c1,_c2,_c3,_c4))
end

function depth_at_midpoint(radius_ratio::Float64, ld::LimbDarkeningParam4thOrder)
    c0 = 1-sum(ld.coeff)
    omega = c0/4+sum(ld.coeff./(5:8))
    ksq = 1-radius_ratio^2
    tmp0 = c0/4*ksq
    tmp1 = ld.coeff[1]/5*ksq^(5//4)
    tmp2 = ld.coeff[2]/6*ksq^(3//2)
    tmp3 = ld.coeff[3]/7*ksq^(7//4)
    tmp4 = ld.coeff[4]/8*ksq^2
    return (tmp0+tmp1+tmp2+tmp3+tmp4)/omega
end

# Cite Mandel & Agol 2002
# See https://arxiv.org/pdf/astro-ph/0210099.pdf and https://faculty.washington.edu/agol/mandel_agol_errata.pdf
# https://github.com/ericagol/ExoJulia/blob/master/ExoJulia/Transit/occultsmall.jl
function depth_small_planet(z::T,p::T,ld::LimbDarkeningParam4thOrder) where T<:Real
  @assert p<=0.2
  @assert sum(c)<=1
  c = ld.coeff

  one_over_Omega_times_4 = 1/(1-c[1]/5-c[2]/3-3*c[3]/7-c[4]/2)
  if z == zero(z)
    F = p*p*one_over_Omega_times_4
  elseif z <= 1-p
    sig1=sqrt(sqrt(1-(z-p)^2))
    sig2=sqrt(sqrt(1-(z+p)^2))
    ave_surface_intensity=(1-c[1]*(1+(sig2^5-sig1^5)/(5*p*z))-c[2]*(1+(sig2^6-sig1^6)/(6*p*z))-c[3]*(1+(sig2^7-sig1^7)/(7*p*z))-c[4]*(p^2+z^2))
    F = p*p*ave_surface_intensity*one_over_Omega_times_4
  elseif z<1+p
    x=1-(z-p)^2
    tmp=1-c[1]*(1-0.8*x^(1//4))-c[2]*(1-2//3*sqrt(x))-c[3]*(1-4//7*x^(3//4))-c[4]*(1-0.5*x)
    F = tmp*(p^2*acos((z-1)/p)-(z-1)*sqrt(p^2-(z-1)^2))*one_over_Omega_times_4/pi
  else
    F = zero(T)
  end
  return F
end


text:
using PlanktonIndividuals.Grids

using MeshArrays

function test_rectilinear_grid()
    grid = RectilinearGrid(size = (4,6,2), x = (0,12), y = (0,12), z = (0,-8), halo = (2,2,2))

    @test grid.Nx == 4
    @test grid.Ny == 6
    @test grid.Nz == 2

    @test grid.Δx == 3.0
    @test grid.Δy == 2.0
    @test grid.dzC == [4.0,4.0,4.0,4.0,4.0,4.0]
    @test grid.dzF == [4.0,4.0,4.0,4.0,4.0,4.0]

    @test length(grid.xC) == 4+2*2
    @test length(grid.yC) == 6+2*2
    @test length(grid.zC) == 2+2*2
    @test length(grid.xF) == 4+2*2
    @test length(grid.yF) == 6+2*2
    @test length(grid.zF) == 2+2*2

    return nothing
end

function test_vertically_stretched_rectilinear_grid()
    zfs = [0.0, -1.0, -3.0, -5.0, -10.0]
    grid = RectilinearGrid(size = (4,6,4), x = (0,12), y = (0,12), z = zfs, halo = (2,2,2))

    @test grid.Nx == 4
    @test grid.Ny == 6
    @test grid.Nz == 4

    @test grid.Δx == 3.0
    @test grid.Δy == 2.0
    @test grid.dzC == [1.0,1.0,1.5,2.0,3.5,5.0,5.0,5.0]
    @test grid.dzF == [1.0,1.0,1.0,2.0,2.0,5.0,5.0,5.0]

    @test grid.zF == [2.0,1.0,-0.0,-1.0,-3.0,-5.0,-10.0,-15.0]
    @test grid.zC == [1.5,0.5,-0.5,-2.0,-4.0,-7.5,-12.5,-17.5]

    @test length(grid.xC) == 4+2*2
    @test length(grid.yC) == 6+2*2
    @test length(grid.zC) == 4+2*2
    @test length(grid.xF) == 4+2*2
    @test length(grid.yF) == 6+2*2
    @test length(grid.zF) == 4+2*2

    return nothing
end

function test_rectilinear_areas_volumes()
    zfs = [0.0, -1.0, -3.0, -5.0, -10.0]
    grid = RectilinearGrid(size = (4,6,4), x = (0,12), y = (0,12), z = zfs, halo = (2,2,2))
    @test ΔxC(1,1,1,grid) == grid.Δx
    @test ΔyC(1,1,1,grid) == grid.Δy
    @test ΔzC(1,1,1,grid) == grid.dzC[1]
    @test ΔxF(1,1,1,grid) == grid.Δx
    @test ΔyF(1,1,1,grid) == grid.Δy
    @test ΔzF(1,1,1,grid) == grid.dzF[1]
    @test Ax(1,1,1,grid) == grid.Δy*grid.dzF[1]
    @test Ay(1,1,1,grid) == grid.Δx*grid.dzF[1]
    @test Az(1,1,1,grid) == grid.Δx*grid.Δy
    @test volume(1,1,1,grid) == grid.Δx*grid.Δy*grid.dzF[1]

    return nothing
end

function test_lat_lon_grid()
    grid = LatLonGrid(size = (360,160,10), lat = (-80,80), lon = (-180,180), z = (0,-20))
    @test grid.Nx == 360
    @test grid.Ny == 160
    @test grid.Nz == 10

    @test length(grid.xC) == 360+2*2
    @test length(grid.yC) == 160+2*2
    @test length(grid.zC) == 10+2*2
    @test length(grid.xF) == 360+2*2
    @test length(grid.yF) == 160+2*2
    @test length(grid.zF) == 10+2*2
end

function test_load_lat_lon_grid()
    MeshArrays.GRID_LL360_download();
	γ = MeshArrays.GridSpec("PeriodicChannel",MeshArrays.GRID_LL360);
	Γ = MeshArrays.GridLoad(γ;option="full");

	grid_info=(XC=Γ.XC[1], YC=Γ.YC[1], RC=Γ.RC,
        	XW=Γ.XW[1], YS=Γ.YS[1], RF=Γ.RF,
            DXC=Γ.DXC[1], DYC=Γ.DYC[1], DRC=Γ.DRC,
            DXG=Γ.DXG[1], DYG=Γ.DYG[1], DRF=Γ.DRF,
            RAW=Γ.RAW[1], RAS=Γ.RAS[1], RAC=Γ.RAC[1],
            hFacC=ones(360,160,50), hFacW=ones(360,160,50), hFacS=ones(360,160,50))

    mask = ones(360,160,50)

    grid = LoadLatLonGrid(;grid_info = grid_info, size = (360,160,50), lat = (-80,80), lon= (-180,180), landmask = mask)

    @test grid.Nx == 360
    @test grid.Ny == 160
    @test grid.Nz == 50

    @test length(grid.xC) == 360+2*2
    @test length(grid.yC) == 160+2*2
    @test length(grid.zC) == 50+2*2
    @test length(grid.xF) == 360+2*2
    @test length(grid.yF) == 160+2*2
    @test length(grid.zF) == 50+2*2

    @test minimum(grid.dzF) == 10.0
    @test maximum(grid.dzF) == 456.5

    return nothing
end

function test_lat_lon_areas_volumes()
    grid = LatLonGrid(size = (360,160,10), lat = (-80,80), lon = (-180,180), z = (0,-20))
    @test ΔxC(1,1,1,grid) == grid.dxC[1,1]
    @test ΔyC(1,1,1,grid) == grid.dyC[1,1]
    @test ΔzC(1,1,1,grid) == grid.dzC[1]
    @test ΔxF(1,1,1,grid) == grid.dxF[1,1]
    @test ΔyF(1,1,1,grid) == grid.dyF[1,1]
    @test ΔzF(1,1,1,grid) == grid.dzF[1]
    @test Ax(1,1,1,grid) == grid.Ax[1,1,1]
    @test Ay(1,1,1,grid) == grid.Ay[1,1,1]
    @test Az(1,1,1,grid) == grid.Az[1,1]
    @test volume(1,1,1,grid) == grid.Vol[1,1,1]

    return nothing
end

@testset "Grids" begin
    @testset "Rectilinear Grid" begin
        test_rectilinear_grid()
        test_vertically_stretched_rectilinear_grid()
        test_rectilinear_areas_volumes()
        grid = RectilinearGrid(size = (4,6,2), x = (0,12), y = (0,12), z = (0,-8), halo = (2,2,2))
        @test try
            show(grid); println()
            true
        catch err
            println("error in show(::RectilinearGrid)")
            println(sprint(showerror, err))
            false
        end
        @test grid isa RectilinearGrid
    end
    @testset "Latitude Longitude Grid" begin
        test_lat_lon_grid()
        test_load_lat_lon_grid()
        test_lat_lon_areas_volumes()
        grid = LatLonGrid(size = (360,160,10), lat = (-80,80), lon = (-180,180), z = (0,-20))
        @test try
            show(grid); println()
            true
        catch err
            println("error in show(::LatLonGrid)")
            println(sprint(showerror, err))
            false
        end
        @test grid isa LatLonGrid
    end
end


text:
# This file was generated by the Julia Swagger Code Generator
# Do not modify this file directly. Modify the swagger specification instead.

mutable struct IoK8sApiRbacV1alpha1Subject <: SwaggerModel
    apiVersion::Any # spec type: Union{ Nothing, String } # spec name: apiVersion
    kind::Any # spec type: Union{ Nothing, String } # spec name: kind
    name::Any # spec type: Union{ Nothing, String } # spec name: name
    namespace::Any # spec type: Union{ Nothing, String } # spec name: namespace

    function IoK8sApiRbacV1alpha1Subject(;apiVersion=nothing, kind=nothing, name=nothing, namespace=nothing)
        o = new()
        validate_property(IoK8sApiRbacV1alpha1Subject, Symbol("apiVersion"), apiVersion)
        setfield!(o, Symbol("apiVersion"), apiVersion)
        validate_property(IoK8sApiRbacV1alpha1Subject, Symbol("kind"), kind)
        setfield!(o, Symbol("kind"), kind)
        validate_property(IoK8sApiRbacV1alpha1Subject, Symbol("name"), name)
        setfield!(o, Symbol("name"), name)
        validate_property(IoK8sApiRbacV1alpha1Subject, Symbol("namespace"), namespace)
        setfield!(o, Symbol("namespace"), namespace)
        o
    end
end # type IoK8sApiRbacV1alpha1Subject

const _property_map_IoK8sApiRbacV1alpha1Subject = Dict{Symbol,Symbol}(Symbol("apiVersion")=>Symbol("apiVersion"), Symbol("kind")=>Symbol("kind"), Symbol("name")=>Symbol("name"), Symbol("namespace")=>Symbol("namespace"))
const _property_types_IoK8sApiRbacV1alpha1Subject = Dict{Symbol,String}(Symbol("apiVersion")=>"String", Symbol("kind")=>"String", Symbol("name")=>"String", Symbol("namespace")=>"String")
Base.propertynames(::Type{ IoK8sApiRbacV1alpha1Subject }) = collect(keys(_property_map_IoK8sApiRbacV1alpha1Subject))
Swagger.property_type(::Type{ IoK8sApiRbacV1alpha1Subject }, name::Symbol) = Union{Nothing,eval(Base.Meta.parse(_property_types_IoK8sApiRbacV1alpha1Subject[name]))}
Swagger.field_name(::Type{ IoK8sApiRbacV1alpha1Subject }, property_name::Symbol) =  _property_map_IoK8sApiRbacV1alpha1Subject[property_name]

function check_required(o::IoK8sApiRbacV1alpha1Subject)
    (getproperty(o, Symbol("kind")) === nothing) && (return false)
    (getproperty(o, Symbol("name")) === nothing) && (return false)
    true
end

function validate_property(::Type{ IoK8sApiRbacV1alpha1Subject }, name::Symbol, val)
end


text:
# Takes two .bench files as input
length(ARGS) >= 2 || error("Please provide two .bench files as input args.")

using Pkg
Pkg.activate(@__DIR__)

using PkgBenchmark

f1 = ARGS[1]
f2 = ARGS[2]

r1 = readresults(f1)
r2 = readresults(f2)

j = judge(r1, r2)

name1 = replace(f1, ".bench" => "")
name2 = replace(f2, ".bench" => "")

export_markdown(string("judge_",name1,"_",name2,".md"), j)


text:
abstract type Curve{F} end

# (short) Weierstrass normal form, y² = x³ + ax + b
@auto_hash_equals struct WNFCurve{F} <: Curve{F}
    a::F
    b::F
end

# generalized Weierstrass normal form, y² + a₁xy + a₃y = x³ + a₂x² + a₄x + a₆
@auto_hash_equals struct GWNFCurve{F} <: Curve{F}
    a₁::F
    a₂::F
    a₃::F
    a₄::F
    a₆::F
end

curve(a, b) = curve(promote(a, b)...)
curve(a::F, b::F) where {F} = begin
    ec = WNFCurve(a, b)
    issingular(ec) && error("curve $(ec) is singular")
    ec
end

curve(a₁, a₂, a₃, a₄, a₆) = curve(promote(a₁, a₂, a₃, a₄, a₆)...)
curve(a₁::F, a₂::F, a₃::F, a₄::F, a₆::F) where {F} = begin
    ec = GWNFCurve(a₁, a₂, a₃, a₄, a₆)
    issingular(ec) && error("curve $(ec) is singular")
    ec
end

issingular(ec::Curve) = discriminant(ec) == 0
issmooth(ec::Curve) = !singular(ec)

discriminant(ec::WNFCurve) = -16 * (4*ec.a^3 + 27*ec.b^2)
discriminant(ec::GWNFCurve) = begin
    b₂ =   ec.a₁^2 + 4*ec.a₂
    b₄ = 2*ec.a₄   +   ec.a₁*ec.a₃
    b₆ =   ec.a₃^2 + 4*ec.a₆
    b₈ = (ec.a₁^2)*ec.a₆ + 4*ec.a₂*ec.a₆ - ec.a₁*ec.a₃*ec.a₄ + ec.a₂*(ec.a₃^2) - ec.a₄^2
    -(b₂^2)*b₈ - 8*b₄^3 - 27*b₆^2 + 9*b₂*b₄*b₆
end

contains(ec::WNFCurve{F}, x::F, y::F) where {F} = y^2 == x^3 + ec.a*x + ec.b
contains(ec::GWNFCurve{F}, x::F, y::F) where {F} = y^2 + ec.a₁*x*y + ec.a₃*y == x^3 + ec.a₂*x^2 + ec.a₄*x + ec.a₆

ring(n::Curve{T}) where {T} = T
ring(ec::WNFCurve{T}) where {T<:RingElem} = parent(ec.a)
ring(ec::GWNFCurve{T}) where {T<:RingElem} = parent(ec.a₁)

field(ec::Curve{T}) where {T<:FieldElem} = ring(ec)

show(io::IO, ec::WNFCurve) = print(io, "{y² = x³ + [$(ec.a)]x + [$(ec.b)]   x, y ∈ $(ring(ec))}")
show(io::IO, ec::GWNFCurve) = print(io, "{y² + [$(ec.a₁)]xy + [$(ec.a₃)]y == x³ + [$(ec.a₂)]x² + [$(ec.a₄)]x + [$(ec.a₆)]   x, y ∈ $(ring(ec))}")


text:
using Turing
using Test

@model function test_ex_rt()
  x ~ Normal(10, 1)
  y ~ Normal(x / 2, 1)
  z = y + 1
  x = x - 1
  x, y, z
end

mf = test_ex_rt()

for alg = [HMC(0.2, 3), PG(20, 2000), SMC(), IS(10000), Gibbs(PG(20, 1, :x), HMC(0.2, 3, :y))]
  chn = sample(mf, alg)
  @test mean(chn[:x]) ≈ 10.0 atol=0.2
  @test mean(chn[:y]) ≈ 5.0 atol=0.2
end


text:
using Mux, MemPool, Distributed, Sockets
import Dagger.Tables

struct LinePlot
    core_key::Symbol
    data_key::Symbol
    xlabel::String
    ylabel::String
end
StructTypes.StructType(::Type{LinePlot}) = StructTypes.CustomStruct()
StructTypes.lower(lp::LinePlot) = "linePlot(svg_container, \"$(lp.core_key)\", \"$(lp.data_key)\", \"$(lp.xlabel)\", \"$(lp.ylabel)\")"

struct GanttPlot
    core_key::Symbol
    id_key::Symbol
    timeline_key::Symbol
    esat_key::Symbol
    psat_key::Symbol
    title::String
end
StructTypes.StructType(::Type{GanttPlot}) = StructTypes.CustomStruct()
StructTypes.lower(lp::GanttPlot) = "ganttPlot(svg_container, \"$(lp.core_key)\", \"$(lp.id_key)\", \"$(lp.timeline_key)\", \"$(lp.esat_key)\", \"$(lp.psat_key)\", \"$(lp.title)\")"

struct GraphPlot
    core_key::Symbol
    id_key::Symbol
    timeline_key::Symbol
    profile_key::Symbol
    title::String
end
StructTypes.StructType(::Type{GraphPlot}) = StructTypes.CustomStruct()
StructTypes.lower(lp::GraphPlot) = "graphPlot(svg_container, \"$(lp.core_key)\", \"$(lp.id_key)\", \"$(lp.timeline_key)\", \"$(lp.profile_key)\", \"$(lp.title)\")"

struct ProfileViewer
    core_key::Symbol
    profile_key::Symbol
    title::String
end
StructTypes.StructType(::Type{ProfileViewer}) = StructTypes.CustomStruct()
StructTypes.lower(lp::ProfileViewer) = "profileViewer(svg_container, \"$(lp.core_key)\", \"$(lp.profile_key)\", \"$(lp.title)\")"

## Core globals
const D3R_RUNNING = Dict{Int,Bool}()
const D3R_LOGS = Dict{Int,Dict{Symbol,Any}}()
const D3R_ACTUAL_PORT = Dict{Int,Int}()

## Client/forwarding globals
# Port -> Worker -> Sockets
# TODO: Concrete socket types
const D3R_CLIENT_SOCKETS = Dict{Int,Dict{Int,Vector{Any}}}()

const D3R_WORKER_HOST_MAP = Dict{Int,IPAddr}()
const D3R_WORKER_PORT_MAP = Dict{Int,Int}()

struct D3Renderer
    port::Int
    port_range::UnitRange{Int}
    config::Vector{Any}
    config_updated::Ref{Bool}
    update_cond::Condition
    seek_store
end
D3Renderer(port::Int, port_range=50000:59999; seek_store=nothing) =
    D3Renderer(port, port_range, Vector{Any}(), Ref{Bool}(false), Condition(), seek_store)
Dagger.init_similar(d3r::D3Renderer) =
    D3Renderer(d3r.port,
               d3r.port_range,
               d3r.config,
               Ref{Bool}(false),
               Condition(),
               d3r.seek_store)

function (d3r::D3Renderer)(logs)
    D3R_LOGS[d3r.port] = logs
    d3r_init(d3r.port, d3r.port_range, d3r.config_updated, d3r.config, d3r.seek_store)
    notify(d3r.update_cond)
end

function Base.push!(d3r::D3Renderer, x)
    push!(d3r.config, x)
    d3r.config_updated[] = true
end

function d3r_init(port::Int, port_range::UnitRange, config_updated::Ref{Bool}, config, seek_store)
    get(D3R_RUNNING, port, false) && return

    # Ensure some variables are setup
    get!(()->Dict{Symbol,Any}(), D3R_LOGS, port)
    get!(()->[], get!(()->Dict{Int,Vector{Any}}(), D3R_CLIENT_SOCKETS, port), myid())

    # Server configuration
    d3_app = (
        Mux.prod_defaults,
        page("/", req->Dict(
            :status=>301,
            :headers=>Dict(:Location=>"/index.html"),
            :body=>"Moved permanently"
        )),
        page("/index.html", req->begin
            html = String(read(joinpath(@__DIR__, "index.html")))
            hdrs = Dict(
                Symbol("Content-Type")=>"text/html",
                # For interactive development
                Symbol("Cache-Control")=>"no-store",
            )
            Dict(
                :body=>html,
                :headers=>hdrs,
            )
        end),
        page("/fs/:path", req->begin
            path = req[:params][:path]
            path = "/tmp/" * path
            if ispath(path)
                Dict(:body=>String(read(path)))
            else
                Dict(:status=>404)
            end
        end),
        page("/worker_id", respond("$(myid())")),
        page("/worker/:id/:location", req->begin
            id = parse(Int, req[:params][:id])
            location = req[:params][:location]
            worker_host, worker_port = worker_host_port(id, port, port_range)
            res = HTTP.get("http://$worker_host:$worker_port/$location")
            return Dict(
                :status=>res.status,
                :body=>String(res.body),
            )
        end),
        Mux.notfound()
    )
    d3_wapp = (
        Mux.wdefaults,
        route("/data_feed", req->begin
            sock = req[:socket]
            client_handler(sock, myid(), port, port_range, config_updated, config, seek_store)
        end),
        route("/worker/:id/data_feed", req->begin
            sock = req[:socket]
            id = parse(Int, req[:params][:id])
            client_handler(sock, id, port, port_range, config_updated, config, seek_store)
        end),
        Mux.wclose,
        Mux.notfound()
    )

    actual_port = get_actual_port(port, port_range)

    # Start D3R server
    Mux.serve(Mux.App(mux(d3_app...)),
              Mux.App(mux(d3_wapp...)),
              actual_port)
    @debug "D3R initialized for port $port at actual port $actual_port"
    D3R_ACTUAL_PORT[port] = actual_port
    D3R_RUNNING[port] = true
end
function client_handler(sock, id, port, port_range, config_updated, config, seek_store)
    fsock = nothing
    if id != myid()
        fsock_ready = Base.Event()
        worker_host, worker_port = worker_host_port(id, port, port_range)
        Mux.HTTP.WebSockets.open("ws://$worker_host:$worker_port/data_feed") do _fsock
            fsock = _fsock
            @debug "D3R forwarder for $id ready"
            notify(fsock_ready)
            while !eof(fsock) && isopen(fsock)
                try
                    bytes = readavailable(fsock)
                    if length(bytes) == 0
                        sleep(0.1)
                        continue
                    end
                    data = String(bytes)
                    #@info "D3R forwarder for $id received data"
                    write(sock, data)
                catch err
                    if err isa Mux.WebSockets.WebSocketClosedError || err isa Base.IOError
                        # Force-close client and forwarder
                        @async close(sock)
                        @async close(fsock)
                        break
                    end
                    @error "D3R forwarder for $id error!" exception=(err,catch_backtrace())
                    rethrow(err)
                end
            end
            @debug "D3R forwarder for $id closed"
        end
        wait(fsock_ready)
    end
    if id == myid()
        @debug "D3R client for $id sending initial config"
        write(sock, JSON3.write((;cmd="data", payload=sanitize(D3R_LOGS[port]))))
        _workers = workers()
        if !(myid() in _workers)
            # FIXME: Get this from the Context
            _workers = vcat(myid(), _workers)
        end
        write(sock, JSON3.write((;cmd="config", payload=sanitize((;myid=myid(),workers=_workers,ctxs=config)))))
    end
    push!(get!(()->[], D3R_CLIENT_SOCKETS[port], id), sock)
    @debug "D3R client for $id ready"
    while !eof(sock) && isopen(sock)
        try
            data = String(read(sock))
            @debug "D3R client for $id received: $data"
            if id == myid()
                #= FIXME
                if config_updated[]
                    config_updated[] = false
                    write(sock, JSON3.write((;cmd="config", payload=sanitize(config))))
                end
                =#
                if seek_store !== nothing
                    if data == "fulldata"
                        raw_logs = seek_store[0:typemax(UInt64)]
                        logs = Dict{Symbol,Vector{Any}}()
                        for (idx,key) in enumerate(Tables.columnnames(raw_logs))
                            logs[key] = Tables.columns(raw_logs)[idx]
                        end
                        write(sock, JSON3.write((;cmd="data", payload=sanitize(logs))))
                        continue
                    end
                    m = match(r"seek\(([0-9]*),([0-9]*)\)", data)
                    if m !== nothing
                        ts_start, ts_stop = parse.(Ref(UInt64), m.captures)
                        raw_logs = seek_store[ts_start:ts_stop]
                        logs = Dict{Symbol,Vector{Any}}()
                        for (idx,key) in enumerate(Tables.columnnames(raw_logs))
                            logs[key] = Tables.columns(raw_logs)[idx]
                        end
                        write(sock, JSON3.write((;cmd="data", payload=sanitize(logs))))
                        continue
                    end
                end
                if data == "data"
                    write(sock, JSON3.write((;cmd="data", payload=sanitize(D3R_LOGS[port]))))
                end
            else
                @debug "D3R client sending to forwarder: $data"
                write(fsock, data)
            end
        catch err
            if err isa Mux.WebSockets.WebSocketClosedError || err isa Base.IOError
                idx = findfirst(x->x==sock, D3R_CLIENT_SOCKETS[port][id])
                if idx !== nothing
                    deleteat!(D3R_CLIENT_SOCKETS[port][id], idx)
                end
                # Force-close client and forwarder
                @async close(sock)
                @async close(fsock)
                break
            end
            @error "D3R client for $id error!" exception=(err,catch_backtrace())
            rethrow(err)
        end
    end
    @debug "D3R client for $id closed"
end
function Dagger.Events.creation_hook(d3r::D3Renderer, log)
    for sock in get!(()->[], get!(()->Dict{Int,Vector{Any}}(), D3R_CLIENT_SOCKETS, d3r.port), myid())
        try
            if isopen(sock)
                write(sock, JSON3.write((;cmd="add", payload=sanitize(log))))
            end
        catch err
            if err isa Mux.WebSockets.WebSocketClosedError
                continue
            end
            rethrow(err)
        end
    end
end
function Dagger.Events.deletion_hook(d3r::D3Renderer, idx)
    for sock in get!(()->[], get!(()->Dict{Int,Vector{Any}}(), D3R_CLIENT_SOCKETS, d3r.port), myid())
        try
            if isopen(sock)
                write(sock, JSON3.write((;cmd="delete", payload=idx)))
            end
        catch err
            if err isa Mux.WebSockets.WebSocketClosedError
                continue
            end
            rethrow(err)
        end
    end
end
function worker_host_port(id::Int, port::Int, port_range::UnitRange)
    worker_host = get!(D3R_WORKER_HOST_MAP, port) do
        worker_host, worker_port = remotecall_fetch(id, port) do port
            (MemPool.host, get_actual_port(port, port_range))
        end
        D3R_WORKER_PORT_MAP[port] = worker_port
        worker_host
    end
    worker_port = D3R_WORKER_PORT_MAP[port]
    return worker_host, worker_port
end
function get_actual_port(port::Int, port_range::UnitRange)
    get!(D3R_ACTUAL_PORT, port) do
        # Check if the requested port is available
        # If not, pick another port in the provided range
        actual_port = port
        while true
            port_available = try
                close(listen(actual_port))
                true
            catch err
                err isa Base.IOError || rethrow(err)
                if err.code == -98
                    false
                else
                    rethrow(err)
                end
            end
            if port_available
                if port != actual_port && myid() == 1
                    @warn "Port $port not available, serving on $actual_port"
                end
                break
            end
            actual_port = rand(port_range)
        end
        actual_port
    end
end


text:
#' # OPTINUM-RAPPORT
#'
#' Nom et Prénom                       Groupe:
#'
#' Date:
using LinearAlgebra
using Optinum

#' ## Exercice 1 - <le sujet de l'exercice>

function question1()

  print("TODO !")

  return

end

question1()

# Un commentaire sur l'exercice 1

max_stars_repo_name: faicaltoubali/ENSEEIHT

max_issues_repo_name: faicaltoubali/ENSEEIHT

max_forks_repo_name: faicaltoubali/ENSEEIHT


text:
using RareEventABC
using Plots; pgfplots()
default(guidefont=font(25), tickfont=font(25), legendfont=font(25), size=(500,500));
using StatPlots
using ProgressMeter
using DataFrames

##ABAKALIKI DATA
size_obs = 30;
popsize = 120;
robs = cumsum([0,13,7,2,3,0,0,1,4,5,3,2,0,2,0,5,3,1,4,0,1,1,1,2,0,1,5,0,5,5]);
duration_obs = maximum(robs);
robs = [robs; [Inf for i in 1:90]];

##MODELS
using RareEventABC
using Distributions

##Markoivan SIR binned to 5 day periods
sir_binned = RareEventABC.SEIRDetails(robs, popsize, θ->RareEventABC.Exponential(1.), θ->RareEventABC.DiracDelta(0.0), θ->Distributions.Exponential(θ[2]), x->floor(x/5.)*5.);

##Gamma infectious period
sir_gamma_I = RareEventABC.SEIRDetails(robs, popsize, θ->RareEventABC.Exponential(1.), θ->RareEventABC.DiracDelta(0.0), θ->Distributions.Gamma(θ[2], θ[3]), x->x);

##Weibull pressure thresholds
sir_weibull_P = RareEventABC.SEIRDetails(robs, popsize, θ->RareEventABC.Weibull(θ[2], 1.), θ->RareEventABC.DiracDelta(0.0), θ->Distributions.Exponential(θ[3]), x->x);

prior_2 = MvIndependent([Exponential(10.), Exponential(10.)]); ##Corresponds to Kypraios et al 2016 prior (n.b. Paper specified Exp(0.1) under rate parameterisation. Julia uses scale parameterisation.)
prior_3 = MvIndependent([Exponential(10.), Exponential(10.), Exponential(10.)]); ##Corresponds to Kypraios et al 2016 prior (n.b. Paper specified Exp(0.1) under rate parameterisation. Julia uses scale parameterisation.)

eps = 15.0;

priors = Distribution[prior_2, prior_3, prior_3];
models = RareEventABC.SEIRDetails[sir_binned, sir_gamma_I, sir_weibull_P];

"
Perform an ABC rejection analysis

Arguments:

* `model` is model to use
* `prior` is the prior distribution
* `N` is number of acceptances required
* `threshold` is threshold to investigate
"
function doABC(model::RareEventABC.SEIRDetails, prior::Distribution, N::Int, threshold::Float64)
    output_sample = zeros(Float64, (N,2))
    nsims = 0
    tic()
    @showprogress for i in 1:N
        while true
            prior_draw = rand(prior)
            next_dist = RareEventABC.getDistance(model, RareEventABC.simulateEpidemic(model, prior_draw))
            nsims += 1
            if next_dist <= threshold
                output_sample[i,:] = prior_draw
                break ##i.e. increment i in for loop
            end ##if
        end ##while
    end ##for
    timeABC = toq()
    time_per_sample = timeABC / N
    return time_per_sample
end ##trialABC

doABC(sir_binned, prior_2, 100, 100.0)

"
Run ABC-MCMC

Arguments:

* `model` is model to use
* `prior` is the prior distribution
* `nits` is number of iterations to perform
* `threshold` is threshold to investigate
* `increment` is random walk increment distribution
* `pars_init` is initial parameter values
"
function doABCMCMC(model::RareEventABC.SEIRDetails, prior::Distribution, nits::Int, threshold::Float64,
                   increment::Distribution, pars_init::Array{Float64, 1})
    tic()
    nacc = 0
    npars = length(pars_init)
    theta = pars_init
    sum_theta = theta
    sum_theta2 = theta .^2
    @showprogress for i in 2:nits
        theta_prop = theta + rand(increment)
        logratio = logpdf(prior, theta_prop) - logpdf(prior, theta)
        if (log(rand()) < logratio) ##1st step of acceptance, based on prior
            next_dist = RareEventABC.getDistance(model, RareEventABC.simulateEpidemic(model, theta_prop))
            if next_dist <= threshold
                theta = theta_prop ##Acceptance
                nacc += 1
            end
        end
        sum_theta .+= theta
        sum_theta2 .+= theta .^2
    end
    t = toq()
    mean_theta = sum_theta / nits
    std_theta = sqrt(max(0.0, sum_theta2./nits - mean_theta .^ 2))
    print("Mean thetas $(round(mean_theta, 2)), sdevs $(round(std_theta, 2)) \n")
    return DataFrame(threshold = threshold, nacc = nacc, accrate = nacc/nits)
end

inc1 = MvNormal([0.,0.], [0.00198571 -0.130246; -0.130246 23.1669]);
start1 = [0.11994, 10.4409];
doABCMCMC(sir_binned, prior_2, 10^7, 15.0, inc1, start1)

inc2 = MvNormal(zeros(3), [0.00139816   -0.113794    0.0590885;
                          -0.113794     19.9754    -10.3812;
                           0.0590885   -10.3812      7.47516]);
start2 = [0.103064, 4.11688, 4.28655];
doABCMCMC(sir_gamma_I, prior_3, 10^7, 15.0, inc2, start2)

inc3 = MvNormal(zeros(3), [0.00107213  0.006198   -0.0642286;
                           0.006198    0.0832468   0.0162699;
                          -0.0642286   0.0162699  10.1916]);
start3 = [0.0555384, 0.773309, 14.1865];
doABCMCMC(sir_gamma_I, prior_3, 10^7, 15.0, inc3, start3)


text:
##### Numeric simplification

"""
    default_rules(expr, context::T)::RuleSet

The `RuleSet` to be used by default for a given expression and the context.
Julia packages defining their own context types should define this method.

By default SymbolicUtils will try to apply appropriate rules for expressions
of symtype Number.
"""
default_rules(x, ctx) = SIMPLIFY_RULES

"""
    simplify(x, ctx=EmptyCtx();
        rules=default_rules(x, ctx),
        fixpoint=true,
        applyall=true,
        recurse=true)

Simplify an expression by applying `rules` until there are no changes.
The second argument, the context is passed to every [`Contextual`](#Contextual)
predicate and can be accessed as `(@ctx)` in the right hand side of `@rule` expression.

By default the context is an `EmptyCtx()` -- which means there is no contextual information.
Any arbitrary type can be used as a context, and packages defining their own contexts
should define `default_rules(ctx::TheContextType)` to return a `RuleSet` that will
be used by default while simplifying under that context.

If `fixpoint=true` this will repeatedly apply the set of rules until there are no changes.
Applies them once if `fixpoint=false`.

The `applyall` and `recurse` keywords are forwarded to the enclosed
`RuleSet`, they are mainly used for internal optimization.
"""
function simplify(x, ctx=EmptyCtx(); rules=default_rules(x, ctx), fixpoint=true, applyall=true, recurse=true)
    if fixpoint
        SymbolicUtils.fixpoint(rules, x, ctx; recurse=recurse, applyall=recurse)
    else
        rules(x, ctx; recurse=recurse, applyall=recurse)
    end
end

Base.@deprecate simplify(x, rules::RuleSet; kwargs...)  simplify(x, rules=rules; kwargs...)

"""
    substitute(expr, dict)

substitute any subexpression that matches a key in `dict` with
the corresponding value.
"""
function substitute(expr, dict)
    RuleSet([@rule ~x::(x->haskey(dict, x)) => dict[~x]])(expr)
end

### Predicates

sym_isa(::Type{T}) where {T} = @nospecialize(x) -> x isa T || symtype(x) <: T
is_operation(f) = @nospecialize(x) -> (x isa Term) && (operation(x) == f)

isliteral(::Type{T}) where {T} = x -> x isa T
isnumber(x) = isliteral(Number)(x)

_iszero(t) = false
_iszero(x::Number) = iszero(x)
_isone(t) = false
_isone(x::Number) = isone(x)

# A total ordering
<ₑ(a::Real,    b::Real) = abs(a) < abs(b)
<ₑ(a::Complex, b::Complex) = (abs(real(a)), abs(imag(a))) < (abs(real(b)), abs(imag(b)))
<ₑ(a::Real,    b::Complex) = true
<ₑ(a::Complex, b::Real) = false

<ₑ(a::Symbolic, b::Number) = false
<ₑ(a::Number,   b::Symbolic) = true

arglength(a) = length(arguments(a))
function <ₑ(a::Sym, b::Term)
    args = arguments(b)
    if length(args) === 2
        n1, n2 = !isnumber(args[1]) , !isnumber(args[2])
        if n1 && n2
            # both subterms are terms, so it's definitely firster
            return true
        elseif n1
            return isequal(a, args[1]) || a <ₑ args[1]
        elseif n2
            return isequal(a, args[2]) || a <ₑ args[2]
        else
            # both arguments are not numbers
            # This case when a <ₑ Term(^, [1,-1])
            # so this term should go to the left.
            return false
        end
    elseif length(args) === 1
        # make sure a < sin(a) < b^2 < b
        if isequal(a, args[1])
            return true # e.g sin(a)*a should become a*sin(a)
        else
            return a<ₑargs[1]
        end
    else
        # variables to the right
        return false
    end
end

<ₑ(a::Symbolic, b::Sym) = !(b <ₑ a)

function <ₑ(a::Symbol, b::Symbol)
    # Enforce the order [+,-,\,/,^,*]
    if b === :*
        a in (:^, :/, :\, :-, :+)
    elseif b === :^
        a in (:/, :\, :-, :+) && return true
    elseif b === :/
        a in (:\, :-, :+) && return true
    elseif b === :\
        a in (:-, :+) && return true
    elseif b === :-
        a === :+ && return true
    elseif a in (:*, :^, :/, :-, :+)
        false
    else
        a < b
    end
end

<ₑ(a::Sym, b::Sym) = a.name < b.name
<ₑ(a::T, b::S) where {T, S} = T===S ? isless(a, b) : nameof(T) < nameof(S)

function <ₑ(a::Term, b::Term)
    if arglength(a) === 0
        return operation(a) <ₑ b
    elseif arglength(b) === 0
        return a <ₑ operation(b)
    end

    na = nameof(operation(a))
    nb = nameof(operation(b))

    if 0 < arglength(a) <= 2 && 0 < arglength(b) <= 2
        # e.g. a < sin(a) < b ^ 2 < b
        @goto compare_args
    end

    if na !== nb
        return na <ₑ nb
    elseif arglength(a) != arglength(b)
        return arglength(a) < arglength(b)
    else
        @label compare_args
        aa, ab = arguments(a), arguments(b)
        if length(aa) !== length(ab)
            return length(aa) < length(ab)
        else
            terms = zip(Iterators.filter(!isnumber, aa), Iterators.filter(!isnumber, ab))

            for (x,y) in terms
                if x <ₑ y
                    return true
                elseif y <ₑ x
                    return false
                end
            end

            # compare the numbers
            nums = zip(Iterators.filter(isnumber, aa),
                       Iterators.filter(isnumber, ab))
            if any(a <ₑ b for (a, b) in nums)
                return true
            end
        end
        return na <ₑ nb # all args are equal, compare the name
    end
end

issortedₑ(args) = issorted(args, lt=<ₑ)

# are there nested ⋆ terms?
function isnotflat(⋆)
    function (args)
        for t in args
            if t isa Term && operation(t) === (⋆)
                return true
            end
        end
        return false
    end
end

function hasrepeats(x)
    length(x) <= 1 && return false
    for i=1:length(x)-1
        if isequal(x[i], x[i+1])
            return true
        end
    end
    return false
end

function merge_repeats(merge, xs)
    length(xs) <= 1 && return false
    merged = Any[]
    i=1

    while i<=length(xs)
        l = 1
        for j=i+1:length(xs)
            if isequal(xs[i], xs[j])
                l += 1
            else
                break
            end
        end
        if l > 1
            push!(merged, merge(xs[i], l))
        else
            push!(merged, xs[i])
        end
        i+=l
    end
    return merged
end

### Simplification rules

pow(x,y) = y==0 ? 1 : y<0 ? inv(x)^(-y) : x^y
pow(x::Symbolic,y) = y==0 ? 1 : Base.:^(x,y)
pow(x, y::Symbolic) = Base.:^(x,y)
pow(x::Symbolic,y::Symbolic) = Base.:^(x,y)

# Numbers to the back
function flatten_term(⋆, args)
    # flatten nested ⋆
    flattened_args = []
    for t in args
        if t isa Term && operation(t) === (⋆)
            append!(flattened_args, arguments(t))
        else
            push!(flattened_args, t)
        end
    end
    Term(⋆, flattened_args)
end

function sort_args(f, args)
    if length(args) < 2
        return Term(f, args)
    elseif length(args) == 2
        x, y = args
        return Term(f, x <ₑ y ? [x,y] : [y,x])
    end
    args = args isa Tuple ? [args...] : args
    Term(f, sort(args, lt=<ₑ))
end


text:
abstract type AbstractBuffer end

struct DataBuffer <: AbstractBuffer
    data_array
    function DataBuffer()
        data_array = Any[]
        new(data_array)
    end
end

function Base.push!(buffer::DataBuffer, args...; kwargs...)
    error("Define specific method `push!`")
end


text:
module Ganymede

using SpaceAnalysis: me, mp, ϵ0, q

export Rg, μ₀, amu, upstream_value, γ
export Index, getIndex

const ionElectronMassRatio = 100 # ion-electron mass ratio in PIC simulation
const ionMass = 14 # single fluid ion mass in Ganymede's simulation
const Rg = 2634000. # Ganymede's radius, [m]
const dx = 1/100 # grid resolution for stream tracing, [Rg]
# Rho Ux Uy Uz Bx By Bz Pe P, NaN for compensating the No index in the 1st col.
const upstream_value = [56., 140., 0., 0., -10., -6., -86., 0.2, 3.4]
const γ = 5/3 # adiabatic index

const Z = 1 # Charge state
const ni = 4e6 # upstream ion number density for G8 flyby, [amu/m^3]
const ω_pi = √(ni*q^2/(ϵ0*ionMass*mp))*Z # ion plasma frequency, [rad/s]

# ion inertial length at Ganymede's upstream, [Rg]
const dᵢ = c / (√(ni*q^2/(ϵ0*ionMass*mp))*Z) / Rg

struct Index
   Rho_::Integer
   Ux_::Integer
   Uy_::Integer
   Uz_::Integer
   Bx_::Integer
   By_::Integer
   Bz_::Integer
   Pe_::Integer
   P_::Integer
   B_::UnitRange{Integer}
   U_::UnitRange{Integer}
end

function getIndex(header::Vector{AbstractString})
   Rho_= findfirst(x->x=="Rho", header)
   P_  = findfirst(x->x=="P", header)
   Pe_ = findfirst(x->x=="Pe", header)
   Bx_ = findfirst(x->x=="Bx", header)
   By_ = Bx_ + 1
   Bz_ = Bx_ + 2
   B_  = Bx_:(Bx_+2)
   Ux_ = findfirst(x->x=="Ux", header)
   Uy_ = Ux_ + 1
   Uz_ = Ux_ + 2
   U_  = Ux_:(Ux_+2)

   id = Index(Rho_, Ux_, Uy_, Uz_, Bx_, By_, Bz_, Pe_, P_, B_, U_)
end

end


text:
module NonhydrostaticModels

export NonhydrostaticModel

using DocStringExtensions

using KernelAbstractions: @index, @kernel, Event, MultiEvent
using KernelAbstractions.Extras.LoopInfo: @unroll

using Oceananigans.Utils: launch!
using Oceananigans.Grids
using Oceananigans.Solvers
using Oceananigans.Distributed: MultiArch, DistributedFFTBasedPoissonSolver, reconstruct_global_grid
using Oceananigans.ImmersedBoundaries: ImmersedBoundaryGrid

import Oceananigans: fields, prognostic_fields

function PressureSolver(arch::MultiArch, local_grid::RegRectilinearGrid)
    global_grid = reconstruct_global_grid(local_grid)
    return DistributedFFTBasedPoissonSolver(global_grid, local_grid)
end

PressureSolver(arch, grid::RegRectilinearGrid)  = FFTBasedPoissonSolver(grid)
PressureSolver(arch, grid::HRegRectilinearGrid) = FourierTridiagonalPoissonSolver(grid)

# *Evil grin*
PressureSolver(arch, ibg::ImmersedBoundaryGrid) = PressureSolver(arch, ibg.underlying_grid)

#####
##### NonhydrostaticModel definition
#####

include("nonhydrostatic_model.jl")
include("pressure_field.jl")
include("show_nonhydrostatic_model.jl")
include("set_nonhydrostatic_model.jl")

#####
##### Time-stepping NonhydrostaticModels
#####

"""
    fields(model::NonhydrostaticModel)

Returns a flattened `NamedTuple` of the fields in `model.velocities` and `model.tracers`.
"""
fields(model::NonhydrostaticModel) = merge(model.velocities, model.tracers)
prognostic_fields(model::NonhydrostaticModel) = fields(model)

include("solve_for_pressure.jl")
include("update_hydrostatic_pressure.jl")
include("update_nonhydrostatic_model_state.jl")
include("pressure_correction.jl")
include("nonhydrostatic_tendency_kernel_functions.jl")
include("calculate_nonhydrostatic_tendencies.jl")
include("correct_nonhydrostatic_immersed_tendencies.jl")

end # module


text:
# Includes all algorithms
# -----------------------

include("policies.jl")
include("biased_random_walk.jl")

include("sarsaL.jl")
include("sarsaL_modulated_particlefilterjump.jl")

include("forward.jl")
include("hybrid_sarsaL.jl")
include("hybrid_glaescher.jl")

include("REINFORCE.jl")
include("REINFORCE_modulated.jl")
include("REINFORCE_baseline.jl")
include("actor_critic.jl")
include("actor_critic_modulatedActor.jl")

include(joinpath("learner", "Tlearner", "tgeneric.jl"))
include(joinpath("learner", "Tlearner", "tleakyglaescher.jl"))
include(joinpath("learner", "Tlearner", "tparticlefilterjump.jl"))

include(joinpath("learner", "Rlearner", "rewardlearners.jl"))

include(joinpath("learner", "prioritizedsweepinglearner.jl"))
include(joinpath("learner", "REINFORCEmodulatedlearner.jl"))
include(joinpath("learner", "REINFORCEbaselinelearner.jl"))
include(joinpath("learner", "REINFORCElearner.jl"))
include(joinpath("learner", "sarsaLmodulatedlearner.jl"))
include(joinpath("learner", "sarsaLlearner.jl"))
include(joinpath("learner", "actorcriticmodulatedlearner.jl"))
include(joinpath("learner", "actorcriticlearner.jl"))
include(joinpath("learner", "hybridconvexlearner.jl"))
include(joinpath("learner", "hybridlearner.jl"))
include(joinpath("learner", "forwardlearnerBF.jl"))
include(joinpath("learner", "forwardlearner.jl"))

include("prioritized_sweeping_particlefilterjump.jl")
include("hybrid_prioritizedsweeping_particlefilterjump.jl")
include("hybrid_actorcritic_FWD.jl")
include("hybrid_actorcritic_FWD_BF.jl")

include("actor_critic_actioncost.jl")
include("actor_critic_modulatedActor_actioncost.jl")


text:
#==== Move ====================================================================#

function locmove(l::Symbol, Y, C::Constants, S::State, X::Auxiliary, U::Counts)
    if l == :hmc
        loc_hmc(Y,C,S,X,U)
    else
        loc_gibbs(Y,C,S,X,U)
    end
end

#==== Gibbs ===================================================================#

function loc_gibbs(Y, C::Constants, S::State, X::Auxiliary, U::Counts)

    k,α,β = S.k, C.α, C.β
    @views locs,hts =S.locs[1:k+2], S.hts[1:k+1]
    if (k == 0)
        return
    end

    @views ix = X.ix[2:k+1] # locations will be
    shuffle!(ix)            # updated in random order.
    for j in ix
        sp,s,sn = locs[j-1], locs[j], locs[j+1]
        hp,h = hts[j-1], hts[j]
        x = sp + rand()*(sn-sp)

        lratio = (x < s) ? llhdratio(Y,x,s,s,h,h,hp) :
                           llhdratio(Y,s,x,x,hp,hp,h)
        lratio += log((sn-x)/(sn-s) * (x-sp)/(s-sp))

        ratio = exp(max(-30,min(0,lratio)))
        if rand() < ratio
            locs[j] = x
            U.laccepts += 1
        end
        U.lprops += 1
    end
    loc_init(Y,S,X)          # leave the auxiliary
    loc_unconstrained(k,C,X) # variables in a
    ix .= 2:k+1              # consistent state.

end

#==== HMC =====================================================================#

function loc_hmc(Y, C::Constants, S::State, X::Auxiliary, U::Counts)

    k = S.k
    ε = .8C.leps + rand()*.4C.leps # base step size ± 20%.
    loc_init(Y,S,X)
    loc_unconstrained(k,C,X)
    loc_momenta(k,X)
    ep = loc_potential(k,C,X) # previous energy; also computes the gradient.
    kp = loc_kinetic(k,X)

    loc_pstep(.5ε,k,X)
    for i = 1:C.lleaps
        loc_qstep(ε,k,Y,C,X,U)
        if i != C.lleaps
            loc_potential(k,C,X)
            loc_pstep(ε,k,X)
        end
    end
    en = loc_potential(k,C,X)
    loc_pstep(.5ε,k,X)
    kn = loc_kinetic(k,X)

    p = rand()
    l = exp(ep+kp-en-kn)

    if p < l
        @views S.locs[1:k+2] .= X.locs[1:k+2]
        @views S.hts[1:k+1] .= X.hts[1:k+1] # order might have changed.
        U.laccepts += 1
    else
        aux_init(Y,C,S,X) # leave the aux. variables in a consistent state.
    end
    U.lprops += 1

end

function loc_pstep(s, k, X::Auxiliary)
    @. @views X.ptlocs[1:k] -= s * X.∇tlocs[1:k]
end

function loc_qstep(s, k, Y, C::Constants, X::Auxiliary, U::Counts)
    @views tlocs,ptlocs = X.tlocs[1:k], X.ptlocs[1:k]
    @. tlocs += s*ptlocs
    if !issorted(tlocs)
        @views ix = X.ix[1:k]
        sortperm!(ix, tlocs, initialized=true)
        permute!(tlocs, ix)
        permute!(ptlocs, ix)
        @views permute!(X.hts[2:k+1], ix)
        @views permute!(X.lhts[2:k+1], ix)
        ix .= 1:k
        U.lsorts += 1
    end
    loc_constrained(k,Y,C,X)
end

function loc_potential(k, C::Constants, X::Auxiliary)
    R,α,β = C.R, C.α, C.β
    @views locs,hts,lhts = X.locs[2:k+1], X.hts[1:k+1], X.lhts[1:k+1]
    @views widths,counts = X.widths[1:k+1], X.counts[1:k+1]

    e = sum(@. (widths+β)*hts - (counts+α)*lhts - log(widths)) -
        sum(@. log(locs*(R-locs)))
    X.∇tlocs[1:k] .= @. begin
        (@view(hts[1:end-1]) - @view(hts[2:end]) -
        1/@view(widths[1:end-1]) + 1/@view(widths[2:end]))*(locs*(R-locs)/R) -
        (1-2locs/R)
    end
    return e
end

loc_kinetic(k, X::Auxiliary) = @views .5 * sum(X.ptlocs[1:k].^2)

#==== Auxiliary ===============================================================#

function loc_init(Y, S::State, X::Auxiliary) # locs, widths, counts.
    k = S.k
    @views X.locs[1:k+2] .= S.locs[1:k+2]
    loc_metrics(k,Y,X)
end

function loc_metrics(k, Y, X::Auxiliary) # widths, counts.
    sp = 0.
    y = 1 # index of the first uncounted datum.
    for i = 1:k+1
        s = X.locs[i+1]
        X.widths[i] = s-sp
        n = searchsortedfirst(@view(Y[y:end]), s) - 1
        X.counts[i] = n
        sp,y = s,y+n
    end
end

function loc_unconstrained(k, C::Constants, X::Auxiliary) # tlocs, texps, dlocs.
    @views locs = X.locs[2:k+1]
    @. X.tlocs[1:k] = log(locs/(C.R-locs))
end

function loc_constrained(k, Y, C::Constants, X::Auxiliary) # locs, metrics.
    @. @views X.locs[2:k+1] = C.R/(1+exp(-X.tlocs[1:k]))
    loc_metrics(k,Y,X)
end

function loc_momenta(k, X::Auxiliary) # ptlocs.
    @views randn!(X.ptlocs[1:k])
end

#==== I/O =====================================================================#

function loc_println(k, X::Auxiliary)
    println("X: locs = $(string(X.locs[1:k+2])),",
              " hts = $(string(X.hts[1:k+1]))\n",
            "   widths = $(string(X.widths[1:k+1])),",
              " counts = $(X.counts[1:k+1])\n",
            "   tlocs = $(string(X.tlocs[1:k])),",
              " ∇tlocs = $(string(X.∇tlocs[1:k]))\n",
            "   ptlocs = $(string(X.ptlocs[1:k])), ix = $(X.ix[1:k])")
end

max_stars_repo_name: kevdur/rjcp

max_issues_repo_name: kevdur/rjcp

max_forks_repo_name: kevdur/rjcp


text:
## Instrumental variables problems

dt = readcsv("iv.csv");

y = dt[:,1];
x = dt[:,2];
z = dt[:,3:7];

f(theta) = z.*(y-x*theta);

gmms2 = GMMEstimator(f, [.0])
initialize!(gmm2s)
estimate!(gmm2s)

srand(1)
using MomentBasedEstimators
using ModelsGenerators
y,x,z = ModelsGenerators.sim_iv_d01(n = 10000, CP = 2000)
xb    = randn(10000,2)
x     = [x xb]
z     = [z xb]

f(theta) = z.*(y-x*theta);

gmm = MomentBasedEstimators.GMMEstimator(f, [.1, 0, 0]);
md = MomentBasedEstimators.MDEstimator(f, [.1, 0, 0]);

MomentBasedEstimators.initialize!(gmm);
MomentBasedEstimators.initialize!(md);

estimate!(gmm);
estimate!(md);

maximum(vcov(gmm)-vcov(md))<1e-07
maximum(vcov(md, false, :Unweighted) - vcov(md, false, :Weighted))<1e-05

MomentBasedEstimators.setparLB!(md, coef(md))
MomentBasedEstimators.setparUB!(md, coef(md))
estimate!(md);

MomentBasedEstimators.setparLB!(md, [0.,0,0])
MomentBasedEstimators.setparUB!(md, [0.,0,0])

## Methods

# 1. Coefficient
coef(gmm)
coef(md)

## 2. Variances
vcov(gmm)
vcov(md)

## 3. Objectiv value

srand(1)
y,x,z=randiv(100)
xb = randn(100,2)
x = [x xb]
z = [z xb]

z=sparse(z)

f(theta) = broadcast(*, z, (y-x*theta))

g = MomentBasedEstimators.GMMEstimator(f, [.1, 0, 0])
m = MomentBasedEstimators.MDEstimator(f, [.1, 0, 0])

MomentBasedEstimators.initialize!(g);
MomentBasedEstimators.initialize!(m);

@time estimate!(g);
@time estimate!(m);

srand(1)
y,x,z=randiv(100)
xb = randn(100,2)
x = [x xb]
z = [z xb]

f(theta) = z.*(y-x*theta);

g = MomentBasedEstimators.GMMEstimator(f, [.1, 0, 0])
m = MomentBasedEstimators.MDEstimator(f, [.1, 0, 0])

MomentBasedEstimators.initialize!(g)
MomentBasedEstimators.initialize!(m)

@time estimate!(g);
@time estimate!(m);

## Before
## elapsed time: 0.129503039 seconds (69486712 bytes allocated, 40.36% gc time)


text:
"""
    EmpiricalCopula{d,MT}

Fields:
  - u::MT - the matrix of observations.

Constructor

    EmpiricalCopula(u;pseudos=true)

The EmpiricalCopula in dimension ``d`` is parameterized by a pseudo-data matrix wich should have shape (d,N). Its expression is given as :

```math
C(\\mathbf x) = \\frac{1}{N}\\sum_{i=1}^n \\mathbf 1_{\\mathbf u_i \\le \\mathbf x}
```

This function is very practical, be be aware that this is not a true copula (since ``\\mathbf u`` are only pseudo-observations). The constructor allows you to pass dirctly pseudo-observations (the default) or will compute them for you. You can then compute the `cdf` of the copula, and sample it through the standard interface.
"""
struct EmpiricalCopula{d,MT} <: Copula{d}
    u::MT
end
Base.eltype(C::EmpiricalCopula{d,MT}) where {d,MT} = Base.eltype(C.u)
function EmpiricalCopula(u;pseudos=true)
    d = size(u,1)
    if !pseudos
        u = pseudos(u)
    else
        @assert all(0 .<= u .<= 1)
    end
    return EmpiricalCopula{d,typeof(u)}(u)
end
function Distributions.cdf(C::EmpiricalCopula{d,MT},u) where {d,MT}
   return mean(all(C.u .<= u,dims=1)) # might not be very efficient implementation.
end
function Distributions._rand!(rng::Distributions.AbstractRNG, C::EmpiricalCopula{d,MT}, x::AbstractVector{T}) where {d,MT,T<:Real}
    x .= C.u[:,Base.rand(rng,axes(C.u,2),1)[1]]
end
function Base.rand(rng::Distributions.AbstractRNG,C::EmpiricalCopula{d,MT}) where {d,MT}
    C.u[:,Base.rand(rng,axes(C.u,2),1)[1]]
end
function Distributions.fit(::Type{CT},u) where {CT <: EmpiricalCopula}
    return EmpiricalCopula(u)
end


text:
include(joinpath(pwd(), "src/models/car.jl"))
include(joinpath(pwd(), "src/constraints/obstacles.jl"))

optimize = true

# Horizon
T = 51

# Time step
tf = 1.0
h = tf / (T - 1)

# Bounds

# ul <= u <= uu
ul, uu = control_bounds(model, T, -3.0, 3.0)

# Initial and final states
x1 = [0.0; 0.0; 0.0]
xT = [1.0; 1.0; 0.0]
xl, xu = state_bounds(model, T, x1 = x1, xT = xT)

# Objective
obj = quadratic_tracking_objective(
        [t < T ? Diagonal(ones(model.n)) : Diagonal(10.0 * ones(model.n)) for t = 1:T],
        [Diagonal(1.0e-1 * ones(model.m)) for t = 1:T-1],
        [xT for t = 1:T], [zeros(model.m) for t = 1:T])

# Constraints
circles = [(0.85, 0.3, 0.1),
           (0.375, 0.75, 0.1),
           (0.25, 0.2, 0.1),
           (0.75, 0.8, 0.1)]

# Constraints
function obstacles!(c, x)
    c[1] = circle_obs(x[1], x[2], circles[1][1], circles[1][2], circles[1][3])
    c[2] = circle_obs(x[1], x[2], circles[2][1], circles[2][2], circles[2][3])
    c[3] = circle_obs(x[1], x[2], circles[3][1], circles[3][2], circles[3][3])
    c[4] = circle_obs(x[1], x[2], circles[4][1], circles[4][2], circles[4][3])
    nothing
end

n_stage = 4
n_con = n_stage * T
con_obstacles = ObstacleConstraints(n_con, (1:n_con), n_stage)

# Problem
prob = trajectory_optimization_problem(model,
           obj,
           T,
           h = h,
           xl = xl,
           xu = xu,
           ul = ul,
           uu = uu,
           con = con_obstacles
           )

# Trajectory initialization
X0 = linear_interp(x1, xT, T) # linear interpolation on state
U0 = random_controls(model, T, 0.001) # random controls

# Pack trajectories into vector
Z0 = pack(X0, U0, prob)

#NOTE: may need to run examples multiple times to get good trajectories
# Solve nominal problem
if optimize
    @time Z̄ = solve(prob, copy(Z0))
    @save joinpath(@__DIR__, "sol.jld2") Z̄
else
    println("Loading solution...")
    @load joinpath(@__DIR__, "sol.jld2") Z̄
end

max_stars_repo_name: zacmanchester/DirectMotionPlanning

max_issues_repo_name: zacmanchester/DirectMotionPlanning

max_forks_repo_name: zacmanchester/DirectMotionPlanning


text:
using VariantsNMF
using Test

#### simulate
V = get_one_simulated_V()

@testset "nmf_FI" begin
    #### params for rank selection
    rc_params = RCParams(
        K_min = 2,
        K_max = 6,
        n_iter = 10
    )

    #### nmf β-divergence, multiplicative updates
    nmf_global_params = NMFParams(
        init          = :random,
        dist          = truncated(Normal(1, 1), 0, Inf),
        max_iter      = 1_000,
        stopping_crit = :conn,
        stopping_iter = 10,
        verbose       = false,
    )

    nmf_local_params = FIParams(
        β         = 1,
        l₁ratio_H = 0,
        α_H       = 0,
        alg       = :mm
    )

    #### NMF struct
    nmf = NMF(
        solver        = nmf_FI,
        global_params = nmf_global_params,
        local_params  = nmf_local_params
    )

    #### get consensus metrics for each rank
    @time rc_results = rank_by_consensus(V, rc_params, nmf)

    @test "rank" in names(rc_results.df_metrics)
    @test "cophenetic" in names(rc_results.df_metrics)
    @test "dispersion" in names(rc_results.df_metrics)
    @test "silhouette" in names(rc_results.df_metrics)
    @test "sparse_W" in names(rc_results.df_metrics)
    @test "sparse_H" in names(rc_results.df_metrics)
end


text:
"""
    ULogarithmic(x)

Represents the positive real number `x` by storing its logarithm.

!!! tip
    If you know `logx=log(x)` then use [`exp(ULogarithmic, logx)`](@ref exp(::Type{ULogarithmic},::Real)) instead.
"""
struct ULogarithmic{T} <: Real
  log::T
  Base.exp(::Type{ULogarithmic{T}}, x::T) where {T<:Real} = new{T}(x)
end

"""
    Logarithmic(x)

Represents the real number `x` by storing its absolute value as a [`ULogarithmic`](@ref) and its sign bit.

!!! tip
    If you know `logx=log(abs(x))` then use [`exp(Logarithmic, logx, signbit)`](@ref exp(::Type{Logarithmic},::Real)) instead.
"""
struct Logarithmic{T} <: Real
  abs::ULogarithmic{T}
  signbit::Bool
  Logarithmic{T}(abs::ULogarithmic{T}, signbit::Bool=false) where {T<:Real} = new{T}(abs, signbit)
end

"""
    CLogarithmic(x)

Represents the complex number `x` by storing its absolute value as a [`ULogarithmic`](@ref) and its angle.

!!! tip
    If you know `logx=log(abs(x))` then use [`exp(ComplexLogarithmic, logx, angle)`](@ref exp(::Type{ComplexLogarithmic},::Real)) instead.
"""
struct CLogarithmic{T} <: Number
  abs :: ULogarithmic{T}
  angle :: T
  CLogarithmic{T}(abs::ULogarithmic{T}, angle::T=zero(T)) where {T<:Real} = new{T}(abs, angle)
end

# for convenience
const AnyLog{T} = Union{ULogarithmic{T}, Logarithmic{T}, CLogarithmic{T}}
const RealLog{T} = Union{ULogarithmic{T}, Logarithmic{T}}


text:
"""
    AliveCallback(analysis_interval=0, alive_interval=analysis_interval÷10)

Inexpensive callback showing that a simulation is still running by printing
some information such as the current time to the screen every `alive_interval`
time steps. If `analysis_interval ≂̸ 0`, the output is omitted every
`analysis_interval` time steps.
"""
mutable struct AliveCallback
  start_time::Float64
  alive_interval::Int
  analysis_interval::Int
end

function AliveCallback(; analysis_interval=0,
                         alive_interval=analysis_interval÷10)

  alive_callback = AliveCallback(0.0, alive_interval, analysis_interval)

  DiscreteCallback(alive_callback, alive_callback, # the first one is the condition, the second the affect!
                   save_positions=(false,false),
                   initialize=initialize!)
end

function Base.show(io::IO, cb::DiscreteCallback{Condition,Affect!}) where {Condition, Affect!<:AliveCallback}
  alive_callback = cb.affect!
  print(io, "AliveCallback(alive_interval=", alive_callback.alive_interval, ")")
end

function Base.show(io::IO, ::MIME"text/plain", cb::DiscreteCallback{Condition,Affect!}) where {Condition, Affect!<:AliveCallback}
  if get(io, :compact, false)
    show(io, cb)
  else
    alive_callback = cb.affect!

    setup = [
             "interval" => alive_callback.alive_interval,
            ]
    summary_box(io, "AliveCallback", setup)
  end
end

function initialize!(cb::DiscreteCallback{Condition,Affect!}, u, t, integrator) where {Condition, Affect!<:AliveCallback}

  alive_callback = cb.affect!
  alive_callback.start_time = time_ns()
  return nothing
end

# this method is called to determine whether the callback should be activated
function (alive_callback::AliveCallback)(u, t, integrator)
  @unpack alive_interval, analysis_interval = alive_callback

  return alive_interval > 0 && (
    (integrator.iter % alive_interval == 0 && (analysis_interval == 0 || integrator.iter % analysis_interval != 0)) ||
    isfinished(integrator))
end

# this method is called when the callback is activated
function (alive_callback::AliveCallback)(integrator)
  @unpack t, dt, iter = integrator

  # Checking for floating point equality is OK here as `DifferentialEquations.jl`
  # sets the time exactly to the final time in the last iteration
  if isfinished(integrator) && mpi_isroot()
    println("─"^100)
    println("Trixi simulation run finished.    Final time: ", integrator.t, "    Time steps: ", integrator.iter)
    println("─"^100)
    println()
  elseif mpi_isroot()
    runtime_absolute = 1.0e-9 * (time_ns() - alive_callback.start_time)
    @printf("#timesteps: %6d │ Δt: %.4e │ sim. time: %.4e │ run time: %.4e s\n",
            iter, dt, t, runtime_absolute)
  end

  # avoid re-evaluating possible FSAL stages
  u_modified!(integrator, false)
  return nothing
end


text:
#An inspect Reward to find a solution in a few number of steps, in the case when one episode corresponds to finind a solution.

mutable struct InspectReward <: SeaPearl.AbstractReward
    value::Float32
end

InspectReward(model::SeaPearl.CPModel) = InspectReward(0)

function SeaPearl.set_reward!(::Type{SeaPearl.StepPhase}, lh::SeaPearl.LearnedHeuristic{SR, InspectReward, O}, model::SeaPearl.CPModel, symbol::Union{Nothing, Symbol}) where {
    SR <: SeaPearl.AbstractStateRepresentation,
    O <: SeaPearl.ActionOutput
}
    nothing
end

function SeaPearl.set_reward!(::Type{SeaPearl.DecisionPhase}, lh::SeaPearl.LearnedHeuristic{SR, InspectReward, O}, model::SeaPearl.CPModel) where {
    SR <: SeaPearl.AbstractStateRepresentation,
    O <: SeaPearl.ActionOutput
}
    lh.reward.value += -1
    nothing
end

function SeaPearl.set_reward!(::Type{SeaPearl.EndingPhase}, lh::SeaPearl.LearnedHeuristic{SR, InspectReward, O}, model::SeaPearl.CPModel, symbol::Union{Nothing, Symbol}) where {
    SR <: SeaPearl.AbstractStateRepresentation,
    O <: SeaPearl.ActionOutput
}
    if symbol == :Feasible || symbol == :FoundSolution
        lh.reward.value += 200
    else
        lh.reward.value-=100
    end
end

max_stars_repo_name: corail-research/SeaPearlZoo

max_issues_repo_name: corail-research/SeaPearlZoo

max_forks_repo_name: corail-research/SeaPearlZoo