text:
# Calibration of Gent-McWilliams to a baroclinic adjustment problem
pushfirst!(LOAD_PATH, joinpath(@__DIR__, ".."))

using ParameterEstimocean
using Oceananigans
using Oceananigans.Units
using Oceananigans.Models.HydrostaticFreeSurfaceModels: SliceEnsembleSize
using Oceananigans.TurbulenceClosures: FluxTapering
using LinearAlgebra, CairoMakie, DataDeps, JLD2

using Oceananigans.Architectures: arch_array

# using ElectronDisplay

architecture = CPU()

# filedir = @__DIR__
# filename = "baroclinic_adjustment_double_Lx_zonal_average.jld2"
# filepath = joinpath(filedir, filename)
# Base.download("https://www.dropbox.com/s/f8zsb33vwwwmjjm/$filename", filepath)

# filepath = "/Users/navid/Research/mesoscale-parametrization-OSM2022/baroclinic_adjustment-double_Lx/short_save_often_run/baroclinic_adjustment_double_Lx_zonal_average.jld2"
filepath = "baroclinic_adjustment_double_Lx_zonal_average_80dayrun.jld2"

file = jldopen(filepath)
coriolis = file["serialized/coriolis"]

# number of grid points
Nx, Ny, Nz = file["grid/Nx"], file["grid/Ny"], file["grid/Nz"]

# Domain
const Lx, Ly, Lz = file["grid/Lx"], file["grid/Ly"], file["grid/Lz"]

close(file)

field_names = (:b, :c, :u, :v)
forward_map_names = (:b, :c)

using ParameterEstimocean.Transformations: Transformation

transformation = (b = ZScore(),
                  c = ZScore(),
                  u = ZScore())

transformation = ZScore()

space_transformation = SpaceIndices(x=:, y=2:4:Ny-1, z=20:2:Nz-1)

transformation = (b = Transformation(space = space_transformation, normalization=ZScore()),
                  c = Transformation(space = space_transformation, normalization=ZScore()),
                  u = Transformation(space = space_transformation, normalization=ZScore()),
                  v = Transformation(space = space_transformation, normalization=RescaledZScore(1e-1)))

transformation = ZScore()

transformation = (b = Transformation(space = space_transformation, normalization=ZScore()),
                  c = Transformation(space = space_transformation, normalization=ZScore()))

times = [40days-12hours, 40days]

observations = SyntheticObservations(filepath; transformation, times, field_names, forward_map_names)

#####
##### Simulation
#####

Nensemble = 50

slice_ensemble_size = SliceEnsembleSize(size=(Ny, Nz), ensemble=Nensemble)

ensemble_grid = RectilinearGrid(architecture,
                                size = slice_ensemble_size,
                                topology = (Flat, Bounded, Bounded),
                                y = (-Ly/2, Ly/2),
                                z = (-Lz, 0),
                                halo=(3, 3))

@show ensemble_grid

file = jldopen(filepath)
closures = file["serialized/closure"]
close(file)

gent_mcwilliams_diffusivity = IsopycnalSkewSymmetricDiffusivity(slope_limiter = FluxTapering(1e-2))

gm_closure_ensemble = arch_array(architecture, [deepcopy(gent_mcwilliams_diffusivity) for _ = 1:Nensemble])

closure_ensemble = (gm_closure_ensemble, closures[1], closures[2])

@show "no convective adjustment"
closure_ensemble = (gm_closure_ensemble, closures[1])

ensemble_model = HydrostaticFreeSurfaceModel(grid = ensemble_grid,
                                             tracers = (:b, :c),
                                             buoyancy = BuoyancyTracer(),
                                             coriolis = coriolis,
                                             closure = closure_ensemble,
                                             free_surface = ImplicitFreeSurface())

Δt = 1.0
simulation = Simulation(ensemble_model; Δt, stop_time=times[end])

priors = (
     κ_skew = ScaledLogitNormal(bounds = (300, 3000)),
     κ_symmetric = ScaledLogitNormal(bounds = (300, 3000))
 )

free_parameters = FreeParameters(priors)

using ParameterEstimocean.Observations: FieldTimeSeriesCollector

simulation_fields = fields(simulation.model)
collected_fields = NamedTuple(name => simulation_fields[name] for name in ParameterEstimocean.Observations.forward_map_names(observations))
time_series_collector = FieldTimeSeriesCollector(collected_fields, observation_times(observations), architecture=CPU())

calibration = InverseProblem(observations, simulation, free_parameters; time_series_collector)

eki = EnsembleKalmanInversion(calibration;
                              noise_covariance = 5e-3,
                              resampler = Resampler(acceptable_failure_fraction=1.0))

iterate!(eki; iterations = 10)

# Last, we visualize few metrics regarding how the EKI calibration went about.

y = observation_map(calibration)
#=
θ̅(iteration) = [eki.iteration_summaries[iteration].ensemble_mean...]
varθ(iteration) = eki.iteration_summaries[iteration].ensemble_var

weight_distances = [norm(θ̅(iter)) for iter in 1:eki.iteration]
output_distances = [norm(forward_map(calibration, θ̅(iter))[:, 1] - y) for iter in 1:eki.iteration]
ensemble_variances = [varθ(iter) for iter in 1:eki.iteration]

f = Figure()
lines(f[1, 1], 1:eki.iteration, weight_distances, color = :red, linewidth = 2,
      axis = (title = "Parameter norm",
              xlabel = "Iteration",
              ylabel="|θ̅ₙ|",
              yscale = log10))
lines(f[1, 2], 1:eki.iteration, output_distances, color = :blue, linewidth = 2,
      axis = (title = "Output distance",
              xlabel = "Iteration",
              ylabel="|G(θ̅ₙ) - y|",
              yscale = log10))
ax3 = Axis(f[2, 1:2], title = "Parameter convergence",
           xlabel = "Iteration",
           ylabel = "Ensemble variance",
           yscale = log10)

for (i, pname) in enumerate(free_parameters.names)
    ev = getindex.(ensemble_variances, i)
    lines!(ax3, 1:eki.iteration, ev / ev[1], label = String(pname), linewidth = 2)
end

axislegend(ax3, valign = :top, halign = :right)

save("summary_bca.png", f); nothing #hide
save("summary_bca.svg", f); nothing #hide

# ![](summary_channel.svg)

# And also we plot the the distributions of the various model ensembles for few EKI iterations to see
# if and how well they converge to the true diffusivity values.

=#

f = Figure()

axtop = Axis(f[1, 1])

axmain = Axis(f[2, 1],
              xlabel = "κ_skew [m² s⁻¹]",
              ylabel = "κ_symmetric [m² s⁻¹]")

axright = Axis(f[2, 2])
scatters = []
labels = String[]

