text:
SRC_DIR = joinpath(Base.source_dir(), "..", "..", "..", "src")
OUT_DIR = joinpath(Base.source_dir(), "..", "..", "..", "output10")
mkpath(OUT_DIR)

#include(joinpath(SRC_DIR, "BridgeSDEInference.jl"))
#using Main.BridgeSDEInference
include(joinpath(SRC_DIR, "BridgeSDEInference_for_tests.jl"))

using StaticArrays, LinearAlgebra, GaussianDistributions
using Distributions
using Random
using Bridge
using CSV, DataFrames
DIR = "auxiliary"
include(joinpath(SRC_DIR, DIR, "data_simulation_fns.jl"))
include(joinpath(SRC_DIR, DIR, "utility_functions.jl"))
include(joinpath(SRC_DIR, DIR, "plotting_fns.jl"))

#=============================================================================
                Routines for setting up the MCMC sampler
=============================================================================#
function _prepare_setup(updt_order, ρ=0.96, num_mcmc_steps=4*10^3,
                        do_blocking=false, b_step=2, num_steps_for_change_pt=10,
                        thin_path=10^0, save_path_every=1*10^3)
    model_setup = DiffusionSetup(Pˣ, P̃, PartObs())
    set_observations!(model_setup, [L for _ in P̃], [Σ for _ in P̃], obs_vals,
                      obs_time)
    set_imputation_grid!(model_setup, 1/500)
    set_x0_prior!(model_setup,
                  GsnStartingPt(x0, @SMatrix [400.0 0.0 0.0;
                                                0.0 20.0 0.0;
                                                0.0 0.0 20.0]), # prior over starting point
                  x0)
    set_auxiliary!(model_setup; skip_for_save=thin_path,
                   adaptive_prop=Adaptation(x0, [0.7, 0.4, 0.2, 0.2, 0.2],
                                                [500, 500, 500, 500, 500],
                                            1))
    initialise!(eltype(x0), model_setup, Vern7(), false, NoChangePt(num_steps_for_change_pt))

    blocks = (NoBlocking(), NoBlocking())
    if do_blocking
        blocks = create_blocks( ChequeredBlocking(), model_setup.P,
                                (knots=collect(1:length(obs_vals)-2)[1:b_step:end],
                                ϵ=10^(-10),
                                change_pt=SimpleChangePt(num_steps_for_change_pt)))
    end
    mcmc_setup = MCMCSetup(
        Imputation(blocks[1], ρ, Vern7()),
        Imputation(blocks[2], ρ, Vern7()),
        ParamUpdate(ConjugateUpdt(), [1,2,3], fill(0.0, 4), nothing,
                    MvNormal(fill(0.0, 3), diagm(0=>fill(1000.0, 3))),
                    UpdtAuxiliary(Vern7(), check_if_recompute_ODEs(P̃, [1,2,3])))
    )
    schedule = MCMCSchedule(num_mcmc_steps, updt_order,
                            (save=save_path_every, verbose=10^2, warm_up=0,
                             readjust=(x->x%100==0), fuse=(x->false)) )
    mcmc_setup, schedule, model_setup
end

#==============================================================================
                        Auxiliary routines for saving
==============================================================================#

function save_paths(tt, paths, filename)
    d = length(paths[1][1])
    xx = copy(tt)
    for i in 1:length(paths)
        for j in 1:d
            xx = hcat(xx, [p[j] for p in paths[i]])
        end
    end
    CSV.write(joinpath(OUT_DIR, filename), DataFrame(xx))
end

