text:
@testset "Ball" begin
    s = Ball([3,4], 5)
    @test s == deepcopy(s)
    @test hash(s) == hash(deepcopy(s))
    @test ndims(s) == 2
    @test [3,9] ∈ s
    @test [3,9.1] ∉ s

    @test ((x,nout) = surfpt_nearby([3,4],s); norm(x-[3,4])≈5 && norm(nout)≈1)  # handle point at center properly
    @test all([surfpt_nearby([3+ρ*5sx,4],s)[1] ≈ [3+5sx,4] for ρ = (one⁻⁻,1,one⁺⁺), sx = (-1,1)])
    @test all([surfpt_nearby([3,4+ρ*5sy],s)[1] ≈ [3,4+5sy] for ρ = (one⁻⁻,1,one⁺⁺), sy = (-1,1)])
    @test all([surfpt_nearby([3+ρ*5sx/√2,4+ρ*5sy/√2],s)[1] ≈ [3+5sx/√2,4+5sy/√2]
               for ρ = (one⁻⁻,one⁺⁺), sx = (-1,1), sy = (-1,1)])
    @test all([(p = [3+5sx,4]; surfpt_nearby(p,s) ≈ (p,[sx,0])) for sx = (-1,1)])  # handle point on boundary properly
    @test all([(p = [3,4+5sy]; surfpt_nearby(p,s) ≈ (p,[0,sy])) for sy = (-1,1)])  # handle point on boundary properly

    @test normal([-1,2],s) == normalize([-1,2] - [3,4])
    @test bounds(s) == ([-2,-1],[8,9])
    @test checkbounds(s)
    @test checkbounds(Ball([1,2,3], 2))

    @test (∆ = rand(2); translate(s,∆) ≈ Ball([3,4]+∆, 5))
end


text:
using Test: @test, @testset, @inferred
using Statistics: mean, std, var
using Random: randn!, bitrand
using LogExpFunctions: logsumexp
using RestrictedBoltzmannMachines: RBM, BinaryRBM, Binary, Spin, Potts, Gaussian, ReLU, dReLU, xReLU, pReLU
using RestrictedBoltzmannMachines: energy, free_energy, sample_from_inputs, sample_v_from_v
using RestrictedBoltzmannMachines: anneal, anneal_zero, ais, aise, raise, log_partition_zero_weight
using RestrictedBoltzmannMachines: logmeanexp, logvarexp, logstdexp
using RestrictedBoltzmannMachines: log_partition

@testset "logmeanexp, logvarexp" begin
    A = randn(5,3,2)
    for dims in (2, (1,2), :)
        @test logmeanexp(A; dims) ≈ log.(mean(exp.(A); dims))
        for corrected in (true, false)
            @test logvarexp(A; dims, corrected) ≈ log.(var(exp.(A); dims, corrected))
            @test logstdexp(A; dims, corrected) ≈ log.(std(exp.(A); dims, corrected))
        end
    end
    @test logvarexp(A; dims=2) ≈ log.(var(exp.(A); dims=2))
    @test logstdexp(A; dims=2) ≈ log.(std(exp.(A); dims=2))
    @test logmeanexp(A) ≈ log.(mean(exp.(A)))
    @test logvarexp(A) ≈ log.(var(exp.(A)))
    @test logstdexp(A) ≈ log.(std(exp.(A)))
end

@testset "logmeanexp properties" begin
    X = randn(1000, 1000)
    @test only(logmeanexp(logmeanexp(X; dims=1); dims=2)) ≈ logmeanexp(X)
    @test only(logmeanexp(-logmeanexp(-X; dims=1); dims=2)) ≤ only(-logmeanexp(-logmeanexp(X; dims=1); dims=2))
    x = randn()
    @test logmeanexp([x]) ≈ x
    X = randn(1000, 1000, 1)
    @test logmeanexp(X; dims=3) ≈ X
    @test -logmeanexp(-X) ≤ logmeanexp(X)
end

@testset "log_partition_zero_weight" begin
    rbm0 = BinaryRBM(randn(3), randn(2), zeros(3,2))
    @test log_partition_zero_weight(rbm0) ≈ log_partition(rbm0)
    rbm1 = RBM(rbm0.visible, rbm0.hidden, randn(3,2))
    @test log_partition_zero_weight(rbm1) ≈ log_partition_zero_weight(rbm0)
end

@testset "anneal layer" begin
    β = 0.3
    N = 11

    init = Binary(randn(N))
    final = Binary(randn(N))
    x = sample_from_inputs(final)
    @test energy(anneal(init, final; β), x) ≈ (1 - β) * energy(init, x) + β * energy(final, x)

    init = Spin(randn(N))
    final = Spin(randn(N))
    x = sample_from_inputs(final)
    @test energy(anneal(init, final; β), x) ≈ (1 - β) * energy(init, x) + β * energy(final, x)

    init = Potts(randn(N))
    final = Potts(randn(N))
    x = sample_from_inputs(final)
    @test energy(anneal(init, final; β), x) ≈ (1 - β) * energy(init, x) + β * energy(final, x)

    init = Gaussian(randn(N), rand(N))
    final = Gaussian(randn(N), rand(N))
    x = sample_from_inputs(final)
    @test energy(anneal(init, final; β), x) ≈ (1 - β) * energy(init, x) + β * energy(final, x)

    init = ReLU(randn(N), rand(N))
    final = ReLU(randn(N), rand(N))
    x = sample_from_inputs(final)
    @test energy(anneal(init, final; β), x) ≈ (1 - β) * energy(init, x) + β * energy(final, x)

    init = dReLU(randn(N), randn(N), rand(N), rand(N))
    final = dReLU(randn(N), randn(N), init.γp, init.γn)
    x = sample_from_inputs(final)
    @test energy(anneal(init, final; β), x) ≈ (1 - β) * energy(init, x) + β * energy(final, x)

    init = pReLU(randn(N), rand(N), randn(N), rand(N) .- 0.5)
    final = pReLU(randn(N), init.γ, randn(N), init.η)
    x = sample_from_inputs(final)
    @test energy(anneal(init, final; β), x) ≈ (1 - β) * energy(init, x) + β * energy(final, x)

    init = xReLU(randn(N), rand(N), randn(N), randn(N))
    final = xReLU(randn(N), init.γ, randn(N), init.ξ)
    x = sample_from_inputs(final)
    @test energy(anneal(init, final; β), x) ≈ (1 - β) * energy(init, x) + β * energy(final, x)
end

@testset "anneal_zero" begin
    β = 0.3
    N = 11
    M = 5
    rbm = BinaryRBM(randn(N), randn(M), randn(N,M))
    init = Binary(randn(N))
    null = BinaryRBM(init.θ, zeros(M), zeros(N,M))
    v = bitrand(N, 10)
    h = bitrand(M, 10)
    @test energy(anneal_zero(init, rbm), v, h) ≈ energy(anneal(null, rbm; β=0), v, h)
    @test iszero(anneal_zero(init, rbm).w)

    layer = Binary(randn(N))
    @test iszero(anneal_zero(layer).θ)

    layer = Spin(randn(N))
    @test iszero(anneal_zero(layer).θ)

    layer = Potts(randn(N))
    @test iszero(anneal_zero(layer).θ)

    layer = Gaussian(randn(N), rand(N))
    @test iszero(anneal_zero(layer).θ)
    @test anneal_zero(layer).γ == layer.γ

    layer = ReLU(randn(N), rand(N))
    @test iszero(anneal_zero(layer).θ)
    @test anneal_zero(layer).γ == layer.γ

    layer = dReLU(randn(N), randn(N), rand(N), rand(N))
    @test iszero(anneal_zero(layer).θp)
    @test iszero(anneal_zero(layer).θn)
    @test anneal_zero(layer).γp == layer.γp
    @test anneal_zero(layer).γn == layer.γn

    layer = pReLU(randn(N), rand(N), randn(N), rand(N) .- 0.5)
    @test iszero(anneal_zero(layer).θ)
    @test iszero(anneal_zero(layer).Δ)
    @test anneal_zero(layer).γ == layer.γ
    @test anneal_zero(layer).η == layer.η

    layer = xReLU(randn(N), rand(N), randn(N), randn(N))
    @test iszero(anneal_zero(layer).θ)
    @test iszero(anneal_zero(layer).Δ)
    @test anneal_zero(layer).γ == layer.γ
    @test anneal_zero(layer).ξ == layer.ξ
end

@testset "anneal" begin
    N = 10
    M = 7
    β = 0.3
    rbm0 = BinaryRBM(randn(N), randn(M), randn(N,M))
    rbm1 = BinaryRBM(randn(N), randn(M), randn(N,M))
    rbm = anneal(rbm0, rbm1; β)
    v = bitrand(N)
    h = bitrand(M)
    @test energy(rbm, v, h) ≈ (1 - β) * energy(rbm0, v, h) + β * energy(rbm1, v, h)
end

@testset "ais" begin
    rbm0 = BinaryRBM(randn(2), randn(1), zeros(2,1))
    rbm1 = BinaryRBM(randn(2), randn(1), randn(2,1))
    v0 = sample_v_from_v(rbm0, bitrand(2, 100000); steps=1)
    data = ais(rbm0, rbm1, v0; nbetas=10)
    @test logmeanexp(data) ≈ log_partition(rbm1) - log_partition(rbm0) rtol=0.1
    @test mean(logmeanexp.(Iterators.partition(data, 10))) ≈ log_partition(rbm1) - log_partition(rbm0) rtol=0.1

    v0 = sample_v_from_v(rbm0, bitrand(2); steps=1)
    data = ais(rbm0, rbm1, v0; nbetas=100000)
    @test only(data) ≈ log_partition(rbm1) - log_partition(rbm0) rtol=0.1
end

@testset "aise" begin
    rbm = BinaryRBM(randn(2), randn(1), randn(2,1))
    lZ = aise(rbm; nbetas=3, nsamples=100000)
    @test logmeanexp(lZ) ≈ log_partition(rbm) rtol=0.1
    lZ = aise(rbm; nbetas=10000, nsamples=10)
    @test logmeanexp(lZ) ≈ log_partition(rbm) rtol=0.1
end

@testset "raise (w = 0)" begin
    rbm = BinaryRBM(randn(2), randn(1), zeros(2,1))
    v = sample_v_from_v(rbm, bitrand(2, 50000); steps=1)
    lZ = raise(rbm; v, nbetas=10)
    @test logmeanexp(lZ) ≈ log_partition(rbm) rtol=0.1
    v = sample_v_from_v(rbm, bitrand(2, 10); steps=1)
    lZ = raise(rbm; v, nbetas=10000)
    @test logmeanexp(lZ) ≈ log_partition(rbm) rtol=0.1
end

@testset "raise" begin
    rbm = BinaryRBM(randn(2), randn(1), randn(2,1))
    v = sample_v_from_v(rbm, bitrand(2, 50000); steps=1000) # need to equilibrate
    lZ = raise(rbm; v, nbetas=10)
    @test logmeanexp(lZ) ≈ log_partition(rbm) rtol=0.1
    @test -logmeanexp(-lZ) ≈ log_partition(rbm) rtol=0.1
end

@testset "aise ≤ lZ ≤ raise" begin
    rbm = BinaryRBM(randn(3), randn(2), randn(3,2))
    lZ = log_partition(rbm)
    v = sample_v_from_v(rbm, bitrand(size(rbm.visible)..., 100_000); steps=500)
    lZf = aise(rbm; nbetas=3, nsamples=size(v)[end])
    lZr = raise(rbm; v, nbetas=3)
    @info @test mean(logmeanexp.(Iterators.partition(lZf, 2))) ≤ lZ ≤ mean(-logmeanexp.(Iterators.partition(-lZr, 2)))
    @test -logmeanexp(-lZr) ≤ logmeanexp(lZr)
end

@testset "ais gaussian" begin
    rbm0 = RBM(Gaussian(randn(20), 1 .+ 5rand(20)), Gaussian(randn(10), 1 .+ rand(10)), randn(20, 10)/10)
    rbm1 = RBM(Gaussian(randn(20), 1 .+ 5rand(20)), Gaussian(randn(10), 1 .+ rand(10)), randn(20, 10)/10)
    lZ0 = log_partition(rbm0)
    lZ1 = log_partition(rbm1)
    v0 = sample_v_from_v(rbm0, randn(20, 10); steps=1)
    R = ais(rbm0, rbm1, v0, nbetas=10000)
    @test logmeanexp(R) ≈ lZ1 - lZ0 rtol=0.1
end

@testset "zero hidden units" begin
    rbm = BinaryRBM(randn(2), randn(0), randn(2,0))
    v = sample_v_from_v(rbm, bitrand(2, 100); steps=1)
    lZ = aise(rbm; nbetas=3, nsamples=100)
    @test logmeanexp(lZ) ≈ log_partition(rbm)
    lZ = raise(rbm; v, nbetas=3)
    @test logmeanexp(lZ) ≈ log_partition(rbm)
end

ext: jl


text:
using PolyChaos
using LinearAlgebra
using CSV
using DelimitedFiles
using DataFrames
using Dates
using ColorSchemes
using StatsBase
using Random

export processInputs,
       getInputNames,
       processInputsBg,
       getInputNamesBg,
       processMissingValues,
       processOutputs,
       buildMultivariateBasis,
       buildCoefficientMatrix,
       solvePCE,
       solveRegPCE,
       performGSA,
       gsaMain,
       gsa,
       processMainEffects,
       getSubsetMI,
       getSubsetIndex,
       processTimeInfo,
       plotMainEffects,
       plotInteractionEffects,
       bootstrapGSA,
       getBootstrapSummary,
       regularized_fit,
       regularized_plot,
       train_test_split

# function to read and process inputs
function processInputs(inputsPath, retainedRunIdx; useBg=true)
    inputRaw = CSV.read(inputsPath, DataFrame)
    if !useBg
        inputsFiltered = inputRaw[!, Not(["BrFactor_ADAPT",
                                    "PoyntingFluxPerBSi",
                                    "LperpTimesSqrtBSi",
                                    "Helicity",
                                    "bgRunIdx",
                                        ])
                                    ]
    else
        inputsFiltered = inputRaw[!, Not(["Helicity",
                                    "bgRunIdx",
                                        ])
                                    ]
    end

    X = Matrix(inputsFiltered[retainedRunIdx, :])
    return X
end

