text:
"""
Pre-computes a nonuniform fast Fourier transform of type `N`.

For best performance, choose the right number of threads by `FFTW.set_num_threads(4)`, for example.
"""
struct NUFFTPlan{N,T,FFT} <: Plan{T}
    U::Matrix{T}
    V::Matrix{T}
    p::FFT
    t::Vector{Int}
    temp::Matrix{T}
    temp2::Matrix{T}
    Ones::Vector{T}
end

"""
Pre-computes a nonuniform fast Fourier transform of type I.
"""
function plan_nufft1(ω::AbstractVector{T}, ϵ::T) where {T<:AbstractFloat}
    N = length(ω)
    ωdN = ω/N
    t = AssignClosestEquispacedFFTpoint(ωdN)
    γ = PerturbationParameter(ωdN, AssignClosestEquispacedGridpoint(ωdN))
    K = FindK(γ, ϵ)
    U = constructU(ωdN, K)
    V = constructV(range(zero(T), stop=N-1, length=N), K)
    p = plan_bfft!(V, 1)
    temp = zeros(Complex{T}, N, K)
    temp2 = zeros(Complex{T}, N, K)
    Ones = ones(Complex{T}, K)

    NUFFTPlan{1, eltype(U), typeof(p)}(U, V, p, t, temp, temp2, Ones)
end

"""
Pre-computes a nonuniform fast Fourier transform of type II.
"""
function plan_nufft2(x::AbstractVector{T}, ϵ::T) where T<:AbstractFloat
    N = length(x)
    t = AssignClosestEquispacedFFTpoint(x)
    γ = PerturbationParameter(x, AssignClosestEquispacedGridpoint(x))
    K = FindK(γ, ϵ)
    U = constructU(x, K)
    V = constructV(range(zero(T), stop=N-1, length=N), K)
    p = plan_fft!(U, 1)
    temp = zeros(Complex{T}, N, K)
    temp2 = zeros(Complex{T}, N, K)
    Ones = ones(Complex{T}, K)

    NUFFTPlan{2, eltype(U), typeof(p)}(U, V, p, t, temp, temp2, Ones)
end

"""
Pre-computes a nonuniform fast Fourier transform of type III.
"""
function plan_nufft3(x::AbstractVector{T}, ω::AbstractVector{T}, ϵ::T) where T<:AbstractFloat
    N = length(x)
    s = AssignClosestEquispacedGridpoint(x)
    t = AssignClosestEquispacedFFTpoint(x)
    γ = PerturbationParameter(x, s)
    K = FindK(γ, ϵ)
    u = constructU(x, K)
    v = constructV(ω, K)

    p = plan_nufft1(ω, ϵ)

    D1 = Diagonal(1 .- (s .- t .+ 1)./N)
    D2 = Diagonal((s .- t .+ 1)./N)
    D3 = Diagonal(exp.(-2 .* im .* T(π) .* ω ))
    U = hcat(D1*u, D2*u)
    V = hcat(v, D3*v)

    temp = zeros(Complex{T}, N, 2K)
    temp2 = zeros(Complex{T}, N, 2K)
    Ones = ones(Complex{T}, 2K)

    NUFFTPlan{3, eltype(U), typeof(p)}(U, V, p, t, temp, temp2, Ones)
end

function (*)(p::NUFFTPlan{N,T}, c::AbstractArray{V}) where {N,T,V}
    mul!(zeros(promote_type(T,V), size(c)), p, c)
end

function mul!(f::AbstractVector{T}, P::NUFFTPlan{1,T}, c::AbstractVector{T}) where {T}
    U, V, p, t, temp, temp2, Ones = P.U, P.V, P.p, P.t, P.temp, P.temp2, P.Ones

    broadcast!(*, temp, c, U)
    conj!(temp)
    fill!(temp2, zero(T))
    recombine_rows!(temp, t, temp2)
    p*temp2
    conj!(temp2)
    broadcast!(*, temp, V, temp2)
    mul!(f, temp, Ones)

    f
end

function mul!(F::Matrix{T}, P::NUFFTPlan{N,T}, C::Matrix{T}) where {N,T}
    for J = 1:size(F, 2)
        mul_col_J!(F, P, C, J)
    end
    F
end

function broadcast_col_J!(f, temp::Matrix, C::Matrix, U::Matrix, J::Int)
    N = size(C, 1)
    COLSHIFT = N*(J-1)
    @inbounds for j = 1:size(temp, 2)
        for i = 1:N
            temp[i,j] = f(C[i+COLSHIFT],U[i,j])
        end
    end
    temp
end

function mul_col_J!(F::Matrix{T}, P::NUFFTPlan{1,T}, C::Matrix{T}, J::Int) where {T}
    U, V, p, t, temp, temp2, Ones = P.U, P.V, P.p, P.t, P.temp, P.temp2, P.Ones

    broadcast_col_J!(*, temp, C, U, J)
    conj!(temp)
    fill!(temp2, zero(T))
    recombine_rows!(temp, t, temp2)
    p*temp2
    conj!(temp2)
    broadcast!(*, temp, V, temp2)
    COLSHIFT = size(C, 1)*(J-1)
    mul_for_col_J!(F, temp, Ones, 1+COLSHIFT, 1)

    F
end

function mul!(f::AbstractVector{T}, P::NUFFTPlan{2,T}, c::AbstractVector{T}) where {T}
    U, V, p, t, temp, temp2, Ones = P.U, P.V, P.p, P.t, P.temp, P.temp2, P.Ones

    broadcast!(*, temp, c, V)
    p*temp
    reindex_temp!(temp, t, temp2)
    broadcast!(*, temp, U, temp2)
    mul!(f, temp, Ones)

    f
end

function mul_col_J!(F::Matrix{T}, P::NUFFTPlan{2,T}, C::Matrix{T}, J::Int) where {T}
    U, V, p, t, temp, temp2, Ones = P.U, P.V, P.p, P.t, P.temp, P.temp2, P.Ones

    broadcast_col_J!(*, temp, C, V, J)
    p*temp
    reindex_temp!(temp, t, temp2)
    broadcast!(*, temp, U, temp2)
    COLSHIFT = size(C, 1)*(J-1)
    mul_for_col_J!(F, temp, Ones, 1+COLSHIFT, 1)

    F
end

function mul!(f::AbstractVector{T}, P::NUFFTPlan{3,T}, c::AbstractVector{T}) where {T}
    U, V, p, t, temp, temp2, Ones = P.U, P.V, P.p, P.t, P.temp, P.temp2, P.Ones

    broadcast!(*, temp2, c, V)
    mul!(temp, p, temp2)
    reindex_temp!(temp, t, temp2)
    broadcast!(*, temp, U, temp2)
    mul!(f, temp, Ones)

    f
end

function mul_col_J!(F::Matrix{T}, P::NUFFTPlan{3,T}, C::Matrix{T}, J::Int) where {T}
    U, V, p, t, temp, temp2, Ones = P.U, P.V, P.p, P.t, P.temp, P.temp2, P.Ones

    broadcast_col_J!(*, temp2, C, V, J)
    mul!(temp, p, temp2)
    reindex_temp!(temp, t, temp2)
    broadcast!(*, temp, U, temp2)
    COLSHIFT = size(C, 1)*(J-1)
    mul_for_col_J!(F, temp, Ones, 1+COLSHIFT, 1)

    F
end

mul_for_col_J!(y::AbstractVecOrMat{T}, A::AbstractMatrix{T}, x::AbstractVecOrMat{T}, istart::Int, jstart::Int) where T =
    mul_for_col_J!(y, A, x, istart, jstart, 1, 1)

function mul_for_col_J!(y::AbstractVecOrMat{T}, A::AbstractMatrix{T}, x::AbstractVecOrMat{T}, istart::Int, jstart::Int, INCX::Int, INCY::Int) where T
    m, n = size(A)
    ishift, jshift = istart-INCY, jstart-INCX
    @inbounds for i = 1:m
        y[ishift+i*INCY] = zero(T)
    end
    @inbounds for j = 1:n
        xj = x[jshift+j*INCX]
        for i = 1:m
            y[ishift+i*INCY] += A[i,j]*xj
        end
    end

    y
end

function reindex_temp!(temp::Matrix{T}, t::Vector{Int}, temp2::Matrix{T}) where {T}
    @inbounds for j = 1:size(temp, 2)
        for i = 1:size(temp, 1)
            temp2[i, j] = temp[t[i], j]
        end
    end
    temp2
end

function recombine_rows!(temp::Matrix{T}, t::Vector{Int}, temp2::Matrix{T}) where {T}
    @inbounds for j = 1:size(temp, 2)
        for i = 1:size(temp, 1)
            temp2[t[i], j] += temp[i, j]
        end
    end
    temp2
end

"""
Computes a nonuniform fast Fourier transform of type I:

```math
f_j = \\sum_{k=0}^{N-1} c_k e^{-2\\pi{\\rm i} \\frac{j}{N} \\omega_k},\\quad{\\rm for}\\quad 0 \\le j \\le N-1.
```
"""
nufft1(c::AbstractVector, ω::AbstractVector{T}, ϵ::T) where {T<:AbstractFloat} = plan_nufft1(ω, ϵ)*c

"""
Computes a nonuniform fast Fourier transform of type II:

```math
f_j = \\sum_{k=0}^{N-1} c_k e^{-2\\pi{\\rm i} x_j k},\\quad{\\rm for}\\quad 0 \\le j \\le N-1.
```
"""
nufft2(c::AbstractVector, x::AbstractVector{T}, ϵ::T) where {T<:AbstractFloat}  = plan_nufft2(x, ϵ)*c

"""
Computes a nonuniform fast Fourier transform of type III:

```math
f_j = \\sum_{k=0}^{N-1} c_k e^{-2\\pi{\\rm i} x_j \\omega_k},\\quad{\\rm for}\\quad 0 \\le j \\le N-1.
```
"""
nufft3(c::AbstractVector, x::AbstractVector{T}, ω::AbstractVector{T}, ϵ::T) where {T<:AbstractFloat} = plan_nufft3(x, ω, ϵ)*c

const nufft = nufft3
const plan_nufft = plan_nufft3

"""
Pre-computes a 2D nonuniform fast Fourier transform.

For best performance, choose the right number of threads by `FFTW.set_num_threads(4)`, for example.
"""
struct NUFFT2DPlan{T,P1,P2} <: Plan{T}
    p1::P1
    p2::P2
    temp::Vector{T}
end

"""
Pre-computes a 2D nonuniform fast Fourier transform of type I-I.
"""
function plan_nufft1(ω::AbstractVector{T}, π::AbstractVector{T}, ϵ::T) where T<:AbstractFloat
    p1 = plan_nufft1(ω, ϵ)
    p2 = plan_nufft1(π, ϵ)
    temp = zeros(Complex{T}, length(π))

    NUFFT2DPlan(p1, p2, temp)
end

"""
Pre-computes a 2D nonuniform fast Fourier transform of type II-II.
"""
function plan_nufft2(x::AbstractVector{T}, y::AbstractVector{T}, ϵ::T) where T<:AbstractFloat
    p1 = plan_nufft2(x, ϵ)
    p2 = plan_nufft2(y, ϵ)
    temp = zeros(Complex{T}, length(y))

    NUFFT2DPlan(p1, p2, temp)
end

function (*)(p::NUFFT2DPlan{T}, C::Matrix{V}) where {T,V}
    mul!(zeros(promote_type(T,V), size(C)), p, C)
end

function mul!(F::Matrix{T}, P::NUFFT2DPlan{T}, C::Matrix{T}) where {T}
    p1, p2, temp = P.p1, P.p2, P.temp

    # Apply 1D x-plan to all columns
    mul!(F, p1, C)

    # Apply 1D y-plan to all rows
    for i = 1:size(C, 1)
        @inbounds for j = 1:size(F, 2) temp[j] = F[i,j] end
        mul!(temp, p2, temp)
        @inbounds for j = 1:size(F, 2) F[i,j] = temp[j] end
    end

    F
end

"""
Computes a 2D nonuniform fast Fourier transform of type I-I:

```math
F_{i,j} = \\sum_{k=0}^{M-1}\\sum_{\\ell=0}^{N-1} C_{k,\\ell} e^{-2\\pi{\\rm i} (\\frac{i}{M} \\omega_k + \\frac{j}{N} \\pi_{\\ell})},\\quad{\\rm for}\\quad 0 \\le i \\le M-1,\\quad 0 \\le j \\le N-1.
```
"""
nufft1(C::Matrix, ω::AbstractVector{T}, π::AbstractVector{T}, ϵ::T) where {T<:AbstractFloat} = plan_nufft1(ω, π, ϵ)*C

"""
Computes a 2D nonuniform fast Fourier transform of type II-II:

```math
F_{i,j} = \\sum_{k=0}^{M-1}\\sum_{\\ell=0}^{N-1} C_{k,\\ell} e^{-2\\pi{\\rm i} (x_i k + y_j \\ell)},\\quad{\\rm for}\\quad 0 \\le i \\le M-1,\\quad 0 \\le j \\le N-1.
```
"""
nufft2(C::Matrix, x::AbstractVector{T}, y::AbstractVector{T}, ϵ::T) where {T<:AbstractFloat} = plan_nufft2(x, y, ϵ)*C

FindK(γ::T, ϵ::T) where {T<:AbstractFloat} = γ ≤ ϵ ? 1 : Int(ceil(5*γ*exp(lambertw(log(10/ϵ)/γ/7))))

(AssignClosestEquispacedGridpoint(x::AbstractVector{T})::AbstractVector{T}) where {T<:AbstractFloat} = round.([Int], size(x, 1)*x)
(AssignClosestEquispacedFFTpoint(x::AbstractVector{T})::Array{Int,1}) where {T<:AbstractFloat} = mod.(round.([Int], size(x, 1)*x), size(x, 1)) .+ 1
(PerturbationParameter(x::AbstractVector{T}, s_vec::AbstractVector{T})::AbstractFloat) where {T<:AbstractFloat} = norm(size(x, 1)*x - s_vec, Inf)

