text:
import MLJBase
using DataFrames
import DataFramesMeta
using BioSequences
using ShapML
using OrderedCollections
using ProgressMeter
using ColorSchemes
using GLMakie
using CairoMakie
using CategoricalArrays
using MLJ
using NamedArrays
using Statistics
using Serialization
using Distributions
using StatsBase: weights
using StatsBase: sample
using StatsBase: countmap

d_base2int = Dict(
    # DNA_N => 0,
    DNA_A => 1,
    DNA_C => 2,
    DNA_G => 3,
    DNA_T => 4,
)
d_int2base = Dict(i => base for (base, i) in d_base2int)
d_int2string = Dict(i => string(Char(base)) for (base, i) in d_base2int)
all_bases = sort(collect(keys(d_base2int)))
base_levels = Int8.(sort(collect(values(d_base2int))))

function convert_DNA_to_int(bases, datatype = Int8)
    bases_int = zeros(datatype, length(bases))
    for (i, base) in enumerate(bases)
        bases_int[i] = d_base2int[base]
    end
    return bases_int
end

function convert_DNAs_to_categorical(bases, datatype = Int8)
    bases_int = convert_DNA_to_int(bases)
    return CategoricalArray{datatype,1,UInt8}(bases_int, levels = base_levels)
end

function convert_DNA_dataframe_to_categorical(df)
    df = copy(df)
    for column in names(df)
        df[!, column] = convert_DNAs_to_categorical(df[!, column])
    end
    return df
end

# function bases2dataframe(bases, prefix)
#     ohe = all_bases .== permutedims(bases)
#     base_names = string.(Char.(all_bases))
#     column_names = [Symbol("$(prefix)_$(base)") for base in base_names]
#     return DataFrame(table(permutedims(ohe); names = column_names))
# end

# function dataframe2ohe(df)
#     hcat([bases2dataframe(df[!, column], column) for column in names(df)]...)
# end

function get_Xy(filename, N_rows = 1000)

    df = deserialize(filename).df
    if 0 < N_rows && N_rows < size(df, 1)
        sample_rows = sample(1:nrow(df), N_rows, replace = false)
        df = df[sample_rows, :]
    end

    X = convert_DNA_dataframe_to_categorical(select(df, Not(:y)))
    y = CategoricalArray{Int8,1,Int8}(df.y)
    return X, y
end

function get_matching_models(X, y)
    hot_model = OneHotEncoder()
    hot = machine(hot_model, X)
    fit!(hot)
    Xt = MLJ.transform(hot, X)
    # models(matching(X, y))
    # models(matching(Xt, y))
    models() do model
        matching(model, Xt, y) && model.prediction_type == :probabilistic #&&
        # model.is_pure_julia
    end
end

#%%

function transpose(df::DataFrame)
    df_T = DataFrame([[names(df)]; collect.(eachrow(df))], [:column; Symbol.(axes(df, 1))])
    return df_T
end

function get_df_LR_long(mach::Machine)
    params_raw = fitted_params(mach)
    params = hcat(
        # DataFrame(intercept = params_raw.logistic_classifier.intercept),
        DataFrame(params_raw.logistic_classifier.coefs),
    )
    params_T = rename(transpose(params), :2 => :values)
    return params_T
end

