text:
# Testing a normalizing flow for seismic denosing
# Authors: Ali Siahkoohi, alisk@gatech.edu
#          Gabrio Rizzuti, rizzuti.gabrio@gatech.edu
#          Philipp Witte, pwitte3@gatech.edu
# Date: September 2020
"""
Example:
julia test_warm-start_unsupervised_hint.jl
"""

using DrWatson
@quickactivate :FastApproximateInference

using InvertibleNetworks
using JLD
using JLD2
using Random
using Statistics
using ArgParse
using BSON
using Printf
using PyPlot
using Seaborn
using Flux: gpu, cpu
using Flux.Data: DataLoader

set_style("whitegrid")
rc("font", family="serif", size=13)
font_prop = matplotlib.font_manager.FontProperties(
    family="serif",
    style="normal",
    size=13
)
sfmt=matplotlib.ticker.ScalarFormatter(useMathText=true)
sfmt.set_powerlimits((0, 0))

s = ArgParseSettings()
@add_arg_table! s begin
    "--max_epoch"
        help = "Maximum # of epochs"
        arg_type = Int
        default = 15
    "--epoch"
        help = "Epoch to load net params"
        arg_type = Int
        default = 10
    "--lr"
        help = "starting learning rate"
        arg_type = Float32
        default = 1f-4
    "--lr_step"
        help = "lr scheduler applied at evey lr_step epoch"
        arg_type = Int
        default = 5
    "--batchsize"
        help = "batch size"
        arg_type = Int
        default = 4
    "--n_hidden"
        help = "# hidden channels"
        arg_type = Int
        default = 64
    "--depth"
        help = "depth of the network"
        arg_type = Int
        default = 12
    "--sigma"
        help = "noise std"
        arg_type = Float32
        default = 2f-1
    "--sim_name"
        help = "simulation name"
        arg_type = String
        default = "warm-start-unsupervised-hint-seismic-denoising-256x256-0.3CS"
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"]
sigma = parsed_args["sigma"]
sim_name = parsed_args["sim_name"]

if parsed_args["epoch"] == -1
    parsed_args["epoch"] = parsed_args["max_epoch"]
end
epoch = parsed_args["epoch"]

# Loading the experiment—only network weights and training loss
Params, fval, exp_path = load_experiment(parsed_args; return_path=true)

Y_fixed = wload(exp_path)["Array(Y_obs)"]
X_fixed = wload(exp_path)["X_true"]
Zy_fixed = wload(exp_path)["Array(Zy_obs)"]
Zx = wload(exp_path)["Array(Z_train)"]
A = wload(exp_path)["A"]
pretrained_args = wload(exp_path)["pretrained_args"]

nx, ny, n_in = size(X_fixed)[1:3]

CH = NetworkConditionalHINT(n_in, n_hidden, depth)
CH_rev = NetworkConditionalHINT(n_in, n_hidden, depth)

put_params!(CH_rev, convert(Array{Any,1}, Params))
CH_rev = CH_rev |> gpu

Params0, fval0 = load_experiment(pretrained_args)
put_params!(CH, convert(Array{Any,1}, Params0))

# Testing data
test_size = size(Zx, 4)

# Now select single fixed sample from all Ys
Y_fixed = Y_fixed |> gpu
Zy_fixed = Zy_fixed |> gpu

CH_rev = reverse(CH_rev)

# Draw new Zx, while keeping Zy fixed
test_size = 1000
X_post0 = zeros(Float32, nx, ny, n_in, test_size)
X_post = zeros(Float32, nx, ny, n_in, test_size)

CH = CH |> gpu

test_loader = DataLoader(randn(Float32, nx, ny, n_in, test_size) |> gpu, batchsize=8, shuffle=false)

for (itr, zx) in enumerate(test_loader)
    counter = (itr - 1)*size(zx)[4] + 1

    zx = zx |> gpu

    X_post0[:, :, :, counter:(counter+size(zx)[4]-1)] = (CH.inverse(zx, repeat(Zy_fixed, 1, 1, 1, size(zx, 4)))[1] |> cpu)
    X_post[:, :, :, counter:(counter+size(zx)[4]-1)] = (CH_rev.forward(zx, repeat(Zy_fixed, 1, 1, 1, size(zx, 4)))[1] |> cpu)
    @printf("Sampling: [%d/%d]\r", counter, test_size)
end

X_post = wavelet_unsqueeze(X_post)
X_post0 = wavelet_unsqueeze(X_post0)

X_fixed = wavelet_unsqueeze(X_fixed)
Y_fixed = wavelet_unsqueeze(Y_fixed)

# Some stats
X_post_mean = mean(X_post; dims=4)
X_post_std = std(X_post; dims=4)
X_post_mean0 = mean(X_post0; dims=4)
X_post_std0 = std(X_post0; dims=4)

save_dict = @strdict max_epoch epoch lr lr_step batchsize n_hidden depth sigma
save_path = plotsdir(sim_name, savename(save_dict; digits=6))

# Plot the conditional mean estimate
fig = figure("x_cm", figsize=(5, 5))
imshow(
    X_fixed[:, :, 1, 1], vmin=-4.0, vmax=4.0, aspect=1, cmap="Greys", resample=true,
    interpolation="lanczos", filterrad=1
)
title("True model")
colorbar(fraction=0.047, pad=0.01, format=sfmt)
grid(false)
ax = plt.gca()
ax.axes.xaxis.set_visible(false)
ax.axes.yaxis.set_visible(false)
safesave(joinpath(save_path, "true_model.png"), fig)
close(fig)