function constructU(x::AbstractVector{T}, K::Int) where {T<:AbstractFloat}
    # Construct a low rank approximation, using Chebyshev expansions
    # for AK = exp(-2*pi*1im*(x[j]-j/N)*k):
    N = length(x)
    s = AssignClosestEquispacedGridpoint(x)
    er = N*x - s
    γ = norm(er, Inf)
    Diagonal(exp.(-im*(π*er)))*ChebyshevP(K-1, er/γ)*Bessel_coeffs(K, γ)
end

function constructV(ω::AbstractVector{T}, K::Int) where {T<:AbstractFloat}
    complex(ChebyshevP(K-1, ω.*(two(T)/length(ω)) .- 1))
end

function Bessel_coeffs(K::Int, γ::T) where {T<:AbstractFloat}
    # Calculate the Chebyshev coefficients of exp(-2*pi*1im*x*y) on [-gam,gam]x[0,1]
    cfs = zeros(Complex{T}, K, K)
    arg = -γ*π/two(T)
    for p = 0:K-1
     	for q = mod(p,2):2:K-1
    		cfs[p+1,q+1] = 4*(1im)^q*besselj((p+q)/2,arg).*besselj((q-p)/2,arg)
    	end
    end
    cfs[1,:] = cfs[1,:]/two(T)
    cfs[:,1] = cfs[:,1]/two(T)
    return cfs
end

function ChebyshevP(n::Int, x::AbstractVector{T}) where T<:AbstractFloat
    # Evaluate Chebyshev polynomials of degree 0,...,n at x:
    N = size(x, 1)
    Tcheb = Matrix{T}(undef, N, n+1)

    # T_0(x) = 1.0
    One = convert(eltype(x),1.0)
    @inbounds for j = 1:N
        Tcheb[j, 1] = One
    end
    # T_1(x) = x
    if ( n > 0 )
        @inbounds for j = 1:N
            Tcheb[j, 2] = x[j]
        end
    end
    # 3-term recurrence relation:
    twoX = 2x
    @inbounds for k = 2:n
        @simd for j = 1:N
            Tcheb[j, k+1] = twoX[j]*Tcheb[j, k] - Tcheb[j, k-1]
        end
    end
    return Tcheb
end


text:
using DSP

"""
    y = undbv(x)

Convert `x` from dB to a voltage
"""
function undbv(x)
    return db2amp.(x)
end

ext: jl

max_stars_repo_name: Mahmoud-Kharsa/DeltaSigma

max_issues_repo_name: Mahmoud-Kharsa/DeltaSigma

max_forks_repo_name: Mahmoud-Kharsa/DeltaSigma


text:
using Test
using Recurrences

include("examples.jl")
include("empirical.jl")

for x in keys(cf)
    @info "$(string(x))"
    a, b = trace(eval(x))
    @test a == b
end

# for x in keys(cf)
#     lrs = lrs_sequential(eval(x), :n)
#     @test solve(lrs) == cf[x]
# end


text:
module DynamicLinearModels

export DLMPlot
export extract
export to_matrix
export from_matrix
export simulate
export compute_prior
export check_dimensions
export kfilter
export ksmoother
export evolutional_covariances
export fitted
export forecast
export estimate

using Distributions
using LinearAlgebra
using RecipesBase

"""
    CovMat

Alias for symmetric, real valued, dense matrices.
"""
const CovMat{RT <: Real} = Symmetric{RT, Matrix{RT}}

"""
    DLMPlot

Indicator type for plotting recipe.
"""
struct DLMPlot end

"""
    plot(::DLMPlot, Y, f, Q[, fh, Qh; factor = 1.64, index = 1])

Recipe for easily plotting the results obtained by the package routines, where
`Y` is the observations matrix, `f` and `Q` are results from `fitted`, and
`fh` and `Qh` are the results from `forecast`. Factor implies the credibility
or the credibility intervals interval, e.g. a factor of 1.64 implies a
credibility of 90%. Index indicates which observational index is to be plotted.
"""
@recipe function plot(::DLMPlot,
                      Y::Matrix{RT},
                      f::Matrix{RT},
                      Q::Vector{CovMat{RT}},
                      fh::Union{Matrix{RT}, Nothing} = nothing,
                      Qh::Union{Vector{CovMat{RT}}, Nothing} = nothing;
                      factor = 1.64,
                      index = 1,
                      nreps = 1) where RT <: Real

    n = size(f[1], 1)
    T = size(Y, 1)

    if isnothing(fh) != isnothing(Qh)
        throw(ArgumentError("Missingness of fh and Qh must match."))
    end

    x::Vector{UnitRange{Int}} = []
    y::Vector{Vector{RT}} = []
    st = Matrix{Symbol}(undef, 1, 0)
    lt = Matrix{Symbol}(undef, 1, 0)
    co = Matrix{Symbol}(undef, 1, 0)
    lb = Matrix{String}(undef, 1, 0)

    for i = 1:nreps
        local_index = n * (i-1) + index
        push!(x, 1:T)
        push!(y, Y[:,local_index])
        st = hcat(st, :scatter)
        lt = hcat(lt, :auto)
        co = hcat(co, :grey)
        lb = hcat(lb, i == 1 ? "Observations" : "")
    end

    yhat = f[:,index]
    Qhat = extract(Q, index)

    push!(x, 1:T, 1:T, 1:T)
    push!(y,
          yhat,
          yhat + factor * sqrt.(Qhat),
          yhat - factor * sqrt.(Qhat))
    st = hcat(st, :line, :line, :line)
    lt = hcat(lt, :solid, :dot, :dot)
    co = hcat(co, :orange, :orange, :orange)
    lb = hcat(lb, "Fitted", "", "")

    if !isnothing(fh)
        h = size(fh, 1)
        fend = f[end,index]
        Qend = Q[end][index,index]
        yfor = fh[:,index]
        Qfor = extract(Qh, index)

        push!(x, T:T+h, T:T+h, T:T+h)
        push!(y,
              vcat(fend, yfor),
              vcat(fend, yfor) + factor * sqrt.(vcat(Qend, Qfor)),
              vcat(fend, yfor) - factor * sqrt.(vcat(Qend, Qfor)))
        st = hcat(st, :line, :line, :line)
        lt = hcat(lt, :solid, :dot, :dot)
        co = hcat(co, :skyblue, :skyblue, :skyblue)
        lb = hcat(lb, "Forecast", "", "")
    end

    seriestype := st
    linestyle := lt
    color --> co
    label --> lb
    (x, y)
end

"""
    extract(x, index)

Utility function that facilitates fetching the time series as a vector for one
the objects used in the package.

If `x` is a `Vector{Vector{RT}}` it returns `[x[t][index] for t = 1:T]`;
If `x` is a `Vector{CovMat{RT}}` it returns `[x[t][index,index] for t = 1:T]`.
"""
function extract(x::Vector{<: Union{Vector{RT}, CovMat{RT}}},
                 index::Union{<: Integer, BitVector}) where RT <: Real

    if typeof(x) == Vector{Vector{RT}}
        return [x[t][index] for t = 1:size(x,1)]
    else
        return [x[t][index,index] for t = 1:size(x,1)]
    end
end

