text:
using LibGit2

export ArchiveSource, FileSource, GitSource, DirectorySource

"""
An `AbstractSource` is something used as source to build the package.  Sources
are installed to `\${WORKSPACE}/srcdir` in the build environment.

Concrete subtypes of `AbstractSource` are:

* [`ArchiveSource`](@ref): a remote archive to download from the Internet;
* [`FileSource`](@ref): a remote file to download from the Internet;
* [`GitSource`](@ref): a remote Git repository to clone;
* [`DirectorySource`](@ref): a local directory to mount.
"""
abstract type AbstractSource end

"""
    ArchiveSource(url::String, hash::String; unpack_target::String = "")

Specify a remote archive in one of the supported archive formats (e.g., TAR or
ZIP balls) to be downloaded from the Internet from `url`.  `hash` is the
64-character SHA256 checksum of the file.

In the builder environment, the archive will be automatically unpacked to
`\${WORKSPACE}/srcdir`, or in its subdirectory pointed to by the optional
keyword `unpack_target`, if provided.
"""
struct ArchiveSource <: AbstractSource
    url::String
    hash::String
    unpack_target::String
end
ArchiveSource(url::String, hash::String; unpack_target::String = "") =
    ArchiveSource(url, hash, unpack_target)

# List of optionally compressed TAR archives that we know how to deal with
const tar_extensions = [".tar", ".tar.gz", ".tgz", ".tar.bz", ".tar.bz2",
                        ".tar.xz", ".tar.Z", ".txz", ".tar.zst"]
# List of general archives that we know about
const archive_extensions = vcat(tar_extensions, ".zip")

"""
    FileSource(url::String, hash::String; filename::String = basename(url))

Specify a remote file to be downloaded from the Internet from `url`.  `hash` is
the 64-character SHA256 checksum of the file.

In the builder environment, the file will be saved under `\${WORKSPACE}/srcdir`
with the same name as the basename of the originating URL, unless the the
keyword argument `filename` is specified.
"""
struct FileSource <: AbstractSource
    url::String
    hash::String
    filename::String
end
FileSource(url::String, hash::String; filename::String = basename(url)) =
    FileSource(url, hash, filename)

"""
    GitSource(url::String, hash::String; unpack_target::String = "")

Specify a remote Git repository to clone form `url`.  `hash` is the 40-character
SHA1 revision to checkout after cloning.

The repository will be cloned in `\${WORKSPACE}/srcdir`, or in its subdirectory
pointed to by the optional keyword `unpack_target`, if provided.
"""
struct GitSource <: AbstractSource
    url::String
    hash::String
    unpack_target::String
end
GitSource(url::String, hash::String; unpack_target::String = "") =
    GitSource(url, hash, unpack_target)

"""
    DirectorySource(path::String; target::String = basename(path), follow_symlinks=false)

Specify a local directory to mount from `path`.

The content of the directory will be mounted in `\${WORKSPACE}/srcdir`, or in
its subdirectory pointed to by the optional keyword `target`, if provided.
Symbolic links are replaced by a copy of the target when `follow_symlinks` is
`true`.
"""
struct DirectorySource <: AbstractSource
    path::String
    target::String
    follow_symlinks::Bool
end
# When setting up the source, by default we won't follow symlinks.  However,
# there are cases where this is necessary, for example when we have symlink
# patchsets across multiple versions of GCC, etc...
DirectorySource(path::String; target::String = "", follow_symlinks::Bool=false) =
    DirectorySource(path, target, follow_symlinks)

# This is not meant to be used as source in the `build_tarballs.jl` scripts but
# only to set up the source in the workspace.
struct SetupSource{T<:AbstractSource}
    path::String
    hash::String
    target::String
    follow_symlinks::Bool
end
# `follow_symlinks` is used only for DirectorySource, let's have a method without it.
SetupSource{T}(path::String, hash::String, target::String) where {T} =
    SetupSource{T}(path, hash, target, false)
# This is used in wizard/obtain_source.jl to automatically guess the parameter
# of SetupSource from the URL
function SetupSource(url::String, path::String, hash::String, target::String)
    if endswith(url, ".git")
        return SetupSource{GitSource}(path, hash, target)
    elseif any(endswith(url, ext) for ext in archive_extensions)
        return SetupSource{ArchiveSource}(path, hash, target)
    else
        return SetupSource{FileSource}(path, hash, target)
    end
end

struct PatchSource
    name::String
    patch::String
end

function download_source(source::T; verbose::Bool = false, downloads_dir = storage_dir("downloads")) where {T<:Union{ArchiveSource,FileSource}}
    gettarget(s::ArchiveSource) = s.unpack_target
    gettarget(s::FileSource) = s.filename
    if isfile(source.url)
        # Immediately abspath() a src_url so we don't lose track of
        # sources given to us with a relative path
        src_path = abspath(source.url)

        # And if this is a locally-sourced tarball, just verify
        verify(src_path, source.hash)
    else
        # Otherwise, download and verify
        src_path = joinpath(downloads_dir, string(source.hash, "-", basename(source.url)))
        download_verify(source.url, source.hash, src_path)
    end
    return SetupSource{T}(src_path, source.hash, gettarget(source))
end

function download_source(source::GitSource; verbose::Bool = false, downloads_dir = storage_dir("downloads"))
    src_path = joinpath(downloads_dir, basename(source.url))

    # If this git repository already exists, ensure that its origin remote actually matches
    if isdir(src_path)
        origin_url = LibGit2.with(LibGit2.GitRepo(src_path)) do repo
            LibGit2.url(LibGit2.get(LibGit2.GitRemote, repo, "origin"))
        end

        # If the origin url doesn't match, wipe out this git repo.  We'd rather have a
        # thrashed cache than an incorrect cache.
        if origin_url != source.url
            rm(src_path; recursive=true, force=true)
        end
    end

    if isdir(src_path)
        @info "Cached repository found in $(src_path)"
        # If we didn't just mercilessly obliterate the cached git repo, use it!
        LibGit2.with(LibGit2.GitRepo(src_path)) do repo
            LibGit2.fetch(repo)
        end
    else
        @info "Cloning $(source.url) to $(src_path)..."
        # If there is no src_path yet, clone it down.
        repo = LibGit2.clone(source.url, src_path; isbare=true)
    end
    return SetupSource{GitSource}(src_path, source.hash, source.unpack_target)
end

function download_source(source::DirectorySource; verbose::Bool = false)
    if !isdir(source.path)
        error("Could not find directory \"$(source.path)\".")
    end
    @info "Directory \"$(source.path)\" found"
    return SetupSource{DirectorySource}(abspath(source.path), "", source.target, source.follow_symlinks)
end

"""
    download_source(source::AbstractSource; verbose::Bool = false)

Download the given `source`.  All downloads are cached within the
BinaryBuilder `downloads` storage directory.
"""
download_source

# Add JSON serialization to sources
JSON.lower(fs::ArchiveSource) = Dict("type" => "archive", extract_fields(fs)...)
JSON.lower(fs::FileSource) = Dict("type" => "file", extract_fields(fs)...)
JSON.lower(gs::GitSource) = Dict("type" => "git", extract_fields(gs)...)
JSON.lower(ds::DirectorySource) = Dict("type" => "directory", extract_fields(ds)...)

# When deserialiasing the JSON file, the sources are in the form of
# dictionaries.  This function converts the dictionary back to the appropriate
# AbstractSource.
function sourcify(d::Dict)
    if d["type"] == "directory"
        return DirectorySource(d["path"])
    elseif d["type"] == "git"
        return GitSource(d["url"], d["hash"])
    elseif d["type"] == "file"
        return FileSource(d["url"], d["hash"])
    elseif d["type"] == "archive"
        return ArchiveSource(d["url"], d["hash"])
    else
        error("Cannot convert to source")
    end
end

# XXX: compatibility functions.  These are needed until we support old-style
# Pair/String sources specifications.
coerce_source(source::AbstractSource) = source
function coerce_source(source::AbstractString)
    @warn "Using a string as source is deprecated, use DirectorySource instead"
    return DirectorySource(source)
end
function coerce_source(source::Pair)
    src_url, src_hash = source
    if endswith(src_url, ".git")
        @warn "Using a pair as source is deprecated, use GitSource instead"
        return GitSource(src_url, src_hash)
    elseif any(endswith(src_url, ext) for ext in archive_extensions)
        @warn "Using a pair as source is deprecated, use ArchiveSource instead"
        return ArchiveSource(src_url, src_hash)
    else
        @warn "Using a pair as source is deprecated, use FileSource instead"
        return FileSource(src_url, src_hash)
    end
end


text:
using HTTPStatus
include("HTTPStatus.jl")

ext: jl

max_stars_repo_name: kracekumar/HTTPStatus

max_issues_repo_name: kracekumar/HTTPStatus

max_forks_repo_name: kracekumar/HTTPStatus


text:
abstract type AbstractGrid{dim, T} end

"""
```
struct RectilinearGrid{dim, T, N, M} <: AbstractGrid{dim, T}
    grid::JuAFEM.Grid{dim, N, T, M}
    nels::NTuple{dim, Int}
    sizes::NTuple{dim, T}
    corners::NTuple{2, JuAFEM.Vec{dim, T}}
end
```

`dim`: dimension of the problem

`T`: number type for computations and coordinates

`N`: number of nodes in a cell of the grid

`M`: number of faces in a cell of the grid

`grid`: a JuAFEM.Grid struct

`nels`: number of elements in every dimension

`sizes`: dimensions of each rectilinear cell

`corners`: 2 corner points of the rectilinear grid

API:

`function RectilinearGrid(nels::NTuple{dim,Int}, sizes::NTuple{dim,T}) where {dim,T}`

Example:

`rectgrid = RectilinearGrid((60,20), (1.0,1.0))`
"""
struct RectilinearGrid{dim, T, N, M} <: AbstractGrid{dim, T}
    grid::JuAFEM.Grid{dim, N, T, M}
    nels::NTuple{dim, Int}
    sizes::NTuple{dim, T}
    corners::NTuple{2, JuAFEM.Vec{dim, T}}
end
nnodespercell(::RectilinearGrid{dim,T,N,M}) where {dim, T, N, M} = N
nfacespercell(::RectilinearGrid{dim,T,N,M}) where {dim, T, N, M} = M

left(rectgrid::RectilinearGrid, x) = x[1] ≈ rectgrid.corners[1][1]
right(rectgrid::RectilinearGrid, x) = x[1] ≈ rectgrid.corners[2][1]
bottom(rectgrid::RectilinearGrid, x) = x[2] ≈ rectgrid.corners[1][2]
top(rectgrid::RectilinearGrid, x) = x[2] ≈ rectgrid.corners[2][2]
back(rectgrid::RectilinearGrid, x) = x[3] ≈ rectgrid.corners[1][3]
front(rectgrid::RectilinearGrid, x) = x[3] ≈ rectgrid.corners[2][3]
middlex(rectgrid::RectilinearGrid, x) = x[1] ≈ (rectgrid.corners[1][1] + rectgrid.corners[2][1]) / 2
middley(rectgrid::RectilinearGrid, x) = x[2] ≈ (rectgrid.corners[1][2] + rectgrid.corners[2][2]) / 2
middlez(rectgrid::RectilinearGrid, x) = x[3] ≈ (rectgrid.corners[1][3] + rectgrid.corners[2][3]) / 2

nnodes(cell::Type{JuAFEM.Cell{dim,N,M}}) where {dim, N, M} = N
nnodes(cell::JuAFEM.Cell) = nnodes(typeof(cell))

function RectilinearGrid(nels::NTuple{dim,Int}, sizes::NTuple{dim,T}) where {dim,T}
    geoshape = dim === 2 ? Quadrilateral : Hexahedron
    corner1 = JuAFEM.Vec{dim}(fill(T(0), dim))
    corner2 = JuAFEM.Vec{dim}((nels .* sizes))
    grid = generate_grid(geoshape, nels, corner1, corner2);

    N = nnodes(geoshape)
    M = JuAFEM.nfaces(geoshape)
    ncells = prod(nels)
    return RectilinearGrid{dim, T, N, M}(grid, nels, sizes, (corner1, corner2))
end


text:
module SmoothAPI

using ImageBase
using ImageBase.ImageCore: GenericGrayImage, GenericImage

"""
    AbstractImageAlgorithm

The root of image algorithms type system
"""
abstract type AbstractImageAlgorithm end

