text:
# Different ways to simulate molecules

export
    VelocityVerlet,
    StormerVerlet,
    simulate!

# Forces are often expressed per mol but this dimension needs removing for use in the integrator
function remove_molar(x)
    fx = first(x)
    if dimension(fx) == u"𝐋 * 𝐍^-1 * 𝐓^-2"
        T = typeof(ustrip(fx))
        return x / T(Unitful.Na)
    else
        return x
    end
end

"""
    VelocityVerlet(; <keyword arguments>)

The velocity Verlet integrator.

# Arguments
- `dt::T`: the time step of the simulation.
- `coupling::C=NoCoupling()`: the coupling which applies during the simulation.
"""
struct VelocityVerlet{T, C}
    dt::T
    coupling::C
end

VelocityVerlet(; dt, coupling=NoCoupling()) = VelocityVerlet(dt, coupling)

"""
    StormerVerlet(; <keyword arguments>)

The Störmer-Verlet integrator.
In this case the `velocities` given to the simulator act as the previous step
coordinates for the first step.
Does not currently work with units or thermostats.

# Arguments
- `dt::T`: the time step of the simulation.
- `coupling::C=NoCoupling()`: the coupling which applies during the simulation.
"""
struct StormerVerlet{T, C}
    dt::T
    coupling::C
end

StormerVerlet(; dt, coupling=NoCoupling()) = StormerVerlet(dt, coupling)

"""
    simulate!(system, simulator, n_steps; parallel=true)
    simulate!(system, simulator; parallel=true)

Run a simulation on a system according to the rules of the given simulator.
Custom simulators should implement this function.
"""
function simulate!(sys::System{D, S, false},
                    sim::VelocityVerlet,
                    n_steps::Integer;
                    parallel::Bool=true) where {D, S}
    # See https://www.saylor.org/site/wp-content/uploads/2011/06/MA221-6.1.pdf for
    #   integration algorithm - used shorter second version
    neighbors = find_neighbors(sys, sys.neighbor_finder; parallel=parallel)
    accels_t = accelerations(sys, neighbors; parallel=parallel)
    accels_t_dt = zero(accels_t)

    @showprogress for step_n in 1:n_steps
        run_loggers!(sys, neighbors, step_n)

        # Update coordinates
        for i in 1:length(sys)
            sys.coords[i] += sys.velocities[i] * sim.dt + remove_molar(accels_t[i]) * (sim.dt ^ 2) / 2
            sys.coords[i] = wrap_coords.(sys.coords[i], sys.box_size)
        end

        accels_t_dt = accelerations(sys, neighbors; parallel=parallel)

        # Update velocities
        for i in 1:length(sys)
            sys.velocities[i] += remove_molar(accels_t[i] + accels_t_dt[i]) * sim.dt / 2
        end

        apply_coupling!(sys, sim, sim.coupling)
        neighbors = find_neighbors(sys, sys.neighbor_finder, neighbors, step_n; parallel=parallel)

        accels_t = accels_t_dt
    end
    return sys
end

function simulate!(sys::System{D, S, true},
                    sim::VelocityVerlet,
                    n_steps::Integer;
                    parallel::Bool=true) where {D, S}
    if length([inter for inter in values(sys.general_inters) if !inter.nl_only]) > 0
        neighbors_all = all_neighbors(length(sys))
    else
        neighbors_all = nothing
    end
    neighbors = find_neighbors(sys, sys.neighbor_finder)
    accels_t = accelerations(sys, sys.coords, sys.atoms, neighbors, neighbors_all)
    accels_t_dt = zero(accels_t)

    for step_n in 1:n_steps
        run_loggers!(sys, neighbors, step_n)

        sys.coords += sys.velocities .* sim.dt .+ (remove_molar.(accels_t) .* sim.dt ^ 2) ./ 2
        sys.coords = wrap_coords_vec.(sys.coords, (sys.box_size,))
        accels_t_dt = accelerations(sys, sys.coords, sys.atoms, neighbors, neighbors_all)
        sys.velocities += remove_molar.(accels_t .+ accels_t_dt) .* sim.dt / 2

        apply_coupling!(sys, sim, sim.coupling)
        neighbors = find_neighbors(sys, sys.neighbor_finder, neighbors, step_n)

        accels_t = accels_t_dt
    end
    return sys
end

function simulate!(sys::System,
                    sim::StormerVerlet,
                    n_steps::Integer;
                    parallel::Bool=true)
    neighbors = find_neighbors(sys, sys.neighbor_finder; parallel=parallel)

    @showprogress for step_n in 1:n_steps
        run_loggers!(sys, neighbors, step_n)

        accels_t = accelerations(sys, neighbors; parallel=parallel)

        # Update coordinates
        coords_copy = sys.coords
        for i in 1:length(sys)
            sys.coords[i] = sys.coords[i] + vector(sys.velocities[i], sys.coords[i], sys.box_size) + remove_molar(accels_t[i]) * sim.dt ^ 2
            sys.coords[i] = wrap_coords.(sys.coords[i], sys.box_size)
        end
        sys.velocities = coords_copy

        apply_coupling!(sys, sim, sim.coupling)
        neighbors = find_neighbors(sys, sys.neighbor_finder, neighbors, step_n; parallel=parallel)
    end
    return sys
end


text:
# %% accessing model
@testset "Model Access" begin
@testset "Accessing model" begin
    model = ABM(Agent0; properties = Dict(:a => 2, :b => "test"))
    # add 5 aents
    for i in 1:5
        add_agent!(model)
    end
    @test model.scheduler == fastest
    @test typeof(model.agents) <: Dict
    @test model.space == nothing
    @test model.properties == Dict(:a => 2, :b => "test")
    a = model[1]
    @test a isa Agent0
    @test a.id == 1
    @test model.a == 2
    @test model.b == "test"

    newa = Agent0(6)
    model[6] = newa
    @test model[6] == newa
    @test_throws ArgumentError model[7] = newa

    prop2 = Agent2(1, 0.5)
    model2 = ABM(Agent0; properties = prop2)
    @test model2.weight == 0.5

end

@testset "model access typestability" begin
    prop1 = Dict(:a => 0.5)
    prop2 = Agent2(1, 0.5)
    model1 = ABM(Agent0; properties = prop1)
    model2 = ABM(Agent0; properties = prop2)

    test1(model) = model.a
    test2(model) = model.weight
    @test_nowarn @inferred test1(model1)
    @test_nowarn @inferred test2(model2)
end

@testset "Core methods" begin
    model = ABM(Agent3, GridSpace((5,5)))
    @test Agents.agenttype(model) == Agent3
    @test Agents.spacetype(model) <: GridSpace
    @test ne(model) == 40
    @test nv(model) == 25
    @test Agents.agent_positions(model) == Agents.agent_positions(model.space)
    @test size(model.space) == (5,5)
    @test all(isempty(n, model) for n in nodes(model))
end

@testset "Display" begin
    model = ABM(Agent3, GridSpace((5,5)))
    @test sprint(show, model) == "AgentBasedModel with 0 agents of type Agent3\n space: GridSpace with 25 nodes and 40 edges\n scheduler: fastest"
    @test sprint(show, model.space) == "GridSpace with 25 nodes and 40 edges"
    model = ABM(Agent6, ContinuousSpace(2))
    @test sprint(show, model.space) == "2-dimensional periodic ContinuousSpace"
end

end

ext: jl


text:
@testset "IO" begin
    c8 = Chip()
    @test_nowarn CHIP8.loadApplication(c8, "pong2.c8")
end


text:
#
#TODO: As far as we know `q`, `q_x`, `q_xx`, we can move
# the rest into a function, which is less error prone, and
# will save a lot of space
#
@doc """
### EulerEquationMod.applySourceTerm

  This function updates eqn.res with the source term.

  Inputs:
    mesh
    sbp
    eqn
    opts
    src_func:  the functor that returns the value of the source term at a node
               This functor must have the signature:
               src_func(q, coords, params, t)
               where coords is a vector of length 2 or 3 containing the x and y
               coordinates of the node, params is the ParamType, t is the
               current time, and q is the vector to be populated with the
               source term values.

  Outputs: none

"""->
function applySourceTerm(mesh, sbp, eqn, opts, src_func::SRCType)
#TODO: check that the k loop vectorizes
  weights = sbp.w
  q_vals = zeros(eltype(eqn.q), mesh.numDofPerNode)
  t = eqn.params.t

  for i=1:mesh.numEl
    jac_i = ro_sview(mesh.jac, :, i)
    res_i = sview(eqn.res, :, :, i)
    for j=1:mesh.numNodesPerElement
      coords_j = ro_sview(mesh.coords, :, j, i)
      src_func(q_vals, coords_j, eqn.params, t)
      fac = weights[j]/jac_i[j]
      for k=1:mesh.numDofPerNode
        res_i[k, j] += fac*q_vals[k]
      end
    end
  end

  return nothing
end

mutable struct SRC0 <: SRCType  # dummy source functor, it should nevery actually be called
end

@doc """
### AdvectionEquationMod.SRC0

  This is the zero source term.  This is the default of source term
  is specified
"""->
mutable struct SRCExp <: SRCType
end

function (obj::SRCExp)(q::AbstractVector, coords::AbstractVector, params::ParamType{2}, t)
  x = coords[1]
  y = coords[2]
  gamma_1 = params.gamma_1
  # a and b are parameters determine the scale and offset of the solution
  # the must be the same here and in the boundary condition
  af = 1/5  # a = 1/af
  b = 0.01

  q[1] = 2*y*af*exp(2*x*y*af + b) + 3*x*af*exp(3*x*y*af + b)
  q[2] = 3*y*af*exp(3*x*y*af + b) + 5*y*af*exp(5*x*y*af + b) + 4*x*af*exp(4*x*y*af+b)
  q[3] = 4*y*af*exp(4*x*y*af + b) + 10*x*af*exp(5*x*y*af + b)
  q[4] = 6*y*af*(1/gamma_1 + 1.5)*exp(6*x*y*af + b) + 2*y*af*exp(4*x*y*af + b) + 7*x*af*(1/gamma_1 + 1.5)*exp(7*x*y*af + b) + 2.5*x*af*exp(5*x*y*af + b)

  # make this work for other variables
  convertFromNaturalToWorkingVars(params, q, q)
end

# declare some global variables
# this is somewhat namespace polluting
global const MMSExp_a = 1/500
global const MMSExp_b = 0.01
global const MMSExp_c1 = 1
global const MMSExp_c2 = 2
global const MMSExp_c3 = 3
global const MMSExp_c4 = 4
global const MMSExp_c5 = 20
global const MMSExp_d1 = 1
global const MMSExp_d2 = 0.05
global const MMSExp_d3 = 0.15
global const MMSExp_d4 = 0.25
global const MMSExp_d5 = 1

function (obj::SRCExp)(q::AbstractVector, coords::AbstractVector, params::ParamType{3}, t)

  x = coords[1]
  y = coords[2]
  z = coords[3]

  # constant parameters
  gamma_1 = params.gamma_1
  a = MMSExp_a
  b = MMSExp_b
  c1 = MMSExp_c1
  c2 = MMSExp_c2
  c3 = MMSExp_c3
  c4 = MMSExp_c4
  c5 = MMSExp_c5
  d1 = MMSExp_d1
  d2 = MMSExp_d2
  d3 = MMSExp_d3
  d4 = MMSExp_d4
  d5 = MMSExp_d5

  t2 = exp(b);
  t3 = a*c2*x*y*z;
  t4 = exp(t3);
  t5 = a*c4*x*y*z;
  t6 = exp(t5);
  t7 = c4*d4*t6*x*y;
  t8 = a*c3*x*y*z;
  t9 = exp(t8);
  t10 = c3*d3*t9*x*z;
  t11 = a*c5*x*y*z;
  t12 = exp(t11);
  t13 = 1.0/d1;
  t16 = a*c1*x*y*z;
  t14 = exp(-t16);
  t15 = c2*d2*t4*y*z;
  t17 = b-t16;
  t18 = exp(t17);
  t19 = d2*d2;
  t20 = a*c2*x*y*z*2.0;
  t21 = exp(t20);
  t22 = d3*d3;
  t23 = a*c3*x*y*z*2.0;
  t24 = exp(t23);
  t25 = d4*d4;
  t26 = a*c4*x*y*z*2.0;
  t27 = exp(t26);
  t28 = b+t11;
  t29 = exp(t28);
  t30 = c2*t19*t21;
  t31 = c3*t22*t24;
  t32 = c4*t25*t27;
  t33 = t30+t31+t32;
  t34 = t19*t21;
  t35 = t22*t24;
  t36 = t25*t27;
  t37 = t34+t35+t36;
  t38 = 1.0/gamma_1;
  t39 = c1-c4;
  t40 = exp(-a*t39*x*y*z);
  t41 = c1-c3;
  t42 = exp(-a*t41*x*y*z);
  t43 = d5*t29;
  t44 = d5*t29*t38;
  t45 = t13*t18*t37*(1.0/2.0);
  t46 = t43+t44+t45;
  t47 = c1-c2;
  t48 = exp(-a*t47*x*y*z);
  q[1] = a*t2*(t7+t10+t15);

  q[2] = a*c5*d5*t2*t12*y*z + a*d2*t2*t4*t13*t14*(t7+t10 - c1*d4*t6*x*y + c2*d4*t6*x*y - c1*d3*t9*x*z + c2*d3*t9*x*z - c1*d2*t4*y*z + c2*d2*t4*y*z*2.0);

  q[3] = a*c5*d5*t2*t12*x*z + a*d3*t2*t9*t13*t14*( (t7+t15 - c1*d4*t6*x*y + c3*d4*t6*x*y - c1*d3*t9*x*z) + (c3*d3*t9*x*z*2.0 - c1*d2*t4*y*z + c3*d2*t4*y*z) );

  q[4] = a*c5*d5*t2*t12*x*y + a*d4*t2*t6*t13*t14*( (t10+t15 - c1*d4*t6*x*y + c4*d4*t6*x*y*2.0 - c1*d3*t9*x*z) + (c4*d3*t9*x*z - c1*d2*t4*y*z + c4*d2*t4*y*z) );

  q[5] = d4*t13*t40*(a*t13*t18*t33*x*y + a*c5*d5*t29*x*y + a*c5*d5*t29*t38*x*y - a*c1*t13*t18*t37*x*y*(1.0/2.0)) + d3*t13*t42*(a*t13*t18*t33*x*z + a*c5*d5*t29*x*z + a*c5*d5*t29*t38*x*z - a*c1*t13*t18*t37*x*z*(1.0/2.0)) + d2*t13*t48*(a*t13*t18*t33*y*z + a*c5*d5*t29*y*z + a*c5*d5*t29*t38*y*z - a*c1*t13*t18*t37*y*z*(1.0/2.0)) - a*d4*t13*t39*t40*t46*x*y - a*d3*t13*t41*t42*t46*x*z - a*d2*t13*t46*t47*t48*y*z;

  return nothing
