text:
struct Mesh
    nx::Int                     # Number of DOFs in x (Longitude)
    ny::Int                     # Number of DOFs in y (Latitude)
    dof::Int                    # Total number of DOFs: nx * ny
    h::Float64                  # Grid size [radians]: π/(ny-1)
    geom::Float64               # Geometrical parameter at poles
    csc2::Array{Float64,1}      # Metric term for the transformation to spherical coordinates: cosc²(θ), where θ is the colatitude angle
    cot::Array{Float64,1}       # Metric term for the transformation to spherical coordinates: cot(θ), where θ is the colatitude angle
    area::Array{Float64,1}      # Area of each cell on the surface of the sphere (*only* depends on latitude)
end

function Mesh(nx=128)
    ny = Int(nx/2+1) # 65
    dof = nx*ny

    h = pi / (ny-1) # Uniform grid size in radians

    # Trigonometric definitions for inner degrees of freedom
    csc2 = zeros(Float64,ny) # csc^2(theta), where theta is the colatitude angle (from the pole)
    cot  = zeros(Float64,ny) #   cot(theta), where theta is the colatitude angle (from the pole)
    for j=2:ny-1
        theta = h*(j-1)
        sintheta = sin(theta)
        csc2[j] = 1/sintheta^2
        cot[j] = cos(theta)/sintheta
    end

    area = zeros(Float64,ny) # Fractional area for the cells depending on latitude
    # Grid point fractional area for inner DOFs
    for j=2:ny-1
        area[j] =  sin(0.5*h)*sin(h*(j-1))/nx
    end

    # Fractional area for the poles
    area[1]  = 0.5*(1 - cos(0.5*h))
    area[ny] = area[1]

    geom = sin(0.5*h)/area[1]

    return Mesh(nx,ny,dof,h,geom,csc2,cot,area)
end

Base.size(mesh::Mesh) = (mesh.nx, mesh.ny)

function Base.show(io::IO, mesh::Mesh)
  nx, ny = size(mesh)
  print(io, "Mesh() with ", nx, "×", ny, " degrees of freedom")
end


text:
module RocketLoggerActorTest

using Test
using Rocket

@testset "LoggerActor" begin

    println("Testing: actor LoggerActor")

    @testset begin
        actor = logger(Int)

        @test !isempty(actor.name)
        @test actor.io === nothing
    end

    @testset begin
        buffer = IOBuffer()
        actor  = logger(Int, buffer)

        next!(actor, 1)

        output = String(take!(buffer))
        @test occursin("Data", output)
        @test occursin("1", output)

        next!(actor, 42)

        output = String(take!(buffer))
        @test occursin("Data", output)
        @test occursin("42", output)

        error!(actor, "err")

        output = String(take!(buffer))
        @test occursin("Error", output)
        @test occursin("err", output)

        complete!(actor)

        output = String(take!(buffer))
        @test occursin("Completed", output)

        @test_throws InconsistentSourceActorDataTypesError{Int64,String} next!(actor, "string")
    end

    @testset begin
        buffer = IOBuffer()
        actor  = logger(buffer, "CustomName")

        subscribe!(of(1), actor)
        output = String(take!(buffer))
        @test occursin("CustomName", output)
        @test occursin("Data", output)
        @test occursin("1", output)

        subscribe!(of(42), actor)
        output = String(take!(buffer))
        @test occursin("CustomName", output)
        @test occursin("Data", output)
        @test occursin("42", output)

        subscribe!(faulted("error"), actor)
        output = String(take!(buffer))
        @test occursin("CustomName", output)
        @test occursin("Error", output)
        @test occursin("error", output)

        subscribe!(completed(), actor)
        output = String(take!(buffer))
        @test occursin("CustomName", output)
        @test occursin("Completed", output)
    end

    @testset begin
        @test logger(Int) isa LoggerActor{Int, Nothing}
        @test logger()    isa Rocket.LoggerActorFactory{Nothing}
    end
end

end

ext: jl


text:
@testset "VariableTimeZone" begin
    @testset "equality" begin
        warsaw = first(compile("Europe/Warsaw", tzdata["europe"]))
        another_warsaw = first(compile("Europe/Warsaw", tzdata["europe"]))

        @test warsaw == warsaw
        @test warsaw === warsaw
        @test warsaw == another_warsaw
        @test warsaw !== another_warsaw
        @test isequal(warsaw, another_warsaw)
        @test hash(warsaw) == hash(another_warsaw)
    end

    @testset "broadcastable" begin
        warsaw = first(compile("Europe/Warsaw", tzdata["europe"]))

        # Validate that VariableTimeZone is treated as a scalar during broadcasting
        @test size(warsaw .== warsaw) == ()
    end

    @testset "links" begin
        # "Arctic/Longyearbyen" is a link to "Europe/Oslo"
        oslo = first(compile("Europe/Oslo", tzdata["europe"]))
        longyearbyen = first(compile("Arctic/Longyearbyen", tzdata["europe"]))

        @test oslo.name != longyearbyen.name
        @test oslo.transitions == longyearbyen.transitions
        @test oslo.cutoff == longyearbyen.cutoff

        @test oslo != longyearbyen
        @test oslo !== longyearbyen
        @test !isequal(oslo, longyearbyen)
        @test hash(oslo) != hash(longyearbyen)
    end

    @testset "cutoff differs" begin
        a = first(compile("Europe/Warsaw", tzdata["europe"]))
        b = VariableTimeZone(a.name, a.transitions, nothing)

        @test a.name == b.name
        @test a.transitions == b.transitions
        @test a.cutoff != b.cutoff

        @test a == b
        @test a !== b
        @test !isequal(a, b)
        @test hash(a) != hash(b)
    end
end


text:
# Internal helper functions for the show output of our structs.

const SHOWINDENTION = 4  # Number of spaces used for each indention level

function showfield(io::IO, key; indention=SHOWINDENTION)
    indentstr = " " ^ indention
    @printf io "%s%-20s %s " indentstr key (isempty(key) ? " " : ":")
end

function showfield(io::IO, key, value...; kwargs...)
    showfield(io, key; kwargs...)
    print(io, value...)
end

function showfieldln(io::IO, args...; kwargs...)
    showfield(io, args...; kwargs...)
    println(io)
end


text:
using DashBootstrapComponents

progress = dbc_progress([
    dbc_progress(value = 20, color = "success", bar = true),
    dbc_progress(value = 30, color = "warning", bar = true),
    dbc_progress(value = 20, color = "danger", bar = true),
],);

max_stars_repo_name: glsdown/dash-bootstrap-components

max_issues_repo_name: glsdown/dash-bootstrap-components

max_forks_repo_name: glsdown/dash-bootstrap-components


text:
function kfoldperm(N,k)
    #function from Dan Getz  https://stackoverflow.com/questions/37989159/how-to-divide-my-data-into-distincts-mini-batches-randomly-julia
    #because my brain was shutoff on a monday afternoon
    n,r = divrem(N,k)
    b = collect(1:n:N+1)
    for i in 1:length(b)
        b[i] += i > r ? r : i-1
    end
    p = randperm(N)
    return [p[r] for r in [b[i]:b[i+1]-1 for i=1:k]]
end