"""
    AbstractImageFilter <: AbstractImageAlgorithm

Filters are image algorithms whose input and output are both images
"""
abstract type AbstractImageFilter <: AbstractImageAlgorithm end

include("smooth.jl")

end # module SmoothAPI


text:
"""
    solve(problem::IntervalConstantBTPDE, gradient; callbacks = AbstractCallback[])

Solve the Bloch-Torrey partial differential equation using P1 finite elements.
This function uses a manual time stepping scheme (theta-rule), that requires a degree of
implicitness `θ` and a time step `Δt`.
    `θ = 0.5`: Crank-Nicolson (second order)
    `θ = 1.0`: Implicit Euler (first order)
"""
function solve(
    problem::IntervalConstantBTPDE{T},
    gradient::ScalarGradient;
    callbacks = AbstractCallback[],
) where {T}
    (; θ, timestep, model, matrices) = problem
    (; γ) = model
    (; M, S, R, Mx, Q) = matrices

    isconstant(gradient.profile) || error("Time profile must be interval-wise constant")
    ivals = intervals(gradient.profile)

    ρ = initial_conditions(model)

    Jac!(J, g⃗) =
        @.(J = -(S + Q + R + im * γ * (g⃗[1] * Mx[1] + g⃗[2] * Mx[2] + g⃗[3] * Mx[3])))

    t = 0.0
    ξ = copy(ρ)
    Ey = copy(ξ)
    J = -(S + Q + im * sum(Mx)) # Allocate sparsity pattern

    for cb ∈ callbacks
        initialize!(cb, problem, gradient, ξ, t)
    end

    # Crank-Nicolson time stepping
    for i = 1:(length(ivals)-1)
        @debug "Solving for interval [%g, %g]\n" ivals[i] ivals[i+1]

        # Adjust time step to divide interval uniformly
        ival_length = ivals[i+1] - ivals[i]
        nt = round(Int, ival_length / timestep)
        Δt = ival_length / nt

        # Compute gradient at midpoint of interval
        tmid = (ivals[i] + ivals[i+1]) / 2
        g⃗ = gradient(tmid)

        # Build matrices for interval
        Jac!(J, g⃗)
        F = lu(M .- Δt .* θ .* J)
        # F = factorize(M .- Δt .* θ .* J)
        E = @. complex(M) + Δt * (1 - θ) * J

        # Step to end of interval
        for _ = 1:nt
            mul!(Ey, E, ξ)
            ldiv!(ξ, F, Ey)
            t += Δt
            for cb ∈ callbacks
                update!(cb, problem, gradient, ξ, t)
            end
        end
    end

    finalize!.(callbacks)

    ξ
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 IoK8sApiCoreV1EndpointAddress <: SwaggerModel
    hostname::Any # spec type: Union{ Nothing, String } # spec name: hostname
    ip::Any # spec type: Union{ Nothing, String } # spec name: ip
    nodeName::Any # spec type: Union{ Nothing, String } # spec name: nodeName
    targetRef::Any # spec type: Union{ Nothing, IoK8sApiCoreV1ObjectReference } # spec name: targetRef

    function IoK8sApiCoreV1EndpointAddress(;hostname=nothing, ip=nothing, nodeName=nothing, targetRef=nothing)
        o = new()
        validate_property(IoK8sApiCoreV1EndpointAddress, Symbol("hostname"), hostname)
        setfield!(o, Symbol("hostname"), hostname)
        validate_property(IoK8sApiCoreV1EndpointAddress, Symbol("ip"), ip)
        setfield!(o, Symbol("ip"), ip)
        validate_property(IoK8sApiCoreV1EndpointAddress, Symbol("nodeName"), nodeName)
        setfield!(o, Symbol("nodeName"), nodeName)
        validate_property(IoK8sApiCoreV1EndpointAddress, Symbol("targetRef"), targetRef)
        setfield!(o, Symbol("targetRef"), targetRef)
        o
    end
end # type IoK8sApiCoreV1EndpointAddress

const _property_map_IoK8sApiCoreV1EndpointAddress = Dict{Symbol,Symbol}(Symbol("hostname")=>Symbol("hostname"), Symbol("ip")=>Symbol("ip"), Symbol("nodeName")=>Symbol("nodeName"), Symbol("targetRef")=>Symbol("targetRef"))
const _property_types_IoK8sApiCoreV1EndpointAddress = Dict{Symbol,String}(Symbol("hostname")=>"String", Symbol("ip")=>"String", Symbol("nodeName")=>"String", Symbol("targetRef")=>"IoK8sApiCoreV1ObjectReference")
Base.propertynames(::Type{ IoK8sApiCoreV1EndpointAddress }) = collect(keys(_property_map_IoK8sApiCoreV1EndpointAddress))
Swagger.property_type(::Type{ IoK8sApiCoreV1EndpointAddress }, name::Symbol) = Union{Nothing,eval(Base.Meta.parse(_property_types_IoK8sApiCoreV1EndpointAddress[name]))}
Swagger.field_name(::Type{ IoK8sApiCoreV1EndpointAddress }, property_name::Symbol) =  _property_map_IoK8sApiCoreV1EndpointAddress[property_name]

function check_required(o::IoK8sApiCoreV1EndpointAddress)
    (getproperty(o, Symbol("ip")) === nothing) && (return false)
    true
end

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


text:
using GCN
using Test

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

max_stars_repo_name: alok/GCN

max_issues_repo_name: alok/GCN

max_forks_repo_name: alok/GCN


text:
using OrdinaryDiffEq
using Trixi

###############################################################################
# semidiscretization of the compressible Euler equations
gamma = 5/3
equations = CompressibleEulerEquations2D(gamma)

"""
    initial_condition_blob(x, t, equations::CompressibleEulerEquations2D)

The blob test case taken from
- Agertz et al. (2006)
  Fundamental differences between SPH and grid methods
  [arXiv: astro-ph/0610051](https://arxiv.org/abs/astro-ph/0610051)
"""
function initial_condition_blob(x, t, equations::CompressibleEulerEquations2D)
  # blob test case, see Agertz et al. https://arxiv.org/pdf/astro-ph/0610051.pdf
  # other reference: https://arxiv.org/pdf/astro-ph/0610051.pdf
  # change discontinuity to tanh
  # typical domain is rectangular, we change it to a square
  # resolution 128^2, 256^2
  # domain size is [-20.0,20.0]^2
  # gamma = 5/3 for this test case
  R = 1.0 # radius of the blob
  # background density
  dens0 = 1.0
  Chi = 10.0 # density contrast
  # reference time of characteristic growth of KH instability equal to 1.0
  tau_kh = 1.0
  tau_cr = tau_kh/1.6 # crushing time
  # determine background velocity
  velx0 = 2*R*sqrt(Chi)/tau_cr
  vely0 = 0.0
  Ma0 = 2.7 # background flow Mach number Ma=v/c
  c = velx0/Ma0 # sound speed
  # use perfect gas assumption to compute background pressure via the sound speed c^2 = gamma * pressure/density
  p0 = c*c*dens0/equations.gamma
  # initial center of the blob
  inicenter = [-15,0]
  x_rel = x-inicenter
  r = sqrt(x_rel[1]^2 + x_rel[2]^2)
  # steepness of the tanh transition zone
  slope = 2
  # density blob
  dens = dens0 + (Chi-1) * 0.5*(1+(tanh(slope*(r+R)) - (tanh(slope*(r-R)) + 1)))
  # velocity blob is zero
  velx = velx0 - velx0 * 0.5*(1+(tanh(slope*(r+R)) - (tanh(slope*(r-R)) + 1)))
  return prim2cons(SVector(dens, velx, vely0, p0), equations)
end
initial_condition = initial_condition_blob

surface_flux = flux_lax_friedrichs
volume_flux  = flux_ranocha
basis = LobattoLegendreBasis(4)

indicator_sc = IndicatorHennemannGassner(equations, basis,
                                         alpha_max=0.4,
                                         alpha_min=0.0001,
                                         alpha_smooth=true,
                                         variable=pressure)

volume_integral = VolumeIntegralShockCapturingHG(indicator_sc;
                                                 volume_flux_dg=volume_flux,
                                                 volume_flux_fv=surface_flux)

solver = DGSEM(basis, surface_flux, volume_integral)

coordinates_min = (-20.0, -20.0)
coordinates_max = ( 20.0,  20.0)

mesh = TreeMesh(coordinates_min, coordinates_max,
                initial_refinement_level=6,
                n_cells_max=100_000,)

semi = SemidiscretizationHyperbolic(mesh, equations, initial_condition, solver)

###############################################################################
# ODE solvers, callbacks etc.

tspan = (0.0, 8.0)
ode = semidiscretize(semi, tspan)

summary_callback = SummaryCallback()

analysis_interval = 100
analysis_callback = AnalysisCallback(semi, interval=analysis_interval)

alive_callback = AliveCallback(analysis_interval=analysis_interval)

save_solution = SaveSolutionCallback(interval=100,
                                     save_initial_solution=true,
                                     save_final_solution=true,
                                     solution_variables=cons2prim)

amr_indicator = IndicatorHennemannGassner(semi,
                                          alpha_max=1.0,
                                          alpha_min=0.0001,
                                          alpha_smooth=false,
                                          variable=Trixi.density)
amr_controller = ControllerThreeLevelCombined(semi, amr_indicator, indicator_sc,
                                              base_level=4,
                                              med_level =0, med_threshold=0.0003, # med_level = current level
                                              max_level =7, max_threshold=0.003,
                                              max_threshold_secondary=indicator_sc.alpha_max)
amr_callback = AMRCallback(semi, amr_controller,
                           interval=1,
                           adapt_initial_condition=true,
                           adapt_initial_condition_only_refine=true)

stepsize_callback = StepsizeCallback(cfl=0.25)

callbacks = CallbackSet(summary_callback,
                        analysis_callback, alive_callback,
                        save_solution,
                        amr_callback, stepsize_callback)

###############################################################################
# run the simulation

sol = solve(ode, CarpenterKennedy2N54(williamson_condition=false),
            dt=1.0, # solve needs some value here but it will be overwritten by the stepsize_callback
            save_everystep=false, callback=callbacks);
summary_callback() # print the timer summary


text:
# # generic fallback, probably slow. TODO: can we have a fast version of this?
# # Is a macro more suitable here?
# function Tensors.tomandel!(v::Vector{T}, r::Residuals) where T
#     # TODO check for length of v
#     start = 1
#     for fn in fieldnames(typeof(r))
#         start = _set_mandel_values!(v, getfield(r, fn), start)
#     end
#     return v
# end

# function _set_mandel_values!(v::Vector{T}, r::SymmetricTensor{order, dim, T, M}, start::Int) where {order, dim, T, M}
#     tomandel!(v[start:end], r)
#     return start + M
# end
# function _set_mandel_values!(v::Vector{T}, r::T, start::Int) where T
#     v[start] = r
#     return start + 1
# end

function vector_residual!(R::Function, r_vector::Vector{T}, x_vector::Vector{T}, m) where T
    # construct residuals with type T
    x_tensor = frommandel(Tensors.get_base(typeof(m)), x_vector)
    r_tensor = R(x_tensor)
    tomandel!(r_vector, r_tensor)
    return r_vector
end

function update_cache!(cache, f)
    cache.f = f
    chunk = ForwardDiff.Chunk(cache.x_f)
    jac_cfg = ForwardDiff.JacobianConfig(cache.f, cache.F, cache.x_f, chunk)
    ForwardDiff.checktag(jac_cfg, cache.f, cache.x_f)

    F2 = copy(cache.F)
    function j_forwarddiff!(J, x)
        ForwardDiff.jacobian!(J, cache.f, F2, x, jac_cfg, Val{false}())
    end
    function fj_forwarddiff!(F, J, x)
        jac_res = DiffResults.DiffResult(F, J)
        ForwardDiff.jacobian!(jac_res, cache.f, F2, x, jac_cfg, Val{false}())
        DiffResults.value(jac_res)
    end

    cache.df = j_forwarddiff!
    cache.fdf = fj_forwarddiff!
    return cache
end


text:
"""
A two dimensional nonlinear discrete time (DT) system.
"""

Base.@kwdef struct TwoDimensionalNonlinearDTSystem <: AbstractEnv
    c = 0.5
    d = 1.0
end

function State(env::TwoDimensionalNonlinearDTSystem)
    return function (x1, x2)
        ComponentArray(x1=x1, x2=x2)
    end
end