function get_df_LR_wide(df::DataFrame)
    N_rows = Int(size(df, 1) // 4)
    m = permutedims(reshape(df.values, (4, N_rows)))
    return DataFrame(m, [:A, :C, :G, :T])
end

#%%

function logodds2prob(logodds::Float64)
    return exp(logodds) / (1 + exp(logodds))
end

#%%

function get_glm_fitresult(mach::Machine)
    fitresult, decode = mach.fitresult.predict.machine.fitresult
    return fitresult
end

#%%

function print_performance(eval, model_name, multiplier = 100, sigdigits = 3)
    μ = round(multiplier * eval.measurement[2], sigdigits = sigdigits)
    per_fold = eval.per_fold[2]
    σ = round(multiplier * std(per_fold) / sqrt(length(per_fold)), digits = sigdigits)
    println("$model_name = ($μ ± $σ) %")
end;

#%%

function plot_LR_fit_coefficients(df)

    max_pos = size(df, 1)
    positions = 1:max_pos
    bases = [:A, :C, :G, :T]

    f = Figure()
    ax = Axis(
        f[1, 1],
        title = "LR fit coefficients as a function of position",
        xlabel = "Position",
        ylabel = "Fit coefficient (LR)",
        limits = (0.5, max_pos + 0.5, nothing, nothing),
        # xticks = 1:2:seq_length,
    )

    colormap = [x for x in ColorSchemes.Set1_9.colors]
    for (i, base) in enumerate(bases)
        scatterlines!(
            ax,
            positions,
            df[:, base],
            color = colormap[i],
            markercolor = colormap[i],
            label = String(base),
        )
    end

    axislegend(position = :cb, nbanks = 4)
    return f
end

#%%
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
#                     ██   ██  ██████  ██████
#                      ██ ██  ██       ██   ██
#                       ███   ██   ███ ██████
#                      ██ ██  ██    ██ ██   ██
#                     ██   ██  ██████  ██████
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -

# XGB = @load XGBoostClassifier
# pipe_xgb = @pipeline(OneHotEncoder, XGB(), name = "pipeline_xgb",);
# machine_xgb = machine(pipe_xgb, X, y);
# fit!(machine_xgb)
# yhat_xgb = predict(machine_xgb, X);
# # predict_mean(machine_xgb, X)
# acc_xgb = accuracy(predict_mode(machine_xgb, X), y)
# println("Accuracy, XGB, ", round(acc_xgb, digits = 3))

# evaluate(
#     pipe_xgb,
#     X,
#     y,
#     resampling = CV(nfolds = 5, shuffle = true, rng = 1234),
#     measures = [LogLoss(), Accuracy()],
#     verbosity = 0,
# )

#%%

function convert_type(df::DataFrame, type::DataType)
    mapcols(col -> int(col; type = type) .- 1, df)
end

#%%

function MLJBase.accuracy(yhat::MLJBase.UnivariateFiniteArray, y::MLJ.CategoricalVector)
    return accuracy(mode.(yhat), y)
end

function MLJBase.confusion_matrix(
    yhat::MLJBase.UnivariateFiniteArray,
    y::MLJ.CategoricalVector,
)
    confusion_matrix(coerce(mode.(yhat), OrderedFactor), coerce(y, OrderedFactor))
end

#%%

function is_nested(x)
    if eltype(x) <: AbstractArray
        return true
    else
        return false
    end
end

#%%

function _plot_roc_check(y_hats, names_roc)
    if length(y_hats) != length(names_roc)
        throw(ArgumentError("'y_hats' and 'names_roc' is not the same length."))
    end
    if !is_nested(y_hats)
        y_hats = (y_hats,)
    end
    if isa(names_roc, String)
        names_roc = (names_roc,)
    end
    return y_hats, names_roc
end

function plot_roc(y_hats, y, names_roc)
    y_hats, names_roc = _plot_roc_check(y_hats, names_roc)

    f = Figure()
    ax = Axis(f[1, 1], xlabel = "FPR", ylabel = "TPR", title = "ROC Curve")
    for (y_hat, name) in zip(y_hats, names_roc)
        fprs, tprs, ts = roc_curve(y_hat, y)
        auc_ = round(100 * area_under_curve(y_hat, y), digits = 1)
        legend_name = "$(name), AUC: $(auc_)%"
        lines!(ax, fprs, tprs, label = legend_name)
    end
    axislegend(ax, position = :rb)
    f
end

#%%

# A wrapper function for ShapML.
function shap_predict_function(model, data)
    y_pred = int(predict_mode(model, data), type = Int64) .- 1
    data_pred = DataFrame(y_pred = y_pred)
    return data_pred
end

function shap_compute(; mach, X, N = 100, sample_size = 100, use_reference = false)

    # Compute Shapley feature-level predictions for N instances.
    explain = copy(X[1:N, :])

    # An optional reference population to compute the baseline prediction.
    reference = use_reference ? copy(X) : nothing

    # sample_size  # Number of Monte Carlo samples.

    # Compute stochastic Shapley values.
    data_shap = ShapML.shap(
        explain = explain,
        reference = reference,
        model = mach,
        predict_function = shap_predict_function,
        sample_size = sample_size,
        seed = 1,
    )

    return data_shap

end

#%%

function get_shap_global_plot_data(data_shap; sort = false)

    data_plot = combine(
        groupby(data_shap, :feature_name),
        [:shap_effect] => (x -> mean(abs.(x))) => :mean_effect,
        [:shap_effect] => (x -> std(abs.(x))) => :std_effect,
        [:shap_effect] => (x -> std(abs.(x)) / sqrt(length(x))) => :sdom_effect,
    )

    if sort
        data_plot = sort(data_plot, order(:mean_effect, rev = true))
    end

    baseline = round(data_shap.intercept[1], digits = 1)

    tbl = (
        x = 1:size(data_plot)[1],
        y = data_plot.mean_effect,
        y_std = data_plot.std_effect,
        y_sdom = data_plot.sdom_effect,
        name = data_plot.feature_name,
        baseline = baseline,
        df = data_plot,
    )

    return tbl
end

function plot_shap_global(
    data_shap;
    sort = false,
    add_errorbars = false,
    savefig = true,
    filename = "fig_shap_global",
)

    tbl = get_shap_global_plot_data(data_shap; sort = sort)

    f = Figure()
    ax = Axis(
        f[1, 1],
        yticks = (tbl.x, tbl.name),
        title = "Feature Importance",
        xlabel = "|Shapley effect| (baseline = $(tbl.baseline))",
        ylabel = "Feature",
        limits = (0, nothing, 0, nothing),
        yreversed = true,
    )
    barplot!(
        ax,
        tbl.x,
        tbl.y,
        direction = :x,
        # bar_labels = :y,
        # flip_labels_at = 0.3,
        # label_formatter = x-> "Flip at $(x)?",
    )
    if add_errorbars
        errorbars!(ax, tbl.y, tbl.x, tbl.y_sdom, whiskerwidth = 12, direction = :x)
    end

    if !savefig
        return f
    else
        save("./figures/$filename.pdf", f)
    end
end

function plot_shap_variable(data_shap; feature = "x1", xlim = nothing, ylim = nothing)

    data_plot = DataFramesMeta.@chain data_shap begin
        DataFramesMeta.@rsubset :feature_name == feature
    end
    # data_plot = data_shap[data_shap.feature_name .== "x1", :]

    baseline = round(data_shap.intercept[1], digits = 1)

    tbl = (
        x = int(data_plot.feature_value, type = Int64),
        y = data_plot.shap_effect,
        xticks = 1:4,
        name = string.(levels(data_plot.feature_value)),
        baseline = baseline,
    )

    f = Figure()
    ax = Axis(
        f[1, 1],
        xticks = (tbl.xticks, tbl.name),
        title = "Feature Importance for $feature",
        xlabel = feature,
        ylabel = "|Shapley effect| (baseline = $(tbl.baseline))",
    )
    violin!(ax, tbl.x, tbl.y, show_median = true)
    if xlim !== nothing
        xlims!(ax, xlim) # as vector
    end
    if ylim !== nothing
        ylims!(ax, ylim) # as vector
    end
    return f
end

function plot_shap_variables(
    data_shap;
    xlim = nothing,
    ylim = nothing,
    savefig = true,
    filename = "fig_shap_global",
)

    if xlim === nothing
        xlim = (0.5, 4.5)
    end
    if ylim === nothing
        ylim = (minimum(data_shap.shap_effect) - 0.1, maximum(data_shap.shap_effect) + 0.1)
    end

    features = unique(data_shap.feature_name)
    d_figures = OrderedDict{String,Figure}()

    for feature in features
        d_figures[feature] =
            plot_shap_variable(data_shap; feature = feature, xlim = xlim, ylim = ylim)
    end

    if !savefig
        return d_figures

    else

        files = String[]
        for (feature, fig) in d_figures
            file = "./figures/$(filename)__$feature.pdf"
            save(file, fig)
            push!(files, file)
        end
        run(`/usr/local/bin/pdftk $files cat output ./figures/$filename.pdf`)

        # delete files
        for file in files
            rm(file)
        end

    end

end

#current_figure()

#%%

import Random.seed!;

function get_accuracies_pr_base(X, y, train, test)

    half_seq_length = Int((size(X, 2)) / 2)
    N_positions_vec = 1:half_seq_length

    accs_LR = Float64[]
    accs_LGB = Float64[]

    X_org = copy(X)

    p = Progress(sum(N_positions_vec), 1)   # minimum update interval: 1 second
    position = 0
    for N_positions in N_positions_vec

        ProgressMeter.update!(p, position)
        position += N_positions

        X_cut = hcat(X_org[:, 1:N_positions], X_org[:, end-N_positions+1:end])

        seed!(42)
        mach_LR = machine(pipe_LR, selectrows(X_cut, train), selectrows(y, train))
        fit!(mach_LR, verbosity = 0)
        acc_LR = accuracy(predict(mach_LR, selectrows(X_cut, test)), selectrows(y, test))
        push!(accs_LR, acc_LR)

        seed!(42)
        mach_LGB = machine(pipe_LGB, selectrows(X_cut, train), selectrows(y, train))
        fit!(mach_LGB, verbosity = 0)
        acc_LGB = accuracy(predict(mach_LGB, selectrows(X_cut, test)), selectrows(y, test))
        push!(accs_LGB, acc_LGB)

    end

    accs = [
        ("N_positions_vec", N_positions_vec),
        ("Logistic Regression", accs_LR),
        ("LightGBM", accs_LGB),
    ]

    return accs
end

function add_missing_bases!(d)
    k = keys(d)
    for i in 1:4
        if !(i in k)
            d[i] = 0
        end
    end
end

function predict_custom(X_row::DataFrameRow)
    d = countmap(X_row)
    add_missing_bases!(d)

    AT = d[1]+d[4]
    GC = d[3]+d[2]
    if AT > GC
        return 1
    elseif AT == GC
        return rand([0, 1])
    else
        return 0
    end
end

function predict_custom(X::DataFrame, N=-1)
    if N < 0
        N = size(X, 1)
    end
    return [predict_custom(X[i, :]) for i in 1:N]
end

function get_accuracies_pr_base_centered(X, y, train, test; add_analytical = true, add_custom=true)

    half_seq_length = Int((size(X, 2)) / 2)
    N_positions_vec = 0:half_seq_length-1

    accs_LR = Float64[]
    accs_LGB = Float64[]
    accs_custom = Float64[]

    X_org = copy(X)

    p = Progress(sum(N_positions_vec), 1)   # minimum update interval: 1 second
    position = 0
    for N_positions in N_positions_vec

        ProgressMeter.update!(p, position)
        position += N_positions

        middle_idxs = half_seq_length-N_positions:half_seq_length+N_positions+1
        middle_X = X_org[:, middle_idxs]

        seed!(42)
        mach_LR = machine(pipe_LR, selectrows(middle_X, train), selectrows(y, train))
        fit!(mach_LR, verbosity = 0)
        acc_LR = accuracy(predict(mach_LR, selectrows(middle_X, test)), selectrows(y, test))
        push!(accs_LR, acc_LR)

        seed!(42)
        mach_LGB = machine(pipe_LGB, selectrows(middle_X, train), selectrows(y, train))
        fit!(mach_LGB, verbosity = 0)
        acc_LGB =
            accuracy(predict(mach_LGB, selectrows(middle_X, test)), selectrows(y, test))
        push!(accs_LGB, acc_LGB)

        if add_custom
            seed!(42)
            yhat_custom = predict_custom(selectrows(middle_X, test));
            acc_custom = accuracy(yhat_custom, selectrows(y, test))
            push!(accs_custom, acc_custom)
        end

    end

    accs = [
        ("N_positions_vec", N_positions_vec),
        ("Logistic Regression", accs_LR),
        ("LightGBM", accs_LGB),
    ]

    if add_analytical
        analytical_accuracies = compute_analytical_accuracies(
            N_half = half_seq_length,
            p_GC_sig = 0.4,
            p_GC_bkg = 0.5,
            weight = [1, 1],
        )
        push!(accs, ("Analytical", analytical_accuracies))
    end

    if add_custom
        push!(accs, ("Custom", accs_custom))
    end

    return accs
end

function plot_accuracy_function_of_bases(accuracies; ylimits = (nothing, nothing))

    N_positions_vec = accuracies[1][2]
    seq_length = maximum(N_positions_vec)

    colormap = [x for x in ColorSchemes.Set1_9.colors]

    f = Figure()
    ax = Axis(
        f[1, 1],
        title = "Accuracy as a function of number of bases included (symmetric)",
        xlabel = "# bases included (symmetric)",
        ylabel = "Accuracy",
        # limits = (0.5, seq_length + 0.5, 0.634, 0.701),
        limits = (0.5, seq_length + 0.5, ylimits...),
        xticks = 1:2:seq_length,
    )

    for (i, acc) in enumerate(accuracies[2:end])
        scatterlines!(
            ax,
            N_positions_vec,
            acc[2],
            color = colormap[i],
            markercolor = colormap[i],
            label = acc[1],
        )
    end

    axislegend(position = :rb)
    return f
end

function plot_accuracy_function_of_bases_centered(accuracies; ylimits = (nothing, nothing))

    N_positions_vec = accuracies[1][2] .+ 1
    half_seq_length = maximum(N_positions_vec)

    colormap = [x for x in ColorSchemes.Set1_9.colors]

    f = Figure()
    ax = Axis(
        f[1, 1],
        title = "Accuracy as a function of number of bases included (centered)",
        xlabel = "# bases included (centered)",
        ylabel = "Accuracy",
        # limits = (0.5, half_seq_length + 0.5, 0.634, 0.701),
        limits = (0.5, half_seq_length + 0.5, ylimits...),
        xticks = 1:2:half_seq_length,
    )

    for (i, acc) in enumerate(accuracies[2:end])
        scatterlines!(
            ax,
            N_positions_vec,
            acc[2],
            color = colormap[i],
            markercolor = colormap[i],
            label = acc[1],
        )
    end

    axislegend(position = :rb)
    return f
end

#%%

function get_scores(yhat, label = 1)
    return pdf.(yhat, Int8(label))
end

function plot_score_density(yhat, y_test, title = "")

    yscores = get_scores(yhat)

    mask_signal = (y_test .== 1)
    mask_background = .!mask_signal
    yscores_signal = yscores[mask_signal]
    yscores_background = yscores[mask_background]

    f = Figure()
    ax = Axis(
        f[1, 1],
        title = title,
        xlabel = "Score",
        ylabel = "Density",
        limits = (0, 1, 0, nothing),
    )

    density!(
        ax,
        yscores_signal,
        color = (:red, 0.1),
        strokecolor = :red,
        strokewidth = 3,
        strokearound = false,
        label = "Signal",
    )
    density!(
        ax,
        yscores_background,
        color = (:blue, 0.1),
        strokecolor = :blue,
        strokewidth = 3,
        strokearound = false,
        label = "Background",
    )
    axislegend(position = :lt, nbanks = 4)

    return f
end

#%%

function tostring(base::DNA)
    return string(Char(base))
end

function get_firstbase_lastbase_mask(X, firstbase, lastbase)
    return (X[:, 1] .== d_base2int[firstbase]) .&& (X[:, end] .== d_base2int[lastbase])
end

function get_base_stratified_measure(X, yhat, y_test, measure_func = area_under_curve)

    # acc_table["A", "C"]
    # acc_table["Base 1" => "A", "Base 76" => "C"]

    base_stratified_measure = Float64[]
    for firstbase in all_bases
        for lastbase in all_bases
            mask_first_last = get_firstbase_lastbase_mask(X, firstbase, lastbase)
            # acc = measure(mode.(yhat[mask_first_last]), y_test[mask_first_last])
            measure = measure_func(yhat[mask_first_last], y_test[mask_first_last])
            push!(base_stratified_measure, measure)
        end
    end

    base_stratified_measure_2d = permutedims(reshape(base_stratified_measure, (4, 4)))
    measure_table =
        100 .* NamedArray(
            base_stratified_measure_2d,
            (tostring.(all_bases), tostring.(all_bases)),
            ("Base 1", "Base 76"),
        )
    return measure_table
end

function Statistics.mean(x::CategoricalVector)
    return mean(int(x) .- 1)
end

function Base.sum(x::CategoricalVector)
    return sum(int(x) .- 1)
end

function get_base_stratified_signal_proportion(X, y_test, func = mean, datatype = Float64)
    base_stratified_means = datatype[]
    for firstbase in all_bases
        for lastbase in all_bases
            mask_first_last = get_firstbase_lastbase_mask(X, firstbase, lastbase)
            push!(base_stratified_means, func(y_test[mask_first_last]))
        end
    end
    base_stratified_means_2d = permutedims(reshape(base_stratified_means, (4, 4)))
    acc_table =
        100 .* NamedArray(
            base_stratified_means_2d,
            (tostring.(all_bases), tostring.(all_bases)),
            ("Base 1", "Base 76"),
        )
    return acc_table
end

#%%

function get_base_counts_pr_position(X; normalise = true)
    df = vcat(
        [DataFrame(Dict("$i" => sum(X[:, col] .== i) for i = 1:4)) for col in names(X)]...,
    )
    rename!(df, d_int2string)
    if normalise
        df = df ./ size(X, 1)
    end

    return df
end

function make_base_fraction_plot(base_counts_signal, base_counts_background)

    N = size(base_counts_signal, 1)

    f = Figure()
    ax = Axis(
        f[1, 1],
        title = "Base fraction for signal (ancient) and background (modern)",
        xlabel = "Read position",
        ylabel = "Base fraction",
        limits = (0.5, N + 0.5, 0.1, 0.4),
        xticks = [1; collect(5:5:N)],
    )

    colormap = [color for color in ColorSchemes.Set1_9.colors]

    for (i, column) in enumerate(names(base_counts_signal))
        lines!(
            ax,
            1:N,
            base_counts_signal[:, column],
            color = colormap[i],
            markercolor = colormap[i],
            label = "Signal, Base: $column",
        )
    end

    for (i, column) in enumerate(names(base_counts_background))
        lines!(
            ax,
            1:N,
            base_counts_background[:, column],
            color = colormap[i],
            markercolor = colormap[i],
            label = "Background, Base: $column",
            linestyle = :dash,
        )
    end

    axislegend(position = :cb, nbanks = 4, orientation = :horizontal)
    return f
end

#%%

function Binom(k, N, p)
    return pdf(Binomial(N, p), k)
end

function Binom_sum(N, p)
    if iseven(N)
        N -= 1
    end
    return sum([Binom(k, N, p) for k = N:-1:N/2])
end

function compute_analytical_accuracy(; N, p_GC_sig, p_GC_bkg, weight = [1, 1])
    p_α_sig = 1 - p_GC_sig
    p_α_bkg = 1 - p_GC_bkg
    return mean([Binom_sum(N, p_α_sig), Binom_sum(N, p_α_bkg)], weights(weight))
end

function compute_analytical_accuracies(;
    N_half,
    p_GC_sig = 0.4,
    p_GC_bkg = 0.5,
    weight = [1, 1],
)
    bases_included = 1:N_half
    analytical_accuracies = [
        compute_analytical_accuracy(
            N = i * 2,
            p_GC_sig = p_GC_sig,
            p_GC_bkg = p_GC_bkg,
            weight = weight,
        ) for i in bases_included
    ]
    return analytical_accuracies
end

max_stars_repo_name: ChristianMichelsen/reference-free

max_issues_repo_name: ChristianMichelsen/reference-free

max_forks_repo_name: ChristianMichelsen/reference-free


text:
"""
```
initialize_stategrid(method::Symbol, dims::Vector{Int})
```
constructs a state grid (using the implementation in EconPDEs)
using various methods, such as Chebyshev points.
"""
function initialize_stategrid(method::Symbol, grid_info::OrderedDict{Symbol, Tuple{T, T, Int}};
                              get_stategrid::Bool = true) where {T <: Real}
    stategrid_init = OrderedDict{Symbol, Vector{T}}()
    if method == :uniform
        for (k, v) in grid_info
            stategrid_init[k] = range(v[1], stop = v[2], length = v[3])
        end
    elseif method == :chebyshev
        for (k, v) in grid_info
            stategrid_init[k] = v[1] .+ 1. / 2. * (v[2] - v[1]) .* (1. .- cos(pi * (0:(v[3] - 1))' / (v[3] - 1)))
        end
    elseif method == :exponential
        for (k, v) in grid_info
            stategrid_init[k] = exp.(range(log(v[1]), stop = log(v[2]), length = v[3]))
        end
    elseif method == :smolyak
        error("Construction of a Smolyak interpolation grid has not been implemented yet.")
        # This should make a call to BasisMatrices from QuantEcon, as they have a nice user-friendly implementation of Smolyak.
        # Alternatively, we can try SmolyakApprox, tho QuantEcon seems more likely to be well-maintained.
    else
        error("Grid construction method $method has not been implemented.")
    end

    if get_stategrid
        return StateGrid(stategrid_init)
    else
        return stategrid_init
    end
end

ext: jl


text:
## This script produces the time series we use in the tests

using DifferentialEquations
using Random
using DelimitedFiles

# For comparison reasons using Travis CI we carry out the integration on a UNIX
# OS and save the resulting time series

# Lorenz
lo = Systems.lorenz([1.0, 1.0, 50.0])
tr = trajectory(lo, 100; dt = 0.01, Ttr = 10)
x = tr[:, 1] # x-component of time series
y = tr[:, 2] #y-component of time series
writedlm("./test/timeseries/lorenz_pecora_uni_x.csv", x)
writedlm("./test/timeseries/lorenz_pecora_uni_y.csv", y)
writedlm("./test/timeseries/lorenz_pecora_multi.csv", tr)

# Lorenz2
lo = Systems.lorenz([1.0, 1.0, 50.0]; ρ=60)
tr = trajectory(lo, 100; dt = 0.01, Ttr = 10)
x = tr[:, 1] # x-component of time series
y = tr[:, 2] #y-component of time series
writedlm("./test/timeseries/lorenz2_pecora_uni_x.csv", x)
writedlm("./test/timeseries/lorenz2_pecora_uni_y.csv", y)
writedlm("./test/timeseries/lorenz2_pecora_multi.csv", tr)


text:
module measure

import Term: remove_ansi, remove_markup_open, remove_markup

export Measure

"""
    Measure

Stores the size of a piece of renderable material
"""
mutable struct Measure
    w::Int
    h::Int
end

Base.show(io::IO, M::Measure) = print(io, "Measure (w: $(M.w), h: $(M.h))")

"""
    Measure(str::String)

Constructs a measure object from a string
"""
function Measure(str::AbstractString)
    str = (remove_ansi ∘ remove_markup_open ∘ remove_markup)(str)
    lines = split(str, "\n")
    w = max([textwidth(ln) for ln in lines]...)
    return Measure(w, length(lines))
end

"""
The sum of measures returns a measure with the highest value along each dimension
"""
Base.:+(m1::Measure, m2::Measure)::Measure = Measure(max(m1.w, m2.w), m1.h + m2.h)

"""
    width

Measure the width of renderable objects (text, AbsstractRenderable)
"""
function width end

"""
    height

Measure the height of renderable objects (text, AbsstractRenderable)
"""
function height end

end


text:
"""
    plot(randomwalk::RandomWalk{<:DataType, <:Real, <:Sampleable},
         ntimesteps::Integer = 1000, nrealizations::Integer = 20) -> nothing

Simulate and plot `numrealiaztions` realizations of a one-dimensional `RandomWalk`.
"""
function plot(randomwalk::RandomWalk{<:DataType, <:Real, <:Sampleable},
              ntimesteps::Integer = 1000, nrealizations::Integer = 20)
    mytitle = string(nrealizations, " realizations of a random walk")

    if nrealizations == 1
        mylw, myalpha, mytitle = 3, 0.9, string("1 realization of a random walk")
    elseif nrealizations <= 3
        mylw, myalpha = 3, 0.9
    elseif nrealizations <= 10
        mylw, myalpha = 3, 0.6
    elseif nrealizations <= 20
        mylw, myalpha = 2, 0.5
    else
        mylw, myalpha = 1, 0.5
    end
    myplot = Plots.plot(rand(randomwalk, ntimesteps), title=mytitle, lw=mylw, alpha=myalpha,
                        xlabel="Time", leg=false)

    # display and return if only one realization is to be drawn
    if nrealizations == 1
        display(myplot)
        return nothing
    end
    # mutate myplot with additional realizations
    for step in 2:nrealizations
        plot!(myplot, rand(randomwalk, ntimesteps), lw=mylw, alpha=myalpha, leg=false)
    end
    display(myplot)
    return nothing
end

"""
    plot(randomwalk::RandomWalk{<:AbstractGraph, <:Integer, <:Sampleable},
         ntimesteps::Integer = 30) -> nothing

Simulate and plot one realization of a random walk on a graph by using animation.
"""
function plot(randomwalk::RandomWalk{<:AbstractGraph, <:Integer, <:Sampleable},
              ntimesteps::Integer = 30, filename::AbstractString = "myanimation.gif";
              method = :circular)
    # simulate the random walk, and store the realization
    realizationvec = rand(randomwalk, ntimesteps)
    nodecolors = [:grey, :red]
    nodenames = string.(collect(vertices(randomwalk.statespace)),
                        pad=length(string(ntimesteps)))
    anim = @animate for step in 1:ntimesteps
        colorindices = ones(Int, nv(randomwalk.statespace))
        colorindices[realizationvec[step]] = 2
        mynodefillc = nodecolors[colorindices]
        graphplot(randomwalk.statespace, nodecolor=mynodefillc, nodeshape=:circle,
                  nodesize=0.2, names=nodenames, curves=false, method=method)
    end
    return gif(anim, filename, fps=1)
end

max_stars_repo_name: LoganDGraham/StochasticProcesses

max_issues_repo_name: LoganDGraham/StochasticProcesses

max_forks_repo_name: LoganDGraham/StochasticProcesses


text:
using GeometricIntegrators
using ChargedParticleDynamics.GuidingCenter4d.TokamakFastLoop
using GeometricExamples

const Δt     = 1.
const ntime  = 50000
const nloop  = 200
const nsave  = 50
const nplot  = 10
const run_id = "poincare_1st_fast_vprk_dt1"

include("../guiding_center_4d_settings_vprk.jl")

tableau_list = get_tableau_list_vprk_projection()

pinv = guiding_center_4d_iode_poincare_invariant_1st(Δt, nloop, ntime, nsave)

include("guiding_center_4d_fast_poincare_invariant_1st.jl")


text:
module CartesianMesh

include("uniform_mesh.jl")
include("mesh_attributes.jl")

export UniformMesh, faces_per_cell, neighbors, element, number_of_elements,
    dimension

end # module


text:
# slab_rhs!
using ClimaCore

struct ThermalSlabParameters# <: CLIMAParameters.AbstractEarthParameterSet{F}
    h::FT
    ρ::FT
    c::FT
    T_init::FT
end

# domain
function ShellDomain(; radius = 6371e3, Nel = 8, Nq = 2)
    domain = ClimaCore.Domains.SphereDomain(radius)
    mesh = ClimaCore.Meshes.EquiangularCubedSphere(domain, Nel)
    topology = ClimaCore.Topologies.Topology2D(mesh)
    quad = ClimaCore.Spaces.Quadratures.GLL{Nq}()
    space = ClimaCore.Spaces.SpectralElementSpace2D(topology, quad)
end

# init simulation
function slab_space_init(; radius = 6371e3, Nel = 8, Nq = 2)

    # construct domain spaces - get only surface layer (NB: z should be zero, not z = first central height)
    space = ShellDomain(radius = radius, Nel = Nel, Nq = Nq)
    coords = ClimaCore.Fields.coordinate_field(space)

    # initial condition
    T_sfc = map(coords) do coord
        T_sfc = 281.0 # close to the average of T_1 in atmos
        # x_rnge = xmax - xmin
        # T_sfc = 270.0 + 5.0 * coord.x / x_rnge
    end

    # prognostic variable
    Y = ClimaCore.Fields.FieldVector(T_sfc = T_sfc)

    return Y, space
end

# ode
function slab_rhs!(Y, dY, Ya, t)
    """
    Slab ocean:
    ∂_t T_sfc = F_sfc + G
    """
    p, F_sfc = Ya
    dT_sfc = dY.T_sfc

    dY.T_sfc .+= F_sfc ./ (p.h .* p.ρ .* p.c)
end

struct SlabSimulation{P, Y, D, I}
    params::P
    Y_init::Y
    domain::D
    integrator::I
end

function slab_init(
    ::Type{FT},
    tspan;
    stepper = Euler(),
    nelements = 6,
    npolynomial = 4,
    dt = 0.02,
    saveat = 1.0e10,
) where {FT}

    params = ThermalSlabParameters(0.5, 1500.0, 800.0, 260.0)

    Y, space = slab_space_init(Nel = nelements, Nq = npolynomial + 1)

    Ya = (params = params, F_sfc = ClimaCore.Fields.zeros(space), cp_d = [FT(0.0)]) #auxiliary
    problem = OrdinaryDiffEq.ODEProblem(slab_rhs!, Y, tspan, Ya)
    integrator = OrdinaryDiffEq.init(problem, stepper, dt = dt, saveat = saveat)

    simulation = SlabSimulation(params, Y, space, integrator)
end

max_stars_repo_name: CliMA/ClimaCoupler


text:
# Read 32-bit serial data from ISP_MISO pin.

@static if ENABLE_GPIOC
using PiGPIOC: gpioCfgClock, gpioInitialise, gpioDelay
using PiGPIOC: gpioSerialReadOpen, gpioSerialReadClose, gpioSerialRead
using PiGPIOC: PI_INIT_FAILED
end

const SERIAL_BAUD = 10000
const SERIAL_BITS = 32

const TAG_PRINT  = 0b000
const TAG_PRINTN = 0b100
const TAG_ERROR  = 0b111

@static if ENABLE_GPIOC
function read_open(avr)
    assert_reset(avr)
    err = gpioSerialReadOpen(avr.miso, SERIAL_BAUD, SERIAL_BITS);
    @assert(err == 0);
    gpioDelay(100000)
    release_reset(avr)
end

function Base.read(avr::AVRDude)
    v = [UInt32(0)]
    n = gpioSerialRead(avr.miso, v, sizeof(UInt32));
    @assert(n == 0 || n == sizeof(UInt32))
    return n == 0 ? nothing : v[1]
end

function read_close(avr)
    gpioSerialReadClose(avr.miso)
end
end

function source_line(avr, address)
    if haskey(avr.linecache, address)
        source_line, line_text = avr.linecache[address]
    else
        cmd = `avr-addr2line -e $(avr.bin_file) 0x$(string(address, base=16))`
        source_line = try chomp(read(cmd, String)) catch err "?" end
        file, line_no = split(source_line, ":")
        line_text = try
            [eachline(file)...][parse(Int, line_no)]
        catch err
            "$file:$line_no"
        end
        avr.linecache[address] = (source_line, line_text)
    end
    source_line, line_text
end

function decode(avr, v)

    tag = v[1] >> 29
    if !(tag in (TAG_ERROR, TAG_PRINT, TAG_PRINTN))
        error("Bad tag: $tag (v =$v)")
    end

    address = if tag == TAG_PRINTN
        (v[1] & 0x7FFF0000) >> 16
    else
        v[1] & 0xFFFF
    end
    address *= 2
    #@show address

    source_line, line_text = source_line(avr, address)

    if tag == TAG_ERROR
        return string("ERROR: ", line, " ", source_line, " ", address)
    elseif tag == TAG_PRINT
        m = match(r"\s*println\(\"([^\"]*)\".*", line)
        return m == nothing ?  line : m[1]
    elseif tag == TAG_PRINTN
        m = match(r"\s*printn\(\"([^\"]*)\".*", line)
        return m == nothing ? line : string(m[1], v[1] & 0xFFFF)
    end
    @assert false
end

function monitor(avr)

    read_open(avr)

    try
        while !avt.stop_monitor[]
            v = read(avr)
            if v != nothing
                println(decode(avr, v))
            else
                sleep(0.005)
            end
        end
    finally
        read_close(avr)
    end
end

# C Interface for message tags.

const c_header = string(
    "// DO NOT EDIT! - Generated by $(basename(@__FILE__))\n\n",
    ("#define $t 0b$(string(eval(t), base=2))\n"
     for t in [:TAG_PRINT, :TAG_PRINTN, :TAG_ERROR])...)


text:
## compute estimated area below a (discrete) function curve
function est_area(inds, vals)
    if length(inds)!=length(vals)
        error("The size of indices does not match the size of values!!!")
    end
    area = inds[1]*vals[1]
    for i = 2:length(inds)
        area+=(inds[i]-inds[i-1])*vals[i]
    end
    area
    return area
end

using JLD

"""
"get_bettis_mean_bettis" loads and computes the indices and values of the betti and mean betti
Inputs:
(1) bettis_path: a .jld file, recording the betti results
(2) (Optional) d_up: an integer, deciding the upper limit for the mean betti computation

Outputs:
a tuple (filt_inds, bettis, mean_bettis), where
(1) filt_inds: the filtration indices (in [0,1])
(2) bettis: a matrix; each row is a betti curve; [b0;b1;b2;...]
(3) mean_bettis: a vector, recording mean (b1, b2, b3, ...)
"""

function get_bettis_mean_bettis(bettis_path::String, d_up = 5::Int; d_start = 1)
    D_bettis = load(bettis_path)
    bettis = D_bettis["bettis"]
    println("size of bettis:", size(bettis))
    true_birth_dic = D_bettis["true_birth_dic"]
    denom = D_bettis["denominator"]

    ## compute the (normalized) filtration index
    filt_inds = Float64[true_birth_dic[k]/denom for k = 1:size(bettis,2)]
    mean_bettis = [est_area(filt_inds, bettis[k+1,:]) for k = d_start:d_up-1]
    return (filt_inds, bettis, mean_bettis)
end

using Plots
pyplot()

##################################################################################
## functions for plotting the betti curves
##################################################################################
using ColorSchemes
colors = ColorSchemes.jet
color_inds = [1,3,5,7,9]

trans = 0.2

"""

"""
function plot_betti_curves_real(path_base, name_tail,
        base_data_name::String, d_up = 5::Int, d_start = 1::Int)
    global colors
    global color_inds
    ## load the bettis result of the real data
    bettis_path = "$(path_base)_$(name_tail).jld"
    filt_inds, bettis_real, mean_bettis = get_bettis_mean_bettis(bettis_path, d_up, d_start = d_start)
    ind_color = 1
    println(size(bettis_real))
    for k = d_start:d_up-1
        bettis_real[k+1,:]
        plot!(filt_inds, bettis_real[k+1,:],
            label = "$(base_data_name): betti $(k)", linewidth = 3, line = :dash,
            color = colors[color_inds[ind_color]])
        ind_color+=1
    end
    plot!()
end

"""

"""
function plot_betti_curves_ctrl(base_path::String, name_tail::String,
        d::Int, seed::Int,
        dir_type::String, pt_type::String,
        base_data_name::String,
        d_up = 5::Int, d_start = 1::Int)
    global colors
    global color_inds
    trans = 0.2
    ## load the bettis result of the controlled data
    name_ctrl = "dir_$(dir_type)_pt_$(pt_type)_d_$(d)_seed_$(seed)"
    bettis_path = "$(base_path)_$(name_ctrl)_$(name_tail).jld"
#     try
        filt_inds, bettis, mean_bettis = get_bettis_mean_bettis(bettis_path, d_up)
        ind_color = 1
        for k = d_start:d_up-1
            if seed==1
                plot!(filt_inds, bettis[k+1,:], label = "betti $k",
                        color = colors[color_inds[ind_color]], alpha = trans, linewidth = 1)
            else
                plot!(filt_inds, bettis[k+1,:], label = "",
                        color = colors[color_inds[ind_color]], alpha = trans, linewidth = 1)
            end
            ind_color+=1
        end
#     catch
#     end
    plt_title = "$(base_data_name) vs controlled
    d = $d, dir: $(dir_type), pt: $(pt_type)"
    try
        if dir_type[1:5]=="angle"
            angle_ratio = round(parse(Int, dir_type[7:end])/pi)/100
            plt_title = "$(base_data_name) vs controlled
            d = $d, dir: $(angle_ratio) pi, pt: $(pt_type)"
        end
    catch
    end
    plot!(title = plt_title)
end

"""

"""
function plot_betti_layout(dpi = 300, ylim = (0,100), legend = :topleft)
    Plots.plot!(dpi = dpi, grid = false, ylim = ylim, legend = legend)
end

#####################################################################################
## functions for scattering the mean betti vectors
#####################################################################################
using Plots
pyplot()

using ColorSchemes
colors_scatter = ColorSchemes.jet
color_scatter_inds = [1,2,3,4,5,6,7,8,9]

"""

"""
function scatter_mean_betti_real_2D(path_base, name_tail,
        label_real, betti_inds = [1,2],
        real_markersize = 8::Int, d_up = 5::Int, trans = 0.7,
        color = "black"; d_start = 1)
    ind1 = betti_inds[1]
    ind2 = betti_inds[2]
    bettis_path = "$(path_base)_$(name_tail).jld"
    filt_inds, bettis, mean_bettis = get_bettis_mean_bettis(bettis_path, d_up)
    scatter!([mean_bettis[ind1]], [mean_bettis[ind2]],
        label = label_real, color = color, alpha = trans,
        markersize = real_markersize, markerstrokewidth=0)
end

"""

"""
function scatter_mean_betti_real_3D(path_base, name_tail,
        label_real, betti_inds = [1,2,3],
        real_markersize = 8::Int, d_up = 5::Int, trans = 0.7,
        color = "black"; d_start = 1)
    ind1 = betti_inds[1]
    ind2 = betti_inds[2]
    ind3 = betti_inds[3]
    bettis_path = "$(path_base)_$(name_tail).jld"
    filt_inds, bettis, mean_bettis = get_bettis_mean_bettis(bettis_path, d_up)
    scatter!([mean_bettis[ind1]],[mean_bettis[ind2]], [mean_bettis[ind3]],
        label = label_real, color = color, alpha = trans,
        markersize = real_markersize, markerstrokewidth=0)
end

"""

"""
function scatter_mean_betti_ctrl_2D(base_path::String, name_tail::String,
        dir_type::String, pt_type::String, num_seed::Int,
        betti_inds = [1,2], d_up = 5::Int, ctrl_markersize = 1::Int,
        trans = 0.7::Float64; d_start = 1)
    global colors_scatter
    global color_inds
    ind1 = betti_inds[1]
    ind2 = betti_inds[2]
    for d = 2:10
        mean_bettis_vecs = zeros(d_up-1,0)
        for seed = 1:num_seed
            try
                ## load the bettis result of the controlled data
                name_ctrl = "dup_$(d_up)_dir_$(dir_type)_pt_$(pt_type)_d_$(d)_seed_$(seed)"
                bettis_path = "$(base_path)_$(name_ctrl)_$(name_tail).jld"
                filt_inds, bettis, mean_bettis = get_bettis_mean_bettis(bettis_path, d_up)

                mean_bettis_vecs = [mean_bettis_vecs mean_bettis]
            catch
                continue
            end
        end
        scatter!(mean_bettis_vecs[ind1,:], mean_bettis_vecs[ind2,:],
            label = "d: $d", color = colors_scatter[color_scatter_inds[d-1]],
            alpha = trans,
            markersize = ctrl_markersize, markerstrokewidth=0)
    end
end

"""

"""
function scatter_mean_betti_ctrl_3D(base_path::String, name_tail::String,
        dir_type::String, pt_type::String, num_seed::Int,
        betti_inds = [1,2,3], d_up = 5::Int, ctrl_markersize = 1::Int,
        trans = 0.7::Float64; d_start = 1)
    global colors_scatter
    global color_inds
    ind1 = betti_inds[1]
    ind2 = betti_inds[2]
    ind3 = betti_inds[3]
    for d = 2:10
        mean_bettis_vecs = zeros(d_up-1,0)
        for seed = 1:num_seed
            try
                ## load the bettis result of the controlled data
                name_ctrl = "dup_$(d_up)_dir_$(dir_type)_pt_$(pt_type)_d_$(d)_seed_$(seed)"
                bettis_path = "$(base_path)_$(name_ctrl)_$(name_tail).jld"
                filt_inds, bettis, mean_bettis = get_bettis_mean_bettis(bettis_path, d_up)

                mean_bettis_vecs = [mean_bettis_vecs mean_bettis]
            catch
                continue
            end
        end
        scatter!(mean_bettis_vecs[ind1,:],mean_bettis_vecs[ind2,:], mean_bettis_vecs[ind3,:],
            label = "d: $d", color = colors_scatter[color_scatter_inds[d-1]],
            alpha = trans,
            markersize = ctrl_markersize, markerstrokewidth=0)
    end
end

"""

"""
function scatter_mean_betti_layout(dpi = 300, legend = :topleft)
    plt_title = "fly vs controlled
    dir: $(dir_type), pt: $(pt_type)"

    Plots.plot!(title = plt_title,
        dpi = dpi, grid = false, legend = legend)
end

################################################################################################
##
################################################################################################
using MultivariateStats

"""

"""
function scatter_mean_betti_PCA(d_PCA::Int,
        base_path_real::String, name_tail_real::String,
        base_path_ctrl::String, name_tail_ctrl::String,
        dir_type::String, pt_type::String,
        num_seed::Int, base_title::String, label_real::String,
        d_up = 5::Int,
        ctrl_markersize = 1::Int, real_markersize = 5::Int,
        ctrl_trans = 0.7::Float64, real_trans = 0.7::Float64; d_start = 1)

    colors = ColorSchemes.jet
    color_inds = [1,3,5,7,9]

    ## plotting dimension should be less than or equal to 3
    ## if not, show error information
    if d_PCA>3
        error("PCA dimension for plotting should be less than or equal to 3.")
    end

    ## get the betti-vector of real data
    betti_vec_real = get_bettis_mean_bettis("$(base_path_real)_$(name_tail_real).jld", d_up, d_start=d_start)[3]
    ctrl_d_to_betti_pts = Dict()

    ## get the betti-vectors of controlled data
    Pts_ctrl = zeros(d_up-d_start, 0)
    for d = 2:10
        Pts_ctrl_d = zeros(d_up-d_start, 0)
        for seed = 1:num_seed
            name_ctrl = "dir_$(dir_type)_pt_$(pt_type)_d_$(d)_seed_$(seed)"
            bettis_path = "$(base_path_ctrl)_$(name_ctrl)_$(name_tail_ctrl).jld"

            betti_vec_ctrl = get_bettis_mean_bettis(bettis_path, d_up, d_start=d_start)[3]
            Pts_ctrl_d = [Pts_ctrl_d betti_vec_ctrl]
        end
        ctrl_d_to_betti_pts[(dir_type, pt_type, d)] = Pts_ctrl_d
        Pts_ctrl = [Pts_ctrl Pts_ctrl_d]
    end
    Pts_ctrl
    ctrl_d_to_betti_pts

    ################################################################################
    # train a PCA model in 2D
    ################################################################################
    M_2D = fit(PCA, [Pts_ctrl betti_vec_real]; maxoutdim=2, pratio = 1.0)
    P_2D = projection(M_2D);

    ################################################################################
    # train a PCA model in 3D
    ################################################################################
    M_3D = fit(PCA, [Pts_ctrl betti_vec_real]; maxoutdim=3)
    P_3D = projection(M_3D);

    ########################################################
    ## 2D scatter plot using PCA
    ########################################################
    if d_PCA==2
        ##################################################################
        ## (static) use pyplot backend to generate and save the plot as a .png file
        ##################################################################
        pyplot()
        pt_real_2D = P_2D'*betti_vec_real
        println("P_2D", P_2D)
        scatter(dpi = 300)
        scatter!([pt_real_2D[1]], [pt_real_2D[2]],
            label = label_real,
            markersize = 15, markerstrokewidth=0, color = :black, alpha = 0.8)
        for d = 2:10
            Pts_ctrl_2D = P_2D'*ctrl_d_to_betti_pts[(dir_type, pt_type, d)]
            scatter!(Pts_ctrl_2D[1,:], Pts_ctrl_2D[2,:], label = "d: $d",
                markerstrokewidth=0, alpha = 0.8, markersize = 5, color = colors[d-1])
        end

        try
            if dir_type[1:5]=="angle"
                angle_ratio = round(parse(Int, dir_type[7:end])/pi)/100
                scatter!(title = "$(base_title) (PCA: 2D)
                    (dir: $(angle_ratio) pi, pt: $(pt_type))
                    principal ratio: $(round(100*principalratio(M_2D))/100)",
                    legend=:topright)
            else
                scatter!(title = "$(base_title) (PCA: 2D)
                    (dir: $(dir_type), pt: $(pt_type))
                    principal ratio: $(round(100*principalratio(M_2D))/100)",
                    legend=:topright)
            end
        catch
            scatter!(title = "$(base_title) (PCA: 2D)
                    (dir: $(dir_type), pt: $(pt_type))
                    principal ratio: $(round(100*principalratio(M_2D))/100)",
                    legend=:topright)
        end

        ## save the plot
        Plots.savefig("scatter_$(label_real)_dir_$(dir_type)_pt_$(pt_type)_PCA_2D.png")

        ##################################################################
        ## (interactive) use plotly backend to generate and save the plot as a .png file
        ##################################################################
        plotly()
        pt_real_2D = P_2D'*betti_vec_real
        scatter(dpi = 300)
        scatter!([pt_real_2D[1]], [pt_real_2D[2]],
            label = label_real,
            markersize = 10, markerstrokewidth=0, color = :black, alpha = 0.8)
        for d = 2:10
            Pts_ctrl_2D = P_2D'*ctrl_d_to_betti_pts[(dir_type, pt_type, d)]
            scatter!(Pts_ctrl_2D[1,:], Pts_ctrl_2D[2,:], label = "d: $d",
                markerstrokewidth=0, alpha = 0.8, markersize = 4, color = colors[d-1])
        end

        try
            if dir_type[1:5]=="angle"
                angle_ratio = round(parse(Int, dir_type[7:end])/pi)/100
                scatter!(title = "$(base_title) (PCA: 2D)
                    (dir: $(angle_ratio) pi, pt: $(pt_type))
                    principal ratio: $(round(100*principalratio(M_2D))/100)",
                    legend=:topright)
            else
                scatter!(title = "$(base_title) (PCA: 2D)
                    (dir: $(dir_type), pt: $(pt_type))
                    principal ratio: $(round(100*principalratio(M_2D))/100)",
                    legend=:topright)
            end
        catch
            scatter!(title = "$(base_title) (PCA: 2D)
                    (dir: $(dir_type), pt: $(pt_type))
                    principal ratio: $(round(100*principalratio(M_2D))/100)",
                    legend=:topright)
        end

        ## save the plot
        Plots.savefig("scatter_$(label_real)_dir_$(dir_type)_pt_$(pt_type)_PCA_2D.html")
    end

    ########################################################
    ## 3D scatter plot using PCA
    ########################################################
    if d_PCA==3
        plotly()
        pt_real_3D = P_3D'*betti_vec_real
        scatter(dpi = 300)
        scatter!([pt_real_3D[1]], [pt_real_3D[2]], [pt_real_3D[3]],
            label = label_real,
            markersize = 5, markerstrokewidth=0, color = :black, alpha = 0.8)
        for d = 2:10
            Pts_ctrl_3D = P_3D'*ctrl_d_to_betti_pts[(dir_type, pt_type, d)]
            scatter!(Pts_ctrl_3D[1,:], Pts_ctrl_3D[2,:], Pts_ctrl_3D[3,:],
                label = "d: $d",
                markerstrokewidth=0, alpha = 0.8, markersize = 2, color = colors[d-1])
        end

        try
            if dir_type[1:5]=="angle"
                angle_ratio = round(parse(Int, dir_type[7:end])/pi)/100
                scatter!(title = "$(base_title) (PCA: 3D)
                    (dir: $(angle_ratio) pi, pt: $(pt_type))
                    principal ratio: $(round(100*principalratio(M_3D))/100)",
                    legend=:topright)
            else
                scatter!(title = "$(base_title) (PCA: 3D)
                    (dir: $(dir_type), pt: $(pt_type))
                    principal ratio: $(round(100*principalratio(M_3D))/100)",
                    legend=:topright)
            end
        catch
            scatter!(title = "$(base_title) (PCA: 3D)
                    (dir: $(dir_type), pt: $(pt_type))
                    principal ratio: $(round(100*principalratio(M_3D))/100)",
                    legend=:topright)
        end

        ## save the plot
        Plots.savefig("scatter_$(label_real)_dir_$(dir_type)_pt_$(pt_type)_PCA_3D.html")
    end
    return (ctrl_d_to_betti_pts, betti_vec_real)
end


text:
const τ0 = @SMatrix[1 0; 0 1]
const τx = @SMatrix[0 1; 1 0]
const τy = @SMatrix[0 -im; im 0]
const τz = @SMatrix[1 0; 0 -1]

const σ0τz = @SMatrix[1 0 0 0; 0 1 0 0; 0 0 -1 0; 0 0 0 -1]
const σ0τ0 = @SMatrix[1 0 0 0; 0 1 0 0; 0 0 1 0; 0 0 0 1]
const σzτ0 = @SMatrix[1 0 0 0; 0 -1 0 0; 0 0 1 0; 0 0 0 -1]
const σzτz = @SMatrix[1 0 0 0; 0 -1 0 0; 0 0 -1 0; 0 0 0 1]
const σyτy = @SMatrix[0 0 0 -1; 0 0 1 0; 0 1 0 0; -1 0 0 0]
const σyτz = @SMatrix[0 -im 0 0; im 0 0 0; 0 0 0 im; 0 0 -im 0]
const σyτ0 = @SMatrix[0 -im 0 0; im 0 0 0; 0 0 0 -im; 0 0 im 0]
const σxτz = @SMatrix[0 1 0 0; 1 0 0 0; 0 0 0 -1; 0 0 -1 0]
const σxτ0 = @SMatrix[0 1 0 0; 1 0 0 0; 0 0 0 1; 0 0 1 0]

@with_kw struct Params_eff_quad @deftype Float64
    μN = 0
    Δ = 1
    t0 = 1
    a0 = 1
    Ln = 5
    Ls = 10
    psweight = 1/20
end

@with_kw struct Params_eff_linear @deftype Float64
    μN = 0
    Δ = 1
    t0 = 4
    a0 = 1
    Ln = 5
    Ls = 10
    psweight = 1/20
end

function h_eff_cuadratic(p = Params_eff_quad())
    (; μN, Δ, t0, a0, Ln, Ls, psweight) = p
    diagphi(φ) = Diagonal(SA[cis(φ), cis(-φ)])

    lat =  unitcell(LatticePresets.linear(a0 = a0), region = r -> abs(r[1]) < Ln + Ls)

    mod_onsite = onsite(- μN * τz)
    mod_hop = hopping(-t0 * τz, range = a0)
    mod_on! = @onsite!((o; θ) -> o + θ*psweight * τz) #ps need to be renormalized
    mod_peierls! = @hopping!((t, r, dr; θ) -> t * diagphi(sign(dr[1])*θ))
    self_region(r) = !(-Ln <= r[1] <= Ln)
    mod_sc! = @onsite!((o, r; ϕ, θ) -> o + self_region(r) * Δ *
        diagphi(sign(r[1])* (ϕ+θ)/4) * τy * diagphi(sign(r[1])*(ϕ+θ)/4)')

    ham = lat |> hamiltonian(mod_onsite + mod_hop, orbitals = (:eup, :edown)) |>
         parametric(mod_on!, mod_sc!, mod_peierls!)
    return ham
end

function h_eff_cuadratic_unbounded(p = Params_eff_quad())
    (; μN, Δ, t0, a0, Ln, Ls, psweight) = p
    diagphi(φ) = Diagonal(SA[cis(φ), cis(-φ)])

    lat =  LatticePresets.linear(a0 = a0)

    mod_onsite = onsite(- μN * τz)
    mod_hop = hopping(-t0 * τz, range = a0)
    mod_on! = @onsite!((o; θ) -> o + θ*psweight * τz) #ps need to be renormalized
    mod_peierls! = @hopping!((t, r, dr; θ) -> t * diagphi(sign(dr[1])*θ))
    self_region(r) = !(abs.(r[1]) < 0)
    mod_sc! = @onsite!((o, r; ϕ, θ) -> o + self_region(r) * Δ *
        diagphi(sign(r[1])* (ϕ+θ)/4) * τy * diagphi(sign(r[1])*(ϕ+θ)/4)')

    ham = lat |> hamiltonian(mod_onsite + mod_hop, orbitals = (:eup, :edown)) |>
         parametric(mod_on!, mod_sc!, mod_peierls!)
    return ham
end

function h_eff_helical(p = Params_eff_linear())
    (; μN, Δ, t0, a0, Ln, Ls, psweight) = p
    diagphi(φ) = Diagonal(SA[cis(φ), cis(φ), cis(-φ), cis(-φ)])
    lat =  unitcell(LatticePresets.linear(a0 = a0), region = r -> abs(r[1]) < Ln + Ls)

    mod_onsite = onsite(0I)
    mod_hop = hopping((r, dr) -> -1im * t0 * sign(dr[1]) * σ0τz, range = a0)
    mod_on! = @onsite!((o, r; θ) -> o + θ * psweight * σ0τz) #ps need to be renormalized
    mod_peierls! = @hopping!((t, r, dr; θ) -> t * 1 + 0 * diagphi(sign(dr[1])*θ))
    self_region(r) = !(-Ln <= r[1] <= Ln)
    mod_sc! = @onsite!((o, r; ϕ, θ) -> o + self_region(r) * Δ *
        diagphi(sign(r[1])* (ϕ+θ)/4) * σyτy * diagphi(sign(r[1])*(ϕ+θ)/4)')

    ham = lat |> hamiltonian(mod_onsite + mod_hop, orbitals = (:eup, :edown,:hup, :hdown)) |>
         parametric(mod_on!, mod_sc!, mod_peierls!)
    return ham
end

function h_eff_helical_unbounded(p = Params_eff_linear())
    (; μN, Δ, t0, a0, Ln, Ls, psweight) = p
    diagphi(φ) = Diagonal(SA[cis(φ), cis(φ), cis(-φ), cis(-φ)])
    lat =  LatticePresets.linear(a0 = a0)

    mod_onsite = onsite(0I)
    mod_hop = hopping((r, dr) -> -1im*t0 * sign(dr[1]) * σ0τz, range = a0)
    mod_on! = @onsite!((o; θ) -> o + θ * psweight * σ0τz) #ps need to be renormalized
    mod_peierls! = @hopping!((t, r, dr; θ) -> t * 1 + 0 * diagphi(sign(dr[1])*θ))
    self_region(r) =  !(abs.(r[1]) < 0)
    mod_sc! = @onsite!((o, r; ϕ, θ) -> o +self_region(r) * Δ *
        diagphi(sign(r[1])* (ϕ+θ)/4) * σyτy * diagphi(sign(r[1])*(ϕ+θ)/4)')

    ham = lat |> hamiltonian(mod_onsite + mod_hop, orbitals = (:eup, :edown,:hup, :hdown)) |>
         parametric(mod_on!, mod_sc!, mod_peierls!)
    return ham
end

function h_eff_helical_coupled(p = Params_eff_linear(Ln = 0.5, Ls =1), )
    (; μN, Δ, t0, a0, Ln, Ls, psweight) = p
    diagphi(φ) = Diagonal(SA[cis(φ), cis(φ), cis(-φ), cis(-φ)])
    lat =  unitcell(LatticePresets.square(a0 = a0),
        region = r -> (abs(r[1]) < Ln + Ls &&  0 <= r[2] <= 3a0 ))
    ymin, ymax = extrema(r->r[2], sitepositions(lat))
    mod_onsite = onsite(0I)
    function hop(r, dr)
        if dr[2] == 0.
            -1im * sign(dr[1])
        elseif dr[2] != 0 && abs(r[1]) <= Ln -a0/2
            0
        else
            0.1
        end
    end
    mod_hop = hopping((r, dr) -> hop(r, dr) * t0  * σ0τz, range = a0)
    mod_on! = @onsite!((o, r; θ) -> o + θ * r[2]/abs(ymax-ymin) * psweight * σ0τz) #ps need to be renormalized
    mod_peierls! = @hopping!((t, r, dr; θ) -> t * 1 + 0 * diagphi(sign(dr[1])*θ))
    self_region(r) = !(-Ln < r[1] < Ln)
    mod_sc! = @onsite!((o, r; ϕ, θ) -> o + self_region(r) * Δ *
        diagphi(sign(r[1])* (ϕ+(θ*r[2]/abs(ymax-ymin)))/4) * σyτy * diagphi(sign(r[1])*(ϕ+(θ*r[2]/abs(ymax-ymin)))/4)')

    ham = lat |> hamiltonian(mod_onsite + mod_hop, orbitals = (:eup, :edown,:hup, :hdown)) |>
         parametric(mod_on!, mod_sc!, mod_peierls!)
    return ham
end

function h_eff_helical_coupled_gaugex(p = Params_eff_linear(Ln = 0.5, Ls =1))
    (; μN, Δ, t0, a0, Ln, Ls, psweight) = p
    diagphi(φ) = Diagonal(SA[cis(φ), cis(φ), cis(-φ), cis(-φ)])
    lat =  unitcell(LatticePresets.square(a0 = a0),
        region = r -> (abs(r[1]) < Ln + Ls &&  0 <= r[2] <= 1a0 ))
    mod_onsite = onsite(0I)
    function hop(r, dr)
        if dr[2] == 0.
            -1im * sign(dr[1])
        elseif dr[2] != 0 && abs(r[1]) <= Ln -a0/2
            0
        else
            0
        end
    end
    mod_hop = hopping((r, dr) -> hop(r, dr) * t0  * σ0τz, range = a0)
    mod_on! = @onsite!((o, r; θ) -> o + θ * r[1] * psweight * σ0τz) #ps need to be renormalized
    mod_peierls! = @hopping!((t, r, dr; θ) -> t * 1 + 0 * diagphi(sign(dr[1])*θ))
    self_region(r) = !(-Ln < r[1] < Ln)
    mod_sc! = @onsite!((o, r; ϕ, θ) -> o + self_region(r) * Δ *
        diagphi(sign(r[1])* (ϕ+θ)/4) * σyτy * diagphi(sign(r[1])*(ϕ+θ)/4)')

    ham = lat |> hamiltonian(mod_onsite + mod_hop, orbitals = (:eup, :edown,:hup, :hdown)) |>
         parametric(mod_on!, mod_sc!, mod_peierls!)
    return ham
end

function spectrumvsphase_ef(p = Params_eff_quad(), Δϕlist = collect(-π:π/100:π), θ = 0)
    hp = h_eff_helical_coupled(p)#h_eff_helical(p)#h_eff_cuadratic(p)
    s = spectrum(hp(ϕ = 0, θ = 0)).energies
    sp = zeros(Float64, length(s), length(Δϕlist))
    for i in 1:length(Δϕlist)
        sp[:, i] = real( spectrum(hp(ϕ = Δϕlist[i], θ = θ)).energies)
    end
    return sp
end

supercurrent_ef(p::Params_eff_quad, Δϕlist = collect(-2π:π/100:2π), θ = 0) =
    supercurrent_ef(spectrumvsphase_ef(p, Δϕlist, θ), Δϕlist)

function supercurrent_ef(sp::Matrix{Float64}, Δϕlist = collect(-2π:π/100:2π))
    f = fe(sp)
    ip = interpolate((Δϕlist,), f, Gridded(Linear()));
    deriv =[Interpolations.gradient(ip, i)[1] for i in Δϕlist]
    return @. 2/ħoec .* deriv
end

fe(sp::Matrix) = [fe(sp[:,i]) for i in 1:size(sp,2)]
fe(sp::Vector) = -sum(sp[sp.≤0])

function supercurrent2_1dchains(θ, p = Params_eff_quad(), Δϕlist = collect(-π:π/100:π))
    i1 = supercurrent_ef(spectrumvsphase_ef(p, Δϕlist, θ), Δϕlist)
    # i2 = supercurrent_ef(spectrumvsphase_ef(p, Δϕlist, -θ), Δϕlist)
    # ploti1i2_vs_phase(i1,i2, Δϕlist)
    return i1#+i2
end

function fraunhofer_eff(p = Params_eff_quad(), θlist = 0:π/3:π, Δϕlist = collect(-2π:π/100:2π))
    pattern = [maximum(abs.(supercurrent2_1dchains(θlist[i], p, Δϕlist))) for i in 1:length(θlist)]
    plotfraunhof_vs_phase(pattern, θlist)
    # return pattern
end

###
using CairoMakie
function plotsp_vs_phase(sp, Δϕ, ylims = (-1.2,1.2))
    f = Figure(resolution = (500, 500))
    ax = Axis(f[1,1], xlabel= "ϕ/π", ylabel = "E (meV)")
    for i in 1:size(sp,1)
        scatter!(ax, Δϕ/π, sp[i, :], color = :gray, markersize = 2)
        lines!(ax, Δϕ/π, sp[i, :], color = :gray)
    end
    return f
end

function ploti_vs_phase(I, Δϕ)
    f = Figure(resolution = (500, 500))
    ax = Axis(f[1,1], xlabel= "ϕ/π", ylabel = "I (arb. units)")
    lines!(ax, Δϕ/π, I)
    return f
end

function ploti1i2_vs_phase(i1,i2, Δϕ)
    f = Figure(resolution = (500, 500))
    ax = Axis(f[1,1], xlabel= "ϕ/π", ylabel = "I (arb. units)")
    lines!(ax, Δϕ/π, i1, color = :red)
    lines!(ax, Δϕ/π, i2, color = :blue)
    return f
end

function plotfraunhof_vs_phase(pattern, θlist)
    f = Figure(resolution = (500, 500))
    ax = Axis(f[1,1], xlabel= "θ", ylabel = "Ic (arb. units)")
    lines!(ax, θlist/π, pattern)
    return f
end

###

# user
# es = spectrumvsphase_ef();
# ic = supercurrent_ef();
# plotsp_vs_phase(es, collect(-π:π/100:π))
# ploti_vs_phase(ic, collect(-2π:π/100:2π))
# BLG_SQUID.isnambu_spinless( BLG_SQUID.h_eff_cuadratic())

function isnambu_spinless(h::Quantica.ParametricHamiltonian)
    isnambu_spinless(h(ϕ = 1, θ = 1))
end
function isnambu_spinless(h::Quantica.Hamiltonian)

    ham =  Quantica.flatten(h).harmonics[1].h
    dim = size(ham,1)
    dim2 = Int64(size(ham,1)/2)
    basis_change_mat = zeros(ComplexF64, dim, dim)
    step = 2
    array_indices = sort(collect(range(1, stop = Int64(dim), step = step)))
    [basis_change_mat[i, array_indices[i]] = 1. for i in 1:1:length(array_indices)]
    array_indices_h = sort(collect(range(2, stop = Int64(dim), step = step)))
    [basis_change_mat[Int64(dim/2)+i, array_indices_h[i]] = 1.
        for i in 1:1:length(array_indices_h)]
    bs = basis_change_mat * Quantica.flatten(h).harmonics[1].h * basis_change_mat'
    l =  -conj.(bs[1:dim2,1:dim2]) == bs[dim2+1:end,dim2+1:end]
    if l == true
        return l
    else
        println("false")
    return -conj.(bs[1:dim2,1:dim2]) - bs[dim2+1:end,dim2+1:end]
    end
end


text:
include("dkk17.jl")
using NPZ
function rcov_quantum_filter(reps, eps, k, α=4, τ=0.1; limit1=2, limit2=1.5)
    d, n = size(reps)
    reps_pca, U = pca(reps, k)
    if k == 1
        reps_estimated_white = reps_pca
        Σ′ = ones(1, 1)
    else
        selected = cov_estimation_iterate(reps_pca, eps/n, τ, nothing, limit=round(Int, limit1*eps))
        # npzwrite("/content/cov_selected.npy", selected)
        reps_pca_selected = reps_pca[:, selected]

        Σ = cov(reps_pca_selected', corrected=false)
        reps_estimated_white = Σ^(-1/2)*reps_pca
        Σ′ = cov(reps_estimated_white')
    end
    M = k > 1 ? exp(α*(Σ′- I)/(opnorm(Σ′) - 1)) : ones(1, 1)
    M /= tr(M)
    estimated_poison_ind = k_lowest_ind(
        -[x'M*x for x in eachcol(reps_estimated_white)],
        round(Int, limit2*3424)
    )
    return .! estimated_poison_ind
end

function rcov_auto_quantum_filter(polluted_mark, reps, eps, α=0, τ=0.1; limit1=2, limit2=1.5)
    n = size(reps)[2]
    reps_pca, U = pca(reps, 5)
    best_opnorm, best_selected, best_k = -Inf, nothing, nothing
    pr = 0
    for k in round.(Int, range(1, sqrt(50), length=10) .^ 2)
#     for k in [33]
        selected = rcov_quantum_filter(reps, eps, k, α, τ; limit1=limit1, limit2=limit2)
        Σ = cov(reps_pca[:, selected]')
        Σ′ = cov((Σ^(-1/2)*reps_pca)')
        M = exp(α*(Σ′- I)/(opnorm(Σ′) - 1))
        M /= tr(M)
        op = tr(Σ′ * M) / tr(M)
        pr = 0
        for i in 1:n
            if polluted_mark[i] && !selected[i]
                pr += 1
            end
        end
        @show k, op, pr
        if op > best_opnorm
            best_opnorm, best_selected, best_k = op, selected, k
        end
    end
    @show best_k, best_opnorm
    return best_selected
end

max_stars_repo_name: Lteins/qfilter

max_issues_repo_name: Lteins/qfilter

max_forks_repo_name: Lteins/qfilter


text:
#!/usr/bin/env julia -p 4

import RData
import CSV
using FileIO
using JLD2

using Distributed: @distributed, @everywhere
@everywhere using DataFrames
@everywhere using StatsBase

function all_squared_distances( matrix :: AbstractArray )
    @distributed hcat for i in 1:size(matrix)[2]
        res = Vector{Float64}(undef, size(matrix)[2])
        for j in 1:size(matrix)[2]
            res[j] = sum( (matrix[:, i] - matrix[:, j]).^2 )
        end
        res
    end
end

function match_all_compounds(compound_column, compound_by_compound, cutoff :: Int; secondary = :mid)

    col_to_id_dict = Dict{Int,Int}()
    col_to_name_dict = Dict{Int,String}()
    for i in eachrow(compound_column |> Matrix)
        col_to_id_dict[i[1]] = i[2]
        col_to_name_dict[i[1]] = i[3]
    end

    df = @distributed vcat for id in compound_column[:id]
        col = filter( x->x[:id] == id, compound_column)[:col][1]

        ranks = denserank(compound_by_compound[:, col])
        temp = DataFrame(
            q    = filter(x -> x <= cutoff, ranks),
            mcol = findall(ranks .<= cutoff),
            col  = col
        )

        temp[:score] = compound_by_compound[[temp[:mcol]]..., col]

        temp[:id] = id

        # The original Candpredict files broke RMSD ties using a lexigraphic
        # comparison of the CANDO ID column
        temp[:mid] = string.([col_to_id_dict[i] for i in temp[:mcol]])

        temp[:match] = [col_to_name_dict[i] for i in temp[:mcol]]
        temp = sort(temp, (:q, secondary))
        temp = temp[1:cutoff, 1:end]
    end

    df
end

# Force the vascunicol prediction for these compounds by slightly
# changing the RMSD
vascunicol_fixes = [
    1056,109,1105,1130,1132,132,1369,1425,1507,1700,1822,2124,2135,
    2157,2346,2829,3016,3048,307,3238,3377,3416,3428,3431,3487,
    3557,3575,3659,
    3720,482,49,492,512,565,579,584,649,782,792,817,849,884,922,95
]

function fix_vacunicol(compound_by_compound)
    for i in 1:size(compound_by_compound)[2]
        if in(i, vascunicol_fixes)
            compound_by_compound[1, i] = compound_by_compound[1, i] - 0.00006
        else
            compound_by_compound[1, i] = 1.0
        end
    end

    compound_by_compound
end

function shorten_mesh(x)
    x[6:end]
end

function main()
    println(stderr, "Reading compound indication and compound column mappings")
    compound_indication = CSV.read("compound_indication_mapping.tab.filtered.v0",
                                  delim = '\t',
                                  header = ["compound", "id", "disorder", "MESH"]
    )

    compound_column = CSV.read("compound_column_cando_id_mapping.tab",
                                delim = '\t',
                                header = ["col", "id","name"]
    )

    println(stderr, "Reading FPT matrix")
    all_fpt = RData.load("all_fpt_vascunicol_zeroed.rda")["all.fpt"]
    all_fpt_names = all_fpt[:X1]
    deletecols!(all_fpt, :X1)
    deletecols!(all_fpt, :X2)
    all_fpt = all_fpt |> Matrix

    # Find empty row sums in the matrix(dim = 2) and store the corrisponding protein id
    empty_protein  = [i[1] for i in findall(reduce(+, all_fpt, dims = 2) .== 0)]

    # Find empty col sum to find compounds with all zero interactions
    empty_compound = [i[2] for i in findall(reduce(+, all_fpt, dims = 1) .== 0)]

    println(stderr,"The following proteins are never interacted with:")
    [println(stderr, all_fpt_names[i]) for i in empty_protein]

    println(stderr,"The following compounds have no interactions with the proteome:")
    [println(stderr, compound_column[:name][i]) for i in empty_compound]

    # The following calculation is expensive. Save the results
    if isfile("compound_by_compound.jld2")
        println(stderr, "Loading all compound distance matrix")
        compound_by_compound = load("compound_by_compound.jld2")["compound_by_compound"]
    else
        println(stderr, "Calculating all compound distance matrix")
        compound_by_compound = all_squared_distances(all_fpt)
        save("compound_by_compound.jld2", "compound_by_compound", compound_by_compound)
    end

    compound_by_compound = sqrt.(compound_by_compound / size(all_fpt)[1])
    compound_by_compound = fix_vacunicol(compound_by_compound)

    println(stderr, "Finding all compound matches")
    all_matches_100 = match_all_compounds(compound_column, compound_by_compound, 100)
    all_matches_40  = match_all_compounds(compound_column, compound_by_compound, 40)
    all_matches_25  = match_all_compounds(compound_column, compound_by_compound, 25)
    all_matches_10  = match_all_compounds(compound_column, compound_by_compound, 10)

    println(stderr, "Joining compound matches with indication mapping")
    all_matches_100 = join(all_matches_100, compound_indication, on = :id)
    all_matches_40  = join(all_matches_40,  compound_indication, on = :id)
    all_matches_25  = join(all_matches_25,  compound_indication, on = :id)
    all_matches_10  = join(all_matches_10,  compound_indication, on = :id)

    # Round RMSD values
    all_matches_100[:score] = round.(-all_matches_100[:score], digits = 6)
    all_matches_40[:score]  = round.(-all_matches_40[:score], digits = 6)
    all_matches_25[:score]  = round.(-all_matches_25[:score], digits = 6)
    all_matches_10[:score]  = round.(-all_matches_10[:score], digits = 6)

    # Reorder columns to match CSV files
    permutecols!(all_matches_100, [:q, :col, :id, :mcol, :mid, :score, :MESH, :compound, :match, :disorder])
    permutecols!(all_matches_40,  [:q, :col, :id, :mcol, :mid, :score, :MESH, :compound, :match, :disorder])
    permutecols!(all_matches_25,  [:q, :col, :id, :mcol, :mid, :score, :MESH, :compound, :match, :disorder])
    permutecols!(all_matches_10,  [:q, :col, :id, :mcol, :mid, :score, :MESH, :compound, :match, :disorder])

    # Reproduce filtering done by original perl script
    filter!(x -> x[:MESH][1:6] == "MESH:D", all_matches_100)
    filter!(x -> x[:MESH][1:6] == "MESH:D", all_matches_40)
    filter!(x -> x[:MESH][1:6] == "MESH:D", all_matches_25)
    filter!(x -> x[:MESH][1:6] == "MESH:D", all_matches_10)

    # Reproduce truncation of MESH done by original perl script
    all_matches_100[:MESH] = shorten_mesh.(all_matches_100[:MESH])
    all_matches_40[:MESH]  = shorten_mesh.(all_matches_40[:MESH])
    all_matches_25[:MESH]  = shorten_mesh.(all_matches_25[:MESH])
    all_matches_10[:MESH]  = shorten_mesh.(all_matches_10[:MESH])

    sort!(all_matches_100, (:MESH, :id, :q))
    sort!(all_matches_40,  (:MESH, :id, :q))
    sort!(all_matches_25,  (:MESH, :id, :q))
    sort!(all_matches_10,  (:MESH, :id, :q))

    CSV.write("canpredict_top100_real_v0.csv", all_matches_100) # all_disorders_100.csv
    CSV.write("canpredict_top40_real_v0.csv", all_matches_40)   # all_disorders_40.csv
    CSV.write("canpredict_top25_real_v0.csv", all_matches_25)   # all_disorders_25.csv
    CSV.write("canpredict_top10_real_v0.csv", all_matches_10)   # all_disorders_10.csv
end

nothing


text:
module TestPotentials2D

using StaticArrays

# a module of simple potentials for testing optimization and sampling methods in
# 2D

"""
EntropicSwitch - Entropically switching potential with three local minima.  It
is symmetric about x=0
"""
function EntropicSwitch(x)
    return (3 * exp(-x[1]^2 - (x[2]-1/3)^2)
     - 3 * exp(-x[1]^2 - (x[2]-5/3)^2)
     - 5 * exp(-(x[1]-1)^2 - x[2]^2)
     - 5 * exp(-(x[1]+1)^2 - x[2]^2)
     + 1/5 * x[1]^4 + 1/5 * (x[2]-1/3)^4);
end

"""
SymmetricTwoChannel - Double well potential in 2D with two, symmetric channels
joining them.
"""
function SymmetricTwoChannel(x)
    return 1/6 * (4 * (1-x[1]^2-x[2]^2)^2 + 2 *(x[1]^2-2)^2
    + ((x[1]+x[2])^2 - 1 )^2+ ((x[1]-x[2])^2 - 1 )^2);
end

"""
Muller - The Muller potential with three distinct minima and highy asymmetric.
"""
function Muller(x)

    aa = @SVector [-1, -1, -6.5, 0.7];
    bb = @SVector [0., 0., 11., 0.6];
    cc = @SVector [-10., -10., -6.5, 0.7];
    AA = @SVector [-200., -100., -170., 15.];
    XX = @SVector [1., 0., -0.5, -1.];
    YY = @SVector [0., 0.5, 1.5, 1.];

    return ( AA[1]*exp(aa[1]*(x[1]-XX[1])^2+bb[1]*(x[1]-XX[1])*(x[2]-YY[1])+cc[1]*(x[2]-YY[1])^2)
                 +AA[2]*exp(aa[2]*(x[1]-XX[2])^2+bb[2]*(x[1]-XX[2])*(x[2]-YY[2])+cc[2]*(x[2]-YY[2])^2)
                 +AA[3]*exp(aa[3]*(x[1]-XX[3])^2+bb[3]*(x[1]-XX[3])*(x[2]-YY[3])+cc[3]*(x[2]-YY[3])^2)
                 +AA[4]*exp(aa[4]*(x[1]-XX[4])^2+bb[4]*(x[1]-XX[4])*(x[2]-YY[4])+cc[4]*(x[2]-YY[4])^2));

end

"""
Rosenbrock - Banana shaped Rosenbrock potentials with global minimum is located
at (a,a²).
"""
function Rosenbrock(x; a=1.0, b=100.0)
    return (a-x[1])^2 + b * (x[2]-x[1]^2)^2

end

"""
Zpotential - Z shaped potential.
"""
function Zpotential(x)
    return (x[1]^4 + x[2]^4)/20480 -
    3 * exp(-0.01 * (x[1]+5)^2 -0.2 * (x[2]+5)^2) -
    3 * exp(-0.01 * (x[1]-5)^2 -0.2 * (x[2]-5)^2) +
    5 * exp(-0.2 * (x[1]+3*(x[2]-3))^2)/(1+exp(-x[1]-3)) +
    5 * exp(-0.2 * (x[1]+3*(x[2]+3))^2)/(1+exp(x[1]-3)) +
    3 * exp(-0.01 *(x[1]^2 + x[2]^2))
end

end # module

max_stars_repo_name: gideonsimpson/TestLandscapes

max_issues_repo_name: gideonsimpson/TestLandscapes

max_forks_repo_name: gideonsimpson/TestLandscapes


text:
function azeotrope_pressure(model::EoSModel, T; v0 =nothing)
    TYPE = promote_type(eltype(T),eltype([0.5,0.5]))
#     lb_v = lb_volume(model,x)
    ts = T_scales(model,[0.5,0.5])
    pmix = p_scale(model,[0.5,0.5])
    if v0 === nothing
        v0 = x0_bubble_pressure(model,T,[0.5,0.5])
    end
    len = length(v0[1:end-1])
    #xcache = zeros(eltype(x0),len)
    Fcache = zeros(eltype(v0[1:end-1]),len)
    f! = (F,z) -> Obj_az_pressure(model, F, T, exp10(z[1]), exp10(z[2]), z[3:end],z[3:end],ts,pmix)
    r  =Solvers.nlsolve(f!,v0[1:end-1],LineSearch(Newton()))
    sol = Solvers.x_sol(r)
    v_l = exp10(sol[1])
    v_v = exp10(sol[2])
    y = FractionVector(sol[3:end])
    P_sat = pressure(model,v_l,T,y)
    return (P_sat, v_l, v_v, y)
end

function Obj_az_pressure(model::EoSModel, F, T, v_l, v_v, x, y,ts,ps)
    x   = FractionVector(x)
    y   = FractionVector(y) #julia magic, check misc.jl
    μ_l = VT_chemical_potential(model,v_l,T,x)
    μ_v = VT_chemical_potential(model,v_v,T,y)
    p_l = pressure(model,v_l,T,x)
    p_v = pressure(model,v_v,T,y)
    for i in 1:length(x)
        F[i] = (μ_l[i]-μ_v[i])/(R̄*ts[i])
    end
    F[end] = (p_l-p_v)/ps
    return F
end

function azeotrope_temperature(model,p)
    f(z) = Obj_azeotrope_temperature(model,z,p)
    pure = split_model(model)
    sat = saturation_temperature.(pure,p)
    Ti   = zeros(length(pure),1)
    for i ∈ 1:length(pure)
        if isnan(sat[i][1])
            Tc,pc,vc = crit_pure(pure[i])
            g(x) = p-pressure(pure[i],vc,x,[1.])
            Ti[i] = Roots.find_zero(g,(Tc))
        else
            Ti[i] = sat[i][1]
        end
    end
    T = Roots.find_zero(f,(minimum(Ti)*0.9,maximum(Ti)*1.1))
    p,v_l,v_v,y = azeotrope_pressure(model,T)
    return T,v_l,v_v,y
end

function Obj_azeotrope_temperature(model,T,p)
    p̃,v_l,v_ll,xx = azeotrope_pressure(model,T)
    return p̃-p
end


text:
const UNIVERSAL_PLURALS = [
    "afreet" => "afreets",
    "agency" => "agencies",
    "bacterium" => "bacteria",
    "bias" => "biases",
    "bison" => "bison",
    "child" => "children",
    "child-in-law" => "children-in-law",
    "corps" => "corps",
    "datum" => "data",
    "deer" => "deer",
    "diagnosis" => "diagnoses",
    "die" => "dice",
    "elf" => "elves",
    "fish" => "fish",
    "fizz" => "fizzes",
    "foot" => "feet",
    "gas" => "gases",
    "goose" => "geese",
    "graffito" => "graffiti",
    "iron" => "irons",
    "louse" => "lice",
    "man" => "men",
    "moose" => "moose",
    "mother-in-law" => "mothers-in-law",
    "macro" => "macros",
    "neo" => "neos",
    "phenomenon" => "phenomena",
    "pig" => "pigs",
    "sex" => "sexes",
    "suffix" => "suffixes",    # "suffices" not standard Latin
    "tomato" => "tomatoes",
    "topaz" => "topazes",
    "unicorn" => "unicorns",
    "quiz" => "quizzes",
    "Javanese" => "Javanese",
    "Major General" => "Major Generals",
    "Postmaster General" => "Postmasters General"]

const CLASSICAL_PLURALS = [
    "beau" => "beaux",
    "cactus" => "cacti",
    "cow" => "kine",
    "focus" => "foci",
    "formula" => "formulae",
    "genius" => "genii",
    "iris" => "irides",
    "larynx" => "larynges",
    "lemma" => "lemmata",
    "lumen" => "lumina",
    "matrix" => "matrices",
    "milieu" => "milieux",
    "money" => "monies",
    "octopus" => "octopodes",  # "octopi" not standard Latin
    "seraph" => "seraphim",
    "soprano" => "soprani",
    "vertex" => "vertices"]

# classical plurals not expected to be handled by singularize
const EXTRA_CLASSICAL_PLURALS = [
    "thee" => "ye",
    "thou" => "ye"
]

const MODERN_PLURALS = [
    "beau" => "beaus",
    "cactus" => "cactuses",
    "cow" => "cows",
    "focus" => "focuses",
    "formula" => "formulas",
    "genius" => "geniuses",
    "iris" => "irises",
    "larynx" => "larynxes",
    "lemma" => "lemmas",
    "lumen" => "lumens",
    "matrix" => "matrixes",
    "milieu" => "milieus",
    "money" => "moneys",
    "octopus" => "octopuses",
    "seraph" => "seraphs",
    "soprano" => "sopranos",
    "vertex" => "vertexes"]

@testset "Pluralize" begin
    @testset for (singular, plural) in UNIVERSAL_PLURALS
        @test pluralize(singular, classical=true) == plural
        @test pluralize(singular, classical=false) == plural
    end
    @testset for (singular, plural) in [CLASSICAL_PLURALS;
                                        EXTRA_CLASSICAL_PLURALS]
        @test pluralize(singular, classical=true) == plural
    end
    @testset for (singular, plural) in MODERN_PLURALS
        @test pluralize(singular, classical=false) == plural
    end

    # invalid arguments currently throw — but we need to decide if this is
    # actually a good idea — it might be safer just to return the string
    # untouched
    @test_throws ArgumentError pluralize("")

    # currently we tolerate trailing whitespace, but what we return does not
    # have it — this is again a decision we should rethink
    @test pluralize("bar ") == "bars"

    # check that it works with a different string type
    # in this case, SubString
    @test pluralize(split("Julia language")[2]) == "languages"
end

@testset "Singularize" begin
    @testset for (singular, plural) in [UNIVERSAL_PLURALS;
                                        CLASSICAL_PLURALS;
                                        MODERN_PLURALS]
        @test singularize(plural) == singular
    end

    # we can't sensibly singularize words without singular forms; the most
    # sensible thing is to return the argument unchanged
    @test singularize("sand") == "sand"

    # check that it works with a different string type
    @test singularize(split("programming languages")[2]) == "language"
end


text:
# The package contains tools to solve an N-body problem in Newtonian gravity,
# and plot and animate the orbits and the consverved quantities.
# This package is based on the DifferentialEquations package tutorial:
# https://github.com/JuliaDiffEq/DiffEqTutorials.jl/blob/master/PhysicalModels/Outer-Solar-System.ipynb

module PlanetarySystem

using DifferentialEquations
using RecursiveArrayTools
using Plots

const G = 2.95912208286e-4 # Newton's constant, units: [L]= AU, [M]= M_sun, [t]= Earth days
const ∑ = sum # symbol for sum function
const N = 5 #Number of bodies (central star included)

#############################################################

# Physical model:

# Gravitational newtonian potential for the NBsolution() solver.
# p is a parameter, t the time, the coordinates x, y, z are arrays running over the bodies,
# M is the array of the bodies' masses
potential(p, t, x, y, z, M) = -G*∑(i->∑(j->(M[i]*M[j])/sqrt((x[i]-x[j])^2 + (y[i]-y[j])^2 + (z[i]-z[j])^2), 1:i-1), 2:N)

# Kinetic energy
# the velocity components  vx, vy, vz are 2D arrays running over the bodies and the time steps,
# M is the array of the bodies' masses
kinetic(vx, vy, vz, M) = 1/2 * ( ∑(M.*vx.*vx,dims=1) + ∑(M.*vy.*vy,dims=1) + ∑(M.*vz.*vz,dims=1) )
# potential, to compute the Hamiltonian
# the coordinates x, y, z are 2D arrays running over the bodies and the time steps,
# M is the array of the bodies' masses
pot(x, y, z, M)= - G*∑(i->∑(j->(M[i]*M[j])./sqrt.((x[i,:]-x[j,:]).^2+(y[i,:]-y[j,:]).^2+(z[i,:]-z[j,:]).^2), 1:i-1), 2:N)

# Hamiltonian (total energy) - variables as above
H(x, y, z, vx, vy, vz, M) = kinetic(vx, vy, vz, M) + pot(x, y, z, M)'

# Angular momentum components  - variables as above
Lx(x, y, z, vx, vy, vz, M) =  ∑(i-> M[i] .*  (vy[i,:] .* x[i,:] - vx[i,:] .* y[i,:]), 1:N)'
Ly(x, y, z, vx, vy, vz, M) =  ∑(i-> M[i] .*  (-vz[i,:] .* x[i,:] + vx[i,:] .* z[i,:]), 1:N)'
Lz(x, y, z, vx, vy, vz, M) =  ∑(i-> M[i] .*  (vz[i,:] .* y[i,:] - vy[i,:] .* z[i,:]), 1:N)'

export H, Lx, Ly, Lz

#############################################################

# Main
export SolveAndPlot
function SolveAndPlot(M::Array{Float64,1}, vel, pos, tspan::Tuple{Float64,Float64}, planets::Array{String,1}, plotfilename::String, animfilename::String, fifilename::String)

sol=NBsolution(M, vel, pos, tspan)
myplot(sol,plotfilename,planets)
animation(sol,animfilename,planets)
plot_first_integrals(sol, M, fifilename, planets)

end

#############################################################

# Solves the N-body problem with a sympletic method, Yoshida6.

# M is the array of the bodies' masses,
# vel contains the initial velocity vectors of all bodies (see example in Readme)
# pos contains the initial position of all bodies (see example in Readme)
# tspan is the time interval

export NBsolution
function NBsolution(M::Array{Float64,1}, vel, pos, tspan::Tuple{Float64,Float64})
    nprob = NBodyProblem(potential, M, vel, pos, tspan)
    sol = solve(nprob, Yoshida6(), dt=10) # can change the method or the time step dt
end

#############################################################

# Plots the orbits in coordinate space and saves to file
# sol is the output of NBsolution()
# filename is the name of the file where the plot is saved
# planets is the array of the planets' names

export myplot
function myplot(sol, filename::String, planets::Array{String,1})
    plot(size=(500, 400),
         title  = "Planetary orbits",
         legend=:bottomleft,
         xlim=(-30,30), ylim=(-30,30), zlim=(-30,30)
         ,markerstrokewidth=0
         #,bg=:black
         #,xlabel = "x", ylabel = "y", zlabel = "z"
         )

    for i in 1:N
    scatter!(sol,
              vars=(3*N+i,4*N+i,5*N+i),
              label=planets[i],
              markersize=1,
              markerstrokewidth=0)
    end

    savefig(filename)
end

#############################################################

# Animates the planets and saves to file
# see above

export animation
function animation(sol, filename::String, planets::Array{String,1})

    nplots=500
    element(i)=round(Int, length(sol.t)/nplots-1)*i

    anim = @animate for i=1:nplots #
        scatter((sol[3*N+1,element(i)],sol[4*N+1,element(i)],sol[5*N+1,element(i)]),
                #xlabel = "x", ylabel = "y", zlabel = "z",
                xlim=(-30,30), ylim=(-30,30), zlim=(-30,30),
                label=planets[1],
                markerstrokewidth=0,
                markercolor=:yellow,
                marker=10,
                title  = "Planetary system",
                size=(500, 400))
        if(N>1)
            for j in 2:(N)
        scatter!((sol[3*N+j,element(i)],sol[4*N+j,element(i)],sol[5*N+j,element(i)]),
                 marker = 5,
                 markerstrokewidth=0,
                 label=planets[j])
            end
        end
    end

    gif(anim, filename, fps = 17) #change fps to adjust the speed of the animation

end

#############################################################

# Plots the fractional variation of conserved quantities (energy and angular momentum)
# see above

export plot_first_integrals
function plot_first_integrals(sol, M::Array{Float64,1}, filename::String, planets::Array{String,1})

    vx=sol[1:N,:]
    vy=sol[N+1:2*N,:]
    vz=sol[2*N+1:3*N,:]
    x= sol[3*N+1:4*N,:]
    y= sol[4*N+1:5*N,:]
    z= sol[5*N+1:6*N,:]

    plot(sol.t, (H(x, y, z, vx, vy, vz, M)[1] .- H(x, y, z, vx, vy, vz, M)')./H(x, y, z, vx, vy, vz, M)[1],
         lab="Energy variation",
         title="Conserved Quantities",
         xlabel="t",
         size=(500, 400))
    plot!(sol.t, (Lx(x, y, z, vx, vy, vz, M)[1].-Lx(x, y, z, vx, vy, vz, M)')./Lx(x, y, z, vx, vy, vz, M)[1], lab="Angular momentum variation, x")
    plot!(sol.t, (Ly(x, y, z, vx, vy, vz, M)[1].-Ly(x, y, z, vx, vy, vz, M)')./Ly(x, y, z, vx, vy, vz, M)[1], lab="Angular momentum variation, y")
    plot!(sol.t, (Lz(x, y, z, vx, vy, vz, M)[1].-Lz(x, y, z, vx, vy, vz, M)')./Lz(x, y, z, vx, vy, vz, M)[1], lab="Angular momentum variation, z")

    savefig(filename)

end

end # module


text:
using Test, Sandbox

if Sys.islinux()
    @test isa(Sandbox.get_kernel_version(), VersionNumber)
    @test Sandbox.check_kernel_version()
end

if executor_available(UnprivilegedUserNamespacesExecutor)
    @testset "UnprivilegedUserNamespacesExecutor" begin
        @test_logs (:info, "Testing Unprivileged User Namespaces Executor (read-only, read-write)") begin
            with_executor(UnprivilegedUserNamespacesExecutor) do exe
                @test probe_executor(exe; test_read_only_map=true, test_read_write_map=true, verbose=true)
            end
        end
    end
else
    @error("Skipping Unprivileged tests, as it does not seem to be available")
end

# Only test privileged runner if sudo doesn't require a password
if success(`sudo -k -n true`)
    if executor_available(PrivilegedUserNamespacesExecutor)
        @testset "PrivilegedUserNamespacesExecutor" begin
            @test_logs (:info, "Testing Privileged User Namespaces Executor (read-only, read-write)") begin
                with_executor(PrivilegedUserNamespacesExecutor) do exe
                    @test probe_executor(exe; test_read_only_map=true, test_read_write_map=true, verbose=true)
                end
            end
        end
    else
        @error("Skipping Privileged tests, as it does not seem to be available")
    end
end


text:
using .CUDA

cuda_available = true

function compute_integral(m, qx, qy, xs, nx::NamedTuple{(:nx, :type, :batchsize)})
    integrate_cuda(m, nx.type, qx, qy, xs, nx.batchsize)
end

function integrate_cuda_inner!(m::Softcore, A::AbstractArray{T}, x, y, xy::AbstractVector{<:Tuple}, stride) where T
    dist2((x1, x2), (y1, y2)) = (x1-y1)^2 + (x2-y2)^2
    k0 = threadIdx().x + (blockIdx().x - 1) * blockDim().x
    j0 = threadIdx().y + (blockIdx().y - 1) * blockDim().y
    if k0 <= stride[1] && j0 <= stride[2]
        for k in k0:stride[1]:length(x)
            for j = j0:stride[2]:length(y)
                z = zero(T)
                for l in 1:length(xy)
                    d = dist2((x[k],y[j]),xy[l]) :: T
                    d = sqrt(d) :: T
                    z += m.kernel_integral(d)
                    A[k0, j0, l] += exp(-z)
                end
            end
        end
    end
    nothing
end
function integrate_cuda(m, T::Type, x,y,xy, batchsize)
    A = CUDA.zeros(T, (cld(length(x), batchsize), cld(length(y), batchsize), length(xy)))
    x = convert(CuVector{T}, collect(x))
    y = convert(CuVector{T}, collect(y))
    xy = convert(CuVector{Tuple{T,T}}, xy)
    num_threads = (256,1)
    num_blocks = cld.(size(A)[1:2], num_threads)
    stride = size(A)[1:2]
    @cuda blocks=num_blocks threads=num_threads integrate_cuda_inner!(m, A, x, y, xy, stride)
    z = sum(A, dims=(1,2))
    a = collect(z)
    CUDA.unsafe_free!(A)
    CUDA.unsafe_free!(z)
    a
end

max_stars_repo_name: mikkoku/SequentialPointProcesses

max_issues_repo_name: mikkoku/SequentialPointProcesses

max_forks_repo_name: mikkoku/SequentialPointProcesses


text:
export PekerisRayModel

"""
    struct PekerisRayModel{T} <: PropagationModel{T}

A fast differentiable ray model that only supports isovelocity constant depth
environments.

---

    PekerisRayModel(env, rays)

Create a Pekeris ray propagation model with a maximum of `rays` ray arrivals.
"""
struct PekerisRayModel{T} <: PropagationModel{T}
  env::T
  rays::Int
  function PekerisRayModel(env, rays)
    rays > 0 || throw(ArgumentError("Number of rays should be more than 0"))
    new{typeof(env)}(check(PekerisRayModel, env), rays)
  end
end

### interface functions

function check(::Type{PekerisRayModel}, env::Union{<:UnderwaterEnvironment,Missing})
  if env !== missing
    altimetry(env) isa FlatSurface || throw(ArgumentError("PekerisRayModel only supports environments with flat sea surface"))
    bathymetry(env) isa ConstantDepth || throw(ArgumentError("PekerisRayModel only supports constant depth environments"))
    ssp(env) isa IsoSSP || throw(ArgumentError("PekerisRayModel only supports isovelocity environments"))
  end
  env
end

function arrivals(model::PekerisRayModel, tx1::AcousticSource, rx1::AcousticReceiver)
  # based on Chitre (2007)
  c = soundspeed(ssp(model.env), 0.0, 0.0, 0.0)
  h = depth(bathymetry(model.env), 0.0, 0.0)
  f = nominalfrequency(tx1)
  p1 = location(tx1)
  p2 = location(rx1)
  R² = abs2(p1[1] - p2[1]) + abs2(p1[2] - p2[2])
  R = R² == 0 ? R² : √R²   # ForwardDiff compatible version of √R²
  d1 = -p1[3]
  d2 = -p2[3]
  [arrival(j, model, R, R², d1, d2, h, c, f) for j ∈ 1:model.rays]
end

function eigenrays(model::PekerisRayModel, tx1::AcousticSource, rx1::AcousticReceiver)
  # based on Chitre (2007)
  c = soundspeed(ssp(model.env), 0.0, 0.0, 0.0)
  h = depth(bathymetry(model.env), 0.0, 0.0)
  f = nominalfrequency(tx1)
  p1 = location(tx1)
  p2 = location(rx1)
  R² = abs2(p1[1] - p2[1]) + abs2(p1[2] - p2[2])
  R = R² == 0 ? R² : √R²   # ForwardDiff compatible version of √R²
  d1 = -p1[3]
  d2 = -p2[3]
  [arrival(j, model, R, R², d1, d2, h, c, f, p1, p2) for j ∈ 1:model.rays]
end

function rays(model::PekerisRayModel, tx1::AcousticSource, θ::Real, rmax)
  -π/2 < θ < π/2 || throw(ArgumentError("θ must be between -π/2 and π/2"))
  c = soundspeed(ssp(model.env), 0.0, 0.0, 0.0)
  h = depth(bathymetry(model.env), 0.0, 0.0)
  f = nominalfrequency(tx1)
  p1 = location(tx1)
  d1 = -p1[3]
  d2 = d1 - rmax * tan(θ)
  s = 0
  b = 0
  while true
    if d2 < 0
      s += 1
      d2 = -d2
    elseif d2 > h
      b += 1
      d2 = 2h - d2
    else
      break
    end
  end
  if 4s - 2 == 4b + 2
    j = 4s - 2
  elseif 4s + 3 == 4b - 1
    j = 4s + 3
  else
    @assert s == b
    j = θ > 0 ? 4s : 4s + 1
  end
  p2 = (p1[1] + rmax, p1[2], -d2)
  arrival(j, model, rmax, rmax*rmax, d1, d2, h, c, f, p1, p2)
end

### helper functions

# memoized version of absorption
const cached_absorption = Ref{Tuple{Float64,Float64,Float64}}()
function fast_absorption(f::Float64, D, S::Float64)
  if cached_absorption[][1:2] == (f, S)
    db2amp(cached_absorption[][3] * D)
  else
    dBperm = amp2db(absorption(f, 1.0, S))
    cached_absorption[] = (f, S, dBperm)
    db2amp(dBperm * D)
  end
end

# fallback
fast_absorption(f, D, S) = absorption(f, D, S)

# complex ForwardDiff friendly version of x^n
ipow(x, n::Int) = prod(x for _ ∈ 1:n)

function arrival(j, model, R, R², d1, d2, h, c, f, p1=missing, p2=missing)
  upward = iseven(j)
  s1 = 2*upward - 1
  n = div(j, 2)
  s = div(n + upward, 2)
  b = div(n + (1-upward), 2)
  s2 = 2*iseven(n) - 1
  dz = 2*b*h + s1*d1 - s1*s2*d2
  D = √(R² + abs2(dz))
  θ = atan(R, dz)
  t = D/c
  A = Complex(1.0, 0.0) / D * fast_absorption(f, D, salinity(model.env))
  s > 0 && (A *= ipow(reflectioncoef(seasurface(model.env), f, θ), s))
  b > 0 && (A *= ipow(reflectioncoef(seabed(model.env), f, θ), b))
  λ = π/2 - θ
  if typeof(p1) === Missing
    RayArrival(t, conj(A), s, b, s1*λ, -s1*s2*λ)    # conj(A) needed to match with Bellhop
  else
    raypath = Array{typeof(p1)}(undef, 2+s+b)
    raypath[1] = p1
    if s + b > 0
      dx = p2[1] - p1[1]
      dy = p2[2] - p1[2]
      z = (1-upward) * h
      r = abs(z-d1) * tan(θ)
      raypath[2] = (p1[1] + r/R * dx, p1[2] + r/R * dy, -z)
      for i ∈ 3:length(raypath)-1
        r += h * tan(θ)
        z = h - z
        raypath[i] = (p1[1] + r/R * dx, p1[2] + r/R * dy, -z)
      end
    end
    raypath[end] = p2
    RayArrival(t, conj(A), s, b, s1*λ, -s1*s2*λ, raypath)
  end
end


text:
"""
    reindex(a::AbstractIndex, inds::AbstractVector{Integer}) -> AbstractIndex

Returns and index of the same type as `a` where the keys the new keys are
constructed by indexing into the keys of `a` with `inds` and the values have the
same starting value but a length matching `inds`.

## Examples
```jldoctest
julia> x, y, z = Index(1:10, 2:11), Index(1:10), SimpleIndex(1:10);

julia> reindex(x, collect(1:2:10))
Index([1, 3, 5, 7, 9] => 2:6)

julia> reindex(y, collect(1:2:10))
Index([1, 3, 5, 7, 9] => Base.OneTo(5))

julia> reindex(z, collect(1:2:10))
SimpleIndex(1:5)
```
"""
function reindex(a::AbstractIndex, inds::AbstractVector{T}) where {T<:Integer}
    @boundscheck checkbounds(a, inds)
    return unsafe_reindex(a, inds)
end

"""
    unsafe_reindex(a::AbstractIndex, inds::AbstractVector) -> AbstractIndex

Similar to `reindex` this function returns an index of the same type as `a` but
doesn't check that `inds` is inbounds. New subtypes of `AbstractIndex` must
implement a unique `unsafe_reindex` method.

See also: [`reindex`](@ref)
"""
function unsafe_reindex(a::AbstractIndex, inds)
    error("New subtypes of `AbstractIndex` must implement a unique `unsafe_reindex` method.")
end
function unsafe_reindex(a::Index{name}, inds) where {name}
    return Index{name}(
        @inbounds(keys(a)[inds]),
        _reindex(values(a), inds),
        AllUnique,
        LengthChecked
       )
end
function unsafe_reindex(a::SimpleIndex{name}, inds) where {name}
    return SimpleIndex{name}(_reindex(values(a), inds))
end

_reindex(a::OneTo{T}, inds) where {T} = OneTo{T}(length(inds))
_reindex(a::OneToMRange{T}, inds) where {T} = OneToMRange{T}(length(inds))
_reindex(a::OneToSRange{T}, inds) where {T} = OneToSRange{T}(length(inds))
_reindex(a::T, inds) where {T<:AbstractUnitRange} = T(first(a), first(a) + length(inds) - 1)