text:
#!/usr/bin/env julia

# This file contains data and functions helping with the translation of symbols for the use of RHEOS without unicode symbols.

# translation table for symbols. left = symbol use by user, right = symbol used by model declaraction or RHEOS structs.
const symbol_convertion_table=(
    time = :t,
    t_col = :t,
    strain = :ϵ,
    epsilon = :ϵ,
    ϵ_col = :ϵ,
    ε = :ϵ,
    stress = :σ,
    sigma = :σ,
    σ_col = :σ,
    omega = :ω,
    ω_col = :ω,
    Gp_col = :Gp,
    Gpp_col = :Gpp,
    eta = :η,
    alpha = :α,
    beta = :β,
    gamma = :γ,
    c_alpha = :cₐ,
    c_beta = :cᵦ,
    c_gamma = :cᵧ,
    k_1 = :k₁,
    k_2 = :k₂,
    k_a = :kₐ,
    k_beta = :kᵦ,
    k_gamma = :kᵧ,
    eta_1 = :η₁,
    eta_2 = :η₂,
    eta_a = :ηₐ,
    eta_beta = :ηᵦ,
    eta_gamma = :ηᵧ,
    tau = :τ,
    G_phi = :Gᵩ,
    G_0 = :G₀,
    G_1 = :G₁,
    G_2 = :G₂
    )

function symbol_to_unicode(s::Symbol)
    s in keys(symbol_convertion_table) ? symbol_convertion_table[s] : s
end

function symbol_to_unicode(nt::NamedTuple)
    NamedTuple{Tuple([ symbol_to_unicode(s) for s in keys(nt) ])}( values(nt) )
end

# Edge case useful for function taking option tuple parameters, e.g. modelfit
function symbol_to_unicode(nt::Nothing)
    nothing
end

# When called from Python, it is useful to accept Dict rather than NamedTuple which are not supported.
function symbol_to_unicode(d::Dict)
    symbol_to_unicode( NamedTuple{([Symbol(e) for e in keys(d)]...,)}(values(d)) )
end

function unicode_to_text(s::Symbol)
    r=findfirst(e->e==s, symbol_convertion_table)
    isnothing(r) ? s : r
end

function unicode_to_text(nt::NamedTuple)
    NamedTuple{Tuple([ unicode_to_text(s) for s in keys(nt) ])}( values(nt) )
end

#
#   Convenience function for python users.
#

"""
    function namedtuple(d::Dict)
    function namedtuple(;kwargs...)

`namedtuple` returns the named tuple formed using keyword arguments or a dictionary passed as parameter.

# Example
```@example
julia> nt = namedtuple(x=1,y=2)
(x = 1, y = 2)
```
"""
function namedtuple(d::Dict)
    NamedTuple{([Symbol(e) for e in keys(d)]...,)}(values(d))
end

function namedtuple(;kwargs...)
    kwargs.data
end

"""
    function dict(nt::NamedTuple)
    function dict(;kwargs...)

`dict` returns the dictionary formed using keyword arguments or a named tuple passed as parameter.

# Example
```@example
julia> d = dict(x=1,y=2)
Dict{Symbol, Int64} with 2 entries:
  :y => 2
  :x => 1
```
"""
function dict(nt::NamedTuple)
    Dict(zip(keys(nt),values(nt)))
end

function dict(;kwargs...)
    dict(kwargs.data)
end


text:
module PandorasBoxDashToggleBlock

using ..Ahorn, Maple

@mapdef Entity "pandorasBox/dashToggleBlock" DashTogleBlock(x::Integer, y::Integer, index::String="0", divisor::Integer=2)

const placements = Ahorn.PlacementDict(
    "Dash Toggle Block (Pandora's Box)" => Ahorn.EntityPlacement(
        DashTogleBlock,
        "rectangle"
    ),
)

Ahorn.minimumSize(entity::DashTogleBlock) = 8, 8
Ahorn.resizable(entity::DashTogleBlock) = true, true

Ahorn.selection(entity::DashTogleBlock) = Ahorn.getEntityRectangle(entity)

function Ahorn.render(ctx::Ahorn.Cairo.CairoContext, entity::DashTogleBlock, room::Maple.Room)
    x = Int(get(entity.data, "x", 0))
    y = Int(get(entity.data, "y", 0))

    width = Int(get(entity.data, "width", 8))
    height = Int(get(entity.data, "height", 8))

    Ahorn.drawRectangle(ctx, 0, 0, width, height, (0.8, 0.3, 1.0, 0.4), (0.8, 0.3, 1.0, 1.0))
end

end

ext: jl

max_stars_repo_name: DemoJameson/PandorasBox

max_issues_repo_name: DemoJameson/PandorasBox

max_forks_repo_name: DemoJameson/PandorasBox


text:
include("../../src/MutualInformation/mutual-information.jl")
include("reader_MI_tests.jl")

function get_layer_bins_vector(vector; neuron_nbins = 5, nbins = 5, lower_lim=-1, upper_lim=1)
    columns = separate_columns(vector)
    hists = create_bin_histogram.(columns, [neuron_nbins], [lower_lim], [upper_lim])
    bounds = (get_weights_length.(hists)...,)
    indices = [getindex.(get_bin.(hists, v),1) for v in vector]
    indices = [correct_upper_bounds(ind, bounds) for ind in indices]
    combination_values = get_bin_from_indices.(indices, [bounds])
    max_value = foldr(*, bounds)
    create_bin_vector(combination_values, nbins, 1, max_value)
end

function calculate_joint_probability_MI(X, Y, lim_X, lim_Y, nbins)
    number_neurons_X = length(X[1])
    bins_X = nbins^(number_neurons_X)
    if number_neurons_X > 1
        linearized_X = get_layer_bins_vector(X, neuron_nbins=bins_X, nbins=nbins, lower_lim=lim_X[1], upper_lim=lim_X[2])
        limits_X = (minimum(linearized_X), maximum(linearized_X))
    else
        linearized_X = getindex.(X, 1)
        limits_X = lim_X
    end

    number_neurons_Y = length(Y[1])
    bins_Y = nbins^(number_neurons_Y)
    if number_neurons_Y > 1
        linearized_Y = get_layer_bins_vector(Y, neuron_nbins=nbins, nbins=bins_Y, lower_lim=lim_Y[1], upper_lim=lim_Y[2])
        limits_Y = (minimum(linearized_Y), maximum(linearized_Y))
    else
        linearized_Y = getindex.(Y, 1)
        limits_Y = lim_Y
    end

    calculate_joint_probability(linearized_X, linearized_Y, limits_X, limits_Y, (bins_X+1, round(Int, (bins_Y))))
end

function get_layer_probabilities(vector; neuron_nbins = 5, nbins = 5, lower_lim=-1, upper_lim=1)
    columns = separate_columns(vector)
    hists = create_bin_histogram.(columns, [neuron_nbins], [lower_lim], [upper_lim])
    probabilities = map(h -> h.weights/sum(h.weights), hists)
    probability = fit(Histogram, (colunms...), nbins = neuron_nbins)
end

#mutual_information(calculate_joint_probability()...)

# dt2 = get_numerical_dataset_test2()
#p2 = calculate_joint_probability_MI(dt2[1], dt2[2], (-2, 1.5), (-1, 1), 3)
#MI = mutual_information(p2...)

#=
Cálculo da informação mútua (caso 1):
-> Valor teórico calculado: I(X;T) = 0.047056227
-> Valor obtido pela aproximação: I(X;T) = 0.04700065437122643 (3 bins)
-> Valor obtido pela aproximação: I(X;T) = 0.04510842519666598 (2 bins)

Cálculo da informação mútua utilizando os valores teóricos (caso 2):
-> Valor teórico calculado: I(X;T) = 0.1685846799340601
-> Valor calculado pela aproximação: I(X;T) = 0.16843457465541672 (9 bins)
-> Valor calculado pela aproximação: I(X;T) = 0.05633855652784133 (4 bins)
=#

max_stars_repo_name: Lucas-Muniz/MNEAT

max_issues_repo_name: Lucas-Muniz/MNEAT

max_forks_repo_name: Lucas-Muniz/MNEAT


text:
#Unscented Kalman Inversion: specific structures and function definitions

"""
    Unscented{FT<:AbstractFloat, IT<:Int} <: Process

An unscented Kalman Inversion process.

# Fields

$(TYPEDFIELDS)

# Constructors

    Unscented(
        u0_mean::AbstractVector{FT},
        uu0_cov::AbstractMatrix{FT};
        α_reg::FT = 1.0,
        update_freq::IT = 1,
        modified_unscented_transform::Bool = true,
        prior_mean::Union{AbstractVector{FT}, Nothing} = nothing,
    ) where {FT <: AbstractFloat, IT <: Int}

Construct an Unscented Inversion Process.

Inputs:

  - `u0_mean`: Mean at initialization.
  - `uu0_cov`: Covariance at initialization.
  - `α_reg`: Hyperparameter controlling regularization toward the prior mean (0 < `α_reg` ≤ 1),
  default should be 1, without regulariazion.
  - `update_freq`: Set to 0 when the inverse problem is not identifiable,
  namely the inverse problem has multiple solutions, the covariance matrix
  will represent only the sensitivity of the parameters, instead of
  posterior covariance information; set to 1 (or anything > 0) when
  the inverse problem is identifiable, and the covariance matrix will
  converge to a good approximation of the posterior covariance with an
  uninformative prior.
  - `modified_unscented_transform`: Modification of the UKI quadrature given
    in Huang et al (2021).
  - `prior_mean`: Prior mean used for regularization.

$(METHODLIST)
"""
mutable struct Unscented{FT <: AbstractFloat, IT <: Int} <: Process
    "an interable of arrays of size `N_parameters` containing the mean of the parameters (in each `uki` iteration a new array of mean is added)"
    u_mean::Any  # ::Iterable{AbtractVector{FT}}
    "an iterable of arrays of size (`N_parameters x N_parameters`) containing the covariance of the parameters (in each `uki` iteration a new array of `cov` is added)"
    uu_cov::Any  # ::Iterable{AbstractMatrix{FT}}
    "an iterable of arrays of size `N_y` containing the predicted observation (in each `uki` iteration a new array of predicted observation is added)"
    obs_pred::Any # ::Iterable{AbstractVector{FT}}
    "weights in UKI"
    c_weights::AbstractVector{FT}
    mean_weights::AbstractVector{FT}
    cov_weights::AbstractVector{FT}
    "covariance of the artificial evolution error"
    Σ_ω::AbstractMatrix{FT}
    "covariance of the artificial observation error"
    Σ_ν_scale::FT
    "regularization parameter"
    α_reg::FT
    "regularization vector"
    r::AbstractVector{FT}
    "update frequency"
    update_freq::IT
    "current iteration number"
    iter::IT
end

function Unscented(
    u0_mean::AbstractVector{FT},
    uu0_cov::AbstractMatrix{FT};
    α_reg::FT = 1.0,
    update_freq::IT = 1,
    modified_unscented_transform::Bool = true,
    prior_mean::Union{AbstractVector{FT}, Nothing} = nothing,
) where {FT <: AbstractFloat, IT <: Int}

    N_par = size(u0_mean, 1)
    # ensemble size
    N_ens = 2 * N_par + 1

    c_weights = zeros(FT, N_par)
    mean_weights = zeros(FT, N_ens)
    cov_weights = zeros(FT, N_ens)

    # set parameters λ, α
    α = min(sqrt(4 / N_par), 1.0)
    λ = α^2 * N_par - N_par

    c_weights[1:N_par] .= sqrt(N_par + λ)
    mean_weights[1] = λ / (N_par + λ)
    mean_weights[2:N_ens] .= 1 / (2 * (N_par + λ))
    cov_weights[1] = λ / (N_par + λ) + 1 - α^2 + 2.0
    cov_weights[2:N_ens] .= 1 / (2 * (N_par + λ))

    if modified_unscented_transform
        mean_weights[1] = 1.0
        mean_weights[2:N_ens] .= 0.0
    end

    u_mean = Vector{FT}[]  # array of Vector{FT}'s
    push!(u_mean, u0_mean) # insert parameters at end of array (in this case just 1st entry)
    uu_cov = Matrix{FT}[]  # array of Matrix{FT}'s
    push!(uu_cov, uu0_cov) # insert parameters at end of array (in this case just 1st entry)

    obs_pred = Vector{FT}[]  # array of Vector{FT}'s

    Σ_ω = (2 - α_reg^2) * uu0_cov
    Σ_ν_scale = 2.0

    r = isnothing(prior_mean) ? u0_mean : prior_mean
    iter = 0

    Unscented(
        u_mean,
        uu_cov,
        obs_pred,
        c_weights,
        mean_weights,
        cov_weights,
        Σ_ω,
        Σ_ν_scale,
        α_reg,
        r,
        update_freq,
        iter,
    )
end

# outer constructors
function EnsembleKalmanProcess(
    obs_mean::AbstractVector{FT},
    obs_noise_cov::Union{AbstractMatrix{FT}, UniformScaling{FT}},
    process::Unscented{FT, IT};
    Δt = FT(1),
    rng::AbstractRNG = Random.GLOBAL_RNG,
    failure_handler_method::FM = IgnoreFailures(),
    localization_method::LM = NoLocalization(),
) where {FT <: AbstractFloat, IT <: Int, FM <: FailureHandlingMethod, LM <: LocalizationMethod}

    #initial parameters stored as columns
    init_params = [DataContainer(update_ensemble_prediction!(process, Δt), data_are_columns = true)]
    # Dimensionality
    N_par = length(process.u_mean[1])
    N_obs = length(obs_mean)
    # ensemble size
    N_ens = 2N_par + 1 #stored with data as columns
    #store for model evaluations
    g = []
    # error store
    err = FT[]
    # timestep store
    Δt = Array([Δt])
    # failure handler
    fh = FailureHandler(process, failure_handler_method)
    # localizer
    loc = Localizer(localization_method, N_par, N_obs, FT)

    EnsembleKalmanProcess{FT, IT, Unscented}(
        init_params,
        obs_mean,
        obs_noise_cov,
        N_ens,
        g,
        err,
        Δt,
        process,
        rng,
        fh,
        loc,
    )
end

function FailureHandler(process::Unscented, method::IgnoreFailures)
    function failsafe_update(uki, u, g, failed_ens)
        #perform analysis on the model runs
        update_ensemble_analysis!(uki, u, g)
        #perform new prediction output to model parameters u_p
        u = update_ensemble_prediction!(uki.process, uki.Δt[end])
        return u
    end
    return FailureHandler{Unscented, IgnoreFailures}(failsafe_update)
end