function V1(env::TwoDimensionalNonlinearDTSystem)
    return function (x)
        if x > 0
            ans = x^(1/4)
        else
            ans = x^4
        end
        ans
    end
end

function Dynamics!(env::TwoDimensionalNonlinearDTSystem)
    @unpack c, d = env
    @Loggable function dynamics!(dx, x, p, t; u)
        @unpack x1, x2 = x
        @log state = x
        @log input = u
        dx.x1 = c*(x1 + x2)  # x1_next = dx.x1
        dx.x2 = c*x2 + u[1]  # x2_next = dx.x2
        @onlylog next_state = dx
    end
end

function CubicSolution(env::TwoDimensionalNonlinearDTSystem)
    @unpack c, d = env
    return function (x2)
        a = (
             (cbrt(sqrt(3)*sqrt(27*c^2*x2^2 + 8) + 9*c*x2) / 3^(2/3))
             - (2 / (cbrt(3) * cbrt(sqrt(3)*sqrt(27*c^2*x2^2 + 8) + 9*c*x2)))
            )  # solve a^3 + 2a - 2*c*x2 for a ∈ ℝ
    end
end

function RunningCost(env::TwoDimensionalNonlinearDTSystem)
    @unpack c, d = env
    return function (x, u)
        @unpack x1, x2 = x
        dx = copy(x)  # will be next state
        Dynamics!(env)(dx, x, nothing, nothing; u=u)
        a = CubicSolution(env)(x2)
        r̄ = (
             (3/4)*a^4 + a^2
             - d*V1(env)(dx.x1) + d*V1(env)(x1)
             + (1-c^2)*x2^2
            )  # for V(x) = d*sqrt(abs(x1)) + x2^2
        r = r̄ + (1/4)*u[1]^4
    end
end

function OptimalValue(env::TwoDimensionalNonlinearDTSystem)
    @unpack c, d = env
    return function (x)
        @unpack x1, x2 = x
        d*V1(env)(x1) + x2^2
    end
end

function OptimalQValue(env::TwoDimensionalNonlinearDTSystem)
    return function (x, u)
        dx = copy(x)
        Dynamics!(env)(dx, x, nothing, nothing; u=u)
        OptimalValue(env)(dx) + RunningCost(env)(x, u)  # V(x_next) + r(x, u)
    end
end

function OptimalControl(env::TwoDimensionalNonlinearDTSystem)
    return function (x)
        @unpack x2 = x
        a = CubicSolution(env)(x2)
        -a
    end
end


text:
export MachOHandle, FatMachOHandle

struct MachOHandle{T <: IO} <: ObjectHandle
    # Backing IO and start point within the IOStream of this MachO object
    io::T
    start::Int

    # The parsed-out header of the MachO object
    header::MachOHeader

    # The path of the file this was created with, if it exists
    path::String
end

function readmeta(io::IO,::Type{MachOHandle})
    start = position(io)

    # Peek at the magic
    magic = read(io,UInt32)
    seek(io, start)

    # See which magic it is:
    header_type = macho_header_type(magic)
    endianness = macho_endianness(magic)

    # Unpack the header
    header = unpack(io, header_type, endianness)
    return MachOHandle(io, start, header, path(io))
end

## IOStream-like operations:
start(oh::MachOHandle) = oh.start
iostream(oh::MachOHandle) = oh.io

## Format-specific properties:
header(oh::MachOHandle) = oh.header
endianness(oh::MachOHandle) = macho_endianness(header(oh).magic)
is64bit(oh::MachOHandle) = macho_is64bit(header(oh).magic)
isrelocatable(oh::MachOHandle) = header(oh).filetype == MH_OBJECT
isexecutable(oh::MachOHandle) = header(oh).filetype == MH_EXECUTE
islibrary(oh::MachOHandle) = header(oh).filetype == MH_DYLIB
mangle_section_names(oh::MachOHandle, name) = string("__", name)
mangle_symbol_name(oh::MachOHandle, name::AbstractString) = string("_", name)
format_string(::Type{H}) where {H <: MachOHandle} = "MachO"

# Section information
section_header_size(oh::MachOHandle) = sizeof(section_header_type(oh))
function section_header_type(oh::H) where {H <: MachOHandle}
    if is64bit(oh)
        return MachOSection64{H}
    else
        return MachOSection32{H}
    end
end

# Misc. stuff
path(oh::MachOHandle) = oh.path


text:
# Plot  recipes for choropleth maps

# Map shape coordinates for Plots
function mapshapecoords(P::Vector{<:Union{Missing,AbstractPolygon,AbstractMultiPolygon}})::Tuple{Vector{Float64}, Vector{Float64}}
    scoords = shapecoords.(P)

    x = map(a -> a[1], scoords)
    y = map(a -> a[2], scoords)

    x = reduce(vcat, x)
    y = reduce(vcat, y)

    x, y
end

function mapshapecoords(P::Vector{<:Union{Missing,AbstractPoint}})::Tuple{Vector{Float64}, Vector{Float64}}
    scoords = shapecoords.(P)

    x = map(a -> a[1][1], scoords)
    y = map(a -> a[1][2], scoords)

    x, y
end

# Choropleth Map
@recipe function f(A::Any, colorvar::Union{AbstractVector,Symbol}, mcr::AbstractMapClassificator)

    if istable(A)
        # If Table or DataFrame
        (:geometry in propertynames(A)) || throw(ArgumentError("table does not have :geometry information"))
        P = A.geometry

        if isa(colorvar, Symbol)
            cvar = A[!, colorvar]
        elseif isa(colorvar, Vector{Symbol})
            cvar = [A[!, c] for c in colorvar]
        else
            cvar = colorvar
        end
        #(isa(colorvar, Symbol) || isa(colorvar, Vector{Symbol})) ? cvar = A[!, colorvar] : cvar = colorvar

    else
        P = A
        cvar = colorvar
    end

    # Check is geometry
    isa(P, Vector{<:Union{Missing, AbstractGeometry}}) || throw(ArgumentError("invalid geometry"))

    isCoLocation = false
    if isa(mcr, AbstractCoLocationMapClassificator)
        isCoLocation = true
    else
        length(P) == length(cvar) || throw(ArgumentError("dimensions must match: A has ($(length(P))), colorvar has ($(length(cvar)))"))
    end

    gtype = geotype(P[1])

    # Default color palette
    defpalette, defrev = defaultpalette(mcr)

    # Get plot attributes
    lower  = get(plotattributes, :lower, :open)
    upper  = get(plotattributes, :upper, :closed)
    digits = get(plotattributes, :digits, 2)
    sep    = get(plotattributes, :sep, ", ")
    counts = get(plotattributes, :counts, true)
    rev    = get(plotattributes, :rev, defrev)

    # Classify values and get labels
    if isa(mcr, AbstractGraduatedMapClassificator)
        mc = mapclassify(mcr, cvar, lower = lower, upper = upper)
        labels = maplabels(mc, digits = digits, sep = sep, counts = counts)
    elseif isa(mcr, AbstractUniqueMapClassificator)
        mc = mapclassify(mcr, cvar)
        labels = maplabels(mc, counts = counts)
    elseif isa(mcr, AbstractCoLocationMapClassificator)
        mc = mapclassify(mcr, cvar, lower = lower, upper = upper)
        labels = maplabels(mc, counts = counts)
    end

    k = length(mc)
    group = assignments(mc)

    # If palette is :Paired adjust for number of categories
    if defpalette == :Paired
        if k <= 12 defpalette = Symbol(string(defpalette) * "_" * string(k)) end
    end

    # Default attributes for maps
    legend --> true
    grid --> false
    showaxis --> false
    ticks --> false
    aspect_ratio --> :equal
    color_palette --> defpalette

    # Get colors from pallete
    mappal = palette(plotattributes[:color_palette], k - isCoLocation, rev = rev)

    # For each category in the map
    for i in 1:k

        # Subset of geometries for the category
        subset = P[group .== i]

        # color for the category
        if isCoLocation
            # Assing grey color for CoLocation when there is no match
            if i > 1
                catcolor = mappal[i - 1]
            else
                catcolor = :lightgrey
            end
        else
            catcolor = mappal[i]
        end

        # If empty subset, plot a NaN point to display the text in the legend but not the color
        if length(subset) == 0
            @series begin
                seriestype := :scatter
                seriescolor := catcolor
                label := labels[i]
                [NaN], [NaN]
            end
            continue
        end

        # Get shape map coordinates and plot polygon or point
        x, y = mapshapecoords(P[group .== i])

        if gtype == :Polygon || gtype == :MultiPolygon
            @series begin
                seriestype := :shape
                seriescolor := catcolor
                label := labels[i]
                (x, y)
            end
        elseif gtype == :Point
            @series begin
                seriestype := :scatter
                seriescolor := catcolor
                label := labels[i]
                x, y
            end
        end
    end

end

# LISA Cluster Map
@recipe function f(A::Any, lisavar::AbstractLocalSpatialAutocorrelation)

    if istable(A)
        # If Table or DataFrame
        (:geometry in propertynames(A)) || throw(ArgumentError("table does not have :geometry information"))
        P = A.geometry
    else
        P = A
    end

    # Check is geometry
    isa(P, Vector{<:Union{Missing, AbstractGeometry}}) || throw(ArgumentError("invalid geometry"))

    gtype = geotype(P[1])

    # Get plot attributes
    counts = get(plotattributes, :counts, true)
    pthreshold = get(plotattributes, :sig, 0.05)
    adjust = get(plotattributes, :adjust, :none)

   q = assignments(lisavar, pthreshold, adjust = adjust)

    mc = mapclassify(Unique(labelsorder(lisavar)), q)
    labels = maplabels(mc, counts = counts)

    k = length(mc)
    group = assignments(mc)
    grouplabs = mc.grouplabs

    # Default attributes for maps
    legend --> true
    grid --> false
    showaxis --> false
    ticks --> false
    aspect_ratio --> :equal

    # For each category in the map
    for i in 1:k

        # Subset of geometries for the category
        subset = P[group .== i]

        if grouplabs[i] == "ns"
            catcolor = :lightgrey
            alpha = 0.4
        end

        catcolor, alpha = labelcolor(lisavar, grouplabs[i] )

        # If empty subset, plot a NaN point to display the text in the legend but not the color
        if length(subset) == 0
            @series begin
                seriestype := :scatter
                seriescolor := catcolor
                alpha := alpha
                label := labels[i]
                [NaN], [NaN]
            end
            continue
        end

        # Get shape map coordinates and plot polygon or point
        x, y = mapshapecoords(P[group .== i])

        if gtype == :Polygon || gtype == :MultiPolygon
            @series begin
                seriestype := :shape
                seriescolor := catcolor
                alpha := alpha
                label := labels[i]
                (x, y)
            end
        elseif gtype == :Point
            @series begin
                seriestype := :scatter
                seriescolor := catcolor
                alpha := alpha
                label := labels[i]
                x, y
            end
        end
    end

end


text:
cholupper(a) = cholesky(a).U
cholupper(a::Number) = sqrt(a)
chollower(a) = cholesky(a).L
chollower(a::Number) = sqrt(a)
export cholupper, chollower
cholupper!(a) = cholesky!(a).U
cholupper!(a::Number) = sqrt(a)
chollower!(a) = cholesky!(a).L
chollower!(a::Number) = sqrt(a)
export cholupper!, chollower!