"""
Function to process inputs for background solar wind runs alone.
Assume that input file is supplied in the standardized 0-1 scaling, with following
columns:
BrMin, BrFactor_ADAPT, nChromoSiAWSoM, PoyntingFluxPerBSi, LperpTimesSqrtBSi,
StochasticExponent, rMinWaveReflection

"""
function processInputsBg(inputsPath, retainedRunIdx)
    inputRaw = CSV.read(inputsPath, DataFrame)
    X = Matrix(inputRaw[retainedRunIdx, :])
    return X
end

"""
Fetch parameter names for background solar wind runs
"""
function getInputNamesBg(inputsPath)
    inputRaw = CSV.read(inputsPath, DataFrame)
    return names(inputRaw)
end

function getInputNames(inputsPath; useBg=true)
    inputRaw = CSV.read(inputsPath, DataFrame)
    if !useBg
        inputsFiltered = inputRaw[!, Not(["BrFactor_ADAPT",
                                    "PoyntingFluxPerBSi",
                                    "LperpTimesSqrtBSi",
                                    "Helicity",
                                    "bgRunIdx",
                                        ])
                                    ]
    else
        inputsFiltered = inputRaw[!, Not(["Helicity",
                                    "bgRunIdx",
                                        ])
                                    ]
    end
    return names(inputsFiltered)
end

function processMissingValues(outputRaw)
    missingIdx = [findall(ismissing, outputRaw[:, i]) for i in 1:size(outputRaw, 2)]
    mostMissingsCol = argmax(length.(missingIdx))
    # Identify column with greatest number of missings
    dMissing = diff(missingIdx[mostMissingsCol])
    # if elements are missing at beginning and end
    if maximum(dMissing) > 1
        trimmedRange = (argmax(dMissing) + 1):(argmax(dMissing) + maximum(dMissing) - 1)
    # if elements are missing only at beginning or only at end
    elseif maximum(dMissing) == 1
        if missingIdx[mostMissingsCol][1] > 1 # missing only at end
            trimmedRange = 1:(missingIdx[mostMissingsCol][1] - 1)
        else # missing only at beginning
            trimmedRange = (missingIdx[mostMissingsCol][end] + 1):size(outputRaw, 1)
        end
    end
    return trimmedRange
end