function save_marginals(tt, paths, filename, indices)
    d = length(paths[1][1])
    xx = copy(tt[indices])
    N = length(paths)
    stride = div(N, 100)
    for i in 1:length(paths)
        if i % stride == 0
            print(div(i, 100), "% done...\n")
        end
        for j in 1:d
            xx = hcat(xx, [p[j] for p in paths[i][indices]])
        end
    end
    CSV.write(joinpath(OUT_DIR, filename), DataFrame(xx'))
end

function save_param_chain(chain, filename)
    d = length(chain[1])
    xx = [c[1] for c in chain]
    for i in 2:d
        xx = hcat(xx, [c[i] for c in chain])
    end
    CSV.write(joinpath(OUT_DIR, filename), DataFrame(xx))
end

function save_obs(tt, paths, filename)
    d = length(paths[1][1])
    xx = copy(tt)
    for i in 1:length(paths)
        for j in 1:d
            xx = hcat(xx, [p[j] for p in paths[i]])
        end
    end
    CSV.write(joinpath(OUT_DIR, filename), DataFrame(xx))
end

function save_history(to_save, filename)
    n, k = length(to_save), length(to_save[1])
    xx = zeros(n, k)
    for i in 1:n
        xx[i,1:end] .= to_save[i]
    end
    CSV.write(joinpath(OUT_DIR, filename), DataFrame(xx))
end

#==============================================================================
                        Generate densely observed process
==============================================================================#

Random.seed!(4)
θˣ = [10.0, 28.0, 8.0/3.0, 3.0]
#Pˣ = LorenzCV(θˣ...)
Pˣ = LorenzHypo(θˣ...)

x0, dt, T = ℝ{3}(1.5, -1.5, 25.0), 1/5000, 2.0
tt = 0.0:dt:T
XX, _ = simulate_segment(ℝ{2}(0.0, 0.0), x0, Pˣ, tt)

Σdiagel = 5*10^0
Σ = SMatrix{2,2}(1.0I)*Σdiagel
L = @SMatrix[0.0 1.0 0.0;
             0.0 0.0 1.0]

skip_me = 50
obs_time, obs_vals = XX.tt[1:skip_me:end], [rand(Gaussian(L*x, Σ)) for x in XX.yy[1:skip_me:end]]

θ_init = [5.0, 15.0, 6.0, 3.0]
aux_flag = Val{(false,true,true)}()
P̃ = [LorenzCVAux(θ_init..., t₀, u, T, v, aux_flag, x0[3]) for (t₀, T, u, v)
     in zip(obs_time[1:end-1], obs_time[2:end], obs_vals[1:end-1], obs_vals[2:end])]
Pˣ = LorenzCV(θ_init...)

P̃ = [LorenzHypoAux(θ_init..., t₀, u, T, v, aux_flag, x0[3]) for (t₀, T, u, v)
     in zip(obs_time[1:end-1], obs_time[2:end], obs_vals[1:end-1], obs_vals[2:end])]
Pˣ = LorenzHypo(θ_init...)

print("mid-point time: ", XX.tt[div(length(XX.yy),2)+1], ", mid-point value: ",
      XX.yy[div(length(XX.yy),2)+1], "\n")
print("three-quarters time: ", XX.tt[3*div(length(XX.yy),4)+1], ", three-quarters value: ",
      XX.yy[3*div(length(XX.yy),4)+1], "\n")

save_paths(obs_time, [obs_vals], "many_obs_data.csv")

#==============================================================================
                    Run the expriment: very large number of blocks
==============================================================================#
setup = _prepare_setup([[1,3],[2,3]], 0.1, 1*10^5, true, 1, 6, 1, 10^2)
Random.seed!(4)
out, elapsed = @timeit mcmc(setup...)

#plot_acceptance([out[2].updates[1],out[2].updates[2]], [500, 500, 500, 500, 500])
plot_chains(out[2]; truth=[10.0, 28.0, 8.0/3.0, 3.0],
            ylims=[nothing, (25,30), (2,5), (0,10)])
plot_paths(out[1], out[2], setup[2]; obs=(times=obs_time[2:end],
                     vals=[[v[1] for v in obs_vals[2:end]],
                           [v[2] for v in obs_vals[2:end]]], indices=[2,3]))

save_paths(out[1].time, out[1].paths, "many_obs_many_blocks_paths.csv")
save_param_chain(out[2].θ_chain, "many_obs_many_blocks_chain.csv")
save_history(out[2].updates[1].accpt_history, "many_obs_many_blocks_accpt_hist_1.csv")
save_history(out[2].updates[2].accpt_history, "many_obs_many_blocks_accpt_hist_2.csv")
save_history(out[2].updates[1].ρs_history, "many_obs_many_blocks_rho_hist_1.csv")
save_history(out[2].updates[2].ρs_history, "many_obs_many_blocks_rho_hist_2.csv")

# repeat, this time save marginals
setup = _prepare_setup([[1,3],[2,3]], 0.1, 1*10^5, true, 1, 6, 1000, 1)
Random.seed!(4)
out, elapsed = @timeit mcmc(setup...)
display(out.accpt_tracker)
#_temp = 0
#save_marginals(out[1].time, out[1].paths[1+_temp*10^4:(_temp+1)*10^4],
#               join(["many_obs_many_blocks_marginals_", string(_temp), ".csv"]),
#               [1,50,100,150,200])
save_marginals(out[1].time, out[1].paths, "many_obs_many_blocks_marginals.csv"),
               [1,50,100,150,200])
#==============================================================================
                    Run the expriment: longer blocks (less of them)
==============================================================================#
setup = _prepare_setup([[1,3],[2,3]], 0.5, 1*10^5, true, 4, 6, 1, 10^2)
Random.seed!(4)
out, elapsed = @timeit mcmc(setup...)

plot_chains(out[2]; truth=[10.0, 28.0, 8.0/3.0, 3.0],
            ylims=[nothing, (25,30), (2,5), (0,10)])
plot_paths(out[1], out[2], setup[2]; obs=(times=obs_time[2:end],
                     vals=[[v[1] for v in obs_vals[2:end]],
                           [v[2] for v in obs_vals[2:end]]], indices=[2,3]))

save_paths(out[1].time, out[1].paths, "many_obs_medium_blocks_paths.csv")
save_param_chain(out[2].θ_chain, "many_obs_medium_blocks_chain.csv")
save_history(out[2].updates[1].accpt_history, "many_obs_medium_blocks_accpt_hist_1.csv")
save_history(out[2].updates[2].accpt_history, "many_obs_medium_blocks_accpt_hist_2.csv")
save_history(out[2].updates[1].ρs_history, "many_obs_medium_blocks_rho_hist_1.csv")
save_history(out[2].updates[2].ρs_history, "many_obs_medium_blocks_rho_hist_2.csv")

setup = _prepare_setup([[1,3],[2,3]], 0.5, 1*10^5, true, 4, 6, 1000, 1)
Random.seed!(4)
out, elapsed = @timeit mcmc(setup...)
save_marginals(out[1].time, out[1].paths, "many_obs_medium_blocks_marginals.csv", [1,50,100,150,200])
#_temp = 9
#save_marginals(out[1].time, out[1].paths[1+_temp*10^4:(_temp+1)*10^4],
#               join(["many_obs_medium_blocks_marginals_", string(_temp), ".csv"]),
#               [1,50,100,150,200])
#==============================================================================
                    Run the expriment: very long blocks
==============================================================================#
setup = _prepare_setup([[1,3]], 0.5, 1*10^5, false, 40, 6, 1, 10^2)
Random.seed!(4)
out, elapsed = @timeit mcmc(setup...)
display(out.accpt_tracker)

plot_chains(out[2]; truth=[10.0, 28.0, 8.0/3.0, 3.0],
            ylims=[nothing, (25,30), (2,5), (0,10)])
plot_paths(out[1], out[2], setup[2]; obs=(times=obs_time[2:end],
                     vals=[[v[1] for v in obs_vals[2:end]],
                           [v[2] for v in obs_vals[2:end]]], indices=[2,3]))

save_paths(out[1].time, out[1].paths, "many_obs_long_blocks_paths.csv")
save_param_chain(out[2].θ_chain, "many_obs_long_blocks_chain.csv")
save_history(out[2].updates[1].accpt_history, "many_obs_long_blocks_accpt_hist_1.csv")
save_history(out[2].updates[1].ρs_history, "many_obs_long_blocks_rho_hist_1.csv")

setup = _prepare_setup([[1,3]], 0.5, 1*10^5, false, 40, 6, 1000, 1)
Random.seed!(4)
out, elapsed = @timeit mcmc(setup...)
save_marginals(out[1].time, out[1].paths, "many_obs_long_blocks_marginals.csv", [1,50,100,150,200])
#_temp = 9
#save_marginals(out[1].time, out[1].paths[1+_temp*10^4:(_temp+1)*10^4],
#               join(["many_obs_long_blocks_marginals_", string(_temp), ".csv"]),
#               [1,50,100,150,200])

#==============================================================================
                        Generate sparsely observed process
==============================================================================#

Random.seed!(4)
θˣ = [10.0, 28.0, 8.0/3.0, 3.0]
Pˣ = LorenzCV(θˣ...)

x0, dt, T = ℝ{3}(1.5, -1.5, 25.0), 1/5000, 2.0
tt = 0.0:dt:T
XX, _ = simulate_segment(ℝ{3}(0.0, 0.0, 0.0), x0, Pˣ, tt)

Σdiagel = 5*10^(-2)
Σ = SMatrix{2,2}(1.0I)*Σdiagel
L = @SMatrix[0.0 1.0 0.0;
             0.0 0.0 1.0]

skip_me = 1000
obs_time, obs_vals = XX.tt[1:skip_me:end], [rand(Gaussian(L*x, Σ)) for x in XX.yy[1:skip_me:end]]

θ_init = [5.0, 15.0, 6.0, 3.0]
aux_flag = Val{(false,true,true)}()
P̃ = [LorenzCVAux(θ_init..., t₀, u, T, v, aux_flag, x0[3]) for (t₀, T, u, v)
     in zip(obs_time[1:end-1], obs_time[2:end], obs_vals[1:end-1], obs_vals[2:end])]
Pˣ = LorenzCV(θ_init...)
print("mid-point time: ", XX.tt[div(length(XX.yy),2)+1], ", mid-point value: ",
      XX.yy[div(length(XX.yy),2)+1], "\n")
print("three-quarters time: ", XX.tt[3*div(length(XX.yy),4)+1], ", three-quarters value: ",
      XX.yy[3*div(length(XX.yy),4)+1], "\n")
save_paths(obs_time, [obs_vals], "few_obs_data.csv")

#==============================================================================
                    Run the expriment: very long blocks
==============================================================================#
setup = _prepare_setup([[1,3]], 0.96, 1*10^5, false, 40, 15, 1, 10^2)
Random.seed!(4)
out, elapsed = @timeit mcmc(setup...)

plot_chains(out[2]; truth=[10.0, 28.0, 8.0/3.0, 3.0],
            ylims=[nothing, nothing, nothing, (0,10)])
plot_paths(out[1], out[2], setup[2]; obs=(times=obs_time[2:end],
                     vals=[[v[1] for v in obs_vals[2:end]],
                           [v[2] for v in obs_vals[2:end]]], indices=[2,3]))

save_paths(out[1].time, out[1].paths, "few_obs_long_blocks_paths.csv")
save_param_chain(out[2].θ_chain, "few_obs_long_blocks_chain.csv")
save_history(out[2].updates[1].accpt_history, "few_obs_long_blocks_accpt_hist_1.csv")
save_history(out[2].updates[1].ρs_history, "few_obs_long_blocks_rho_hist_1.csv")

setup = _prepare_setup([[1,3]], 0.96, 1*10^5, false, 4, 15, 1000, 1)
Random.seed!(4)
out, elapsed = @timeit mcmc(setup...)
save_marginals(out[1].time, out[1].paths, "few_obs_long_blocks_marginals.csv", [1,3,6,9,10])
_temp = 9
save_marginals(out[1].time, out[1].paths[1+_temp*10^4:(_temp+1)*10^4],
               join(["few_obs_long_blocks_marginals_", string(_temp), ".csv"]),
               [1,3,6,9,10])

#==============================================================================
                    Run the expriment: very large number of blocks
==============================================================================#
setup = _prepare_setup([[1,3],[2,3]], 0.9, 1*10^5, true, 1, 100, 1, 10^2)
Random.seed!(4)
out, elapsed = @timeit mcmc(setup...)

plot_chains(out[2]; truth=[10.0, 28.0, 8.0/3.0, 3.0],
            ylims=[nothing, (25,30), (2,5), (0,10)])
plot_paths(out[1], out[2], setup[2]; obs=(times=obs_time[2:end],
                     vals=[[v[1] for v in obs_vals[2:end]],
                           [v[2] for v in obs_vals[2:end]]], indices=[2,3]))

save_paths(out[1].time, out[1].paths, "few_obs_many_blocks_paths.csv")
save_param_chain(out[2].θ_chain, "few_obs_many_blocks_chain.csv")
save_history(out[2].updates[1].accpt_history, "few_obs_many_blocks_accpt_hist_1.csv")
save_history(out[2].updates[2].accpt_history, "few_obs_many_blocks_accpt_hist_2.csv")
save_history(out[2].updates[1].ρs_history, "few_obs_many_blocks_rho_hist_1.csv")
save_history(out[2].updates[2].ρs_history, "few_obs_many_blocks_rho_hist_2.csv")

# repeat, this time save marginals
setup = _prepare_setup([[1,3],[2,3]], 0.9, 1*10^5, true, 1, 100, 1000, 1)
Random.seed!(4)
out, elapsed = @timeit mcmc(setup...)
save_marginals(out[1].time, out[1].paths, "few_obs_many_blocks_marginals.csv", [1,3,6,9,10])
_temp = 9
save_marginals(out[1].time, out[1].paths[1+_temp*10^4:(_temp+1)*10^4],
               join(["few_obs_many_blocks_marginals_", string(_temp), ".csv"]),
               [1,3,6,9,10])
#==============================================================================
            Run the expriment: very large number of blocks and adaptive
#NOTE because adaptation is currently implemented in a rather inefficient way
# the code responsible for it is by default **COMMENTED OUT**. Consequently,
# to run the code with adaptation one needs to manually remove `#` from the file
# `src/examples/lorenz_system_const_vola.jl` in line 104--105. Run the code
# below and comment it out again.
==============================================================================#
setup = _prepare_setup([[1,3],[2,3]], 0.9, 1*10^5, true, 1, 100, 1, 10^2)
Random.seed!(4)
out, elapsed = @timeit mcmc(setup...)

plot_chains(out[2]; truth=[10.0, 28.0, 8.0/3.0, 3.0],
            ylims=[nothing, (25,30), (2,5), (0,10)])
plot_paths(out[1], out[2], setup[2]; obs=(times=obs_time[2:end],
                     vals=[[v[1] for v in obs_vals[2:end]],
                           [v[2] for v in obs_vals[2:end]]], indices=[2,3]))

save_paths(out[1].time, out[1].paths, "few_obs_many_blocks_and_adpt_paths.csv")
save_param_chain(out[2].θ_chain, "few_obs_many_blocks_and_adpt_chain.csv")
save_history(out[2].updates[1].accpt_history, "few_obs_many_blocks_and_adpt_accpt_hist_1.csv")
save_history(out[2].updates[2].accpt_history, "few_obs_many_blocks_and_adpt_accpt_hist_2.csv")
save_history(out[2].updates[1].ρs_history, "few_obs_many_blocks_and_adpt_rho_hist_1.csv")
save_history(out[2].updates[2].ρs_history, "few_obs_many_blocks_and_adpt_rho_hist_2.csv")

setup = _prepare_setup([[1,3],[2,3]], 0.9, 1*10^5, true, 1, 100, 1000, 1)
Random.seed!(4)
out, elapsed = @timeit mcmc(setup...)
save_marginals(out[1].time, out[1].paths, "few_obs_many_blocks_and_adpt_marginals.csv", [1,3,6,9,10])
_temp = 9
save_marginals(out[1].time, out[1].paths[1+_temp*10^4:(_temp+1)*10^4],
               join(["few_obs_many_blocks_and_adpt_marginals_", string(_temp), ".csv"]),
               [1,3,6,9,10])

#==============================================================================
            Run the expriment: very large number of blocks and adaptive
#NOTE needs to be done by manually uncommenting the definitions of B₀ and β₀
==============================================================================#
setup = _prepare_setup([[1,3]], 0.96, 1*10^5, false, 40, 15, 1, 10^2)
Random.seed!(4)
out, elapsed = @timeit mcmc(setup...)

plot_chains(out[2]; truth=[10.0, 28.0, 8.0/3.0, 3.0],
            ylims=[nothing, nothing, nothing, (0,10)])
plot_paths(out[1], out[2], setup[2]; obs=(times=obs_time[2:end],
                     vals=[[v[1] for v in obs_vals[2:end]],
                           [v[2] for v in obs_vals[2:end]]], indices=[2,3]))

save_paths(out[1].time, out[1].paths, "few_obs_long_blocks_zero_drift_paths.csv")
save_param_chain(out[2].θ_chain, "few_obs_long_blocks_zero_drift_chain.csv")
save_history(out[2].updates[1].accpt_history, "few_obs_long_blocks_zero_drift_accpt_hist_1.csv")
save_history(out[2].updates[1].ρs_history, "few_obs_long_blocks_zero_drift_rho_hist_1.csv")

setup = _prepare_setup([[1,3]], 0.96, 1*10^5, false, 4, 15, 1000, 1)
Random.seed!(4)
out, elapsed = @timeit mcmc(setup...)
save_marginals(out[1].time, out[1].paths, "few_obs_long_blocks_zero_drift_marginals.csv", [1,3,6,9,10])
_temp = 9
save_marginals(out[1].time, out[1].paths[1+_temp*10^4:(_temp+1)*10^4],
               join(["few_obs_long_blocks_zero_drift_marginals_", string(_temp), ".csv"]),
               [1,3,6,9,10])


text:
function partwise(A::SparseMatrixCSC{Tv, Ti}, Π::MapPartition{Ti}) where {Tv, Ti}
    @inbounds begin
        m, n = size(A)
        N = nnz(A)
        pos = A.colptr
        idx = A.rowval
        val = A.nzval
        K = length(Π)

        Πos = zeros(Ti, K + 1)
        πos = zeros(Ti, K + 1)
        hst = zeros(Ti, K)
        idx′ = undefs(Ti, N)
        val′ = undefs(Tv, N)

        for j = 1:n
            for q in pos[j] : pos[j + 1] - 1
                i = idx[q]
                k = Π.asg[i]
                πos[k + 1] += hst[k] != j
                hst[k] = j
                k₀ = k
                Πos[k + 1] += 1
            end
        end

        q = 1
        j′ = 1
        for k = 1:(K + 1)
            (Πos[k], q) = (q, q + Πos[k])
            (πos[k], j′) = (j′, j′ + πos[k])
        end
        n′ = j′ - 1

        pos′ = undefs(Ti, n′ + 1)
        prm = zeros(Ti, n′)
        zero!(hst)

        for j = 1:n
            for q in pos[j] : pos[j + 1] - 1
                i = idx[q]
                k = Π.asg[i]
                q′ = Πos[k + 1]
                idx′[q′] = i
                val′[q′] = val[q]
                Πos[k + 1] = q′ + 1
                if hst[k] != j
                    j′ = πos[k + 1]
                    pos′[j′] = q′
                    prm[j′] = j
                    πos[k + 1] = j′ + 1
                end
                hst[k] = j
            end
        end
        pos′[n′ + 1] = N + 1

        return ((n, K, n′, πos, prm), SparseMatrixCSC(m, n′, pos′, idx′, val′))
    end
end

struct PartwiseCount{Ti <: Integer, Ocl}
    n::Int
    K::Int
    n′::Int
    πos::Vector{Ti}
    prm::Vector{Ti}
    ocl::Ocl
end

Base.size(arg::PartwiseCount) = (arg.n + 1, arg.n + 1, arg.K)

partwisecount!(args...; kwargs...) = partwisecount!(NoHint(), args...; kwargs...)
partwisecount!(::AbstractHint, args...; kwargs...) = @assert false
partwisecount!(hint::AbstractHint, n, K, n′, πos, prm, ocl; kwargs...) =
    PartwiseCount(hint, n, K, n′, πos, prm, ocl; kwargs...)

PartwiseCount(hint::AbstractHint, n, K, n′, πos::Vector{Ti}, prm::Vector{Ti}, ocl::Ocl; kwargs...) where {Ti, Ocl} =
    PartwiseCount{Ti, Ocl}(hint, n, K, n′, πos, prm, ocl; kwargs...)

PartwiseCount{Ti, Ocl}(hint::AbstractHint, n, K, n′, πos::Vector{Ti}, prm::Vector{Ti}, ocl::Ocl; kwargs...) where {Ti, Ocl} =
    PartwiseCount{Ti, Ocl}(n, K, n′, πos, prm, ocl)

Base.getindex(arg::PartwiseCount{Ti}, j::Integer, j′::Integer, k::Integer) where {Ti} = arg(j, j′, k)
function (arg::PartwiseCount{Ti})(j::Integer, j′::Integer, k::Integer) where {Ti}
    @inbounds begin
        tmp = @view arg.prm[arg.πos[k] : arg.πos[k + 1] - 1]
        rnk_j = arg.πos[k] + searchsortedfirst(tmp, j) - 1
        rnk_j′ = arg.πos[k] + searchsortedfirst(tmp, j′) - 1
        return arg.ocl(rnk_j, rnk_j′)
    end
end

mutable struct PartwiseCountStepper{Ti, Ocl} <: AbstractArray{Ti, 3}
    rnk_j::Ti
    rnk_j′::Ti
    cst::PartwiseCount{Ti, Ocl}
end

Base.size(arg::PartwiseCountStepper) = size(arg.cst)

partwisecount!(hint::StepHint, n, K, n′, πos, prm, ocl; kwargs...) =
    PartwiseCountStepper(hint, n, K, n′, πos, prm, ocl; kwargs...)

PartwiseCountStepper(hint::AbstractHint, n, K, n′, πos::Vector{Ti}, prm::Vector{Ti}, ocl::Ocl; kwargs...) where {Ti, Ocl} =
    PartwiseCountStepper{Ti, Ocl}(hint, n, K, n′, πos, prm, ocl; kwargs...)

function PartwiseCountStepper{Ti, Ocl}(hint::AbstractHint, n, K, n′, πos::Vector{Ti}, prm::Vector{Ti}, ocl::Ocl; kwargs...) where {Ti, Ocl}
    rnk_j = 0
    rnk_j′ = 0
    cst = PartwiseCount{Ti, Ocl}(hint, n, K, n′, πos::Vector{Ti}, prm::Vector{Ti}, ocl::Ocl; kwargs...)
    return PartwiseCountStepper{Ti, Ocl}(rnk_j, rnk_j′, cst)
end

Base.getindex(arg::PartwiseCountStepper{Ti}, j::Integer, j′::Integer, k::Integer) where {Ti} = arg(j, j′, k)
function (arg::PartwiseCountStepper{Ti})(j::Integer, j′::Integer, k::Integer) where {Ti}
    @inbounds begin
        tmp = @view arg.cst.prm[arg.cst.πos[k] : arg.cst.πos[k + 1] - 1]
        rnk_j = arg.cst.πos[k] + searchsortedfirst(tmp, j) - 1
        arg.rnk_j = rnk_j
        rnk_j′ = arg.cst.πos[k] + searchsortedfirst(tmp, j′) - 1
        arg.rnk_j′ = rnk_j′
        return arg.cst.ocl(rnk_j, rnk_j′)
    end
end

@propagate_inbounds function (stp::Step{Net})(_j::Same, _j′::Next, _k::Same) where {Ti, Net <: PartwiseCountStepper{Ti}}
    begin
        arg = stp.ocl
        j = destep(_j)
        j′ = destep(_j′)
        k = destep(_k)
        rnk_j = arg.rnk_j
        rnk_j′ = arg.rnk_j′
        πos = arg.cst.πos
        prm = arg.cst.prm
        n′ = arg.cst.n′

        if rnk_j′ < πos[k + 1] && prm[rnk_j′] < j′
            rnk_j′ += 1
            arg.rnk_j′ = rnk_j′
            return Step(arg.cst.ocl)(Same(rnk_j), Next(rnk_j′))
        else
            return Step(arg.cst.ocl)(Same(rnk_j), Same(rnk_j′))
        end
    end
end

@propagate_inbounds function (stp::Step{Net})(_j::Next, _j′::Same, _k::Same) where {Ti, Net <: PartwiseCountStepper{Ti}}
    begin
        arg = stp.ocl
        j = destep(_j)
        j′ = destep(_j′)
        k = destep(_k)
        rnk_j = arg.rnk_j
        rnk_j′ = arg.rnk_j′
        πos = arg.cst.πos
        prm = arg.cst.prm
        n′ = arg.cst.n′

        if rnk_j < πos[k + 1] && prm[rnk_j] < j
            rnk_j += 1
            arg.rnk_j = rnk_j
            return Step(arg.cst.ocl)(Next(rnk_j), Same(rnk_j′))
        else
            return Step(arg.cst.ocl)(Same(rnk_j), Same(rnk_j′))
        end
    end
end

ext: jl


text:
using Parameters

@with_kw struct MDInput

  #
  # Atomistic simulation parameters
  #

  # Number of particles
  n :: Int64 = 1000

  # Simulation time
  dt :: Float64 = 0.1
  tmax :: Float64 = 1000

  # Size of the box
  side :: Float64 = 1000.
  cutoff :: Float64 = 10.

  # Printing properties
  iprint :: Int64 = 100

  # Save how many points of the trajectory:
  nsave :: Int64 = 200

  # Temperature
  kavg_target :: Float64 = 5.
  kavg_tol :: Float64 = 1e-3

  # Parameters of Berendsen bath
  tau :: Float64 = 50.
  nequil :: Int64 = -1

  # Lennard-Jones parameters
  eps :: Float64 = 1.
  sig :: Float64 = 2.

  # Parameters of optimization method for initial point
  tol :: Float64 = 1e-2

  # Fraction of initially sick individuals
  Si :: Float64 = 0.01

  # Cross-section and probability of contamination
  xsec :: Float64 = 4.
  pcont :: Float64 = 0.3

  # chance of dieing or getting immune (per step)
  pdie :: Float64 = 0.0001
  pimmune :: Float64 = 0.0005

  # Energy function (defaults to QUAD inside codes)
  uf = "uf_QUAD!"

end


text:
"""
Iceberg(n_timesteps)

Driver routine for iceberg simulation using WAVI.
Defines parameters, initialises state, and runs n_timesteps.
"""
function iceberg_test(;end_time=1000.)

#Number of timesteps between plots
n_steps_plot=Inf

# Iceberg
nx = 20
ny = 20
nσ = 4
x0 = -50000.0
y0 = -50000.0
dx = 5000.0
dy = 5000.0
#Homogenous Dirichlet boundary conditions
u_iszero=falses(nx+1,ny)
u_iszero[div(nx,2)+1,:].=true
v_iszero=falses(nx,ny+1)
v_iszero[:,div(ny,2)+1].=true
grid = Grid(nx = nx,
            ny = ny,
            nσ = nσ,
            x0 = x0,
            y0 = y0,
            dx = dx,
            dy = dy,
            u_iszero = u_iszero,
            v_iszero = v_iszero)

h_mask=sqrt.(grid.xxh.^2+grid.yyh.^2) .< 45000
h_mask = convert(Array{Bool,2}, h_mask)
grid = @set grid.h_mask = h_mask

#Model domain mask on h-grid
starting_thickness=zeros(nx,ny)
starting_thickness[h_mask] .= 200.0
initial_conditions = InitialConditions(initial_thickness = starting_thickness)

accumulation_rate=0.3
params = Params(accumulation_rate = accumulation_rate)

bed_elevation=-500.0.*ones(nx,ny)

maxiter_picard=1 #No need for Picard iteration if runninng to steady state
solver_params = SolverParams(maxiter_picard = maxiter_picard)

#make the model
model = Model(grid = grid,
              bed_elevation = bed_elevation,
              params = params,
              solver_params = solver_params,
              initial_conditions = initial_conditions)

timestepping_params = TimesteppingParams(dt = 0.1, end_time = end_time)
simulation = Simulation(model = model,
                        timestepping_params = timestepping_params)
run_simulation!(simulation)
return simulation
end


text:
function conjugate(r::Union{Regression,Regression!}, Y, Γ0::AbstractArray, ts)
  Γ0_ = Matrix(Γ0)
  conjugate(r, Y, Mitosis.Gaussian{(:F,:Γ)}(zeros(eltype(Γ0_), size(Γ0_, 1)), Γ0_), ts)
end

function conjugate(r::Union{Regression,Regression!}, Y::RODESolution, Γ0::AbstractArray)
  Γ0_ = Matrix(Γ0)
  Y_ = Y.u
  ts = Y.t
  conjugate(r, Y_, Mitosis.Gaussian{(:F,:Γ)}(zeros(eltype(Γ0_), size(Γ0_, 1)), Γ0_), ts)
end

function conjugate(r::Union{Regression,Regression!}, YY::EnsembleSolution, Γ0::AbstractArray)
  Γ0_ = Matrix(Γ0)
  G = Mitosis.Gaussian{(:F,:Γ)}(zeros(eltype(Γ0_), size(Γ0_, 1)), Γ0_)
  for Y in YY
    Y_ = Y.u
    ts = Y.t
    G = conjugate(r, Y_, G, ts)
  end
  G
end

function conjugate(r::Regression, Y, prior::Gaussian, ts)
  @unpack k, fjac, ϕ0func, θ, dy, isscalar = r
  @unpack g, p = k

  t = ts[1]
  y = Y[1]

  ϕ = calc_J(y, r, θ, t)
  μ = copy(prior.F)
  Γ = copy(prior.Γ)

  for i in 1:length(Y)-1
    ϕ = calc_J(y, r, θ, t)
    if isinplace(r)
      g(dy, y, p, t)
    else
      dy = g(y, p, t)
    end
    if !isscalar
      Gϕ = pinv(outer_(dy))*ϕ
    else
      Gϕ = pinv(outer_(reshape(dy, :, 1)))*ϕ
    end
    zi = ϕ'*Gϕ
    t2 = ts[i + 1]
    y2 = Y[i + 1]

    ds = t2 - t

    # compute intercept
    if ϕ0func !== nothing
      ϕ0 = ϕ0func(y, p, t)
      μ = μ .+ Gϕ'*((y2 - y) .- ϕ0*ds)
    else
      μ = μ .+ Gϕ'*(y2 - y)
    end

    t = t2
    y = y2
    Γ = Γ .+ zi*ds

  end
  return Mitosis.Gaussian{(:F,:Γ)}(μ, Γ)
end

function conjugate(r::Regression!, Y, prior::Gaussian, ts)
  @unpack k, fjac!, ϕ0func!, ϕ, ϕ0, y, y2, θ, dy, isscalar = r
  @unpack g, p = k

  t = ts[1]
  copyto!(y, Y[1])

  calc_J!(ϕ, r, θ, t)
  μ = prior.F
  Γ = prior.Γ

  for i in 1:length(Y)-1
    calc_J!(ϕ, r, θ, t)
    if isinplace(r)
      g(dy, y, p, t)
    else
      dy = g(y, p, t)
    end
    if !isscalar
      Gϕ = pinv(outer_(dy))*ϕ
    else
      Gϕ = pinv(outer_(reshape(dy, :, 1)))*ϕ
    end
    zi = ϕ'*Gϕ
    t2 = ts[i + 1]
    copyto!(y2, Y[i + 1])

    ds = t2 - t

    # compute intercept
    if ϕ0func! !== nothing
      ϕ0func!(ϕ0, y, p, t)
      μ .= μ + Gϕ'*((y2 - y) - ϕ0*ds)
    else
      μ .= μ + Gϕ'*(y2 - y)
    end

    t = t2
    copyto!(y, y2)
    Γ .= Γ + zi*ds

  end
  return Mitosis.Gaussian{(:F,:Γ)}(μ, Γ)
end


text:
using testPQPolygonSDK
using Dates
using DataFrames

# build a user model -
options = Dict{String,Any}()
options["email"] = "yp392@cornell.edu"
options["apikey"] = "abc1234" # do _not_ check in a real API key

# build the user model -
user_model = model(PQPolygonSDKUserModel, options)

# now that we have the user_model, let's build an endpoint model -
endpoint_options = Dict{String,Any}()
endpoint_options["asset_class"] = "stocks"
endpoint_options["data_type"] = "trade"
endpoint_options["id"] = "N/A"
endpoint_options["sip"] = "CTA"
endpoint_options["order"] = "asc"
endpoint_options["limit"] = "10"
endpoint_options["sort"] = "id"

endpoint_model = model(PolygonConditionsEndpointModel, user_model, endpoint_options)

# build the url -
base = "https://api.polygon.io"
my_url_string = url(base, endpoint_model)

# make the call -
(header, data) = api(PolygonConditionsEndpointModel, my_url_string)


text:
################################################################################
# setup
################################################################################
using Revise
using DataBench, FastGroupBy, ShortStrings, CategoricalArrays, SortingLab,
    SortingAlgorithms, Feather, BenchmarkTools, DataFrames, DataFramesMeta,
    Lazy, Query

srand(1);
const N = 10_000_000; const K = 100
@time df = createSynDataFrame(N, K); #31 #40

# bench results
benchresults = []
bch2sec(bchres) = mean(bchres.times)/1e9

################################################################################
# DT[, sum(v1), keyby=id1]
# short string & string
# Status: SLOWER but if converted to Categorical then it's faster
################################################################################
# test String15
@time df[:id1_ss] = ShortString7.(df[:id1]);
# @time a = fastby(sum, df[:id1_ss], df[:v1]);
# @time a = fastby(sum, df[:id1_ss], df[:v1]);

# @time a = fastby(sum, df[:id1], df[:v1]);
# @time a = fastby(sum, df[:id1], df[:v1]);

sumv1id1_ss = @benchmark fastby(sum, df[:id1_ss], df[:v1])
sumv1id1 = @benchmark fastby(sum, df[:id1], df[:v1])

sumv1id1_dfm_fn(df) = @> begin
    df
    @by(:id1, sumx=sum(:v1))
end

sumv1id1_dfm = @benchmark sumv1id1_dfm_fn(df)

sumv1id1_query_fn(df) = @from i in df begin
    @group i by i.id1 into g
    @select {r=sum(g..v1)}
    @collect DataFrame
end
sumv1id1_query = @benchmark sumv1id1_query_fn(df)

# sumv1id1_df_fn(df) = aggregate(df[[:id1,:v1]], :id1, sum)
# sumv1id1_df = @benchmark sumv1id1_df_fn(df)

push!(benchresults,
     ("1 @by(:id1, sumx=sum(:v1))", bch2sec(sumv1id1_ss),    "FastGroupBy.jl\n fastby()", "ShortStrings.jl")
    ,("1 @by(:id1, sumx=sum(:v1))", bch2sec(sumv1id1),       "FastGroupBy.jl\n fastby()", "String")
    ,("1 @by(:id1, sumx=sum(:v1))", bch2sec(sumv1id1_dfm),   "DataFramesMeta.jl", "String")
    ,("1 @by(:id1, sumx=sum(:v1))", bch2sec(sumv1id1_query), "Query.jl", "String")
)

################################################################################
# DT[, sum(v1), keyby=id1]
# test categorical
# Status: FASTER 9x
################################################################################
@time df[:id1_cate] = categorical(df[:id1]); #7
@time df[:id1_cate] = compress(df[:id1_cate]); # 0.5
# @time sumby(df[:id1_cate], df[:v1]);
# @time fastby(sum, df[:id1_cate], df[:v1]);
# @time fastby(sum, df[:id1_cate], df[:v1]);

# @time fgroupreduce(+, df[:id1_cate], df[:v1]);
# @time fgroupreduce(+, df[:id1_cate], df[:v1]);

sumv1id1_cate = @benchmark fastby(sum, df[:id1_cate], df[:v1])
sumv1id1_cate_groupreduce = @benchmark fgroupreduce(+, df[:id1_cate], df[:v1])

sumv1id1_cate_dfm_fn(df) = @> begin
    df
    @by(:id1_cate, sumx=sum(:v1))
end

sumv1id1_cate_dfm = @benchmark sumv1id1_cate_dfm_fn(df)

umv1id1_cate_query_fn(df) = @from i in df begin
    @group i by i.id1_cate into g
    @select {r=sum(g..v1)}
    @collect DataFrame
end

umv1id1_cate_query = @benchmark umv1id1_cate_query_fn(df)

push!(benchresults,
     ("1 @by(:id1, sumx=sum(:v1))", bch2sec(sumv1id1_cate),    "FastGroupBy.jl\n fastby()", "CategoricalArrays.jl")
    ,("1 @by(:id1, sumx=sum(:v1))", bch2sec(sumv1id1_cate_groupreduce),    "FastGroupBy.jl\n fgroupreduce()", "CategoricalArrays.jl")
    ,("1 @by(:id1, sumx=sum(:v1))", bch2sec(sumv1id1_cate_dfm),   "DataFramesMeta.jl", "CategoricalArrays.jl")
    ,("1 @by(:id1, sumx=sum(:v1))", bch2sec(umv1id1_cate_query), "Query.jl", "CategoricalArrays.jl")
)

################################################################################
# DT[, sum(v1), keyby="id1,id2"]
# test categorical
# Status: FASTER 9x
################################################################################
@time df[:id2_cate] = categorical(df[:id2]); #7
@time df[:id2_cate] = compress(df[:id2_cate]); # 0.5
# @time fgroupreduce(+, (df[:id1_cate], df[:id2_cate]), df[:v1]);
# @time fgroupreduce(+, (df[:id1_cate], df[:id2_cate]), df[:v1]);

fn, byveccv, val = +, (df[:id1_cate], df[:id2_cate]), df[:v1]

@time a = fgroupreduce(+, (df[:id1_cate], df[:id2_cate]), df[:v1])

sumv1id1id2_groupreduce = @benchmark fgroupreduce(+, (df[:id1_cate], df[:id2_cate]), df[:v1])

@code_warntype fgroupreduce(+, (df[:id1_cate], df[:id2_cate]), df[:v1])

sumv1id1id2_dfm_fn(df) = @> begin
    df
    @by([:id1_cate,:id2_cate], r=sum(:v1))
end

sumv1id1id2_dfm = @benchmark sumv1id1id2_dfm_fn(df)

sumv1id1id2_query_fn(df) =
@from i in df begin
    @group i by (i.id1,i.id2) into g
    @select {r=sum(g..v1)}
    @collect DataFrame
end

sumv1id1id2_query = @benchmark sumv1id1id2_query_fn(df)

push!(benchresults,
     ("2 @by([:id1_cate,:id2_cate], r=sum(:v1))", bch2sec(sumv1id1id2_groupreduce),    "FastGroupBy.jl\n fgroupreduce()", "CategoricalArrays.jl")
    ,("2 @by([:id1_cate,:id2_cate], r=sum(:v1))", bch2sec(sumv1id1id2_dfm),   "DataFramesMeta.jl", "CategoricalArrays.jl")
    ,("2 @by([:id1_cate,:id2_cate], r=sum(:v1))", bch2sec(sumv1id1id2_query), "Query.jl", "CategoricalArrays.jl")
)

################################################################################
# DT[, list(sum(v1),mean(v3)), keyby=id3]
# test categorical
# Status: SLOWER
# TODO: make a fast fsortandperm for integers
# TODO: make this faster than data.table
# TODO: look for mean and make it into a reduce
# TO BEAT: 0.6s for 10m; 14.11s for 100m;
#
# NOTES:
# The key issue here is that we need to sort two vectors in order to perform the
# groupby operation. R's seems to have a really fast way to doing that
# Potential make the group into radix sort instead of counting sort
#
# What about making an RLE vector type?
#
# I can radixsort the 10m refs in 0.08 seconds so still 0.52 seconds to distribute
#
# I can sortperm the 10m refs in 0.28 seconds so still
################################################################################
@time df[:id3_cate] = categorical(df[:id3]) |> compress;
# import FastGroupBy: BaseRadixSortSafeTypes, fastby

sumv1meanv3id3_cate = @benchmark fastby((sum, mean), df[:id3_cate], (df[:v1], df[:v3]))
# BenchmarkTools.Trial:
#   memory estimate:  435.87 MiB
#   allocs estimate:  1000315
#   --------------
#   minimum time:     1.084 s (0.00% GC)
#   median time:      1.130 s (0.00% GC)  mean time:        1.133 s (0.00% GC)
#   maximum time:     1.170 s (0.00% GC)
#   --------------
#   samples:          5
#   evals/sample:     1

# BenchmarkTools.Trial:
#   memory estimate:  430.21 MiB
#   allocs estimate:  977648
#   --------------
#   minimum time:     801.169 ms (0.00% GC)
#   median time:      912.131 ms (0.00% GC)
#   mean time:        945.560 ms (0.00% GC)
#   maximum time:     1.200 s (0.00% GC)
#   --------------
#   samples:          6
#   evals/sample:     1

sumv1meanv3id3_cate_dfm_fn(df) = @> begin
    df
    @by([:id3], sumv1=sum(:v1), meanv2=mean(:v2))
end

sumv1meanv3id3_cate_dfm = @benchmark sumv1meanv3id3_cate_dfm_fn(df)
# BenchmarkTools.Trial:  memory estimate:  1022.60 MiB
#   allocs estimate:  7395590  --------------
#   minimum time:     7.985 s (0.00% GC)
#   median time:      7.985 s (0.00% GC)
#   mean time:        7.985 s (0.00% GC)
#   maximum time:     7.985 s (0.00% GC)
#   --------------
#   samples:          1
#   evals/sample:     1

sumv1meanv3id3_cate_ql_fn(df) = @from i in df begin
    @group i by i.id3 into g
    @select {s=sum(g..v1),m=mean(g..v3)}
    @collect DataFrame
end

sumv1meanv3id3_cate_ql = @benchmark sumv1meanv3id3_cate_ql_fn(df)

push!(benchresults,
     ("3 @by([:id1_cate,:id2_cate], r=sum(:v1))", bch2sec(sumv1meanv3id3_cate),    "FastGroupBy.jl\n fastby()", "CategoricalArrays.jl")
    ,("3 @by([:id1_cate,:id2_cate], r=sum(:v1))", bch2sec(sumv1meanv3id3_cate_dfm),   "DataFramesMeta.jl", "CategoricalArrays.jl")
    ,("3 @by([:id1_cate,:id2_cate], r=sum(:v1))", bch2sec(sumv1meanv3id3_cate_ql), "Query.jl", "CategoricalArrays.jl")
)

################################################################################
# DT[, lapply(.SD, mean), keyby=id4, .SDcols=7:9]
# TO BEAT: 0.6s for 10m; 2.76s for 100m; 6.45s for 250m
# TODO: make CountingSort
# Status: MUCH SLOWER if the number of elements is large and saturating memory
################################################################################
# df[:id4_32] = Int32.(df[:id4])
# df[:id4_8] = Int8.(df[:id4])

# @time a = fastby((mean, mean, mean), df[:id4], (df[:v1], df[:v2], df[:v3]));

# @time a = fastby((mean, mean, mean), df[:id4_32], (df[:v1], df[:v2], df[:v3]));
# @time a = fastby((mean, mean, mean), df[:id4_32], (df[:v1], df[:v2], df[:v3]));

# @time a = fastby((mean, mean, mean), df[:id4_8], (df[:v1], df[:v2], df[:v3]));
# @time a = fastby((mean, mean, mean), df[:id4_8], (df[:v1], df[:v2], df[:v3]));

mean79id4 = @benchmark fastby((mean, mean, mean), df[:id4], (df[:v1], df[:v2], df[:v3]))
# BenchmarkTools.Trial:
#   memory estimate:  540.57 MiB
#   allocs estimate:  2470
#   --------------
#   minimum time:     456.440 ms (0.00% GC)
#   median time:      494.681 ms (0.00% GC)
#   mean time:        587.366 ms (16.06% GC)
#   maximum time:     1.310 s (64.82% GC)
#   --------------
#   samples:          9
#   evals/sample:     1

mean79id4_dfm_fn(df) = @> df begin
    @by(:id4, mean1=mean(:v1), mean2=mean(:v2), mean3=mean(:v3))
end

mean79id4_dfm = @benchmark mean79id4_dfm_fn(df)

mean79id4_ql_fn(df) = @from i in df begin
    @group i by i.id4 into g
    @select {m7=mean(g..v1),m8=mean(g..v2),m9=mean(g..v3)}
    @collect DataFrame
end

mean79id4_ql = @benchmark mean79id4_ql_fn(df)

push!(benchresults,
     ("4 @by(:id4, mean1=mean(:v1), mean2=mean(:v2), mean3=mean(:v3))", bch2sec(mean79id4),    "FastGroupBy.jl\n fastby()", "Int")
    ,("4 @by(:id4, mean1=mean(:v1), mean2=mean(:v2), mean3=mean(:v3))", bch2sec(mean79id4_dfm),   "DataFramesMeta.jl", "Int")
    ,("4 @by(:id4, mean1=mean(:v1), mean2=mean(:v2), mean3=mean(:v3))", bch2sec(mean79id4_ql), "Query.jl", "Int")
)

################################################################################
# DT[, lapply(.SD, sum), keyby=id6, .SDcols=7:9]
# TO BEAT: 0.4s for 10m; 11.33s for 100m; 31.31 for 250m
################################################################################
df[:id6] = Int64.(df[:id6]); # this is needed not because it was `Int64` but `Float64`
# @time a = fastby((sum, sum, sum), df[:id6], (df[:v1], df[:v2], df[:v3]));
# @time a = fastby((sum, sum, sum), df[:id6], (df[:v1], df[:v2], df[:v3]));

sum79id6 = @benchmark fastby((sum, sum, sum), df[:id6], (df[:v1], df[:v2], df[:v3]))
# BenchmarkTools.Trial:
#   memory estimate:  598.48 MiB
#   allocs estimate:  1330706
#   --------------
#   minimum time:     950.109 ms (0.00% GC)
#   median time:      986.441 ms (0.00% GC)
#   mean time:        997.222 ms (0.00% GC)
#   maximum time:     1.095 s (0.00% GC)
#   --------------
#   samples:          6
#   evals/sample:     1

sum79id6_dfm_fn(df) = @> df begin
    @by(:id6, sum1=sum(:v1), sum2=sum(:v2), sum3=sum(:v3))
end

sum79id6_dfm = @benchmark sum79id6_dfm_fn(df)

sum79id6_ql_fn(df) = @from i in df begin
    @group i by i.id6 into g
    @select {m7=sum(g..v1),m8=sum(g..v2),m9=sum(g..v3)}
    @collect DataFrame
end

sum79id6_ql = @benchmark sum79id6_ql_fn(df)

push!(benchresults,
     ("5 @by(:id6, sum1=sum(:v1), sum2=sum(:v2), sum3=sum(:v3))", bch2sec(sum79id6),    "FastGroupBy.jl\n fastby()", "Int")
    ,("5 @by(:id6, sum1=sum(:v1), sum2=sum(:v2), sum3=sum(:v3))", bch2sec(sum79id6_dfm),   "DataFramesMeta.jl", "Int")
    ,("5 @by(:id6, sum1=sum(:v1), sum2=sum(:v2), sum3=sum(:v3))", bch2sec(sum79id6_ql), "Query.jl", "Int")
)

if false
    df[:id6_cate] = categorical(df[:id6])
    @benchmark fgroupreduce.((+,), (df[:id6_cate],), (df[:v1], df[:v2], df[:v3]))
    # BenchmarkTools.Trial:
    # memory estimate:  123.80 MiB
    # allocs estimate:  600971
    # --------------
    # minimum time:     347.807 ms (0.00% GC)
    # median time:      364.181 ms (0.00% GC)
    # mean time:        363.018 ms (0.00% GC)
    # maximum time:     385.564 ms (0.00% GC)
    # --------------
    # samples:          14
    # evals/sample:     1
end

################################################################################
# generate r data
################################################################################
using RCall

rres=R"""
memory.limit(2^31-1)
library(data.table)
N=$N; K=$K
set.seed(1)
DT <- data.table(
  id1 = sample(sprintf("id%03d",1:K), N, TRUE),      # large groups (char)
  id2 = sample(sprintf("id%03d",1:K), N, TRUE),      # large groups (char)
  id3 = sample(sprintf("id%010d",1:(N/K)), N, TRUE), # small groups (char)
  id4 = sample(K, N, TRUE),                          # large groups (int)
  id5 = sample(K, N, TRUE),                          # large groups (int)
  id6 = sample(N/K, N, TRUE),                        # small groups (int)
  v1 =  sample(5, N, TRUE),                          # int in range [1,5]
  v2 =  sample(5, N, TRUE),                          # int in range [1,5]
  v3 =  sample(round(runif(100,max=100),4), N, TRUE) # numeric e.g. 23.5749
)
cat("GB =", round(sum(gc()[,2])/1024, 3), "\n")
replicate(5, list(
    system.time( DT[, sum(v1), keyby=id1] ),
    system.time( DT[, sum(v1), keyby="id1,id2"] ),
    system.time( DT[, list(sum(v1),mean(v3)), keyby=id3] ),
    system.time( DT[, lapply(.SD, mean), keyby=id4, .SDcols=7:9] ),
    system.time( DT[, lapply(.SD, sum), keyby=id6, .SDcols=7:9] )
))
"""

rres1 = [mean(x->x[3], getindex.(rres,j,1:5)) for j=1:5]

################################################################################
# Combine Julia Data with R Data
################################################################################
dfres = DataFrame()
for (i, n) in enumerate([:benchmark,:seconds,:method,:eltype])
    dfres[n] = [j[i] for j in benchresults]
end
dfres

# prepare the R results
dfres_r = DataFrame(
    benchmark = unique(dfres[:benchmark]) |> sort,
    seconds = rres1,
    method = ["R\n data.table" for i=1:5],
    eltype = ["String", "String", "String", "Int", "Int"]
)

dfres_all = vcat(dfres, dfres_r)

function getstr(str)
    ss = split.(str,"\n")
    n = length(ss)
    if n == 2
        return ss[1], "", ss[2]
    elseif n == 1
        return ss[1], "", ""
    else
        return ss[1], ss[2], ss[3]
    end
end

getstrres = getstr.(dfres_all[:method])

dfres_all[:pkg] = [getstrres1[1] for getstrres1 in getstrres]
dfres_all[:func] = [getstrres1[3] for getstrres1 in getstrres]

dfres_all[:label] = lstrip.(((func, pkg)->ifelse(func=="",pkg, func)).(dfres_all[:func], dfres_all[:pkg]))

################################################################################
# Convert to IndexedTables.jl
################################################################################
using IndexedTables, JuliaDB
@time dfit = table(df) # 11s
@time dfit2 = reindex(dfit, (:id1,:id2,:id3,:id4,:id5,:id6)) #26

# see https://github.com/JuliaComputing/JuliaDB.jl/issues/139
@benchmark groupreduce(+, dfit2, :id6, select = :v1)

################################################################################
# plotting of results
################################################################################
using Plots
gr(size=(760, 500))
# gr()

isquery(x) = x[1:5] == "Query"

function plotresult_groupby(benchmark, dfres_all = dfres_all, addtitle="")
    ms = maximum(dfres_all[:seconds])
    dfres1 = @where(dfres_all, :benchmark .==  benchmark )
    dfres1 = @where(dfres1, .!isquery.(:method))
    sort!(dfres1, cols = [:pkg, :func])

    bar(
        dfres1[:pkg].*"\n".*dfres1[:func].*" ".*dfres1[:eltype],
        # log.(dfres1[:seconds]),
        (dfres1[:seconds]),
        # orientation = :h,
        # horizontal = true,
        # group = dfres1[:pkg],
        fillcolor = categorical(dfres1[:pkg]).refs,
        # bar_position = :dodge,
        ylabel = "seconds",
        label="",
        # labels = map(string, categorical(dfres1[:pkg]).refs),
        title = "Bench $benchmark $addtitle \n 10m rows ",
        ylim = (0, 4)
    )
    # @df dfres1 groupedbar(
    #     :pkg,
    #     :seconds,
    #     # orientation = :h,
    #     # horizontal = true,
    #    orientation = :horizontal,
    #     group = :label,
    #     bar_position = :dodge,
    #     xlabel = "seconds",
    #     title = "Bench $benchmark\n 10m rows"
    # )

    # barplot(
    #     x = dfres1[:method] .* "\n" .* dfres1[:eltype],
    #     y = dfres1[:seconds],
    #     horizontal = true,
    #     group=collect(1:length(dfres1[:seconds]))#,
    #     # label = "seconds",
    #     # title = "Bench $benchmark\n 10m rows"
    # )
end

plotresult_groupby(dfres_all[1,:benchmark], @where(dfres_all, :eltype .== "String"), "(String type)")
savefig("groupby1a.png")

plotresult_groupby(dfres_all[1,:benchmark], @where(dfres_all, (:eltype .== "CategoricalArrays.jl") .| (:pkg .== "R")), "(Categorical type)")
savefig("groupby1b.png")

plotresult_groupby(unique(dfres_all[:benchmark])[2])
savefig("groupby2.png")

plotresult_groupby(unique(dfres_all[:benchmark])[3])
savefig("groupby3.png")

plotresult_groupby(unique(dfres_all[:benchmark])[4])
savefig("groupby4.png")

plotresult_groupby(unique(dfres_all[:benchmark])[5])
savefig("groupby5.png")


text:
"""
    augmented_trip_meta_graph_fws_dists(tg::TG, dist_fn::Function,
                                        n_depots::Int64, n_sites::Int64,
                                        depot_sites::Vector{LOC}, drone_params::DroneParams) where {LOC, TG <: TransitGraph}

Get distance matrix from Floyd-Warshall Shortest Paths alg on the augmented
trip metagraph. See Appendix II-B for details.
"""
function augmented_trip_meta_graph_fws_dists(tg::TG, dist_fn::Function,
                                             n_depots::Int64, n_sites::Int64,
                                             depot_sites::Vector{LOC}, drone_params::DroneParams) where {LOC, TG <: TransitGraph}

    n_trips = length(tg.transit_trips)
    n_verts = n_depots + n_sites + n_trips

    # Create the FWS dists matrix
    fws_dists = Inf*ones(n_verts, n_verts)

    # First do depots-sites
    for i = 1:n_depots + n_sites
        for j = 1:n_depots + n_sites

            if i != j
                dist = dist_fn(depot_sites[i], depot_sites[j])
                if dist < drone_params.max_distance
                    fws_dists[i, j] = dist
                    fws_dists[j, i] = dist
                end
            end
        end
    end

    # Now do depot-sites to trips
    for i = 1:n_depots + n_sites
        for j = n_depots+n_sites+1:n_verts
            trip = tg.transit_trips[j - n_depots - n_sites]
            mindist = Inf

            for rwp in trip

                dist = dist_fn(depot_sites[i], tg.stop_to_location[rwp.stop_id])
                mindist = (dist < mindist) ? dist : mindist

            end

            fws_dists[i, j] = mindist
            fws_dists[j, i] = mindist
        end
    end

    # Finally, do trip-trip
    for i = n_depots + n_sites + 1:n_verts
        trip1 = tg.transit_trips[i - n_depots - n_sites]
        for j = n_depots + n_sites + 1:n_verts
            trip2 = tg.transit_trips[j - n_depots - n_sites]
            mindist = Inf

            # Iterate over both trips and compute min distance
            for rwp1 in trip1
                for rwp2 in trip2

                    dist = dist_fn(tg.stop_to_location[rwp1.stop_id], tg.stop_to_location[rwp2.stop_id])

                    if drone_params.avg_speed * (rwp2.arrival_time - rwp1.arrival_time) > dist && dist < mindist
                        mindist = dist
                    end
                end
            end

            fws_dists[i, j] = mindist

        end
    end

    # Run FWS to get distances
    floyd_warshall!(fws_dists)

    return fws_dists
end

"""
    get_stop_idx_to_trip_ids(tg::TG) where {TG <: TransitGraph}

One-to-many mapping from stop IDs to trip IDs passing through them
"""
function get_stop_idx_to_trip_ids(tg::TG) where {TG <: TransitGraph}

    stop_idx_to_trips = Dict{Int64,Set{Int64}}()

    for (trip_id, trip) in enumerate(tg.transit_trips)

        for rp in trip
            @assert haskey(tg.stop_to_location, rp.stop_id)
            if ~(haskey(stop_idx_to_trips, rp.stop_id))
                stop_idx_to_trips[rp.stop_id] = Set{Int64}()
            end
            push!(stop_idx_to_trips[rp.stop_id], trip_id)
        end

    end

    return stop_idx_to_trips
end

"""
    true_stop_to_locations(stop_to_location::Dict{Int64,LOC}, stop_idx_to_trips::Dict) where {LOC}

Filter out any stops that are not actually used by any trip.
"""
function true_stop_to_locations(stop_to_location::Dict{Int64,LOC}, stop_idx_to_trips::Dict) where {LOC}

    keys_to_keep = Set{Int64}()

    true_stop_to_locs = Dict{Int64,LOC}()

    for stlkey in collect(keys(stop_to_location))

        if haskey(stop_idx_to_trips, stlkey)
            push!(keys_to_keep, stlkey)
        end
    end

    # Now copy over
    for kk in keys_to_keep
        true_stop_to_locs[kk] = stop_to_location[kk]
    end

    return true_stop_to_locs

end

"""
    generate_city_halton_nn(city_params::CityParams; n_points::Int64 = 100,
                        bases::Vector{Int64} = [2, 3], discard::Int64 = 0)

Sample a halton sequence of points in a city and create a NNTree with them.
"""
function generate_city_halton_nn(city_params::CityParams; n_points::Int64 = 100,
                                 bases::Vector{Int64} = [2, 3], discard::Int64 = 0)

    lower = [city_params.lat_start, city_params.lon_start]
    upper = [city_params.lat_end, city_params.lon_end]

    # Generate the halton sequence
    city_halton_points = get_halton_sequence(n_points, bases, lower, upper, discard)

    # Generate the NN tree
    # nn_tree.data has the map from idx to point
    halton_nn_tree = BallTree(city_halton_points, EuclideanLatLong())

    return halton_nn_tree, city_halton_points
end

"""
    get_travel_time_estimate(halton_nn_tree::BallTree, loc1::LOC, loc2::LOC,
                             estimate_matrix::Matrix{Float64}) where LOC

Using the surrogate travel time estimate matrix, compute the
estimate between any two locations by looking up their corresponding nearest rep. location.
"""
function get_travel_time_estimate(halton_nn_tree::BallTree, loc1::LOC, loc2::LOC,
                                  estimate_matrix::Matrix{Float64}) where LOC

    # Get the nearest neighbor for each location
    # IMP - If they are the same, return 0.0
    loc1_vect = convert_to_vector(loc1)
    loc1_idxs, _ = knn(halton_nn_tree, loc1_vect, 1)
    loc1_idx = loc1_idxs[1]

    loc2_vect = convert_to_vector(loc2)
    loc2_idxs, _ = knn(halton_nn_tree, loc2_vect, 1)
    loc2_idx = loc2_idxs[1]

    # Returns 0.0 if same!
    return estimate_matrix[loc1_idx, loc2_idx]
end

"""
    compute_all_pairs_estimates(env::MAPFTransitEnv, n_halton_points::Int64, weight::Float64)

Preprocessing. Compute the pairwise travel times between pairs of Voronoi sites. See Appendix III
"""
function compute_all_pairs_estimates(env::MAPFTransitEnv, n_halton_points::Int64, weight::Float64)

    travel_time_estimates = zeros(n_halton_points, n_halton_points)

    @sync for i = 1:n_halton_points
        println("From point $i:")
        for j = 1:n_halton_points

            if i != j
                orig_str = string("d-", i)
                goal_str = string("d-", j)

                envcp = deepcopy(env)
                Threads.@spawn travel_time_estimates[i, j] = get_depot_to_site_travel_time!(envcp, weight, orig_str, goal_str)
            end
        end
    end

    return travel_time_estimates
end


text:
using Dorayaki
using Random
using Plots

const lr = 0.2
const iters = 10000

function main(x, y)
    I, H, O = 1, 10, 1
    w1 = Var(0.01 * randn(H, I))
    b1 = Var(zeros(H))
    w2 = Var(0.01 * randn(O, H))
    b2 = Var(zeros(O))

    function predict(p)
        q = linear(w1, p, b1)
        q = sigmoid(q)
        q = linear(w2, q, b2)
        return q
    end

    for i in 1:iters
        y_pred = predict(x)
        loss = mean_squared_error(y, y_pred)

        cleargrad!(w1, b1, w2, b2)

        gradient!(loss)

        w1.data .-= lr .* w1.grad.data
        b1.data .-= lr .* b1.grad.data
        w2.data .-= lr .* w2.grad.data
        b2.data .-= lr .* b2.grad.data

        if i % 1000 == 0
            println(loss.data)
        end
    end
    return predict
end

x = rand(1, 100)
y = sin.(2 * pi .* x) .+ rand(1, 100)

ret = @time main(x, y)
fig = plot(x', y', st=:scatter)
f(x) = ret(x).data[1]
plot!(fig, f)


text:
module Accumulo

using Thrift

import Base: close, isready, show, eof, start, next, done
import Thrift: flush

export AccumuloSession, AccumuloAuth, AccumuloAuthSASLPlain
export close

# export useful enums
#export PartialKey, TablePermission, SystemPermission, ScanType, ScanState, CompactionType, CompactionReason, TimeType
export IteratorScope, Durability, ConditionalStatus

# export useful types
export Key, KeyValue, IteratorSetting

# export table administration commands
# TODO: namespaces
export table_clone, table_create, table_delete, table_rename, table_du, table_config, table_config!, table_versions!, table_export, table_import, table_offline, table_online, table_exists, tables

# export table control commands
# TODO: locality groups, compaction
export flush, constraints, add_constraints, remove_constraints, table_splits, table_split, table_merge

# export table iterators
export iter, iters, add_iter, remove_iter, check_iter

# export writer functions
export batch, conditional_batch, batch_writer, conditional_batch_writer, where, update, delete, close, flush

# export scanner functions
export scanner, scanner_key, records, close, eof, start, next, done

## enable logging only during debugging
#using Logging
#const logger = Logging.configure(level=DEBUG)
##const logger = Logging.configure(filename="/tmp/accumulo$(getpid()).log", level=DEBUG)
#macro logmsg(s)
#    quote
#        debug($(esc(s)))
#    end
#end
macro logmsg(s)
end

include("proxy/proxy.jl")
using .proxy
import .proxy: update
include("types.jl")
include("sess.jl")
include("table_admin.jl")
include("table_control.jl")
include("iter.jl")
include("writers.jl")
include("scanners.jl")

end # module


text:
using CUDA

function matmul(a, b, c, d, conf;
                transform_global_to_shared_a = Transform.Elementwise(),
                transform_global_to_shared_b = Transform.Elementwise(),
                transform_global_to_shared_c = Transform.Elementwise(),
                transform_shared_to_global_d = Transform.Elementwise(),
                transform_shared_to_regs_a = Transform.Elementwise(),
                transform_shared_to_regs_b = Transform.Elementwise(),
                transform_shared_to_regs_c = Transform.Elementwise(),
                transform_regs_to_shared_d = Transform.Elementwise(),
                epilogue = Epilogue.Default(),
                kernel = Kernel.matmul_singlestage)

    args = [a, b, c, d,
            transform_global_to_shared_a, transform_global_to_shared_b, transform_global_to_shared_c, transform_shared_to_global_d,
            transform_shared_to_regs_a, transform_shared_to_regs_b, transform_shared_to_regs_c, transform_regs_to_shared_d,
            epilogue,
            conf]

    hostkernel = @cuda launch=false kernel(args...)
    attributes(hostkernel.fun)[CUDA.FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES] = 64 * 1024
    hostkernel(args...; conf.launch_args...)
end


text:
struct VolumePartsIter{
    TM<:AbstractMatrix,
    TV<:AbstractVector,
    Tg <: AbstractVector,
    Tts<:AbstractVector,
    Texpand<:NamedTuple,
    T<:Real} <: AbstractVolumePartsIter{TM, TV, T}

  #parameters
  A::TM
  Ã::TM #factored A, used in regressions
  G::TM

  g::Tg #vector version by referene of G, should be equivelent

  ts::Tts
  tsM1::Tts
  expand::Texpand

  dims::NamedTuple

  #inner constructor needed to identify elTΠ
  function VolumePartsIter(
    A::TM, Ã::TM, G::TM, g::Tg, ts::Tts, tsM1::Tts, expand::Texpand,
    ::Type{TV} = typeof(similar(A,1))) where {
      TM, TV, Tg, Tts, Texpand}

    return new{TM, TV, Tg, Tts, Texpand, eltype(TM)}(A,Ã,G,g,ts,tsM1, expand)
  end
end

#pre-allocates for the vpexpand
function VolumePartsIter(A::TA,  Ã::TA, G::TA, ts::Tts, tsM1::Tts=ts.-1,#=
  ::Type{TM}, ::Type{TV}, ::Type{T}=eltype(TM)=#) where {
    TA<:AbstractMatrix, TM, TV, T, Tts<:AbstractVector}

  g = vec(G)

  expand = (
    A = view(A, :, ts),
    LA = view(A, :, tsM1),
    Ã = view(Ã, :, ts),
    LÃ = view(Ã, :, tsM1),
    G = view(G, :, tsM1),
    index  = (
      A = to_indices(A, (:, ts)),
      LA = to_indices(A, (:, tsM1)),
      Ã = to_indices(Ã, (:, ts)),
      LÃ = to_indices(Ã, (:, tsM1)),
      G = to_indices(G, (:, tsM1)),
    ))

  Θ = VolumePartsIter(A, Ã, G, g, ts, tsM1, expand)
  return Θ
end

#basic constructor from dimension arguments
function VolumePartsIter(T::Int, K::Int, ts::Vector{Int}, ::Type{Tvp}=PARAM[:itertype],
  ::Type{TM}=PARAM[:itergpu] ? CUDA.CuMatrix{Tvp} : Matrix{Tvp},
  ::Type{TV}=PARAM[:itergpu] ? CUDA.CuVector{Tvp} : Vector{Tvp},
    ) where {Tvp<:Real, TV<:AbstractVector{Tvp}, TM<:AbstractMatrix{Tvp}}

  A = TM(abs.(randn(Tvp, K,T)))
  Ã = TM(abs.(randn(Tvp, K,T)))
  G = TM(ones(Tvp, K, T-1))

  return VolumePartsIter(A, Ã, G, ts, #=TM, TV, Tvp=#)
end

#needed due to the segemented array structure
function Base.deepcopy(Θ::VolumePartsIter)
  A  = deepcopy(Θ.A)
  Ã = deepcopy(Θ.Ã)
  G = deepcopy(Θ.G)

  return VolumePartsIter(A,Ã,G,Θ.ts,Θ.tsM1)
end

#invididual functions for computing volume constriubtions
#genabsμₖ(Aᵢ, LAᵢ, wsᵢ, Rtwsᵢ, RLwsᵢ) = abs(LAᵢ*(Rtwsᵢ - RLwsᵢ)) + abs(Aᵢ*wsᵢ-LAᵢ*Rtwsᵢ)
genabsμₖ(Aᵢ, LAᵢ, wsᵢ, RLwsᵢ) = abs(Aᵢ * wsᵢ - LAᵢ * RLwsᵢ)
genabsμₖ(Aᵢ, LAᵢ, wsᵢ, Rtwsᵢ, RLwsᵢ)  = genabsμₖ(Aᵢ, LAᵢ, wsᵢ, RLwsᵢ)
signabsμₖ(Aᵢ, LAᵢ, wsᵢ, Rtwsᵢ, RLwsᵢ) = sign(Aᵢ*wsᵢ-LAᵢ*Rtwsᵢ) #derivitive

#project the volume at a particular time t
#WARNING- the below only works when sum(Lwsₜ)==sum(wsₜ). If this is not true, need to carry
#an extra vector for the previous weights
function projectt(Aᵢ, LAᵢ, wsₜ, Rtwsₜ, RLwsₜ,  Mₜ)

  return Mₜ * genabsμₖ.(Aᵢ', LAᵢ', wsₜ, Rtwsₜ, RLwsₜ,)
end
function projectt(Aᵢ, LAᵢ, wsₜ, Rtwsₜ, RLwsₜ, ::Nothing)
  return genabsμₖ.(Aᵢ', LAᵢ', wsₜ, Rtwsₜ, RLwsₜ,)
end

#projectt(Aᵢ, LAᵢ, wsₜ, RLwsₜ, ::Nothing) = genabsμₖ.(Aᵢ', LAᵢ', wsₜ, RLwsₜ, 1.0 +sum(RLwsₜ)-sum(wsₜ))

#project the volume for each time t
function projectbyt(A, xsection)
  dims = (K = size(A,1), T = size(A,2))
  @unpack xws, xRtws, xRLws, xM = xsection

  #project the volume at each time t, then concatenate the results
  factored =  reduce(vcat, broadcast(1:(dims.T-1), xws, xRtws,xRLws, xM) do t, xwsₜ, xRtwsₜ, xRLwsₜ, xMₜ
    projectt(A[:,t+1], A[:,t], xwsₜ, xRtwsₜ, xRLwsₜ, xMₜ)
  end)

  return factored
end

#the below is used to get the active volume
genμₖ(Aᵢ, LAᵢ, wsᵢ, RLwsᵢ) = Aᵢ * wsᵢ - LAᵢ * RLwsᵢ
genμₖ(Aᵢ, LAᵢ, wsᵢ, Rtwsᵢ, RLwsᵢ)  = genμₖ(Aᵢ, LAᵢ, wsᵢ, RLwsᵢ)
function projectpurchasest(Aᵢ, LAᵢ, wsₜ, Rtwsₜ, RLwsₜ,  Mₜ)

  return Mₜ * genμₖ.(Aᵢ', LAᵢ', wsₜ, Rtwsₜ, RLwsₜ,)
end
function projectpurchasest(Aᵢ, LAᵢ, wsₜ, Rtwsₜ, RLwsₜ, ::Nothing)
  return genμₖ.(Aᵢ', LAᵢ', wsₜ, Rtwsₜ, RLwsₜ,)
end
function projectpurchasesbyt(A, xsection)
  dims = (K = size(A,1), T = size(A,2))
  @unpack xws, xRtws, xRLws, xM = xsection

  #project the volume at each time t, then concatenate the results
  factored =  reduce(vcat, broadcast(1:(dims.T-1), xws, xRtws,xRLws, xM) do t, xwsₜ, xRtwsₜ, xRLwsₜ, xMₜ
    projectpurchasest(A[:,t+1], A[:,t], xwsₜ, xRtwsₜ, xRLwsₜ, xMₜ)
  end)

  return factored
end

#this just runs the regression for the purposes of extracting the coefficients
function coef(Θ::AbstractVolumePartsIter{TM,TV,T}, Xv::XYIterControl,
  RegressionType::Val) where {TM, TV, T}

  @unpack G,Ã,A = Θ
  dims = (K = size(A,1), T = size(A,2))
  initializeÃandG!(Θ, Xv)
  RHS::TM = hcat(projectbyt(Ã, Xv.xsection), Xv.W̃)

  β = reg(RHS, Xv.ṽ, RegressionType)
  return β
end

#version that performs cross-sectional regressions as needed
function (Θ::AbstractVolumePartsIter{TM,TV,T})(Xv::XYIterControl,
  RegressionType::Val) where {TM,TV,T}
  @unpack G,Ã,A = Θ

  dims = (K = size(A,1), T = size(A,2))

  #set Ã to the cumulative product of G
  initializeÃandG!(Θ, Xv)
  #@info "type of Xv: $(typeof(Xv))"

  RHS::TM = hcat(projectbyt(Ã, Xv.xsection), Xv.W̃)

  #run the regression for A₁
  β = reg(RHS, Xv.ṽ, RegressionType)
  #below shoudl be equivelent to A[:,1] .= β[1:dims.K]
  #β = (b->ifelse(b>0, b, b-100.)).(β) |> TV

  copyto!(A, 1, β, 1, dims.K)
  #A[:,1] .= parent(β)[1:dims.K]

  #absμ .= vec(sum(absμₖ,dims=2))
  #@eval Main controlbeta = $(β |> Vector)
  absμ = RHS * β

  return absμ
end

function (Θ::AbstractVolumePartsIter{TM,TV,T})(Xv::XYIterControl,
  RegressionType::Val{:none},
  noregressionbase::T = T(PARAM[:iternoregressionbase])) where {TM,TV,T}

  @unpack G,Ã,A = Θ

  dims = (K = size(A,1), T = size(A,2))

  #set Ã to the cumulative product of G
  initializeÃandG!(Θ, Xv)
  #@info "type of Xv: $(typeof(Xv))"

  RHS::TM = hcat(projectbyt(Ã, Xv.xsection), Xv.W̃)

  #run the regression for A₁

  βW̃ = (size(Xv.W̃,2) > 0) ? reg(Xv.W̃, Xv.ṽ, Val{:cholesky}()) : TV()
  β = [TV(ones(T, dims.K) .* noregressionbase); βW̃]
  #below shoudl be equivelent to A[:,1] .= β[1:dims.K]
  copyto!(A, 1, β, 1, dims.K)
  #A[:,1] .= parent(β)[1:dims.K]

  #absμ .= vec(sum(absμₖ,dims=2))
  #@eval Main controlbeta = $(β |> Vector)
  absμ = RHS * β

  return absμ
end

#we can probably dispatch all of this just from Xv
#NOTE: this is mainly for ZYGOTE
function (Θ::AbstractVolumePartsIter{TM,TV,T})(G::Matrix{T},
    Xv::XYIterControl,
    RegressionType::Val) where {TM,TV,T}

  dims = (K=size(Θ.A,1), T = size(Θ.A,2))

  someones = ones(T, dims.K, 1)
  prodmatG = cumprodbyrow(G)
  Ã = hcat(someones, prodmatG) |> TM

  RHS::TM = hcat(projectbyt(Ã, Xv.xsection), Xv.W̃)

  #run the regression for A₁
  β = reg(RHS, Xv.ṽ, RegressionType)
  #β = (b->ifelse(b>0, b, b-100.)).(β) |> TV

  #β = ones(size(RHS,2))
  absμ = RHS * β

  return absμ
end

#we can probably dispatch all of this just from Xv
#NOTE: this is mainly for ZYGOTE
function (Θ::AbstractVolumePartsIter{TM,TV,T})(G::Matrix{T},
    Xv::XYIterControl,
    RegressionType::Val{:none},
    noregressionbase::T = T(PARAM[:iternoregressionbase])) where {TM,TV,T}

  dims = (K=size(Θ.A,1), T = size(Θ.A,2))

  someones = ones(T, dims.K, 1)
  prodmatG = cumprodbyrow(G)
  Ã = hcat(someones, prodmatG) |> TM

  RHS::TM = hcat(projectbyt(Ã, Xv.xsection), Xv.W̃)

  #run the regression for A₁
  βW̃ = (size(Xv.W̃,2) > 0) ? reg(Xv.W̃, Xv.ṽ, Val{:choleskyzygote}()) : TV()
  β = [TV(ones(T, dims.K) .* noregressionbase); βW̃]

  absμ = RHS * β

  return absμ
end

#NOTE: this is mainly for ZYGOTE
function (Θ::AbstractVolumePartsIter{TM,TV,T})(G::Matrix{T},
    Xv::XYIterControl,
    RegressionType::Val{:zygotefallbackreg}) where {TM,TV,T}

  dims = (K=size(Θ.A,1), T = size(Θ.A,2))

  someones = ones(T, dims.K, 1)
  prodmatG = cumprodbyrow(G, limitA=true)
  Ã = hcat(someones, prodmatG) |> TM

  RHS::TM = hcat(projectbyt(Ã, Xv.xsection), Xv.W̃)

  #run the regression for A₁
  β = reg(RHS, Xv.ṽ, RegressionType)
  #β = (b->ifelse(b>0, b, b-100.)).(β) |> TV

  absμ = RHS * β

  return absμ
end

function lossforzygote(G::Matrix{T}, Θ::AbstractVolumePartsIter{TM,TV,T}, Xv::XYIterControl,;
  RegressionType::Val=Val{PARAM[:iterregressiontypezygote]}()) where {TM,TV,T}

  #throw("hello")

  #ε = factored * Aₜ .- Xv.ṽ
  Σε2 = sum((Θ(G, Xv, RegressionType) .- Xv.ṽ).^2) #0.3ms

  return Σε2
end

function ∇lossforzygote(Θ, Xv::AbstractXYIter{TM,TV,T},  GrowthType::Val;
    RegressionType=Val(PARAM[:iterregressiontypezygote]),
    GradType=Val(PARAM[:itergradtype])) where {TM,TV,T}

  #the actual gradient function

  function ∇loss(G::Matrix, GradType::Val{:zygote})
    local grad

    try
      grad = gradient((G)->lossforzygote(G, Θ, Xv,
        RegressionType=RegressionType), G)[1]
    catch err
      (err == InterruptException()) && throw(err)
      errmessage = "$err"
      errmessage = length(errmessage)  > 1000 ? errmessage[1:1000] : errmessage
      @warn "Gradient failed with error $errmessage\nAttempting fallback method"
      try
        grad = gradient((G)->lossforzygote(G, Θ, Xv,
          RegressionType=Val{:zygotefallbackreg}()), G)[1]
      catch err2
        @error "Gradient failed with error $errmessage" exception=(err, catch_backtrace())
        errmessage2 = "$err2"
        errmessage2 = length(errmessage2)  > 1000 ? errmessage2[1:1000] : errmessage2
        @error "Gradient fallback also failed with error $errmessage2" exception=(err2, catch_backtrace())
      end
    end

  end

  function ∇loss(G::Matrix, GradType::Val{:ag})
   return ∇lossag(G, Θ, Xv)
  end

  ∇loss(G::Matrix) = ∇loss(G::Matrix, GradType)

  function ∇vec!(out, x::TV, ::Val{:identity}) where TV <: Vector

    if !(x===Θ.g) #if they are the same array, no need to update
      copyto!(Θ.g, x)
    end

      copyto!(out, ∇loss(Θ.G)) #converts to TV

    return nothing
  end

  function ∇vec!(out, x::TV, ::Val{:identity}) where TV <: CuVector

    copyto!(Θ.g, x)
    G = Θ.G |> Matrix
    copyto!(out, ∇loss(G)) #converts to TV

    return nothing
  end

  function ∇vec!(out, x::TV, ::Val{:log}) where TV <: Vector
    #@info "∇vec! log"
    (x ≡ Θ.g) && throw("x and g should not be the same vector if growthtype≠identity")
    copyto!(Θ.g, exp.(x))

    copyto!(out, ∇loss(Θ.G) .* Θ.G) #adjust by G since d(expy)/dx=expy*dy/dx

    return nothing
  end

  function ∇vec!(out, x::TV, ::Val{:log}) where TV <: CuVector
    (x ≡ Θ.g) && throw("x and g should not be the same vector if growthtype≠identity")
    copyto!(Θ.g, exp.(x))
    G = Θ.G |> Matrix
    copyto!(out, ∇loss(G) .* G)
    #@info "sumabsout: $(sum(abs.(out)))"
  end

  ∇vec!(out, x) = ∇vec!(out, x |> TV, GrowthType) #this version performs the conversion
  ∇vec!(out, x::TV) = ∇vec!(out, x, GrowthType)

  return ∇vec!
end

#Updates G based on values of A
#of course, this changes the RHS side of teh regressions, so any
#factored cross-sections need to be recalculated
updateGfromA!(args...;kwargs...) = throw("I dropped updateGfromA!. If needed, just replace
  with Θ.A[:,2:end] ./ Θ.A[:,1:(end-1)] (though be careful about Ã vs A)")

#uses a base value of A and the G matrix to update all other values of A
function updateAfromGAₜ!(Θ::VolumePartsIter, Xv::AbstractXYIter{TM,TV,T},t::Int) where {TM,TV,T}

  local A₀::TV
  local Aₜ::TV = Θ.A[:,t]
  dims = (T=size(Θ.A,2),K=size(Θ.A,1))

  #compute the cumulative growth
  ΠG::TM = similar(Θ.A)
  ΠG[:,1] .= T(1.0)
  ΠG[:,2:end] .= cumprodbyrow(Θ.G)

  A₀ = Aₜ ./ ΠG[:,t] #identify the intial value of assets

  Θ.A .= A₀ .* ΠG #use the growth to compute A
  (vec(Θ.A[:,t]) ≈ Aₜ) || throw("vec(Θ.A[:,t]) ($(vec(Θ.A[:,t]))) ≠ Aₜ ($(Aₜ))! ") #quick integrity check
  Θ.A[:,t] .= Aₜ #preserve exact values to partly mitigate roundoff error

  return nothing
end

function updateAfromGAₜ!(Θ::VolumePartsIter, Xv::AbstractXYIter{TM,TV,T},
  t::Int) where {TM<:CuMatrix,TV<:CuVector,T}

  local Aₜ::TV = Θ.A[:,t]
  dims = (T=size(Θ.A,2),K=size(Θ.A,1))

  #compute the cumulative growth
  ΠG::TM = similar(Θ.A)
  ΠG[:,1] .= T(1.0)
  ΠG[:,2:end] .= cumprodbyrow(Θ.G)
  #ΠG .= hcat(CUDA.ones(T, dims.K, 1), cumprod(Θ.G,dims=2))

  #A₀ = Aₜ ./ ΠG[:,t] #identify the intial value of assets
  #Θ.A .= A₀ .* ΠG
  A = (Aₜ ./ ΠG[:,t]) .* ΠG
  copyto!(Θ.A, A) #use the growth to compute A
  #@assert vec(Θ.A[:,t]) ≈ Aₜ #quick integrity check
  Θ.A[:,t] .= Aₜ #preserve exact values to partly mitigate roundoff error

  return nothing
end

#computs the loss, but does so
#given the current values of A irrespective of G
function lossfromA(Θ::AbstractVolumePartsIter{TM,TV,T}, Xv::XYIterControl, RegressionType::Val
    ) where {TM,TV,T}

  @unpack G,Ã,A = Θ
  dims = (K = size(A,1), T = size(A,2))
  RHS::TM = hcat(projectbyt(A, Xv.xsection), Xv.W̃) #NOTE the use of A, not Ã
  β = reg(RHS, Xv.ṽ, RegressionType)
  absμ = RHS * β

  loss = sum((Xv.ṽ .- absμ) .^2 )

  return loss
end

#NOTE: could be used in some future GPU version
cudaGᵢ(Aᵢ,LAᵢ)= Aᵢ/LAᵢ
cudafactorAᵢ(wsᵢ, RLwsᵢ, Gᵢ) = CUDA.abs(wsᵢ - RLwsᵢ/Gᵢ)
cudafactorLAᵢ(wsᵢ, RLwsᵢ, Gᵢ) = CUDA.abs(Gᵢ*wsᵢ - RLwsᵢ)

#takes a volume part object and writes the results to a df
function formresults(panel::AbstractDataFrame, ms::MeasureSpec, Θ::AbstractVolumeParts{TM,TV,T},
    Xv::XYIterControl,  ::Val{:iter}) where {TM,TV,T}

  grossexposures = (Θ.A |> Matrix)' .* grossleverages(ms)'
  growthrates = vcat(zeros(eltype(grossexposures), size(grossexposures,2))', (Θ.G|>Matrix!)')

  Fcoefs::Vector{Symbol} = ms.Fξs
  Fgcoefs::Vector{Symbol} = (s->replace(string(s),r"^Z"=>"G_") |> Symbol).(Fcoefs)
  FWcoefs::Vector{Symbol} = [ms.FW...;]

  if length(Fcoefs) ≠ size(Θ.A,1)
    throw("Coeffcients not in expected form. Fcoefs: $Fcoefs")
  end

  resultsgrossexposures = DataFrame(grossexposures,Fcoefs)
  resultsgrowthrates = DataFrame(growthrates,Fgcoefs)

  tidx::Dict = gettidx(panel, ms)
  tidxinv::Vector{Date} = sort(collect(keys(tidx)))

  results::DataFrame = DataFrame(date=tidxinv)

  results = [results resultsgrossexposures resultsgrowthrates]

  return results

end

max_stars_repo_name: clintonTE/CCA

max_issues_repo_name: clintonTE/CCA

max_forks_repo_name: clintonTE/CCA


text:
module TKTDbayes

end # module

max_stars_repo_name: virgile-baudrot/TKTDbayes

max_issues_repo_name: virgile-baudrot/TKTDbayes

max_forks_repo_name: virgile-baudrot/TKTDbayes


text:
# This file solve the flap reponses of blades
function sflapre(ψ,λ_α,θ_cp,θ_lat,θ_lon) #staticflapre
  # 其中θ应当输入平均值，例如mean(θ_cp)
  γ_ = (ρ*8.2*(0.06)*R^4)/Iβ
  μ = μ_air
  F0 = 1+μ^2/2
  F_B1 = μ
  F_λ = 1
  A0 = 2*μ
  A_B1 = 1+3*μ^2/4
  A_λ = μ
  A_a1 = 1-μ^2/4
  B_β0 = μ
  B_A1 = 1+μ^2/4
  B_b1 = 1+μ^2/4
  γ_B = γ_/2*B_β0/B_b1
  θ0_ = θcp[1]
  θ_lat_ = θ_lat[1]
  θ_lon_ = θ_lon[1]
  Matrix_β = [γ_/2*F0 0 γ_/2*F_B1 γ_/2*F_λ;
              A0/A_a1 0 A_B1/A_a1 A_λ/A_a1;
              γ_B*F0 -B_A1/B_b1  γ_B*F_B1 γ_B*F_λ]*[θ0_,θ_lat_,θ_lon_,λ_α]
  β0 = Matrix_β[1]
  β_lon = Matrix_β[2]
  β_lat = Matrix_β[3]
  β = β0-β_lon*cos(ψ)-β_lat*sin(ψ)
  dβ = Ω*(β_lon*sin(ψ)-β_lat*cos(ψ))
  ddβ = Ω*(β_lon*cos(ψ)+β_lat*cos(ψ))
  return β,dβ,ddβ,β_lon,β_lat,β0
end

function dflapre(betanow,betaxnow,Mbeta_aero)	#动态挥舞响应
  # 输入：当前挥舞角、挥舞角速度
  # 输出：步进挥舞角、挥舞角速度
	# 常规参数
	Mcen = ecut*R*Ω^2*m_*1/2*(R-ecut*R)^2
	MG = m_*g*1/2*(R-ecut*R)^2
  nite = 10 # 迭代步数
  h = dψ/nite # 龙格库塔求解间隔
  ite = 1
  betanext = 0.0
  betaxnext = 0.0

	# 挥舞方程
	funcbetax(xxbeta) = 1/(Iβ*Ω^2)*(Mbeta_aero-MG*cos(xxbeta)-
											Mcen*sin(xxbeta))-
											cos(xxbeta)*sin(xxbeta)

	# 龙格库塔求解
  while ite<=nite
  	k1 = betaxnow
  	kx1 = funcbetax(betanow)
  	k2 = betaxnow+h/2*k1
  	kx2 = funcbetax(betanow+h/2*kx1)
  	k3 = betaxnow+h/2*k2
  	kx3 = funcbetax(betanow+h/2*kx2)
  	k4 = betaxnow+h*k3
  	kx4 = funcbetax(betanow+h*kx3)
  	betanext = betanow + h/6*(k1+2*k2+2*k3+k4)
  	betaxnext = betaxnow + h/6*(kx1+2*kx2+2*kx3+kx4)
    betanow = betanext
    betaxnow = betaxnext
    ite = ite+1
  end

	return betanext,betaxnext
end

max_stars_repo_name: dingzhw/RotorTrim

max_issues_repo_name: helicopterliu/RotorTrim

max_forks_repo_name: helicopterliu/RotorTrim


text:
using Documenter, ArviZ
using MCMCChains: MCMCChains # make `from_mcmcchains` available for API docs
using SampleChains: SampleChains # make `from_samplechains` available for API docs

makedocs(;
    modules=[ArviZ],
    sitename="ArviZ.jl",
    pages=[
        "Home" => "index.md",
        "Quickstart" => "quickstart.md",
        "API" => "api.md",
        "Example Gallery" =>
            ["Matplotlib" => "mpl_examples.md", "Bokeh" => "bokeh_examples.md"],
        "Reference" => "reference.md",
    ],
    checkdocs=:exports,
    format=Documenter.HTML(;
        prettyurls=haskey(ENV, "CI"),
        assets=["assets/favicon.ico"],
        sidebar_sitename=false,
        canonical="stable",
    ),
    linkcheck=true,
    analytics="G-W1G68W77YV",
)

deploydocs(; repo="github.com/arviz-devs/ArviZ.jl.git", devbranch="main", push_preview=true)


text:
struct UnitCell{Dimensions, T <: AbstractFloat}
    tag::Symbol # currently unused
    # could be used to speed up neighbor search
    basis::Vector{Point{Dimensions, T}}
    dirs::NTuple{Dimensions, Vec{Dimensions, T}}
end

function UnitCell(
        pos::Point{D, T},
        dirs::NTuple{D, Vec{D, T}}
    ) where {D, T <: AbstractFloat}
    UnitCell([pos], dirs)
end

function UnitCell(
        basis::Vector{Point{D, T}},
        dirs::NTuple{D, Vec{D, T}}
    ) where {D, T <: AbstractFloat}
    UnitCell(:None, basis, dirs)
end

"""
    get_pos(uc, a, u)

Returns R = Rₐ + ∑ᵢ uᵢaᵢ for a UnitCell uc.
"""
function get_pos(
        uc::UnitCell{D, T},
        alpha::Int64,
        uvw::NTuple{D, Int64}
    ) where {D, T}
    uc.basis[alpha] + reduce(+, map((i, v) -> i*v, uvw, uc.dirs))
    # B.pos + reduce(+, map((i, v) -> i * v, uvw, B.dirs))
end

# Returns the dimensionality of a UnitCell
dims(uc::UnitCell{D, T}) where {D, T} = D

function merge(ucs::UnitCell{D, T}...) where {D, T}
    @assert all(
        uc -> all(
            v -> v in ucs[1].dirs,
            uc.dirs
        ),
        ucs
    ) "Lattice vectors must match!"

    UnitCell(
        ucs[1].tag,
        [p for uc in ucs for p in uc.basis],
        ucs[1].dirs
    )
end

# TODO
# center! for more safety in get_neighbors

# Literally just constructors from here on out

################################################################################
### 2D Lattices
################################################################################

"""
    square()

Returns a UnitCell representing a square lattice.

Keyword arguments include
- `pos::Point{2, T} = Point{2}(0.)`: Starting position for this lattice.
- `a::T = 1.0`: Scaling along the first lattice vector.
- `b::T = 1.0`: Scaling along the second lattice vector.
- `isa::Symbol = :cubic`: Lattice shape to use (:cubic or :primitive)
"""
function square(;
        pos::Point{2, T} = Point{2}(0.),
        a::T = 1.0,
        b::T = 1.0,
        isa::Symbol = :cubic
    ) where T <: AbstractFloat
    UnitCell(:cubic, pos, (a * Vec(1., 0.), b * Vec(0., 1.)))
end

"""
    triangle()

Returns a UnitCell representing a triangular lattice.

Keyword arguments include
- `pos::Point{2, T} = Point{2}(0.)`: Starting position for this lattice.
- `a::T = 1.0`: Scaling along the first lattice vector.
- `b::T = 1.0`: Scaling along the second lattice vector.
- `isa::Symbol = :primitive`: Lattice shape to use (:cubic or :primitive)
"""
function triangle(;
        pos::Point{2, T} = Point{2}(0.),
        a::T = 1.0,
        b::T = 1.0,
        isa::Symbol = :primitive
    ) where T <: AbstractFloat

    if isa == :cubic
        return UnitCell(
            :cubic,
            [pos, pos + Point(a*sind(30.), b*cosd(30.))],
            (a * Vec(1., 0.), b * Vec(0., 1.))
        )
    else
        return UnitCell(
            :primitive,
            pos,
            (a*Vec(0., 1.), b*Vec(cosd(30.0), sind(30.0)))
        )
    end
end

"""
    honeycomb()

Returns a UnitCell representing a honeycomb lattice.

Keyword arguments include
- `pos::Point{2, T} = Point{2}(0.)`: Starting position for this lattice.
- `a::T = 1.0`: Scaling along the first lattice vector.
- `b::T = 1.0`: Scaling along the second lattice vector.
- `isa::Symbol = :primitive`: Lattice shape to use (:cubic or :primitive)
"""
function honeycomb(;
        pos::Point{2, T} = Point{2}(0.),
        a::T = 1.0,
        b::T = 1.0,
        isa::Symbol = :primitive
    ) where T <: AbstractFloat

    if isa == :cubic
        return UnitCell(
            :cubic,
            [pos, pos + Point(0.0, b)],
            (2a * cosd(30) * Vec(1., 0.), 2b * cosd(30) * Vec(0., 1.))
        )
    else
        return UnitCell(
            :primitive,
            [pos, pos + Point(a, 0.)],
            (
                2a * cosd(30) * Vec(0., 1.),
                2b * cosd(30) * Vec(cosd(30.0), sind(30.0))
            )
        )
    end
end

################################################################################
### 3D Lattices
################################################################################

"""
    sc()

Returns a UnitCell representing a simple cubic lattice.

Keyword arguments include
- `pos::Point{3, T} = Point{3}(0.)`: Starting position for this lattice.
- `a::T = 1.0`: Scaling along the first lattice vector.
- `b::T = 1.0`: Scaling along the second lattice vector.
- `c::T = 1.0`: Scaling along the third lattice vector.
- `isa::Symbol = :cubic`: Lattice shape to use (:cubic or :primitive)
"""
function sc(;
        pos::Point{3, T} = Point{3}(0.),
        a::T = 1.0,
        b::T = 1.0,
        c::T = 1.0,
        isa::Symbol = :cubic
    ) where T <: AbstractFloat
    UnitCell(
        :cubic,
        [pos],
        (
            a * Vec(1., 0., 0.),
            b * Vec(0., 1., 0.),
            c * Vec(0., 0., 1.)
        )
    )
end

"""
    bcc()

Returns a UnitCell representing a body centered cubic lattice.

Keyword arguments include
- `pos::Point{3, T} = Point{3}(0.)`: Starting position for this lattice.
- `a::T = 1.0`: Scaling along the first lattice vector.
- `b::T = 1.0`: Scaling along the second lattice vector.
- `c::T = 1.0`: Scaling along the third lattice vector.
- `isa::Symbol = :primitive`: Lattice shape to use (:cubic or :primitive)
"""
function bcc(;
        pos::Point{3, T} = Point{3}(0.),
        a::T = 1.0,
        b::T = 1.0,
        c::T = 1.0,
        isa::Symbol = :primitive
    ) where T <: AbstractFloat

    if isa == :cubic
        return UnitCell(
            :cubic,
            [pos, pos + Point{3}(0.5)],
            (
                a * Vec(1., 0., 0.),
                b * Vec(0., 1., 0.),
                c * Vec(0., 0., 1.)
            )
        )
    else
        return UnitCell(
            :primitive,
            [pos],
            (
                a * 0.5 * Vec(1., 1., -1.),
                b * 0.5 * Vec(1., -1., 1.),
                c * 0.5 * Vec(-1., 1., 1.)
            )
        )
    end
end

"""
    fcc()

Returns a UnitCell representing a face centered cubic lattice.

Keyword arguments include
- `pos::Point{3, T} = Point{3}(0.)`: Starting position for this lattice.
- `a::T = 1.0`: Scaling along the first lattice vector.
- `b::T = 1.0`: Scaling along the second lattice vector.
- `c::T = 1.0`: Scaling along the third lattice vector.
- `isa::Symbol = :primitive`: Lattice shape to use (:cubic or :primitive)
"""
function fcc(;
        pos::Point{3, T} = Point{3}(0.),
        a::T = 1.0,
        b::T = 1.0,
        c::T = 1.0,
        isa::Symbol = :primitive
    ) where T <: AbstractFloat

    if isa == :cubic
        return UnitCell(
            :cubic,
            [
                pos,
                pos + Point(0.5, 0.5, 0.0),
                pos + Point(0.5, 0.0, 0.5),
                pos + Point(0.0, 0.5, 0.5)
            ],
            (
                a * Vec(1., 0., 0.),
                b * Vec(0., 1., 0.),
                c * Vec(0., 0., 1.)
            )
        )
    else
        return UnitCell(
            :primitive,
            [pos],
            (
                a * 0.5 * Vec(1., 1., 0.),
                b * 0.5 * Vec(1., 0., 1.),
                c * 0.5 * Vec(0., 1., 1.)
            )
        )
    end
end

"""
    diamond()

Returns a UnitCell representing a diamond lattice.

Keyword arguments include
- `pos::Point{3, T} = Point{3}(0.)`: Starting position for this lattice.
- `a::T = 1.0`: Scaling along the first lattice vector.
- `b::T = 1.0`: Scaling along the second lattice vector.
- `c::T = 1.0`: Scaling along the third lattice vector.
- `isa::Symbol = :primitive`: Lattice shape to use (:cubic or :primitive)
"""
function diamond(;
        pos::Point{3, T} = Point{3}(0.),
        a::T = 1.0,
        b::T = 1.0,
        c::T = 1.0,
        isa::Symbol = :primitive
    ) where T <: AbstractFloat

    merge(
        fcc(pos = pos, a=a, b=b, c=c, isa=isa),
        fcc(pos = pos + Point{3}(0.25), a=a, b=b, c=c, isa=isa)
    )
end


text:
# This file is a part of JuliaFEM.
# License is MIT: see https://github.com/JuliaFEM/JuliaFEM.jl/blob/master/LICENSE.md

using JuliaFEM, Test
using AsterReader: RMEDFile, aster_read_nodes, aster_read_data

## Two rings
# RING1 = inner, RING2 = outer, RINGS combined mesh. Set T=1.0 for
# inner ring and T=2.0 for outer ring, measure temperature from middle of ring.
# Results are calculated using Code Aster for comparison.

# test 3d heat, two rings, and compare to CA solution
meshfile = @__DIR__() * "/testdata/primitives.med"
mesh = aster_read_mesh(meshfile, "RINGS_UNION")

rings = Problem(Heat, "RINGS", 1)
rings_elements_1 = create_elements(mesh, "RING1")
rings_elements_2 = create_elements(mesh, "RING2")
update!(rings_elements_1, "thermal conductivity", 1.0)
update!(rings_elements_2, "thermal conductivity", 1.0)
add_elements!(rings, rings_elements_1)
add_elements!(rings, rings_elements_2)

bc_inner = Problem(Dirichlet, "INNER SURFACE", 1, "temperature")
bc_inner_elements = create_elements(mesh, "RING1_INNER")
update!(bc_inner_elements, "temperature 1", 1.0)
add_elements!(bc_inner, bc_inner_elements)

bc_outer = Problem(Dirichlet, "OUTER SURFACE", 1, "temperature")
bc_outer_elements = create_elements(mesh, "RING2_OUTER")
update!(bc_outer_elements, "temperature 1", 2.0)
add_elements!(bc_outer, bc_outer_elements)

analysis = Analysis(Linear)
add_problems!(analysis, rings, bc_inner, bc_outer)
run!(analysis)

temp_jf = rings("temperature", 0.0)

fn = @__DIR__() * "/testdata/rings.rmed"
results = RMEDFile(fn)
nodes = aster_read_nodes(results)
temp_ca = aster_read_data(results, "TEMP")

rtol = [norm(temp_jf[j]-temp_ca[j]) / max(temp_jf[j],temp_ca[j]) for j in keys(temp_jf)]
@test maximum(rtol) < 1.0e-12


text:
using LabClimateDashboard
using Test
using Dates
using DelimitedFiles
using HTTP
using Revise

include(joinpath(@__DIR__, "utils.jl"))

@testset "LabClimateDashboard.jl" begin
    logdata_path = generate_logdata()

    @async LabClimateDashboard.run(logdata_path; port=8050)

    r = HTTP.request("GET", "http://127.0.0.1:8050"; verbose=2, connect_timeout=60)
    # test if we receive http status code 200 (Ok)
    @test r.status == 200

    println("Press ENTER to terminate!")
    readline()
end


text:
using LinearAlgebra

function run_test_bench(model::C, pixel_accuracy::Float64 = 1e-5, ray_accuracy::Float64 = 1e-4) where C <: CameraModel

    @testset "Check basics and interface implementation for $(C)." begin

        forward = lookdirection(model)
        up = updirection(model)
        right = cross(forward, up)

        @testset "Test model axes" begin
            @test norm(forward) ≈ 1.0
            @test norm(up) ≈ 1.0
            @test norm(right)  ≈ 1.0
        end

        pixelsize = sensorsize(model)

        # Get a pixel location somewhere near the image centre.
        image_centre = (pixelsize .+ 1) ./ 2
        slight_offset = 0.042 .* pixelsize
        some_pixel_location = image_centre + slight_offset

        # Get the ray that belongs to that pixel.
        ray = pixel2ray(model, some_pixel_location)

        # Generate a 3D point along that ray.
        point = origin(ray) + 4.2 .* direction(ray)

        # Some models might not implement point2pixel.
        if canreproject(model)
            reprojection = point2pixel(model, point)
            @test some_pixel_location ≈ reprojection atol = pixel_accuracy
        else
            @info "point2pixel not implemented for $(C)."
        end

    end

end