text:
let

    # fitting noisy data to an exponential model
    model(x, p) = @. p[1] * exp(-x * p[2])
    #model(x,p) = [p[1],100*(p[2]-p[1]^2)]
    # some example data
    Random.seed!(12345)
    xdata = range(0, stop=10, length=50_000)
    ydata = model(xdata, [1.0, 2.0]) + 0.01*randn(length(xdata))
    p0 = [.5, .5]

    function jacobian_model(x,p)
        J = Array{Float64}(undef, length(x), length(p))
        @. J[:,1] = exp(-x*p[2])     #dmodel/dp[1]
        @. @views J[:,2] = -x*p[1]*J[:,1]
        J
    end

    #Generating the "automatic" avv
    model_inplace(out, p) =  @. out = p[1] * exp(-xdata * p[2])

    hessians = Array{Float64}(undef, length(xdata)*length(p0), length(p0))
    h! = make_hessian(model_inplace, xdata, p0)
    auto_avv! = Avv(h!, length(p0), length(xdata))

    # a couple notes on the Avv function:
    # - the basic idea is to see the model output as simply a collection of functions: f1...fm
    # - then Avv return an array of size m, where each elements corresponds to
    # v'H(p)v, with H an n*n Hessian matrix of the m-th function, with n the size of p
    function manual_avv!(dir_deriv,p,v)
        v1 = v[1]
        v2 = v[2]
        for i=1:length(xdata)
            #compute all the elements of the H matrix
            h11 = 0
            h12 = (-xdata[i] * exp(-xdata[i] * p[2]))
            #h21 = h12
            h22 = (xdata[i]^2 * p[1] * exp(-xdata[i] * p[2]))

            # manually compute v'Hv. This whole process might seem cumbersome, but
            # allocating temporary matrices quickly becomes REALLY expensive and might even
            # render the use of geodesic acceleration terribly inefficient
            dir_deriv[i] = h11*v1^2 + 2*h12*v1*v2 + h22*v2^2

        end
    end

    curve_fit(model, jacobian_model, xdata, ydata, p0; maxIter=1); #warmup
    curve_fit(model, jacobian_model, xdata, ydata, p0; maxIter=1, avv! = manual_avv!,lambda=0, min_step_quality = 0); #lambda = 0 to match Mark's code
    curve_fit(model, jacobian_model, xdata, ydata, p0; maxIter=10, avv! = auto_avv!,lambda=0, min_step_quality = 0)

    println("--------------\nPerformance of curve_fit vs geo")

    println("\t Non-inplace")
    fit = @time curve_fit(model, jacobian_model, xdata, ydata, p0; maxIter=1000)
    @test fit.converged

    println("\t Geodesic")
    fit_geo = @time curve_fit(model, jacobian_model, xdata, ydata, p0; maxIter=10, avv! = manual_avv!,lambda=0, min_step_quality = 0)
    @test fit_geo.converged

    println("\t Geodesic - auto avv!")
    fit_geo_auto = @time curve_fit(model, jacobian_model, xdata, ydata, p0; maxIter=10, avv! = auto_avv!,lambda=0, min_step_quality = 0)
    @test fit_geo_auto.converged

    @test maximum(abs.(fit.param-[1.0, 2.0])) < 1e-1
    @test maximum(abs.(fit.param-fit_geo.param)) < 1e-6
    @test maximum(abs.(fit.param-fit_geo_auto.param)) < 1e-6

    #with noise
    yvars = 1e-6*rand(length(xdata))
    ydata = model(xdata, [1.0, 2.0]) + sqrt.(yvars) .* randn(length(xdata))

    #warm up
    curve_fit(model, jacobian_model, xdata, ydata, 1 ./ yvars, p0; maxIter=1)
    curve_fit(model, jacobian_model, xdata, ydata,  1 ./ yvars, p0; maxIter=1, avv! = manual_avv!,lambda=0, min_step_quality = 0)
    curve_fit(model, jacobian_model, xdata, ydata,  1 ./ yvars, p0; maxIter=1, avv! = auto_avv!,lambda=0, min_step_quality = 0)

    println("--------------\nPerformance of curve_fit vs geo with weights")

    println("\t Non-inplace")
    fit_wt = @time curve_fit(model, jacobian_model, xdata, ydata, 1 ./ yvars, p0; maxIter=100)
    @test fit_wt.converged

    println("\t Geodesic")
    fit_geo_wt = @time curve_fit(model, jacobian_model, xdata, ydata,  1 ./ yvars, p0; maxIter=100, avv! = manual_avv!,lambda=0, min_step_quality = 0)
    @test fit_geo_wt.converged

    println("\t Geodesic - auto avv!")
    fit_geo_auto_wt = @time curve_fit(model, jacobian_model, xdata, ydata,  1 ./ yvars, p0; maxIter=100, avv! = auto_avv!,lambda=0, min_step_quality = 0)
    @test fit_geo_auto_wt.converged

    @test maximum(abs.(fit_wt.param-[1.0, 2.0])) < 1e-1
    @test maximum(abs.(fit_wt.param-fit_geo_wt.param)) < 1e-6
    @test maximum(abs.(fit_wt.param-fit_geo_auto_wt.param)) < 1e-6

end


text:
const superscripts = Dict(0=>"⁰",1=>"¹",2=>"²",3=>"³",4=>"⁴",5=>"⁵",6=>"⁶",7=>"⁷",8=>"⁸",9=>"⁹")

## JacobiQ and JacobiQWeight

function Base.show(io::IO,S::JacobiQ)
    S.a == S.b == 0 ? print(io,"LegendreQ(") : print(io,"JacobiQ($(S.a),$(S.b),")
    show(io,domain(S))
    print(io,")")
end

function Base.show(io::IO,s::JacobiQWeight)
    d=domain(s)
    #TODO: Get shift and weights right
    if s.α==s.β
        print(io,"(x^2-1)^$(s.α)[")
    elseif s.α==0
        print(io,"(x-1)^$(s.β)[")
    elseif s.β==0
        print(io,"(x+1)^$(s.α)[")
    else
        print(io,"(x+1)^$(s.α)*(x-1)^$(s.β)[")
    end

    show(io,s.space)
    print(io,"]")
end

## CauchyWeight

function Base.show(io::IO,s::CauchyWeight)
    #TODO: Get shift and weights right
    if order(s) == 0
        print(io,"π⁻¹log|y-x|[")
    elseif order(s) ≥ 1
        print(io,"π⁻¹|y-x|⁻"*mapreduce(x->superscripts[x],*,reverse!(digits(order(s))))*"[")
    else
        print(io,"π⁻¹|y-x|"*mapreduce(x->superscripts[x],*,reverse!(digits(-order(s))))*"[")
    end
    show(io,s.space)
    print(io,"]")
end

## GreensFun

function Base.show(io::IO,G::GreensFun)
    print(io,"GreensFun with kernels:\n{")
    Kernels = kernels(G)
    for i in 1:length(G)
        print(io,"\n ")
        show(io,Kernels[i])
    end
    print(io,"\n}")
end


text:
#Loading necessary modules
using PorousMaterials

PorousMaterials.set_path_to_data("$data_path")
path = PorousMaterials.PATH_TO_DATA
working_dir = pwd() * "/"
mkdir(working_dir * "Output")

temperature = $temperature
ljff = LJForceField("$ff", cutoffradius=$cutoff, mixing_rules="$mixing")
framework = Framework(working_dir * "$framework", check_charge_neutrality=false)
rep_factor = replication_factors(framework.box, ljff)
framework = replicate(framework, rep_factor)
density = crystal_density(framework)

for adsorbate in $adsorbates
    result = open("Output/Ev_vdw_${frameworkname}_"*adsorbate*"_"*adsorbate*".csv","w")
    write(result,"!!!Generated results using aiida-porousmaterials plugin!!!\n")
    write(result,"Framework Density\n")
    write(result,string(density),"\n")
    write(result,"Temperature(K)\n")
    write(result,string(temperature),"\n")
    write(result, "Ev_K,boltzmann_factor,weighted_energy_K,Rv_A,x,y,z,adsorbate\n")
    posfile = open(working_dir * "${frameworkname}_"*adsorbate*".voro_accessible")
    lines = readlines(posfile)
    n_nodes = parse(Int, lines[1])
    for k = 1:n_nodes
        xyz = split(lines[2+k])[2:4]
        r = split(lines[2+k])[5]
        x = parse.(Float64, xyz)
        molecule = Molecule(adsorbate)
        set_fractional_coords!(molecule, framework.box)
        translate_to!(molecule,framework.box.c_to_f * x)
        energy = (vdw_energy(framework, molecule, ljff))
        boltzmann_factor = exp(-energy / temperature)
        wtd_energy = boltzmann_factor * energy
        saveresults = energy, boltzmann_factor, wtd_energy, r, x[1], x[2], x[3], adsorbate
        write(result, join(saveresults,","), "\n")
    end
    close(posfile)
    close(result)
end

for adsorbate in $adsorbates
    result = open("Output/Ev_vdw_${frameworkname}_PLD_"*adsorbate*".csv","w")
    write(result,"!!!Generated results using aiida-porousmaterials plugin!!!\n")
    write(result,"Framework Density\n")
    write(result,string(density),"\n")
    write(result,"Temperature(K)\n")
    write(result,string(temperature),"\n")
    write(result, "Ev_K,boltzmann_factor,weighted_energy_K,Rv_A,x,y,z,adsorbate\n")
    posfile = open(working_dir * "${frameworkname}_PLD.voro_accessible")
    lines = readlines(posfile)
    n_nodes = parse(Int, lines[1])
    for k = 1:n_nodes
        xyz = split(lines[2+k])[2:4]
        r = split(lines[2+k])[5]
        x = parse.(Float64, xyz)
        molecule = Molecule(adsorbate)
        set_fractional_coords!(molecule, framework.box)
        translate_to!(molecule,framework.box.c_to_f * x)
        energy = (vdw_energy(framework, molecule, ljff))
        boltzmann_factor = exp(-energy / temperature)
        wtd_energy = boltzmann_factor * energy
        saveresults = energy, boltzmann_factor, wtd_energy, r, x[1], x[2], x[3], adsorbate
        write(result, join(saveresults,","), "\n")
    end
    close(posfile)
    close(result)
end

max_stars_repo_name: pzarabadip/aiida-porousmaterials

max_issues_repo_name: pzarabadip/aiida-porousmaterials

max_forks_repo_name: pzarabadip/aiida-porousmaterials


text:
using Test
using EMIRT
using EMIRT.AffinityMaps

@testset "test affinitymap" begin
    aff = rand(Float32, 128,128,16,3)

    println("transform affinitymap to edge list...")
    @time aff2edgelist(aff)

    println("segment affinity map using connectivity analysis ...")
    @time aff2seg(aff)

    println("transform to uniform distribution ...")
    @time aff2uniform(aff)

    println("exchange x and z axis of affinity map...")
    @time exchangeaffxz!(aff)
end # end of test set


text:
input = readlines()

numbers = parse.(Int, split(input[1], ","))

function check_victory(m, drawn)
    if maximum(sum(drawn; dims=1)) == 5 || maximum(sum(drawn; dims=2)) == 5
        (@. m * !drawn) |> sum
    else
        0
    end
end

function solve()
    best = []
    for i in 1:100
        m = [
            parse.(Int, split(input[3 + (i-1)*6]));;
            parse.(Int, split(input[4 + (i-1)*6]));;
            parse.(Int, split(input[5 + (i-1)*6]));;
            parse.(Int, split(input[6 + (i-1)*6]));;
            parse.(Int, split(input[7 + (i-1)*6]))
        ]

        drawn = falses(size(m))

        for (j, n) in enumerate(numbers)
            pos = findnext(==(n), m, CartesianIndex(1, 1))
            if !isnothing(pos)
                drawn[pos] = true
            end
            score = check_victory(m, drawn)
            if score > 0
                score *= n
                push!(best, (j, score))
                break
            end
        end
    end
    sort!(best)
    best[begin], best[end]
end

# Part 1 and 2
solve() .|> println

max_stars_repo_name: pLOPeGG/JuliaOfCode

max_issues_repo_name: pLOPeGG/JuliaOfCode

max_forks_repo_name: pLOPeGG/JuliaOfCode


text:
function SelectionSort(x::AbstractVector)
    for i = 1:length(x)
        min = i

        for j = i+1:length(x)
            if x[j] < x[min]
                min = j
            end
        end

        temp = x[i]
        x[i] = x[min]
        x[min] = temp
    end

    return x
end


text:
# Invertible CNN layer from Putzky & Welling (2019)
# Author: Philipp Witte, pwitte3@gatech.edu
# Date: January 2020

using InvertibleNetworks, LinearAlgebra, Test, Random

Random.seed!(1);

# Input
nx = 28
ny = 28
n_in = 8
n_hidden = 8
batchsize = 2

# Input images
X = randn(Float32, nx, ny, n_in, batchsize)
X0 = randn(Float32, nx, ny, n_in, batchsize)
dX = X - X0

# Invertible layers
L = CouplingLayerIRIM(n_in, n_hidden)
L01 = CouplingLayerIRIM(n_in, n_hidden)
L02 = CouplingLayerIRIM(n_in, n_hidden)

###################################################################################################
# Test invertibility

X_ = L.inverse(L.forward(X))
@test isapprox(norm(X - X_)/norm(X), 0f0; atol=1e-2)

X_ = L.forward(L.inverse(X))
@test isapprox(norm(X - X_)/norm(X), 0f0; atol=1e-2)

###################################################################################################
# Gradient tests

