text:
using Dates
using ArgParse
using Printf
using Random
using GaussianProcesses

# before_all = Dates.now()

include("../../btg.jl")
s = ArgParseSettings()
# The defaut setting: --test: multiple length scale, QMC
@add_arg_table! s begin
    "--test"
        help = "test or not"
        action = :store_true
    "--validate"
        help = "do cross validation or not"
        action = :store_true
    "--fast"
        help = "use fast validation or not"
        action = :store_true
    "--GP"
        help = "test GP model"
        action = :store_true
    "--logGP"
        help = "test log-GP model"
        action = :store_true
    "--p"
        help = "degree of covariate function"
        arg_type = Int
        default = 1
    "--singletest"
        help = "write log to single test"
        action = :store_true
    "--ntrain"
        help = "number of training points"
        arg_type = Int
        default = 101
    "--ntest"
        help = "number of testing points"
        arg_type = Int
        default = 401
    "--noiselevel"
        help = "noise level in observation"
        arg_type = Int
        default = 1
    "--lmin"
        help = "minimal length scale"
        arg_type = Float64
        default = 0.5
    "--lmax"
        help = "maximum length scale"
        arg_type = Float64
        default = 2.
    "--randseed"
        help = "noise level in observation"
        arg_type = Int
        default = 1234
end
parsed_args = parse_args(ARGS, s)

function covariate_fun(x, p)
    n = size(x, 1)
    d = size(x, 2)
    if p == 1
        return ones(n, 1)
    elseif p == 1 + d
        return hcat(ones(n), x)
    else
        throw(ArgumentError("Only support constant or linear convariate."))
    end
end

true_fun(x, s) = sin(x) + s*randn() + 3
true_fun_noise(x) = max(true_fun(x, 1e-1), 0)^1/3
true_fun_0(x) = true_fun(x, 0.)^1/3

# training set
randseed = parsed_args["randseed"]; rng = MersenneTwister(randseed)
n_train = parsed_args["ntrain"]
data = range(-pi, stop=pi, length=n_train)
# data = 2pi .* rand(n_train) .- pi
target = (sin.(data) .+ 0.1 * parsed_args["noiselevel"].* randn(rng, n_train) .+ 10).^(1/3)
x = reshape(data, n_train, 1)
Fx = covariate_fun(x, parsed_args["p"])
max_train = maximum(target)
y = target ./ max_train
trainingData0 = trainingData(x, Fx, y)
d = getDimension(trainingData0); n = getNumPts(trainingData0); p = getCovDimension(trainingData0)
# testing set
n_test = parsed_args["ntest"]
x_test = range(-pi, stop=pi, length=n_test)
x_test = reshape(x_test, n_test, 1)
Fx_test = covariate_fun(x_test, parsed_args["p"])
y_test_true = (sin.(x_test) .+ 10).^(1/3)
elapsedmin = 0

#parameter setting
# myquadtype = parsed_args["sparse"] ? ["SparseCarlo", "SparseCarlo"] : ["QuasiMonteCarlo", "QuasiMonteCarlo"]
myquadtype = ["Gaussian", "Gaussian"]
rangeλ = [-1.5 1.]
lmin = parsed_args["lmin"]
lmax = parsed_args["lmax"]
rangeθ = [1/lmax^2 1/lmin^2]
@info "rangeθ, rangel:" rangeθ, [lmin lmax]
# rangeθ = [0.111 25]
# build btg model
btg0 = btg(trainingData0, rangeθ, rangeλ; quadtype = myquadtype)
(pdf0_raw, cdf0_raw, dpdf0_raw, quantInfo0_raw) = solve(btg0);