end

"""
  Functor for source term corresponding to ICPeriodicMMS
"""
mutable struct SRCPeriodicMMS <: SRCType
end

function (obj::SRCPeriodicMMS)(q::AbstractVector, coords::AbstractVector,
              params::ParamType{2}, t)

  x = coords[1]
  y = coords[2]
  gamma_1 = params.gamma_1

  t4 = t*2.0;
  t2 = -t4+x+y;
  t3 = 3.141592653589793*t2;
  t5 = cos(t3);
  t6 = sin(t3);
  t7 = t6+1.5E1;
  t8 = 3.141592653589793*gamma_1*t5*t7*(1.0/5.0E1);
  q[1] = 0
  q[2] = t8;
  q[3] = t8;
  q[4] = 3.141592653589793*gamma_1*t5*t7*(1.0/2.5E1);

  return nothing
end

function (obj::SRCPeriodicMMS)(q::AbstractVector, coords::AbstractVector,
              params::ParamType{3}, t)

  x = coords[1]
  y = coords[2]
  z = coords[3]
  gamma_1 = params.gamma_1
  gamma = params.gamma
#=
  t2 = x+y+z;
  t3 = 3.141592653589793*t2;
  t4 = cos(t3);
  t5 = sin(t3);
  t6 = t5*(1.0/1.0E1);
  t7 = t6+2.0;
  t8 = 3.141592653589793*gamma_1*t4*t7*(1.0/5.0);
=#
#  tmp = pi*gamma_1*cos(pi*(x + y + z))*(sin(pi*(x + y + z))/10 + 2)/5
#=
  q[1] = 0;
  q[2] = 0
  q[3] = 0
  q[4] = 0
  q[5] = 0;
=#

  t5 = t*2.0;
  t2 = -t5+x+y+z;
  t3 = 3.141592653589793*t2;
  t4 = cos(t3);
  t6 = gamma_1*2.5E1;
  t7 = sin(t3);
  t8 = gamma_1*t7*2.0;
  t9 = t6+t8+1.0E1;
  t10 = 3.141592653589793*t4*t9*(1.0/1.0E2);
  q[1] = 3.141592653589793*t4*(1.0/1.0E1);
  q[2] = t10;
  q[3] = t10;
  q[4] = t10;
  q[5] = 3.141592653589793*t4*(gamma_1*7.5E1+t7*2.0+gamma_1*t7*6.0+4.0E1)*(1.0/1.0E2);

  return nothing
end

"""
  Source term for testing shock capturing terms
"""
mutable struct SRCLaplace <: SRCType
end

function (obj::SRCLaplace)(q::AbstractVector, coords::AbstractVector,
              params::ParamType{2}, t)

  x = coords[1]; y = coords[2]
  val = 5*sin(x)*sin(2*y)
  for i=1:length(q)
    q[i] = val
  end
end

@doc """
### EulerEquationMod.SRCDict

  Stores the functors that evaluate source terms at a node.  Every new
  functor should be added to this dictonary

  All functors must have the signature:

  src_func(q, coords, params::ParamType, t)

  where coords is the vector of length 2 containing the x and y coordinates
  of the node, t is the current time, and q is the vector to be populated with
  the source term values.
"""->
global const SRCDict = Dict{String, SRCType}(
"SRCExp" => SRCExp(),
"SRCPeriodicMMS" => SRCPeriodicMMS(),
"SRCLaplace" => SRCLaplace(),
"SRC0" => SRC0(),
)

@doc """
### EulerEquationMod.getSRCFunctors

  This function gets the functor specified by opts["SRCname"] and stores
  it to the equation object.  Currently one 1 source functor is allowed.

"""->
function getSRCFunctors(mesh::AbstractMesh, sbp::AbstractOperator,
                        eqn::EulerData, opts)

  # currently we only allow 1 source functor
  eqn.src_func = SRCDict[opts["SRCname"]]
  return nothing
end


text:
import ..SingleStackUtils: single_stack_diagnostics, NodalStack
using ..ODESolvers

# Convenience wrapper
single_stack_diagnostics(solver_config; kwargs...) = single_stack_diagnostics(
    solver_config.dg.grid,
    solver_config.dg.balance_law,
    gettime(solver_config.solver),
    solver_config.dg.direction;
    prognostic = solver_config.Q,
    auxiliary = solver_config.dg.state_auxiliary,
    diffusive = solver_config.dg.state_gradient_flux,
    hyperdiffusive = solver_config.dg.states_higher_order[2],
    kwargs...,
)

# Convenience wrapper
NodalStack(solver_config; kwargs...) = NodalStack(
    solver_config.dg.balance_law,
    solver_config.dg.grid;
    prognostic = solver_config.Q,
    auxiliary = solver_config.dg.state_auxiliary,
    diffusive = solver_config.dg.state_gradient_flux,
    hyperdiffusive = solver_config.dg.states_higher_order[2],
    kwargs...,
)


text:
module HeavyChain

type Chain
  q::Array{Float64, 1} #generalised coordinates. Contains [x0, alpha_0, alpha_1, alpha_n], the pos of the cart and all link-angles
  u::Array{Float64, 1} #time derivative of q
  massmat::Array{Float64, 2}
  fvec::Array{Float64, 1}
end

function create_chain()
  q = zeros(21) #cart position and all angles are zero
  u = zeros(21) #all velocities are zero
  massmat = zeros(21, 21) #mass matrix, memory is overwritten at each forward simulation step
  fvec = zeros(21) #force vector, memory is overwritten at each forward simulation step
  return Chain(q, u, massmat, fvec)
end

export Chain, create_chain

include("src/simulate.jl")
export simulate!

#include("src/plot.jl")
#export heavychain_plot!

include("src/misc.jl")
export getendpoint

end

max_stars_repo_name: NikonPic/REINFORCEMENT-JULIA

max_issues_repo_name: NikonPic/REINFORCEMENT-JULIA

max_forks_repo_name: NikonPic/REINFORCEMENT-JULIA


text:
using Test
using StaticArrays
using CellListMap

using MolecularMinimumDistances

@testset "Initalization functions" begin

    x = [ rand(SVector{3,Float64}) for _ in 1:12 ]
    #
    # Initialization functions
    #
    @test length(init_list(x, i -> mol_index(i,4))) == 3
    @test length(init_list(Float64, 3)) == 3
    inds = [ 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2, 2 ]
    @test length(init_list(x, i -> inds[i])) == 2

end

@testset "Single set" begin
    #
    # Single set of molecules
    #
    x = [ rand(SVector{3,Float64}) for _ in 1:135 ]

    for box in [ Box([1,1,1],0.2),
                 Box([1.0 0.2 0.0
                      0.2 1.0 0.2
                      0.0 0.0 1.0], 0.2) ]

        x_list_naive = MolecularMinimumDistances.naive_md(x,3,box)

        x_list = minimum_distances(x,3,box,parallel=false)
        @test x_list ≈ x_list_naive

        x_list = minimum_distances(x,3,box,parallel=true)
        @test x_list ≈ x_list_naive

        cl = CellList(x,box,parallel=false)
        x_list = init_list(x, i -> mol_index(i,3))
        minimum_distances!(i -> mol_index(i,3), x_list,box,cl,parallel=false)
        @test x_list ≈ x_list_naive

        cl = CellList(x,box,parallel=true)
        minimum_distances!(i -> mol_index(i,3),x_list,box,cl,parallel=true)
        @test x_list ≈ x_list_naive

    end

end

@testset "Disjoint: one list" begin

    #
    # Disjoint sets: return only one list
    #
    x = [ rand(SVector{3,Float64}) for _ in 1:100 ]
    y = [ rand(SVector{3,Float64}) for _ in 1:90 ]

    for box in [ Box([1,1,1],0.2),
                 Box([1.0 0.2 0.0
                      0.2 1.0 0.2
                      0.0 0.0 1.0], 0.2) ]

        x_list_naive = MolecularMinimumDistances.naive_md(x,y,5,box)

        x_list = minimum_distances(x,y,5,box,parallel=false)
        @test x_list ≈ x_list_naive

        x_list = minimum_distances(x,y,5,box,parallel=true)
        @test x_list ≈ x_list_naive

        cl = CellList(x,y,box,parallel=false)
        x_list = init_list(x, i -> mol_index(i,5))
        minimum_distances!(
            i -> mol_index(i,5),
            x_list,box,cl;
            parallel=false
        )
        @test x_list ≈ x_list_naive

        cl = CellList(x,y,box,parallel=true)
        minimum_distances!(
            i -> mol_index(i,5),
            x_list,box,cl;
            parallel=true
        )
        @test x_list ≈ x_list_naive

    end

end

@testset "Disjoint: two lists" begin

    #
    # Disjoint sets of molecules
    #
    x = [ rand(SVector{3,Float64}) for _ in 1:100 ]
    y = [ rand(SVector{3,Float64}) for _ in 1:90 ]

    for box in [ Box([1,1,1],0.2),
                 Box([1.0 0.2 0.0
                      0.2 1.0 0.2
                      0.0 0.0 1.0], 0.2) ]

        x_list_naive, y_list_naive = MolecularMinimumDistances.naive_md(x,y,5,3,box)

        x_list, y_list = minimum_distances(x,y,5,3,box,parallel=false)
        @test x_list ≈ x_list_naive
        @test y_list ≈ y_list_naive

        x_list, y_list = minimum_distances(x,y,5,3,box,parallel=true)
        @test x_list ≈ x_list_naive
        @test y_list ≈ y_list_naive

        cl = CellList(x,y,box,parallel=false)
        x_list = init_list(x, i -> mol_index(i,5))
        y_list = init_list(y, i -> mol_index(i,3))
        minimum_distances!(
            i -> mol_index(i,5),
            i -> mol_index(i,3),
            x_list,y_list,box,cl;
            parallel=false
        )
        @test x_list ≈ x_list_naive
        @test y_list ≈ y_list_naive

        cl = CellList(x,y,box,parallel=true)
        minimum_distances!(
            i -> mol_index(i,5),
            i -> mol_index(i,3),
            x_list,y_list,box,cl;
            parallel=true
        )
        @test x_list ≈ x_list_naive
        @test y_list ≈ y_list_naive

    end

end # @testset


text:
using LinearAlgebra: logabsdet

@testset "injective" begin

    @inj function foo()
        x = @read(:x)
        y = @read(:y)
        z = @read(:z)
        a = x + y
        b = x - y
        c = x * y * z
        @write(a, :a)
        @write(b, :b)
        @write(c, :c)
    end

    x = 10.123
    y = -1.213
    z = 1.095

    input = DynamicAssignment()
    input[:x] = x
    input[:y] = y
    input[:z] = z

    J = [1.0 1.0 (y*z); 1.0 -1.0 (x*z); 0.0 0.0 (x*y)]
    (expected,) = logabsdet(J)

    (output, logdet) = apply(foo, (), input)
    @test isapprox(logdet, expected)

end

max_stars_repo_name: BoltzmannBrain/Gen

max_issues_repo_name: BoltzmannBrain/Gen

max_forks_repo_name: BoltzmannBrain/Gen


text:
"""Defines the SymOp, Crystal types"""

"""
    SymOp

Defines a symmetry operation belonging to a 3D space group, operating on fractional coordinates.
"""
struct SymOp
    R :: Mat3
    T :: Vec3
end

function Base.show(io::IO, ::MIME"text/plain", s::SymOp)
    for i in 1:3
        terms = []
        for (j, a) in enumerate(["x", "y", "z"])
            Rij = s.R[i,j]
            if Rij == 1
                push!(terms, "+"*a)
            elseif Rij == -1
                push!(terms, "-"*a)
            elseif Rij != 0
                push!(terms, format_number_or_fraction(Rij) * a)
            end
        end
        Ti = s.T[i]
        if Ti != 0
            push!(terms, format_number_or_fraction(Ti))
        end
        if isempty(terms)
            push!(terms, "0")
        end
        if terms[1][1] == '+'
            terms[1] = terms[1][2:end]
        end
        foreach(t -> print(io, t), terms)
        if i < 3
            print(io, ",")
        end
    end
end

function Base.isapprox(s1::SymOp, s2::SymOp; atol)
    return isapprox(s1.R, s2.R; atol) && isapprox(s1.T, s2.T; atol)
end

function transform(s::SymOp, r::Vec3)
    return s.R*r + s.T
end

# Wrap each coordinate of position r into the range [0,1). To account for finite
# precision, wrap 1-ϵ to -ϵ, where ϵ=symprec is a tolerance parameter.
function wrap_to_unit_cell(r::Vec3; symprec=1e-5)
    return @. mod(r+symprec, 1) - symprec
end

function is_same_position(x, y; symprec=1e-5)
    return norm(rem.(x-y, 1, RoundNearest)) < symprec
end

struct SiteSymmetry
    symbol       :: String
    multiplicity :: Int
    wyckoff      :: Char
end

"""
    Crystal

A type holding all geometry and symmetry information needed to represent
 a three-dimensional crystal.
"""
struct Crystal
    lat_vecs       :: Mat3                                 # Lattice vectors as columns
    positions      :: Vector{Vec3}                         # Positions in fractional coords
    types          :: Vector{String}                       # Types
    classes        :: Vector{Int}                          # Class indices
    sitesyms       :: Union{Nothing, Vector{SiteSymmetry}} # Optional site symmetries
    symops         :: Vector{SymOp}                        # Symmetry operations
    spacegroup     :: String                               # Description of space group
    symprec        :: Float64                              # Tolerance to imperfections in symmetry
end

"""
    nbasis(crystal::Crystal)

Number of basis positions (sublattices) in the unit cell.
"""
nbasis(cryst::Crystal) = length(cryst.positions)

"""
    cell_volume(crystal::Crystal)

Volume of the crystal unit cell.
"""
cell_volume(cryst::Crystal) = abs(det(cryst.lat_vecs))

"""
    equiv_sites(crystal::Crystal, b::Int)

Returns a list of all basis indices in the same symmetry equivalency class
as the provided index `b`.
"""
equiv_sites(cryst::Crystal, b::Int) = findall(==(cryst.classes[b]), cryst.classes)

cell_type(cryst::Crystal) = cell_type(cryst.lat_vecs)
lattice_vectors(cryst::Crystal) = cryst.lat_vecs
lattice_params(cryst::Crystal) = lattice_params(cryst.lat_vecs)