# Loss Function
function loss(L, X, Y)
    Y_ = L.forward(X)
    ΔY = Y_ - Y

    f = .5f0*norm(ΔY)^2
    ΔX = L.backward(ΔY, Y_)[1]

    # Pass back gradients w.r.t. input X and from the residual block and 1x1 conv. layer
    return f, ΔX, L.C.v1.grad, L.C.v2.grad, L.C.v3.grad, L.RB.W1.grad, L.RB.W2.grad, L.RB.W3.grad
end

# Gradient test w.r.t. input X0
Y = L.forward(X)
f0, ΔX = loss(L, X0, Y)[1:2]
h = 0.1f0
maxiter = 5
err1 = zeros(Float32, maxiter)
err2 = zeros(Float32, maxiter)

print("\nGradient test invertible layer\n")
for j=1:maxiter
    f = loss(L, X0 + h*dX, Y)[1]
    err1[j] = abs(f - f0)
    err2[j] = abs(f - f0 - h*dot(dX, ΔX))
    print(err1[j], "; ", err2[j], "\n")
    global h = h/2f0
end

@test isapprox(err1[end] / (err1[1]/2^(maxiter-1)), 1f0; atol=1f1)
@test isapprox(err2[end] / (err2[1]/4^(maxiter-1)), 1f0; atol=1f1)

# Gradient test w.r.t. weights of residual block
Y = L.forward(X)
Lini = deepcopy(L02)
dW1 = L.RB.W1.data - L02.RB.W1.data
dW2 = L.RB.W2.data - L02.RB.W2.data
dW3 = L.RB.W3.data - L02.RB.W3.data

f0, ΔX, Δv1, Δv2, Δv3, ΔW1, ΔW2, ΔW3 = loss(L02, X, Y)
h = 0.1f0
maxiter = 5
err3 = zeros(Float32, maxiter)
err4 = zeros(Float32, maxiter)

print("\nGradient test invertible layer\n")
for j=1:maxiter
    L02.RB.W1.data = Lini.RB.W1.data + h*dW1
    L02.RB.W2.data = Lini.RB.W2.data + h*dW2
    L02.RB.W3.data = Lini.RB.W3.data + h*dW3
    f = loss(L02, X, Y)[1]
    err3[j] = abs(f - f0)
    err4[j] = abs(f - f0 - h*dot(dW1, ΔW1) - h*dot(dW2, ΔW2) - h*dot(dW3, ΔW3))
    print(err3[j], "; ", err4[j], "\n")
    global h = h/2f0
end

@test isapprox(err3[end] / (err3[1]/2^(maxiter-1)), 1f0; atol=1f1)
@test isapprox(err4[end] / (err4[1]/4^(maxiter-1)), 1f0; atol=1f1)

# Gradient test w.r.t. 1x1 conv weights
Y = L.forward(X)
Lini = deepcopy(L01)
dv1 = L.C.v1.data - L01.C.v1.data
dv2 = L.C.v2.data - L01.C.v2.data
dv3 = L.C.v3.data - L01.C.v3.data

f0, ΔX, Δv1, Δv2, Δv3, ΔW1, ΔW2, ΔW3 = loss(L01, X, Y)
h = 0.1f0
maxiter = 4
err5 = zeros(Float32, maxiter)
err6 = zeros(Float32, maxiter)

print("\nGradient test invertible layer\n")
for j=1:maxiter
    L01.C.v1.data = Lini.C.v1.data + h*dv1
    L01.C.v2.data = Lini.C.v2.data + h*dv2
    L01.C.v3.data = Lini.C.v3.data + h*dv3
    f = loss(L01, X, Y)[1]
    err5[j] = abs(f - f0)
    err6[j] = abs(f - f0 - h*dot(dv1, Δv1) - h*dot(dv2, Δv2) - h*dot(dv3, Δv3))
    print(err5[j], "; ", err6[j], "\n")
    global h = h/2f0
end

@test isapprox(err5[end] / (err5[1]/2^(maxiter-1)), 1f0; atol=1f1)
@test isapprox(err6[end] / (err6[1]/4^(maxiter-1)), 1f0; atol=1f1)

###################################################################################################
# Jacobian-related tests

# Gradient test

# Initialization
L = CouplingLayerIRIM(n_in, n_hidden)
θ = deepcopy(get_params(L))
L0 = CouplingLayerIRIM(n_in, n_hidden)
θ0 = deepcopy(get_params(L0))
X = randn(Float32, nx, ny, n_in, batchsize)

# Perturbation (normalized)
dθ = θ-θ0; dθ .*= norm.(θ0)./(norm.(dθ).+1f-10)
dX = randn(Float32, nx, ny, n_in, batchsize); dX *= norm(X)/norm(dX)

# Jacobian eval
dY, Y = L.jacobian(dX, dθ, X)

# Test
print("\nJacobian test\n")
h = 0.1f0
maxiter = 5
err7 = zeros(Float32, maxiter)
err8 = zeros(Float32, maxiter)
for j=1:maxiter
    set_params!(L, θ+h*dθ)
    Y_loc = L.forward(X+h*dX)
    err7[j] = norm(Y_loc - Y)
    err8[j] = norm(Y_loc - Y - h*dY)
    print(err7[j], "; ", err8[j], "\n")
    global h = h/2f0
end

@test isapprox(err7[end] / (err7[1]/2^(maxiter-1)), 1f0; atol=1f1)
@test isapprox(err8[end] / (err8[1]/4^(maxiter-1)), 1f0; atol=1f1)

# Adjoint test

set_params!(L, θ)
dY, Y = L.jacobian(dX, dθ, X)
dY_ = randn(Float32, size(dY))
dX_, dθ_ = L.adjointJacobian(dY_, Y)
a = dot(dY, dY_)
b = dot(dX, dX_)+dot(dθ, dθ_)
@test isapprox(a, b; rtol=1f-3)


text:
# shared functionality among TMJets implementations

using TaylorModels: TaylorModelN
using TaylorModels: fp_rpa, remainder, initialize!

# =================================
# Defaut values for the parameters
# =================================

const DEFAULT_MAX_STEPS_TMJETS = 2000
const DEFAULT_ABS_TOL_TMJETS = 1e-10
const DEFAULT_ORDER_T_TMJETS = 8
const DEFAULT_ORDER_Q_TMJETS = 2

"""
    TMJets

The algorithm TMJets defaults to `TMJets21b`.
"""
const TMJets = TMJets21b

# =======================================================================
# Initialization funtions to prepare the input for validated integration
# =======================================================================

# fallback
function _initialize(X0::LazySet, orderQ, orderT)
    return _initialize(box_approximation(X0), orderQ, orderT)
end

# taylor model representations
function _initialize(X0::TaylorModelReachSet, orderQ, orderT)
    return set(X0)
end

_initialize(X0::Vector{TaylorModel1{TaylorN{T}, T}}, orderQ, orderT) where {T} = X0

# hyperrectangular sets
_initialize(X0::AbstractHyperrectangle, orderQ, orderT) = convert(IntervalBox, box_approximation(X0))
_initialize(X0::IntervalBox, orderQ, orderT) = X0
_initialize(X0::IntervalArithmetic.Interval, orderQ, orderT) = IntervalBox(X0)

# zonotopic sets
function _initialize(X0::AbstractZonotope, orderQ, orderT)
    X0z = convert(Zonotope, X0)
    if order(X0z) == 1
        X = overapproximate(X0z, TaylorModelReachSet, orderQ=orderQ, orderT=orderT)
    else
        X = overapproximate(X0z, TaylorModelReachSet, orderQ=orderQ, orderT=orderT, box_reduction=true)
    end
    return set(X)
end

function _initialize(X0::CartesianProduct{N, <:AbstractHyperrectangle, <:AbstractHyperrectangle}, orderQ, orderT) where {N}
    return convert(IntervalBox, convert(Hyperrectangle, X0))
end

function _initialize(X0::CartesianProduct{N, <:Zonotope, <:Interval}, orderQ, orderT) where {N}
    Z = X0.X
    G = Z.generators
    ord = order(Z)
    n = dim(Z)

    if ord == 1
        # "exact"
        X = _overapproximate_structured(X0, TaylorModelReachSet, orderQ=orderQ, orderT=orderT)

    elseif (size(G) == (n, 2n + 1)) && isdiag(view(G, :, (n+2):(2n+1)))
        X = _overapproximate_structured_full(X0, TaylorModelReachSet, orderQ=orderQ, orderT=orderT)

    elseif (size(G) == (n, 2(n + 1))) && isdiag(view(G, :, (n+2):(2n+1))) && iszero(view(G, :, 2n+2))
        aux = Zonotope(Z.center, view(G, :, 1:(2n+1)))
        X0z = convert(Zonotope, aux × X0.Y)
        X = _overapproximate_structured(X0z, TaylorModelReachSet, orderQ=orderQ, orderT=orderT)

    else # otherwise, resort to a box overapproximation
        X0z = convert(Zonotope, X0)
        X = overapproximate(X0z, TaylorModelReachSet, orderQ=orderQ, orderT=orderT, box_reduction=true)
    end

    return set(X)
end

# =================================
# External solver function
# =================================

function _solve_external(f!, X0, t0, T, orderQ, orderT, abstol, maxsteps, Δt0; kwargs...)
    N = eltype(X0)

    # box overapproximation of the initial states
    X0box = convert(IntervalBox, box_approximation(X0))

    # extract solver name and options
    parse_eqs = get(kwargs, :parse_eqs, false)
    solver_name = get(kwargs, :solver_name, TM.validated_integ)
    solver_kwargs = get(kwargs, :solver_kwargs, Dict(:maxsteps=>maxsteps))

    # call external solver
    tv, xv, xTM1v = solver_name(f!, X0box, t0, T, orderQ, orderT, abstol; parse_eqs, solver_kwargs...)

    # build flowpipe
    F = Vector{TaylorModelReachSet{N}}()
    sizehint!(F, maxsteps)
    @inbounds for i in eachindex(tv)
        δt = TimeInterval(tv[i] + domain(xTM1v[1, i]))
        Ri = TaylorModelReachSet(xTM1v[:, i], δt + Δt0)
        push!(F, Ri)
    end
    ext = Dict{Symbol, Any}(:tv => tv, :xv => xv, :xTM1v => xTM1v)
    return Flowpipe(F, ext)
end


text:
function fun2(a)
    println(a+b)
end

function fun1()
    fun2(4)
end

fun1()


text:
using MonteCarloX
using Random

include("../examples/contact_process.jl")

"""
    test_gillespie_fast(;verbose=false)

Testing if sampling single event correctly

there is clearly a drift in the precision of sum_rates (needs to be compensated
in system implementation of course
"""
function test_gillespie(;verbose = false)
    return pass
end


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

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

using StaticArrays
using Distributions
using Random
# Let's generate the data
# -----------------------
using Bridge
DIR = "auxiliary"
include(joinpath(SRC_DIR, DIR, "data_simulation_fns.jl"))
include(joinpath(SRC_DIR, DIR, "utility_functions.jl"))
Random.seed!(4)
#θ₀ = [10.0, 28.0, 8.0/3.0, 3.0, 3.0, 3.0]
θˣ = [10.0, 28.0, 8.0/3.0, 3.0]
Pˣ = LorenzCV(θˣ...)

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

θ_init = [5.0, 15.0, 6.0, 8.0]
Pˣ = LorenzCV(θ_init...)

skip = 1000

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

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

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

setup = MCMCSetup(Pˣ, P̃, PartObs())
set_observations!(setup, [L for _ in P̃], [Σ for _ in P̃], obs_vals, obs_time) # uses default fpt
set_imputation_grid!(setup, 1/2000)
set_transition_kernels!(setup,
                        [RandomWalk([], []),
                         RandomWalk([2.0, 1.0, 0.64, 0.5], 4)],
                        0.96, true, [[1,2,3],[4]],
                        (ConjugateUpdt(),
                         MetropolisHastingsUpdt()
                        ),                           # update types
                        Adaptation(x0,
                                   [0.7, 0.4, 0.2, 0.2, 0.2],
                                   [500, 500, 500, 500, 500],
                                   1)
                        )
set_priors!(setup,
            Priors((MvNormal([0.0,0.0,0.0], diagm(0=>[1000.0, 1000.0, 1000.0])),
                    ImproperPrior())),               # priors over parameters
            GsnStartingPt(x0, @SMatrix [20.0 0.0 0.0;
                                        0.0 20.0 0.0;
                                        0.0 0.0 400.0]), # prior over starting point
            x0
            )
