text:
using CUDA

using Oceananigans.Fields: AbstractField, compute_at!
using Oceananigans.LagrangianParticleTracking: LagrangianParticles
using Oceananigans.Diagnostics: slice_parent

fetch_output(output, model, field_slicer) = output(model)

# Needed to support `fetch_output` with `model::Nothing`.
time(model) = model.clock.time
time(::Nothing) = nothing

function fetch_output(field::AbstractField, model, field_slicer)
    compute_at!(field, time(model))
    return slice_parent(field_slicer, field)
end

function fetch_output(lagrangian_particles::LagrangianParticles, model, field_slicer)
    particle_properties = lagrangian_particles.properties
    names = propertynames(particle_properties)
    return NamedTuple{names}([getproperty(particle_properties, name) for name in names])
end

convert_output(output, writer) = output
convert_output(output::AbstractArray, writer) = CUDA.@allowscalar writer.array_type(output)

# Need to broadcast manually because of https://github.com/JuliaLang/julia/issues/30836
convert_output(outputs::NamedTuple, writer) =
    NamedTuple{keys(outputs)}(writer.array_type.(values(outputs)))

fetch_and_convert_output(output, model, writer) =
    convert_output(fetch_output(output, model, writer.field_slicer), writer)


text:
d0=d1=d2=0
"""
We demonstrate core capabilities of EquivariantOperators.jl: particle mesh, finite differences, and machine learning. In a particle mesh context, we place 2 +1 point atoms in a grid creating a scalar field. We then do a finite difference calculation of the electric field and potential. Switching gears to machine learning, we train our equivariant neural network to learn this transformation from the charge distribution (scalar field) to the electric potential (scalar field) and field (vector field), essentially learning the Green's function solution to Poisson's Equation. Finally, we treat our point atoms as proton nuclei and train another neural network to predict the ground state electronic density of H2.
"""

repo = ".."
include("$repo/src/neural_networks.jl")
include("$repo/src/atoms.jl")
# include("../src/plotutils.jl")
# include("$repo/src/EquivariantOperators.jl")
# using .EquivariantOperators

using FileIO
using BSON: @save,@load
using IterTools
using Functors

using Plots
using Random
using Optim
using Flux, CUDA

Random.seed!(1)
# ENV["CUDA_VISIBLE_DEVICES"] = -1
mygpu = identity

# name = "AAACQNRQYWKRAR-SCLBCKFNSA-N"
# name = "benzene"
# name = "h2"
# name="dsgdb9nsd_000001"
# case = load("..\\density-prediction\\data\\$name.jld2", "cases")[1]

folder = "..\\density-prediction\\data"
names = [split(x, ".")[1] for x in readdir(folder) if x == "h2.jld2"]
systems = [load("$folder\\$name.jld2", "cases") for name in names]
# file = "$folder\\$name.jld2"
# case = load(file,"case")

data = []
grid=nothing
for (cases, name) in zip(systems, names)
    for case in cases
        @unpack resolution,
        density,
        core_density,
        charges,
        core_charges,
        # atoms,
        # core_atoms,
        positions = case
        atoms=charges

        origin = ones(3)
         cell = resolution*Matrix(I,3,3)
      sz = size(density)
global        grid = Grid(cell, sz; origin)

        ρe = cat(density, dims = 4)
        # ρe *=total/(sum(ρe*dV))

        # x = ρp |> mygpu
        # y = ρe |> mygpu

        push!(data, Frame(atoms,positions,grid;ρe))
    end
end

@unpack ρp,grid,ρe,ops=data[1]
@unpack dv,cell=grid
@unpack ϕ,E=ops
ϕp = ϕ(ρp)
ϕe = ϕ(ρe)
# E=dv*sum(ρe.*ϕp)
# E1=dv*sum(ρe.*ϕe)
##
using GLMakie
GLMakie.inline!(false)
# fig = volume(E.kernel[:, :, :, 1])
# fig = volume(ϕ.kernel[:, :, :, 1])
fig = volume(d4[:, :, :, 1])
display(fig)
#
# ##
# @unpack dx=grid
# rmax=3.
# m=DensityPredictor(dx,rmax)
# ps=Flux.params(m.L,m.D,m.D1)
# # ps=Flux.params(m)
# yhat =0
# function loss(x)
#     y=x.ρe
#     global yhat=m(x)
#     sumy = sum(y)
#     @show l = nae(yhat, y; sumy)
#     l
# end
#
# data0 = data
# data = vcat([fill(x, 5) for x in data]...)
# loss(data[1])
# opt = ADAM(0.1)
#
# for i = 1:10
#     Flux.train!(loss, ps, data, opt)
# end
# @save "m.bson" m
#
# ##
# @load "m.bson" m
# function fg!(F, G, x)
#     ρe_hat, ρp = predict_density(x, atoms, grid)
#     ρhat = -ρe_hat + ρp
#     @show forces_hat = calcforces(x, atoms, ρhat, grid)
#     if G != nothing
#         G .= -forces_hat
#     end
#     norm(forces_hat) / sum(atoms)
# end;
#
# x0 = positions
#
# xres = optimize(
#     Optim.only_fg!(fg!),
#     x0,
#     LBFGS(),
#     Optim.Options(show_trace = true, f_tol = tol),
# )
# xmin = Optim.minimizer(xres)
# dmin = norm(xmin[:, 1] - xmin[:, 2])
# @printf "\nOptimal bond length for Ecut=%.2f: %.3f Bohr\n" Ecut dmin
#
# ##
# using GLMakie
# include("$repo/src/plotutils.jl")
#
# # volume(x[:, :, :, 1])
# # volume(y[:, :, :, 1])
# GLMakie.inline!(false)
# # fig=GLMakie.contour(10yhat[:, :, :, 1])
# # fig=volume(10yhat[:, :, :, 1],algorithm = :iso,)
# # fig=volume(10y[:, :, :, 1],algorithm = :iso,)
# # fig=volume(10yhat[:, :, :, 1],algorithm = :absorption,)
# fig = volume(yhat[:, :, :, 1])
# display(fig)
# # volume(10z[:, :, :, 3])
#
# # vis(L)

ext: jl


text:
module BenchmarkPlots
using RecipesBase
using BenchmarkTools: Trial, BenchmarkGroup

@recipe function f(::Type{Trial}, t::Trial)
    seriestype --> :violin
    legend --> false
    yguide --> "t / ns"
    xticks --> false
    t.times
end

@recipe function f(g::BenchmarkGroup, keys=keys(g))
    legend --> false
    yguide --> "t / ns"
    for k in keys
        @series begin
            label --> string(k)
            xticks --> true
            [string(k)], g[k]
        end
    end
end

end


text:
@testset "update CIR Tests" begin
    tmp = DualOptimalFiltering.update_CIR_params([0.5, 0.5], 1., 1.5, 1., [0, 1], [5])
    @test tmp[1] == 2.5
    @test tmp[2] == [5, 6]
    @test tmp[3][1] ≈ 0.131579 atol=10.0^(-5)
    @test tmp[3][2] ≈ 0.868421 atol=10.0^(-5)
end;

@testset "predict CIR Tests" begin
    tmp = DualOptimalFiltering.predict_CIR_params([1.], 1., 2., 1., 1., [5], 1.)
    @test tmp[1] ≈ 1.072578883495754 atol=10.0^(-5)
    @test tmp[2] == 0:5
    [@test tmp[3][k] ≈ [0.686096, 0.268465, 0.0420196, 0.0032884, 0.000128673, 2.01396e-6][k] atol=10.0^(-5) for k in 1:6]
end

@testset "CIR filtering tests" begin
    Random.seed!(0)
    X = DualOptimalFiltering.generate_CIR_trajectory(range(0, stop = 2, length = 20), 3, 3., 0.5, 1);
    Y = map(λ -> rand(Poisson(λ),10), X);
    data = Dict(zip(range(0, stop = 2, length = 20), Y))
    Λ_of_t, wms_of_t, θ_of_t = DualOptimalFiltering.filter_CIR(3., 0.5, 1.,1.,data);
    [@test isreal(k) for k in Λ_of_t |> keys]
    [@test isinteger(sum(k)) for k in Λ_of_t |> values]
    [@test isreal(sum(k)) for k in wms_of_t |> values]
    [@test isreal(k) for k in θ_of_t |> values]
end;

max_stars_repo_name: konkam/DualOptimalFiltering

max_issues_repo_name: konkam/DualOptimalFiltering

max_forks_repo_name: konkam/DualOptimalFiltering


text:
using DemoPackageJCM
using Test

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


text:
"""
    zs = smoothconv(z, nas)

Smoothen field `z(ns,ns)` with a circular kernel of diameter `ns/nas` using convolution.
Takes into account missing values.
"""
function smoothconv(zi, nas)

    iinan = findall(isnan.(zi))
    iinotnan = findall(.~isnan.(zi))
    zi[iinan] .= 0.

    nss = size(zi)
    ns = nss[1]
    sdim = div(div(ns, nas), 2) # the smoothing sigma is half large scale pixel wide

    mask=zeros(ns, ns)
    for i = 1:ns
        for j = 1:ns
            kx = i-1
            ky = j-1
            if (i>ns/2+1)
                kx = i-ns-1
            end
            if (j>ns/2+1)
                ky = j-ns-1
            end
            r2 = kx*kx+ky*ky
            mask[i, j] = exp(-(r2/(sdim*sdim))/2)
        end
    end
    fm = fft(mask)
    zf = real(ifft(fm.*fft(zi)))/sum(mask)
    if length(iinan)>0
        zi1 = deepcopy(zi)
        zi1[iinotnan] .= 1.0
        zf = zf./(real(ifft(fm.*fft(zi1)))/sum(mask))
    end
    zf[iinan] .= NaN
    return zf
end


text:
function render(cur_size, cur_data, cur_layout, cur_config)
  if is_repl
    render_unicode(cur_size, cur_data, cur_layout, cur_config)
  else
    render_html(cur_size, cur_data, cur_layout, cur_config)
  end
end

# render_unicode in repl.jl

function render_html(cur_size, cur_data, cur_layout, cur_config)
  cur_plot_div = "js-plot-" * string(UUIDs.uuid4())

  cur_html = """

  """

  cur_html *= _render_html_script(
    cur_plot_div, cur_data, cur_layout, cur_config
  )

  HTML(cur_html)
end

function _render_html_script(cur_plot_div, cur_data, cur_layout, cur_config)
  """

  """
end


text:
"""
    abstract type AbstractCosmoGeometry end

Abstract type for (multi-dimensional) geometry volumes, particularly for filtering with [`filter`](@ref).

Any subtypes of `AbstractCosmoGeometry` have to implement the following methods:
- [`CosmoParticles.geometry_enclosing_corners`](@ref)
- [`CosmoParticles.geometry_enclosing_center`](@ref)
- [`CosmoParticles.mask_in!`](@ref)
"""
abstract type AbstractCosmoGeometry end

"""
    geometry_enclosing_corners(geo::AbstractCosmoGeometry)

Return the lower left and upper right corners of the enclosing box of the geometry as a tuple of vectors.

The enclosing box is not necessarily the tightest fitting box.

This is not exported.
"""
function geometry_enclosing_corners(::AbstractCosmoGeometry) end

"""
    geometry_enclosing_center(geo::AbstractCosmoGeometry)

Return the center of the enclosing box of the geometry as a vector.

This is not exported.
"""
function geometry_enclosing_center(::AbstractCosmoGeometry) end

@doc raw"""
    mask_in!(mask::BitVector, pos::AbstractMatrix{<:Number}, geo::AbstractCosmoGeometry)

Return the `BitVector` mask of the positions (``\mathrm{dims} × N``) located within the geometry.

This is not exported.
"""
function mask_in!(::BitVector, ::AbstractMatrix{<:Number}, ::AbstractCosmoGeometry) end

@doc raw"""
    mask_in(pos::AbstractMatrix{<:Number}, geo::AbstractCosmoGeometry)

Return the `BitArray` mask of the positions (``\mathrm{dims} × N``) located within the geometry.

Calls [`mask_in!`](@ref) by default.

This is not exported.
"""
function mask_in(pos::AbstractMatrix{<:Number}, geo::AbstractCosmoGeometry)
    mask = BitVector(undef, size(pos, 2))
    return mask_in!(mask, pos, geo)
end

"""
    struct CosmoHyperrectangle{T,N} <: AbstractCosmoGeometry where {T<:Number}
        lowerleft::Vector{T}
        upperright::Vector{T}
    end

`N`-dimensional hyperrectangle aligned with the coordinate system axes given by its lower left and upper right corners.

The dimensions of space are given by `N`.
"""
struct CosmoHyperrectangle{T,N} <: AbstractCosmoGeometry where {T<:Number}
    lowerleft::Vector{T}
    upperright::Vector{T}

    function CosmoHyperrectangle{T,N}(
        lowerleft::AbstractVector{<:Number},
        upperright::AbstractVector{<:Number},
    ) where {T<:Number,N}
        @assert length(lowerleft) == length(upperright) == N
        new{T,N}(convert(Vector{T}, lowerleft), convert(Vector{T}, upperright))
    end
end

"""
    CosmoHyperrectangle(center::AbstractVector{<:Number}, sidelengths::Number...)

Return a `CosmoHyperrectangle` centered around `center`, with the given `sidelengths`.

The side lengths are given in order of the axes, i.e., x, y or x, y, z for 2D and 3D, respectively.
"""
function CosmoHyperrectangle(center::AbstractVector{<:Number}, sidelengths::Number...)
    halflengths = [sidelengths...]
    halflengths *= 1 // 2
    return CosmoHyperrectangle(center .- halflengths, center .+ halflengths)
end