"""
    Crystal(lat_vecs, positions; types=nothing, symprec=1e-5)

Constructs a crystal from the complete list of atom positions `positions`,
representing fractions (between 0 and 1) of the lattice vectors `lat_vecs`. All
symmetry information is automatically inferred.
"""
function Crystal(lat_vecs, positions; types::Union{Nothing,Vector{String}}=nothing, symprec=1e-5)
    lat_vecs = convert(Mat3, lat_vecs)
    positions = [convert(Vec3, p) for p in positions]
    if isnothing(types)
        types = fill("", length(positions))
    end
    return crystal_from_inferred_symmetry(lat_vecs, positions, types; symprec)
end

"""
    Crystal(lat_vecs, positions, symbol::String; types=nothing, symprec=1e-5)

Builds a crystal by applying the symmetry operators for a given spacegroup
symbol.
"""
function Crystal(lat_vecs::Mat3, positions, symbol::String; types::Union{Nothing,Vector{String}}=nothing, setting=nothing, symprec=1e-5)
    positions = [convert(Vec3, p) for p in positions]
    if isnothing(types)
        types = fill("", length(positions))
    end
    crystal_from_symbol(lat_vecs, positions, types, symbol; setting, symprec)
end

"""
    Crystal(lat_vecs, positions, spacegroup_number; types=nothing, symprec=1e-5)

Builds a crystal by applying symmetry operators for a given international
spacegroup number.
"""
function Crystal(lat_vecs::Mat3, positions, spacegroup_number::Int; types::Union{Nothing,Vector{String}}=nothing, setting=nothing, symprec=1e-5)
    positions = [convert(Vec3, p) for p in positions]
    if isnothing(types)
        types = fill("", length(positions))
    end
    symbol = international_short_names[spacegroup_number]
    crystal_from_symbol(lat_vecs, positions, types, symbol; setting, symprec)
end

const n_hall_numbers = 530
const n_space_groups = 230

const international_short_names = begin
    ret = fill("", n_space_groups)
    for n in 1:n_hall_numbers
        sgt = Spglib.get_spacegroup_type(n)
        ret[sgt.number] = sgt.international_short
    end
    ret
end

function spacegroup_name(hall_number::Int)
    # String representation of space group
    sgt = Spglib.get_spacegroup_type(hall_number)
    return "HM symbol '$(sgt.international)' ($(sgt.number))"
end

function symops_from_spglib(rotations, translations)
    Rs = Mat3.(transpose.(eachslice(rotations, dims=3)))
    Ts = Vec3.(eachcol(translations))
    symops = SymOp.(Rs, Ts)
end

# Sort the sites according to class and fractional coordinates.
function sort_sites!(cryst::Crystal)
    function less_than(i, j)
        ci = cryst.classes[i]
        cj = cryst.classes[j]
        if ci != cj
            return ci < cj
        end
        ri = cryst.positions[i]
        rj = cryst.positions[j]
        for k = 3:-1:1
            if !isapprox(ri[k], rj[k], atol=cryst.symprec)
                return ri[k] < rj[k]
            end
        end
        error("Positions $i and $j cannot be distinguished.")
    end
    perm = sort(eachindex(cryst.positions), lt=less_than)
    cryst.positions .= cryst.positions[perm]
    cryst.classes .= cryst.classes[perm]
    cryst.types .= cryst.types[perm]
end

function crystal_from_inferred_symmetry(lat_vecs::Mat3, positions::Vector{Vec3}, types::Vector{String}; symprec=1e-5)
    for i in 1:length(positions)
        for j in i+1:length(positions)
            ri = positions[i]
            rj = positions[j]
            if is_same_position(ri, rj; symprec)
                error("Positions $ri and $rj are symmetry equivalent.")
            end
        end
    end

    positions = wrap_to_unit_cell.(positions; symprec)

    cell = Spglib.Cell(lat_vecs, hcat(positions...), types)
    d = Spglib.get_dataset(cell, symprec)
    classes = d.crystallographic_orbits
    # classes = d.equivalent_atoms
    symops = symops_from_spglib(d.rotations, d.translations)
    spacegroup = spacegroup_name(d.hall_number)

    # renumber class indices so that they go from 1:max_class
    classes = [findfirst(==(c), unique(classes)) for c in classes]
    @assert unique(classes) == 1:maximum(classes)

    # multiplicities for the equivalence classes
    multiplicities = map(classes) do c
        # indices that belong to class c
        idxs = findall(==(c), classes)
        # indices in the primitive cell that belong to class c
        prim_idxs = unique(d.mapping_to_primitive[idxs])
        # number of atoms in the standard cell that correspond to each primitive index for class c
        counts = [count(==(i), d.std_mapping_to_primitive) for i in prim_idxs]
        # sum over all equivalent atoms in the primitive cell
        sum(counts)
    end

    sitesyms = SiteSymmetry.(d.site_symmetry_symbols, multiplicities, d.wyckoffs)

    ret = Crystal(lat_vecs, positions, types, classes, sitesyms, symops, spacegroup, symprec)
    validate(ret)
    return ret
end

# Build Crystal using the space group denoted by a unique Hall number. The complete
# list is given at http://pmsl.planet.sci.kobe-u.ac.jp/~seto/?page_id=37&lang=en
function crystal_from_hall_number(lat_vecs::Mat3, positions::Vector{Vec3}, types::Vector{String}, hall_number::Int; symprec=1e-5)
    cell = cell_type(lat_vecs)
    hall_cell = cell_type(hall_number)
    allowed_cells = all_compatible_cells(hall_cell)
    @assert cell in allowed_cells "Hall number $hall_number requires a $hall_cell cell, but found $cell."

    if hall_cell == Sunny.monoclinic
        is_compatible = is_compatible_monoclinic_cell(lat_vecs, hall_number)
        @assert is_compatible "Lattice vectors define a monoclinic cell that is incompatible with Hall number $hall_number."
    end

    symops = symops_from_spglib(Spglib.get_symmetry_from_database(hall_number)...)
    spacegroup = spacegroup_name(hall_number)

    return crystal_from_symops(lat_vecs, positions, types, symops, spacegroup; symprec)
end

function crystal_from_symbol(lat_vecs::Mat3, positions::Vector{Vec3}, types::Vector{String}, symbol::String; setting=nothing, symprec=1e-5)
    hall_numbers = Int[]
    crysts = Crystal[]

    for hall_number in 1:n_hall_numbers
        sgt = Spglib.get_spacegroup_type(hall_number)

        if (replace(symbol, " "=>"") == sgt.international_short ||
            symbol in [sgt.hall_symbol, sgt.international, sgt.international_full])

            # Some Hall numbers may be incompatible with unit cell of provided
            # lattice vectors; skip them.
            is_compatible = true

            cell = cell_type(lat_vecs)
            hall_cell = cell_type(hall_number)
            allowed_cells = all_compatible_cells(hall_cell)

            # Special handling of trigonal space groups
            if Sunny.is_trigonal_symmetry(hall_number)
                # Trigonal symmetry must have either hexagonal or rhombohedral
                # cell, according to the Hall number.
                is_latvecs_valid = cell in [Sunny.rhombohedral, Sunny.hexagonal]
                if !is_latvecs_valid
                    error("Symbol $symbol requires a rhomobohedral or hexagonal cell, but found $cell.")
                end
                is_compatible = cell in allowed_cells
            else
                # For all other symmetry types, there is a unique cell for each Hall number
                if !(cell in allowed_cells)
                    error("Symbol $symbol requires a $hall_cell cell, but found $cell.")
                end
            end

            if hall_cell == Sunny.monoclinic
                is_compatible = is_compatible_monoclinic_cell(lat_vecs, hall_number)
            end

            if is_compatible
                cryst = crystal_from_hall_number(lat_vecs, positions, types, hall_number; symprec)
                push!(hall_numbers, hall_number)
                push!(crysts, cryst)
            end
        end
    end

    if length(crysts) == 0
        error("Could not find symbol '$symbol' in database.")
    elseif length(crysts) == 1
        return first(crysts)
    else
        if !isnothing(setting)
            i = findfirst(hall_numbers) do hall_number
                sgt = Spglib.get_spacegroup_type(hall_number)
                setting == sgt.choice
            end
            if isnothing(i)
                error("The symbol '$symbol' is ambiguous, and the specified setting '$setting' is not valid.")
            else
                return crysts[i]
            end
        end

        println("The symbol '$symbol' is ambiguous! Returning all crystals:")
        for (i, (hall_number, c)) in enumerate(zip(hall_numbers, crysts))
            sgt = Spglib.get_spacegroup_type(hall_number)
            hm_symbol = sgt.international
            choice = sgt.choice
            n_atoms = length(c.positions)
            i_str = @sprintf "%2d" i
            natoms_str = @sprintf "%2d" n_atoms
            println("   $i_str. \"$hm_symbol\", setting=\"$choice\", generates $natoms_str atoms")
        end
        println()
        println("Note: To disambiguate, you may pass a named parameter, setting=\"...\".")
        println()
        return crysts
    end
end

"Builds a crystal from an explicit set of symmetry operations and a minimal set of positions "
function crystal_from_symops(lat_vecs::Mat3, positions::Vector{Vec3}, types::Vector{String}, symops::Vector{SymOp}, spacegroup::String; symprec=1e-5)
    all_positions = Vec3[]
    all_types = String[]
    classes = Int[]

    for i = eachindex(positions)
        for s = symops
            x = wrap_to_unit_cell(transform(s, positions[i]); symprec)

            idx = findfirst(y -> is_same_position(x, y; symprec), all_positions)
            if isnothing(idx)
                push!(all_positions, x)
                push!(all_types, types[i])
                push!(classes, i)
            else
                j = classes[idx]
                if i != j
                    error("Reference positions $(positions[i]) and $(positions[j]) are symmetry equivalent.")
                end
            end
        end
    end

    # Atoms are sorted by contiguous equivalence classes: 1, 2, ..., n
    @assert unique(classes) == 1:maximum(classes)

    # Ask Spglib to infer the spacegroup for the given positions and types
    inferred = crystal_from_inferred_symmetry(lat_vecs, all_positions, all_types; symprec)

    # Compare the inferred symops to the provided ones
    is_subgroup = all(symops) do s
        any(inferred.symops) do s′
            isapprox(s, s′; atol=symprec)
        end
    end
    is_supergroup = all(inferred.symops) do s
        any(symops) do s′
            isapprox(s, s′; atol=symprec)
        end
    end

    if !is_subgroup
        println("WARNING: User provided symmetry operation could not be inferred by Spglib,")
        println("which likely indicates a non-conventional unit cell.")
    end

    # If the inferred symops match the provided ones, then we use the inferred
    # Crystal. Otherwise we must construct a new Crystal and do not have site
    # symmetry information.
    ret = if is_subgroup && is_supergroup
        inferred
    else
        Crystal(lat_vecs, all_positions, all_types, classes, nothing, symops, spacegroup, symprec)
    end
    sort_sites!(ret)
    validate(ret)
    return ret
end

"""
    subcrystal(cryst, types) :: Crystal

Filters sublattices of a `Crystal` by atom `types`, keeping the space group
unchanged.
"""
function subcrystal(cryst::Crystal, types::Vararg{String, N}) where {N}
    for s in types
        if !(s in cryst.types)
            error("types string '$s' is not present in crystal.")
        end
    end
    idxs = findall(in(types), cryst.types)
    classes = unique(cryst.classes[idxs])
    return subcrystal(cryst, classes...)
end

"""
    subcrystal(cryst, classes) :: Crystal

Filters sublattices of `Crystal` by equivalence `classes`, keeping the space
group unchanged.
"""
function subcrystal(cryst::Crystal, classes::Vararg{Int, N}) where {N}
    for c in classes
        if !(c in cryst.classes)
            error("Class '$c' is not present in crystal.")
        end
    end

    idxs = findall(in(classes), cryst.classes)
    new_positions = cryst.positions[idxs]
    new_types = cryst.types[idxs]
    new_classes = cryst.classes[idxs]
    new_sitesyms = cryst.sitesyms[idxs]

    if idxs != 1:maximum(idxs)
        println("WARNING, atoms are being renumbered.")
    end

    ret = Crystal(cryst.lat_vecs, new_positions, new_types, new_classes, new_sitesyms,
                  cryst.symops, cryst.spacegroup, cryst.symprec)
    return ret
end

function Base.show(io::IO, ::MIME"text/plain", cryst::Crystal)
    printstyled(io, "Crystal\n"; bold=true, color=:underline)
    println(io, cryst.spacegroup)

    if is_standard_form(cryst.lat_vecs)
        (a, b, c, α, β, γ) = lattice_params(cryst.lat_vecs)
        @printf io "Lattice params a=%.4g, b=%.4g, c=%.4g, α=%.4g°, β=%.4g°, γ=%.4g°\n" a b c α β γ
    else
        println(io, "Lattice vectors:")
        for a in eachcol(cryst.lat_vecs)
            @printf io "   [%.4g %.4g %.4g]\n" a[1] a[2] a[3]
        end
    end

    @printf io "Volume %.4g\n" cell_volume(cryst)

    for c in unique(cryst.classes)
        i = findfirst(==(c), cryst.classes)
        print(io, "Class $c")
        if cryst.types[i] != ""
            print(io, ", Type '$(cryst.types[i])'")
        end
        if !isnothing(cryst.sitesyms)
            # TODO: simplify in Julia 1.7
            symbol = cryst.sitesyms[i].symbol
            multiplicity = cryst.sitesyms[i].multiplicity
            wyckoff = cryst.sitesyms[i].wyckoff
            print(io, ", Site sym '$symbol', Wyckoff $multiplicity$wyckoff")
        end
        println(io, ":")

        for i in findall(==(c), cryst.classes)
            r = cryst.positions[i]
            r = round.(r, digits=12) # very close to zero is probably exactly zero
            @printf io "   %d. [%.4g, %.4g, %.4g]\n" i r[1] r[2] r[3]
        end
    end
end