set_mcmc_params!(setup,
                 4*10^3,            # number of mcmc steps
                 1*10^3,            # save path every ... iteration
                 10^2,              # print progress message every ... iteration
                 10^0,              # thin the path imputatation points for save
                 100                # number of first iterations without param update
                 )
#set_blocking!(setup, ChequeredBlocking(),
#              (collect(1:length(obs_vals)-2)[1:2:end], 10^(-10), SimpleChangePt(100)))
set_solver!(setup, Vern7(), NoChangePt())
initialise!(eltype(x0), setup)

Random.seed!(4)
out, elapsed = @timeit mcmc(setup)
display(out.accpt_tracker)

include(joinpath(SRC_DIR, DIR, "plotting_fns.jl"))
plot_chains(out; truth=[10.0, 28.0, 8.0/3.0, 3.0],
            ylims=[nothing, (25,30), (2,5), (0,10)])
plot_paths(out; obs=(times=obs_time[2:end],
                     vals=[[v[1] for v in obs_vals[2:end]],
                           [v[2] for v in obs_vals[2:end]]], indices=[2,3]))
ws.accpt_tracker


text:
using Oceananigans.Models.NonhydrostaticModels: calculate_pressure_source_term_fft_based_solver!
using Oceananigans.Architectures: device_event
import Oceananigans.Models.NonhydrostaticModels: solve_for_pressure!

function solve_for_pressure!(pressure, solver::DistributedFFTBasedPoissonSolver, Δt, U★)
    rhs = first(solver.storage)
    arch = solver.architecture
    grid = solver.my_grid

    rhs_event = launch!(arch, grid, :xyz, calculate_pressure_source_term_fft_based_solver!,
                        rhs, grid, Δt, U★, dependencies = device_event(arch))

    wait(device(arch), rhs_event)

    # Solve pressure Poisson equation for pressure, given rhs
    solve!(pressure, solver)

    return pressure
end


text:
module ScuffTools

using Plots; plotly()
using DataFrames
using MultiLayerNFRHT

# constants
const w0 = 3.0e14

abstract type FileType end
struct SIFlux       <: FileType end
struct TotalFlux{T} <: FileType end

# export of functions
export FileType,SIFlux,TotalFlux,
       import_data,simulation_data,
       plot_scuff,benchmark_freq,
       scuff_job,scuff_parallel,
       transfer_w,total_transfer,frequency_file

# Files
include("plotting.jl")
include("importing.jl")
include("run_on_cluster.jl")
include("utils.jl")

end #module


text:
import PowerSimulations:
    SimulationStoreParams,
    ModelStoreParams,
    get_problem_exports,
    should_export_dual,
    should_export_parameter,
    should_export_variable,
    OptimizationContainerMetadata

function _make_params()
    sim = Dict(
        "initial_time" => Dates.DateTime("2020-01-01T00:00:00"),
        "step_resolution" => Dates.Hour(24),
        "num_steps" => 2,
    )
    problem_defs = OrderedDict(
        :ED => Dict(
            "execution_count" => 24,
            "horizon" => 12,
            "interval" => Dates.Hour(1),
            "resolution" => Dates.Hour(1),
            "base_power" => 100.0,
            "system_uuid" => Base.UUID("4076af6c-e467-56ae-b986-b466b2749572"),
        ),
        :UC => Dict(
            "execution_count" => 1,
            "horizon" => 24,
            "interval" => Dates.Hour(1),
            "resolution" => Dates.Hour(24),
            "base_power" => 100.0,
            "system_uuid" => Base.UUID("4076af6c-e467-56ae-b986-b466b2749572"),
        ),
    )
    container_metadata = OptimizationContainerMetadata(
        Dict(
            "ActivePowerVariable__ThermalStandard" =>
                PSI.VariableKey(ActivePowerVariable, ThermalStandard),
            "EnergyVariable__HydroEnergyReservoir" =>
                PSI.VariableKey(EnergyVariable, HydroEnergyReservoir),
            "OnVariable__ThermalStandard" =>
                PSI.VariableKey(OnVariable, ThermalStandard),
        ),
    )
    problems = OrderedDict{Symbol, ModelStoreParams}()
    for problem in keys(problem_defs)
        problem_params = ModelStoreParams(
            problem_defs[problem]["execution_count"],
            problem_defs[problem]["horizon"],
            problem_defs[problem]["interval"],
            problem_defs[problem]["resolution"],
            problem_defs[problem]["base_power"],
            problem_defs[problem]["system_uuid"],
            container_metadata,
        )

        problems[problem] = problem_params
    end

    return SimulationStoreParams(
        sim["initial_time"],
        sim["step_resolution"],
        sim["num_steps"],
        problems,
        # Emulation Problem Params. Export not implemented yet
        OrderedDict{Symbol, ModelStoreParams}(),
    )
end

@testset "Test export from JSON" begin
    params = _make_params()
    exports = SimulationResultsExport(joinpath(DATA_DIR, "results_export.json"), params)

    valid = Dates.DateTime("2020-01-01T06:00:00")
    valid2 = Dates.DateTime("2020-01-02T23:00:00")
    invalid = Dates.DateTime("2020-01-01T02:00:00")
    invalid2 = Dates.DateTime("2020-01-03T00:00:00")

    @test should_export_variable(
        exports,
        valid,
        :ED,
        PSI.VariableKey(ActivePowerVariable, ThermalStandard),
    )
    @test should_export_variable(
        exports,
        valid2,
        :ED,
        PSI.VariableKey(ActivePowerVariable, ThermalStandard),
    )
    @test !should_export_variable(
        exports,
        invalid,
        :ED,
        PSI.VariableKey(ActivePowerVariable, ThermalStandard),
    )
    @test !should_export_variable(
        exports,
        invalid2,
        :ED,
        PSI.VariableKey(ActivePowerVariable, ThermalStandard),
    )
    @test !should_export_variable(
        exports,
        valid,
        :ED,
        PSI.VariableKey(ActivePowerVariable, RenewableFix),
    )
    @test should_export_parameter(
        exports,
        valid,
        :ED,
        PSI.ParameterKey(ActivePowerTimeSeriesParameter, ThermalStandard),
    )
    @test !should_export_dual(
        exports,
        valid,
        :ED,
        PSI.ConstraintKey(ActivePowerVariableLimitsConstraint, RenewableFix),
    )

    @test should_export_variable(
        exports,
        valid,
        :UC,
        PSI.VariableKey(OnVariable, ThermalStandard),
    )
    @test !should_export_variable(
        exports,
        valid,
        :UC,
        PSI.VariableKey(ActivePowerVariable, RenewableFix),
    )
    @test should_export_parameter(
        exports,
        valid,
        :UC,
        PSI.ParameterKey(ActivePowerTimeSeriesParameter, ThermalStandard),
    )
    @test should_export_dual(
        exports,
        valid,
        :UC,
        PSI.ConstraintKey(ActivePowerVariableLimitsConstraint, RenewableFix),
    )

    @test exports.path == "export_path"
    @test exports.format == "csv"
    @test "csv" in list_supported_formats(SimulationResultsExport)
end

@testset "Invalid exports" begin
    params = _make_params()
    valid = Dates.DateTime("2020-01-01T00:00:00")
    invalid = Dates.DateTime("2020-01-03T00:00:00")

    # Invalid start_time
    @test_throws IS.InvalidValue SimulationResultsExport(
        Dict("start_time" => invalid, "models" => [Dict("name" => "ED")]),
        params,
    )

    # Invalid end_time
    @test_throws IS.InvalidValue SimulationResultsExport(
        Dict("end_time" => invalid, "models" => [Dict("name" => "ED")]),
        params,
    )

    # Invalid format
    @test_throws IS.InvalidValue SimulationResultsExport(
        Dict("format" => "invalid", "models" => [Dict("name" => "ED")]),
        params,
    )

    # Missing name
    @test_throws IS.InvalidValue SimulationResultsExport(
        Dict("models" => [Dict("variables" => ["ActivePowerVariable__ThermalStandard"])]),
        params,
    )
end


text:
module ADMM

using Printf
using LinearAlgebra
using SparseArrays
using ..Concepts
using ..Estimator

using ..ModelFitting
using ..Utilities
using ..Utilities.FastEigen
using ..Utilities.PrettyPrinter
using ..Losses
using ..MathLib
using Logging

import  StatsBase
import LinearAlgebra:norm

# export complete

mutable struct RunHistory
    primfeas::Array{Float64,1}
    dualfeas::Array{Float64,1}

    function RunHistory()
        new(Array{Float64,1}(),Array{Float64,1}())
    end
end

import KrylovKit

# function project(v,e)
#     return e * diagm(0 => v) * e';
# end

const header_list = ["Iter",
                     " R(dual)",
                     "R(primal)",
                     "ℒ(Gaussian)",
                     "ℒ(Bernoulli)",
                     "ℒ(Poisson)",
                     "ℒ(NegBin) ",
                     " ℒ(Gamma) ",
                     "λ‖diag(Z)‖ᵢ",
                     " μ⟨I, X⟩ ",
                     " ‖Z₁₂‖ᵢ "]

function l1BallProjection(v,b)
    if (norm(v,1) <= b);
        return v
    end
    n = length(v);
    nvec = hcat(1:n...)'[:];
    vv = sort(abs.(v),lt = (x,y) -> !isless(x,y));
    idxsort = sortperm(abs.(v),lt = (x,y) -> !isless(x,y));
    vsum = cumsum(vv);
    tmp  = vv .-(vsum .- b)./nvec;
    idx  = findall(x->x>0,tmp);
    if !isempty(idx);
        k = maximum(idx);
    else
        println("something is wrong")
    end
    lam = (vsum[k] .- b) ./ k;
    xx = zeros(length(idxsort));
    xx[idxsort,1] = max.(vv .- lam,0);
    x = sign.(v).*xx;
    return x;
end

@private
function initialize_warmup(tracker, A)
    warmup = Dict{Symbol, Array{Float64}}()
    if haskey(tracker, :Bernoulli) && length(tracker[:Bernoulli][:Observed]) > 0
        warmup[:Bernoulli] = rand(length(tracker[:Bernoulli][:Observed]), 1)
    end
    if haskey(tracker, :Gaussian) && length(tracker[:Gaussian][:Observed]) > 0
        warmup[:Gaussian] = rand(length(tracker[:Gaussian][:Observed]), 1)
    end
    if haskey(tracker, :Poisson) && length(tracker[:Poisson][:Observed]) > 0
        warmup[:Poisson] = rand(length(tracker[:Poisson][:Observed]), 1)
    end
    if haskey(tracker, :Gamma) && length(tracker[:Gamma][:Observed]) > 0
        warmup[:Gamma] = rand(length(tracker[:Gamma][:Observed]), 1)
    end
    if haskey(tracker, :NegativeBinomial) && length(tracker[:NegativeBinomial][:Observed]) > 0
        warmup[:NegativeBinomial] = rand(length(tracker[:NegativeBinomial][:Observed]), 1)
    end
    return warmup
end

@private
function ensure_feasible_data_type(set_of_types)
    @info(set_of_types)
    if !issubset(set_of_types, Set([:Gaussian,
                                    :Poisson,
                                    :NegativeBinomial,
                                    :Gamma,
                                    :Bernoulli]))
        @error("Unrecognized data type.")
        return nothing
    end
end

@private
function ensure_valid_param_estimates(::Type{Val{:Gaussian}}, value::T) where T
    return true
end

@private
function ensure_valid_param_estimates(::Type{Val{:Bernoulli}}, value::T) where T
    return true
end

@private
function ensure_valid_param_estimates(::Type{Val{:Gamma}}, value::T) where T
    return true
end

@private
function ensure_valid_param_estimates(::Type{Val{:Posson}}, value::T) where T
    return true
end

@private
function ensure_valid_param_estimates(::Type{Val{:NegativeBinomial}}, value::T) where T
    if keys(value) != Set((:p, :r))
        @error("Unrecognized parameter name for Negative Binomial. Expected (:p, :r)")
        return false
    end
    if !(0 < value[:p] < 1)
        @error("Estimator for p in Negative Binomial distribution should be in (0, 1)")
        return false
    end
    if value[:r] < 0
        @error("Estimator for r in Negative Binomial distribution should be in (0, ∞)")
        return false
    end
    return true
end