for iter in [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
    ## Make parameter matrix
    parameters = eki.iteration_summaries[iter].parameters
    Nensemble = length(parameters)
    Nparameters = length(first(parameters))
    parameter_ensemble_matrix = [parameters[i][j] for i=1:Nensemble, j=1:Nparameters]

    label = iter == 0 ? "Initial ensemble" : "Iteration $iter"
    push!(labels, label)
    push!(scatters, scatter!(axmain, parameter_ensemble_matrix))
    density!(axtop, parameter_ensemble_matrix[:, 1])
    density!(axright, parameter_ensemble_matrix[:, 2], direction = :y)
end

colsize!(f.layout, 1, Fixed(300))
colsize!(f.layout, 2, Fixed(200))

rowsize!(f.layout, 1, Fixed(200))
rowsize!(f.layout, 2, Fixed(300))

Legend(f[1, 2], scatters, labels, valign = :bottom, halign = :left)

hidedecorations!(axtop, grid = false)
hidedecorations!(axright, grid = false)

xlims!(axright, 0, 0.025)
ylims!(axtop, 0, 0.025)

save("distributions_bca_2.png", f); nothing #hide
save("distributions_bca_2.svg", f); nothing #hide

# ![](distributions_channel.svg)


text:
@testset "Inverse" begin
    for z ∈ Zs
        @test z * inv(z) ≈ inv(z) * z ≈ 1v
    end
end

ext: jl


text:
module CellValues

using Test
using Gridap
using Gridap.Helpers

export CellValue
export IterCellValue
export IndexCellValue

import Base: iterate
import Base: length
import Base: eltype
import Base: size
import Base: getindex
import Base: IndexStyle
import Base: show

export test_iter_cell_value
export test_index_cell_value

abstract type IterCellValue{T} end

function iterate(::IterCellValue{T})::Union{Nothing,Tuple{T,Any}} where T
  @abstractmethod
end

function iterate(::IterCellValue{T},state)::Union{Nothing,Tuple{T,Any}} where T
  @abstractmethod
end

length(::IterCellValue)::Int = @abstractmethod

eltype(::Type{C}) where C <: IterCellValue{T} where T = T

# Indexable cell Values

abstract type IndexCellValue{T,N} <: AbstractArray{T,N} end

function getindex(::IndexCellValue{T,N}, ::Vararg{<:Integer,N})::T where {T,N}
  @abstractmethod
end

size(x::IndexCellValue) = @abstractmethod

IndexStyle(::Type{<:IndexCellValue{T,N}} where {T,N}) = IndexLinear()

# Cell Values

const CellValue{T} = Union{IterCellValue{T},IndexCellValue{T}}

function show(io::IO,self::CellValue)
  for (i, a) in enumerate(self)
    println(io,"$i -> $a")
  end
end

function show(io::IO,::MIME"text/plain",self::CellValue)
  if length(self) < 15
    show(io,self)
  else
    for (i, a) in enumerate(self)
      if i >= 15
        print(io,"... (total length $(length(self)))")
        break
      end
      println(io,"$i -> $a")
    end
  end
end

# Testers

function test_iter_cell_value(icv::CellValue{T},a::AbstractArray{T}) where T
  _test_iter_cell_value(icv,a)
end

function _test_iter_cell_value(icv,a)

  @test length(icv) == length(a)

  i = 0
  for v in icv
    i += 1
    @assert _eq(v,a[i])
  end

  @test i == length(a)

  for v in icv
    @assert eltype(icv) == typeof(v)
  end

end

_eq(a,b) = a == b

_eq(a::Real,b::Real) = a ≈ b

function test_index_cell_value(icv::IndexCellValue{T},a::AbstractArray{T}) where T
  _test_index_cell_value(icv,a)
end

function _test_index_cell_value(icv,a)

  @test size(icv) == size(a)

  for i in eachindex(icv)
    v = icv[i]
    @assert v == a[i]
  end

  @test IndexStyle(icv) == IndexStyle(a)

end

end # module CellValues


text:
# This file was generated, do not modify it. # hide
using RDatasets, DataFrames
boston = dataset("MASS", "Boston")
first(boston, 3)


text:
# Closable objects:
ContextManagers.maybeenter(io::IO) = io
ContextManagers.maybeenter(ch::AbstractChannel) = ch

if isdefined(Base, :AbstractLock)
    const AbstractLock = Base.AbstractLock
else
    const AbstractLock = ReentrantLock
end
ContextManagers.maybeenter(l::AbstractLock) = (lock(l); l)
ContextManagers.exit(l::AbstractLock) = unlock(l)

"""
    ContextManagers.opentemp([parent]; kwargs...) -> tf

Create and open a temporary file. The path and the `IO` object can be accessed
through the properties `.path` and `.io` of the returned object `tf`
respectively. The positional and named arguments are passed to `mktemp`.  The
file is automatically removed and the IO object is automatically closed when
used with `@with` or `with`.
"""
ContextManagers.opentemp

struct TemporaryFile
    path::String
    io::IOStream
end

ContextManagers.opentemp(; kwargs...) = TemporaryFile(mktemp(; kwargs...)...)
ContextManagers.opentemp(parent; kwargs...) = TemporaryFile(mktemp(parent; kwargs...)...)

ContextManagers.maybeenter(tf::TemporaryFile) = tf
function ContextManagers.exit(tf::TemporaryFile)
    (path, io) = tf
    rm(path; force = true)
    close(io)
end

Base.IteratorSize(::Type{TemporaryFile}) = Base.HasLength()
Base.length(::TemporaryFile) = 2
Base.iterate(tf::TemporaryFile) = (tf.path, tf.io)
Base.iterate(::TemporaryFile, io::IOStream) = (io, nothing)
Base.iterate(::TemporaryFile, ::Nothing) = nothing

Base.IteratorEltype(::Type{TemporaryFile}) = Base.HasEltype()
Base.eltype(::Type{TemporaryFile}) = Union{String,IOStream}

"""
    ContextManagers.opentempdir(parent=tempdir(); kwargs...) -> td

Create aa temporary directory. The path can be accessed through the property
`.path`  of the returned object `td`. When this is used with `@with` or `with`,
`td` is unwrapped to a `path` automatically. For example, in the do block of
`@with(paath = opentempdir()) do; ...; end`, a string `path` is available.
"""
ContextManagers.opentempdir

struct TemporaryDirectory
    path::String
end

ContextManagers.opentempdir(; kwargs...) = TemporaryDirectory(mktempdir(; kwargs...))
ContextManagers.opentempdir(parent; kwargs...) =
    TemporaryDirectory(mktempdir(parent; kwargs...))

ContextManagers.maybeenter(td::TemporaryDirectory) = td
ContextManagers.value(td::TemporaryDirectory) = td.path
function ContextManagers.exit(td::TemporaryDirectory)
    path = td.path
    rm(path; force = true, recursive = true)
end

"""
    ContextManagers.IgnoreError()

Ignore an error (if any).

# Example

```jldoctest
julia> using ContextManagers: @with, IgnoreError

julia> @with(
           IgnoreError(),
       ) do
           error("error")
       end
```

# Extended help

Note that `IgnoreError` only ignores the error from the "inner" code.  That is
to say, in the following code, the error from the `doblock` and the context
managers `c` and `d` are ignored but not the error from the context managers `a`
and `b`.

```
@with(
    a,
    b,
    IgnoreError(),
    c,
    d,
) do
    doblock
end
```
"""
ContextManagers.IgnoreError

ContextManagers.maybeenter(ie::ContextManagers.IgnoreError) = ie
ContextManagers.exit(::ContextManagers.IgnoreError, _) = ContextManagers.Handled()


text:
using Test
using Tables: schema, Schema
using Cleaner: CleanTable, reinfer_schema

@testset "reinfer_schema is working as expected" begin
    testCT = CleanTable([:A, :B, :C], [[1, 2, 3], Any[4, missing, "z"], Any["5", "6", "9"]])
    testCT2 = CleanTable(
        [:A, :B, :C], [[1, 2, 3, 4], Any[5, missing, "z", 2.0], Any["6", "7", "8", "9"]]
    )

    @test schema(reinfer_schema(testCT)) ==
          Schema([:A, :B, :C], [Int, Union{Missing,String,Int}, String])
    @test schema(reinfer_schema(testCT2)) ==
          Schema([:A, :B, :C], [Int, Union{Missing,Real,String}, String])
    @test schema(reinfer_schema(testCT2; max_types=2)) ==
          Schema([:A, :B, :C], [Int, Any, String])

    if Threads.nthreads() > 1
        three_typed = append!([4, missing, "z"], collect(1:999_997))
        testCT = CleanTable(
            [:A, :B, :C],
            [collect(1:1_000_000), three_typed, Any[repeat(["a"], 1_000_000)...]],
        )

        @test schema(reinfer_schema(testCT)) ==
              Schema([:A, :B, :C], [Int, Union{Missing,String,Int}, String])
    end
end

max_stars_repo_name: TheRoniOne/Cleaner

max_issues_repo_name: TheRoniOne/Cleaner

max_forks_repo_name: TheRoniOne/Cleaner


text:
module TestLloyd

using ParallelKMeans
using Test
using Random

@testset "basic kmeans" begin
    X = [1. 2. 4.;]
    res = kmeans(X, 1; n_threads = 1, tol = 1e-6, verbose = false)
    @test res.assignments == [1, 1, 1]
    @test res.centers[1] ≈ 2.3333333333333335
    @test res.totalcost ≈ 4.666666666666666
    @test res.converged

    res = kmeans(X, 2; n_threads = 1, init = [1.0 4.0], tol = 1e-6, verbose = false)
    @test res.assignments == [1, 1, 2]
    @test res.centers ≈ [1.5 4.0]
    @test res.totalcost ≈ 0.5
    @test res.converged
end

@testset "no convergence yield last result" begin
    X = [1. 2. 4.;]
    res = kmeans(X, 2; n_threads = 1, init = [1.0 4.0], tol = 1e-6, max_iters = 1, verbose = false)
    @test !res.converged
    @test res.totalcost ≈ 0.5
end

@testset "singlethread linear separation" begin
    Random.seed!(2020)

    X = rand(3, 100)
    res = kmeans(X, 3; n_threads = 1, tol = 1e-6, verbose = false)

    @test res.totalcost ≈ 14.16198704459199
    @test res.converged
    @test res.iterations == 11
end

@testset "multithread linear separation quasi two threads" begin
    Random.seed!(2020)

    X = rand(3, 100)
    res = kmeans(X, 3; n_threads = 2, tol = 1e-6, verbose = false)

    @test res.totalcost ≈ 14.16198704459199
    @test res.converged
end

end # module


text:
#=
   https://dtai.cs.kuleuven.be/problog/tutorial/basic/10_inhibitioneffects.html
   """
   Example 2: Social Network

   An infectious disease spreads through a population as follows: when- ever two people are in regular
   contact with each other and one is infected, there is a probability of 0.6 of the infection spreading
   also to the other person. Given a set of initially infected people and a graph of connections between
   individuals in the population, the goal is to predict the spread of the disease.
   """

   ProbLog model
   """
   person(a).
   person(b).

   0.1::initialInf(X) :- person(X).
   0.1::contact(X,Y) :- person(X), person(Y).

   inf(X)      :- initialInf(X).
   0.6::inf(X) :- contact(X, Y), inf(Y).

   query(inf(_)). % inf(a): 0.1054 inf(b): 0.1054
   """

   Cf ~/webppl/inhibition_infection.wppl

=#

using Turing, StatsPlots, DataFrames
include("jl_utils.jl")

@model function inhibition_infection()
    n = 3
    people = 1:n

    initialInf = tzeros(n)
    for p in 1:n
        initialInf[p] ~ flip(0.1)
    end

    contact = tzeros(n,n)
    for x in 1:n, y in 1:n
        if x == y
            contact[x,y] ~ flip(0.0)
        else
            contact[x,y] ~ flip(0.1)
        end
    end

    a_inf = tzeros(n)
    for p in 1:n
        if initialInf[p]==1
            a_inf[p] ~ flip(1.0)
        else
            check = any(q -> q != p && contact[p,q]==1 && a_inf[q]==1,people)
            if check==1
                a_inf[p] ~ flip(0.6)
            else
                a_inf[p] ~ flip(0.0)
            end
        end
    end

    return a_inf[1] # a_inf

end

model = inhibition_infection()

num_chns = 4

# chns = sample(model, Prior(), 10_000)

# chns = sample(model, MH(), MCMCThreads(), 10_000, num_chns)
# chns = sample(model, MH(), 10_000)

# chns = sample(model, PG(15), MCMCThreads(), 1_000, num_chns)
# chns = sample(model, PG(15), 1_000)

# chns = sample(model, SMC(1000), MCMCThreads(), 10_000, num_chns)
chns = sample(model, SMC(1000), 10_000)

# chns = sample(model, IS(), 1_000)

#
display(chns)
show_var_dist_pct(chns,:len,1000)

chns_params = Turing.MCMCChains.get_sections(chns, :parameters)
genq = generated_quantities(model, chns_params)
show_var_dist_pct(genq)

max_stars_repo_name: tias/hakank

max_issues_repo_name: tias/hakank

max_forks_repo_name: tias/hakank


text:
module TestLength
    using Base.Test
    using CategoricalArrays

    pool = CategoricalPool([1, 2, 3])
    @test length(pool) == 3

    pool = CategoricalPool([1, 2, 3], [3, 2, 1])
    @test length(pool) == 3
end


text:
# ComponentsGDPbyIndustry.jl
#= This script pulls the data from:
- Industry Economic Account Data
	- GDP by Industry Components of Value Added
		- Value Added (A)
		- Compensation of employees (A)
		- Wages and Salaries (A)
		- Taxes on production and imports, less subsidies (A)
		- Gross operating surplus (A)
For all industries for years 2008 - 2019
Computes the NAICS2 sector code blow-up factors to be applied to salary and wages

## Source: https://apps.bea.gov/iTable/iTable.cfm?reqid=150&step=2&isuri=1&categories=compbyind
## Note: Endpoint is not available for consuming it from the API. We use the spreadsheet file.
=#
using HTTP: HTTP, request, URI
using JSON3: JSON3
using XLSX: XLSX, readxlsx, readtable, readdata, sheetnames
using CSV: CSV
using DataFrames: DataFrames, DataFrame, combine, groupby, ByRow, dropmissing!,
				  transform!, select!, select, leftjoin
const BEA_API_KEY = get(ENV, "api_bea_token", "")

#=
Download the information and format it.
=#
# Download the spreadsheet file for the components of value added
if !isfile(joinpath("data", "ComponentsOfVa.xlsx"))
	download(string(URI(scheme = "https",
						host = "apps.bea.gov",
						path = "/industry/Release/XLS/CompByInd/ComponentsOfVa.xlsx")),
		 	 joinpath("data", "ComponentsOfVa.xlsx"))
end
# Read the spreadsheet file
components_of_va = readxlsx(joinpath("data", "ComponentsOfVa.xlsx"))
# The first sheet has the name of the series to sheets mapping
tblnames = Dict(components_of_va[1]["C5:C17"] .=> sheetnames(components_of_va)[2:end])
# These are the cell range for all industries for years 2008 - 2019
range_for_industries_2008_2019 = "N11:Y96"
# Get the industry names
all_industries = convert(Vector{String}, vec(components_of_va[2]["B11:B96"]))
# Read the data for each variable and format it together with the industry and year
components_of_va_data =
	DataFrame(reduce(hcat,
					 vec([isa(x, Number) ? x : missing
					 	  for x in components_of_va[tblnames[k]][range_for_industries_2008_2019]])
						  for k in sort(collect(keys(tblnames)))),
			  sort!(collect(keys(tblnames))))
components_of_va_data[!,:industry] .= repeat(all_industries, outer = length(2008:2019))
components_of_va_data[!,:year] .= repeat(2008:2019, inner = length(all_industries))
components_of_va_data =
	components_of_va_data[!,union([:industry, :year], propertynames(components_of_va_data))]
transform!(components_of_va_data, :industry => ByRow(strip) => :industry)

#=
Find the mapping to of industry to NAICS2 (sector) through the APIs
=#
# url = URI(scheme = "https", host = "apps.bea.gov", path = "/api/data",
#           query = ["UserID" => BEA_API_KEY,
#                    "method" => "GetParameterList",
#                    "DatasetName" => "UnderlyingGDPbyIndustry",
#                    "ResultFormat" => "JSON"
#                   ])
# response = request("GET", url)
# @assert response.status == 200
# json = JSON3.read(response.body)
# json.BEAAPI.Results.Parameter
#=
Using the code above I found the ParameterName

Using a similar table, I can get the `Industry` list with the NAICS
TableID: 6 -- Components of Value Added by Industry (A)
Frequency: "A" -- Annual
=#
# url = URI(scheme = "https", host = "apps.bea.gov", path = "/api/data",
#           query = ["UserID" => BEA_API_KEY,
#                    "method" => "GetParameterValues",
#                    "DatasetName" => "UnderlyingGDPbyIndustry",
#                 #    "TableID" => 210,
#                    "ParameterName" => "Industry",
#                    "ResultFormat" => "JSON"
#         ])
# response = request("GET", url)
# @assert response.status == 200
# json = JSON3.read(response.body)
# industries = DataFrame(json.BEAAPI.Results.ParamValue)
# CSV.write(joinpath("data", "BEA_Industry_Naics.csv"), industries)
#=
We obtain a mapping from BEA industry to NAICS2.
=#
industries = CSV.read(joinpath("data", "BEA_Industry_Naics.csv"), DataFrame)
transform!(industries, :Key => ByRow(x -> match(r"\d+\-?\d+", x)) => :naics)
transform!(industries, :naics => ByRow(x -> isnothing(x) ? missing : x.match) => :naics)
dropmissing!(industries)
transform!(industries, :Desc => ByRow(x -> replace(x, " (A,Q)" => "")) => :industry)
select!(industries, [:naics, :industry])
industries = combine(groupby(industries, :naics), first)
#=
Join the data so we can obtain the relevant information at the NAICS2 level.
=#
naics = leftjoin(components_of_va_data, industries, on = :industry)
select!(naics,
		["year", "naics", "industry",
		 "Wages and salaries",
		 "Compensation of employees",
		 "Value Added",
		 "Gross operating surplus",
		 "Taxes on production and imports, less subsidies"])
dropmissing!(naics)
#=
We find that we need to handle these four cases:
- "31-33" use 31
- "44-45" use 44
- "48-49" use 48
- "99" do not include
=#
bea_oews = Dict("31" => "31-33", "44" => "44-45", "48" => "48-49")
#=
Compute simple salary and wages to resource cost factors per industry (NAICS2).
=#
comp_of_va = select(naics, ["year", "naics", "Wages and salaries", "Compensation of employees", "Value Added"])
transform!(comp_of_va, ["Compensation of employees", "Wages and salaries"] => ByRow(/) => :l_w)
transform!(comp_of_va,
		   :naics => ByRow(x -> get(bea_oews, string(x), x)),
		   renamecols = false)
select!(comp_of_va, [:year, :naics, :l_w])
# setdiff(oews_industries, resource_cost_factors[!,:naics])
CSV.write(joinpath("data", "comp_of_va.csv"), comp_of_va)

max_stars_repo_name: Nosferican/Estimating-Resource-Costs

max_issues_repo_name: Nosferican/Estimating-Resource-Costs

max_forks_repo_name: Nosferican/Estimating-Resource-Costs


text:
#=
Copyright (c) 2018-2022 Chris Coey, Lea Kapelevich, and contributors

This Julia package Hypatia.jl is released under the MIT license; see LICENSE
file in the root directory or at https://github.com/chriscoey/Hypatia.jl

interior point stepping routines for algorithms based on the
homogeneous self dual embedding
=#

# update the RHS for prediction direction
function update_rhs_pred(
    solver::Solver{T},
    rhs::Point{T},
    ) where {T <: Real}
    rhs.x .= solver.x_residual
    rhs.y .= solver.y_residual
    rhs.z .= solver.z_residual
    rhs.tau[] = solver.tau_residual

    for (s_k, d_k) in zip(rhs.s_views, solver.point.dual_views)
        @. s_k = -d_k
    end

    rhs.kap[] = -solver.point.kap[]

    return rhs
end

# update the prediction RHS with adjustment
function update_rhs_predadj(
    solver::Solver{T},
    rhs::Point{T},
    dir::Point{T},
    ) where {T <: Real}
    rhs.vec .= 0

    rteps = sqrt(eps(T))
    irtrtmu = inv(sqrt(sqrt(solver.mu)))
    for (k, cone_k) in enumerate(solver.model.cones)
        Cones.use_dder3(cone_k) || continue
        H_prim_dir_k = cone_k.vec1
        prim_k_scal = cone_k.vec2
        prim_dir_k = dir.primal_views[k]

        @. prim_k_scal = irtrtmu * prim_dir_k
        Cones.hess_prod_slow!(H_prim_dir_k, prim_dir_k, cone_k)
        dder3_k = Cones.dder3(cone_k, prim_k_scal)

        # only use third order deriv if it nearly satisfies an identity
        dot1 = dot(dder3_k, cone_k.point)
        dot2 = irtrtmu * dot(prim_k_scal, H_prim_dir_k)
        dder3_viol = abs(dot1 - dot2) / (rteps + abs(dot2))
        if dder3_viol < T(1e-4) # TODO tune
            @. rhs.s_views[k] = H_prim_dir_k + dder3_k
        end
    end

    taubar = solver.point.tau[]
    tau_dir_tau = dir.tau[] / taubar
    rhs.kap[] = tau_dir_tau * solver.mu / taubar * (1 + tau_dir_tau)

    return rhs
end

# update the RHS for centering direction
function update_rhs_cent(
    solver::Solver{T},
    rhs::Point{T},
    ) where {T <: Real}
    rhs.x .= 0
    rhs.y .= 0
    rhs.z .= 0
    rhs.tau[] = 0

    rtmu = sqrt(solver.mu)
    for (k, cone_k) in enumerate(solver.model.cones)
        duals_k = solver.point.dual_views[k]
        grad_k = Cones.grad(cone_k)
        @. rhs.s_views[k] = -duals_k - rtmu * grad_k
    end

    # kap
    rhs.kap[] = -solver.point.kap[] + solver.mu / solver.point.tau[]

    return rhs
end

# update the centering RHS with adjustment
function update_rhs_centadj(
    solver::Solver{T},
    rhs::Point{T},
    dir::Point{T},
    ) where {T <: Real}
    rhs.vec .= 0

    rteps = sqrt(eps(T))
    irtrtmu = inv(sqrt(sqrt(solver.mu)))
    for (k, cone_k) in enumerate(solver.model.cones)
        Cones.use_dder3(cone_k) || continue
        H_prim_dir_k_scal = cone_k.vec1
        prim_k_scal = cone_k.vec2
        prim_dir_k = dir.primal_views[k]

        @. prim_k_scal = irtrtmu * prim_dir_k
        Cones.hess_prod_slow!(H_prim_dir_k_scal, prim_k_scal, cone_k)
        dder3_k = Cones.dder3(cone_k, prim_k_scal)

        # only use third order deriv if it nearly satisfies an identity
        dot1 = dot(dder3_k, cone_k.point)
        dot2 = dot(prim_k_scal, H_prim_dir_k_scal)
        dder3_viol = abs(dot1 - dot2) / (rteps + abs(dot2))
        if dder3_viol < T(1e-4) # TODO tune
            rhs.s_views[k] .= dder3_k
        end
    end

    taubar = solver.point.tau[]
    tau_dir_tau = dir.tau[] / taubar
    rhs.kap[] = tau_dir_tau * solver.mu / taubar * tau_dir_tau

    return rhs
end

include("predorcent.jl")
include("combined.jl")


text:
using SpectralEmulator
using Test

@testset "SpectralEmulator.jl" begin
    # Write your own tests here.
end


text:
"""

```math
\\tilde X = \\hat X_{t+1|t} = A\\hat X_{t|t}
```
"""
function KalmanFilter(M::LGSSM, y)
    RR = M.R * M.R'
    SS = M.S * M.S'
    ydim, qdim = size(M.B)
    xdim, pdim = size(M.R)
    n = size(y, 1)

    ϵ = zeros(ydim, 1, n)
    Γ = zeros(ydim, ydim, n)
    Γ⁻¹ = zeros(ydim, ydim, n)
    K = zeros(xdim, ydim, n) # Kalman filter gain
    H = zeros(xdim, ydim, n) # Kalman prediction gain
    Λ = zeros(xdim, xdim, n)

    Xp = zeros(xdim, 1, n)
    Σp = zeros(xdim, xdim, n)
    Xf = zeros(xdim, 1, n)
    Σf = zeros(xdim, xdim, n)

    ℓ = 0.0

    for t in 1:n
        # Prediction
        if t == 1
            Xp[:,:,t] = 0.0
            Σp[:,:,t] = M.Σᵥ
        else
            Xp[:,:,t] = M.A * Xf[:,:,t-1]
            Σp[:,:,t] = M.A * Σf[:,:,t-1] * M.A' + RR
        end

        # Correction
        ϵ[:,:,t] = y[t] - M.B * Xp[:,:,t]
        Γ[:,:,t] = M.B*Σp[:,:,t]*M.B' + SS
        Γ⁻¹[:,:,t] = inv(Γ[:,:,t])
        K[:,:,t] = Σp[:,:,t]*M.B'*Γ⁻¹[:,:,t]
        H[:,:,t] = M.A*K[:,:,t]
        Λ[:,:,t] = M.A - H[:,:,t]*M.B
        Xf[:,:,t] = Xp[:,:,t] + K[:,:,t]*ϵ[:,:,t]
        Σf[:,:,t] = Σp[:,:,t] - K[:,:,t]*M.B*Σp[:,:,t]
        ℓ += logdet(Γ[:,:,t]) + ϵ[:,:,t]'*Γ⁻¹[:,:,t]*ϵ[:,:,t]
    end
    return Dict(:Xp => Xp,
                :Xf => Xf,
                :Σp => Σp,
                :Σf => Σf,
                :ϵ => ϵ,
                :Γ => Γ,
                :Γ⁻¹ => Γ⁻¹,
                :K => K,
                :H => H,
                :Λ => Λ,
                :ℓ => -ℓ/2)
end

"""

    Disturbanve Smoother

    Algorithm 5.2.15 Cappe 2005
"""
function DisturbanceSmoother(M::LGSSM, y)
    n = size(y, 1)
    pdim = size(M.B, 2)
    qdim = size(y, 2)
    RR = M.R*M.R'

    p = zeros(pdim, 1, n-1)
    C = zeros(pdim, pdim, n-1)
    U = zeros(pdim, pdim, n-1)
    Ξ = zeros(pdim, pdim, n-1)

    Xs = zeros(pdim, 1, n)
    Σs = zeros(pdim, pdim, n)

    KF = KalmanFilter(M, y)

    # Smoothed disturbances
    for t in n-1:-1:1
        if t == n-1
            p[:,:,t] = M.B'*KF[:Γ⁻¹][:,:,n]*KF[:ϵ][:,:,n]
            C[:,:,t] = M.B'*KF[:Γ⁻¹][:,:,n]*M.B
        else
            p[:,:,t] = M.B'*KF[:Γ⁻¹][:,:,t+1]*KF[:ϵ][:,:,t+1] + KF[:Λ][:,:,t+1]'*p[:,:,t+1]
            C[:,:,t] = M.B'*KF[:Γ⁻¹][:,:,t+1]*M.B + KF[:Λ][:,:,t+1]'*C[:,:,t+1]*KF[:Λ][:,:,t+1]
        end
        U[:,:,t] = M.R'*p[:,:,t]
        Ξ[:,:,t] = I - M.R'*C[:,:,t]*M.R
    end

    # Smoothed states
    Xs[:,:,1] = M.Σᵥ*(M.B'*KF[:Γ⁻¹][:,:,1]*KF[:ϵ][:,:,1] + KF[:Λ][:,:,1]*p[:,:,1])
    Σs[:,:,1] = M.Σᵥ - M.Σᵥ*(M.B'*KF[:Γ⁻¹][:,:,1]*M.B + KF[:Λ][:,:,1]'*C[:,:,1]*KF[:Λ][:,:,1])*M.Σᵥ
    for t in 1:n-1
        Xs[:,:,t+1] = M.A*Xs[:,:,t] + M.R*U[:,:,t]
        Σs[:,:,t+1] = M.A*Σs[:,:,t]*M.A' + M.R*Ξ[:,:,t]*M.R' -
                        M.A*KF[:Σp][:,:,t]*KF[:Λ][:,:,t]'*C[:,:,t]*RR -
                        RR*C[:,:,t]*KF[:Λ][:,:,t]*KF[:Σp][:,:,t]*M.A'
    end
    return Dict(:Xp => KF[:Xp],
                :Xf => KF[:Xf],
                :Xs => Xs,
                :Σp => KF[:Σp],
                :Σf => KF[:Σf],
                :Σs => Σs,
                :ϵ => KF[:ϵ],
                :Γ => KF[:Γ],
                :Γ⁻¹ => KF[:Γ⁻¹],
                :K => KF[:K],
                :H => KF[:H],
                :Λ => KF[:Λ],
                :p => p,
                :C => C,
                :U => U,
                :Ξ => Ξ,
                :ℓ => KF[:ℓ])
end

"""
    Backward Information Recursion

    Proposition 5.2.21 Cappe (2005)
"""
function BackwardInformationRecursion(M::LGSSM, y)
    n = size(y, 1)
    pdim = size(M.B, 2)
    qdim = size(y, 2)
    RR = M.R*M.R'
    SS = M.S*M.S'
    SS⁻¹ = inv(SS)

    κ = zeros(pdim, 1, n)
    κ̃ = zeros(pdim, 1, n)
    Π = zeros(pdim, pdim, n)
    Π̃ = zeros(pdim, pdim, n)

    Xs = zeros(pdim, 1, n)
    Σs = zeros(pdim, pdim, n)

    KF = KalmanFilter(M, y)

    for t in n-1:-1:1
        κ̃[:,:,t+1] = M.B'*SS⁻¹*y[t+1] + κ[:,:,t+1]
        Π̃[:,:,t+1] = M.B'*SS⁻¹*M.B + Π[:,:,t+1]
        V = inv(I + Π̃[:,:,t+1]*RR)
        κ[:,:,t] = M.A'V*κ̃[:,:,t+1]
        Π[:,:,t] = M.A'V*Π̃[:,:,t+1]*M.A
    end
    Dict(:Xp => KF[:Xp],
                :Xf => KF[:Xf],
                :Σp => KF[:Σp],
                :Σf => KF[:Σf],
                :κ => κ,
                :Π => Π)
end

"""
    Forward-backward smoother

    Algorithm 5.2.22 Cappe (2005)
"""
function ForwardBackwardSmoother(M::LGSSM, y)
    n = size(y, 1)
    pdim = size(M.B, 2)
    qdim = size(y, 2)
    RR = M.R*M.R'
    SS = M.S*M.S'
    SS⁻¹ = inv(SS)

    κ = zeros(pdim, 1, n)
    κ̃ = zeros(pdim, 1, n)
    Π = zeros(pdim, pdim, n)
    Π̃ = zeros(pdim, pdim, n)

    Xs = zeros(pdim, 1, n)
    Σs = zeros(pdim, pdim, n)

    KF = KalmanFilter(M, y)

    W = inv(I + Π[:,:,n]*KF[:Σf][:,:,n])
    Xs[:,:,n] = KF[:Xf][:,:,n] + KF[:Σf][:,:,n]*W*(κ[:,:,n] - Π[:,:,n]*KF[:Xf][:,:,n])
    Σs[:,:,n] = KF[:Σf][:,:,n] - KF[:Σf][:,:,n]*W*Π[:,:,n]*KF[:Σf][:,:,n]

    for t in n-1:-1:1
        κ̃[:,:,t+1] = M.B'*SS⁻¹*y[t+1] + κ[:,:,t+1]
        Π̃[:,:,t+1] = M.B'*SS⁻¹*M.B + Π[:,:,t+1]
        V = inv(I + Π̃[:,:,t+1]*RR)
        κ[:,:,t] = M.A'V*κ̃[:,:,t+1]
        Π[:,:,t] = M.A'V*Π̃[:,:,t+1]*M.A
        W = inv(I + Π[:,:,t]*KF[:Σf][:,:,t])
        Xs[:,:,t] = KF[:Xf][:,:,t] + KF[:Σf][:,:,t]*W*(κ[:,:,t] - Π[:,:,t]*KF[:Xf][:,:,t])
        Σs[:,:,t] = KF[:Σf][:,:,t] - KF[:Σf][:,:,t]*W*Π[:,:,t]*KF[:Σf][:,:,t]
    end
    Dict(:Xp => KF[:Xp],
                :Xf => KF[:Xf],
                :Xs => Xs,
                :Σp => KF[:Σp],
                :Σf => KF[:Σf],
                :Σs => Σs,
                :κ => κ,
                :Π => Π)
end

"""
    Backward Information Recursion

    Algorithm 6.1.2 Cappe (2005)
"""
function BackwardMarkovianStateSampling(M::LGSSM, y, N::Int)
    n = size(y, 1)
    KF = KalmanFilter(M, y)
    pdim = size(M.B, 2)
    qdim = size(y, 2)

    X = zeros(pdim, n, N)
    X[:,1,:] = rand(MvNormal(reshape(KF[:Xf][:,:,1],1), KF[:Σf][:,:,1]), N)
    for t in 2:n
        X[:,t,:] = rand(MvNormal(reshape(KF[:Xf][:,:,t],1), KF[:Σf][:,:,t]), N)
    end
    return X
end

"""
    Sampling with Dual Smoothing

    Algorithm 6.1.3 Cappe (2005)
"""
function SamplingWithDualSmoothing(M::LGSSM, y)
    n = size(y, 1)
    pdim = size(M.B, 2)
    qdim = size(y, 2)

    XY′ = Generate(M, n)
    KF = DisturbanceSmoother(M, y)
    KF′ = DisturbanceSmoother(M, XY′[:Y])
    return KF[:Xf] + XY′[:X] - KF′[:Xf]
end


text:
# for sum(xs, weights) (#522)
Base.sum(xs::AbstractArray, weights::AbstractArray) = dot(xs, weights)
struct SumRuleConfig <: RuleConfig{Union{HasReverseMode}} end

@testset "Maps and Reductions" begin
    @testset "sum(x; dims=$dims)" for dims in (:, 2, (1, 3))
        # Forward
        test_frule(sum, rand(5); fkwargs=(; dims=dims))
        test_frule(sum, rand(ComplexF64, 2, 3, 4); fkwargs=(; dims=dims))

        # Reverse
        test_rrule(sum, rand(5); fkwargs=(; dims=dims))
        test_rrule(sum, rand(ComplexF64, 2, 3, 4); fkwargs=(; dims=dims))

        # Structured matrices
        test_rrule(sum, rand(5)'; fkwargs=(; dims=dims))
        y, back = rrule(sum, UpperTriangular(rand(5, 5)); dims=dims)
        unthunk(back(y * (1 + im))[2]) isa UpperTriangular{Float64}
        @test_skip test_rrule(
            sum,
            UpperTriangular(rand(5, 5)) ⊢ randn(5, 5);
            fkwargs=(; dims=dims),
            check_inferred=false,
        ) # Problem: in add!!  Evaluated: isapprox

        # Boolean -- via @non_differentiable
        test_rrule(sum, randn(5) .> 0; fkwargs=(; dims=dims))

        # Function allowing for 2nd derivatives
        for x in (rand(5), rand(2, 3, 4))
            dy = maximum(x; dims=dims)
            test_frule(ChainRules._unsum, x, dy, dims)
            test_rrule(ChainRules._unsum, x, dy, dims)
        end

        # Arrays of arrays
        for x in
            ([rand(ComplexF64, 3) for _ in 1:4], [rand(3) for _ in 1:2, _ in 1:3, _ in 1:4])
            test_rrule(sum, x; fkwargs=(; dims=dims), check_inferred=false)

            dy = sum(x; dims=dims)
            ddy = rrule(ChainRules._unsum, x, dy, dims)[2](x)[3]
            @test size(ddy) == size(dy)
        end
    end

    @testset "sum abs2" begin
        sizes = (3, 4, 7)
        @testset "dims = $dims" for dims in (:, 1)
            @testset "Array{$N, $T}" for N in eachindex(sizes), T in (Float64, ComplexF64)
                x = randn(T, sizes[1:N]...)
                test_frule(sum, abs2, x; fkwargs=(; dims=dims))
                test_rrule(sum, abs2, x; fkwargs=(; dims=dims))
            end

            # Boolean -- via @non_differentiable, test that this isn't ambiguous
            test_rrule(sum, abs2, randn(5) .> 0; fkwargs=(; dims=dims))
        end
    end  # sum abs2

    @testset "sum(f, xs)" begin
        # This calls back into AD
        test_rrule(sum, abs, [-4.0, 2.0, 2.0])
        test_rrule(sum, cbrt, randn(5))
        test_rrule(sum, Multiplier(2.0), [2.0, 4.0, 8.0])

        # Complex numbers
        test_rrule(sum, sqrt, rand(ComplexF64, 5))
        test_rrule(sum, abs, rand(ComplexF64, 3, 4))  # complex -> real

        # inference fails for array of arrays
        test_rrule(sum, sum, [[2.0, 4.0], [4.0, 1.9]]; check_inferred=false)

        # dims kwarg
        test_rrule(sum, abs, [-2.0 4.0; 5.0 1.9]; fkwargs=(; dims=1))
        test_rrule(sum, abs, [-2.0 4.0; 5.0 1.9]; fkwargs=(; dims=2))
        test_rrule(sum, sqrt, rand(ComplexF64, 3, 4); fkwargs=(; dims=(1,)))

        test_rrule(sum, abs, @SVector [1.0, -3.0])

        # covectors
        x = [-4.0 2.0; 2.0 -1.0]
        test_rrule(sum, inv, x[1, :]')
        test_rrule(sum, inv, x[1:1, :]')
        test_rrule(sum, inv, transpose(view(x, 1, :)))

        # Make sure we preserve type for StaticArrays
        ADviaRuleConfig = ChainRulesTestUtils.ADviaRuleConfig
        _, pb = rrule(ADviaRuleConfig(), sum, abs, @SVector [1.0, -3.0])
        @test pb(1.0) isa Tuple{NoTangent,NoTangent,SVector{2,Float64}}

        # make sure we preserve type for Diagonal
        _, pb = rrule(ADviaRuleConfig(), sum, abs, Diagonal([1.0, -3.0]))
        @test pb(1.0)[3] isa Diagonal

        # Boolean -- via @non_differentiable, test that this isn't ambiguous
        test_rrule(sum, sqrt, randn(5) .> 0)
        test_rrule(sum, sqrt, randn(5, 5) .> 0; fkwargs=(; dims=1))
        # ... and Bool produced by function
        @test_skip test_rrule(sum, iszero, randn(5))  # DimensionMismatch("second dimension of A, 1, does not match length of x, 0")
    end

    # https://github.com/JuliaDiff/ChainRules.jl/issues/522
    @testset "sum(xs, weights) (#522)" begin
        xs = rand(5)
        weights = rand(5)

        @test rrule(SumRuleConfig(), Base.sum, xs, weights) isa Nothing
    end

    @testset "prod" begin
        @testset "Array{$T}" for T in [Float64, ComplexF64]
            @testset "size = $sz, dims = $dims" for (sz, dims) in [
                ((12,), :),
                ((12,), 1),
                ((3, 4), 1),
                ((3, 4), 2),
                ((3, 4), :),
                ((3, 4), [1, 2]),
                ((3, 4, 1), 1),
                ((3, 2, 2), 3),
                ((3, 2, 2), 2:3),
            ]
                x = randn(T, sz)
                test_rrule(prod, x; fkwargs=(dims=dims,), check_inferred=true)
                x[1] = 0
                test_rrule(prod, x; fkwargs=(dims=dims,), check_inferred=true)
                x[5] = 0
                test_rrule(prod, x; fkwargs=(dims=dims,), check_inferred=true)
                x[3] = x[7] = 0  # two zeros along some slice, for any dims
                test_rrule(prod, x; fkwargs=(dims=dims,), check_inferred=true)

                if ndims(x) == 3
                    xp = PermutedDimsArray(x, (3, 2, 1))  # not a StridedArray
                    test_rrule(prod, xp; fkwargs=(dims=dims,), check_inferred=true)
                end
            end

            @testset "structured wrappers" begin
                # Adjoint -- like PermutedDimsArray this may actually be used
                xa = adjoint(rand(T, 4, 4))
                test_rrule(prod, xa)
                test_rrule(prod, xa; fkwargs=(dims=2,))
                # use Adjoint for tangent
                test_rrule(prod, xa ⊢ rand(T, 4, 4)')
                test_rrule(prod, xa ⊢ rand(T, 4, 4)'; fkwargs=(dims=2,))
                @test unthunk(rrule(prod, adjoint(rand(T, 3, 3)))[2](1.0)[2]) isa Matrix
                @test unthunk(
                    rrule(prod, adjoint(rand(T, 3, 3)); dims=1)[2](ones(1, 3))[2]
                ) isa Matrix

                # Diagonal -- a stupid thing to do, product of zeros! Shouldn't be an error though:
                @test iszero(unthunk(rrule(prod, Diagonal(rand(T, 3)))[2](1.0)[2]))
                @test iszero(
                    unthunk(rrule(prod, Diagonal(rand(T, 3)); dims=1)[2](ones(1, 3))[2])
                )
                # does a division for the complex case, so is not necessarily exact
                @test isapprox(
                    unthunk(rrule(prod, Diagonal(rand(T, 1)))[2](1.0)[2]), # 1x1 sparse matrix
                    hcat(1);
                    rtol=T <: Complex ? 2eps() : 0.0,
                )
                @test unthunk(
                    rrule(prod, Diagonal(ones(T, 2)); dims=1)[2](ones(1, 2))[2]
                ) == Diagonal([0.0 1; 1 0])

                # Triangular -- almost equally stupud
                @test iszero(
                    unthunk(rrule(prod, UpperTriangular(rand(T, 3, 3)))[2](1.0)[2])
                )
                @test unthunk(rrule(prod, UpperTriangular(ones(T, 2, 2)))[2](1.0)[2]) ==
                    UpperTriangular([0.0 0; 1 0])

                # Symmetric -- at least this doesn't have zeros, still an unlikely combination
                xs = Symmetric(rand(T, 4, 4))
                @test unthunk(rrule(prod, Symmetric(T[1 2; -333 4]))[2](1.0)[2]) ==
                    [16 8; 8 4]
                # TODO debug why these fail  https://github.com/JuliaDiff/ChainRules.jl/issues/475
                @test_skip test_rrule(prod, xs)
                @test_skip test_rrule(prod, xs; fkwargs=(dims=2,))
            end
        end
        @testset "Array{Float32}, no zero entries" begin
            v = [1f-5, 1f-10, 1f-15, 1f-20]
            @test prod(v) == 0
            @test unthunk(rrule(prod, v)[2](1.0f0)[2]) == zeros(4)
            test_rrule(prod, v)
        end
    end # prod
end

@testset "Accumulations" begin
    @testset "cumprod" begin
        v = round.(10 .* randn(9), sigdigits=3)
        test_rrule(cumprod, v)
        v[3] = 0
        test_rrule(cumprod, v)
        v[6] = 0
        test_rrule(cumprod, v)

        @testset "higher dimensions, dims=$dims" for dims in (1, 2, 3)
            m = round.(10 .* randn(4, 5), sigdigits=3)
            test_rrule(cumprod, m; fkwargs=(; dims=dims), atol=0.1)
            m[2, 2] = 0
            m[2, 4] = 0
            test_rrule(cumprod, m; fkwargs=(; dims=dims))

            t = round.(10 .* randn(3, 3, 3), sigdigits=3)
            test_rrule(cumprod, t; fkwargs=(; dims=dims))
            t[2, 2, 2] = 0
            t[2, 3, 3] = 0
            test_rrule(cumprod, t; fkwargs=(; dims=dims))
        end

        @testset "types" begin
            back = rrule(cumprod, [1, 2, 3])[2]  # rule allows integer input, but test_rrule does not
            @test unthunk(back(fill(0.5, 3))[2]) == [9 / 2, 2, 1]

            back = rrule(cumprod, PermutedDimsArray([1 2; 3 4], (2, 1)); dims=1)[2]
            @test unthunk(back(ones(Float32, 2, 2))[2]) == [3 5; 1 3]

            @test_throws Exception cumprod(Symmetric([1 2; 3 4]); dims=1) # forward pass fails, so can't test gradient

            back = rrule(cumprod, Diagonal([1, 2]); dims=1)[2]
            @test unthunk(back(fill(0.5, 2, 2))[2]) ≈ [1/2 0; 0 0]  # ProjectTo'd to Diagonal now
        end
    end
end


text:
"""

Recursively parse the model to construct command line for chain id.

$(SIGNATURES)

### Required arguments
```julia
* `m::CmdStanSampleModel`              : CmdStanSampleModel
* `id::Int`                            : Chain id
```
Not exported
"""
function cmdline(m::Union{SampleModel, Sample, Adapt, Hmc, Engine}, id)

  #=
  `./bernoulli3 sample num_samples=1000 num_warmup=1000
    save_warmup=0 thin=1 adapt engaged=1 gamma=0.05 delta=0.8 kappa=0.75
    t0=10.0 init_buffer=75 term_buffer=50 window=25 algorithm=hmc engine=nuts
    max_depth=10 metric=diag_e stepsize=1.0 stepsize_jitter=1.0 random
    seed=-1 init=bernoulli3_1.init.R id=1 data file=bernoulli3_1.data.R
    output file=bernoulli3_samples_1.csv refresh=100`,
  =#
  cmd = ``
  if isa(m, SampleModel)
    # Handle the model name field for unix and windows
    cmd = `$(m.exec_path)`

    # Sample() specific portion of the model
    cmd = `$cmd $(cmdline(getfield(m, :method), id))`

    # Common to all models
    cmd = `$cmd random seed=$(getfield(m, :seed).seed)`

    # Init file required?
    if length(m.init_file) > 0 && isfile(m.init_file[id])
      cmd = `$cmd init=$(m.init_file[id])`
    else
      cmd = `$cmd init=$(m.init.bound)`
    end

    # Data file required?
    if length(m.data_file) > 0 && isfile(m.data_file[id])
      cmd = `$cmd id=$(id) data file=$(m.data_file[id])`
    end

    # Output files
    cmd = `$cmd output`
    if length(m.sample_file[id]) > 0
      cmd = `$cmd file=$(m.sample_file[id])`
    end
    if length(m.diagnostic_file) > 0
      cmd = `$cmd diagnostic_file=$(m.diagnostic_file[id])`
    end

    # Refresh rate
    cmd = `$cmd refresh=$(string(getfield(m, :output).refresh))`

  else

    # The 'recursive' part
    if isa(m, SamplingAlgorithm)
      cmd = `$cmd algorithm=$(split(lowercase(string(typeof(m))), '.')[end])`
    elseif isa(m, Engine)
      cmd = `$cmd engine=$(split(lowercase(string(typeof(m))), '.')[end])`
    else
      cmd = `$cmd $(split(lowercase(string(typeof(m))), '.')[end])`
    end
    for name in fieldnames(typeof(m))
      if  isa(getfield(m, name), String) || isa(getfield(m, name), Tuple)
        cmd = `$cmd $(name)=$(getfield(m, name))`
      elseif length(fieldnames(typeof(getfield(m, name)))) == 0
        if isa(getfield(m, name), Bool)
          cmd = `$cmd $(name)=$(getfield(m, name) ? 1 : 0)`
        else
          if name == :metric || isa(getfield(m, name), DataType)
            cmd = `$cmd $(name)=$(split(lowercase(string(typeof(getfield(m, name)))), '.')[end])`
          else
            if name == :algorithm && typeof(getfield(m, name)) == StanSample.Fixed_param
              cmd = `$cmd $(name)=fixed_param`
            else
              cmd = `$cmd $(name)=$(getfield(m, name))`
            end
          end
        end
      else
        cmd = `$cmd $(cmdline(getfield(m, name), id))`
      end
    end
  end

  cmd

end


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

mutable struct ParameterValuesValue <: SwaggerModel
    value::Any # spec type: Union{ Nothing, Any } # spec name: value

    function ParameterValuesValue(;value=nothing)
        o = new()
        validate_property(ParameterValuesValue, Symbol("value"), value)
        setfield!(o, Symbol("value"), value)
        o
    end
end # type ParameterValuesValue

const _property_map_ParameterValuesValue = Dict{Symbol,Symbol}(Symbol("value")=>Symbol("value"))
const _property_types_ParameterValuesValue = Dict{Symbol,String}(Symbol("value")=>"Any")
Base.propertynames(::Type{ ParameterValuesValue }) = collect(keys(_property_map_ParameterValuesValue))
Swagger.property_type(::Type{ ParameterValuesValue }, name::Symbol) = Union{Nothing,eval(Base.Meta.parse(_property_types_ParameterValuesValue[name]))}
Swagger.field_name(::Type{ ParameterValuesValue }, property_name::Symbol) =  _property_map_ParameterValuesValue[property_name]

function check_required(o::ParameterValuesValue)
    true
end

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


text:
# Create and return VTK grids that are ready to be filled with data (vtu version)
function build_vtk_grids(::Val{:vtu}, mesh::TreeMesh, n_visnodes, verbose,
                         output_directory, is_datafile, filename)
  coordinates, levels, center_level_0, length_level_0 = extract_mesh_information(mesh)

  # Extract number of spatial dimensions
  ndims_ = size(coordinates, 1)

  # Calculate VTK points and cells
  verbose && println("| Preparing VTK cells...")
  if is_datafile
    @timeit "prepare VTK cells (node data)" begin
      vtk_points, vtk_cells = calc_vtk_points_cells(Val(ndims_), coordinates, levels,
                                                    center_level_0, length_level_0, n_visnodes)
    end
  end

  # Prepare VTK points and cells for celldata file
  @timeit "prepare VTK cells (cell data)" begin
    vtk_celldata_points, vtk_celldata_cells = calc_vtk_points_cells(Val(ndims_), coordinates,
                                                                    levels, center_level_0,
                                                                    length_level_0, 1)
  end

  # Determine output file names
  base, _ = splitext(splitdir(filename)[2])
  vtk_filename = joinpath(output_directory, base)
  vtk_celldata_filename = vtk_filename * "_celldata"

  # Open VTK files
  verbose && println("| Building VTK grid...")
  if is_datafile
    @timeit "build VTK grid (node data)" vtk_nodedata = vtk_grid(vtk_filename, vtk_points,
                                                                 vtk_cells)
  else
    vtk_nodedata = nothing
  end
  @timeit "build VTK grid (cell data)" vtk_celldata = vtk_grid(vtk_celldata_filename,
                                                                vtk_celldata_points,
                                                                vtk_celldata_cells)

  return vtk_nodedata, vtk_celldata
end

# Create and return VTK grids that are ready to be filled with data (vti version)
function build_vtk_grids(::Val{:vti}, mesh, n_visnodes, verbose,
                         output_directory, is_datafile, filename)

  coordinates, levels, center_level_0, length_level_0 = extract_mesh_information(mesh)

  # Extract number of spatial dimensions
  ndims_ = size(coordinates, 1)

  # Prepare VTK points and cells for celldata file
  @timeit "prepare VTK cells" vtk_celldata_points, vtk_celldata_cells = calc_vtk_points_cells(
      Val(ndims_), coordinates, levels, center_level_0, length_level_0, 1)

  # Determine output file names
  base, _ = splitext(splitdir(filename)[2])
  vtk_filename = joinpath(output_directory, base)
  vtk_celldata_filename = vtk_filename * "_celldata"

  # Open VTK files
  verbose && println("| Building VTK grid...")
  if is_datafile
    # Determine level-wise resolution
    max_level = maximum(levels)
    resolution = n_visnodes * 2^max_level

    Nx = Ny = resolution + 1
    dx = dy = length_level_0/resolution
    origin = center_level_0 .- 1/2 * length_level_0
    spacing = [dx, dy]
    @timeit "build VTK grid (node data)" vtk_nodedata = vtk_grid(vtk_filename, Nx, Ny,
                                                        origin=origin,
                                                        spacing=spacing)
  else
    vtk_nodedata = nothing
  end
  @timeit "build VTK grid (cell data)" vtk_celldata = vtk_grid(vtk_celldata_filename,
                                                      vtk_celldata_points,
                                                      vtk_celldata_cells)

  return vtk_nodedata, vtk_celldata
end

# Create and return VTK grids that are ready to be filled with data
# (StructuredMesh/UnstructuredMesh2D/P4estMesh version)
function build_vtk_grids(::Val{:vtu},
                         mesh::Union{StructuredMesh, UnstructuredMesh2D, P4estMesh},
                         n_visnodes, verbose, output_directory, is_datafile, filename)

  @timeit "prepare coordinate information" node_coordinates = calc_node_coordinates(mesh, n_visnodes)

  # Calculate VTK points and cells
  verbose && println("| Preparing VTK cells...")
  if is_datafile
    @timeit "prepare VTK cells (node data)" begin
      vtk_points, vtk_cells = calc_vtk_points_cells(node_coordinates)
    end
  end

  # Prepare VTK points and cells for celldata file
  @timeit "prepare VTK cells (cell data)" begin
    vtk_celldata_points, vtk_celldata_cells = calc_vtk_points_cells(node_coordinates)
  end

  # Determine output file names
  base, _ = splitext(splitdir(filename)[2])
  vtk_filename = joinpath(output_directory, base)
  vtk_celldata_filename = vtk_filename * "_celldata"

  # Open VTK files
  verbose && println("| Building VTK grid...")
  if is_datafile
    @timeit "build VTK grid (node data)" vtk_nodedata = vtk_grid(vtk_filename, vtk_points,
                                                                 vtk_cells)
  else
    vtk_nodedata = nothing
  end
  @timeit "build VTK grid (cell data)" vtk_celldata = vtk_grid(vtk_celldata_filename,
                                                                vtk_celldata_points,
                                                                vtk_celldata_cells)

  return vtk_nodedata, vtk_celldata
end

function calc_node_coordinates(mesh::StructuredMesh, n_visnodes)
  # Extract number of spatial dimensions
  ndims_ = ndims(mesh)
  n_elements = prod(size(mesh))

  nodes = range(-1, 1, length=n_visnodes)
  f(args...; kwargs...) = Base.invokelatest(mesh.mapping, args...; kwargs...)

  node_coordinates = Array{Float64, ndims_+2}(undef, ndims_, ntuple(_ -> n_visnodes, ndims_)..., n_elements)

  return calc_node_coordinates!(node_coordinates, f, nodes, mesh)
end

function calc_node_coordinates(mesh::UnstructuredMesh2D, n_visnodes)
  # Extract number of spatial dimensions
  ndims_ = ndims(mesh)
  n_elements = length(mesh)

  # get the uniform nodes
  nodes = range(-1, 1, length=n_visnodes)

  # intialize the container for the node coordinates
  node_coordinates = Array{Float64, ndims_+2}(undef, ndims_, ntuple(_ -> n_visnodes, ndims_)..., n_elements)

  # work container for the corners of elements
  four_corners = zeros(eltype(mesh.corners), 4, 2)

  # loop through all elements and call the correct node coordinate constructor based on curved
  for element = 1:n_elements
    if mesh.element_is_curved[element]
      Trixi.calc_node_coordinates!(node_coordinates, element, nodes, view(mesh.surface_curves, :, element))
    else # straight sided element
      for i in 1:4, j in 1:2
        # pull the (x,y) values of these corners out of the global corners array
        four_corners[i, j] = mesh.corners[j, mesh.element_node_ids[i, element]]
      end
      Trixi.calc_node_coordinates!(node_coordinates, element, nodes, four_corners)
    end
  end

  return node_coordinates
end

function calc_node_coordinates(mesh::P4estMesh, n_visnodes)
  # Extract number of spatial dimensions
  ndims_ = ndims(mesh)

  nodes = range(-1, 1, length=n_visnodes)

  node_coordinates = Array{Float64, ndims_+2}(undef, ndims_,
                                              ntuple(_ -> n_visnodes, ndims_)...,
                                              Trixi.ncells(mesh))

  return Trixi.calc_node_coordinates!(node_coordinates, mesh, nodes)
end

# Calculation of the node coordinates for `StructuredMesh` in 2D
function calc_node_coordinates!(node_coordinates::AbstractArray{<:Any, 4}, f, nodes, mesh)
  linear_indices = LinearIndices(size(mesh))

  # Get cell length in reference mesh
  dx = 2 / size(mesh, 1)
  dy = 2 / size(mesh, 2)

  for cell_y in 1:size(mesh, 2), cell_x in 1:size(mesh, 1)
    element = linear_indices[cell_x, cell_y]

    # Calculate node coordinates of reference mesh
    cell_x_offset = -1 + (cell_x-1) * dx + dx/2
    cell_y_offset = -1 + (cell_y-1) * dy + dy/2

    for j in eachindex(nodes), i in eachindex(nodes)
      # node_coordinates are the mapped reference node_coordinates
      node_coordinates[:, i, j, element] .= f(cell_x_offset + dx/2 * nodes[i],
                                              cell_y_offset + dy/2 * nodes[j])
    end
  end

  return node_coordinates
end

# Calculation of the node coordinates for `StructuredMesh` in 3D
function calc_node_coordinates!(node_coordinates::AbstractArray{<:Any, 5}, f, nodes, mesh)
  linear_indices = LinearIndices(size(mesh))

  # Get cell length in reference mesh
  dx = 2 / size(mesh, 1)
  dy = 2 / size(mesh, 2)
  dz = 2 / size(mesh, 3)

  for cell_z in 1:size(mesh, 3), cell_y in 1:size(mesh, 2), cell_x in 1:size(mesh, 1)
    element = linear_indices[cell_x, cell_y, cell_z]

    # Calculate node coordinates of reference mesh
    cell_x_offset = -1 + (cell_x-1) * dx + dx/2
    cell_y_offset = -1 + (cell_y-1) * dy + dy/2
    cell_z_offset = -1 + (cell_z-1) * dz + dz/2

    for k in eachindex(nodes), j in eachindex(nodes), i in eachindex(nodes)
      # node_coordinates are the mapped reference node_coordinates
      node_coordinates[:, i, j, k, element] .= f(cell_x_offset + dx/2 * nodes[i],
                                                 cell_y_offset + dy/2 * nodes[j],
                                                 cell_z_offset + dz/2 * nodes[k])
    end
  end

  return node_coordinates
end

# Determine and return filenames for PVD fiels
function pvd_filenames(filenames, pvd, output_directory)
  # Determine pvd filename
  if !isnothing(pvd)
    # Use filename if given on command line
    filename = pvd

    # Strip of directory/extension
    filename, _ = splitext(splitdir(filename)[2])
  else
    filename = get_pvd_filename(filenames)

    # If filename is empty, it means we were not able to determine an
    # appropriate file thus the user has to supply one
    if filename == ""
      error("could not auto-detect PVD filename (input file names have no common prefix): " *
            "please provide a PVD filename name with the keyword argument `pvd=...`")
    end
  end

  # Get full filenames
  pvd_filename = joinpath(output_directory, filename)
  pvd_celldata_filename = pvd_filename * "_celldata"

  return pvd_filename, pvd_celldata_filename
end

# Determine filename for PVD file based on common name
function get_pvd_filename(filenames)
  filenames = getindex.(splitdir.(filenames), 2)
  bases = getindex.(splitext.(filenames), 1)
  pvd_filename = longest_common_prefix(bases)
  return pvd_filename
end

# Determine longest common prefix
function longest_common_prefix(strings)
  # Return early if array is empty
  if isempty(strings)
    return ""
  end

  # Count length of common prefix, by ensuring that all strings are long enough
  # and then comparing the next character
  len = 0
  while all(length.(strings) .> len) && all(getindex.(strings, len+1) .== strings[1][len+1])
    len +=1
  end

  return strings[1][1:len]
end

# Convert coordinates and level information to a list of points and VTK cells (2D version)
function calc_vtk_points_cells(::Val{2}, coordinates::AbstractMatrix{Float64},
                               levels::AbstractVector{Int},
                               center_level_0::AbstractVector{Float64},
                               length_level_0::Float64,
                               n_visnodes::Int=1)
  ndim = 2
  # Create point locator
  pl = PointLocator{2}(center_level_0, length_level_0, 1e-12)

  # Create arrays for points and cells
  n_elements = length(levels)
  points = Vector{SVector{2, Float64}}()
  vtk_cells = Vector{MeshCell}(undef, n_elements * n_visnodes^ndim)
  point_ids = Vector{Int}(undef, 2^ndim)

  # Reshape cell array for easy-peasy access
  reshaped = reshape(vtk_cells, n_visnodes, n_visnodes, n_elements)

  # Create VTK cell for each Trixi element
  for element_id in 1:n_elements
    # Extract cell values
    cell_x = coordinates[1, element_id]
    cell_y = coordinates[2, element_id]
    cell_dx = length_level_0 / 2^levels[element_id]

    # Adapt to visualization nodes for easy-to-understand loops
    dx = cell_dx / n_visnodes
    x_lowerleft = cell_x - cell_dx/2 - dx/2
    y_lowerleft = cell_y - cell_dx/2 - dx/2

    # Create cell for each visualization node
    for j = 1:n_visnodes
      for i = 1:n_visnodes
        # Determine x and y
        x = x_lowerleft + i * dx
        y = y_lowerleft + j * dx

        # Get point id for each vertex
        point_ids[1] = insert!(pl, points, (x - dx/2, y - dx/2))
        point_ids[2] = insert!(pl, points, (x + dx/2, y - dx/2))
        point_ids[3] = insert!(pl, points, (x - dx/2, y + dx/2))
        point_ids[4] = insert!(pl, points, (x + dx/2, y + dx/2))

        # Add cell
        reshaped[i, j, element_id] = MeshCell(VTKCellTypes.VTK_PIXEL, copy(point_ids))
      end
    end
  end

  # Convert array-of-points to two-dimensional array
  vtk_points = Matrix{Float64}(undef, ndim, length(points))
  for point_id in 1:length(points)
    vtk_points[1, point_id] = points[point_id][1]
    vtk_points[2, point_id] = points[point_id][2]
  end

  return vtk_points, vtk_cells
end

# Convert coordinates and level information to a list of points and VTK cells (3D version)
function calc_vtk_points_cells(::Val{3}, coordinates::AbstractMatrix{Float64},
                               levels::AbstractVector{Int},
                               center_level_0::AbstractVector{Float64},
                               length_level_0::Float64,
                               n_visnodes::Int=1)
  ndim = 3
  # Create point locator
  pl = PointLocator{3}(center_level_0, length_level_0, 1e-12)

  # Create arrays for points and cells
  n_elements = length(levels)
  points = Vector{SVector{3, Float64}}()
  vtk_cells = Vector{MeshCell}(undef, n_elements * n_visnodes^ndim)
  point_ids = Vector{Int}(undef, 2^ndim)

  # Reshape cell array for easy-peasy access
  reshaped = reshape(vtk_cells, n_visnodes, n_visnodes, n_visnodes, n_elements)

  # Create VTK cell for each Trixi element
  for element_id in 1:n_elements
    # Extract cell values
    cell_x = coordinates[1, element_id]
    cell_y = coordinates[2, element_id]
    cell_z = coordinates[3, element_id]
    cell_dx = length_level_0 / 2^levels[element_id]

    # Adapt to visualization nodes for easy-to-understand loops
    dx = cell_dx / n_visnodes
    x_lowerleft = cell_x - cell_dx/2 - dx/2
    y_lowerleft = cell_y - cell_dx/2 - dx/2
    z_lowerleft = cell_z - cell_dx/2 - dx/2

    # Create cell for each visualization node
    for k = 1:n_visnodes, j = 1:n_visnodes, i = 1:n_visnodes
      # Determine x and y
      x = x_lowerleft + i * dx
      y = y_lowerleft + j * dx
      z = z_lowerleft + k * dx

      # Get point id for each vertex
      point_ids[1] = insert!(pl, points, (x - dx/2, y - dx/2, z - dx/2))
      point_ids[2] = insert!(pl, points, (x + dx/2, y - dx/2, z - dx/2))
      point_ids[3] = insert!(pl, points, (x - dx/2, y + dx/2, z - dx/2))
      point_ids[4] = insert!(pl, points, (x + dx/2, y + dx/2, z - dx/2))
      point_ids[5] = insert!(pl, points, (x - dx/2, y - dx/2, z + dx/2))
      point_ids[6] = insert!(pl, points, (x + dx/2, y - dx/2, z + dx/2))
      point_ids[7] = insert!(pl, points, (x - dx/2, y + dx/2, z + dx/2))
      point_ids[8] = insert!(pl, points, (x + dx/2, y + dx/2, z + dx/2))

      # Add cell
      reshaped[i, j, k, element_id] = MeshCell(VTKCellTypes.VTK_VOXEL, copy(point_ids))
    end
  end

  # Convert array-of-points to two-dimensional array
  vtk_points = Matrix{Float64}(undef, ndim, length(points))
  for point_id in 1:length(points)
    vtk_points[1, point_id] = points[point_id][1]
    vtk_points[2, point_id] = points[point_id][2]
    vtk_points[3, point_id] = points[point_id][3]
  end

  return vtk_points, vtk_cells
end

# Convert coordinates and level information to a list of points and VTK cells for `StructuredMesh` (2D version)
function calc_vtk_points_cells(node_coordinates::AbstractArray{<:Any,4})
  n_elements = size(node_coordinates, 4)
  size_ = size(node_coordinates)
  n_points = prod(size_[2:end])
  # Linear indices to access points by node indices and element id
  linear_indices = LinearIndices(size_[2:end])

  # Use lagrange nodes as VTK points
  vtk_points = reshape(node_coordinates, (2, n_points))
  vtk_cells = Vector{MeshCell}(undef, n_elements)

  # Create cell for each element
  for element in 1:n_elements
    vertices = [linear_indices[1, 1, element],
                linear_indices[end, 1, element],
                linear_indices[end, end, element],
                linear_indices[1, end, element]]

    edges = vcat(linear_indices[2:end-1, 1, element],
                 linear_indices[end, 2:end-1, element],
                 linear_indices[2:end-1, end, element],
                 linear_indices[1, 2:end-1, element])

    faces = vec(linear_indices[2:end-1, 2:end-1, element])

    point_ids = vcat(vertices, edges, faces)
    vtk_cells[element] = MeshCell(VTKCellTypes.VTK_LAGRANGE_QUADRILATERAL, point_ids)
  end

  return vtk_points, vtk_cells
end

# Convert coordinates and level information to a list of points and VTK cells for `StructuredMesh` (3D version)
function calc_vtk_points_cells(node_coordinates::AbstractArray{<:Any,5})
  n_elements = size(node_coordinates, 5)
  size_ = size(node_coordinates)
  n_points = prod(size_[2:end])
  # Linear indices to access points by node indices and element id
  linear_indices = LinearIndices(size_[2:end])

  # Use lagrange nodes as VTK points
  vtk_points = reshape(node_coordinates, (3, n_points))
  vtk_cells = Vector{MeshCell}(undef, n_elements)

  # Create cell for each element
  for element in 1:n_elements
    vertices = [linear_indices[1, 1, 1, element],
                linear_indices[end, 1, 1, element],
                linear_indices[end, end, 1, element],
                linear_indices[1, end, 1, element],
                linear_indices[1, 1, end, element],
                linear_indices[end, 1, end, element],
                linear_indices[end, end, end, element],
                linear_indices[1, end, end, element]]

    # This order doesn't make any sense. This is completely different
    # from what is shown in
    # https://blog.kitware.com/wp-content/uploads/2018/09/Source_Issue_43.pdf
    # but this is the way it works.
    edges = vcat(linear_indices[2:end-1, 1, 1, element],
                 linear_indices[end, 2:end-1, 1, element],
                 linear_indices[2:end-1, end, 1, element],
                 linear_indices[1, 2:end-1, 1, element],
                 linear_indices[2:end-1, 1, end, element],
                 linear_indices[end, 2:end-1, end, element],
                 linear_indices[2:end-1, end, end, element],
                 linear_indices[1, 2:end-1, end, element],
                 linear_indices[1, 1, 2:end-1, element],
                 linear_indices[end, 1, 2:end-1, element],
                 linear_indices[1, end, 2:end-1, element],
                 linear_indices[end, end, 2:end-1, element])

    # See above
    faces = vcat(vec(linear_indices[1, 2:end-1, 2:end-1, element]),
                 vec(linear_indices[end, 2:end-1, 2:end-1, element]),
                 vec(linear_indices[2:end-1, 1, 2:end-1, element]),
                 vec(linear_indices[2:end-1, end, 2:end-1, element]),
                 vec(linear_indices[2:end-1, 2:end-1, 1, element]),
                 vec(linear_indices[2:end-1, 2:end-1, end, element]))

    volume = vec(linear_indices[2:end-1, 2:end-1, 2:end-1, element])

    point_ids = vcat(vertices, edges, faces, volume)
    vtk_cells[element] = MeshCell(VTKCellTypes.VTK_LAGRANGE_HEXAHEDRON, point_ids)
  end

  return vtk_points, vtk_cells
end


text:
function MH(data::Data, prior::Prior, candiate::Candidate, draws = 1_000)

    x, n = data.x, data.n
    J = length(x)

    # Prepare posterior
    post = Chain(zeros(draws), zeros(draws), zeros(draws, N), zeros(draws, 2))

    # Starting values (random in the unit area)
    post.μ[1] = rand()
    post.σ[1] = rand()

    post.acceptance[1, :] = 1

    # Old posterior
    lp = logposterior(prior::Prior, post.μ[1], post.σ[1], x, n)

    for i in 2:draws

        ### MU ###

        # Candidate value for μ
        μ = rand(Normal(post.μ[i - 1], 0.5))

        # Candidate log-posterior
        lpc = logposterior(prior::Prior, μ, post.σ[i - 1], x, n)

        if (lpc - lp) > log(rand())
            post.μ[i] = μ
            post.acceptance[:, 1] = 1
        else
            post.μ[i] = post.μ[i - 1]
        end

        ### SIGMA ###

        # Candidate value for σ
        σ = rand(Normal(post.σ[i - 1], 0.5))

        # Candidate log-posterior
        lpc = logposterior(prior::Prior, post.μ[i - 1], σ, x, n)

        if (lpc - lp) > log(rand())
            post.σ[i] = σ
            post.acceptance[:, 2] = 1
        else
            post.σ[i] = post.σ[i - 1]
        end

        ### THETA ###
        for j in 1:J
            θ[i, j] = rand(Beta((1 - post.σ[i])/post.σ[i] * post.μ[i] + x[j], (1 - post.σ[i])/post.σ[i] * (1 - post.μ[i]) + n[j] - x[j]))
        end

    end
    return post
end

function logprior(prior::Beta, value::AbstractFloat)
    μ, ϕ = mu_and_phi(prior.a, prior.b)
    (μ - 1)*log(value) + (μ - 1)*log(1 - value)
end

function loglikelihood(μ, σ, x, n)
    l = logpdf(Beta(μ * (1 - σ)/σ + x), (1 - μ)*(1 - σ)/σ + n - x)
    l = l - logpdf(Beta(μ * (1 - σ)/σ, (1 - μ)*(1 - σ)/σ)) * length(x)
    return l
end

function logposterior(prior::Prior, μ, σ, x, n)
    l = loglikelihood(μ, σ, x, n)
    p = logprior(prior.μ, μ) + logprior(prior.σ, σ)
    return l+p
end

function alpha_and_beta(μ, ϕ)
    α = μ * (1 - ϕ)/ϕ
    β = (1 - μ) * (1 - ϕ)/ϕ
    return (α = α, β = β)
end

function mu_and_phi(α, β)
    μ = α / (α + β)
    ϕ = 1 / (1 + α + β)
end


text:
# source: https://github.com/FluxML/Flux.jl/blob/master/src/optimise/optimisers.jl

abstract type AbstractOptimiser end

const ϵ = 1e-8

"""
    Descent(η = 0.1)
Classic gradient descent optimiser with learning rate `η`.
For each parameter `p` and its gradient `δp`, this runs `p -= η*δp`
# Parameters
- Learning rate (`η`): Amount by which gradients are discounted before updating
                       the weights.
"""
mutable struct Descent <: AbstractOptimiser
  eta::Float64
end

Descent() = Descent(0.1)

function apply!(o::Descent, x, Δ)
  Δ .*= o.eta
end

"""
    Momentum(η = 0.01, ρ = 0.9)
Gradient descent optimizer with learning rate `η` and momentum `ρ`.
# Parameters
- Learning rate (`η`): Amount by which gradients are discounted before updating
                       the weights.
- Momentum (`ρ`): Controls the acceleration of gradient descent in the
                  prominent direction, in effect dampening oscillations.
# Examples
```julia
opt = Momentum()
opt = Momentum(0.01, 0.99)
```
"""
mutable struct Momentum <: AbstractOptimiser
  eta::Float64
  rho::Float64
  velocity::IdDict
end

Momentum(η = 0.01, ρ = 0.9) = Momentum(η, ρ, IdDict())

function apply!(o::Momentum, x, Δ)
  η, ρ = o.eta, o.rho
  v = get!(() -> zero(x), o.velocity, x)::typeof(x)
  @. v = ρ * v - η * Δ
  @. Δ = -v
end

"""
    Nesterov(η = 0.001, ρ = 0.9)
Gradient descent optimizer with learning rate `η` and Nesterov momentum `ρ`.
# Parameters
- Learning rate (`η`): Amount by which gradients are discounted before updating
                       the weights.
- Nesterov momentum (`ρ`): Controls the acceleration of gradient descent in the
                           prominent direction, in effect dampening oscillations.
# Examples
```julia
opt = Nesterov()
opt = Nesterov(0.003, 0.95)
```
"""
mutable struct Nesterov <: AbstractOptimiser
  eta::Float64
  rho::Float64
  velocity::IdDict
end

Nesterov(η = 0.001, ρ = 0.9) = Nesterov(η, ρ, IdDict())

function apply!(o::Nesterov, x, Δ)
  η, ρ = o.eta, o.rho
  v = get!(() -> zero(x), o.velocity, x)::typeof(x)
  d = @. ρ^2 * v - (1+ρ) * η * Δ
  @. v = ρ*v - η*Δ
  @. Δ = -d
end

"""
    RMSProp(η = 0.001, ρ = 0.9)
Optimizer using the
[RMSProp](https://www.cs.toronto.edu/~tijmen/csc321/slides/lecture_slides_lec6.pdf)
algorithm. Often a good choice for recurrent networks. Parameters other than learning rate
generally don't need tuning.
# Parameters
- Learning rate (`η`): Amount by which gradients are discounted before updating
                       the weights.
- Momentum (`ρ`): Controls the acceleration of gradient descent in the
                  prominent direction, in effect dampening oscillations.
# Examples
```julia
opt = RMSProp()
opt = RMSProp(0.002, 0.95)
```
"""
mutable struct RMSProp <: AbstractOptimiser
  eta::Float64
  rho::Float64
  acc::IdDict
end

RMSProp(η = 0.001, ρ = 0.9) = RMSProp(η, ρ, IdDict())

function apply!(o::RMSProp, x, Δ)
  η, ρ = o.eta, o.rho
  acc = get!(() -> zero(x), o.acc, x)::typeof(x)
  @. acc = ρ * acc + (1 - ρ) * Δ^2
  @. Δ *= η / (√acc + ϵ)
end

"""
    ADAM(η = 0.001, β::Tuple = (0.9, 0.999))
[ADAM](https://arxiv.org/abs/1412.6980) optimiser.
# Parameters
- Learning rate (`η`): Amount by which gradients are discounted before updating
                       the weights.
- Decay of momentums (`β::Tuple`): Exponential decay for the first (β1) and the
                                   second (β2) momentum estimate.
# Examples
```julia
opt = ADAM()
opt = ADAM(0.001, (0.9, 0.8))
```
"""
mutable struct ADAM <: AbstractOptimiser
  eta::Float64
  beta::Tuple{Float64,Float64}
  state::IdDict
end

ADAM(η = 0.001, β = (0.9, 0.999)) = ADAM(η, β, IdDict())

function apply!(o::ADAM, x, Δ)
  η, β = o.eta, o.beta

  mt, vt, βp = get!(o.state, x) do
      (zero(x), zero(x), Float64[β[1], β[2]])
  end :: Tuple{typeof(x),typeof(x),Vector{Float64}}

  @. mt = β[1] * mt + (1 - β[1]) * Δ
  @. vt = β[2] * vt + (1 - β[2]) * Δ^2
  @. Δ =  mt / (1 - βp[1]) / (√(vt / (1 - βp[2])) + ϵ) * η
  βp .= βp .* β

  return Δ
end

"""
    RADAM(η = 0.001, β::Tuple = (0.9, 0.999))
[Rectified ADAM](https://arxiv.org/abs/1908.03265) optimizer.
# Parameters
- Learning rate (`η`): Amount by which gradients are discounted before updating
                       the weights.
- Decay of momentums (`β::Tuple`): Exponential decay for the first (β1) and the
                                   second (β2) momentum estimate.
# Examples
```julia
opt = RADAM()
opt = RADAM(0.001, (0.9, 0.8))
```
"""
mutable struct RADAM <: AbstractOptimiser
  eta::Float64
  beta::Tuple{Float64,Float64}
  state::IdDict
end

RADAM(η = 0.001, β = (0.9, 0.999)) = RADAM(η, β, IdDict())

function apply!(o::RADAM, x, Δ)
  η, β = o.eta, o.beta
  ρ∞ = 2/(1-β[2])-1

  mt, vt, βp, t = get!(o.state, x) do
      (zero(x), zero(x), Float64[β[1], β[2]], Ref(1))
  end :: Tuple{typeof(x),typeof(x),Vector{Float64},Ref{Int}}

  @. mt = β[1] * mt + (1 - β[1]) * Δ
  @. vt = β[2] * vt + (1 - β[2]) * Δ^2
  ρ = ρ∞ - 2t[] * βp[2] / (1 - βp[2])
  if ρ > 4
    r = sqrt((ρ-4)*(ρ-2)*ρ∞/((ρ∞-4)*(ρ∞-2)*ρ))
    @. Δ =  mt / (1 - βp[1]) / (√(vt / (1 - βp[2])) + ϵ) * η * r
  else
    @. Δ =  mt / (1 - βp[1]) * η
  end
  βp .= βp .* β
  t[] += 1

  return Δ
end

"""
    AdaMax(η = 0.001, β::Tuple = (0.9, 0.999))
[AdaMax](https://arxiv.org/abs/1412.6980) is a variant of ADAM based on the ∞-norm.
# Parameters
- Learning rate (`η`): Amount by which gradients are discounted before updating
                       the weights.
- Decay of momentums (`β::Tuple`): Exponential decay for the first (β1) and the
                                   second (β2) momentum estimate.
# Examples
```julia
opt = AdaMax()
opt = AdaMax(0.001, (0.9, 0.995))
```
"""
mutable struct AdaMax <: AbstractOptimiser
  eta::Float64
  beta::Tuple{Float64,Float64}
  state::IdDict
end

AdaMax(η = 0.001, β = (0.9, 0.999)) = AdaMax(η, β, IdDict())

function apply!(o::AdaMax, x, Δ)
  η, β = o.eta, o.beta

  mt, ut, βp = get!(o.state, x) do
      (zero(x), zero(x), Float64[β[1], β[2]])
  end :: Tuple{typeof(x),typeof(x),Vector{Float64}}

  @. mt = β[1] * mt + (1 - β[1]) * Δ
  @. ut = max(β[2] * ut, abs(Δ))
  @. Δ = (η/(1 - βp[1])) * mt/(ut + ϵ)
  βp .= βp .* β

  return Δ
end

"""
    OADAM(η = 0.0001, β::Tuple = (0.5, 0.9))
[OADAM](https://arxiv.org/abs/1711.00141) (Optimistic ADAM)
is a variant of ADAM adding an "optimistic" term suitable for adversarial training.
# Parameters
- Learning rate (`η`): Amount by which gradients are discounted before updating
                       the weights.
- Decay of momentums (`β::Tuple`): Exponential decay for the first (β1) and the
                                   second (β2) momentum estimate.
# Examples
```julia
opt = OADAM()
opt = OADAM(0.001, (0.9, 0.995))
```
"""
mutable struct OADAM <: AbstractOptimiser
  eta::Float64
  beta::Tuple{Float64,Float64}
  state::IdDict
end

OADAM(η = 0.001, β = (0.5, 0.9)) = OADAM(η, β, IdDict())

function apply!(o::OADAM, x, Δ)
  η, β = o.eta, o.beta

  mt, vt, Δ_, βp = get!(o.state, x) do
      (zero(x), zero(x), zero(x), Float64[β[1], β[2]])
  end :: Tuple{typeof(x),typeof(x),typeof(x),Vector{Float64}}

  @. mt = β[1] * mt + (1 - β[1]) * Δ
  @. vt = β[2] * vt + (1 - β[2]) * Δ^2
  @. Δ = -Δ_
  @. Δ_ = η * mt / (1 - βp[1]) / (√(vt / (1 - βp[2])) + ϵ)
  @. Δ += 2Δ_
  βp .= βp .* β

  return Δ
end

"""
    ADAGrad(η = 0.1)
[ADAGrad](http://www.jmlr.org/papers/volume12/duchi11a/duchi11a.pdf) optimizer. It has
parameter specific learning rates based on how frequently it is updated.
Parameters don't need tuning.
# Parameters
- Learning rate (`η`): Amount by which gradients are discounted before updating
                       the weights.
# Examples
```julia
opt = ADAGrad()
opt = ADAGrad(0.001)
```
"""
mutable struct ADAGrad <: AbstractOptimiser
  eta::Float64
  acc::IdDict
end

ADAGrad(η = 0.1) = ADAGrad(η, IdDict())

function apply!(o::ADAGrad, x, Δ)
  η = o.eta
  acc = get!(() -> fill!(similar(x), ϵ), o.acc, x)::typeof(x)
  @. acc += Δ^2
  @. Δ *= η / (√acc + ϵ)
end

"""
    ADADelta(ρ = 0.9)
[ADADelta](https://arxiv.org/abs/1212.5701) is a version of ADAGrad adapting its learning
rate based on a window of past gradient updates.
Parameters don't need tuning.
# Parameters
- Rho (`ρ`): Factor by which the gradient is decayed at each time step.
# Examples
```julia
opt = ADADelta()
opt = ADADelta(0.89)
```
"""
mutable struct ADADelta <: AbstractOptimiser
  rho::Float64
  state::IdDict
end

ADADelta(ρ = 0.9) = ADADelta(ρ, IdDict())

function apply!(o::ADADelta, x, Δ)
  ρ = o.rho
  acc, Δacc = get!(() -> (zero(x), zero(x)), o.state, x)::NTuple{2,typeof(x)}
  @. acc = ρ * acc + (1 - ρ) * Δ^2
  # DON'T remove epsilon from numerator
  # or even out of the square roots
  @. Δ *= √(Δacc + ϵ) / √(acc + ϵ)
  @. Δacc = ρ * Δacc + (1 - ρ) * Δ^2
  return Δ
end

"""
    AMSGrad(η = 0.001, β::Tuple = (0.9, 0.999))
The [AMSGrad](https://openreview.net/forum?id=ryQu7f-RZ) version of the ADAM
optimiser. Parameters don't need tuning.
# Parameters
- Learning rate (`η`): Amount by which gradients are discounted before updating
                       the weights.
- Decay of momentums (`β::Tuple`): Exponential decay for the first (β1) and the
                                   second (β2) momentum estimate.
# Examples
```julia
opt = AMSGrad()
opt = AMSGrad(0.001, (0.89, 0.995))
```
"""
mutable struct AMSGrad <: AbstractOptimiser
  eta::Float64
  beta::Tuple{Float64, Float64}
  state::IdDict
end

AMSGrad(η = 0.001, β = (0.9, 0.999)) = AMSGrad(η, β, IdDict())

function apply!(o::AMSGrad, x, Δ)
  η, β = o.eta, o.beta

  mt, vt, v̂t = get!(o.state, x) do
    (fill!(similar(x), ϵ), fill!(similar(x), ϵ), fill!(similar(x), ϵ))
  end :: NTuple{3,typeof(x)}

  @. mt = β[1] * mt + (1 - β[1]) * Δ
  @. vt = β[2] * vt + (1 - β[2]) * Δ ^ 2
  @. v̂t = max(v̂t, vt)
  @. Δ = η * mt / (√v̂t + ϵ)
end

"""
    NADAM(η = 0.001, β::Tuple = (0.9, 0.999))
[NADAM](https://openreview.net/forum?id=OM0jvwB8jIp57ZJjtNEZ) is a Nesterov variant of ADAM.
Parameters don't need tuning.
# Parameters
- Learning rate (`η`): Amount by which gradients are discounted before updating
                       the weights.
- Decay of momentums (`β::Tuple`): Exponential decay for the first (β1) and the
                                   second (β2) momentum estimate.
# Examples
```julia
opt = NADAM()
opt = NADAM(0.002, (0.89, 0.995))
```
"""
mutable struct NADAM <: AbstractOptimiser
  eta::Float64
  beta::Tuple{Float64, Float64}
  state::IdDict
end

NADAM(η = 0.001, β = (0.9, 0.999)) = NADAM(η, β, IdDict())

function apply!(o::NADAM, x, Δ)
  η, β = o.eta, o.beta

  mt, vt, βp = get!(o.state, x) do
    (zero(x), zero(x), Float64[o.beta[1], o.beta[2]])
  end :: Tuple{typeof(x),typeof(x),Vector{Float64}}
  β1p, β2p = βp

  @. mt = β[1] * mt + (1 - β[1]) * Δ
  @. vt = β[2] * vt + (1 - β[2]) * Δ^2
  @. Δ = (β[1] * mt / (1 - β[1] * β1p) + (1 - β[1]) * Δ / (1 - β1p)) / (√(vt * β[2] / (1 - β2p)) + ϵ) * η
  βp .= βp .* β

  return Δ
end

"""
    ADAMW(η = 0.001, β::Tuple = (0.9, 0.999), decay = 0)
[ADAMW](https://arxiv.org/abs/1711.05101) is a variant of ADAM fixing (as in repairing) its
weight decay regularization.
# Parameters
- Learning rate (`η`): Amount by which gradients are discounted before updating
                       the weights.
- Decay of momentums (`β::Tuple`): Exponential decay for the first (β1) and the
                                   second (β2) momentum estimate.
- `decay`: Decay applied to weights during optimisation.
# Examples
```julia
opt = ADAMW()
opt = ADAMW(0.001, (0.89, 0.995), 0.1)
```
"""
ADAMW(η = 0.001, β = (0.9, 0.999), decay = 0) =
  Optimiser(ADAM(1, β), WeightDecay(decay), Descent(η))

"""
    AdaBelief(η = 0.001, β::Tuple = (0.9, 0.999))
The [AdaBelief](https://arxiv.org/abs/2010.07468) optimiser is a variant of the well-known
ADAM optimiser.
# Parameters
- Learning rate (`η`): Amount by which gradients are discounted before updating
                       the weights.
- Decay of momentums (`β::Tuple`): Exponential decay for the first (β1) and the
                                   second (β2) momentum estimate.
# Examples
```julia
opt = AdaBelief()
opt = AdaBelief(0.001, (0.9, 0.8))
```
"""
mutable struct AdaBelief
  eta::Float64
  beta::Tuple{Float64,Float64}
  state::IdDict
end

AdaBelief(η = 0.001, β = (0.9, 0.999)) = AdaBelief(η, β, IdDict())

function apply!(o::AdaBelief, x, Δ)
  η, β = o.eta, o.beta
  mt, st = get!(() -> (zero(x), zero(x)), o.state, x)::NTuple{2,typeof(x)}
  @. mt = β[1] * mt + (1 - β[1]) * Δ
  @. st = β[2] * st + (1 - β[2]) * (Δ - mt)^2
  @. Δ =  η * mt / (√(st) + ϵ)
  return Δ
end


text:
struct BCache{T<:Union{UInt,Nothing}}
    p::Ptr{Cvoid}
    i::T
end