"""
    FailureHandler(process::Unscented, method::SampleSuccGauss)

Provides a failsafe update that
 - computes all means and covariances over the successful sigma points,
 - rescales the mean weights and the off-center covariance weights of the
    successful particles to sum to the same value as the original weight sums.
"""
function FailureHandler(process::Unscented, method::SampleSuccGauss)
    function succ_gauss_analysis!(uki, u_p, g, failed_ens)

        obs_mean = uki.obs_mean
        Σ_ν = uki.process.Σ_ν_scale * uki.obs_noise_cov
        successful_ens = filter(x -> !(x in failed_ens), collect(1:size(g, 2)))

        ############# Prediction step
        u_p_mean = construct_successful_mean(uki, u_p, successful_ens)
        uu_p_cov = construct_successful_cov(uki, u_p, u_p_mean, successful_ens)

        ###########  Analysis step
        g_mean = construct_successful_mean(uki, g, successful_ens)
        gg_cov = construct_successful_cov(uki, g, g_mean, successful_ens) + Σ_ν / uki.Δt[end]
        ug_cov = construct_successful_cov(uki, u_p, u_p_mean, g, g_mean, successful_ens)

        # Localize
        ug_cov = ug_cov .* uki.localizer.kernel

        tmp = ug_cov / gg_cov

        u_mean = u_p_mean + tmp * (obs_mean - g_mean)
        uu_cov = uu_p_cov - tmp * ug_cov'

        ########### Save results
        push!(uki.process.obs_pred, g_mean) # N_ens x N_data
        push!(uki.process.u_mean, u_mean) # N_ens x N_params
        push!(uki.process.uu_cov, uu_cov) # N_ens x N_data

        push!(uki.g, DataContainer(g, data_are_columns = true))

        compute_error!(uki)
    end
    function failsafe_update(uki, u, g, failed_ens)
        #perform analysis on the model runs
        succ_gauss_analysis!(uki, u, g, failed_ens)
        #perform new prediction output to model parameters u_p
        u = update_ensemble_prediction!(uki.process, uki.Δt[end])
        return u
    end
    return FailureHandler{Unscented, SampleSuccGauss}(failsafe_update)
end

"""
    construct_sigma_ensemble(
        process::Unscented,
        x_mean::Array{FT},
        x_cov::AbstractMatrix{FT},
    ) where {FT <: AbstractFloat, IT <: Int}

Construct the sigma ensemble based on the mean `x_mean` and covariance `x_cov`.
"""
function construct_sigma_ensemble(
    process::Unscented,
    x_mean::AbstractVector{FT},
    x_cov::AbstractMatrix{FT},
) where {FT <: AbstractFloat, IT <: Int}

    N_x = size(x_mean, 1)

    c_weights = process.c_weights

    # compute cholesky factor L of x_cov
    local chol_xx_cov
    try
        chol_xx_cov = cholesky(Hermitian(x_cov)).L
    catch
        _, S, Ut = svd(x_cov)
        # find the first singular value that is smaller than 1e-8
        ind_0 = searchsortedfirst(S, 1e-8, rev = true)
        S[ind_0:end] .= S[ind_0 - 1]
        chol_xx_cov = (qr(sqrt.(S) .* Ut).R)'
    end

    x = zeros(FT, N_x, 2 * N_x + 1)
    x[:, 1] = x_mean
    for i in 1:N_x
        x[:, i + 1] = x_mean + c_weights[i] * chol_xx_cov[:, i]
        x[:, i + 1 + N_x] = x_mean - c_weights[i] * chol_xx_cov[:, i]
    end

    return x
end