@private
function preprocess_user_input_estimators(input::Optional{Dict{Symbol, Dict{Symbol, T}}}) where T<:Real
    if isnothing(input)
        return Dict{Symbol, Any}()
    end
    ensure_feasible_data_type(keys(input))
    processed_input = Dict{Symbol, Dict{Symbol, Any}}()

    for dist in [:NegativeBinomial, :Gaussian, :Poisson, :Bernoulli, :Gamma]
        if haskey(input, dist)
            # @show(dist)
            processed_input[dist] = Dict{Symbol, Any}()
            processed_input[dist] = merge(processed_input[dist],
                                          preprocess_user_input_param_estimates(Val{dist}, input[dist]))
        end
    end
    return processed_input
end

@private
function preprocess_user_input_param_estimates(::Type{T}, value::S) where {T, S}
    ensure_valid_param_estimates(T, value)
    return Dict{Symbol, Any}(:user_input => value)
end

@private
function initialize_estimators(tracker, A, user_input_estimators)
    estimators = Dict{Symbol, Dict{Symbol, Any}}()
    for sym in [:Gaussian, :Bernoulli, :Poisson, :Gamma, :NegativeBinomial]
        estimators[sym] = Dict{Symbol, Any}()
    end
    if haskey(tracker, :NegativeBinomial) && length(tracker[:NegativeBinomial][:Observed]) > 0
        @info("Found negative binomial items. Use MOM for r and p")
        estimators[:NegativeBinomial][:MOM] = Concepts.estimator(MOM{AbstractNegativeBinomial}(),
                                                                 convert(Array{Float64},
                                                                         A[tracker[:NegativeBinomial][:Observed]]))
        @info(estimators[:NegativeBinomial][:MOM])
    end
    return merge(estimators, preprocess_user_input_estimators(user_input_estimators))
end

@private
function update(::Type{Val{:Gaussian}},
                A            ::Array{MaybeMissing{Float64}},
                Y12          ::Array{Float64},
                tracker,
                ρ            ::Float64,
                gd_iter      ::Int64,
                warmup       ::Dict{Symbol, Array{Float64}},
                γ            ::Float64,
                use_autodiff ::Bool,
                closed_form  ::Bool)
    if closed_form == true
        if haskey(tracker, :Gaussian) && length(tracker[:Gaussian][:Observed]) > 0
            Y12[tracker[:Gaussian][:Observed]] .= (1 / (1 + ρ)) * (A[tracker[:Gaussian][:Observed]] + ρ * Y12[tracker[:Gaussian][:Observed]])
        end
    else
        if haskey(tracker, :Gaussian) && length(tracker[:Gaussian][:Observed]) > 0
            Y12[tracker[:Gaussian][:Observed]] = train(use_autodiff ? provide(Loss{AbstractGaussian}()) : Loss{AbstractGaussian}(),
                                                       fx   = warmup[:Gaussian],
                                                       y    = A[tracker[:Gaussian][:Observed]],
                                                       c    = Y12[tracker[:Gaussian][:Observed]],
                                                       ρ    = ρ,
                                                       iter = gd_iter,
                                                       γ    = 0.2)
            warmup[:Gaussian] = Y12[tracker[:Gaussian][:Observed]]
        end
    end
end

@private
function update(::Type{Val{:Bernoulli}},
                A            ::Array{MaybeMissing{Float64}},
                Y12          ::Array{Float64},
                tracker,
                ρ            ::Float64,
                gd_iter      ::Int64,
                warmup       ::Dict{Symbol, Array{Float64}},
                γ            ::Float64,
                use_autodiff ::Bool)
    if haskey(tracker, :Bernoulli) && length(tracker[:Bernoulli][:Observed]) > 0
        Y12[tracker[:Bernoulli][:Observed]] = train(use_autodiff ? provide(Loss{AbstractBernoulli}()) : Loss{AbstractBernoulli}(),
                                                    fx   = warmup[:Bernoulli],
                                                    y    = A[tracker[:Bernoulli][:Observed]],
                                                    c    = Y12[tracker[:Bernoulli][:Observed]],
                                                    ρ    = ρ,
                                                    iter = gd_iter,
                                                    γ    = 0.2)
        warmup[:Bernoulli] = Y12[tracker[:Bernoulli][:Observed]]
    end
end

@private
function update(::Type{Val{:Poisson}},
                A            ::Array{MaybeMissing{Float64}},
                Y12          ::Array{Float64},
                tracker,
                ρ            ::Float64,
                gd_iter      ::Int64,
                warmup       ::Dict{Symbol, Array{Float64}},
                γ            ::Float64,
                use_autodiff ::Bool)
    if haskey(tracker, :Poisson) && length(tracker[:Poisson][:Observed]) > 0
        Y12[tracker[:Poisson][:Observed]] = train(use_autodiff ? provide(Loss{AbstractPoisson}()) : Loss{AbstractPoisson}(),
                                                  fx   = warmup[:Poisson],
                                                  y    = A[tracker[:Poisson][:Observed]],
                                                  c    = Y12[tracker[:Poisson][:Observed]],
                                                  ρ    = ρ,
                                                  iter = gd_iter,
                                                  γ    = 0.1)

        warmup[:Poisson] = Y12[tracker[:Poisson][:Observed]]
    end
end

@private
function update(::Type{Val{:Gamma}},
                A       ::Array{MaybeMissing{Float64}},
                Y12     ::Array{Float64},
                tracker,
                ρ       ::Float64,
                gd_iter ::Int64,
                warmup  ::Dict{Symbol, Array{Float64}},
                γ       ::Float64, use_autodiff::Bool)
    if haskey(tracker, :Gamma) && length(tracker[:Gamma][:Observed]) > 0
        Y12[tracker[:Gamma][:Observed]] = train(use_autodiff ? provide(Loss{AbstractGamma}()) : Loss{AbstractGamma}(),
                                                fx   = warmup[:Gamma],
                                                y    = A[tracker[:Gamma][:Observed]],
                                                c    = Y12[tracker[:Gamma][:Observed]],
                                                ρ    = ρ,
                                                iter = gd_iter,
                                                γ    = 0.2)
        warmup[:Gamma] = Y12[tracker[:Gamma][:Observed]]
    end
end

@private
function update(::Type{Val{:NegativeBinomial}},
                A            ::Array{MaybeMissing{Float64}},
                Y12          ::Array{Float64},
                tracker,
                ρ            ::Float64,
                gd_iter      ::Int64,
                warmup       ::Dict{Symbol, Array{Float64}},
                γ            ::Float64,
                use_autodiff ::Bool,
                estimator    ::Dict{Symbol, Any})
    if haskey(tracker, :NegativeBinomial) && length(tracker[:NegativeBinomial][:Observed]) > 0
        local r_est = nothing
        if haskey(estimator, :user_input)
            r_est = estimator[:user_input][:r]
        else
            r_est = estimator[:MOM][:r]
        end
        Y12[tracker[:NegativeBinomial][:Observed]] = negative_binomial_train(fx         = warmup[:NegativeBinomial],
                                                                             y          = A[tracker[:NegativeBinomial][:Observed]],
                                                                             c          = Y12[tracker[:NegativeBinomial][:Observed]],
                                                                             ρ          = ρ,
                                                                             iter       = gd_iter,
                                                                             γ          = 0.2,
                                                                             r_estimate = r_est)
        warmup[:NegativeBinomial] = Y12[tracker[:NegativeBinomial][:Observed]]
    end
end

@private
function update(::Type{Val{:ρ}}, primal_feasibility::Float64, dual_feasibility::Float64, current::Float64 = nothing)
    if (primal_feasibility < 0.5 * dual_feasibility)
        return 0.7 * current
    elseif (primal_feasibility > 2 * dual_feasibility)
        return 1.3 * current
    end
    return current
end

@private
function balance_gap(ρ::Float64, primal_feasibility::Float64, dual_feasibility::Float64)
    if (primal_feasibility < 0.5 * dual_feasibility)
        return 0.7 * ρ
    elseif (primal_feasibility > 2 * dual_feasibility)
        return 1.3 * ρ
    end
    return ρ
end

# update
const update(arg::Symbol, args...) = update(Val{arg}, args...)

@private
function calculate_primal_and_dual_residual(X, Z, W, C, Xinput, ρ)
    Fnorm = x -> norm(x,2);
    Xdual    = -ρ * (X - Xinput)
    Zdual    = ρ * (Z - C)
    normX    = 1 + Fnorm(X)
    primfeas = Fnorm(X - Z) / normX
    err1     = 1/ρ * Fnorm(W - Xdual)
    err2     = 1/ρ * Fnorm(W - Zdual)
    dualfeas = maximum([err1, err2]) / normX
    return primfeas, dualfeas
end

function set_diagonal(mat::Array{Float64, 2}, val::Array{Float64, 1})
    [mat[i, i] = val[i] for i in 1:length(val)]
end

function calculate_Z12_update(A, C,tracker, ρ, α, warmup, use_autodiff, gd_iter, estimators, closed_form)
    d1, d2 = size(A)
    Z12 = C[1:d1, (d1+1):(d1+d2)]
    update(:Gaussian,         A, Z12, tracker, ρ, gd_iter, warmup, 0.2,   use_autodiff, closed_form)
    update(:Bernoulli,        A, Z12, tracker, ρ, gd_iter, warmup, 0.2,   use_autodiff)
    update(:Poisson,          A, Z12, tracker, ρ, gd_iter, warmup, 0.2,  use_autodiff)
    update(:Gamma,            A, Z12, tracker, ρ, gd_iter, warmup, 0.005, use_autodiff)
    update(:NegativeBinomial, A, Z12, tracker, ρ, gd_iter, warmup, 0.005, use_autodiff, estimators[:NegativeBinomial])
    project!(ClosedInterval{Float64}(-α, α), Z12)
    return Z12
end

function set_block_12(mat, d1, d2, val)
    @. mat[1:d1, (d1+1):(d1+d2)] = val
end

function set_block_21(mat, d1, d2, val)
    @. mat[(d1+1):(d1+d2), 1:d1] = val
end

@private
function initialize_trackers(A::Array{MaybeMissing{Float64}}, type_assignment)
    type_tracker = Utilities.IndexTracker{Symbol}()
    if type_assignment == nothing
        disjoint_join(type_tracker, Concepts.fit(DistributionDeduction(), A))
    else
        disjoint_join(type_tracker, type_assignment)
    end
    disjoint_join(type_tracker, Concepts.fit(ObservedOrMissing(), A))
    return groupby(type_tracker, [:Observed, :Missing]), type_tracker
end

@private
function ensure_feasible(A::Array{MaybeMissing{Float64}})
    if isnothing(A)
        @error(io, "please provide data matrix")
        throw(MethodError())
    end
end

function format_log_data(x)
    if x == -10000000
        return "N/A"
    end
    return @sprintf("%3.2e", x)
end

@private
function print_optimization_log(iter,A, X, Z, Z12, W, II, Rp, Rd, ρ, λ, μ, tracker, estimators, io)
    R = abs.(maximum(diag(Z)))
    local gaussian_loss          = -10000000
    local bernoulli_loss         = -10000000
    local poisson_loss           = -10000000
    local gamma_loss             = -10000000
    local negative_binomial_loss = -10000000
    if haskey(tracker, :Gaussian)
        # gaussian_loss = norm(Z12[tracker[:Gaussian][:Observed]] -  A[tracker[:Gaussian][:Observed]])^2
        gaussian_loss = evaluate(Loss{AbstractGaussian}(),
                                 Z12[tracker[:Gaussian][:Observed]],
                                 A[tracker[:Gaussian][:Observed]],
                                 similar(Z12[tracker[:Gaussian][:Observed]]),
                                 0)
    end
    if haskey(tracker, :Bernoulli)
        bernoulli_loss = evaluate(Loss{AbstractBernoulli}(),
                                  Z12[tracker[:Bernoulli][:Observed]],
                                  A[tracker[:Bernoulli][:Observed]],
                                  similar(Z12[tracker[:Bernoulli][:Observed]]),
                                  0)
    end

    if haskey(tracker, :NegativeBinomial)
        local r_est = nothing
        if haskey(estimators[:NegativeBinomial], :user_input)
            r_est = estimators[:NegativeBinomial][:user_input][:r]
        else
            r_est = estimators[:NegativeBinomial][:MOM][:r]
        end
        negative_binomial_loss = evaluate(Loss{AbstractNegativeBinomial}(),
                                          Z12[tracker[:NegativeBinomial][:Observed]],
                                          A[tracker[:NegativeBinomial][:Observed]],
                                          similar(Z12[tracker[:NegativeBinomial][:Observed]]),
                                          0,
                                          r_estimate = r_est)
    end
    if haskey(tracker, :Poisson)
        poisson_loss = evaluate(Loss{AbstractPoisson}(),
                                Z12[tracker[:Poisson][:Observed]],
                                A[tracker[:Poisson][:Observed]],
                                similar(Z12[tracker[:Poisson][:Observed]]),
                                0)
    end
    if haskey(tracker, :Gamma)
        gamma_loss = evaluate(Loss{AbstractGamma}(),
                              Z12[tracker[:Gamma][:Observed]],
                              A[tracker[:Gamma][:Observed]],
                              similar(Z12[tracker[:Gamma][:Observed]]),
                              0)
        # gamma_loss = 0
    end

    data = [iter,
            Rp,
            Rd,
            gaussian_loss,
            bernoulli_loss,
            poisson_loss,
            negative_binomial_loss,
            gamma_loss,
            maximum(abs.(Z12)),
            μ * tr(II * X),
            abs.(maximum(diag(Z)))
            ]
    new_data = map(x -> format_log_data(x) ,data)
    new_data[1] = string(iter)
    add_row(header_list, data=new_data, io = io)