fig = figure("x_cm", figsize=(5, 5))
imshow(
    Y_fixed[:, :, 1, 1], vmin=-4.0, vmax=4.0, aspect=1, cmap="Greys", resample=true,
    interpolation="lanczos", filterrad=1
)
title("Observed data")
colorbar(fraction=0.047, pad=0.01, format=sfmt)
grid(false)
ax = plt.gca()
ax.axes.xaxis.set_visible(false)
ax.axes.yaxis.set_visible(false)
safesave(joinpath(save_path, "observed_data.png"), fig)
close(fig)

# Plot the conditional mean estimate
fig = figure("x_cm", figsize=(5, 5))
imshow(
    X_post_mean[:, :, 1, 1], vmin=-4.0, vmax=4.0, aspect=1, cmap="Greys", resample=true,
    interpolation="lanczos", filterrad=1
)
title("Conditional mean — Equation (6)")
colorbar(fraction=0.047, pad=0.01, format=sfmt)
grid(false)
ax = plt.gca()
ax.axes.xaxis.set_visible(false)
ax.axes.yaxis.set_visible(false)
safesave(joinpath(save_path, "conditional_mean.png"), fig)
close(fig)

# Plot the conditional mean estimate
fig = figure("x_cm", figsize=(5, 5))
imshow(
    X_post_mean0[:, :, 1, 1], vmin=-4.0, vmax=4.0, aspect=1, cmap="Greys", resample=true,
    interpolation="lanczos", filterrad=1
)
title("Conditional mean — Equation (3)")
colorbar(fraction=0.047, pad=0.01, format=sfmt)
grid(false)
ax = plt.gca()
ax.axes.xaxis.set_visible(false)
ax.axes.yaxis.set_visible(false)
safesave(joinpath(save_path, "pretrained-conditional_mean.png"), fig)
close(fig)

# Plot the pointwise standard deviation
fig = figure("x_std", figsize=(5, 5))
imshow(
    X_post_std0[:, :, 1, 1], vmin=.30,
    vmax=.42, aspect=1, cmap="OrRd", resample=true,
    interpolation="lanczos", filterrad=1
)
title("STD — Equation (3)")
colorbar(fraction=0.047, pad=0.01, format=sfmt)
grid(false)
ax = plt.gca()
ax.axes.xaxis.set_visible(false)
ax.axes.yaxis.set_visible(false)
safesave(joinpath(save_path, "pretrained-pointwise_std.png"), fig)
close(fig)

# Plot the pointwise standard deviation
fig = figure("x_std", figsize=(5, 5))
imshow(
    X_post_std[:, :, 1, 1], vmin=.13,
    vmax=.19, aspect=1, cmap="OrRd", resample=true,
    interpolation="lanczos", filterrad=1
)
title("STD — Equation (6)")
colorbar(fraction=0.047, pad=0.01, format=sfmt)
grid(false)
ax = plt.gca()
ax.axes.xaxis.set_visible(false)
ax.axes.yaxis.set_visible(false)
safesave(joinpath(save_path, "pointwise_std.png"), fig)
close(fig)

# Plot the conditional mean estimate
fig = figure("err", figsize=(5, 5))
imshow(
    2f0.*(X_fixed[:, :, 1, 1] - X_post_mean[:, :, 1, 1])./(abs.(X_fixed[:, :, 1, 1])  + abs.(X_post_mean[:, :, 1, 1])), vmin=-2.0, vmax=2.0, aspect=1, cmap="twilight_shifted", resample=true,
    interpolation="lanczos", filterrad=1
)
title("Relative error — Equation (6)")
colorbar(fraction=0.047, pad=0.01, format=sfmt)
grid(false)
ax = plt.gca()
ax.axes.xaxis.set_visible(false)
ax.axes.yaxis.set_visible(false)
safesave(joinpath(save_path, "conditional_mean_err.png"), fig)
close(fig)

# Plot the conditional mean estimate
fig = figure("err", figsize=(5, 5))
imshow(
    2f0.*(X_fixed[:, :, 1, 1] - X_post_mean0[:, :, 1, 1])./(abs.(X_fixed[:, :, 1, 1])  + abs.(X_post_mean0[:, :, 1, 1])), vmin=-2.0, vmax=2.0, aspect=1, cmap="twilight_shifted", resample=true,
    interpolation="lanczos", filterrad=1
)
title("Relative error — Equation (3)")
colorbar(fraction=0.047, pad=0.01, format=sfmt)
grid(false)
ax = plt.gca()
ax.axes.xaxis.set_visible(false)
ax.axes.yaxis.set_visible(false)
safesave(joinpath(save_path, "pretrained-conditional_mean_err.png"), fig)
close(fig)


text:
using YaoExtensions
using Yao, Test

@testset "Mod" begin
    @test_throws AssertionError Mod{4}(4,10)
    @test_throws AssertionError Mod{2}(3,10)
    m = Mod{4}(3,10)
    @test mat(m) ≈ applymatrix(m)
    @test isunitary(m)
    @test isunitary(mat(m))
    @test m' == Mod{4}(7,10)
end

