text:
#####
##### Run in 1LM
#####

test_vgg() = AutoTM.Experiments.test_vgg()
conventional_inception()    = AutoTM.Experiments.conventional_inception()
conventional_resnet()       = AutoTM.Experiments.conventional_resnet()
conventional_vgg()          = AutoTM.Experiments.conventional_vgg()
conventional_densenet()     = AutoTM.Experiments.conventional_densenet()

"""
    common_ratios()

Ratios of PMM to DRAM used for the AutoTM paper.
"""
common_ratios() = [
    1 // 0,
    8 // 1,
    4 // 1,
    1 // 1,
    0 // 1,
]

conventional_functions() = [
    conventional_inception(),
    conventional_resnet(),
    conventional_vgg(),
    conventional_densenet(),
]

#####
##### Standard Routine
#####

"""
    run_conventional(fn, optimizers, ratios)

Run experiment for a conventional workload.

Argument Description
--------------------

* `fn`: The network to run. Options are
    - `test_vgg()`
    - `conventional_vgg()`
    - `conventional_inception()`
    - `conventional_resnet()`
    - `conventional_densenet()`

* `optimizers`: Iteratable of optimizers to use. Options:
    - `AutoTM.Optimizer.Static`
    - `AutoTM.Optimizer.Synchronous`
    - `AutoTM.Optimizer.Numa`

* `ratios`: Iterable of ratios (`Rational{Int}`) of ratios of PMM and DRAM to use.
    Defaults to `conventional_ratios() = [1 // 0, 8 // 1, 4 // 1, 1 // 1, 0 // 1]`

Results will be saved to `Benchmarker/data/cpu`.
"""
function run_conventional(fn, optimizers, ratios = common_ratios())
    ratios = common_ratios()

    opt = Iterators.flatten(o.(ratios) for o in optimizers)
    cache = AutoTM.Experiments.CPU_CACHE
    execute(fn, opt, cache, nGraph.Backend("CPU"))
end

############################################################################################
# Plots
############################################################################################

#####
##### Front Plot
#####

function plot_front()
    f = conventional_inception()
    ratio = 4 // 1
    cache = AutoTM.Experiments.CPU_CACHE
    formulations = ("numa", "synchronous",)

    plot_front(f, ratio, cache; formulations = formulations)
end

#####
##### PMM Speedup Plot
#####