"""
Function to read and process shifted / unshifted outputs. For background solar
wind runs, only works with unshifted runs.
"""
function processOutputs(outputPath, retainedRunIdx; QoI="Ur", useShifted=false)
    if useShifted
        outputRaw = readdlm(joinpath(outputPath, QoI * "SimTrimmed.txt"))
        replace!(x -> x == "missing" ? missing : x, outputRaw)
        # trim off missing values!
        trimmedRange = processMissingValues(outputRaw)
        Y = outputRaw[trimmedRange, :]
    else
        outputRaw = readdlm(joinpath(QOIS_PATH, QoI * "Sim_earth.txt"))
        Y = outputRaw[:, retainedRunIdx]
    end

    return Array{Float64, 2}(Y')
end

# function to build multivariate basis for PCE
function buildMultivariateBasis(D, P)
    # D is dim of stochastic space
    # P is degree of PCE

    univariateBasis = Uniform01OrthoPoly(P) # Monic case
    # univariateBasis = Legendre01OrthoPoly(P)

    return MultiOrthoPoly([univariateBasis for i in 1:D], P)
end

# function to build PCE
function buildCoefficientMatrix(X; pceDegree=2)
    nd = size(X, 2) # dimension of stochastic space
    multivariateBasis = buildMultivariateBasis(nd, pceDegree)
    # monic is false, but normalizing is true, alternately both are true!
    # revise this later to add it as an argument.
    A = evaluate(X, multivariateBasis, true, true)'
    # A = evaluate(X, multivariateBasis, false, true)'
    return Matrix(A)
end

# function to regress PCE
function solvePCE(A, y; regularize=false)
    β = A \ y
    return β
end

"""
Function to solve regularized linear system for PCE. By default, regularization parameter
λ is 0 i.e. it performs OLS.
"""
function solveRegPCE(A, y; λ = 0)
    β = (A'A + λ * I) \ (A' * y)
    return β
end

# function to return subset of multi-index for particular effect
function getSubsetMI(multiIndex, effectIdx)
    subsetMI = []
    m, n = size(multiIndex)
    for i in 1:m
        # maybe set notation would simplify this operation?
        push!(subsetMI, (all(multiIndex[i, effectIdx] .> 0)) & (multiIndex[i, setdiff(1:n, effectIdx)] == zeros((n - length(unique(effectIdx))))))
    end
    return Array{Bool, 1}(subsetMI)
end

function getSubsetIndex(multiIndex, effectIdx)
    subsetMI = []
    m, n = size(multiIndex)
    for i in 1:m
        # maybe set notation would simplify this operation?
        push!(subsetMI, (all(multiIndex[i, effectIdx] .> 0)) & (multiIndex[i, setdiff(1:n, effectIdx)] == zeros((n - length(unique(effectIdx))))))
    end
    return multiIndex[Array{Bool, 1}(subsetMI), :]
end

function performGSA(inputsPath, outputPath; gsa_kwargs...)
    # process inputs
    X = processInputs(inputsPath, retainedRunIdx; useBg=useBg)
    # A = buildCoefficientMatrix(inputsPath, retainedRunIdx; useBg=useBg, pceDegree=pceDegree)
    # process outputs
    Y = processOutputs(outputPath, retainedRunIdx;
                    QoI=QoI,
                    useShifted=useShifted
                    )

    # use X and Y to get Sobol Indices
    gsaIndices = gsa(X, Y;
                    regularize=regularize,
                    pceDegree = pceDegree
                    )

    return gsaIndices
end

# function to calculate _only_ main effects!
function gsaMain(β::AbstractArray; nEffects = 7, pceDegree=2, multiIndex=nothing)
    nTimePts = size(β, 2)
    mainEffects = zeros(nEffects, nTimePts);
    # Diagonal elements contain main effects
    # Off diagonal elements contain interactions
    if isnothing(multiIndex)
        multiIndex = buildMultivariateBasis(nEffects, pceDegree).ind
    end
    for t in 1:nTimePts
        varF = sum(β[2:end, t].^2)
        for i in 1:nEffects
            subsetMultiIndex = getSubsetMI(multiIndex, [i, i])
            varSubset = sum(β[findall(subsetMultiIndex), t].^2)
            mainEffects[i, t] = varSubset / varF
        end
    end
    return mainEffects
end

# function to calculate tensor of main and joint effects
function gsa(X, Y; regularize=false, pceDegree=2)

    nEffects = size(X, 2)
    nTimePts = size(Y, 2)

    gsaIndices = zeros(nEffects, nEffects, nTimePts);
    # Diagonal elements contain main effects
    # Off diagonal elements contain interactions

    A = buildCoefficientMatrix(X; pceDegree=pceDegree)
    β = solvePCE(A, Y; regularize=regularize)
    multiIndex = buildMultivariateBasis(nEffects, pceDegree).ind
    for t in 1:nTimePts
        varF = sum(β[2:end, t].^2)
        for i in 1:nEffects
            for j in 1:nEffects
                subsetMultiIndex = getSubsetMI(multiIndex, [i, j])
                varSubset = sum(β[findall(subsetMultiIndex), t].^2)
                gsaIndices[i, j, t] = varSubset / varF
            end
        end
    end
    return gsaIndices
end

# calculate with β as input!
function gsa(β::AbstractArray; nEffects=7, pceDegree=2)
    @assert factorial(nEffects + pceDegree) / (factorial(nEffects) * factorial(pceDegree)) == size(β, 1)
    nTimePts = size(β, 2)
    multiIndex = buildMultivariateBasis(nEffects, pceDegree).ind
    gsaIndices = zeros(nEffects, nEffects, nTimePts);
    for t in 1:nTimePts
        varF = sum(β[2:end, t].^2)
        for i in 1:nEffects
            for j in 1:nEffects
                subsetMultiIndex = getSubsetMI(multiIndex, [i, j])
                varSubset = sum(β[findall(subsetMultiIndex), t].^2)
                gsaIndices[i, j, t] = varSubset / varF
            end
        end
    end
    return gsaIndices
end

function processMainEffects(gsaIndices)
    nEffects = size(gsaIndices, 1)
    nTimePts = size(gsaIndices, 3)
    mainEffects = zeros(nEffects, nTimePts)
    for i in 1:nTimePts
        mainEffects[:, i] = diag(gsaIndices[:, :, i])
    end
    return mainEffects
end

# function to make plots of everything (this will require time information)
"""
Process time to enable plotting of time information on graph.
Problematic: Trim index information. Need to reorganize data so it is
easier to deal with.
"""
function processTimeInfo(timeData; trimIndices = (39, 141))
    if timeData isa Vector{String}
        obsTimes = DateTime.(chomp.(timeData))
    else
        obsTimes = DateTime.(timeData)
    end
    startTime = obsTimes[1]
    obsTimesTrimmed = obsTimes[trimIndices[1]: trimIndices[2]]
    return startTime, obsTimesTrimmed
end

function plotMainEffects(mainEffects, timeData, inputNames;
                        palette=palette(:tab10, rev=true),
                        trimIndices=(1, 720),
                        tickStep=72,
                        tickFormat="dd-u",
                        title="Title",
                        lineWidth=0.0,
                        barWidth=2,
                        dpi=1000
                        )
    # we will first permute these so they show up with correct convention in
    # our bar plot

    # only columns are reversed after transpose, we are not changing the time series.
    mainEffectsReversed = reverse(mainEffects', dims=2)
    startTime, obsTimesTrimmed = processTimeInfo(timeData; trimIndices=trimIndices)
    obsTimeTicks = range(obsTimesTrimmed[1],
                        obsTimesTrimmed[end],
                        step=Hour(tickStep)
                        )
    xTicks = findall(in(obsTimeTicks), obsTimesTrimmed)
    labels = Dates.format.(obsTimeTicks, tickFormat)
    groupedbar(
            mainEffectsReversed,
            bar_position=:stack,
            bar_width=barWidth,
            legend=:outertopright,
            label=permutedims(reverse(inputNames)),
            xticks=(xTicks, labels),
            xminorticks=12,
            figsize=(1000, 600),
            color=[palette[i] for i in 1:size(mainEffectsReversed, 2)]',
            line=(lineWidth, :black),
            title=title,
            xlims=(1, length(obsTimesTrimmed)),
            ylims=(0, 1),
            dpi=dpi,
            framestyle=:box,
            )
    # plot!(xticks=(ticks, labels))
    plot!(xlabel = "Start Time: $(Dates.format(startTime, "dd-u-yy HH:MM:SS"))")
    plot!(ylabel = "Main Effects")

end

"""
Replicate Python function. Plot ONLY one half of the Matrix
and include text labels to annotate.
"""
function plotInteractionEffects(gsaIndices,
                            timeData,
                            inputNames;
                            trimIndices=(1, 720),
                            timeIdx=nothing,
                            dpi=1000,
                            color=:viridis
                            # summaryPlot="mean"
                            )
    if !isnothing(timeIdx)
        interactions = LowerTriangular(gsaIndices[:, :, timeIdx])
        _, obsTimes = processTimeInfo(timeData; trimIndices=trimIndices)
        plotTime = obsTimes[timeIdx]
        plotTitle = "Time:  $(Dates.format(plotTime, "dd-u-yy HH:MM:SS"))"
    else
        interactions = LowerTriangular(mean(gsaIndices, dims=3)[:, :, 1])
        plotTitle = "Mean Interaction Effects"
    end

    # interactions[maskElements] .= NaN

    # Tick labels
    xTickLabels = permutedims(reverse(inputNames))

    # Plot heatmap with annotations
    interactionPlot = Plots.heatmap(
                              inputNames,
                              inputNames,
                              interactions,
                              yflip=true,
                              c=color,
                              xrot=20,
                              fillalpha=0.7,
                              grid=false,
                              framestyle=:box,
                              dpi=dpi
                            )

    # #
    # return interactionPlot
    return interactionPlot
end

# function to perform bootstrapping

# function drawInputSamples(inputsPath, retainedRunIdx;
#                         useBg=true,
#                         NSamples = 100
#                         )

#     return inputSamples, trainIdx
# end

"""
Function to perform bootstrapping for a given number of replicates and sample size.
The bootstrapping process is given in:
Storlie, Curtis B., et al. Implementation and Evaluation of Nonparametric Regression Procedures for Sensitivity Analysis of Computationally
Demanding Models.
Reliability Engineering & System Safety, vol. 94, no. 11, Nov. 2009, pp. 1735–63. DOI.org (Crossref), https://doi.org/10.1016/j.ress.2009.05.007.
"""
function bootstrapGSA(X, Y;
                    regularize=false,
                    pceDegree=2,
                    NReplications = 1000,
                    replace=true
                    # NSamples = 100
                    )

    A = buildCoefficientMatrix(X; pceDegree=pceDegree)
    # bootstrap sampling
    nSamplesEnd = 100
    nSamplesStart = 20
    samplingRange = range(nSamplesStart, nSamplesEnd; step=10)

    nEffects = size(X, 2)
    nTimePts = size(Y, 2)

    mainEffectsBootstrap = zeros(nEffects, nTimePts, NReplications, length(samplingRange));
    multiIndex = buildMultivariateBasis(nEffects, pceDegree).ind

    for replication in 1:NReplications
        for (iSample, n) in enumerate(samplingRange)
            # println("New sample size: ", n)
            sampleIndex = sample(1:size(X, 1), n, replace=replace)
            ASamples = A[sampleIndex, :]
            YSamples = Y[sampleIndex, :]
            # A = buildCoefficientMatrix(XSamples; pceDegree=pceDegree)
            β = solvePCE(ASamples, YSamples; regularize=regularize)

            # can save some time here by extending method for supplied multiindex, and calculating multiindex outside of both our for loops.
            mainEffectsBootstrap[:, :, replication, iSample] = gsaMain(β; nEffects=nEffects, pceDegree=pceDegree, multiIndex=multiIndex)
        end
    end

    return mainEffectsBootstrap
    # meanEffects = mean(mainEffectsBootstrap; dims=3)[:, :, 1, :]; # take mean across replications and squeeze out dimension
    # stdEffects  = std(mainEffectsBootstrap; dims=3)[:, :, 1, :];  # take std across replications and squeeze out dimension
    # return meanEffects, stdEffects
end

"""
Function to obtain summary of bootstrap procedure results.
"""
function getBootstrapSummary(bootstrapEffects)
    # note specifying the dimension needs knowledge of which index encodes replications!. Not robust!!
    meanEffects = mean(mainEffectsBootstrap; dims=3)[:, :, 1, :]; # take mean across replications and squeeze out singleton dimension
    stdEffects = std(mainEffectsBootstrap; dims=3)[:, :, 1, :];
    return meanEffects, stdEffects
end

## SECTION TO TEST AND PLOT PREDICTIONS OF SURROGATE

"""
Take in input matrix X and output matrix Y and split it in suitable fashion into train and test sets.
The `ratio` parameter specifies sizes i.e. test size / train size ~ 0.2
"""
function train_test_split(X, Y; ratio=0.2, seed=20220405)
    # Random.seed!(seed)
    shuffledIdx = shuffle(1:size(X, 1))
    # perform splitting
    testLength = floor(Int, ratio * length(runs_to_keep))
    testIdx = shuffledIdx[1:testLength]
    trainIdx = shuffledIdx[(testLength + 1):end]

    XTrain = X[trainIdx, :]
    XTest  = X[testIdx,  :]

    YTrain = Y[trainIdx, :]
    YTest  = Y[testIdx, :]

    return XTrain, YTrain, XTest, YTest, trainIdx, testIdx
end

"""
Function to do L2 fit and return predictions as well as original
simulation runs.
"""
function regularized_fit(X, Y, lambda;
                                ratio=0.2,
                                seed=20220405,
                                pceDegree=2
                                )

    XTrain, YTrain, XTest, YTest, trainIdx, testIdx = train_test_split(X, Y;
                                                            ratio=ratio,
                                                            seed=seed
                                                            )
    ATrain  = buildCoefficientMatrix(XTrain; pceDegree=pceDegree)
    ATest   = buildCoefficientMatrix(XTest; pceDegree=pceDegree)
    βTrain  = solveRegPCE(ATrain, YTrain; λ = lambda)
    YPred   = ATest * βTrain

    return XTest, YTest, YPred, testIdx
end

"""
Extend regularized fit to use supplied XTrain, XTest, etc
"""
function regularized_fit(XTrain, YTrain, XTest, lambda;
                        pceDegree=2)
    ATrain  = buildCoefficientMatrix(XTrain; pceDegree=pceDegree)
    ATest   = buildCoefficientMatrix(XTest; pceDegree=pceDegree)
    βTrain  = solveRegPCE(ATrain, YTrain; λ = lambda)
    YPred   = ATest * βTrain

    return YPred
end

"""
Function to make plots for all test runs across different lambda values.
"""
function regularized_plot(YTest, YPred, testIdx, plotIdx, lambda)
    pReg = plot(YTest[plotIdx, :], line=(:red, 2), label="Truth")
    plot!(YPred[plotIdx, :], line=(:blue, 2), label="Prediction")
    plot!(ylims=(200, 900))
    rmse = sqrt(mean((YTest[plotIdx, :] - YPred[plotIdx, :]).^2))
    plot!(title = "RMSE:$(round(rmse, digits=2)) " * " λ:$(lambda)  Idx:$(testIdx[plotIdx])")
    plot!(titlefontsize=6)
    plot!(legendfontsize=4, fg_legend=false)
    return pReg, rmse
end

"""
Function to plot confidence intervals based on constructed PCE!
"""


text:
####
#### prognostic_to_primitive! and primitive_to_prognostic!
####

####
#### Wrappers (entry point)
####

"""
    prognostic_to_primitive!(atmos::AtmosModel, prim::Vars, prog::Vars, aux::Vars)

Convert prognostic variables `prog` to primitive
variables `prim` for the atmos model `atmos`.

!!! note
    The only field in `aux` required for this
    method is the geo-potential.
"""
function prognostic_to_primitive!(
    atmos::AtmosModel,
    prim::Vars,
    prog::Vars,
    aux,
)
    energy_model(atmos) isa TotalEnergyModel ||
        error("TotalEnergyModel only supported")
    prognostic_to_primitive!(atmos, moisture_model(atmos), prim, prog, aux)
    prognostic_to_primitive!(
        turbconv_model(atmos),
        atmos,
        moisture_model(atmos),
        prim,
        prog,
    )
end

"""
    primitive_to_prognostic!(atmos::AtmosModel, prog::Vars, prim::Vars, aux::Vars)

Convert primitive variables `prim` to prognostic
variables `prog` for the atmos model `atmos`.

!!! note
    The only field in `aux` required for this
    method is the geo-potential.
"""
function primitive_to_prognostic!(
    atmos::AtmosModel,
    prog::Vars,
    prim::Vars,
    aux,
)
    primitive_to_prognostic!(atmos, moisture_model(atmos), prog, prim, aux)
    primitive_to_prognostic!(
        turbconv_model(atmos),
        atmos,
        moisture_model(atmos),
        prog,
        prim,
    )
end

####
#### prognostic to primitive
####

function prognostic_to_primitive!(
    atmos,
    moist::DryModel,
    prim::Vars,
    prog::Vars,
    aux::Vars,
)
    ts = new_thermo_state(atmos, prog, aux)
    prim.ρ = air_density(ts) # Needed for recovery of energy, not prog.ρ in anelastic1d
    prim.u = prog.ρu ./ density(atmos, prog, aux)
    prim.p = pressure(atmos, ts, aux)
end

function prognostic_to_primitive!(
    atmos,
    moist::EquilMoist,
    prim::Vars,
    prog::Vars,
    aux::Vars,
)
    ts = new_thermo_state(atmos, prog, aux)
    prim.ρ = air_density(ts) # Needed for recovery of energy, not prog.ρ in anelastic1d
    prim.u = prog.ρu ./ density(atmos, prog, aux)
    prim.p = pressure(atmos, ts, aux)
    prim.moisture.q_tot = PhasePartition(ts).tot
end

function prognostic_to_primitive!(
    atmos,
    moist::NonEquilMoist,
    prim::Vars,
    prog::Vars,
    aux::Vars,
)
    ts = new_thermo_state(atmos, prog, aux)
    prim.ρ = air_density(ts) # Needed for recovery of energy, not prog.ρ in anelastic1d
    prim.u = prog.ρu ./ density(atmos, prog, aux)
    prim.p = pressure(atmos, ts, aux)
    prim.moisture.q_tot = PhasePartition(ts).tot
    prim.moisture.q_liq = PhasePartition(ts).liq
    prim.moisture.q_ice = PhasePartition(ts).ice
end

####
#### primitive to prognostic
####

function primitive_to_prognostic!(
    atmos,
    moist::DryModel,
    prog::Vars,
    prim::Vars,
    aux::Vars,
)
    energy_model(atmos) isa TotalEnergyModel ||
        error("TotalEnergyModel only supported")
    ts = PhaseDry_ρp(parameter_set(atmos), prim.ρ, prim.p)
    e_kin = prim.u' * prim.u / 2
    ρ = density(atmos, prim, aux)
    e_pot = gravitational_potential(atmos.orientation, aux)

    prog.ρ = ρ
    prog.ρu = ρ .* prim.u
    prog.energy.ρe = ρ * total_energy(e_kin, e_pot, ts)
end

function primitive_to_prognostic!(
    atmos,
    moist::EquilMoist,
    prog::Vars,
    prim::Vars,
    aux::Vars,
)
    energy_model(atmos) isa TotalEnergyModel ||
        error("TotalEnergyModel only supported")
    ts = PhaseEquil_ρpq(
        parameter_set(atmos),
        prim.ρ,
        prim.p,
        prim.moisture.q_tot,
        true,
    )
    e_kin = prim.u' * prim.u / 2
    ρ = density(atmos, prim, aux)
    e_pot = gravitational_potential(atmos.orientation, aux)

    prog.ρ = ρ
    prog.ρu = ρ .* prim.u
    prog.energy.ρe = ρ * total_energy(e_kin, e_pot, ts)
    prog.moisture.ρq_tot = ρ * PhasePartition(ts).tot
end

function primitive_to_prognostic!(
    atmos,
    moist::NonEquilMoist,
    prog::Vars,
    prim::Vars,
    aux::Vars,
)
    energy_model(atmos) isa TotalEnergyModel ||
        error("TotalEnergyModel only supported")
    q_pt = PhasePartition(
        prim.moisture.q_tot,
        prim.moisture.q_liq,
        prim.moisture.q_ice,
    )
    ts = PhaseNonEquil_ρpq(parameter_set(atmos), prim.ρ, prim.p, q_pt)
    e_kin = prim.u' * prim.u / 2
    ρ = density(atmos, prim, aux)
    e_pot = gravitational_potential(atmos.orientation, aux)

    prog.ρ = ρ
    prog.ρu = ρ .* prim.u
    prog.energy.ρe = ρ * total_energy(e_kin, e_pot, ts)
    prog.moisture.ρq_tot = ρ * PhasePartition(ts).tot
    prog.moisture.ρq_liq = ρ * PhasePartition(ts).liq
    prog.moisture.ρq_ice = ρ * PhasePartition(ts).ice
end

function construct_face_auxiliary_state!(
    bl::AtmosModel,
    aux_face::AbstractArray,
    aux_cell::AbstractArray,
    Δz::FT,
) where {FT <: Real}
    param_set = parameter_set(bl)
    _grav = FT(grav(param_set))
    var_aux = Vars{vars_state(bl, Auxiliary(), FT)}
    aux_face .= aux_cell

    if !(bl.orientation isa NoOrientation)
        var_aux(aux_face).orientation.Φ =
            var_aux(aux_cell).orientation.Φ + _grav * Δz / 2
    end
end


text:
# Use
#
# DOCUMENTER_DEBUG=true GKSwstype=100 julia --color=yes make.jl local [linkcheck] [fixdoctests]
#
# for local builds.

using Pkg
using Documenter
using Literate
using Plots
using SolidStateDetectors
using SolidStateDetectors.ConstructiveSolidGeometry

function fix_literate_output(content)
    content = replace(content, "EditURL = \"@__REPO_ROOT_URL__/\"" => "")
    return content
end

gen_content_dir = joinpath(@__DIR__, "src", "tutorials")
for tut_lit_fn in filter(fn -> endswith(fn, "_lit.jl"), readdir(gen_content_dir))
    lit_src_fn = joinpath(gen_content_dir, tut_lit_fn)
    tut_basename = tut_lit_fn[1:end-7] # remove "_lit.jl"
    Literate.markdown(lit_src_fn, gen_content_dir, name = tut_basename, documenter = true, credit = true, postprocess = fix_literate_output)
    Literate.notebook(lit_src_fn, gen_content_dir, execute = false, name = tut_basename, documenter = true, credit = true)
    Literate.script(lit_src_fn, gen_content_dir, keep_comments = false, name = tut_basename, documenter = true, credit = false)
end

makedocs(
    sitename = "SolidStateDetectors.jl",
    modules = [SolidStateDetectors, SolidStateDetectors.ConstructiveSolidGeometry],
    pages = [
        "Home" => "index.md",
        "Manual" => Any[
            "Installation" => "man/installation.md",
            "Configuration Files" => "man/config_files.md",
            "Constructive Solid Geometry (CSG)" => "man/csg.md",
            "Grids" => "man/Grids.md",
            "Electric Potential" => "man/electric_potential.md",
            "Electric Field" => "man/electric_field.md",
            "Charge Drift" => "man/charge_drift.md",
            "Weighting Potentials" => "man/weighting_potentials.md",
            "Capacitances" => "man/capacitances.md",
            "IO" => "man/IO.md",
            "Plotting" => "man/plotting.md",
        ],
        "Tutorials" => [
            "tutorials/complete_simulation_chain_IVC.md",
            "tutorials/custom_impurity_density_pn_junction.md",
        ],
        "API" => "api.md",
        "LICENSE" => "LICENSE.md",
    ],
    doctest = ("fixdoctests" in ARGS) ? :fix : true,
    format = Documenter.HTML(canonical = "https://JuliaPhysics.github.io/SolidStateDetectors.jl/stable/", prettyurls = !("local" in ARGS)),
    linkcheck = ("linkcheck" in ARGS),
    strict = !("local" in ARGS),
)

deploydocs(
    repo = "github.com/JuliaPhysics/SolidStateDetectors.jl.git",
    devbranch = "main",
    devurl = "main",
    forcepush = true,
    push_preview = true,
)


text:
#=
Copyright (c) 2018-2022 Chris Coey, Lea Kapelevich, and contributors

This Julia package Hypatia.jl is released under the MIT license; see LICENSE
file in the root directory or at https://github.com/chriscoey/Hypatia.jl

TODO improve description of the model

see e.g.:
- "The Power of Convex Relaxation" by Emmanuel J. Candes and Terence Tao
- http://www.mit.edu/~parrilo/pubs/talkfiles/ISMP2009.pdf
- https://www.cvxpy.org/examples/dgp/pf_matrix_completion.html (includes geomean)

extended formulations to (u, W) in EpiNormSpectral(use_dual = true) uses:
min 1/2(tr(W1) + tr(W2))
[W1 X; X' W2] ⪰ 0
=#

struct MatrixCompletionNative{T <: Real} <: ExampleInstanceNative{T}
    m::Int
    n::Int
    geomean_constr::Bool # add a constraint on the geomean of unknown values
    nuclearnorm_obj::Bool # use nuclear norm in the objective, else spectral norm
    use_hypogeomean::Bool # use hypogeomean cone, else power cone formulation
    use_epinormspectral::Bool # use epinormspectral cone (or dual), else PSD cone
end

function build(inst::MatrixCompletionNative{T}) where {T <: Real}
    (m, n) = (inst.m, inst.n)
    @assert m <= n
    mn = m * n
    rt2 = sqrt(T(2))

    num_known = round(Int, mn * 0.8)
    known_rows = rand(1:m, num_known)
    known_cols = rand(1:n, num_known)
    known_vals = 2 * rand(T, num_known) .- 1

    is_known = fill(false, mn)
    # h for the rows that X (the matrix, not epigraph variable) participates in
    h_norm_x = zeros(T, m * n)
    for (k, (i, j)) in enumerate(zip(known_rows, known_cols))
        known_idx = (j - 1) * m + i
        # if not using the epinorminf cone, indices relate to X'
        h_norm_x[known_idx] = known_vals[k]
        is_known[known_idx] = true
    end

    num_known = sum(is_known) # if randomly generated, some indices may repeat
    num_unknown = m * n - num_known

    # epinormspectral cone or its dual- get vec(X) in G and h
    if inst.use_epinormspectral
        c = vcat(one(T), zeros(T, num_unknown))
        G_norm = zeros(T, mn, num_unknown)
        total_idx = 1
        unknown_idx = 1
        for j in 1:n, i in 1:m
            if !is_known[total_idx]
                G_norm[total_idx, unknown_idx] = -1
                unknown_idx += 1
            end
            total_idx += 1
        end

        # add first row and column for epigraph variable
        G_norm = [
            -one(T)    zeros(T, 1, num_unknown);
            zeros(T, mn)    G_norm;
            ]
        h_norm_x = vcat(zero(T), h_norm_x)
        h_norm = h_norm_x

        cones = Cones.Cone{T}[Cones.EpiNormSpectral{T, T}(m, n,
            use_dual = inst.nuclearnorm_obj)]
    else
        # build an extended formulation for the norm used in the objective
        if inst.nuclearnorm_obj
            # extended formulation for nuclear norm
            # X, W_1, W_2
            num_W1_vars = Cones.svec_length(m)
            num_W2_vars = Cones.svec_length(n)
            num_vars = num_W1_vars + num_W2_vars + num_unknown

            A = zeros(T, 0, num_vars)
            # unknown entries in X' unlike above
            num_rows = num_W1_vars + num_W2_vars + mn
            G_norm = zeros(T, num_rows, num_unknown + num_W1_vars + num_W2_vars)
            h_norm = zeros(T, num_rows)
            # first block W_1
            G_norm[1:num_W1_vars, (num_unknown + 1):(num_unknown + num_W1_vars)] =
                -Matrix{T}(I, num_W1_vars, num_W1_vars)

            offset = num_W1_vars
            # index to count rows in the bottom half of the large to-be-PSD matrix
            idx = 0
            # index only in X
            X_var_idx = 0
            W2_var_idx = 0
            # index of unknown vars (the x variables in the standard from), can
            # increment as we are moving row wise in X' (i.e. columnwise in X)
            unknown_idx = 0
            # fill bottom `n` rows
            for i in 1:n
                # X'
                for j in 1:m
                    idx += 1
                    X_var_idx += 1
                    if !is_known[X_var_idx]
                        unknown_idx += 1
                        G_norm[offset + idx, unknown_idx] = -1
                    else
                        h_norm[offset + idx] = h_norm_x[X_var_idx]
                    end
                end
                # second block W_2
                for j in 1:i
                    idx += 1
                    W2_var_idx += 1
                    G_norm[offset + idx,
                        num_unknown + num_W1_vars + W2_var_idx] = -1
                end
            end
            Cones.scale_svec!(G_norm, rt2)
            Cones.scale_svec!(h_norm, rt2)
            cones = Cones.Cone{T}[Cones.PosSemidefTri{T, T}(num_rows)]
            c_W1 = Cones.smat_to_svec!(zeros(T, num_W1_vars),
                Diagonal(one(T) * I, m), rt2)
            c_W2 = Cones.smat_to_svec!(zeros(T, num_W2_vars),
                Diagonal(one(T) * I, n), rt2)
            c = vcat(zeros(T, num_unknown), c_W1, c_W2) / 2
        else
            # extended formulation for spectral norm
            num_rows = Cones.svec_length(m) + m * n + Cones.svec_length(n)
            G_norm = zeros(T, num_rows, num_unknown + 1)
            h_norm = zeros(T, num_rows)
            # first block epigraph variable * I
            for i in 1:m
                G_norm[sum(1:i), 1] = -1
            end
            offset = Cones.svec_length(m)
            # index to count rows in the bottom half of the large to-be-PSD matrix
            idx = 1
            # index only in X
            var_idx = 1
            # index of unknown vars (the x variables in the standard from), can
            # increment it because we are moving row wise in X'
            unknown_idx = 1
            # fill bottom `n` rows
            for i in 1:n
                # X'
                for j in 1:m
                    if !is_known[var_idx]
                        G_norm[offset + idx, 1 + unknown_idx] = -1
                        unknown_idx += 1
                    else
                        h_norm[offset + idx] = h_norm_x[var_idx]
                    end
                    idx += 1
                    var_idx += 1
                end
                # second block epigraph variable * I
                # skip `i` rows which will be filled with zeros
                idx += i
                G_norm[offset + idx - 1, 1] = -1
            end
            Cones.scale_svec!(G_norm, rt2)
            Cones.scale_svec!(h_norm, rt2)
            cones = Cones.Cone{T}[Cones.PosSemidefTri{T, T}(num_rows)]
            c = vcat(one(T), zeros(T, num_unknown))
        end
    end

    if inst.geomean_constr
        if inst.use_hypogeomean
            # hypogeomean for values to be filled
            G_geo = vcat(zeros(T, 1, num_unknown), Matrix{T}(-I,
                num_unknown, num_unknown))
            h = vcat(h_norm, one(T), zeros(T, num_unknown))

            # if using extended with spectral objective G_geo needs to be
            # prepadded with an epigraph variable
            if inst.nuclearnorm_obj
                if inst.use_epinormspectral
                    prepad = zeros(T, num_unknown + 1, 1)
                    postpad = zeros(T, num_unknown + 1, 0)
                else
                    prepad = zeros(T, num_unknown + 1, 0)
                    postpad = zeros(T, num_unknown + 1,
                        size(G_norm, 2) - num_unknown)
                end
            else
                prepad = zeros(T, num_unknown + 1, 1)
                postpad = zeros(T, num_unknown + 1, 0)
            end
            G = [
                G_norm;
                prepad  G_geo  postpad
                ]
            push!(cones, Cones.HypoGeoMean{T}(1 + num_unknown))
        else
            # number of 3-dimensional power cones needed is num_unknown - 1,
            # number of new variables is num_unknown - 2
            @assert num_unknown > 3 # power cone formulation minimum

            # first num_unknown columns overlap with G_norm, column for the
            # epigraph variable of the spectral cone added later
            len_power = 3 * (num_unknown - 1)
            G_geo_unknown = zeros(T, len_power, num_unknown)
            G_geo_newvars = zeros(T, len_power, num_unknown - 2)
            # first cone is a special case since two of the original variables
            # participate in it
            G_geo_unknown[1, 1] = -1
            G_geo_unknown[2, 2] = -1
            G_geo_newvars[3, 1] = -1
            push!(cones, Cones.GeneralizedPower{T}(fill(inv(T(2)), 2), 1))
            offset = 4
            # loop over new vars
            for i in 1:(num_unknown - 3)
                G_geo_newvars[offset + 2, i + 1] = -1
                G_geo_newvars[offset + 1, i] = -1
                G_geo_unknown[offset, i + 2] = -1
                push!(cones, Cones.GeneralizedPower{T}([inv(T(i + 2)),
                    T(i + 1) / T(i + 2)], 1))
                offset += 3
            end

            # last row also special because hypograph variable is fixed
            G_geo_unknown[offset, num_unknown] = -1
            G_geo_newvars[offset + 1, num_unknown - 2] = -1
            push!(cones, Cones.GeneralizedPower{T}([inv(T(num_unknown)),
                T(num_unknown - 1) / T(num_unknown)], 1))
            h = vcat(h_norm, zeros(T, 3 * (num_unknown - 2)), T[0, 0, 1])

            # if using extended with spectral objective G_geo needs to be
            # prepadded with an epigraph variable
            if inst.nuclearnorm_obj
                if inst.use_epinormspectral
                    prepad = zeros(T, len_power, 1)
                    postpad = zeros(T, len_power, 0)
                else
                    prepad = zeros(T, len_power, 0)
                    postpad = zeros(T, len_power, size(G_norm, 2) - num_unknown)
                end
            else
                prepad = zeros(T, len_power, 1)
                postpad = zeros(T, len_power, 0)
            end
            G = [
                G_norm  zeros(T, size(G_norm, 1), num_unknown - 2);
                prepad  G_geo_unknown  postpad  G_geo_newvars
                ]

            c = vcat(c, zeros(T, num_unknown - 2))
        end
    else
        G = G_norm
        h = h_norm
    end

    A = zeros(T, 0, size(G, 2))
    b = T[]

    model = Models.Model{T}(c, A, b, G, h, cones)
    return model
end


text:
module MoreMatrices

const ɛ = ([0 0 0; 0 0 1; 0 -1 0], [0 0 -1; 0 0 0; 1 0 0], [0 1 0; -1 0 0; 0 0 0])

export
    ɛ,

    # operations.jl

    kron_shuffle!,
    kron_shuffle,
    kron_unshuffle!,
    kron_unshuffle,
    kron_factor,
    kron_sum,
    join_blocks,
    split_blocks,
    rotation!,
    rotation,
    euler_zyz!,
    euler_zyz,
    nearest_orthogonal!,
    nearest_orthogonal,
    nearest_special_orthogonal!,
    nearest_special_orthogonal,
    cayley,
    rotation2skew,
    vec2skew!,
    vec2skew,
    skew2vec!,
    skew2vec,

    # random_matrices.jl

    random_orthogonal!,
    random_orthogonal,
    random_rotation,
    random_permutation

include("operations.jl")
include("random_matrices.jl")

end


text:
export @exfiltrate, setfield!!

# original version by Stefan Karpinski in slack (2.8.2020)
"""
    setfieled!!(module, var, value)

Set a globla variable in an arbitrary module
"""
setfield!!(m::Module, var::Symbol, val::Any) = m.eval(:($var = $val))

"""
    @exfiltrate

Usage example:
julia> function f(x, y)
           a = 1
           @exfiltrate
           b = "foo"
           return x, b
       end
f (generic function with 1 method)
julia> f(3, 4)
(3, "foo")
julia> x
3
julia> y
4
julia> a
1
julia> b
ERROR: UndefVarError: b not defined
"""
macro exfiltrate()
    quote
        for (var, val) in Base.@locals
            setfield!!(Main, var, val)
        end
    end
end


text:
function calculate_kinetics(X, data_dictionary)

	# load all species dictionary
	all_species_dict = data_dictionary["all_species_dict"]  # String-->Int

	# generate kinetics for signaling, assume saturation kinetics
	# load Monod affinity constants
	MonodK = data_dictionary["MonodAffinityConstantDict"]  # String-->Float
	# load reaction kinetic constants
	kcat = data_dictionary["kcat_signaling"]  # in order of rnx

	# write reaction rate equations
	rnx_rate_vector = zeros(7)
	rnx_rate_vector[1] = kcat[1]*X[12]*X[2]/(MonodK["MonodK~rnx1~m_A_e"]+X[2])
	rnx_rate_vector[2] = kcat[2]*X[13]*X[6]/(MonodK["MonodK~rnx2~m_B_c"]+X[6])
	rnx_rate_vector[3] = kcat[3]*X[14]*X[7]/(MonodK["MonodK~rnx3~m_C_c"]+X[7])
	rnx_rate_vector[4] = kcat[4]*X[8]*X[5]/(MonodK["MonodK~rnx4~m_A_c"]+X[5])
	rnx_rate_vector[5] = kcat[5]*X[11]*X[6]/(MonodK["MonodK~rnx5~m_B_c"]+X[6])
	rnx_rate_vector[6] = kcat[6]*X[9]*X[5]/(MonodK["MonodK~rnx6~m_A_c"]+X[5])
	rnx_rate_vector[7] = kcat[7]*X[10]*X[7]/(MonodK["MonodK~rnx7~m_C_c"]+X[7])

	# generate kinetics & control terms for TXTL
	# load data from data dictionary
	W_value_dict = data_dictionary["W_value_dict"]  #String-->Float
	coop = data_dictionary["cooperativity"]
	disassociation_constant_dict = data_dictionary["transferFunctionDisassociationConstantDict"]
	background_control_term = data_dictionary["backgroundGeneExpressionControlTermDict"]
	kcatTX = data_dictionary["kcatTranscription"]
	RNAPconc = data_dictionary["RNAPConcentration"]
	geneConc = data_dictionary["avgGeneConcentration"]
	saturationTX = data_dictionary["transcriptionSaturationConstant"]
	kcatTL = data_dictionary["kcatTranslation"]
	RIBOconc = data_dictionary["RIBOConcentration"]
	saturationTL = data_dictionary["translationSaturationConstant"]

	# write control terms, transcription rate equations, and translation rate equations
	TX_control_term = Dict{String, Float64}()  # control term
	TX_rate_vector = Dict{String, Float64}()  # transcription rate
	TL_rate_vector = Dict{String, Float64}()  # translation rate
	# mRNA_E4_c and p_E4_c
	mRNA_E4_c_m_B_c = W_value_dict["W~m_B_c~mRNA_E4_c"] * X[6]^coop/((disassociation_constant_dict["KD~m_B_c"])^coop+X[6]^coop)
	up_action_on_mRNA_E4_c = mRNA_E4_c_m_B_c
	TX_control_term["mRNA_E4_c"] = (background_control_term["mRNA_E4_c"] + up_action_on_mRNA_E4_c)/(1 + background_control_term["mRNA_E4_c"] + up_action_on_mRNA_E4_c)
	TX_rate_vector["mRNA_E4_c"] = TX_control_term["mRNA_E4_c"]*kcatTX*RNAPconc*geneConc/(saturationTX+geneConc)
	TL_rate_vector["p_E4_c"] = kcatTL*RIBOconc*X[18]/(saturationTL+X[18])
	# mRNA_E2_c and p_E2_c
	mRNA_E2_c_m_A_c = W_value_dict["W~m_A_c~mRNA_E2_c"] * X[5]^coop/((disassociation_constant_dict["KD~m_A_c"])^coop+X[5]^coop)
	up_action_on_mRNA_E2_c = mRNA_E2_c_m_A_c
	TX_control_term["mRNA_E2_c"] = (background_control_term["mRNA_E2_c"] + up_action_on_mRNA_E2_c)/(1 + background_control_term["mRNA_E2_c"] + up_action_on_mRNA_E2_c)
	TX_rate_vector["mRNA_E2_c"] = TX_control_term["mRNA_E2_c"]*kcatTX*RNAPconc*geneConc/(saturationTX+geneConc)
	TL_rate_vector["p_E2_c"] = kcatTL*RIBOconc*X[16]/(saturationTL+X[16])
	# mRNA_E3_c and p_E3_c
	mRNA_E3_c_m_C_c = W_value_dict["W~m_C_c~mRNA_E3_c"] * X[7]^coop/((disassociation_constant_dict["KD~m_C_c"])^coop+X[7]^coop)
	up_action_on_mRNA_E3_c = mRNA_E3_c_m_C_c
	TX_control_term["mRNA_E3_c"] = (background_control_term["mRNA_E3_c"] + up_action_on_mRNA_E3_c)/(1 + background_control_term["mRNA_E3_c"] + up_action_on_mRNA_E3_c)
	TX_rate_vector["mRNA_E3_c"] = TX_control_term["mRNA_E3_c"]*kcatTX*RNAPconc*geneConc/(saturationTX+geneConc)
	TL_rate_vector["p_E3_c"] = kcatTL*RIBOconc*X[17]/(saturationTL+X[17])
	# mRNA_TA_c and p_TA_c
	mRNA_TA_c_m_A_e = W_value_dict["W~m_A_e~mRNA_TA_c"] * X[2]^coop/((disassociation_constant_dict["KD~m_A_e"])^coop+X[2]^coop)
	up_action_on_mRNA_TA_c = mRNA_TA_c_m_A_e
	TX_control_term["mRNA_TA_c"] = (background_control_term["mRNA_TA_c"] + up_action_on_mRNA_TA_c)/(1 + background_control_term["mRNA_TA_c"] + up_action_on_mRNA_TA_c)
	TX_rate_vector["mRNA_TA_c"] = TX_control_term["mRNA_TA_c"]*kcatTX*RNAPconc*geneConc/(saturationTX+geneConc)
	TL_rate_vector["p_TA_c"] = kcatTL*RIBOconc*X[19]/(saturationTL+X[19])
	# mRNA_TB_c and p_TB_c
	mRNA_TB_c_m_B_c = W_value_dict["W~m_B_c~mRNA_TB_c"] * X[6]^coop/((disassociation_constant_dict["KD~m_B_c"])^coop+X[6]^coop)
	up_action_on_mRNA_TB_c = mRNA_TB_c_m_B_c
	TX_control_term["mRNA_TB_c"] = (background_control_term["mRNA_TB_c"] + up_action_on_mRNA_TB_c)/(1 + background_control_term["mRNA_TB_c"] + up_action_on_mRNA_TB_c)
	TX_rate_vector["mRNA_TB_c"] = TX_control_term["mRNA_TB_c"]*kcatTX*RNAPconc*geneConc/(saturationTX+geneConc)
	TL_rate_vector["p_TB_c"] = kcatTL*RIBOconc*X[20]/(saturationTL+X[20])
	# mRNA_E1_c and p_E1_c
	mRNA_E1_c_m_A_c = W_value_dict["W~m_A_c~mRNA_E1_c"] * X[5]^coop/((disassociation_constant_dict["KD~m_A_c"])^coop+X[5]^coop)
	up_action_on_mRNA_E1_c = mRNA_E1_c_m_A_c
	TX_control_term["mRNA_E1_c"] = (background_control_term["mRNA_E1_c"] + up_action_on_mRNA_E1_c)/(1 + background_control_term["mRNA_E1_c"] + up_action_on_mRNA_E1_c)
	TX_rate_vector["mRNA_E1_c"] = TX_control_term["mRNA_E1_c"]*kcatTX*RNAPconc*geneConc/(saturationTX+geneConc)
	TL_rate_vector["p_E1_c"] = kcatTL*RIBOconc*X[15]/(saturationTL+X[15])
	# mRNA_TC_c and p_TC_c
	mRNA_TC_c_m_C_c = W_value_dict["W~m_C_c~mRNA_TC_c"] * X[7]^coop/((disassociation_constant_dict["KD~m_C_c"])^coop+X[7]^coop)
	up_action_on_mRNA_TC_c = mRNA_TC_c_m_C_c
	TX_control_term["mRNA_TC_c"] = (background_control_term["mRNA_TC_c"] + up_action_on_mRNA_TC_c)/(1 + background_control_term["mRNA_TC_c"] + up_action_on_mRNA_TC_c)
	TX_rate_vector["mRNA_TC_c"] = TX_control_term["mRNA_TC_c"]*kcatTX*RNAPconc*geneConc/(saturationTX+geneConc)
	TL_rate_vector["p_TC_c"] = kcatTL*RIBOconc*X[21]/(saturationTL+X[21])

	return (rnx_rate_vector, TX_rate_vector, TL_rate_vector)
end


text:
using DataFrames

#Create Dataframe
myDataFrame = DataFrame(Name = ["Paul", "Bryan", "Julia", "Helmut"],
                        Age = [22, 90, 38, 67])

#add a new column
myDataFrame[!,:name_new_column] .= 1

println(myDataFrame)

max_stars_repo_name: RainerNeu/juliaExamplesWithJuliaPro

max_issues_repo_name: RainerNeu/juliaExamplesWithJuliaPro

max_forks_repo_name: RainerNeu/juliaExamplesWithJuliaPro


text:
# Autogenerated wrapper script for Expat_jll for x86_64-linux-musl
export libexpat, xmlwf

JLLWrappers.@generate_wrapper_header("Expat")
JLLWrappers.@declare_library_product(libexpat, "libexpat.so.1")
JLLWrappers.@declare_executable_product(xmlwf)
function __init__()
    JLLWrappers.@generate_init_header()
    JLLWrappers.@init_library_product(
        libexpat,
        "lib/libexpat.so",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_executable_product(
        xmlwf,
        "bin/xmlwf",
    )

    JLLWrappers.@generate_init_footer()
end  # __init__()


text:
using CUDA
using PicoQuant
using BenchmarkTools

let
    r = 4
    c = 4
    n = 16
    circ = create_ghz_preparation_circuit(n)
    #circ = create_RQC(r, c, 8)

    function sim_circ_full(circ, gpu=false)
        backend = gpu ? InteractiveBackend{CuArray{ComplexF32}}() : InteractiveBackend{Array{ComplexF32}}()
       	tng = convert_qiskit_circ_to_network(circ, backend)
	add_input!(tng, "0"^n)
        if gpu
	    CUDA.@time output_node = full_wavefunction_contraction!(tng)
        else
	    @time output_node = full_wavefunction_contraction!(tng)
        end
    end

    function sim_circ_mps(circ, gpu=false)
        backend = gpu ? InteractiveBackend{CuArray{ComplexF32}}() : InteractiveBackend{Array{ComplexF32}}()
       	tng = convert_qiskit_circ_to_network(circ, decompose=true, transpile=true, backend)
	add_input!(tng, "0"^n)
        if gpu
	    CUDA.@profile output_node = contract_mps_tensor_network_circuit!(tng)
        else
	    @time output_node = contract_mps_tensor_network_circuit!(tng)
        end
    end

    #sim_circ_full(circ)
    #sim_circ_full(circ)

    #sim_circ_full(circ, true)
    #sim_circ_full(circ, true)

    sim_circ_mps(circ)
    sim_circ_mps(circ)

    sim_circ_mps(circ, true)
    sim_circ_mps(circ, true)

end


text:
struct NoValue end

"""
    liftr(f) -> f′

Try to lift `X -> Y` to `X -> Union{Some{Y},Nothing}`; return `f`
as-is if not defined.
"""
liftr(f::F) where {F} = f

function define_liftr_impl()
    exprs = Expr[]
    for n in names(Maybe; all = true)
        @debug "`define_liftr_impl()`" n isdefined(Base, n)
        isdefined(Base, n) || continue
        f0 = getfield(Base, n)
        f1 = getfield(Maybe, n)
        @debug "`define_liftr_impl()`" f0 f1 f0 isa Function
        f0 isa Function || continue
        ex = quote
            liftr(::$(typeof(f0))) = $f1
        end
        push!(exprs, ex)
    end
    return Expr(:block, exprs...)
end

define_liftr() = Base.eval(Maybe.Implementations, define_liftr_impl())

function finalize_implementations()
    define_liftr()
end

# See ThreadsX.jl
function define_docstrings()
    docstrings = Pair[:Maybe=>joinpath(dirname(@__DIR__), "README.md")]
    docsdir = joinpath(@__DIR__, "docs")
    for filename in readdir(docsdir)
        stem, ext = splitext(filename)
        ext == ".md" || continue
        name = Symbol(stem)
        if name in names(Maybe, all = true)
            push!(docstrings, name => joinpath(docsdir, filename))
        elseif name in names(Maybe.Extras, all = true)
            push!(docstrings, :(Extras.$name) => joinpath(docsdir, filename))
        end
    end
    for (name, path) in docstrings
        include_dependency(path)
        doc = read(path, String)
        doc = replace(doc, r"^```julia"m => "```jldoctest $name")
        doc = replace(doc, "<kbd>TAB</kbd>" => "_TAB_")
        @eval Maybe $Base.@doc $doc $name
    end
end

function finalize_package()
    @eval Maybe begin
        const T = Implementations.MaybeType
        const $(Symbol("@?")) = Implementations.$(Symbol("@?"))
        export $(Symbol("@?"))
    end
    define_docstrings()
end


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

using Markdown
using InteractiveUtils

# ╔═╡ 0a57947c-3a56-11ec-1341-534834850daf
md"""
# Chapter 1
## Section 3
"""

# ╔═╡ 8fcfe29e-049c-4df3-9d4b-325d3f4e171e
function newton(f, fp, x0, tol)
	x = x0
	while abs(f(x)) > tol
		x = x - f(x) / fp(x)
	end
	x
end

# ╔═╡ 74555744-e6a7-47bd-96b9-7e8909664f42
newton(x -> x^2 - 5., x -> 2. * x, 1., 10e-6)

# ╔═╡ 40d3098f-4b1b-4ca1-be89-5eb9a18ef3a8
function secant(f, x0, x1, tol)
	x = x0
	xx = x1
	while abs(f(xx)) > tol
		tmp = xx
		xx = xx - f(xx) * (xx - x) / (f(xx) - f(x))
		x = xx
	end
	xx
end

# ╔═╡ 7d023bfd-53fb-46a1-84dc-1c34519387dc
secant(x -> x^2 - 5., 2, 2.3, 10e-6)
# Well ... I'm definitly not got at finding roots

# ╔═╡ Cell order:
# ╟─0a57947c-3a56-11ec-1341-534834850daf
# ╠═8fcfe29e-049c-4df3-9d4b-325d3f4e171e
# ╠═74555744-e6a7-47bd-96b9-7e8909664f42
# ╠═40d3098f-4b1b-4ca1-be89-5eb9a18ef3a8
# ╠═7d023bfd-53fb-46a1-84dc-1c34519387dc

max_stars_repo_name: zazbone/ComputationalPhysics

max_issues_repo_name: zazbone/ComputationalPhysics

max_forks_repo_name: zazbone/ComputationalPhysics


text:
using Documenter, FastTransforms

makedocs(
			doctest = false,
			format = Documenter.HTML(),
			sitename = "FastTransforms.jl",
			authors = "Richard Mikael Slevinsky",
			pages = Any[
					"Home" => "index.md"
					]
			)

deploydocs(
    repo   = "github.com/JuliaApproximation/FastTransforms.jl.git",
    )


text:
export wcc, wcc_dt

"""

    wcc(ref, cur, t, window, fs, window_length, window_step)

Measure dv/v using the windowed cross-correlation following Snieder et al., 2012

# Arguments
- `ref::AbstractArray`: Reference time series
- `cur::AbstractArray`: Current time series
- `t::AbstractArray`: Time axis common to ref and cur
- `window::AbstractArray`: Indices for a subset of t over which to measure phase lags
- `fs::Float64`: Sampling frequency
- `window_length::Float64`: Length of time window within which to find time shifts
- `window_step::Float64`: Time step to advance window

# Returns
- `dvv::Float64`: dv/v for current correlation for a range of frequencies
- `dvv_err::Float64`: Error for calculation of dv/v for a range of frequencies
- `int::Float64`: Intercept for regression calculation
- `int_err::Float64`: Error for intercept
- `dvv0::Float64`: dv/v for current correlation forced through the origin
- `dvv0_err::Float64`: Error for calculation of dvv0 for a range of frequencies
"""
function wcc(ref::AbstractArray, cur::AbstractArray, t::AbstractArray, window::AbstractArray, fs::Float64,
             window_length::Float64, window_step::Float64)
    # measure phase shifts
    time_axis, delta_t, cc_max = wcc_dt(ref, cur, t, window, fs, window_length, window_step)
    # perform linear regression of dt/t=-dv/v
    dvv, dvv_err, int, int_err, dvv0, dvv0_err = dvv_lstsq(time_axis, delta_t)

    return dvv, dvv_err, int, int_err, dvv0, dvv0_err
end

"""

    wcc(ref, cur, t, window, fs, window_length, window_step, freqbands; norm=true)

Measure dv/v over a range of frequencies using the windowed cross-correlation following Snieder et al., 2012

# Arguments
- `ref::AbstractArray`: Reference time series
- `cur::AbstractArray`: Current time series
- `t::AbstractArray`: Time axis common to ref and cur
- `window::AbstractArray`: Indices for a subset of t over which to measure phase lags
- `fs::Float64`: Sampling frequency
- `window_length::Float64`: Length of time window within which to find time shifts
- `window_step::Float64`: Time step to advance window
- `freqbands::AbstractArray`: Frequency bands over which to compute dv/v
- `norm::Bool`: Whether or not to normalize signals before dv/v

# Returns
- `freqbands::AbstractArray`: Array of frequencies where dv/v was measured
- `dvv::Float64`: dv/v for current correlation for a range of frequencies
- `dvv_err::Float64`: Error for calculation of dv/v for a range of frequencies
- `int::Float64`: Intercept for regression calculation
- `int_err::Float64`: Error for intercept
- `dvv0::Float64`: dv/v for current correlation forced through the origin
- `dvv0_err::Float64`: Error for calculation of dvv0 for a range of frequencies
"""
function wcc(ref::AbstractArray, cur::AbstractArray, t::AbstractArray, window::AbstractArray, fs::Float64,
             window_length::Float64, window_step::Float64, freqbands::AbstractArray; norm::Bool=true)
     # define sampling interval
     dt = 1/fs

     # calculate the CWT of the time series, using identical parameters for both calculations
     cwt1, sj, freqs, coi = cwt(ref, dt, minimum(freqbands), maximum(freqbands))
     cwt2, sj, freqs, coi = cwt(cur, dt, minimum(freqbands), maximum(freqbands))

     # if a frequency window is given (instead of a set of frequency bands), we assume
     # dv/v should be calculated for each frequency. We construct a 2D array of the
     # form [f1 f1; f2 f2; ...], which can be treated the same as a 2D array of frequency bands
     if ndims(freqbands)==1
         freqbands = hcat(freqs, freqs)
     end
     # number of frequency bands
     (nbands,_) = size(freqbands)

     # initialize dvv and err arrays
     dvv = zeros(nbands)
     dvv_err = similar(dvv)
     int = similar(dvv)
     int_err = similar(dvv)
     dvv0 = similar(dvv)
     dvv0_err = similar(dvv)

     # loop over frequency bands
     for iband=1:nbands
         (fmin, fmax) = freqbands[iband, :]

         # get current frequencies over which we apply icwt
         # frequency checks
         if fmax < fmin
             println("Error: please ensure columns 1 and 2 are right frequency limits in freqbands!")
         else
             freq_ind = findall(f->(f>=fmin && f<=fmax), freqs)
         end

         # perform icwt
         icwt1 = icwt(cwt1[:,freq_ind], sj[freq_ind], dt)
         icwt2 = icwt(cwt2[:,freq_ind], sj[freq_ind], dt)
         wcwt1 = real.(icwt1)
         wcwt2 = real.(icwt2)

         # normalize both signals, if appropriate
         if norm
             ncwt1 = ((wcwt1 .- mean(wcwt1)) ./ std(wcwt1))
             ncwt2 = ((wcwt2 .- mean(wcwt2)) ./ std(wcwt2))
         else
             ncwt1 = wcwt1
             ncwt2 = wcwt2
         end

         time_axis, delta_t, cc_max = wcc_dt(ncwt1, ncwt2, t, window, fs, window_length, window_step)
         dvv[iband], dvv_err[iband], int[iband], int_err[iband], dvv0[iband], dvv0_err[iband] = dvv_lstsq(time_axis, delta_t)
     end

     return freqbands, dvv, dvv_err, int, int_err, dvv0, dvv0_err
end

"""

    wcc_dt(ref, cur, fs, tmin, window_length, window_step)

Measure phase shifts using the windowed cross-correlation following Snieder et al., 2012

# Arguments
- `ref::AbstractArray`: Reference time series
- `cur::AbstractArray`: Current time series
- `t::AbstractArray`: Time axis common to ref and cur
- `window::AbstractArray`: Indices for a subset of t over which to measure phase lags
- `fs::Float64`: Sampling frequency
- `window_length::Float64`: Length of time window within which to find time shifts
- `window_step::Float64`: Time step to advance window

# Returns
- `time_axis::AbstractArray`: Center times of windows
- `dt::AbstractArray`: Time shifts
- `cc_max::AbstractArray`: Maximum correlation coefficient for each window
"""
function wcc_dt(ref::AbstractArray, cur::AbstractArray, t::AbstractArray, window::AbstractArray, fs::Float64,
             window_length::Float64, window_step::Float64)
     # create time axis for mwcs
     time_axis = Array(t[window[1]] + window_length / 2. : window_step :
                       t[window[end]] - window_length / 2.)

     window_length_samples = convert(Int,window_length * fs)
     window_step_samples = convert(Int,window_step * fs)
     minind = window[1]:window_step_samples:window[end] - window_length_samples
     padd = convert(Int,2 ^ (ceil(log2(abs(window_length_samples))) + 2))

     N = length(minind)
     dt = zeros(N)
     err = zeros(N)

     cci = zeros(window_length_samples,N)
     cri = zeros(window_length_samples,N)

     # fill matrices
     for ii = 1:N
         cci[:,ii] = cur[minind[ii]:minind[ii]+window_length_samples-1]
         cri[:,ii] = ref[minind[ii]:minind[ii]+window_length_samples-1]
     end

     # preprocess
     SeisNoise.demean!(cci)
     SeisNoise.detrend!(cci)
     SeisNoise.taper!(cci,fs,max_percentage=0.85)
     SeisNoise.demean!(cri)
     SeisNoise.detrend!(cri)
     SeisNoise.taper!(cri,fs,max_percentage=0.85)

     xcorr = zeros(2*window_length_samples-1, N)
     for ii=1:N
         xcorr[:,ii] = DSP.xcorr(cci[:,ii], cri[:,ii])
     end

     # get maximum correlation coefficient and its index
     cc_max, cc_max_ind = findmax(xcorr, dims=1)
     cc_max_ind = [cc_max_ind[i][1] for i=1:length(cc_max_ind)]
     # get time shift
     dt = (cc_max_ind .- window_length_samples)/fs

     return time_axis, dt, cc_max
end


text:
# To-do: derive high-fidelity gradient
using DwsimOpt
using JuMP
using Ipopt
using FiniteDifferences
using LinearAlgebra
# using GAMS
# using MadNLP

## SIMULATION OPTIMIZATION PROBLEM DEFINITION
path2sim = "c:/Users/lfsfr/Desktop/DwsimOpt.jl/examples/PRICO/PRICO_composite.dwxmz"
path2dwsim = "C:/Users/lfsfr/AppData/Local/DWSIM7/"

include("c:/Users/lfsfr/Desktop/DwsimOpt.jl/src/SimOpt.jl")
sim_jl = py"sim"

# Declare the simulation optimization problem
py"""
# Import dwsim-python data exchange interface (has to be after Automation2)
from dwsimopt.py2dwsim import create_pddx, assign_pddx

# Assign DoF:
create_pddx( ["MR-1", "CompoundMassFlow", "Nitrogen", "kg/s"],    sim, element="dof" )
create_pddx( ["MR-1", "CompoundMassFlow", "Methane", "kg/s"],     sim, element="dof" )
create_pddx( ["MR-1", "CompoundMassFlow", "Ethane", "kg/s"],      sim, element="dof" )
create_pddx( ["MR-1", "CompoundMassFlow", "Propane", "kg/s"],     sim, element="dof" )
create_pddx( ["VALV-01", "OutletPressure", "Mixture", "Pa"],      sim, element="dof" )
create_pddx( ["COMP-4", "OutletPressure", "Mixture", "Pa"],       sim, element="dof" )

# Assign F
create_pddx( ["Sum_W", "EnergyFlow", "Mixture", "kW"], sim, element="fobj" )

# adding constraints (g_i <= 0):
g1 = create_pddx( ["MITA1-Calc", "OutputVariable", "mita", "°C"], sim, element="constraint", assign=False )
r = 1.0/3
assign_pddx( lambda: r*(3-g1[0]()) , ["MITA1-Calc", "OutputVariable", "mita", "°C"], sim, element="constraint" )
create_pddx( ["MSTR-27", "MassFraction", "Liquid", "x"], sim, element="constraint" )
create_pddx( ["MR-1", "MassFraction", "Liquid", "x"], sim, element="constraint" )
create_pddx( ["MSTR-03", "MassFraction", "Liquid", "x"], sim, element="constraint" )
create_pddx( ["MSTR-05", "MassFraction", "Liquid", "x"], sim, element="constraint" )

# decision variables bounds
# x0 = np.array( [1.07249112e-04, 1.93057244e-04, 2.16881626e-04, 7.44968925e-04, 3.43386624e+05, 6.39810411e+06] )
# x0 = np.array( [1.12083333e-04, 1.30420113e-04, 2.57333445e-04, 8.00266697e-04, 3.45000000e+05, 1.53000000e+02] )
x0 = np.array( [1.12083333e-04, 2.20416667e-04, 1.42216680e-04, 7.54165016e-04, 3.45000000e+05, 7.20000000e+06] )
x0 = np.array( [1.3*1.12083333e-04, 1.3*2.20416667e-04, 1.3*1.42216680e-04, 1.3*7.54165016e-04, 3.45000000e+05, 7.20000000e+06] )
bounds_raw = np.array( [0.5*np.asarray(x0), 1.5*np.asarray(x0)] )   # 50 % around base case

# regularizer calculation
regularizer = np.zeros(x0.size)
regularizer_b = np.zeros(x0.size)
for i in range(len(regularizer)):
    regularizer[i] = 10**(-1*math.floor(math.log(x0[i],10))) # regularizer for magnitude order of 1e0

# bounds regularized
bounds_reg = regularizer*bounds_raw + regularizer_b

# objective and constraints lambda definitions
f = lambda x: sim.calculate_optProblem(np.asarray(x)/regularizer)[0:sim.n_f]
g = lambda x: sim.calculate_optProblem(np.asarray(x)/regularizer)[sim.n_f:(sim.n_f+sim.n_g)]
"""

f = py"f"
g = py"g"
x0 = py"x0*regularizer"
searchSpace = py"bounds_reg"
dim = sim_jl.n_dof
op1 = optProblem(f, g, x0, searchSpace, dim, sim_jl, py"sim.n_dof", py"sim.n_f", py"sim.n_g")
op1.sim_jl.verbose = false;
save_sim() = py"""sim.interface.SaveFlowsheet(sim.flowsheet,$pwd()+"/examples/PRICO/PRICO_composite2.dwxmz",True)"""

# external functions and jacobian calculations
x0 = 1.0 * (op1.searchSpace[1, :] + op1.searchSpace[2, :]) ./ 2
# x0 = py"x0*regularizer"
iter = 0
function f_block(x)
    global iter
    iter += 1
    return [op1.f(x)[1]; op1.g(x)[:]]
end
first_time = true
h = 0.01
flag_diff_method = "2nd_central"
flag_FiniteDifference = false
function jac_block(x)
    # println("hi")
    global x0, ∇f_bkp, ∇g1_bkp, ∇g2_bkp, ∇g3_bkp, ∇g4_bkp, ∇g5_bkp, first_time
    if norm(x0 - x) > 1e-6 || first_time == true
        f0 = f_block(x)
        println("Calculating jacobian of black-box functions at x=", x, " and f^{sim}=", f0)
        # flag_diff_method = "1st_forward"
        global h, flag_FiniteDifference, flag_diff_method
        tmp = zeros(6, length(x))
        if flag_FiniteDifference
            tmp = jacobian(central_fdm(2, 1), f_block, x)[1]
        else
            for i = 1:length(x)
                if flag_diff_method == "1st_forward"
                    tmp[:, i] = (f_block([x[1:i-1]; x[i] + h; x[i+1:end]]) - f0) / h
                elseif flag_diff_method == "2nd_central"
                    tmp[:, i] = (f_block([x[1:i-1]; x[i] + h; x[i+1:end]]) - f_block([x[1:i-1]; x[i] - h; x[i+1:end]])) / (2 * h)
                end
            end
        end
        println("∇f_block=", tmp)
        # tmp = jacobian(forward_fdm(2, 1), f_block, x)[1]
        ∇f_bkp = tmp[1, :]
        ∇g1_bkp = tmp[2, :]
        ∇g2_bkp = tmp[3, :]
        ∇g3_bkp = tmp[4, :]
        ∇g4_bkp = tmp[5, :]
        ∇g5_bkp = tmp[6, :]
        x0 = x
        first_time = false
        return tmp
    end
end

function setup_model()
    ## JuMP MODEL STUFF
    model = Model(optimizer_with_attributes(Ipopt.Optimizer))
    # model = Model(Ipopt.Optimizer)
    # model = Model(GAMS.Optimizer)
    # set_optimizer_attribute(model, GAMS.ModelType(), "NLP")
    # set_optimizer_attribute(model, "NLP", "CONOPT")
    # set_optimizer_attribute(model, "resLIM", 1200.0)
    # model = Model(() -> MadNLP.Optimizer(print_level = MadNLP.INFO, max_iter = 100))

    # VARIABLES
    i = 1
    @variable(model, op1.searchSpace[1, i]' <= x1 <= op1.searchSpace[2, i]', start = (op1.searchSpace[1, i] + op1.searchSpace[2, i]) ./ 2)
    i = 2
    @variable(model, op1.searchSpace[1, i]' <= x2 <= op1.searchSpace[2, i]', start = (op1.searchSpace[1, i] + op1.searchSpace[2, i]) ./ 2)
    i = 3
    @variable(model, op1.searchSpace[1, i]' <= x3 <= op1.searchSpace[2, i]', start = (op1.searchSpace[1, i] + op1.searchSpace[2, i]) ./ 2)
    i = 4
    @variable(model, op1.searchSpace[1, i]' <= x4 <= op1.searchSpace[2, i]', start = (op1.searchSpace[1, i] + op1.searchSpace[2, i]) ./ 2)
    i = 5
    @variable(model, op1.searchSpace[1, i]' <= x5 <= op1.searchSpace[2, i]', start = (op1.searchSpace[1, i] + op1.searchSpace[2, i]) ./ 2)
    i = 6
    @variable(model, op1.searchSpace[1, i]' <= x6 <= op1.searchSpace[2, i]', start = (op1.searchSpace[1, i] + op1.searchSpace[2, i]) ./ 2)

    # OBJECTIVE FUNCTION
    fobj_bb(x1, x2, x3, x4, x5, x6) = op1.f([x1, x2, x3, x4, x5, x6])[1]
    function ∇fobj_bb(g::AbstractVector{T}, x1::T, x2::T, x3::T, x4::T, x5::T, x6::T) where {T}
        # println("I am in ∇fobj_bb with x=", x1)
        global x0, ∇f_bkp, ∇g1_bkp, ∇g2_bkp, ∇g3_bkp, ∇g4_bkp, ∇g5_bkp
        jac_block([x1, x2, x3, x4, x5, x6])
        ∇f_tmp = ∇f_bkp
        # ∇f_tmp = grad(forward_fdm(2, 1), op1.f, [x1, x2, x3, x4, x5, x6, x7, x8])[1]
        for i = 1:length([x1, x2, x3, x4, x5, x6])
            g[i] = ∇f_tmp[i]
        end
    end
    register(model, :fobj_bb, op1.dim, fobj_bb, ∇fobj_bb; autodiff = false)
    @NLobjective(model, Min, fobj_bb(x1, x2, x3, x4, x5, x6))

    # CONSTRAINTS
    g1(x) = op1.g(x)[1]
    g2(x) = op1.g(x)[2]
    g3(x) = op1.g(x)[3]
    g4(x) = op1.g(x)[4]
    g5(x) = op1.g(x)[5]
    g1_bb(x1, x2, x3, x4, x5, x6) = op1.g([x1, x2, x3, x4, x5, x6])[1]
    g2_bb(x1, x2, x3, x4, x5, x6) = op1.g([x1, x2, x3, x4, x5, x6])[2]
    g3_bb(x1, x2, x3, x4, x5, x6) = op1.g([x1, x2, x3, x4, x5, x6])[3]
    g4_bb(x1, x2, x3, x4, x5, x6) = op1.g([x1, x2, x3, x4, x5, x6])[4]
    g5_bb(x1, x2, x3, x4, x5, x6) = op1.g([x1, x2, x3, x4, x5, x6])[5]
    function ∇g1_bb(g::AbstractVector{T}, x1::T, x2::T, x3::T, x4::T, x5::T, x6::T) where {T}
        # println("I am in ∇g1_bb with x=", x1)
        global x0, ∇f_bkp, ∇g1_bkp, ∇g2_bkp, ∇g3_bkp, ∇g4_bkp, ∇g5_bkp
        jac_block([x1, x2, x3, x4, x5, x6])
        ∇g1_tmp = ∇g1_bkp
        for i = 1:length([x1, x2, x3, x4, x5, x6])
            g[i] = ∇g1_tmp[i]
        end
    end
    function ∇g2_bb(g::AbstractVector{T}, x1::T, x2::T, x3::T, x4::T, x5::T, x6::T) where {T}
        # println("I am in ∇g2_bb with x=", x1)
        global x0, ∇f_bkp, ∇g1_bkp, ∇g2_bkp, ∇g3_bkp, ∇g4_bkp, ∇g5_bkp
        jac_block([x1, x2, x3, x4, x5, x6])
        ∇g2_tmp = ∇g2_bkp
        for i = 1:length([x1, x2, x3, x4, x5, x6])
            g[i] = ∇g2_tmp[i]
        end
    end
    function ∇g3_bb(g::AbstractVector{T}, x1::T, x2::T, x3::T, x4::T, x5::T, x6::T) where {T}
        # println("I am in ∇g2_bb with x=", x1)
        global x0, ∇f_bkp, ∇g1_bkp, ∇g2_bkp, ∇g3_bkp, ∇g4_bkp, ∇g5_bkp
        jac_block([x1, x2, x3, x4, x5, x6])
        ∇g3_tmp = ∇g3_bkp
        for i = 1:length([x1, x2, x3, x4, x5, x6])
            g[i] = ∇g3_tmp[i]
        end
    end
    function ∇g4_bb(g::AbstractVector{T}, x1::T, x2::T, x3::T, x4::T, x5::T, x6::T) where {T}
        # println("I am in ∇g2_bb with x=", x1)
        global x0, ∇f_bkp, ∇g1_bkp, ∇g2_bkp, ∇g3_bkp, ∇g4_bkp, ∇g5_bkp
        jac_block([x1, x2, x3, x4, x5, x6])
        ∇g4_tmp = ∇g4_bkp
        for i = 1:length([x1, x2, x3, x4, x5, x6])
            g[i] = ∇g4_tmp[i]
        end
    end
    function ∇g5_bb(g::AbstractVector{T}, x1::T, x2::T, x3::T, x4::T, x5::T, x6::T) where {T}
        # println("I am in ∇g2_bb with x=", x1)
        global x0, ∇f_bkp, ∇g1_bkp, ∇g2_bkp, ∇g3_bkp, ∇g4_bkp, ∇g5_bkp
        jac_block([x1, x2, x3, x4, x5, x6])
        ∇g5_tmp = ∇g5_bkp
        for i = 1:length([x1, x2, x3, x4, x5, x6])
            g[i] = ∇g5_tmp[i]
        end
    end
    register(model, :g1_bb, op1.dim, g1_bb, ∇g1_bb; autodiff = false)
    @NLconstraint(model, g1_bb(x1, x2, x3, x4, x5, x6) <= 0)
    register(model, :g2_bb, op1.dim, g2_bb, ∇g2_bb; autodiff = false)
    @NLconstraint(model, g2_bb(x1, x2, x3, x4, x5, x6) == 0)
    register(model, :g3_bb, op1.dim, g2_bb, ∇g2_bb; autodiff = false)
    @NLconstraint(model, g3_bb(x1, x2, x3, x4, x5, x6) == 0)
    register(model, :g4_bb, op1.dim, g2_bb, ∇g2_bb; autodiff = false)
    @NLconstraint(model, g4_bb(x1, x2, x3, x4, x5, x6) == 0)
    register(model, :g5_bb, op1.dim, g2_bb, ∇g2_bb; autodiff = false)
    @NLconstraint(model, g5_bb(x1, x2, x3, x4, x5, x6) == 0)

    return model
end

model = setup_model()
function set_opt_params(model)
    set_optimizer_attribute(model, "tol", 1e-2)
    set_optimizer_attribute(model, "dual_inf_tol", 1e-2)
    set_optimizer_attribute(model, "constr_viol_tol", 1e-3)
    set_optimizer_attribute(model, "compl_inf_tol", 1e-2)
    set_optimizer_attribute(model, "acceptable_tol", 1e-2)
    set_optimizer_attribute(model, "acceptable_iter", 3)
    # set_optimizer_attribute(model, "mu_target", 1e-6)

    set_optimizer_attribute(model, "hessian_approximation", "limited-memory")

    set_optimizer_attribute(model, "max_iter", 50)
    set_optimizer_attribute(model, "max_cpu_time", 3600.0)
    set_optimizer_attribute(model, "print_level", 5)
    return model
end
model = set_opt_params(model)
JuMP.optimize!(model)

iter
println("Model used h=", h, " and tol=", get_optimizer_attribute(model, "tol"))
println("Initial value x0=", (op1.searchSpace[1, :] + op1.searchSpace[2, :]) ./ 2)
op1.searchSpace

@show JuMP.termination_status(model)
@show JuMP.primal_feasibility_report(model)
@show JuMP.dual_status(model)

# retrieve the objective value, corresponding x values and the status
println([value(model[:x1]), value(model[:x2]), value(model[:x3]), value(model[:x4]), value(model[:x5]), value(model[:x6])])
println(JuMP.objective_value(model))
println(JuMP.termination_status(model))

f_block(JuMP.value.([value(model[:x1]), value(model[:x2]), value(model[:x3]), value(model[:x4]), value(model[:x5]), value(model[:x6])]))


text:
using Revise
using Distributed
# addprocs(2)

# @everywhere begin
    using Random
    using POMDPs
    using Spot
    using POMDPModelChecking
    import Cairo
    using POMDPModelTools
    using POMDPModels
    using POMDPSimulators
    using POMDPGifs
    using ProgressMeter
    using BeliefUpdaters
    using SARSOP
    using FileIO
    using JLD2

    include("../test/blind_gridworld.jl")

    const LABELLED_STATES = Dict(GWPos(4,3) => :crash, GWPos(4,6)=>:crash,  GWPos(9,3)=>:a, GWPos(8,8)=>:b)

    # const LABELLED_STATES = Dict(GWPos(4,3) => :crash, GWPos(4,6)=>:crash)

    function POMDPModelChecking.labels(mdp::SimpleGridWorld, s, a)
        if haskey(LABELLED_STATES, s)
            return tuple(LABELLED_STATES[s])
        else
            return ()
        end
    end

    include("mc_simulation.jl")

# end

prop = ltl"!crash U a"
# prop = ltl"!crash U a & !crash U b"
# prop = ltl"G !crash"

mdp = SimpleGridWorld(rewards=Dict(GWPos(9,3)=>10.0), terminate_from=Set([]))

@everywhere POMDPModelChecking.labels(pomdp::BlindGridWorld, s, a) = labels(pomdp.simple_gw, s, a)
pomdp = BlindGridWorld(exit=GWPos(-1,-1), simple_gw = mdp)

@show length(states(pomdp))
@show length(actions(pomdp))
@show length(observations(pomdp))

# overwriding discount, to trigger SARSOP iterations
POMDPs.discount(problem::Union{ProductMDP, ProductPOMDP}) = 0.999
sarsop = SARSOPSolver(precision=1e-2, timeout=30)
solver = ModelCheckingSolver(solver=sarsop, property=prop, verbose=true)

# run(`rm model.pomdpx`)
policy = solve(solver, pomdp);

# these are the bounds printed by the SARSOP output
# (unfortunately no way to get them programmatically with the current SARSOP implementation)
UBOUND = 0.914258
LBOUND = 0.900412

successes, mu, sig = many_sims(policy.problem, policy, 100);

using PGFPlots

p = Plots.Linear(1:length(mu), mu, legendentry="MC estimate", style="blue, thick", mark="none");
p2 = Plots.Linear(1:length(mu), LBOUND*ones(length(mu)), style="ultra thick, black, dotted", mark="none", legendentry="SARSOP lower bound");
p3 = Plots.Linear(1:length(mu), UBOUND*ones(length(mu)), style="ultra thick, black, dashed", mark="none", legendentry="SARSOP upper bound");
ax = Axis([p, p2, p3], ylabel="Probability of success", xlabel="Number of simulations", legendPos="south east", style="grid=both")

PGFPlots.save("succ.pdf", ax)

#=

up = DiscreteUpdater(policy.problem)
b0 = initialize_belief(up, initialstate_distribution(policy.problem))
@show value(policy, b0)

policy.memory = 1
ctx1 = POMDPModels.render(mdp, Dict(), valuecolor=s->value(policy, deterministic_belief(policy.problem, ProductState(s, policy.memory))),
                                      action=s->action(policy, deterministic_belief(policy.problem, ProductState(s, policy.memory))),
                                      landmark = s -> s ∈ keys(LABELLED_STATES));
ctx1 |> PNG("gw.png")

## simulation

up = DiscreteUpdater(policy.problem)
b0 = initialize_belief(up, initialstate_distribution(policy.problem))
hr = HistoryRecorder(max_steps=50)
hist = simulate(hr, policy.problem, policy, up, b0)
prod_state_hist = hist.state_hist
state_hist = [s.s for s in hist.state_hist];
hist2 = POMDPHistory(state_hist, [getfield(hist, f) for f in fieldnames(POMDPHistory) if f != :state_hist]...);
makegif(pomdp, hist2, filename="bgw.gif", spec="(s,a)")

pomdp = BlindGridWorld()

sarsop = SARSOPSolver(precision=1e-3)
solver = ModelCheckingSolver(solver=sarsop, property=prop, verbose=true)

solve(solver, pomdp)

automata = DeterministicRabinAutomata(prop)

label_to_array(prop)

translator = LTLTranslator(deterministic=true, buchi=true, state_based_acceptance=true)
aut = translate(translator, prop)

aut = SpotAutomata(aut)

get_rabin_acceptance(aut)

dra = to_generalized_rabin(aut)

get_rabin_acceptance(dra)

dra.a.to_str()
get_rabin_acceptance(dra)

states = 1:num_states(dra)
initial_state = get_init_state_number(dra)
APs = atomic_propositions(dra)
edgelist, labels = get_edges_labels(dra)

conditions = label_to_array.(labels)

label_to_array(labels[1])

=#


text:
import RRSTExperiments: InputEdge

function solve_rec_st_model(n::Int, E::Vector{InputEdge}, k::Int)::Float64

  V = collect(1:n) # set of nodes
  Vminus1 = setdiff(V, [1]) # commodity nodes
  A = [(e.i, e.j) for e ∈ E] ∪ [(e.j, e.i) for e ∈ E]
  L = n-1-k

  # Model
  model = Model(CPLEX.Optimizer)
  set_silent(model)
  set_optimizer_attribute(model, "CPX_PARAM_EPINT", 1e-8)

  # Variables
  @variable(model, fx[A, Vminus1] ≥ 0)
  @variable(model, wx[A] ≥ 0)
  @variable(model, x[E] ≥ 0)

  @variable(model, fy[A, Vminus1] ≥ 0)
  @variable(model, wy[A] ≥ 0)
  @variable(model, y[E] ≥ 0)
  @variable(model, z[E] ≥ 0)

  # Objective
  @objective(model, Min, sum(e.C * x[e] for e ∈ E) + sum((e.c + e.d) * y[e] for e ∈ E))

  #Constraints
  #Constraints
  for k ∈ Vminus1 # sources
      @constraint(model, sum(fx[(j,1),k] for j ∈ filter(j -> (j,1) ∈ A, V)) -
      sum(fx[(1,j),k] for j ∈ filter(j -> (1,j) ∈ A,V)) == -1)
      @constraint(model, sum(fy[(j,1),k] for j ∈ filter(j -> (j,1) ∈ A, V)) -
      sum(fy[(1,j),k] for j ∈ filter(j -> (1,j) ∈ A,V)) == -1)
  end

  for k ∈ Vminus1, i ∈ Vminus1 # balances
      if i ≠ k
          @constraint(model, sum(fx[(j,i),k] for j ∈ filter(j -> (j,i) ∈ A, V)) -
          sum(fx[(i,j),k] for j ∈ filter(j -> (i,j) ∈ A, V)) == 0)
          @constraint(model, sum(fy[(j,i),k] for j ∈ filter(j -> (j,i) ∈ A, V)) -
          sum(fy[(i,j),k] for j ∈ filter(j -> (i,j) ∈ A, V)) == 0)
      end
  end

  for k ∈ Vminus1 # sinks
      @constraint(model, sum(fx[(j,k),k] for j ∈ filter(j -> (j,k) ∈ A,V)) -
      sum(fx[(k,j),k] for j ∈ filter(j -> (k,j) ∈ A,V)) == 1)
      @constraint(model, sum(fy[(j,k),k] for j ∈ filter(j -> (j,k) ∈ A,V)) -
      sum(fy[(k,j),k] for j ∈ filter(j -> (k,j) ∈ A,V)) == 1)
  end

  for k ∈ Vminus1, a ∈ A # capacity
      @constraint(model, fx[a,k] ≤ wx[a])
      @constraint(model, fy[a,k] ≤ wy[a])
  end

  # tree
  @constraint(model, sum(wx[a] for a ∈ A) == n - 1)
  @constraint(model, sum(wy[a] for a ∈ A) == n - 1)

  for e ∈ E
      @constraint(model, x[e] == wx[(e.i, e.j)] + wx[(e.j, e.i)])
      @constraint(model, y[e] == wy[(e.i, e.j)] + wy[(e.j, e.i)])
  end

  for e ∈ E
      @constraint(model, x[e] ≥ z[e])
      @constraint(model, y[e] ≥ z[e])
  end

  @constraint(model, sum(z[e] for e ∈ E) ≥ L)

  # Solve
  optimize!(model)

  return objective_value(model)
end


text:
function Persa.train!(model::MatrixFactorization,
                dataset::Persa.Dataset;
                γ = 0.001,
                λ = 0.02,
                max_epochs = 1000)

    err = Inf
    p = Progress(max_epochs)
    for epoch = 1:max_epochs
        update!(model, dataset, γ, λ)
        e = objective(model, dataset, 0.1)

        if e > err
            break
        end

        err = e

        ProgressMeter.next!(p; showvalues = [(:epoch, epoch), (:error, err)])
    end

    return nothing
end


text:
# Autogenerated wrapper script for SuiteSparse32_jll for x86_64-linux-musl
export libamd_a, libbtf_a, libcolamd_a, libklu_a, libsuitesparseconfig_a

JLLWrappers.@generate_wrapper_header("SuiteSparse32")
JLLWrappers.@declare_file_product(libamd_a)
JLLWrappers.@declare_file_product(libbtf_a)
JLLWrappers.@declare_file_product(libcolamd_a)
JLLWrappers.@declare_file_product(libklu_a)
JLLWrappers.@declare_file_product(libsuitesparseconfig_a)
function __init__()
    JLLWrappers.@generate_init_header()
    JLLWrappers.@init_file_product(
        libamd_a,
        "lib/libamd.a",
    )

    JLLWrappers.@init_file_product(
        libbtf_a,
        "lib/libbtf.a",
    )

    JLLWrappers.@init_file_product(
        libcolamd_a,
        "lib/libcolamd.a",
    )

    JLLWrappers.@init_file_product(
        libklu_a,
        "lib/libklu.a",
    )

    JLLWrappers.@init_file_product(
        libsuitesparseconfig_a,
        "lib/libsuitesparseconfig.a",
    )

    JLLWrappers.@generate_init_footer()
end  # __init__()


text:
function get_grid_points(obj::FieldInterpolationGrid, arg0::RealFieldElement, arg1::RealFieldElement)
    return jcall(obj, "getGridPoints", Vector{RealFieldElement}, (RealFieldElement, RealFieldElement), arg0, arg1)
end


text:
include("../les/OceananigansData.jl")
include("../les/LESbraryData.jl")

prefix = pwd()*"/src/les/data/"

directories=Dict(
    "free_convection"          => "2daySuite/three_layer_constant_fluxes_Qu0.0e+00_Qb1.0e-07_f1.0e-04_Nh128_Nz128/three_layer_constant_fluxes_Qu0.0e+00_Qb1.0e-07_f1.0e-04_Nh128_Nz128_averaged_statistics.jld2",
    "strong_wind"              => "2daySuite/three_layer_constant_fluxes_Qu1.0e-03_Qb0.0e+00_f1.0e-04_Nh128_Nz128/three_layer_constant_fluxes_Qu1.0e-03_Qb0.0e+00_f1.0e-04_Nh128_Nz128_averaged_statistics.jld2",
    "strong_wind_no_coriolis"  => "2daySuite/three_layer_constant_fluxes_Qu2.0e-04_Qb0.0e+00_f0.0e+00_Nh128_Nz128/three_layer_constant_fluxes_Qu2.0e-04_Qb0.0e+00_f0.0e+00_Nh128_Nz128_averaged_statistics.jld2",
    "weak_wind_strong_cooling" => "2daySuite/three_layer_constant_fluxes_Qu2.0e-04_Qb1.0e-07_f1.0e-04_Nh128_Nz128/three_layer_constant_fluxes_Qu2.0e-04_Qb1.0e-07_f1.0e-04_Nh128_Nz128_averaged_statistics.jld2",
    "strong_wind_weak_cooling" => "2daySuite/three_layer_constant_fluxes_Qu8.0e-04_Qb5.0e-09_f1.0e-04_Nh128_Nz128/three_layer_constant_fluxes_Qu8.0e-04_Qb5.0e-09_f1.0e-04_Nh128_Nz128_averaged_statistics.jld2",
    "strong_wind_weak_heating" => "2daySuite/three_layer_constant_fluxes_Qu8.0e-04_Qb-1.0e-08_f1.0e-04_Nh128_Nz128/three_layer_constant_fluxes_Qu8.0e-04_Qb-1.0e-08_f1.0e-04_Nh128_Nz128_averaged_statistics.jld2",
)

# directories=Dict(
#
#     "free_convection" => Dict(
#         2 => "2daySuite/three_layer_constant_fluxes_Qu0.0e+00_Qb1.0e-07_f1.0e-04_Nh128_Nz128/three_layer_constant_fluxes_Qu0.0e+00_Qb1.0e-07_f1.0e-04_Nh128_Nz128_averaged_statistics.jld2",
#     ),
#
#     "strong_wind" => Dict(
#         2 => "2daySuite/three_layer_constant_fluxes_Qu1.0e-03_Qb0.0e+00_f1.0e-04_Nh128_Nz128/three_layer_constant_fluxes_Qu1.0e-03_Qb0.0e+00_f1.0e-04_Nh128_Nz128_averaged_statistics.jld2",
#     ),
#
#     "strong_wind_no_coriolis" => Dict(
#         2 => "2daySuite/three_layer_constant_fluxes_Qu2.0e-04_Qb0.0e+00_f0.0e+00_Nh128_Nz128/three_layer_constant_fluxes_Qu2.0e-04_Qb0.0e+00_f0.0e+00_Nh128_Nz128_averaged_statistics.jld2",
#     ),
#
#     "weak_wind_strong_cooling" => Dict(
#         2 => "2daySuite/three_layer_constant_fluxes_Qu2.0e-04_Qb1.0e-07_f1.0e-04_Nh128_Nz128/three_layer_constant_fluxes_Qu2.0e-04_Qb1.0e-07_f1.0e-04_Nh128_Nz128_averaged_statistics.jld2",
#     ),
#
#     "strong_wind_weak_cooling" => Dict(
#         2 => "2daySuite/three_layer_constant_fluxes_Qu8.0e-04_Qb5.0e-09_f1.0e-04_Nh128_Nz128/three_layer_constant_fluxes_Qu8.0e-04_Qb5.0e-09_f1.0e-04_Nh128_Nz128_averaged_statistics.jld2",
#     ),
#
#     "strong_wind_weak_heating" => Dict(
#         2 => "2daySuite/three_layer_constant_fluxes_Qu8.0e-04_Qb-1.0e-08_f1.0e-04_Nh128_Nz128/three_layer_constant_fluxes_Qu8.0e-04_Qb-1.0e-08_f1.0e-04_Nh128_Nz128_averaged_statistics.jld2",
#     ),
# )

"""
read_les_output(filename::String; days=0)
"""
function read_les_output(filename::String)

    if filename[end-2:end]==".nc"
        # println("Collecting NetCDF data.")
        # include("../les/convert_netcdf_to_data.jl")
        filename=prefix*"data_netcdf/$(filename[1:end-3])/$(filename)"
        return ReadNetCDF_OceananigansData(filename)
    elseif filename[1:13]=="general_strat"
        # println("Collecting JLD2 data.")
        # include("../les/convert_jld2_to_data.jl")
        filename=prefix*"general_strat_sandreza/$(filename[1:end-5])/$(filename)"
        return ReadJLD2_OceananigansData(filename)
    else
        filename=prefix*directories[filename]
        return ReadJLD2_LESbraryData(filename)
    end
end

# read_les_output("strong_wind_weak_heating"; days=2)

max_stars_repo_name: adelinehillier/LearnConvection

max_issues_repo_name: adelinehillier/LearnConvection

max_forks_repo_name: adelinehillier/LearnConvection


text:
"""
    ImplicitFunction{F,C,L}

Differentiable wrapper for an implicit function `x -> ŷ(x)` whose output is defined by explicit conditions `F(x,ŷ(x)) = 0`.

We can obtain the Jacobian of `ŷ` with the implicit function theorem:
```
∂₁F(x,ŷ(x)) + ∂₂F(x,ŷ(x)) * ∂ŷ(x) = 0
```
If `x ∈ ℝⁿ`, `y ∈ ℝᵐ` and `F(x,y) ∈ ℝᶜ`, this amounts to solving the linear system `A * J = B`, where `A ∈ ℝᶜᵐ`, `B ∈ ℝᶜⁿ` and `J ∈ ℝᵐⁿ`.

# Fields:
- `forward::F`: callable of the form `x -> ŷ(x)`
- `conditions::C`: callable of the form `(x,y) -> F(x,y)`
- `linear_solver::L`: callable of the form `(A,b) -> u` such that `A * u = b`
"""
struct ImplicitFunction{F,C,L}
    forward::F
    conditions::C
    linear_solver::L
end

"""
    ImplicitFunction(forward, conditions)

Construct an `ImplicitFunction` with `Krylov.gmres` as the default linear solver.

# See also
- [`ImplicitFunction{F,C,L}`](@ref)
"""
function ImplicitFunction(forward::F, conditions::C) where {F,C}
    return ImplicitFunction(forward, conditions, Krylov.gmres)
end

struct SolverFailureException <: Exception
    msg::String
end

"""
    implicit(x)

Make [`ImplicitFunction`](@ref) callable by applying `implicit.forward`.
"""
(implicit::ImplicitFunction)(x) = implicit.forward(x)

"""
    frule(rc, (_, dx), implicit, x)

Custom forward rule for [`ImplicitFunction`](@ref).

We compute the Jacobian-vector product `Jv` by solving `Au = Bv` and setting `Jv = u`.
"""
function ChainRulesCore.frule(
    rc::RuleConfig, (_, dx), implicit::ImplicitFunction, x::AbstractVector
)
    (; forward, conditions, linear_solver) = implicit

    y = forward(x)
    n, m = length(x), length(y)

    conditions_x(x̃) = conditions(x̃, y)
    conditions_y(ỹ) = -conditions(x, ỹ)

    pushforward_A(dỹ) = frule_via_ad(rc, (NoTangent(), dỹ), conditions_y, y)[2]
    pushforward_B(dx̃) = frule_via_ad(rc, (NoTangent(), dx̃), conditions_x, x)[2]

    mul_A!(res, v) = res .= pushforward_A(v)
    mul_B!(res, v) = res .= pushforward_B(v)

    A = LinearOperator(Float64, m, m, false, false, mul_A!)
    B = LinearOperator(Float64, m, n, false, false, mul_B!)

    dx_vec = Vector(unthunk(dx))
    b = B * dx_vec
    dy_vec, stats = linear_solver(A, b)
    if !stats.solved
        throw(SolverFailureException("The linear solver failed to converge"))
    end
    return y, dy_vec
end

"""
    rrule(rc, implicit, x)

Custom reverse rule for [`ImplicitFunction`](@ref).

We compute the vector-Jacobian product `Jᵀv` by solving `Aᵀu = v` and setting `Jᵀv = Bᵀu`.
"""
function ChainRulesCore.rrule(rc::RuleConfig, implicit::ImplicitFunction, x::AbstractVector)
    (; forward, conditions, linear_solver) = implicit

    y = forward(x)
    n, m = length(x), length(y)

    conditions_x(x̃) = conditions(x̃, y)
    conditions_y(ỹ) = -conditions(x, ỹ)

    pullback_Aᵀ = last ∘ rrule_via_ad(rc, conditions_y, y)[2]
    pullback_Bᵀ = last ∘ rrule_via_ad(rc, conditions_x, x)[2]

    mul_Aᵀ!(res, v) = res .= pullback_Aᵀ(v)
    mul_Bᵀ!(res, v) = res .= pullback_Bᵀ(v)

    Aᵀ = LinearOperator(Float64, m, m, false, false, mul_Aᵀ!)
    Bᵀ = LinearOperator(Float64, n, m, false, false, mul_Bᵀ!)

    function implicit_pullback(dy)
        dy_vec = Vector(unthunk(dy))
        u, stats = linear_solver(Aᵀ, dy_vec)
        if !stats.solved
            throw(SolverFailureException("The linear solver failed to converge"))
        end
        dx_vec = Bᵀ * u
        return (NoTangent(), dx_vec)
    end

    return y, implicit_pullback
end


text:
"""
    Momentum(::Vector{XVariable}; lr=1e-4, inertia=0.9, L1decay=0.001, L2decay=0.01)

Implements stochastic gradient descent with momentum
"""
mutable struct Momentum <: Optimizer
    xparams::Vector{XVariable}
    v::Vector
    lr::AbstractFloat
    inertia::AbstractFloat
    L1decay::AbstractFloat
    L2decay::AbstractFloat
    name::String
    function Momentum(xparams::Vector{XVariable}; lr=1e-4, inertia=0.9, L1decay=0.001, L2decay=0.01)
        num = length(xparams)
        vel = Vector(undef,num)
        for i = 1:num
            c , θ = xparams[i]
            vel[i] = Zeros(typeof(ᵛ(θ)), θ.shape)
        end
        new(xparams, vel, lr, inertia, L1decay, L2decay, "Momentum")
    end
end

function Base.show(io::IO, O::Momentum)
    print("Momentum(lr=$(O.lr), inertia=$(O.inertia), L1decay=$(O.L1decay), L2decay=$(O.L2decay))")
end

function update!(O::Momentum; clipfn::Function=LPInfNormClip, clipvalue=10.0)
    vel = O.v
    μ = - O.lr
    ρ = O.inertia
    λ₁ = O.L1decay
    λ₂ = O.L2decay

    for i = 1:length(O.xparams)
        c , θ = O.xparams[i]
        ∇ = clipfn(setNanInfZero(δ(θ)), clipvalue)
        𝒗 = ᵛ(θ)
        @. vel[i] = ρ * vel[i] + ∇
        if c == 'w'
            if λ₁==0 && λ₂==0
                @. 𝒗 += μ * vel[i]
            elseif λ₁==0 && λ₂!=0
                @. 𝒗 += μ * (vel[i] + λ₂ * 𝒗)
            elseif λ₁!=0 && λ₂==0
                @. 𝒗 += μ * (vel[i] + λ₁ * sign(𝒗))
            else  # λ₁!=0 && λ₂!=0
                @. 𝒗 += μ * (vel[i] + λ₁ * sign(𝒗) + λ₂ * 𝒗)
            end
        else
            @. 𝒗 += μ * vel[i]
        end
    end
end

max_stars_repo_name: sonosole/Delta

max_issues_repo_name: sonosole/Delta

max_forks_repo_name: sonosole/Delta