@testset "KMod" begin
    @test_throws AssertionError KMod{6, 2}(4,10)
    @test_throws AssertionError KMod{4, 2}(3,10)
    m = KMod{6, 2}(3,10)
    @test mat(m) ≈ applymatrix(m)
    @test isunitary(m)
    @test isunitary(mat(m))
    @test m' == KMod{6, 2}(7,10)
end

ext: jl


text:
module BasicReservoirComputing
"""Basic implementation of a Reservoir computer

Author: Jason Platt
Email: 1763platt@gmail.com
"""

module RC

export rc, train_RC, test_RC, forecast_RC, TLM!, Global_LEs

using SparseArrays
using LinearAlgebra
using Random
using ArnoldiMethod
using DynamicalSystems
using Parameters
using Distributions

if (Threads.nthreads() > 1) & !Sys.isapple()
    using MKL
end

struct rc
    """Structure holding the rc parameters and matrixes

    Args:
        D::Int : Input system Dimension
        N::Int : Reservoir Dimension
        SR::Float : Spectral Radius
        ρA::Float : Density of A
        α::Float : Leak Rate
        σ::Float : Input Strength
        random_state::Int : Random seed
        σb::Float : Bias
        β::Float : Tikhonov regularization
        (params)::NamedTuple : Named tuple with fields SR, α, σ, σb, β
    """
    D::Int # Input System Dimension
    N::Int # Reservoir Dimension
    SR::Float64 # Spectral Radius of A
    ρA::Float64 # Density of A
    α::Float64 # Leak Rate
    σ::Float64 # Input Strength
    σb::Float64 # Input Bias
    β::Float64 # Tikhonov Regularization
    A::SparseMatrixCSC{Float64, Int64} # Adjacency Matrix
    Win::Array{Float64,2} # Input Matrix
    Wout::Array{Float64,2} # Output Matrix
    W::Array{Float64,2} # A + Win Wout
    rng::MersenneTwister # Random Number Generator
    climstd::Vector{Float64} # climatological standard deviation
    Δt::Float64 # Time step of the data/RC
    function rc(D, N, ρA, Δt, SR, α, σ, σb, β;
                 random_state=111111)
        @assert SR > 0
        # Random number generator
        rng = MersenneTwister(random_state)

        # Initialize the connectivity matrix/Win and allocate memory
        A, Win, Wout, W, climstd = initialize_weights(D, N, SR, ρA, σ, rng)
        new(D, N, SR, ρA, α, σ, σb, β, A, Win, Wout, W, rng, climstd, Δt)
    end
    function rc(D, N, ρA, Δt, params::NamedTuple; random_state=111111)
        rng = MersenneTwister(random_state)
        @unpack SR, α, σ, σb, β = params
        @assert SR > 0
        A, Win, Wout, W, climstd = initialize_weights(D, N, SR, ρA, σ, rng)
        new(D, N, SR, ρA, α, σ, σb, β, A, Win, Wout, W, rng, climstd, Δt)
    end
end

function initialize_weights(D, N, SR, ρA, σ, rng)
	"""Initialize the adjacency matrix, Win and allocate memory"""
    A = get_connection_matrix(N, ρA, SR, rng)
    Win = rand(rng, Uniform(-σ, σ), N, D)
    Wout = Array{Float64}(undef, (D, N))
    W = Array{Float64}(undef, (N, N))
    climstd = Vector{Float64}(undef, D)
    return A, Win, Wout, W, climstd
end

function get_connection_matrix(N, ρA, SR, rng)
	"""Initialize A and rescale maximal eigenvalue"""
    uni = Uniform(-1.0, 1.0)
    rf(rng, N) = rand(rng, uni, N)
    A = sprand(rng, N, N, ρA, rf)
    λs, _ = partialeigen(partialschur(A, nev=1, tol=1e-3, which=LM())[1])
    maxeig = abs(λs[1])
    A = A.*(SR/maxeig)
    return A
end

function train_RC(rc::rc, u; spinup = 100, readout = false)
    """Train the rc

    This function sets the matrix rc.Wout
    u ∈ D×T
    r ∈ N×T

    Args:
        rc : reservoir computer
        u : Array of data of shape D×Time
        spinup : Amount of data to use for spinup
    Returns:
        if readout=true: Wout*r, the readout of the rc training data
        else: return the last state of r, can be used to predict
              forward from the training data
    """
    @assert size(u)[1] == rc.D
    r = generate(rc, u)
    @views compute_Wout(rc, r[:, 1+spinup:end], u[:, 1+spinup:end])
    rc.W[:, :] .= rc.Win*rc.Wout .+ rc.A
    rc.climstd[:] .= vec(std(u, dims=2))
    if readout == true return rc.Wout*r end
    return r[:, end]
end

function forecast_RC(rc::rc, nsteps; uspin=nothing, r0 = nothing)
    """Forecast nsteps forward in time from the end of uspin or from rc state r0

    Make sure to train the RC before forecasting.  Requires either uspin or r0.
    If both are provided will use uspin to set r0.
    """
    if !isnothing(uspin)
    	@assert size(uspin)[1] == rc.D
        rspin = generate(rc, uspin)
        r0 = rspin[:, end]
    end
    @assert !isnothing(r0)

    rfc = Array{Float64}(undef, (rc.N, nsteps))
    rfc[:, 1] = r0

    for t in 1:nsteps-1
        @views @inbounds @fastmath auto_rc!(rfc[:, t+1], rc, rfc[:, t])
    end
    return rc.Wout*rfc