function setγ(γ::R,
    ρ::R,
    T::UniformScaling{R},
    Cwwf::PartialHermitianEigen{R,R},
    y::Array{R,1},
    wb::Array{R,1}) where {R<:Real}
    while true
        rhs = ρ*sqrt((y-wb)'*T*(y-wb))
        tmp = wbinv(T/γ, Cwwf[:vectors], Diagonal(convert(R,1.0)./Cwwf[:values]), y-wb)
        lhs = γ*sqrt(tmp'*tmp/T.λ)
        if lhs < rhs
            γ *= 2
        else
            break
        end
    end
    γ
end


text:
#=
# 1D advection-diffusion using FV with indexed equations:
#   flux_ij = a_i.n * u_ij + D_j * grad(u_ij).n
# Indices cover a range of advection speeds and diffusion rates.
=#

### If the Finch package has already been added, use this line #########
using Finch # Note: to add the package, first do: ]add "https://github.com/paralab/Finch.git"

### If not, use these four lines (working from the examples directory) ###
# if !@isdefined(Finch)
#     include("../Finch.jl");
#     using .Finch
# end
##########################################################################

init_finch("addiff1dindexed");

useLog("addiff1dindexedlog", level=3)

# Configuration setup
domain(1)
solverType(FV)
timeStepper(RK4)

# Mesh
n = 60 # number of elements
mesh(LINEMESH, elsperdim=n, bids=2)

# Set up indices and indexed values
nspeeds = 5
ndiffs = 4
speed = index("speed", range = [1,nspeeds])
diff = index("diff", range = [1,ndiffs])

speeds = zeros(nspeeds);
diffs = zeros(ndiffs);
bdrystr = [];
initstr = [];
astr = [];
dstr = [];

for j=1:ndiffs
    diffs[j] = 0.06 * (j-1)/(ndiffs-1);
    push!(dstr, string(diffs[j]));
end
for i=1:nspeeds
    speeds[i] = 0.5 + i/nspeeds;
    ts = 0.2/speeds[i];
    push!(astr, string(speeds[i]));
    for j=1:ndiffs
        push!(bdrystr, 0);
        push!(initstr, "x<=0.3&&x>0.1 ? 1 : 0");
    end
end

# Variables and BCs
u = variable("u", type=VAR_ARRAY, location=CELL, index = [speed, diff])
boundary(u, 1, FLUX, bdrystr)
boundary(u, 2, NO_BC)

# Time interval and initial condition
T = 0.3;
timeInterval(T)
initial(u, initstr)

# The flux and source terms of the conservation equation
coefficient("a", astr, VAR_ARRAY) # advection velocity
coefficient("d", dstr, VAR_ARRAY) # diffusion rate
flux(u, "upwind(a[speed],u[speed, diff]) - d[diff] * dot(grad(u[speed, diff]),normal())")

# Arrange the loops as desired. Outermost first.
assemblyLoops(u, ["elements", speed, diff])
# assemblyLoops(u, [speed, diff, "elements"])

solve(u)

finalize_finch()

##### Uncomment below to compare to plot solution

# # The exact solution with constant velocity v
# a = 1;
# exact = zeros(n);
# x = Finch.fv_info.cellCenters[:]
# for i=1:n
#     xt = x[i] - a*T;
#     if xt < 0.3 && xt > 0.1
#         exact[i] = 1
#     end
# end

# using Plots
# pyplot();
# ps = Array{Any,1}(undef, ndiffs);
# for j=1:ndiffs
#     global xs = x;
#     global labels = "exact a=1,d=0";
#     global vals = exact;
#     for i=1:nspeeds
#         global xs;
#         global vals;
#         global labels;
#         xs = [xs x];
#         vals = [vals u.values[i+nspeeds*(j-1),:]];
#         labels = [labels "a="*string(speeds[i])*", d="*string(diffs[j])];
#     end
#     ps[j] = plot(xs, vals, markershape=:circle, label=labels, reuse=false)
# end
# display(plot(ps[1], ps[2], ps[3], ps[4], layout=4))

max_stars_repo_name: paralab/Finch

max_issues_repo_name: paralab/Finch

max_forks_repo_name: paralab/Finch


text:
function reflection_coefficient(ω::T, dwave::DiscretePlaneWaveMode{T},
    source::PlaneSource{T,2,1,Acoustic{T,2}}, material::Material{2,Halfspace{T,2}};
     # medium::Acoustic{T,2}, specie::Specie{T,2}; θin::T = 0.0,
     scheme::Symbol = :trapezoidal) where T <: AbstractFloat

    k = ω / source.medium.c
    M = dwave.basis_order
    σ = k .* integration_scheme(dwave.x; scheme=scheme) # multiple by k to be the same as the non-dimensional version

    θin = transmission_angle(source,material)

    R = T(2)*number_density(material.species) / (cos(θin)*k^2)*sum(
        im^T(m)*exp(im*k*dwave.x[j]*cos(θin) - im*θin*m)*dwave.amplitudes[j,m+M+1,1]*σ[j]
    for m=-M:M, j in eachindex(dwave.x))

    return R
end

max_stars_repo_name: arturgower/EffectiveWaves


text:
using Dash, DashHtmlComponents, DashCoreComponents, DataFrames

df7 = DataFrame(a = [1, 2, 3],
               b = [4, 1, 4],
               c = ["x", "y", "z"])

app = dash()

app.layout = html_div() do
    dcc_dropdown(id="dropdown",
                 options = [
                    (label = i, value = i) for i in df7[:, "c"]
                 ],
                 value="x"
    ),
    html_div(id="output-2")
end

callback!(
    app,
    Output("output-2", "children"),
    Input("dropdown", "value"),
) do value
    # Here, `df7` is an example of a variable that is
    # 'outside the scope of this function'.
    # It is not safe to modify or reassign this variable
    # inside this callback.
    # do not do this, this is not safe!
    global df7 = filter!(row -> row.c == value, df7)
    return length(df7[1, :])
end

run_server(app, "0.0.0.0", debug=true)

max_stars_repo_name: joelostblom/dash-docs

max_issues_repo_name: joelostblom/dash-docs

max_forks_repo_name: joelostblom/dash-docs


text:
#=
Visualization of the gravities of different bodies around the solar system inspired by Dr. James O'Donoghue visualization here: https://twitter.com/physicsJ/status/1414233142861262855

Planetary data from NASA: https://nssdc.gsfc.nasa.gov/planetary/factsheet/ =#

using Javis # v0.7.1

diameters = Dict(
    "Mercury" => 4879,
    "Venus" => 12104,
    "Earth" => 12756,
    "Moon" => 3475,
    "Mars" => 6792,
    "Jupiter" => 142984,
    "Saturn" => 120536,
    "Uranus" => 51118,
    "Neptune" => 49528,
    "Pluto" => 2370,
)

struct Planet
    position::Int # Position from the sun
    name::String # Name of the planet
    gravity::Real # planets gravity in m/s^2
    color::String # color of the planet
    radius::Real # whatever fake units make the viz pretty
end

# Automatically calculate radius for us. Calculation is totally arbritrary
Planet(pos, name, grav, color) = Planet(pos, name, grav, color, log2(diameters[name]) * 2)

planets = [
    Planet(1, "Mercury", 3.7, "snow4")
    Planet(2, "Venus", 8.9, "navajowhite")
    Planet(3, "Earth", 9.8, "lightskyblue")
    Planet(4, "Moon", 1.6, "gainsboro")
    Planet(5, "Mars", 3.7, "orangered")
    Planet(6, "Jupiter", 23.1, "olive")
    Planet(7, "Saturn", 9, "burlywood")
    Planet(8, "Uranus", 8.7, "cyan3")
    Planet(9, "Neptune", 11, "dodgerblue")
    Planet(10, "Pluto", 0.7, "rosybrown4")
]

framerate = 30 # frames / second
height = 1000 # meters | also the amount of pixels the balls will fall

frames = let
    # Get planet with slowest acceleration
    slowest = [p.gravity for p in planets] |> minimum

    # Kinematic equation to determine seconds to fall from height
    time = 2 * height / slowest |> sqrt

    # Calculate total frames and add a few seconds at the end to display final result
    ceil(Int, time * framerate) + framerate * 5
end

function ground(args...)
    background("black")
    sethue("white")
end

font_height = 25

# Give room for header and footer
top_padding = font_height
bottom_padding = font_height * 5 + 50 + maximum(ceil.(Int, [p.radius * 4 for p in planets]))
total_height = height + top_padding + bottom_padding

# Height that planets start at
start_height = total_height / 2 - top_padding * 4

# Make video square by using total_height for height and width
width = total_height
myvideo = Video(width, total_height)
Background(1:frames, ground)

# Determine spacing and make some functions to help space the planets and text
spacing = width / (length(planets) + 1)
x_calc(planet::Planet) = (planet.position) * spacing - width / 2 + spacing / 2
y_height(row) = total_height / 2 - font_height / 2 - font_height * row

# 1km and 0km lines have to be drawn first to be under everything
Object(
    1:frames,
    JLine(
        Point(-width / 2, -start_height),
        Point(width, -start_height),
        color = "darkgray",
        linewidth = 5,
    ),
)
Object(
    1:frames,
    JLine(
        Point(-width / 2, -start_height + height),
        Point(width, -start_height + height),
        color = "darkgray",
        linewidth = 5,
    ),
)

# Make a list of tuples for each planet object and the planets stats
#     to make it easier to iterate over each planet and its object.
planet_objects = [
    (
        p,
        Object(
            1:frames,
            JCircle(O, p.radius, color = p.color, action = :fill),
            Point(x_calc(p), -start_height),
        ),
    ) for p in planets
]

function gravity_force(p::Planet, args)

    # unpack args provided by `Action`
    video, obj, action, frame = args

    # Calculate current time and position
    time = frame / framerate
    position = 0.5 * p.gravity * time^2

    # Keep the planet from falling further than the height specified at the beginning
    y_position = min(position, height)

    # Update the position of the planet
    obj.change_keywords[:center] = Point(0, y_position)

    # Leave trail to give an idea of the acceleration after planet has finished falling
    if frame % framerate * 8 == 0
        Object(
            frame:frames,
            JCircle(O, p.radius / 5, color = p.color, action = :fill),
            Point(x_calc(p), y_position - start_height),
        )
    end

    # Update text for planets current state

    # Determine if planet is finished moving.
    if y_position == height
        # Set all the text for after the planet is done moving

        # Calculate final time and velocity
        t_final = sqrt(height / (0.5 * p.gravity))
        v_final = p.gravity * t_final

        if t_final < time
            Object(
                frame:frame,
                @JShape begin
                    fontsize(font_height)
                    sethue("springgreen4")
                    text(
                        string(round(t_final, digits = 1), "s"),
                        Point(x_calc(p), y_height(1)),
                        halign = :center,
                    )
                    text(
                        string(round(Int, v_final), "m/s"),
                        Point(x_calc(p), y_height(0)),
                        halign = :center,
                    )
                end
            )
        end
    else
        # update text while planet is falling
        Object(
            frame:frame,
            @JShape begin
                fontsize(font_height)
                text(
                    string(round(time, digits = 1), "s"),
                    Point(x_calc(p), y_height(1)),
                    halign = :center,
                )
                text(
                    string(round(p.gravity * time, digits = 2), "m/s"),
                    Point(x_calc(p), y_height(0)),
                    halign = :center,
                )
            end
        )
    end
end

for (p, obj) in planet_objects
    act!(obj, Action(1:frames, (args...) -> gravity_force(p, args)))

    # Set text that is static during entire animation for each planet

    # Set planet info
    Object(
        1:frames,
        @JShape begin
            fontsize(font_height)
            sethue(p.color)
            text(p.name, Point(x_calc(p), y_height(3)), halign = :center)
            text(
                string(p.gravity, "m/s^2"),
                Point(x_calc(p), y_height(2)),
                halign = :center,
            )
        end
    )

    # Place a static copy of each planet in its column
    Object(
        1:frames,
        JCircle(O, p.radius, color = p.color, action = :fill),
        Point(x_calc(p), y_height(4) - p.radius),
    )
end

# Set the legend text
Object(
    1:frames,
    @JShape begin
        fontsize(font_height)
        x_pt = -width / 2 + 10
        text("1km", Point(x_pt, -start_height - 10), halign = :left)
        text("0km", Point(x_pt, -start_height + height - 10), halign = :left)

        fontsize(font_height * 0.75)
        text("Planet:", Point(x_pt, y_height(3)), halign = :left)
        text("Acceleration:", Point(x_pt, y_height(2)), halign = :left)

        text("Time:", Point(x_pt, y_height(1)), halign = :left)
        text("Velocity:", Point(x_pt, y_height(0)), halign = :left)

        fontsize(font_height * 2)
        sethue("royalblue")
        text(
            "Ball Falling 1km on Bodies in the Solar System",
            Point(0, -total_height / 2 + font_height * 2.5),
            halign = :center,
        )
    end
)

render(myvideo; liveview = true, framerate = framerate)
# render(myvideo; pathname = "gravities.gif", framerate = framerate)


text:
#=
--------------------------------------------------------------------------------------------------------------------------------
Types and constructors
--------------------------------------------------------------------------------------------------------------------------------
=#

"""
    VMASettings(...)

Define an immutable structure used to initialise the estimation routine for VMA(r) models.

# Arguments
- `Y`: observed measurements (`nxT`)
- `n`: Number of series
- `T`: Number of observations
- `r`: Order of the moving average polynomial
- `nr`: n*r
- `m`: nr+n*1_{q=0}
- `λ`: overall shrinkage hyper-parameter for the elastic-net penalty
- `α`: weight associated to the LASSO component of the elastic-net penalty
- `β`: additional shrinkage for distant lags (p>1)
- `Γ`: Diagonal matrix used to input the hyperparameters in the estimation - see Pellegrino (2022) for details
- `ε`: Small number (default: 1e-4)
- `tol`: tolerance used to check convergence (default: 1e-4)
- `max_iter`: maximum number of iterations for the estimation algorithm (default: 1000)
- `prerun`: number of iterations prior the actual estimation algorithm (default: 2)
- `check_quantile`: check the quantile of the relative change of the parameters for convergence purposes (default: false)
- `verb`: Verbose output (default: true)
"""
struct VMASettings <: EstimSettings
    Y::Union{FloatMatrix, JMatrix{Float64}}
    n::Int64
    T::Int64
    r::Int64
    nr::Int64
    m::Int64
    λ::Float64
    α::Float64
    β::Float64
    Γ::DiagMatrix
    ε::Float64
    tol::Float64
    max_iter::Int64
    prerun::Int64
    check_quantile::Bool
    verb::Bool
end

# VMASettings constructor
function VMASettings(Y::Union{FloatMatrix, JMatrix{Float64}}, r::Int64, λ::Float64, α::Float64, β::Float64; ε::Float64=1e-4, tol::Float64=1e-4, max_iter::Int64=1000, prerun::Int64=2, check_quantile::Bool=false, verb::Bool=true)

    # Compute missing inputs
    n, T = size(Y);
    m = n*r+n;
    Γ = build_Γ(n, r, λ, β);

    # Return VMASettings
    return VMASettings(Y, n, T, r, n*r, m, λ, α, β, Γ, ε, tol, max_iter, prerun, check_quantile, verb);
end

#=
--------------------------------------------------------------------------------------------------------------------------------
Initialisation
--------------------------------------------------------------------------------------------------------------------------------
=#

"""
    check_model_bounds(estim::VMASettings)

Check whether the arguments for the estimation of a VMA(r) model are correct.
"""
function check_model_bounds(estim::VMASettings)
    check_bounds(estim.r, 0);
    check_bounds(estim.λ, 0);
    check_bounds(estim.α, 0, 1);
    check_bounds(estim.β, 1);
    check_bounds(estim.max_iter, 3);
    check_bounds(estim.max_iter, estim.prerun);
    check_bounds(estim.n, 2); # It supports only multivariate models (for now ...)
end

"""
    initialise(estim::VMASettings)

Initialise the ECM algorithm for VMA(r) models.

# Arguments
- `estim`: settings used for the estimation

# References
Pellegrino (2022)
"""
function initialise(estim::VMASettings)

    # Initialise C_star
    C_star_init = zeros(estim.n, estim.m)

    # Interpolate data
    data = interpolate_series(estim.Y, estim.n, estim.T);

    # VAR(sqrt(T)) data
    sqrt_T = floor(sqrt(estim.T)) |> Int64;
    Y, Y_lagged = lag(data, sqrt_T);

    # Corresponding coefficients
    Γ_sqrt_T = build_Γ(estim.n, sqrt_T, estim.λ, estim.β);
    VAR_coefficients = Y*Y_lagged'/Symmetric(Y_lagged*Y_lagged' + Γ_sqrt_T);
    enforce_causality_and_invertibility!(VAR_coefficients);

    # Residuals
    residuals = Y - VAR_coefficients*Y_lagged;

    # VMA(r) data
    Y_VMA, _ = lag(Y, estim.r);
    U, U_VMA = lag(residuals, estim.r);

    # Initialise VMA
    B_star_init = Y_VMA*U_VMA'/Symmetric(U_VMA*U_VMA' + estim.Γ);
    enforce_causality_and_invertibility!(B_star_init);
    Q_init = Symmetric((U*U')/size(U,2));

    sspace = KalmanSettings(estim.Y, sspace_representation(estim, B_star_init, C_star_init, Q_init)..., compute_loglik=false);

    # Coordinates
    coordinates_measurement_states = collect(1:estim.m);
    coordinates_transition_current = collect(1:estim.n);
    coordinates_transition_P0 = LinearIndices(sspace.P0)[:];
    coordinates_free_params_B = CartesianIndices(sspace.B)[:, 1+estim.n:estim.m][:];

    # Useful views
    B_star = @view sspace.B[:, coordinates_measurement_states];
    Q_view = @view sspace.DQD.data[1:sspace.Y.n, 1:sspace.Y.n];

    # Return output
    return sspace, B_star, nothing, Q_view, coordinates_measurement_states, coordinates_transition_current, nothing, nothing, coordinates_transition_P0, coordinates_free_params_B, nothing;
end

#=
--------------------------------------------------------------------------------------------------------------------------------
CM-step
--------------------------------------------------------------------------------------------------------------------------------
=#

"""
    cm_step!(estim::VMASettings, sspace::KalmanSettings, B_star::SubArray{Float64}, C_star::Nothing, Q_view::SubArray{Float64}, F_raw::FloatMatrix, G::Nothing, H::Nothing, M::FloatArray, N::Vector{SparseMatrixCSC{Float64, Int64}}, O::Array{VectorsArray{Float64},1}, coordinates_free_params_B::CoordinatesVector, coordinates_free_params_C::Nothing)

Run the CM-step for elastic-net VMA(r) models.
"""
function cm_step!(estim::VMASettings, sspace::KalmanSettings, B_star::SubArray{Float64}, C_star::Nothing, Q_view::SubArray{Float64}, F_raw::FloatMatrix, G::Nothing, H::Nothing, M::FloatArray, N::Vector{SparseMatrixCSC{Float64, Int64}}, O::Array{VectorsArray{Float64},1}, coordinates_free_params_B::CoordinatesVector, coordinates_free_params_C::Nothing)

    # Make sure F is symmetric
    F = Symmetric(F_raw)::SymMatrix;

    # Companion forms: initial settings
    B_companion_params = @view sspace.B[:, 1+estim.n:estim.m];
    B_companion_old = companion_form(B_companion_params, extended=false);

    # Loop over the free parameters
    for ij in coordinates_free_params_B

        # Compute remaining sufficient statistics
        sum_numerator, sum_denominator, i, j = cm_step_time_loop(sspace, B_star, ij, N, O);

        # Compute new value for sspace.B[ij]
        B_star[ij] = soft_thresholding(M[ij] - sum_numerator, 0.5*estim.α*estim.ε*estim.Γ[j-estim.n, j-estim.n]);
        B_star[ij] /= sum_denominator + (1-estim.α)*estim.ε*estim.Γ[j-estim.n, j-estim.n];
    end

    # Companion forms: final settings
    B_companion_new = companion_form(B_companion_params, extended=false);

    # Adjust coefficients (if needed)
    enforce_causality_and_invertibility!(B_companion_new, B_companion_old, B_companion_params);

    # CM-step for Q
    for ij in eachindex(F)
        Q_view[ij] = F[ij]/estim.T;
    end
end

#=
--------------------------------------------------------------------------------------------------------------------------------
Model validation
--------------------------------------------------------------------------------------------------------------------------------
=#

"""
    VMASettings(Y::Union{FloatMatrix, JMatrix{Float64}}, model_args::Tuple, model_kwargs::Nothing, p::Int64, λ::Float64, α::Float64, β::Float64)
    VMASettings(Y::Union{FloatMatrix, JMatrix{Float64}}, model_args::Tuple, model_kwargs::NamedTuple, p::Int64, λ::Float64, α::Float64, β::Float64)

Generate the relevant VMASettings structure by using the candidate hyperparameters `p`, `λ`, `α`, `β`.
"""
VMASettings(Y::Union{FloatMatrix, JMatrix{Float64}}, model_args::Tuple, model_kwargs::Nothing, p::Int64, λ::Float64, α::Float64, β::Float64) = VMASettings(Y, p, λ, α, β, verb=false);
VMASettings(Y::Union{FloatMatrix, JMatrix{Float64}}, model_args::Tuple, model_kwargs::NamedTuple, p::Int64, λ::Float64, α::Float64, β::Float64) = VMASettings(Y, p, λ, α, β; model_kwargs...);


text:
module FileTools

__precompile__(true)

export @f_str

abstract AbstractFile

macro f_str(filename)
  name,format = splitext(filename)
  support(filename, Val{symbol(format)})
end

include("jsonfile.jl")
include("txtfile.jl")

end # module


text:
function get_os_info()
    if Sys.iswindows()
        return "win-x64"
    elseif Sys.isapple()
        return "macos-universal"
    elseif Sys.islinux()
        return "linux-x64"
    else
        error("This OS is not supported.")
    end
end

function get_archive_file_ext()
    if Sys.iswindows()
        return ".zip"
    else
        return ".tar.gz"
    end
end

function get_lib_ext()
    if Sys.iswindows()
        return ".dll"
    elseif Sys.isapple()
        return ".dylib"
    elseif Sys.islinux()
        return ".so"
    end
end

function get_version()
    proj_toml_path = joinpath(@__DIR__, "..", "Project.toml")
    open(proj_toml_path, "r") do fp
        for line in eachline(fp)
            println(line)
            if startswith(line, "version = ")
                return string(strip(replace(line, "version = " => ""), '\"'))
            end
        end
        error("Version is not found in Project.toml.")
    end
end

function extract_zip(zipfile)
    lib_ext = get_lib_ext()
    run(`unzip $zipfile bin/"*"$lib_ext`)
    cp("bin", joinpath(@__DIR__, "..", "src", "bin"); force=true)
    rm("bin"; force=true, recursive=true)
    rm(zipfile; force=true)
end

function extract_targz(tarfile)
    run(`tar -xvf $tarfile`)
    cp("bin", joinpath(@__DIR__, "..", "src", "bin"); force=true)
    rm("bin"; force=true, recursive=true)
    rm(tarfile; force=true)
end

function replace_latest_binary(version::String)
    local base_url = "https://github.com/shinolab/autd3-library-software/releases/download/"
    module_path = joinpath(@__DIR__, "..", "src")
    ext = get_archive_file_ext()
    os_info = get_os_info()
    url = base_url * "/v" * version * "/autd3-" * "v" * version * "-" * os_info * ext
    tmp_archive_path = joinpath(module_path, "tmp" * ext)
    file = download(url, tmp_archive_path)
    if startswith(os_info, "win")
        extract_zip(tmp_archive_path)
    elseif startswith(os_info, "macos")
        extract_targz(tmp_archive_path)
    elseif startswith(os_info, "linux")
        extract_targz(tmp_archive_path)
    end
end

_version = get_version()
replace_latest_binary(_version)


text:
# Integer Functions

# TODO: vector types
const generic_integer_types = [Cchar, Cuchar, Cshort, Cushort, Cint, Cuint, Clong, Culong]

# generically typed

for gentype in generic_integer_types
@eval begin

@device_override Base.abs(x::$gentype) = @builtin_ccall("abs", $gentype, ($gentype,), x)
@device_function abs_diff(x::$gentype, y::$gentype) = @builtin_ccall("abs_diff", $gentype, ($gentype, $gentype), x, y)

@device_function add_sat(x::$gentype, y::$gentype) = @builtin_ccall("add_sat", $gentype, ($gentype, $gentype), x, y)
@device_function hadd(x::$gentype, y::$gentype) = @builtin_ccall("hadd", $gentype, ($gentype, $gentype), x, y)
@device_function rhadd(x::$gentype, y::$gentype) = @builtin_ccall("rhadd", $gentype, ($gentype, $gentype), x, y)

@device_override Base.clamp(x::$gentype, minval::$gentype, maxval::$gentype) = @builtin_ccall("clamp", $gentype, ($gentype, $gentype, $gentype), x, minval, maxval)

@device_function clz(x::$gentype) = @builtin_ccall("clz", $gentype, ($gentype,), x)
@device_function ctz(x::$gentype) = @builtin_ccall("ctz", $gentype, ($gentype,), x)

@device_function mad_hi(a::$gentype, b::$gentype, c::$gentype) = @builtin_ccall("mad_hi", $gentype, ($gentype, $gentype, $gentype), a, b, c)
@device_function mad_sat(a::$gentype, b::$gentype, c::$gentype) = @builtin_ccall("mad_sat", $gentype, ($gentype, $gentype, $gentype), a, b, c)

# XXX: these definitions introduce ambiguities
#@device_override Base.max(x::$gentype, y::$gentype) = @builtin_ccall("max", $gentype, ($gentype, $gentype), x, y)
#@device_override Base.min(x::$gentype, y::$gentype) = @builtin_ccall("min", $gentype, ($gentype, $gentype), x, y)

@device_function mul_hi(x::$gentype, y::$gentype) = @builtin_ccall("mul_hi", $gentype, ($gentype, $gentype), x, y)

@device_function rotate(v::$gentype, i::$gentype) = @builtin_ccall("rotate", $gentype, ($gentype, $gentype), v, i)

@device_function sub_sat(x::$gentype, y::$gentype) = @builtin_ccall("sub_sat", $gentype, ($gentype, $gentype), x, y)

@device_function popcount(x::$gentype) = @builtin_ccall("popcount", $gentype, ($gentype,), x)

@device_function mad24(x::$gentype, y::$gentype, z::$gentype) = @builtin_ccall("mad24", $gentype, ($gentype, $gentype, $gentype), x, y, z)
@device_function mul24(x::$gentype, y::$gentype) = @builtin_ccall("mul24", $gentype, ($gentype, $gentype), x, y)

end
end

# specifically typed

@device_function upsample(hi::Cchar, lo::Cuchar) = @builtin_ccall("upsample", Cshort, (Cchar, Cuchar), hi, lo)
upsample(hi::Cuchar, lo::Cuchar) = @builtin_ccall("upsample", Cushort, (Cuchar, Cuchar), hi, lo)
upsample(hi::Cshort, lo::Cushort) = @builtin_ccall("upsample", Cint, (Cshort, Cushort), hi, lo)
upsample(hi::Cushort, lo::Cushort) = @builtin_ccall("upsample", Cuint, (Cushort, Cushort), hi, lo)
upsample(hi::Cint, lo::Cuint) = @builtin_ccall("upsample", Clong, (Cint, Cuint), hi, lo)
upsample(hi::Cuint, lo::Cuint) = @builtin_ccall("upsample", Culong, (Cuint, Cuint), hi, lo)


text:
module DSP

using Interpolation
using Conv
using Distributions
using DistributedArrays
using DSP # from julia
using SparseArrays
using LinearAlgebra
using FFTW

function chirp!(x; tgrid=nothing,
	fmin=nothing, # minimum fraction of Nyquist sampling
	fmax=nothing, # maximum fraction of Nyquist
	)

	fs=inv(step(tgrid))
	freqmin=(fmin===nothing) ? 0.0 : fmin*fs*0.5
	freqmax=(fmax===nothing) ? fs*0.5 : fmax*fs*0.5

	k=(freqmax-freqmin)*inv(tgrid[end]-tgrid[1])

	for it in 1:length(tgrid)
		t=(it-1)*step(tgrid)
		x[it] = sin(2. * π *(freqmin*t+k/2.0*t*t))
	end

end
"""
Random source signal models for Param.
Generate a band-limited
random signal of a particular maximum time, `tmax`.
Sinc function in the frequency domain corresponds to
a box car function is applied in the time domain.

This method has some extra allocations.
"""
function freq_rand!(x;
	tgrid=nothing,
	fmin=nothing, # minimum fraction of Nyquist sampling
	fmax=nothing, # maximum fraction of Nyquist
	nevent_fracs::Array{Float64}=[0.0],
	rng=Normal(),
	verbose=false,
	)

	nt=size(x,1)
	(tgrid===nothing) && (tgrid=range(0.0, stop=(nt-1)*1.0, step=1.0))  # create dummy tgrid
	fgrid=DSP.rfftfreq(length(tgrid),inv(step(tgrid))) # corresponding fgrid
	(nt≠length(tgrid)) && error("dimension")

	tmax=abs(tgrid[nt]-tgrid[1])

	xfreq=complex.(zeros(length(fgrid))) # allocation in freq domain
	xx=zeros(length(tgrid)) # allocation in time domain

	fs=inv(step(tgrid))
	freqmin=(fmin===nothing) ? 0.0 : fmin*fs*0.5
	freqmax=(fmax===nothing) ? fs*0.5 : fmax*fs*0.5
	(freqmax ≤ freqmin) && error("fmin and fmax")

	Δf = inv(tmax) # resolution in the frequency domain
	(Δf ≤ step(fgrid)) ? error("desired resolution smaller than grid sampling") :
	(Δf ≥ (freqmax-freqmin)) ? error("need to increase tmax") :
	fvec = [f for f in freqmin:Δf:freqmax]
	ifvec = fill(0, size(fvec))

	verbose && println("number of frequencies added to signal:\t", size(fvec,1))
	verbose && println("interval between random variable:\t", Δf)
	verbose && println("minimum frequency added to signal\t",minimum(fvec))
	verbose && println("maximum frequency added to signal\t",maximum(fvec))
	for iff in eachindex(fvec)
		ifvec[iff] =  Interpolation.indminn(fgrid, fvec[iff])[1]
	end

	for iff in eachindex(fvec)
		# translated sinc function
		fshift=fgrid[ifvec[iff]]  # shift for translated sinc function
		X=rand(rng) * exp(im*rand(Uniform(-pi, pi))) # uniformly distributed phase and random amplitude
		Xe=complex(1.0) # intialization for events
		for tfrac in nevent_fracs
			Xe += exp(im*fshift*2*π*tfrac*tgrid[nt]) # phase translations for each event
		end
		X *= Xe # add events to X
		for ifff in 1:length(fgrid)
			α=tmax*(abs(fgrid[ifff] - fshift)) # argument of sinc function
			xfreq[ifff] += (X * sinc(α)) # go
		end
	end
	xx=irfft(xfreq, nt)
	normalize!(xx)
	copyto!(x, xx)
	return x

end

"""
Tapering is necessary to be able to input random signal into finite-difference code
Filtering tapering are applied only if the length of the time series is greater than 10
"""
function get_tapered_random_tmax_signal(tgrid;
					fmin=nothing,
					fmax=nothing,
					tmaxfrac::Float64=1.0,
					dist=Uniform(-2.0, 2.0),
					sparsep=1.0,
					taperperc=20.
					)
	filt_flag=(length(tgrid) > 5) && (!(fmin === nothing)) && (!(fmax===nothing))

	fs = inv(step(tgrid));
	if(filt_flag)
		designmethod = Butterworth(6);
		filtsource = Bandpass(fmin, fmax; fs=fs);
	end

	itind = argmin(abs.(tgrid.-abs(tmaxfrac)*tgrid[end]))
	if(tmaxfrac>0.0)
		its=1:itind
	elseif(tmaxfrac<0.0)
		its=itind:length(tgrid)
	end
	# 20% taper window
	twin = taper(ones(length(its)),taperperc)
	X = zeros(length(its))
	wavsrc = zeros(length(tgrid))
	if(filt_flag)
		X[:] = rand(dist, length(its)) .* twin
	else
		X[:] = rand(dist, length(its))
	end
	if(sparsep ≠ 1.0)
		Xs=sprandn(length(X), sparsep)
		X[findall(Xs.==0.0)].=0.0
	end
	# band limit
	(filt_flag) && (filt!(X, digitalfilter(filtsource, designmethod), X))

	(length(X) ≠ 1) && normalize!(X)
	wavsrc[its] = X
	return wavsrc
end

"""
* `x` : first dimension is time
* `p` : sparsity
* `rng` : Random Number Generator
* `btperc` : Taper Perc
* `etperc` : Taper Perc
* `bfrac` : zeros at the beginning of each coloumn
* `efrac` : zeros fraction at the end of each coloumn
"""
function tapered_rand!(x;p=1.0,rng=Normal(),btperc=0., etperc=0.,bfrac=0.0, efrac=0.0)
	x[:]=0.0
	nt=size(x,1)

	a=1+round(Int,bfrac*nt)
	b=nt-round(Int,efrac*nt)

	dd=size(x)[2:end]
	for i in CartesianRange(dd)
		xx=view(x,a:b,i)
		rand!(rng, xx)
		taper!(xx, bperc=btperc, eperc=etperc)
		nxx=size(xx,1)
		if(p ≠ 1.0)
			xs=sprandn(nxx, p)
			xx[findn(xs.==0.0)]=0.0
		end
	end
	scale!(x, inv(vecnorm(x)))
	return x
end

"""
randomly shift and add x to itself
using circshift; not memory efficient
"""
function cshift_add_rand!(x::AbstractVector; ntimes=1, tminfrac=0.0, tmaxfrac=1.0, rand_amp_flag=true)
	nt=size(x,1)
	a=1+round(Int,tminfrac*nt)
	b=round(Int,tmaxfrac*nt)

	for it in 1:ntimes
		if(rand_amp_flag)
			amp=randn()
		else
			amp=1.0
		end
		# a random shift
		its=rand(DiscreteUniform(a,b))
		xx = circshift(x,(its,)) # pre-allocate for performance?
		for i in eachindex(x)
			x[i] += xx[i] * amp
		end
	end
end

"""
Construct Toy Green's functions
Decaying peaks, control number of events, and their positions, depending on bfrac and efrac.
* `afrac` : control amplitude of peaks
"""
function toy_green!(x;nevents=1,afrac=[inv(2^(ie-1)) for ie in 1:nevents],bfrac=0.0,efrac=0.0)
	nt=size(x,1)
	nr=size(x,2)
	x[:]=0.0
	a=1+round(Int,bfrac*nt)
	b=nt-round(Int,efrac*nt)
	(b-a<nevents+2) && error("not enough samples")

	itvec=zeros(Int64,nevents+1)
	for ir in 1:nr
		itvec[:]=0
		for ie in 1:nevents
			if(ie==1)
				# first event, direct arrival
				itvec[ie]=rand(DiscreteUniform(a,b-nevents-1))
			else
				# all other events after first event
				itvec[ie]=rand(DiscreteUniform(itvec[1]+1,b-1))
			end
			sgn=1.
			# randomly choose sign for all other events
			if(ie≠1)
				sgn=(rand(Bernoulli())==1) ? 1. : -1.
			end
			x[itvec[ie],ir]=sgn*afrac[ie]
		end
	end
	return x
end

function findfreq(
		  x::Array{Float64, ND},
		  tgrid;
		  attrib::Symbol=:peak,
		  threshold::Float64=-50.
		  ) where {ND}

	cx=rfft(x,[1]);
	fgrid=DSP.rfftfreq(length(tgrid),inv(step(tgrid))) # corresponding fgrid

	ax = (abs.(cx).^2); # power spectrum in dB

	if(maximum(ax) == 0.0)
		warn("x is zero"); return 0.0
	else
		ax /= maximum(ax);
		ax = 10. .* log10.(ax)
	end

	if(attrib == :max)
		iii=findlast(ax .>= threshold)
	elseif(attrib == :min)
		iii=findfirst(ax .>= threshold)
	elseif(attrib == :peak)
		iii=argmax(ax)
	end
	ii=CartesianIndices(size(ax))[iii][1]
	return fgrid[ii]

end

"""
Cosine taper a N-dimensional array along its first dimension.

# Arguments
* `x::Array{Float64,N}` :
* `perc::Float64` : taper percentage
"""
function taper(x::AbstractArray,perc::Float64)
	xout=copy(x);
	taper!(xout,perc)
	return xout
end

function taper!(x, perc=0.0; bperc=perc,eperc=perc)

	nt=size(x,1)
	nttb=min(round(Int,nt*bperc*0.01), nt)
	ntte=min(round(Int,nt*eperc*0.01), nt)
	kb=inv(2.0*round(Int,nt*bperc*0.01)-1)*pi
	ke=inv(2.0*round(Int,nt*eperc*0.01)-1)*pi
	for i in CartesianIndices(size(x))
		if(1≤i[1]≤nttb)
			x[i] *= sin((i[1]-1)*kb)
		end
		if(nt-ntte+1≤i[1]≤nt)
			x[i] *= sin((-i[1]+nt)*ke)
		end
	end
end

include("Filters.jl")

end # module


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

using Markdown
using InteractiveUtils

# ╔═╡ 96c26596-9fef-11ec-351a-11ad73639906
begin

	const PROJECT_ROOT = dirname(@__DIR__)
	const TEST_DIR = joinpath(PROJECT_ROOT, "test")

	import Pkg

	new_env = mktempdir()
	cp(TEST_DIR, new_env; force=true)

	copy!(LOAD_PATH, ["@", PROJECT_ROOT])
	Pkg.activate(new_env)
	Pkg.instantiate()

	# development dependencies
	import Downloads
	import JSON
	using BenchmarkTools
	using PlutoTest

	using HypertextLiteral
	import AbstractPlutoDingetjes

	import Pluto
	using PlutoHooks
end

# ╔═╡ 72204279-d847-4bc6-9365-83a0299c8a38
s = Pluto.ServerSession(;options=Pluto.Configuration.from_flat_kwargs(;
	launch_browser=false,
	disable_writing_notebook_files=true,
	auto_reload_from_file=true,
	lazy_workspace_creation=true,
))

# ╔═╡ 1ee7d783-88c1-418f-ae18-d1a25049b467
const nb_path = joinpath(PROJECT_ROOT, "src", "notebook.jl")

# ╔═╡ a0deb59c-9fb1-4e3c-9761-690e3c42af76
begin
	for n in values(s.notebooks)
		Pluto.SessionActions.shutdown(s, n; keep_in_session=false)
	end
	const nb_original = Pluto.SessionActions.open(
		s, nb_path;
		run_async=false,
	)
end

# ╔═╡ 06c8d69c-16bd-4dc1-97ef-f3410c30e756

# ╔═╡ c7f642e6-0660-4bc0-948e-77fadd128b59

# ╔═╡ d6fe479d-0e08-4e1d-af57-deb1e78ce300

# ╔═╡ bf9e7649-6489-4464-adaf-a7e7b9d7c9c0
md"""
### Making sure all cells ran

Some cells skipped running because they were rendered server-side (markdown-only cells). Run them manually.
"""

# ╔═╡ c4d5f03c-3787-4d8d-afff-a324a76ad899
rendered_original = Pluto.WorkspaceManager.eval_fetch_in_workspace(
	(s, nb_original),
	:(PlutoRunner.cell_results |> keys |> collect)
)

# ╔═╡ 0d6a1a94-44a7-454f-ae42-8fde9c3c0977
begin
	to_render = setdiff([c.cell_id for c in nb_original.cells], rendered_original)
	Pluto.update_save_run!(s, nb_original, [nb_original.cells_dict[id] for id in to_render]; run_async=false)

	nb = nb_original
end;

# ╔═╡ 40212e56-3085-4425-a043-1d43fd9ded2c

# ╔═╡ 32a28dac-0ea8-4ba9-a708-37a30241e700
md"""
### Choosing the cells that we want to render
"""

# ╔═╡ 2f251c46-0c96-482b-acbb-2c98f342deff
selected_cell_indices = let
	match = findfirst(nb.cells) do c
		occursin("# Implementation", c.code)
	end

	match === nothing ? eachindex(nb.cells) : (1:match-1)
end

# ╔═╡ cbf5e7ce-8a58-4a83-a36a-2f2fb7830c0b
selected_cells = nb.cells[selected_cell_indices]

# ╔═╡ 40100b17-f1c3-410c-b5bd-572ee3587afa
length(selected_cells)

# ╔═╡ c6f9f70d-a0fd-4c7f-bfbe-f26e4df054cd

# ╔═╡ a3e2bf13-a0ed-4b58-a503-c6c81652db3c
md"""
### Render to Markdown functions
"""

# ╔═╡ 3aa0f01d-7cf9-45c2-af4d-2db45098c635
function cell_repr(cell::Pluto.Cell)::Tuple{String,MIME}
	if Pluto.ends_with_semicolon(cell.code)
		return "", MIME"text/plain"()
	end
	ws = Pluto.WorkspaceManager.get_workspace((s, nb_original))
	Pluto.WorkspaceManager.eval_fetch_in_workspace(
		(s, nb_original),
		quote
			result = PlutoRunner.cell_results[$(cell.cell_id)]

			if result isa CapturedException
				val = result
				stack = [s for (s, _) in val.processed_bt]

			    function_wrap_index = findlast(f -> occursin("#==#", String(f.file)), stack)

			    if function_wrap_index === nothing
			        for _ in 1:2
			            until = findfirst(b -> b.func == :eval, reverse(stack))
			            stack = until === nothing ? stack : stack[1:end - until]
			        end
			    else
			        stack = stack[1:function_wrap_index]
			    end

				return (
					sprint(showerror, result.ex, stack),
					MIME"text/plain"()
				)
			end

			if result isa Nothing
				return ("", MIME"text/plain"())
			end

			mime = if result isa Markdown.MD
				MIME"text/markdown"()
			else
				MIME"text/plain"()
			end

			repr(mime, result; context=IOContext(devnull,
				:color => false, :limit => true, :displaysize => (18, 88),
				:module => getfield(Main, $(QuoteNode(ws.module_name)))
			)), mime
		end
	)
end

# ╔═╡ 3293c088-653e-4ca4-9432-e01c12a1ed18
:(a + $(QuoteNode(:s)))

# ╔═╡ 699ebe9c-ff49-4a27-952a-e72eddd4dfa8
function cell_md_data(cell::Pluto.Cell)::String
	data, mime = cell_repr(cell)
	result = if cell.code_folded
		let
			prefix = let
				r = cell.output.rootassignee
				if r isa Symbol
					"$(r) = "
				else
					""
				end
			end

			if mime isa MIME"text/markdown" || mime isa MIME"text/html"
				prefix * data
			else
				isempty(data) ? prefix : "```\n$(prefix * data)\n```"
			end
		end
	else
		if isempty(data)
			"```julia\n$(cell.code)\n```"
		else
			if mime isa MIME"text/markdown" || mime isa MIME"text/html"
				"```julia\n$(cell.code)\n```\n$(data)"
			else
				multiline = occursin("\n", data)
				"```julia\n$(cell.code)\n\n#=>  $(multiline ? "\n" : "")$(data)\n```"
			end
		end
	end

	if isempty(result)
		"\n
\n\n"
	else
		result * "\n"
	end
end

# ╔═╡ 13f47c9d-a2e6-4ebf-b483-001c6d2fa25e

# ╔═╡ 5c4a510e-c717-480e-9207-66059f206194
md"""
### Let's go!
"""

# ╔═╡ b8716356-2e56-4c0c-bf25-1d9bbe63c9c5
md_data = replace(
	join(cell_md_data.(selected_cells), "\n"),
	r"^@skip_as_script "m => ""
)

# ╔═╡ 2988531d-e94d-4cd9-aa87-f88a115d5f25
Text(md_data)

# ╔═╡ cf83c853-ba25-40db-a635-902f64779ffc
@htl """

Preview

$(Markdown.parse(md_data))

"""

# ╔═╡ 013e190a-e119-414f-8898-972c994e1dcf
md"""
### Write to `README.md`
"""

# ╔═╡ d7a2b1da-95ae-408b-9531-0b0a4ce46054
const REAMDE_PATH = joinpath(PROJECT_ROOT, "README.md")

# ╔═╡ b81b8c8e-ba08-4d58-a502-663d93c32964
write(REAMDE_PATH, md_data)

# ╔═╡ Cell order:
# ╠═96c26596-9fef-11ec-351a-11ad73639906
# ╠═72204279-d847-4bc6-9365-83a0299c8a38
# ╠═1ee7d783-88c1-418f-ae18-d1a25049b467
# ╠═a0deb59c-9fb1-4e3c-9761-690e3c42af76
# ╠═06c8d69c-16bd-4dc1-97ef-f3410c30e756
# ╠═c7f642e6-0660-4bc0-948e-77fadd128b59
# ╟─d6fe479d-0e08-4e1d-af57-deb1e78ce300
# ╟─bf9e7649-6489-4464-adaf-a7e7b9d7c9c0
# ╠═c4d5f03c-3787-4d8d-afff-a324a76ad899
# ╠═0d6a1a94-44a7-454f-ae42-8fde9c3c0977
# ╟─40212e56-3085-4425-a043-1d43fd9ded2c
# ╟─32a28dac-0ea8-4ba9-a708-37a30241e700
# ╠═2f251c46-0c96-482b-acbb-2c98f342deff
# ╟─cbf5e7ce-8a58-4a83-a36a-2f2fb7830c0b
# ╠═40100b17-f1c3-410c-b5bd-572ee3587afa
# ╟─c6f9f70d-a0fd-4c7f-bfbe-f26e4df054cd
# ╟─a3e2bf13-a0ed-4b58-a503-c6c81652db3c
# ╠═3aa0f01d-7cf9-45c2-af4d-2db45098c635
# ╠═3293c088-653e-4ca4-9432-e01c12a1ed18
# ╠═699ebe9c-ff49-4a27-952a-e72eddd4dfa8
# ╟─13f47c9d-a2e6-4ebf-b483-001c6d2fa25e
# ╟─5c4a510e-c717-480e-9207-66059f206194
# ╟─b8716356-2e56-4c0c-bf25-1d9bbe63c9c5
# ╠═2988531d-e94d-4cd9-aa87-f88a115d5f25
# ╟─cf83c853-ba25-40db-a635-902f64779ffc
# ╠═013e190a-e119-414f-8898-972c994e1dcf
# ╠═d7a2b1da-95ae-408b-9531-0b0a4ce46054
# ╠═b81b8c8e-ba08-4d58-a502-663d93c32964


text:
# -------------------------------------------------------------------------------------------------------------------------------------------------------------
# Script to estimate Caliendo and Parro (2015) trade elasticities
# -------------------------------------------------------------------------------------------------------------------------------------------------------------

using DataFrames, RData, XLSX, LinearAlgebra, Statistics, CSV

dir = "X:/VIVES/1-Personal/Florian/git/IOTables/src/"
include(dir * "model/transform_data.jl") # Script with functions to import and transform raw data
include(dir * "elasticity/elasticities_functions.jl") # Script with functions to compute statistics for estimating sectoral trade elasticities (method of Caliendo and Parro, 2015)
# include(dir * "elasticity/wto_tariffs_2018.jl") # Script to prepare WTO tariff data from 2018

# -------------------------------------------------------------------------------------------------------------------------------------------------------------

# Data specification
dir_raw = "C:/Users/u0148308/data/raw/" # location of raw data

# WIOD rev. 2013
source = "WIOD"
revision = "2013"
year = 2011 # specified year
N = 41 # number of countries
S = 35 # number of industries

# # WIOD rev. 2016
# source = "WIOD"
# revision = "2016"
# year = 2014 # specified year
# N = 44 # number of countries
# S = 56 # number of industries

# # OECD rev. 2021
# source = "OECD"
# revision = "2021"
# year = 1995 # specified year
# N = 71 # number of countries
# S = 45 # number of industries

# -------------------------------------------------------------------------------------------------------------------------------------------------------------

Z, F, Y, F_ctry, TB_ctry, VA_ctry, VA_coeff, γ, α, π_Z, π_F = transform_data(dir_raw, source, revision, year, N, S)

# -------------------------------------------------------------------------------------------------------------------------------------------------------------

df_tariffs = DataFrame(XLSX.readtable(dir_raw * "WTO/tariff_matrix.xlsx", "Sheet1")...)
τ_Z = Matrix(convert.(Float64, df_tariffs[:,2:end])) # NS×N
τ_Z = 1.0 .+ τ_Z ./ 100 # NS×N

# -------

# CP statistic is computed using absorption (X), i.e. imports not exports
# therefore, we use the transpose of Z, so we have destination as rows and origin as columns as in τ_Z
M = copy(transpose(Z))

lhs_Z, rhs_Z = elasticity_data(M, τ_Z, "intermediate", N, S)

# -------------------------------------------------------------------------------------------------------------------------------------------------------------
# export regression data with fixed effects included
# -------------------------------------------------------------------------------------------------------------------------------------------------------------

combinations = Int64(sum([n*(n+1)/2 for n in 1:N-2]))

# sector FE
industries = lpad.(1:S, 2, '0') # left pad with leading zeros so we can sort later on
FE_ind = repeat(industries, inner=combinations) # S*combinations×1, sectors are already sorted!

# exporter/importer FE
FE_ctry = zeros(combinations, N)
e = 1
for i in 1:N-2 # i -> j
    for j in i+1:N-1 # j -> n
        for n in j+1:N # n -> i
            FE_ctry[e,i] = 1
            FE_ctry[e,n] = 1
            FE_ctry[e,j] = 1
            e += 1
        end
    end
end

FE_ctry = Int64.(repeat(FE_ctry, S)) # S*combinations×N

# combine data in a DataFrame
ctry_names_2013 = ["AUS", "AUT", "BEL", "BGR", "BRA", "CAN", "CHN", "CYP", "CZE", "DEU", "DNK", "ESP", "EST", "FIN", "FRA", "GBR",
    "GRC", "HUN", "IDN", "IND", "IRL", "ITA", "JPN", "KOR", "LTU", "LUX", "LVA", "MEX", "MLT", "NLD", "POL", "PRT", "ROM", "RUS",
    "SVK", "SVN", "SWE", "TUR", "TWN", "USA"]
col_names = [["industry", "lhs_Z", "rhs_Z"]; ctry_names_2013; "RoW"]

df_reg = DataFrames.DataFrame([FE_ind lhs_Z rhs_Z FE_ctry], col_names)

# CSV.write(dir_raw * "df_reg.csv", df_reg)
# XLSX.writetable(dir_raw * "df_reg.xlsx", df_reg, overwrite=true) # to see results better

# -------------------------------------------------------------------------------------------------------------------------------------------------------------

using GLM

df_reg_Z = subset(df_reg, :lhs_Z => ByRow( x -> !(ismissing(x) || isnothing(x) || isnan(x)) ))
df_reg_Z = df_reg_Z[:, [:industry, :lhs_Z, :rhs_Z]]
transform!(df_reg_Z, :industry => ByRow(string), [:lhs_Z, :rhs_Z] .=> ByRow(Float64), renamecols=false) # columns need to have the right type!

XLSX.writetable(dir_raw * "df_reg_Z.xlsx", df_reg_Z, overwrite=true) # to see results better

ols_Z = GLM.lm(@formula(lhs_Z ~ 0 + rhs_Z), df_reg_Z)

# -----

for i in unique(df_reg_Z.industry)
    gdf = subset(df_reg_Z, :industry => ByRow(x-> x == i))
    ols_Z = GLM.lm(@formula(lhs_Z ~ 0 + rhs_Z), gdf)
    println("Result for industry: $i \n $ols_Z \n")
end


text:
# __precompile__(true)
module Pkg

import Random
import REPL
using REPL.TerminalMenus

depots() = Base.DEPOT_PATH
logdir() = joinpath(depots()[1], "logs")
devdir() = get(ENV, "JULIA_PKG_DEVDIR", joinpath(homedir(), ".julia", "dev"))
const UPDATED_REGISTRY_THIS_SESSION = Ref(false)

have_warned_session = false
function print_first_command_header()
    global have_warned_session
    have_warned_session && return
    isinteractive() || return
    if !PKG3_IS_PRECOMPILED && !haskey(ENV, "JULIA_PKG3_DISABLE_PRECOMPILE_WARNING")
        @info """
        Pkg is running without precompile statements, first action will be slow.
        Rebuild julia with the environment variable `JULIA_PKG3_PRECOMPILE` set to enable precompilation of Pkg.
        This message can be disabled by setting the env variable `JULIA_PKG3_DISABLE_PRECOMPILE_WARNING`.
        """
    end
    have_warned_session = true
end

# load snapshotted dependencies
include("../ext/TOML/src/TOML.jl")

include("GitTools.jl")
include("PlatformEngines.jl")
include("Types.jl")
include("Display.jl")
include("Pkg2/Pkg2.jl")
include("GraphType.jl")
include("Resolve.jl")
include("Operations.jl")
include("API.jl")
include("REPLMode.jl")

# Define new variables so tab comleting Pkg. works.
const add          = API.add
const rm           = API.rm
const up           = API.up
const test         = API.test
const gc           = API.gc
const init         = API.init
const build        = API.build
const installed    = API.installed
const pin          = API.pin
const free         = API.free
const checkout     = API.checkout
const develop      = API.develop
const generate     = API.generate
const instantiate  = API.instantiate
const resolve      = API.resolve
const status       = Display.status
const update       = up
const activate     = API.activate
const setprotocol! = API.setprotocol!

# legacy CI script support
import .API: clone, dir

import .REPLMode: @pkg_str
export @pkg_str

#function __init__()
    if isdefined(Base, :active_repl)
        REPLMode.repl_init(Base.active_repl)
    else
        atreplinit() do repl
            if isinteractive() && repl isa REPL.LineEditREPL
                isdefined(repl, :interface) || (repl.interface = REPL.setup_interface(repl))
                REPLMode.repl_init(repl)
            end
        end
    end
#end

using .Types
using UUIDs
import LibGit2
import Dates
# This crashes low memory systems and some of Julia's CI
# so keep it disabled by default for now.
if haskey(ENV, "JULIA_PKG3_PRECOMPILE")
    const PKG3_IS_PRECOMPILED = true
    include("precompile.jl")
else
    const PKG3_IS_PRECOMPILED = false
end

METADATA_compatible_uuid(pkg::String) = Types.uuid5(Types.uuid_package, pkg)

end # module


text:
# Statistical model for analyzing Radial Velocity observations
# Assumes existance of data in arrays times, obs, sigma_obs
# Assumes existance of functions: TODO

export loglikelihood, logprior_noise, is_valid_noise

num_noise_param() = 0    # Set to 1 for using jitter, 0 for none
const Jitter0 = 1.0  # units of m/s
const max_jitter = 10000.0
const jitter_prior_norm = -log(log1p(max_jitter/Jitter0))
jitter_offset(theta::Vector) = noise_param_offset(theta)

function set_jitter(theta::Vector, jitter)
  if num_noise_param() == 1
     theta[jitter_offset(theta)] = log1p(jitter/Jitter0)
  end
  return
end

function extract_jitter(theta::Vector)
  if num_noise_param() == 1
     return Jitter0*(exp(theta[jitter_offset(theta)])-1)
  else
     return 0
  end
end

function is_valid_noise(p::Vector)
  param = extract_noise_param(p)
  @assert length(param) == num_noise_param()
  if num_noise_param() == 1
     jitter = extract_jitter(p)
     if !(0<=jitter<=max_jitter)
       return false
     end
  end
  return true
end

function loglikelihood(p::Vector)
  num_pl = num_planets(p)
  if !is_valid(p) return -Inf end  # prempt model evaluation
  # Set t, o, and so to point to global arrays with observational data, while enforcing types
  t::Array{Float64,1} = times
  o::Array{Float64,1} = obs
  so::Array{Float64,1} = sigma_obs
  @assert length(times) == length(obs) == length(sigma_obs)
  jitter_sq = num_noise_param() ==1 ? extract_jitter(p)^2 : 0.
  chisq = zero(eltype(p))
  log_normalization = -0.5*length(t)*log(2pi)
  for i in 1:length(t)
    model::eltype(p) = calc_model_rv(p,t[i])
    sigma_eff_sq = so[i]^2+jitter_sq
    chisq += abs2(model-o[i])/sigma_eff_sq
    log_normalization -= 0.5*log(sigma_eff_sq)
  end
  return -0.5*(chisq)+log_normalization
end

function logprior_noise(p::Vector)
  logp = zero(eltype(p))
  if num_noise_param() == 1
    if !is_valid_noise(p)
      logp -=Inf          # prempt model evaluation
      return logp
    end
    jitter::eltype(p) = extract_jitter(p)
    logp += -log1p(jitter/Jitter0) + jitter_prior_norm
  end
  return logp::eltype(p)
end


text:
# Example plotting using PyPlot

using COMTRADE
using Glob
using PyCall
@pyimport matplotlib.backends.backend_pdf as pdf
using PyPlot

# Plot the first six analog channels in each of the test data files
function process_files()
    p = pdf.PdfPages("waves.pdf")
    ioff()
    for fn in glob("../test/data/*.cfg")
        println(fn)
        fn = replace(fn, ".cfg", "")
        z = read_comtrade(fn)
        f, ax = subplots(6, 1, sharex = true, figsize = [12, 12])
        nm = names(z.dat)
        for i in 1:6
            ax[i,1][:plot](z.dat[:time], z.dat[i+2])
            ax[i,1][:set_ylabel](z.cfg.A[i, :ch_id])
        end
        ax[1,1][:set_title](fn)
        p[:savefig]()
    end
    p[:close]()
    ion()
end

process_files()


text:
_jl_libRmath = dlopen("libRmath-julia")

## c function signatures:
## double dnchisq(double x, double df, double ncp, int give_log)
## double pnchisq(double x, double df, double ncp, int lower_tail, int log_p)
## double qnchisq(double p, double df, double ncp, int lower_tail, int log_p)
## double rnchisq(double df, double lambda)

immutable NChiSq <: ContinuousUnivariateDistribution
    ν::Float64
    λ::Float64
    function NChiSq(ν::Float64, λ::Float64)
        if ν <= 0
            error("ν of NChiSq must be larger than zero.")
        end
        if λ <= 0
            error("λ of NChiSq must be larger than zero.")
        end

        return new(ν, λ)
    end
end

function NChiSq(params::Array{Float64, 1})
    return NChiSq(params...)
end

import Base.Multimedia.display
function display(ncs::NChiSq)
    nDigits = 3
    println("\nUnivariate noncentral chi-squared distribution:")
    println("    ν: $(round(ncs.ν, nDigits))")
    println("    λ: $(round(ncs.λ, nDigits))")
end

## access parameters
##------------------

import Distributions.dof
function dof(ncs::NChiSq)
    return ncs.ν
end

## getParams
##----------

function getParams(ncs::NChiSq)
    return (ncs.ν, ncs.λ)
end

## rand
##-----

import Distributions.rand
function rand(ncs::NChiSq)
    return ccall(dlsym(_jl_libRmath,:rnchisq), Float64,
                 (Float64,Float64),
                 ncs.ν, ncs.λ)
end

function rand(ncs::NChiSq, nObs::Int)
    simVals = Array(Float64, nObs)
    for ii=1:nObs
        simVals[ii] = rand(ncs)
    end
    return simVals
end

## pdf
##----

import Distributions.pdf
function pdf(d::NChiSq, x::Float64)
    return ccall(dlsym(_jl_libRmath,:dnchisq), Float64,
                 (Float64,Float64,Float64,Int32),
                 x, d.ν, d.λ, false)
end

function pdf(d::NChiSq, x::Array{Float64, 1})
    nVals = length(x)
    pdfVals = Array(Float64, nVals)
    for ii=1:nVals
        pdfVals[ii] = pdf(d, x[ii])
    end
    return pdfVals
end

## cdf
##----

import Distributions.cdf
function cdf(d::NChiSq, x::Float64)
    return ccall(dlsym(_jl_libRmath,:pnchisq), Float64,
                 (Float64,Float64,Float64,Int32,Int32),
                 x, d.ν, d.λ, true, false)
end

function cdf(d::NChiSq, x::Array{Float64, 1})
    nVals = length(x)
    cdfVals = Array(Float64, nVals)
    for ii=1:nVals
        cdfVals[ii] = cdf(d, x[ii])
    end
    return cdfVals
end

## quantile
##---------

import Distributions.quantile
function quantile(d::NChiSq, x::Float64)
    return ccall(dlsym(_jl_libRmath,:qnchisq), Float64,
                 (Float64,Float64,Float64,Int32,Int32),
                 x, d.ν, d.λ, true, false)
end

function quantile(d::NChiSq, x::Array{Float64, 1})
    nVals = length(x)
    pVals = Array(Float64, nVals)
    for ii=1:nVals
        pVals[ii] = quantile(d, x[ii])
    end
    return pVals
end

## nllh
##-----

function nllh(d::NChiSq, data::Array{Float64, 1})
    return -sum(log(pdf(d, data)))
end

function nllh_nchis(params::Array{Float64, 1}, data::Array{Float64, 1})
    d = NChiSq(params...)
    return -sum(log(pdf(d, data)))
end

## fit
##----

import Distributions.fit
function fit(dt::Type{NChiSq}, data::Array{Float64, 1})

    opt = Opt(:LN_BOBYQA, 2)
    function objFun(parameters::Vector, grad::Vector)
        ## objective function calculating portfolio variance
        if length(grad) > 0
            ## no partial derivative given
        end

        ## calculate portfolio variance
        nllhVal = nllh_nchis(parameters, data)
        return nllhVal
    end

    min_objective!(opt, objFun)
    xtol_rel!(opt, 1e-4)

    lower_bounds!(opt, [0.1, 0.1])
    upper_bounds!(opt, [Inf, Inf])

    initNu = 3.0
    initMu = mean(data)
    initParams = [initNu, initMu]

    (minf, minx, ret) = optimize(opt, initParams)
    return NChiSq(minx)
end


text:
@doc raw"""
    CostProblem <: Problem

speficy a problem for solvers just based on cost functions, i.e.
gradient free ones.

# Fields

* `M`            – a manifold $\mathcal M$
* `cost` – a function $F\colon\mathcal M\to\mathbb R$ to minimize

# See also
[`NelderMead`](@ref)
"""
struct CostProblem{mT <: Manifold} <: Problem
    M::mT
    cost::Function
end

@doc raw"""
    NelderMeadOptions <: Options

Describes all parameters and the state of a Nealer-Mead heuristic based
optimization algorithm.

# Fields

The naming of these parameters follows the [Wikipedia article](https://en.wikipedia.org/wiki/Nelder–Mead_method)
of the Euclidean case. The default is given in brackets, the required value range
after the description

* `population` – an `Array{`point`,1}` of $n+1$ points $x_i$, $i=1,\ldots,n+1$, where $n$ is the
  dimension of the manifold.
* `stopping_criterion` – ([`StopAfterIteration`](@ref)`(2000)`) a [`StoppingCriterion`](@ref)
* `retraction` – (`exp`) the rectraction to use
* `α` – (`1.`) reflection parameter ($\alpha > 0$)
* `γ` – (`2.`) expansion parameter ($\gamma>0$)
* `ρ` – (`1/2`) contraction parameter, $0 < \rho \leq \frac{1}{2}$,
* `σ` – (`1/2`) shrink coefficient, $0 < \sigma \leq 1$
* `x` – (`p[1]`) - a field to collect the current best value

# Constructors

    NelderMead(M,stop, retr; α=1. , γ=2., ρ=1/2, σ=1/2)

construct a Nelder-Mead Option with a set of `dimension(M)+1` random points.

    NelderMead(p, stop retr; α=1. , γ=2., ρ=1/2, σ=1/2)

construct a Nelder-Mead Option with a set `p` of points
"""
mutable struct NelderMeadOptions{T} <: Options
    population::Array{T,1}
    stop::StoppingCriterion
    α::Real
    γ::Real
    ρ::Real
    σ::Real
    x::T
    costs::Array{Float64,1}
    function NelderMeadOptions(
        M::MT;
        stop::StoppingCriterion = StopAfterIteration(2000),
        α = 1.,
        γ = 2.,
        ρ=1/2,
        σ = 1/2
    ) where {MT <: Manifold }
        p = [random_point(M) for i=1:(manifold_dimension(M)+1) ]
        new{eltype(p)}(p, stop, α, γ, ρ, σ, p[1],[])
    end
    function NelderMeadOptions(population::Array{T,1},
        stop::StoppingCriterion = StopAfterIteration(2000);
        α = 1., γ = 2., ρ=1/2, σ = 1/2
    ) where {T}
        return new{T}(population, stop, α, γ, ρ, σ, population[1],[])
    end
end
get_solver_result(o::NelderMeadOptions) = o.x


text:
const DOCUMENTER_UUID = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
const STANDARD_KWS = [:modules, :format, :pages, :repo, :sitename, :authors, :assets]

"""
Add a `Documenter` subtype to a template's plugins to add support for documentation
generation via [Documenter.jl](https://github.com/JuliaDocs/Documenter.jl).

By default, the plugin generates a minimal index.md and a make.jl file. The make.jl
file contains the Documenter.makedocs command with predefined values for `modules`,
`format`, `pages`, `repo`, `sitename`, and `authors`.

The subtype is expected to include the following fields:
* `assets::Vector{AbstractString}`, a list of filenames to be included as the `assets`
kwarg to `makedocs`
* `gitignore::Vector{AbstractString}`, a list of files to be added to the `.gitignore`

It may optionally include the field `additional_kwargs::Union{AbstractDict, NamedTuple}`
to allow additional kwargs to be added to `makedocs`.
"""
abstract type Documenter <: CustomPlugin end

function gen_plugin(p::Documenter, t::Template, pkg_name::AbstractString)
    path = joinpath(t.dir, pkg_name)
    docs_dir = joinpath(path, "docs")
    mkpath(docs_dir)

    # Create the documentation project.
    proj = Base.current_project()
    try
        Pkg.activate(docs_dir)
        Pkg.add(PackageSpec(; name="Documenter", uuid=DOCUMENTER_UUID))
    finally
        proj === nothing ? Pkg.activate() : Pkg.activate(proj)
    end

    tab = repeat(" ", 4)
    assets_string = if !isempty(p.assets)
        mkpath(joinpath(docs_dir, "src", "assets"))
        for file in p.assets
            cp(file, joinpath(docs_dir, "src", "assets", basename(file)))
        end

        # We want something that looks like the following:
        # [
        #         assets/file1,
        #         assets/file2,
        #     ]
        s = "[\n"
        for asset in p.assets
            s *= """$(tab^2)"assets/$(basename(asset))",\n"""
        end
        s *= "$tab]"

        s
    else
        "String[]"
    end

    kwargs_string = if :additional_kwargs in fieldnames(typeof(p)) &&
        fieldtype(typeof(p), :additional_kwargs) <: Union{AbstractDict, NamedTuple}
        # We want something that looks like the following:
        #     key1="val1",
        #     key2="val2",
        #
        kws = [keys(p.additional_kwargs)...]
        valid_keys = filter(k -> !in(Symbol(k), STANDARD_KWS), kws)
        if length(p.additional_kwargs) > length(valid_keys)
            invalid_keys = filter(k -> Symbol(k) in STANDARD_KWS, kws)
            @warn string(
                "Ignoring predefined Documenter kwargs ",
                join(map(repr, invalid_keys), ", "),
                " from additional kwargs"
            )
        end
        join(map(k -> string(tab, k, "=", repr(p.additional_kwargs[k]), ",\n"), valid_keys))
    else
        ""
    end

    make = """
        using Documenter, $pkg_name

        makedocs(;
            modules=[$pkg_name],
            format=Documenter.HTML(),
            pages=[
                "Home" => "index.md",
            ],
            repo="https://$(t.host)/$(t.user)/$pkg_name.jl/blob/{commit}{path}#L{line}",
            sitename="$pkg_name.jl",
            authors="$(t.authors)",
            assets=$assets_string,
        $kwargs_string)
        """
    docs = """
    # $pkg_name.jl

    ```@index
    ```

    ```@autodocs
    Modules = [$pkg_name]
    ```
    """

    gen_file(joinpath(docs_dir, "make.jl"), make)
    gen_file(joinpath(docs_dir, "src", "index.md"), docs)
end

function Base.show(io::IO, p::Documenter)
    spc = "  "
    println(io, nameof(typeof(p)), ":")

    n = length(p.assets)
    s = n == 1 ? "" : "s"
    print(io, spc, "→ $n asset file$s")
    if n == 0
        println(io)
    else
        println(io, ": ", join(map(a -> replace(a, homedir() => "~"), p.assets), ", "))
    end

    n = length(p.gitignore)
    s = n == 1 ? "" : "s"
    print(io, "$spc→ $n gitignore entrie$s")
    n > 0 && print(io, ": ", join(map(repr, p.gitignore), ", "))
end


text:
export Dimer, dimer

type Dimer <: Model
  S2 :: Int
  Jz :: Float64
  Jxy :: Float64
  coeff :: BondCoeffs
  function Dimer(S2, Jz, Jxy)
    c = bond_coeffs(S2)
    new(S2, Jz, Jxy, c)
  end
end
dimer(S2::Integer, Jz::Real, Jxy::Real) = Dimer(S2,Jz, Jxy)
dimer(S2::Integer, J::Real) = dimer(S2,J,J)
num_sites(::Dimer) = 2
num_bonds(::Dimer) = 1

function interall(dimer::Dimer, t::Real = 0.0)
  const io = open("interall.def", "w")
  header_interall(io, dimer)
  bond(io, 0, 1, dimer.Jz, dimer.Jxy*cis(t), dimer.coeff)

  close(io)
end

max_stars_repo_name: yomichi/HPhi-samples

max_issues_repo_name: yomichi/HPhi-samples

max_forks_repo_name: yomichi/HPhi-samples


text:
# This file was generated, do not modify it. # hide
Random.seed!(1515)

@load KMeans
@pipeline SPCA2(std = Standardizer(),
                pca = PCA(),
                km = KMeans(k=3))

spca2_mdl = SPCA2()
spca2 = machine(spca2_mdl, X)
fit!(spca2)

assignments = report(spca2).reports[1].assignments
mask1 = assignments .== 1
mask2 = assignments .== 2
mask3 = assignments .== 3;


text:
# Train an normalizing flow for seismic interpolation
# Authors: Rajiv Kumar
#          Ali Siahkoohi, alisk@gatech.edu
# Date: February 2021
"""
Example:
julia train_hint_interpolation.jl
"""

using DrWatson
@quickactivate :WavefieldRecoveryUQ

using InvertibleNetworks
using JLD2
using HDF5
using Random
using Statistics
using ArgParse
using ProgressMeter
using Flux: gpu, cpu
using Flux.Data: DataLoader
using Flux.Optimise: Optimiser, ExpDecay, update!, ADAM

# Random seed
Random.seed!(66)

s = ArgParseSettings()
@add_arg_table! s begin
    "--max_epoch"
        help = "Maximum # of epochs"
        arg_type = Int
        default = 12
    "--lr"
        help = "starting learning rate"
        arg_type = Float32
        default = 1f-3
    "--lr_step"
        help = "lr scheduler applied at evey lr_step epoch"
        arg_type = Int
        default = 1
    "--batchsize"
        help = "batch size"
        arg_type = Int
        default = 8
    "--n_hidden"
        help = "# hidden channels"
        arg_type = Int
        default = 64
    "--depth"
        help = "depth of the network"
        arg_type = Int
        default = 8
    "--subsample_fac"
        help = "subsampling factor"
        arg_type = Float32
        default = 7.5f-1
    "--sim_name"
        help = "simulation name"
        arg_type = String
        default = "wavefield_recovery_and_UQ"
end
parsed_args = parse_args(s)

max_epoch = parsed_args["max_epoch"]
lr = parsed_args["lr"]
lr_step = parsed_args["lr_step"]
batchsize = parsed_args["batchsize"]
n_hidden = parsed_args["n_hidden"]
depth = parsed_args["depth"]
subsample_fac = parsed_args["subsample_fac"]
sim_name = parsed_args["sim_name"]

# Define raw data directory
mkpath(datadir("dataset"))
data_path = datadir("dataset", "frequency_slices.h5")

# Download the dataset into the data directory if it does not exist
if isfile(data_path) == false
    run(`wget https://www.dropbox.com/s/wcqucrlye615fhi/'
        'frequency_slices.h5 -q -O $data_path`)
end

# Load seismic images and create training and testing data
h5open(data_path, "r") do file
    X_orig = file["dataset"][:, :, :, 1:2:end]

    X_orig = permutedims(X_orig, (2, 3, 1, 4))
    global nx, ny, nc, nsamples = size(X_orig)

    # Whiten data
    AN = ActNorm(nsamples)
    X_orig = AN.forward(X_orig)
    global AN_params = get_params(AN)

    # Forward operator
    global M = Mask([nx, ny], subsample_fac)

    # Split in training/testing
    ntrain = Int(floor((nsamples*.75)))
    ntest = nsamples - ntrain

    rand_per = randperm(nsamples)
    global train_idx = rand_per[1:ntrain]
    global test_idx = rand_per[ntrain+1:end]

    global X_train = X_orig[:, :, :, train_idx]
    global Y_train = M*X_train

end

# Create network
CH = NetworkConditionalHINT(
    nc, n_hidden, depth, k1=3, k2=3, p1=1, p2=1
) |> gpu

# Training
# Batch extractor
train_loader = DataLoader(
    range(1, length(train_idx), step=1),batchsize=batchsize, shuffle=true
)
num_batches = length(train_loader)

# Optimizer
opt = Optimiser(ExpDecay(lr, .9f0, num_batches*lr_step, 0f0), ADAM(lr))

# Training log keeper
fval = zeros(Float32, num_batches*max_epoch)
p = Progress(num_batches*max_epoch)

for epoch=1:max_epoch
    for (itr, idx) in enumerate(train_loader)
        Base.flush(Base.stdout)

        X = X_train[:, :, :, idx]
        Y = Y_train[:, :, :, idx]

        X = X |> gpu
        Y = Y |> gpu

        fval[(epoch-1)*num_batches + itr] = loss_supervised(CH, X, Y)[1]

        ProgressMeter.next!(
            p;
            showvalues = [
                (:Epoch, epoch),
                (:Itreration, itr),
                (:NLL, fval[(epoch-1)*num_batches + itr])
            ]
        )

        # Update params
        for p in get_params(CH)
            update!(opt, p.data, p.grad)
        end
        clear_grad!(CH)
    end

    # Saving parameters and logs
    Params = get_params(CH) |> cpu
    save_dict = @strdict epoch max_epoch lr lr_step batchsize n_hidden depth subsample_fac sim_name M.M AN_params Params fval train_idx test_idx
    @tagsave(
        datadir(sim_name, savename(save_dict, "jld2"; digits=6)),
        save_dict;
        safe=true
    )

end


text:
workspace()
ENV["JuBrain_Mode"]="MPI"
using JuBrain

model=networkModel()
model.parameters="taupre=20;taupost=20;tmax=50"
model.equations="dapre/dt=-apre/taupre;dapost/dt=-apost/taupost;dW/dt=0"

model.groups=Groups(groupName=[1,2],groupSize=[100,100])
initialize(model,groups=1,expr="tspike=(i-1)*tmax/(100-1)")
initialize(model,groups=2,expr="tspike=(100-i)*tmax/(100-1)")
Connect(model,expr="apre=0.0")
Connect(model,expr="apost=0.0")
Connect(model,expr="W=0.0")
Connect(model,pre=1,post=2,condition="i==j",expr="Apre=0.01")
Connect(model,pre=1,post=2,condition="i==j",expr="Apost=-0.01*1.05")

model.spike="abs(t-tspike).<=0.5/2"
model.plasticity="apre[j,i]+=Apre[j,i];W[j,i]+=apost[j,i];apost[i,j]+=Apost[i,j];W[i,j]+=apre[i,j]"
model.dt=0.5;model.saveVars=["W","tspike"]

model.postProcs="""
@load("Hebbian_code.jld")
using MatlabPlot
index1=model.groups[1][:localindex];index2=model.groups[2][:localindex]
index=sub2ind(size(W),index2,index1)
mplot(tspike[index2]-tspike[index1],W[index])"""

buildnet(name="Hebbian_code",model=model,duration=100)


text:
#################### Chains ####################

#################### Constructors ####################
"""
    Chains(iters::Integer, params::Integer;
    start::Integer=1, thin::Integer=1, chains::Integer=1,
    names::Vector{T}=AbstractString[]) where {T<:AbstractString}
"""
function Chains(
    iters::Integer,
    params::Integer;
    start::Integer = 1,
    thin::Integer = 1,
    chains::Integer = 1,
    names::Vector{T} = AbstractString[],
    ntrees::Integer = 1,
    tree_names::Vector{Symbol} = Symbol[],
) where {T<:AbstractString}
    value = Array{Float64}(undef, length(start:thin:iters), params, chains)
    value2 = Array{AbstractString}(undef, length(start:thin:iters), ntrees, chains)
    fill!(value, NaN)

    Chains(
        value,
        value2,
        start = start,
        thin = thin,
        names = names,
        tree_names = string.(tree_names),
    )
end

"""
    Chains(value::Array{T, 3},
    start::Integer=1, thin::Integer=1,
    names::Vector{W}=AbstractString[], chains::Vector{V}=Int[], moves::Vector{V}=Int[0])
    where {T<:Real, U<:AbstractString, V<:Integer, W <: AbstractString}
"""
function Chains(
    value::Array{T,3},
    start::Integer = 1,
    thin::Integer = 1,
    names::Vector{W} = AbstractString[],
    chains::Vector{V} = Int[],
) where {T<:Real,U<:AbstractString,V<:Integer,W<:AbstractString}
    Chains(
        value,
        Array{String,3}(undef, size(value)),
        start = start,
        thin = thin,
        names = names,
        chains = chains,
    )
end
"""
    Chains(value::Array{T, 3},
    value2::Array{U,3};
    start::Integer=1, thin::Integer=1,
    names::Vector{W}=AbstractString[], chains::Vector{V}=Int[], moves::Vector{V}=Int[0])
    where {T<:Real, U<:AbstractString, V<:Integer, W <: AbstractString}
"""
function Chains(
    value::Array{T,3},
    value2::Array{U,3};
    start::Integer = 1,
    thin::Integer = 1,
    names::Vector{W} = AbstractString[],
    chains::Vector{V} = Int[],
    tree_names::Vector{X} = AbstractString[],
) where {T<:Real,U<:AbstractString,V<:Integer,W<:AbstractString,X<:AbstractString}
    n, p, m = size(value)

    if isempty(names)
        names = map(i -> "Param$i", 1:p)
    elseif length(names) != p
        throw(DimensionMismatch("size(value, 2) and names length differ"))
    end

    if isempty(chains)
        chains = collect(1:m)
    elseif length(chains) != m
        throw(DimensionMismatch("size(value, 3) and chains length differ"))
    end

    v = convert(Array{Float64,3}, value)
    Chains(
        v,
        range(start, step = thin, length = n),
        AbstractString[names...],
        Int[chains...],
        value2,
        tree_names,
    )
end
"""
    Chains(value::Matrix{T};
    start::Integer=1, thin::Integer=1,
    names::Vector{U}=AbstractString[], chains::Integer=1)
    where {T<:Real, U<:AbstractString}
"""
function Chains(
    value::Matrix{T};
    start::Integer = 1,
    thin::Integer = 1,
    names::Vector{U} = AbstractString[],
    chains::Integer = 1,
) where {T<:Real,U<:AbstractString}
    cont_vals = reshape(value, size(value, 1), size(value, 2), 1)
    tree_vals = Array{String,3}(undef, size(cont_vals))
    Chains(
        cont_vals,
        tree_vals,
        start = start,
        thin = thin,
        names = names,
        chains = Int[chains],
    )
end
"""
    Chains(value::Vector{T};
    start::Integer=1, thin::Integer=1,
    names::U="Param1", chains::Integer=1) where {T<:Real, U <: AbstractString}

Construct a `Chains` object that stores MCMC sampler output.

Returns an object of type `Chains`.

* `iters`: total number of iterations in each sampler run, of which `length(start:thin:iters)` outputted iterations will be stored in the object.

* `params`: number of parameters to store.

* `value`: array whose first, second (optional), and third (optional) dimensions index outputted iterations, parameter elements, and runs of an MCMC sampler, respectively.

* `start`: number of the first iteration to be stored.

* `thin`: number of steps between consecutive iterations to be stored.

* `chains`: number of simulation runs for which to store output, or indices to the runs (default: 1, 2, …).

* `names`: names to assign to the parameter elements (default: `"Param1"`, `"Param2"`, …).
"""
function Chains(
    value::Vector{T};
    start::Integer = 1,
    thin::Integer = 1,
    names::U = "Param1",
    chains::Integer = 1,
) where {T<:Real,U<:AbstractString}
    Chains(
        reshape(value, length(value), 1, 1),
        start = start,
        thin = thin,
        names = U[names],
        chains = Int[chains],
    )
end

#################### Indexing ####################
"""
    Base.getindex(c::Chains, window, names, chains)
Subset MCMC sampler output. The syntax `c[i, j, k]` is converted to `getindex(c, i, j, k)`.

Subsetted sampler output stored in the same type of object as that supplied in the call.

* `c` : sampler output to subset.

* `window` : indices of the form `start:stop` or `start:thin:stop` can be used to subset iterations, where `start` and `stop` define a range for the subset and `thin` will apply additional thinning to existing sampler output.

* `names` : indices for subsetting of parameters that can be specified as strings, integers, or booleans identifying parameters to be kept. `ModelChains` may additionally be indexed by model node symbols.

* `chains` : indices for chains can be integers or booleans.

A value of `:` can be specified for any of the dimensions to indicate no subsetting.
"""
function Base.getindex(c::Chains, window, names, chains)
    inds1 = window2inds(c, window)
    inds2 = names2inds(c, names)
    if !isdefined(c.trees, 1)
        newsize = size(c.value[inds1, inds2, chains])
        Chains(
            c.value[inds1, inds2, chains],
            Array{String,length(newsize)}(undef, newsize),
            start = first(c) + (first(inds1) - 1) * step(c),
            thin = step(inds1) * step(c),
            names = c.names[inds2],
            chains = c.chains[chains],
        )
    else
        Chains(
            c.value[inds1, inds2, chains],
            c.trees[inds1, :, chains],
            start = first(c) + (first(inds1) - 1) * step(c),
            thin = step(inds1) * step(c),
            names = c.names[inds2],
            chains = c.chains[chains],
        )
    end
end

Base.lastindex(c::AbstractChains, i) = size(c, i)
"""
    Base.setindex!(c::AbstractChains, value, iters, names, chains)

Store MCMC sampler output at a given index. The syntax `c[i, j, k] = value` is converted to `setindex!(c, value, i, j, k)`.

Returns an object of the same type as c with the sampler output stored in the specified indices.

* `c` : object within which to store sampler output.

* `value` : sampler output.

* `iters` : iterations can be indexed as a `start:stop` or `start:thin:stop` range, a single numeric index, or a vector of indices; and are taken to be relative to the index range store in the `c.range` field.

* `names` : indices for subsetting of parameters can be specified as strings, integers, or booleans.

* `chains` : indices for chains can be integers or booleans.

A value of `:` can be specified for any of the dimensions to indicate no subsetting.
"""
function Base.setindex!(c::AbstractChains, value, iters, names, chains)
    setindex!(c.value, value, iters2inds(c, iters), names2inds(c, names), chains)
end

macro mapiters(iters, c)
    return esc(quote
        ($iters .- first($c)) ./ step($c) .+ 1.0
    end)
end

window2inds(c::AbstractChains, window) =
    throw(ArgumentError("$(typeof(window)) iteration indexing is unsupported"))
window2inds(c::AbstractChains, ::Colon) = window2inds(c, 1:size(c, 1))
window2inds(c::AbstractChains, window::Int) = window2inds(c, window:window)
window2inds(c::AbstractChains, window::AbstractRange) = begin
    range = @mapiters(window, c)
    a = max(ceil(Int, first(range)), 1)
    b = step(window)
    c = min(floor(Int, last(range)), size(c.value, 1))
    a:b:c
end

iters2inds(c::AbstractChains, iters) = iters
iters2inds(c::AbstractChains, ::Colon) = 1:size(c.value, 1)
iters2inds(c::AbstractChains, iters::AbstractRange) =
    convert(StepRange{Int,Int}, @mapiters(iters, c))
iters2inds(c::AbstractChains, iter::Real) = Int(@mapiters(iter, c))
iters2inds(c::AbstractChains, iters::Vector{T}) where {T<:Real} =
    Int[@mapiters(i, c) for i in iters]

names2inds(c::AbstractChains, names) = names
names2inds(c::AbstractChains, ::Colon) = 1:size(c.value, 2)
names2inds(c::AbstractChains, name::Real) = [name]
names2inds(c::AbstractChains, name::AbstractString) = names2inds(c, [name])
names2inds(c::AbstractChains, names::Vector{T}) where {T<:AbstractString} =
    indexin(names, c.names)

#################### Concatenation ####################
"""
    Base.cat(c1::AbstractChains, args::AbstractChains...; dims::Integer)

Concatenate input MCMC chains along a specified dimension. For dimensions other than the specified one, all input chains must have the same sizes, which will also be the sizes of the output chain. The size of the output chain along the specified dimension will be the sum of the sizes of the input chains in that dimension. vcat concatenates vertically along dimension 1, and has the alternative syntax [chain1; chain2; ...]. hcat concatenates horizontally along dimension 2, and has the alternative syntax [chain1 chain2 ...].

Returns a `Chains` object containing the concatenated input.

* `dim` : dimension (1, 2, or 3) along which to concatenate the input chains.

* `c1`, `args...` : Chains to concatenate.

"""
function Base.cat(c1::AbstractChains, args::AbstractChains...; dims::Integer)
    dims == 1 ? cat1(c1, args...) :
    dims == 2 ? cat2(c1, args...) :
    dims == 3 ? cat3(c1, args...) :
    throw(ArgumentError("cannot concatenate along dimension $dim"))
end

function cat1(c1::A, args::A...) where {A<:AbstractChains}

    range = c1.range
    for c in args
        last(range) + step(range) == first(c) ||
            throw(ArgumentError("noncontiguous chain iterations"))
        step(range) == step(c) || throw(ArgumentError("chain thinning differs"))
        range = first(range):step(range):last(c)
    end

    names = c1.names
    tree_names = c1.tree_names
    all(c -> c.names == names, args) || throw(ArgumentError("chain names differ"))
    all(c -> c.tree_names == tree_names, args) ||
        throw(ArgumentError("chain tree names differ"))

    chains = c1.chains
    all(c -> c.chains == chains, args) || throw(ArgumentError("sets of chains differ"))

    value = cat(c1.value, map(c -> c.value, args)..., dims = 1)

    if isassigned(c1.trees, 1)
        trees = cat(c1.trees, map(c -> c.trees, args)..., dims = 1)
    else
        trees = c1.trees
    end

    Chains(
        value,
        trees,
        start = first(range),
        thin = step(range),
        names = names,
        chains = chains,
        tree_names = tree_names,
    )
end

function cat2(c1::A, args::A...) where {A<:AbstractChains}
    range = c1.range
    all(c -> c.range == range, args) || throw(ArgumentError("chain ranges differ"))

    names = c1.names
    n = length(names)
    for c in args
        names = union(names, c.names)
        n += length(c.names)
        n == length(names) || throw(ArgumentError("non-unique chain names"))
    end

    chains = c1.chains
    all(c -> c.chains == chains, args) || throw(ArgumentError("sets of chains differ"))

    value = cat(c1.value, map(c -> c.value, args)..., dims = 2)
    if isassigned(c1.trees, 1)
        trees = cat(c1.trees, map(c -> c.trees, args)..., dims = 2)
    else
        trees = c1.trees
    end

    Chains(
        value,
        trees,
        start = first(range),
        thin = step(range),
        names = names,
        chains = chains,
    )
end

function cat3(c1::A, args::A...) where {A<:AbstractChains}
    range = c1.range
    all(c -> c.range == range, args) || throw(ArgumentError("chain ranges differ"))

    names = c1.names
    all(c -> c.names == names, args) || throw(ArgumentError("chain names differ"))
    tree_names = c1.tree_names
    all(c -> c.tree_names == tree_names, args) || throw(ArgumentError("chain names differ"))

    value = cat(c1.value, map(c -> c.value, args)..., dims = 3)

    if isassigned(c1.trees, 1)
        value2 = cat(c1.trees, map(c -> c.trees, args)..., dims = 3)
    else
        value2 = c1.trees
    end

    Chains(
        value,
        value2,
        start = first(range),
        thin = step(range),
        names = names,
        tree_names = tree_names,
    )
end

Base.hcat(c1::AbstractChains, args::AbstractChains...) = cat(c1, args..., dims = 2)

Base.vcat(c1::AbstractChains, args::AbstractChains...) = cat(c1, args..., dims = 1)

#################### Base Methods ####################
"""
    Base.keys(c::AbstractChains)

Returns names of parameter elements.

* `c` : Chain to return names of.

"""
function Base.keys(c::AbstractChains)
    c.names
end
"""
    Base.show(io::IO, c::AbstractChains)

Prints header and values of Chain.

* `io` : IO stream on which to print.

* `AbstractChains` : Chain to print.
"""
function Base.show(io::IO, c::AbstractChains)
    print(io, "Object of type \"$(summary(c))\"\n\n")
    print(io, header(c))
    isa(c, ModelChains) && show(io, c.sim_params; short = true)
    println()
    show(io, c.value)
end
"""
    Base.size(c::AbstractChains)

Returns Tuple containing last iteration of MCMC sampler output and dimensions of Chain dimensions of c.value.

* `c` : Chain object of interest.

"""
function Base.size(c::AbstractChains)
    dim = size(c.value)
    last(c), dim[2], dim[3]
end
"""
    Base.size(c::AbstractChains, ind)

Returns last iteration of MCMC sampler output, or dimension derived from C, according to value of ind.

* `c` : Chain object of interest.

* `ind` : index of tuple to return; 1 returns last iteration of MCMC sampler output, 2 and 3 return dimensions of c.value.
"""
function Base.size(c::AbstractChains, ind)
    size(c)[ind]
end

Base.first(c::AbstractChains) = first(c.range)
Base.step(c::AbstractChains) = step(c.range)
Base.last(c::AbstractChains) = last(c.range)

#################### Auxilliary Functions ####################

function combine(c::AbstractChains)
    n, p, m = size(c.value)
    value = Array{Float64}(undef, n * m, p)
    for j = 1:p
        idx = 1
        for i = 1:n, k = 1:m
            value[idx, j] = c.value[i, j, k]
            idx += 1
        end
    end
    value
end

function header(c::AbstractChains)
    string(
        "Iterations = $(first(c)):$(last(c))\n",
        "Thinning interval = $(step(c))\n",
        "Chains = $(join(map(string, c.chains), ","))\n",
        "Samples per chain = $(length(c.range))\n",
        #,
        #"NNI moves = $(c.moves)\n"
    )
end

function indiscretesupport(c::AbstractChains, bounds::Tuple{Real,Real} = (0, Inf))
    nrows, nvars, nchains = size(c.value)
    result = Array{Bool}(undef, nvars * (nrows > 0))
    for i = 1:nvars
        result[i] = true
        for j = 1:nrows, k = 1:nchains
            x = c.value[j, i, k]
            if !isinteger(x) || x < bounds[1] || x > bounds[2]
                result[i] = false
                break
            end
        end
    end
    result
end

function link(c::AbstractChains)
    cc = copy(c.value)
    for j = 1:length(c.names)
        x = cc[:, j, :]
        if minimum(x) > 0.0
            cc[:, j, :] = maximum(x) < 1.0 ? logit.(x) : log.(x)
        end
    end
    cc
end

max_stars_repo_name: erathorn/MCPhylo

max_issues_repo_name: erathorn/MCPhylo

max_forks_repo_name: erathorn/MCPhylo


text:
# Test that the way we specify boundary conditions works as expected
using MPI
using OrderedCollections
using StaticArrays
using Statistics
using Test

using CLIMAParameters
struct EarthParameterSet <: AbstractEarthParameterSet end
const param_set = EarthParameterSet()

using ClimateMachine
using ClimateMachine.Land
using ClimateMachine.Land.SoilWaterParameterizations
using ClimateMachine.Mesh.Topologies
using ClimateMachine.Mesh.Grids
using ClimateMachine.DGMethods
using ClimateMachine.DGMethods.NumericalFluxes
using ClimateMachine.DGMethods: BalanceLaw, LocalGeometry
using ClimateMachine.MPIStateArrays
using ClimateMachine.GenericCallbacks
using ClimateMachine.ODESolvers
using ClimateMachine.VariableTemplates
using ClimateMachine.SingleStackUtils
using ClimateMachine.BalanceLaws:
    BalanceLaw, Prognostic, Auxiliary, Gradient, GradientFlux, vars_state

@testset "Boundary condition functions" begin
    ClimateMachine.init()

    FT = Float64

    function init_soil_water!(land, state, aux, localgeo, time)
        myfloat = eltype(state)
        state.soil.water.ϑ_l = myfloat(land.soil.water.initialϑ_l(aux))
        state.soil.water.θ_i = myfloat(land.soil.water.initialθ_i(aux))
    end

    soil_param_functions =
        SoilParamFunctions{FT}(porosity = 0.75, Ksat = 1e-7, S_s = 1e-3)
    bottom_flux_amplitude = FT(-3.0)
    f = FT(pi * 2.0 / 300.0)
    # nota bene: the flux is -K∇h
    bottom_flux =
        (aux, t) -> bottom_flux_amplitude * sin(f * t) * aux.soil.water.K
    surface_state = (aux, t) -> eltype(aux)(0.2)
    ϑ_l0 = (aux) -> eltype(aux)(0.2)
    bc = LandDomainBC(
        bottom_bc = LandComponentBC(soil_water = Neumann(bottom_flux)),
        surface_bc = LandComponentBC(soil_water = Dirichlet(surface_state)),
    )
    soil_water_model = SoilWaterModel(FT; initialϑ_l = ϑ_l0)
    soil_heat_model = PrescribedTemperatureModel()

    m_soil = SoilModel(soil_param_functions, soil_water_model, soil_heat_model)
    sources = ()
    m = LandModel(
        param_set,
        m_soil;
        boundary_conditions = bc,
        source = sources,
        init_state_prognostic = init_soil_water!,
    )

    N_poly = 5
    nelem_vert = 50

    # Specify the domain boundaries
    zmax = FT(0)
    zmin = FT(-1)

    driver_config = ClimateMachine.SingleStackConfiguration(
        "LandModel",
        N_poly,
        nelem_vert,
        zmax,
        param_set,
        m;
        zmin = zmin,
        numerical_flux_first_order = CentralNumericalFluxFirstOrder(),
    )

    t0 = FT(0)
    timeend = FT(300)
    dt = FT(0.05)

    solver_config = ClimateMachine.SolverConfiguration(
        t0,
        timeend,
        driver_config,
        ode_dt = dt,
    )
    mygrid = solver_config.dg.grid
    Q = solver_config.Q
    aux = solver_config.dg.state_auxiliary
    grads = solver_config.dg.state_gradient_flux
    K∇h_vert_ind =
        varsindex(vars_state(m, GradientFlux(), FT), :soil, :water)[3]
    K_ind = varsindex(vars_state(m, Auxiliary(), FT), :soil, :water, :K)
    n_outputs = 30

    every_x_simulation_time = ceil(Int, timeend / n_outputs)

    dons_arr = Dict([k => Dict() for k in 1:n_outputs]...)

    iostep = [1]
    callback = GenericCallbacks.EveryXSimulationTime(
        every_x_simulation_time,
    ) do (init = false)
        t = ODESolvers.gettime(solver_config.solver)
        K = aux[:, K_ind, :]
        K∇h_vert = grads[:, K∇h_vert_ind, :]
        all_vars = Dict{String, Array}(
            "t" => [t],
            "K" => K,
            "K∇h_vert" => K∇h_vert,
        )
        dons_arr[iostep[1]] = all_vars
        iostep[1] += 1
        nothing
    end

    ClimateMachine.invoke!(solver_config; user_callbacks = (callback,))
    t = ODESolvers.gettime(solver_config.solver)
    K = aux[:, K_ind, :]
    K∇h_vert = grads[:, K∇h_vert_ind, :]
    all_vars = Dict{String, Array}("t" => [t], "K" => K, "K∇h_vert" => K∇h_vert)
    dons_arr[n_outputs] = all_vars

    computed_bottom_∇h =
        [dons_arr[k]["K∇h_vert"][1] for k in 1:n_outputs] ./ [dons_arr[k]["K"][1] for k in 1:n_outputs]

    t = [dons_arr[k]["t"][1] for k in 1:n_outputs]
    # we need a -1 out in front here because the flux BC is on -K∇h
    prescribed_bottom_∇h = t -> FT(-1) * FT(-3.0 * sin(pi * 2.0 * t / 300.0))

    MSE = mean((prescribed_bottom_∇h.(t) .- computed_bottom_∇h) .^ 2.0)
    @test MSE < 1e-7
end


text:
function _random_invertible_matrix(n, B)
  A = identity_matrix(FlintZZ, n)
  if n == 1
    return A
  end
  for k in 1:10
    i = rand(1:(n - 1))
    j = rand((i + 1):n)
    A[i, j] += rand(B)
  end
  @assert det(A) == 1
  C = identity_matrix(FlintZZ, n)
  for k in 1:10
    i = rand(1:(n - 1))
    j = rand((i + 1):n)
    C[i, j] += rand(B)
  end
  @assert det(C) == 1
  return transpose(C) * A
end

const lattices_and_aut_order = [
(([[2]]), 2),
# 2
(([[1, 0], [0, 2]]), 4),
(([[2, -1], [-1, 2]]), 12),
# 3
(([[2, 1, 0],
   [1, 2, 0],
   [0, 0, 26]]), 24),
# 4
(([[1, 0, 0, 0],
   [0, 2, -1, 1],
   [0, -1, 3, -1],
   [0, 1, -1, 3]]), 16),
# 5
(([[2, 1, 1, 1, -1],
   [1, 2, 1, 1, 0],
   [1, 1, 2, 1, -1],
   [1, 1, 1, 2, -1],
   [-1, 0, -1, -1, 2]]), 3840),
(([[1, 0, 0, 0, 0],
   [0, 1, 0, 0, 0],
   [0, 0, 1, 0, 0],
   [0, 0, 0, 2, 1],
   [0, 0, 0, 1, 3]]), 192),
# 6
(([[2, -1, 0, 0, 0, 0],
   [-1, 2, -1, 0, 0, 0],
   [0, -1, 2, -1, 0, -1],
   [0, 0, -1, 2, -1, 0],
   [0, 0, 0, -1, 2, 0],
   [0, 0, -1, 0, 0, 2]]), 103680),
# 6
(([[1, 0, 0, 0, 0, 0],
   [0, 1, 0, 0, 0, 0],
   [0, 0, 2, 1, 0, 1],
   [0, 0, 1, 3, 1, 1],
   [0, 0, 0, 1, 2, 1],
   [0, 0, 1, 1, 1, 3]]), 512),
#(([[8, 4, 4, 4, 4, 4, 4, 2, 4, 4, 4, 2, 4, 2, 2, 2, 4, 2, 2, 2, 0, 0, 0, -3],
#   [4, 4, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 1, 1, 2, 1, 1, 2, 1, 0, 0, -1],
#   [4, 2, 4, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 1, 2, 1, 2, 2, 1, 1, 1, 0, 0, -1],
#   [4, 2, 2, 4, 2, 2, 2, 2, 2, 2, 2, 2, 2, 1, 1, 2, 2, 1, 2, 1, 1, 0, 0, -1],
#   [4, 2, 2, 2, 4, 2, 2, 2, 2, 2, 2, 1, 2, 2, 2, 2, 2, 2, 2, 2, 1, 0, 0, -1],
#   [4, 2, 2, 2, 2, 4, 2, 2, 2, 2, 2, 1, 2, 2, 1, 1, 2, 1, 2, 1, 0, 0, 0, -1],
#   [4, 2, 2, 2, 2, 2, 4, 2, 2, 2, 2, 1, 2, 1, 2, 1, 2, 1, 1, 2, 0, 0, 0, -1],
#   [2, 2, 2, 2, 2, 2, 2, 4, 1, 1, 1, 2, 1, 2, 2, 2, 1, 2, 2, 2, 2, 0, 0, 1],
#   [4, 2, 2, 2, 2, 2, 2, 1, 4, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 1, 1, 1, -1],
#   [4, 2, 2, 2, 2, 2, 2, 1, 2, 4, 2, 2, 2, 2, 1, 1, 2, 2, 1, 1, 0, 1, 0, -1],
#   [4, 2, 2, 2, 2, 2, 2, 1, 2, 2, 4, 2, 2, 1, 2, 1, 2, 1, 2, 1, 0, 0, 1, -1],
#   [2, 2, 2, 2, 1, 1, 1, 2, 2, 2, 2, 4, 1, 2, 2, 2, 1, 2, 2, 2, 2, 1, 1, 1],
#   [4, 2, 2, 2, 2, 2, 2, 1, 2, 2, 2, 1, 4, 2, 2, 2, 2, 1, 1, 1, 1, 1, 1, -1],
#   [2, 2, 1, 1, 2, 2, 1, 2, 2, 2, 1, 2, 2, 4, 2, 2, 1, 2, 2, 2, 2, 2, 1, 1],
#   [2, 1, 2, 1, 2, 1, 2, 2, 2, 1, 2, 2, 2, 2, 4, 2, 1, 2, 2, 2, 2, 1, 2, 1],
#   [2, 1, 1, 2, 2, 1, 1, 2, 2, 1, 1, 2, 2, 2, 2, 4, 1, 2, 2, 2, 2, 1, 1, 1],
#   [4, 2, 2, 2, 2, 2, 2, 1, 2, 2, 2, 1, 2, 1, 1, 1, 4, 2, 2, 2, 1, 1, 1, -1],
#   [2, 1, 2, 1, 2, 1, 1, 2, 2, 2, 1, 2, 1, 2, 2, 2, 2, 4, 2, 2, 2, 2, 1, 1],
#   [2, 1, 1, 2, 2, 2, 1, 2, 2, 1, 2, 2, 1, 2, 2, 2, 2, 2, 4, 2, 2, 1, 2, 1],
#   [2, 2, 1, 1, 2, 1, 2, 2, 2, 1, 1, 2, 1, 2, 2, 2, 2, 2, 2, 4, 2, 1, 1, 1],
#   [0, 1, 1, 1, 1, 0, 0, 2, 1, 0, 0, 2, 1, 2, 2, 2, 1, 2, 2, 2, 4, 2, 2, 2],
#   [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 1, 1, 2, 1, 1, 1, 2, 1, 1, 2, 4, 2, 2],
#   [0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 1, 1, 1, 1, 2, 1, 1, 1, 2, 1, 2, 2, 4, 2],
#   [-3, -1, -1, -1, -1, -1, -1, 1, -1, -1, -1, 1, -1, 1, 1, 1, -1, 1, 1, 1, 2,
#    2, 2, 4]]), 8315553613086720000)]
(([[2, -1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
   [-1, 2, -1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
   [0, -1, 2, -1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
   [0, 0, -1, 2, -1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
   [0, 0, 0, -1, 2, -1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
   [0, 0, 0, 0, -1, 2, -1, 0, 0, 0, 0, 0, 0, 0, 0, 0],
   [0, 0, 0, 0, 0, -1, 2, -1, 0, 0, 0, 0, 0, 0, 0, 0],
   [0, 0, 0, 0, 0, 0, -1, 2, -1, 0, 0, 0, 0, 0, 0, 0],
   [0, 0, 0, 0, 0, 0, 0, -1, 2, -1, 0, 0, 0, 0, 0, 0],
   [0, 0, 0, 0, 0, 0, 0, 0, -1, 2, -1, 0, 0, 0, 0, 0],
   [0, 0, 0, 0, 0, 0, 0, 0, 0, -1, 2, -1, 0, 0, 0, 0],
   [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, -1, 2, -1, 0, 0, 0],
   [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, -1, 2, -1, 0, 0],
   [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, -1, 2, -1, 0],
   [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, -1, 2, -1],
   [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, -1, 2]]), 711374856192000)]

function test_automorphisms(L, G, ambient_representation)
  if ambient_representation
    @test all(g * gram_matrix(ambient_space(L)) * transpose(g)  == gram_matrix(ambient_space(L)) for g in G)
  else
    @test all(g * gram_matrix(L) * transpose(g) == gram_matrix(L) for g in G)
  end
end

@testset "Z-lattice automorphisms" begin
  for (m, o) in lattices_and_aut_order
    n = length(m[1])
    G = matrix(FlintZZ, n, n, reduce(vcat, m))
    L = Zlattice(gram = G)
    Ge = automorphism_group_generators(L, ambient_representation = true)
    test_automorphisms(L, Ge, true)
    Ge = automorphism_group_generators(L, ambient_representation = false)
    test_automorphisms(L, Ge, false)
    @test automorphism_group_order(L) == o
  end

  D = lattice_database()
  l = number_of_lattices(D)
  for i in 1:min(l, 200)
    L = lattice(D, i)
    L = Zlattice(gram = gram_matrix(L)) # to avoid caching
    @test automorphism_group_order(L) == lattice_automorphism_group_order(D, i)
  end
end

@testset "Z-lattice isomorphisms" begin
  for (m, o) in lattices_and_aut_order
    n = length(m[1])
    G = matrix(FlintZZ, n, n, reduce(vcat, m))
    L = Zlattice(gram = G)
    X = _random_invertible_matrix(n, -3:3)
    @assert abs(det(X)) == 1
    L2 = Zlattice(gram = X * G * X')
    b, T = isisometric(L, L2, ambient_representation = false)
    @test b
    @test T * gram_matrix(L2) * T' == gram_matrix(L)
    L2 = Zlattice(X, gram = G)
    b, T = isisometric(L, L2, ambient_representation = false)
    @test b
    @test T * gram_matrix(L2) * T' == gram_matrix(L)
    b, T = isisometric(L, L2, ambient_representation = true)
    @test b
    @test T * gram_matrix(ambient_space(L2)) * T' ==
            gram_matrix(ambient_space(L))
  end

  D = lattice_database()
  l = number_of_lattices(D)
  for i in 1:min(l, 200)
    L = lattice(D, i)
    L = Zlattice(gram = gram_matrix(L)) # to avoid caching
    n = rank(L)
    X = change_base_ring(FlintQQ, _random_invertible_matrix(n, -3:3))
    @assert abs(det(X)) == 1
    L2 = Zlattice(gram = X * gram_matrix(L) * X')
    b, T = isisometric(L, L2, ambient_representation = false)
    @test b
    @test T * gram_matrix(L2) * T' == gram_matrix(L)
  end
end