function Base.:(==)(r1::CosmoHyperrectangle, r2::CosmoHyperrectangle)
    return r1.lowerleft == r2.lowerleft && r1.upperright == r2.upperright
end

"""
    CosmoCuboid{T} = CosmoHyperrectangle{T,3}

Alias for a 3D [`CosmoHyperrectangle`](@ref).
"""
const CosmoCuboid{T} = CosmoHyperrectangle{T,3} where {T<:Number}

function CosmoCuboid(center::AbstractVector{<:Number}, x::Number, y::Number, z::Number)
    halflengths = [x, y, z]
    halflengths *= 1 // 2
    return CosmoCuboid(center .- halflengths, center .+ halflengths)
end

"""
    CosmoRectangle{T} = CosmoHyperrectangle{T,2}

Alias for a 2D [`CosmoHyperrectangle`](@ref).
"""
const CosmoRectangle{T} = CosmoHyperrectangle{T,2} where {T<:Number}

function CosmoRectangle(center::AbstractVector{<:Number}, x::Number, y::Number)
    halflengths = [x, y]
    halflengths *= 1 // 2
    return CosmoRectangle(center .- halflengths, center .+ halflengths)
end

"""
    CosmoHypercube(center::AbstractVector{<:Number}, radius::Number)

Return a cubic `CosmoCuboid` centered around `center`, with equal sidelengths.
"""
function CosmoHypercube(center::AbstractVector{<:Number}, radius::Number)
    return CosmoHyperrectangle(center .- radius, center .+ radius)
end

"""
    CosmoCube(center::AbstractVector{<:Number}, radius::Number)

Return a cubic `CosmoCuboid` centered around `center`, with equal sidelengths.
"""
function CosmoCube(center::AbstractVector{<:Number}, radius::Number)
    return CosmoCuboid(center .- radius, center .+ radius)
end

"""
    CosmoSquare(center::AbstractVector{<:Number}, radius::Number)

Return a square `CosmoRectangle` centered around `center`, with equal sidelengths.
"""
function CosmoSquare(center::AbstractVector{<:Number}, radius::Number)
    return CosmoRectangle(center .- radius, center .+ radius)
end

for (name, N) in zip([:CosmoHyperrectangle, :CosmoCuboid, :CosmoRectangle], [:(length(lowerleft)), 3, 2])
    quote
        function $name(
            lowerleft::AbstractVector{T1},
            upperright::AbstractVector{T2},
        ) where {T1<:Number,T2<:Number}
            T = promote_type(T1, T2) |> float
            return CosmoHyperrectangle{T,$N}(lowerleft, upperright)
        end
    end |> eval
end

geometry_enclosing_corners(r::CosmoHyperrectangle) = r.lowerleft, r.upperright

function geometry_enclosing_center(r::CosmoHyperrectangle)
    return 1 // 2 .* (r.lowerleft .+ r.upperright)
end

function mask_in!(mask::BitVector, pos::AbstractMatrix{<:Number}, r::CosmoCuboid{T}) where {T}
    @assert size(pos, 1) == 3
    @assert size(pos, 2) == length(mask)
    @inbounds @views @. mask = (
        (r.lowerleft[1] ≤ pos[1, :] ≤ r.upperright[1]) &
        (r.lowerleft[2] ≤ pos[2, :] ≤ r.upperright[2]) &
        (r.lowerleft[3] ≤ pos[3, :] ≤ r.upperright[3])
    )
end

function mask_in!(mask::BitVector, pos::AbstractMatrix{<:Number}, r::CosmoRectangle{T}) where {T}
    @assert size(pos, 1) == 2
    @assert size(pos, 2) == length(mask)
    @inbounds @views @. mask =
        ((r.lowerleft[1] ≤ pos[1, :] ≤ r.upperright[1]) & (r.lowerleft[2] ≤ pos[2, :] ≤ r.upperright[2]))
end

function mask_in!(mask::BitVector, pos::AbstractMatrix{<:Number}, r::CosmoHyperrectangle{T,N}) where {T,N}
    @assert size(pos, 1) == N
    @assert size(pos, 2) == length(mask)
    @inbounds @views @. mask = r.lowerleft[1] ≤ pos[1, :] ≤ r.upperright[1]
    for i in 2:N
        mask .&= @inbounds @views @. r.lowerleft[i] ≤ pos[i, :] ≤ r.upperright[i]
    end
    return mask
end

"""
    struct CosmoHypersphere{T,N} <: AbstractCosmoGeometry where {T<:Number}
        center::Vector{T}
        radius::T
    end

`N`-dimensional hypersphere given by its `center` and `radius`.

If different types are passed to the constructor, the types will be promoted without throwing an error.
"""
struct CosmoHypersphere{T,N} <: AbstractCosmoGeometry where {T<:Number}
    center::Vector{T}
    radius::T

    function CosmoHypersphere{T,N}(center::AbstractVector{<:Number}, radius::Number) where {T<:Number,N}
        @assert length(center) == N
        new{T,N}(convert(Vector{T}, center), convert(T, radius))
    end
end

"""
    CosmoSphere{T} = CosmoHypersphere{T,3}

Alias for a 2D [`CosmoHypersphere`](@ref).
"""
const CosmoSphere{T} = CosmoHypersphere{T,3} where {T<:Number}

"""
    CosmoCircle{T} = CosmoHypersphere{T,2}

Alias for a 2D [`CosmoHypersphere`](@ref).
"""
const CosmoCircle{T} = CosmoHypersphere{T,2} where {T<:Number}

for (name, N) in zip([:CosmoHypersphere, :CosmoSphere, :CosmoCircle], [:(length(center)), 3, 2])
    quote
        function $name(center::AbstractVector{T1}, radius::T2) where {T1<:Number,T2<:Number}
            T = promote_type(T1, T2) |> float
            return CosmoHypersphere{T,$N}(center, radius)
        end
    end |> eval
end

"""
    CosmoHypersphere(N::Integer, radius::T)

Returns an `N`-dimensional hypersphere around the origin.
"""
function CosmoHypersphere(N::Integer, radius::T) where {T<:Number}
    Tfloat = float(T)
    return CosmoHypersphere{Tfloat,N}(zeros(Tfloat, N), radius)
end

"""
    CosmoSphere(radius::T)

Returns a sphere around the origin.
"""
function CosmoSphere(radius::T) where {T<:Number}
    Tfloat = float(T)
    return CosmoHypersphere{Tfloat,3}(zeros(Tfloat, 3), radius)
end

"""
    CosmoCircle(radius::T)

Returns a circle around the origin.
"""
function CosmoCircle(radius::T) where {T<:Number}
    Tfloat = float(T)
    return CosmoHypersphere{Tfloat,2}(zeros(Tfloat, 2), radius)
end

function Base.:(==)(s1::CosmoHypersphere, s2::CosmoHypersphere)
    return s1.center == s2.center && s1.radius == s2.radius
end

function geometry_enclosing_corners(s::CosmoHypersphere)
    return s.center .- s.radius, s.center .+ s.radius
end

geometry_enclosing_center(s::CosmoHypersphere) = s.center

function mask_in!(mask::BitVector, pos::AbstractMatrix{<:Number}, s::CosmoSphere)
    @assert size(pos, 1) == 3
    @assert size(pos, 2) == length(mask)
    r² = s.radius^2
    if all(==(0), s.center)
        @inbounds @views @. mask = pos[1, :]^2 + pos[2, :]^2 + pos[3, :]^2 ≤ r²
    else
        @inbounds @views @. mask =
            (pos[1, :] - s.center[1])^2 + (pos[2, :] - s.center[2])^2 + (pos[3, :] - s.center[3])^2 ≤ r²
    end
end

function mask_in!(mask::BitVector, pos::AbstractMatrix{<:Number}, circle::CosmoCircle)
    @assert size(pos, 1) == 2
    @assert size(pos, 2) == length(mask)
    r² = circle.radius^2
    if all(==(0), circle.center)
        @inbounds @views @. mask = pos[1, :]^2 + pos[2, :]^2 ≤ r²
    else
        @inbounds @views @. mask = (pos[1, :] - circle.center[1])^2 + (pos[2, :] - circle.center[2])^2 ≤ r²
    end
end

function mask_in!(mask::BitVector, pos::AbstractMatrix{<:Number}, s::CosmoHypersphere{T,N}) where {T,N}
    @assert size(pos, 1) == N
    @assert size(pos, 2) == length(mask)
    r² = s.radius^2
    center = s.center
    if all(==(0), center)
        @inbounds for i in eachindex(mask)
            s = (pos[1, i] - center[1])^2
            for j in 2:length(center)
                s += pos[j, i]^2
            end
            mask[i] = s ≤ r²
        end
    else
        @inbounds for i in eachindex(mask)
            s = (pos[1, i] - center[1])^2
            for j in 2:length(center)
                s += (pos[j, i] - center[j])^2
            end
            mask[i] = s ≤ r²
        end
    end
    return mask
end

"""
    struct CosmoCylinder{T} <: AbstractCosmoGeometry where {T<:Number}
        startpos::Vector{T}
        endpos::Vector{T}
        radius::T
    end

Cylinder given by its end points `startpos` and `endpos` and `radius`.

If different types are passed to the constructor, the types will be promoted without throwing an error.
"""
struct CosmoCylinder{T} <: AbstractCosmoGeometry where {T<:Number}
    startpos::Vector{T}
    endpos::Vector{T}
    radius::T

    function CosmoCylinder{T}(startpos::AbstractVector, endpos::AbstractVector, radius) where {T<:Number}
        @assert length(startpos) == 3 && length(endpos) == 3
        return new{T}(convert(Vector{T}, startpos), convert(Vector{T}, endpos), convert(T, radius))
    end
end

function CosmoCylinder(
    startpos::AbstractVector{T1},
    endpos::AbstractVector{T2},
    radius::T3,
) where {T1<:Number,T2<:Number,T3<:Number}
    T = promote_type(T1, T2, T3) |> float
    return CosmoCylinder{T}(startpos, endpos, radius)
end

function Base.:(==)(c1::CosmoCylinder, c2::CosmoCylinder)
    return c1.startpos == c2.startpos && c1.endpos == c2.endpos && c1.radius == c2.radius
end

function geometry_enclosing_corners(c::CosmoCylinder)
    # https://iquilezles.org/www/articles/diskbbox/diskbbox.htm
    a = c.endpos .- c.startpos
    @inbounds a² = a[1]^2 + a[2]^2 + a[3]^2
    e = @. c.radius * sqrt(1 - a^2 / a²)
    return (min.(c.startpos .- e, c.endpos .- e), max.(c.startpos .+ e, c.endpos .+ e))
end

geometry_enclosing_center(c::CosmoCylinder) = 1 // 2 .* (c.startpos .+ c.endpos)

function mask_in!(mask::BitVector, pos::AbstractMatrix{<:Number}, c::CosmoCylinder)
    # https://stackoverflow.com/questions/47932955/how-to-check-if-a-3d-point-is-inside-a-cylinder
    @assert size(pos, 1) == 3
    @assert size(pos, 2) == length(mask)
    pos1 = @view pos[1, :]
    pos2 = @view pos[2, :]
    pos3 = @view pos[3, :]
    p1 = c.startpos
    p2 = c.endpos
    Δ = p2 .- p1
    norm2Δ = Δ[1]^2 + Δ[2]^2 + Δ[3]^2
    r² = c.radius^2
    @inbounds @. mask =
        ((pos1 - p1[1]) * Δ[1] + (pos2 - p1[2]) * Δ[2] + (pos3 - p1[3]) * Δ[3] ≥ 0) &
        ((pos1 - p2[1]) * Δ[1] + (pos2 - p2[2]) * Δ[2] + (pos3 - p2[3]) * Δ[3] ≤ 0) &
        (
            (
                ((pos2 - p1[2]) * Δ[3] - (pos3 - p1[3]) * Δ[2])^2 +
                ((pos3 - p1[3]) * Δ[1] - (pos1 - p1[1]) * Δ[3])^2 +
                ((pos1 - p1[1]) * Δ[2] - (pos2 - p1[2]) * Δ[1])^2
            ) / norm2Δ ≤ r²
        )
end

"""
    struct CosmoStandingCylinder{T} <: AbstractCosmoGeometry where {T<:Number}
        center::Vector{T}
        height::T
        radius::T
    end

Standing cylinder given by its `center`, `height`, and `radius`.

Standing means that the cylinder is oriented such that its axis is aligned with the z axis.
If different types are passed to the constructor, the types will be promoted without throwing an error.
"""
struct CosmoStandingCylinder{T} <: AbstractCosmoGeometry where {T<:Number}
    center::Vector{T}
    height::T
    radius::T

    function CosmoStandingCylinder{T}(center::AbstractVector, height, radius) where {T<:Number}
        @assert length(center) == 3
        return new{T}(convert(Vector{T}, center), convert(T, height), convert(T, radius))
    end
end

function CosmoStandingCylinder(
    center::AbstractVector{T1},
    height::T2,
    radius::T3,
) where {T1<:Number,T2<:Number,T3<:Number}
    T = promote_type(T1, T2, T3) |> float
    return CosmoStandingCylinder{T}(center, height, radius)
end

function Base.:(==)(c1::CosmoStandingCylinder, c2::CosmoStandingCylinder)
    return c1.center == c2.center && c1.height == c2.height && c1.radius == c2.radius
end