####################################
############### Test ###############
####################################
if parsed_args["test"]
    @info "Start Test"
    before = Dates.now()
    count_test = 0
    error_abs = 0.
    error_sq = 0.
    nlpd = 0.
    id_fail = []
    id_nonproper = []
    CI_set = zeros(n_test, 2)
    median_set = zeros(n_test)
    for i in 1:n_test
        global error_abs, error_sq, nlpd, count_test
        # mod(i, 20) == 0 ? (@info i) : nothing
        # @info "i" i
        x_test_i = reshape(x_test[i, :], 1, d)
        Fx_test_i = reshape(Fx_test[i, :], 1, p)
        try
            pdf_test_i, cdf_test_i, dpdf_test_i, quantbound_test_i, support_test_i = pre_process(x_test_i, Fx_test_i, pdf0_raw, cdf0_raw, dpdf0_raw, quantInfo0_raw)
            y_test_i_true = y_test_true[i]
            median_test_i = max_train * quantile(cdf_test_i, quantbound_test_i, support_test_i)[1]
            # @info "True, median " y_test_i_true, median_test_i
            median_set[i] = median_test_i
            try
                CI_test_i = max_train .* credible_interval(cdf_test_i, quantbound_test_i, support_test_i; mode=:equal, wp=.95)[1]
                count_test += (y_test_i_true >= CI_test_i[1])&&(y_test_i_true <= CI_test_i[2]) ? 1 : 0
                # @info "CI" CI_test_i
                CI_set[i, :] = CI_test_i
            catch err
                append!(id_fail, i)
            end
            error_abs += abs(y_test_i_true - median_test_i)
            error_sq += (y_test_i_true - median_test_i)^2
            nlpd += log(pdf_test_i(y_test_i_true/max_train))
        # @info "Count, id_fail" count_test, id_fail
        catch err
            append!(id_nonproper, i)
        end
        count_test /= n_test - length(id_fail) - length(id_nonproper)
        error_abs  /= n_test - length(id_nonproper)
        error_sq   /= n_test - length(id_nonproper)
        nlpd       /= -n_test - length(id_nonproper)
        after = Dates.now()
        elapsedmin = round(((after - before) / Millisecond(1000))/60, digits=5)

        # Plot
        PyPlot.close("all") #close existing windows
        PyPlot.plot(x_test, median_set, label = "BTG median")
        PyPlot.plot(x_test, y_test_true, label = "true")
        PyPlot.fill_between(dropdims(x_test; dims = 2), CI_set[:, 1], CI_set[:, 2], alpha = 0.3, label = "95% confidence interval")
        PyPlot.scatter(x, target, s = 10, c = "k", marker = "*")
        PyPlot.legend(fontsize=8)
        PyPlot.grid()
        PyPlot.title("BTG $myquadtype", fontsize=10)
        PyPlot.savefig("exp_synthetic1_btg_range_$(Int(10*lmin))_$(Int(10*lmax))_noise$(parsed_args["noiselevel"])_randseed_$(randseed)_p$(parsed_args["p"]).pdf")
    end

    if parsed_args["GP"]
        @info "Start GP"
        global error_abs_GP, error_sq_GP, CI_test_GP, count_test_GP, nlpd_GP
        # training set
        x = reshape(x, 1, n_train)
        y = target
        # build and fit a GP
        # mymean = MeanLin(zeros(d))
        mymean = MeanZero()
        kern = SE(zeros(d),0.0)
        gp = GP(x, y, mymean, kern)
        optimize!(gp)
        # predict
        x_test = reshape(x_test, 1, n_test)
        μ, σ² = predict_y(gp, x_test); stdv = sqrt.(σ²)
        error_GP = abs.(μ .- y_test_true)
        error_abs_GP = mean(error_GP)
        error_sq_GP = mean(error_GP.^2)
        CI_test_GP = hcat( -1.96 .* stdv .+ μ, 1.96 .* stdv .+ μ)
        count_test_GP = sum((y_test_true .>= CI_test_GP[:, 1]) .* (y_test_true .<= CI_test_GP[:,2]))/n_test
        nlpd_GP = -mean(log.(pdf.(Normal(), (y_test_true.-μ)./stdv)./stdv))
        # Plot
        # PyPlot.close("all") #close existing windows
        # PyPlot.plot(dropdims(x_test; dims = 1), μ, label = "GP median")
        # PyPlot.scatter(x, target, s = 10, c = "k", marker = "*")
        # PyPlot.fill_between(dropdims(x_test; dims = 1), CI_test_GP[:, 1], CI_test_GP[:, 2], alpha = 0.3, label = "95% credibel interval")
        # PyPlot.legend(fontsize=8)
        # PyPlot.grid()
        # PyPlot.title("GP", fontsize=10)
        # PyPlot.savefig("exp_synthetic1_GP_range_$(Int(10*lmin))_$(Int(10*lmax))_noise$(parsed_args["noiselevel"]).pdf")
    end

    if parsed_args["logGP"]
        @info "Start logGP"
        global error_abs_logGP, error_sq_logGP, CI_test_logGP, count_test_logGP, nlpd_logGP
        x = reshape(x, 1, n_train)
        y = target
        trans = BoxCox()
        g_fixed(x) = trans(x, 0.); dg(x) = partialx(trans, x, 0.)
        invg(x) = inverse(trans, x, 0.)
        gy = g_fixed.(y)
        # build and fit a GP
        # mymean = MeanLin(zeros(d))
        mymean = MeanZero()
        kern = SE(zeros(d),0.0)
        loggp = GP(x, gy, mymean, kern)
        optimize!(loggp)
        # predict
        x_test = reshape(x_test, 1, n_test)
        μ, σ² = predict_y(loggp, x_test); stdv = sqrt.(σ²)
        CI_test_logGP = invg.(hcat(-1.96.*stdv .+ μ, 1.96.*stdv .+ μ))
        count_test_logGP = sum((y_test_true .>= CI_test_logGP[:, 1]) .* (y_test_true .<= CI_test_logGP[:,2]))/n_test
        y_pred = invg.(μ)
        error_logGP = abs.(y_pred .- y_test_true)
        error_abs_logGP = mean(error_logGP)
        error_sq_logGP = mean(error_logGP.^2)
        nlpd_logGP = -mean(log.( dg.(y_test_true) .* pdf.(Normal(), (g_fixed.(y_test_true).-μ)./stdv) ./stdv ))
        # Plot
        # PyPlot.close("all") #close existing windows
        # PyPlot.plot(dropdims(x_test; dims = 1), invg.(μ), label = "logGP median")
        # PyPlot.scatter(x, target, s = 10, c = "k", marker = "*")
        # PyPlot.fill_between(dropdims(x_test; dims = 1), CI_test_logGP[:, 1], CI_test_logGP[:, 2], alpha = 0.3, label = "95% credibel interval")
        # PyPlot.legend(fontsize=8)
        # PyPlot.grid()
        # PyPlot.title("logGP", fontsize=10)
        # PyPlot.savefig("exp_synthetic1_logGP_range_$(Int(10*lmin))_$(Int(10*lmax))_noise$(parsed_args["noiselevel"]).pdf")
    end

    # plot and Compare - full comparison between btg, GP, and logGP
    if parsed_args["GP"] && parsed_args["logGP"] && parsed_args["test"]
        x_test = reshape(x_test, n_test)
        y_test_true = reshape(y_test_true, n_test)
        median_set = reshape(median_set, n_test)
        PyPlot.close("all") #close existing windows
        PyPlot.scatter(x, target, s = 10, c = "k", marker = "*")
        PyPlot.plot(x_test, y_test_true, label = "true")
        PyPlot.plot(x_test, median_set, label = "BTG median")
        # PyPlot.fill_between(x_test, CI_set[:, 1], CI_set[:, 2], alpha = 0.3, label = "95% CI BTG")
        # PyPlot.fill_between(x_test, CI_test_GP[:, 1], CI_test_GP[:, 2], alpha = 0.3, label = "95% CI GP")
        # PyPlot.fill_between(x_test, CI_test_logGP[:, 1], CI_test_logGP[:, 2], alpha = 0.3, label = "95% CI logGP")
        PyPlot.plot(x_test, CI_set[:, 1], "b-", label = "95% CI BTG")
        PyPlot.plot(x_test, CI_set[:, 2], "b-")

        PyPlot.plot(x_test, CI_test_GP[:, 1], "k:", label = "95% CI GP")
        PyPlot.plot(x_test, CI_test_GP[:, 2], "k:")

        PyPlot.plot(x_test, CI_test_logGP[:, 1], "c-.", label = "95% CI logGP")
        PyPlot.plot(x_test, CI_test_logGP[:, 2], "c-.")

        PyPlot.title("Compare BTG, GP and logGP", fontsize=10)
        PyPlot.legend(fontsize=8)
        PyPlot.grid()
        PyPlot.savefig("exp_synthetic1_compare_range_$(Int(10*lmin))_$(Int(10*lmax))_noise$(parsed_args["noiselevel"])_randseed_$(randseed)_p$(parsed_args["p"]).pdf")
    end

    io1 = open("Exp_synthetic1_test.txt", "a")
    write(io1, "\n$(Dates.now()) \n" )
    write(io1, "Data set: exp_synthetic1
        train size:  $n_train;  test size:  $n_test;  dimension: $d;   degree of covariate functions: $(parsed_args["p"])
        randseed: $randseed;    noise level in training data: $(0.1*parsed_args["noiselevel"]) \n")
    write(io1, "BTG model:
        $myquadtype  ;  rangeλ: $rangeλ;  lengthscale: [$lmin, $lmax];  rangeθ: [$(@sprintf("%2.3f", rangeθ[1])), $(@sprintf("%2.3f", rangeθ[2]))] \n")
    if parsed_args["GP"] && parsed_args["logGP"]
        write(io1, "Compare test results: ")
        write(io1, "                               BTG               GP               logGP
        credible intervel accuracy percentage:   $(@sprintf("%11.8f", count_test))       $(@sprintf("%11.8f", count_test_GP))       $(@sprintf("%11.8f", count_test_logGP))
        mean absolute error:                     $(@sprintf("%11.8f", error_abs))       $(@sprintf("%11.8f", error_abs_GP))       $(@sprintf("%11.8f", error_abs_logGP))
        mean squared error:                      $(@sprintf("%11.8f", error_sq))       $(@sprintf("%11.8f", error_sq_GP))       $(@sprintf("%11.8f", error_sq_logGP))
        mean negative log predictive density:    $(@sprintf("%11.8f", nlpd))       $(@sprintf("%11.8f", nlpd_GP))       $(@sprintf("%11.8f", nlpd_logGP))
        Time cost by prediction: $elapsedmin
        BTG: Failed index in credible intervel:   $id_fail
        BTG: Failed index in pdf computation:     $id_nonproper\n")
    else
        write(io1, "BTG test results:
        credible intervel accuracy percentage:   $(@sprintf("%11.8f", count_test))
        mean absolute error:                     $(@sprintf("%11.8f", error_abs))
        mean squared error:                      $(@sprintf("%11.8f", error_sq))
        mean negative log predictive density:    $(@sprintf("%11.8f", nlpd))
        Time cost by prediction: $elapsedmin
        Failed index in credible intervel:       $id_fail
        BTG: Failed index in pdf computation:     $id_nonproper\n")
    end
    close(io1)

end

####################################
############ Validation ############
####################################
function lootd(td::AbstractTrainingData, i::Int64)
    x = getPosition(td)
    Fx = getCovariates(td)
    z = getLabel(td)
    x_minus_i = x[[1:i-1;i+1:end], :]
    Fx_minus_i = Fx[[1:i-1;i+1:end], :]
    z_minus_i = z[[1:i-1;i+1:end]]
    x_i = x[i:i, :]
    Fx_i = Fx[i:i, :]
    z_i = z[i:i, :]
    return trainingData(x_minus_i, Fx_minus_i, z_minus_i), x_i, Fx_i, z_i
end

mutable struct validation_plot_buffer
    ygrid::Array{Float64, 1}
    median::Float64
    CI::Array{Float64, 1}
    pdf_ytrue::Float64
    pdf_median_i::Float64
    pdf_CI::Array{Float64, 1}
    function validation_plot_buffer(ygrid_i, median_i, CI_i, pdf_ytrue_i, pdf_median_i, pdf_CI_i)
        return new(ygrid_i, median_i, CI_i, pdf_ytrue_i, pdf_median_i, pdf_CI_i)
    end
end

function unpack(a::validation_plot_buffer)
    return (a.ygrid, a.median, a.CI, a.pdf_ytrue, a.pdf_median_i, a.pdf_CI)
end

if parsed_args["validate"]
    @info "Start validation"
    before = Dates.now()
    id_fail_val = []
    id_nonproper_val = []
    count_val = 0
    error_abs_val = 0.
    error_sq_val = 0.
    nlpd_val = 0.
    # prepare space for plotting
    n_row = 6; n_col = 8
    xgrid = range(.01, stop=1.3, length=100)
    validation_plot_buffer_dict = Dict{T where T<:Int64, validation_plot_buffer}()
    for i = 1:n_train
        global count_val, error_abs_val, error_sq_val, nlpd_val
        mod(i, 10) == 0 ? (@info i) : nothing
        if parsed_args["fast"]
            (pdf_val_i_raw, cdf_val_i_raw, dpdf_val_i_raw, quantInfo_val_i_raw) = solve(btg0, validate = i)
        else
            (trainingdata_minus_i, x_i, Fx_i, z_i) = lootd(trainingData0, i)
            btg_val_naive = btg(trainingdata_minus_i, rangeθ, rangeλ; quadtype = myquadtype)
            (pdf_val_i_raw, cdf_val_i_raw, dpdf_val_i_raw, quantInfo_val_i_raw) = solve(btg_val_naive)
        end
        try
            x_i = x[i:i, :]; Fx_i = Fx[i:i, :]
            pdf_val_i, cdf_val_i, dpdf_val_i, quantbound_val_i, support_val_i = pre_process(x_i, Fx_i, pdf_val_i_raw, cdf_val_i_raw, dpdf_val_i_raw, quantInfo_val_i_raw)
            median_val_i = quantile(cdf_val_i, quantbound_val_i, support_val_i)[1]
            y_val_i_true = getLabel(btg0.trainingData)[i]
            pdf_ytrue_i = pdf_val_i(y_val_i_true)
            try
                CI_val_i = credible_interval(cdf_val_i, quantbound_val_i, support_val_i; mode=:equal, wp=.95)[1]
                count_val += (y_val_i_true >= CI_val_i[1])&&(y_val_i_true <= CI_val_i[2]) ? 1 : 0
                if length(validation_plot_buffer_dict) < n_row*n_col
                    ygrid_i = pdf_val_i.(xgrid)
                    pdf_CI_i = pdf_val_i.(CI_val_i)
                    pdf_median_i = pdf_val_i.(median_val_i)
                    push!(validation_plot_buffer_dict, i => validation_plot_buffer(ygrid_i, median_i, CI_val_i, pdf_ytrue_i, pdf_median_i, pdf_CI_i))
                end
            catch err
                append!(id_fail_val, i)
            end
            error_abs_val += abs(y_val_i_true - median_val_i)
            error_sq_val += (y_val_i_true - median_val_i)^2
            nlpd_val -= log(pdf_ytrue_i)
        # @info "Count, id_fail" count_test, id_fail
        catch err
            append!(id_nonproper_val, i)
        end
    end
    count_val      /= n_train - length(id_fail_val) - length(id_nonproper_val)
    error_abs_val  /= n_train - length(id_nonproper_val)
    error_sq_val   /= n_train - length(id_nonproper_val)
    nlpd_val       /= n_train - length(id_nonproper_val)
    after = Dates.now()
    elapsedmin = round(((after - before) / Millisecond(1000))/60, digits=5)

    # save results
    io2 = open("Exp_exp_synthetic1_LOOCV.txt", "a")
    write(io2, "\n$(Dates.now()), randseed: $randseed \n" )
    write(io2, "Data set: exp_synthetic1
        id_train:  $id_train;   posx: $posx;   posc: $posc\n")
    write(io2, "BTG model:
        $myquadtype  ;  rangeλ: $rangeλ;   rangeθ: $rangeθs (single length-scale: $(parsed_args["single"]))
        Fast LOOCV: $(parsed_args["fast"]) \n")
    write(io2, "LOOCV on training set results:
    credible intervel accuracy percentage:   $(@sprintf("%11.8f", count_val))
    mean absolute error:                     $(@sprintf("%11.8f", error_abs_val))
    mean squared error:                      $(@sprintf("%11.8f", error_sq_val))
    mean negative log predictive density:    $(@sprintf("%11.8f", nlpd_val))
    Time validation took: $elapsedmin
    Failed index in credible intervel:       $id_fail_val
    BTG: Failed index in pdf computation:     $id_nonproper_val\n")
    close(io2)

    # Plot the first 48
    PyPlot.close("all") #close existing windows
    plt, axs = PyPlot.subplots(n_row, n_col)
    PyPlot.suptitle("Partial Plots of LOOCV", fontsize=10)
    id = 0
    for i in keys(validation_plot_buffer_dict)
        id += 1
        global id
        ind1 = Int64(ceil(id/n_col))
        ind2 = Int64(id - n_col*(floor((id-.1)/n_col)))
        y_val_i_true = getLabel(btg0.trainingData)[i]
        ygrid_i, median_i, CI_val_i, pdf_ytrue_i, pdf_median_i, pdf_CI_i = unpack(validation_plot_buffer_dict[i])
        axs[ind1, ind2].plot(xgrid, ygrid_i, linewidth = 1.0, linestyle = "-")
        # axs[ind1, ind2].plot(xgrid, cdf1.(xgrid), color = "orange", linewidth = 1.0, linestyle = "-")
        axs[ind1, ind2].vlines(y_val_i_true, 0, pdf_ytrue_i, label = "true value")
        axs[ind1, ind2].vlines(median_val_i, 0, pdf_median_i, label = "median")
        CI_id = (xgrid .> CI_val_i[1]) .* (xgrid .< CI_val_i[2])
        CI_xrange = vcat(CI_val_i[1], xgrid[CI_id], CI_val_i[2])
        CI_yrange = vcat(pdf_CI_i[1], ygrid_i[CI_id], pdf_CI_i[2]) # utilize previous evaluation results
        axs[ind1, ind2].fill_between(CI_xrange, 0, CI_yrange, alpha = 0.3, label = "95% confidence interval")
    end
    # Hide x labels and tick labels for top plots and y ticks for right plots.
    for ax in axs
        ax.label_outer()
    end
    PyPlot.savefig("figure/exp_exp_synthetic1_$(id_train[1])_$(id_train[2])_$(myquadtype[1])$(myquadtype[2])_posc$(posc[2])_rθ_$(Int(rangeθs[1]))_$(Int(rangeθs[2]))_rλ_$(Int(rangeλ[1]))_$(Int(rangeλ[2]))_$(parsed_args["fast"]).pdf")

end

# after_all = Dates.now()
# elapsedmin_all = round(((after_all - before_all) / Millisecond(1000))/60, digits=5)
# @info "Total time" elapsedmin_all

max_stars_repo_name: dbindel/btg

max_issues_repo_name: dbindel/btg

max_forks_repo_name: dbindel/btg


text:
using WaveProp.Geometry
using WaveProp.Integration
using Test

# # Integration

# The `WaveProp.Integration` module implements the various quadrature rules and
# integration routines required throughout the package. The interface of
# the `Integration` module is given through `AbstractQuadratureRule{D}`, where the
# type parameter `D` is a subtype of `AbstractReferenceShape`.

# For a given a domain `D`, an $N$-point quadrature rule over `D` is a set of nodes $\{x_i\}$
# and weights $\{w_i\}$, $1 \leq i \leq N$, which can be used to integrate a function over `D`. That is
# ```math
# \int_D f(x) dx \approx \sum_{i=1}^{N} f(x_i)w_i
# ```

# There are several concrete subtypes of `AbstractQuadratureRule`. To see a list
# of available implementations you can type:
subtypes(Integration.AbstractQuadratureRule)

# Different quadrature rules may require different input arguments in order to
# be instantiate. For example, the 1d quadrature rules of `GaussLegendre` type
# require as argument `n`, the number of quadrature nodes. You can instantiate a
# 10-point `GaussLegendre` rule as
q = GaussLegendre(10)

# To get the nodes and weights of `q` you can simply type:
x̂,ŵ = q()

# Similarly, you can `integrate` a callable object usign the syntax
I = integrate(x->cos(x)*exp(x),q)

# By default, calling an `AbstractQuadrature` object with no arguments returns
# the underlying quadrature rule for integration over `domain(q)`. The basic
# strategy for integration over an arbitrarily shaped element `el` is simple:
# first, develop a quadrature rules on reference shape `D == domain(el)`, and
# then *lift* the quadrature nodes `xᵢ` and weights `wᵢ` to the element `el`
# using its parametrization. Notice that this requires `el` to implement `el(u)`
# and `jacobian(el,u)` for `u ∈ D`.

# We demonstrate how to do that in a simple element: a line segment
# connecting `a=(0,0)` to `b=(2,2)`:
a,b = (0,0), (2,2)
el = line(a,b)
x,w = q(el)

# It is easy to check that `sum(w) == √8`, i.e. the length of the line
@test sum(w) ≈  √8

ext: jl

max_stars_repo_name: WaveProp/WaveProp

max_issues_repo_name: WaveProp/WaveProp

max_forks_repo_name: WaveProp/WaveProp


text:
function OCT_gradient(ω1, TRF, TR, ω0, B1, m0s, R1f, R2f, Rx, R1s, T2s, R2slT, grad_list, weights)
    (E, dEdω1, dEdTRF) = calculate_propagators_ω1(ω1, TRF, TR, ω0, B1, m0s, R1f, R2f, Rx, R1s, T2s, R2slT, grad_list)
    Q = calcualte_cycle_propgator(E)
    Y = propagate_magnetization(Q, E)
    (CRB, d) = dCRBdm(Y, weights)
    P = calculate_adjoint_state(d, Q, E)
    (grad_ω1, grad_TRF) = calculate_gradient_inner_product(P, Y, E, dEdω1, dEdTRF)
    return (CRB, grad_ω1, grad_TRF)
end

function OCT_TV_gradient(ω1, TRF, TR, ω0, B1, m0s, R1f, R2f, Rx, R1s, T2s, R2slT, grad_list, weights, λ_TV)
    (E, dEdω1, dEdTRF) = calculate_propagators_ω1(ω1, TRF, TR, ω0, B1, m0s, R1f, R2f, Rx, R1s, T2s, R2slT, grad_list)
    Q = calcualte_cycle_propgator(E)
    Y = propagate_magnetization(Q, E)

    (CRB, d) = dCRBdm(Y, weights)
    P = calculate_adjoint_state(d, Q, E)
    (grad_ω1, grad_TRF) = calculate_gradient_inner_product(P, Y, E, dEdω1, dEdTRF)

    (TV, d) = dthetaTVdm(Y, λ_TV)
    P = calculate_adjoint_state(d, Q, E)
    (grad_ω1_TV, grad_TRF_TV) = calculate_gradient_inner_product(P, Y, E, dEdω1, dEdTRF)

    grad_ω1  .-= grad_ω1_TV
    grad_TRF .-= grad_TRF_TV

    return (CRB+TV, grad_ω1, grad_TRF)
end

function calculate_propagators_ω1(ω1, TRF, TR, ω0, B1, m0s, R1f, R2f, Rx, R1s, T2s, R2slT, grad_list::AbstractArray{T}; rfphase_increment=[π]) where T <: grad_param
    E      = Array{SMatrix{11,11,Float64,121}}(undef, length(ω1), length(rfphase_increment), length(grad_list))
    dEdω1  = Array{SMatrix{11,11,Float64,121}}(undef, length(ω1), length(rfphase_increment), length(grad_list))
    dEdTRF = Array{SMatrix{11,11,Float64,121}}(undef, length(ω1), length(rfphase_increment), length(grad_list))

    cache = ExponentialUtilities.alloc_mem(zeros(22, 22), ExpMethodHigham2005Base())
    dH = zeros(Float64, 22, 22)

    for r in eachindex(rfphase_increment)
        u_rot = z_rotation_propagator(rfphase_increment[r], grad_m0s())
        for g in eachindex(grad_list)
            grad = grad_list[g]

            u_fp = xs_destructor(grad_list[g]) * exp(hamiltonian_linear(0, B1, ω0, TR / 2, m0s, R1f, R2f, Rx, R1s, 0, 0, 0, grad_list[g]))
            u_pl = propagator_linear_inversion_pulse(ω1[1], TRF[1], B1,
            R2slT[1](TRF[1], ω1[1] * TRF[1], B1, T2s),
            R2slT[2](TRF[1], ω1[1] * TRF[1], B1, T2s),
            R2slT[3](TRF[1], ω1[1] * TRF[1], B1, T2s),
            grad_list[g])
            E[1,r,g] = u_fp * u_pl * u_rot * u_fp
            dEdω1[1,r,g] = @SMatrix zeros(11, 11)
            dEdTRF[1,r,g] = @SMatrix zeros(11, 11)

            for t in 2:length(ω1)
                calculte_propagator!(E, dEdω1, dEdTRF, t, r, g, ω1[t], TRF[t], TR, ω0, B1, m0s, R1f, R2f, Rx, R1s, T2s, R2slT, grad, u_rot, dH, cache)
            end
        end
    end
    return (E, dEdω1, dEdTRF)
end

function calculte_propagator!(E, dEdω1, dEdTRF, t, r, g, ω1, TRF, TR, ω0, B1, m0s, R1f, R2f, Rx, R1s, T2s, R2slT, grad, u_rot, dH, cache)

    H_fp = hamiltonian_linear(0.0, B1, ω0, 1.0, m0s, R1f, R2f, Rx, R1s, 0.0, 0.0, 0.0, grad)
    ux = xs_destructor(grad)
    u_fp = ux * exp(H_fp * ((TR - TRF) / 2))

    H_pl = hamiltonian_linear(ω1, B1, ω0, 1, m0s, R1f, R2f, Rx, R1s,
        R2slT[1](TRF, ω1 * TRF, B1, T2s),
        R2slT[2](TRF, ω1 * TRF, B1, T2s),
        R2slT[3](TRF, ω1 * TRF, B1, T2s),
        grad)

    dHdω1 = d_hamiltonian_linear_dω1(B1, 1,
        R2slT[4](TRF, ω1 * TRF, B1, T2s),
        R2slT[6](TRF, ω1 * TRF, B1, T2s),
        R2slT[7](TRF, ω1 * TRF, B1, T2s),
        grad)

    @views dH[1:11,1:11]   .= H_pl
    @views dH[12:22,12:22] .= H_pl
    @views dH[1:11,12:22]  .= 0
    @views dH[12:22,1:11]  .= dHdω1
    dH .*= TRF
    E_pl = exponential!(dH, ExpMethodHigham2005Base(), cache)

    E_pl1 = SMatrix{11,11}(@view E_pl[1:11,1:11])
    E_pl2 = SMatrix{11,11}(@view E_pl[12:end,1:11])
    E[t,r,g]     = u_fp * E_pl1 * u_rot * u_fp
    dEdω1[t,r,g] = u_fp * E_pl2 * u_rot * u_fp

    # TRF
    dHdTRF = H_pl + d_hamiltonian_linear_dTRF_add(TRF,
        R2slT[5](TRF, ω1 * TRF, B1, T2s),
        R2slT[8](TRF, ω1 * TRF, B1, T2s),
        R2slT[9](TRF, ω1 * TRF, B1, T2s),
        grad)
    H_pl *= TRF
    @views dH[1:11,1:11]   .= H_pl
    @views dH[12:22,12:22] .= H_pl
    @views dH[1:11,12:22]  .= 0
    @views dH[12:22,1:11]  .= dHdTRF
    E_pl = exponential!(dH, ExpMethodHigham2005Base(), cache)

    E_pl1 = SMatrix{11,11}(@view E_pl[1:11,1:11])
    E_pl2 = SMatrix{11,11}(@view E_pl[12:end,1:11])
    dEdTRF[t,r,g] = u_fp * ((E_pl2 - (1 / 2 * H_fp * E_pl1)) * u_rot - (1 / 2 * E_pl1 * u_rot * H_fp)) * u_fp
    return nothing
end

function calcualte_cycle_propgator(E)
    Q = Array{SMatrix{11,11,Float64}}(undef, size(E, 2), size(E, 3))

    for g in 1:size(E, 3), r in 1:size(E, 2)
        A = E[1,r,g]
        for t = size(E, 1):-1:2
            A = A * E[t,r,g]
        end
        Q[r,g] = A0(grad_m0s()) - A
    end
    return Q
end

function propagate_magnetization(Q, E)
    Y = similar(E, SVector{11,Float64})

    for r in 1:size(E, 2), g in 1:size(E, 3)
        m = Q[r,g] \ C(grad_m0s())

        Y[1,r,g] = m
        for t = 2:size(E, 1)
            m = E[t,r,g] * m
            Y[t,r,g] = m
        end
    end
    return Y

    return S
end

# the commented line in this function are required for the adjoint state to be correct; but since these entries are not used for the OCT algorithm, we skip calculating them.
function calculate_adjoint_state(d, Q, E)
    P = similar(E, Array{Float64})
    λ = @view P[end,:,:]

    for g in 1:size(E, 3), r in 1:size(E, 2)
        λ[r,g] = d(size(E, 1), r, g)
    end

    for r = 1:size(E, 2)
        for t = size(E, 1) - 1:-1:1
            λ[r,1] = transpose(E[t + 1,r,1]) * λ[r,1]
            λ[r,1] += d(t, r, 1)
            for g = 2:size(E, 3)
                λ[r,g] = transpose(E[t + 1,r,g][6:10,:]) * λ[r,g][6:10]
                λ[r,1][1:5] .+= λ[r,g][1:5]
                # λ[r,1][end]  += λ[r,g][end]
                λ[r,g] .+= d(t, r, g)
            end
        end

    # for g = 2:length(grad_list)
    #     λ[r,g][1:5] .= λ[r,1][1:5]
    #     λ[r,g][end]  = λ[r,1][end]
    # end

        P[end,r,1] = inv(transpose(Q[r,1])) * λ[r,1]
        for g = 2:size(E, 3)
            λ[r,g] = inv(transpose(Q[r,g]))[:,6:10] * λ[r,g][6:10]
            λ[r,1][1:5] .+= λ[r,g][1:5]
            # P[end,r,1][end]  += P[end,r,g][end]
        end

        # step 5: propagate adjoint state
        for t in size(E, 1):-1:2
            for g = 1:size(E, 3)
                _E = E[mod(t, size(E, 1)) + 1,r,g]
                if g == 1
                    P[t - 1,r,g] = transpose(_E) * P[t,r,g]
                else
                    P[t - 1,r,g] = transpose(_E[6:10,:]) * P[t,r,g][6:10]
                    P[t - 1,r,1][1:5] .+= P[t - 1,r,g][1:5]
                    P[t - 1,r,1][end]  += P[t - 1,r,g][end]
                end
                P[t - 1,r,g] .+= d(t, r, g)
            end

        # for g = 2:length(grad_list)
        #     P[t - 1,r,g][1:5] .= P[t - 1,r,1][1:5]
        #     P[t - 1,r,g][end]  = P[t - 1,r,1][end]
        # end
        end
    end
    return P
end

function dCRBdm(Y, w)
    _dCRBdx = Array{Float64}(undef, size(Y, 1), size(Y, 2), size(Y, 3) + 1)
    _dCRBdy = similar(_dCRBdx)
    F = zeros(ComplexF64, size(Y, 3) + 1, size(Y, 3) + 1)
    dFdy = similar(F)
    tmp  = similar(F)

    for g2 in 0:size(Y, 3), g1 in 0:size(Y, 3), r in 1:size(Y, 2), t in 1:size(Y, 1)
        if g1 == 0
            s1 = Y[t,r,1][1] - 1im * Y[t,r,1][2]
        else
            s1 = Y[t,r,g1][6] - 1im * Y[t,r,g1][7]
        end
        if g2 == 0
            s2 = Y[t,r,1][1] + 1im * Y[t,r,1][2]
        else
            s2 = Y[t,r,g2][6] + 1im * Y[t,r,g2][7]
        end
        F[g1 + 1,g2 + 1] += s1 * s2
    end
    Fi = inv(F)
    CRB = real.(w * diag(Fi))

    for g1 in 0:size(Y, 3), r in 1:size(Y, 2), t in 1:size(Y, 1)
        # derivative wrt. x
        dFdy .= 0
        dFdy[g1 + 1,1] = Y[t,r,1][1] + 1im * Y[t,r,1][2]
        dFdy[1,g1 + 1] = Y[t,r,1][1] - 1im * Y[t,r,1][2]
        for g2 in 1:size(Y, 3)
            dFdy[g1 + 1,g2 + 1] = Y[t,r,g2][6] + 1im * Y[t,r,g2][7]
            dFdy[g2 + 1,g1 + 1] = Y[t,r,g2][6] - 1im * Y[t,r,g2][7]
        end
        dFdy[g1 + 1,g1 + 1] = 2 * real(dFdy[g1 + 1,g1 + 1])
        mul!(dFdy, Fi, mul!(tmp, dFdy, Fi))
        _dCRBdx[t,r,g1 + 1] = real.(w * diag(dFdy))

        # derivative wrt. y
        dFdy .= 0
        dFdy[g1 + 1,1] = Y[t,r,1][2] - 1im * Y[t,r,1][1]
        dFdy[1,g1 + 1] = Y[t,r,1][2] + 1im * Y[t,r,1][1]
        for g2 in 1:size(Y, 3)
            dFdy[g1 + 1,g2 + 1] = Y[t,r,g2][7] - 1im * Y[t,r,g2][6]
            dFdy[g2 + 1,g1 + 1] = Y[t,r,g2][7] + 1im * Y[t,r,g2][6]
        end
        dFdy[g1 + 1,g1 + 1] = 2 * real(dFdy[g1 + 1,g1 + 1])
        mul!(dFdy, Fi, mul!(tmp, dFdy, Fi))
        _dCRBdy[t,r,g1 + 1] = real.(w * diag(dFdy))
    end

    d(t, r, g) = @SVector [_dCRBdx[t,r,1], _dCRBdy[t,r,1],0,0,0,_dCRBdx[t,r,g + 1], _dCRBdy[t,r,g + 1],0,0,0,0]

    return (CRB, d)
end

function dthetaTVdm(Y, λ)
    # ϑ = vec([atan(Y[t,r,1][1] / Y[t,r,1][3]) for t=1:size(Y,1), r=1:size(Y,2)])
    # TV = (ϑ[1] - ϑ[end])^2
    # for t=2:length(ϑ)
    #     TV += (ϑ[t] - ϑ[t-1])^2
    # end

    T = size(Y, 1)
    R = size(Y, 2)
    @inline x(t, r) = Y[t,r,1][1]
    @inline z(t, r) = Y[t,r,1][3]
    @inline ϑ(t, r) = atan(x(t, r) / z(t, r))

    TV = zero(eltype(Y[1]))
    _dthetaTVdx = ones(size(Y, 1), size(Y, 2))
    _dthetaTVdz = ones(size(Y, 1), size(Y, 2))

    ϑₙ₋₁ = ϑ(T, R)
    ϑₙ₋₂ = ϑ(T - 1, R)
    tₙ₋₁ = T
    rₙ₋₁ = R
    for r ∈ 1:R
        for t ∈ 1:T
            ϑₙ = ϑ(t, r)
            TV += (ϑₙ - ϑₙ₋₁)^2

            dₙ  = x(t, r)^2 + z(t, r)^2
            _dthetaTVdx[t,r] *= z(t, r) / dₙ
            _dthetaTVdz[t,r] *= -x(t, r) / dₙ
            _dthetaTVdx[tₙ₋₁,rₙ₋₁] *= 2λ * (2ϑₙ₋₁ - ϑₙ - ϑₙ₋₂)
            _dthetaTVdz[tₙ₋₁,rₙ₋₁] *= 2λ * (2ϑₙ₋₁ - ϑₙ - ϑₙ₋₂)

            tₙ₋₁ = t
            rₙ₋₁ = r
            ϑₙ₋₂ = ϑₙ₋₁
            ϑₙ₋₁ = ϑₙ
        end
    end

    d(t, r, _) = @SVector [_dthetaTVdx[t,r], 0, _dthetaTVdz[t,r],0,0,0,0,0,0,0,0]
    return (λ * TV, d)
end

function calculate_gradient_inner_product(P, Y, E, dEdω1, dEdTRF)
    grad_ω1 = zeros(size(E, 1))
    grad_TRF = zeros(size(E, 1))

    for g in 1:size(E, 3), r in 1:size(E, 2), t in 2:size(E, 1)
        if g == 1
            grad_ω1[t]  -= transpose(P[t - 1,r,g]) * (dEdω1[t,r,g] * Y[t - 1,r,g])
            grad_TRF[t] -= transpose(P[t - 1,r,g]) * (dEdTRF[t,r,g] * Y[t - 1,r,g])
        else
            a = dEdω1[t,r,g] * Y[t - 1,r,g]
            b = dEdTRF[t,r,g] * Y[t - 1,r,g]
            @inbounds for i = 6:10
                grad_ω1[t]  -= P[t - 1,r,g][i] * a[i]
                grad_TRF[t]  -= P[t - 1,r,g][i] * b[i]
            end
        end
    end
    return (grad_ω1, grad_TRF)
end


text:
import Base

mutable struct Add <: DiffableFunction
    grad_field::GradField
end

mutable struct Sub <: DiffableFunction
    grad_field::GradField
end

mutable struct Neg <: DiffableFunction
    grad_field::GradField
end

mutable struct Mul <: DiffableFunction
    grad_field::GradField
end

mutable struct Div <: DiffableFunction
    grad_field::GradField
end

mutable struct Pow <: DiffableFunction
    grad_field::GradField
    c
end

@inline forward(::Add, x1, x2)  = x1 + x2
# case: a - b
@inline forward(::Sub, x1, x2) = x1 - x2
# case: -a
@inline forward(::Sub, x) = -x
@inline forward(::Neg, x) = -x
@inline forward(::Mul, x1, x2) = x1 * x2
@inline forward(::Div, x1, x2) = x1 / x2
@inline forward(f::Pow, x1) = x1^f.c

function backward(::Add, gy::Variable{T}) where {T <: Real}
    return (gy, gy)
end

function backward(::Add, gy::Variable{T}) where {T <: AbstractArray}
    @. return (gy, gy)
end

function backward(::Sub, gy::Variable{T}) where {T <: Real}
    return (gy, -gy)
end

function backward(::Sub, gy::Variable{T}) where {T <: AbstractArray}
    return (gy, -gy)
end

function backward(::Neg, gy::Variable{T})  where {T <: Real}
    return -gy
end

function backward(::Neg, gy::Variable{T})  where {T <: AbstractArray}
    @. return -gy
end

function backward(f::Mul, gy::Variable{T}) where {T <: Real}
    x1, x2 = f.grad_field.inputs
    return (x2 * gy, x1 * gy)
end

function backward(f::Mul, gy::Variable{T}) where {T <: AbstractArray}
    x1, x2 = f.grad_field.inputs
    @. return (x2 .* gy, x1 .* gy)
end

function backward(f::Div, gy::Variable{T})  where {T <: Real}
    x1, x2 = f.grad_field.inputs
    return (1 / x2) * gy, (-x1 / (x2 * x2)) * gy
end

function backward(f::Div, gy::Variable{T})  where {T <: AbstractArray}
    x1, x2 = f.grad_field.inputs
    @. return (1 / x2) * gy, (-x1 / (x2 * x2)) * gy
end

function backward(f::Pow, gy::Variable{T})  where {T <: Real}
    x = f.grad_field.inputs[1]
    c = f.c
    return (c * (x^(c - 1))) * gy
end

function backward(f::Pow, gy::Variable{T})  where {T <: AbstractArray}
    x = f.grad_field.inputs[1]
    c = f.c
    @. return (c * (x^(c - 1))) * gy
end

const normal_operators = Dict(
    :+ => Add,
    :- => Sub,
    :* => Mul,
    :/ => Div,
)

Base.:-(x::Variable) = Neg(GradField())(x)
Base.:^(x::Variable, c) = Pow(GradField(), c)(x)
is_support(::typeof(^)) = true
get_jt_struct(::typeof(^)) = Pow

for (op, jt_func) in normal_operators
    @eval is_support(::typeof(Base.$op)) = true
    @eval is_support_broadcast(::typeof(Base.$op)) = true
    @eval get_jt_struct(::typeof(Base.$op)) = $jt_func
    @eval Base.$op(x1::Variable, x2::Real) = Base.$op(promote(x1, x2)...)
    @eval Base.$op(x1::Real, x2::Variable) = Base.$op(promote(x1, x2)...)
    @eval Base.$op(x1::Variable, x2::Variable) = $jt_func(GradField())(x1, x2)
end


text:
"""
    qqmap(obs::ClimGrid, ref::ClimGrid, fut::ClimGrid; method="Additive", detrend=true, window::Int=15, rankn::Int=50, thresnan::Float64=0.1, keep_original::Bool=false, interp::Function = Linear(), extrap::Function = Flat())

Quantile-Quantile mapping bias correction. For each julian day of the year (+/- **window** size), a transfer function is estimated through an empirical quantile-quantile mapping.

The quantile-quantile transfer function between **ref** and **obs** is etimated on a julian day (and grid-point) basis with a moving window around the julian day. Hence, for a given julian day, the transfer function is then applied to the **fut** dataset for a given julian day.

**Options**

**method::String = "Additive" (default) or "Multiplicative"**. Additive is used for most climate variables. Multiplicative is usually bounded variables such as precipitation and humidity.

**detrend::Bool = true (default)**. A 4th order polynomial is adjusted to the time series and the residuals are corrected with the quantile-quantile mapping.

**window::Int = 15 (default)**. The size of the window used to extract the statistical characteristics around a given julian day.

**rankn::Int = 50 (default)**. The number of bins used for the quantile estimations. The quantiles uses by default 50 bins between 0.01 and 0.99. The bahavior between the bins is controlled by the interp keyword argument. The behaviour of the quantile-quantile estimation outside the 0.01 and 0.99 range is controlled by the extrap keyword argument.

**thresnan::Float64 = 0.1 (default)**. The fraction is missing values authorized for the estimation of the quantile-quantile mapping for a given julian days. If there is more than **treshnan** missing values, the output for this given julian days returns NaNs.

**keep_original::Bool = false (default)**. If **keep_original** is set to true, the values are set to the original values in presence of too many NaNs.

**interp = Interpolations.Linear() (default)**. When the data to be corrected lies between 2 quantile bins, the value of the transfer function is linearly interpolated between the 2 closest quantile estimation. The argument is from Interpolations.jl package.

**extrap = Interpolations.Flat() (default)**. The bahavior of the quantile-quantile transfer function outside the 0.01-0.99 range. Setting it to Flat() ensures that there is no "inflation problem" with the bias correction. The argument is from Interpolation.jl package.
"""
function qqmap(obs::ClimGrid, ref::ClimGrid, fut::ClimGrid; method::String="Additive", detrend::Bool=true, window::Int=15, rankn::Int=50, thresnan::Float64=0.1, keep_original::Bool=false, interp=Linear(), extrap=Flat())

    # Consistency checks # TODO add more checks for grid definition
    @argcheck size(obs[1], 1) == size(ref[1], 1) == size(fut[1], 1)
    @argcheck size(obs[1], 2) == size(ref[1], 2) == size(fut[1], 2)

    # Modify dates (e.g. 29th feb are dropped/lost by default)
    obs = ClimateTools.dropfeb29(obs)
    ref = ClimateTools.dropfeb29(ref)
    fut = ClimateTools.dropfeb29(fut)

    # Remove trend if specified
    if detrend == true
        # Obs
        obs_polynomials = ClimateTools.polyfit(obs)
        obs = obs - ClimateTools.polyval(obs, obs_polynomials)
        # Ref
        ref_polynomials = ClimateTools.polyfit(ref)
        ref = ref - ClimateTools.polyval(ref, ref_polynomials)
        # Fut
        fut_polynomials = ClimateTools.polyfit(fut)
        poly_values = ClimateTools.polyval(fut, fut_polynomials)
        fut = fut - poly_values
    end

    #Get date vectors
    datevec_obs = get_timevec(obs)
    datevec_ref = get_timevec(ref)
    datevec_fut = get_timevec(fut)

    # Julian days vectors
    obs_jul = Dates.dayofyear.(datevec_obs)
    ref_jul = Dates.dayofyear.(datevec_ref)
    fut_jul = Dates.dayofyear.(datevec_fut)

    # Prepare output array (replicating data type of fut ClimGrid)
    dataout = fill(convert(typeof(fut[1].data[1]), NaN), (size(fut[1], 1), size(fut[1],2), size(fut[1], 3)))::Array{typeof(fut[1].data[1]), T} where T

    if minimum(ref_jul) == 1 && maximum(ref_jul) == 365
        days = 1:365
    else
        days = minimum(ref_jul)+window:maximum(ref_jul)-window
        start = Dates.monthday(minimum(datevec_obs))
        finish = Dates.monthday(maximum(datevec_obs))
    end

    # Reshape. Allows multi-threading on grid points.
    obsin = reshape(obs[1].data, (size(obs[1].data, 1)*size(obs[1].data, 2), size(obs[1].data, 3)))
    refin = reshape(ref[1].data, (size(ref[1].data, 1)*size(ref[1].data, 2), size(ref[1].data, 3)))
    futin = reshape(fut[1].data, (size(fut[1].data, 1)*size(fut[1].data, 2), size(fut[1].data, 3)))
    dataoutin = reshape(dataout, (size(dataout, 1)*size(dataout, 2), size(dataout, 3)))

    # Looping over grid points using multiple-dispatch calls to qqmap
    Threads.@threads for k = 1:size(obsin, 1)

        obsvec = obsin[k,:]
        refvec = refin[k,:]
        futvec = futin[k,:]

        dataoutin[k, :] = qqmap(obsvec, refvec, futvec, days, obs_jul, ref_jul, fut_jul, method=method, detrend=detrend, window=window, rankn=rankn, thresnan=thresnan, keep_original=keep_original, interp=interp, extrap=extrap)

    end

    # Apply mask
    dataout = applymask(dataout, obs.msk)

    # Rebuilding ClimGrid
    lonsymbol = Symbol(fut.dimension_dict["lon"])
    latsymbol = Symbol(fut.dimension_dict["lat"])

    dataout2 = AxisArray(dataout, Axis{lonsymbol}(fut[1][Axis{lonsymbol}][:]), Axis{latsymbol}(fut[1][Axis{latsymbol}][:]),Axis{:time}(datevec_fut))

    timeattrib = fut.timeattrib
    timeattrib["calendar"] = "365_day"

    globalattribs = fut.globalattribs
    if haskey(globalattribs, "history")
        globalattribs["history"] = string(globalattribs["history"], " - ", "Bias-corrected with ClimateTools.jl")
    else
        globalattribs["history"] = string("Bias-corrected with ClimateTools.jl")
    end

    C = ClimGrid(dataout2; longrid=fut.longrid, latgrid=fut.latgrid, msk=fut.msk, grid_mapping=fut.grid_mapping, dimension_dict=fut.dimension_dict, timeattrib=timeattrib, model=fut.model, frequency=fut.frequency, experiment=fut.experiment, run=fut.run, project=fut.project, institute=fut.institute, filename=fut.filename, dataunits=fut.dataunits, latunits=fut.latunits, lonunits=fut.lonunits, variable=fut.variable, typeofvar=fut.typeofvar, typeofcal="365_day", varattribs=fut.varattribs, globalattribs=globalattribs)

    if detrend == true
        C = C + poly_values
    end

    return C

end

"""
    qqmap(obsvec::Array{N,1}, refvec::Array{N,1}, futvec::Array{N,1}, days, obs_jul, ref_jul, fut_jul; method::String="Additive", detrend::Bool=true, window::Int64=15, rankn::Int64=50, thresnan::Float64=0.1, keep_original::Bool=false, interp=Linear(), extrap=Flat())

Quantile-Quantile mapping bias correction for single vector. This is a low level function used by qqmap(A::ClimGrid ..), but can work independently.

"""
function qqmap(obsvec::Array{N,1} where N, refvec::Array{N,1} where N, futvec::Array{N,1} where N, days, obs_jul, ref_jul, fut_jul; method::String="Additive", detrend::Bool=true, window::Int64=15, rankn::Int64=50, thresnan::Float64=0.1, keep_original::Bool=false, interp=Linear(), extrap=Flat())

    # range over which quantiles are estimated
    P = range(0.01, stop=0.99, length=rankn)
    obsP = similar(obsvec, length(P))
    refP = similar(refvec, length(P))
    sf_refP = similar(refvec, length(P))

    # Prepare output array
    dataout = similar(futvec, (size(futvec)))

    # LOOP OVER ALL DAYS OF THE YEAR
    # TODO Define "days" instead of 1:365
    for ijulian in days

        # idx for values we want to correct
        idxfut = (fut_jul .== ijulian)

        # Find all index of moving window around ijulian day of year
        idxobs = ClimateTools.find_julianday_idx(obs_jul, ijulian, window)
        idxref = ClimateTools.find_julianday_idx(ref_jul, ijulian, window)

        obsval = obsvec[idxobs] .+ eps(0.0) # values to use as ground truth
        refval = refvec[idxref] .+ eps(0.0)# values to use as reference for sim
        futval = futvec[idxfut] .+ eps(0.0) # values to correct

        obsvalnan = isnan.(obsval)
        refvalnan = isnan.(refval)
        futvalnan = isnan.(futval)

        if (sum(obsvalnan) < (length(obsval) * thresnan)) & (sum(refvalnan) < (length(refval) * thresnan)) & (sum(futvalnan) < (length(futval) * thresnan))

            # Estimate quantiles for obs and ref for ijulian
            quantile!(obsP, obsval[.!obsvalnan], P, sorted=false)
            quantile!(refP, refval[.!refvalnan], P, sorted=false)

            if lowercase(method) == "additive" # used for temperature
                sf_refP .= obsP .- refP
                itp = interpolate((refP,), sf_refP, Gridded(interp))
                itp = extrapolate(itp, extrap) # add extrapolation
                # dataout[idxfut] .= itp(futval) .+ futval
                futval .= itp(futval) .+ futval

                futval[futvalnan] .= 0.0

            elseif lowercase(method) == "multiplicative" # used for precipitation
                sf_refP .= obsP ./ refP
                sf_refP[sf_refP .< 0] .= eps(1.0)
                sf_refP[isnan.(sf_refP)] .= 1.0#eps(1.0)#1.0
                itp = interpolate((refP,), sf_refP, Gridded(interp))
                itp = extrapolate(itp, extrap) # add extrapolation
                futval .= itp(futval) .* futval

                futval[futvalnan] .= 0.0

            else
                error("Wrong method")
            end
            # Replace values with new ones
            dataout[idxfut] .= futval
        else

            if keep_original
                # # Replace values with original ones (i.e. too may NaN values for robust quantile estimation)
                dataout[idxfut] .= futval
            else
                # DO NOTHING (e.g. if there is no reference, we want NaNs and not original values)
                dataout[idxfut] .= NaN
            end
        end
    end

    return dataout

end


text:
using DimensionalData
using DelimitedFiles
using Dates
using SciMLBase
using DiffEqBase

Base.@kwdef mutable struct Process # Not ensemble
    process = nothing
    parameter_profile::Union{Function, Tuple, Array, NTuple} = x->0.0 # Can be a tuple of functions, if the system has more than one parameter
    parameter_profile_parameters = [] # Can be a tuple of tuples
    X0::Vector = [nothing]
    t0::Union{Float64, Int64} = 0.0
    transient_t0::Union{Float64, Int64} = t0 # t0 will always take precedence
    dt::Union{Float64, Int64} = 0.001
    savedt::Union{Float64, Int64} = 0.01
    tmax::Union{Float64, Int64} = 100.0
    alg::Union{SciMLBase.SciMLAlgorithm, Function, Nothing} = nothing
    solver_opts::Dict = Dict(:adaptive=>false)
    #parameter_rng::UInt64 = seed()
    solver_rng::Int64 = seed()
    id::Int64 = abs(rand(Int64)) # Just a unique number for this simulation
    date::String = string(Dates.now())
    solution = nothing
    varnames::Vector{Symbol} = defaultvars(length(X0))
end

function subshow(io, P)
    namelen = maximum(length.(string.((fieldnames∘typeof)(P)))) + 2
    fillfun = x -> repeat(" ", namelen-length(string(x)))
    showfields = (fieldnames∘typeof)(P)
    rows = ["   $x:$(fillfun(x))"*string(getfield(P, x))*"\n" for x ∈ showfields]
    issolved = !isnothing(getsolution(P))
    rows[findfirst(showfields .== :solution)] = "   solution:$(fillfun(:solution))$issolved\n"
    print(io, reduce(*, rows))
end
function Base.show(io::IO, P::Process)
    print(io, "Process with fields:\n")
    subshow(io, P)
end
function Base.show(io::IO, m::MIME"text/plain", P::Process)
    printstyled(io, "Process", color=:red, bold=true)
    printstyled(io, " with fields:\n", bold=true)
    subshow(io, P)
end

function defaultvars(x)
    vars = Symbol.(Char.([(120:122)..., 119, 117, 118]))
    (length(x) <= length(vars) ? vars[1:x] : 1:x)
end

function Process(D::Dict)
    for s ∈ setdiff(keys(D), (:date, :solution))
        if D[s] isa String
            if s == :process
                try
                    D[s] = eval(Meta.parse(D[s])) # Maybe we can't parse the process as a function
                catch
                    D[s] = D[s] # It must be a string. You won't be able to solve this process, so the timeseries better be supplied.
                end
            else
                if s == :gitcommit
                    pop!(D, s) # We don't need this
                else
                    D[s] = eval(Meta.parse(D[s]))
                end
            end
        end
    end
    Process(;D...)
end

function (P::Process)(;kwargs...)
    # Can use field aliases here
    kwargs = Dict(kwargs)
    repalias!(kwargs, process_aliases)

    P2 = deepcopy(P)
    [setfield!(P2, x, y) for (x, y) in kwargs]
    setfield!(P2, :solution, nothing) # You've changed some parameters, so the solution is no longer valid
    setfield!(P2, :id, abs(rand(Int64))) # New id, yeah?
    setfield!(P2, :solver_rng, seed()) # New random seed, yeah?
    setfield!(P2, :date, string(Dates.now())) # New datetime, yeah?
    return P2
end
export Process

"""Field Aliases for Process constructors"""
process_aliases = Dict(
    :process =>                     [:sim, :system, :processes],
    :parameter_profile =>           [:profile, :profiles, :𝑝, :𝑃, :parameter_profiles],
    :parameter_profile_parameters =>[:parameters, :ps, :params, :param, :parameter,
                                     :profile_parameters, :parameterprofileparameters,
                                     :profileparameters, :𝔓, :𝔭],
    :X0 =>                          [:initial_conditions, :X, :X_0, :X₀, :𝑥₀, :𝑋₀, :𝑥0,
                                     :𝑋0],
    :transient_t0 =>                [:transient, :cutoff, :tₜ, :𝑡ₜ, :tt],
    :t0 =>                          [:tstart, :t₀, :𝑡₀],
    :dt =>                          [:δt, :𝛿t, :δ𝑡, :𝛿𝑡],
    :savedt =>                      [:save_dt, :save_Δt, :save_δt, :save_𝛥t, :save_𝛿t, :save_Δ𝑡,
                                     :save_δ𝑡, :save_𝛥𝑡, :save_𝛿𝑡, :Δt, :𝛥t, :Δ𝑡, :Δ𝑡],
    :tmax =>                        [:t_max, :T, :Tmax, :T_max, :𝑇],
    :alg =>                         [:algorithm, :solver],
    :solver_opts =>                 [:opts, :solopts, :sol_opts, :solveropts],
    :solver_rng =>                  [:rng, :rngseed, :rng_seed, :solverrng],
    :id =>                          [:identifier, :inventory_id],
    :date =>                        [:time, :datetime],
    :solution =>                    [:sol, :result, :output],
    :varnames =>                    [:variables, :variablenames, :variable_names]
)
function repalias!(D, aliai::Dict)
    for d ∈ keys(D)
        for a ∈ keys(aliai)
            if d ∈ aliai[a]
                D[a] = pop!(D, d)
            end
        end
    end
end

"""A function to handle simulations that are specified with a Process type"""
function solution!(P::Process) # vars::Tuple=Tuple(1:size(P.X0)[1])
    if isnothing(P.solution)
        @debug "Solving for the $(getprocess(P)) process ($(getid(P)))"
        P.solution = P.process(P)
    end
    x = P.solution
end
export solution
function simulate(P::Process)
    P2 = deepcopy(P)
    P2.solution = solution!(P2)
    return P2
end
simulate!(P::Process) = begin P.solution = solution!(P); return nothing end
export simulate
export simulate!

# * Might have various solution types, so timeseries gets all of them as an array
timeseries(s::SciMLBase.AbstractTimeseriesSolution, dim::Real) = s[dim, :]
timeseries(s::SciMLBase.AbstractTimeseriesSolution, dim::Union{Vector, UnitRange}=1:size(s.u[1], 1)) = copy(s[dim, :]')
function timeseries(s::AbstractArray, dim::Union{Vector, UnitRange, Real}=1:size(s, 2))
    if s isa Vector
        if length(dim) != 1 || dim[1] != 1
            error("Cannot index the second dimension of the input, which is a vector")
        end
        s[:]
    else
        s[:, dim]
    end
end
function timeseries!(P::Process, dim=1:length(getX0(P)); transient::Bool=false)
    x = timeseries(solution!(P), dim)
    if transient
        idxs = 1:length(times(P, transient=true))
    else
        idxs = (length(P.transient_t0:P.savedt:P.t0)):1:length(times(P, transient=true))
    end
    saveTimes = (P.transient_t0:P.savedt:P.tmax)[idxs]
    namevars = P.varnames[dim]
    if size(x, 2) > 1
        x = DimArray(x[idxs, :], (Ti(saveTimes), Dim{:Variable}(namevars)))
    else
        x = DimArray(x[idxs], (Ti(saveTimes),))
    end
end
timeDims(T::DimArray) = dims(T, Ti).val
variableDims(T::DimArray) = dims(T, :Variable).val
# function timeseries(s::Tuple, dim::Union{Real, Vector, Tuple}=1)
#     timeseries(s[1], dim) # You gave the metadata as well
# end
export timeseries

"""Convenient access to some solution info"""
function times(P::Process; transient::Bool=false)
    if transient
        P.transient_t0:P.savedt:P.tmax
    else
        P.t0:P.savedt:P.tmax
    end
end
times(T::DimArray) = dims(T, Ti).val
export times

parameter_function(P::Process) = tuplef2ftuple(P.parameter_profile, P.parameter_profile_parameters)
export parameter_function

function parameter_functions(P::Process)
    if all(isempty.(getps(P)))
        getprofiles(P)
    else
        if P.parameter_profile isa Union{Tuple, Vector}
            [P.parameter_profile[x](P.parameter_profile_parameters[x]...) for x in 1:length(P.parameter_profile)]
        else
            P.parameter_profile(P.parameter_profile_parameters...)
        end
    end
end
export parameter_functions

function parameterseries(P::Process; p=nothing, kwargs...)
    ps = hcat(parameter_function(P).(times(P; kwargs...))...)
    if size(ps, 1) == 1 # This 1 × N array, which should be a vector
        ps = ps[:]
    end
    if isnothing(p) || (ps isa Vector)
        return ps
    else
        return ps[p, :]
    end
end
export parameterseries

# Access the fields of a process with functions
for field ∈ keys(process_aliases)
    f = Symbol(:get, field)
    eval(quote
        $f(P::Process) = P.$field; export $f
    end)
    for field_alias ∈ process_aliases[field]
        fa = Symbol(:get, field_alias)
        eval(quote
            $fa = $f; #export fa
        end)
    end
end

""" Which parameters are not constant?"""
function getvaryingparameters(P::Process)
    fs = getparameter_profile(P)
    ps = typeof.(fs) .!= typeof(constantParameter)
    findall(ps)
end
export getvaryingparameters

function forcevec(x)
    if !(x isa Union{AbstractArray, Tuple})
        x = [x]
    else
        x
    end
end
export forcevec

function forcemat(x)
    if x isa AbstractVector
        x = reshape(x, :, 1)
    end
    return x
end
export forcemat

function trimtransient(P::Process)
    if !isempty(P.solution)
        P.solution = timeseries(P, transient=false)
        P.X0 = P.solution[1, :] # If you want to resimulate, start at this point. Assumes the entire description of the system is contained in time and position, which is not unreasonable
    end
    P.transient_t0 = P.t0
    return P
end

function saveTimeseries!(P::Process, folder::String="./", delim::Char=','; transient::Bool=false, fileroot="timeseries")
    X = timeseries(P, transient=transient)
    mkpath(folder)
    if !transient
        P = trimtransient(P)
    end
    filename = joinpath(folder, fileroot*"_"*string(getid(P))*".csv")
    P.solution = abspath(folder)
    #P.solution = nothing
    # @info "Saving time-series data to $filename"
    writedlm(filename, X, delim)
end
export saveTimeseries!

function gettimeseriesfile(P::Process, folder::String)
    filename = filter(x->occursin("timeseries_"*string(getid(P)), x), readdir(folder))
end

function timeseries(P::Process, dim=1:length(getX0(P)); folder::Union{String, Bool}=(getsolution(P) isa String), kwargs...)
    if folder isa Bool && folder
        if getsolution(P) isa String
            folder = getsolution(P)
        else
            folder = "./"
        end
    end
    if folder isa String
        filename = gettimeseriesfile(P, folder)
    else
        return timeseries!(P, dim; kwargs...)
    end
    if isempty(filename)
        return timeseries!(P, dim; kwargs...)
    end
    filename = filename[1]
    @debug "Loading time-series data from $filename"
    filename = joinpath(folder, filename)
    P.solution = readdlm(filename, ',', Float64)
    return timeseries!(P, dim; kwargs...)
end

function updateparam(P::Process, p::Integer, profile, value)
    profiles = [getparameter_profile(P)...]
    values = [getparameter_profile_parameters(P)...]
    if length(profiles) == 1
        @assert p == 1
        profiles = [profile]
        values = value
    else
        profiles = [profiles[1:p-1]..., profile, profiles[p+1:end]...] # 🤮
        values = [values[1:p-1]..., value, values[p+1:end]...]
    end
    return P(parameter_profile=Tuple(profiles), parameter_profile_parameters=values)
end
function updateparam(P::Process, p::Integer, profile::Function)
    value = getparameter_profile_parameters(P)
    profiles = [getparameter_profile(P)...]
    value = (length(profiles) == 1 && p == 1) ? value : value[p]
    return updateparam(P, p, profile, value)
end
function updateparam(P::Process, p::Integer, value::Union{Number, Tuple, Vector})
    profile = getparameter_profile(P)
    profile = (length(profile) == 1 && p == 1) ? profile : profile[p]
    return updateparam(P, p, profile, value)
end
export updateparam


text:
function _moment(θ)
    (β,σ) = θ;
    util = Utility(σ);
    dutil = (s) -> Tracker.gradient(util,s)[1].data;
    dudc  = dutil.(a_t);
    dudc1 = dutil.(a_t1);
    dtrans=ddc.dtrans;
    dtrans_s = (a,s) -> dtrans(a,s)[2].data;

    return(dudc - β .* dtrans_s.(a_t1,s_t1) .* dudc1)
end

mutable struct Estimation
    data
    Vhat::ApproxFn
    Transition::Transition
    exWeight::Array{Float64,2}
    ltp1::Array{Int,1}
    invl::Array{Int,1}
    #  -ltp1[j] = location in sivec of observation for same firm at
    #  time t+1, it is < 0 if there is no such observation
    #  -invl[j] = location in sivec of observation for same firm at
    #   time t-1, it is < 0 if there is no such observation
    function Estimation(data)
        a=data.action,s=data.state;
        nT=size(a,2);nM=size(a,1);
        a_t = []; a_t1 = []; s_t = []; s_t1=[];
        a_t = [];
        a_t1 = [];
        s_t = [];
        for t = 1:(nT-1)
            a_t  = vcat(a_t,data.action[:,t]);
            s_t  = vcat(s_t,data.state[:,t]);
            a_t1 = vcat(a_t1,data.action[:,t+1]);
            s_t1 = vcat(s_t1,data.state[:,t+1])
        end
        # First estimate transition density

    end
        # for t = 1:(nT-1)
        #     global a_t  = vcat(a_t,data.action[:,t]);
        #     global s_t  = vcat(s_t,data.state[:,t]);
        #     global a_t1 = vcat(a_t1,data.action[:,t+1])
        # end

    end

end


text:
function get_colors(scene::Scene, ego_id::Int64 = 1)
    colors = Dict{Int64, Colorant}()
    for veh in scene
        colors[veh.id] = veh.id == ego_id ?  COLOR_CAR_EGO : COLOR_CAR_OTHER
    end
    return colors
end

function animate_record(env, rec::SceneRecord; sim_dt::Float64=0.1, cam=FitToContentCamera(0.), overlays = SceneOverlay[])
    duration =rec.nframes*sim_dt
    fps = Int(1/sim_dt)
    function render_rec(t, dt)
        frame_index = Int(floor(t/dt)) + 1
        return AutoViz.render(rec[frame_index-nframes(rec)], env, overlays, cam=cam)
    end
    return duration, fps, render_rec
end

function animate_scenes(scenes::Vector{Scene}, env; overlays::Vector{SceneOverlay} = SceneOverlay[], cam=FitToContentCamera(0.),  sim_dt=0.1)
    duration =length(scenes)*sim_dt
    fps = Int(1/sim_dt)
    # cam = FitToContentCamera(0.)
    function render_rec(t, dt)
        frame_index = Int(floor(t/dt)) + 1
        disp = SceneOverlay[overlays...]
        push!(disp, TextOverlay(text = ["Step: $frame_index"],
                                                           font_size=20,
                                                           pos=VecE2(-15,-10),
                                                           incameraframe=true))
        return AutoViz.render(scenes[frame_index], env, disp, cam=cam)
    end
    return duration, fps, render_rec
end

function animate_scenes(scenes::Vector{Scene},
                        actions::Vector{Float64},
                        env;
                        overlays::Vector{SceneOverlay} = SceneOverlay[],
                        cam=FitToContentCamera(0.),
                        sim_dt=0.1)
    duration =length(scenes)*sim_dt
    fps = Int(1/sim_dt)
    function render_rec(t, dt)
        frame_index = Int(floor(t/dt)) + 1
        return AutoViz.render(scenes[frame_index],
                              env,
                              cat(1, overlays, TextOverlay(text = ["Acc: $(actions[frame_index])"],
                                                           font_size=20,
                                                           pos=VecE2(-15,-10),
                                                           incameraframe=true)),
                              cam=cam,
                              car_colors=get_colors(scenes[frame_index]))
    end
    return duration, fps, render_rec
end

max_stars_repo_name: MaximeBouton/AutomotivePOMDPs

max_issues_repo_name: MaximeBouton/AutomotivePOMDPs

max_forks_repo_name: MaximeBouton/AutomotivePOMDPs


text:
#---------------------------------------------------------
# Tables and AG-Grid Utility Functions
#---------------------------------------------------------

module GridUtils

import ..Glimmer
import Tables
using OrderedCollections
import Dates
import Markdown
import REPL
import JSON, JSONTables

export  gridCell,
        gridButtonCell,
        gridTreeViewCell,
        gridRowIndexCell,
        gridSelectCell

function _gridCommonCell(;kwargs...)
    res = Dict()

    for a in kwargs
        key = string(a[1])
        if (key == "header")
            key = "headerName"
        elseif (key == "eventName")
            key = "juliaEventName"
        end
        res[key] = a[2]
    end

    return res;
end

function gridCell(;
    kwargs...
)
    res = _gridCommonCell(;kwargs...)
    return res
end

function gridRowIndexCell(;
    kwargs...
)
    res = _gridCommonCell(;kwargs...)

    res["valueGetter"] = "node.rowIndex + 1"

    res["editable"] = false;
    res["sortable"] = false;
    res["floatingFilter"] = false;
    res["filter"] = false

    return res
end

function gridButtonCell(;
    kwargs...
)
    res = _gridCommonCell(;kwargs...)

    res["cellRenderer"] = "buttonCellRenderer"

    res["editable"] = false;
    res["sortable"] = false;
    res["floatingFilter"] = false;
    res["filter"] = ""

    return res
end

function gridTreeViewCell(;
    kwargs...
)
    res = _gridCommonCell(;kwargs...)

    res["cellRenderer"] = "treeViewCellRenderer"

    res["editable"] = false;
    res["sortable"] = false;
    res["floatingFilter"] = false;
    res["filter"] = ""

    return res
end

function gridSelectCell(;
    kwargs...
)
    res = _gridCommonCell(;kwargs...)

    res["cellRenderer"] = "selectCellRenderer"
    res["cellEditor"] = "selectCellEditor"

    return res
end

function _is_javascript_safe(x::Integer)
    min_safe_int = -(Int64(2)^53-1)
    max_safe_int = Int64(2)^53-1
    min_safe_int < x < max_safe_int
end

function _is_javascript_safe(x::AbstractFloat)
    min_safe_float = -(Float64(2)^53-1)
    max_safe_float = Float64(2)^53-1
    min_safe_float < x < max_safe_float
end

function table2json(schema, rows)
    res = []
    for (i, row) in enumerate(rows)
        r = OrderedDict()
        Tables.eachcolumn(schema, row) do val, ind, name
            if val isa Real && isfinite(val) && _is_javascript_safe(val)
                # JSON.show_pair(columnwriter, ser, name, val)
                r[name] = val
            elseif val === nothing || val === missing
                # JSON.show_pair(columnwriter, ser, name, repr(val))
                r[name] = repr(val)
            else
                # JSON.show_pair(columnwriter, ser, name, sprint(print, val))
                r[name] = sprint(print, val)
            end
        end
        push!(res, r)
    end
    return res
end

"""
    table2agGrid(table)

Takes a `Tables.jl` compatable table and return the agGrid represantation of it with default colDefs.
"""
function table2agGrid(table)
    res = OrderedDict()

    rows = Tables.rows(table)
    schema = Tables.schema(rows)

    names = schema.names
    types = schema.types

    # create the column definitions
    res["columnDefs"] = [
        OrderedDict(
            "headerName" => string(n),
            # "editable" => cell_changed !== nothing,
            "headerTooltip" => string(types[i]),
            "field" => string(n),
            "sortable" => true,
            "resizable" => true,
            "type" => types[i] <: Union{Missing, T where T <: Number} ? "numericColumn" : nothing,
            "filter" => types[i] <: Union{Missing, T where T <: Dates.Date} ? "agDateColumnFilter" :
                    types[i] <: Union{Missing, T where T <: Number} ? "agNumberColumnFilter" : true,
        ) for (i, n) in enumerate(schema.names)
    ]

    res["defaultColDef"] = OrderedDict(
        "width" => 150,
        "editable" => true,
        "filter" => "agTextColumnFilter",
        "floatingFilter" => true,
        "resizable" => true,
    )

    res["rowData"] = table2json(schema, rows)

    res["_metaData"] = OrderedDict(
        "ag_grid_materializer" => Tables.materializer(table),
        "ag_grid_schema" => schema,
    )

    return res;
end

"""
    table2agGridTree(table, id::String, parentId::String)

Takes a `Tables.jl` compatable table and return the agGrid represantation of it with default colDefs.
The table data is assumed to represent an hericical data structure where the `id` parameter is the field containing a unique id, and the
`parentId` parameter is the field containing the parent id of the specific row.
"""
function table2agGridTree(table, id::String, parentId::String)

    #
    # build a dict representing the tree from the flat table
    #
    rows = Tables.rows(table)
    id_2_node_map = Dict()
    for (i, row) in enumerate(rows)
        # node = create_node(row)
        node = Dict([Symbol(col) => Tables.getcolumn(row, col) for col in Tables.columnnames(row)])

        # each node contains a treeview specific information in a _treeview field
        node[:_treeview] = Dict{Symbol, Any}(
            :children => [],
            :expanded => false,
            :filtered => false,
            :level => 0,
        )

        id_2_node_map[Tables.getcolumn(row, Symbol(id))] = node
    end

    # move children nodes under their PrepareTableOptions
    # looping on the grid instead of the id_2_node_map in order to keep the original order of items in the table
    # without the need to use OrderedDict
    final_tree = []
    for (i, row) in enumerate(rows)
        k = Tables.getcolumn(row, Symbol(id))
        v = id_2_node_map[k]
        parent = v[ Symbol(parentId)]
        if (parent === missing || parent == "" || parent == -1)
            push!(final_tree, v)
        else
            parent_node = id_2_node_map[parent]
            push!(parent_node[:_treeview][:children], v)
        end
    end

    res = OrderedDict()

    rows = Tables.rows(table)
    schema = Tables.schema(rows)

    names = schema.names
    types = schema.types

    # create the column definitions
    res["columnDefs"] = [
        OrderedDict(
            "headerName" => string(n),
            # "editable" => cell_changed !== nothing,
            "headerTooltip" => string(types[i]),
            "field" => string(n),
            "sortable" => true,
            "resizable" => true,
            "type" => types[i] <: Union{Missing, T where T <: Number} ? "numericColumn" : nothing,
            "filter" => types[i] <: Union{Missing, T where T <: Dates.Date} ? "agDateColumnFilter" :
                    types[i] <: Union{Missing, T where T <: Number} ? "agNumberColumnFilter" : true,
        ) for (i, n) in enumerate(schema.names)
    ]

    res["defaultColDef"] = OrderedDict(
        "width" => 150,
        "editable" => true,
        "filter" => "agTextColumnFilter",
        "floatingFilter" => true,
        "resizable" => true,
    )

    # res["rowData"] = table2json(schema, rows)
    res["rowData"] = final_tree

    res["_metaData"] = OrderedDict(
        "ag_grid_materializer" => Tables.materializer(table),
        "ag_grid_schema" => schema,
    )

    return res;
end

"""
    getTypes(m::Any, arr::Vector{Any})

This function has no place in this package and is just a utility function that fills an array with information of types in the given module.
The function is here in order to test the agGrid tree view capabilities but it has nothing to do with UI or Visualization.
"""
function getTypes(m::Any, arr::Vector{Any})
    # key = string(m)

    item = Dict{Any, Any}(
        :name => string(m),
        :type => string(m), # should be just m
        :desc => "",
        :children => Vector{Any}(),
        :file => nothing,
        :line => nothing,
        :docs => "",
        :hasDocs => false,
        :module => "?",
        :symbol => "?",
    )

    if (m isa Module)
        sym = names(m; all=true)
        for s in sym
            if (occursin("#", string(s)))
                continue
            end
            if (string(s) in ["eval", "include"])
                continue
            end

            local type = nothing
            try
                type = getfield(m, s)
            catch e
                @warn "Can't process [$m].[$s]\n$e\nskipping..."
                continue
            end

            # filter the actual module we are scanning - it appears in the list of names
            if (type === m)
                continue
            end

            stat = getTypes(type, item[:children])

            # special treatment for names of primitie types
            if (stat)
                if (isprimitivetype(typeof(type)) || type isa AbstractString || type isa Function)
                    item[:children][end][:name] = "$m.$s"
                    item[:children][end][:type] = "$m.$s"
                end
                item[:children][end][:module] = string(m)
                item[:children][end][:symbol] = string(s)

                # try to get documentation
                binding = Docs.Binding(m, s)
                documentation = Markdown.plain(REPL.doc(binding))
                # if item has documentation we want to add the summary of it anyway
                if (!startswith(documentation, "No documentation found."))

                    item[:children][end][:hasDocs] = true
                    alias = Docs.aliasof(binding)
                    extra_doc = Markdown.plain(REPL.summarize(alias, Union{}))
                    # @info documentation extra_doc[24:end]

                    documentation = documentation * "-------------------------------" * extra_doc[25:end]
                end

                item[:children][end][:docs] = documentation

            end

        end
        len = length(item[:children])
        item[:desc] = "Module containing $len items"
        push!(arr, item)
    elseif (m isa DataType)
        item[:desc] = "DataType Super[$(supertype(m))] Mutable[$(getfield(m, :mutable))]]"
        push!(arr, item)
    elseif (isprimitivetype(typeof(m)))
        item[:desc] = "PrimitiveType Type[$(typeof(m))] value[$m]"
        push!(arr, item)
    elseif (isprimitivetype(typeof(m)))
        item[:desc] = "PrimitiveType Type[$(typeof(m))] value[$m]"
        push!(arr, item)
    elseif (m isa AbstractString)
        item[:desc] = "AbstractString Type[$(typeof(m))] value[$m]"
        push!(arr, item)
    elseif (m isa Function)
        # item[:desc] = "Function Super[$(supertype(m))] Mutable[$(getfield(m, :mutable))]]"
        meths = methods(m)
        if (length(meths) == 1)
            item[:desc] = "Function Signature[$(meths.ms[1].sig)]"
            item[:file] = string(meths.ms[1].file)
            item[:line] = meths.ms[1].line
        else
            item[:desc] = "Function Signature[$(meths)]"
        end
        push!(arr, item)
    elseif (m isa UnionAll)
        item[:desc] = "UnionAll Body[$(getfield(m, :body))]"
        push!(arr, item)
    else
        @warn "Don't know how to handle [$m] [$(typeof(m))]"
        return false
    end
    return true
end

function toJSONTableRow(item, id, parent_id)
    # shorten the name to the part after the last period

    short_name = item[:name]
    pos = findlast(".", short_name)
    if (pos !== nothing)
        short_name = short_name[pos[1]+1:end]
    end

    return Dict(
        "_id" => id,
        "_parentId" => parent_id,
        "name" => short_name,
        "type" => item[:type],
        "desc" => item[:desc],
        "file" => item[:file],
        "line" => item[:line],
        "docs" => item[:docs],
        "hasDocs" => item[:hasDocs],
        "module" => item[:module],
        "symbol" => item[:symbol],
    )
end

function toJSONTableRecursive(data::Vector{Any}, parent_id::Int64, json_data::Vector{Any})
    for item in data
        id = length(json_data)
        row = toJSONTableRow(item, id, parent_id)
        push!(json_data, row)

        toJSONTableRecursive(item[:children], id, json_data)
    end
end

"""
    toJSONTable(data::Vector{Any}, parent_id::Int64 = -1)

Utility function to convert the result of `getTypes` in to a JSON table (Tables.jl compatiable data structure)
"""
function toJSONTable(data::Vector{Any}, parent_id::Int64 = -1)
    json_data = Vector{Any}()

    toJSONTableRecursive(data, parent_id, json_data)

    # return JSONTables.jsontable(json_data)
    # return json_data
    JSONTables.jsontable(JSON.json(json_data))
end

function toJSONTable(jsonStr::String)
    JSONTables.jsontable(jsonStr)
end

end # module GridUtils


text:
# Libraries
using Distributed;
@everywhere include("./../src/ElasticNetVAR.jl");
@everywhere using Main.ElasticNetVAR;
using Random;
using FileIO, XLSX, DataFrames, BSON;

# Load data
data_raw = DataFrame(XLSX.readtable("./data/data.xlsx", "FX")...);

# Data
data = data_raw[3:end, 3:end] |> JArray{Float64,2};

# Store dates
date = data_raw[3:end,2];

# Skip EA19 and New Zealand
data = data[:, [1;2;4;collect(6:size(data,2))]] |> JArray{Float64,2};

# Transpose data
data = permutedims(data);

# End presample
end_ps = 80; # 1979-12-31

# End validation sample
end_vs = 160; # 1999-12-31

# Data used for validation
data_validation = data[:, 1:end_vs];
n, T = size(data_validation);

# Select optimal d
@info "Running optimal_d(n, T)";
d = optimal_d(n, T);

# Set options for the selection problem
p_grid=[1, 5]; λ_grid=[1e-4, 5]; α_grid=[0, 1]; β_grid=[1, 5];
vs = ValidationSettings(4, data_validation, t0=end_ps, subsample=d/(n*T), max_samples=5000, log_folder_path=".");
hg = HyperGrid(p_grid, λ_grid, α_grid, β_grid, 1000);

# Run algorithm
Random.seed!(1);
@info "Running select_hyperparameters(vs, hg)";
candidates, errors = select_hyperparameters(vs, hg);

# Save to file
save("./res.bson", Dict("data" => data,
                        "date" => date,
                        "end_ps" => end_ps,
                        "end_vs" => end_vs,
                        "d" => d,
                        "vs" => vs,
                        "hg" => hg,
                        "candidates" => candidates,
                        "errors" => errors));


text:
function u1_unit(nbit::Int, i::Int, j::Int)
    chain(nbit, put(nbit, i=>Rz(0)),
    put(nbit, j=>Rz(0)),
    put(nbit, (i,j)=>rot(SWAP, 0))
    )
end

function u1_circuit(nbit_measure::Int, nbit_virtual::Int, nlayer::Int, nrepeat::Int, entangler_pairs)
    circuit = sequence()
    nbit_used = nbit_measure + nbit_virtual

    for i=1:nrepeat
        unit = chain(nbit_used)
        for j=1:nlayer
            push!(unit, chain(nbit_used, u1_unit(nbit_used, i, j) for (i,j) in entangler_pairs))
            for k=1:(i==nrepeat ? nbit_used : nbit_measure)
                put(nbit_used, k=>Rz(0))
            end
        end
        push!(circuit, unit)
    end
    dispatch!(circuit, :random)
end

function model(::Val{:u1}; nbit = 20, V=4, B=4096, nlayer=5, pairs)
    nrepeat = (nbit - V)
    c = u1_circuit(1, V, nlayer, nrepeat, pairs) |> autodiff(:QC)
    chem = QuantumMPS(1, V, 0, c, zero_state(V+1, B), [i%2 for i=1:nbit])
    chem
end

max_stars_repo_name: Roger-luo/QuantumMPS

max_issues_repo_name: Roger-luo/QuantumMPS

max_forks_repo_name: Roger-luo/QuantumMPS


text:
using Test, Astro

print("25.a Sun position, low precision")
# Meeus page 165
# Calculate the Sun's position on 1992 October 13 at 0h TD
# should be geometric long mean equinox of date 199degrees  54 minutes 26.18 seconds
#          apparent longitude 199d54m23.56s
#          apparent latitude +0.72 seconds
#          radius vector 0.99760853
#          apparent RA 13h13m30.749s
#          apparent decl -1d47m01.74s
L, R = longitude_radius_low(2448908.5)
@test rad2deg(L) ≈  199.90988 atol= 1e-2# should be 1e-5 ?
@test isapprox(R, 0.99766, atol=1e-5)
L = apparent_longitude_low(2448908.5, L)
println("\nL is $(rad2deg(L)) degrees, R is $R au")
@test rad2deg(L) ≈  199.90895 atol= 1e-2# should be 1e-5 ?
println(" passed\n\n")

print("28.a Equation of time")
# Find equation of time on 1992 October 13 at 0h
# should be 13.70949 minutes, 13m42.6s
E = equation_time(cal_to_jd(1992,10,13))
@test E ≈  13.7 atol= .5# half a second?
E = equation_time(cal_to_jd(1980,7,27.5))
@test isapprox(E, -6.4, atol=.1)
println(" passed")

#=

map(rad2deg,sun_dimension3(cal_to_jd(1992,10,13)))
 2014-10-08T16:59:16 loaded data for Earth 2.4489085e6 (1992.0,10.0,13.0)
(199.90729724204937,0.00020664594475074667,57.158757815790885)

=#

print(" Duffett-Smith: RA and DEC of the sun of July 27 1980")
# Duffett-Smith page 83 Practical Astronomy with your calculator
# What were the RA and DEC of the sun of July 27 1980
# sun's longitude should be 124.108828
# > ra should be 08h25m44s
# dec should be 19h13m53s
jd = cal_to_jd(1980,7,27)
Lsun, radius = longitude_radius_low(cal_to_jd(1980,7,27))
# ecliptic latitude is 0
Bsun = 0
set_obliquity(obliquity_high(jd))
ra, decl = ecl_to_equ(Lsun, Bsun, currentobliquity)

#println(" For July 27, 1980, the Sun's RA and Dec were")
#println(" RA: $(d_to_dms(rad2deg(ra))), Decl: $(d_to_dms(rad2deg(decl)))")
# RA is (126.0,26.0,8.426670651592705), Decl is (19.0,13.0,52.04805574268332)

#=

Lsun, radius = longitude_radius_low(cal_to_jd(1980,7,27))
> (2.166121184301187,1.0155158031473217)

rad2deg(2.166121184301187)
> 124.10960177433758        - sun's longitude looks good

Bsun = 0
obliquity = obliquity_high(jd)

ra, decl = ecl_to_equ(Lsun, Bsun, obliquity)

ecl_to_equ returns:
        Right ascension in radians (2 pi is 24 hours)
        Declination in radians

> (2.206718804590345,0.335646440686569)

map(rad2deg, (ra,decl))
> (126.435674075181,19.231124459928523)
θθ

=#

print("\n\n25.b Sun position, high precision, page 169")

jd_to_cal(2448908.5)

L, B, R = sun_dimension3(cal_to_jd(1992,10,13))

# meeus gives L as 19°.907372, B as -0.000179, R as 0.99760775
# L, B, R = 199.91598259307455, 0.00020664594475074667, 0.9976085202355934

@test isapprox(rad2deg(L), 199.907372, atol=0.1)
@test isapprox(rad2deg(B) , -0.000179, atol=.1)
@test isapprox(R, 0.9976, atol=.1)

L, B = vsop_to_fk5(2448908.5, L, B)
@test isapprox(rad2deg(L), 199.907347, atol=.1)

aberration = aberration_low(R)
@test isapprox(rad2deg(aberration) * 3600, -20.539, atol=0.001)
println(" passed")

print("\n\n26.a 1 rectangular coordinates of the sun, relative to the mean equinox of the date")
X, Y, Z = rectangular_md(2448908.5)
@test isapprox(X, -0.9379952, atol=.1) #
@test isapprox(Y, -0.3116544, atol=0.1) # ..
@test isapprox(Z, -0.1351215, atol=0.1) # ..
println(" passed")


text:
function plot_stockout(fig, gp,ax,x,y,smoothing,ticks)
    xpos = fill(0, length(unique(x)))
    ypos = zeros(length(xpos))
    yall = zero(ypos)
    xall = zero(xpos)
    last = first(x)
    idx = 1
    for (new, b) in zip(x, y)
        idx += new !== last
        last = new
        if  b > 0
            ypos[idx] += b
            xpos[idx] += 1
        end
        yall[idx] += b
        xall[idx] += 1
    end
    ypos ./= [n == 0 ? one(n) : n for n in xpos]
    yall ./= xall
    if smoothing
        ypos = accumulate!((o,n) -> o*0.8+n*0.2, ypos,ypos)
        yall = accumulate!((o,n) -> o*0.8+n*0.2, yall,ypos)
    end
    l1 = lines!(ax,ticks, ypos, color = :red)
    l2 = lines!(ax,ticks, yall, color = :green)
    gp[2,1] = Legend(fig, [l2,l1], ["mean stockout", "mean ⊕ stockout"], tellwidth = false, tellhight = false, valign = :top , halign = :right)
    nothing
end

function plot_fillrate(fig, gp,ax,x,y,smoothing,ticks)
    ax.yticks = 0:0.1:1
    xone = fill(0, length(unique(x)))
    yone = zeros(length(xone))
    yall = zero(yone)
    xall = zero(xone)
    last = first(x)
    idx = 1
    for (new, b) in zip(x, y)
        idx += new !== last
        last = new
        if b == 1
            yone[idx] += 1
        end
        yall[idx] += b
        xall[idx] += 1
    end
    yall ./= xall
    yone ./= xall
    if smoothing
        yall = accumulate!((o,n) -> o*0.8+n*0.2, yall,yall)
    end
    b = barplot!(ax,ticks, yone, color = :gray80)
    l = lines!(ax,ticks, yall, color = :red)
    gp[2,1] = Legend(fig, [b, l], ["α service", "β service"], tellwidth = false, tellhight = false, valign = :top , halign = :right)
    nothing
end


text:
# Reference: Table 3 in Roquet et. al., "Defining a Simplified yet 'Realistic'
# Equation of State for Seawater", Journal of Physical Oceanography (2015).
optimized_roquet_coeffs = Dict(
                  :linear => (R₀₁₀ = -1.775e-1, R₁₀₀ = 7.718e-1, R₀₂₀ = 0,         R₀₁₁ = 0,          R₂₀₀ = 0,         R₁₀₁ = 0,         R₁₁₀ = 0),
               :cabbeling => (R₀₁₀ = -0.844e-1, R₁₀₀ = 7.718e-1, R₀₂₀ = -4.561e-3, R₀₁₁ = 0,          R₂₀₀ = 0,         R₁₀₁ = 0,         R₁₁₀ = 0),
:cabbeling_thermobaricity => (R₀₁₀ = -0.651e-1, R₁₀₀ = 7.718e-1, R₀₂₀ = -5.027e-3, R₀₁₁ = -2.5681e-5, R₂₀₀ = 0,         R₁₀₁ = 0,         R₁₁₀ = 0),
                :freezing => (R₀₁₀ = -0.491e-1, R₁₀₀ = 7.718e-1, R₀₂₀ = -5.027e-3, R₀₁₁ = -2.5681e-5, R₂₀₀ = 0,         R₁₀₁ = 0,         R₁₁₀ = 0),
            :second_order => (R₀₁₀ =  0.182e-1, R₁₀₀ = 8.078e-1, R₀₂₀ = -4.937e-3, R₀₁₁ = -2.4677e-5, R₂₀₀ = -1.115e-4, R₁₀₁ = -8.241e-6, R₁₁₀ = -2.446e-3)
)

"""
    RoquetIdealizedNonlinearEquationOfState{F, C, T} <: AbstractNonlinearEquationOfState

Parameters associated with the idealized nonlinear equation of state proposed by
Roquet et al., "Defining a Simplified yet 'Realistic' Equation of State for Seawater",
Journal of Physical Oceanography (2015).
"""
struct RoquetIdealizedNonlinearEquationOfState{F, C, FT} <: AbstractNonlinearEquationOfState
                   ρ₀ :: FT
    polynomial_coeffs :: C
end

type_convert_roquet_coeffs(FT, coeffs) = NamedTuple{propertynames(coeffs)}(Tuple(FT(R) for R in coeffs))

"""
    RoquetIdealizedNonlinearEquationOfState([FT=Float64,] flavor, ρ₀=1024.6,
                                            polynomial_coeffs=optimized_roquet_coeffs[flavor])

Returns parameters for the idealized polynomial nonlinear equation of state with
reference density `ρ₀` and `polynomial_coeffs` proposed by Roquet et al., "Defining a
Simplified yet 'Realistic' Equation of State for Seawater", Journal of Physical
Oceanography (2015). The default reference density is `ρ₀ = 1024.6 kg m⁻³`, the average
surface density of seawater in the world ocean.

The `flavor` of the nonlinear equation of state is a symbol that selects a set of optimized
polynomial coefficients defined in Table 2 of Roquet et al., "Defining a Simplified yet
'Realistic' Equation of State for Seawater", Journal of Physical Oceanography (2015), and
further discussed in the text surrounding equations (12)--(15). The optimization minimizes
errors in estimated horizontal density gradient estiamted from climatological temperature
and salinity distributions between the 5 simplified forms chosen by Roquet et. al
and the full-fledged [TEOS-10](http://www.teos-10.org) equation of state.

The equations of state define the density anomaly `ρ′`, and have
the polynomial form

    `ρ′(T, S, D) = Σᵢⱼₐ Rᵢⱼₐ Tⁱ Sʲ Dᵃ`,

where `T` is conservative temperature, `S` is absolute salinity, and `D` is the
geopotential depth, currently just `D = -z`. The `Rᵢⱼₐ` are constant coefficients.

Flavors of idealized nonlinear equations of state
=================================================

    - `:linear`: a linear equation of state, `ρ′ = R₁₀₀ * T + R₀₁₀ * S`

    - `:cabbeling`: includes quadratic temperature term,
                    `ρ′ = R₁₀₀ * T + R₀₁₀ * S + R₀₂₀ * T^2`

    - `:cabbeling_thermobaricity`: includes 'thermobaricity' term,
                                   `ρ′ = R₁₀₀ * T + R₀₁₀ * S + R₀₂₀ * T^2 + R₀₁₁ * T * D`

    - `:freezing`: same as `:cabbeling_thermobaricity` with modified constants to increase
                   accuracy near freezing

    - `:second_order`: includes quadratic salinity, halibaricity, and thermohaline term,
                       `ρ′ = R₁₀₀ * T + R₀₁₀ * S + R₀₂₀ * T^2 + R₀₁₁ * T * D`
                             + R₂₀₀ * S^2 + R₁₀₁ * S * D + R₁₁₀ * S * T`

Example
=======

julia> using Oceananigans

julia> eos = Oceananigans.RoquetIdealizedNonlinearEquationOfState(:cabbeling);

julia> eos.polynomial_coeffs
(R₀₁₀ = -0.0844, R₁₀₀ = 0.7718, R₀₂₀ = -0.004561, R₀₁₁ = 0.0, R₂₀₀ = 0.0, R₁₀₁ = 0.0, R₁₁₀ = 0.0)

References
==========

    - Roquet et al., "Defining a Simplified yet 'Realistic' Equation of State for
      Seawater", Journal of Physical Oceanography (2015).

    - "Thermodynamic Equation of State for Seawater" (TEOS-10), http://www.teos-10.org
"""
function RoquetIdealizedNonlinearEquationOfState(FT, flavor=:cabbeling_thermobaricity;
                                                 polynomial_coeffs=optimized_roquet_coeffs[flavor], ρ₀=1024.6)
    typed_coeffs = type_convert_roquet_coeffs(FT, polynomial_coeffs)
    return RoquetIdealizedNonlinearEquationOfState{flavor, typeof(typed_coeffs), FT}(ρ₀, typed_coeffs)
end

RoquetIdealizedNonlinearEquationOfState(flavor::Symbol=:cabbeling_thermobaricity; kwargs...) =
    RoquetIdealizedNonlinearEquationOfState(Float64, flavor; kwargs...)

@inline ρ′(Θ, Sᴬ, D, eos::RoquetIdealizedNonlinearEquationOfState) =
    @inbounds (   eos.polynomial_coeffs.R₁₀₀ * Sᴬ
                + eos.polynomial_coeffs.R₀₁₀ * Θ
                + eos.polynomial_coeffs.R₀₂₀ * Θ^2
                + eos.polynomial_coeffs.R₀₁₁ * Θ * D
                + eos.polynomial_coeffs.R₂₀₀ * Sᴬ^2
                + eos.polynomial_coeffs.R₁₀₁ * Sᴬ * D
                + eos.polynomial_coeffs.R₁₁₀ * Sᴬ * Θ )

@inline thermal_expansion(Θ, Sᴬ, D, eos::RoquetIdealizedNonlinearEquationOfState) =
    @inbounds (       eos.polynomial_coeffs.R₀₁₀
                + 2 * eos.polynomial_coeffs.R₀₂₀ * Θ
                +     eos.polynomial_coeffs.R₀₁₁ * D
                +     eos.polynomial_coeffs.R₁₁₀ * Sᴬ )

@inline haline_contraction(Θ, Sᴬ, D, eos::RoquetIdealizedNonlinearEquationOfState) =
    @inbounds (       eos.polynomial_coeffs.R₁₀₀
                + 2 * eos.polynomial_coeffs.R₂₀₀ * Sᴬ
                    + eos.polynomial_coeffs.R₁₀₁ * D
                    + eos.polynomial_coeffs.R₁₁₀ * Θ )


text:
"""
    lambdipole(U, R, grid::TwoDGrid; center=(mean(g.x), mean(g.y))

Return the 2D vorticity field of the Lamb dipole with strength `U` and radius `R`, centered on
`center=(xc, yc)` and on the `grid`. The default value of `center` is the middle of the grid.
"""
function lambdipole(U, R, grid::TwoDGrid{T, A}; center=(mean(grid.x), mean(grid.y))) where {T, A}
  firstzero = 3.8317059702075123156
  k = firstzero / R # dipole wavenumber for radius R in terms of first zero of besselj
  q0 = -2U * k / besselj(0, k * R) # dipole amplitude for strength U and radius R
  x, y = gridpoints(grid)
  xc, yc = center
  r = @. sqrt( (x - xc)^2 + (y - yc)^2 )
  CUDA.@allowscalar besselj1 = A([besselj(1, k * r[i, j]) for i=1:grid.nx, j=1:grid.ny])
  q = @. q0 * besselj1 * (y - yc) / r
  @. q = ifelse(r >= R, 0, q)
  return q
end

"""
    peakedisotropicspectrum(g, kpeak, E0; mask=mask, allones=false)
Generate a real and random two-dimensional vorticity field q(x, y) with
a Fourier spectrum peaked around a central non-dimensional wavenumber kpeak and
normalized so that its total energy is E0.
"""
function peakedisotropicspectrum(g::TwoDGrid{T, A}, kpeak::Real, E0::Real; mask=ones(size(g.Krsq)), allones=false) where {T, A}
  if g.Lx !== g.Ly
    error("the domain is not square")
  else
    k0 = kpeak * 2π / g.Lx
    modk = sqrt.(g.Krsq)
    psik = A(zeros(T, (g.nk, g.nl)))
    psik = @. (modk^2 * (1 + (modk / k0)^4))^(-0.5)
    CUDA.@allowscalar psik[1, 1] = 0.0
    phases = randn(Complex{T}, size(g.Krsq))
    phases_real, phases_imag = real.(phases), imag.(phases)
    phases = A(phases_real) + im * A(phases_imag)
    psih = @. phases * psik
    if allones; psih = psik; end
    psih = psih .* A(mask)
    Ein = real(sum(g.Krsq .* abs2.(psih) / (g.nx * g.ny)^2))
    psih = psih * sqrt(E0 / Ein)
    CUDA.@allowscalar q = A(-irfft(g.Krsq .* psih, g.nx))
  end
  return q
end