end

function test_RC(rc::rc, test_arr; ϵ=0.3, spinup=100)
    """Test the RC on all the forecasts in test_arr

    Args:
        rc : reservoir computer
        test_arr : array of forecasts DxT

    Returns:
        valid prediction time of all the tests
    """
    vpt = Float64[]
    for (i,test) in enumerate(test_arr)
    	@assert size(test)[1] == rc.D
        utest = test[:, spinup:end]
        nsteps = size(utest)[2]
        upred = forecast_RC(rc, nsteps, uspin = test[:, 1:spinup])

        diff = sqrt.(mean(((utest .- upred)./rc.climstd).^2, dims=1))

        thresh = vec(diff .> ϵ)
        idx = findfirst(thresh)
        if isnothing(idx) idx = length(thresh) end
        push!(vpt, idx*rc.Δt)
    end
    return vpt
end

function auto_rc!(rtp1, rc::rc, rt)
    rtp1[:] .= rc.α.*tanh.(rc.A*rt .+ rc.Win*rc.Wout*rt .+ rc.σb).+(1 .- rc.α).*rt
end

function driven_rc!(rtp1, rc::rc, rt, ut)
    rtp1[:] .= rc.α.*tanh.(rc.A*rt .+ rc.Win*ut .+ rc.σb).+(1 .- rc.α).*rt
end

function generate(rc::rc, u)
    """Generate hidden RC states on input data u"""
    T = size(u)[2]
    r = Array{Float64}(undef, (rc.N, T))
    r[:, 1] = zeros(rc.N)
    for t in 2:T
        @views @inbounds @fastmath driven_rc!(r[:, t], rc, r[:, t-1], u[:, t-1])
    end
    return r
end

function compute_Wout(rc::rc, r, u)
	"""Linear regression to find Wout"""
    try
        rc.Wout[:, :] .= (Symmetric(r*r'+rc.β*I)\(r*u'))'
    catch e
        println("Training Failed")
        println(e)
    end
end

function TLM(rc, r)
    """The tangent linear model maps how perturbations
    to the model evolve in time.
    """
    t0 = rc.α.*(1 .- tanh.(rc.W*r .+ rc.σb).^2)
    Q = t0.*rc.W
    Q[diagind(Q)].+=(1 - rc.α)
    return Q
end

function TLM!(Q, rc, r)
    """The tangent linear model maps how perturbations
    to the model evolve in time.
    """
    t0 = rc.α.*(1 .- tanh.(rc.W*r .+ rc.σb).^2)
    Q[:, :] .= t0.*rc.W
    Q[diagind(Q)].+=(1 - rc.α)
end

function Global_LEs(rc::rc, u, nspin; num_exp = 1, renorm_steps=10)
    """Find the Lyapunov Exponents of the rc

    Args:
        rc : reservoir computer
        u : input data, must be enough for the estimate. length(u)>renorm_steps*num_evals+nspin
        nspin : spinup steps to use
        num_exp : number of exponents to compute out of N
        renorm_steps : How many steps before renormalizing using QR decomposition
        num_evals : Number of renormalization steps to conduct
    Returns:
        Lyapunov Exponents : Array
    """
    function f!(xnew, x, p, t)
        driven_rc!(xnew, rc, x, @view(p[1][:, t+1]))
    end
    function fJac!(Jnew, x, p, t)
        TLM!(Jnew, rc, x)
    end
    num_evals = trunc(Int, (size(u)[2]-nspin)/renorm_steps)
    DS = DiscreteDynamicalSystem(f!, zeros(rc.N), [u], fJac!)
    return lyapunovspectrum(DS, num_evals, num_exp, Δt = renorm_steps, Ttr = nspin)/rc.Δt
end
end

module Opt

export opt, opt_lower_bounds, opt_upper_bounds

using ..RC
using LinearAlgebra
using Parameters
using CMAEvolutionStrategy

if (Threads.nthreads() > 1) & !Sys.isapple()
    using MKL
end

opt_lower_bounds = @with_kw (SR_lb = 0.01, α_lb = 0.0, σ_lb = 1e-12, σb_lb = 0.0, β_lb = 1e-12)
opt_upper_bounds = @with_kw (SR_ub = 2.0, α_ub = 1.0, σ_ub = 1.0, σb_ub = 4.0, β_ub = 2.0)

struct opt
    """Optimization object

    Args:
        train_data, valid_data_list, spinup_steps, N, ρA
        lb::NamedTuple : use opt_lower_bounds to set lower bounds on params
        ub::NamedTuple : use opt_upped_bounds to set upper bounds on params
    """
    lb::Array{Float64, 1} # lower bounds
    ub::Array{Float64, 1} # upped bounds
    train_data::Array{Float64, 2} # Training data D×T
    valid_data_list::Array{AbstractArray{Float64,2},1} # List of validation data
    spinup_steps::Int64 # Number of spinup steps
    N::Int64 # rc size
    ρA::Float64 # density of A
    function opt(train_data, valid_data_list, spinup_steps, N, ρA;
                 lb = opt_lower_bounds(), ub = opt_upper_bounds())
        @unpack SR_lb, α_lb, σ_lb, σb_lb, β_lb = lb
        @unpack SR_ub, α_ub, σ_ub, σb_ub, β_ub = ub
        lb = [SR_lb, α_lb, log(σ_lb), σb_lb, log(β_lb)]
        ub = [SR_ub, α_ub, log(σ_ub), σb_ub, log(β_ub)]
        @assert lb < ub
        @assert ρA <= 1.0
        @assert ρA >= 0.0
        new(lb, ub, train_data, valid_data_list, spinup_steps, N, ρA)
    end