end

#TODO
# function standardize(A, tracker, estimators)
#   A[tracker[:Gaussian][:Observed]] .= (A[tracker[:Gaussian][:Observed]] .- estimators[:Gaussian][:μ]) ./ estimators[:Gaussian][:σ]
# end

#TODO
# function destandardize(A, type_tracker, estimators)
#   A[tracker[:Gaussian]]] .= (A[tracker] .* estimators[:Gaussian][:σ]) .+ estimators[:Gaussian][:μ]
# end
struct OneShotADMM end

function Concepts.complete(model::OneShotADMM;
                           A::Array{MaybeMissing{Float64}} = nothing,
                           α::Float64                       = maximum(A[findall(x -> !ismissing(x),A)]),
                           λ::Float64                       = 5e-1,
                           μ::Float64                       = 5e-4,
                           ρ::Float64                       = 0.3,
                           τ::Float64                       = 1.618,
                           maxiter::Int64                   = 200,
                           stoptol::Float64                 = 1e-5,
                           use_autodiff::Bool               = false,
                           gd_iter::Int64                   = 50,
                           debug_mode::Bool                 = false,
                           interactive_plot                 = false,
                           type_assignment                  = nothing,
                           warmup                           = nothing,
                           start_var                        = nothing,
                           dynamic_ρ                        = true,
                           user_input_estimators            = nothing,
                           project_rank                     = nothing,
                           io::IO                           = Base.stdout,
                           eigen_solver                     = KrylovMethods(),
                           closed_form_update               = false)
    return complete(;A = A,
                    α = α,
                    λ = λ,
                    μ = μ,
                    ρ = ρ,
                    τ = τ,
                    maxiter = maxiter,
                    stoptol = stoptol,
                    use_autodiff = use_autodiff,
                    gd_iter = gd_iter,
                    debug_mode = debug_mode,
                    interactive_plot = interactive_plot,
                    type_assignment = type_assignment,
                    warmup = warmup,
                    start_var = start_var,
                    dynamic_ρ,
                    user_input_estimators = user_input_estimators,
                    project_rank = project_rank,
                    io = io,
                    eigen_solver = eigen_solver,
                    closed_form_update = closed_form_update)
end

function Concepts.complete(;A::Array{MaybeMissing{Float64}} = nothing,
                           α::Float64                       = maximum(A[findall(x -> !ismissing(x),A)]),
                           λ::Float64                       = 5e-1,
                           μ::Float64                       = 5e-4,
                           ρ::Float64                       = 0.3,
                           τ::Float64                       = 1.618,
                           maxiter::Int64                   = 200,
                           stoptol::Float64                 = 1e-5,
                           use_autodiff::Bool               = false,
                           gd_iter::Int64                   = 50,
                           debug_mode::Bool                 = false,
                           interactive_plot                 = false,
                           type_assignment                  = nothing,
                           warmup                           = nothing,
                           start_var                        = nothing,
                           dynamic_ρ                        = true,
                           user_input_estimators            = nothing,
                           project_rank                     = nothing,
                           io::IO                           = Base.stdout,
                           eigen_solver                     = KrylovMethods(),
                           closed_form_update               = false)
    logger = SimpleLogger(io)
    global_logger(logger)
    if isnothing(project_rank)
        @info("Using Full Eigen Decomposition.")
    else
        @info("Using Fast Eigen")
    end
    ensure_feasible(A)
    d1, d2                               = size(A);
    Z::Array{Float64, 2}                 = zeros(d1 + d2,  d1 + d2)
    X::Array{Float64, 2}                 = zeros(d1 + d2,  d1 + d2)
    W::Array{Float64, 2}                 = zeros(d1 + d2,  d1 + d2)
    C::Array{Float64, 2}                 = zeros(d1 + d2,  d1 + d2)
    if !isnothing(start_var)
        Z = start_var[:Z]
        X = start_var[:X]
        W = start_var[:W]
        C = start_var[:C]
    end
    Xinput::Array{Float64, 2}            = zeros(d1 + d2,  d1 + d2)
    II                                   = sparse(1.0I, d1 + d2, d1 + d2)
    tracker, type_tracker                = initialize_trackers(A, type_assignment)
    # initialize warmup input for various gradient descent procedures
    if isnothing(warmup)
        warmup::Dict{Symbol, Array{Float64}} = initialize_warmup(tracker, A)
    end
    # initialize various estimators
    estimators::Dict{Symbol, Any}        = initialize_estimators(tracker, A, user_input_estimators)
    # print optimization path table header
    table_header(header_list, io = io)
    for iter = 1:maxiter
        @. Xinput = Z + W/ρ
        # step 1
        # try
        if isnothing(project_rank)
            X = project(SemidefiniteCone(), Z + W / ρ - (μ / ρ) * II)
        else
            X = project(SemidefiniteCone(rank = project_rank), Z + W / ρ - (μ / ρ) * II,
                        eigs_implementation = eigen_solver)
        end
        # catch
        #   @warn("Manual fix for numerical instability")
        #   fix =  Z + W / ρ - (μ / ρ) * II
        #   fix[findall(x -> x== Inf || isnan(x), fix)] .= rand()
        #   X = project(SemidefiniteCone(rank = 20), fix)
        # end
        # Step 2
        @. C = X - 1/ρ * W; @. Z = C
        Z12 = calculate_Z12_update(A, C, tracker, ρ, α, warmup, use_autodiff, gd_iter, estimators, closed_form_update)
        set_block_12(Z, d1, d2, Z12)
        set_block_21(Z, d1, d2, Z12')
        set_diagonal(Z, diag(C) - (λ / ρ) * l1BallProjection(diag(C) * ρ / λ, 1))
        # step 3
        @. W = W + τ * ρ * (Z - X)
        if rem(iter, 10)==1
            primfeas, dualfeas = calculate_primal_and_dual_residual(X, Z, W, C, Xinput, ρ)
            print_optimization_log(iter, A, X, Z, Z12, W, II, primfeas, dualfeas, ρ, λ, μ, tracker, estimators, io)
            if dynamic_ρ
                ρ = balance_gap(ρ, primfeas, dualfeas)
            end
        end
    end
    completedMatrix = C[1:d1, (d1+1):(d1+d2)]
    last_info = Dict(:Z => Z,
                     :X => X,
                     :W => W,
                     :C => C,
                     :warmup => warmup)
    return completedMatrix, type_tracker, tracker, last_info
end

include("./ChainedADMM.jl")

end


text:
using Distributed
using SharedArrays

@everywhere function BoyerMoore(A::AbstractArray{T, 1}, k::Int=2)::Dict{T, Int} where T
    candidates = Dict{T, Int}()
    for a in A
        if length(candidates) < k - 1 || haskey(candidates, a)
            candidates[a] = get!(candidates, a, 0) + 1
        else
            to_del = Vector{T}()
            for key in keys(candidates)
                candidates[key] -= 1
                if candidates[key] <= 0
                    append!(to_del, key)
                end
            end
            for key in to_del
                pop!(candidates, key)
            end
        end
    end
    return candidates
end

function majority_element(A::Vector{T}, k::Int=2)::Vector{T} where T
    @assert k >= 2 "k must be an integer no less than 2"

    candidates = BoyerMoore(A, k)
    for key in keys(candidates)
        candidates[key] = 0
    end
    for a in A
        if haskey(candidates, a)
            candidates[a] += 1
        end
    end

    bar = div(length(A), k) + 1
    return [key for (key, v) in candidates if v >= bar]
end

@everywhere function merge_candidates!(X::Dict{T, Int}, Y::Dict{T, Int}, k::Int=2) where T
    for (key, v) in Y
        if length(X) < k - 1 || haskey(X, key)
            X[key] = get!(X, key, 0) + v
        else
            min_v = min(minimum(values(X)), v)
            to_del = Vector{T}()
            for a in keys(X)
                X[a] -= min_v
                if X[a] <= 0
                    append!(to_del, a)
                end
            end
            for a in to_del
                pop!(X, a)
            end
            if v > min_v
                X[key] = v - min_v
            end
        end
    end
    return X
end

function distributed_majority_element(A::Vector{T}, p::Int, k::Int=2)::Vector{T} where T
    @assert k >= 2 "k must be an integer no less than 2"

    n = length(A)
    step = n ÷ p

    A = SharedVector(A)
    candidates = @distributed merge_candidates! for i = 1:p
        left = (i - 1) * step + 1
        right = i == p ? n : i * step
        BoyerMoore(view(A, left:right), k)
    end

    global_counter = @distributed mergewith(+) for i = 1:p
        counter = Dict(key => 0 for (key, v) in candidates)
        left = (i - 1) * step + 1
        right = i == p ? n : i * step
        for a in view(A, left:right)
            if haskey(counter, a)
                counter[a] += 1
            end
        end
        counter
    end

    bar = n ÷ k + 1
    return [key for (key, v) in global_counter if v >= bar]
end

function parallel_majority_element(A::Vector{T}, p::Int, k::Int=2)::Vector{T} where T
    @assert k >= 2 "k must be an integer no less than 2"

    n = length(A)
    step = n ÷ p

    pool = Vector{Dict{T, Int}}(undef, p)
    Threads.@threads for i = 1:p
        left = (i - 1) * step + 1
        right = i == p ? n : i * step
        pool[i] = BoyerMoore(view(A, left:right), k)
    end

    candidates = reduce(merge_candidates!, pool)

    Threads.@threads for i = 1:p
        pool[i] = Dict(key => 0 for (key, v) in candidates)
        left = (i - 1) * step + 1
        right = i == p ? n : i * step
        for a in view(A, left:right)
            if haskey(pool[i], a)
                pool[i][a] += 1
            end
        end
    end
    counter = reduce(mergewith(+), pool)
    bar = n ÷ k + 1
    return [key for (key, v) in counter if v >= bar]
end


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

Testrunner for the LSH module

================================================================#

#==============
Helper functions and constants for LSH
==============#

include("utils.jl")

#========================
Tests
========================#

include("test_intervals.jl")
include("test_similarities.jl")

include(joinpath("hashes", "test_simhash.jl"))
include(joinpath("hashes", "test_minhash.jl"))
include(joinpath("hashes", "test_lphash.jl"))
include(joinpath("hashes", "test_sign_alsh.jl"))
include(joinpath("hashes", "test_lshfunction.jl"))

include(joinpath("function_hashing", "test_monte_carlo.jl"))


text:
using PICDataStructures, Test
using LaserTypes
using StaticArrays
using UnitfulAtomic
using PhysicalConstants.CODATA2018: c_0

λ = 15105.659329029153
c = austrip(c_0)
m = 2
p = 2
a₀ = 2.0
w₀ = 75 * λ

laser = setup_laser(LaguerreGaussLaser, :atomic;
    λ,
    m,
    p,
    a₀,
    w₀,
    profile = GaussProfile,
    τ = 10λ/(2π*c), z₀ = 0.
)

x_grid = y_grid = z_grid = -30:30

grid = SparseAxisGrid(x_grid, y_grid, y_grid)

t = 0.

electricfield = vectorfield(grid) do (x,y,z)
    r = SVector{3}(x, y, z).*w₀
    E(r, t, laser)
end

@test ndims(electricfield) == 3

E_slice = selectdim(electricfield, :z, 0.)

using CairoMakie

plotdata(E_slice)