"""
construct_mean `x_mean` from ensemble `x`.
"""
function construct_mean(
    uki::EnsembleKalmanProcess{FT, IT, Unscented},
    x::AbstractVecOrMat{FT};
    mean_weights = uki.process.mean_weights,
) where {FT <: AbstractFloat, IT <: Int}

    if isa(x, AbstractMatrix{FT})
        @assert size(x, 2) == length(mean_weights)
        return Array((mean_weights' * x')')
    else
        @assert length(mean_weights) == length(x)
        return mean_weights' * x
    end
end

"""
    construct_successful_mean(
        uki::EnsembleKalmanProcess{FT, IT, Unscented},
        x::AbstractVecOrMat{FT},
        successful_indices::Union{AbstractVector{IT}, AbstractVector{Any}},
    ) where {FT <: AbstractFloat, IT <: Int}

Constructs mean over successful particles by rescaling the quadrature
weights over the successful particles. If the central particle fails
in a modified unscented transform, the mean is computed as the
ensemble mean over all successful particles.
"""
function construct_successful_mean(
    uki::EnsembleKalmanProcess{FT, IT, Unscented},
    x::AbstractVecOrMat{FT},
    successful_indices::Union{AbstractVector{IT}, AbstractVector{Any}},
) where {FT <: AbstractFloat, IT <: Int}

    mean_weights = deepcopy(uki.process.mean_weights)
    # Check if modified
    if sum(mean_weights[2:end]) ≈ 0 && !(1 in successful_indices)
        mean_weights .= 1 / length(successful_indices)
    else
        mean_weights = mean_weights ./ sum(mean_weights[successful_indices])
    end
    x_succ = isa(x, AbstractMatrix) ? x[:, successful_indices] : x[successful_indices]
    return construct_mean(uki, x_succ; mean_weights = mean_weights[successful_indices])
end

"""
construct_cov `xx_cov` from ensemble `x` and mean `x_mean`.
"""
function construct_cov(
    uki::EnsembleKalmanProcess{FT, IT, Unscented},
    x::AbstractVecOrMat{FT},
    x_mean::Union{FT, AbstractVector{FT}, Nothing} = nothing;
    cov_weights = uki.process.cov_weights,
) where {FT <: AbstractFloat, IT <: Int}

    x_mean = isnothing(x_mean) ? construct_mean(uki, x) : x_mean

    if isa(x, AbstractMatrix{FT})
        @assert isa(x_mean, AbstractVector{FT})
        N_x, N_ens = size(x)
        xx_cov = zeros(FT, N_x, N_x)

        for i in 1:N_ens
            xx_cov .+= cov_weights[i] * (x[:, i] - x_mean) * (x[:, i] - x_mean)'
        end
    else
        @assert isa(x_mean, FT)
        N_ens = length(x)
        xx_cov = FT(0)

        for i in 1:N_ens
            xx_cov += cov_weights[i] * (x[i] - x_mean) * (x[i] - x_mean)
        end
    end
    return xx_cov
end

"""
    construct_successful_cov(
        uki::EnsembleKalmanProcess{FT, IT, Unscented},
        x::AbstractVecOrMat{FT},
        x_mean::Union{AbstractVector{FT}, FT},
        successful_indices::Union{AbstractVector{IT}, AbstractVector{Any}},
    ) where {FT <: AbstractFloat, IT <: Int}

Constructs variance of `x` over successful particles by rescaling the
off-center weights over the successful off-center particles.
"""
function construct_successful_cov(
    uki::EnsembleKalmanProcess{FT, IT, Unscented},
    x::AbstractVecOrMat{FT},
    x_mean::Union{FT, AbstractVector{FT}, Nothing},
    successful_indices::Union{AbstractVector{IT}, AbstractVector{Any}},
) where {FT <: AbstractFloat, IT <: Int}

    cov_weights = deepcopy(uki.process.cov_weights)

    # Rescale non-center sigma weights to sum to original value
    orig_weight_sum = sum(cov_weights[2:end])
    sum_indices = filter(x -> x > 1, successful_indices)
    succ_weight_sum = sum(cov_weights[sum_indices])
    cov_weights[2:end] = cov_weights[2:end] .* (orig_weight_sum / succ_weight_sum)

    x_succ = isa(x, AbstractMatrix) ? x[:, successful_indices] : x[successful_indices]
    return construct_cov(uki, x_succ, x_mean; cov_weights = cov_weights[successful_indices])
end

"""
construct_cov `xy_cov` from ensemble x and mean `x_mean`, ensemble `obs_mean` and mean `y_mean`.
"""
function construct_cov(
    uki::EnsembleKalmanProcess{FT, IT, Unscented},
    x::AbstractMatrix{FT},
    x_mean::AbstractVector{FT},
    obs_mean::AbstractMatrix{FT},
    y_mean::AbstractVector{FT};
    cov_weights = uki.process.cov_weights,
) where {FT <: AbstractFloat, IT <: Int, P <: Process}

    N_x, N_ens = size(x)
    N_y = length(y_mean)
    xy_cov = zeros(FT, N_x, N_y)

    for i in 1:N_ens
        xy_cov .+= cov_weights[i] * (x[:, i] - x_mean) * (obs_mean[:, i] - y_mean)'
    end

    return xy_cov
end

"""
    construct_successful_cov(
        uki::EnsembleKalmanProcess{FT, IT, Unscented},
        x::AbstractMatrix{FT},
        x_mean::AbstractArray{FT},
        obs_mean::AbstractMatrix{FT},
        y_mean::AbstractArray{FT},
        successful_indices::Union{AbstractVector{IT}, AbstractVector{Any}},
    ) where {FT <: AbstractFloat, IT <: Int}

Constructs covariance of `x` and `obs_mean - y_mean` over successful particles by rescaling
the off-center weights over the successful off-center particles.
"""
function construct_successful_cov(
    uki::EnsembleKalmanProcess{FT, IT, Unscented},
    x::AbstractMatrix{FT},
    x_mean::AbstractVector{FT},
    obs_mean::AbstractMatrix{FT},
    y_mean::AbstractVector{FT},
    successful_indices::Union{AbstractVector{IT}, AbstractVector{Any}},
) where {FT <: AbstractFloat, IT <: Int}
    N_ens, N_x, N_y = uki.N_ens, length(x_mean), length(y_mean)

    cov_weights = deepcopy(uki.process.cov_weights)

    # Rescale non-center sigma weights to sum to original value
    orig_weight_sum = sum(cov_weights[2:end])
    sum_indices = filter(x -> x > 1, successful_indices)
    succ_weight_sum = sum(cov_weights[sum_indices])
    cov_weights[2:end] = cov_weights[2:end] .* (orig_weight_sum / succ_weight_sum)

    x_succ = isa(x, AbstractMatrix) ? x[:, successful_indices] : x[successful_indices]
    obs_mean_succ = isa(x, AbstractMatrix) ? obs_mean[:, successful_indices] : obs_mean[successful_indices]
    return construct_cov(uki, x_succ, x_mean, obs_mean_succ, y_mean; cov_weights = cov_weights[successful_indices])
end

"""
uki prediction step : generate sigma points
"""
function update_ensemble_prediction!(process::Unscented, Δt::FT) where {FT <: AbstractFloat}

    process.iter += 1
    # update evolution covariance matrix
    if process.update_freq > 0 && process.iter % process.update_freq == 0
        process.Σ_ω = (2 - process.α_reg^2) * process.uu_cov[end]
    end

    u_mean = process.u_mean[end]
    uu_cov = process.uu_cov[end]

    α_reg = process.α_reg
    r = process.r
    Σ_ω = process.Σ_ω

    N_par = length(process.u_mean[1])
    ############# Prediction step:

    u_p_mean = α_reg * u_mean + (1 - α_reg) * r
    uu_p_cov = α_reg^2 * uu_cov + Σ_ω * Δt

    ############ Generate sigma points
    u_p = construct_sigma_ensemble(process, u_p_mean, uu_p_cov)
    return u_p
end

"""
uki analysis step
g is the predicted observations  `Ny x N_ens` matrix
"""
function update_ensemble_analysis!(
    uki::EnsembleKalmanProcess{FT, IT, Unscented},
    u_p::AbstractMatrix{FT},
    g::AbstractMatrix{FT},
) where {FT <: AbstractFloat, IT <: Int}

    obs_mean = uki.obs_mean
    Σ_ν = uki.process.Σ_ν_scale * uki.obs_noise_cov

    ############# Prediction step:

    u_p_mean = construct_mean(uki, u_p)
    uu_p_cov = construct_cov(uki, u_p, u_p_mean)

    ###########  Analysis step

    g_mean = construct_mean(uki, g)
    gg_cov = construct_cov(uki, g, g_mean) + Σ_ν / uki.Δt[end]
    ug_cov = construct_cov(uki, u_p, u_p_mean, g, g_mean)

    # Localize
    ug_cov = ug_cov .* uki.localizer.kernel

    tmp = ug_cov / gg_cov

    u_mean = u_p_mean + tmp * (obs_mean - g_mean)
    uu_cov = uu_p_cov - tmp * ug_cov'

    ########### Save results
    push!(uki.process.obs_pred, g_mean) # N_ens x N_data
    push!(uki.process.u_mean, u_mean) # N_ens x N_params
    push!(uki.process.uu_cov, uu_cov) # N_ens x N_data

    push!(uki.g, DataContainer(g, data_are_columns = true))

    compute_error!(uki)
end

"""
    update_ensemble!(
        uki::EnsembleKalmanProcess{FT, IT, Unscented},
        g_in::AbstractMatrix{FT};
        Δt_new = nothing,
        failed_ens = nothing,
    ) where {FT <: AbstractFloat, IT <: Int}

Updates the ensemble according to an Unscented process.

Inputs:
 - uki :: The EnsembleKalmanProcess to update.
 - g_in :: Model outputs, they need to be stored as a `N_obs × N_ens` array (i.e data are columms).
 - Δt_new :: Time step to be used in the current update.
 - failed_ens :: Indices of failed particles. If nothing, failures are computed as columns of `g`
    with NaN entries.
"""
function update_ensemble!(
    uki::EnsembleKalmanProcess{FT, IT, Unscented},
    g_in::AbstractMatrix{FT};
    Δt_new = nothing,
    failed_ens = nothing,
) where {FT <: AbstractFloat, IT <: Int}
    #catch works when g_in non-square
    if !(size(g_in)[2] == uki.N_ens)
        throw(
            DimensionMismatch(
                "ensemble size in EnsembleKalmanProcess and g_in do not match, try transposing or check ensemble size",
            ),
        )
    end

    u_p_old = get_u_final(uki)

    set_Δt!(uki, Δt_new)
    fh = uki.failure_handler

    if isnothing(failed_ens)
        _, failed_ens = split_indices_by_success(g_in)
    end
    if !isempty(failed_ens)
        @info "$(length(failed_ens)) particle failure(s) detected. Handler used: $(nameof(typeof(fh).parameters[2]))."
    end

    u_p = fh.failsafe_update(uki, u_p_old, g_in, failed_ens)

    push!(uki.u, DataContainer(u_p, data_are_columns = true))

    return u_p
end

function get_u_mean_final(uki::EnsembleKalmanProcess{FT, IT, Unscented}) where {FT <: AbstractFloat, IT <: Int}
    return uki.process.u_mean[end]
end

function compute_error!(uki::EnsembleKalmanProcess{FT, IT, Unscented}) where {FT <: AbstractFloat, IT <: Int}
    mean_g = uki.process.obs_pred[end]
    diff = uki.obs_mean - mean_g
    X = uki.obs_noise_cov \ diff # diff: column vector
    newerr = dot(diff, X)
    push!(uki.err, newerr)
end

function get_error(uki::EnsembleKalmanProcess{FT, IT, Unscented}) where {FT <: AbstractFloat, IT <: Int}
    return uki.err
end

function Gaussian_2d(
    u_mean::AbstractVector{FT},
    uu_cov::AbstractMatrix{FT},
    Nx::IT,
    Ny::IT;
    xx = nothing,
    yy = nothing,
) where {FT <: AbstractFloat, IT <: Int}
    # 2d Gaussian plot
    u_range = [min(5 * sqrt(uu_cov[1, 1]), 5); min(5 * sqrt(uu_cov[2, 2]), 5)]

    if xx === nothing
        xx = Array(LinRange(u_mean[1] - u_range[1], u_mean[1] + u_range[1], Nx))
    end
    if yy == nothing
        yy = Array(LinRange(u_mean[2] - u_range[2], u_mean[2] + u_range[2], Ny))
    end
    X, Y = repeat(xx, 1, Ny), Array(repeat(yy, 1, Nx)')
    Z = zeros(FT, Nx, Ny)

    det_uu_cov = det(uu_cov)

    for ix in 1:Nx
        for iy in 1:Ny
            Δxy = [xx[ix] - u_mean[1]; yy[iy] - u_mean[2]]
            Z[ix, iy] = exp(-0.5 * (Δxy' / uu_cov * Δxy)) / (2 * pi * sqrt(det_uu_cov))
        end
    end

    return xx, yy, Z
end


text:
### TANGENT OPERATORS ###

"""
    compute_stress_tangent(ϵ, r::Rheology{T,::Nothing,::Nothing,::Elasticity,::DruckerPrager}, s::MaterialState ; elastic_only = false) where {T,D,V<:Nothing,E::Elasticity,P::DruckerPrager}

Computes the consistent tangent operator associated with a Drucker-Prager return mapping to the smooth part of the yield function cone based on current strain `ϵ`. Updates the material states and returns the corrected stress as well as the fourth order consistent tangent tensorial operator.
"""
function compute_stress_tangent_old(ϵ,
                                r::Rheology{T,Nothing,Nothing,E,P},
                                s::MaterialState,
                                clock::Clock;
                                noplast = false) where {T,E<:Elasticity,P<:DruckerPrager}

    # unpack elastic stiffness tensor
    Dᵉ = r.elasticity.Dᵉ

    ##### Evaluation of trial values
    ϵᵉ_trial = ϵ - s.ϵᵖ # trial-elastic_strain
    σ_trial = Dᵉ ⊡ ϵᵉ_trial # trial-stress

    if noplast # during first newton iteration to homogenize strain field
        s.temp_ϵ = ϵ
        s.temp_σ = σ_trial
        return s.temp_σ, Dᵉ
    end

    # unpack plastic properties
    C = r.plasticity.C
    H = r.plasticity.H
    η = r.plasticity.η
    ξ = r.plasticity.ξ

    p_trial = 1/3 * tr(σ_trial) # trial pressure, negative in compression
    s_trial = dev(σ_trial) # trial deviatoric stress
    τ_trial = get_τ(s_trial,r.plasticity) # effetive trial-stress (2nd invariant of the deviatoric stress tensor)
    τ_yield = -η*p_trial + ξ*(C + H*s.ϵ̅ᵖ)
    F_trial  = τ_trial - τ_yield # Trial-value of the yield surface

    if F_trial < 0.0 # elastic loading
        s.temp_ϵ = ϵ
        s.temp_σ = σ_trial
        return s.temp_σ, Dᵉ
    else # plastic loading
        # unpack
        G = r.elasticity.G
        K = r.elasticity.K
        ηᵛᵖ = r.plasticity.ηᵛᵖ
        η̅ = r.plasticity.η̅

        # Plastic flow potential gradient
        I2D = SymmetricTensor{2,3}(δ) # second order identity tensor
        ∇Q = ((1/(2*τ_trial)) * s_trial + η̅/3 * I2D)
        #∇Q_scalar = sqrt(2/3 * ∇Q ⊡ ∇Q)

        ##### Compute incremental plastic multiplier
        Δλ_factor = 1/(G + K*η*η̅ + ηᵛᵖ/clock.Δt + ξ^2*H) #ξ*H*∇Q_scalar
        Δλ = Δλ_factor * F_trial

        ##### Compute unit deviatoric flow vector
        ϵᵉdev_trial = dev(ϵᵉ_trial) #ϵᵉ_trial .- [ϵᵉvol_d3_trial, ϵᵉvol_d3_trial, ϵᵉvol_d3_trial, 0.0]
        norm_ϵᵉdev_t = norm(ϵᵉdev_trial)#sqrt(ϵᵉdev_trial[1]^2 + ϵᵉdev_trial[2]^2 + ϵᵉdev_trial[3]^2 + 2*ϵᵉdev_trial[4]^2) # get the norm of the deviatoric elastic trial strain

        # prevent division by zero if ϵᵉdev_trial is zero valued
        if norm_ϵᵉdev_t != 0.0
            inv_norm = 1.0/norm_ϵᵉdev_t
        else
            inv_norm = 0.0
        end
        uni_dev = ϵᵉdev_trial*inv_norm # unit deviatoric flow vector

        ##### assemble tangent
        Isymdev = SymmetricTensor{4,3}(Isymdev_func) # fourth order deviatoric symmetric tensor
        A = Δλ_factor
        A_fact = 2G * (1.0 - Δλ/(sqrt(2)*norm_ϵᵉdev_t))
        A_factd3 = A_fact/3
        B_fact = 2G * (Δλ/(sqrt(2)*norm_ϵᵉdev_t) - G*A)
        C_fact = -sqrt(2)*G*A*K
        D_fact = K*(1.0 - K*η*η̅*A)
        D = A_fact * Isymdev +
            B_fact * uni_dev ⊗ uni_dev +
            C_fact * (η * uni_dev⊗I2D + η̅ * I2D⊗uni_dev) +
            D_fact * I2D ⊗ I2D #TODO check whether this form of the D term or the one from the book's code is the good one => (D_fact - A_factd3) or just D_fact.

        Δϵᵖ = Δλ*∇Q # plastic strain increment
        Δϵ̅ᵖ = Δλ*ξ#Δλ*∇Q_scalar # accumulated plastic strain increment
        s.temp_ϵ = ϵ
        s.temp_ϵᵖ = s.ϵᵖ + Δϵᵖ # plastic strain
        s.temp_ϵ̅ᵖ = s.ϵ̅ᵖ + Δϵ̅ᵖ # accumulated plastic strain
        s.temp_σ = σ_trial - Δλ * Dᵉ⊡∇Q
        return s.temp_σ, D
    end
end

#TODO put back this function if the version iterating over Δλ doesn't work.
function compute_stress_tangent_prev(ϵ,
                                r::Rheology{T,Nothing,Nothing,E,P},
                                s::MaterialState,
                                clock::Clock;
                                noplast = false) where {T,E<:Elasticity,P<:DruckerPrager}

    # unpack elastic stiffness tensor
    Dᵉ = r.elasticity.Dᵉ

    ##### Evaluation of trial values
    ϵᵉ_trial = ϵ - s.ϵᵖ # trial-elastic_strain
    σ_trial = Dᵉ ⊡ ϵᵉ_trial # trial-stress

    if noplast # during first newton iteration to homogenize strain field
        s.temp_ϵ = ϵ
        s.temp_σ = σ_trial
        return s.temp_σ, Dᵉ
    end

    # unpack plastic properties
    C = r.plasticity.C
    H = r.plasticity.H
    η = r.plasticity.η
    ξ = r.plasticity.ξ

    p_trial = 1/3 * tr(σ_trial) # trial pressure, negative in compression
    s_trial = dev(σ_trial) # trial deviatoric stress
    τ_trial = get_τ(s_trial,r.plasticity) # effetive trial-stress (2nd invariant of the deviatoric stress tensor)
    τ_yield = -η*p_trial + ξ*(C + H*s.ϵ̅ᵖ)
    F_trial  = τ_trial - τ_yield # Trial-value of the yield surface

    if F_trial < 0.0 # elastic loading
        s.temp_ϵ = ϵ
        s.temp_σ = σ_trial
        return s.temp_σ, Dᵉ
    else # plastic loading
        # unpack
        G = r.elasticity.G
        K = r.elasticity.K
        ηᵛᵖ = r.plasticity.ηᵛᵖ
        η̅ = r.plasticity.η̅

        # Plastic flow potential gradient
        I2D = SymmetricTensor{2,3}(δ) # second order identity tensor
        ∇Q = ((1/(2*τ_trial)) * s_trial + η̅/3 * I2D)
        #∇Q_scalar = sqrt(2/3 * ∇Q ⊡ ∇Q)

        ##### Compute incremental plastic multiplier
        Δλ_factor = 1/(G + K*η*η̅ + ηᵛᵖ/clock.Δt + ξ^2*H) # + ηᵛᵖ/clock.Δt pour visco-plasticité
        Δλ = Δλ_factor * F_trial

        ##### Test the validity of the return mapping to the smooth portion of the cone
        is_valid = (τ_trial - G*Δλ >= 0)

        ##### Use relevant return mapping algorithm

        if is_valid # return to the smooth portion of the cone

            ##### Compute unit deviatoric flow vector
            ϵᵉdev_trial = dev(ϵᵉ_trial) #ϵᵉ_trial .- [ϵᵉvol_d3_trial, ϵᵉvol_d3_trial, ϵᵉvol_d3_trial, 0.0]
            norm_ϵᵉdev_t = norm(ϵᵉdev_trial)#sqrt(ϵᵉdev_trial[1]^2 + ϵᵉdev_trial[2]^2 + ϵᵉdev_trial[3]^2 + 2*ϵᵉdev_trial[4]^2) # get the norm of the deviatoric elastic trial strain

            # prevent division by zero if ϵᵉdev_trial is zero valued
            if norm_ϵᵉdev_t != 0.0
                inv_norm = 1.0/norm_ϵᵉdev_t
            else
                inv_norm = 0.0
            end
            uni_dev = ϵᵉdev_trial*inv_norm # unit deviatoric flow tensor

            ##### assemble tangent
            Isymdev = SymmetricTensor{4,3}(Isymdev_func) # fourth order deviatoric symmetric tensor
            A = Δλ_factor
            A_fact = 2G * (1.0 - Δλ/(sqrt(2)*norm_ϵᵉdev_t))
            A_factd3 = A_fact/3
            B_fact = 2G * (Δλ/(sqrt(2)*norm_ϵᵉdev_t) - G*A)
            C_fact = -sqrt(2)*G*A*K
            D_fact = K*(1.0 - K*η*η̅*A)
            D = A_fact * Isymdev +
                B_fact * uni_dev ⊗ uni_dev +
                C_fact * (η * uni_dev⊗I2D + η̅ * I2D⊗uni_dev) +
                D_fact * I2D ⊗ I2D #TODO check whether this form of the D term or the one from the book's code is the good one => (D_fact - A_factd3) or just D_fact.

            Δϵᵖ = Δλ*∇Q # plastic strain increment
            Δϵ̅ᵖ = Δλ*ξ#Δλ*∇Q_scalar # accumulated plastic strain increment
            s.temp_σ = σ_trial - Δλ * Dᵉ⊡∇Q

        else # return to the apex
            α = ξ/η
            β = ξ/η̅
            Δϵᵖ_vol = (p_trial - (C + H*s.ϵ̅ᵖ)*β) / (H*α*β + K)
            Δϵᵖ = 1/3 * Δϵᵖ_vol * I2D
            Δϵ̅ᵖ = α * Δϵᵖ_vol
            s.temp_σ = (p_trial - K*Δϵᵖ_vol)*I2D

            D = K * (1 - K/(K + α*β*H)) * I2D ⊗ I2D
        end

        ### update strains
        s.temp_ϵ = ϵ
        s.temp_ϵᵖ = s.ϵᵖ + Δϵᵖ # plastic strain
        s.temp_ϵ̅ᵖ = s.ϵ̅ᵖ + Δϵ̅ᵖ # accumulated plastic strain

        return s.temp_σ, D
    end
end

function compute_stress_tangent(ϵ,
                                r::Rheology{T,Nothing,Nothing,E,P},
                                s::MaterialState,
                                clock::Clock;
                                noplast = false) where {T,E<:Elasticity,P<:DruckerPrager}

    # unpack elastic stiffness tensor
    Dᵉ = r.elasticity.Dᵉ

    ##### Evaluation of trial values
    ϵᵉ_trial = ϵ - s.ϵᵖ # trial-elastic_strain
    σ_trial = Dᵉ ⊡ ϵᵉ_trial # trial-stress

    if noplast # during first newton iteration to homogenize strain field
        s.temp_ϵ = ϵ
        s.temp_σ = σ_trial
        return s.temp_σ, Dᵉ
    end

    # unpack plastic properties
    C = r.plasticity.C
    H = r.plasticity.H
    η = r.plasticity.η
    ξ = r.plasticity.ξ
    ηᵛᵖ = r.plasticity.ηᵛᵖ
    η̅ = r.plasticity.η̅
    # unpack elastic properties
    G = r.elasticity.G
    K = r.elasticity.K

    p_trial = 1/3 * tr(σ_trial) # trial pressure, negative in compression
    s_trial = dev(σ_trial) # trial deviatoric stress
    τ_trial = get_τ(s_trial,r.plasticity) # effetive trial-stress (2nd invariant of the deviatoric stress tensor)
    τ_yield = -η*p_trial + ξ*(C + H*s.ϵ̅ᵖ)
    F_trial_0  = τ_trial - τ_yield # Trial-value of the yield surface

    converged_Δλ = false
    Δλ₀ = 0.0
    Δλ = 0.0
    Δλ_factor = 0.0
    F_trial = 0.0

    while !converged_Δλ # iterate until Δλ converge
        # compute F_trial with current Δλ₀
        F_trial = F_trial_0 - ηᵛᵖ*Δλ₀/clock.Δt
        ##### Compute incremental plastic multiplier
        Δλ_factor = 1/(G + K*η*η̅ + ηᵛᵖ/clock.Δt + ξ^2*H) # + ηᵛᵖ/clock.Δt pour visco-plasticité
        Δλ = Δλ_factor * F_trial
        # Check error
        (abs(Δλ-Δλ₀)/Δλ₀ <= 1e-8) && (converged_Δλ = true)
        Δλ₀ = Δλ
    end

    if F_trial < 0.0 # elastic loading
        s.temp_ϵ = ϵ
        s.temp_σ = σ_trial
        return s.temp_σ, Dᵉ
    else # plastic loading

        # Plastic flow potential gradient
        I2D = SymmetricTensor{2,3}(δ) # second order identity tensor
        ∇Q = ((1/(2*τ_trial)) * s_trial + η̅/3 * I2D)
        #∇Q_scalar = sqrt(2/3 * ∇Q ⊡ ∇Q)

        ##### Test the validity of the return mapping to the smooth portion of the cone
        is_valid = (τ_trial - G*Δλ >= 0)

        ##### Use relevant return mapping algorithm

        if is_valid # return to the smooth portion of the cone

            ##### Compute unit deviatoric flow vector
            ϵᵉdev_trial = dev(ϵᵉ_trial) #ϵᵉ_trial .- [ϵᵉvol_d3_trial, ϵᵉvol_d3_trial, ϵᵉvol_d3_trial, 0.0]
            norm_ϵᵉdev_t = norm(ϵᵉdev_trial)#sqrt(ϵᵉdev_trial[1]^2 + ϵᵉdev_trial[2]^2 + ϵᵉdev_trial[3]^2 + 2*ϵᵉdev_trial[4]^2) # get the norm of the deviatoric elastic trial strain

            # prevent division by zero if ϵᵉdev_trial is zero valued
            if norm_ϵᵉdev_t != 0.0
                inv_norm = 1.0/norm_ϵᵉdev_t
            else
                inv_norm = 0.0
            end
            uni_dev = ϵᵉdev_trial*inv_norm # unit deviatoric flow tensor

            ##### assemble tangent
            Isymdev = SymmetricTensor{4,3}(Isymdev_func) # fourth order deviatoric symmetric tensor
            A = Δλ_factor
            A_fact = 2G * (1.0 - Δλ/(sqrt(2)*norm_ϵᵉdev_t))
            A_factd3 = A_fact/3
            B_fact = 2G * (Δλ/(sqrt(2)*norm_ϵᵉdev_t) - G*A)
            C_fact = -sqrt(2)*G*A*K
            D_fact = K*(1.0 - K*η*η̅*A)
            D = A_fact * Isymdev +
                B_fact * uni_dev ⊗ uni_dev +
                C_fact * (η * uni_dev⊗I2D + η̅ * I2D⊗uni_dev) +
                D_fact * I2D ⊗ I2D #TODO check whether this form of the D term or the one from the book's code is the good one => (D_fact - A_factd3) or just D_fact.

            Δϵᵖ = Δλ*∇Q # plastic strain increment
            Δϵ̅ᵖ = Δλ*ξ#Δλ*∇Q_scalar # accumulated plastic strain increment
            s.temp_σ = σ_trial - Δλ * Dᵉ⊡∇Q

        else # return to the apex
            α = ξ/η
            β = ξ/η̅
            Δϵᵖ_vol = (p_trial - (C + H*s.ϵ̅ᵖ)*β) / (H*α*β + K)
            Δϵᵖ = 1/3 * Δϵᵖ_vol * I2D
            Δϵ̅ᵖ = α * Δϵᵖ_vol
            s.temp_σ = (p_trial - K*Δϵᵖ_vol)*I2D

            D = K * (1 - K/(K + α*β*H)) * I2D ⊗ I2D
        end

        ### update strains
        s.temp_ϵ = ϵ
        s.temp_ϵᵖ = s.ϵᵖ + Δϵᵖ # plastic strain
        s.temp_ϵ̅ᵖ = s.ϵ̅ᵖ + Δϵ̅ᵖ # accumulated plastic strain

        return s.temp_σ, D
    end
end

function compute_stress_tangent(ϵ,
                                r::Rheology{T,Nothing,V,E,Nothing},
                                s::MaterialState,
                                clock::Clock;
                                noplast = false) where {T,V<:Viscosity,E<:Elasticity}

    #@info "good stress tangent function"

    ##### viscoelastic tangent operator
    Gᵛᵉ = 1/(1/r.elasticity.G + clock.Δt/r.viscosity.η)
    λᵛᵉ = λ_from_KG(r.elasticity.K,Gᵛᵉ)
    Dᵛᵉ = get_elastic_stiffness_tensor(Gᵛᵉ,λᵛᵉ)

    ##### stress and strain update
    Gratio = Gᵛᵉ/r.elasticity.G
    I2D = SymmetricTensor{2,3}(δ) # second order identity tensor
    σ⁰ = (1/3)*tr(s.σ)*I2D + Gratio*dev(s.σ)
    s.temp_σ = σ⁰ + Dᵛᵉ ⊡ (ϵ - s.ϵ)
    s.temp_ϵ = ϵ

    return s.temp_σ, Dᵛᵉ
end

function compute_stress_tangent(ϵ,
                                r::Rheology{T,Nothing,V,E,P},
                                s::MaterialState,
                                clock::Clock;
                                noplast = false) where {T,V<:Viscosity,E<:Elasticity,P<:DruckerPrager}

    ##### viscoelastic tangent operator
    Gᵛᵉ = 1/(1/r.elasticity.G + clock.Δt/r.viscosity.η)
    λᵛᵉ = λ_from_KG(r.elasticity.K,Gᵛᵉ)
    Dᵛᵉ = get_elastic_stiffness_tensor(Gᵛᵉ,λᵛᵉ)

    # trial stress
    Gratio = Gᵛᵉ/r.elasticity.G
    I2D = SymmetricTensor{2,3}(δ) # second order identity tensor
    σ⁰ = (1/3)*tr(s.σ)*I2D + Gratio*dev(s.σ)
    # TODO find good expression for ϵᵉ_trial
    #ϵᵉ_trial = ϵ - s.ϵᵖ # should be bad because doesn't take viscous deformation into account
    ϵᵉ_trial = ϵ - s.ϵ

    σ_trial = σ⁰ + Dᵛᵉ ⊡ (ϵ - s.ϵ)

    if noplast # during first newton iteration to homogenize strain field
        s.temp_ϵ = ϵ
        s.temp_σ = σ_trial
        return s.temp_σ, Dᵛᵉ
    end

    # unpack plastic properties
    C = r.plasticity.C
    H = r.plasticity.H
    η = r.plasticity.η
    ξ = r.plasticity.ξ

    p_trial = 1/3 * tr(σ_trial) # trial pressure, negative in compression
    s_trial = dev(σ_trial) # trial deviatoric stress
    τ_trial = get_τ(s_trial,r.plasticity) # effetive trial-stress (2nd invariant of the deviatoric stress tensor)
    τ_yield = -η*p_trial + ξ*(C + H*s.ϵ̅ᵖ)
    F_trial  = τ_trial - τ_yield # Trial-value of the yield surface

    if F_trial < 0.0 # elastic loading
        s.temp_ϵ = ϵ
        s.temp_σ = σ_trial
        return s.temp_σ, Dᵛᵉ
    else # plastic loading
        # unpack
        K = r.elasticity.K
        ηᵛᵖ = r.plasticity.ηᵛᵖ
        η̅ = r.plasticity.η̅

        # Plastic flow potential gradient
        ∇Q = ((1/(2*τ_trial)) * s_trial + η̅/3 * I2D)
        #∇Q_scalar = sqrt(2/3 * ∇Q ⊡ ∇Q)

        ##### Compute incremental plastic multiplier
        Δλ_factor = 1/(Gᵛᵉ + K*η*η̅ + ηᵛᵖ/clock.Δt + ξ^2*H) #ξ*H*∇Q_scalar
        Δλ = Δλ_factor * F_trial

        ##### Test the validity of the return mapping to the smooth portion of the cone
        is_valid = (τ_trial - Gᵛᵉ*Δλ >= 0)

        ##### Use relevant return mapping algorithm

        if is_valid # return to the smooth portion of the cone
            ##### Compute unit deviatoric flow vector
            ϵᵉdev_trial = dev(ϵᵉ_trial) #ϵᵉ_trial .- [ϵᵉvol_d3_trial, ϵᵉvol_d3_trial, ϵᵉvol_d3_trial, 0.0]
            norm_ϵᵉdev_t = norm(ϵᵉdev_trial)#sqrt(ϵᵉdev_trial[1]^2 + ϵᵉdev_trial[2]^2 + ϵᵉdev_trial[3]^2 + 2*ϵᵉdev_trial[4]^2) # get the norm of the deviatoric elastic trial strain

            # prevent division by zero if ϵᵉdev_trial is zero valued
            if norm_ϵᵉdev_t != 0.0
                inv_norm = 1.0/norm_ϵᵉdev_t
            else
                inv_norm = 0.0
            end
            uni_dev = ϵᵉdev_trial*inv_norm # unit deviatoric flow vector

            ##### assemble tangent
            Isymdev = SymmetricTensor{4,3}(Isymdev_func) # fourth order deviatoric symmetric tensor
            A = Δλ_factor
            A_fact = 2Gᵛᵉ * (1.0 - Δλ/(sqrt(2)*norm_ϵᵉdev_t))
            A_factd3 = A_fact/3
            B_fact = 2Gᵛᵉ * (Δλ/(sqrt(2)*norm_ϵᵉdev_t) - Gᵛᵉ*A)
            C_fact = -sqrt(2)*Gᵛᵉ*A*K
            D_fact = K*(1.0 - K*η*η̅*A)
            D = A_fact * Isymdev +
                B_fact * uni_dev ⊗ uni_dev +
                C_fact * (η * uni_dev⊗I2D + η̅ * I2D⊗uni_dev) +
                D_fact * I2D ⊗ I2D #TODO check whether this form of the D term or the one from the book's code is the good one => (D_fact - A_factd3) or just D_fact.

            Δϵᵖ = Δλ*∇Q # plastic strain increment
            Δϵ̅ᵖ = Δλ*ξ # accumulated plastic strain increment
            s.temp_σ = σ_trial - Δλ * Dᵛᵉ⊡∇Q
        else
            α = ξ/η
            β = ξ/η̅
            Δϵᵖ_vol = (p_trial - (C + H*s.ϵ̅ᵖ)*β) / (H*α*β + K)
            Δϵᵖ = 1/3 * Δϵᵖ_vol * I2D
            Δϵ̅ᵖ = α * Δϵᵖ_vol
            s.temp_σ = (p_trial - K*Δϵᵖ_vol)*I2D
            D = K * (1 - K/(K + α*β*H)) * I2D ⊗ I2D
        end
        s.temp_ϵ = ϵ
        s.temp_ϵᵖ = s.ϵᵖ + Δϵᵖ # plastic strain
        s.temp_ϵ̅ᵖ = s.ϵ̅ᵖ + Δϵ̅ᵖ # accumulated plastic strain
        s.temp_σ = σ_trial - Δλ * Dᵛᵉ⊡∇Q
        return s.temp_σ, D
    end
end

loss_of_convexity(u,t,integrator) = (u[1] == 1.0 ? 0 : 1) # must return zero to be triggered

function compute_damaged_stiffness_tensor(r::Rheology,ϵij,D)

    # unpack
    G = r.elasticity.G
    ν = r.elasticity.ν

    # Damage constants
    c1, c2, c3 = compute_c1c2c3(r,D)
    A, B = compute_AB(r,c1,c2,c3)
    A₁ = compute_A1(r,A)
    B₁ = compute_B1(r,B)
    Γ = compute_Γ(r,A₁,B₁)

    # strain invariants
    ϵ = tr(ϵij)
    e = dev(ϵij)
    γ = sqrt(2.0 * e ⊡ e)

    @assert !isnan(G)
    @assert !isnan(c1)
    @assert !isnan(c2)
    @assert !isnan(c3)
    @assert !isnan(A)
    @assert !isnan(B)
    @assert !isnan(A₁)
    @assert !isnan(B₁)
    @assert !isnan(Γ)
    @assert !isnan(ϵ)
    @assert !isnan(γ)
    @assert !isnan((3*(1-2ν))/(2*(1+ν)))
    @assert !isnan(A₁^2/2)
    # println("A₁*B₁*ϵ : ", A₁*B₁*ϵ)
    # println("(2*γ) : ", (2*γ))
    #@assert !isnan(A₁*B₁*ϵ/(2*γ)) # returns NaN, because 0/0

    (γ == 0) && (γ += 1e-9)#zero(typeof(ϵij))) # TODO remove redundancy
    ϵ̂ = ϵij/γ

    # get stiffness factors
    if G/Γ == 0
        Cμ = 0.0
        Cλ = 0.0
        Cσ = 0.0
        Cσσ = 0.0
    else
        Cμ = G/Γ * ( (3*(1-2ν))/(2*(1+ν)) + A₁^2/2 - A₁*B₁*ϵ/(2*γ) )
        Cλ = G/Γ * ( 3*ν/(1+ν) + B₁^2/2 - A₁^2/3 + A₁*B₁*ϵ/γ + 2A₁*B₁*ϵ^3/(9γ^3) )
        Cσ = - G/Γ * ( A₁*B₁ + 2*A₁*B₁*ϵ^2/(3*γ^2) )
        Cσσ = G/Γ * (2A₁*B₁*ϵ/γ)
    end

    if isnan(Cμ)
        println("Cμ = ", Cμ)
        @assert !isnan(G/Γ)
        println("G/Γ = ", G/Γ)
        @assert !isnan( (3*(1-2ν))/(2*(1+ν)) + A₁^2/2 - A₁*B₁*ϵ/(2*γ) )
        println("second term Cμ = ", (3*(1-2ν))/(2*(1+ν)) + A₁^2/2 - A₁*B₁*ϵ/(2*γ) )
        @assert !isnan(G/Γ * ( (3*(1-2ν))/(2*(1+ν)) + A₁^2/2 - A₁*B₁*ϵ/(2*γ) ))
    end
    if isnan(Cλ)
        println("Cλ = ", Cλ)
    end
    @assert !isnan(Cσ)
    @assert !isnan(Cσσ)
    # functional form of the stiffness tensor
    C_func(i,j,k,l) = Cμ * ( δ(k,i)*δ(l,j) + δ(l,i)*δ(k,j) ) +
                      Cλ * ( δ(i,j)*δ(k,l) ) +
                      Cσ * ( ϵ̂[i,j]*δ(k,l) + δ(i,j)*ϵ̂[k,l] ) +
                      Cσσ * (ϵ̂[i,j]*ϵ̂[k,l])

    # assemble the tensor
    return SymmetricTensor{4,3}(C_func)
end

function compute_stress_tangent(ϵij,
                                r::Rheology{T,TD,Nothing,TE,Nothing},
                                s::MaterialState,
                                clock::Clock;
                                noplast = false) where {T,TD<:Damage,TE<:Elasticity}

    # unpack elastic stiffness tensor
    Cᵉ = r.elasticity.Dᵉ

    ##### Evaluation of trial values
    ϵᵉ_trial = ϵij - s.ϵᵖ # trial-elastic_strain
    σ_trial = Cᵉ ⊡ ϵᵉ_trial # trial-stress

    if noplast # during first newton iteration to homogenize strain field
        s.temp_ϵ = ϵij
        s.temp_σ = σ_trial
        return s.temp_σ, Cᵉ
    end

    Δϵ = ϵij - s.ϵ # strain increment
    p_trial = 1/3 * tr(σ_trial) # trial pressure, negative in compression
    s_trial = dev(σ_trial) # trial deviatoric stress
    τ_trial = get_τ(s_trial,r.damage) # effective trial-stress (2nd invariant of the deviatoric stress tensor)
    KI_trial = compute_KI(r,p_trial,τ_trial,s.D) # KI at the end of the timestep if the loading is purely elastic

    if KI_trial < 0.0 # elastic loading over Δt
        s.temp_ϵ = ϵij
        s.temp_σ = σ_trial
        return s.temp_σ, Cᵉ
    else # damage growth over Δt
        # Converged stresses and strains invariants
        ϵ0 = tr(s.ϵ)
        e0 = dev(s.ϵ)
        γ0 = sqrt(2.0 * e0 ⊡ e0)

        σ0 = 1/3 * tr(s.σ)
        s0 = dev(s.σ)
        τ0 = sqrt(0.5 * s0 ⊡ s0)

        # Initial conditions
        u0 = [s.D, ϵ0, γ0, σ0, τ0]
        # Time span
        tspan = (0.0,clock.Δt)
        sub_Δt_ini = tspan[2]/100 # initial subtimestep for RK4 interation algorithm
        damage_reltol = 1e-6
        # strain invariants rates
        dϵdt = tr(Δϵ)/clock.Δt

        e = dev(ϵij)
        γ = sqrt(2.0 * e ⊡ e)
        Δγ = γ - γ0
        dγdt = Δγ/clock.Δt
        # OLD AND WRONG :
        #Δe = dev(Δϵ)
        #dγdt = sqrt(2.0 * Δe ⊡ Δe)/clock.Δt

        # ODE system parameters
        params = [r, dϵdt, dγdt]

        # Hand made RK4
        u_vec, t_vec, cohesion_loss_flag = compute_RK4_adaptative(u0,params,state_system!,tspan,sub_Δt_ini,damage_reltol)

        # TEST CODE
        #println("t_vec : ", t_vec)
        #println("D_vec : ", [u[1] for u in u_vec])
        #lineplot(t_vec,u_vec[1])

        # DiffEq solve alternatively
        # # Problem instantiation
        # state_problem = ODEProblem(state_system!,u0,tspan,params)
        #
        # # the ODE system is only solved until D reaches 1.0 => loss of convexity of the free energy function.
        # termination_callback = ContinuousCallback(loss_of_convexity,terminate!)
        # sol = OrdinaryDiffEq.solve(state_problem, Tsit5(), callback=termination_callback)

        # TODO for damaged-plastic material, the idea is to check if the integration has been done over the whole Δt. If it stopped before, it means loss of cohesion of the material. A plastic update is therefore required in that case.

        # update states
        s.temp_D = u_vec[end][1]#sol[1,end]

        if s.temp_D > s.D
            s.temp_σ = compute_σij(r,s.temp_D,ϵij) # changed from ϵᵉ_trial to ϵ

            # TEST
            σ_temp = 1/3 * tr(s.temp_σ) # trial pressure, negative in compression
            s_temp = dev(s.temp_σ) # trial deviatoric stress
            τ_temp = get_τ(s_temp,r.damage)

            #println("sigma diff : ", σ_temp - u_vec[end][4])
            #println("tau diff : ", τ_temp - u_vec[end][5])
            invCᵉ = get_elastic_compliance_tensor(r)
            Δϵᵖ = Δϵ - invCᵉ⊡(s.temp_σ - s.σ) # damage strain increment
            s.temp_ϵ = ϵij
            s.temp_ϵᵖ = s.ϵᵖ + Δϵᵖ # plastic strain
            # Construction of the tangent operator ∂σ_n/∂ϵ_n at {D = D_n+1}:
            C = compute_damaged_stiffness_tensor(r,s.ϵ,s.D) #TODO s.D or S.temp_D
            return s.temp_σ, C
        else
            s.temp_ϵ = ϵij
            s.temp_σ = σ_trial
            return s.temp_σ, Cᵉ
        end
        # EQUALITY TEST TODO investigate that
        #Δσ = C ⊡ Δϵ
        # println("Δϵ : ", Δϵ)
        # println("tangent : ", C)
        # println("Δσ from tangent : ", Δσ)
        # println("Δσ from constitutive relationships : ", s.temp_σ - s.σ)
        #Δσ = s.temp_σ - s.σ
        #println("Δσ : ", Δσ)
        #println("Δϵ : ", Δϵ)
        #C2 = Δσ ⊗ pinv(Δϵ)
        #Δσ_C2 = C2 ⊡ Δϵ
        #println("Δσ const rel : ", Δσ)
        #println("Δσ from Δσ ⊗ inv(Δϵ) : ", Δσ_C2)

        #Δϵ̅ᵖ = Δλ*ξ#Δλ*∇Q_scalar # accumulated plastic strain increment

        #s.temp_ϵ̅ᵖ = s.ϵ̅ᵖ + Δϵ̅ᵖ # accumulated plastic strain
        #s.temp_σ = σ_trial - Δλ * Dᵉ⊡∇Q
    end
end

function compute_stress_tangent(ϵij,
                                r::Rheology{T,TD,Nothing,TE,TP},
                                s::MaterialState,
                                clock::Clock;
                                noplast = false) where {T,TD<:Damage,TE<:Elasticity,TP}

    # unpack elastic stiffness tensor
    Cᵉ = r.elasticity.Dᵉ

    ##### Evaluation of trial values
    ϵᵉ_trial = ϵij - s.ϵᵖ # trial-elastic_strain
    σ_trial = Cᵉ ⊡ ϵᵉ_trial # trial-stress

    if noplast # during first newton iteration to homogenize strain field
        s.temp_ϵ = ϵij
        s.temp_σ = σ_trial
        return s.temp_σ, Cᵉ
    end

    Δϵ = ϵij - s.ϵ # strain increment
    p_trial = 1/3 * tr(σ_trial) # trial pressure, negative in compression
    s_trial = dev(σ_trial) # trial deviatoric stress
    τ_trial = get_τ(s_trial,r.damage) # effetive trial-stress (2nd invariant of the deviatoric stress tensor)
    KI_trial = compute_KI(r,p_trial,τ_trial,s.D) # KI at the end of the timestep if the loading is purely elastic

    if KI_trial < 0.0 # elastic loading over Δt
        s.temp_ϵ = ϵij
        s.temp_σ = σ_trial
        return s.temp_σ, Cᵉ
    else # damage growth over Δt
        # Converged stresses and strains invariants
        ϵ0 = tr(s.ϵ)
        e0 = dev(s.ϵ)
        γ0 = sqrt(2.0 * e0 ⊡ e0)

        σ0 = 1/3 * tr(s.σ)
        s0 = dev(s.σ)
        τ0 = sqrt(0.5 * s0 ⊡ s0)

        # Initial conditions
        u0 = [s.D, ϵ0, γ0, σ0, τ0]
        # Time span
        tspan = (0.0,clock.Δt)
        sub_Δt_ini = tspan[2]/100 # initial subtimestep for RK4 interation algorithm
        damage_reltol = 1e-6
        # strain invariants rates
        dϵdt = tr(Δϵ)/clock.Δt

        e = dev(ϵij)
        γ = sqrt(2.0 * e ⊡ e)
        Δγ = γ - γ0
        dγdt = Δγ/clock.Δt
        #Δe = dev(Δϵ)
        #dγdt = sqrt(2.0 * Δe ⊡ Δe)/clock.Δt

        # ODE system parameters
        params = [r, dϵdt, dγdt]

        # Hand made RK4
        u_vec, t_vec, cohesion_loss_flag = compute_RK4_adaptative(u0,params,state_system!,tspan,sub_Δt_ini,damage_reltol)

        # TEST CODE
        #println("t_vec : ", t_vec)
        #println("D_vec : ", [u[1] for u in u_vec])
        #lineplot(t_vec,u_vec[1])

        # DiffEq solve alternatively
        # # Problem instantiation
        # state_problem = ODEProblem(state_system!,u0,tspan,params)
        #
        # # the ODE system is only solved until D reaches 1.0 => loss of convexity of the free energy function.
        # termination_callback = ContinuousCallback(loss_of_convexity,terminate!)
        # sol = OrdinaryDiffEq.solve(state_problem, Tsit5(), callback=termination_callback)

        # TODO for damaged-plastic material, the idea is to check if the integration has been done over the whole Δt. If it stopped before, it means loss of cohesion of the material. A plastic update is therefore required in that case.

        # Handle loss of cohesion => branching to a plastic behavior
        if (TP <: Plasticity) & cohesion_loss_flag
            r_nodamage = Rheology(damage = Nothing,
                                  viscosity = r.viscosity,
                                  elasticity = r.elasticity,
                                  plasticity = r.plasticity)
                                  # TODO tester avec ou sans reset de plastic strain
            compute_stress_tangent(ϵij, r_nodamage, s, clock)
        end

        # update state
        s.temp_D = u_vec[end][1]#sol[1,end]

        if s.temp_D > s.D
            s.temp_σ = compute_σij(r,s.temp_D,ϵij) # changed from ϵᵉ_trial to ϵ

            ## TEST ##
            if s.temp_D - s.D > 0.0001
                p_temp = 1/3 * tr(s.temp_σ) # pressure, negative in compression
                s_temp = dev(s.temp_σ) # deviatoric stress
                τ_temp = get_τ(s_temp,r.damage)
                println("p diff : ", p_temp - u_vec[end][4])
                println("tau diff : ", τ_temp - u_vec[end][5])
            end
            invCᵉ = get_elastic_compliance_tensor(r)
            Δϵᵖ = Δϵ - invCᵉ⊡(s.temp_σ - s.σ) # damage strain increment
            s.temp_ϵ = ϵij
            s.temp_ϵᵖ = s.ϵᵖ + Δϵᵖ # plastic strain
            # Construction of the tangent operator ∂σ_n/∂ϵ_n at {D = D_n+1}:
            C = compute_damaged_stiffness_tensor(r,s.ϵ,s.D) #TODO check s.D or temp_D
            return s.temp_σ, C
        else
            s.temp_ϵ = ϵij
            s.temp_σ = σ_trial
            return s.temp_σ, Cᵉ
        end

        # EQUALITY TEST TODO investigate that
        #Δσ = C ⊡ Δϵ
        # println("Δϵ : ", Δϵ)
        # println("tangent : ", C)
        # println("Δσ from tangent : ", Δσ)
        # println("Δσ from constitutive relationships : ", s.temp_σ - s.σ)
        #Δσ = s.temp_σ - s.σ
        #println("Δσ : ", Δσ)
        #println("Δϵ : ", Δϵ)
        #C2 = Δσ ⊗ pinv(Δϵ)
        #Δσ_C2 = C2 ⊡ Δϵ
        #println("Δσ const rel : ", Δσ)
        #println("Δσ from Δσ ⊗ inv(Δϵ) : ", Δσ_C2)

        #Δϵ̅ᵖ = Δλ*ξ#Δλ*∇Q_scalar # accumulated plastic strain increment

        #s.temp_ϵ̅ᵖ = s.ϵ̅ᵖ + Δϵ̅ᵖ # accumulated plastic strain
        #s.temp_σ = σ_trial - Δλ * Dᵉ⊡∇Q
    end
end

function compute_stress_tangent(ϵij,
                                r::Rheology{T,TD,TV,TE,TP},
                                s::MaterialState,
                                clock::Clock;
                                noplast = false) where {T,TD<:Damage,TV<:Viscosity,TE<:Elasticity,TP<:Plasticity}

    ##### viscoelastic tangent operator
    Gᵛᵉ = 1/(1/r.elasticity.G + clock.Δt/r.viscosity.η)
    λᵛᵉ = λ_from_KG(r.elasticity.K,Gᵛᵉ)
    Cᵛᵉ = get_elastic_stiffness_tensor(Gᵛᵉ,λᵛᵉ)

    # trial stress
    Gratio = Gᵛᵉ/r.elasticity.G
    I2D = SymmetricTensor{2,3}(δ) # second order identity tensor
    σ⁰ = (1/3)*tr(s.σ)*I2D + Gratio*dev(s.σ)
    ϵᵉ_trial = ϵij - s.ϵᵖ # trial-elastic_strain
    Δϵ = ϵij - s.ϵ # strain increment
    σ_trial = σ⁰ + Cᵛᵉ ⊡ (ϵij - s.ϵ) # trial stress

    ##### Evaluation of trial values OLD
    # ϵᵉ_trial = ϵ - s.ϵᵖ # trial-elastic_strain
    # σ_trial = Cᵉ ⊡ ϵᵉ_trial # trial-stress

    if noplast # during first newton iteration to homogenize strain field
        s.temp_ϵ = ϵij
        s.temp_σ = σ_trial
        return s.temp_σ, Cᵛᵉ
    end

    p_trial = 1/3 * tr(σ_trial) # trial pressure, negative in compression
    s_trial = dev(σ_trial) # trial deviatoric stress
    τ_trial = get_τ(s_trial,r.damage) # effetive trial-stress (2nd invariant of the deviatoric stress tensor)
    KI_trial = compute_KI(r,p_trial,τ_trial,s.D) # KI at the end of the timestep if the loading is purely elastic

    if KI_trial < 0.0 # elastic loading over Δt
        s.temp_ϵ = ϵij
        s.temp_σ = σ_trial
        return s.temp_σ, Cᵛᵉ
    else # damage growth over Δt
        # Converged stresses and strains invariants
        ϵ0 = tr(s.ϵ)
        e0 = dev(s.ϵ)
        γ0 = sqrt(2.0 * e0 ⊡ e0)

        σ0 = 1/3 * tr(s.σ)
        s0 = dev(s.σ)
        τ0 = sqrt(0.5 * s0 ⊡ s0)

        # Initial conditions
        u0 = [s.D, ϵ0, γ0, σ0, τ0]
        # Time span
        tspan = (0.0,clock.Δt)
        sub_Δt_ini = tspan[2]/100 # initial subtimestep for RK4 interation algorithm
        damage_reltol = 1e-6
        # strain invariants rates
        dϵdt = tr(Δϵ)/clock.Δt
        e = dev(ϵij)
        γ = sqrt(2.0 * e ⊡ e)
        Δγ = γ - γ0
        dγdt = Δγ/clock.Δt
        # OLD AND WRONG :
        #Δe = dev(Δϵ)
        #dγdt = sqrt(2.0 * Δe ⊡ Δe)/clock.Δt

        # ODE system parameters
        params = [r, dϵdt, dγdt]

        # Hand made RK4
        u_vec, t_vec, cohesion_loss_flag = compute_RK4_adaptative(u0,params,state_system!,tspan,sub_Δt_ini,damage_reltol)

        # TEST CODE
        #println("t_vec : ", t_vec)
        #println("D_vec : ", [u[1] for u in u_vec])
        #lineplot(t_vec,u_vec[1])

        # DiffEq solve alternatively
        # # Problem instantiation
        # state_problem = ODEProblem(state_system!,u0,tspan,params)
        #
        # # the ODE system is only solved until D reaches 1.0 => loss of convexity of the free energy function.
        # termination_callback = ContinuousCallback(loss_of_convexity,terminate!)
        # sol = OrdinaryDiffEq.solve(state_problem, Tsit5(), callback=termination_callback)

        # Handle loss of cohesion => branching to a plastic behavior
        if u_vec[end][1] >= 1
            r_nodamage = Rheology(damage = Nothing,
                                  viscosity = r.viscosity,
                                  elasticity = r.elasticity,
                                  plasticity = r.plasticity)
            compute_stress_tangent(ϵij, r_nodamage,s,clock ; noplast = false)
        end

        # update temporary state
        s.temp_D = u_vec[end][1]

        if s.temp_D > s.D
            s.temp_σ = compute_σij(r,s.temp_D,ϵij) # changed from ϵᵉ_trial to ϵ
            #invCᵉ = get_elastic_compliance_tensor(r) # TODO add a viscoelastic compliance tensor getter
            Δϵᵖ = Δϵ - inv(Cᵛᵉ)⊡(s.temp_σ - s.σ) # damage strain increment
            s.temp_ϵ = ϵij
            s.temp_ϵᵖ = s.ϵᵖ + Δϵᵖ # plastic strain
            # Construction of the tangent operator ∂σ_n/∂ϵ_n at {D = D_n+1}:
            C = compute_damaged_stiffness_tensor(r,s.ϵ,s.temp_D)
            return s.temp_σ, C
        else
            s.temp_ϵ = ϵij
            s.temp_σ = σ_trial
            return s.temp_σ, Cᵛᵉ
        end
        # EQUALITY TEST TODO investigate that
        #Δσ = C ⊡ Δϵ
        # println("Δϵ : ", Δϵ)
        # println("tangent : ", C)
        # println("Δσ from tangent : ", Δσ)
        # println("Δσ from constitutive relationships : ", s.temp_σ - s.σ)
        #Δσ = s.temp_σ - s.σ
        #println("Δσ : ", Δσ)
        #println("Δϵ : ", Δϵ)
        #C2 = Δσ ⊗ pinv(Δϵ)
        #Δσ_C2 = C2 ⊡ Δϵ
        #println("Δσ const rel : ", Δσ)
        #println("Δσ from Δσ ⊗ inv(Δϵ) : ", Δσ_C2)

        #Δϵ̅ᵖ = Δλ*ξ#Δλ*∇Q_scalar # accumulated plastic strain increment

        #s.temp_ϵ̅ᵖ = s.ϵ̅ᵖ + Δϵ̅ᵖ # accumulated plastic strain
        #s.temp_σ = σ_trial - Δλ * Dᵉ⊡∇Q
    end
end

function compute_jacobian(f!, u ; eps=1e-9)
    Iv = Int[]
    Jv = Int[]
    Vv = Float64[]
    len_u = length(u)
    δu = Vector{Float64}(undef,len_u)
    half_δu = similar(δu)
    ∂f∂u_j = similar(δu)
    f_up = similar(δu)
    f_down = similar(δu)
    @inbounds for j in 1:len_u
        fill!(δu,0.0)
        δu[j] = eps
        half_δu .= δu./2
        f!(f_up, u .+ half_δu)
        f!(f_down, u .- half_δu)
        @. ∂f∂u_j = ( f_up - f_down ) /  δu[j]
        for i in 1:len_u
            if ∂f∂u_j[i] != 0.0
                push!(Iv,i)
                push!(Jv,j)
                push!(Vv,∂f∂u_j[i])
            end
        end
    end
    J = sparse(Iv,Jv,Vv,len_u,len_u)
    return J
end


text:
#=
This file is part of the replication code for: Hasenzagl, T., Pellegrino, F., Reichlin, L., & Ricco, G. (2020). A Model of the Fed's View on Inflation.
Please cite the paper if you are using any part of the code for academic work (including, but not limited to, conference and peer-reviewed papers).
=#

include("./quarterly2monthly.jl");

function read_data(data_path)

     date_monthly   = [];
     data_monthly   = [];
     date_quarterly = [];
     data_quarterly = [];

     # Monthly data
     try
          f_monthly    = DataFrame(XLSX.readtable(data_path, "monthly")...);
          date_monthly = f_monthly[18:end, 2] |> Array{Date,1};
          data_monthly = f_monthly[18:end, 3:end] |> Array{Union{Float64, Missing},2};

     catch
          print("There are no monthly variables in $data_path \n");
     end

     # Quarterly data
     try
          f_quarterly    = DataFrame(XLSX.readtable(data_path, "quarterly")...);
          date_quarterly = f_quarterly[18:end, 2] |> Array{Date,1};;
          data_quarterly = f_quarterly[18:end, 3:end] |> Array{Union{Float64, Missing},2};

     catch
          print("There are no quarterly variables in $data_path \n");
     end

     # Merge
     if size(data_monthly)[1] == 0 && size(data_quarterly)[1] > 0
          date = date_quarterly;
          data = data_quarterly;
          nM   = 0;
          nQ   = size(data)[2];

     elseif size(data_monthly)[1] > 0 && size(data_quarterly)[1] == 0
          date = date_monthly;
          data = data_monthly;
          nM   = size(data)[2];
          nQ   = 0;

     elseif size(data_monthly)[1] == 0 && size(data_quarterly)[1] == 0
          date     = [];
          data     = [];
          nM       = 0;
          nQ       = 0;
          MNEMONIC = [];
          print("There is no data in $data_path \n");

     else
          Xq   = quarterly2monthly(data_quarterly, size(data_monthly)[1]);
          date = date_monthly;
          data = [[data_monthly; missing.*ones(size(Xq,1)-size(data_monthly,1), size(data_monthly,2))] Xq];
          nM   = size(data_monthly)[2];
          nQ   = size(data_quarterly)[2];
     end

     data = data |> Array{Union{Float64, Missing},2};

     # MNEMONIC
     if nM != 0 || nQ != 0
          info_data = DataFrame(XLSX.readtable(data_path, "transf")...);
          MNEMONIC  = info_data[1:end, 2] |> Array{String,1};
     end

     return data, date, nM, nQ, MNEMONIC;
end


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

const Δt     = 10.
const ntime  = 5000
const nloop  = 200
const nsave  = 5
const nplot  = 10
const run_id = "poincare_1st_fast_vprk_dt10"

include("../guiding_center_4d_settings_vprk.jl")

tableau_list = get_tableau_list_vprk_projection()

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

include("guiding_center_4d_fast_poincare_invariant_1st.jl")


text:
"""
     invertDict(D::OrderedDict; sorted=false)

For a given injective dictionary `D` with key set `K` and value set `V`, this
function returns a dict with key set `V` and values `Array{K}`.
"""
function invertDict(D::OrderedDict; sorted=false)
    uniqueV = unique(values(D))
    invD = OrderedDict{valtype(D),Array{keytype(D),1}}(e => Array{valtype(D),1}() for e  in uniqueV)
    for (k,v) in D
		!(v == k) && push!(invD[v], k)
    end
    if sorted
        for k in keys(invD)
            sort!(invD[k])
        end
    end
    return invD
end

"""
    check_for_equivalence_relation(f::Predicate, dom)

Given a Predicate `f` over a domain `dom`, return `true` if `f`
is a equivalence relation, `false` otherwise.
One can create a symmetric equivalency relation with [`make_symmetric(f)`](@ref)
TODO: check transitivity.

# Examples
```
julia> check_for_equivalence_relation(Predicate((x,y) -> x == -y), -5:10)
true
julia> check_for_equivalence_relation(Predicate((x,y) -> x == y+2), -1:10)
false
```
"""
function check_for_equivalence_relation(f::Predicate, dom)
  for (i,el1) in enumerate(dom)
    for el2 in dom
      !(f(el1,el2) == f(el2,el1)) && return false
    end
  end
  #TODO: check transitivity
  return true
end

function minimal_set(parents::Dict{T,T}, vl::Array{T,1}) where T
    vl_min = Array{T,1}()
    map = Dict{Int64,Int64}()
    used_i = 1
    for (i,el) in enumerate(vl)
        if parents[el] == el
            push!(vl_min, el)
            map[i] = used_i
            used_i += 1
        end
    end
    return  map,vl_min
end


text:
module test_server

using HTTP, Sockets, Test

function testget(url, m=1)
    r = []
    @sync for i in 1:m
        l = rand([0,0,10,1000,10000])
        body = Vector{UInt8}(rand('A':'Z', l))
        # println("sending request...")
        @async push!(r, HTTP.request("GET", "$url/$i", [], body))
    end
    return r
end

@testset "HTTP.listen" begin
    port = 8087 # rand(8000:8999)

    # echo response
    handler = (http) -> begin
        request::HTTP.Request = http.message
        request.body = read(http)
        closeread(http)
        request.response::HTTP.Response = HTTP.Response(request.body)
        request.response.request = request
        startwrite(http)
        write(http, request.response.body)
    end

    server = Sockets.listen(Sockets.InetAddr(parse(IPAddr, "127.0.0.1"), port))
    tsk = @async HTTP.listen(handler, "127.0.0.1", port; server=server)
    sleep(3.0)
    @test !istaskdone(tsk)
    r = testget("http://127.0.0.1:$port")
    @test r[1].status == 200
    close(server)
    sleep(0.5)
    @test istaskdone(tsk)

    server = Sockets.listen(Sockets.InetAddr(parse(IPAddr, "127.0.0.1"), port))
    tsk = @async HTTP.listen(handler, "127.0.0.1", port; server=server)

    handler2 = HTTP.Handlers.RequestHandlerFunction(req->HTTP.Response(200, req.body))

    server2 = Sockets.listen(Sockets.InetAddr(parse(IPAddr, "127.0.0.1"), port+100))
    tsk2 = @async HTTP.serve(handler2, "127.0.0.1", port+100; server=server2)
    sleep(0.5)
    @test !istaskdone(tsk)
    @test !istaskdone(tsk2)

    r = testget("http://127.0.0.1:$port")
    @test r[1].status == 200

    r = testget("http://127.0.0.1:$(port+100)")
    @test r[1].status == 200

    rs = testget("http://127.0.0.1:$port/", 20)
    foreach(rs) do r
        @test r.status == 200
    end

    r = HTTP.get("http://127.0.0.1:$port/"; readtimeout=30)
    @test r.status == 200
    @test String(r.body) == ""

    # large headers
    tcp = Sockets.connect(ip"127.0.0.1", port)
    x = "GET / HTTP/1.1\r\n$(repeat("Foo: Bar\r\n", 10000))\r\n";
    @show length(x)
    write(tcp, "GET / HTTP/1.1\r\n$(repeat("Foo: Bar\r\n", 10000))\r\n")
    sleep(0.1)
    @test occursin(r"HTTP/1.1 413 Request Entity Too Large", String(read(tcp)))

    # invalid HTTP
    tcp = Sockets.connect(ip"127.0.0.1", port)
    write(tcp, "GET / HTP/1.1\r\n\r\n")
    sleep(0.1)
    @test occursin(r"HTTP/1.1 400 Bad Request", String(readavailable(tcp)))

    # no URL
    tcp = Sockets.connect(ip"127.0.0.1", port)
    write(tcp, "SOMEMETHOD HTTP/1.1\r\nContent-Length: 0\r\n\r\n")
    sleep(0.1)
    r = String(readavailable(tcp))
    @test occursin(r"HTTP/1.1 400 Bad Request", r)

    # Expect: 100-continue
    tcp = Sockets.connect(ip"127.0.0.1", port)
    write(tcp, "POST / HTTP/1.1\r\nContent-Length: 15\r\nExpect: 100-continue\r\n\r\n")
    sleep(0.1)
    client = String(readavailable(tcp))
    @test client == "HTTP/1.1 100 Continue\r\n\r\n"

    write(tcp, "Body of Request")
    sleep(0.1)
    client = String(readavailable(tcp))

    println("client:")
    println(client)
    @test occursin("HTTP/1.1 200 OK\r\n", client)
    @test occursin("Transfer-Encoding: chunked\r\n", client)
    @test occursin("Body of Request", client)

    hello = (http) -> begin
        request::HTTP.Request = http.message
        request.body = read(http)
        closeread(http)
        request.response::HTTP.Response = HTTP.Response("Hello")
        request.response.request = request
        startwrite(http)
        write(http, request.response.body)
    end
    close(server)
    close(server2)

    # keep-alive vs. close: issue #81
    port += 1
    server = Sockets.listen(Sockets.InetAddr(parse(IPAddr, "127.0.0.1"), port))
    tsk = @async HTTP.listen(hello, "127.0.0.1", port; server=server, verbose=true)
    sleep(0.5)
    @test !istaskdone(tsk)
    tcp = Sockets.connect(ip"127.0.0.1", port)
    write(tcp, "GET / HTTP/1.0\r\n\r\n")
    sleep(0.5)
    client = String(readavailable(tcp))
    @show client
    @test client == "HTTP/1.1 200 OK\r\n\r\nHello"

    # SO_REUSEPORT
    println("Testing server port reuse")
    t1 = @async HTTP.listen(hello, "127.0.0.1", 8089; reuseaddr=true)
    sleep(0.5)
    @test !istaskdone(t1)

    println("Starting second server listening on same port")
    t2 = @async HTTP.listen(hello, "127.0.0.1", 8089; reuseaddr=true)
    sleep(0.5)
    @test Sys.iswindows() ? istaskdone(t2) : !istaskdone(t2)

    println("Starting server on same port without port reuse (throws error)")
    try
        HTTP.listen(hello, "127.0.0.1", 8089)
    catch e
        @test e isa Base.IOError
        @test startswith(e.msg, "listen")
        @test e.code == Base.UV_EADDRINUSE
    end

    # test automatic forwarding of non-sensitive headers
    # this is a server that will "echo" whatever headers were sent to it
    t1 = @async HTTP.listen("127.0.0.1", 8090) do http
        request::HTTP.Request = http.message
        request.body = read(http)
        closeread(http)
        request.response::HTTP.Response = HTTP.Response(200, request.headers)
        request.response.request = request
        startwrite(http)
        write(http, request.response.body)
    end

    sleep(0.5)
    @test !istaskdone(t1)

    # test that an Authorization header is **not** forwarded to a domain different than initial request
    r = HTTP.get("http://httpbin.org/redirect-to?url=http://127.0.0.1:8090", ["Authorization"=>"auth"])
    @test !HTTP.hasheader(r, "Authorization")

    # test that an Authorization header **is** forwarded to redirect in same domain
    r = HTTP.get("http://httpbin.org/redirect-to?url=https://httpbin.org/response-headers?Authorization=auth")
    @test HTTP.hasheader(r, "Authorization")
end # @testset

end # module


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

using Markdown
using InteractiveUtils

# ╔═╡ a0a65df0-1e4f-11ec-0c0c-87a3e3c07e73
using RemoteS, GMT, SatelliteToolbox

# The AQUA orbits TLE contents for the month of September 2021
tle1 = "1 27424U 02022A   21245.83760660  .00000135  00000-0  39999-4 0  9997";
tle2 = "2 27424  98.2123 186.0654 0002229  67.6025 313.3829 14.57107527 28342";
orb = sat_tracks(tle=[tle1; tle2], duration="1D");

# The orbit track can be visualized with
imshow(orb,  proj=:Robinson, region=:global, coast=true)

# ╔═╡ 491c9bd0-1e63-11ec-0eca-33e88b5d9b74
using Dates

orb = sat_tracks(tle=[tle1; tle2], start=DateTime("2021-09-02T13:30:00"), stop=DateTime("2021-09-02T13:40:00"), step="1m");

Dsc = RemoteS.sat_scenes(orb, RemoteS.SCENE_HALFW["AQUA"], "AQUA");

# ╔═╡ 8295ae60-1e4f-11ec-1e88-35a4e3c1ee58
md"# Plot AQUA satellite tracks

Compute 1 day of AQUA satellite orbits starting at current local time. Note, this
will be accurate for the month of September 2021. For other dates it needs an updated TLE. Note also that repeating the commands below will produce different results since the current local time is used as the starting point.

Orbits calculation rely on the [SatelliteToolbox](https://github.com/JuliaSpace/SatelliteToolbox.jl) package, which is not a direct dependency (meaning, it's not loaded automatically) but need to be installed for this to work. Orbits are calculated with the help of the so called Two Line Element files that unfortunatelly have _accuracy validity_ quite short (around one month). They can be obtainded from the [Space Track](https://www.space-track.org) site.
"

# ╔═╡ ff8c1ec0-1e51-11ec-1605-f542f9898cfd
md"There is no point in plotting orbits for a lengthier duration because they would clutter the figure but we can compute them and display over a specific region on Earth.
"

# ╔═╡ 536ac320-1e52-11ec-343e-05e3a18e5b93
# Compute orbits during 2 days at 15 seconds intervals
orb = sat_tracks(tle=[tle1; tle2], duration="2D", step=15);

# and plot them on my "vicinity"
D = RemoteS.within_BB(orb, [-20, 10, 30, 50]);

imshow(D, proj=:guess, coast=true, dpi=200)

# ╔═╡ bbb879e0-1e57-11ec-32f6-ad3bc0a15aac
md"And imagine we would like to know the names of the AQUA scene files that cover a certain location? We can do it with the ``findscenes`` function to which we must provide the location of interest, the satellite (currently only AQUA or TERRA), the duration of the search and if we want Chlorophyl-a concentration or Sea Surface Temperature (see function help for details).
"

# ╔═╡ 63d520b0-1e58-11ec-23b8-c586b28ab99d
# Get the scene names of data with Chlorophyl-a concentration covering
# the point (-8,36) and for two days before "2021-09-07T17:00:00"
tle1 = "1 27424U 02022A   21245.83760660  .00000135  00000-0  39999-4 0  9997";
tle2 = "2 27424  98.2123 186.0654 0002229  67.6025 313.3829 14.57107527 28342";
findscenes(-8,36, start="2021-09-07T17:00:00", sat=:aqua, day=true, duration=-2, oc=1, tle=[tle1, tle2])

# ╔═╡ 6d7213e0-1e5c-11ec-18e4-99322dfd3273
md"
```
	2-element Vector{String}:
	A2021251125500.L2_LAC_OC.nc
	A2021252134000.L2_LAC_OC.nc
```
"

# ╔═╡ 12c0adf0-1e5a-11ec-024e-61169e4c5e09
md"Now you can download them from [```https://oceandata.sci.gsfc.nasa.gov/ob/getfile/A2021251125500.L2_LAC_OC.nc```](https://oceandata.sci.gsfc.nasa.gov/ob/getfile/A2021251125500.L2_LAC_OC.nc ) (but you will need to register first) and play with the ```grid_at_sensor``` function for interpolation and visualization.
"

# ╔═╡ 1dd36980-1e64-11ec-0753-617fe170d8ef
md"Another cool thing we can do is to plot the area extent of the AQUA or TERRA scenes. We do it by first computing and orbit with a `start` and `stop` times and with increments of 1 minute (attention, this is an important factor). With that orbit we use the `sat_tracks` function to compute the polygons delimiting those areas.
"

# ╔═╡ 2d9c2220-1e65-11ec-0b33-1f76ba219cb7
md"The scene names are stored in the `header` field of the `Dsc` GMTdatset

julia> Dsc[1].header

`AQUA_MODIS.20210902T133001.L2.SST.NRT.nc`

julia> Dsc[2].header

`AQUA_MODIS.20210902T133501.L2.SST.NRT.nc`
"

# ╔═╡ Cell order:
# ╠═8295ae60-1e4f-11ec-1e88-35a4e3c1ee58
# ╠═a0a65df0-1e4f-11ec-0c0c-87a3e3c07e73
# ╠═ff8c1ec0-1e51-11ec-1605-f542f9898cfd
# ╠═536ac320-1e52-11ec-343e-05e3a18e5b93
# ╠═bbb879e0-1e57-11ec-32f6-ad3bc0a15aac
# ╠═63d520b0-1e58-11ec-23b8-c586b28ab99d
# ╠═6d7213e0-1e5c-11ec-18e4-99322dfd3273
# ╠═12c0adf0-1e5a-11ec-024e-61169e4c5e09
# ╠═1dd36980-1e64-11ec-0753-617fe170d8ef
# ╠═491c9bd0-1e63-11ec-0eca-33e88b5d9b74
# ╠═2d9c2220-1e65-11ec-0b33-1f76ba219cb7

max_stars_repo_name: GenericMappingTools/RemoteS

max_issues_repo_name: GenericMappingTools/RemoteS

max_forks_repo_name: GenericMappingTools/RemoteS


text:
struct GLMTrait{D<:ResponseDistribution, L<:InverseLinkFunction}
formula::String
mu::Vector{Float64}
dist:: D
link:: L
end

function GLMTrait(mu::Number, df, dist::D, link::L) where {D, L}
    return(GLMTrait{D, L}(string(mu), repeat([mu], size(df, 1)), dist, link))
end

function GLMTrait(formula::String, df, dist::D, link::L) where {D, L}
    mu = mean_formula(formula, df)
    return(GLMTrait{D, L}(formula, mu, dist, link))
end

function Multiple_GLMTraits(formulas, df, dist::ResponseDistribution, link::InverseLinkFunction)
  vec = [GLMTrait(formulas[i], df, dist, link) for i in 1:length(formulas)] #vector of GLMTrait objects
  return(vec)
end

# we put type of the dist vector as Any since we want to allow for any ResponseType{Poisson(), LogLink()}, ResponseType{Normal(), IdentityLink()}
function Multiple_GLMTraits(formulas::Vector{String}, df::DataFrame, dist::Vector, link::Vector)
  vec = [GLMTrait(formulas[i], df, dist[i], link[i]) for i in 1:length(formulas)]
  return(vec)
end

# lmm: multiple traits (MVN)

struct LMMTrait{T}
formulas::Vector{String}
mu::Matrix{Float64}
vc::T
  function LMMTrait(formulas, df, vc::T) where T
    n_traits = length(formulas)
    n_people = size(df)[1]
    mu = zeros(n_people, n_traits)
    for i in 1:n_traits
      #calculate the mean vector
      mu[:, i] += mean_formula(formulas[i], df)
    end
    return(new{T}(formulas, mu, vc))
  end
end

max_stars_repo_name: sarah-ji/Tutorials

max_issues_repo_name: sarah-ji/Tutorials

max_forks_repo_name: sarah-ji/Tutorials


text:
using Pkg
Pkg.activate(".")
using Soss

myModel = @model begin
    N ~ Poisson(100)
    K ~ Poisson(2.5)
    p ~ Dirichlet(K, 1.0)
    μ ~ Normal(0,1.5) |> iid(K)
    σ ~ HalfNormal(1)
    θ = [(m,σ) for m in μ]
    x ~ MixtureModel(Normal, θ, p) |> iid(N)
end

using Parameters
using Plots
pyplot()

using Random
using StatsPlots

m = @model begin
    p ~ Uniform()
    μ ~ Normal(0, 1.5) |> iid(2)
    σ ~ HalfNormal(1)
    θ = [(m, σ) for m in μ]
    x ~ MixtureModel(Normal, θ, [p,1-p]) |> iid(100)
end

Random.seed!(20)
grt = rand(m)
mplot(grt)

function mplot(data)
    @unpack μ,σ,p,x = data
    p = [p, 1-p]
    lo,hi = extrema(x)
    lo=min(lo,-4)
    hi=max(hi,4)
    xs = range(lo,hi,length=1000)
    dist = MixtureModel(Normal, tuple.(μ,σ), p)
    plt = plot(xs, p[1] .* pdf.(Normal(μ[1],σ), xs)
        , legend=false, fill = (0, 0.1, :red), linecolor=:red
        , axis=nothing)
    plot!(plt, xs, p[2] .* pdf.(Normal(μ[2],σ), xs)
        , fill = (0, 0.1, :blue), linecolor=:blue)
    plot!(plt,xs, pdf.(dist, xs), legend=false, linewidth=2, linecolor=:black)
    scatter!(plt, x, zeros(size(x))
        , marker=:vline, markercolor=:black, markersize=10)
    xlims!((-4,4))
end

r = makeRand(m)
anim = @animate for i=1:100
    mplot(r())
end
gif(anim, "m.gif", fps = 1)

#############################
# Inverse model (posterior distribution)

m_fwd = m(:p,:μ,:σ)
r = makeRand(m_fwd)
anim = @animate for i=1:100
    v = merge(grt, r(; grt...))
    mplot(v)
end
gif(anim, "m_fwd.gif", fps = 3)

m_inv = m(:x)
post = nuts(m_inv; grt...)
post = post.samples

anim = @animate for par in post[1:100]
    par = merge(grt, par)
    mplot(par)
end
gif(anim, "m_inv.gif", fps = 1)

####################################
# Posterior predictive checks

using StatsBase
xs  = range(-4,4,length=100)
ppc = repeat(grt.x, 1, 1000)
r = makeRand(m_fwd)

plt = plot(xs,ecdf(grt.x).(xs)
    ,linewidth=3,legend=false,size=(1200,800), axis=nothing)
# Each column is a posterior prediction
for j in 1:1000
    ppc = r(;post[j]...).x
    plot!(plt, xs, (rand(100) ./ 100 .+ ecdf(ppc).(xs)) ./ 1.01,linewidth=3, linecolor=:black, linealpha=0.01)
end
plot!(xs,ecdf(grt.x).(xs),linewidth=5,legend=false,linecolor=:white)
plot!(xs,ecdf(grt.x).(xs),linewidth=4,legend=false,linecolor=:orange)
plt
savefig("ppc.png")


text:
using UnionFind

# edges must be pre-sorted according to weight.
function kruskal{T <: Integer}(nodes :: T, edges :: Array{(T, T)})
    uf = UnionFinder(nodes)
    mst = Array{(T, T)}

    for i in 1:length(edges)
        (u, v) = edges[i]
        if find!(uf, u) != find!(uf, v)
            union!(uf, u, v)
            push!(mst, (u, v))
        end
    end
end


text:
#!/usr/local/bin/julia
#
const PKGS = ["BenchmarkTools","StatsBase","Nullables","SQLite","Feather","MySQL","LibPQ","DataStreams",
              "JuliaDB","IndexedTables","OnlineStats","StatsPlots","StaticArrays","ShiftedArrays","StatsPlots"]

println("Installing required packages ...")

using Pkg
ipks = Pkg.installed()
for p in PKGS
    println("> $p")
    !haskey(ipks,p) && Pkg.add(p)
end
println("Done!")


text:
function SimplePointChecker(arg0::jdouble, arg1::jdouble)
    return SimplePointChecker((jdouble, jdouble), arg0, arg1)
end

function SimplePointChecker(arg0::jdouble, arg1::jdouble, arg2::jint)
    return SimplePointChecker((jdouble, jdouble, jint), arg0, arg1, arg2)
end

function converged(obj::SimplePointChecker, arg0::jint, arg1::Pair, arg2::Pair)
    return jcall(obj, "converged", jboolean, (jint, Pair, Pair), arg0, arg1, arg2)
end


text:
#---------------------------------------------------------------------------------------------------
# Helper Funtions
#
# 1. mat_sqrt computes matrix square root.
# 2. canonical_gauging regauges (sometimes with trucation) the MPS.
# 3. vals_group groups values into different degenerate classes.
#---------------------------------------------------------------------------------------------------
function mat_sqrt(mat::Hermitian; ztol::Real=1e-20)
    vals_all, vecs_all = eigen(mat)
    pos = vals_all .> ztol
    vals_sqrt = sqrt.(vals_all[pos])
    vecs = vecs_all[:, pos]
    L = vecs * Diagonal(vals_sqrt)
    R = Diagonal(1 ./ vals_sqrt) * vecs'
    L, R
end
#---------------------------------------------------------------------------------------------------
function canonical_gauging(Γ, L, R)
    @tensor Γ2[:] := R[-1,1] * Γ[1,-2,2] * L[2,-3]
    Γ2
end
#---------------------------------------------------------------------------------------------------
function vals_group(vals::Vector; sorttol::Real=1e-3)
    pos = Vector{Vector{Int64}}(undef, 0)
    temp = zeros(Int64, 0)
    current_val = vals[1]
    for i=1:length(vals)
        if abs(vals[i]-current_val) .< sorttol
            push!(temp, i)
        else
            push!(pos, temp)
            temp = [i]
            current_val = vals[i]
        end
    end
    push!(pos, temp)
    pos
end

#---------------------------------------------------------------------------------------------------
# Right Canonical Form
#
# 1. The algorithm will tensor that is right-normalized.
# 2. Automatically trim null-space.
# 3. Is NOT guaranteed that the outputs are non-degenerate.
# 4. The algorithm is UNSTABLE if the transfer matrix has huge condition number.
#---------------------------------------------------------------------------------------------------
export right_canonical
function right_canonical(
    Γ::AbstractArray{<:Number, 3};
    itr::Integer=100,
    ztol::Real=1e-20
)
    ρ = steady_mat(Γ, itr)
    L, R = mat_sqrt(ρ, ztol=ztol)
    canonical_gauging(Γ, L, R)
end

#---------------------------------------------------------------------------------------------------
# Block Decomposition

# 1. Further decompose the right-renormalized tensor into right-renormalized tensors.
# 2. Ensure the outputs are non-degenerate.
# 3. block_trim only return one block when degeneracy is encountered.
#---------------------------------------------------------------------------------------------------
function block_decomp(
    Γ::AbstractArray{<:Number, 3};
    itr::Integer=100,
    sorttol::Real=1e-3
)
    vals, vecs = rand_fixed_mat(Γ, itr) |> eigen
    vgroup = vals_group(vals, sorttol=sorttol)
    res = begin
        num = length(vgroup)
        ctype = promote_type(eltype(Γ), eltype(vecs))
        Vector{Array{ctype, 3}}(undef, num)
    end
    for i=1:num
        p = vecs[:, vgroup[i]]
        Γc = canonical_gauging(Γ, p, p')
        res[i] = Γc
    end
    res
end
#---------------------------------------------------------------------------------------------------
function block_trim(
    Γ::AbstractArray{<:Number, 3};
    itr::Integer=100,
    sorttol::Real=1e-3
)
    vals, vecs = rand_fixed_mat(Γ, itr) |> eigen
    p = begin
        vgroup = vals_group(vals, sorttol=sorttol)
        i = length.(vgroup) |> argmin
        vecs[:, vgroup[i]]
    end
    canonical_gauging(Γ, p, p')
end

#---------------------------------------------------------------------------------------------------
# Block Right Canonical Form
#
# 1. All-at-once method.
# 2. Return multiple non-degenerate right-canonical form.
#---------------------------------------------------------------------------------------------------
export block_canonical
function block_canonical(
    Γ::AbstractArray{<:Number, 3};
    itr::Integer=100,
    ztol::Real=1e-20,
    sorttol::Real=1e-3,
    trim::Bool=false
)
    Γ_RC = right_canonical(Γ, itr=itr, ztol=ztol)
    res = if trim
        block_trim(Γ_RC, itr=itr, sorttol=sorttol)
    else
        block_decomp(Γ_RC, itr=itr, sorttol=sorttol)
    end
    res
end

#---------------------------------------------------------------------------------------------------
# Schmidt Canonical Form
#
# 1. Given a right canonical form, return a Schmidt canonical form.
# 2. This algorithm assume there is no degeneracy.
#---------------------------------------------------------------------------------------------------
export schmidt_canonical
function schmidt_canonical(
    Γ::AbstractArray{<:Number,3};
    itr::Integer=100,
    bound::Integer=100,
    cutoff::Real=1e-14,
    renormalize::Bool=true
)
    X, Yt = begin
        R = steady_mat(Γ, itr)
        L = steady_mat(Γ, itr, dir=:l)
        R_res = cholesky(R)
        L_res = cholesky(L)
        R_res.L, L_res.U
    end
    U, S, V = svd_trim(Yt * X, bound, cutoff, renormalize)
    R_mat = inv(Yt) * U * Diagonal(S)
    L_mat = V * inv(X)
    Γ_new = canonical_gauging(Γ, R_mat, L_mat)
    Γ_new, S
end


text:
# load external packages -
using Distributions
using Plots

# load my codes -
include("Library.jl")


text:
function i_plus_x(i, x, N)
	i = (i+N+x) % N
	(i == 0) && (i = N)
	return i
end

function mean(vet::AbstractVector)
	(length(vet) == 0) && (return 0)
	return sum(vet) / length(vet)
end

function issymmetric(mat::AbstractMatrix)
	(size(mat, 1) != size(mat, 2)) && (return false)
	for i in 1:size(mat, 1), j in i+1:size(mat, 2)
		(mat[i,j] != mat[j,i]) && (return false)
	end
	return true
end

function hasnodeOrError(network::AbstractNetwork, idx_node::Integer)
	hasnode(network, idx_node) ||
		throw(ArgumentError("Invalid node index! Node $idx_node doesn't exist in the network!"))
end


text:
# using Revise
using Test
using Logging
using LarSurf
# Logging.configure(level==Logging.Debug)

# include("../src/LarSurf.jl")
# include("../src/block.jl")

@testset "Boundary matrix shape test" begin
    grid_size = [2,3,4]
    b3, larmodel = LarSurf.calculate_boundary3(grid_size)
#     println(size(b3))
#     @test size(b3) =
#     slides0 = LarSurf.random_image([7, 7, 7], [1,2,2], [3, 4, 5], 2)
#     Matrix(slides0)
#     @test all(slides0 .== [1, 5, 6, 10, 11, 15])

    # slides0 = LarSurf.data_sub_from_block_sub([5, 5, 5], 1, [1,2,3])
    # @test all(slides0 .== [0, 6, 5, 11, 10, 16])
    Cnum = prod(grid_size)
    Fnum = LarSurf.grid_number_of_faces(grid_size)
    @test size(b3, 1) == Cnum
    @test size(b3, 2) == Fnum

    V, (VV, EV, FV, CV) = larmodel
    @test size(b3, 1) == length(CV)
    @test size(b3, 2) == length(FV)

    @test size(b3) == (24, 98)
    FV = larmodel[2][3]
    # display(FV)
    # println(Fnum)

end

@testset "reagular brick boundary matrix to image file" begin
    # using ImageView
    sz = 4
    grid_size = [sz, sz, sz]
    b3, larmodel = LarSurf.calculate_boundary3(grid_size)
    barr = Matrix(b3 .> 0)
    nvoxels = sz^3
    nfaces = 3 * (sz^2 * (sz + 1) )
    @test size(barr, 1) == nvoxels
    @test size(barr, 2) == nfaces
    # imshow(barr)
end

@testset "Tests get boundary" begin
    grid_size = [2,3,4]
    LarSurf.set_param(
        boundary_allow_memory=false,
        boundary_allow_read_files=false,
        boundary_allow_write_files=false
        )
    b3a, larmodel_a = LarSurf.get_boundary3(grid_size)

    LarSurf.set_param(boundary_allow_memory=true)
    b3b, larmodel_b = LarSurf.get_boundary3(grid_size)

    @test size(b3a) == (24, 98)
    @test size(b3b) == (24, 98)

end


text:
using MPIMeasurements

scanner = MPIScanner("HeadScanner.toml")
daq = getDAQ(scanner)

params = toDict(daq.params)

params["studyName"]="TestTobi"
params["studyDescription"]="A very cool measurement"
params["scannerOperator"]="Tobi"
params["dfStrength"]=[1e-3]
params["acqNumFrames"]=100
params["acqNumAverages"]=10

x = linspace(0,1,5)
#params["acqFFValues"] = [0.0]
params["acqFFValues"] = repeat( cat(1,x,reverse(x[2:end-1])),inner=5)
params["acqNumPeriodsPerFrame"]=length(params["acqFFValues"])

#u = measurement(daq, params, controlPhase=false)
filename = measurement(daq, params, MDFStore, controlPhase=true)


text:
using DecisionTrees
using CSV
using DataFrames
using Profile

pkgdir = Pkg.dir("DecisionTrees") # how long will this remain supported?
testdir = joinpath(pkgdir, "test")
datadir = joinpath(pkgdir, "data")

elt = [Int,	Float64,	Float64,	Float64,	Float64,	Float64,	Int,	String,	String,	String,	Int,	String,	String,	String,	String,	String,	String,	String,	String,	String,	String,	Int,	String,	Int,	String,	String,	Int,	String,	String,	Int,	String,	String,	Int,	String,	String,	Int,	String,	Int,	String,	String,	Int,	Int,	String,	Int,	Int,	String,	Int,	Int,	String,	Int,	String,	String,	String,	Int,	String,	String,	String,	String,	Int,	Int,	Int,	String,	String,	String,	String,	String,	Int,	Int,	Int,	Int,	Int,	Int,	Int,	Int,	Int,	Int,	Int,	Int,	Int,	Float64,	Int,	Float64,	Float64,	Float64,	Float64,	Int,	Int,	Int,	Int,	Int,	Int]

thisfile = joinpath(datadir, "GermanMotorPremiums", "data1Small.csv")
@assert isfile(thisfile)
@time df_tmp = CSV.read(thisfile, allowmissing=:none, types=elt, categorical=false, rows_for_type_detect=10000);

selected_explanatory_vars = ["PLZ_WOHNORT","ART_DES_WOHNEIGENTUM","GEBURTSDATUM","FAMILIENSTAND","NATIONALITAET","GESCHLECHT","FINANZIERUNGSART","STADT","KENNZEICHEN"]

# keep 10 largest PLZ only
vorig = deepcopy(df_tmp[:PLZ_WOHNORT])
counts, freqs, vals, keep = getCounts(vorig, threshold=10)
vnew = df_tmp[:PLZ_WOHNORT]
mapToOther!(vnew,keep,9999999)

##################################################
# Prepare the data
##################################################

dtmtable, sett, df_prepped = prepare_dataframe_for_dtm!(df_tmp, treat_as_categorical_variable=["PLZ_WOHNORT"], weightcol="EXPOSURE", numcol="LOSS20HALF", denomcol="PREMIUM66", independent_vars=selected_explanatory_vars);
sett.minWeight = -.2

# selected_explanatory_vars=[    "VORSCHAEDEN_ANZAHL",    "MALLORCA_POLICE",	"SCHUTZBRIEF_INKL",	"FREIE_WERKSTATTWAHL",	"AUTOMOBILCLUB_MITGLIED_SEIT",	"BAHNCARD",	"ZAHLUNGSWEISE",	"JAHRESKARTE_OEPNV",	"MOTORRAD_BESITZER",	"AUTOMOBILCLUB",	"SFKLASSE_VOLLKASKO",	"SFKLASSE_HAFTPFLICHT",	"STELLPLATZ_ABSCHLIESSBAR",	"NAECHTLICHER_STELLPLATZ",	"NUTZUNGSWEISE",	"JAEHRLICHE_FAHRLEISTUNG",	"TSN",	"ERSTZULASSUNG",	"HSN",	"FINANZIERUNGSART",	"ZULASSUNG_AUF_VERSICHERUNGSNEHM",	"STADT",	"KENNZEICHEN",	"PLZ_DES_HALTER",	"SELBSTGENUTZTES_WOHNEIGENTUM",	"ART_DES_WOHNEIGENTUM",	"GEBURTSDATUM",	"FAMILIENSTAND",	"NATIONALITAET",	"GESCHLECHT",	"FUEHRERSCHEIN_ERWORBEN_AM",	"VORSCHAEDEN0_typeKH",	"VORSCHAEDEN0_typetk",	"VORSCHAEDEN0_month",	"VORSCHAEDEN0_year",	"VORSCHAEDEN1_typetk",	"VORSCHAEDEN1_month",	"VORSCHAEDEN1_year",	"VORSCHAEDEN2_typevk",	"VORSCHAEDEN2_month",	"VORSCHAEDEN2_year",	"adacid",	"name",	"marke",	"modell",	"preis",	"getriebeart",	"antriebsart",	"Fahrzeugklasse",	"co2klasse",	"kw",	"ps",	"tueranzahl",	"Motorart",	"Kraftstoffart",	"Motorbauart",	"Schadstoffklasse",	"Karosserie",	"Sitzanzahl",	"typklasseh_num",	"typklassetk_num",	"typklassevk_num",	"hubraum2",	"drehmoment2",	"breite2",	"radstand2",	"laenge2",	"hoehe2",	"leergewicht2",	"gesamtgewicht2",	"zuladung2",	"kofferraumvolumen_num",	"hoechstgeschwindigkeit2",	"verbrauchgesamt2",	"verbrauchausserorts2",	"verbrauchinnerorts2",	"beschleunigung2",	"tank2",	"kfzsteuer2",	"anzahlgaenge2",	"anzahlzylinder2",	"co2_wert",	"modellstart_y"]
##################################################
# run tree
##################################################
strs, resm = dtm(dtmtable, sett)

##################################################
# run boosting
##################################################

sett.iterations = 10
sett.model_type = "boosted_tree"
strs, resm2 = dtm(dtmtable, sett)

Profile.clear_malloc_data()

strs, resm2 = dtm(dtmtable, sett)
@info("allocation run done.")
quit();


text:
# 2D LSRTM/FWI parallelization test with 2 sources and 2 vintages
#
# Ziyi Yin, ziyi.yin@gatech.edu
# August 2021

parsed_args = parse_commandline()

nlayer = parsed_args["nlayer"]
tti = parsed_args["tti"]
fs =  parsed_args["fs"]

### Model
model, model0, dm = setup_model(tti, 4)
q, srcGeometry, recGeometry, info = setup_geom(model; nsrc=4)
q1 = q[[1,4]]
q2 = q[[2,3]]
srcGeometry1 = subsample(srcGeometry,[1,4])
srcGeometry2 = subsample(srcGeometry,[2,3])
recGeometry1 = subsample(recGeometry,[1,4])
recGeometry2 = subsample(recGeometry,[2,3])
info1 = subsample(info,[1,4])
info2 = subsample(info,[2,3])

opt = Options(sum_padding=true, free_surface=fs)
F1 = judiModeling(info1, model, srcGeometry1, recGeometry1; options=opt)
F2 = judiModeling(info2, model, srcGeometry2, recGeometry2; options=opt)

# Observed data
dobs1 = F1*q1
dobs2 = F2*q2

@testset "FWI/LSRTM objective multi-level parallelization test with $(nlayer) layers and tti $(tti) and freesurface $(fs)" begin

    ftol = 1f-5
    dm1 = 2f0*circshift(dm, 10)
    dm2 = 2f0*circshift(dm, 30)

    Jm0, grad = lsrtm_objective([model0, model0], [q1, q2], [dobs1, dobs2], [dm1, dm2]; options=opt, nlind=true)
    Jm01, grad1 = lsrtm_objective(model0, q1, dobs1, dm1; options=opt, nlind=true)
    Jm02, grad2 = lsrtm_objective(model0, q2, dobs2, dm2; options=opt, nlind=true)
    @test isapprox(Jm01+Jm02, Jm0; rtol=ftol)
    @test isapprox(grad1, grad[1]; rtol=ftol)
    @test isapprox(grad2, grad[2]; rtol=ftol)

    _Jm0, _grad = fwi_objective([model0, model0], [q1, q2], [dobs1, dobs2]; options=opt)
    _Jm01, _grad1 = fwi_objective(model0, q1, dobs1; options=opt)
    _Jm02, _grad2 = fwi_objective(model0, q2, dobs2; options=opt)
    @test isapprox(_Jm01+_Jm02, _Jm0; rtol=ftol)
    @test isapprox(_grad1, _grad[1]; rtol=ftol)
    @test isapprox(_grad2, _grad[2]; rtol=ftol)

end


text:
module Pollen

using ANSIColoredPrinters
using AbstractTrees
using Base.Docs
using CommonMark
using CommonMark: Document, Item, Text, Paragraph, List, Heading,
    Emph, SoftBreak, Link, Code, Node, AbstractContainer, CodeBlock, ThematicBreak,
    BlockQuote, Admonition, Attributes, Image, Citation
using FilePathsBase
using DataStructures: DefaultDict, OrderedDict
import Gumbo
using JuliaFormatter
using Mustache
using LiveServer
using IJulia
import LiveServer
using HTTP
using TOML
using IOCapture
using JSON3
using Revise

include("xtree.jl")
include("selectors.jl")
include("catas.jl")
include("folds.jl")

#include("xexpr.jl")
include("reflectionutils.jl")
include("references.jl")
include("files.jl")

include("formats.jl")
include("markdown.jl")
include("html.jl")
include("jupyter.jl")
#include("json.jl")

include("rewriters.jl")
include("project.jl")
include("builders.jl")

include("serve/events.jl")
include("serve/server.jl")
include("serve/servefiles.jl")

include("rewriters/documentfolder.jl")
include("rewriters/referencer.jl")
include("rewriters/documenttree.jl")
include("rewriters/basic.jl")
include("rewriters/assets.jl")
include("rewriters/templater.jl")
include("rewriters/coderunner.jl")
include("rewriters/inserter.jl")
include("rewriters/toc.jl")
include("rewriters/packagewatcher.jl")

#=
include("serve.jl")
include("servelazy.jl")
=#
include("projects.jl")

export select,
    XTree, XNode, XLeaf,
    cata, catafirst, replace, replacefirst, fold, catafold,
    SelectNode,
    Replacer, Inserter, HTMLTemplater, ExecuteCode,
    NthChild, FirstChild, Before, After,
    Project, build,
    SelectTag, SelectOr, XExpr, ChangeTag, htmlify, AddSlugID, AddTableOfContents, SelectAttrEq,
    Selector, parse, HTML, Markdown, resolveidentifier, serve,
    # rewriters
    AddID, HTMLify, ChangeLinkExtension, FormatCode, AddTableOfContents, Referencer, DocumentFolder,
    documentationproject, Server, runserver, ServeFiles, ServeFilesLazy,
    PackageWatcher,
    RelativeLinks

end


text:
using Rasters, Test
using Rasters.LookupArrays, Rasters.Dimensions

@testset "set"  begin
    A = [missing 7; 2 missing]
    ga = Raster(A, (Y(-20:40:20), X(50:10:60)); missingval=missing)

    # Use the Projected lookup, with crs
    lu = Projected(; crs=EPSG(2024))
    ga = set(ga, Y=lu, X=lu)
    @test crs(ga) == EPSG(2024)
    @test mappedcrs(ga) == nothing

    # Set it with usercrs as well
    lu = Projected(; crs=EPSG(2024), mappedcrs=EPSG(4326))
    ga = set(ga, Y=lu, X=lu)
    @test crs(ga) == EPSG(2024)
    @test mappedcrs(ga) == EPSG(4326)

    # Make them intervals
    ga = set(ga, Y=Intervals(Start()), X=Intervals(Start()))
    @test sampling(ga) == (Intervals(Start()), Intervals(Start()))

    # Change the parent array
    ga = set(ga, [1 2; missing 5])
    @test all(parent(ga) .=== [1 2; missing 5])

    # TODO: Add more tests once the lookups are cleaned up
end