function validate(cryst::Crystal)
    # Atoms of the same class must have the same type
    for i in eachindex(cryst.positions)
        for j in eachindex(cryst.positions)
            if cryst.classes[i] == cryst.classes[j]
                @assert cryst.types[i] == cryst.types[j]
            end
        end
    end

    # Rotation matrices in global coordinates must be orthogonal
    for s in cryst.symops
        R = cryst.lat_vecs * s.R * inv(cryst.lat_vecs)
        # Due to possible imperfections in the lattice vectors, only require
        # that R is approximately orthogonal
        @assert norm(R*R' - I) < cryst.symprec "Lattice vectors and symmetry operations are incompatible."
    end

    # TODO: Check that space group is closed and that symops have inverse?
end

#= Definitions of common crystals =#

function cubic_crystal(; a=1.0)
    lat_vecs = lattice_vectors(a, a, a, 90, 90, 90)
    basis_vecs = [[0, 0, 0]]
    Crystal(lat_vecs, basis_vecs)
end

function fcc_crystal(; a=1.0)
    lat_vecs = lattice_vectors(a, a, a, 90, 90, 90)
    basis_vecs = [[0, 0, 0]/2,
                  [1, 1, 0]/2,
                  [1, 0, 1]/2,
                  [0, 1, 1]/2]
    cryst = Crystal(lat_vecs, basis_vecs)
    sort_sites!(cryst)
    cryst
end

function fcc_primitive_crystal(; a=1.0)
    lat_vecs = [1 1 0; 1 0 1; 0 1 1]' * a/2
    basis_vecs = [[0, 0, 0]]
    Crystal(lat_vecs, basis_vecs)
end

function bcc_crystal(; a=1.0)
    lat_vecs = lattice_vectors(a, a, a, 90, 90, 90)
    basis_vecs = [[0, 0, 0]/2,
                  [1, 1, 1]/2,]
    Crystal(lat_vecs, basis_vecs)
end

function bcc_primitive_crystal(; a=1.0)
    lat_vecs = [1 1 -1; 1 -1 1; -1 1 1]' * a/2
    basis_vecs = [[0, 0, 0]]
    Crystal(lat_vecs, basis_vecs)
end

function diamond_crystal(; a=1.0)
    lat_vecs = lattice_vectors(a, a, a, 90, 90, 90)
    basis_vecs = [
        [0, 0, 0]/4,
        [2, 2, 0]/4,
        [1, 1, 1]/4,
        [3, 3, 1]/4,
        [2, 0, 2]/4,
        [0, 2, 2]/4,
        [3, 1, 3]/4,
        [1, 3, 3]/4,
    ]
    cryst = Crystal(lat_vecs, basis_vecs)
    sort_sites!(cryst)
    cryst
end

function diamond_primitive_crystal(; a=1.0)
    lat_vecs = [1 1 0; 1 0 1; 0 1 1]' * a/2
    basis_vecs = [
        [0, 0, 0]/4,
        [1, 1, 1]/4,
    ]
    Crystal(lat_vecs, basis_vecs)
end

function pyrochlore_lattice(; a=1.0)
    lat_vecs = [1 1 0; 1 0 1; 0 1 1]' * a/2
    basis_vecs = [
        [5, 5, 5]/8,
        [1, 5, 5]/8,
        [5, 5, 1]/8,
        [5, 1, 5]/8
    ]
    cryst = Crystal(lat_vecs, basis_vecs)
    sort_sites!(cryst)
    cryst
end


text:
# This file is a part of BAT.jl, licensed under the MIT License (MIT).

using BAT
using Test

using ValueShapes, Distributions, ForwardDiff

@testset "test_generic_transforms" begin
    @testset "uniform_cdf" begin
        a = -1.0
        b = 3.0
        x = 2
        u = cdf(Uniform(a, b), x)

        @test @inferred(BAT.uniform_cdf(x, a, b)) ≈ u
        @test @inferred(BAT.uniform_invcdf(u, a, b)) ≈ x
        @test BAT.uniform_invcdf_ladj(u, a, b) ≈ log(abs(ForwardDiff.derivative(u -> BAT.uniform_invcdf(u, a, b), u)))

        @test @inferred(BAT.UniformCDFTrafo(a, b)) isa BAT.VariateTransform{Univariate,BAT.UnitSpace,BAT.MixedSpace}
        trafo = BAT.UniformCDFTrafo(a, b)
        @test @inferred(varshape(trafo)) == ScalarShape{Real}()
        @test BAT.target_space(trafo) == BAT.UnitSpace()
        @test BAT.source_space(trafo) == BAT.MixedSpace()

        @test @inferred(BAT.apply_vartrafo(trafo, x, 0)).v ≈ u
        @test BAT.apply_vartrafo(trafo, x, 0).ladj ≈ log(abs(ForwardDiff.derivative(x -> cdf(Uniform(a, b), x), x)))
        @test isnan(BAT.apply_vartrafo(trafo, x, NaN).ladj)
        @test @inferred(BAT.apply_vartrafo(inv(trafo), u, 0)).v ≈ x
        @test BAT.apply_vartrafo(inv(trafo), u, 0).ladj ≈ log(abs(ForwardDiff.derivative(u -> quantile(Uniform(a, b), u), u)))
        @test isnan(BAT.apply_vartrafo(inv(trafo), u, NaN).ladj)

        @test @inferred(inv(trafo)) isa BAT.VariateTransform{Univariate,BAT.MixedSpace,BAT.UnitSpace}
        @test @inferred(BAT.apply_vartrafo(inv(trafo) ∘ trafo, x, 0)).v ≈ x
        @test BAT.apply_vartrafo(inv(trafo) ∘ trafo, x, 0.79).ladj ≈ 0.79

        @test @inferred(BAT.apply_vartrafo(inv(trafo ∘ inv(trafo)), u, 0)).v ≈ u
        @test BAT.apply_vartrafo(inv(trafo ∘ inv(trafo)), u, 0.79).ladj ≈ 0.79
    end

    @testset "logistic_cdf" begin
        mu = 2.0
        theta = 3.0
        x = 4
        u = cdf(Logistic(mu, theta), x)

        @test @inferred(BAT.logistic_cdf(x, mu, theta)) ≈ u
        @test @inferred(BAT.logistic_invcdf(u, mu, theta)) ≈ x
        @test BAT.logistic_invcdf_ladj(u, mu, theta) ≈ log(abs(ForwardDiff.derivative(u -> BAT.logistic_invcdf(u, mu, theta), u)))

        @test @inferred(BAT.LogisticCDFTrafo(mu, theta)) isa BAT.VariateTransform{Univariate,BAT.UnitSpace,BAT.InfiniteSpace}
        trafo = BAT.LogisticCDFTrafo(mu, theta)
        @test @inferred(varshape(trafo)) == ScalarShape{Real}()
        @test BAT.target_space(trafo) == BAT.UnitSpace()
        @test BAT.source_space(trafo) == BAT.InfiniteSpace()

        @test @inferred(BAT.apply_vartrafo(trafo, x, 0)).v ≈ u
        @test BAT.apply_vartrafo(trafo, x, 0).ladj ≈ log(abs(ForwardDiff.derivative(x -> cdf(Logistic(mu, theta), x), x)))
        @test isnan(BAT.apply_vartrafo(trafo, x, NaN).ladj)
        @test @inferred(BAT.apply_vartrafo(inv(trafo), u, 0)).v ≈ x
        @test BAT.apply_vartrafo(inv(trafo), u, 0).ladj ≈ log(abs(ForwardDiff.derivative(u -> quantile(Logistic(mu, theta), u), u)))
        @test isnan(BAT.apply_vartrafo(inv(trafo), u, NaN).ladj)

        @test @inferred(inv(trafo)) isa BAT.VariateTransform{Univariate,BAT.InfiniteSpace,BAT.UnitSpace}
        @test @inferred(BAT.apply_vartrafo(inv(trafo) ∘ trafo, x, 0.79)).v ≈ x
        @test BAT.apply_vartrafo(inv(trafo) ∘ trafo, x, 0.79).ladj ≈ 0.79

        @test @inferred(BAT.apply_vartrafo(inv(trafo ∘ inv(trafo)), u, 0.79)).v ≈ u
        @test BAT.apply_vartrafo(inv(trafo ∘ inv(trafo)), u, 0.79).ladj ≈ 0.79
    end

    @testset "exponential_cdf" begin
        theta = 3.0
        x = 5
        u = cdf(Exponential(theta), x)

        @test @inferred(BAT.exponential_cdf(x, theta)) ≈ u
        @test @inferred(BAT.exponential_invcdf(u, theta)) ≈ x
        @test BAT.exponential_cdf_ladj(x, theta) ≈ log(abs(ForwardDiff.derivative(x -> BAT.exponential_cdf(x, theta), x)))

        @test @inferred(BAT.ExponentialCDFTrafo(theta)) isa BAT.VariateTransform{Univariate,BAT.UnitSpace,BAT.MixedSpace}
        trafo = BAT.ExponentialCDFTrafo(theta)
        @test @inferred(varshape(trafo)) == ScalarShape{Real}()
        @test BAT.target_space(trafo) == BAT.UnitSpace()
        @test BAT.source_space(trafo) == BAT.MixedSpace()

        @test @inferred(BAT.apply_vartrafo(trafo, x, 0)).v ≈ u
        @test BAT.apply_vartrafo(trafo, x, 0).ladj ≈ log(abs(ForwardDiff.derivative(x -> cdf(Exponential(theta), x), x)))
        @test isnan(BAT.apply_vartrafo(trafo, x, NaN).ladj)
        @test @inferred(BAT.apply_vartrafo(inv(trafo), u, 0)).v ≈ x
        @test BAT.apply_vartrafo(inv(trafo), u, 0).ladj ≈ log(abs(ForwardDiff.derivative(u -> quantile(Exponential(theta), u), u)))
        @test isnan(BAT.apply_vartrafo(inv(trafo), u, NaN).ladj)

        @test @inferred(inv(trafo)) isa BAT.VariateTransform{Univariate,BAT.MixedSpace,BAT.UnitSpace}
        @test @inferred(BAT.apply_vartrafo(inv(trafo) ∘ trafo, x, 0.79)).v ≈ x
        @test BAT.apply_vartrafo(inv(trafo) ∘ trafo, x, 0.79).ladj ≈ 0.79

        @test @inferred(BAT.apply_vartrafo(inv(trafo ∘ inv(trafo)), u, 0.79)).v ≈ u
        @test BAT.apply_vartrafo(inv(trafo ∘ inv(trafo)), u, 0.79).ladj ≈ 0.79
    end

    @testset "scaled_log" begin
        theta = 3.0
        x = 4
        y = log(x / theta)

        @test @inferred(BAT.ScaledLogTrafo(theta)) isa BAT.VariateTransform{Univariate,BAT.InfiniteSpace,BAT.MixedSpace}
        trafo = BAT.ScaledLogTrafo(theta)
        @test @inferred(varshape(trafo)) == ScalarShape{Real}()
        @test BAT.target_space(trafo) == BAT.InfiniteSpace()
        @test BAT.source_space(trafo) == BAT.MixedSpace()

        @test @inferred(BAT.apply_vartrafo(trafo, x, 0)).v ≈ y
        @test BAT.apply_vartrafo(trafo, x, 0).ladj ≈ log(abs(ForwardDiff.derivative(trafo, x)))
        @test isnan(BAT.apply_vartrafo(trafo, x, NaN).ladj)
        @test @inferred(BAT.apply_vartrafo(inv(trafo), y, 0)).v ≈ x
        @test BAT.apply_vartrafo(inv(trafo), y, 0).ladj ≈ log(abs(ForwardDiff.derivative(inv(trafo), y)))
        @test isnan(BAT.apply_vartrafo(inv(trafo), y, NaN).ladj)
    end
end


text:
const USE_GPU = false  # Use GPU? If this is set false, then no GPU needs to be available
using ParallelStencil
using ParallelStencil.FiniteDifferences2D
@static if USE_GPU
    @init_parallel_stencil(CUDA, Float64, 2)
    macro pow(args...)  esc(:(CUDA.pow($(args...)))) end
    macro tanh(args...) esc(:(CUDA.tanh($(args...)))) end
else
    @init_parallel_stencil(Threads, Float64, 2)
    pow(x,y) = x^y
    macro pow(args...)  esc(:(pow($(args...)))) end
    macro tanh(args...) esc(:(Base.tanh($(args...)))) end
end
using Plots, Printf, Statistics, LinearAlgebra

@parallel function update_old!(Phi_o::Data.Array, ∇V_o::Data.Array, Phi::Data.Array, ∇V::Data.Array)
    @all(Phi_o) = @all(Phi)
    @all(∇V_o)  = @all(∇V)
    return
end

@parallel function compute_params_∇!(EtaC::Data.Array, K_muf::Data.Array, Rog::Data.Array, ∇V::Data.Array, ∇qD::Data.Array, Phi::Data.Array, Pf::Data.Array, Pt::Data.Array, Vx::Data.Array, Vy::Data.Array, qDx::Data.Array, qDy::Data.Array, μs::Data.Number, η2μs::Data.Number, R::Data.Number, λPe::Data.Number, k_μf0::Data.Number, ϕ0::Data.Number, nperm::Data.Number, θ_e::Data.Number, θ_k::Data.Number, ρfg::Data.Number, ρsg::Data.Number, ρgBG::Data.Number, dx::Data.Number, dy::Data.Number)
    @all(EtaC)  = (1.0-θ_e)*@all(EtaC)  + θ_e*( μs/@all(Phi)*η2μs*(1.0+0.5*(1.0/R-1.0)*(1.0+@tanh((@all(Pf)-@all(Pt))/λPe))) )
    @all(K_muf) = (1.0-θ_k)*@all(K_muf) + θ_k*( k_μf0*@pow((@all(Phi)/ϕ0), nperm) )
    @all(Rog)   = ρfg*@all(Phi) + ρsg*(1.0-@all(Phi)) - ρgBG
    @all(∇V)    = @d_xa(Vx)/dx  + @d_ya(Vy)/dy
    @all(∇qD)   = @d_xa(qDx)/dx + @d_ya(qDy)/dy
    return
end

@parallel function compute_RP!(dτPf::Data.Array, RPt::Data.Array, RPf::Data.Array, K_muf::Data.Array, ∇V::Data.Array, ∇qD::Data.Array, Pt::Data.Array, Pf::Data.Array, EtaC::Data.Array, Phi::Data.Array, Pfsc::Data.Number, Pfdmp::Data.Number, min_dxy2::Data.Number, dx::Data.Number, dy::Data.Number)
    @inn(dτPf) = min_dxy2/@maxloc(K_muf)/4.1/Pfsc
    @all(RPt)  =                 - @all(∇V)  - (@all(Pt) - @all(Pf))/(@all(EtaC)*(1.0-@all(Phi)))
    @all(RPf)  = @all(RPf)*Pfdmp - @all(∇qD) + (@all(Pt) - @all(Pf))/(@all(EtaC)*(1.0-@all(Phi)))
    return
end

@parallel function compute_P_τ!(Pt::Data.Array, Pf::Data.Array, τxx::Data.Array, τyy::Data.Array, σxy::Data.Array, RPt::Data.Array, RPf::Data.Array, dτPf::Data.Array, Vx::Data.Array, Vy::Data.Array, ∇V::Data.Array, dτPt::Data.Number, μs::Data.Number, β_n::Data.Number, dx::Data.Number, dy::Data.Number)
    @all(Pt)  = @all(Pt) +      dτPt *@all(RPt)
    @all(Pf)  = @all(Pf) + @all(dτPf)*@all(RPf)
    @all(τxx) = 2.0*μs*( @d_xa(Vx)/dx - 1.0/3.0*@all(∇V) - β_n*@all(RPt) )
    @all(τyy) = 2.0*μs*( @d_ya(Vy)/dy - 1.0/3.0*@all(∇V) - β_n*@all(RPt) )
    @all(σxy) = 2.0*μs*(0.5*( @d_yi(Vx)/dy + @d_xi(Vy)/dx ))
    return
end

@parallel function compute_res!(Rx::Data.Array, Ry::Data.Array, dVxdτ::Data.Array, dVydτ::Data.Array, τxx::Data.Array, τyy::Data.Array, σxy::Data.Array, Pt::Data.Array, Rog::Data.Array, dampX::Data.Number, dampY::Data.Number, dx::Data.Number, dy::Data.Number)
    @all(Rx)    = @d_xi(τxx)/dx + @d_ya(σxy)/dy - @d_xi(Pt)/dx
    @all(Ry)    = @d_yi(τyy)/dy + @d_xa(σxy)/dx - @d_yi(Pt)/dy - @av_yi(Rog)
    @all(dVxdτ) = dampX*@all(dVxdτ) + @all(Rx)
    @all(dVydτ) = dampY*@all(dVydτ) + @all(Ry)
    return
end

@parallel function compute_update!(Vx::Data.Array, Vy::Data.Array, qDx::Data.Array, qDy::Data.Array, Phi::Data.Array, dVxdτ::Data.Array, dVydτ::Data.Array, K_muf::Data.Array, Pf::Data.Array, Phi_o::Data.Array, ∇V::Data.Array, ∇V_o::Data.Array, dτV::Data.Number, ρfg::Data.Number, ρgBG::Data.Number, CN::Data.Number, dt::Data.Number, dx::Data.Number, dy::Data.Number)
    @inn(Vx)  =  @inn(Vx) + dτV*@all(dVxdτ)
    @inn(Vy)  =  @inn(Vy) + dτV*@all(dVydτ)
    @inn(qDx) = -@av_xi(K_muf)*(@d_xi(Pf)/dx)
    @inn(qDy) = -@av_yi(K_muf)*(@d_yi(Pf)/dy + (ρfg - ρgBG))
    @all(Phi) =  @all(Phi_o) + (1.0-@all(Phi))*(CN*@all(∇V_o) + (1.0-CN)*@all(∇V))*dt
    return
end

@parallel_indices (ix,iy) function bc_x!(A::Data.Array)
    A[1  , iy] = A[2    , iy]
    A[end, iy] = A[end-1, iy]
    return
end

@parallel_indices (ix,iy) function bc_y!(A::Data.Array)
    A[ix, 1  ] = A[ix, 2    ]
    A[ix, end] = A[ix, end-1]
    return
end

##################################################
@views function HydroMech2D()
    # Physics - scales
    ρfg      = 1.0             # fluid rho*g
    k_μf0    = 1.0             # reference permeability
    ηC0      = 1.0             # reference bulk viscosity
    # Physics - non-dimensional parameters
    η2μs     = 10.0            # bulk/shear viscosity ration
    R        = 500.0           # Compaction/decompaction strength ratio for bulk rheology
    nperm    = 3.0             # Carman-Kozeny exponent
    ϕ0       = 0.01            # reference porosity
    ra       = 2               # radius of initil porosity perturbation
    λ0       = 1.0             # standard deviation of initial porosity perturbation
    t_tot    = 0.02            # total time
    # Physics - dependent scales
    ρsg      = 2.0*ρfg         # solid rho*g
    lx       = 20.0            # domain size x
    ly       = ra*lx           # domain size y
    ϕA       = 2*ϕ0            # amplitude of initial porosity perturbation
    λPe      = 0.01            # effective pressure transition zone
    dt       = 1e-5            # physical time-step
    # Numerics
    CN       = 0.5             # Crank-Nicolson CN=0.5, Backward Euler CN=0.0
    res      = 128
    nx, ny   = res-1, ra*res-1 # numerical grid resolutions; should be a mulitple of 32-1 for optimal GPU perf
    ε        = 1e-5            # non-linear tolerance
    iterMax  = 5e3             # max nonlinear iterations
    nout     = 200             # error checking frequency
    β_n      = 1.0             # numerical compressibility
    Vdmp     = 5.0             # velocity damping for momentum equations
    Pfdmp    = 0.8             # fluid pressure damping for momentum equations
    Vsc      = 2.0             # reduction of PT steps for velocity
    Ptsc     = 2.0             # reduction of PT steps for total pressure
    Pfsc     = 4.0             # reduction of PT steps for fluid pressure
    θ_e      = 9e-1            # relaxation factor for non-linear viscosity
    θ_k      = 1e-1            # relaxation factor for non-linear permeability
    dt_red   = 1e-3            # reduction of physical timestep
    # Derived physics
    μs       = ηC0*ϕ0/η2μs                       # solid shear viscosity
    λ        = λ0*sqrt(k_μf0*ηC0)                # initial perturbation width
    ρgBG     = ρfg*ϕ0 + ρsg*(1.0-ϕ0)             # Background density
    # Derived numerics
    dx, dy   = lx/(nx-1), ly/(ny-1)              # grid step in x, y
    min_dxy2 = min(dx,dy)^2
    dτV      = min_dxy2/μs/(1.0+β_n)/4.1/Vsc     # PT time step for velocity
    dτPt     = 4.1*μs*(1.0+β_n)/max(nx,ny)/Ptsc
    dampX    = 1.0-Vdmp/nx
    dampY    = 1.0-Vdmp/ny
    # Array allocations
    Phi_o    = @zeros(nx  ,ny  )
    Pt       = @zeros(nx  ,ny  )
    Pf       = @zeros(nx  ,ny  )
    Rog      = @zeros(nx  ,ny  )
    ∇V       = @zeros(nx  ,ny  )
    ∇V_o     = @zeros(nx  ,ny  )
    ∇qD      = @zeros(nx  ,ny  )
    dτPf     = @zeros(nx  ,ny  )
    RPt      = @zeros(nx  ,ny  )
    RPf      = @zeros(nx  ,ny  )
    τxx      = @zeros(nx  ,ny  )
    τyy      = @zeros(nx  ,ny  )
    σxy      = @zeros(nx-1,ny-1)
    dVxdτ    = @zeros(nx-1,ny-2)
    dVydτ    = @zeros(nx-2,ny-1)
    Rx       = @zeros(nx-1,ny-2)
    Ry       = @zeros(nx-2,ny-1)
    Vx       = @zeros(nx+1,ny  )
    Vy       = @zeros(nx  ,ny+1)
    qDx      = @zeros(nx+1,ny  )
    # Initial conditions
    qDy      =   zeros(nx  ,ny+1)
    Phi      = ϕ0*ones(nx  ,ny  )
    Radc     =   zeros(nx  ,ny  )
    Radc    .= [(((ix-1)*dx-0.5*lx)/λ/4.0)^2 + (((iy-1)*dy-0.25*ly)/λ)^2 for ix=1:size(Radc,1), iy=1:size(Radc,2)]
    Phi[Radc.<1.0] .= Phi[Radc.<1.0] .+ ϕA
    EtaC     = μs./Phi.*η2μs
    K_muf    = k_μf0.*(Phi./ϕ0)
    ϕ0bc     = mean.(Phi[:,end])
    qDy[:,[1,end]] .= (ρsg.-ρfg).*(1.0.-ϕ0bc).*k_μf0.*(ϕ0bc./ϕ0).^nperm
    Phi      = Data.Array(Phi)
    EtaC     = Data.Array(EtaC)
    K_muf    = Data.Array(K_muf)
    qDy      = Data.Array(qDy)
    t        = 0.0
    it       = 1
    # Preparation of visualisation
    ENV["GKSwstype"]="nul"; if isdir("viz2D_out")==false mkdir("viz2D_out") end; loadpath = "./viz2D_out/"; anim = Animation(loadpath,String[])
    println("Animation directory: $(anim.dir)")
    X, Y, Yv = 0:dx:lx, 0:dy:ly, (-dy/2):dy:(ly+dy/2)
    # Time loop
    while t<t_tot
        @parallel update_old!(Phi_o, ∇V_o, Phi, ∇V)
        err=2*ε; iter=1; niter=0
        while err > ε && iter <= iterMax
            if (iter==11)  global wtime0 = Base.time()  end
            @parallel compute_params_∇!(EtaC, K_muf, Rog, ∇V, ∇qD, Phi, Pf, Pt, Vx, Vy, qDx, qDy, μs, η2μs, R, λPe, k_μf0, ϕ0, nperm, θ_e, θ_k, ρfg, ρsg, ρgBG, dx, dy)
            @parallel compute_RP!(dτPf, RPt, RPf, K_muf, ∇V, ∇qD, Pt, Pf, EtaC, Phi, Pfsc, Pfdmp, min_dxy2, dx, dy)
            @parallel (1:size(dτPf,1), 1:size(dτPf,2))  bc_x!(dτPf)
            @parallel (1:size(dτPf,1), 1:size(dτPf,2))  bc_y!(dτPf)
            @parallel compute_P_τ!(Pt, Pf, τxx, τyy, σxy, RPt, RPf, dτPf, Vx, Vy, ∇V, dτPt, μs, β_n, dx, dy)
            @parallel compute_res!(Rx, Ry, dVxdτ, dVydτ, τxx, τyy, σxy, Pt, Rog, dampX, dampY, dx, dy)
            @parallel compute_update!(Vx, Vy, qDx, qDy, Phi, dVxdτ, dVydτ, K_muf, Pf, Phi_o, ∇V, ∇V_o, dτV, ρfg, ρgBG, CN, dt, dx, dy)
            @parallel (1:size(Vx,1),  1:size(Vx,2))  bc_y!(Vx)
            @parallel (1:size(Vy,1),  1:size(Vy,2))  bc_x!(Vy)
            @parallel (1:size(qDx,1), 1:size(qDx,2)) bc_y!(qDx)
            @parallel (1:size(qDy,1), 1:size(qDy,2)) bc_x!(qDy)
            if mod(iter,nout)==0
                global norm_Ry, norm_RPf
                norm_Ry = norm(Ry)/length(Ry); norm_RPf = norm(RPf)/length(RPf); err = max(norm_Ry, norm_RPf)
                # @printf("iter = %d, err = %1.3e [norm_Ry=%1.3e, norm_RPf=%1.3e] \n", iter, err, norm_Ry, norm_RPf)
            end
            iter+=1; niter+=1
        end
        # Performance
        wtime    = Base.time()-wtime0
        A_eff    = (8*2)/1e9*nx*ny*sizeof(Data.Number)  # Effective main memory access per iteration [GB] (Lower bound of required memory access: Te has to be read and written: 2 whole-array memaccess; Ci has to be read: : 1 whole-array memaccess)
        wtime_it = wtime/(niter-10)                     # Execution time per iteration [s]
        T_eff    = A_eff/wtime_it                       # Effective memory throughput [GB/s]
        @printf("it = %d, time = %1.3e sec (@ T_eff = %1.2f GB/s) \n", it, wtime, round(T_eff, sigdigits=2))
        # Visualisation
        default(size=(500,700))
        if mod(it,5)==0
            p1 = heatmap(X, Y,  Array(Phi)'  , aspect_ratio=1, xlims=(X[1],X[end]), ylims=(Y[1],Y[end]), c=:viridis, title="porosity")
            p2 = heatmap(X, Y,  Array(Pt-Pf)', aspect_ratio=1, xlims=(X[1],X[end]), ylims=(Y[1],Y[end]), c=:viridis, title="effective pressure")
            p3 = heatmap(X, Yv, Array(qDy)'  , aspect_ratio=1, xlims=(X[1],X[end]), ylims=(Yv[1],Yv[end]), c=:viridis, title="vertical Darcy flux")
            p4 = heatmap(X, Yv, Array(Vy)'   , aspect_ratio=1, xlims=(X[1],X[end]), ylims=(Yv[1],Yv[end]), c=:viridis, title="vertical velocity")
            display(plot(p1, p2, p3, p4)); frame(anim)
        end
        # Time
        dt = dt_red/(1e-10+maximum(abs.(∇V)))
        t  = t + dt
        it+=1
    end
    gif(anim, "HydroMech2D.gif", fps = 15)
    return
end

HydroMech2D()


text:
# This is an extended version of exact_time_inference.jl. It combines it with
# Optim + ParameterHandling + Zygote to learn the kernel parameters.
# Each of these other packages know nothing about TemporalGPs, they're just general-purpose
# packages which play nicely with TemporalGPs (and AbstractGPs).

using AbstractGPs
using TemporalGPs

# Load up the separable kernel from TemporalGPs.
using TemporalGPs: RegularSpacing

# Load standard packages from the Julia ecosystem
using Optim # Standard optimisation algorithms.
using ParameterHandling # Helper functionality for dealing with model parameters.
using Zygote # Algorithmic Differentiation

# Declare model parameters using `ParameterHandling.jl` types.
# var_kernel is the variance of the kernel, λ the inverse length scale, and var_noise the
# variance of the observation noise. Note that they're all constrained to be positive.
flat_initial_params, unflatten = ParameterHandling.flatten((
    var_kernel = positive(0.6),
    λ = positive(0.1),
    var_noise = positive(2.0),
));

# Construct a function to unpack flattened parameters and pull out the raw values.
unpack = ParameterHandling.value ∘ unflatten;
params = unpack(flat_initial_params);

function build_gp(params)
    k = params.var_kernel * Matern52Kernel() ∘ ScaleTransform(params.λ)
    return to_sde(GP(k), SArrayStorage(Float64))
end

# Specify a collection of inputs. Must be increasing.
T = 1_000_000;
x = RegularSpacing(0.0, 1e-4, T);

# Generate some noisy synthetic data from the GP.
y = rand(f(x, params.var_noise));

# Specify an objective function for Optim to minimise in terms of x and y.
# We choose the usual negative log marginal likelihood (NLML).
function objective(params)
    f = build_gp(params)
    return -logpdf(f(x, params.var_noise), y)
end

# Optimise using Optim. Zygote takes a little while to compile.
training_results = Optim.optimize(
    objective ∘ unpack,
    θ -> only(Zygote.gradient(objective ∘ unpack, θ)),
    flat_initial_params + randn(3), # Perturb the parameters to make learning non-trivial
    BFGS(
        alphaguess = Optim.LineSearches.InitialStatic(scaled=true),
        linesearch = Optim.LineSearches.BackTracking(),
    ),
    Optim.Options(show_trace = true);
    inplace=false,
);

# Extracting the final values of the parameters. Should be moderately close to truth.
final_params = unpack(training_results.minimizer);

# Construct the posterior as per usual.
f_final = build_gp(final_params)
f_post = posterior(f_final(x, final_params.var_noise), y);

# Specify some locations at which to make predictions.
T_pr = 1_200_000;
x_pr = RegularSpacing(0.0, 1e-4, T_pr);

# Compute the exact posterior marginals at `x_pr`.
f_post_marginals = marginals(f_post(x_pr));
m_post_marginals = mean.(f_post_marginals);
σ_post_marginals = std.(f_post_marginals);

# Generate a few posterior samples. Not fantastically-well optimised at present.
f_post_samples = [rand(f_post(x_pr)) for _ in 1:5];

# Visualise the posterior. The if block is just to prevent it running in CI.
if get(ENV, "TESTING", "FALSE") == "FALSE"
    using Plots
    plt = plot();
    scatter!(plt, x, y; label="", markersize=0.1, alpha=0.1);
    plot!(plt, f_post(x_pr); ribbon_scale=3.0, label="");
    plot!(x_pr, f_post_samples; color=:red, label="");
    savefig(plt, "posterior.png");
end


text:
using QuantumComputing

# Quantum circuit of Deutsch-Jozsa algorithm
#
#      +---+  +---+  +---+  +---+
# |0⟩ -+ X +--+ H +--+   +--+ H +--
#      +---+  +---+  |   |  +---+
#                    | U |
#             +---+  |   |
# |0⟩ --------+ H +--+   +---------
#             +---+  +---+
#

psi = Ket(0,0,0,0,0)
n = length(psi)

Uf = reduce(*,[Hadamard(n,i) for i=1:n])

G1 = Not(n,5)
G2 = reduce(*,[Hadamard(n,i) for i=1:n])
G3 = reduce(*,[Hadamard(n,i) for i=1:n-1])

psi = G3*Uf*G2*G1*psi

println(psi)
println(extract(psi,5))


text:
# This file is a part of TypeDBClient.  License is MIT: https://github.com/Humans-of-Julia/TypeDBClient.jl/blob/main/LICENSE

struct EntityType <: AbstractEntityType
    label::Label
    is_root::Bool
end


text:
using Reexport

@reexport module CrossMappings

using Statistics
using Distributions
using StatsBase

include("utils.jl")
include("ccm.jl")
include("pairwise_asymmetric_inference.jl")

end # module


text:
for op in [:filepicker, :opendialog, :savedialog, :datepicker, :colorpicker, :timepicker, :spinbox,
           :autocomplete, :input, :dropdown, :checkbox, :toggle, :togglecontent,
           :textbox, :textarea, :button, :slider, :rangeslider, :rangepicker, :entry,
           :radiobuttons, :checkboxes, :toggles, :togglebuttons, :tabs, :tabulator, :accordion,
           :wdglabel, :latex, :alert, :highlight, :notifications, :mask, :tooltip!, :confirm]
    @eval begin
        function $op(args...; kwargs...)
            length(args) > 0 && args[1] isa AbstractBackend &&
                error("Function " * string($op) * " was about to overflow: check the signature")
            $op(get_backend(), args...; kwargs...)
        end

        widget(::Val{$(Expr(:quote, op))}, args...; kwargs...) = $op(args...; kwargs...)
    end
end

function defaultlayout(::AbstractBackend)
    ui -> div(values(components(ui))..., observe(ui))
end

input(::Type{<:Bool}, args...; kwargs...) = toggle(args...; kwargs...)
input(::Type{<:AbstractString}, args...; kwargs...) = textbox(args...; kwargs...)
input(::Type{<:Dates.Date}, args...; kwargs...) = datepicker(args...; kwargs...)
input(::Type{<:Dates.Time}, args...; kwargs...) = timepicker(args...; kwargs...)
input(::Type{<:Colorant}, args...; kwargs...) = colorpicker(args...; kwargs...)
function input(::Type{T}, args...; value=Observable{Union{Nothing, T}}(nothing), kwargs...) where {T<:Real}
    spinbox(args...; value=value, kwargs...)
end

"""
`widget(args...; kwargs...)`

Automatically convert Julia types into appropriate widgets. `kwargs` are passed to the
more specific widget function.

## Examples

```julia
map(display, [
    widget(1:10),                 # Slider
    widget(false),                # Checkbox
    widget("text"),               # Textbox
    widget(1.1),                  # Spinbox
    widget([:on, :off]),          # Toggle Buttons
    widget(Dict("π" => float(π), "τ" => 2π)),
    widget(colorant"red"),        # Color picker
    widget(Dates.today()),        # Date picker
    widget(Dates.Time()),         # Time picker
    ]);
```
"""
function widget end

widget(x; kwargs...) = x
widget(x::Observable; label = nothing) = widget(get_backend(), x; label = label)
widget(x::AbstractRange; kwargs...) = slider(x; kwargs...)
widget(x::AbstractVector; kwargs...) = togglebuttons(x; kwargs...)
widget(x::AbstractVector{<:Real}; kwargs...) = slider(x; kwargs...)
widget(x::AbstractVector{Bool}; kwargs...) = togglebuttons(x; kwargs...)
widget(x::Tuple; kwargs...) = togglebuttons(collect(x); kwargs...)
widget(x::AbstractDict; kwargs...) = togglebuttons(x; kwargs...)

function widget(x::T; kwargs...) where T<:Union{Bool, AbstractString, Real, Dates.Date, Dates.Time, Colorant}
    input(T; value=x, kwargs...)
end


text:
mutable struct TimeStamps # seconds
    ts0::Float64 # model construction start
    ts1::Float64 # ts1 ≤ _init!
    ts2::Float64 # _init! ≤ ts2 ≤ @threads
    ts3::Float64 # @threads ≤ ts3 ≤ remote_start
    ts4::Float64 # remote_start ≤ ts4 ≤ main_run
    ts5::Float64 # main_run ≤ ts5 ≤ remote_stop
    ts6::Float64 # remote_stop ≤ ts6
end

function TimeStamps(ts::Float64 = zero(Float64))
    t = zero(Float64)
    return TimeStamps(ts, t, t, t, t, t, t)
end
TimeStamps(model::_Model) = TimeStamps(get_time_stamp(model))

add_time!(stamps, ts, ::Val{1}) = stamps.ts1 = ts
add_time!(stamps, ts, ::Val{2}) = stamps.ts2 = ts
add_time!(stamps, ts, ::Val{3}) = stamps.ts3 = ts
add_time!(stamps, ts, ::Val{4}) = stamps.ts4 = ts
add_time!(stamps, ts, ::Val{5}) = stamps.ts5 = ts
add_time!(stamps, ts, ::Val{6}) = stamps.ts6 = ts

add_time!(stamps, i) = add_time!(stamps, time(), Val(i))

get_time(stamps, ::Val{0}) = stamps.ts1
get_time(stamps, ::Val{1}) = stamps.ts1
get_time(stamps, ::Val{2}) = stamps.ts2
get_time(stamps, ::Val{3}) = stamps.ts3
get_time(stamps, ::Val{4}) = stamps.ts4
get_time(stamps, ::Val{5}) = stamps.ts5
get_time(stamps, ::Val{6}) = stamps.ts6

get_time(stamps, i) = get_time(stamps, Val(i))

function time_info(stamps)
    info = Dict([
        :model => get_time(stamps, 1) - get_time(stamps, 0),
        :init => get_time(stamps, 2) - get_time(stamps, 1),
        :threads_start => get_time(stamps, 3) - get_time(stamps, 2),
        :remote_start => get_time(stamps, 4) - get_time(stamps, 3),
        :local_run => get_time(stamps, 5) - get_time(stamps, 4),
        :remote_stop => get_time(stamps, 6) - get_time(stamps, 5),
        :total_run => get_time(stamps, 6) - get_time(stamps, 1),
        :model_and_run => get_time(stamps, 6) - get_time(stamps, 0),
    ])
    return info
end


text:
function test_reading_vcf()
    vcf = vcf_read(joinpath(dirname(@__FILE__),"..","testdata","vcf","floats.vcf"))
    if version(vcf) != v"4.1.0"
        return false
    end
    return true
end

function test_vcf_operation()
    v1 = vcf_read(joinpath(dirname(@__FILE__),"..","testdata","vcf","cmp-test-a.vcf"))
    v2 = vcf_read(joinpath(dirname(@__FILE__),"..","testdata","vcf","cmp-test-b.vcf"))
    minus = v1 - v2
    intersect = v1 * v2
    return length(minus) == 1 && length(intersect) == length(v1)-1
end

@test test_reading_vcf()
@test test_vcf_operation()


text:
import Base: convert
import SymEngine: convert, free_symbols

"""
   is_linearcombination(L::Basic)

Return whether the expression `L` is a linear combination of its symbols.

### Input

- `L` -- expression

### Output

`true` if `L` is a linear combination or false otherwise.

### Examples

```jldoctest
julia> using SpaceExParser: is_linearcombination

julia> is_linearcombination(:(2*x1 - 4))
true

julia> is_linearcombination(:(6.1 - 5.3*f - 0.1*g))
true

julia> is_linearcombination(:(2*x1^2))
false

julia> is_linearcombination(:(x1^2 - 4*x2 + x3 + 2))
false
```
"""
function is_linearcombination(L::Basic)
    all(isempty.(free_symbols.(diff.(L, free_symbols(L)))))
end

is_linearcombination(L::Expr) = is_linearcombination(convert(Basic, L))

"""
    is_halfspace(expr::Expr)

Return wheter the given expression corresponds to a halfspace.

### Input

- `expr` -- a symbolic expression

### Output

`true` if `expr` corresponds to a halfspace or `false` otherwise.

### Examples

```jldoctest
julia> using SpaceExParser: is_halfspace

julia> all(is_halfspace.([:(x1 <= 0), :(x1 < 0), :(x1 > 0), :(x1 >= 0)]))
true

julia> is_halfspace(:(2*x1 <= 4))
true

julia> is_halfspace(:(6.1 <= 5.3*f - 0.1*g))
true

julia> is_halfspace(:(2*x1^2 <= 4))
false

julia> is_halfspace(:(x1^2 > 4*x2 - x3))
false

julia> is_halfspace(:(x1 > 4*x2 - x3))
true
```
"""
function is_halfspace(expr::Expr)::Bool

    # check that there are three arguments
    # these are the comparison symbol, the left hand side and the right hand side
    if (length(expr.args) != 3) || !(expr.head == :call)
        return false
    end

    # check that this is an inequality
    if !(expr.args[1] in [:(<=), :(<), :(>=), :(>)])
        return false
    end

    # convert to symengine expressions
    lhs, rhs = convert(Basic, expr.args[2]), convert(Basic, expr.args[3])

    # check if the expression defines a halfspace
    return is_linearcombination(lhs) && is_linearcombination(rhs)
end

"""
    is_hyperplane(expr::Expr)

Return wheter the given expression corresponds to a hyperplane.

### Input

- `expr` -- a symbolic expression

### Output

`true` if `expr` corresponds to a halfspace or `false` otherwise.

### Examples

```jldoctest
julia> using SpaceExParser: is_hyperplane

julia> is_hyperplane(:(x1 = 0))
true

julia> is_hyperplane(:(2*x1 = 4))
true

julia> is_hyperplane(:(6.1 = 5.3*f - 0.1*g))
true

julia> is_hyperplane(:(2*x1^2 = 4))
false

julia> is_hyperplane(:(x1^2 = 4*x2 - x3))
false

julia> is_hyperplane(:(x1 = 4*x2 - x3))
true
```
"""
function is_hyperplane(expr::Expr)::Bool

    # check that there are three arguments
    # these are the comparison symbol, the left hand side and the right hand side
    if (length(expr.args) != 2) || !(expr.head == :(=))
        return false
    end

    # convert to symengine expressions
    lhs = convert(Basic, expr.args[1])

    if :args in fieldnames(typeof(expr.args[2]))
        # treats the 4 in :(2*x1 = 4)
        rhs = convert(Basic, expr.args[2].args[2])
    else
        rhs = convert(Basic, expr.args[2])
    end

    # check if the expression defines a hyperplane
    return is_linearcombination(lhs) && is_linearcombination(rhs)
end

"""
    convert(::Type{HalfSpace{N}}, expr::Expr; vars=nothing) where {N}

Return a `LazySet.HalfSpace` given a symbolic expression that represents a halfspace.

### Input

- `expr` -- a symbolic expression
- `vars` -- (optional, default: `nothing`): set of variables with respect to which
            the gradient is taken; if nothing, it takes the free symbols in the given expression

### Output

A `HalfSpace`, in the form `ax <= b`.

### Examples

```jldoctest convert_halfspace
julia> using LazySets: HalfSpace

julia> convert(HalfSpace, :(x1 <= -0.03))
HalfSpace{Float64, Vector{Float64}}([1.0], -0.03)

julia> convert(HalfSpace, :(x1 < -0.03))
HalfSpace{Float64, Vector{Float64}}([1.0], -0.03)

julia> convert(HalfSpace, :(x1 > -0.03))
HalfSpace{Float64, Vector{Float64}}([-1.0], 0.03)

julia> convert(HalfSpace, :(x1 >= -0.03))
HalfSpace{Float64, Vector{Float64}}([-1.0], 0.03)

julia> convert(HalfSpace, :(x1 + x2 <= 2*x4 + 6))
HalfSpace{Float64, Vector{Float64}}([1.0, 1.0, -2.0], 6.0)
```

You can also specify the set of "ambient" variables, even if not
all of them appear:

```jldoctest convert_halfspace
julia> using SymEngine: Basic

julia> convert(HalfSpace, :(x1 + x2 <= 2*x4 + 6), vars=Basic[:x1, :x2, :x3, :x4])
HalfSpace{Float64, Vector{Float64}}([1.0, 1.0, 0.0, -2.0], 6.0)
```
"""
function convert(::Type{HalfSpace{N}}, expr::Expr; vars::Vector{Basic}=Basic[]) where {N}

    @assert is_halfspace(expr) "the expression :(expr) does not correspond to a halfspace"

    # check sense of the inequality, assuming < or <= by default
    got_geq = expr.args[1] in [:(>=), :(>)]

    # get sides of the inequality
    lhs, rhs = convert(Basic, expr.args[2]), convert(Basic, expr.args[3])

    # a1 x1 + ... + an xn + K [cmp] 0 for cmp in <, <=, >, >=
    eq = lhs - rhs
    if isempty(vars)
        vars = free_symbols(eq)
    end
    K = subs(eq, [vi => zero(N) for vi in vars]...)
    a = convert(Basic, eq - K)

    # convert to numeric types
    K = convert(N, K)
    a = convert(Vector{N}, diff.(a, vars))

    if got_geq
        return HalfSpace(-a, K)
    else
        return HalfSpace(a, -K)
    end
end

# type-less default halfspace conversion
convert(::Type{HalfSpace}, expr::Expr; vars::Vector{Basic}=Basic[]) = convert(HalfSpace{Float64}, expr; vars=vars)

"""
    convert(::Type{Hyperplane{N}}, expr::Expr; vars=nothing) where {N}

Return a `LazySet.Hyperplane` given a symbolic expression that represents a hyperplane.

### Input

- `expr` -- a symbolic expression
- `vars` -- (optional, default: `nothing`): set of variables with respect to which
            the gradient is taken; if nothing, it takes the free symbols in the given expression

### Output

A `Hyperplane`, in the form `ax = b`.

### Examples

```jldoctest convert_hyperplane
julia> using LazySets: Hyperplane

julia> convert(Hyperplane, :(x1 = -0.03))
Hyperplane{Float64, Vector{Float64}}([1.0], -0.03)

julia> convert(Hyperplane, :(x1 + 0.03 = 0))
Hyperplane{Float64, Vector{Float64}}([1.0], -0.03)

julia> convert(Hyperplane, :(x1 + x2 = 2*x4 + 6))
Hyperplane{Float64, Vector{Float64}}([1.0, 1.0, -2.0], 6.0)
```

You can also specify the set of "ambient" variables in the hyperplane, even if not
all of them appear:

```jldoctest convert_hyperplane
julia> using SymEngine: Basic

julia> convert(Hyperplane, :(x1 + x2 = 2*x4 + 6), vars=Basic[:x1, :x2, :x3, :x4])
Hyperplane{Float64, Vector{Float64}}([1.0, 1.0, 0.0, -2.0], 6.0)
```
"""
function convert(::Type{Hyperplane{N}}, expr::Expr; vars::Vector{Basic}=Basic[]) where {N}

    @assert is_hyperplane(expr) "the expression :(expr) does not correspond to a Hyperplane"

    # get sides of the inequality
    lhs = convert(Basic, expr.args[1])

    # treats the 4 in :(2*x1 = 4)
    rhs = :args in fieldnames(typeof(expr.args[2])) ? convert(Basic, expr.args[2].args[2]) : convert(Basic, expr.args[2])

    # a1 x1 + ... + an xn + K = 0
    eq = lhs - rhs
    if isempty(vars)
        vars = free_symbols(eq)
    end
    K = subs(eq, [vi => zero(N) for vi in vars]...)
    a = convert(Basic, eq - K)

    # convert to numeric types
    K = convert(N, K)
    a = convert(Vector{N}, diff.(a, vars))

    return Hyperplane(a, -K)
end

# type-less default Hyperplane conversion
convert(::Type{Hyperplane}, expr::Expr; vars::Vector{Basic}=Basic[]) = convert(Hyperplane{Float64}, expr; vars=vars)

"""
    free_symbols(expr::Expr, set_type::Type{LazySet})

Return the free symbols in an expression that represents a given set type.

### Input

- `expr` -- symbolic expression

### Output

A list of symbols, in the form of SymEngine `Basic` objects.

### Examples

```jldoctest
julia> using SpaceExParser: free_symbols

julia> using LazySets: HalfSpace

julia> free_symbols(:(x1 <= -0.03), HalfSpace)
1-element Vector{SymEngine.Basic}:
 x1

julia> free_symbols(:(x1 + x2 <= 2*x4 + 6), HalfSpace)
3-element Vector{SymEngine.Basic}:
 x2
 x1
 x4
```
"""
function free_symbols(expr::Expr, set_type::Type{<:HalfSpace})
    # get sides of the inequality
    lhs, rhs = convert(Basic, expr.args[2]), convert(Basic, expr.args[3])
    return free_symbols(lhs - rhs)
end

function free_symbols(expr::Expr, set_type::Type{<:Hyperplane})
    # get sides of the inequality
    lhs = convert(Basic, expr.args[1])

    # treats the 4 in :(2*x1 = 4)
    rhs = :args in fieldnames(typeof(expr.args[2])) ? convert(Basic, expr.args[2].args[2]) : convert(Basic, expr.args[2])
    return free_symbols(lhs - rhs)
end

function free_symbols(expr::Expr)
    if is_hyperplane(expr)
        return free_symbols(expr, HyperPlane)
    elseif is_halfspace(expr)
        return free_symbols(expr, Halfspace)
    else
        error("the free symbols for the expression $(expr) is not implemented")
    end
end


text:
"""
    TTVFaster

Computes first order eccentricity transit timing variations (TTVs) with respect to the following initial conditions: the planet-star mass ratio [μ], the initial transit time (of the averaged orbit) [t0], the mean orbital period [Per], the eccentricity [e], and the longitude of periastron [barred ω].
"""

module TTVFaster

# include("ttv_nplanet.jl")
include("ttv_wrapper.jl")

export Planet_plane_hk, compute_ttv!
export Planet_plane
# export ttv_nplanet, ttv_wrapper, ttv_chisquare

end

max_stars_repo_name: bmlindor/TTVFaster

max_issues_repo_name: bmlindor/TTVFaster

max_forks_repo_name: bmlindor/TTVFaster


text:
module OceanGrids

using Unitful, Interpolations
using LinearAlgebra, SparseArrays

abstract type OceanGrid end

"""
    OceanRectilinearGrid

Ocean rectilinear grid.
Here I use the slightly abuse the term rectilinear,
because the grid supposedly represents layers around the (curved) globe.
Here "rectilinear" means that the grid is defined entirely by just its
`lat`, `lon`, `depth`, `δlat`, `δlon`, and `δdepth` vectors.
For example, the Ocean Circulation Inverse Model (OCIM) version 1
grid is rectilinear in that sense.
"""
struct OceanRectilinearGrid <: OceanGrid
    lat          # °
    lon          # °
    depth        # m
    δlat         # °
    δlon         # °
    δdepth       # m
    lat_3D       # °
    lon_3D       # °
    depth_3D     # m
    δy           # m
    δx_3D        # m
    δy_3D        # m
    δz_3D        # m
    volume_3D    # m³
    depth_top    # m
    depth_top_3D # m
    A_2D         # m²
    wet3D::BitArray
    nlon
    nlat
    ndepth
    nboxes
end

"""
    OceanCurvilinearGrid

Ocean curvilinear grid.
The grid is curvilinear if it can be accessed using Cartesian Indices.
I.e., it maps to a rectilinear grid.
This type of grid typically cannot be inferred from only
`lat`, `lon`, `depth`, `δlat`, `δlon`, and `δdepth` vectors.
Instead, the full 3-dimensional arrays
`lat_3D`, `lon_3D`, `D`epth_3D`, `δx_3D`, `δy_3D`, and `δz_3D`
are required.
For exaample, the displaced-pole grid of the POP model is curcvilinear.
"""
struct OceanCurvilinearGrid <: OceanGrid
    lat_3D       # °
    lon_3D       # °
    depth_3D     # m
    δx_3D        # m
    δy_3D        # m
    δz_3D        # m
    volume_3D    # m³
    depth_top    # m
    depth_top_3D # m
    A_2D         # m²
    wet3D::BitArray
    nlon
    nlat
    ndepth
    nboxes
end

const TU = AbstractArray{<:Quantity}
const T1 = AbstractArray{<:Number}

"""
    OceanGrid(elon::T, elat::U, edepth::V; R=6371.0u"km") where {T<:TU, U<:TU, V<:TU}

Returns an `OceanRectilinearGrid` with boxes whose edges are defined by the
`elon`, `elat`, and `edepth` vectors.
The globe radius can be changed with the keyword `R` (default value 6371 km)
The 3D array of wet boxes, `wet3D` is true everywhere by default.
"""
function OceanGrid(elon::TU, elat::TU, edepth::TU; R=6371.0u"km")
    unitful_data, unitless_data = generate_grid_data_no_wet3D(elon, elat, edepth, R)
    #
    nlat, nlon, ndepth = unitless_data[[2,1,3]]
    wet3D = trues(nlat, nlon, ndepth)
    return OceanRectilinearGrid(unitful_data..., wet3D, unitless_data...)
end

"""
    OceanGrid(elon::T, elat::U, edepth::V, wet3D::BitArray; R=6371.0u"km")

Same as `OceanGrid(elon, elat, edepth; R)` but with specified `wet3D`.
"""
function OceanGrid(elon::TU, elat::TU, edepth::TU, wet3D::BitArray; R=6371.0u"km")
    unitful_data, unitless_data = generate_grid_data_no_wet3D(elon, elat, edepth, R)
    return OceanRectilinearGrid(unitful_data..., wet3D, unitless_data...)
end

function generate_grid_data_no_wet3D(elon::TU, elat::TU, edepth::TU, R)
    R = R |> u"m"
    edepth = edepth .|> u"m"
    elat = elat .|> u"°"
    edon = elon .|> u"°"
    # δ objects
    nlon, nlat, ndepth = length(elon) - 1, length(elat) - 1, length(edepth) - 1
    nboxes = nlon * nlat * ndepth
    lon = edges_to_centers(elon)
    lat = edges_to_centers(elat)
    depth = edges_to_centers(edepth)
    δlon, δlat, δdepth = diff(elon), diff(elat), diff(edepth)
    # depth objects
    δz_3D = repeat(reshape(δdepth, (1,1,ndepth)), outer=(nlat,nlon,1))
    depth_3D = repeat(reshape(depth, (1,1,ndepth)), outer=(nlat,nlon,1))
    depth_top = cumsum(δdepth) - δdepth
    depth_top_3D = repeat(reshape(depth_top, (1,1,ndepth)), outer=(nlat,nlon,1))
    # lat objects
    lat_3D = repeat(reshape(lat, (nlat,1,1)), outer=(1,nlon,ndepth))
    δy = R * δlat ./ 360u"°"
    δy_3D = repeat(reshape(δy, (nlat,1,1)), outer=(1,nlon,ndepth))
    # lon objects
    lon_3D = repeat(reshape(lon, (1,nlon,1)), outer=(nlat,1,ndepth))
    # For δx_3D, first calculate the area and then infer the equivalent distance in meters
    A_2D = R^2 * abs.(sin.(elat[1:end-1]) - sin.(elat[2:end])) * ustrip.(δlon .|> u"rad")'
    A_3D = repeat(A_2D, outer=(1, 1, ndepth))
    δx_3D = A_3D ./ δy_3D
    # volume
    volume_3D = δx_3D .* δy_3D .* δz_3D
    unitful_data = (lat, lon, depth, δlat, δlon, δdepth, lat_3D, lon_3D, depth_3D,
                    δy, δx_3D, δy_3D, δz_3D, volume_3D, depth_top, depth_top_3D, A_2D)
    unitless_data = (nlon, nlat, ndepth, nboxes)
    return unitful_data, unitless_data
end

# Convert units in case not provided
function OceanGrid(elon::TU, elat::TU, edepth::T1; R=6371.0u"km")
    @warn "`edepth` provided without units — assuming meters"
    return OceanGrid(elon, elat, edepth * u"m"; R=R)
end
function OceanGrid(elon::T1, elat::TU, edepth::TU; R=6371.0u"km")
    @warn "`elon` provided without units — assuming degrees"
    return OceanGrid(elon * u"°", elat, edepth; R=R)
end
function OceanGrid(elon::TU, elat::T1, edepth::TU; R=6371.0u"km")
    @warn "`elat` provided without units — assuming degrees"
    return OceanGrid(elon, elat * u"°", edepth; R=R)
end
function OceanGrid(elon::T1, elat::T1, edepth::TU; R=6371.0u"km")
    @warn "`elat` and `elon` provided without units — assuming degrees"
    return OceanGrid(elon * u"°", elat * u"°", edepth; R=R)
end
function OceanGrid(elon::T1, elat::T1, edepth::T1; R=6371.0u"km")
    @warn "`elat`, `elon`, and `edepth` provided without units — assuming degrees and meters"
    return OceanGrid(elon * u"°", elat * u"°", edepth * u"m"; R=R)
end
function OceanGrid(elon::TU, elat::T1, edepth::T1; R=6371.0u"km")
    @warn "`elat` and `edepth` provided without units — assuming and meters"
    return OceanGrid(elon, elat * u"°", edepth * u"m"; R=R)
end
function OceanGrid(elon::T1, elat::TU, edepth::T1; R=6371.0u"km")
    @warn "`elon` and `edepth` provided without units — assuming degrees and meters"
    return OceanGrid(elon * u"°", elat, edepth * u"m"; R=R)
end

edges_to_centers(x::Vector) = 0.5 * (x[1:end-1] + x[2:end])
edges_to_centers(x::AbstractRange) = x[1:end-1] .+ 0.5step(x)

"""
    OceanGrid(nlat::Int, nlon::Int, ndepth::Int)

Returns a regularly spaced `OceanRectilinearGrid` with size `nlat`, `nlon`, and `ndepth`.
"""
function OceanGrid(nlat::Int, nlon::Int, ndepth::Int)
    elat = range(-90,90,length=nlat+1) * u"°"
    elon = range(0,360,length=nlon+1) * u"°"
    edepth = range(0,3682,length=ndepth+1) * u"m"
    return OceanGrid(elon, elat, edepth)
end

function Base.show(io::IO, g::OceanGrid)
    println("OceanGrid of size $(g.nlat)×$(g.nlon)×$(g.ndepth) (lat×lon×depth)")
end

"""
    OceanGridBox

Ocean grid box.
Each grid can be looped over, where the `OceanGridBox` are the elements of the grid.
This useful to investigate specific boxes in the grid.
"""
struct OceanGridBox
    I
    lat          # °
    lon          # °
    depth        # m
    δlat         # °
    δlon         # °
    δdepth       # m
    δx           # m
    δy           # m
    δz           # m
    volume       # m³
    depth_top    # m
    A            # m²
    iswet::Bool
end

"""
    box(g::OceanGrid, i, j, k)

Accesses the individual box of `g::OceanGrid` at index `(i,j,k)`.
Each grid can be looped over, where the `OceanGridBox` are the elements of the grid.
This useful to investigate specific boxes in the grid.
"""
function box(g::OceanGrid, i, j, k)
    return OceanGridBox(
                        CartesianIndex(i,j,k),
                        g.lat[i],
                        g.lon[j],
                        g.depth[k],
                        g.δlat[i],
                        g.δlon[j],
                        g.δdepth[k],
                        g.δx_3D[i,j,k],
                        g.δy_3D[i,j,k],
                        g.δz_3D[i,j,k],
                        g.volume_3D[i,j,k],
                        g.depth_top[k],
                        g.A_2D[i,j],
                        g.wet3D[i,j,k]
                       )
end

"""
    box(g::OceanGrid, I)

Accesses the individual box of `g::OceanGrid` at index `I`.
Each grid can be looped over, where the `OceanGridBox` are the elements of the grid.
This useful to investigate specific boxes in the grid.
"""
function box(g::OceanGrid, I)
    i,j,k = CartesianIndices((g.nlat,g.nlon,g.ndepth))[I].I
    return box(g::OceanGrid, i, j, k)
end

Base.getindex(g::OceanGrid, I) = box(g,I)

"""
    Base.size(g::OceanGrid)

Size of the grid.
"""
Base.size(g::OceanGrid) = g.nlat, g.nlon, g.ndepth

"""
    Base.length(g::OceanGrid)

Size of the grid.
"""
Base.length(g::OceanGrid) = g.nlat * g.nlon * g.ndepth

function Base.show(io::IO, b::OceanGridBox)
    println("$(wet_print(b.iswet)) OceanGridBox at $(b.I):")
    println("  location: $(round(b.lat,digits=1))N, $(round(b.lon,digits=1))E")
    println("  depth: $(round(b.depth,digits=1))")
    println("  size: $(round(b.δx |> u"km",digits=1)) × $(round(b.δy |> u"km",digits=1)) × $(round(b.δz,digits=1)) (δx × δy × δz)")
end

wet_print(iswet) = iswet ? "Wet" : "Dry"

Base.round(q::Quantity; digits=0) = round(q |> ustrip, digits=digits) * unit(q)

area(b::OceanGridBox) = b.A |> u"km^2"
volume(b::OceanGridBox) = b.volume

Base.iterate(g::OceanGrid) = box(g::OceanGrid, 1), 1
function Base.iterate(g::OceanGrid, i)
    if i == g.nboxes
        return nothing
    else
        return box(g::OceanGrid, i+1), i+1
    end
end

Base.isequal(g₁::T, g₂::T) where {T<:OceanGrid} = prod([getfield(g₁,f) == getfield(g₂,f) for f in fieldnames(T)])
Base.:(==)(g₁::T, g₂::T) where {T<:OceanGrid} = isequal(g₁, g₂)

export OceanGrid, OceanCurvilinearGrid, OceanRectilinearGrid, box, OceanGridBox

"""
    iswet(grid)

Returns the 3D BitArray of wet boxes of the grid.
"""
iswet(g::T) where {T<:OceanGrid} = g.wet3D
export iswet

"""
    latvec(grid)

Returns the vector (wet points) of latitudes (units stripped).
"""
latvec(g::T) where {T<:OceanGrid} = ustrip.(g.lat_3D[iswet(g)])

"""
    lonvec(grid)

Returns the vector (wet points) of longitudes (units stripped).
"""
lonvec(g::T) where {T<:OceanGrid} = ustrip.(g.lon_3D[iswet(g)])

"""
    depthvec(grid)

Returns the vector (wet points) of depths (units stripped).
"""
depthvec(g::T) where {T<:OceanGrid} = ustrip.(g.depth_3D[iswet(g)])

"""
    latlondepthvecs(g)

Returns `latvec(g), lonvec(g), depthvec(g)`.
"""
latlondepthvecs(g::T) where {T<:OceanGrid} = latvec(g), lonvec(g), depthvec(g)

export latvec, lonvec, depthvec, latlondepthvecs

#===================================
1D Vector <-> 3D array conversions
===================================#
"""
    rearrange_into_3Darray(x, wet3D::BitArray)

Returns a 3D array of `x` rearranged to the `true` entries of `wet3D`.
Entries where `wet3D` is `false` are filled with `NaN`s.
"""
function rearrange_into_3Darray(x, wet3D::BitArray)
    iwet = findall(wet3D)
    x3d = fill(NaN, size(wet3D))
    x3d[iwet] .= x
    return x3d
end
"""
    rearrange_into_3Darray(x, grid)

Returns a 3D array of `x` rearranged to the wet boxes of the grid.
"""
rearrange_into_3Darray(x, grid) = rearrange_into_3Darray(x, grid.wet3D)

export rearrange_into_3Darray

#============
Interpolation
============#
"""
    interpolate(x, g)

Returns an "interpolation object" for fast interpolation.

It must be queried within the range of the grid in longitude...
"""
function Interpolations.interpolate(x, g::T) where {T<:OceanGrid}
    knots = (ustrip.(g.lat), periodic_longitude(ustrip.(g.lon)), ustrip.(g.depth))
    A = periodic_longitude(rearrange_into_3Darray(x, g))
    itp = interpolate(knots, A, Gridded(Constant()))
    return extrapolate(itp, (Flat(), Periodic(), Flat())) # extrapolate periodically
end
function Interpolations.interpolate(x, g::OceanGrid, lat, lon, depth; itp=interpolate(x,g))
    return itp(ustrip(lat), ustrip.(lon), ustrip.(depth))
end
function Interpolations.interpolate(x, g::OceanGrid, lats::Vector, lons::Vector, depths::Vector; itp=interpolate(x,g))
    return [itp(y,x,z) for (y,x,z) in zip(ustrip.(lats), ustrip.(lons), ustrip.(depths))]
end
function Interpolations.interpolate(x, g::OceanGrid, MD::NamedTuple; itp=interpolate(x,g))
    return interpolate(x, g, MD.lat, MD.lon, MD.depth; itp=itp)
end
function Interpolations.interpolate(x, g::OceanGrid, obs; itp=interpolate(x,g))
    return interpolate(x, g, obs.metadata; itp=itp)
end
Interpolations.interpolate(x, g::OceanGrid, ::Missing; itp=nothing) = missing
export interpolate

"""
    iswet(g, lat, lon, depth)
    iswet(g, lats, lons, depths)
    iswet(g, metadata)

Returns the indices of the provided locations that are "in" a wet box of `g`.

Technically, this uses a nearest neighbour interpolation algorithm, so the bounding boxes
will not work perfectly.
If AIBECS will solely rely on this nearest neighbour interpolation,
then it might be a good idea to replace Interpolations.jl with NearestNeighbors.jl.
"""
function iswet(g, args...; itp=interpolate(1:count(iswet(g)),g))
    J = interpolate(1:count(iswet(g)), g, args...; itp=itp)
    return findall((!isnan).(J))
end
iswet(g, ::Missing; itp=nothing) = missing

"""
    interpolationmatrix(g, lats, lons, depths)
    interpolationmatrix(g, metadata)

Returns the sparse matrix `M` such that `M*x` is the model vector `x` interpolated
onto the provided locations/metadata.

Technically, this requires a linear interpolation, in the sense that
`interpolation(x) = M * x` for some `M`, which seems to require a nearest neighbor
interpolation.
If AIBECS will solely rely on this nearest neighbour interpolation,
then it might be a good idea to replace Interpolations.jl with NearestNeighbors.jl.
Also, the few observations that lie in model land (instead of water) are discarded.

In the future, a more generic approach will use an sparse-aware AD for any type of
interpolation, and use an interpolation that somehow does not discard observations
made at locations not inside an ocean grid cell.
"""
function interpolationmatrix(g, args...; itp=interpolate(1:count(iswet(g)),g))
    J = interpolate(1:count(iswet(g)), g, args...; itp=itp)
    iwet = findall((!isnan).(J))
    J = Int.(J[iwet])
    I = collect(1:length(iwet))
    return sparse(I, J, true, length(I), count(iswet(g)))
end
interpolationmatrix(g, ::Missing; itp=nothing) = missing
export interpolationmatrix

#=
    lats = range(ustrip(grd.lat[1]), ustrip(grd.lat[end]), length=length(grd.lat))
    lons = range(ustrip(grd.lon[1]), ustrip(grd.lon[1])+360, length=length(grd.lon)+1)
    stp = Interpolations.scale(itp, lats, lons, 1:ndepths) # re-scale to the actual domain
    etp = extrapolate(stp, (Line(), Periodic(), Line())) # periodic longitude

    xs = range(0, 2π, length=11)
    ys = range(-π/2, π/2, length=6)
    A = f.(xs,ys')
    itp = interpolate(A, BSpline(Linear())) # interpolate linearly between the data points
    stp = scale(itp, xs, ys) # re-scale to the actual domain
    etp = extrapolate(stp, (Periodic(), Throw())) # extrapolate periodically
=#

periodic_longitude(x3D::Array{T,3}) where T = view(x3D,:,[1:size(x3D,2); 1],:)
periodic_longitude(v::Vector{T}) where T = [v; v[1]+360]

#===================================
Off-diaongal matrices
===================================#

"""
    buildIbelow(wet3D, iwet)

Builds the shifted-diagonal sparse matrix of the indices of below neighbours.

Ibelow[i,j] = 1 if the box represented by the linear index i
lies directly below the box represented by the linear index j.
"""
function buildIbelow(wet3D, iwet)
    nlat, nlon, ndepth = size(wet3D)
    n = nlon * nlat * (ndepth + 1)
    In = sparse(I, n, n)
    idx = zeros(Int, nlat, nlon, ndepth + 1)
    idx[:] = 1:n
    idx .= idx[:, :, [2:ndepth + 1; 1]]      # downward shift
    return In[idx[:], :][iwet, iwet]
end
"""
    buildIbelow(grid)

Builds the shifted-diagonal sparse matrix of the indices of below neighbours for `grid`.
See `buildIbelow(wet3D, iwet)`.
"""
buildIbelow(grid) = buildIbelow(grid.wet3D, indices_of_wet_boxes(grid))
export buildIbelow

"""
    buildIabove(wet3D, iwet)

Builds the shifted-diagonal sparse matrix of the indices of above neighbours.

Iabove[i,j] = 1 if the box represented by the linear index i
lies directly above the box represented by the linear index j.
"""
buildIabove(wet3D, iwet) = copy(transpose(buildIbelow(wet3D, iwet)))
"""
    buildIabove(grid)

Builds the shifted-diagonal sparse matrix of the indices of above neighbours for `grid`.
See `buildIabove(wet3D, iwet)`.
"""
buildIabove(grid) = copy(transpose(buildIbelow(grid)))
export buildIabove

#===================================
Functions returning useful arrays,
vectors, and constants
===================================#

"""
    array_of_volumes(grid)

Returns the 3D array of volumes of grid boxes.
"""
array_of_volumes(grid) = grid.volume_3D
export array_of_volumes

"""
    indices_of_wet_boxes(wet3D::BitArray)

Returns the vector of the indices of wet grid boxes.
"""
indices_of_wet_boxes(wet3D::BitArray) = findall(vec(wet3D))
"""
    indices_of_wet_boxes(grid)

Returns the vector of the indices of wet grid boxes.
"""
indices_of_wet_boxes(grid) = indices_of_wet_boxes(grid.wet3D)
export indices_of_wet_boxes

"""
    number_of_wet_boxes(wet3D::BitArray)

Returns the number of wet grid boxes.
"""
number_of_wet_boxes(wet3D::BitArray) = length(indices_of_wet_boxes(wet3D))
"""
    number_of_wet_boxes(grid)

Returns the number of wet grid boxes.
"""
number_of_wet_boxes(grid) = number_of_wet_boxes(grid.wet3D)
export number_of_wet_boxes

"""
    vector_of_volumes(grid)

Returns the vector of volumes of wet boxes.
"""
vector_of_volumes(grid) = array_of_volumes(grid)[indices_of_wet_boxes(grid.wet3D)]
export vector_of_volumes

"""
    vector_of_depths(grid)

Returns the vector of depths of the center of wet boxes.
"""
vector_of_depths(grid) = grid.depth_3D[indices_of_wet_boxes(grid.wet3D)]
export vector_of_depths

"""
    vector_of_depths(grid)

Returns the vector of depths of the top of wet boxes.
"""
vector_of_top_depths(grid) = grid.depth_top_3D[indices_of_wet_boxes(grid.wet3D)]
export vector_of_top_depths

#===================================
Functions return weighted norms
===================================#

"""
    weighted_norm²(x, w)

Returns the square of the weighted norm of `x` using weights `w`.
"""
weighted_norm²(x, w) = transpose(x) * Diagonal(w) * x

"""
    weighted_norm(x, v)

Returns the weighted norm of `x` using weights `w`.
"""
weighted_norm(x, w) = sqrt(weighted_norm²(x, w))

"""
    weighted_mean(x, w)

Returns the weighted mean of `x` using weights `w`.
"""
weighted_mean(x, w) = transpose(w) * x / sum(w)

"""
    vmean(x, w, I)

Returns the weighted mean of `x` using weights `w`, but only over indices `I`.
This is useful to get the observed mean.
(Because there are some grid boxes without observations.)
"""
weighted_mean(x, w, I) = transpose(w[I]) * x[I] / sum(w[I])
export weighted_norm²

end # module


text:
# log scale tests

scale_x_data = 10.0 .^ (-2:2)
scale_y_data = rand(5)

plot_1 = plot(scale_x_data, scale_y_data, title="plot 1")
yscale!(:log10)
xscale!(:log)

plot_2 = plot(
  scale_x_data, scale_y_data,
  xscale=:log10, yscale=:log
)

title!(plot_2, "plot 2")

test_is_same_object(plot_1, plot_2; except=["title"])

# lin scale tests

plot_A = plot(scale_x_data, scale_y_data, title="plot A")
xscale!(:log) ; yscale!(:linear)

plot_B = plot(
  scale_x_data, scale_y_data, xscale=:log10
)

title!("plot B")

test_is_same_object(plot_A, plot_B; except=["title"])
test_is_not_same_object(plot_A, plot_1; except=["title"])

# non-positive tests

plot(1:5, 1:5)
plot(1:5, 1:5, xscale=:log10, yscale=:log10)

plot(0:4, 1:5, xscale=:lin, yscale=:log10)
@test_throws AssertionError plot(0:4, 1:5, xscale=:log, yscale=:log10)

plot(1:5, -1:3, xscale=:log10, yscale=:identity)
@test_throws AssertionError plot(1:5, -1:3, xscale=:log10, yscale=:log10)