"""
    to_matrix(x)

Converts a vector of vectors to a matrix.
"""
function to_matrix(x::Vector{Vector{RT}}) where RT <: Real

    return collect(hcat(x...)')
end

"""
    from_matrix(x)

Converts a matrix to a vector of vector.
"""
function from_matrix(x::Matrix{RT}) where RT <: Real

    return [x[t,:] for t = 1:size(x, 1)]
end

"""
    check_dimensions(F, G[; V, W, Y])

Utility function that checks for dimension mistmatches in the given arguments
and returns the dimensions for the observations and for the state-space.
"""
function check_dimensions(F::Matrix{RT},
                          G::Matrix{RT};
                          V::Union{CovMat{RT}, Nothing} = nothing,
                          W::Union{CovMat{RT}, Nothing} = nothing,
                          Y::Union{Matrix{RT}, Nothing} = nothing,
                          nreps::Integer = 1) where RT <: Real

    n = size(F, 1)
    p = size(F, 2)

    if size(G, 1) != size(G, 2)
        throw(DimensionMismatch("G must be a square matrix."))
    end

    if size(G) != (p, p)
        throw(DimensionMismatch("Dimension of G dos not match that of F."))
    end

    if !isnothing(V)
        # Only need to check one dimension since Symmetric guarantees square
        if size(V, 1) != n
            throw(DimensionMismatch("Dimensions of V are a mismatch."))
        end
    end

    if !isnothing(W)
        # Only need to check one dimension since Symmetric guarantees square
        if size(W, 1) != p
            throw(DimensionMismatch("Dimensions of W are a mismatch."))
        end
    end

    if !isnothing(Y)
        if size(Y, 2) != n * nreps
            throw(DimensionMismatch("Dimension of Y incompatible with F."))
        end
    end

    return n, p
end

"""
    simulate(F, G, V, W, θ₀, T[, nreps])

Simulates a time-series Dynamic Linear Model specified by the quadruple
(`F`, `G`, `V`, `W`) with a starting state of `θ₁ <- Nₚ(G θ₀, W)`, with `T`
observations. A parameter `nreps` may be passed indicating the number of
replicates to be generated. Returns the generated `θ` and `y`.

Note that the parametrizations being considered in this package is such that
    `y[t] = F * y[t-1] + ϵ`
and not the notation from West and Harrison (1996) where
    `y[t] = F' * y[t-1] + ϵ`.
"""
function simulate(F::Matrix{RT},
                  G::Matrix{RT},
                  V::CovMat{RT},
                  W::CovMat{RT},
                  θ₀::Vector{RT},
                  T::Integer,
                  nreps::Integer = 1) where RT <: Real

    n, p = check_dimensions(F, G, V=V, W=W)

    θ = Matrix{RT}(undef, T, p)
    y = Matrix{RT}(undef, T, n * nreps)

    ω = MultivariateNormal(zeros(p), W)
    ϵ = MultivariateNormal(zeros(n), V)

    θ[1,:] = G * θ₀ + rand(ω)
    y[1,:] = repeat(F * θ[1,:], nreps) + reshape(rand(ϵ, nreps), nreps * n)
    for t = 2:T
        θ[t,:] = G * θ[t-1,:] + rand(ω)
        y[t,:] = repeat(F * θ[t,:], nreps) + reshape(rand(ϵ, nreps), nreps * n)
    end

    return θ, y
end

"""
    compute_prior(Y, F, m₀, C₀)

Utility function for computing a smart prior that's not informative and at the
same does not lead to numerical or visualization issues. Only takes effect if
one of `m₀` is `C₀` is `nothing`, otherwise it just returns `m₀` and `C₀`.
"""
function compute_prior(Y::Matrix{RT},
                       F::Matrix{RT},
                       m₀::Union{Vector{RT}, Nothing},
                       C₀::Union{CovMat{RT}, Nothing}) where RT <: Real

     p = size(F, 2)

     if isnothing(m₀)
         m₀ = zeros(p)
     end

     if isnothing(C₀)
         magic_sdev = maximum(abs.(Y[1,:] - F * m₀))
         C₀ = Symmetric(diagm(repeat([magic_sdev^2], p)))
     end

     return m₀, C₀
end

"""
    exclude_low_weights(Y, η, ϵ)

Internal function which returns a version of the original `Y` and `η` variables,
containing only the observations with weights above `ϵ`.
"""
@inline function exclude_low_weights(Y::Matrix{RT},
                                     η::Vector{RT},
                                     ϵ::RT) where RT <: Real

    # TODO: Make sure I'm returning Views

    good_mask = η .> ϵ
    if all(good_mask)
        return Y, η
    end

    n = Int(size(Y[1,:], 1) / size(η, 1))
    return Y[:,repeat(good_mask, inner=n)], η[good_mask]
end

"""
    kfilter_core(y, F, G, V, W, a, R)

Internal function which actually performs the computation step.
"""
@inline function kfilter_core(y::Vector{RT},
                              F::Matrix{RT},
                              G::Matrix{RT},
                              V::CovMat{RT},
                              W::CovMat{RT},
                              m::Vector{RT},
                              C::CovMat{RT}) where RT <: Real

    a = G * m
    R = Symmetric(G * C * G') + W
    f = F * a
    Q = Symmetric(F * R * F') + V
    A = R * F' * inv(Q)
    m = a + A * (y - f)
    C = R - Symmetric(A * Q * A')

    return a, R, m, C
end

"""
    kfilter_core(y, F, G, V, W, a, R)

Internal function which actually performs the computation step.
"""
@inline function kfilter_core(y::Vector{RT},
                              F::Matrix{RT},
                              G::Matrix{RT},
                              V::CovMat{RT},
                              δ::RT,
                              m::Vector{RT},
                              C::CovMat{RT}) where RT <: Real

    a = G * m
    R = Symmetric(G * C * G') / δ
    f = F * a
    Q = Symmetric(F * R * F') + V
    A = R * F' * inv(Q)
    m = a + A * (y - f)
    C = R - Symmetric(A * Q * A')

    return a, R, m, C
end

"""
    kfilter(Y, F, G, V, W[, m₀, C₀])

Filtering routine for the Dynamic Linear Model (`F`, `G`) where the
observational and evolutional covariance matrices `V` and `W` are known and
constant. `Y` is the matrix of observations with `T` rows and `n` columns.
Prior parameters `m₀` and `C₀` may be omitted, in which case `compute_prior`
kicks in to assign a prior.

Returns one-step ahead prior means and covariances `a` and `R`, and online
means and covariances `m` and `C`.
"""
function kfilter(Y::Matrix{RT},
                 F::Matrix{RT},
                 G::Matrix{RT},
                 V::CovMat{RT},
                 W::CovMat{RT},
                 m₀::Union{Vector{RT}, Nothing} = nothing,
                 C₀::Union{CovMat{RT}, Nothing} = nothing) where RT <: Real

    n, p = check_dimensions(F, G, V=V, W=W, Y=Y)
    T = size(Y, 1)

    m₀, C₀ = compute_prior(Y, F, m₀, C₀)

    a = Matrix{RT}(undef, T, p)
    m = Matrix{RT}(undef, T, p)
    R = Vector{CovMat{RT}}(undef, T)
    C = Vector{CovMat{RT}}(undef, T)

    a[1,:], R[1], m[1,:], C[1] = kfilter_core(Y[1,:], F, G, V, W, m₀, C₀)
    for t = 2:T
        a[t,:], R[t], m[t,:], C[t] = kfilter_core(Y[t,:], F, G, V, W, m[t-1,:], C[t-1])
    end

    return a, R, m, C
end

"""
    kfilter(Y, F, G, V, δ[, m₀, C₀])

Filtering routine for a discount factor Dynamic Linear Model (`F`, `G`) where
the observational covariance matrix `V` is known and constant and evolutional
covariance matrices `W[1], ..., W[T]` are indirectly modelled through a discount
factor `δ`. See West & Harrison (1996) for further information of the discount
factor apporach. `Y` is the matrix of observations with `T` rows and `n`
columns.  Prior parameters `m₀` and `C₀` may be omitted, in which case
`compute_prior` kicks in to assign a prior.

Returns one-step ahead prior means and covariances `a` and `R`, and online
means and covariances `m` and `C`.
"""
function kfilter(Y::Matrix{RT},
                 F::Matrix{RT},
                 G::Matrix{RT},
                 V::CovMat{RT},
                 δ::RT,
                 m₀::Union{Vector{RT}, Nothing} = nothing,
                 C₀::Union{CovMat{RT}, Nothing} = nothing) where RT <: Real

    n, p = check_dimensions(F, G, V=V, Y=Y)
    T = size(Y, 1)

    m₀, C₀ = compute_prior(Y, F, m₀, C₀)

    a = Matrix{RT}(undef, T, p)
    m = Matrix{RT}(undef, T, p)
    R = Vector{CovMat{RT}}(undef, T)
    C = Vector{CovMat{RT}}(undef, T)

    a[1,:], R[1], m[1,:], C[1] = kfilter_core(Y[1,:], F, G, V, δ, m₀, C₀)
    for t = 2:T
        a[t,:], R[t], m[t,:], C[t] = kfilter_core(Y[t,:], F, G, V, δ, m[t-1,:], C[t-1])
    end

    return a, R, m, C
end

"""
    kfilter(Y, F, G, V, δ[, m₀, C₀])

Filtering routine for a discount factor Dynamic Linear Model (`F`, `G`) where
the observational covariance matrices `V[1], …, V[T]` are known and evolutional
covariance matrices `W[1], ..., W[T]` are indirectly modelled through a discount
factor `δ`. See West & Harrison (1996) for further information of the discount
factor apporach. `Y` is the matrix of observations with `T` rows and `n`
columns. Prior parameters `m₀` and `C₀` may be omitted, in which case
`compute_prior` kicks in to assign a prior.

Returns one-step ahead prior means and covariances `a` and `R`, and online
means and covariances `m` and `C`.
"""
function kfilter(Y::Matrix{RT},
                 F::Matrix{RT},
                 G::Matrix{RT},
                 V::Vector{CovMat{RT}},
                 δ::RT,
                 m₀::Union{Vector{RT}, Nothing} = nothing,
                 C₀::Union{CovMat{RT}, Nothing} = nothing) where RT <: Real

    n, p = check_dimensions(F, G, Y=Y)
    T = size(Y, 1)

    m₀, C₀ = compute_prior(Y, F, m₀, C₀)

    a = Matrix{RT}(undef, T, p)
    m = Matrix{RT}(undef, T, p)
    R = Vector{CovMat{RT}}(undef, T)
    C = Vector{CovMat{RT}}(undef, T)

    a[1,:], R[1], m[1,:], C[1] = kfilter_core(Y[1,:], F, G, V[1], δ, m₀, C₀)
    for t = 2:T
        a[t,:], R[t], m[t,:], C[t] = kfilter_core(Y[t,:], F, G, V[t], δ, m[t-1,:], C[t-1])
    end

    return a, R, m, C
end

"""
    kfilter(Y, F, G, V, η, δ[, m₀, C₀])

Filtering routine for a discount factor Dynamic Linear Model (`F`, `G`) where
the observational covariance matrix `V` is known and constants and evolutional
covariance matrices `W[1], ..., W[T]` are indirectly modelled through a discount
factor `δ` and observations have replications. `Y` is the matrix of observations
with `T` rows and `n * nreps` columns, each replicate with a weight`η[i]`.
Prior parameters `m₀` and `C₀` may be omitted, in which case `compute_prior`
kicks in to assign a prior.

Returns one-step ahead prior means and covariances `a` and `R`, and online
means and covariances `m` and `C`.
"""
function kfilter(Y::Matrix{RT},
                 F::Matrix{RT},
                 G::Matrix{RT},
                 V::CovMat{RT},
                 η::Vector{RT},
                 δ::RT,
                 m₀::Union{Vector{RT}, Nothing} = nothing,
                 C₀::Union{CovMat{RT}, Nothing} = nothing;
                 exclude_under::Union{RT, Nothing} = nothing) where RT <: Real

    if !isnothing(exclude_under)
        Y, η = exclude_low_weights(Y, η, exclude_under)
    end

    nreps = size(η, 1)
    n, p = check_dimensions(F, G, V=V, Y=Y, nreps=nreps)
    T = size(Y, 1)

    η = copy(η)
    η /= sum(η)

    F = repeat(F, nreps)
    V = Symmetric(kron(diagm(1 ./ η), V))

    m₀, C₀ = compute_prior(Y, F, m₀, C₀)

    a = Matrix{RT}(undef, T, p)
    m = Matrix{RT}(undef, T, p)
    R = Vector{CovMat{RT}}(undef, T)
    C = Vector{CovMat{RT}}(undef, T)

    a[1,:], R[1], m[1,:], C[1] = kfilter_core(Y[1,:], F, G, V, δ, m₀, C₀)
    for t = 2:T
        a[t,:], R[t], m[t,:], C[t] = kfilter_core(Y[t,:], F, G, V, δ, m[t-1,:], C[t-1])
    end

    return a, R, m, C
end

"""
    kfilter(Y, F, G, V, η, δ[, m₀, C₀])

Filtering routine for a discount factor Dynamic Linear Model (`F`, `G`) where
the observational covariance matrix `V` is known and constants and evolutional
covariance matrices `W[1], ..., W[T]` are indirectly modelled through a discount
factor `δ` and observations have replications. `Y` is the matrix of observations
with `T` rows and `n * nreps` columns, each replicate with dynamic weights
`η[t,i]`.  Prior parameters `m₀` and `C₀` may be omitted, in which case
`compute_prior` kicks in to assign a prior.

Returns one-step ahead prior means and covariances `a` and `R`, and online
means and covariances `m` and `C`.
"""
function kfilter(Y::Matrix{RT},
                 F::Matrix{RT},
                 G::Matrix{RT},
                 V::CovMat{RT},
                 η::Matrix{RT},
                 δ::RT,
                 m₀::Union{Vector{RT}, Nothing} = nothing,
                 C₀::Union{CovMat{RT}, Nothing} = nothing;
                 exclude_under::Union{RT, Nothing} = nothing) where RT <: Real

    if !isnothing(exclude_under)
        Y, η = exclude_low_weights(Y, η, exclude_under)
    end

    nreps = size(η, 2)
    n, p = check_dimensions(F, G, V=V, Y=Y, nreps=nreps)
    T = size(Y, 1)

    η = copy(η)
    for t = 1:T
        η[t,:] /= sum(η[t,:])
    end

    F = repeat(F, nreps)

    m₀, C₀ = compute_prior(Y, F, m₀, C₀)

    a = Matrix{RT}(undef, T, p)
    m = Matrix{RT}(undef, T, p)
    R = Vector{CovMat{RT}}(undef, T)
    C = Vector{CovMat{RT}}(undef, T)

    VV = Symmetric(kron(diagm(1 ./ η[1,:]), V))
    a[1,:], R[1], m[1,:], C[1] = kfilter_core(Y[1,:], F, G, VV, δ, m₀, C₀)
    for t = 2:T
        VV = Symmetric(kron(diagm(1 ./ η[t,:]), V))
        a[t,:], R[t], m[t,:], C[t] = kfilter_core(Y[t,:], F, G, VV, δ, m[t-1,:], C[t-1])
    end

    return a, R, m, C
end

"""
    ksmoother(G, a, R, m, C)

Filtering routine for a Dynamic Linear Model ( ⋅, `G`), where `a` and `R` are
the filtered one-step ahead prior means and covariances, and `m` and `C` are
the filtered online means and covariances.

Returns the posterior means and covariances `s` and `S`.
"""
function ksmoother(G::Matrix{RT},
                   a::Matrix{RT},
                   R::Vector{CovMat{RT}},
                   m::Matrix{RT},
                   C::Vector{CovMat{RT}}) where RT <: Real

    p = size(m, 2)
    T = size(R, 1)

    s = similar(m, T, p)
    S = similar(C, T)

    s[T,:] = m[T,:]
    S[T] = C[T]

    for t = T-1:-1:1
        B::Matrix{RT} = C[t] * G' * inv(R[t+1])
        s[t,:] = m[t,:] + B * (s[t+1,:] - a[t+1,:])
        S[t] = C[t] - Symmetric(B * (R[t+1] - S[t+1]) * B')
    end

    return s, S
end

"""
    evolutional_covariances(Y, F, G, V, δ[, m₀, C₀])

Compute the implied values of the evolutional covariances W[1], ..., W[T] when
considering a discount factor approach.
"""
function evolutional_covariances(Y::Matrix{RT},
                                 F::Matrix{RT},
                                 G::Matrix{RT},
                                 V::CovMat{RT},
                                 δ::RT,
                                 m₀::Union{Vector{RT}, Nothing} = nothing,
                                 C₀::Union{CovMat{RT}, Nothing} = nothing) where RT <: Real

    T = size(Y, 1)
    m, C = compute_prior(Y, F, m₀, C₀)

    W = Vector{CovMat{RT}}(undef, T)

    for t = 1:T
        _, _, m, C = kfilter_core(Y[t,:], F, G, V, δ, m, C)
        W[t] = Symmetric(G * C * G') * ((1. - δ) / δ)
    end

    return W
end

"""
    evolutional_covariances(Y, F, G, V, η, δ[, m₀, C₀])

Compute the implied values of the evolutional covariances W[1], ..., W[T] when
considering a discount factor approach, and weighted replicates.
"""
function evolutional_covariances(Y::Matrix{RT},
                                 F::Matrix{RT},
                                 G::Matrix{RT},
                                 V::CovMat{RT},
                                 η::Vector{RT},
                                 δ::RT,
                                 m₀::Union{Vector{RT}, Nothing} = nothing,
                                 C₀::Union{CovMat{RT}, Nothing} = nothing) where RT <: Real

    nreps = size(η, 1)
    F = repeat(F, nreps)
    V = Symmetric(kron(diagm(1 ./ η), V))
    return evolutional_covariances(Y, F, G, V, δ, m₀, C₀)
end

"""
    evolutional_covariances(Y, F, G, V, δ[, m₀, C₀])

Compute the implied values of the evolutional covariances W[1], ..., W[T] when
considering a discount factor approach, and dynamically weighted replicates.
"""
function evolutional_covariances(Y::Matrix{RT},
                                 F::Matrix{RT},
                                 G::Matrix{RT},
                                 V::CovMat{RT},
                                 η::Matrix{RT},
                                 δ::RT,
                                 m₀::Union{Vector{RT}, Nothing} = nothing,
                                 C₀::Union{CovMat{RT}, Nothing} = nothing) where RT <: Real

    #TODO: We need a proper time-by-time exclude_low_weights in this case.

    nreps = size(η, 2)
    T = size(Y, 1)
    F = repeat(F, nreps)
    m, C = compute_prior(Y, F, m₀, C₀)

    W = Vector{CovMat{RT}}(undef, T)

    for t = 1:T
        VV = Symmetric(kron(diagm(1 ./ η[t,:]), V))
        _, _, m, C = kfilter_core(Y[t,:], F, G, VV, δ, m, C)
        W[t] = Symmetric(G * C * G') * ((1. - δ) / δ)
    end

    return W
end

"""
    fitted(F, V, m, C)

Computes the fitted values for the data, for a model with observational matrix
`F`, evolutional covariance matrix `V`, and state means `m` and `C`. Note that
this can be done with the one-step ahead priors, online parameters or, more
appropriately, smoother results.

Returns observational means and covariances `f` and `Q`.
"""
function fitted(F::Matrix{RT},
                V::CovMat{RT},
                m::Matrix{RT},
                C::Vector{CovMat{RT}}) where RT <: Real

    n = size(F, 1)
    T = size(m, 1)

    f = Matrix{RT}(undef, T, n)
    Q = Vector{CovMat{RT}}(undef, T)

    for t = 1:T
        f[t,:] = F * m[t,:]
        Q[t] = Symmetric(F * C[t] * F') + V
    end

    return f, Q
end

"""
    forecast(F, G, V, W, μ, Σ, h)

Filtering routine for the Dynamic Linear Model (`F`, `G`) where the
observational and evolutional covariance matrices `V` and `W` are known and
constant. `μ` and `Σ` are the mean and covariance matrix for the last state
given the most recent information, and `h` is the forecasting window.

Returns observational means and covariances `f` and `Q`.
"""
function forecast(F::Matrix{RT},
                  G::Matrix{RT},
                  V::CovMat{RT},
                  W::CovMat{RT},
                  μ::Vector{RT},
                  Σ::CovMat{RT},
                  h::Integer) where RT <: Real

    n = size(F, 1)

    f = Matrix{RT}(undef, h, n)
    Q = Vector{CovMat{RT}}(undef, h)

    a, R = μ, Σ
    for t in 1:h
        a = G * a
        R = Symmetric(G * R * G') + W
        f[t,:] = F * a
        Q[t] = Symmetric(F * R * F') + V
    end

    return f, Q
end

"""
    forecast(F, G, V, δ, μ, Σ, h)

Filtering routine for the Dynamic Linear Model (`F`, `G`) where the
observational covariance matrix `V` is known and constants, and evolutional
covariance matrices W[1], ..., W[T] are indirectly modeled through a discount
factor `δ`. `μ` and `Σ` are the mean and covariance matrix for the last state
given the most recent information, and `h` is the forecasting window.

Returns observational means and covariances `f` and `Q`.
"""
function forecast(F::Matrix{RT},
                  G::Matrix{RT},
                  V::CovMat{RT},
                  δ::RT,
                  μ::Vector{RT},
                  Σ::CovMat{RT},
                  h::Integer) where RT <: Real

    return forecast(F, G, V, Σ / δ, μ, Σ, h)
end

"""
    estimate(Y, F, G, δ[, m₀, C₀; maxit, ϵ])

Obtains maximum a posteriori estimates for the states and observational
covariance matrix for a Dynamic Linear Model (`F`, `G`), considering a discount
factor `δ` for the evolutional covariance matrices. Prior parameters `m₀` and
`C₀` may be omitted, in which case `compute_prior` kicks in to assign a prior.
Parameters `maxit` and `ϵ` control the maximum number of iterations and the
numerical precision for early convergence for the Coordinate Descent algorithm,
respectively.

Returns point estimates for the states `θ` and point estimate for the
covariance matrix `V`. Also returns the number of iterations until convergence.
If negative, it means the algorithm stopped from reaching the maximum number
of iterations.
"""
function estimate(Y::Matrix{RT},
                  F::Matrix{RT},
                  G::Matrix{RT},
                  δ::RT,
                  m₀::Union{Vector{RT}, Nothing} = nothing,
                  C₀::Union{CovMat{RT}, Nothing} = nothing;
                  maxit::Integer = 50,
                  ϵ::RT = 1e-8) where RT <: Real

    n, p = check_dimensions(F, G, Y=Y)
    T = size(Y, 1)

    # Initialize values
    ϕ = ones(n)
    θ = zeros(T, p)

    it_count = zero(maxit)

    # Coordinate descent algorithm: Iterate on conditional maximums
    for it = 1:maxit
        prev_θ = copy(θ)

        # Conditional maximum for the states is the mean from the normal
        # distributions resulting from Kalman smoothing
        a, R, m, C = kfilter(Y, F, G, Symmetric(diagm(ϕ)), δ, m₀, C₀)
        θ, _ = ksmoother(G, a, R, m, C)

        # Conditional maximum for variance comes from the Gamma distribution
        # which is InverseGamma(T - 1, sum((y - θ)^2)) for which the mode is
        # given by the sum((y - F θ)^2) / T.
        ϕ = zero(ϕ)
        for t = 1:T
            ϕ += (Y[t,:] - F * θ[t,:]) .^ 2 / T
        end

        # Check for early convergence condition
        it_count = it
        if sum((θ - prev_θ) .^ 2) < p * T * ϵ^2
            break
        end
    end

    return θ, Symmetric(diagm(ϕ)), it_count < maxit ? it_count : -1
end

"""
    estimate(Y, F, G, η, δ[, m₀, C₀; maxit, ϵ])

Obtains maximum a posteriori estimates for the states and observational
covariance matrix for a Dynamic Linear Model (`F`, `G`), considering a discount
factor `δ` for the evolutional covariance matrices, and weighted replicates.
Prior parameters `m₀` and `C₀` may be omitted, in which case `compute_prior`
kicks in to assign a prior. Parameters `maxit` and `ϵ` control the maximum
number of iterations and the numerical precision for early convergence for the
Coordinate Descent algorithm, respectively.

Returns point estimates for the states `θ` and point estimate for the
covariance matrix `V`. Also returns the number of iterations until convergence.
If negative, it means the algorithm stopped from reaching the maximum number
of iterations.
"""
function estimate(Y::Matrix{RT},
                  F::Matrix{RT},
                  G::Matrix{RT},
                  η::Vector{RT},
                  δ::RT,
                  m₀::Union{Vector{RT}, Nothing} = nothing,
                  C₀::Union{CovMat{RT}, Nothing} = nothing;
                  maxit::Integer = 50,
                  ϵ::RT = 1e-8,
                  exclude_under::Union{RT, Nothing} = nothing) where RT <: Real

    if !isnothing(exclude_under)
        Y, η = exclude_low_weights(Y, η, exclude_under)
    end

    nreps = size(η, 1)
    n, p = check_dimensions(F, G, Y=Y, nreps=nreps)
    T = size(Y, 1)
    FF = repeat(F, nreps)

    η = copy(η)
    η /= sum(η)

    index_map = [(n * (l - 1) + 1):(n * (l - 1) + n) for l = 1:nreps]

    # Initialize values
    ϕ = ones(n)
    θ = zeros(T, p)

    it_count = zero(maxit)

    # Coordinate descent algorithm: Iterate on conditional maximums
    for it = 1:maxit
        prev_θ = copy(θ)

        # Conditional maximum for the states is the mean from the normal
        # distributions resulting from Kalman smoothing
        V = Symmetric(kron(diagm(1 ./ η), diagm(ϕ)))
        a, R, m, C = kfilter(Y, FF, G, V, δ, m₀, C₀)
        θ, _ = ksmoother(G, a, R, m, C)

        # Conditional maximum for variance comes from the Gamma distribution
        ϕ = zero(ϕ)
        for t = 1:T
            for i = 1:nreps
                ϕ += η[i] * (Y[t,index_map[i]] - F * θ[t,:]) .^ 2 / T
            end
        end

        # Check for early convergence condition
        it_count = it
        if sum((θ - prev_θ) .^ 2) < p * T * ϵ^2
            break
        end
    end

    return θ, Symmetric(diagm(ϕ)), it_count < maxit ? it_count : -1
end

"""
    estimate(Y, F, G, η, δ[, m₀, C₀; maxit, ϵ])

Obtains maximum a posteriori estimates for the states and observational
covariance matrix for a Dynamic Linear Model (`F`, `G`), considering a discount
factor `δ` for the evolutional covariance matrices, and dynamically weighted
replicates. Prior parameters `m₀` and `C₀` may be omitted, in which case
`compute_prior` kicks in to assign a prior. Parameters `maxit` and `ϵ` control
the maximum number of iterations and the numerical precision for early
convergence for the Coordinate Descent algorithm, respectively.

Returns point estimates for the states `θ` and point estimate for the
covariance matrix `V`. Also returns the number of iterations until convergence.
If negative, it means the algorithm stopped from reaching the maximum number
of iterations.
"""
function estimate(Y::Matrix{RT},
                  F::Matrix{RT},
                  G::Matrix{RT},
                  η::Matrix{RT},
                  δ::RT,
                  m₀::Union{Vector{RT}, Nothing} = nothing,
                  C₀::Union{CovMat{RT}, Nothing} = nothing;
                  maxit::Integer = 50,
                  ϵ::RT = 1e-8,
                  exclude_under::Union{RT, Nothing} = nothing) where RT <: Real

    #TODO: We need a proper time-by-time exclude_low_weights in this case.

    nreps = size(η, 2)
    n, p = check_dimensions(F, G, Y=Y, nreps=nreps)
    T = size(Y, 1)
    FF = repeat(F, nreps)

    η = copy(η)
    for t = 1:T
        η[t,:] /= sum(η[t,:])
    end

    index_map = [(n * (l - 1) + 1):(n * (l - 1) + n) for l = 1:nreps]

    # Initialize values
    ϕ = ones(n)
    θ = zeros(T, p)

    it_count = zero(maxit)

    # Coordinate descent algorithm: Iterate on conditional maximums
    for it = 1:maxit
        prev_θ = copy(θ)

        # Conditional maximum for the states is the mean from the normal
        # distributions resulting from Kalman smoothing
        V = [Symmetric(kron(diagm(1 ./ η[t,:]), diagm(ϕ))) for t = 1:T]
        a, R, m, C = kfilter(Y, FF, G, V, δ, m₀, C₀)
        θ, _ = ksmoother(G, a, R, m, C)

        # Conditional maximum for variance comes from the Gamma distribution
        ϕ = zero(ϕ)
        for t = 1:T
            for i = 1:nreps
                ϕ += η[i] * (Y[t,index_map[i]] - F * θ[t,:]) .^ 2 / T
            end
        end

        # Check for early convergence condition
        it_count = it
        if sum((θ - prev_θ) .^ 2) < p * T * ϵ^2
            break
        end
    end

    return θ, Symmetric(diagm(ϕ)), it_count < maxit ? it_count : -1
end

end # module


text:
module SimpleDamage

using Base: @kwdef
using StaticArrays
import DamagedShearBand as DSB
using IfElse

using Reexport
@reexport using DamagedShearBand: Rheology

include("types.jl")
include("damage_functions.jl")
include("functions.jl")
include("data_simulation.jl")

export Model, NumericalSetup, TriaxialSetup, DeformationControl, ConstantStrainRate, ConstantStress, ConstitutiveModel, LinearWeakening, AsymptoticWeakening, EnergyBasedWeakening
export ConstantYieldStress, CoulombYieldStress, SimpleCharlesLaw, InvariantsKICharlesLaw, PrincipalKICharlesLaw, MicroMechanicalCharlesLaw, Elasticity, IncompressibleElasticity, Geom2D, Geom3D
export Plasticity, MinViscosityThreshold, DamageThreshold
export update_derivatives, update_derivatives!
export init_variables, init_params

end#module

max_stars_repo_name: Leooop/SimpleDamage

max_issues_repo_name: Leooop/SimpleDamage

max_forks_repo_name: Leooop/SimpleDamage


text:
using TimerOutputs

"""
    block_sparse_tucker(input_vec::CompressedTuckerState, cluster_ops, clustered_ham;
        max_iter    = 20,
        max_iter_pt = 200,
        nbody       = 4,
        H0          = "Hcmf",
        thresh_var  = 1e-4,
        thresh_foi  = 1e-6,
        thresh_pt   = 1e-5,
        tol_ci      = 1e-5,
	resolve_ss  = true,
        do_pt       = true,
        tol_tucker  = 1e-6)

# Arguments
- `input_vec::CompressedTuckerState`: initial state
- `cluster_ops`: local cluster operators
- `clustered_ham::ClusteredOperator`: hamiltonian
- `max_iter = 20`: max number of iterations
- `max_iter_pt = 200`: max number of iterations to solve for the 1st order wavefunction
- `nbody = 4`: include up to `nbody` terms when forming the FOIS to search
- `H0 = "Hcmf"`: zeroth-order hamiltonian for computing 1st order wavefunction ["Hcmf", "H"]
- `nbody`: max number of nbody terms in the Hamiltonian used for creating FOIS
- `thresh_var`: Compression threshold for the variational solution
- `thresh_foi`: Compression threshold for the FOIS
- `thresh_pt`: Compression threshold for the first-order wavefunction (if used)
- `tol_ci`:     Convergence threshold for the CI (norm of residual)
- `resolve_ss`:  After compressing previous variational state, should we resolve in new subspace?
- `do_pt = true`: Compute pt1 wavefunction for finding updated compression basis?
- `tol_tucker`: Convergence threshold for Tucker iterations (energy change)
# Returns
- `e_var::Float64`: the final variational energy
- `v_var::CompressedTuckerState`: the final variational state

See also: [`CompressedTuckerState`](@ref), [`Tucker`](@ref)
"""
function block_sparse_tucker(input_vec::CompressedTuckerState, cluster_ops, clustered_ham;
        max_iter    = 20,
        max_iter_pt = 200, # max number of iterations for solving for PT1
        nbody       = 4,
        H0          = "Hcmf",
        thresh_var  = 1e-4,
        thresh_foi  = 1e-6,
        thresh_pt   = 1e-5,
        tol_ci      = 1e-5,
	resolve_ss  = true,
        do_pt       = true,
        tol_tucker  = 1e-6)
    #={{{=#
    e_last = 0.0
    e0     = 0.0
    e_var  = 0.0
    e_pt2  = 0.0
    ref_vec = deepcopy(input_vec)
    clustered_S2 = extract_S2(input_vec.clusters)

    to = TimerOutput()

    for iter in 1:max_iter
        println(" --------------------------------------------------------------------")
        println(" Iterate PT-Var:       Iteration #: ",iter)
        println(" --------------------------------------------------------------------")

        #
        # Compress Variational Wavefunction
        dim1 = length(ref_vec)
        norm1 = orth_dot(ref_vec, ref_vec)
        ref_vec = compress(ref_vec, thresh=thresh_var)
        normalize!(ref_vec)
        dim2 = length(ref_vec)
        norm2 = orth_dot(ref_vec, ref_vec)
        @printf(" Compressed Ref state from: %8i → %8i (thresh = %8.1e)\n", dim1, dim2, thresh_var)
        @printf(" Norm of compressed state: %12.8f \n", norm2)

        #
        # Solve variationally in reference space
        if resolve_ss
            println()
            @printf(" Solve zeroth-order problem. Dimension = %10i\n", length(ref_vec))
            @timeit to "CI small" e0, ref_vec = tucker_ci_solve(ref_vec, cluster_ops, clustered_ham, tol=tol_ci)
        end
        #       sig = deepcopy(ref_vec)
        #       zero!(sig)
        #       build_sigma!(sig, ref_vec, cluster_ops, clustered_ham)
        #       e0 = orth_dot(ref_vec, sig)
        if iter == 1
            e_last = e0
        end

        tmp = deepcopy(ref_vec)
        zero!(tmp)
        @timeit to "S2" build_sigma!(tmp, ref_vec, cluster_ops, clustered_S2)
        @printf(" <S^2> = %12.8f\n", orth_dot(tmp,ref_vec))

        #
        # Get First order wavefunction
        println()
        println(" Compute first order wavefunction. Reference space dim = ", length(ref_vec))
        @timeit to "FOIS" pt1_vec  = build_compressed_1st_order_state(ref_vec, cluster_ops, clustered_ham, nbody=nbody, thresh=thresh_foi)

        #
        # Compress FOIS
        norm1 = orth_dot(pt1_vec, pt1_vec)
        dim1 = length(pt1_vec)

        if length(pt1_vec) == 0
            error("zero length FOIS?")
        end
        @timeit to "compress" pt1_vec = compress(pt1_vec, thresh=thresh_foi)
        norm2 = orth_dot(pt1_vec, pt1_vec)
        dim2 = length(pt1_vec)
        @printf(" FOIS Compressed from:     %8i → %8i (thresh = %8.1e)\n", dim1, dim2, thresh_foi)
        @printf(" Norm of |1>:              %12.8f \n", norm2)
        @printf(" Overlap between <1|0>:    %8.1e\n", nonorth_dot(pt1_vec, ref_vec, verbose=0))

        if do_pt
            #
            #
            println()
            println(" Compute PT vector. Reference space dim = ", length(ref_vec))
            @timeit to "PT1" pt1_vec, e_pt2= hylleraas_compressed_mp2(pt1_vec, ref_vec, cluster_ops, clustered_ham; tol=tol_ci, do_pt=do_pt, max_iter=max_iter_pt, H0=H0)
            #
            # Compress first order wavefunction
            norm1 = orth_dot(pt1_vec, pt1_vec)
            dim1 = length(pt1_vec)
            @timeit to "compress" pt1_vec = compress(pt1_vec, thresh=thresh_pt)
            norm2 = orth_dot(pt1_vec, pt1_vec)
            dim2 = length(pt1_vec)
            @printf(" PT   Compressed from:     %8i → %8i (thresh = %8.1e)\n", dim1, dim2, thresh_pt)
            @printf(" Norm of |1>:              %12.8f \n", norm2)
            @printf(" Overlap between <1|0>:    %8.1e\n", nonorth_dot(pt1_vec, ref_vec, verbose=0))
        end

        #
        # Solve variationally in compressed FOIS
        # CI
        println()
        var_vec = deepcopy(ref_vec)
        #zero!(pt1_vec)
        @timeit to "add" nonorth_add!(var_vec, pt1_vec)
        norm1 = orth_dot(var_vec, var_vec)
        dim1 = length(var_vec)
        @timeit to "compress" var_vec = compress(var_vec, thresh=thresh_pt)
        dim2 = length(var_vec)
        @printf(" Var  Compressed from:     %8i → %8i (thresh = %8.1e)\n", dim1, dim2, thresh_pt)
        normalize!(var_vec)
        @printf(" Solve in compressed FOIS. Dimension =   %10i\n", length(var_vec))
        @timeit to "CI big" e_var, var_vec = tucker_ci_solve(var_vec, cluster_ops, clustered_ham, tol=tol_ci)

        tmp = deepcopy(var_vec)
        zero!(tmp)
        @timeit to "S2" build_sigma!(tmp, var_vec, cluster_ops, clustered_S2)
        @printf(" <S^2> = %12.8f\n", orth_dot(tmp,var_vec))

        ref_vec = var_vec

        @printf(" E(Ref)      = %12.8f\n", e0[1])
        do_pt == false || @printf(" E(PT2) tot  = %12.8f\n", e_pt2)
        @printf(" E(var) tot  = %12.8f\n", e_var[1])
    	show(to)
        println("")

        if abs(e_last[1] - e_var[1]) < tol_tucker
            println("*Converged")
            return e_var, ref_vec
            break
        end
        e_last = e_var

    end
    println(" Not converged")
    show(to)
    println("")
    return e_var,ref_vec
end
#=}}}=#

max_stars_repo_name: nmayhall-vt/FermiCG

max_issues_repo_name: nmayhall-vt/FermiCG

max_forks_repo_name: nmayhall-vt/FermiCG


text:
function SecantSolver()
    return SecantSolver(())
end

function SecantSolver(arg0::jdouble)
    return SecantSolver((jdouble,), arg0)
end

function SecantSolver(arg0::jdouble, arg1::jdouble)
    return SecantSolver((jdouble, jdouble), arg0, arg1)
end


text:
using Breadth
using Test

@testset "Baseline" begin
    @test 1 == 1
    x = Variable("x", Sets)
    println(x)
    println(x ∨ x)
    println(~x)
    println(ConditionSet(x, x))
end


text:
using RecordedArrays
using RecordedArrays.ResizingTools
using Random
using LoopVectorization
using EvolutionaryModelingTools
using EvolutionaryModelingTools: sample

const T = 10.0
const r = 0.5
const K = 100
const μ = 0.01
const v = [10;]
const m = ones(1, 1)

# simulate with this package
## define reactions

# define growth with @reaction_eq, loop vectorization is disabled but inbounds, fastmath and  offset are enable
@reaction_eq growth r * (1 - μ) v[i] → 2v[i] avx=false inbounds=true fastmath=true offset=true

# define mutation with @reaction and @quickloop with given index
@reaction mutation begin
    @quickloop r * μ * v[i]
    begin
        push!(v, 1)
        n = length(v)
        resize!(m, (n, n))
        # rand competition matrix
        m[:, n] = rand(rng, n)
        m[n, 1:n-1] = rand(rng, 1, n - 1)
    end
end

# define competition with @reaction and @quickloop
@reaction competition begin
    @quickloop (r / K) * v[i] * m[i, j] * v[j] turbo # force enable loop vectorization
    begin
        i = ind[1]
        v[i] -= 1
        if v[i] == 0 # check extinction
            deleteat!(v, i)
            resize!(m, (not(i), not(i)))
        end
    end
end

# run functions for test and benchmark
function run_gillespie(rng=Random.GLOBAL_RNG, T=T, r=r, K=K, μ=μ, v=v, m=m)
    ps = (; r, K, μ, v, m = SimpleRDArray(m))
    return gillespie(rng, T, ps, (growth, mutation, competition))
end

function run_gillespie_record(rng=Random.GLOBAL_RNG, T=T, r=r, K=K, μ=μ, v=v, m=m)
    clock = ContinuousClock(T)
    ps = (;
        r,
        K,
        μ,
        v=recorded(DynamicEntry, clock, v),
        m=recorded(StaticEntry, clock, m),
    )
    return gillespie(rng, clock, ps, (growth, mutation, competition))
end

# simulate manually
function run_manually(rng=Random.GLOBAL_RNG, T=T, r=r, K=K, μ=μ, v=v, m=m)
    t = 0.0
    v′ = deepcopy(v)
    m′ = SimpleRDArray(m)
    while t <= T
        growth_rate = r * (1 - μ) * v′
        mutation_rate = r * μ * v′
        c = r / K # baseline competition coefficient
        competition_rate = @. c * v′ * m′ * v′'
        rates_sum = (sum(growth_rate), sum(mutation_rate), sum(competition_rate))
        rates_sum_acc = accumulate(+, rates_sum)
        rates_sum_sum = rates_sum_acc[end]
        t += -log(rand(rng)) / rates_sum_sum
        rn = rand(rng) * rates_sum_sum
        ind = searchsortedfirst(collect(rates_sum_acc), rn)
        if ind == 1
            i = sample(growth_rate, rn)
            v′[i] += 1
        elseif ind == 2
            i = sample(mutation_rate, rn - rates_sum_acc[1])
            push!(v′, 1)
            n = length(v′)
            resize!(m′, (n, n))
            # rand competition matrix
            m′[:, n] = rand(rng, n)
            m′[n, 1:n-1] = rand(rng, 1, n - 1)
        elseif ind == 3
            i = sample(competition_rate, rn - rates_sum_acc[2])[1]
            v′[i] -= 1
            if v′[i] == 0 # check extinction
                deleteat!(v′, i)
                resize!(m′, (not(i), not(i)))
            end
        end
        length(v′) == 0 && break
    end
    return (; v=v′, m=m′)
end
function run_manually_record(rng=Random.GLOBAL_RNG, T=T, r=r, K=K, μ=μ, v=v, m=m)
    clock = ContinuousClock(T)
    v′ = recorded(DynamicEntry, clock, v)
    m′ = recorded(StaticEntry, clock, m)
    for _ in clock
        growth_rate = r * (1 - μ) * v′
        mutation_rate = r * μ * v′
        c = r / K # baseline competition coefficient
        competition_rate = @. c * v′ * m′ * v′'
        rates_sum = (sum(growth_rate), sum(mutation_rate), sum(competition_rate))
        rates_sum_acc = accumulate(+, rates_sum)
        rates_sum_sum = rates_sum_acc[end]
        τ = -log(rand(rng)) / rates_sum_sum
        increase!(clock, τ)
        rn = rand(rng) * rates_sum_sum
        ind = searchsortedfirst(collect(rates_sum_acc), rn)
        if ind == 1
            i = sample(growth_rate, rn)
            v′[i] += 1
        elseif ind == 2
            i = sample(mutation_rate, rn - rates_sum_acc[1])
            push!(v′, 1)
            n = length(v′)
            resize!(m′, (n, n))
            # rand competition matrix
            m′[:, n] = rand(rng, n)
            m′[n, 1:n-1] = rand(rng, 1, n - 1)
        elseif ind == 3
            i = sample(competition_rate, rn - rates_sum_acc[2])[1]
            v′[i] -= 1
            if v′[i] == 0 # check extinction
                deleteat!(v′, i)
                resize!(m′, (not(i), not(i)))
            end
        end
        length(v′) == 0 && break
    end
    return (; v=v′, m=m′)
end


text:
using ModelingToolkitStandardLibrary.Electrical, ModelingToolkit, OrdinaryDiffEq, Test
using ModelingToolkitStandardLibrary.Blocks: Step, Constant, Sine, Cosine, ExpSine, Ramp, Square, Triangular
using ModelingToolkitStandardLibrary.Blocks: square, triangular
# using Plots

@parameters t

@testset "sensors" begin
    @named source = Sine(offset=1, amplitude=10, frequency=5)
    @named voltage = Voltage()
    @named resistor = Resistor(R=1)
    @named capacitor = Capacitor(C=1)
    @named ground = Ground()

    @named voltage_sensor = VoltageSensor()
    @named current_sensor = CurrentSensor()
    @named power_sensor = PowerSensor()

    connections = [
        connect(source.output, voltage.V)
        connect(voltage.p, resistor.p)
        connect(resistor.n, current_sensor.p)
        connect(current_sensor.n, power_sensor.pc)
        connect(power_sensor.nc, capacitor.p)
        connect(capacitor.n, voltage.n, ground.g)
        connect(capacitor.p, voltage_sensor.p)
        connect(capacitor.n, voltage_sensor.n)
        connect(capacitor.p, power_sensor.pv)
        connect(capacitor.n, power_sensor.nv)
    ]

    @named model = ODESystem(connections, t; systems=[resistor, capacitor, source, voltage, ground, voltage_sensor, current_sensor, power_sensor])
    sys = structural_simplify(model)
    prob = ODAEProblem(sys, Pair[], (0.0, 10.0))
    sol = solve(prob, Tsit5())

    # Plots.plot(sol; vars=[capacitor.v, voltage_sensor.v])
    # Plots.plot(sol; vars=[power_sensor.power, capacitor.i * capacitor.v])
    # Plots.plot(sol; vars=[resistor.i, current_sensor.i])
    @test sol[capacitor.v] ≈ sol[voltage_sensor.v] atol=1e-3
    @test sol[power_sensor.power] ≈ sol[capacitor.i * capacitor.v] atol=1e-3
    @test sol[resistor.i] ≈ sol[current_sensor.i] atol=1e-3
end

# simple voltage divider
@testset "voltage divider" begin
    @named source = Constant(k=10)
    @named voltage = Voltage()
    @named R1 = Resistor(R=1e3)
    @named R2 = Resistor(R=1e3)
    @named ground = Ground()

    connections = [
        connect(source.output, voltage.V)
        connect(voltage.p, R1.p)
        connect(R1.n, R2.p)
        connect(R2.n, voltage.n, ground.g)
    ]

    @named model = ODESystem(connections, t, systems=[R1, R2, source,voltage, ground])
    sys = structural_simplify(model)
    prob = ODEProblem(sys, Pair[], (0, 2.0))
    sol = solve(prob, Rodas4()) # has no state; does not work with Tsit5
    @test sol[R1.p.v][end] ≈ 10 atol=1e-3
    @test sol[R1.n.v][end] ≈ 5 atol=1e-3
    @test sol[R2.n.v][end] ≈ 0 atol=1e-3
end

# simple RC
@testset "RC" begin
    @named source = Constant(k=10)
    @named voltage = Voltage()
    @named resistor = Resistor(R=1)
    @named capacitor = Capacitor(C=1)
    @named ground = Ground()

    connections = [
        connect(source.output, voltage.V)
        connect(voltage.p, resistor.p)
        connect(resistor.n, capacitor.p)
        connect(capacitor.n, voltage.n, ground.g)
    ]

    @named model = ODESystem(connections, t; systems=[resistor, capacitor, source, voltage, ground])
    sys = structural_simplify(model)
    prob = ODAEProblem(sys, [capacitor.v => 0.0], (0.0, 10.0))
    sol = solve(prob, Tsit5())

    # Plots.plot(sol; vars=[source.v, capacitor.v])
    @test sol[capacitor.v][end] ≈ 10 atol=1e-3
end

# simple RL
@testset "RL" begin
    @named source = Constant(k=10)
    @named voltage = Voltage()
    @named resistor = Resistor(R=1)
    @named inductor = Inductor(L=1.0)
    @named ground = Ground()

    connections = [
        connect(source.output, voltage.V)
        connect(voltage.p, resistor.p)
        connect(resistor.n, inductor.p)
        connect(inductor.n, voltage.n, ground.g)
    ]

    @named model = ODESystem(connections, t; systems=[resistor, inductor, source, voltage, ground])
    sys = structural_simplify(model)
    prob = ODAEProblem(sys, [inductor.i => 0.0], (0.0, 10.0))
    sol = solve(prob, Tsit5())

    # Plots.plot(sol; vars=[inductor.i, inductor.i])
    @test sol[inductor.i][end] ≈ 10 atol=1e-3
end

@testset "RC with voltage sources" begin
    R, C = 1, 1
    @named voltage = Voltage()
    @named source_const = Constant(k=10)
    @named source_sin = Sine(offset=1, amplitude=10, frequency=2, start_time=0.5, phase=0)
    @named source_step = Step(offset=1, height=10, start_time=0.5)
    @named source_tri = Triangular(offset=1, start_time=0.5, amplitude=10, frequency=2)
    @named source_dsin = ExpSine(offset=1, amplitude=10, frequency=2, start_time=0.5, phase=0, damping=0.5)
    @named source_ramp = Ramp(offset=1, height=10, start_time=0.5, duration=1)
    sources = [source_const, source_sin, source_step, source_tri, source_dsin, source_ramp]

    @named resistor = Resistor(; R)
    @named capacitor = Capacitor(; C)
    @named ground = Ground()

    for source in sources
        connections = [
            connect(source.output, voltage.V)
            connect(voltage.p, resistor.p)
            connect(resistor.n, capacitor.p)
            connect(capacitor.n, voltage.n, ground.g)
        ]

        @named model = ODESystem(connections, t; systems=[resistor, capacitor, source, ground, voltage])
        sys = structural_simplify(model)
        prob = ODAEProblem(sys, [capacitor.v => 10.0], (0.0, 10.0))
        @test_nowarn sol = solve(prob, Rodas5())
        @test_nowarn sol = solve(prob, Tsit5())

        # Plots.plot(sol; vars=[voltage.v, capacitor.v])
    end
end

# RC with current sources
@testset "RC with current sources" begin
    start_time = 2
    @named current = Current()
    @named source = Step(start_time=2)
    @named resistor = Resistor(R=1)
    @named capacitor = Capacitor(C=1)
    @named ground = Ground()

    connections = [
        connect(source.output, current.I)
        connect(current.p, resistor.n)
        connect(capacitor.n, resistor.p)
        connect(capacitor.p, current.n, ground.g)
    ]

    @named model = ODESystem(connections, t; systems=[ground, resistor,  current, capacitor, source])
    sys = structural_simplify(model)
    prob = ODAEProblem(sys, [capacitor.v => 0.0], (0.0, 10.0))
    sol = solve(prob, Tsit5())
    y(x, st) = (x .> st) .* abs.(collect(x) .- st)
    @test sum(reduce(vcat, sol.u) .- y(sol.t, start_time)) ≈ 0 atol=1e-2

end

@testset "Integrator" begin
    R=1e3
    f=1
    Vin=5
    @named ground = Ground()
    @named R1 = Resistor(R=R)
    @named R2 = Resistor(R=100*R)
    @named C1 = Capacitor(C=1/(2 * pi * f * R))
    @named opamp = IdealOpAmp()
    @named square = Square(amplitude=Vin)
    @named voltage = Voltage()
    @named sensor = VoltageSensor()

    connections = [
        connect(square.output, voltage.V)
        connect(voltage.p, R1.p)
        connect(R1.n, C1.n, R2.p, opamp.n1)
        connect(opamp.p2, C1.p, R2.n)
        connect(opamp.p1, ground.g, opamp.n2, voltage.n)
        connect(opamp.p2, sensor.p)
        connect(sensor.n, ground.g)
    ]
    @named model = ODESystem(connections, t, systems = [R1, R2, opamp, square, voltage, C1, ground, sensor])
    sys = structural_simplify(model)
    u0 = [
        C1.v => 0.0
        R1.v => 0.0
    ]
    prob = ODEProblem(sys, u0, (0, 100.0))
    sol = solve(prob, Rodas4())
    @test sol[opamp.v2] == sol[C1.v] # Not a great one however. Rely on the plot
    @test sol[opamp.p2.v] == sol[sensor.v]

    # plot(sol, vars=[sensor.v, square.v, C1.v])
end

_step(x, h, st) = ifelse(x < st, 0, h)
_cos_wave(x, f, A, st, ϕ) = A*cos(2*π*f*(x-st) + ϕ)
_ramp(x, st, d, h) = ifelse(x < st, 0,
                                  ifelse(x < (st + d), (x - st) * h / d, h))
_sine_wave(x, f, A, st, ϕ) = A*sin(2*π*f*(x - st) + ϕ)
_damped_sine_wave(x, f, A, st, ϕ, d) = exp((st-x)*d)*A*sin(2*π*f*(x-st) + ϕ)

@testset "Voltage function generators" begin
    st, o, h, f, A, et, ϕ, d, δ = 0.7, 1.25, 3, 2, 2.5, 2, π/4, 0.1, 0.0001

    @named res = Resistor(R=1)
    @named cap = Capacitor(C=1)
    @named ground = Ground()
    @named voltage = Voltage()
    @named voltage_sensor = VoltageSensor()
    @named step = Step(start_time=st, offset=o, height=h)
    @named cosine = Cosine(offset=o, amplitude=A, frequency=f, start_time=st, phase=ϕ)
    @named sine = Sine(offset=o, amplitude=A, frequency=f, start_time=st, phase=ϕ)
    @named damped_sine = ExpSine(offset=o, amplitude=A, frequency=f, start_time=st, phase=ϕ, damping=d)
    @named ramp   = Ramp(offset=o, start_time=st, duration=et-st, height=h)
    @named vsquare = Square(offset=o, start_time=st, amplitude=A, frequency=f)
    @named tri = Triangular(offset=o, start_time=st, amplitude=A, frequency=f)
    # @named vsawtooth = SawTooth(amplitude=A, start_time=st, frequency=f, offset=o)

    sources = [step, cosine, sine, damped_sine, ramp, tri, vsquare] #, vsawtooth]
    waveforms(i, x) = getindex([o .+  _step.(x, h, st),
                                o .+ (x .> st) .* _cos_wave.(x, f, A, st, ϕ),
                                o .+ (x .> st) .* _sine_wave.(x, f, A, st, ϕ),
                                o .+ (x .> st) .* _damped_sine_wave.(x, f, A, st, ϕ, d),
                                o .+ _ramp.(x, st, (et-st), h),
                                triangular.(x, f, A, o, st),
                                square.(x, f, A, o, st)], i)
                                # o .+ (x .> st). * _sawtooth_wave.(x, δ, f, A, st),

    for i in 1:length(sources)
        source = sources[i]
        @info "Testing Voltage with $(source.name) source"
        eqs = [
            connect(source.output, voltage.V)
            connect(voltage.p, voltage_sensor.p, res.p)
            connect(res.n, cap.p)
            connect(ground.g, voltage_sensor.n, voltage.n, cap.n)
        ]
        @named vmodel = ODESystem(eqs, t, systems = [voltage_sensor, res, cap, source, voltage, ground])
        vsys = structural_simplify(vmodel)

        u0 = []

        prob = ODAEProblem(vsys, u0, (0, 10.0))
        sol = solve(prob, dt=0.1, Tsit5())

        @test sol[voltage.V.u] ≈ waveforms(i, sol.t) atol=1e-1
        @test sol[voltage.p.v] ≈ sol[voltage.V.u]
        # For visual inspection
        # plt = plot(sol; vars=[voltage.v])
        # savefig(plt, "test_voltage_$(source.name)")
    end
end

@testset "Current function generators" begin
    st, o, h, f, A, et, ϕ, d, δ = 0.7, 1.25, 3, 2, 2.5, 2, π/4, 0.1, 0.0001

    @named ground = Ground()
    @named res = Resistor(R=1.0)
    @named cap = Capacitor(C=1)
    @named current_sensor = CurrentSensor()
    @named current = Current()
    @named step = Step(start_time=st, offset=o, height=h)
    @named cosine = Cosine(offset=o, amplitude=A, frequency=f, start_time=st, phase=ϕ)
    @named sine = Sine(offset=o, amplitude=A, frequency=f, start_time=st, phase=ϕ)
    @named damped_sine = ExpSine(offset=o, amplitude=A, frequency=f, start_time=st, phase=ϕ, damping=d)
    @named ramp = Ramp(offset=o, start_time=st, duration=et-st, height=h)
    @named vsquare = Square(offset=o, start_time=st, amplitude=A, frequency=f)
    @named tri = Triangular(offset=o, start_time=st, amplitude=A, frequency=f)
    # @named isawtooth = SawTooth(amplitude=A, start_time=st, frequency=f, offset=o)

    sources = [step, cosine, sine, damped_sine, ramp, tri, vsquare] #, idamped_sine]
    waveforms(i, x) = getindex([o .+  _step.(x, h, st),
                                o .+ (x .> st) .* _cos_wave.(x, f, A, st, ϕ),
                                o .+ (x .> st) .* _sine_wave.(x, f, A, st, ϕ),
                                o .+ (x .> st) .* _damped_sine_wave.(x, f, A, st, ϕ, d),
                                o .+ _ramp.(x, st, (et-st), h),
                                triangular.(x, f, A, o, st),
                                square.(x, f, A, o, st)], i)
                                # # o .+ (x .> st). * _sawtooth_wave.(x, δ, f, A, st)

    for i in 1:length(sources)
        source = sources[i]
        @info "Testing Current with $(source.name) source"
        eqs = [
            connect(source.output, current.I)
            connect(current.p, current_sensor.n)
            connect(current_sensor.p, res.p)
            connect(res.n, cap.p)
            connect(current.n, ground.g, cap.n)
        ]
        @named model = ODESystem(eqs, t, systems = [current_sensor, source, current, res, cap, ground])
        isys = structural_simplify(model)

        u0 = []

        prob = ODAEProblem(isys, u0, (0, 10.0))
        sol = solve(prob, dt=0.1, Tsit5())

        @test sol[current.I.u] ≈ waveforms(i, sol.t) atol=1e-1
        @test sol[current.I.u] ≈ sol[current.p.i] atol=1e-1
        # For visual inspection
        # plt = plot(sol)
        # savefig(plt, "test_current_$(source.name)")
    end
end


text:
# License for this file: MIT (expat)
# Copyright 2017-2018, DLR Institute of System Dynamics and Control
#
# This file is part of module
#   Modia3D.Composition (Modia3D/Composition/_module.jl)
#
# Functions that have to be provided from an object used for contact pair response

#contactStart(material::Modia3D.AbstractContactMaterial, obj1, obj2, rContact, contactNormal, elasticContactReductionFactor) =
#             error("Function contactStart not defined for material ", String(typeof(material)))

#contactEnd(material::Modia3D.AbstractContactMaterial, obj1, obj2) =
#             error("Function contactEnd not defined for material ", String(typeof(material)))

#responseCalculation(material::Modia3D.AbstractContactMaterial, obj1, obj2, rContact, contactNormal, distanceWithHysteresis, time) = error("Function responseCalculation not defined for material ", String(typeof(material)))

# Utility functions

"""
    responseMaterial = contactStart(obj1, obj2, rContact, contactNormal)

Return a response material object at contact start.
"""
function contactStart(obj1::Object3D, obj2::Object3D, rContact::SVector{3,Float64},
                      contactNormal::SVector{3,Float64}, elasticContactReductionFactor::Float64)
    name1 = obj1.feature.contactMaterial
    name2 = obj2.feature.contactMaterial
    material = Shapes.getContactPairMaterial(name1, name2)
    return contactStart(material, obj1, obj2, rContact, contactNormal, elasticContactReductionFactor)
end

"""
    delta_dot = normalRelativeVelocityAtContact(obj1, obj2, rContact, contactNormal)

Return the relative velocity in normal direction `contactNormal` at
contact point `rContact` of the two penetrating objects `obj1, obj2`.
"""
function normalRelativeVelocityAtContact(obj1::Object3D, obj2::Object3D,
                                         rContact::Frames.Vector3D, e_n::Frames.Vector3D)
  r_rel1 = rContact - obj1.r_abs
  r_rel2 = rContact - obj2.r_abs
  w1     = obj1.R_abs'*obj1.w
  w2     = obj2.R_abs'*obj2.w
  v1     = obj1.v0 + cross(w1,r_rel1)
  v2     = obj2.v0 + cross(w2,r_rel2)
  v_rel  = v2 - v1
  delta_dot = dot(v_rel,e_n)
  return delta_dot
end


text:
using Test, MonteCarloZoo, Random, Plots, SpecialFunctions
import Distributions: Kolmogorov, cdf

function ks_test(samples, target_cdf)
    sorted_samples = sort(samples, dims=2)
    D = 0
    for (i, (s, t)) in enumerate(zip(sorted_samples[:, 1:end - 1],
                                     sorted_samples[:, 1:end]))
        D = max(D,
                abs(target_cdf(s) - i / length(samples)),
                abs(target_cdf(t) - i / length(samples)))
    end
    K = D * sqrt(length(samples))
    p = 1 - cdf(Kolmogorov(), K)
    return p
end

## Rejection Sampling

f(x) = exp(-x^2 / 2) / sqrt(2π)
g(x) = 1 / (π * (1 + x^2))
proposal_sampler = InverseTransformSampler(u -> tan(π * (u - 0.5)))
M = sqrt(2π / ℯ)
N = 1000

s = RejectionSampler(f, g, proposal_sampler, M)
Random.seed!(1729)
samples = sample(s, N)

@test size(samples) == (1, N)
p = ks_test(samples, x -> 1 / 2 * (1 + erf(x / sqrt(2))))
@test p > 0.05

## Inverse Transform Sampling

F_inv(u) = -log(u)  # using symmetry of U (see notes)
N = 1000

s = InverseTransformSampler(F_inv)
Random.seed!(1729)
samples = sample(s, N)

@test size(samples) == (1, N)
p = ks_test(samples, x -> 1 - exp(-x))
@test p > 0.05

## Box-Muller Transform Sampling

function test_box_muller()
    for N in (1000, 1001)
        Random.seed!(1729)
        s = BoxMullerTransformSampler()
        samples = sample(s, N)

        @test size(samples) == (1, N)
        p = ks_test(samples, x -> 1 / 2 * (1 + erf(x / sqrt(2))))
        @test p > 0.05
    end
end

test_box_muller()

## Linear Congruential Generator

a = 1664525
c = 1013904223
m = 2^32
seed = 1729
N = 1000

s = LinearCongruentialGenerator(a, c, m, seed)
samples = sample(s, N)

@test size(samples) == (1, N)
p = ks_test(samples, x -> x * (0 ≤ x ≤ 1))
@test p > 0.05

## Normal-Normal Transformer

μ₁ = 1.0
σ² = 2.0
normal_sampler = BoxMullerTransformSampler()
N = 1000

t1 = NormalNormalTransformer(normal_sampler, (μ₁, σ²))
Random.seed!(1729)
samples = sample(t1, N)

@test size(samples) == (1, N)
p = ks_test(samples, x -> 1 / 2 * (1 + erf((x - μ₁) / sqrt(2 * σ²))))
@test p > 0.05

μ₂ = 2.0

t2 = NormalNormalTransformer(t1, (μ₂, σ²), (μ₁, σ²))
Random.seed!(1729)
samples = sample(t2, N)

@test size(samples) == (1, N)
p = ks_test(samples, x -> 1 / 2 * (1 + erf((x - μ₂) / sqrt(2 * σ²))))
@test p > 0.05


text:
@testset "tracked: accessors" begin

    # TODO: move to properties
    tape = Tape()
    r = rand()
    tr = tracked(tape, r)

    @test gettape(tr) == tape
    @test getid(tr) == -1
    @test getvalue(tr) == r
    setvalue!(tr, r + 1.0)
    @test getvalue(tr) == r + 1.0

    a = rand(3,4)
    ta = tracked(tape, a)

    @test gettape(ta) == tape
    @test getid(ta) == -1
    @test getvalue(ta) == a
    setvalue!(ta, a .+ 1)
    @test getvalue(ta) == a .+ 1

end

@testset "tracked: ops" begin

    tape = Tape()
    xv = rand(); yv = rand()
    x = tracked(tape, xv)
    y = tracked(tape, yv)
    z = 2x^2 + y
    @test getvalue(z) == 2xv^2 + yv
    @test length(tape) == 5

    tape = Tape()
    xv = rand(3,4); yv = rand(4,3)
    x = tracked(tape, xv); y = tracked(tape, yv)
    z = x * y .+ 1.0
    @test getvalue(z) == xv * yv .+ 1.0
    @test length(tape) == 3

    # tracked * untracked
    tape = Tape()
    x = tracked(tape, rand(3,3))
    y = rand(3,3)
    z = x * y
    @test length(tape) == 2
    @test z.val == x.val * y

end


text:
using Plots

#graphs a cylindrical aprx with n planes, radius =r, height = h
function cylinder(n,r,h)
	p = plot(xlims=[-h,h],ylims=[-2*r,2*r],zlims=[-2*r,2*r],legend=false)
	x = range(0,h,length=10)
	alpha = range(0,2*pi,length=n)

	for theta in alpha
		if theta == alpha[end]
			break
		end
		y=r*cos(theta)*ones(length(x))
		z=r*sin(theta)*ones(length(x))
		plot!(x,y,z,seriescolor=:black)
	end
	display(p)
end

#graphs a function to be revolved, at angles theta_i in alpha
#use at least n=10,15
function PolarLines(a,b,n)

	#f is the function to revolve
	f(u) = 1+sin(u)cos(u)
	p = plot(legend=false)
	x=range(a,b,length=10)
	alpha = range(0,2*pi,length=n)
	for theta in alpha
		if theta == alpha[end]
			break
		end
		y=cos(theta)*f.(x)
		z=sin(theta)*f.(x)
		plot!(x,y,z,seriescolor=:black)
	end
	savefig(p,"fig.html")
	p
end

#graphs a function to be revolved, at vertical slices, and with PolarLines
#use at least n=10
function Rings(a,b,n,m)

	#f is the function to revolve
	f(u) = 1+sin(u)cos(u)
	q = plot([f,0],a,b)

	p = plot(legend=false)
	x=range(a,b,length=n)
	alpha = range(0,2*pi,length=m)

	for xi in x
		y = Float64[]
		z = Float64[]
		xp = xi*ones(m)
		for theta in alpha
			push!(y,cos(theta)*f(xi))
			push!(z,sin(theta)*f(xi))

		end
		plot!(p,xp,y,z,seriescolor=:black)

	end
	p1 = PolarLines(a,b,2*m)

	f = plot(
	p,p1,q,
	layout=@layout([a b;c]),
	legend=false,
	titlelocation=:left,
	bottom_margin=10px,
	window_title="Solid of revolution")
	savefig(f,"figures.html")
	f
end


text:
function Random.Sampler(T::Type{<:Random.AbstractRNG}, x::Consistent, ::Val{Inf})
    # sampler for sampling indices of predictions
    predictions = x.predictions
    idxsampler = Random.Sampler(T, eachindex(predictions))
    return Random.SamplerSimple(x, (idxsampler, make_samplers(T, x)))
end

function make_samplers(T::Type{<:Random.AbstractRNG}, x::Consistent{<:Real})
    return Random.Sampler(T, float(eltype(x.predictions)))
end

# set up alias samplers
function make_samplers(
    T::Type{<:Random.AbstractRNG}, x::Consistent{<:AbstractVector{<:Real}}
)
    predictions = x.predictions
    prediction = first(predictions)
    floatsampler = Random.Sampler(T, float(eltype(prediction)))
    integersampler = Random.Sampler(T, Base.OneTo(length(prediction)))
    samplers = let floatsampler = floatsampler, integersampler = integersampler
        map(predictions) do prediction
            return AliasSampler(prediction, floatsampler, integersampler)
        end
    end
    return samplers
end

function make_samplers(T::Type{<:Random.AbstractRNG}, x::Consistent)
    samplers = let T = T
        map(x.predictions) do prediction
            return Random.Sampler(T, prediction)
        end
    end
    return samplers
end

# out-of-place sampling
function Base.rand(
    rng::Random.AbstractRNG, sampler::Random.SamplerSimple{<:Consistent{<:Real}}
)
    idxsampler, floatsampler = sampler.data
    idx = rand(rng, idxsampler)
    @inbounds prediction = sampler[].predictions[idx]
    @inbounds target = rand(rng, floatsampler) < prediction
    return (prediction, target)
end
function Base.rand(rng::Random.AbstractRNG, sampler::Random.SamplerSimple{<:Consistent})
    idxsampler, samplers = sampler.data
    idx = rand(rng, idxsampler)
    @inbounds prediction = sampler[].predictions[idx]
    @inbounds target = rand(rng, samplers[idx])
    return (prediction, target)
end

# in-place sampling
function Random.rand!(
    rng::Random.AbstractRNG,
    out::AbstractArray{<:Tuple{<:Any,<:Any}},
    sampler::Random.SamplerSimple{<:Consistent{<:Real}},
)
    predictions = sampler[].predictions
    idxsampler, floatsampler = sampler.data
    for i in eachindex(out)
        idx = rand(rng, idxsampler)
        @inbounds prediction = predictions[idx]
        @inbounds target = rand(rng, floatsampler) < prediction
        @inbounds out[i] = (prediction, target)
    end
    return out
end
function Random.rand!(
    rng::Random.AbstractRNG,
    out::AbstractArray{<:Tuple{<:Any,<:Any}},
    sampler::Random.SamplerSimple{<:Consistent},
)
    predictions = sampler[].predictions
    idxsampler, samplers = sampler.data
    for i in eachindex(out)
        idx = rand(rng, idxsampler)
        @inbounds prediction = predictions[idx]
        @inbounds target = rand(rng, samplers[idx])
        @inbounds out[i] = (prediction, target)
    end
    return out
end


text:
# Body for `cell.output.body` of
# using PlutoUI
# f(x) = Base.inferencebarrier(x);
# with_terminal() do
#     @code_warntype f(1)
# end
body = raw"""

    """

txt = strip(PlutoStaticHTML._extract_txt(body))
@test startswith(txt, "MethodInstance")
@test endswith(txt, "return %2")
@test contains(txt, '\n')

patched = PlutoStaticHTML._patch_with_terminal(body)
@test contains(patched, """
""")


text:
# Adapted from https://discourse.julialang.org/t/expression-parser/41880/7
# code by Alan R. Rogers, Professor of Anthropology, University of Utah

# Adapted from https://github.com/x3042/Exact-reduction-of-ODE-systems/blob/main/src/myeval.jl
# by Alexander D. , Higher School of Economics

function myeval(e::Union{Expr,Symbol,Number}, map::Dict{Symbol,MPoly})
    try
        return _myeval(e, map)
    catch ex
        println("Can't parse \"$e\"")
        rethrow(ex)
    end
end

function _myeval(s::Symbol, map::Dict{Symbol,MPoly})
    if haskey(map, s)
        return map[s]
    else
        throw(UndefVarError(s))
    end
end

# Numbers are converted to QQ.
function _myeval(x::Number, map::Dict{Symbol,MPoly})
    return QQ(x)
end

# a helper definition for floats
function _myeval(x::Float64, map::Dict{Symbol, MPoly})
    result = QQ(0)

    # Getting the result from the string representation in order
    # to avoid approximations caused by the float representation
    s = string(x)
    denom = 1
    extra_num = 1
    if occursin(r"[eE]", s)
        s, exp = split(s, r"[eE]")
        if exp[1] == "-"
            denom = QQ(10)^(-parse(Int, exp))
        else
            extra_num = QQ(10)^(parse(Int, exp))
        end
    end
    frac = split(s, ".")
    if length(frac) == 1
        result = QQ(parse(fmpz, s)) * extra_num // denom
    else
        result = QQ(parse(fmpz, frac[1] * frac[2])) * extra_num // (denom * 10^(length(frac[2])))
    end

    # too verbose for now
    # @warn "a possibility of inexact float conversion" from=x to=result
    return result
end

# To parse an expression, convert the head to a singleton
# type, so that Julia can dispatch on that type.
function _myeval(e::Expr, map::Dict{Symbol,MPoly})
    return _myeval(Val(e.head), e.args, map)
end

# Call the function named in args[1]
function _myeval(::Val{:call}, args, map::Dict{Symbol,MPoly})
    return _myeval(Val(args[1]), args[2:end], map)
end

# Addition
function _myeval(::Val{:+}, args, map::Dict{Symbol,MPoly})
    x = 0
    for arg ∈ args
        x += _myeval(arg, map)
    end
    return x
end

# Subtraction and negation
function _myeval(::Val{:-}, args, map::Dict{Symbol,MPoly})
    len = length(args)
    if len == 1
        return -_myeval(args[1], map)
    else
        return _myeval(args[1], map) - _myeval(args[2], map)
    end
end

# Multiplication
function _myeval(::Val{:*}, args, map::Dict{Symbol,MPoly})
    x = 1
    for arg ∈ args
        x *= _myeval(arg, map)
    end
    return x
end

# Division
function _myeval(::Val{:/}, args, map::Dict{Symbol,MPoly})
    return _myeval(args[1], map) // _myeval(args[2], map)
end

# Exponentiation
function _myeval(::Val{:^}, args, map::Dict{Symbol, MPoly})
    if 1 != denominator(_myeval(args[2], map))
        @warn "chto-to strannoe happened"
    end

    return _myeval(args[1], map) ^ Int(numerator(_myeval(args[2], map)))
end


text:
include("../opti_test_functions/testfunctions.jl")
include("../../numder/src/numder_module.jl")

module opti
using ..numder
using LinearAlgebra

include("../src/line_search.jl")

struct ToleranceOptions
    step_size_tol::Float64
    f_abs_tol::Float64
    f_rel_tol::Float64
    g_abs_tol::Float64
    max_iter::Int
end
function ToleranceOptions(;
    step_size_tol = 1e-15,
    f_abs_tol = 1e-15,
    f_rel_tol = 1e-15,
    g_abs_tol = 1e-15,
    max_iter = 100)

    ToleranceOptions(
        step_size_tol,
        f_abs_tol,
        f_rel_tol,
        g_abs_tol,
        max_iter)
end

struct BfgsOptions

end

function bfgs(fcn, x0, grad_fcn, options = ToleranceOptions() )

    n_dim = length(x0)
    H     = Matrix{Float64}(I, n_dim, n_dim)
    eye   = Matrix{Float64}(I, n_dim, n_dim)

    # first step is the steepest gradient descent
    g0        = grad_fcn(fcn, x0)
    step_size = line_search(fcn, x0, -g0)
    x1        = x0 - step_size * g0
    g1        = grad_fcn(fcn, x1)

    y            = g1 - g0
    s            = step_size * g0
    scale_factor = dot(y, s) / dot(y, y)
    for ii = 1:n_dim
        H[ii, ii] *= scale_factor
    end

    # main loop
    iter = 0
    for outer iter = 1:options.max_iter

        # update of the matrix
        s       = x1 - x0
        y       = g1 - g0
        sdy     = dot(s, y)
        rho     = 1 / sdy
        rhos    = rho * s
        Hy      = H * y
        Hyst    = Hy * rhos'
        sytHyst = rhos * (y' * Hyst)
        sst     = rhos * s'
        Hnew    = -(Hyst + Hyst') + sytHyst + sst
        Htest   = (eye - rho * s * y') * H * (eye - rho * y * s') + rho * s * s'

        dH = norm(Htest - Hnew)
        println(dH)

        x0        = x1
        p         = -(Hnew * g1)
        step_size = line_search(fcn, x0, p)
        dx        = p * step_size
        x1        = x0 + dx

        if norm(dx) <= options.step_size_tol
            break
        end

        # options.f_abs_tol
        # options.f_rel_tol

        g0        = g1
        g1        = grad_fcn(fcn, x1)

        if norm(g1) <= options.g_abs_tol
            break
        end

        H = Hnew
    end

    return (x1, iter)
end

end

module otest
using ..opti
using ..numder

grad_fcn(fcn, x) = numder.gradient_fw(fcn, x)

w = [1.0 3.0]
fcn(x) = sum(x .* x .* w)

x0 = [-2.0, -3.0]
t_sol = opti.bfgs(fcn, x0, grad_fcn)

end

max_stars_repo_name: HomoModelicus/julia

max_issues_repo_name: HomoModelicus/julia

max_forks_repo_name: HomoModelicus/julia


text:
module TestGradients
    verbose = false

    using Base.Test

    import IntertemporalChoiceHeuristics:
        load,
        fit,
        models,
        make_closures,
        initial_parameters,
        training_losses

    import ForwardDiff

    import Optim: minimizer

    inputs, weights = load()

    λ = 0.0

    @testset "Analytic gradients match AD gradients" begin
        for model in models
            for loss in training_losses
                res = fit(model, inputs, weights, loss)

                Θ = initial_parameters(model)
                Θ′ = minimizer(res)

                cost, cost_gr! = make_closures(model, loss, λ, inputs, weights)

                @test isa(cost(Θ), Float64)

                gr = similar(Θ)

                cost_gr!(Θ, gr)
                fd_gr = ForwardDiff.gradient(cost, Θ)
                @test norm(gr .- fd_gr, Inf) < 1e-10

                cost_gr!(Θ′, gr)
                fd_gr = ForwardDiff.gradient(cost, Θ′)
                @test norm(gr .- fd_gr, Inf) < 1e-10

                if verbose
                    println((model, loss, norm(gr .- fd_gr, Inf)))
                end
            end
        end
    end
end


text:
### Notation
#   small letters => vectors
# capital letters => matrices

function build_A(k, u, η, D, y, params)

    N = params.Ny
    g = params.g

    k2 = k^2
    ik = 1im*k

    # This removed the free-surface deformation in the basic state,
    # which is more like what is done in the Hydrostatic case

    h = 0*η .+ params.H

     U = diagm(0 => u)
     H = diagm(0 => h)
    dU = diagm(0 => D*u)
    dH = diagm(0 => D*h)
     F = diagm(0 => params.f₀ .+ 0*y)
     I = diagm(0 => ones(N+1))

    # Form 1L-RSW Matrix
    #   [u1, v1, h1]

    A = [    k*U         1im*(F-dU)[:, 2:N]    g*k*I;
          -1im*F[2:N,:]     k*U[2:N,2:N]    -1im*g*D[2:N,:];
             k*H         -1im*D*H[:,  2:N]       k*U];

    return A
end


text:
#################### Empirical Move ####################

#################### Types and Constructors ####################

mutable struct EmpiricalTune <: SamplerTune
    logf::Function
    samplerfun::Union{Function, Missing}
    width::Int64
    replacement::Bool

end

EmpiricalTune(x::Vector, logf::Function, samplerfun::Function, width::Int64, replacement::Bool=false)=
    EmpiricalTune(logf, samplerfun, width, replacement)

const EmpiricalVariate = Sampler{EmpiricalTune, T} where T

#################### Sampler Constructor ####################
"""

function Empirical(params::ElementOrVector{Symbol}, width::Int64; replace::Bool=false)

This Sampler generates Samples from a distribution by selecting a random value
from this distribution. The move will always be accepted. Sampling can be done with or without replacement.
"""
function Empirical(params::ElementOrVector{Symbol},
                width::Int64;
                replace::Bool=false)

    tune = EmpiricalTune(logpdf!, rand, width, replace)

    Sampler(Float64[], asvec(params), tune, Symbol[], false)
end

function sample!(v::EmpiricalVariate{T}, lpdf::Function; model::Model=model, kwargs...) where T
    width = v.tune.width
    samfun = v.tune.sampler_fun
    params = v.params
    sam = samfun(model, params, x)
    if width > 1 && v.tune.replacement
        while length(Set(sam)) != width
            sam = samfun(model, params, x)
        end
    end
    v[:] = sam
    v
end

max_stars_repo_name: erathorn/MCPhylo

max_issues_repo_name: erathorn/MCPhylo

max_forks_repo_name: erathorn/MCPhylo


text:
module HWconstrained

	using JuMP, NLopt, DataFrames, Ipopt
	using LinearAlgebra

	export data, table_NLopt, table_JuMP, max_JuMP, obj, constr, max_NLopt

	function data(a=0.5)
		na = 3
		p = ones(3)
		e = [2.0, 0.0, 0.0]
		z = zeros(16,3)
		z_2 = [0.72, 0.92, 1.12, 1.32]
		z_3 = [0.86, 0.96, 1.06, 1.16]
		z[:,1] .= 1.0
		for i=1:4
		z[1+(4*(i-1)):4+(4*(i-1)), 2] .= z_2[i]
		end
		for j = 1:4
			z[j,3] = z_3[j]
			z[j+4,3] = z_3[j]
			z[j+4*2,3] = z_3[j]
			z[j+4*3,3] = z_3[j]
		end
		ns = length(z_2)
		nss = length(z_2) * length(z_3)
		pi = 1/nss
		nc = 1
		return Dict("a"=>a,"na"=>na,"nc"=>nc,"ns"=>ns,"nss"=>nss,"e"=>e,"p"=>p,"z"=>z,"pi"=>pi)
	end

	function max_JuMP(a=0.5)
		m = Model(with_optimizer(Ipopt.Optimizer))
		@variable(m, omega[1:3])
		@variable(m, c >= 0.0)
		@NLconstraint(m, c + sum(data()["p"][i] * (omega[i] - data()["e"][i]) for i in 1:3) == 0.0)
		@NLobjective(m, Max,
			-exp(-a*c) + data()["pi"] * sum(-exp(-a *
			sum(data()["z"][j,i] * omega[i] for i in 1:3)) for j in 1:data()["nss"]))
		JuMP.optimize!(m)
		return Dict("obj"=>objective_value(m),"c"=>value(c),"omegas"=>[value(omega[i]) for i in 1:length(omega)])
	end

	function table_JuMP()
		d = DataFrame(a=[0.5;1.0;5.0],c = zeros(3),omega1=zeros(3),omega2=zeros(3),omega3=zeros(3),fval=zeros(3))
		for i in 1:nrow(d)
			xx = max_JuMP(d[i,:a])
			d[i,:c] = xx["c"]
			d[i,:omega1] = xx["omegas"][1]
			d[i,:omega2] = xx["omegas"][2]
			d[i,:omega3] = xx["omegas"][3]
			d[i,:fval] = xx["obj"]
		end
		return d
	end

## Does not work
	function obj(x::Vector,grad::Vector,data::Dict)
		if length(grad) > 0
			# grad wrt x1
			a = 0.5
			x=[1.0, 1.0, 1.0]
			grad[1] = (-1) * sum(data["pi"] * (-a * data["z"][i,1]) * (- exp(-a * sum(data["z"][i,1] * x[1] +
																				data["z"][i,2] * x[2] +
																				data["z"][i,3] * x[3]))) for i in 1:data["nss"])
			# grad wrt x2
			grad[2] = (-1) * sum(data["pi"] * (-a * data["z"][i,2]) * (- exp(-a * sum(data["z"][i,1] * x[1] +
																				data["z"][i,2] * x[2] +
																				data["z"][i,3] * x[3]))) for i in 1:data["nss"])
			# grad wrt x3
			grad[3] = (-1) * sum(data["pi"] * (-a * data["z"][i,3]) * (- exp(-a * sum(data["z"][i,1] * x[1] +
																				data["z"][i,2] * x[2] +
																				data["z"][i,3] * x[3]))) for i in 1:data["nss"])
			# grad wrt c
			grad[4] = (-1) *  (-a) * (- exp(-a * c))
		end
		return (-1) * (-exp(-a*c) + sum(data["pi"] * (- exp(-a * sum(data["z"][i,1] * x[1] +
										data["z"][i,1] * x[2] +
										data["z"][i,1] * x[3]))) for i=1:data["nss"]))
	end

	function constr(x::Vector,grad::Vector,data::Dict)
		if length(grad) > 0
			grad[1] = data["p"][1]
			grad[2] = data["p"][2]
			grad[3] = data["p"][3]
			grad[4] = 1
		end
		return (c +
			data["p"][1] * (x[1] - data["e"][1]) +
			data["p"][2] * (x[2] - data["e"][2]) +
			data["p"][3] * (x[3] - data["e"][3]))

	end

### BOUND ISSUE --> problem on this function
	function max_NLopt(a=0.5)
		opt = Opt(:LD_MMA, 4)
		lower_bounds!(opt, [0., -Inf, -Inf, -Inf])
		xtol_rel!(opt,1e-4)
		min_objective!(opt,(x,grad) -> obj(x,grad,data()))
		inequality_constraint!(opt, (x,grad) -> constr(x,grad,data()))
		ftol_rel!(opt,1e-9)
		vector_init = vcat(0.0, data()["e"])
		NLopt.optimize(opt, vector_init)
end
#Bounds error

	function table_NLopt()
		d = DataFrame(a=[0.5;1.0;5.0],c = zeros(3),omega1=zeros(3),omega2=zeros(3),omega3=zeros(3),fval=zeros(3))
		for i in 1:nrow(d)
			xx = max_NLopt(d[i,:a])
			for j in 2:ncol(d)-1
				d[i,j] = xx[2][j-1]
			end
			d[i,end] = xx[1]
		end
		return d
	end
table_NLopt()

end # module


text:
using Distributed
addprocs(50)

@everywhere include("ProblemGenerator.jl")
@everywhere include("VirusCallBacks.jl")

########################################################################
#Virus simulation (Figure 7)
########################################################################

#Define the percentages of cells producing IFN
#Each cell has its own unique parameter set
@everywhere function VirusStoch(rep,percentIFN,prob)
    #Make a copy of the parameter container and the initial condition
    θ = deepcopy(prob.p)
	u₀ = deepcopy(prob.u0)

	#Assign cells to be initially infected
	probDistInfected = Poisson(moi)
	u₀[:,:,2] = @. m2c(1e3*rand(probDistInfected,N,N))

	#Reset what cells are initially infected and dead
	θ.cellsInfected = fill(Inf,N,N)
	θ.cellsInfected[findall(u₀[:,:,2] .> 0.0), 1] .= 0.0
	θ.cellsDead = fill(Inf,N,N)
	θ.infectFirstAttempt = trues(N,N)

    #The 11th parameter (kcat7) controls IFN production, if zero then no IFN
    if percentIFN == 0.0
        θ.par[11] .= 0.0 #Set all parameters to zero
    else
		θ.par[11] .= 47639.70295 .* rand(Bernoulli(percentIFN),N,N)
    end

    #Assign the new parameters to the model
	#println(θ)
	@show rep, percentIFN
    probStoch = remake(prob; p=θ,u0=u₀)
    sol = @time solve(probStoch,CVODE_BDF(linear_solver=:GMRES),saveat=0.1,callback=cb)

    #Calculate the cell state dynamics over time
	#Percent, Sample, Time, Healthy, Infected, Dead
	numTime = length(saveTimePoints)
    allStates = zeros(numTime,6)

	allStates[:,1] = fill(percentIFN,numTime)
	allStates[:,2] = fill(rep,numTime)
	allStates[:,3] = saveTimePoints

    for (i,t) in enumerate(saveTimePoints)
        allStates[i,4:6] = cellStates(t,probStoch.p)
    end

    return allStates
end

#Define a new problem
@everywhere prob = ModelSetup(:Virus,:Stochastic,:Hetero)

#Time points to save during simulation
@everywhere const saveTimePoints = range(prob.tspan[1],prob.tspan[2],step=0.1)

#What percent of the cell population will be capable of producing IFNb?
percentIFN = 0.0:0.1:1.0
#Number of repititions for each IFN percentage
reps = 10
#Vector to save the simulations
#states = Vector(undef,reps)

iter = Iterators.product(1:reps,percentIFN)

#Simulate all of the different percentages in parallel and write data
states = pmap(x -> VirusStoch(x[1],x[2],prob),iter)

#Saving the Data
#Need to create a DataFrame with the following columns:
#Percent, Sample, Time, Healthy, Infected, Dead
statesLong = vcat(states...)
statesLong[:,4:6] = statesLong[:,4:6]./nCells #Get cell Percents

VirusSim = convert(DataFrame,statesLong)
rename!(VirusSim,[:Percent, :Sample, :Time, :Healthy, :Infected, :Dead])

#Save the simulation to a CSV to plot
CSV.write("./ServerSimulations/VirusFigure7DataHetero.csv",VirusSim)