function plot_speedup(;
        models = conventional_functions(),
        ratios = common_ratios(),
        formulations = ("numa", "static", "synchronous"),
        cache = AutoTM.Experiments.CPU_CACHE
    )

    # Get rid of the all PMEM and all DRAM case
    deleteat!(ratios, findfirst(isequal(0 // 1), ratios))
    deleteat!(ratios, findfirst(isequal(1 // 0), ratios))

    pgf_speedup(
        models,
        ratios,
        cache;
        formulations = formulations,
    )
end

#####
##### Error Plot
#####

function plot_conventional_error(;
        models = [
            conventional_inception(),
            conventional_resnet(),
            conventional_vgg(),
            conventional_densenet()
        ],
        ratios = common_ratios(),
        legacy = false,
        formulations = ("static", "synchronous")
    )

    caches = [
        AutoTM.Experiments.CPU_CACHE
    ]

    suffix = nothing
    return pgf_error_plot(
        models,
        ratios,
        caches;
        formulations = formulations,
        suffix = suffix,
        legacy = legacy
    )
end

#####
##### Cost Plot
#####

function plot_costs(;
        pairs = [
            conventional_vgg() => "synchronous",
            conventional_densenet() => "synchronous",
            conventional_resnet() => "synchronous",
            conventional_inception() => "synchronous",
        ],
        ratios = common_ratios(),
        cost_ratio = 2.1,
    )

    # Get rid of the all PMEM and all DRAM case
    deleteat!(ratios, findfirst(isequal(0 // 1), ratios))
    cache = AutoTM.Experiments.CPU_CACHE

    pgf_cost(pairs, ratios, cache; cost_ratio = cost_ratio)
end

############################################################################################
# Case Study - Inception
############################################################################################

"""
    case_study_ratios()

Return PMM to DRAM ratios used for the Inception case study.
"""
function case_study_ratios()
    # Start out with the ratios of the whole memory we want to devote to DRAM
    ratios = [ 1 // i for i in 1:10 ]

    # Add some more fractions for the higher end of DRAM and for the case of all PMEM
    push!(ratios,
        1 // 100,
        1 // 80,
        1 // 60,
        1 // 40,
        1 // 20,
        2 // 3,
        3 // 4,
        0 // 1
    )
    sort!(ratios)

    # Convert these ratios into the PMEM // DRAM ratios
    pmem_to_dram_ratios = (one(eltype(ratios)) .- ratios) ./ ratios
    return pmem_to_dram_ratios
end

"""
    inception_case_study()

Run the `conventional_inception()` workload for a large number of different PMM to DRAM ratios.
"""
function inception_case_study()
    f = conventional_inception()

    ratios = case_study_ratios()

    optimizers = Iterators.flatten((
        [Optimizer.Static(r) for r in ratios],
        [Optimizer.Synchronous(r) for r in ratios],
    ))

    cache = SINGLE_KERNEL_PATH
    suffix = "study"

    # Run the workload - don't perform the ratio search because that's not as important
    # for this experiment
    execute(f, optimizers, cache, nGraph.Backend("CPU"), suffix; search_ratio = false)
end

"""
    inception_case_study_plots()

Generate Figure 10a, 10b, and 10c of the paper.
"""
function inception_case_study_plots()
    f = conventional_inception()
    cache = AutoTM.Experiments.CPU_CACHE
    suffix = "study"
    formulations = (
        "static",
        "synchronous",
    )

    # Performance line graph
    a = pgf_plot_performance(f, cache, suffix;
        file = joinpath(FIGDIR, "inception_perf.pdf"),
        formulations = ("static", "synchronous"),
    )

    # Input/output tensor graph
    b = pgf_io_plot(f, cache, suffix;
        file = joinpath(FIGDIR, "inception_io.pdf"),
        formulations = formulations,
    )

    # Movement statistics
    c = pgf_movement_plot(f, cache, suffix;
        file = joinpath(FIGDIR, "inception_movement.pdf"),
        formulation = "synchronous"
    )

    return (
        performance = a,
        io = b,
        movement = c,
    )
end

############################################################################################
# Large Models
############################################################################################

large_inception() = AutoTM.Experiments.large_inception()
large_resnet() = AutoTM.Experiments.large_resnet()
large_densenet() = AutoTM.Experiments.large_densenet()
large_vgg() = AutoTM.Experiments.large_vgg()

"""
    kernel_profile(fns; recache = false)

Perform kernel profiling for each function in `fns`.
Possible values for argument `fns` are:
- `test_vgg()`
- `conventional_vgg()`
- `conventional_inception()`
- `conventional_resnet()`
- `conventional_densenet()`
- `large_vgg()`
- `large_inception()`
- `large_resnet()`
- `large_densenet()`

If `recache = true`, delete previously profiled kernels for each function.
"""
function kernel_profile(fns; recache = false)
    backend = nGraph.Backend("CPU")
    dummy_opt = AutoTM.Optimizer.Static(0)

    for f in wrap(fns)
        AutoTM.Optimizer.factory(
            backend,
            f,
            dummy_opt;
            cache = AutoTM.Profiler.CPUKernelCache(AutoTM.Experiments.CPU_CACHE),
            just_profile = true,
            recache = recache
        )
    end
end

# Default the limit to 185 GB, slightly smaller than the 192 GB DRAM size.
function run_large(f, opt; limit = 185_000_000_000)
    optimizer = opt(limit)
    backend = nGraph.Backend("CPU")
    cache = SINGLE_KERNEL_PATH

    execute(f, optimizer, cache, nGraph.Backend("CPU"); skip_base_check = true, defrag = true)
end

#####
##### 2LM
#####

# Just execute the function directly - nothing too fancy here
function run_2lm(fn)
    opt = Optimizer.Optimizer2LM()
    cache = "nocache"
    backend = nGraph.Backend("CPU")

    savefile = canonical_path(fn, opt, cache, backend)
    stats = ispath(savefile) ? deserialize(savefile) : AutoTM.Profiler._base_stats()
    fex = Optimizer.factory(backend, fn, opt)
    runtime = gettime(fex)

    run = Dict(
        :actual_runtime => runtime,
    )
    push!(stats.runs, run)
    push!(stats.runs, run)
    serialize(savefile, stats)
end

# Generate the plot for the large workloads
function plot_large()
    fns = (large_vgg(), large_inception(), large_resnet(), large_densenet())
    cache = AutoTM.Experiments.CPU_CACHE
    cache_2lm = "nocache"
    formulations = ("static", "synchronous")

    pgf_large_performance(fns, cache, cache_2lm; formulations = formulations)
end

############################################################################################
# GPU Benchmarks
############################################################################################

const GPU_MAX_MEMORY = Ref(11_000_000_000)

# Setup the memory allocator in CuArrays to be as conservative as possible.

# Got this number from looking at nvidia-smi after all the GPU initialization code
# in ngraph runs.
#
# Probably worth double-checking
#const GPU_MEMORY_OVERHEAD = 561_000_000
const GPU_MEMORY_OVERHEAD = Ref(800_000_000)
gpu_adjusted_memory() = GPU_MAX_MEMORY[] - GPU_MEMORY_OVERHEAD[]

"""
    gpu_fns()

Collection of benchmarks used to test the GPU version of AutoTM.
"""
gpu_fns() = (
    AutoTM.Experiments.Inception_v4(64),
    AutoTM.Experiments.Inception_v4(128),
    AutoTM.Experiments.Inception_v4(256),
    AutoTM.Experiments.Resnet200(32),
    AutoTM.Experiments.Resnet200(64),
    AutoTM.Experiments.Resnet200(128),
    AutoTM.Experiments.DenseNet(32),
    AutoTM.Experiments.DenseNet(64),
    AutoTM.Experiments.DenseNet(128),
    AutoTM.Experiments.Vgg19(64),
    AutoTM.Experiments.Vgg19(128),
)

function gpu_profile(; recache = false)
    # Switch to a better allocator since we're allocating a lot.
    CuArrays.pool[] = CuArrays.BinnedPool
    CuArrays.pool[].init()

    fns = gpu_fns()
    opt = AutoTM.Optimizer.Synchronous(gpu_adjusted_memory())
    cache = AutoTM.Experiments.GPU_CACHE
    backend = nGraph.Backend("GPU")

    for f in fns
        @show name(f)
        try
            execute(
                f,
                opt,
                cache,
                backend;
                # kwargs
                just_profile = true,
                skip_base_check = true,
                recache = recache,
            )
        catch e
            isa(e, AutoTM.Profiler.CompilerExit) || rethrow(e)
        end
    end
end

"""
    gpu_go(fn; limit = gpu_adjusted_memory())

Run the `i`th GPU benchmark.

The code is set up this way to facilitate restarting Julia between sessions.
"""
run_gpu(i::Integer; kw...) = run_gpu(gpu_fns()[i]; kw...)
function run_gpu(
        fn;
        limit = gpu_adjusted_memory(),
        optimizers = (AutoTM.Optimizer.Synchronous, AutoTM.Optimizer.Asynchronous),
    )

    # Manually invoke GC to clean up anything that might be still holding data on the GPU.
    CuArrays.reclaim()
    GC.gc()

    @info "Running $(name(fn))"
    opt = [o(limit) for o in wrap(optimizers)]
    cache = AutoTM.Experiments.GPU_CACHE
    backend = nGraph.Backend("GPU")

    execute(fn, opt, cache, backend; adjust_io = true)
    return nothing
end

"""
    gpu_benchmarks(itr = 1:length(gpu_fns()))

Run all GPU benchmarks.
A subset can be run by passing a custom iterator.
"""
gpu_benchmarks(itr = 1:length(gpu_fns())) = run_gpu.(itr)

"""
    plot_gpu_performance(fns = gpu_fns(), formulations = ("synchronous", "asynchronous"))

Generate the GPU performance plot, comparing AutoTM with `cudaMallocManaged`.
"""
function plot_gpu_performance(;
        fns = gpu_fns(),
        formulations = ("synchronous", "asynchronous")
    )

    return pgf_gpu_performance_plot(
        fns,
        AutoTM.Experiments.GPU_CACHE;
        formulations = formulations
    )
end

max_stars_repo_name: darchr/AutoTM

max_issues_repo_name: darchr/AutoTM

max_forks_repo_name: darchr/AutoTM


text:
# NOTE: these are useful routines that are only loaded when the GMT package is already loaded in the REPL

using .GMT

export ImportTopo

"""
    Topo = ImportTopo(limits; file::String="@earth_relief_01m.grd")

Uses GMT to download the topography of a certain region, specified with limits=[lon_min, lon_max, lat_min, lat_max]

Note:
====
- latitude values in the southern hemisphere should have a minus sign (e.g., -2.8)
- longitude values that are "west" should *either* come with a minus sign *or* are defined by values >180

| Dataset                 |   Resolution |   Description                                               |
|:----------------        | ------------ | ----------------------------------------------------------- |
| "@earth\\_relief\\_01s" |	1 arc sec 	 | SRTM tiles (14297 tiles, land only, 60S-60N) [NASA/USGS]    |
| "@earth\\_relief\\_03s"	|   3 arc sec	 | SRTM tiles (14297 tiles, land only, 60S-60N) [NASA/USGS]    |
| "@earth\\_relief\\_15s"	|  15 arc sec	 | SRTM15+ [David Sandwell, SIO/UCSD]                          |
| "@earth\\_relief\\_30s"	|  30 arc sec	 | SRTM30+ [Becker et al., 2009, SIO/UCSD]                     |
| "@earth\\_relief\\_01m"	|   1 arc min	 | ETOPO1 Ice surface [NEIC/NOAA]                              |
| "@earth\\_relief\\_02m"	|   2 arc min	 | ETOPO2v2 Ice surface [NEIC/NOAA]                            |
| "@earth\\_relief\\_03m"	|   3 arc min	 | ETOPO1 after Gaussian spherical filtering (5.6 km fullwidth)|
| "@earth\\_relief\\_04m"	|   4 arc min	 | ETOPO1 after Gaussian spherical filtering (7.5 km fullwidth)|
| "@earth\\_relief\\_05m"	|   5 arc min	 | ETOPO1 after Gaussian spherical filtering (9 km fullwidth)  |
| "@earth\\_relief\\_06m"	|   6 arc min	 | ETOPO1 after Gaussia30n spherical filtering (10 km fullwidth) |
| "@earth\\_relief\\_10m"	|  10 arc min	 | ETOPO1 after Gaussian spherical filtering (18 km fullwidth) |
| "@earth\\_relief\\_15m"	|  20 arc min	 | ETOPO1 after Gaussian spherical filtering (28 km fullwidth) |
| "@earth\\_relief\\_20m"	|  20 arc min	 | ETOPO1 after Gaussian spherical filtering (37 km fullwidth) |
| "@earth\\_relief\\_30m"	|  30 arc min	 | ETOPO1 after Gaussian spherical filtering (55 km fullwidth) |
| "@earth\\_relief\\_60m"	|  60 arc min	 | ETOPO1 after Gaussian spherical filtering (111 km fullwidth)|

*Note*: this routine is only available once the GMT.jl package is loaded on the REPL

# Example
```julia-repl
julia> Topo = ImportTopo([4,20,37,49]);
GeoData
  size  : (960, 720, 1)
  lon   ϵ [ 4.0 : 19.983333333333334]
  lat   ϵ [ 37.0 : 48.983333333333334]
  depth ϵ [ -3.8725 km : 4.2495 km]
  fields: (:Topography,)
```
And you can save this to Paraview with
```julia
julia> Write_Paraview(Topo,"Topo_Alps")
1-element Vector{String}:
 "Topo_Alps.vts"
```
"""
function ImportTopo(limits; file::String="@earth_relief_01m.grd")

    # Correct if negative values are given (longitude coordinates that are west)
    ind = findall(limits[1:2] .< 0);

    if (limits[1] < 0) && (limits[2] < 0)
      limits[ind] .= 360 .+ limits[ind];
      limits[1:2] = sort(limits[1:2])
    end

    # Download topo file
    G               =   gmtread(file, limits=limits, grid=true);

    # Transfer to GeoData
    nx,ny           =   size(G.z,2), size(G.z,1)
    Lon,Lat,Depth   =   LonLatDepthGrid(G.x[1:nx],G.y[1:ny],0);
    Depth[:,:,1]    =   1e-3*G.z';
    Topo            =   GeoData(Lon, Lat, Depth, (Topography=Depth*km,))

    return Topo

end

"""
  ImportTopo(; lat::Vector{2}, lon::Vector{2}, file::String="@earth_relief_01m.grd")

Imports topography (using GMT), by specifying keywords for latitude and longitude ranges

# Example
=========
```julia
julia> Topo = ImportTopo(lat=[30,40], lon=[30, 50] )
```
The values can also be defined as tuples:
```julia
julia> Topo = ImportTopo(lon=(-50, -40), lat=(-10,-5), file="@earth_relief_30s.grd")
```

"""
ImportTopo(; lat=[37,49], lon=[4,20], file::String="@earth_relief_01m.grd") = ImportTopo([lon[1],lon[2], lat[1], lat[2]], file=file)

ext: jl


text:
using DiffEqOperators, Test

@testset "Gradient Operation on a 2-dimensional function" begin

    # For testing the faster implementation for a 2-dim function
    s = x, y = (-5:1.25:5, -5:1.25:5)
    dx = dy = x[2] - x[1]

    f(x::T, y::T) where T = x^2 + y^2

    u0 = [f(X, Y) for X in x, Y in y]

    # Analytic Gradient of the function is given by u_analytic = 2x ê₁ + 2y ê₂

    u_analytic = zeros(Float64,(size(u0).-2)...,2)

    for j in 1:length(y)-2, i in 1:length(x)-2
        u_analytic[i,j,1] = 2*x[i+1]
        u_analytic[i,j,2] = 2*y[j+1]
    end

    A = Gradient(4,(dx,dy),size(u0).-2)

    u = A*u0

    for I in CartesianIndices(u)
        @test u[I] ≈ u_analytic[I] atol=1e-3
    end

    # check for non-uniform grid
    dx = dy = 1.25*ones(10)

    A = Gradient(4,(dx,dy),size(u0).-2)

    u = A*u0

    for I in CartesianIndices(u)
        @test u[I] ≈ u_analytic[I] atol=1e-3
    end
end

@testset "Gradient Operation on a 3-dimensional function" begin

    # For testing the faster implementation for a 3-dim function
    s = x, y, z = (-5:1.25:5, -5:1.25:5, -5:1.25:5)
    dx = dy = dz = x[2] - x[1]

    f(x::T, y::T, z::T) where T = x^2 + y^2 + z^2

    u0 = [f(X, Y, Z) for X in x, Y in y, Z in z]

    # Analytic Gradient of the function is given by u_analytic = 2x ê₁ + 2y ê₂  + 2z ê₃

    u_analytic = zeros(Float64,(size(u0).-2)...,3)

    for k in 1:length(z)-2, j in 1:length(y)-2, i in 1:length(x)-2
        u_analytic[i,j,k,1] = 2*x[i+1]
        u_analytic[i,j,k,2] = 2*y[j+1]
        u_analytic[i,j,k,3] = 2*z[k+1]
    end

    A = Gradient(4,(dx,dy,dz),size(u0).-2)

    u = A*u0

    for I in CartesianIndices(u)
        @test u[I] ≈ u_analytic[I] atol=1e-3
    end

    # check for multiplication with constant

    u1 = 2*A*u0

    for I in CartesianIndices(u)
        @test u1[I] ≈ 2*u[I] atol=1e-3
    end

    # check for non-uniform grid
    dx = dy = dz = 1.25*ones(10)

    A = Gradient(4,(dx,dy,dz),size(u0).-2)

    u = A*u0

    for I in CartesianIndices(u)
        @test u[I] ≈ u_analytic[I] atol=1e-3
    end
end

@testset "Gradient Operation on a 4-dimensional function" begin

    # For Testing the fallback implementation

    s = x, y, z, w = (-5:1.25:5, -5:1.25:5, -5:1.25:5, -5:1.25:5)
    dx = dy = dz = dw = x[2] - x[1]

    f(x::T, y::T, z::T, w::T) where T = x^2 + y^2 + z^2 + w^2

    u0 = [f(X, Y, Z, W) for X in x, Y in y, Z in z, W in w]

    # Analytic Gradient of the function is given by u_analytic = 2x ê₁ + 2y ê₂  + 2z ê₃ + 2w ê₄

    u_analytic = zeros(Float64,(size(u0).-2)...,4)

    for l in 1:length(w)-2, k in 1:length(z)-2, j in 1:length(y)-2, i in 1:length(x)-2
        u_analytic[i,j,k,l,1] = 2*x[i+1]
        u_analytic[i,j,k,l,2] = 2*y[j+1]
        u_analytic[i,j,k,l,3] = 2*z[k+1]
        u_analytic[i,j,k,l,4] = 2*w[l+1]
    end

    A = Gradient(4,(dx,dy,dz,dw),size(u0).-2)

    u = A*u0

    for I in CartesianIndices(u)
        @test u[I] ≈ u_analytic[I] atol=1e-3
    end
end


text:
"""
Building blocks and high level functions for filtering a collection randomly.

The typical use case would be having an iterable (eg rows from a table) for which it is
desirable to keep a subsample based on some characteristic (eg a person identifier).

`random_subsample` provides a filtered iterable, while `RandomSubSampler` and `randkeep!`
expose the building blocks.
"""
module RandomSubSamplers

export RandomSubSampler, randkeep!, random_subsample

using ArgCheck: @argcheck
using DocStringExtensions: FUNCTIONNAME, SIGNATURES
using Random: AbstractRNG, GLOBAL_RNG, rand

struct RandomSubSampler{T, R <: Real}
    table::Dict{T,Bool}
    ratio::R
    function RandomSubSampler{T}(ratio::R) where {T, R <: Real}
        @argcheck 0 ≤ ratio ≤ 1
        new{T,R}(Dict{T,Bool}(), ratio)
    end
end

"""
$(FUNCTIONNAME)([rng], rs, elt)
"""
randkeep!(rng::AbstractRNG, rs::RandomSubSampler, elt) =
    get!(() -> rand(rng) < rs.ratio, rs.table, elt)

randkeep!(rs::RandomSubSampler, elt) = randkeep!(GLOBAL_RNG, rs, elt)

"""
$(SIGNATURES)

Try to infer a type for the dictionary keys from `itr`.
"""
_inferred_keytype(key, itr) = Core.Compiler.return_type(key, (eltype(itr), ))

"""
$(FUNCTIONNAME)([rng], key, ratio, itr)

Return a (lazy) iterator that randomly keeps `ratio` of `itr` by `key`.

`keytype` can be provided for the dictionary.
"""
function random_subsample(rng::AbstractRNG, key, ratio::Real, itr;
                          keytype = _inferred_keytype(key, itr))
    rs = RandomSubSampler{keytype}(ratio)
    Iterators.filter(elt -> randkeep!(rng, rs, key(elt)), itr)
end

random_subsample(key, ratio::Real, itr;
                 keytype = _inferred_keytype(key, itr)) =
    random_subsample(GLOBAL_RNG, key, ratio, itr; keytype = keytype)

end # module

max_stars_repo_name: tpapp/RandomSubSamplers

max_issues_repo_name: tpapp/RandomSubSamplers

max_forks_repo_name: tpapp/RandomSubSamplers


text:
using EarlyStopping, Dates, Test, InteractiveUtils

@testset "criteria.jl" begin
    include("criteria.jl")
end

@testset "disjunction.jl" begin
    include("disjunction.jl")
end

@testset "object_oriented_api.jl" begin
    include("object_oriented_api.jl")
end

# to complete code coverage:
@testset "api.jl" begin
    include("api.jl")
end


text:
module Billard16Balls

using ModiaLang

import Modia3D
using  Modia3D.ModiaInterface

vmatBalls   = Modia3D.VisualMaterial(color="AntiqueWhite", transparency=0.0)
vmatTable   = Modia3D.VisualMaterial(color="DarkGreen", transparency=0.5)
vmatCushion = Modia3D.VisualMaterial(color="Grey20", transparency=0.5)

diameter = 0.06
radius = diameter/2
distance_balls = sqrt(3)/2*diameter
dist = 0.0001
TableX = 2.24
TableY = 1.12
TableZ = 0.05
LyBox = 0.05
LzBox = 0.05
hz = 7*radius/5
hzz = hz/3

Table = Model(
    table = Object3D(parent=:world,
                     translation=[0.0, 0.0, -TableZ/2],
                     feature=Solid(shape=Box(lengthX=TableX, lengthY=TableY, lengthZ=TableZ),
                                   solidMaterial="BilliardTable",
                                   visualMaterial=vmatTable,
                                   collision=true,
                                   collisionSmoothingRadius=0.001))
)

Cushion = Model(
    cushion11 = Object3D(parent=:world,
                         translation=[TableX/2-LyBox/2, 0.0, hz/2],
                         feature=Solid(shape=Box(lengthX=LyBox, lengthY=TableY, lengthZ=hz),
                                       solidMaterial="BilliardCushion",
                                       visualMaterial=vmatCushion)),
    cushion12 = Object3D(parent=:cushion11,
                         translation=[-LyBox/4, 0.0, hz/2+hzz/2],
                         feature=Solid(shape=Box(lengthX=1.5*LyBox, lengthY=TableY, lengthZ=hzz),
                                       solidMaterial="BilliardCushion",
                                       visualMaterial=vmatCushion,
                                       collision=true)),
    cushion31 = Object3D(parent=:world,
                         translation=[-TableX/2+LyBox/2, 0.0, hz/2],
                         feature=Solid(shape=Box(lengthX=LyBox, lengthY=TableY, lengthZ=hz),
                                       solidMaterial="BilliardCushion",
                                       visualMaterial=vmatCushion)),
    cushion32 = Object3D(parent=:cushion31,
                         translation=[LyBox/4, 0.0, hz/2+hzz/2],
                         feature=Solid(shape=Box(lengthX=1.5*LyBox, lengthY=TableY, lengthZ=hzz),
                                       solidMaterial="BilliardCushion",
                                       visualMaterial=vmatCushion,
                                       collision=true)),
    cushion21 = Object3D(parent=:world,
                         translation=[0.0, TableY/2-LyBox/2, hz/2],
                         feature=Solid(shape=Box(lengthX=TableX, lengthY=LyBox, lengthZ=hz),
                                       solidMaterial="BilliardCushion",
                                       visualMaterial=vmatCushion)),
    cushion22 = Object3D(parent=:cushion21,
                         translation=[0.0, -LyBox/4, hz/2+hzz/2],
                         feature=Solid(shape=Box(lengthX=TableX, lengthY=1.5*LyBox, lengthZ=hzz),
                                       solidMaterial="BilliardCushion",
                                       visualMaterial=vmatCushion,
                                       collision=true)),
    cushion41 = Object3D(parent=:world,
                         translation=[0.0, -TableY/2+LyBox/2, hz/2],
                         feature=Solid(shape=Box(lengthX=TableX, lengthY=LyBox, lengthZ=hz),
                                       solidMaterial="BilliardCushion",
                                       visualMaterial=vmatCushion)),
    cushion42 = Object3D(parent=:cushion41,
                         translation=[0.0, LyBox/4, hz/2+hzz/2],
                         feature=Solid(shape=Box(lengthX=TableX, lengthY=1.5*LyBox, lengthZ=hzz),
                                       solidMaterial="BilliardCushion",
                                       visualMaterial=vmatCushion,
                                       collision=true))
)

Ball = Model(
    ball = Object3D(feature=Solid(shape=Sphere(diameter=diameter),
                                  solidMaterial="BilliardBall",
                                  visualMaterial=vmatBalls,
                                  collision=true))
)

Billard = Model(
    world = Object3D(feature=Scene(gravityField=UniformGravityField(g=9.81, n=[0, 0, -1]),
                                   enableContactDetection=true,
                                   nominalLength=0.15*TableX,
                                   animationFile="Billard16Balls.json")),
    table = Table,
    cushion = Cushion,
    ball0 = Ball,
    joint0 = FreeMotion(obj1=:world, obj2=:(ball0.ball),
                        r=Var(init=[-0.8, -0.1, diameter/2]),
                        v=Var(init=[3.0, 0.1, 0.0]),
                        rot=Var(init=[pi/2, 0.0, 0.0])),
    ball1 = Ball,
    joint1 = FreeMotion(obj1=:world, obj2=:(ball1.ball),
                        r=Var(init=[TableX/6, 0.0, diameter/2]),
                        rot=Var(init=[pi/2, 0.0, 0.0])),
    ball2 = Ball,
    joint2 = FreeMotion(obj1=:world, obj2=:(ball2.ball),
                        r=Var(init=[TableX/6+1*distance_balls+dist, 1/2*(diameter+dist), diameter/2]),
                        rot=Var(init=[pi/2, 0.0, 0.0])),
    ball3 = Ball,
    joint3 = FreeMotion(obj1=:world, obj2=:(ball3.ball),
                        r=Var(init=[TableX/6+1*distance_balls+dist, -1/2*(diameter+dist), diameter/2]),
                        rot=Var(init=[pi/2, 0.0, 0.0])),
    ball4 = Ball,
    joint4 = FreeMotion(obj1=:world, obj2=:(ball4.ball),
                        r=Var(init=[TableX/6+2*(distance_balls+dist), (diameter+dist), diameter/2]),
                        rot=Var(init=[pi/2, 0.0, 0.0])),
    ball5 = Ball,
    joint5 = FreeMotion(obj1=:world, obj2=:(ball5.ball),
                        r=Var(init=[TableX/6+2*(distance_balls+dist), 0.0, diameter/2]),
                        rot=Var(init=[pi/2, 0.0, 0.0])),
    ball6 = Ball,
    joint6 = FreeMotion(obj1=:world, obj2=:(ball6.ball),
                        r=Var(init=[TableX/6+2*(distance_balls+dist), -(diameter+dist), diameter/2]),
                        rot=Var(init=[pi/2, 0.0, 0.0])),
    ball7 = Ball,
    joint7 = FreeMotion(obj1=:world, obj2=:(ball7.ball),
                        r=Var(init=[TableX/6+3*(distance_balls+dist), 3/2*(diameter+dist), diameter/2]),
                        rot=Var(init=[pi/2, 0.0, 0.0])),
    ball8 = Ball,
    joint8 = FreeMotion(obj1=:world, obj2=:(ball8.ball),
                        r=Var(init=[TableX/6+3*(distance_balls+dist), 1/2*(diameter+dist), diameter/2]),
                        rot=Var(init=[pi/2, 0.0, 0.0])),
    ball9 = Ball,
    joint9 = FreeMotion(obj1=:world, obj2=:(ball9.ball),
                        r=Var(init=[TableX/6+3*(distance_balls+dist), -1/2*(diameter+dist), diameter/2]),
                        rot=Var(init=[pi/2, 0.0, 0.0])),
    ball10 = Ball,
    joint10 = FreeMotion(obj1=:world, obj2=:(ball10.ball),
                         r=Var(init=[TableX/6+3*(distance_balls+dist), -3/2*(diameter+dist), diameter/2]),
                         rot=Var(init=[pi/2, 0.0, 0.0])),
    ball11 = Ball,
    joint11 = FreeMotion(obj1=:world, obj2=:(ball11.ball),
                         r=Var(init=[TableX/6+4*(distance_balls+dist), 2*(diameter+dist), diameter/2]),
                         rot=Var(init=[pi/2, 0.0, 0.0])),
    ball12 = Ball,
    joint12 = FreeMotion(obj1=:world, obj2=:(ball12.ball),
                         r=Var(init=[TableX/6+4*(distance_balls+dist), (diameter+dist), diameter/2]),
                         rot=Var(init=[pi/2, 0.0, 0.0])),
    ball13 = Ball,
    joint13 = FreeMotion(obj1=:world, obj2=:(ball13.ball),
                         r=Var(init=[TableX/6+4*(distance_balls+dist), 0.0, diameter/2]),
                         rot=Var(init=[pi/2, 0.0, 0.0])),
    ball14 = Ball,
    joint14 = FreeMotion(obj1=:world, obj2=:(ball14.ball),
                         r=Var(init=[TableX/6+4*(distance_balls+dist), -(diameter+dist), diameter/2]),
                         rot=Var(init=[pi/2, 0.0, 0.0])),
    ball15 = Ball,
    joint15 = FreeMotion(obj1=:world, obj2=:(ball15.ball),
                         r=Var(init=[TableX/6+4*(distance_balls+dist), -2*(diameter+dist), diameter/2]),
                         rot=Var(init=[pi/2, 0.0, 0.0]))
)

billard = @instantiateModel(buildModia3D(Billard), unitless=true, log=false, logStateSelection=false, logCode=false)

stopTime = 5.0
testTime = 1.5
tolerance = 1e-7
interval = 0.001
requiredFinalStates = [0.9336946243629118, -0.18365520174670474, 0.029998496991356247, -0.7011746198909901, 0.6311295647736829, 1.1287610685608916e-6, -13.537433737049488, -0.3858748206785683, -51.74321001761849, 10.647446137520763, -27.86255858188846, -9.957728295667122, 0.32509152711043904, 0.027780203284739517, 0.029997667930081795, -0.04542743305030694, 0.02500170283086513, 7.062399855287814e-8, 0.9625529100121986, -0.5890320882437217, 1.5160600432975857, -0.8636537095656692, -0.04474632782683813, 1.4954319791715307, 0.39301955928437254, 0.31437311545523294, 0.029997559985557257, -0.0295450598735267, 0.25844118748219774, 6.543232669438406e-8, -7.838021100027451, -0.22760408431777734, 4.052129880389039e-5, -8.613971169221236, -0.015367180725959132, 0.9844568913082455, 0.3526376383583441, -0.14575202949834226, 0.029998721807484585, -0.05704334595170738, -0.09087467006012945, -8.648550230091023e-8, 6.463720232117719, 0.5613406381152121, -0.6619441310051621, 3.0283371525415186, -0.011031122590753485, 1.901031014930388, 0.47669958004172114, 0.30988266338557435, 0.029997421574535003, -0.0006779043504506098, 0.22297266655607906, 2.9010253957924976e-8, -6.65506642882509, -0.004145772945978028, 0.0028226662531805786, -7.4316293203593204, -7.298765565151683e-5, 0.022599728211138887, 0.3783007120971173, -0.012766185812833935, 0.029998759358176016, -0.04494536278891231, -0.018632373536778858, -7.215437557767047e-8, 1.6943216570588908, 0.19414146703780452, 3.298010373792468, -0.3471933000846414, -0.13201026537226748, 1.577730264308742, 0.46331163648978807, -0.09389324485553788, 0.029998896770933176, -0.0016188042316731567, -0.0037739593426910373, -2.5503675848459563e-9, 2.6598529884755653, 0.23804363229547723, 0.3872083559790523, 0.1248900591032485, 0.00033968828158966596, 0.05356980525109636, 0.5432807165029493, 0.2953230757760231, 0.02999730653311251, -0.13624873106020685, -0.18205262515847373, 1.9858018516413702e-7, -5.40642875525656, -0.19653356814243264, 0.09370138552443177, 6.331724950526496, -3.5125692661867474, 2.2385800870275046, 0.5280895349790748, 0.02944513078495663, 0.029997345184240368, -3.671897262472209e-15, -3.6295581814805492e-12, -1.2427691187347466e-15, 1.5884011771883515, 0.00047788925223401856, 0.045046404635795685, 8.53157003717906e-11, -3.8442455469304e-12, 1.2714782913701774e-13, 0.5298393427066304, -0.03176029891751437, 0.029999002794942116, 1.297965371269723e-12, -7.63423453026827e-13, 3.722976278613608e-18, 1.6241627081677217, -0.00011555321024985877, -0.0050541325338022604, 1.7967580704217367e-11, -1.5382972866683138e-12, -3.0499570729309535e-11, 0.5382383412758586, -0.09514048819597622, 0.02999893580898303, 1.3155505360211643e-12, -8.821797446529922e-13, -7.043094683356967e-16, 1.7309170721863256, -0.022807794119219644, -0.28238846622282227, 2.063596477748024e-11, 8.485564273716919e-13, -3.1023757752800775e-11, 0.592498670205686, 0.3681019267809226, 0.029997236420195104, 0.15727976960938742, -0.019558871461817822, -2.644518715680437e-7, -6.163065667728738, 0.37981142806747936, -0.058040028743833724, 0.5349948316934143, 5.243907614414163, -0.34169516938674743, 0.6344294601093406, 0.0669304999587428, 0.029997172347095737, -0.7614223443828724, 0.006309955741959036, 1.2183649383003493e-6, 1.558857816712138, 0.0042029171671443275, -0.7644036324349628, -0.0191002690729473, -0.4380898296432406, 25.377516875983996, 0.7646288540762318, -0.10696914113095653, 0.029998770268321002, 0.15911209364889162, -0.09298081830712687, -2.5596316103907114e-7, 1.9073899574938322, -0.10513813819059165, -6.883833037985919, 3.0989440441527623, 3.5085557038592837e-6, -5.30302146889589, 0.7873220523004107, -0.17884965224773064, 0.029998733802261777, 0.18239453271826178, -0.10527370713656753, -2.9340175118245415e-7, 2.1544182053207317, -0.43508216020191737, -7.5519516937515325, 3.508723323843255, -9.174029482795626e-7, -6.0791204764983995, 0.7976238621443168, -0.24490123678673625, 0.02999869521225786, 0.19303854447667773, -0.11142219870901685, -1.791306576517481e-7, 2.1859628103222946, -0.6263874799512511, -7.873271866785609, 3.7136784821321944, -4.152815543679789e-7, -6.433930426076936]
simulate!(billard, stopTime=testTime, tolerance=tolerance, interval=interval, log=true, logStates=true, logEvents=true, requiredFinalStates=requiredFinalStates)

@usingModiaPlot
plot(billard, ["joint0.r" "joint0.rot"; "joint0.v" "joint0.w"], figure=1)

end


text:
using Crayons

include("shape.jl")
include("display_obj.jl")
include("display_stack.jl")


text:
module PathDirections
export PathDirection
@enum PathDirection::Bool begin
    up # going towards the root of the tree
    down # going away from the root of the tree
end
end
using .PathDirections

struct TreePath{V, E}
    source::V
    target::V
    edges::Vector{E} # in order
    directions::Vector{PathDirection}
    indexmap::Vector{Int}
end

source(path::TreePath) = path.source
target(path::TreePath) = path.target

@inline Base.findfirst(path::TreePath{<:Any, E}, edge::E) where {E} = path.indexmap[edge_index(edge)] # TODO: remove, use equalto
@inline Base.in(edge::E, path::TreePath{<:Any, E}) where {E} = findfirst(path, edge) != 0 # TODO: findfirst update
@inline directions(path::TreePath) = path.directions
@inline direction(edge::E, path::TreePath{<:Any, E}) where {E} = path.directions[findfirst(path, edge)] # TODO: findfirst update

Base.iterate(path::TreePath) = iterate(path.edges)
Base.iterate(path::TreePath, state) = iterate(path.edges, state)
Base.eltype(path::TreePath) = eltype(path.edges)
Base.length(path::TreePath) = length(path.edges)

function Base.show(io::IO, path::TreePath)
    println(io, "Path from $(path.source) to $(path.target):")
    for edge in path
        directionchar = ifelse(direction(edge, path) == PathDirections.up, '↑', '↓')
        print(io, "$directionchar ")
        show(IOContext(io, :compact => true), edge)
        println(io)
    end
end

function TreePath(src::V, target::V, tree::SpanningTree{V, E}) where {V, E}
    source_to_lca = E[]
    lca_to_target = E[]
    source_current = src
    target_current = target
    while source_current != target_current
        if tree_index(source_current, tree) > tree_index(target_current, tree)
            edge = edge_to_parent(source_current, tree)
            push!(source_to_lca, edge)
            source_current = source(edge, tree)
        else
            edge = edge_to_parent(target_current, tree)
            push!(lca_to_target, edge)
            target_current = source(edge, tree)
        end
    end
    reverse!(lca_to_target)

    edges = collect(E, flatten((source_to_lca, lca_to_target)))
    directions = collect(PathDirection, flatten(((PathDirections.up for e in source_to_lca), (PathDirections.down for e in lca_to_target))))
    indexmap = Vector(sparsevec(Dict(edge_index(e) => i for (i, e) in enumerate(edges)), num_edges(tree)))
    TreePath{V, E}(src, target, edges, directions, indexmap)
end


text:
const BaseType = Union{Float32, Float64,
                       Complex{Float32}, Complex{Float64},
                       Quaternion{Float32}, Quaternion{Float64}}


text:
using Random, Distributions, Roots

"""
    simu_u_ped(nid, ng)
---

# Description
Simulate a simple population of 0:`ng` generations:

- Each generation has `nid` breeders
- Half of them are males.
- Random mating
- Each pair has two offspring.

# Parameters
- `nid`: even (or, +1), positive (or, error)
- `ng`: positive integer

# Results
Return a two column matrix, which are `pa` and `ma`.
Row number = ID number.
Unknown parent = 0.

"""
function simu_u_ped(nid::Int, ng::Int)
    nid < 1 && @error "nid (= $nid) should > 0"
    ng  < 1 && @error "ng (= $ng) should > 0"
    nid % 2 == 1 && begin
        nid += 1
        @warn "nid increased by 1, = $nid now"
    end

    # generation 0
    ped = zeros(Int, 2nid)

    # generation 1:ng
    @inbounds for g in 1:ng
        fra = (g - 1) * nid + 1
        id = shuffle(fra : fra + nid - 1)
        for i in 1:2:nid
            append!(ped, [id[i], id[i+1], id[i], id[i+1]])
        end
    end
    reshape(ped, 2, :)'
end

"""
    simu_p_bvy(ped, h2; μ = .0, Vp = 1.0)
---
Simulate breeding values (`bv`) and phenotypes (`y`) with polygenic method.
Given a pedigree `ped`, heritability `h2`, this function return two vectors, breeding values and phenotypes.

The population mean μ is 0.0, variance of phenotypes is 1.0 by default.
"""
function simu_p_bvy(ped, h2; μ = .0, Vp = 1.0)
    ! (0.0 ≤ h2 ≤ 1.0) && throw(DomainError("h2 should be in [0, 1]"))
    Vp ≤ .0 && throw(DomainError("vp should be in (0, ∞)"))

    Va = Vp * h2
    RndAdd = Normal(.0, sqrt(Va/2.)) # when one parent is unknown
    Mendel = Normal(.0, sqrt(Va)/2.) # when one parent is known, add some Mendelian sampling error
    sde = sqrt(Vp * (1. - h2))

    nid = size(ped)[1]
    y   = rand(Normal(.0, sde), nid) # residuals first for phenotypes
    bv  = zeros(nid)
    @inbounds for i in 1:nid
        pa, ma = ped[i, :]
        bv[i] += pa > 0 ? bv[pa] / 2. + rand(Mendel) : rand(RndAdd)
        bv[i] += ma > 0 ? bv[ma] / 2. + rand(Mendel) : rand(RndAdd)
    end
    y = y + bv .+ μ
    return bv, y
end

"""
    simu_q_founder(nID, h2; μ = .0, Vp = 1.0, nQTL = 1000, shape = 0.25, MAF = 0.01)
---
Simulate a founder population of size = `nid`.  This population will have
- Va = h2
= Vp = 1.0

unless specified otherwize.  This also assume the QTL
- are independent
- frequencies follow a β(0.25) distribution
"""
function simu_q_founder(nID, h2; μ = .0, Vp = 1.0, nQTL = 1000, Shape = 0.25, MAF = 0.01)
    # sample frequencies for QTL
    freq = Float64[]            # pop frequency of (MAF, 1-MAF), inclusive
    while length(freq) < nQTL   # sizeof is equivalent to capacity of C++
        f = rand(Beta(Shape))
        MAF ≤ f ≤ 1-MAF && push!(freq, f)
    end

    allele = begin
        tmp = Bool[]
        for f in freq
            append!(tmp, rand(Bernoulli(f), 2nID))
        end
        reshape(tmp, 2nID, :)
    end

    eQTL = randn(nQTL)
    M = begin
        tmp = Float64[]
        for i in 1:2:2nID
            append!(tmp, convert.(Float64, allele[i, :] + allele[i+1, :]))
        end
        reshape(tmp, nQTL, :)
    end

    for _ in 1:3                # repeat 3 times to reach μ=0., var=h2*Vp
        fv(f) = begin
            tmp = eQTL .* f
            var(M'tmp) - h2*Vp
        end
        eQTL .*= fzero(fv, 1.0)

        fm(m) = begin
            tmp = eQTL .- m
            mean(M'tmp)
        end
        eQTL .-= fzero(fm, 0.)
    end
    allele, eQTL
end

"""
    simu_q_bvy(ped, sde, allele, eQTL)
---
1. Drop `allele` throw founders into `ped`.
2. Calculate `bv`
3. Add noise to `bv` → `y`
4. Return `bv`, `y`
"""
function simu_q_bvy(ped, sde, allele, eQTL)
    nFdr = size(allele)[1] ÷ 2
    nQTL = length(eQTL)
    nID  = size(ped)[1]
    y = randn(nID) .* sde

    for id in nFdr+1:nID        # check pedigree
        if ped[id, 1] == 0 || ped[id, 2] == 0
            throw(DomainError("ID $id's parent(s) unknown"))
        elseif ped[id, 1] >= id || ped[id, 2] >= id
            throw(DomainError("ID $id's parents not before this ID"))
        end
    end

    mat = zeros(Bool, 2nID, nQTL)
    mat[1:2nFdr, :] = allele

    bv = begin
        # Drop alleles
        @inbounds for id in nFdr+1:nID
            pa, ma = ped[id, :]
            ix = CartesianIndex.(2pa .- rand(0:1, nQTL), 1:nQTL)
            mat[2id-1, :] = mat[ix]
            ix = CartesianIndex.(2ma .- rand(0:1, nQTL), 1:nQTL)
            mat[2id, :] = mat[ix]
        end
        M = begin
            tmp = Float64[]
            for i in 1:2:2nID
                append!(tmp, convert.(Float64, mat[i, :] + mat[i+1, :]))
            end
            reshape(tmp, nQTL, :)
        end
        M'eQTL
    end
    y += bv
    bv, y
end


text:
abstract type AbstractSelector; end

struct ColumnSelector{T} <: AbstractSelector
    table::T
    splitby
    compare
    across
    x
    y
    kw::Dict{Symbol, Any}
end

ColumnSelector(df) = ColumnSelector(df, Symbol[], nothing, nothing, nothing, nothing, Dict{Symbol, Any}(:compare => false))

struct Selector{T, F1<:Function, F2<:Function, F3<:Function, F4<:Function, F5<:Function} <: AbstractSelector
    table::T
    splitby::F1
    compare::F2
    across::F3
    x::F4
    y::F5
    kw::Dict{Symbol, Any}
end

Selector(df) = Selector(df, t -> ("y1",), t -> (),t -> 0.0, t -> 0.0, t -> NaN, Dict{Symbol, Any}(:compare => false))

Base.convert(::Type{Selector}, a::Selector) = a
Base.convert(::Type{Selector}, a) = Selector(a)

struct Table2Process{T}
    table::T
    kw::Dict{Symbol, Any}
end

Table2Process(df) = Table2Process(df, Dict{Symbol, Any}())

tuplify(x::Tuple, args...) = (x..., args...)
tuplify(x, args...) = (x, args...)

_get_tup(t) = map(_get, t)

_get(t) = t
_get(t::DataValue) = get(t)

_nafree_tup(t::Tuple) = all(_nafree, t)

_nafree(t) = true
_nafree(t::DataValue) = !isna(t)

function Table2Process(s::Selector)
    enumerable = TableTraits.getiterator(s.table)
    T = eltype(enumerable)
    if s.kw[:compare]
        select_func = t -> tuplify(s.splitby(t), s.compare(t), s.across(t), s.x(t), s.y(t))
    else
        select_func = t -> tuplify(s.splitby(t), s.across(t), s.x(t), s.y(t))
    end
    tuple_vec = map(_get_tup, Iterators.filter(_nafree_tup, (select_func(i) for i in enumerable)))
    Table2Process(columns(convert(Columns, tuple_vec)), s.kw)
end

getcolumn(t::IndexedTables.AbstractIndexedTable, s) = columns(t, s)
getcolumn(t, s) = getindex(t, s)

getlength(t::IndexedTables.AbstractIndexedTable) = length(t)
getlength(t) = size(t, 1)

function Table2Process(s::ColumnSelector)
    if s.splitby == Symbol[]
        splitter = [fill("y1", getlength(s.table))]
    elseif isa(s.splitby, Symbol)
        splitter = [getcolumn(s.table, s.splitby)]
    else
        splitter = map(t -> getcolumn(s.table, t), s.splitby)
    end
    across_col = s.across == nothing ? fill(0.0, getlength(s.table)) : getcolumn(s.table, s.across)
    y_col = s.y == nothing ? fill(NaN, getlength(s.table)) : getcolumn(s.table, s.y)
    x = isa(s.x, Symbol) ? getcolumn(s.table, s.x) : fill(s.x, getlength(s.table))
    if s.compare == nothing
        columns = tuple(splitter..., across_col, x, y_col)
    else
        columns = tuple(splitter..., getcolumn(s.table, s.compare), across_col, x, y_col)
    end
    Table2Process(columns, s.kw)
end

struct ProcessedTable{T}
    table::T
    kw::Dict{Symbol, Any}
end

ProcessedTable(t::Table2Process) = pipeline(t)
ProcessedTable(s::AbstractSelector) = ProcessedTable(Table2Process(s))

function Base.filter(f::Function, p::ProcessedTable)
    filtered_table = filter(f, p.table)
    ProcessedTable(filtered_table, p.kw)
end

nsplits(t::Union{Table2Process, ProcessedTable}) = count(t -> startswith(t, "s"), string.(colnames(t.table)))
listsplits(t::Union{Table2Process, ProcessedTable}) = [Symbol(:s, i) for i = 1:nsplits(t)]


text:
"""
    StringCollision

    string collision between contact points

    origin_parent: position of contact on parent body relative to center of mass
    origin_child: position of contact on parent body relative to center of mass
    length: maximum distance between contact point
"""
mutable struct StringCollision{T,O,I,OI} <: Collision{T,O,I,OI}
    origin_parent::SVector{I,T}
    origin_child::SVector{I,T}
    length::T
end

# function Base.show(io::IO, mime::MIME{Symbol("text/plain")}, collision::StringCollision)
#     summary(io, collision)
#     println(io, "")
#     println(io, "origin_parent: "*string(collision.origin_parent))
#     println(io, "origin_child:  "*string(collision.origin_child))
#     println(io, "length:  "*string(collision.length))
# end

# distance
function distance(collision::StringCollision, xp, qp, xc, qc)
    # contact origin points
    cop = contact_point_origin(xp, qp, collision.origin_parent)
    coc = contact_point_origin(xc, qc, collision.origin_child)

    # distance between contact origins
    d = norm(cop - coc, 2)

    # minimum distance between spheres
    return collision.length - d
end

function ∂distance∂x(gradient::Symbol, collision::StringCollision, xp, qp, xc, qc)
    # contact origin points
    cop = contact_point_origin(xp, qp, collision.origin_parent)
    coc = contact_point_origin(xc, qc, collision.origin_child)

    # distance between contact origins
    # d = norm(cop - coc, 2)
    ∂norm∂d = ∂norm∂x(cop - coc)

    if gradient == :parent
        D = 1.0 * ∂norm∂d * 1.0 * ∂contact_point_origin∂x(xp, qp, collision.origin_parent)
    elseif gradient == :child
        D = 1.0 * ∂norm∂d * -1.0 * ∂contact_point_origin∂x(xc, qc, collision.origin_child)
    end

    # if gradient == :parent
    #     FD = FiniteDiff.finite_difference_jacobian(x -> distance(collision, x, qp, xc, qc), xp)
    # elseif gradient == :child
    #     FD = FiniteDiff.finite_difference_jacobian(x -> distance(collision, xp, qp, x, qc), xc)
    # end

    # @assert norm(-D - FD, Inf) < 1.0e-5

    return -D
end

function ∂distance∂q(gradient::Symbol, collision::StringCollision, xp, qp, xc, qc)
    # contact origin points
    cop = contact_point_origin(xp, qp, collision.origin_parent)
    coc = contact_point_origin(xc, qc, collision.origin_child)

    # distance between contact origins
    # d = norm(cop - coc, 2)
    ∂norm∂d = ∂norm∂x(cop - coc)

    if gradient == :parent
        D = 1.0 * ∂norm∂d *  1.0 * ∂contact_point_origin∂q(xp, qp, collision.origin_parent)
    elseif gradient == :child
        D = 1.0 * ∂norm∂d * -1.0 * ∂contact_point_origin∂q(xc, qc, collision.origin_child)
    end

    # if gradient == :parent
    #     FD = FiniteDiff.finite_difference_jacobian(q -> distance(collision, xp, Quaternion(q..., false), xc, qc), vector(qp))
    # elseif gradient == :child
    #     FD = FiniteDiff.finite_difference_jacobian(q -> distance(collision, xp, qp, xc, Quaternion(q..., false)), vector(qc))
    # end

    # @assert norm(-D - FD, Inf) < 1.0e-5

    return -D
end

# contact point in world frame
function contact_point(relative::Symbol, collision::StringCollision, xp, qp, xc, qc)
    # contact origin points
    cop = contact_point_origin(xp, qp, collision.origin_parent)
    coc = contact_point_origin(xc, qc, collision.origin_child)

    # direction of minimum distance (child to parent)
    # d = cop - coc
    # dir = normalize(d)

    # contact point
    if relative == :parent
        return cop
    elseif relative == :child
        return coc
    end
end

function ∂contact_point∂x(relative::Symbol, jacobian::Symbol, collision::StringCollision, xp, qp, xc, qc)
    # contact origin points
    cop = contact_point_origin(xp, qp, collision.origin_parent)
    coc = contact_point_origin(xc, qc, collision.origin_child)

    # direction of minimum distance (child to parent)
    # d = cop - coc
    # dir = normalize(d)

    if relative == :parent
        # cop
        if jacobian == :parent
            ∂c∂x = ∂contact_point_origin∂x(xp, qp, collision.origin_parent)
            X = ∂c∂x
        elseif jacobian == :child
            X = szeros(eltype(xc), 3, 3)
        end
    elseif relative == :child
        # coc
        if jacobian == :parent
            X = szeros(eltype(xp), 3, 3)
        elseif jacobian == :child
            ∂c∂x = ∂contact_point_origin∂x(xc, qc, collision.origin_child)
            X = ∂c∂x
        end
    end

    # if jacobian == :parent
    #     FD =  FiniteDiff.finite_difference_jacobian(x -> contact_point(relative, collision, x, qp, xc, qc), xp)
    # elseif jacobian == :child
    #     FD = FiniteDiff.finite_difference_jacobian(x -> contact_point(relative, collision, xp, qp, x, qc), xc)
    # end

    # @assert norm(X - FD, Inf) < 1.0e-5

    return X
end

function ∂contact_point∂q(relative::Symbol, jacobian::Symbol, collision::StringCollision, xp, qp, xc, qc)
    # contact origin points
    cop = contact_point_origin(xp, qp, collision.origin_parent)
    coc = contact_point_origin(xc, qc, collision.origin_child)

    # direction of minimum distance (child to parent)
    # d = cop - coc
    # dir = normalize(d)

    if relative == :parent
        # cop
        if jacobian == :parent
            ∂c∂q = ∂contact_point_origin∂q(xp, qp, collision.origin_parent)
            Q = ∂c∂q
        elseif jacobian == :child
            Q = szeros(eltype(xc), 3, 4)
        end
    elseif relative == :child
        # coc
        if jacobian == :parent
            Q = szeros(eltype(xp), 3, 4)
        elseif jacobian == :child
            ∂c∂q = ∂contact_point_origin∂q(xc, qc, collision.origin_child)
            Q = ∂c∂q
        end
    end

    # if jacobian == :parent
    #     FD = FiniteDiff.finite_difference_jacobian(q -> contact_point(relative, collision, xp, Quaternion(q..., false), xc, qc), vector(qp))
    # elseif jacobian == :child
    #     FD = FiniteDiff.finite_difference_jacobian(q -> contact_point(relative, collision, xp, qp, xc, Quaternion(q..., false)), vector(qc))
    # end

    # @assert norm(Q - FD, Inf) < 1.0e-5

    return Q
end

function contact_normal(collision::StringCollision, xp, qp, xc, qc)
    # contact points
    cop = contact_point(:parent, collision, xp, qp, xc, qc)
    coc = contact_point(:child,  collision, xp, qp, xc, qc)

    # unnormalized direction
    dir = cop - coc

    # distance
    dis = distance(collision, xp, qp, xc, qc)

    # normalized direction
    # if dis >= 0.0
    return -normalize(dir)'
    # else
    #     return normalize(dir)'
    # end
end

function ∂contact_normal_transpose∂x(jacobian::Symbol, collision::StringCollision, xp, qp, xc, qc)
    # contact origin points
    cop = contact_point(:parent, collision, xp, qp, xc, qc)
    coc = contact_point(:child,  collision, xp, qp, xc, qc)

    # unnormalized direction
    dir = cop - coc

    # Jacobians
    X = ∂normalize∂x(dir) * (∂contact_point∂x(:parent, jacobian, collision, xp, qp, xc, qc) - ∂contact_point∂x(:child, jacobian, collision, xp, qp, xc, qc))

    # distance
    # dis = distance(collision, xp, qp, xc, qc)

    # normalized direction
    # if dis >= 0.0
    return -X
    # else
    #     return 1.0 * X
    # end
end

function ∂contact_normal_transpose∂q(jacobian::Symbol, collision::StringCollision, xp, qp, xc, qc)
    # contact origin points
    cop = contact_point(:parent, collision, xp, qp, xc, qc)
    coc = contact_point(:child,  collision, xp, qp, xc, qc)

    # unnormalized direction
    dir = cop - coc

    Q = ∂normalize∂x(dir) * (∂contact_point∂q(:parent, jacobian, collision, xp, qp, xc, qc) - ∂contact_point∂q(:child, jacobian, collision, xp, qp, xc, qc))

    # Jacobians
    # if jacobian == :parent
    #     FD = FiniteDiff.finite_difference_jacobian(q -> contact_normal(collision, xp, Quaternion(q..., false), xc, qc)', vector(qp))
    # elseif jacobian == :child
    #     FD = FiniteDiff.finite_difference_jacobian(q -> contact_normal(collision, xp, qp, xc, Quaternion(q..., false))', vector(qc))
    # end

    # distance
    dis = distance(collision, xp, qp, xc, qc)

    # @assert norm((dis >= 0.0 ? 1.0 : -1.0) * Q - FD, Inf) < 1.0e-2

    # normalized direction
    # if dis >= 0.0
    return -Q
    # else
    #     return Q
    # end
end


text:
# ---
# title: 645. Set Mismatch
# id: problem645
# author: Indigo
# date: 2021-07-04
# difficulty: Easy
# categories: Hash Table, Math
# link: <https://leetcode.com/problems/set-mismatch/description/>
# hidden: true
# ---
#
# The set `S` originally contains numbers from 1 to `n`. But unfortunately, due
# to the data error, one of the numbers in the set got duplicated to **another**
# number in the set, which results in repetition of one number and loss of
# another number.
#
# Given an array `nums` representing the data status of this set after the
# error. Your task is to firstly find the number occurs twice and then find the
# number that is missing. Return them in the form of an array.
#
# **Example 1:**
#
#
#
#     Input: nums = [1,2,2,4]
#     Output: [2,3]
#
#
# **Note:**
#
#   1. The given array size will in the range [2, 10000].
#   2. The given array's numbers won't have any order.
#
#
## @lc code=start
using LeetCode

function find_error_nums(nums::Vector{Int})
    cnt = fill(0x0, length(nums))
    for num in nums
        cnt[num] += 1
    end
    return findfirst(==(0x2), cnt), findfirst(==(0x0), cnt)
end
## @lc code=end


text:
### A Pluto.jl notebook ###
# v0.15.1

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : missing
        el
    end
end

# ╔═╡ 0bc12fe2-4d08-4840-b651-b58df23f0277
begin
	using Pkg;
	Pkg.activate(".")
	using IntervalLinearAlgebra, LazySets, Plots, Images, StaticArrays, PlutoUI
end

# ╔═╡ f403b050-e2e6-11eb-1c44-911de0d748bd
html"<button onclick='present()'>present</button>"

# ╔═╡ c9e05bd9-1d7b-41c0-99dc-29b1c1e20090
html"""

IntervalLinearAlgebra.jl

Linear algebra done rigorously

Luca Ferranti, David Sanders, Marcelo Forets

Google Summer of Code (GSoC) 2021

Presentation for JuliaCon 2021

"""

# ╔═╡ c3a5fd72-9a9d-45f3-aaf3-bd61637d286c
md"""
## Interval matrices

An interval matrix $\mathbf{A}\in\mathbb{I}\mathbb{R}^{m\times n}$ is defined as

$\mathbf{A} = \{A \in \mathbb{R}^{m\times n} | A_{ij}\in\mathbf{A}_{ij}\quad i=1,\ldots,m\quad j=1\ldots,n\}$

for example
"""

# ╔═╡ f2cff36e-9136-4eb1-acff-c97af31c2450
AA = [1..2 3..4; 5..6 7..8]

# ╔═╡ 286f40a2-840d-431b-a531-fe9e09dc55f1
md"""
##

 $\mathbf{A}$ can also be represented in *midpoint-radius* notation $A_c\pm A_\Delta$, where $A_c$ is the *midpoint matrix* and $A_\Delta$ is the *radius matrix*.
"""

# ╔═╡ 97e498f5-2e0c-4139-9774-86f58d59d511
AA

# ╔═╡ bb878728-d4ca-4c01-a4f7-bf1811a8f5c2
mid.(AA)

# ╔═╡ a89ee6e9-565a-46d5-82db-bafb251a2200
radius.(AA)

# ╔═╡ 4848561e-2336-40fb-a5c6-8705da85c2ee
md"""
##

### Regular matrices
We say that an interval matrix $\mathbf{A}$ is **regular** if all $A\in\mathbf{A}$ are invertible. Otherwise, the interval matrix is **singular**.

- In general, checking for regularity or singularity is computationally expensive.
"""

# ╔═╡ 5b111f71-4d41-43cf-a3cf-e2a6f4c2adba
html"

    real
    interval
	check property

    invertible
    regular
	coNP-complete

    singular
    singular
	NP-complete

"

# ╔═╡ 395b75e5-e76f-4f3a-b51f-57eaa3de2d4e
md"""
## Interval linear systems

We are now ready to study the square interval linear system $\mathbf{Ax}=\mathbf{b}$ with $\mathbf{A}\in\mathbb{IR}^{n\times n}$ and $\mathbf{b}\in\mathbb{IR}^n$.
"""

# ╔═╡ f39f94f7-59ff-4dd0-93dd-4f9f158c673c
md"""
The solution set $\mathbf{x}$ is defined as

$\mathbf{x}=\{x\in\mathbb{R}^n | Ax=b\text{ for some } A\in\mathbf{A}, b\in\mathbf{b}\}$
"""

# ╔═╡ 3f5a6fed-6175-4459-8be2-9a330cedc938
md"""
- If $\mathbf{A}$ is regular, then the solution set is non-empty and bounded.
"""

# ╔═╡ c223f95a-9885-48ae-b939-4f3fa54a7987
md"""
##
### Finding the solution set

- How do we solve the square interval linear system $\mathbf{A}\mathbf{x}=\mathbf{b}$?
- A naive approach might be to use Monte Carlo, i.e. generate several random real instances from the interval system.
- We will use the following as running example
"""

# ╔═╡ 32f0ff22-d42f-410f-a1a6-b8ea77b696ff
A = [2..4 -2..1;-1..2 2..4]

# ╔═╡ a79e27c9-399c-44e1-ac79-f65480fe5178
b = [-2..2, -2..2]

# ╔═╡ ba3a1ca6-04dc-45db-b686-a036519f2bc9
md"""
##
"""

# ╔═╡ 38883c77-d506-4e81-9d25-0ac78323241a
begin
	Ns = 100_000
	Astatic = @SMatrix [2..4 -2..1;-1..2 2..4]
	bstatic = @SVector [-2..2, -2..2]
	xs = [rand.(A)\rand.(b) for _ in 1:Ns]
	x = [xs[i][1] for i in 1:Ns]
	y = [xs[i][2] for i in 1:Ns]
	histogram2d(x, y, ratio=1)
end

# ╔═╡ 4549ca50-615c-4fba-afde-303ad231f648
md"""
- Pictures obtained by uniformly samplying 100_000 times each interval.
- **Did we cover the whole set?**
"""

# ╔═╡ 0f196850-f4b1-4830-9198-d72a281da05c
md"""
## Solution set characterization

An important theorem, known as **Oettli-Präger theorem**, helps us describing the solution set $\mathbf{x}$.

###### Oettli-Präger theorem
Given the square interval linear system $\mathbf{Ax}=\mathbf{b}$, then

$
x\in\mathbf{x}\Leftrightarrow|A_cx-b_c|\le A_\Delta|x|+b_\Delta$

- The absolute value is taken elementwise.

- We can remove the absolute values by considering one orthant at the time, obtaining $2^n$ systems of linear inequalities.

- Alternatively, the non-linear inequalities can be solved directly using `IntervalConstraintProgramming.jl`

- Both impractical for higher dimensions!
"""

# ╔═╡ d5e3856d-07e7-4e50-9e4d-c0afb7091d8b
md"""
##
"""

# ╔═╡ 95edbbbf-2e35-446d-9eeb-a53f5c6b10ec
polytopes = solve(A, b, LinearOettliPrager())

# ╔═╡ 4880c8a6-de84-4e55-be90-f73256b0888a
begin
	plot(polytopes, ratio=1)
	histogram2d!(x, y)
end

# ╔═╡ 75a66d2a-01b9-4d6a-8fed-8d8159ba6354
md"""
##
### A 3D example
Below the solution of the interval system with

$
\begin{bmatrix}
[4.5, 4.5]&[0, 2]&[0, 2]\\
[0, 2]&[4.5, 4.5]&[0, 2]\\
[0, 2]&[0, 2]& [4.5, 4.5]
\end{bmatrix}\mathbf{x}=\begin{bmatrix}[-1, 1]\\
[-1, 1]\\
[-1, 1]\end{bmatrix}$

"""

# ╔═╡ cd558a80-2825-4953-870c-93551f60c10e
load("./3dexample.png")

# ╔═╡ c5aa282c-c82b-4521-994d-235cb7934c83
md"""
##
### Exact solution, conclusions
* In general, the solution set of an interval linear system is a non-convex polytope.
  - however it is convex in each orthant.
* The computational complexity to find the solution set grows exponentially with the dimension,
* An alternative more feasible approach is to find a *tight enclosure* of the solution.
"""

# ╔═╡ 32481e35-1a98-4b65-b925-264d51d24743
md"""
## Enclosures of interval linear systems

- We say that an interval box $\mathbf{\Sigma}$ is an enclosure of the set $\mathbf{x}$ if $\mathbf{x}\subseteq\mathbf{\Sigma}$

- Ideally, we want the *hull* of the solution $\mathbf{\Sigma}_{H}$, that is the tightest interval box.

- However, finding the exact hull is in general NP-hard
"""

# ╔═╡ 58b8b1b8-93e7-46c0-b0a0-28532b152908
begin
	plot(interval_hull(ConvexHullArray(polytopes)), label="hull", α=0.3)
	plot!(polytopes, ratio=1, α=1)
end

# ╔═╡ faa03639-a431-4cbf-ac03-bbb33fb7b5b8
md"""
## Algorithms to find the enclosure of the system

`InteralLinearAlgebra.jl` has several algorithms to compute an enclosure of $\mathbf{Ax}=\mathbf{b}$ and an user friendly interface to choose what algorithm and precondition mechanisms to use.

#### Implemented algorithms
- Gaussian elimination
- Gauss-Seidel
- Jacobi
- Hansen-Bliek-Rohn
- Krawczyk
"""

# ╔═╡ 5e0cf15b-21ab-45bb-81fa-ee962a46004f
Xge = solve(A, b, GaussianElimination(), NoPrecondition())

# ╔═╡ fc08d410-02fa-480b-a999-841476dfaaf6
md"""
##
"""

# ╔═╡ 57a5fb2a-5cb2-412f-a1a0-5d27e0a5b8f7
begin
	plot(interval_hull(ConvexHullArray(polytopes)), label="hull", α=0.3)
	plot!(IntervalBox(Xge), label="Gaussian elimination", α=0.2, legend=:right)
	plot!(polytopes, ratio=1, α=1)
end

# ╔═╡ a70a6137-88cf-4fd5-ad50-8d76a1379a27
md"""
- For some special cases, the algorithm will return the hull, but in general some overestimation will occur.
"""

# ╔═╡ 0308378a-2d36-4c47-add7-0f53dc9011ca
md"""
## Preconditioning

- For the previous algorithms to work and be numerically stable, there are some requirements on the matrix $\mathbf{A}$.
- If those requirements are not met, one can try to precondition the problem with a *real* matrix $C$. that is apply the algorithms to the linear system

$
C\mathbf{Ax}=C\mathbf{b}$
- A particularly good choice is $C\approx A_c^{-1}$
"""

# ╔═╡ 9dba04bd-a8c3-44da-a1ab-855829c824ae
md"""
##
To understand the need for preconditioning, let us consider the following example.
"""

# ╔═╡ 8f68a5f0-4ec9-40d5-89be-cddf0c34a977
@bind N Slider(2:100; show_value=true, default=5)

# ╔═╡ 22abccab-5c93-4aa4-90e4-8209997a9a86
A1 = tril(fill(1..1, N, N))

# ╔═╡ e77a478d-0270-4120-a2d4-36efc455fb2e
b1 = [-2..2, fill(0..0, N-1)...]

# ╔═╡ 1ec87e5e-2f90-490b-a94f-05ea53d4cb1d
md"""
The correct solution is $[[-2, 2], [-2, 2], [0, 0], [0, 0], \ldots]$
"""

# ╔═╡ de8734a4-46b8-45ba-aff1-fa6a73e85d96
md"""
##
"""

# ╔═╡ e3bee176-9b9a-448f-b2d0-fa53cad35555
solve(A1, b1, GaussianElimination(), NoPrecondition())

# ╔═╡ 495e13e3-bf20-4541-ab6c-f6cb9572e9ed
solve(A1, b1, HansenBliekRohn(), NoPrecondition())

# ╔═╡ f223ac43-79cd-43be-b585-907a0bfb8f27
md"""
##
"""

# ╔═╡ acb54294-917d-41b0-b501-364b9228531b
solve(A1, b1, GaussianElimination(), InverseMidpoint())

# ╔═╡ ea231b9c-4cc9-4d45-b171-4c66536b02a1
solve(A1, b1, HansenBliekRohn(), InverseMidpoint())

# ╔═╡ 44813e6f-d7a2-47d6-9862-b9fd3e102ee2
md"""
##
### Downsides of preconditioning

- In general, the solution set of the preconditioned problem is **not** the solution set of of the original problem.
- Let us consider our previous 2D example.
"""

# ╔═╡ 1a1b8edf-bc1a-466e-b714-3eb71aa145ce
begin
	polytopes2 = solve(A, b, LinearOettliPrager(), InverseMidpoint())
	plot(UnionSetArray(polytopes2), ratio=1, label="preconditioned", legend=:right)
	plot!(UnionSetArray(polytopes), label="original", α=1)
end

# ╔═╡ d8735a5e-8444-4724-9796-a59b37efd0e2
md"""
##
### Take-home lesson
- In general preconditioning is needed to achieve numerical stability.
- This may however enlarge the solution set.
- If preconditioning is not specified, the package performs some heuristic checks to decide a precondition strategy.
"""

# ╔═╡ bbd4086e-af3f-4743-b9fb-01559f567560
solve(A1, b1, GaussianElimination())

# ╔═╡ 1e9bf340-2dca-4602-a60c-b2ccdc9cfb85
md"""
## Conclusions

- [IntervalLinearAlgebra.jl](https://github.com/lucaferranti/IntervalLinearAlgebra.jl) offers (will offer) a toolbox to deal with interval linear systems.
- New features (possibly!) coming:
  - Parametric interval systems
  - Eigenvalues computation
  - Determinant computation

- Ultimate goal: **Linear algebra done rigorously!**
- Download the slides: <https://github.com/lucaferranti/ILAjuliacon2021>
"""

# ╔═╡ Cell order:
# ╟─f403b050-e2e6-11eb-1c44-911de0d748bd
# ╟─0bc12fe2-4d08-4840-b651-b58df23f0277
# ╟─c9e05bd9-1d7b-41c0-99dc-29b1c1e20090
# ╟─c3a5fd72-9a9d-45f3-aaf3-bd61637d286c
# ╟─f2cff36e-9136-4eb1-acff-c97af31c2450
# ╟─286f40a2-840d-431b-a531-fe9e09dc55f1
# ╟─97e498f5-2e0c-4139-9774-86f58d59d511
# ╠═bb878728-d4ca-4c01-a4f7-bf1811a8f5c2
# ╠═a89ee6e9-565a-46d5-82db-bafb251a2200
# ╟─4848561e-2336-40fb-a5c6-8705da85c2ee
# ╟─5b111f71-4d41-43cf-a3cf-e2a6f4c2adba
# ╟─395b75e5-e76f-4f3a-b51f-57eaa3de2d4e
# ╟─f39f94f7-59ff-4dd0-93dd-4f9f158c673c
# ╟─3f5a6fed-6175-4459-8be2-9a330cedc938
# ╟─c223f95a-9885-48ae-b939-4f3fa54a7987
# ╟─32f0ff22-d42f-410f-a1a6-b8ea77b696ff
# ╟─a79e27c9-399c-44e1-ac79-f65480fe5178
# ╟─ba3a1ca6-04dc-45db-b686-a036519f2bc9
# ╟─38883c77-d506-4e81-9d25-0ac78323241a
# ╟─4549ca50-615c-4fba-afde-303ad231f648
# ╟─0f196850-f4b1-4830-9198-d72a281da05c
# ╟─d5e3856d-07e7-4e50-9e4d-c0afb7091d8b
# ╠═95edbbbf-2e35-446d-9eeb-a53f5c6b10ec
# ╟─4880c8a6-de84-4e55-be90-f73256b0888a
# ╟─75a66d2a-01b9-4d6a-8fed-8d8159ba6354
# ╟─cd558a80-2825-4953-870c-93551f60c10e
# ╟─c5aa282c-c82b-4521-994d-235cb7934c83
# ╟─32481e35-1a98-4b65-b925-264d51d24743
# ╟─58b8b1b8-93e7-46c0-b0a0-28532b152908
# ╟─faa03639-a431-4cbf-ac03-bbb33fb7b5b8
# ╠═5e0cf15b-21ab-45bb-81fa-ee962a46004f
# ╟─fc08d410-02fa-480b-a999-841476dfaaf6
# ╟─57a5fb2a-5cb2-412f-a1a0-5d27e0a5b8f7
# ╟─a70a6137-88cf-4fd5-ad50-8d76a1379a27
# ╟─0308378a-2d36-4c47-add7-0f53dc9011ca
# ╟─9dba04bd-a8c3-44da-a1ab-855829c824ae
# ╟─8f68a5f0-4ec9-40d5-89be-cddf0c34a977
# ╟─22abccab-5c93-4aa4-90e4-8209997a9a86
# ╟─e77a478d-0270-4120-a2d4-36efc455fb2e
# ╟─1ec87e5e-2f90-490b-a94f-05ea53d4cb1d
# ╟─de8734a4-46b8-45ba-aff1-fa6a73e85d96
# ╠═e3bee176-9b9a-448f-b2d0-fa53cad35555
# ╠═495e13e3-bf20-4541-ab6c-f6cb9572e9ed
# ╟─f223ac43-79cd-43be-b585-907a0bfb8f27
# ╠═acb54294-917d-41b0-b501-364b9228531b
# ╠═ea231b9c-4cc9-4d45-b171-4c66536b02a1
# ╟─44813e6f-d7a2-47d6-9862-b9fd3e102ee2
# ╟─1a1b8edf-bc1a-466e-b714-3eb71aa145ce
# ╟─d8735a5e-8444-4724-9796-a59b37efd0e2
# ╠═bbd4086e-af3f-4743-b9fb-01559f567560
# ╟─1e9bf340-2dca-4602-a60c-b2ccdc9cfb85


text:
using KmerIndexes
using BaseTestNext
using Bio.Seq

@testset "short sequence" begin
    seq = dna"ACGTAGACGT"
    index = KmerIndex(DNAKmer{4}, seq)
    @test count(kmer(dna"ACGT"), index) === 2
    @test locate(kmer(dna"ACGT"), index) == [1, 7]
    @test count(kmer(dna"AGAC"), index) === 1
    @test locate(kmer(dna"AGAC"), index) == [5]
end


text:
using Stheno: cholesky
using TemporalGPs: to_sde, smooth, SArrayStorage, decorrelate
using Plots
using Optim
using Zygote: gradient
using Random

using GPARatScale

# Create function that computes dtc objective that takes in a Stheno GP, and the
# input locations, and y1 locations

# This should use TemporalGPs inside. Replace obs noice cov matrix with the
# decorrelate op in TemporalGPs
# Line 263 on compute_invariants
# Instead of chol_Σy.U' \ (y - mean(f)) do decorrelate(LGSSM, y - mean(f)) - > exactly the same as what woul've happened if computed the covar matrix and then backsolve against y - mean(f)
function get_dtc_objective(f, u, time, y1, pseudo_time, pseudo_y1)
    # Turn the GPs into an LTISDE
    sde = to_sde(gp, SArrayStorage(Float64))
    # Turn the LTISDEs into LGSSMs by indexing
    chol_Σy = cholesky(f.Σy)
    # Instead of chol_Σy.U' \ (y - mean(f)) do decorrelate(LGSSM, y - mean(f))
    # Exactly the same as what woul've happened if computed the covar matrix and
    # then backsolve against y - mean(f)
    A = cholesky(Symmetric(cov(u))).U' \ (chol_Σy.U' \ cov(f, u))'
    Λ_ε = cholesky(Symmetric(A * A' + I))
    δ = chol_Σy.U' \ (y - mean(f))

    tmp = logdet(chol_Σy) + logdet(Λ_ε) + sum(abs2, δ) - sum(abs2, Λ_ε.U' \ (A * δ))
    _dtc = -(length(y) * typeof(tmp)(log(2π)) + tmp) / 2
    return _dtc, chol_Σy, A
end

# # Example function
# function _compute_intermediates(f::FiniteGP, y::AV{<:Real}, u::FiniteGP)
#     consistency_check(f, y, u)
#     chol_Σy = cholesky(f.Σy)
#     # Lower triangular chol decomp of Cuu backsolved with
#     # (Cfu¹ * L)ᵀ^-1 * chol(c)
#     A = cholesky(Symmetric(cov(u))).U' \ (chol_Σy.U' \ cov(f, u))'
#     Λ_ε = cholesky(Symmetric(A * A' + I))
#     δ = chol_Σy.U' \ (y - mean(f))
#
#     tmp = logdet(chol_Σy) + logdet(Λ_ε) + sum(abs2, δ) - sum(abs2, Λ_ε.U' \ (A * δ))
#     _dtc = -(length(y) * typeof(tmp)(log(2π)) + tmp) / 2
#     return _dtc, chol_Σy, A
# end

x, y_obs, _, _ = generate_small_dataset()
# Model X=R²; (x, y1) -> y2; i.e  v=(x, y1), y = y2, and z=(pseudo_x, pseudo_y1)
time_loc = x
outputs = y_obs[1]
# Generate u, the pseudo-points, by taking every third element
every_third = range(1; stop=length(x), step=3)

pseudo_time_loc = time_loc[every_third]
pseudo_outputs = outputs[every_third]
# Get the dimensions
N = length(time_loc)
M = length(pseudo_time_loc)
# TODO: Perfor  m this concatination when extending to D > 1
# Concatenate inputs so we treat them together
# inp = to_ColVecs([time_loc, y1])
# pseudo_inp = to_ColVecs([pseudo_time_loc, pseudo_y1])

# Generate GP priors
kern = Matern52()
f_prior = GP(kern, GPC())
u_prior = GP(kern, GPC())

# Compute Cfu by using finiteGP. This was to test the pairwise function.
f = f_prior(time_loc, 0.1)
u = f_prior(pseudo_time_loc, 0.1)
Cfu_naive = cov(f, u)

# Compute Cfu by using pairwise
Cfu = pairwise(kern, time_loc, pseudo_time_loc)

# TODO: Use this column by column computation when computing B
Cfu_smart = zeros(N, M)
# Compute column by column
for col_index in 1:M
    Cfu_smart[:, col_index] = pairwise(kern, time_loc, [pseudo_time_loc[col_index]])
end

# Create LTISDE from gps
storage = SArrayStorage(Float64)

f_sde = to_sde(f_prior, storage)
u_sde = to_sde(u_prior, storage)
# Turn the LTISDE into LGSSM by indexing
NOISE_SIGMA = 0.05
# TODO: This is where extending D > 1 failed. SDE doesn't support multi-dim inp
f_lgssm = f_sde(time_loc, NOISE_SIGMA)
u_lgssm = u_sde(pseudo_time_loc, NOISE_SIGMA)

# Here we'll naively compute alpha and Beta
L_naive = cholesky(f.Σy).U'
alpha_naive = L_naive \ outputs
beta_naive = L_naive \ Cfu_naive

# Compute alpha
_, alpha = decorrelate(f_lgssm, outputs)

# Compute beta
beta = zeros((N, M)) # zero initialization
# Iterate over the columns of Cfu and call decorrelate to construct beta
for col_index in 1:M
    col = Cfu[:, col_index]
    # Call decorrelate to compute that column in beta
    _, col_beta = decorrelate(f_lgssm, col)
    beta[:, col_index] = col_beta
end

# Compute A from A = cholesky(Symmetric(cov(u))).U' \ (chol_Σy.U' \ cov(f, u))'

max_stars_repo_name: willtebbutt/GPAR-at-scale

max_issues_repo_name: willtebbutt/GPAR-at-scale

max_forks_repo_name: willtebbutt/GPAR-at-scale


text:
using OrdinaryDiffEq
a=100.0
function df(du,u,p,t)
	du[1:end-1] .= a*(u[2:end] .- u[1:end-1])
	du[end] = -a*u[end]
end

u0 = [0.0,0.0,0.0,0.0,1.0]
tspan = (0,0.1)
prob = ODEProblem(df, u0, tspan)
sol = solve(prob, Tsit5(), saveat=0.001)

leg=["\$1^\\text{st}\$ tap", "\$2^\\text{nd}\$ tap", "\$3^\\text{rd}\$ tap", "\$4^\\text{th}\$ tap", "\$5^\\text{th}\$ tap"]

plts = []
for i in 1:5
	push!(plts, @pgf PlotInc(
			{ no_marks, very_thick },
			Table(; x = sol.t, y = sol[end+1-i,:])
		)
	)
	push!(plts, @pgf LegendEntry(leg[i]))
end

p = @pgf TikzPicture(
	Axis({
		height="4cm",
		width=raw"\textwidth",
		# "scale only axis",
		xmin=0,
		xmax=0.1,
		ymin=0,
		ymax=1,
		xlabel="time in seconds",
		ylabel="impulse response",
		"axis background/.style"={fill="none"},
		},
		plts...
	)
)

max_stars_repo_name: jleugeri/phd

max_issues_repo_name: jleugeri/phd

max_forks_repo_name: jleugeri/phd