end

function (f::opt)(;maxtime=200, popsize = 15, multithread = false)
    """Call future forecasting optimization routine

    This routine optimizes the rc by looking at a number of long term forecasts
    of the data and comparing them.
    """
    if multithread
         @assert Threads.nthreads() > 1
    end
    function Loss(data, spinup_steps, rc::rc)
        """Loss function over 1 forecast"""
        uspin = @view data[:, 1:spinup_steps]
        utrue = @view data[:, f.spinup_steps:end]
        Ttrue = size(utrue)[2]
        upred = forecast_RC(rc, Ttrue; uspin=uspin)
        err = upred.-utrue
        weight = exp.(-(1:Ttrue)./Ttrue)
        return norm(err.*weight')
    end

    function obj(θ)
        """Objective function to optimize"""
        θ = untransform(θ)
        SR, α, logσ, σb, logβ = θ
        D, _ = size(f.train_data)
        reservoir = rc(D, f.N, f.ρA, 1.0, SR, α, exp(logσ), σb, exp(logβ))
        train_RC(reservoir, f.train_data)
        n_valid = length(f.valid_data_list)
        loss = zeros(n_valid)
        for i in 1:n_valid
            @inbounds @views loss[i] = Loss(f.valid_data_list[i], f.spinup_steps, reservoir)
        end
        return sum(loss)
    end

    function untransform(θ::AbstractArray{Float64, 1})
        return f.lb + (f.ub - f.lb).*θ/10
    end

    function transform(θ::Array{Float64, 1})
        θ[:] .= clamp.(10*(θ - f.lb)./(f.ub - f.lb), 0.0, 10.0)
    end

    x0 = (f.ub + f.lb)/2
    x0 = transform(x0)
    obj(x0)
    result = minimize(obj, x0, 4;
                      lower = zeros(5),
                      upper = 10.0*ones(5),
                      popsize = popsize,
                      multi_threading = multithread,
                      verbosity = 1,
                      seed = rand(UInt),
                      maxtime = maxtime)

    xopt = xbest(result)
    xopt = untransform(xopt)
    xopt = (SR = xopt[1], α = xopt[2], σ = exp(xopt[3]), σb = xopt[4], β = exp(xopt[5]))
    return xopt, fbest(result)
end
end
using .RC
using .Opt
export rc, train_RC, test_RC, forecast_RC, TLM!, Global_LEs, opt, opt_lower_bounds, opt_upper_bounds

end

max_stars_repo_name: japlatt/BasicReservoirComputing

max_issues_repo_name: japlatt/BasicReservoirComputing

max_forks_repo_name: japlatt/BasicReservoirComputing


text:
module Cpp

export @cpp

# Useful references:
# http://www.agner.org/optimize/calling_conventions.pdf
# http://mentorembedded.github.io/cxx-abi/abi.html#mangling

const SUBSTITUTION = [""; map(string, collect(0:9)); map(string, collect('A':'Z'))]  # FIXME more than 37 arguments?

# Allow calling of functions in C++ shared libraries.
# Usage: a = @cpp ccall((:mysymbol,mylib),...)
macro cpp(ex)
    # If you get "undefined symbol" errors, use nm or readelf to view
    # the names of the symbols in the library. Then use the resulting
    # information to help improve this macro!
    # Note: for global objects without class or namespace qualifiers
    # (e.g., global consts), you do not need to use this macro (it
    # will work with a plain ccall, even though it is a C++ library)
    msg = "@cpp requires a ccall((:mysymbol, mylib),...)  expression"
    if !isa(ex,Expr) || ex.head != :ccall
        error(msg)
    end

    # Parse the library symbol's name
    exlib = ex.args[1]
    if !(isa(exlib,Expr) && exlib.head == :tuple)
        error(msg)
    end
    sym = exlib.args[1]
    fstr = string(eval(sym))
    #GNU3-4 ABI
    fstr = string("_Z",length(fstr),fstr)

    # Parse the arguments to ccall and construct the parameter type string
    exargtypes = ex.args[3]
    if exargtypes.head != :tuple
        error(msg)
    end
    exargs = exargtypes.args
    pstr = ""
    symtable = (:Void,:Bool,:Cchar,:Char,:ASCIIString,:Int,:Int8,:Uint8,:Int16,:Uint16,:Int32,:Cint,:Uint32,:Int64,:Uint64,:Float32,:Float64)
    # GNU3-4 ABI v.3 and v.4
    ptable =   ('v',  'b',  'c',   'w',  "Pc",        'i', 'a',  'h',   's',   't',    'i',   'i',  'j',    'l',   'm',    'f',     'd')
    msub = ASCIIString[]
    for iarg = 1:length(exargs)
        thisarg = exargs[iarg]
        thisargm = ""
        while isa(thisarg,Expr) && thisarg.head == :curly && thisarg.args[1] == :Ptr
            thisargm = string(thisargm,'P')
            thisarg = thisarg.args[2]
        end
        matched = false
        for isym = 1:length(symtable)
            if thisarg == symtable[isym]
                matched = true
                thisargm = string(thisargm, ptable[isym])
                # Cchar is a special notation just for name mangling,
                # convert back to :Int8
                if thisarg == :Cchar
                    ex.args[3].args[iarg] = :Int8
                end
                break
            end
        end
        if matched
            if length(thisargm) > 1
                # Use substitution
                idx = indexin([thisargm], msub)[1]
                if idx != 0
                    pstr = string(pstr, "S"*SUBSTITUTION[idx]*"_")
                else
                    push!(msub, thisargm)
                    pstr = string(pstr, thisargm)
                end
            else
                pstr = string(pstr, thisargm)
            end
        else
            println(thisarg)
            error("@cpp: argument not recognized")
        end
    end
    ex.args[1].args[1] = Expr(:quote, symbol(string(fstr,pstr)))
    ex.args[1].args[2] = esc(ex.args[1].args[2])
    for ival = 4:length(ex.args)
        ex.args[ival] = esc(ex.args[ival])
    end
    ex
end

end


text:
# CONVERSIONS
function Graph(D::Digraph{T}) where {T}
    G = Graph{T}()
    map(a -> add_arc!(G, a), arcs(D))
    return G
end

function Digraph(G::Graph{T}) where {T}
    D = Digraph{T}()
    map(a -> add_arc!(D, a), arcs(G))
    map(a -> add_arc!(D, reverse(a)), arcs(G))
    return D
end


text:
Network = ResNet

netparams = ResNetHP(
  num_filters=64,
  num_blocks=5,
  conv_kernel_size=(3,1),
  num_policy_head_filters=4,
  num_value_head_filters=32,
  batch_norm_momentum=0.3)

self_play = SelfPlayParams(
  num_games=2_000,
  reset_mcts_every=400,
  mcts=MctsParams(
    use_gpu=true,
    num_workers=64,
    num_iters_per_turn=320,
    cpuct=4,
    temperature=StepSchedule(1.),
    dirichlet_noise_ϵ=0))

arena = ArenaParams(
  num_games=150,
  reset_mcts_every=100,
  update_threshold=(2 * 0.58 - 1),
  mcts=MctsParams(self_play.mcts,
    temperature=StepSchedule(0.3),
    dirichlet_noise_ϵ=0.05))

learning = LearningParams(
  samples_weighing_policy=LOG_WEIGHT,
  batch_size=256,
  loss_computation_batch_size=1024,
  gc_every=nothing,
  learning_rate=1e-3,
  l2_regularization=1e-4,
  nonvalidity_penalty=1.,
  checkpoints=[1, 2, 4])

params = Params(
  arena=arena,
  self_play=self_play,
  learning=learning,
  num_iters=40,
  mem_buffer_size=PLSchedule(
    [      0,       20],
    [200_000, 2_000_000]))

benchmark = [
  Benchmark.Duel(
    Benchmark.Full(self_play.mcts),
    Benchmark.MctsRollouts(self_play.mcts),
    num_games=100)]


text:
export pad_constant, pad_repeat, pad_reflect, pad_zeros

"""
    pad_zeros(x, pad::Tuple; [dims])
    pad_zeros(x, pad::Int; [dims])

Pad the array `x` with zeros.
Equivalent to [`pad_constant`](@ref) with the constant equal to 0.
"""
pad_zeros(x::AbstractArray, pad; dims = :) =
  pad_constant(x, pad, 0; dims = dims)

"""
    pad_constant(x, pad::Tuple, val = 0; [dims = :])
    pad_constant(x, pad::Int, val = 0; [dims = :])

Pad the array `x` with the constant value `val`.

`pad` can be a tuple of integers.
If it is of some length `2 * length(dims)` that specifies the left and right padding size
for each of the dimensions in `dims` as `(l1, r1, ..., ln, rn)`.
If supplied with a tuple of length `length(dims)` instead, it applies symmetric padding.
If `dims` is not given, it defaults to all dimensions.

For integer `pad` input, it is applied on both sides
on every dimension in `dims`.

See also [`pad_zeros`](@ref), [`pad_reflect`](@ref) and [`pad_repeat`](@ref).

```jldoctest

julia> r = reshape(1:4, 2, 2)
2×2 reshape(::UnitRange{Int64}, 2, 2) with eltype Int64:
 1  3
 2  4

julia> pad_constant(r, (1, 2, 3, 4), 8)
5×9 Matrix{Int64}:
 8  8  8  8  8  8  8  8  8
 8  8  8  1  3  8  8  8  8
 8  8  8  2  4  8  8  8  8
 8  8  8  8  8  8  8  8  8
 8  8  8  8  8  8  8  8  8

julia> pad_constant(r, 1, 8)
5×9 Matrix{Int64}:
 8  8  8  8
 8  1  3  8
 8  2  4  8
 8  8  8  8

julia> r = reshape(1:27, 3, 3, 3)
3×3×3 reshape(::UnitRange{Int64}, 3, 3, 3) with eltype Int64:
[:, :, 1] =
 1  4  7
 2  5  8
 3  6  9

[:, :, 2] =
 10  13  16
 11  14  17
 12  15  18

[:, :, 3] =
 19  22  25
 20  23  26
 21  24  27

julia> pad_constant(r, (2,1), dims = 1) # assymetric padding
6×3×3 Array{Int64,3}:
[:, :, 1] =
 0  0  0
 0  0  0
 1  4  7
 2  5  8
 3  6  9
 0  0  0

[:, :, 2] =
  0   0   0
  0   0   0
 10  13  16
 11  14  17
 12  15  18
  0   0   0

[:, :, 3] =
  0   0   0
  0   0   0
 19  22  25
 20  23  26
 21  24  27
  0   0   0

julia> pad_constant(r, (2,1, 3), dims = (1,2)) # padding must always be either the same length as dims, or double it
ERROR: ArgumentError: Could not parse padding (2, 1, 3) and dims (1, 2)
Stacktrace:
[...]
```
"""
pad_constant(x::AbstractArray{T,N}, pad::Int, val = 0; dims = :) where {T,N} =
  pad_constant(x, gen_pad(pad, dims isa Colon ? dims : (dims...,), N), val)
pad_constant(x::AbstractArray{T,N}, pad::Tuple, val = 0; dims = :) where {T,N} =
  pad_constant(x, gen_pad(pad, dims isa Colon ? dims : (dims...,), N), val)

function pad_idx(pad, dims, N)
  is = zip( (2 .* dims) .- 1, (2 .* dims))
end

@inline tuplejoin(x) = x
@inline tuplejoin(x, y) = (x..., y...)
@inline tuplejoin(x, y, z...) = tuplejoin(tuplejoin(x, y), z...)

gen_pad(pad::Int, dims, N) = gen_pad(ntuple(_ -> pad, length(dims)), dims, N)
gen_pad(pad::Int, dims::Colon, N) = ntuple(_ -> (pad, pad), N)
gen_pad(pad, dims::Colon, N) = gen_pad(pad, ntuple(identity, N), N)
gen_pad(pad::Int, dims::Int, N) = gen_pad((pad,pad), (dims,), N)
function gen_pad(pad::NTuple{L,Int}, dims::NTuple{D,Int}, N) where {L,D}
  ntuple(N) do d
   if d in dims
     if L == D
       ix = findfirst(==(d), dims)
       (pad[ix], pad[ix])
     elseif L == 2D
       ix = findfirst(==(d), dims)
       (pad[2ix - 1], pad[2ix])
     else
       throw(ArgumentError("Could not parse padding $pad and dims $dims"))
     end
   else
     (0,0)
   end

  end
end

# Expects length(pad) == 2M
function pad_constant(x::AbstractArray{T,M}, pad::NTuple{N,Tuple{Int,Int}}, val = 0) where {T,M,N}
  sz, c = size_and_center(x, pad)
  res = fill!(similar(x, sz...), val)
  res[c...] = x
  res
end

function size_and_center(x, pad::NTuple{N,NTuple{2, Int}}) where N
  sz = ntuple(i -> pad[i][1] + pad[i][2], N) .+ size(x)
  center = broadcast((x,y) -> x .+ y, axes(x), ntuple(i -> pad[i][1], N))
  sz, center
end

function rrule(::typeof(pad_constant), x::AbstractArray{T,N},
               pad, val; dims = :) where {T,N}
  y = pad_constant(x, pad, val; dims = dims)
  function pad_constant_pullback(Δ)
    p = gen_pad(pad, dims, N)
    outsize, center = size_and_center(x, p)
    (NO_FIELDS, @thunk(Δ[center...]), NoTangent(), NoTangent(),)
  end
  return y, pad_constant_pullback
end

"""
    pad_repeat(x, pad::Tuple; [dims])
    pad_repeat(x, pad::Int; [dims])

Pad the array `x` repeating the values on the border.

`pad` can a tuple of integers `(l1, r1, ..., ln, rn)`
of some length `2n` that specifies the left and right padding size
for each of the dimensions in `dims`. If `dims` is not given,
it defaults to the first `n` dimensions.

For integer `pad` input instead, it is applied on both sides
on every dimension in `dims`. In this case, `dims`
defaults to the first `ndims(x)-2` dimensions
(i.e. excludes the channel and batch dimension).

See also [`pad_reflect`](@ref) and [`pad_constant`](@ref).

```jldoctest

julia> r = reshape(1:9, 3, 3)
3×3 reshape(::UnitRange{Int64}, 3, 3) with eltype Int64:
 1  4  7
 2  5  8
 3  6  9

julia> pad_repeat(r, (1,2,3,4))
6×10 Matrix{Int64}:
 1  1  1  1  4  7  7  7  7  7
 1  1  1  1  4  7  7  7  7  7
 2  2  2  2  5  8  8  8  8  8
 3  3  3  3  6  9  9  9  9  9
 3  3  3  3  6  9  9  9  9  9
 3  3  3  3  6  9  9  9  9  9
```
"""
function pad_repeat(x::AbstractArray, pad::NTuple{M,Int};
                    dims = 1:M÷2) where M
  length(dims) == M ÷ 2 ||
    throw(ArgumentError("The number of dims should be equal to the number of padding dimensions"))
  for (i, d) in enumerate(dims)
    x = pad_repeat(x, (pad[2i-1], pad[2i]); dims=d)
  end
  return x
end

function pad_repeat(x::AbstractArray{F,N}, pad::NTuple{2,Int};
                    dims::Int = 1) where {F,N}
  lpad, rpad = pad

  xlborder = selectdim(x, dims, 1:1)
  nrepl = ntuple(i -> i == dims ? lpad : 1, N)
  xl = repeat(xlborder, outer = nrepl)

  n = size(x, dims)
  xrborder = selectdim(x, dims, n:n)
  nrepr = ntuple(i -> i == dims ? rpad : 1, N)
  xr = repeat(xrborder, outer = nrepr)

  return cat(xl, x, xr, dims = dims)
end

"""
    pad_reflect(x, pad::Tuple; [dims])
    pad_reflect(x, pad::Int; [dims])

Pad the array `x` reflecting its values across the border.

`pad` can a tuple of integers `(l1, r1, ..., ln, rn)`
of some length `2n` that specifies the left and right padding size
for each of the dimensions in `dims`. If `dims` is not given,
it defaults to the first `n` dimensions.

For integer `pad` input instead, it is applied on both sides
on every dimension in `dims`. In this case, `dims`
defaults to the first `ndims(x)-2` dimensions
(i.e. excludes the channel and batch dimension).

See also [`pad_repeat`](@ref) and [`pad_constant`](@ref).

```jldoctest

julia> r = reshape(1:9, 3, 3)
3×3 reshape(::UnitRange{Int64}, 3, 3) with eltype Int64:
 1  4  7
 2  5  8
 3  6  9

julia> pad_reflect(r, (1,2,1,2))
6×6 Matrix{Int64}:
 5  2  5  8  5  2
 4  1  4  7  4  1
 5  2  5  8  5  2
 6  3  6  9  6  3
 5  2  5  8  5  2
 4  1  4  7  4  1
```
"""
function pad_reflect(x::AbstractArray, pad::NTuple{M,Int};
                     dims=1:M÷2) where M
  length(dims) == M ÷ 2 ||
    throw(ArgumentError("The number of dims should be equal to the number of padding dimensions"))
  for (i, d) in enumerate(dims)
    x = pad_reflect(x, (pad[2i-1], pad[2i]); dims = d)
  end
  return x
end

function pad_reflect(x::AbstractArray{F,N}, pad::NTuple{2,Int};
                     dims::Int = 1) where {F,N}
  lpad, rpad = pad

  n = size(x, dims)
  xl = selectdim(x, dims, lpad+1:-1:2)
  xr = selectdim(x, dims, n-1:-1:n-rpad)
  # Alternative selection, not sure which is faster...
  # xl = reverse(selectdim(x, dims, 2:lpad+1), dims)
  # xr = reverse(selectdim(x, dims, n-rpad:n-1), dims)
  return cat(xl, x, xr, dims = dims)
end

# convenience methods for symmetric and homogeneous padding
pad_repeat(x::AbstractArray{F,N}, pad::Int; dims=1:N-2) where {F,N} =
  pad_repeat(x, ntuple(_ -> pad, 2length(dims)); dims = dims)
pad_reflect(x::AbstractArray{F,N}, pad::Int; dims=1:N-2) where {F,N} =
  pad_reflect(x, ntuple(_ -> pad, 2length(dims)); dims = dims)


text:
# Keep track of parameter estimates
m_x = zeros(T+1,)
w_x = .1 *ones(T+1,)
a_τ = 1. *ones(T+1,)
b_τ = .5 *ones(T+1,)

for k = 1:T

    # Initialize recognition distribution marginals
    marginals = Dict(:x_kmin1 => ProbabilityDistribution(Univariate, GaussianMeanPrecision, m=m_x[k], w=w_x[k]),
                     :x_k => ProbabilityDistribution(Univariate, GaussianMeanPrecision, m=m_x[k], w=w_x[k]),
                     :τ => ProbabilityDistribution(Univariate, Gamma, a=a_τ[k], b=b_τ[k]))

    # Initialize data
    data = Dict(:y_k => Y[k],
                :m_x_kmin1 => m_x[k],
                :w_x_kmin1 => w_x[k],
                :τ1 => a_τ[k],
                :τ2 => b_τ[k])

    # Iterate updates
    for n = 1:10
        stepx!(data, marginals)
        stepτ!(data, marginals)
    end

    # Store estimates
    m_x[k+1] = mean(marginals[:x_k])
    w_x[k+1] = inv(cov(marginals[:x_k]))
    a_τ[k+1] = marginals[:τ].params[:a]
    b_τ[k+1] = marginals[:τ].params[:b]

end

# Generate plot of state estimates
p1 = scatter(1:T, Y, color="black", markersize=6, label="observations", xlabel="time (t)", ylabel="signal")
plot!(1:T, m_x[2:end], ribbon=[sqrt.(inv.(w_x[2:end])) sqrt.(inv.(w_x[2:end]))], color="blue", label="inferred states")

# Generate plot of process noise precision estimates
m_τ = a_τ ./ b_τ
s_τ = sqrt.(a_τ ./ b_τ.^2)
p2 = plot(1:T, m_τ[2:end], ribbon=[s_τ[2:end] s_τ[2:end]], color="blue", xlabel="time (t)", legend=:bottomright, ylabel="τ", label="inferred precision")

# Actually plot
p3 = plot(p1,p2, layout=(2,1), size=(800,600))

max_stars_repo_name: Yikeru/BMLIP

max_issues_repo_name: Yikeru/BMLIP

max_forks_repo_name: Yikeru/BMLIP