function geometry_enclosing_corners(c::CosmoStandingCylinder)
    c.center .- [c.radius, c.radius, 1 // 2 * c.height], c.center .+ [c.radius, c.radius, 1 // 2 * c.height]
end

geometry_enclosing_center(c::CosmoStandingCylinder) = c.center

function mask_in!(mask::BitVector, pos::AbstractMatrix{<:Number}, c::CosmoStandingCylinder)
    @assert size(pos, 1) == 3
    @assert size(pos, 2) == length(mask)
    r² = c.radius^2
    halfheight = 1 // 2 * c.height
    @inbounds @views @. mask =
        ((pos[1, :] - c.center[1])^2 + (pos[2, :] - c.center[2])^2 ≤ r²) &
        (c.center[3] - halfheight ≤ pos[3, :] ≤ c.center[3] + halfheight)
end

"""
    CosmoCylinder(c::CosmoStandingCylinder)

Create a cylinder from a standing cylinder.
"""
function CosmoCylinder(c::CosmoStandingCylinder)
    halfheight = 1 // 2 * c.height
    startpos = [c.center[1], c.center[2], c.center[3] - halfheight]
    endpos = [c.center[1], c.center[2], c.center[3] + halfheight]
    return CosmoCylinder(startpos, endpos, c.radius)
end

"""
    CosmoStandingCylinder(c::CosmoCylinder)

Create a standing cylinder from a cylinder.

The end positions of the cylinder have to have the same x and y coordinates.
"""
function CosmoStandingCylinder(c::CosmoCylinder)
    @assert c.startpos[1] == c.endpos[1] && c.startpos[2] == c.endpos[2]
    center = geometry_enclosing_center(c)
    height = abs(c.endpos[3] - c.startpos[3])
    return CosmoStandingCylinder(center, height, c.radius)
end

function Base.:(==)(c1::CosmoCylinder, c2::CosmoStandingCylinder)
    return geometry_enclosing_center(c1) == geometry_enclosing_center(c2) &&
           abs.(c1.startpos[3] - c1.endpos[3]) == c2.height &&
           c1.radius == c2.radius &&
           @views c1.startpos[1:2] == c1.endpos[1:2] # c1 is actually standing
end
Base.:(==)(c1::CosmoStandingCylinder, c2::CosmoCylinder) = ==(c2, c1)

function Base.isapprox(c1::CosmoCylinder, c2::CosmoStandingCylinder; kwargs...)
    return isapprox(geometry_enclosing_center(c1), geometry_enclosing_center(c2); kwargs...) &&
           isapprox(abs.(c1.startpos[3] - c1.endpos[3]), c2.height; kwargs...) &&
           isapprox(c1.radius, c2.radius; kwargs...) &&
           @views isapprox(c1.startpos[1:2], c1.endpos[1:2]; kwargs...) # c1 is actually standing
end
Base.isapprox(c1::CosmoStandingCylinder, c2::CosmoCylinder; kwargs...) = isapprox(c2, c1; kwargs...)


text:
using SymbolicUtils
using Metatheory

sym_eq(x, y) = simplify(x == y)

@testset "Basic optimization" begin
    @syms a b c x y z
    ex = 2a + 2b - (a*(a + b))
    res = optimize(ex)
    @test sym_eq(res, (a+b)*(2-a))

    ex = sin(a^2)/cos(a^2)
    res = optimize(ex)
    @test isequal(res, tan(a^2)) # sym_eq does not work
    # TODO report issue

    ex = sin(1/x * (a * b + a * c))^2 + cos((a*(b+c))/x)^2
    res = optimize(ex)
    @test res == 1
end


text:
mutable struct FrictionCalculator{T,M} <: AbstractFrictionCalculator{T,M}
    model::M
    potential::DependentField{T,Matrix{T}}
    derivative::DependentField{Matrix{T},Matrix{T}}
    friction::DependentField{Matrix{T},Matrix{T}}
    stats::Dict{Symbol,Int}
    function FrictionCalculator{T}(model::M, atoms::Integer) where {T,M<:Model}

        potential = zero(T)
        derivative = zeros(ndofs(model), atoms)
        friction = zeros(ndofs(model)*atoms, ndofs(model)*atoms)

        position = fill(NaN, ndofs(model), atoms)

        stats = Dict{Symbol,Int}(
            :potential=>0,
            :derivative=>0,
            :friction=>0,
        )

        new{T,M}(
            model,
            DependentField(potential, copy(position)),
            DependentField(derivative, copy(position)),
            DependentField(friction, copy(position)),
            stats
        )
    end
end

struct RingPolymerFrictionCalculator{T,M} <: AbstractFrictionCalculator{T,M}
    model::M
    potential::DependentField{Vector{T},Array{T,3}}
    derivative::DependentField{Array{T,3},Array{T,3}}
    friction::DependentField{Array{T,3},Array{T,3}}
    stats::Dict{Symbol,Int}
    function RingPolymerFrictionCalculator{T}(model::M, atoms::Integer, beads::Integer) where {T,M<:Model}

        potential = zeros(beads)
        derivative = zeros(ndofs(model), atoms, beads)
        friction = zeros(ndofs(model)*atoms, ndofs(model)*atoms, beads)

        position = fill(NaN, ndofs(model), atoms, beads)

        stats = Dict{Symbol,Int}(
            :potential=>0,
            :derivative=>0,
            :friction=>0,
        )

        new{T,M}(
            model,
            DependentField(potential, copy(position)),
            DependentField(derivative, copy(position)),
            DependentField(friction, copy(position)),
            stats
        )
    end
end

struct DiabaticFrictionCalculator{T,M} <: AbstractDiabaticCalculator{T,M}
    model::M
    potential::DependentField{Hermitian{T,Matrix{T}},Matrix{T}}
    derivative::DependentField{Matrix{Hermitian{T,Matrix{T}}},Matrix{T}}
    eigen::DependentField{LinearAlgebra.Eigen{T,T,Matrix{T},Vector{T}},Matrix{T}}
    adiabatic_derivative::DependentField{Matrix{Matrix{T}},Matrix{T}}
    nonadiabatic_coupling::DependentField{Matrix{Matrix{T}},Matrix{T}}
    friction::DependentField{Matrix{T},Matrix{T}}
    tmp_mat::Matrix{T}
    stats::Dict{Symbol,Int}

    function DiabaticFrictionCalculator{T}(model::M, atoms::Integer) where {T,M<:Model}
        mat = NQCModels.DiabaticModels.matrix_template(model, T)
        vec = NQCModels.DiabaticModels.vector_template(model, T)

        potential = Hermitian(zero(mat))
        derivative = [Hermitian(zero(mat)) for i=1:ndofs(model), j=1:atoms]
        eigen = Eigen(zero(vec), zero(mat)+I)
        adiabatic_derivative = [zero(mat) for i=1:ndofs(model), j=1:atoms]
        nonadiabatic_coupling = [zero(mat) for i=1:ndofs(model), j=1:atoms]
        friction = zeros(ndofs(model)*atoms, ndofs(model)*atoms)
        tmp_mat = zero(mat)

        position = fill(NaN, ndofs(model), atoms)

        stats = Dict{Symbol,Int}(
            :potential=>0,
            :derivative=>0,
            :eigen=>0,
            :adiabatic_derivative=>0,
            :nonadiabatic_coupling=>0,
            :friction=>0
        )

        new{T,M}(
            model,
            DependentField(potential, copy(position)),
            DependentField(derivative, copy(position)),
            DependentField(eigen, copy(position)),
            DependentField(adiabatic_derivative, copy(position)),
            DependentField(nonadiabatic_coupling, copy(position)),
            DependentField(friction, copy(position)),
            tmp_mat,
            stats
        )
    end
end

struct RingPolymerDiabaticFrictionCalculator{T,M} <: AbstractDiabaticCalculator{T,M}
    model::M
    potential::DependentField{Vector{Hermitian{T,Matrix{T}}},Array{T,3}}
    derivative::DependentField{Array{Hermitian{T,Matrix{T}},3},Array{T,3}}
    eigen::DependentField{Vector{LinearAlgebra.Eigen{T,T,Matrix{T},Vector{T}}},Array{T,3}}
    adiabatic_derivative::DependentField{Array{Matrix{T},3},Array{T,3}}
    nonadiabatic_coupling::DependentField{Array{Matrix{T},3},Array{T,3}}
    friction::DependentField{Matrix{T},Array{T,3}}
    tmp_mat::Matrix{T}
    stats::Dict{Symbol,Int}

    function RingPolymerDiabaticFrictionCalculator{T}(model::M, atoms::Integer, beads::Integer) where {T,M<:Model}
        n = nstates(model)
        matrix_template = NQCModels.DiabaticModels.matrix_template(model, T)
        vector_template = NQCModels.DiabaticModels.vector_template(model, T)

        potential = [Hermitian(matrix_template) for _=1:beads]
        derivative = [Hermitian(matrix_template) for _=1:ndofs(model), _=1:atoms, _=1:beads]
        adiabatic_derivative = [matrix_template for _=1:ndofs(model), _=1:atoms, _=1:beads]
        eigen = [Eigen(vector_template, matrix_template + I) for _=1:beads]
        nonadiabatic_coupling = [matrix_template for _=1:ndofs(model), _=1:atoms, _=1:beads]
        friction = zeros(ndofs(model)*atoms*beads, ndofs(model)*atoms*beads)
        tmp_mat = zeros(T, n, n)

        position = fill(NaN, ndofs(model), atoms, beads)

        stats = Dict{Symbol,Int}(
            :potential=>0,
            :derivative=>0,
            :eigen=>0,
            :adiabatic_derivative=>0,
            :nonadiabatic_coupling=>0,
            :friction=>0
        )

        new{T,M}(
            model,
            DependentField(potential, copy(position)),
            DependentField(derivative, copy(position)),
            DependentField(eigen, copy(position)),
            DependentField(adiabatic_derivative, copy(position)),
            DependentField(nonadiabatic_coupling, copy(position)),
            DependentField(friction, copy(position)),
            tmp_mat,
            stats
        )
    end
end

function Calculator(model::AdiabaticFrictionModel, atoms::Integer, t::Type{T}) where {T}
    FrictionCalculator{t}(model, atoms)
end
function Calculator(model::DiabaticFrictionModel, atoms::Integer, t::Type{T}) where {T}
    DiabaticFrictionCalculator{t}(model, atoms)
end
function Calculator(model::DiabaticFrictionModel, atoms::Integer, beads::Integer, t::Type{T}) where {T}
    RingPolymerDiabaticFrictionCalculator{t}(model, atoms, beads)
end
function Calculator(model::AdiabaticFrictionModel, atoms::Integer, beads::Integer, t::Type{T}) where {T}
    RingPolymerFrictionCalculator{t}(model, atoms, beads)
end

function evaluate_friction!(calc::AbstractFrictionCalculator, R::AbstractMatrix)
    calc.stats[:friction] += 1
    NQCModels.friction!(calc.model, calc.friction, R)
end

function evaluate_friction!(calc::AbstractFrictionCalculator, R::AbstractArray{T,3}) where {T}
    calc.stats[:friction] += 1
    @views for i in axes(R, 3)
        NQCModels.friction!(calc.model, calc.friction[:,:,i], R[:,:,i])
    end
end

@doc raw"""
    evaluate_friction!(calc::DiabaticFrictionCalculator, R::AbstractMatrix)

Evaluate the electronic friction for a model given in the diabatic representation.

```math
γ = 2πħ ∑ⱼ <1|dH|j><j|dH|1> δ(ωⱼ) / ωⱼ
```
Note that the delta function is approximated by a normalised gaussian.
"""
function evaluate_friction!(calc::DiabaticFrictionCalculator, r::AbstractMatrix)

    gauss(x, σ) = exp(-0.5 * x^2 / σ^2) / (σ*sqrt(2π))
    adiabatic_derivative = get_adiabatic_derivative(calc, r)
    eigen = get_eigen(calc, r)

    dofs = ndofs(calc.model)
    calc.friction .= 0
    for i in axes(r, 2) # Atoms
        for j in axes(r, 1) # dofs
            for m=2:nstates(calc.model)
                ω = eigen.values[m] - eigen.values[1]
                g = gauss(ω, calc.model.σ) / ω
                calc.friction[j+(i-1)*dofs] += 2π*abs2(adiabatic_derivative[j,i][m,1])*g
            end
        end
    end
end


text:
# The Computer Language Benchmarks Game
# http://shootout.alioth.debian.org/
#
# Contributed by David Campbell

const ITER = 50

function ismandel(z::Complex{Float64})
    c = z
    for n = 1:ITER
        if abs2(z) > 4
            return false
        end
        z = z^2 + c
    end
    return true
end

function draw_mandel(M::Array{UInt8, 2}, n::Int)
    for y = 0:n-1
        ci = 2y/n - 1
        for x = 0:n-1
            c = complex(2x/n - 1.5, ci)
            if ismandel(c)
                M[div(x, 8) + 1, y + 1] |= 1 << UInt8(7 - x%8)
            end
        end
    end
end

function perf_mandelbrot(n::Int=200)
    if n%8 != 0
        error("Error: n of $n is not divisible by 8")
    end

    M = zeros(UInt8, div(n, 8), n)
    draw_mandel(M, n)

    sprint() do output
        write(output, "P4\n$n $n\n")
        write(output, M)
    end
end


text:
module Optimization

using JuMP, Ipopt
using ClimateMARGO.Models
using ClimateMARGO.Utils
using ClimateMARGO.Diagnostics

export optimize_controls!

include("deterministic_optimization.jl")

end


text:
using ForwardDiff
using RootSolvers

function T2_h_i(β, α::FT, i::IT, N::IT) where {FT, IT}
  a = α; b = β
  g = (b+FT(1))/(b-FT(1))
  t = (FT(i)/FT(N) - a)/(FT(1)-a)
  f = ((b+FT(2)*a)*g^t - b + FT(2)*a)/((FT(2)*a+FT(1))*(FT(1)+g^t))
  return f
end

T2_root_Δh_near_hmax(β, hmin::FT,hmax::FT,α::FT,N::IT,Δh::FT) where {FT,IT} =
  T2_h_i(β,α,N,N) - T2_h_i(β,α,N-1,N) - Δh/(hmax-hmin)

T2_root_Δh_near_hmin(β, hmin::FT,hmax::FT,α::FT,N::IT,Δh::FT) where {FT,IT} =
  T2_h_i(β,α,1,N) - T2_h_i(β,α,0,N) - Δh/(hmax-hmin)

"""
    β_Δh_small

Compute `β` such that `N` grid points
between `hmin` and `hmax` stretch so
that the smallest `Δh` matches the one
given.

!!! warn
    If `N` is small, `β->∞`, and the
    given `Δh` may not be satisfied.
    `N` can be increased to fix this.
"""
function β_Δh_small(hmin::FT, hmax::FT, N::IT, Δh::FT) where {FT, IT}
  tol = eps(FT)
  β_initial = FT(1) + tol # Initial guess
  root_closure(β) = T2_root_Δh_near_hmax(β, hmin, hmax, FT(0), N, Δh)
  args = (NewtonsMethodAD(β_initial), CompactSolution(), SolutionTolerance(eps(FT)))
  sol = find_zero(β -> root_closure(β), args...)
  return sol.root
end

"""
    β_Δh_big

Compute `β` such that `N` grid points
between `hmin` and `hmax` stretch so
that the largest `Δh` matches the one
given.

!!! warn
    If `N` is small, `β->∞`, and the
    given `Δh` may not be satisfied.
    `N` can be increased to fix this.
"""
function β_Δh_big(hmin::FT, hmax::FT, N::IT, Δh::FT) where {FT, IT}
  tol = eps(FT)
  β_initial = FT(1) + tol # Initial guess
  root_closure(β) = T2_root_Δh_near_hmin(β, hmin, hmax, FT(0), N, Δh)
  args = (NewtonsMethodAD(β_initial), CompactSolution(), SolutionTolerance(eps(FT)))
  sol = find_zero(β -> root_closure(β), args...)
  return sol.root
end

"""
    β_Δh_both

Compute `β` such that `N` grid points
between `hmin` and `hmax` stretch so
that the smallest `Δh` matches the one
given.

!!! warn
    If `N` is small, `β->∞`, and the
    given `Δh` may not be satisfied.
    `N` can be increased to fix this.
"""
function β_Δh_both(hmin::FT, hmax::FT, N::IT, Δh::FT) where {FT, IT}
  tol = eps(FT)
  β_initial = FT(1) + tol # Initial guess
  root_closure(β) = T2_root_Δh_near_hmax(β, hmin, hmax, FT(1/2), N, Δh)
  args = (NewtonsMethodAD(β_initial), CompactSolution(), SolutionTolerance(eps(FT)))
  sol = find_zero(β -> root_closure(β), args...)
  return sol.root
end


text:
"""
    alluniquesorted(vec)

Check if all values of sorted vector `vec` are unique.
"""
function alluniquesorted(vec)
    for i in 1:length(vec)-1
        vec[i] == vec[i+1] && return false
    end
    true
end

"""
    sumcols!(buff, col1, col2)

Calculate the symmetric difference of two sorted integer vectors. This is
used to sum columns in the reduction matrix (mod 2).

Writes the result to `buff` and also returns an emptied `col1`, which can be
reused as `buff` in the next iteration.
"""
@inline function sumcols!(buff, col1, col2)
    i = j = 1
    n = length(col1)
    m = length(col2)
    @inbounds while i ≤ n && j ≤ m
        if col1[i] < col2[j]
            push!(buff, col1[i])
            i += 1
        elseif col1[i] > col2[j]
            push!(buff, col2[j])
            j += 1
        else
            i += 1
            j += 1
        end
    end
    @inbounds while j ≤ m
        push!(buff, col2[j])
        j += 1
    end
    @inbounds while i ≤ n
        push!(buff, col1[i])
        i += 1
    end
    empty!(col1)
    buff, col1
end

"""
    getlow(vec)

Get the "low one" in a column of the reduction matrix.
"""
function getlow(col::AbstractVector{T}) where T
    isempty(col) ? zero(T) : last(col)
end

"""
    allpaths(G::SimpleWeightedGraph)

Get all shortest paths in graph.
"""
function allpaths(G::SimpleWeightedGraph{I, T}) where {I, T}
    paths = Vector{DijkstraState{T, I}}(nv(G))
    for v in vertices(G)
        paths[v] = dijkstra_shortest_paths(G, v)
    end
    paths
end

"""
    iterlength(itr)

Get the length of an iterator that does not have `length` defined.
"""
function iterlength(itr)
    len = 0
    for _ in itr
        len += 1
    end
    len
end

"""
    addpath!(res_e, res_p, v1, v2, paths, edgeord)

Find the path from `v1` to `v2` in `paths`, convert to edge indices from
`edgeord` and append the result to `res`.
"""
function addpath!(res_e, res_p, v1, v2, paths, edgeord)
    @inbounds begin
        path = enumerate_paths(paths[v1], v2)
        for i in 1:length(path)-1
            push!(res_e, edgeord[path[i], path[i + 1]])
            push!(res_p, path[i])
        end
    end
    res_e, res_p
end

# Rewrite
@inline function islocallyminimal(sides, distmat, s)
    s12, s23, s31 = sides
    l12, l23, l31 = length.(sides)
    for i in 2:l12, j in 2:l23
        if distmat[s12[i], s23[j]] < distmat[s12[i], s23[1]] + distmat[s23[1], s23[j]] - 2s
            return false
        end
    end
    for i in 2:l23, j in 2:l31
        if distmat[s23[i], s31[j]] < distmat[s23[i], s31[1]] + distmat[s31[1], s31[j]] - 2s
            return false
        end
    end
    for i in 2:l31, j in 2:l12
        if distmat[s31[i], s12[j]] < distmat[s31[i], s12[1]] + distmat[s12[1], s12[j]] - 2s
            return false
        end
    end
    true
end

function showprogress(j, r, ntriangles, nskipped, ngenerators, ngenerators_left, nrows)
    print("\rTriangles visited:",
          lpad(round(Int, 100 * j/ntriangles), 3), "%, l ≈ ",
          rpad(round(3r, 5), 7, 0),
          ", skipped:",
          lpad(round(Int, 100 * nskipped/j), 3), "%",
          ". Generators found:",
          lpad(ngenerators, 2),
          ", left: ",
          rpad(ngenerators_left, floor(Int, log10(nrows) + 1)))
    flush(STDOUT::IO)
end

"""
    persistence(points; printprogress, resolution, subset)

Calculate 1-dimensional intrinsic persistence of `points`.
"""
function persistence(pts::AbstractMatrix{T}; kwargs...) where T
    n, m = size(pts)
    _pts = reinterpret(SVector{n, T}, pts, (m,))
    persistence(_pts; kwargs...)
end
function persistence(pts::AbstractVector{SVector{D, T}};
                     show_trace = false,
                     resolution = nothing) where {D, T}

    # Prepare default arguments.
    if resolution == nothing
        s = sampledensity(pairwise(Euclidean(),
                                   reinterpret(T, pts, (D, length(pts)))))
    else
        s = resolution
    end

    # Main algo.
    show_trace && print("# ===================== #\n",
                        "| Intrinsic persistence |\n",
                        "# ===================== #\n\n",
                        "Preprocessing...")

    # Intrinsic metric distance matrix and adjacency graph.
    metric  = IntrinsicMetric(pts, 2s)
    distmat = pairwise(metric)
    graph   = adjgraph(metric)

    # Number of rows in reduction matrix.
    nrows = ne(graph)
    # Edge ordering matrix.
    edgeord = edgetoorder(graph)
    # Shortest paths from every node.
    paths = allpaths(graph)
    # Triangle iterator.
    triangles = equilaterals(distmat, tol = 2s, lower = 0)
    # Reduction matrix columns indexed by getlow.
    L = fill(Int[], nrows)

    # Result storage.
    generators = IntrinsicGenerator{T}[]
    # Number of generators left to discover.
    ngenerators_left = nrows - length(pts) + length(connected_components(graph))
    nskipped = 0

    if show_trace
        ntriangles = iterlength(triangles)
        print("\r")
        println("* Sample density:       $(round(s, 5))")
        println("* Number of triangles:  $ntriangles")
        println("* Number of generators: $(ngenerators_left)")
        println()
    end

    # j-th column and buffer are swapped back and forth to avoid allocation.
    ∂_j   = Int[]
    buff  = Int[]
    # Sides of the triangle as paths.
    sides = (Int[], Int[], Int[])
    currindex = 0
    r = zero(T)
    for Δ in triangles
        currindex += 1
        ngenerators_left > 0 || break

        v1, v2, v3 = Δ.vertices
        r = Δ.r

        if show_trace && currindex % (ntriangles ÷ 100) == 0
            showprogress(currindex, r, ntriangles, nskipped,
                         length(generators), ngenerators_left, nrows)
        end

        # Build column and path.
        empty!(∂_j)
        empty!(sides[1]); empty!(sides[2]); empty!(sides[3])
        addpath!(∂_j, sides[1], v1, v2, paths, edgeord)
        addpath!(∂_j, sides[2], v2, v3, paths, edgeord)
        addpath!(∂_j, sides[3], v3, v1, paths, edgeord)
        sort!(∂_j)

        # Skip
        #!alluniquesorted(∂_j) && continue

        if true
        if (r ≤ 2s && length(∂_j) > 3) ||
           !alluniquesorted(∂_j) ||
           !islocallyminimal(sides, distmat, s)
            nskipped += 1
            continue
        end
        end

        # Matrix reduction.
        low = getlow(∂_j)
        while low ≠ 0 && !isempty(L[low])
            ∂_j, buff = sumcols!(buff, ∂_j, L[low])
            low = getlow(∂_j)
        end

        # Output.
        if low ≠ 0
            L[low] = ∂_j
            ∂_j = Int[]

            if r > 2s
                path = Int[]
                for i in 1:3
                    append!(path, sides[i])
                end
                lengths = similar(path, T)
                for p_i in eachindex(path)
                    p_j = p_i < length(path) ? p_i + 1 : 1
                    lengths[p_i] = distmat[path[p_i], path[p_j]]
                end

                push!(generators, IntrinsicGenerator(Δ, path, lengths))
            end

            ngenerators_left -= 1
            nskipped += 1
        end
    end
    if show_trace
        showprogress(currindex, r, ntriangles, nskipped,
                     length(generators), ngenerators_left, nrows)
        println()
        println()
    end

    IntrinsicPersistenceResult(generators, pts, collect(1:length(pts)))
end


text:
module Commands

using QuantumESPRESSOCommands:
    ParallelizationFlags,
    PwxConfig,
    PhxConfig,
    Q2rxConfig,
    MatdynxConfig,
    DynmatxConfig,
    QuantumESPRESSOConfig,
    pw,
    ph,
    q2r,
    matdyn,
    dynmat,
    makecmd

export pw, ph, q2r, matdyn, dynmat

end


text:
# This file was generated, do not modify it. # hide
Refpath="/tmp/resultJ.txt"
Tarpath="/tmp/resultJ1.txt"
run(`bash -c "cmp --silent $Refpath $Tarpath || echo \"files are different\""`)

using Plots
using DelimitedFiles
img = readdlm(Tarpath,',',Int)
M,N = img |> size
p = heatmap(1:N,1:M,img,aspect_ratio=1, size=(400, 400))
savefig(p, joinpath(@OUTPUT, "mandJ1.png")) #hide


text:
module JuliaSetPkg

export juliaSet

function juliaSet{T<:Number}(R::Function, x::Array{T}, y::Array{T}, n_iter = 200, escape_tol=1e6)
    A = zeros(length(x), length(y))
    for i in 1:length(x)
      for j in 1:length(y)
        z = x[i] + im * y[j]
        for k in 1:n_iter
          z = R(z)
          if abs(z) > escape_tol # diverged
            A[i,j] = k
            break
          end
        end
        if A[i,j] == 0 # didn't diverge
            A[i,j] = n_iter + 1
        end
      end
    end
    return A
end

end # module

max_stars_repo_name: anjandn/JuliaSetPkg

max_issues_repo_name: anjandn/JuliaSetPkg

max_forks_repo_name: anjandn/JuliaSetPkg


text:
# Bandpass fitler example
using MarketCycles
using Gadfly

# Generate dummy data
srand(121)
n = 1000
op = 100.0 + cumsum(randn(n))
hi = op + rand(n)
lo = op - rand(n)
cl = 100.0 + cumsum(randn(n))
index = collect(1:1:length(cl))
for i = 1:n
	if cl[i] > hi[i]
		cl[i] = hi[i]
	elseif cl[i] < lo[i]
		cl[i] = lo[i]
	end
end

# Apply Bandpass Filter
BP_Filter = BandPassFilter(cl,n=30,bandwidth=.3)

# Plot
white_panel = Theme(
	panel_fill="white",
    default_color="blue",
    background_color="white"
)
	p1 = plot(x=index,y=cl,Geom.line,
	Guide.xlabel(nothing), Guide.ylabel("Price"), Guide.title("Dummy Data"),white_panel)
	p2 = plot(x=index,y=BP_Filter,Geom.line,Guide.xlabel("Time Index"),Guide.title("Band Pass Filter"), Guide.ylabel("Band Pass Filter"),white_panel)
	out = vstack(p1,p2)
	# Save Plot
	filename = joinpath(dirname(@__FILE__), "BP_example.png")
	draw(PNG(filename, 1500px, 800px), out)


text:
const JULIALANG_REPOS = MirrorUpdater.Types.SrcDestPair[
    MirrorUpdater.Types.SrcDestPair(
        ;
        source_url = "https://github.com/JuliaLang/METADATA.jl",
        destination_repo_name = "JuliaLang-METADATA.jl",
        ),
    MirrorUpdater.Types.SrcDestPair(
        ;
        source_url = "https://github.com/JuliaLang/Pkg.jl",
        destination_repo_name =
            "JuliaLang-Pkg.jl-44cfe95a-1eb2-52ea-b672-e2afdf69b78f",
        ),
    MirrorUpdater.Types.SrcDestPair(
        ;
        source_url = "https://github.com/JuliaLang/julia",
        destination_repo_name = "JuliaLang-julia",
        ),
    MirrorUpdater.Types.SrcDestPair(
        ;
        source_url = "https://github.com/JuliaLang/libosxunwind",
        destination_repo_name = "JuliaLang-libosxunwind",
        ),
    MirrorUpdater.Types.SrcDestPair(
        ;
        source_url = "https://github.com/JuliaLang/libuv",
        destination_repo_name = "JuliaLang-libuv",
        ),
    MirrorUpdater.Types.SrcDestPair(
        ;
        source_url = "https://github.com/JuliaLang/readline",
        destination_repo_name = "JuliaLang-readline",
        ),
    MirrorUpdater.Types.SrcDestPair(
        ;
        source_url = "https://github.com/JuliaRegistries/Registrator.jl",
        destination_repo_name = "Registrator.jl-4418983a-e44d-11e8-3aec-9789530b3b3e",
        ),
    MirrorUpdater.Types.SrcDestPair(
        ;
        source_url = "https://github.com/JuliaRegistries/TagBot",
        destination_repo_name = "JuliaRegistries-TagBot",
        ),
    ]


text:
using Test, ModelConstructors

# CURRENTLY ONLY TESTS VALUE, PRIOR, FIXED, AND VALUEBOUNDS SWITCHING, NO REGIME SWITCHING IN OTHER CASES
# ALSO TESTS IF TRANSFORMS WORK CORRECTLY WITH REGIME-SWITCHING

@info "The following error 'get_regime_val(), Input Error: No regime 3' is expected."
@testset "Regime switching with parameters" begin
    u = parameter(:bloop, 2.5230, (1e-8, 5.), (1e-8, 5.), ModelConstructors.SquareRoot(); fixed = false)
    ModelConstructors.set_regime_val!(u, 1, 2.5230)

    @test_throws ParamBoundsError ModelConstructors.set_regime_val!(u, 2, 0.)
    ModelConstructors.set_regime_val!(u, 2, 0.; override_bounds = true)
    uvec = ParameterVector{Float64}(undef, 2)
    uvec[1] = u
    uvec[2] = u

    # Test set_regime_val! and regime_val
    @test !isempty(u.regimes)
    @test u.regimes[:value][1] == 2.5230
    @test u.regimes[:value][2] == 0.
    @test u.value == 2.5230
    @test ModelConstructors.regime_val(u, 1) == 2.5230
    @test ModelConstructors.regime_val(u, 2) == 0.
    @test_throws KeyError ModelConstructors.regime_val(u, 3)
    ModelConstructors.toggle_regime!(u, 2)
    @test u.value == 0.
    @test ModelConstructors.regime_val(u, 1) == 2.5230
    @test ModelConstructors.regime_val(u, 2) == 0.
    ModelConstructors.toggle_regime!(u, 1)
    @test u.value == 2.5230
    @test ModelConstructors.regime_val(u, 1) == 2.5230
    @test ModelConstructors.regime_val(u, 2) == 0.

    # Test get_values and toggling regimes
    @test n_parameters_regime_switching(uvec) == 4
    ModelConstructors.toggle_regime!(u, 1)
    @test ModelConstructors.get_values(uvec; regime_switching = false) == [2.5230, 2.5230]
    @test ModelConstructors.get_values(uvec) == [2.5230, 2.5230, 0., 0.]

    # Test set_regime_prior! and regime_prior
    ModelConstructors.set_regime_prior!(u, 1, Uniform(0., 5.))
    ModelConstructors.set_regime_prior!(u, 2, Normal(0., 1.))

    @test get(ModelConstructors.regime_prior(u, 1)) == Uniform(0., 5.)
    @test get(ModelConstructors.regime_prior(u, 2)) == Normal(0., 1.)

    # Test set_regime_fixed! and regime_fixed
    @test !haskey(u.regimes, :valuebounds)
    ModelConstructors.set_regime_fixed!(u, 1, true)
    @test haskey(u.regimes, :valuebounds)
    ModelConstructors.set_regime_fixed!(u, 2, false)
    @show u.regimes[:value]
    @test_throws ErrorException ModelConstructors.set_regime_fixed!(u, 3, true)

    @test ModelConstructors.regime_fixed(u, 1)
    @test !ModelConstructors.regime_fixed(u, 2)
    @test ModelConstructors.regime_valuebounds(u, 1) == (2.5230, 2.5230)
    @test ModelConstructors.regime_valuebounds(u, 2) == (1e-8, 5.)

    # test regime-switching value bounds
    ModelConstructors.set_regime_valuebounds!(u, 1, (1e-6, 5.))
    ModelConstructors.set_regime_valuebounds!(u, 2, (1e-7, 5.))
    @test ModelConstructors.regime_valuebounds(u, 1) == (1e-6, 5.)
    @test ModelConstructors.regime_valuebounds(u, 2) == (1e-7, 5.)

    # test set_regime_fixed w/update_valuebounds
    ModelConstructors.set_regime_fixed!(u, 1, true; update_valuebounds = (10., 11.))
    ModelConstructors.set_regime_fixed!(u, 2, false; update_valuebounds = (10., 11.))
    @test u.regimes[:valuebounds][1] == (10., 11.)
    @test u.regimes[:valuebounds][2] == (10., 11.)

    # test transform_to_real_line
    ModelConstructors.toggle_regime!(uvec, 1)
    set_regime_val!(uvec[1], 2, 1.; override_bounds = true) # break the valuebounds for now
    set_regime_val!(uvec[2], 2, 1.; override_bounds = true)
    values = ModelConstructors.get_values(uvec)
    real_vals_true = similar(values)
    real_vals_true[1] = ModelConstructors.transform_to_real_line(uvec[1], uvec[1].regimes[:value][1])
    real_vals_true[2] = ModelConstructors.transform_to_real_line(uvec[2], uvec[2].regimes[:value][1])
    real_vals_true[3] = ModelConstructors.transform_to_real_line(uvec[1], uvec[1].regimes[:value][2])
    real_vals_true[4] = ModelConstructors.transform_to_real_line(uvec[2], uvec[2].regimes[:value][2])
    real_vals1 = transform_to_real_line(uvec; regime_switching = true)
    real_vals2 = transform_to_real_line(uvec, values; regime_switching = true)
    @test real_vals1 == real_vals2 == real_vals_true

    # test transform_to_model_space
    model_vals = transform_to_model_space(uvec, real_vals1; regime_switching = true)
    @test model_vals == values
end

@testset "Regime switching with parameters when model regimes are different" begin
    u = parameter(:bloop, 2.5230, (1e-8, 5.), (1e-8, 5.), ModelConstructors.SquareRoot(); fixed = false)
    ModelConstructors.set_regime_val!(u, 1, 2.5230)

    d = Dict(2 => 1, 3 => 2, 4 => 3)
    @test_throws ParamBoundsError ModelConstructors.set_regime_val!(u, 3, 0., d)
    ModelConstructors.set_regime_val!(u, 3, 0., d; override_bounds = true)
    uvec = ParameterVector{Float64}(undef, 2)
    uvec[1] = u
    uvec[2] = u

    # Test set_regime_val! and regime_val
    @test !isempty(u.regimes)
    @test u.regimes[:value][1] == 2.5230
    @test u.regimes[:value][2] == 0.
    @test u.value == 2.5230
    @test ModelConstructors.regime_val(u, 2, d) == 2.5230
    @test ModelConstructors.regime_val(u, 3, d) == 0.
    @test_throws KeyError ModelConstructors.regime_val(u, 4, d)
    ModelConstructors.toggle_regime!(u, 3, d)
    @test u.value == 0.
    @test ModelConstructors.regime_val(u, 2, d) == 2.5230
    @test ModelConstructors.regime_val(u, 3, d) == 0.
    ModelConstructors.toggle_regime!(u, 2, d)
    @test u.value == 2.5230
    @test ModelConstructors.regime_val(u, 2, d) == 2.5230
    @test ModelConstructors.regime_val(u, 3, d) == 0.

    # Test get_values and toggling regimes
    @test n_parameters_regime_switching(uvec) == 4
    ModelConstructors.toggle_regime!(u, 2, d)
    @test ModelConstructors.get_values(uvec; regime_switching = false) == [2.5230, 2.5230]
    @test ModelConstructors.get_values(uvec) == [2.5230, 2.5230, 0., 0.]

    # Test set_regime_prior! and regime_prior
    ModelConstructors.set_regime_prior!(u, 2, Uniform(0., 5.), d)
    ModelConstructors.set_regime_prior!(u, 3, Normal(0., 1.), d)

    @test get(ModelConstructors.regime_prior(u, 1)) == Uniform(0., 5.)
    @test get(ModelConstructors.regime_prior(u, 2)) == Normal(0., 1.)
    @test get(ModelConstructors.regime_prior(u, 2, d)) == Uniform(0., 5.)
    @test get(ModelConstructors.regime_prior(u, 3, d)) == Normal(0., 1.)

    # Test set_regime_fixed! and regime_fixed
    ModelConstructors.set_regime_fixed!(u, 2, true, d)
    ModelConstructors.set_regime_fixed!(u, 3, false, d)
    @test_throws ErrorException ModelConstructors.set_regime_fixed!(u, 4, true, d)

    @test ModelConstructors.regime_fixed(u, 1)
    @test !ModelConstructors.regime_fixed(u, 2)
    @test ModelConstructors.regime_fixed(u, 2, d)
    @test !ModelConstructors.regime_fixed(u, 3, d)
    @test ModelConstructors.regime_valuebounds(u, 2, d) == (2.5230, 2.5230)
    @test ModelConstructors.regime_valuebounds(u, 3, d) == (1e-8, 5.)

    # test regime-switching value bounds
    ModelConstructors.set_regime_valuebounds!(u, 2, (1e-6, 5.), d)
    ModelConstructors.set_regime_valuebounds!(u, 3, (1e-7, 5.), d)
    ModelConstructors.set_regime_valuebounds!(u, 4, (1e-8, 5.), d)
    @test ModelConstructors.regime_valuebounds(u, 1) ==
        ModelConstructors.regime_valuebounds(u, 2, d) == (1e-6, 5.)
    @test ModelConstructors.regime_valuebounds(u, 2) ==
        ModelConstructors.regime_valuebounds(u, 3, d) == (1e-7, 5.)
    @test ModelConstructors.regime_valuebounds(u, 3) ==
        ModelConstructors.regime_valuebounds(u, 4, d) == (1e-8, 5.)

    # test set_regime_fixed w/update_valuebounds
    ModelConstructors.set_regime_fixed!(u, 2, true, d; update_valuebounds = (10., 11.))
    ModelConstructors.set_regime_fixed!(u, 3, false, d; update_valuebounds = (10., 11.))
    @test u.regimes[:valuebounds][1] == (10., 11.)
    @test u.regimes[:valuebounds][2] == (10., 11.)
end

nothing

max_stars_repo_name: FRBNY-DSGE/ModelConstructors

max_issues_repo_name: FRBNY-DSGE/ModelConstructors

max_forks_repo_name: FRBNY-DSGE/ModelConstructors


text:
println("===========================================================================")
println("Beginning execution of script GMHPhotoReceptor.jl/scripts/PhotoReceptor1.jl")
println("===========================================================================")

import GeneralizedMetropolisHastings
import GMHPhotoReceptor

println("Number of parallel processes running: ",nprocs())
println("Use addprocs() in the REPL if you want to run on more processes")

@everywhere using GeneralizedMetropolisHastings
@everywhere using GMHPhotoReceptor
using PyPlot

#Standard M-H for nproposals == 1
#Generalized M-H for nproposals > 1
nproposals = 100

#MCMC iteration specifications
nburnin = 100
niterations = 100
ntunerperiod = 10

###Values of the model
numvilli1 = 30000

#specify the values that determine the priors on the parameters
latencylocation = (2.0,3.5) #uniform distribution with (low,high) values
latencyscale = (0.2,0.7) #uniform distribution with (low,high) values
refractorylocation = (4.0,6.0) #uniform distribution with (low,high) values
refractoryscale = (1.5,2.5) #uniform distribution with (low,high) values
bumpamplitude = (3.0,5.0) #uniform distribution with (low,high) values
bumpshape = (log(3.0),0.3) #lognormal distribution with (location,scale) values; variable can vary roughly between 2.0 and 4.0, but becomes increasingly penalized outside
bumpscale = (log(2.5),0.3) #lognormal distribution with (location,scale) values; variable can vary roughly between 1.5 and 3.5, but becomes increasingly penalized outside

photonfilename = joinpath(splitdir(@__FILE__())[1],"../data/naturallight.jld")
photons1 = photonsequence(photonfilename)
current1 = lightinducedcurrent(photonfilename)

modelpolicy1 = policy(:photoreceptor) #4-parameter model with stochastic lognormal latency and refractory parameters and fixed bump parameters
params1 = parameters(:photoreceptor,modelpolicy1,latencylocation,latencyscale,refractorylocation,refractoryscale)

####Variance for the noise model, estimated from previous runs
variance1 = [3600.0]

println("==========================================")
println("Simulation parameters defined successfully")
println("==========================================")

###Create a PhotoReceptor model
model1 = model(:photoreceptor,params1,photons1,current1,variance1,numvilli1,modelpolicy1)

###Show the model
println("==========================")
println("Model defined successfully")
println("==========================")
show(model1)

###Plot the measurement data (simmulated data + noise)
figure("PhotoReceptor1") ; clf()
subplot(221)
plot(dataindex(model1),measurements(model1);label="measured",linewidth=2)
xlabel("Time")
ylabel("Current (nA)")
xlim(dataindex(model1)[1],dataindex(model1)[end])
title("Light-Induced Current")
grid("on")
legend(loc="upper right",fancybox="true")

###Create a Metropolis sampler with a Normal proposal density
propcov = [0.1,0.01,0.1,0.01]
sampler1 = sampler(:mh,:normal,0.1,propcov)
println("============================")
println("Sampler defined successfully")
println("============================")
show(sampler1)

###Create a tuner that scales the proposal density
#tuner1 = tuner(:scale,ntunerperiod,0.5,:erf)
tuner1 = tuner(:scale,ntunerperiod,0.5,:erf)
println("==========================")
println("Tuner defined successfully")
println("==========================")
show(tuner1)

###Create a Generalized Metropolis-Hastings runner (which will default to Standard MH when nproposals=1)
runnerpolicy1 = policy(:mh,nproposals;model=:stochastic,initialize=:prior)
runner1 = runner(runnerpolicy1,niterations,nproposals;numburnin = nburnin)
println("===========================")
println("Runner defined successfully")
println("===========================")
show(runner1)

###Run the MCMC (can take quite a bit of time)
println("=======================")
println("Run the MCMC simulation")
println("=======================")
@time chain1 = run!(runner1,model1,sampler1,tuner1)
println("=========================")
println("Completed MCMC simulation")
println("=========================")

###Show the result of the simulations
show(chain1)

nparas = numparas(model1)
meanparamvals = mean(samples(chain1),2)
stdparamvals = std(samples(chain1),2)

println("Results of the MCMC simulation:")
for i=1:nparas
    println("mean $(parameters(model1)[i].key): $(meanparamvals[i])")
    println("std $(parameters(model1)[i].key): $(stdparamvals[i])")
end

println("================")
println("Plotting results")
println("================")

###Plot the logposterior values across samples
subplot(222)
plot(1:numsamples(chain1),logposterior(chain1,model1))
title("Log-Posterior values across samples")
xlabel("Samples")
ylabel("Log-Posterior")

###Plot the histograms of parameter values
for i=1:nparas
    subplot(244 + i)
    bins = meanparamvals[i]-1.0:0.1:meanparamvals[i]+1.0
    h = PyPlot.plt[:hist](vec(getindex(samples(chain1),i,:)),bins)
    grid("on")
    title("Parameter $(parameters(model1)[i].key)")
    xlabel("Values")
    ylabel("Samples")
end

###Finally, plot the average model results in the data window
subplot(221)
modeldata = evaluate!(model1,vec(meanparamvals))
plot(dataindex(model1),modeldata;label="average model results",linewidth=2)
xlim(dataindex(model1)[1],dataindex(model1)[end])
legend(loc="upper right",fancybox="true")


text:
function UnivariateObjectiveFunction(arg0::UnivariateFunction)
    return UnivariateObjectiveFunction((UnivariateFunction,), arg0)
end

function get_objective_function(obj::UnivariateObjectiveFunction)
    return jcall(obj, "getObjectiveFunction", UnivariateFunction, ())
end


text:
# using Pkg
# Pkg.activate(@__DIR__)

# For linear algebra
using StaticArrays: SVector
using LinearAlgebra
using SparseArrays: SparseMatrixCSC, sparse
# Fast Cholesky
using KoLesky

# multiple dispatch
import IterativeSolvers: cg!
import LinearAlgebra: mul!, ldiv!
import Base: size

# parser
using ArgParse

# logging
using Logging

using ForwardDiff
using FastGaussQuadrature

# profile
using Plots

# solving Burgers: u_t+ u u_x- nu u_xx=0
function parse_commandline()
    s = ArgParseSettings()
    @add_arg_table s begin
        "--nu"
            help = "ν"
            arg_type = Float64
            default = 0.01/π
        "--kernel"
            arg_type = String
            default = "Matern7half"
        "--sigma"
            help = "lengthscale"
            arg_type = Float64
            default = 0.02
        "--h"
            help = "spatial grid size"
            arg_type = Float64
            default = 0.001
        "--dt"
            help = "time step size"
            arg_type = Float64
            default = 0.01
        "--T"
            help = "final time"
            arg_type = Float64
            default = 1.0
        "--nugget"
            arg_type = Float64
            default = 1e-10
        "--GNsteps"
            arg_type = Int
            default = 2
        "--rho_big"
            arg_type = Float64
            default = 3.0
        "--rho_small"
            arg_type = Float64
            default = 3.0
        "--k_neighbors"
            arg_type = Int
            default = 1
        "--compare_exact"
            arg_type = Bool
            default = false
    end
    return parse_args(s)
end

## PDEs type
abstract type AbstractPDEs end
struct Burgers1d{Tν} <: AbstractPDEs
    # eqn: u_t+ u u_x- νu_xx=0, (x,t) in [-1,1]*[0,1]
    ν::Tν
    bdy::Function
    rhs::Function
    init::Function
    init_∂x::Function
    init_∂∂x::Function
end

## sample points
function sample_points_grid(eqn::Burgers1d, h_in)
    X_domain = -1+h_in:h_in:1-h_in
    # X_domain = -1:h_in:1
    X_boundary = [-1,1]
    return collect(X_domain), X_boundary
end

## exact algorithms
# assemby Gram matrices
function get_Gram_matrices(eqn::Burgers1d, cov::KoLesky.AbstractCovarianceFunction, X_domain, X_boundary, dt, sol, sol_x)
    d = 1
    N_domain = length(X_domain)
    N_boundary = length(X_boundary)

    Δδ_coefs = -eqn.ν
    ∇δ_coefs = sol
    δ_coefs_int = 2/dt .+ sol_x # Crank-Nikoson discretization

    # get linearized PDEs correponding measurements
    meas_δ = [KoLesky.PointMeasurement{d}(SVector{d,Float64}([X_boundary[i]])) for i = 1:N_boundary]
    meas_Δ∇δ = [KoLesky.Δ∇δPointMeasurement{Float64,d}(SVector{d,Float64}([X_domain[i]]), Δδ_coefs, [∇δ_coefs[i]], δ_coefs_int[i]) for i = 1:N_domain]
    meas_δ_int = [KoLesky.PointMeasurement{d}(SVector{d,Float64}([X_domain[i]])) for i = 1:N_domain]

    Theta_train = zeros(N_domain+N_boundary,N_domain+N_boundary)

    measurements = Vector{Vector{<:KoLesky.AbstractPointMeasurement}}(undef,2)
    measurements[1] = meas_δ; measurements[2] = meas_Δ∇δ
    cov(Theta_train, reduce(vcat,measurements))

    meas_∇δ = [KoLesky.Δ∇δPointMeasurement{Float64,d}(SVector{d,Float64}([X_domain[i]]), 0.0, [1.0], 0.0) for i = 1:N_domain]
    meas_Δδ = [KoLesky.Δ∇δPointMeasurement{Float64,d}(SVector{d,Float64}([X_domain[i]]), 1.0, [0.0], 0.0) for i = 1:N_domain]
    measurements_test = Vector{Vector{<:KoLesky.AbstractPointMeasurement}}(undef,3)
    measurements_test[1] = meas_δ_int
    measurements_test[2] = meas_∇δ
    measurements_test[3] = meas_Δδ

    Theta_test = zeros(3*N_domain,N_domain+N_boundary)
    cov(Theta_test, reduce(vcat,measurements_test), reduce(vcat, measurements))
    return Theta_train, Theta_test

end

# iterative GPR
function iterGPR_exact(eqn, cov, X_domain, X_boundary, nugget, dt, T, GNsteps)
    N_domain = length(X_domain)
    N_boundary = length(X_boundary)
    Nt = convert(Int,T/dt)
    # get the rhs and bdy data
    # @show X_domain
    sol_u = [eqn.init(X_domain[i]) for i in 1:N_domain]
    sol_ux = [eqn.init_∂x(X_domain[i]) for i in 1:N_domain]
    sol_uxx = [eqn.init_∂∂x(X_domain[i]) for i in 1:N_domain]
    # @show sol_u
    bdy = [eqn.bdy(X_boundary[i]) for i in 1:N_boundary]
    rhs = [eqn.rhs(X_domain[i]) for i in 1:N_domain]

    rhs_CN = zeros(N_boundary+N_domain)
    rhs_CN[1:N_boundary] = bdy

    cur_rhs_CN = copy(rhs_CN)

    # measurements construction only once
    d = 1
    meas_δ_bd = [KoLesky.PointMeasurement{d}(SVector{d,Float64}([X_boundary[i]])) for i = 1:N_boundary]
    measurements = Vector{Vector{<:KoLesky.AbstractPointMeasurement}}(undef,2)
    measurements[1] = meas_δ_bd;

    meas_∇δ = [KoLesky.Δ∇δPointMeasurement{Float64,d}(SVector{d,Float64}([X_domain[i]]), 0.0, [1.0], 0.0) for i = 1:N_domain]
    meas_Δδ = [KoLesky.Δ∇δPointMeasurement{Float64,d}(SVector{d,Float64}([X_domain[i]]), 1.0, [0.0], 0.0) for i = 1:N_domain]
    meas_δ_int = [KoLesky.PointMeasurement{d}(SVector{d,Float64}([X_domain[i]])) for i = 1:N_domain]
    measurements_test = Vector{Vector{<:KoLesky.AbstractPointMeasurement}}(undef,3)
    measurements_test[1] = meas_δ_int
    measurements_test[2] = meas_∇δ
    measurements_test[3] = meas_Δδ
    Theta_train = zeros(N_domain+N_boundary,N_domain+N_boundary)
    Theta_test = zeros(3*N_domain,N_domain+N_boundary)
    @inbounds for it in 1:Nt
        t = it*dt
        @info "[time] t = $t"
        rhs_CN[N_boundary+1:end] .= rhs + 2/dt*sol_u + eqn.ν*sol_uxx - sol_u.*sol_ux
        for _ in 1:GNsteps
            cur_rhs_CN[N_boundary+1:end] .= rhs_CN[N_boundary+1:end] + sol_u.*sol_ux

            d = 1
            Δδ_coefs = -eqn.ν
            ∇δ_coefs = sol_u
            δ_coefs_int = 2/dt .+ sol_ux # Crank-Nikoson discretization
            meas_Δ∇δ = [KoLesky.Δ∇δPointMeasurement{Float64,d}(SVector{d,Float64}([X_domain[i]]), Δδ_coefs, [∇δ_coefs[i]], δ_coefs_int[i]) for i = 1:N_domain]
            measurements[2] = meas_Δ∇δ

            cov(Theta_train, reduce(vcat,measurements))
            cov(Theta_test, reduce(vcat,measurements_test), reduce(vcat, measurements))

            v = Theta_test*(((Theta_train+nugget*diagm(diag(Theta_train))))\cur_rhs_CN)

            @views sol_u .= v[1:N_domain]
            @views sol_ux .= v[N_domain+1:2*N_domain]
            @views sol_uxx .= v[2*N_domain+1:end]
        end
    end
    return sol_u
end

## algorithm using Kolesky and pcg
# struct that stores the factor of Theta_train
abstract type implicit_mtx end
struct approx_Theta_train{Tv,Ti,Tmtx<:SparseMatrixCSC{Tv,Ti}} <: implicit_mtx
    P::Vector{Ti}
    U::Tmtx
    L::Tmtx
    δ_coefs::Vector{Tv}
    ∇δ_coefs::Vector{Tv}
    Δδ_coef::Tv
    N_boundary::Ti
    N_domain::Ti
end
struct precond_Theta_train{Tv,Ti,Tmtx<:SparseMatrixCSC{Tv,Ti}} <: implicit_mtx
    P::Vector{Ti}
    U::Tmtx
    L::Tmtx
end

function size(A::approx_Theta_train, num)
    return size(A.U,num)
end

function mul!(x, Θtrain::approx_Theta_train, b)
    @views temp = vcat(b[1:Θtrain.N_boundary],Θtrain.δ_coefs.*b[Θtrain.N_boundary+1:end],Θtrain.∇δ_coefs.*b[Θtrain.N_boundary+1:end],Θtrain.Δδ_coef*b[Θtrain.N_boundary+1:end])
    temp[Θtrain.P] = Θtrain.L\(Θtrain.U\temp[Θtrain.P])

    @views x[1:Θtrain.N_boundary] = temp[1:Θtrain.N_boundary]
    @views x[Θtrain.N_boundary+1:end] .= Θtrain.δ_coefs.*(temp[Θtrain.N_boundary+1:Θtrain.N_boundary+Θtrain.N_domain]) .+ Θtrain.∇δ_coefs.* temp[Θtrain.N_boundary+Θtrain.N_domain+1:Θtrain.N_boundary+2*Θtrain.N_domain] + Θtrain.Δδ_coef * temp[Θtrain.N_boundary+2*Θtrain.N_domain+1:end]
end

function ldiv!(x, precond::precond_Theta_train, b)
    x[precond.P] = precond.U*(precond.L*b[precond.P])
end

function iterGPR_fast_pcg(eqn, cov, X_domain, X_boundary, dt, T, nugget, GNsteps; ρ_big =4.0, ρ_small=6.0, k_neighbors = 4, lambda = 1.5, alpha = 1.0)

    @info "[Algorithm]: iterative GPR + fast cholesky factorization + pcg"

    N_domain = length(X_domain)
    N_boundary = length(X_boundary)
    Nt = convert(Int,T/dt)
    # get the rhs and bdy data
    # @show X_domain
    sol_u = [eqn.init(X_domain[i]) for i in 1:N_domain]
    sol_ux = [eqn.init_∂x(X_domain[i]) for i in 1:N_domain]
    sol_uxx = [eqn.init_∂∂x(X_domain[i]) for i in 1:N_domain]
    # @show sol_u
    bdy = [eqn.bdy(X_boundary[i]) for i in 1:N_boundary]
    rhs = [eqn.rhs(X_domain[i]) for i in 1:N_domain]

    rhs_CN = zeros(N_boundary+N_domain)
    rhs_CN[1:N_boundary] = bdy

    cur_rhs_CN = copy(rhs_CN)

    # measurements construction only once
    d = 1
    meas_δ_bd = [KoLesky.PointMeasurement{d}(SVector{d,Float64}([X_boundary[i]])) for i = 1:N_boundary]
    meas_∇δ = [KoLesky.Δ∇δPointMeasurement{Float64,d}(SVector{d,Float64}([X_domain[i]]), 0.0, [1.0], 0.0) for i = 1:N_domain]
    meas_Δδ = [KoLesky.Δ∇δPointMeasurement{Float64,d}(SVector{d,Float64}([X_domain[i]]), 1.0, [0.0], 0.0) for i = 1:N_domain]
    meas_δ_int = [KoLesky.PointMeasurement{d}(SVector{d,Float64}([X_domain[i]])) for i = 1:N_domain]
    measurements = Vector{Vector{<:KoLesky.AbstractPointMeasurement}}(undef,4)
    measurements[1] = meas_δ_bd;
    measurements[2] = meas_δ_int
    measurements[3] = meas_∇δ
    measurements[4] = meas_Δδ

    @info "[Big Theta: implicit factorization] time"
    @time implicit_bigΘ = KoLesky.ImplicitKLFactorization_FollowDiracs(cov, measurements, ρ_big, k_neighbors; lambda = lambda, alpha = alpha)

    @info "[Big Theta: explicit factorization] time"
    @time explicit_bigΘ = KoLesky.ExplicitKLFactorization(implicit_bigΘ; nugget = nugget)
    U_bigΘ = explicit_bigΘ.U
    L_bigΘ = sparse(U_bigΘ')
    P_bigΘ = explicit_bigΘ.P

    Θtrain = approx_Theta_train(P_bigΘ, U_bigΘ, L_bigΘ,zeros(N_domain),zeros(N_domain), -eqn.ν, N_boundary,N_domain)

    implicit_factor = nothing

    measurement_train = Vector{Vector{<:KoLesky.AbstractPointMeasurement}}(undef,2)
    measurement_train[1] = meas_δ_bd;

    @inbounds for it in 1:Nt
        t = it*dt
        @info "[time] t = $t"
        rhs_CN[N_boundary+1:end] .= rhs + 2/dt*sol_u + eqn.ν*sol_uxx - sol_u.*sol_ux
        for step in 1:GNsteps
            @info "[Current GN step] $step"
            cur_rhs_CN[N_boundary+1:end] .= rhs_CN[N_boundary+1:end] + sol_u.*sol_ux

            d = 1
            Δδ_coefs = -eqn.ν
            ∇δ_coefs = sol_u
            δ_coefs_int = 2/dt .+ sol_ux # Crank-Nikoson discretization
            meas_Δ∇δ = [KoLesky.Δ∇δPointMeasurement{Float64,d}(SVector{d,Float64}([X_domain[i]]), Δδ_coefs, [∇δ_coefs[i]], δ_coefs_int[i]) for i = 1:N_domain]
            measurement_train[2] = meas_Δ∇δ

            @info "[Theta Train: implicit factorization] time"
            @time if implicit_factor === nothing
                implicit_factor = KoLesky.ImplicitKLFactorization(cov,  measurement_train, ρ_small, k_neighbors; lambda = lambda, alpha = alpha)
            else
                implicit_factor.supernodes.measurements .= reduce(vcat, collect.(measurement_train))[implicit_factor.P]
            end
            @info "[Theta Train: explicit factorization] time"
            @time explicit_factor = KoLesky.ExplicitKLFactorization(implicit_factor; nugget = nugget)

            U = explicit_factor.U
            L = sparse(U')
            P = explicit_factor.P

            # use the approximate solution as the initial point for the pCG iteration
            Θinv_rhs = U*(L*cur_rhs_CN[P])
            Θinv_rhs[P] = Θinv_rhs

            # pcg step for Theta_train\rhs
            Θtrain.δ_coefs .= δ_coefs_int
            Θtrain.∇δ_coefs .= ∇δ_coefs
            precond = precond_Theta_train(P,U,L)

            @info "[pcg started]"
            @time Θinv_rhs, ch = cg!(Θinv_rhs, Θtrain, cur_rhs_CN; Pl = precond, log=true)
            @info "[pcg finished], $ch"

            # get approximated sol_now = Theta_test * Θinv_rhs
            tmp = vcat(Θinv_rhs[1:N_boundary], δ_coefs_int .* Θinv_rhs[N_boundary+1:end], ∇δ_coefs.*Θinv_rhs[N_boundary+1:end], Δδ_coefs *Θinv_rhs[N_boundary+1:end])
            tmp[P_bigΘ] = L_bigΘ\(U_bigΘ\tmp[P_bigΘ])

            @views sol_u = tmp[N_boundary+1:N_domain+N_boundary]
            @views sol_ux .= tmp[N_domain+N_boundary+1:2*N_domain+N_boundary]
            @views sol_uxx .= tmp[2*N_domain+N_boundary+1:3*N_domain+N_boundary]
        end
    end
    @info "[solver finished]"
    return sol_u
end

function SolveBurgers_ColeHopf(x,t,ν)
    Gauss_pts, weights = gausshermite(100)
    temp = x.-sqrt(4*ν*t)*Gauss_pts
    val1 = weights .* sin.(π*temp) .* exp.(-cos.(π*temp)/(2*π*ν))
    val2 = weights .* exp.(-cos.(π*temp)/(2*π*ν))
    return -sum(val1)/sum(val2)
end

function main(args)
    ν = args.nu
    # ground truth solution
    function fun_u(x)
        return -sin(π*x)
    end

    function grad_u(x)
        return ForwardDiff.derivative(fun_u, x)
    end

    function Delta_u(x)
        return ForwardDiff.derivative(grad_u, x)
    end

    # right hand side
    function fun_rhs(x)
        return 0.0
    end

    # boundary value
    function fun_bdy(x)
        return 0.0
    end

    @info "[solver started] Burgers 1d"
    @info "[equation] u_t+ u u_x- $ν u_xx=0"
    eqn = Burgers1d(ν,fun_bdy,fun_rhs,fun_u,grad_u,Delta_u)

    h_in = args.h::Float64
    X_domain, X_boundary = sample_points_grid(eqn, h_in)
    N_domain = length(X_domain)
    N_boundary = length(X_boundary)
    @info "[sample points] grid size $h_in"
    @info "[sample points] N_domain is $N_domain, N_boundary is $N_boundary"

    lengthscale = args.sigma
    if args.kernel == "Matern5half"
        cov = KoLesky.MaternCovariance5_2(lengthscale)
    elseif args.kernel == "Matern7half"
        cov = KoLesky.MaternCovariance7_2(lengthscale)
    elseif args.kernel == "Matern9half"
        cov = KoLesky.MaternCovariance9_2(lengthscale)
    elseif args.kernel == "Matern11half"
        cov = KoLesky.MaternCovariance11_2(lengthscale)
    elseif args.kernel == "Gaussian"
        cov = KoLesky.GaussianCovariance(lengthscale)
    end
    @info "[kernel] choose $(args.kernel), lengthscale $lengthscale\n"

    nugget = args.nugget::Float64
    @info "[nugget] $nugget"

    dt = args.dt::Float64
    T = args.T::Float64

    GNsteps_approximate = args.GNsteps::Int
    @info "[total GN steps] $GNsteps_approximate"

    ρ_big = args.rho_big::Float64
    ρ_small = args.rho_small::Float64
    k_neighbors = args.k_neighbors::Int
    @info "[Fast Cholesky] ρ_big = $ρ_big, ρ_small = $ρ_small, k_neighbors = $k_neighbors"

    fast_solve() = @time iterGPR_fast_pcg(eqn, cov, X_domain, X_boundary, dt, T, nugget, GNsteps_approximate; ρ_big = ρ_big, ρ_small = ρ_small, k_neighbors=k_neighbors);
    sol = fast_solve()

    truth =  [SolveBurgers_ColeHopf(X_domain[i],T,ν) for i in 1:N_domain]
    pts_accuracy = sqrt(sum((truth-sol).^2)/N_domain)
    @info "[L2 accuracy: pCG method] $pts_accuracy"
    pts_max_accuracy = maximum(abs.(truth-sol))
    @info "[Linf accuracy: pCG method] $pts_max_accuracy"

    plot(X_domain,truth)
    plot!(X_domain, sol)
    if args.compare_exact
        GNsteps_exact = 2
        @info "[comparison: exact method]"
        @time sol_exact = iterGPR_exact(eqn, cov, X_domain, X_boundary, nugget, dt, T, GNsteps_exact)
        pts_accuracy_exact = sqrt(sum((truth-sol_exact).^2)/N_domain)
        @info "[L2 accuracy: exact method] $pts_accuracy_exact"
        pts_max_accuracy_exact = maximum(abs.(truth-sol_exact))
        @info "[Linf accuracy: exact method] $pts_max_accuracy_exact"
        plot!(X_domain,sol_exact)
    end

end

args = parse_commandline()
args = (; (Symbol(k) => v for (k,v) in args)...) # Named tuple from dict
main(args)


text:
function Abs()
    return Abs(())
end

function equals(obj::Object, arg0::Object)
    return jcall(obj, "equals", jboolean, (Object,), arg0)
end

function get_class(obj::Object)
    return jcall(obj, "getClass", Class, ())
end

function hash_code(obj::Object)
    return jcall(obj, "hashCode", jint, ())
end

function notify(obj::Object)
    return jcall(obj, "notify", void, ())
end

function notify_all(obj::Object)
    return jcall(obj, "notifyAll", void, ())
end

function to_string(obj::Object)
    return jcall(obj, "toString", JString, ())
end

function value(obj::Abs, arg0::Derivative)
    return jcall(obj, "value", Derivative, (Derivative,), arg0)
end

function value(obj::Abs, arg0::DerivativeStructure)
    return jcall(obj, "value", DerivativeStructure, (DerivativeStructure,), arg0)
end

function value(obj::Abs, arg0::jdouble)
    return jcall(obj, "value", jdouble, (jdouble,), arg0)
end

function wait(obj::Object)
    return jcall(obj, "wait", void, ())
end

function wait(obj::Object, arg0::jlong)
    return jcall(obj, "wait", void, (jlong,), arg0)
end

function wait(obj::Object, arg0::jlong, arg1::jint)
    return jcall(obj, "wait", void, (jlong, jint), arg0, arg1)
end


text:
@testset "MovingWindow" begin
    using Zarr
    a = Array{Union{Float64,Missing}}(rand(40, 20, 10))
    lon = RangeAxis("Lon", 1:40)
    lat = RangeAxis("Lat", 1:20)
    tim = RangeAxis("Time", 1:10)
    c = YAXArray([lon, lat, tim], a)
    d = tempname()
    c = savecube(c, d, chunksize = Dict("Lon" => 7, "Lat" => 9), backend = :zarr)

    indims = InDims("Time",YAXArrays.MovingWindow("Lon",1,1),window_oob_value = -9999.0)
    r1 = mapCube(c, indims=indims, outdims=OutDims("Time")) do xout,xin
        xout[:] = xin[:,1]
    end
    r2 = mapCube(c, indims = indims, outdims = OutDims("Time")) do xout, xin
        xout[:] = xin[:, 2]
    end
    r3 = mapCube(c, indims = indims, outdims = OutDims("Time")) do xout, xin
        xout[:] = xin[:, 3]
    end
    @test r1.data[:, 2:40, :] == permutedims(a[1:39, :, :], (3, 1, 2))
    @test all(==(-9999.0), r1.data[:, 1, :])
    @test r2.data[:, :, :] == permutedims(a[:, :, :], (3, 1, 2))
    @test r3.data[:, 1:39, :] == permutedims(a[2:40, :, :], (3, 1, 2))
    @test all(==(-9999.0), r3.data[:, 40, :])

    indims = InDims("Time",YAXArrays.MovingWindow("Lon",1,1),YAXArrays.MovingWindow("Lat",1,1))
    r1 = mapCube(c, indims=indims, outdims=OutDims("Time")) do xout,xin
        xout[:] = xin[:,1,1]
    end
    r2 = mapCube(c, indims = indims, outdims = OutDims("Time")) do xout, xin
        xout[:] = xin[:, 2, 2]
    end
    r3 = mapCube(c, indims = indims, outdims = OutDims("Time")) do xout, xin
        xout[:] = xin[:, 3, 3]
    end
    @test r1.data[:, 2:40, 2:20] == permutedims(a[1:39, 1:19, :], (3, 1, 2))
    @test all(ismissing, r1.data[:, 1, :])
    @test all(ismissing, r1.data[:, :, 1])
    @test r2.data[:, :, :] == permutedims(a, (3, 1, 2))
    @test r3.data[:, 1:39, 1:19] == permutedims(a[2:40, 2:20, :], (3, 1, 2))
    @test all(ismissing, r3.data[:, end, :])
    @test all(ismissing, r3.data[:, :, end])

    a = Array{Union{Float64,Missing}}(rand(10,4,  40, 20));
    varax = CategoricalAxis("Variable", 'a':'d')
    tim = RangeAxis("Time", 1:10)
    lon = RangeAxis("Lon", 1:40)
    lat = RangeAxis("Lat", 1:20)
    c = YAXArray([tim, varax, lon,lat], a)
    indims = InDims("Time",YAXArrays.MovingWindow("Lon",1,1))
    r1 = mapCube(c, indims=indims, outdims=OutDims("Time")) do xout,xin
        xout[:] = xin[:,1]
    end
    @test all(ismissing,r1[:,:,1,:])
    @test r1[:,:,2:40,:] == a[:,:,1:end-1,:]

end

max_forks_repo_name: meggart/YAXArrays


text:
#Wave equation Differential Equation
function square_grid_KG!(df,f,param,t)

    #Get Operators and variables
    Opps        = param[1]
    vars        = param[2]

    #Get numerical variables
    Nx  = vars.xnodes
    Ny  = vars.ynodes

    μ   = vars.μ

    #Get operators
    x   = Opps.xcoord
    y   = Opps.ycoord
    av  = Opps.aux_vex

    RHS     = Opps.DiffO_1
    BC      = Opps.BC_Mat
    Alpha   = Opps.Abs_Mat
    Lz      = Opps.Lz

    #Fields of our problem (Julia changes the rhs if you change the lhs)
    ψ   = f.x[1]
    ϕ   = f.x[2]

    dψ  = df.x[1]
    dϕ  = df.x[2]

    #Differential Equation
    if (vars.dissipation)

        #If Rotation is on,
        if( vars.Ω != 0.0 )

            #Update absorption matrix
            alpha!(Alpha , vars , Opps , t)

            #First Derivative
            dψ .=  ϕ

            #Second Derivative
            mul!(dϕ , RHS , ψ )
            av =  Alpha * (ϕ  + Lz * ψ)
            dϕ .=  dϕ .- av

        else

            # No need to get absorption matrix
            #alpha!( Alpha , vars , Opps , t)

            #First Derivative
            dψ .=  ϕ

            #Second derivative
            mul!(dϕ , RHS , ψ )
            av =  Alpha * (ϕ  + Lz * ψ)
            dϕ .=  dϕ .- av

        end

    else

        #First Derivative
        dψ .=  ϕ

        #Second Derivative
        mul!(dϕ , RHS , ψ )

    end

    #Boundary Conditions
    if(vars.Boundaries != :periodic )

        mul!(av , BC , dϕ )
       dϕ .= av

    end

end

function solve_wave_equation_2D(p::Param)

    #Welcome screen
    if( !setup_routine(p) )
        #End simulation
        println("Finished!")
        println(" ")
        println("Thank you for choosing KG Solver!")
        println("___________________________________________________________")
        print("Time Taken:  ")

        return
    end

    #Storage file
    println("Creating HDF5 file...")

    fdata   = "./data"
    fname   = fdata * p.folder *  "$(p.fname).h5"
    fid = h5open(fname, "w")

    HDF5.attributes(fid)["xmin"]      = p.xmin
    HDF5.attributes(fid)["xmax"]      = p.xmax
    HDF5.attributes(fid)["xnodes"]    = p.xnodes

    HDF5.attributes(fid)["ymin"]      = p.ymin
    HDF5.attributes(fid)["ymax"]      = p.ymax
    HDF5.attributes(fid)["ynodes"]    = p.ynodes

    if(p.Boundaries == :periodic )
        HDF5.attributes(fid)["Boundaries"]      = "periodic"
    else
        HDF5.attributes(fid)["Boundaries"]      = "radial"
    end

    if(p.dissipation == :true )
        HDF5.attributes(fid)["Dissipation"]     = "true"

        HDF5.attributes(fid)["R1"]          = p.R1
        HDF5.attributes(fid)["R2"]          = p.R2
        HDF5.attributes(fid)["Rorbit"]      = p.R_orbit

        HDF5.attributes(fid)["Omega"]           = p.Ω
        HDF5.attributes(fid)["alpha1"]          = p.alpha1
        HDF5.attributes(fid)["alpha2"]          = p.alpha2
    else
        HDF5.attributes(fid)["Dissipation"]     = "false"
    end

    println("OK")

    #Time variables
    tspan       = ( 0.0 , p.tmax )
    iter        = 0
    iter_save   = 0
    every       = floor(Int , p.out_every_t/p.deltat + 1.0)

    #Create operators structures
    println("Setting up Operators...")
    Operators = set_up_operator(p)
    println("OK")

    #Initial Conditions
    println("Setting up inicial configuration...")
    ψ , dψ = get_gaussian( p , Operators)

    U = ArrayPartition( Operators.BC_Mat * ψ , Operators.BC_Mat *dψ )
    println("OK")

    #Parameters
    my_params = ( Operators , p)

    #Problem ODE and integrator
    prob = ODEProblem(  square_grid_KG! , U , tspan , my_params )

    integrator = init( prob , RK4() , save_everystep=false , dt=p.deltat , adaptive=false )

    #Create a group for the functions
    g1 =  create_group(fid , "ψ");
    g2 =  create_group(fid , "dψ");

    #Save first configuration of field

    #data set names (first iteration is 0)
    dset_name = string("iteration_",iter_save)

    #data set for both displacement and derivative
    data1 = integrator.u.x[1]
    dset1, dtype1  = create_dataset(g1, dset_name , data1 )
    write_dataset(dset1, dtype1 , data1)

    data2 = integrator.u.x[2]
    dset2, dtype2  = create_dataset(g2, dset_name , data2 )
    write_dataset(dset2, dtype2 , data2)

    #Add HDF5.attributes time to both datasets
    HDF5.attributes(dset1)["time"] =  integrator.t
    HDF5.attributes(dset2)["time"] =  integrator.t

    #Progress bar and time evolution
    println("Starting up simulation...")
    prog1 = Progress( Int(floor(p.tmax))*100 )

    for (u,t) in tuples(integrator)

        #Increase iteration counter
        iter += 1

        #Step integrator
        DifferentialEquations.step!(integrator)

        #Update bar
        ProgressMeter.update!(prog1, Int(floor(integrator.t*100)))

        #If save iteration
        if iter % every == 0

            iter_save += 1

            #data set names
            dset_name = string("iteration_",iter_save)

            #data set for both displacement and derivative
            data1 = integrator.u.x[1]
            dset1, dtype1  = create_dataset(g1, dset_name , data1 )
            write_dataset(dset1, dtype1 , data1)

            data2 = integrator.u.x[2]
            dset2, dtype2  = create_dataset(g2, dset_name , data2 )
            write_dataset(dset2, dtype2 , data2)

            #Add HDF5.attributes time to both datasets
            HDF5.attributes(dset1)["time"] =  integrator.t
            HDF5.attributes(dset2)["time"] =  integrator.t

        end

    end

    #Add general information to file
    HDF5.attributes(fid)["max_iter"] = iter_save
    HDF5.attributes(fid)["max_time"] = integrator.t

    close(fid)
    println("Saving Files...")
    println("Finished!")

    #End simulation
    println(" ")
    println("Thank you for choosing KG Solver!")
    println("___________________________________________________________")

end


text:
module MakieRecipes

# using RecipesBase, MakieCore
# using RecipesBase: @recipe
# using MakieCore: Palette, to_color

# using RecipesPipeline
# using Colors

# # ## Palette
# # The default palette is defined here.

# expand_palette(palette, n = 20; kwargs...) = RGBA.(distinguishable_colors(n, palette; kwargs...))

# const wong = copy(wong_colors)
# begin
#     global wong
#     tmp = wong[1]
#     wong[1] = wong[2]
#     wong[2] = tmp
# end
# const rwong = expand_palette(wong, 50)
# const default_palette = Palette(rwong)

# include("bezier.jl")
# include("pipeline_integration.jl")
# include("attribute_table.jl")
# include("recipeplot.jl")

# # TODO FIXME
# RecipesBase.is_key_supported(::Symbol) = false
# plots(la::MakieLayout.LAxis) = plots(la.scene)
# # FIXME TODO

# function tomakie!(sc::AbstractScene, args...; attrs...)
#     RecipesPipeline.recipe_pipeline!(sc, Dict{Symbol, Any}(attrs), args)
# end

# tomakie!(args...; attrs...) = tomakie!(current_scene(), args...; attrs...)

# tomakie(args...; attrs...) = tomakie!(Scene(), args...; attrs...)

# export tomakie, tomakie!, recipeplot, recipeplot!

end


text:
@testset "149.max-points-on-a-line.jl" begin
    @test max_points([[1,1],[3,2],[5,3],[4,1],[2,3],[1,4]]) == 4
    @test max_points([[1,1],[2,2],[3,3]]) == 3
    @test max_points([[0,0],[94911150,94911151],[94911151,94911152]]) == 2
end


text:
@testset "map" begin
    function _mapview(f, d::AbstractDictionary)
        I = keytype(d)
        T = Core.Compiler.return_type(f, Tuple{eltype(d)})

        return MappedDictionary{I, T, typeof(f), Tuple{typeof(d)}}(f, (d,))
    end

    i = Indices([1,2,3,4,5])

    @test map(iseven, i)::Dictionary == dictionary([1=>false, 2=>true, 3=>false, 4=>true, 5=>false])
    @test map(isodd, i)::Dictionary == dictionary([1=>true, 2=>false, 3=>true, 4=>false, 5=>true])

    @test _mapview(iseven, i)::Dictionaries.MappedDictionary == dictionary([1=>false, 2=>true, 3=>false, 4=>true, 5=>false])
    @test _mapview(isodd, i)::Dictionaries.MappedDictionary == dictionary([1=>true, 2=>false, 3=>true, 4=>false, 5=>true])

    d = Dictionary([1,2,3,4,5], [1,3,2,4,5])

    @test map(iseven, d)::Dictionary == dictionary([1=>false, 2=>false, 3=>true, 4=>true, 5=>false])
    @test map(isodd, d)::Dictionary == dictionary([1=>true, 2=>true, 3=>false, 4=>false, 5=>true])
    @test map(+, d, d)::Dictionary == dictionary([1=>2, 2=>6, 3=>4, 4=>8, 5=>10])
    @test map(+, d, d, d)::Dictionary == dictionary([1=>3, 2=>9, 3=>6, 4=>12, 5=>15])

    @test _mapview(iseven, d)::Dictionaries.MappedDictionary == dictionary([1=>false, 2=>false, 3=>true, 4=>true, 5=>false])
    @test _mapview(isodd, d)::Dictionaries.MappedDictionary == dictionary([1=>true, 2=>true, 3=>false, 4=>false, 5=>true])
end


text:
function column_model_1D_kpp(constants, BCs, ICs, times, parameters=OceanTurb.KPP.Parameters())
    f, α, g, Nz, H = constants.f, constants.α, constants.g, constants.Nz, constants.H
    ρ₀ = 1027.0
    cₚ = 4000.0
    β  = 0.0
    kpp_constants = OceanTurb.Constants(Float64, ρ₀=ρ₀, cP=cₚ, f=f, α=α, β=β, g=g)

    model = OceanTurb.KPP.Model(N=Nz, H=H, stepper=:BackwardEuler, constants=kpp_constants, parameters=parameters)

    U₀ = ICs.u
    V₀ = ICs.v
    T₀ = ICs.T

    model.solution.U.data[1:Nz] .= U₀
    model.solution.V.data[1:Nz] .= V₀
    model.solution.T.data[1:Nz] .= T₀

    UW_flux = BCs.uw.top
    VW_flux = BCs.vw.top
    WT_flux = BCs.wT.top

    Δz = H / Nz

    ∂U₀∂z = (U₀[2] - U₀[1]) / Δz
    ∂V₀∂z = (V₀[2] - V₀[1]) / Δz
    ∂T₀∂z = (T₀[2] - T₀[1]) / Δz

    model.bcs.U.top = OceanTurb.FluxBoundaryCondition(UW_flux)
    model.bcs.V.top = OceanTurb.FluxBoundaryCondition(VW_flux)

    if WT_flux isa Number
        model.bcs.T.top = OceanTurb.FluxBoundaryCondition(WT_flux)
    else
        model.bcs.T.top = OceanTurb.FluxBoundaryCondition(WT_flux(0))
    end

    model.bcs.U.bottom = OceanTurb.GradientBoundaryCondition(∂U₀∂z)
    model.bcs.V.bottom = OceanTurb.GradientBoundaryCondition(∂V₀∂z)
    model.bcs.T.bottom = OceanTurb.GradientBoundaryCondition(∂T₀∂z)

    Nt = length(times)
    U = zeros(Nz, Nt)
    V = zeros(Nz, Nt)
    T = zeros(Nz, Nt)

    UW = zeros(Nz+1, Nt)
    VW = zeros(Nz+1, Nt)
    WT = zeros(Nz+1, Nt)

    # loop the model
    Δt = times[2] - times[1]
    for n in 1:Nt
        OceanTurb.run_until!(model, Δt, times[n])

        if !isa(WT_flux, Number)
            model.bcs.T.top.condition = WT_flux(model.clock.time)
        end

        U[:, n] .= model.solution.U[1:Nz]
        V[:, n] .= model.solution.V[1:Nz]
        T[:, n] .= model.solution.T[1:Nz]

        UW[:, n] .= OceanTurb.diffusive_flux(:U, model)[1:Nz+1]
        VW[:, n] .= OceanTurb.diffusive_flux(:V, model)[1:Nz+1]
        WT[:, n] .= OceanTurb.diffusive_flux(:T, model)[1:Nz+1] .+ OceanTurb.KPP.nonlocal_temperature_flux(model)[1:Nz+1]

        UW[Nz+1, n] = UW_flux
        VW[Nz+1, n] = VW_flux

        if WT_flux isa Number
            WT[Nz+1, n] = WT_flux
        else
            WT[Nz+1, n] = WT_flux(model.clock.time)
        end
    end

    return (; U, V, T, UW, VW, WT)
end