text:
# # 5-bus Market simulation with [PowerSimulations.jl](https://github.com/NREL-SIIP/PowerSimulations.jl)

# **Originally Contributed by**: Clayton Barrows

# ## Introduction

# PowerSimulations.jl supports simulations that consist of sequential optimization problems
# where results from previous problems inform subsequent problems in a variety of ways. This
# example demonstrates some of these capabilities to represent electricity market clearing.

# ## Dependencies and Data
# First, let's create `System`s to represent the Day-Ahead and Real-Time market clearing
# process with hourly, and 5-minute time series data, respectively.

# ### Modeling Packages
using PowerSystems
using PowerSimulations
const PSI = PowerSimulations
IS = PowerSystems.IS

# ### Data management packages
using Dates
using DataFrames

# ### Optimization packages
using Cbc #solver
solver = optimizer_with_attributes(Cbc.Optimizer, "logLevel" => 1, "ratioGap" => 0.5)

# ### 5-bus Data
# The five bus system data here includes hourly day-ahead data, 5-minute real-time market
# data, and 6-second actual data. We'll only use the hourly and 5-minute data for the
# example simulations below, but the 6-second data is included for future development.
base_dir = PowerSystems.download(PowerSystems.TestData; branch = "master");
pm_data = PowerSystems.PowerModelsData(joinpath(base_dir, "matpower", "case5_re_uc.m"))

FORECASTS_DIR = joinpath(base_dir, "forecasts", "5bus_ts", "7day")

tsp_da = IS.read_time_series_file_metadata(joinpath(
    FORECASTS_DIR,
    "timeseries_pointers_da_7day.json",
))
tsp_rt = IS.read_time_series_file_metadata(joinpath(
    FORECASTS_DIR,
    "timeseries_pointers_rt_7day.json",
))
tsp_agc = IS.read_time_series_file_metadata(joinpath(
    FORECASTS_DIR,
    "timeseries_pointers_agc_7day.json",
))

sys_DA = System(pm_data)
reserves = [
    VariableReserve{ReserveUp}("REG1", true, 5.0, 0.1),
    VariableReserve{ReserveUp}("REG2", true, 5.0, 0.06),
    VariableReserve{ReserveUp}("REG3", true, 5.0, 0.03),
    VariableReserve{ReserveUp}("REG4", true, 5.0, 0.02),
]
contributing_devices = get_components(Generator, sys_DA)
for r in reserves
    add_service!(sys_DA, r, contributing_devices)
end

add_time_series!(sys_DA, tsp_da)
transform_single_time_series!(sys_DA, 24, Hour(24))

sys_RT = System(pm_data)
add_time_series!(sys_RT, tsp_rt)
transform_single_time_series!(sys_RT, 12, Hour(1))

sys_AGC = System(pm_data)
add_time_series!(sys_AGC, tsp_agc)

# ## `OperationsProblemTemplate`s
template_uc = template_unit_commitment()
devices = Dict(
    :Generators => DeviceModel(ThermalStandard, ThermalDispatch),
    :Ren => DeviceModel(RenewableDispatch, RenewableFullDispatch),
    :Loads => DeviceModel(PowerLoad, StaticPowerLoad),
    :HydroROR => DeviceModel(HydroDispatch, HydroDispatchRunOfRiver),
    :RenFx => DeviceModel(RenewableFix, FixedOutput),
)
template_ed = template_economic_dispatch(devices = devices)

# ### Define the Simulation Sequence

stages_definition = Dict(
    "UC" => Stage(GenericOpProblem, template_uc, sys_DA, solver),
    "ED" => Stage(GenericOpProblem, template_ed, sys_RT, solver),
)

feedforward_chronologies = Dict(("UC" => "ED") => Synchronize(periods = 24))

feedforward = Dict(
    ("ED", :devices, :Generators) => SemiContinuousFF(
        binary_source_stage = PSI.ON,
        affected_variables = [PSI.ACTIVE_POWER],
    ),
)

cache = Dict("UC" => [TimeStatusChange(ThermalStandard, PSI.ON)])

order = Dict(1 => "UC", 2 => "ED")
horizons = Dict("UC" => 24, "ED" => 12)
intervals = Dict("UC" => (Hour(24), Consecutive()), "ED" => (Hour(1), Consecutive()))

DA_RT_sequence = SimulationSequence(
    step_resolution = Hour(24),
    order = order,
    horizons = horizons,
    intervals = intervals,
    ini_cond_chronology = InterStageChronology(),
    feedforward_chronologies = feedforward_chronologies,
    feedforward = feedforward,
    #cache = cache,
)

# ## `Simulation`

file_path = tempdir()
sim = Simulation(
    name = "5bus-test",
    steps = 1,
    stages = stages_definition,
    stages_sequence = DA_RT_sequence,
    simulation_folder = file_path,
)

# ### Build simulation

build!(sim)

# ### Execute simulation

sim_results = execute!(sim)

# ## Results
ed_results = load_simulation_results(sim_results, "ED");


text:
using CSV
using Test
using WorldBankData
using DataFrames

# reset the country and indicator caches from search_wdi.jl tests
# otherwise downloads will fail since the data is not in the cache
WorldBankData.reset_country_cache()
WorldBankData.reset_indicator_cache()

refdf = DataFrame(CSV.File(joinpath(dirname(@__FILE__), "example_data.csv")))
refdf = select!(refdf, [:iso2c,:country,:NY_GNP_PCAP_CD,:year,:AG_LND_ARBL_HA_PC])
refdf = stack(refdf, [:NY_GNP_PCAP_CD, :AG_LND_ARBL_HA_PC])
refdf[!,:variable] = map(x -> replace(String(x), "_" => "."), refdf[!,:variable])
rename!(refdf, Dict(:variable => :indicator))

# download example case from documentation and compare to csv file
# retry 5 times if no data
function try_download(cntr=5)
    dfweb = ""
    while cntr > 0
        dfweb = wdi(["NY.GNP.PCAP.CD", "AG.LND.ARBL.HA.PC"], ["US", "BR"], 1980, 2008, dflong=true, verbose=true)
        if size(dfweb)[1] != 0
            break
        end
        println("      retry download")
        cntr = cntr - 1
        if cntr == 0
            error("unable to download data")
        end
    end
    dfweb
end

dfweb = try_download()

@test sort!(dfweb) == sort!(refdf)

ext: jl


text:
@testset "Plots" begin
    plot(FourierGrid(4))
    plot(FourierGrid(4)× FourierGrid(4) )
    plot(FourierGrid(4)× FourierGrid(4) × FourierGrid(4) )
    plot(FourierGrid(4)× FourierGrid(4) ,rand(4,4))
    plot(FourierGrid(4) ,rand(4))
end


text:
module ExactDiagonalization
using LinearAlgebra
using SparseArrays
#-----------------------------------------------------------------------------------------------------
import Base: +, -, *, /
import Base: eltype, view, length, size
import Base: Array
import LinearAlgebra: norm, mul!
import SparseArrays: sparse
#-----------------------------------------------------------------------------------------------------
export change!, index, operator
#-----------------------------------------------------------------------------------------------------
# Abstract Basis Type
#-----------------------------------------------------------------------------------------------------
"""
    AbstractBasis

Abstract type for basis object. Any concrete implementation should have the following functions:

Functions:
----------
- size(basis)       : Size of the block.
- change!(basis, i) : Change the bit and return normalization.
- index(basis)      : Get normalization and index of new state.

Optional:
---------
- eltype(basis)     : Default is ComplexF64
"""
abstract type AbstractBasis end
eltype(::AbstractBasis) = ComplexF64
length(b::AbstractBasis) = size(b, 2)
#-----------------------------------------------------------------------------------------------------
# INCLUDE
#-----------------------------------------------------------------------------------------------------
include("basis/TensorBasis.jl")
include("basis/ProjectedBasis.jl")
include("basis/TranslationalBasis.jl")
include("Operator.jl")
include("LinearOperation.jl")
include("Spin.jl")

end


text:
parselist(s::AbstractString) = map(x->(x[1], parse(Int, x[2:end])), split(s, ','))

const path1, path2 = map(parselist, eachline("input"))

manhatten((x, y)::Tuple{Int,Int}) = abs(x) + abs(y)  # from the origin

function trace(path::AbstractVector)
    points = Dict{Tuple{Int,Int},Int}()
    i = j = traveled = 0
    for (direction, distance) in path, step in 1:distance
        traveled += 1
        if direction === 'R'
            i += 1
        elseif direction === 'L'
            i -= 1
        elseif direction === 'U'
            j -= 1
        elseif direction === 'D'
            j += 1
        end
        get!(points, (i, j), traveled)
    end
    return points
end

function main(path1, path2)
    points1 = trace(path1)
    points2 = trace(path2)
    crosses = intersect(keys(points1), keys(points2))
    distance = minimum(manhatten, crosses)
    println("Part 1 answer: ", distance)
    steps = minimum(p->points1[p] + points2[p], crosses)
    println("Part 2 answer: ", steps)
end

main(path1, path2)

max_stars_repo_name: ararslan/advent-solutions

max_issues_repo_name: ararslan/advent-solutions

max_forks_repo_name: ararslan/advent-solutions


text:
# This file is part of the StateMachines.jl Julia package
#
# Paul Bayer
# MIT License
# Github: https://github.com/pbayer/StateMachines.jl
#

"""
    operate(A::StateMachine, input::Id, output::Id)

Operate a state machine. If needed, create input and output channels.
"""
function operate(A::StateMachine, input::Id, output::Id)
end

"""
    start!(A::StateMachine, input::Id="", output::Id=""; wait4it::Bool=false)

Start a state machine. Eventually setup input and output channels and start a
task operating the state machine.

# Arguments
- `A::StateMachine`:
- `input::Id=""`:
- `output::Id=""`:
- `wait4it::Bool=false`:
"""
function start!(A::StateMachine, input::Id="", output::Id=""; wait4it::Bool=false)
    @assert A.state == Undefined() "unable to start: $(A.id) is in state $(A.state)!"
    if wait4it
        operate(A, input, output)
    else
        A.task = @async operate(A, input, output)
    end
end


text:
# ---
# title: 977. Squares of a Sorted Array
# id: problem977
# author: Tian Jun
# date: 2020-10-31
# difficulty: Easy
# categories: Array, Two Pointers
# link: <https://leetcode.com/problems/squares-of-a-sorted-array/description/>
# hidden: true
# ---
#
# Given an integer array `nums` sorted in **non-decreasing** order, return _an
# array of **the squares of each number**  sorted in non-decreasing order_.
#
#
#
# **Example 1:**
#
#
#
#     Input: nums = [-4,-1,0,3,10]
#     Output: [0,1,9,16,100]
#     Explanation: After squaring, the array becomes [16,1,0,9,100].
#     After sorting, it becomes [0,1,9,16,100].
#
#
# **Example 2:**
#
#
#
#     Input: nums = [-7,-3,2,3,11]
#     Output: [4,9,9,49,121]
#
#
#
#
# **Constraints:**
#
#   * `1 <= nums.length <= 104`
#   * `-104 <= nums[i] <= 104`
#   * `nums` is sorted in **non-decreasing** order.
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


text:
#=
using DelimiteFiles
=#

# NOTE: These simulation setting are identical to that of test-sims-6-5 EXCEPT
# that W_true is different. The W_true here has low abundance (< 10% each) for
# 2 pairs (i.e. 4) features. Thus, they should be harder to recover, and
# uncertainty in R_i should be greater.

# Simulation Name
simname = basename("$(@__DIR__)")
println(simname); flush(stdout)

# NOTE: write to scratchdir
function outdir_suffix(pmiss, phi, zind)
  return "pmiss$(pmiss)-phi$(phi)-zind$(zind)"
end

# NOTE: modify
settings = let
  Z1 = Matrix{Bool}(readdlm(joinpath(@__DIR__, "Z1.txt"), Int, comments=true))
  Z2 = Matrix{Bool}(readdlm(joinpath(@__DIR__, "Z2.txt"), Int, comments=true))
  Z3 = Matrix{Bool}(readdlm(joinpath(@__DIR__, "Z3.txt"), Int, comments=true))
  Z = [Z1, Z2, Z3]
  W = readdlm(joinpath(@__DIR__, "W.txt"), comments=true)
  Nfac = 2000
  N = [1, 1] *  Nfac
  pthin = 5
  batchprop = 0.05
  dataseed = 1
  mcmcseed = 1
  temperatures = [1.0, 1.003, 1.006, 1.01]

  [Dict(:simname => simname,
        :repfam_dist_scale => phi,
        :N => N,
        :Z => Z[zind],
        :W => W,
        :thin_samps => 1,
        :nburn => 6000,
        :nsamps => 3000,
        :Lmcmc => Dict(0 => 3, 1 => 3),
        :Kmcmc => 10,
        :pthin => pthin,
        :temperatures => temperatures,
        :ntemps => length(temperatures),
        :propmissingscale => propmissingscale,
        :batchprop => batchprop,
        :dataseed => dataseed,
        :mcmcseed => mcmcseed,
        :outdir_suffix => outdir_suffix(propmissingscale, phi, zind))
   for propmissingscale in [0.6, 0.0]
   for phi in [1.0, 0.0, 10.0, 25.0]
   for zind in 1:3]
end


text:
module LDB
export
    # energy map
    EnergyMap,
    TimeFrequency,
    ProbabilityDensity,
    Signatures,
    energy_map,
    # discriminant measures
    DiscriminantMeasure,
    ProbabilityDensityDM,
    SignaturesDM,
    AsymmetricRelativeEntropy,
    SymmetricRelativeEntropy,
    HellingerDistance,
    LpDistance,
    EarthMoverDistance,
    discriminant_measure,
    # discriminant power
    DiscriminantPower,
    BasisDiscriminantMeasure,
    FishersClassSeparability,
    RobustFishersClassSeparability,
    discriminant_power,
    # local discriminant basis
    LocalDiscriminantBasis,
    fit!,
    fitdec!,
    fit_transform,
    transform,
    inverse_transform,
    change_nfeatures

using LinearAlgebra,
      Wavelets,
      Statistics
import Parameters: @with_kw
import StatsBase: mad

using ..Utils,
      ..DWT
import ..BestBasis: bestbasis_treeselection

include("ldb_energymap.jl")
include("ldb_measures.jl")

## LOCAL DISCRIMINANT BASIS
"""
    LocalDiscriminantBasis

Class type for the Local Discriminant Basis (LDB), a feature selection algorithm developed
by N. Saito and R. Coifman in "Local Discriminant Bases and Their Applications" in the
Journal of Mathematical Imaging and Vision, Vol 5, 337-358 (1995). This struct contains the
following field values:

# Parameters and Attributes:
- `wt::DiscreteWavelet`: (Default: `wavelet(WT.haar)`) A discrete wavelet for transform
  purposes.
- `max_dec_level::Union{Integer, Nothing}`: (Default: `nothing`) Max level of wavelet packet
  decomposition to be computed.
- `dm::DiscriminantMeasure`: (Default: `AsymmetricRelativeEntropy()`) the discriminant
    measure for the LDB algorithm. Supported measures are the `AsymmetricRelativeEntropy()`,
    `LpDistance()`, `SymmetricRelativeEntropy()`, and `HellingerDistance()`.
- `en::EnergyMap`: (Default: `TimeFrequency()`) the type of energy map used. Supported maps
    are `TimeFrequency()`, `ProbabilityDensity()`, and `Signatures()`.
- `dp::DiscriminantPower()`: (Default: `BasisDiscriminantMeasure()`) the measure of
    discriminant power among expansion coefficients. Supported measures are
    `BasisDiscriminantMeasure()`, `FishersClassSeparability()`, and
    `RobustFishersClassSeparability()`.
- `top_k::Union{Integer, Nothing}`: (Default: `nothing`) the top-k coefficients used in each
    node to determine the discriminant measure.
- `n_features::Union{Integer, Nothing}`: (Default: `nothing`) the dimension of output after
    undergoing feature selection and transformation.
- `sz::Union{Vector{T} where T<:Integer, Nothing}`: (Default: `nothing`) Size of signal
- `Γ::Union{AbstractArray{T} where T<:AbstractFloat, AbstractArray{NamedTuple{(:coef,
  :weight), Tuple{S₁, S₂}}} where {S₁<:Array{T} where T<:AbstractFloat, S2<:Union{T,
  Array{T}} where T<:AbstractFloat}, Nothing}`: (Default: `nothing`) computed energy map
- `DM::Union{AbstractArray{<:AbstractFloat}, Nothing}`: (Default: `nothing`) computed
  discriminant measure
- `cost::Union{AbstractVector{<:AbstractFloat}, Nothing}`: (Default: `nothing`) computed
    wavelet packet decomposition (WPD) tree cost based on the discriminant measure `DM`.
- `tree::Union{BitVector, Nothing}`: (Default: `nothing`) computed best WPD tree based on
    the discriminant measure `DM`.
- `DP::Union{AbstractVector{<:AbstractFloat}, Nothing}`: (Default: `nothing`) computed
  discriminant power
- `order::Union{AbstractVector{Integer}, Nothing}`: (Default: `nothing`) ordering of `DP` by
  descending order.
"""
@with_kw mutable struct LocalDiscriminantBasis
    # to be declared by user
    wt::DiscreteWavelet = wavelet(WT.haar)
    max_dec_level::Union{Integer, Nothing} = nothing
    dm::DiscriminantMeasure = AsymmetricRelativeEntropy()
    en::EnergyMap = TimeFrequency()
    dp::DiscriminantPower = BasisDiscriminantMeasure()
    top_k::Union{Integer, Nothing} = nothing
    n_features::Union{Integer, Nothing} = nothing
    # to be computed in fit! or fitdec!
    sz::Union{Tuple{Vararg{T,N}} where {T<:Integer,N}, Nothing} = nothing
    Γ::Union{AbstractArray{<:AbstractFloat},
             AbstractArray{NamedTuple{(:coef, :weight), Tuple{S1, S2}}} where
                {S1<:Array{T} where T<:AbstractFloat,
                 S2<:Union{AbstractFloat, Array{<:AbstractFloat}}},
             Nothing} = nothing
    DM::Union{AbstractArray{<:AbstractFloat}, Nothing} = nothing
    cost::Union{AbstractVector{<:AbstractFloat}, Nothing} = nothing
    tree::Union{BitVector, Nothing} = nothing
    DP::Union{AbstractArray{<:AbstractFloat}, Nothing} = nothing
    order::Union{AbstractVector{Integer}, Nothing} = nothing
end

"""
    fit!(f, X, y)

Fits the Local Discriminant Basis feature selection algorithm `f` onto the signals `X` with
labels `y`.

# Arguments
- `f::LocalDiscriminantBasis`: Local discriminant basis object.
- `X::AbstractArray{S} where S<:AbstractFloat`: Group of signals of the same size, arranged
  in the 2nd (1D signals) or 3rd (2D signals) dimension.
- `y::AbstractVector{T} where T`: Labels corresponding to each signal in X.

# Returns
- `f::LocalDiscriminantBasis`: The same LDB object `f` with updated fields.

# Examples
```julia
using Wavelets, WaveletsExt

X, y = generateclassdata(ClassData(:tri, 5, 5, 5))
f = LocalDiscriminantBasis()
fit!(f, X, y)
```

**See also:** [`LocalDiscriminantBasis`](@ref), [`fit_transform`](@ref),
    [`transform`](@ref), [`inverse_transform`](@ref), [`change_nfeatures`](@ref)
"""
function fit!(f::LocalDiscriminantBasis, X::AbstractArray{S}, y::AbstractVector{T}) where
             {S<:AbstractFloat, T}

    # basic summary of data
    @assert 2 ≤ ndims(X) ≤ 3
    sz = size(X)[1:end-1]
    L = maxtransformlevels(min(sz...))
    # change LocalDiscriminantBasis parameters if necessary
    f.max_dec_level = isnothing(f.max_dec_level) ? L : f.max_dec_level
    @assert 1 ≤ f.max_dec_level ≤ L

    # wavelet packet decomposition
    Xw = wpdall(X, f.wt, f.max_dec_level)

    # fit local discriminant basis
    fitdec!(f, Xw, y)
    return nothing
end

"""
    fitdec!(f, Xw, y)

Fits the Local Discriminant Basis feature selection algorithm `f` onto the decomposed
signals `Xw` with labels `y`.

# Arguments
- `f::LocalDiscriminantBasis`: Local discriminant basis object.
- `Xw::AbstractArray{S} where S<:AbstractFloat`: Group of decomposed signals of the same
  size, arranged in the 3rd (1D signals) or 4th (2D signals) dimension.
- `y::AbstractVector{T} where T`: Labels corresponding to each signal in X.

# Returns
- `f::LocalDiscriminantBasis`: The same LDB object `f` with updated fields.

# Examples
```julia
using Wavelets, WaveletsExt

X, y = generateclassdata(ClassData(:tri, 5, 5, 5))
f = LocalDiscriminantBasis()
Xw = wpdall(X, f.wt)
fitdec!(f, Xw, y)
```

**See also:** [`LocalDiscriminantBasis`](@ref), [`fit!`](@ref), [`fit_transform`](@ref),
    [`transform`](@ref), [`inverse_transform`](@ref), [`change_nfeatures`](@ref)
"""
function fitdec!(f::LocalDiscriminantBasis, Xw::AbstractArray{S}, y::AbstractVector{T}) where
                {S<:AbstractFloat, T}
    # basic summary of data
    @assert 3 ≤ ndims(Xw) ≤ 4
    c = unique(y)       # unique classes
    nc = length(c)      # number of classes
    Ny = length(y)
    f.sz = size(Xw)[1:end-2]
    L = size(Xw)[end-1]
    Nx = size(Xw)[end]
    nelem = prod(f.sz)  # total number of elements in each signal

    # change LocalDiscriminantBasis parameters if necessary
    f.top_k = isnothing(f.top_k) ? nelem : f.top_k
    f.n_features = isnothing(f.n_features) ? nelem : f.n_features
    f.max_dec_level = isnothing(f.max_dec_level) ? L-1 : f.max_dec_level

    # parameter checking
    @assert Nx == Ny
    @assert 1 ≤ f.top_k ≤ nelem
    @assert 1 ≤ f.n_features ≤ nelem
    @assert f.max_dec_level+1 == L
    @assert 1 ≤ f.max_dec_level ≤ maxtransformlevels(min(f.sz...))
    @assert nc > 1

    # --- Construct energy map for each class ---
    f.Γ = energy_map(Xw, y, f.en)

    # --- Compute discriminant measure D and obtain tree cost ---
    f.DM = discriminant_measure(f.Γ, f.dm)
    tree_type = ndims(Xw) == 3 ? :binary : :quad
    cost_len = ndims(Xw) == 3 ? gettreelength(1<<L) : gettreelength(1<<L,1<<L)
    f.cost = Vector{S}(undef, cost_len)
    for i in eachindex(f.cost)
        d = getdepth(i, tree_type)    # node level
        # Extract coefficients for current node
        if tree_type == :binary
            θ = i - 1<<d
            nθ = nodelength(f.sz..., d)         # Number of elements in current node
            rng = (θ*nθ+1):((θ+1)*nθ)
            DMθ = f.DM[rng,d+1]
        else
            rng₁ = getrowrange(f.sz[1], i)
            rng₂ = getcolrange(f.sz[2], i)
            nθ = length(rng₁) * length(rng₂)    # Number of elements in current node
            DMθ = f.DM[rng₁,rng₂,d+1] |> vec
        end
        # Compute cost
        if f.top_k < nθ
            sort!(DMθ, rev=true)
            f.cost[i] = sum(DMθ[1:f.top_k])
        else
            f.cost[i] = sum(DMθ)
        end
    end

    # --- Select best tree and best set of expansion coefficients ---
    f.tree = bestbasis_treeselection(f.cost, f.sz..., :max)
    Xc = getbasiscoefall(Xw, f.tree)

    # --- Obtain and order basis functions by power of discrimination ---
    (f.DP, f.order) = f.dp == BasisDiscriminantMeasure() ?
                      discriminant_power(f.DM, f.tree, f.dp) :
                      discriminant_power(Xc, y, f.dp)
    return nothing
end

"""
    transform(f, X)

Extract the LDB features from signals `X` based on the LDB tree constructed from previously
fitted signals.

# Arguments
- `f::LocalDiscriminantBasis`: Local discriminant basis object.
- `X::AbstractArray{S} where S<:AbstractFloat`: Group of signals of the same size, arranged
  in the 2nd (1D signals) or 3rd (2D signals) dimension.

# Returns
- `Xc::Matrix{S}`: A matrix of the extracted LDB features from `X`, where each column
  corresponds to the coefficients extracted from the corresponding signal.

# Examples
```julia
using Wavelets, WaveletsExt

X, y = generateclassdata(ClassData(:tri, 5, 5, 5))
f = LocalDiscriminantBasis()
fit!(f, X, y)
Xc = transform(f, X)
```

**See also:** [`LocalDiscriminantBasis`](@ref), [`fit!`](@ref), [`fit_transform`](@ref),
    [`inverse_transform`](@ref), [`change_nfeatures`](@ref)
"""
function transform(f::LocalDiscriminantBasis, X::AbstractArray{T}) where T
    # check necessary measurements
    sz = size(X)[1:end-1]
    N = size(X)[end]
    nelem = prod(f.sz)  # total number of elements in each signal
    @assert 2 ≤ ndims(X) ≤ 3
    @assert !isnothing(f.max_dec_level)
    @assert !isnothing(f.top_k)
    @assert !isnothing(f.n_features)
    @assert !isnothing(f.sz)
    @assert !isnothing(f.Γ)
    @assert !isnothing(f.DM)
    @assert !isnothing(f.cost)
    @assert !isnothing(f.tree)
    @assert !isnothing(f.DP)
    @assert !isnothing(f.order)
    @assert sz == f.sz
    @assert 1 ≤ f.max_dec_level ≤ maxtransformlevels(min(f.sz...))
    @assert 1 ≤ f.top_k ≤ nelem
    @assert 1 ≤ f.n_features ≤ nelem

    # wpt on X based on given f.tree
    Xw = wptall(X, f.wt, f.tree)
    # Extract best features
    Xc = Array{T,2}(undef, f.n_features, N)
    for (i, xw) in enumerate(eachslice(Xw, dims=ndims(Xw)))
        Xc[:,i] = xw[f.order[1:f.n_features]]
    end
    return Xc
end

"""
    fit_transform(f, X, y)

Fit and transform the signals `X` with labels `y` based on the LDB class `f`. This function
essentially combines both `fit!` and `transform` into one, and is the preferred method when
trying to extract the LDB coefficients of the training dataset. For the test dataset, one
needs to first run `fit!` on the training dataset before running `transform` on the test
dataset.

# Arguments
- `f::LocalDiscriminantBasis`: Local discriminant basis object.
- `Xw::AbstractArray{S} where S<:AbstractFloat`: Group of decomposed signals of the same
  size, arranged in the 3rd (1D signals) or 4th (2D signals) dimension.
- `y::AbstractVector{T} where T`: Labels corresponding to each signal in X.

# Returns
- `Xc::Matrix{S}`: A matrix of the extracted LDB features from `X`, where each column
  corresponds to the coefficients extracted from the corresponding signal.

# Examples
```julia
using Wavelets, WaveletsExt

X, y = generateclassdata(ClassData(:tri, 5, 5, 5))
f = LocalDiscriminantBasis()
Xc = fit_transform(f, X, y)
```

**See also:** [`LocalDiscriminantBasis`](@ref), [`fit!`](@ref), [`transform`](@ref),
    [`inverse_transform`](@ref), [`change_nfeatures`](@ref)
"""
function fit_transform(f::LocalDiscriminantBasis,
                       X::AbstractArray{S},
                       y::AbstractVector{T}) where {S<:AbstractFloat, T}
    # get necessary measurements
    @assert 2 ≤ ndims(X) ≤ 3
    sz = size(X)[1:end-1]
    N = size(X)[end]
    f.max_dec_level = isnothing(f.max_dec_level) ? maxtransformlevels(min(sz...)) : f.max_dec_level
    @assert 1 ≤ f.max_dec_level ≤ maxtransformlevels(min(sz...))

    # wpd on X
    Xw = wpdall(X, f.wt, f.max_dec_level)

    # fit LDB and return best features
    fitdec!(f, Xw, y)
    Xw = getbasiscoefall(Xw, f.tree)
    # Extract best features
    Xc = Array{S,2}(undef, f.n_features, N)
    for (i, xw) in enumerate(eachslice(Xw, dims=ndims(Xw)))
        Xc[:,i] = xw[f.order[1:f.n_features]]
    end
    return Xc
end

"""
    inverse_transform(f, X)

Compute the inverse transform on the feature matrix `X` to form the original signal based on
the LDB class `f`. This function can be used to reconstruct the signals and compare the
reconstructed features between different classes.

# Arguments
- `f::LocalDiscriminantBasis`: Local discriminant basis object.
- `X::Matrix{S} where S<:AbstractFloat`: A matrix of the extracted LDB features from
  original signals, where each column corresponds to the coefficients extracted from the
  corresponding signal.

# Returns
- `Xₜ::Array{S}`: Reconstructed signals.

# Examples
```julia
X, y = generateclassdata(ClassData(:tri, 5, 5, 5))
f = LocalDiscriminantBasis()
Xc = fit_transform(f, X, y)
Xₜ = inverse_transform(f, Xc)
```

**See also:** [`LocalDiscriminantBasis`](@ref), [`fit!`](@ref), [`fit_transform`](@ref),
    [`transform`](@ref), [`change_nfeatures`](@ref)
"""
function inverse_transform(f::LocalDiscriminantBasis, X::AbstractArray{T,2}) where
                           T<:AbstractFloat
    # get necessary measurements
    @assert size(X,1) == f.n_features
    N = size(X,2)

    # insert the features x into the full matrix X padded with 0's
    Xc = zeros(T, (f.sz..., N))
    @views for (xc, x) in zip(eachslice(Xc, dims=ndims(Xc)), eachslice(X, dims=2))
        for (i,j) in enumerate(f.order[1:f.n_features])
            xc[j] = x[i]
        end
    end

    # iwpt on X
    Xₜ = iwptall(Xc, f.wt, f.tree)
    return Xₜ
end

"""
    change_nfeatures(f, x, n_features)

Change the number of features extracted from the signals from `f.n_features` to
`n_features`.

!!! warning If the input `n_features` is larger than `f.n_features`, it results in the
    regeneration of signals based on the current `f.n_features` before reselecting the
    features. This will cause additional features to be less accurate and effective. Often
    times, the "additional" features tend to be zeros.

# Arguments
- `f::LocalDiscriminantBasis`: Local discriminant basis object.
- `x::AbstractMatrix{S} where S<:AbstractFloat`: A matrix of the extracted LDB features from
  original signals, where each column corresponds to the coefficients extracted from the
  corresponding signal.
- `n_features::Integer`: Desired number of features in selected by LDB.

# Returns
- `xₜ::Matrix{S}`: Matrix of extracted LDB features, where the number of rows is now
  `n_features` instead of the previous `f.n_features`.

**See also:** [`LocalDiscriminantBasis`](@ref), [`fit!`](@ref), [`fit_transform`](@ref),
    [`transform`](@ref), [`inverse_transform`](@ref)
"""
function change_nfeatures(f::LocalDiscriminantBasis,
                          x::AbstractArray{T,2},
                          n_features::Integer) where T<:AbstractFloat

    # check measurements
    @assert !isnothing(f.n_features)
    @assert size(x,1) == f.n_features || throw(ArgumentError("f.n_features and number of rows of x do not match!"))
    @assert 1 ≤ n_features ≤ prod(f.sz)

    # change number of features
    if f.n_features ≥ n_features
        f.n_features = n_features
        xₜ = x[1:f.n_features,:]
    else
        @warn "Proposed n_features larger than currently saved n_features. Results will be less accurate since inverse_transform and transform is involved."
        X = inverse_transform(f, x)
        f.n_features = n_features
        xₜ = transform(f, X)
    end
    return xₜ
end

end # end module


text:
""""
`Cochain(simplices::AbstractVector, values::AbstractVector)`

A type representing a cochain, i.e. a way to associate a value to each simplex. Simplices are expressed as an abstract vector `simplices` (generally represented as tuples). Values are stored in `values`. In practice, `values` will often be a sparse vector, in which case `findnz` returns the simplices where the cochian has non-zero value and the list of values on those simplices.

Use `keys` and `values` to access simplices and values respectively. The value at a specific simplex can be accessed with `getindex`.

## Examples

```julia
julia> c = Cochain([(1, 2), (2, 3), (1, 3)], [6, 9, -2])
Int64-valued 1-Cochain

julia> keys(c)
3-element Array{Tuple{Int64,Int64},1}:
 (1, 2)
 (2, 3)
 (1, 3)

julia> values(c)
3-element Array{Int64,1}:
  6
  9
 -2

julia> c[(2, 3)]
9

julia> using SparseArrays

julia> c = Cochain([(1, 2), (2, 3), (1, 3), (1, 4)], sparse([6, 0, -2, 12]))
Sparse Int64-valued 1-Cochain. Non-zero values:
(1, 2) => 6
(1, 3) => -2
(1, 4) => 12
```
"""
struct Cochain{S<:AbstractVector, V<:AbstractVector}
    simplices::S
    values::V
end

Base.keys(c) = c.simplices
Base.values(c::Cochain) = c.values
Base.length(c::Cochain) = length(values(c))
Base.getindex(s::Cochain, t) = values(s)[searchsortedfirst(keys(s), t)]
Base.eltype(c::Type{Cochain{S, V}}) where {S, V} = eltype(V)
Base.eltype(c::Cochain) = eltype(typeof(c))

SparseArrays.issparse(c::Cochain) = issparse(values(c))
function SparseArrays.findnz(c::Cochain)
    issparse(c) || throw(ArgumentError("Argument of findnz must be sparse"))
    inds, vals = findnz(values(c))
    (keys(c)[inds], vals)
end

function Base.show(io::IO, s::Cochain)
    npts = nv(eltype(keys(s)))
    summary = "$(eltype(s))-valued $(npts-1)-Cochain"
    if !issparse(s)
        print(io, summary)
    else
        println(io, "Sparse $summary. Non-zero values:")
        limit = get(io, :limit, true)
        for (ind, (key, val)) in enumerate(zip(findnz(s)...))
            if ind >= 100 && limit
                println("...")
                break
            else
                show(io, key => val)
                println(io, "")
            end
        end
    end
end

function onecochain(::Type{T}, c::AbstractVector, i) where {T}
    v = sparsevec([i], one(T), length(c))
    Cochain(c, v)
end

function facevalue(cochain, simplex, i)
    @inbounds f = face(simplex, i)
    cochain[f]
end

@generated function boundaryvalue(cochain, simplex::T) where T
    N = nv(T)
    faces = [Expr(:call, :facevalue, :cochain, :simplex, i) for i in 1:N]
    expr = faces[1]
    for i in 2:N
        sign = isodd(i) ? (:+) : (:-)
        expr = Expr(:call, sign, expr, faces[i])
    end
    expr
end


text:
using Oceananigans
module Vizinanigans

using Oceananigans
using GLMakie, AbstractPlotting
using ImageTransformations, Colors
using AbstractPlotting.MakieLayout

export visualize

function visualize(model::Oceananigans.AbstractModel)
    vstates = [Array(interior(model.velocities[velocity])) for velocity in keys(model.velocities)]
    vstatenames = [string(velocity) for velocity in keys(model.velocities)]
    tstates = [Array(interior(model.tracers[tracer])) for tracer in keys(model.tracers)]
    tstatenames = [string(tracer) for tracer in keys(model.tracers)]
    states = vcat(vstates, tstates)
    statenames = vcat(vstatenames, tstatenames)
    visualize(states, statenames = statenames)
    return nothing
end

function visualize(states::AbstractArray; statenames = string.(1:length(states)))
    # Create choices and nodes
    stateindex = collect(1:length(states))
    statenode = Node(stateindex[1])

    colorchoices = [:balance, :thermal, :dense, :deep, :curl, :thermometer]
    colornode = Node(colorchoices[1])

    x, y, z = size(states[1]) # basically only determines the aspect ratio

    # Lift Nodes
    state = @lift(states[$statenode])
    clims = @lift(extrema(states[$statenode]))
    cmap_rgb = @lift(to_colormap($colornode))
    titlename = @lift(" "^10 * " Field " * statenames[$statenode] * " "^10) # use padding and appropriate centering

    # Create scene
    scene, layout = layoutscene()
    # Volume Plot (needs to come first)
    lscene = layout[1:4, 2:4] = LScene(scene)
    volume!(lscene, 0..x, 0..y, 0..z, state,
            camera = cam3d!,
            colormap = cmap_rgb,
            colorrange = clims)
    # Title
    supertitle = layout[1,2] = LText(scene, titlename , textsize = 50, color = :black)
    # Menus
    statemenu = LMenu(scene, options = zip(statenames, stateindex))
    on(statemenu.selection) do s
        statenode[] = s
    end

    colormenu = LMenu(scene, options = zip(colorchoices, colorchoices))
    on(colormenu.selection) do s
        colornode[] = s
    end

    layout[1, 1] = vgrid!(
        LText(scene, "State", width = nothing),
        statemenu,
        LText(scene, "Color", width = nothing),
        colormenu,
    )
    display(scene)
    return nothing
end

end

max_stars_repo_name: glwagner/Visualizanigans

max_issues_repo_name: glwagner/Visualizanigans

max_forks_repo_name: glwagner/Visualizanigans


text:
#
# Denoise an SPD Example with Cyclic Proximal Point to minimize the
#
# L2-TV functional with anisotropic TV
#
# where the example is the same data and cost as for SPD_Image_CP_vs_CPPA
#
# This example and its data is used for comparison in Section 6.2 of
#
# > R. Bergmann, R. Herzog, M. Silva Louzeiro, D. Tenbrinck, J. Vidal Núñez:
# > _Fenchel Duality Theory and a Primal-Dual Algorithm on Riemannian Manifolds_,
# > arXiv: [1908.02022](https://arxiv.org/abs/1908.02022)
#
using Manopt, Manifolds
using Images, CSV, DataFrames, LinearAlgebra, JLD2
#
# Settings
experiment_name = "SPD_Image_CPPA"
export_orig = true # export and render input image
export_result = true # export resulting image
export_table = true # export a table of data recorded through the iterations
export_function_value = true # export the final cost function value for ChambollePock
asy_render_detail = 4 # resolution in asymptote, a high resolution is e.g. 4
results_folder = joinpath(@__DIR__, "Image_TV")
!isdir(results_folder) && mkdir(results_folder)

#
# Algorithm Parameters
α = 6.0
maxIterations = 4000

#
# Manifold and Data
f = artificial_SPD_image2(32)
pixelM = SymmetricPositiveDefinite(3)
if export_orig
    orig_filename = joinpath(results_folder, experiment_name * "-orig.asy")
    asymptote_export_SPD(orig_filename; data=f, scaleAxes=(7.5, 7.5, 7.5))
    render_asymptote(orig_filename; render=asy_render_detail)
end

#
# Build Problem for L2-TV & proximal maps
M = PowerManifold(pixelM, NestedPowerRepresentation(), size(f)...)
d = length(size(f))
rep(d) = (d > 1) ? [ones(Int, d)..., d] : d
fidelity(x) = 1 / 2 * distance(M, x, f)^2
Λ(x) = forward_logs(M, x) # on T_xN
prior(x) = norm(norm.(Ref(pixelM), repeat(x, rep(d)...), Λ(x)), 1)
cost(x) = (1 / α) * fidelity(x) + prior(x)
proxes = ((λ, x) -> prox_distance(M, λ, f, x, 2), (λ, x) -> prox_TV(M, α * λ, x, 1))

x0 = f
@time o = cyclic_proximal_point(
    M,
    cost,
    proxes,
    x0;
    λ=i -> 4 / i,
    stopping_criterion=StopAfterIteration(maxIterations),
    debug=[
        :Iteration,
        " | ",
        DebugProximalParameter(),
        " | ",
        :Change,
        " | ",
        :Cost,
        "\n",
        100,
        :Stop,
    ],
    record=[:Iteration, :Iterate, :Cost],
    return_options=true,
)
y = get_solver_result(o)
yRec = get_record(o)
#
# Results
if export_result
    result_filename = joinpath(results_folder, experiment_name * "-result.asy")
    asymptote_export_SPD(result_filename; data=y, scaleAxes=(7.5, 7.5, 7.5))
    render_asymptote(result_filename; render=asy_render_detail)
end
if export_table
    A = cat([y[1] for y in yRec], [y[3] for y in yRec]; dims=2)
    table_filename = joinpath(results_folder, experiment_name * "-recorded-cost.csv")
    CSV.write(table_filename, DataFrame(A); writeheader=false)
end
if export_function_value
    fctval_filename = joinpath(results_folder, experiment_name * "-cost.jld2")
    values = Dict("cost_function_value" => last(yRec)[3], "iterations" => length(yRec) - 1)
    save(fctval_filename, values)
end


text:
module Benchmark
    using DataArrays # TODO: Remove this
    using DataFrames

    export benchmark, benchmarks, compare

    include("benchmarks.jl")
    include("compare.jl")
end


text:
module TestTable

import PublicAPI
using DocumentationOverview
using Test

using ..TestPkgAPI: ExampleModule

function test_default()
    table = DocumentationOverview.table(DocumentationOverview)
    text = sprint(show, "text/plain", table)
    @test occursin("DocumentationOverview.table", text)
    @test occursin("DocumentationOverview.find", text)
    @test occursin("DocumentationOverview.API", text)
end

function test_publicapi()
    table = DocumentationOverview.table(PublicAPI.of(DocumentationOverview))
    text = sprint(show, "text/plain", table)
    @test occursin("DocumentationOverview.table", text)
    @test occursin("DocumentationOverview.find", text)
    @test occursin("DocumentationOverview.API", text)
end

function test_fullnames()
    table = DocumentationOverview.table(
        :[table, DocumentationOverview.API, find],
        namespace = DocumentationOverview,
    )
    text = sprint(show, "text/plain", table)
    @test occursin("DocumentationOverview.table", text)
    @test occursin("DocumentationOverview.find", text)
    @test occursin("DocumentationOverview.API", text)
end

function test_nodoc()
    table = DocumentationOverview.table(
        DocumentationOverview.API.(split("""
        f
        nodoc
        """), namespace = ExampleModule),
    )
    text = sprint(show, "text/plain", table)
    @test occursin("f", text)
    @test occursin("nodoc", text)
end

end  # module


text:
using MultivariateStats
using LinearAlgebra
using Test
import Statistics: mean, cov
import Random

#Test equivalence of eigenvectors/singular vectors taking into account possible
#phase (sign) differences
#This may end up in Base.Test; see JuliaLang/julia#10651
function test_approx_eq_vecs(a::StridedVecOrMat{S},
    b::StridedVecOrMat{T}, error=nothing) where {S<:Real, T<:Real}
    n = size(a, 2)
    @test n==size(b, 2) && size(a, 1)==size(b, 1)
    error==nothing && (error=n^3*(eps(S)+eps(T)))
    for i=1:n
        ev1, ev2 = a[:, i], b[:, i]
        deviation = min(abs(norm(ev1-ev2)), abs(norm(ev1+ev2)))
        if !isnan(deviation)
            @test isapprox.(deviation, 0.0, atol=error)
        end
    end
end

Random.seed!(34568)

## prepare data
let
d = 5
ns = [10, 15, 20]
nc = length(ns)
n = sum(ns)
Xs = Matrix{Float64}[]
ys = Vector{Int}[]
Ss = Matrix{Float64}[]
cmeans = zeros(d, nc)

for k = 1:nc
    R = qr(randn(d, d)).Q
    nk = ns[k]

    Xk = R * Diagonal(2 * rand(d) .+ 0.5) * randn(d, nk) .+ randn(d)
    yk = fill(k, nk)
    uk = vec(mean(Xk, dims=2))
    Zk = Xk .- uk
    Sk = Zk * Zk'

    push!(Xs, Xk)
    push!(ys, yk)
    push!(Ss, Sk)
    cmeans[:,k] .= uk
end

X = hcat(Xs...)
y = vcat(ys...)
mv = vec(mean(X, dims=2))
Sw = zeros(d, d)
for k = 1:nc
    Sw += Ss[k]
end

Sb = zeros(d, d)
for k = 1:nc
    dv = cmeans[:,k] - mv
    Sb += ns[k] * (dv * dv')
end

Z = X .- mv
Sa = Z * Z'
@test Sw + Sb ≈ Sa

@assert size(X) == (d, n)
@assert size(y) == (n,)

## Stats

S = multiclass_lda_stats(nc, X, y)

@test S.dim == d
@test S.nclasses == nc

@test classweights(S) == ns
@test classmeans(S) ≈ cmeans
@test mean(S) ≈ mv

@test withclass_scatter(S) ≈ Sw
@test betweenclass_scatter(S) ≈ Sb

## Solve

emax = maximum(eigvals(Sw))
Sw_c = copy(Sw)
Sb_c = copy(Sb)

lambda = 1.0e-3
Sw_r = Sw + (lambda * emax) * Matrix(I, d, d)

P1 = mclda_solve(Sb, Sw, :gevd, nc-1, lambda)
@test P1' * Sw_r * P1 ≈ Matrix(I, nc-1, nc-1)
U = Sb * P1
V = Sw_r * P1
# test whether U is proportional to V,
# which indicates P is the generalized eigenvectors
@test U ≈ V*Diagonal(vec(mean(U./V, dims=1)))

P2 = mclda_solve(Sb, Sw, :whiten, nc-1, lambda)
@test P2' * Sw_r * P2 ≈ Matrix(I, nc-1, nc-1)

test_approx_eq_vecs(P1, P2)

## LDA

for T in (Float32, Float64)
    M = fit(MulticlassLDA, nc, convert(Matrix{T}, X), y; method=:gevd, regcoef=convert(T, lambda))
    @test indim(M) == d
    @test outdim(M) == nc - 1
    @test projection(M) ≈ P1
    @test M.pmeans ≈ M.proj'cmeans
    @test transform(M, X) ≈ M.proj'X

    M = fit(MulticlassLDA, nc, convert(Matrix{T}, X), y; method=:whiten, regcoef=convert(T, lambda))
    @test indim(M) == d
    @test outdim(M) == nc - 1
    # @test projection(M) P2  # signs may change
    @test M.pmeans ≈ M.proj'cmeans
    @test transform(M, X) ≈ M.proj'X
end

## High-dimensional LDA (subspace LDA)

# Low-dimensional case (no singularities)

centers = [zeros(5) [10.0;zeros(4)] [0.0;10.0;zeros(3)]]

# Case 1: 3 groups of 500
dX = randn(5,1500);
for i = 0:500:1000
    dX[:,(1:500).+i] .= dX[:,(1:500).+i] .- mean(dX[:,(1:500).+i], dims=2)  # make the mean of each 0
end
dX1 = dX
X1 = [dX[:,1:500].+centers[:,1] dX[:,501:1000].+centers[:,2] dX[:,1001:1500].+centers[:,3]]
label1 = [fill(1,500); fill(2,500); fill(3,500)]
# Case 2: 3 groups, one with 1000, one with 100, and one with 10
dX = randn(5,1110);
dX[:,   1:1000] .= dX[:,   1:1000] .- mean(dX[:,   1:1000], dims=2)
dX[:,1001:1100] .= dX[:,1001:1100] .- mean(dX[:,1001:1100], dims=2)
dX[:,1101:1110] .= dX[:,1101:1110] .- mean(dX[:,1101:1110], dims=2)
dX2 = dX
X2 = [dX[:,1:1000].+centers[:,1] dX[:,1001:1100].+centers[:,2] dX[:,1101:1110].+centers[:,3]]
label2 = [fill(1,1000); fill(2,100); fill(3,10)]

for (X, dX, label) in ((X1, dX1, label1), (X2, dX2, label2))
    w = [length(findall(label.==1)), length(findall(label.==2)), length(findall(label.==3))]
    M = fit(SubspaceLDA, X, label)
    @test indim(M) == 5
    @test outdim(M) == 2
    totcenter = vec(sum(centers.*w', dims=2)./sum(w))
    @test mean(M) ≈ totcenter
    @test classmeans(M) ≈ centers
    @test classweights(M) == w
    x = rand(5)
    @test transform(M, x) ≈ projection(M)'*x
    dcenters = centers .- totcenter
    Hb = dcenters.*sqrt.(w)'
    Sb = Hb*Hb'
    Sw = dX*dX'
    # When there are no singularities, we need to solve the "regular" LDA eigenvalue equation
    proj = projection(M)
    @test Sb*proj ≈ Sw*proj*Diagonal(M.λ)
    @test proj'*Sw*proj ≈ Matrix(I,2,2)

    # also check that this is consistent with the conventional algorithm
    Mld = fit(MulticlassLDA, 3, X, label)
    for i = 1:2
        @test abs(dot(normalize(proj[:,i]), normalize(Mld.proj[:,i]))) ≈ 1
    end
end

# High-dimensional case (undersampled => singularities)
X = randn(10^6, 9)
label = rand(1:3, 9); label[1:3] = 1:3
M = fit(SubspaceLDA, X, label)
centers = M.cmeans
for i = 1:3
    flag = label.==i
    @test centers[:,i] ≈ mean(X[:,flag], dims=2)
    @test M.cweights[i] == sum(flag)
end
dcenters = centers .- mean(X, dims=2)
Hb = dcenters*Diagonal(sqrt.(M.cweights))
Hw = X - centers[:,label]
@test (M.projw'*Hb)*(Hb'*M.projw)*M.projLDA ≈ (M.projw'*Hw)*(Hw'*M.projw)*M.projLDA*Diagonal(M.λ)
Gw = projection(M)'*Hw
@test Gw*Gw' ≈ Matrix(I,2,2)

# Test that nothing breaks if one class has no members
label[1:4] .= 1
label[5:9] .= 3  # no data points have label == 2
M = fit(SubspaceLDA, X, label)
centers = M.cmeans
dcenters = centers .- mean(X, dims=2)
Hb = dcenters*Diagonal(sqrt.(M.cweights))
Hw = X - centers[:,label]
@test (M.projw'*Hb)*(Hb'*M.projw)*M.projLDA ≈ (M.projw'*Hw)*(Hw'*M.projw)*M.projLDA*Diagonal(M.λ)

# Test normalized LDA
function gen_ldadata_2(centers, n1, n2)
    d = size(centers, 1)
    X = randn(d, n1+n2)
    X[:,1:n1]       .-= vec(mean(X[:,1:n1], dims=2))
    X[:,n1+1:n1+n2] .-= vec(mean(X[:,n1+1:n1+n2], dims=2))
    dX = copy(X)
    X[:,1:n1]       .+= centers[:,1]
    X[:,n1+1:n1+n2] .+= centers[:,2]
    label = [fill(1,n1); fill(2,n2)]
    X, dX, label
end
centers = zeros(2,2); centers[1,2] = 10
n1 = 100
for n2 in (100, 1000, 10000)
    X, dX, label = gen_ldadata_2(centers, n1, n2)
    M = fit(SubspaceLDA, X, label; normalize=true)
    proj = projection(M)
    Hb = centers .- mean(centers, dims=2)   # not weighted by number in each cluster
    Sb = Hb*Hb'
    Sw = dX*dX'/(n1+n2)
    @test Sb*proj ≈ Sw*proj*Diagonal(M.λ)
end
end


text:
module PredictiveControl

################################################################################
# Main exports from this package
################################################################################
export ConstrainedTimeInvariantLQR,
       PrimalFullCondensing, PrimalNoCondensing, DualFullCondensing

# Analysis functions
export conditionBound

# Preconditioners
export circulantblockpreconditioner, sdpblockpreconditoner

################################################################################
# Packages we need
################################################################################
using BlockArrays
using ControlSystems
using ControlSystems: nstates, ninputs
using COSMO
using FFTW
using JuMP
using LinearAlgebra
using Logging
using MathOptInterface
using MatrixEquations
using OSQP

const MOI = MathOptInterface

# Overload these functions for our own types
import LinearAlgebra: cond

include( "utilities.jl" )

include( "types/problemTypes.jl" )
include( "types/clqr.jl" )

include( "matrices/fullCondensing.jl" )
include( "matrices/preconditioning.jl" )

include( "analysis/conditioning.jl")

# Solvers
include( "solvers/fastGradientMethod/fgm.jl" )
using .FastGradientMethod
const FGM = FastGradientMethod

export FGM, fastgradientmethod

# Solver interfaces
include( "solvers/interfaces/fgm_interface.jl" )
include( "solvers/interfaces/jump_interface.jl" )
include( "solvers/interfaces/osqp_interface.jl" )

end # module PredictiveControl


text:
# This script processes the QIBA Phantom data using non-linear fitting
# for the original and constrained reference region models

# Run this script by launching julia then entering: include("path/to/e02_2a_qibaProcessNLSQ.jl")

# Estimated runtime: ~14 seconds

# First, figure out the current file path
curFilePath = Base.source_path() # Path to current file
workDir = dirname(curFilePath) # Our working directory (directory of current file)
## If the above doesn't work, then user has to manually enter the location of the CLRRM directory

## Load packages
import DCEMRI
using MAT

# Some required files
matFile = workDir * "/data/QIBA-ToftsV6/QIBAv6-Mini.mat"
warmupFile = workDir * "/data/QIBA-ToftsV6/QIBAv6-Mini4Jl-Warmup.mat"
auxCode = workDir * "/jlfiles/auxCode.jl"
refRegCode = workDir * "/jlfiles/refRegionFunction.jl"

outDir = workDir * "/data/QIBA-ToftsV6/Results"

if !(isdir(outDir))
  mkdir(outDir)
end

# Load pre-requisite code
include(auxCode)

## Warmup run to make sure all is well
results = DCEMRI.fitdata(datafile=warmupFile, models=[2])
rm("output.mat") # Delete output since it isn't needed

# If warmup succeeds, then proceed by loading Reference Region Model
include(refRegCode)

# Load the phantom data
mat = matread(matFile)
Ct = mat["concData"]
Ct = permutedims(Ct, [3, 1, 2])
t = vec(mat["T"])
Cp = vec(mat["aif"])
(sT, sX, sY) = size(Ct)
Ct = reshape(Ct,sT,sX*sY)
mask = ones(sX*sY) .> 0

## Using 'refY' parameters
refKTrans = 0.1;
refVe = 0.1;
refKep = refKTrans/refVe
Crr = DCEMRI.toftskety(t, [refKTrans,refKep], Cp)

# Fit the NRRM and CNRRM to the QIBA phantom data
(paramsCN, residCN, estKepRR, paramsN, residN, runtimeN, runtimeCN)=fitCNRRM(Ct, Crr, t, estKepRR=0.0, doTime=true)

# Export to .mat
results = Dict()
  results["pkParamsN"] = paramsN
  results["pkParamsCN"] = paramsCN
  results["residN"] = residN
  results["residCN"] = residCN
  results["estKepRRN"] = estKepRR
  results["refKt"] = refKTrans
  results["refVe"] = refVe
  results["refKep"] = refKep
  matwrite(outDir * "/QIBAv6-refY-NL.mat", results)

max_stars_repo_name: MPUmri/CRRM

max_issues_repo_name: MPUmri/CRRM

max_forks_repo_name: MPUmri/CRRM


text:
#==
Plot ZonedDateTime, on x-axis.
We convert it to DateTimes, in the local timezone,
and we list that timezone in the x-axis label.

This is just one of many options for how to display this info.
See this PR: https://github.com/JuliaTime/TimeZones.jl/pull/251#issuecomment-603806198
for details on the options and their tradeoffs.
==#
@recipe function f(xs::AbstractVector{<:ZonedDateTime}, ys)
    if !isempty(xs)
        tz = timezone(first(xs))  # convert all to same timezone as first one
        new_xs = DateTime.(astimezone.(xs, tz), Local)

        # xguide is the proper name for xlabel
        label = get(plotattributes, :xguide, "")
        label *= isempty(label) ? "Time zone: $tz" : " ($tz)"
        xguide := label
        new_xs, ys
    else
        # Can't check timezone if empty so we just return an empty list
        [], ys
    end
end


text:
module System

using ComponentArrays
import Flight.Plotting: plots

export assemble_x0, init_x0, init_y0, init_u0, init_d0, f_cont!, f_disc!
export AbstractComponent, AbstractSystem, System

no_extend_error(f::Function, ::Type{S}) where {S} = error(
    "Function $f not implemented for type $S or incorrect call signature")
no_extend_warning(f::Function, ::Type{S}) where {S} = println(
    "Warning: Function $f not implemented for type $S or incorrect call signature")

#anything around which we can build a System
abstract type AbstractComponent end #anything that can go in a System

#init_x0 must return either nothing, an AbstractVector{<:Real} or a NT with a
#Union{AbstractVector{<:Real},Nothing}
init_x0(::C) where {C<:AbstractComponent} = nothing
init_y0(::C) where {C<:AbstractComponent} = nothing
init_u0(::C) where {C<:AbstractComponent} = nothing #sytems are not required to have control inputs
init_d0(::C) where {C<:AbstractComponent} = nothing #systems are not required to have discrete states

abstract type AbstractSystem{C<:AbstractComponent} end

#need the C type parameter for dispatch, the rest for type stability
mutable struct System{C, X <: Union{Nothing, AbstractVector{<:Real}},
                    Y, U, D, P, S} <: AbstractSystem{C}
    ẋ::X #continuous state vector derivative
    x::X #continuous state vector (to be used as a buffer for f_cont! evaluation)
    y::Y #output state
    u::U #control inputs
    d::D #discrete state
    t::Base.RefValue{Float64} #this allows propagation of t updates down the subsystem hierarchy
    params::P
    subsystems::S
end

f_cont!(::S, args...) where {S<:AbstractSystem} = no_extend_error(f_cont!, S)
(f_disc!(::S, args...)::Bool) where {S<:AbstractSystem} = no_extend_error(f_disc!, S)

#f_disc! is free to modify a Hybrid system's discrete state, control inputs and
#continuous state. if it modifies the latter, it must return true, false
#otherwise. it is dangerous to provide a default fallback for f_disc!, because
#if the intended f_disc! implementation for the System has the wrong interface,
#the dispatch will revert to the fallback, which may not be obvious at all. it
#is safer to force each concrete System that does not require an actual f_disc!
#to implement a trivial f_disc! that returns false

#when the System constructor for a certain Component is passed no
#parameters for dx and x, it calls the init_x0 method for that Component, which
#may return

assemble_x0(c::AbstractComponent) = assemble_x0(init_x0(c))

assemble_x0(::Nothing) = (s = nothing, ss = nothing)

assemble_x0(x::AbstractVector{<:Real}) = (s = x, ss = nothing)

function assemble_x0(x::ComponentVector)

    x_ss = Dict{Symbol,Any}()

    for id in keys(x)
        x_ss[id] = view(x, id)
    end

    return (s = x, ss = NamedTuple{Tuple(keys(x_ss))}(values(x_ss)))

end

# function assemble_x0(nt::NamedTuple{L,X}
#     ) where {L, X <: NTuple{N, Union{AbstractVector{<:Real}, Nothing}} where {N}}
function assemble_x0(nt::NamedTuple)
    #assembles the System's state vector from the non-empty subsystems' state
    #vectors. returns the System's state vector and the views to be assigned to
    #the subsystems
    x_ss = Dict{Symbol,Any}(pairs(nt))

    #para cada entry del nt, llamamos a assemble_x0, que nos devolvera un x_s y
    #un x_ss

    x_s_blocks = filter(x_ss) do item
        isa(item.second, AbstractVector{<:Real})
    end
    x_s = (!isempty(x_s_blocks) ? ComponentVector{Float64}(x_s_blocks) : nothing)

    for id in keys(x_ss)
        x_ss[id] = (id in keys(x_s_blocks) ? view(x_s, id) : x_ss[id])
    end

    return (s = x_s, ss = NamedTuple{Tuple(keys(x_ss))}(values(x_ss)))

end

function plots(::AbstractVector{<:Real}, ::AbstractVector{T}) where {T}
    no_extend_warning(plots, T) #nothing to plot by default, warn about it
end

end


text:
@doc raw"""
`LODE`: Lagrangian Ordinary Differential Equation

Defines an implicit initial value problem
```math
\begin{aligned}
\dot{q} (t) &= v(t) , &
q(t_{0}) &= q_{0} , \\
\dot{p} (t) &= f(t, q(t), v(t)) , &
p(t_{0}) &= p_{0} , \\
p(t) &= ϑ(t, q(t), v(t))
\end{aligned}
```
with momentum ``p`` and force field ``f``, given by
```math
\begin{aligned}
p &= \frac{\partial L}{\partial v} , &
f &= \frac{\partial L}{\partial q} ,
\end{aligned}
```
initial conditions ``(q_{0}, p_{0})`` and the solution ``(q,p)`` taking values
in ``\mathbb{R}^{d} \times \mathbb{R}^{d}``.
This is a special case of a differential algebraic equation with dynamical
variables ``(q,p)`` and algebraic variable ``v``.

### Parameters

* `DT <: Number`: data type
* `TT <: Real`: time step type
* `AT <: AbstractArray{DT}`: array type
* `ϑType <: Function`: type of `ϑ`
* `fType <: Function`: type of `f`
* `gType <: Function`: type of `g`
* `ωType <: Function`: type of `ω`
* `v̄Type <: Function`: type of `v̄`
* `f̄Type <: Function`: type of `f̄`
* `lagType <: Function`: Lagrangian type
* `invType <: OptionalNamedTuple`: invariants type
* `parType <: OptionalNamedTuple`: parameters type
* `perType <: OptionalArray{AT}`: periodicity type

### Fields

* `d`: dimension of dynamical variables ``q`` and ``p`` as well as the vector fields ``f`` and ``p``
* `ϑ`: function determining the momentum
* `f`: function computing the vector field
* `g`: function determining the projection, given by ``\nabla \vartheta (q) \cdot \lambda``
* `ω`: function computing the symplectic matrix
* `v̄`: function computing an initial guess for the velocity field ``v`` (optional)
* `f̄`: function computing an initial guess for the force field ``f`` (optional)
* `t₀`: initial time (optional)
* `q₀`: initial condition for `q`
* `p₀`: initial condition for `p`
* `λ₀`: initial condition for `λ` (optional)
* `lagrangian`: function computing the Lagrangian ``L``
* `invariants`: either a `NamedTuple` containing the equation's invariants or `nothing`
* `parameters`: either a `NamedTuple` containing the equation's parameters or `nothing`
* `periodicity`: determines the periodicity of the state vector `q` for cutting periodic solutions

The functions `ϑ` and `f` must have the interface
```julia
    function ϑ(t, q, v, p)
        p[1] = ...
        p[2] = ...
        ...
    end
```
and
```julia
    function f(t, q, v, f)
        f[1] = ...
        f[2] = ...
        ...
    end
```
where `t` is the current time, `q` is the current solution vector, `v` is the
current velocity and `f` and `p` are the vectors which hold the result of
evaluating the functions ``f`` and ``ϑ`` on `t`, `q` and `v`.
The funtions `g` and `v` are specified by
```julia
    function g(t, q, λ, g)
        g[1] = ...
        g[2] = ...
        ...
    end
```
and
```julia
    function v(t, q, p, v)
        v[1] = ...
        v[2] = ...
        ...
    end
```

### Constructors

```julia
LODE(ϑ, f, ω, v̄, f̄, t₀, q₀, p₀, λ₀, lagrangian, invariants, parameters, periodicity)

LODE(ϑ, f, l, ω, t₀, q₀::StateVector, p₀::StateVector, λ₀::StateVector=zero(q₀); kwargs...)
LODE(ϑ, f, l, ω, q₀::StateVector, p₀::StateVector, λ₀::StateVector=zero(q₀); kwargs...)
LODE(ϑ, f, l, ω, t₀, q₀::State, p₀::State, λ₀::StateVector=zero(q₀); kwargs...)
LODE(ϑ, f, l, ω, q₀::State, p₀::State, λ₀::StateVector=zero(q₀); kwargs...)
```

### Keyword arguments

* `v̄ = (t,q,v) -> nothing`
* `f̄ = f`
* `invariants = nothing`
* `parameters = nothing`
* `periodicity = nothing`

"""
struct LODE{ϑType <: Callable, fType <: Callable, gType <: Callable,
            ωType <: Callable, v̄Type <: Callable, f̄Type <: Callable,
            lagType <: Callable,
            invType <: OptionalInvariants,
            parType <: OptionalParameters,
            perType <: OptionalPeriodicity} <: AbstractEquationPODE{invType,parType,perType}

    ϑ::ϑType
    f::fType
    g::gType
    ω::ωType

    v̄::v̄Type
    f̄::f̄Type

    lagrangian::lagType
    invariants::invType
    parameters::parType
    periodicity::perType

    function LODE(ϑ, f, g, ω, v̄, f̄, lagrangian, invariants, parameters, periodicity)
        @assert !isempty(methods(ϑ))
        @assert !isempty(methods(f))
        @assert !isempty(methods(g))
        @assert !isempty(methods(ω))
        @assert !isempty(methods(v̄))
        @assert !isempty(methods(f̄))
        @assert !isempty(methods(lagrangian))

        new{typeof(ϑ), typeof(f), typeof(g), typeof(ω), typeof(v̄), typeof(f̄),
            typeof(lagrangian), typeof(invariants), typeof(parameters), typeof(periodicity)}(
                ϑ, f, g, ω, v̄, f̄, lagrangian, invariants, parameters, periodicity)
    end
end

_lode_default_v̄(t,q,v,params) = nothing

LODE(ϑ, f, g, ω, l; invariants=NullInvariants(), parameters=NullParameters(), periodicity=NullPeriodicity(), v̄=_lode_default_v̄, f̄=f) = LODE(ϑ, f, g, ω, v̄, f̄, l, invariants, parameters, periodicity)

GeometricBase.invariants(equation::LODE) = equation.invariants
GeometricBase.parameters(equation::LODE) = equation.parameters
GeometricBase.periodicity(equation::LODE) = equation.periodicity

hasvectorfield(::LODE) = true
haslagrangian(::LODE) = true

function check_initial_conditions(::LODE, ics::NamedTuple)
    haskey(ics, :q) || return false
    haskey(ics, :v) || haskey(ics, :p) || return false
    haskey(ics, :λ) || return false
    eltype(ics.q) == eltype(ics.p) == eltype(ics.λ) || return false
    typeof(ics.q) == typeof(ics.p) == typeof(ics.λ) || return false
    axes(ics.q) == axes(ics.p) == axes(ics.λ) || return false
    return true
end

function check_methods(equ::LODE, tspan, ics, params)
    applicable(equ.ϑ, tspan[begin], ics.q, zero(ics.q), zero(ics.p), params) || return false
    applicable(equ.f, tspan[begin], ics.q, zero(ics.q), zero(ics.p), params) || return false
    applicable(equ.g, tspan[begin], ics.q, zero(ics.q), zero(ics.p), params) || return false
    applicable(equ.lagrangian, tspan[begin], ics.q, zero(ics.q), params) || return false
    return true
end

function datatype(equ::LODE, ics::NamedTuple)
    @assert check_initial_conditions(equ, ics)
    return eltype(ics.q)
end

function arrtype(equ::LODE, ics::NamedTuple)
    @assert check_initial_conditions(equ, ics)
    typeof(ics.q)
end

_get_ϑ(equ::LODE, params) = (t,q,v,ϑ) -> equ.ϑ(t, q, v, ϑ, params)
_get_f(equ::LODE, params) = (t,q,v,f) -> equ.f(t, q, v, f, params)
_get_g(equ::LODE, params) = (t,q,v,g) -> equ.g(t, q, v, g, params)
_get_ω(equ::LODE, params) = (t,q,v,ω) -> equ.ω(t, q, v, ω, params)
_get_v̄(equ::LODE, params) = (t,q,v)   -> equ.v̄(t, q, v, params)
_get_f̄(equ::LODE, params) = (t,q,v,f) -> equ.f̄(t, q, v, f, params)
_get_l(equ::LODE, params) = (t,q,v)   -> equ.lagrangian(t, q, v, params)

_functions(equ::LODE) = (ϑ = equ.ϑ, f = equ.f, g = equ.g, ω = equ.ω, v̄ = equ.v̄, f̄ = equ.f̄, l = equ.lagrangian)
_functions(equ::LODE, params::OptionalParameters) = (ϑ = _get_ϑ(equ, params),
                                                     f = _get_f(equ, params),
                                                     g = _get_g(equ, params),
                                                     ω = _get_ω(equ, params),
                                                     v̄ = _get_v̄(equ, params),
                                                     f̄ = _get_f̄(equ, params),
                                                     l = _get_l(equ, params))


text:
# magma_sgesvd wrapper
function magmaSgesvd(jobu, jobvt,
	 m, n,
	 A, lda,
	 s,
	 U, ldu,
	 VT, ldvt,
	 work, lwork, info)
    ccall((:magma_sgesvd, libmagma),
               Cint,
               (Cint, Cint, 		# jobu , jobvt
			    Cint, Cint, 		# m, n
				PtrOrCuPtr{Cfloat}, Cint, 	# A, lda
				PtrOrCuPtr{Cfloat}, 		# s
				PtrOrCuPtr{Cfloat}, Cint, 	# U, ldu
				PtrOrCuPtr{Cfloat}, Cint, 	# VT, ldvt
				PtrOrCuPtr{Cfloat}, Cint, PtrOrCuPtr{Cint}),# work, lwork, info

                jobu, jobvt,
				m, n,
				A, lda,
				s,
				U, ldu,
				VT, ldvt,
				work, lwork, info)
end


text:
# * existing: the ancillas and input arguments in the local scope.
#       They should be protected to avoid duplicated allocation.
# * deallocated: the ancillas removed.
#       They should be recorded to avoid using after deallocation.
# * unclassified: the variables from global scope.
#       They can not be allocation target.
struct SymbolTable
    existing::Vector{Symbol}
    deallocated::Vector{Symbol}
    unclassified::Vector{Symbol}
end

function SymbolTable()
    SymbolTable(Symbol[], Symbol[], Symbol[])
end

Base.copy(st::SymbolTable) = SymbolTable(copy(st.existing), copy(st.deallocated), copy(st.unclassified))

# remove a variable from a list
function removevar!(lst::AbstractVector, var)
    index = findfirst(==(var), lst)
    deleteat!(lst, index)
end

# replace a variable in a list with target variable
function replacevar!(lst::AbstractVector, var, var2)
    index = findfirst(==(var), lst)
    lst[index] = var2
end

# allocate a new variable
function allocate!(st::SymbolTable, var::Symbol)
    if var ∈ st.existing
        throw(InvertibilityError("Repeated allocation of variable `$(var)`"))
    elseif var ∈ st.deallocated
        removevar!(st.deallocated, var)
        push!(st.existing, var)
    elseif var ∈ st.unclassified
        throw(InvertibilityError("Variable `$(var)` used before allocation."))
    else
        push!(st.existing, var)
    end
    nothing
end

# find the list containing var
function findlist(st::SymbolTable, var::Symbol)
    if var ∈ st.existing
        return st.existing
    elseif var ∈ st.unclassified
        return st.unclassified
    elseif var in st.deallocated
        return st.deallocated
    else
        return nothing
    end
end

# using a variable
function operate!(st::SymbolTable, var::Symbol)
    if var ∈ st.existing || var ∈ st.unclassified
    elseif var ∈ st.deallocated
        throw(InvertibilityError("Operating on deallocate variable `$(var)`"))
    else
        push!(st.unclassified, var::Symbol)
    end
    nothing
end

# deallocate a variable
function deallocate!(st::SymbolTable, var::Symbol)
    if var ∈ st.deallocated
        throw(InvertibilityError("Repeated deallocation of variable `$(var)`"))
    elseif var ∈ st.existing
        removevar!(st.existing, var)
        push!(st.deallocated, var)
    elseif var ∈ st.unclassified
        throw(InvertibilityError("Deallocating an external variable `$(var)`"))
    else
        throw(InvertibilityError("Deallocating an external variable `$(var)`"))
    end
    nothing
end

# check symbol table to make sure there is symbols introduced in the local scope that has not yet deallocated.
# `a` is the symbol table after running local scope, `b` is the symbol table before running the local scope.
function checksyms(a::SymbolTable, b::SymbolTable=SymbolTable())
    diff = setdiff(a.existing, b.existing)
    if !isempty(diff)
        error("Some variables not deallocated correctly: $diff")
    end
end

function swapsyms!(st::SymbolTable, var1::Symbol, var2::Symbol)
    lst1 = findlist(st, var1)
    lst2 = findlist(st, var2)
    if lst1 !== nothing && lst2 !== nothing
        # exchange variables
        i1 = findfirst(==(var1), lst1)
        i2 = findfirst(==(var2), lst2)
        lst2[i2], lst1[i1] = lst1[i1], lst2[i2]
    elseif lst1 !== nothing
        replacevar!(lst1, var1, var2)
        operate!(st, var1)
    elseif lst2 !== nothing
        replacevar!(lst2, var2, var1)
        operate!(st, var2)
    else
        operate!(st, var1)
        operate!(st, var2)
    end
end

function swapsyms_asymetric!(st::SymbolTable, var1s::Vector, var2::Symbol)
    length(var1s) == 0 && return
    lst1 = findlist(st, var1s[1])
    for k=2:length(var1s)
        if findlist(st, var1s[k]) !== lst1
            error("variable status not aligned: $var1s")
        end
    end
    lst2 = findlist(st, var2)
    if lst1 !== nothing
        removevar!.(Ref(lst1), var1s)
        push!(lst1, var2)
    else
        operate!(st, var2)
    end
    if lst2 !== nothing
        removevar!(lst2, var2)
        push!.(Ref(lst2),var1s)
    else
        operate!.(Ref(st), var1s)
    end
end


text:
using AcousticFWI,Seismic,PyPlot

function main()

    nz = 150
    nx = 250
    nt = 1024
    nf = 1024
    dz = 0.01
    dx = 0.01
    dt = 0.002
    f0 = 10.
    fmin = 1.0
    fmax = 30.
    ext = 50
    atten_max = 2.5

    vp = 2.1*ones(nz,nx)
    vp[61:100,:] = 3.0
    vp[101:end,:] = 2.6

    isz = [3]
    isx = [125]
    igz = 3*ones(Int,nx)
    igx = [1:nx;]
    ot = [0.]

    wav = Ricker(f0=f0,dt=dt)
    u = Helmholtz(isz,isx,ot,vp,wav,fmin,fmax,nf,nt,dz,dx,dt,ext,atten_max)
    d = u[:,3,igx]

    figure(1) ; imshow(1000*vp,cmap="YlOrBr",extent=[0.,(nx-1)*1000*dx,(nz-1)*1000*dz,0.])
              ; xlabel("x (m)")
              ; ylabel("z (m)")
              ; colorbar(ticks=[2100,2300,2500,2700,2900])
              ; ax = gca()
              ; ax[:set_xticks]([0,500,1000,1500,2000])
              ; ax[:set_yticks]([0,300,600,900,1200,1500])
    figure(2) ; subplots_adjust(wspace=0.5,hspace=0.5)
              ; subplot(3,3,1) ; SeisPlot(u[10,:,:],pclip=100,dx=1000*dx,dy=1000*dz,xticks=[0,500,1000,1500,2000],yticks=[0,300,600,900,1200,1500],xlabel="x (m)",ylabel="z (m)",title="t = 0.02s",titlesize=14,labelsize=12,cmap="gray",fignum=2)
              ; subplot(3,3,2) ; SeisPlot(u[75,:,:],pclip=100,dx=1000*dx,dy=1000*dz,xticks=[0,500,1000,1500,2000],yticks=[0,300,600,900,1200,1500],xlabel="x (m)",ylabel="z (m)",title="t = 0.15s",titlesize=14,labelsize=12,cmap="gray",fignum=2)
              ; subplot(3,3,3) ; SeisPlot(u[140,:,:],pclip=100,dx=1000*dx,dy=1000*dz,xticks=[0,500,1000,1500,2000],yticks=[0,300,600,900,1200,1500],xlabel="x (m)",ylabel="z (m)",title="t = 0.28s",titlesize=14,labelsize=12,cmap="gray",fignum=2)
              ; subplot(3,3,4) ; SeisPlot(u[205,:,:],pclip=100,dx=1000*dx,dy=1000*dz,xticks=[0,500,1000,1500,2000],yticks=[0,300,600,900,1200,1500],xlabel="x (m)",ylabel="z (m)",title="t = 0.41s",titlesize=14,labelsize=12,cmap="gray",fignum=2)
              ; subplot(3,3,5) ; SeisPlot(u[270,:,:],pclip=100,dx=1000*dx,dy=1000*dz,xticks=[0,500,1000,1500,2000],yticks=[0,300,600,900,1200,1500],xlabel="x (m)",ylabel="z (m)",title="t = 0.54s",titlesize=14,labelsize=12,cmap="gray",fignum=2)
              ; subplot(3,3,6) ; SeisPlot(u[335,:,:],pclip=100,dx=1000*dx,dy=1000*dz,xticks=[0,500,1000,1500,2000],yticks=[0,300,600,900,1200,1500],xlabel="x (m)",ylabel="z (m)",title="t = 0.67s",titlesize=14,labelsize=12,cmap="gray",fignum=2)
              ; subplot(3,3,7) ; SeisPlot(u[400,:,:],pclip=100,dx=1000*dx,dy=1000*dz,xticks=[0,500,1000,1500,2000],yticks=[0,300,600,900,1200,1500],xlabel="x (m)",ylabel="z (m)",title="t = 0.70s",titlesize=14,labelsize=12,cmap="gray",fignum=2)
              ; subplot(3,3,8) ; SeisPlot(u[465,:,:],pclip=100,dx=1000*dx,dy=1000*dz,xticks=[0,500,1000,1500,2000],yticks=[0,300,600,900,1200,1500],xlabel="x (m)",ylabel="z (m)",title="t = 0.83s",titlesize=14,labelsize=12,cmap="gray",fignum=2)
              ; subplot(3,3,9) ; SeisPlot(u[530,:,:],pclip=100,dx=1000*dx,dy=1000*dz,xticks=[0,500,1000,1500,2000],yticks=[0,300,600,900,1200,1500],xlabel="x (m)",ylabel="z (m)",title="t = 0.96s",titlesize=14,labelsize=12,cmap="gray",fignum=2)
    SeisPlot(d,pclip=98,dx=1,dy=dt,xlabel="Trace",ylabel="t (s)",cmap="gray",fignum=3)

end

main()


text:
using Crystalline
using MPBUtils
write_dir = (@__DIR__)*"" # NB: make sure whatever directory this points to actually exists

function init_lattice(sgnum, Dᵛ::Val{D}, cntr::Char;
                maxGs=ntuple(_->2, Val(D)), expon::Real=1.25) where D

    # level-set surface
    flat = normscale!(modulate(levelsetlattice(sgnum, Dᵛ, maxGs)), expon)
    deleteat!(flat.orbits, 1), deleteat!(flat.orbitcoefs, 1) # trivial constant G=0 term

    # make lattice primitivize
    pflat = primitivize(flat, cntr)

    return pflat
end

# --- choices ---
sgnum = 68    # space group number
D = 3         # dimension
cntr = centering(sgnum, D) # centering symbol (e.g, 'F' for face-centered, etc.)
write_to_file = true
has_tr = true # whether we have time-reversal or not
res = 32      # resolution used in MPB
nbands = 12   # number of bands requested from MPB

# --- find out which little groups we need to assess bandreps/symmetry vector ---
brs  = bandreps(sgnum, D, timereversal=has_tr)
lgs  = Crystalline.matching_littlegroups(brs, Val(D))
plgs = primitivize.(lgs, #=modw=# false)

# --- generate a bunch of mpb-input files ---
ids = 1:10 # create 10 different lattices and "input-files"
runtype = "all" # we don't have TE/TM in 3D; do all polarizations
for id in ids
    # direct basis
    Rs  = directbasis(sgnum, Val(D)) # generate a random basis consistent with `sgnum`
    pRs = primitivize(Rs, cntr)

    # generate a isosurface
    pflat = init_lattice(sgnum, Val(D), cntr)

    # generate some dielectric values
    filling = rand(0.25:0.05:0.75)
    epsin = rand((7.0, 9.0, 11.0, 13.0, 15.0))
    epsout = 1.0

    # write to .sh file for mpb
    filename = mpb_calcname(D, sgnum, id, res, runtype)
    filepath = joinpath(write_dir, filename*".sh")
    write_to_file || continue
    open(filepath, "w") do io
        prepare_mpbcalc!(io, sgnum, pflat, pRs, filling, epsin, epsout, runtype;
                                res=res, lgs=plgs, id=id, nbands=nbands)
        # NB: if you want an mpb dispersion calculation, provide a keyword argument `kvs`
        #     instead of `lgs` in the above. You can also supply a filename for `kvs` if
        #     you want.
    end
    id == 1 && println("MPB setup files written to:")
    println("  ∙  ", filepath, ".sh")
end


text:
#
# This file is part of the YAActL.jl Julia package, MIT license
#
# Paul Bayer, 2020
#

"""
    Func(f, args...; kwargs...)

A structure for passing a function `f` and its arguments
to an actor.
"""
struct Func{X,Y,Z}
    f::X
    args::Y
    kwargs::Z

    Func(f::F, args...; kwargs...) where F<:Function =
        new{typeof(f),typeof(args),typeof(kwargs)}(f, args, kwargs)
    Func() = new{typeof(+),typeof(()),typeof(pairs(()))}(+,(),pairs(()))
end

"""
    Args(args...; kwargs...)

A structure for updating arguments to an actor's behavior.
"""
struct Args{A,B}
    args::A
    kwargs::B

    Args(args...; kwargs...) = new{typeof(args),typeof(kwargs)}(args, kwargs)
end

"""
    Become(x::Func)

An asynchronous [`Message`](@ref) to an actor to change its
behavior.
"""
struct Become <: Message
    x::Func
end

"""
    Call(arg, from::Link)

A synchronous [`Message`](@ref) to an actor to execute its
behavior with `arg...` and to send the result as a [`Response`](@ref)
message to `from`.

If the actor is set to `state` dispatch, it updates its internal
state with the result.
"""
struct Call <: Message
    x
    from::Link
end
Call(from) = Call((),from)

"""
    Cast(arg)

An asynchronous [`Message`](@ref) to an actor to execute
its behavior with `arg...` without sending a response.

If the actor is set to `state` dispatch, it updates its internal
state with the result.
"""
struct Cast <: Message
    x
end
Cast() = Cast(())

"""
    Diag(from::Link)

A synchronous [`Message`](@ref) to an actor to send a
`Response` message with its internal `_ACT` variable to `from`.
"""
struct Diag <: Message
    from::Link
end

"""
    Exec(func::Func, from::Link)

A synchronous [`Message`](@ref) to an actor to execute `func`
and to send a `Response` message with the return value to `from`.
"""
struct Exec <: Message
    func
    from::Link
end

"""
    Init(f::Func)

A [`Message`](@ref) to an actor to execute the given
[`Func`](@ref) and to register it in the [`_ACT`](@ref)
variable.
"""
struct Init <: Message
    x::Func
end

"""
    Query(s::Symbol, from::Link)

A [`Message`](@ref) to an actor to send a
`Response` message with one of its internal state
variables `s` to `from`.

- `s::Symbol` can be one of `:sta`, `:res`, `:bhv`, `:dsp`.
"""
struct Query <: Message
    x::Symbol
    from::Link
end

"""
    Request(x, from::Link)

A generic [`Message`](@ref) for user requests.
"""
struct Request <: Message
    x
    from::Link
end

"""
    Response(y, from::Link=self())

A [`Message`](@ref) representing a response to requests.

# Fields
- `y`: response content,
- `from::Link`: sender link.
"""
struct Response <: Message
    y
    from::Link
end
Response(y) = Response(y, self())

"""
    Stop(code=0)

A [`Message`](@ref) causing an actor to stop with an exit
`code`. If present, it calls its [`term!`](@ref) function with
`code` as last argument.
"""
struct Stop <: Message
    x::Int
end
Stop() = Stop(0)

"""
    Term(x::Func)

A [`Message`](@ref) to an actor to save the given [`Func`](@ref)
and to execute it upon termination.
"""
struct Term <: Message
    x::Func
end

"""
    Timeout()

A return value to signal that a timeout has occurred.
"""
struct Timeout <: Message end

"""
    Update(s::Symbol, x)

An asynchronous [`Message`](@ref) to an actor to update its
internal state `s` to `x`.

- `s::Symbol` can be one of `:sta`, `:dsp`, `:arg`, `:lnk`.
"""
struct Update <: Message
    s::Symbol
    x
end


text:
using Documenter, HomotopyContinuation
import LinearAlgebra

makedocs(
    sitename = "Homotopy Continuation",
    pages = [
        "Introduction" => "index.md",
        "Solving polynomial systems" => [
            "solving.md",
            "solver.md"
        ],
        "Solving parametrized systems with monodromy" => "monodromy.md",
        "Tracking paths" => [
            "path_tracker.md",
            "core_tracker.md",
        ],
        "Homotopies" => "homotopies.md",
        "Polynomial systems" => "systems.md",
        "Reference" => "reference.md"
        ],
    strict=false)

deploydocs(
    repo = "github.com/JuliaHomotopyContinuation/HomotopyContinuation.jl.git"
)


text:
module SentimentAnalysis

using Requests, URIParser, Macros, Logger, App

mutable struct SentimentData
  polarity::Int
  polarity_name::String
  subjective::Bool
end

const AYLIEN_API_URL = "https://api.aylien.com/api/v1/sentiment"
const AYLIEN_APP_KEY_HEADER = "X-AYLIEN-TextAPI-Application-Key"
const AYLIEN_APP_ID_HEADER = "X-AYLIEN-TextAPI-Application-ID"

function analyze(text::String) :: SentimentData
  sentiment = Requests.get( AYLIEN_API_URL * "?mode=tweet&text=$(URIParser.escape(text))",
                            headers = Dict{String,String}(AYLIEN_APP_KEY_HEADER => App.AYLIEN_TOKENS[:api_key],
                                                          AYLIEN_APP_ID_HEADER => App.AYLIEN_TOKENS[:app_id])) |> Requests.json |> to_sentimentdata
end

function to_sentimentdata(data::Dict{String,Any}) :: SentimentData
  try
    polarity =  if data["polarity"] == "positive"
                  1
                elseif data["polarity"] == "negative"
                  -1
                else
                  0
                end

    subjective = data["subjectivity"] == "subjective" ? true : false

    SentimentData(polarity, data["polarity"], subjective)
  catch ex
    Logger.log(string(ex), :err)
    Logger.log(@location_in_file, :err)

    SentimentData(0, "N/A", true)
  end
end

end

max_stars_repo_name: essenciary/genie-demo-tweet-stats

max_issues_repo_name: essenciary/genie-demo-tweet-stats

max_forks_repo_name: essenciary/genie-demo-tweet-stats


text:
function Tuple(arg0::Vector{jdouble})
    return Tuple((Vector{jdouble},), arg0)
end

function abs(obj::Tuple)
    return jcall(obj, "abs", Tuple, ())
end

function acos(obj::Tuple)
    return jcall(obj, "acos", Tuple, ())
end

function acosh(obj::Tuple)
    return jcall(obj, "acosh", Tuple, ())
end

function add(obj::Tuple, arg0::Tuple)
    return jcall(obj, "add", Tuple, (Tuple,), arg0)
end

function add(obj::Tuple, arg0::jdouble)
    return jcall(obj, "add", Tuple, (jdouble,), arg0)
end

function asin(obj::Tuple)
    return jcall(obj, "asin", Tuple, ())
end

function asinh(obj::Tuple)
    return jcall(obj, "asinh", Tuple, ())
end

function atan(obj::Tuple)
    return jcall(obj, "atan", Tuple, ())
end

function atan2(obj::Tuple, arg0::Tuple)
    return jcall(obj, "atan2", Tuple, (Tuple,), arg0)
end

function atanh(obj::Tuple)
    return jcall(obj, "atanh", Tuple, ())
end

function cbrt(obj::Tuple)
    return jcall(obj, "cbrt", Tuple, ())
end

function ceil(obj::Tuple)
    return jcall(obj, "ceil", Tuple, ())
end

function copy_sign(obj::Tuple, arg0::Tuple)
    return jcall(obj, "copySign", Tuple, (Tuple,), arg0)
end

function copy_sign(obj::Tuple, arg0::jdouble)
    return jcall(obj, "copySign", Tuple, (jdouble,), arg0)
end

function cos(obj::Tuple)
    return jcall(obj, "cos", Tuple, ())
end

function cosh(obj::Tuple)
    return jcall(obj, "cosh", Tuple, ())
end

function divide(obj::Tuple, arg0::Tuple)
    return jcall(obj, "divide", Tuple, (Tuple,), arg0)
end

function divide(obj::Tuple, arg0::jdouble)
    return jcall(obj, "divide", Tuple, (jdouble,), arg0)
end

function equals(obj::Tuple, arg0::Object)
    return jcall(obj, "equals", jboolean, (Object,), arg0)
end

function exp(obj::Tuple)
    return jcall(obj, "exp", Tuple, ())
end

function expm1(obj::Tuple)
    return jcall(obj, "expm1", Tuple, ())
end

function floor(obj::Tuple)
    return jcall(obj, "floor", Tuple, ())
end

function get_component(obj::Tuple, arg0::jint)
    return jcall(obj, "getComponent", jdouble, (jint,), arg0)
end

function get_components(obj::Tuple)
    return jcall(obj, "getComponents", Vector{jdouble}, ())
end

function get_dimension(obj::Tuple)
    return jcall(obj, "getDimension", jint, ())
end

function get_field(obj::Tuple)
    return jcall(obj, "getField", Field, ())
end

function get_real(obj::Tuple)
    return jcall(obj, "getReal", jdouble, ())
end

function hash_code(obj::Tuple)
    return jcall(obj, "hashCode", jint, ())
end

function hypot(obj::Tuple, arg0::Tuple)
    return jcall(obj, "hypot", Tuple, (Tuple,), arg0)
end

function linear_combination(obj::Tuple, arg0::Tuple, arg1::Tuple, arg2::Tuple, arg3::Tuple)
    return jcall(obj, "linearCombination", Tuple, (Tuple, Tuple, Tuple, Tuple), arg0, arg1, arg2, arg3)
end

function linear_combination(obj::Tuple, arg0::Tuple, arg1::Tuple, arg2::Tuple, arg3::Tuple, arg4::Tuple, arg5::Tuple)
    return jcall(obj, "linearCombination", Tuple, (Tuple, Tuple, Tuple, Tuple, Tuple, Tuple), arg0, arg1, arg2, arg3, arg4, arg5)
end

function linear_combination(obj::Tuple, arg0::Tuple, arg1::Tuple, arg2::Tuple, arg3::Tuple, arg4::Tuple, arg5::Tuple, arg6::Tuple, arg7::Tuple)
    return jcall(obj, "linearCombination", Tuple, (Tuple, Tuple, Tuple, Tuple, Tuple, Tuple, Tuple, Tuple), arg0, arg1, arg2, arg3, arg4, arg5, arg6, arg7)
end

function linear_combination(obj::Tuple, arg0::Vector{Tuple}, arg1::Vector{Tuple})
    return jcall(obj, "linearCombination", Tuple, (Vector{Tuple}, Vector{Tuple}), arg0, arg1)
end

function linear_combination(obj::Tuple, arg0::Vector{jdouble}, arg1::Vector{Tuple})
    return jcall(obj, "linearCombination", Tuple, (Vector{jdouble}, Vector{Tuple}), arg0, arg1)
end

function linear_combination(obj::Tuple, arg0::jdouble, arg1::Tuple, arg2::jdouble, arg3::Tuple)
    return jcall(obj, "linearCombination", Tuple, (jdouble, Tuple, jdouble, Tuple), arg0, arg1, arg2, arg3)
end

function linear_combination(obj::Tuple, arg0::jdouble, arg1::Tuple, arg2::jdouble, arg3::Tuple, arg4::jdouble, arg5::Tuple)
    return jcall(obj, "linearCombination", Tuple, (jdouble, Tuple, jdouble, Tuple, jdouble, Tuple), arg0, arg1, arg2, arg3, arg4, arg5)
end

function linear_combination(obj::Tuple, arg0::jdouble, arg1::Tuple, arg2::jdouble, arg3::Tuple, arg4::jdouble, arg5::Tuple, arg6::jdouble, arg7::Tuple)
    return jcall(obj, "linearCombination", Tuple, (jdouble, Tuple, jdouble, Tuple, jdouble, Tuple, jdouble, Tuple), arg0, arg1, arg2, arg3, arg4, arg5, arg6, arg7)
end

function log(obj::Tuple)
    return jcall(obj, "log", Tuple, ())
end

function log10(obj::Tuple)
    return jcall(obj, "log10", Tuple, ())
end

function log1p(obj::Tuple)
    return jcall(obj, "log1p", Tuple, ())
end

function multiply(obj::Tuple, arg0::Tuple)
    return jcall(obj, "multiply", Tuple, (Tuple,), arg0)
end

function multiply(obj::Tuple, arg0::jdouble)
    return jcall(obj, "multiply", Tuple, (jdouble,), arg0)
end

function multiply(obj::Tuple, arg0::jint)
    return jcall(obj, "multiply", Tuple, (jint,), arg0)
end

function negate(obj::Tuple)
    return jcall(obj, "negate", Tuple, ())
end

function new_instance(obj::Tuple, arg0::jdouble)
    return jcall(obj, "newInstance", Tuple, (jdouble,), arg0)
end

function pow(obj::Tuple, arg0::Tuple)
    return jcall(obj, "pow", Tuple, (Tuple,), arg0)
end

function pow(obj::Tuple, arg0::jdouble)
    return jcall(obj, "pow", Tuple, (jdouble,), arg0)
end

function pow(obj::Tuple, arg0::jint)
    return jcall(obj, "pow", Tuple, (jint,), arg0)
end

function reciprocal(obj::Tuple)
    return jcall(obj, "reciprocal", Tuple, ())
end

function remainder(obj::Tuple, arg0::Tuple)
    return jcall(obj, "remainder", Tuple, (Tuple,), arg0)
end

function remainder(obj::Tuple, arg0::jdouble)
    return jcall(obj, "remainder", Tuple, (jdouble,), arg0)
end

function rint(obj::Tuple)
    return jcall(obj, "rint", Tuple, ())
end

function root_n(obj::Tuple, arg0::jint)
    return jcall(obj, "rootN", Tuple, (jint,), arg0)
end

function scalb(obj::Tuple, arg0::jint)
    return jcall(obj, "scalb", Tuple, (jint,), arg0)
end

function signum(obj::Tuple)
    return jcall(obj, "signum", Tuple, ())
end

function sin(obj::Tuple)
    return jcall(obj, "sin", Tuple, ())
end

function sin_cos(obj::Tuple)
    return jcall(obj, "sinCos", FieldSinCos, ())
end

function sinh(obj::Tuple)
    return jcall(obj, "sinh", Tuple, ())
end

function sqrt(obj::Tuple)
    return jcall(obj, "sqrt", Tuple, ())
end

function subtract(obj::Tuple, arg0::Tuple)
    return jcall(obj, "subtract", Tuple, (Tuple,), arg0)
end

function subtract(obj::Tuple, arg0::jdouble)
    return jcall(obj, "subtract", Tuple, (jdouble,), arg0)
end

function tan(obj::Tuple)
    return jcall(obj, "tan", Tuple, ())
end

function tanh(obj::Tuple)
    return jcall(obj, "tanh", Tuple, ())
end

function to_degrees(obj::Tuple)
    return jcall(obj, "toDegrees", Tuple, ())
end

function to_radians(obj::Tuple)
    return jcall(obj, "toRadians", Tuple, ())
end


text:
function magnetization1()

    M1 = Magnetization()
    M2 = Magnetization{Int}()
    M3 = Magnetization(0, 0, 0.0)

    @test eltype(M1) == eltype(M3) != eltype(M2)
    return M1 == M2 == M3

end

function magnetization2()

    M = Magnetization()
    show(devnull, M)
    show(devnull, "text/plain", M)

    return true

end

function magnetization3()

    M = Magnetization(1, 2, 3)
    Mcopy = copy(M)
    Mcopyto = Magnetization()
    copyto!(Mcopyto, M)

    @test M ≈ Mcopy ≈ Mcopyto
    @test M == Mcopy == Mcopyto
    return signal(M) == 1 + 2im

end

function magnetization4()

    M = Magnetization()
    MInt = convert(Magnetization{Int}, M)
    MF64 = convert(Magnetization{Float64}, M)

    return M == MInt == MF64

end

function magnetization5()

    src = [0, 0, 0]
    dst = Magnetization(3, 2, 1)
    copyto!(dst, src)

    return Vector(dst) == src

end

function magnetization6()

    src = Magnetization(3, 2, 1)
    dst = [0, 0, 0]
    copyto!(dst, src)

    return dst == Vector(src)

end

function magnetizationmc1()

    M1 = MagnetizationMC(3)
    M2 = MagnetizationMC{Int}(3)
    M3 = MagnetizationMC((0, 0, 0), (0, 0, 0), (0, 0, 0))

    @test eltype(M1) != eltype(M2) == eltype(M3)
    return M1 == M2 == M3

end

function magnetizationmc2()

    M = MagnetizationMC(3)
    show(devnull, M)
    show(devnull, "text/plain", M)

    return true

end

function magnetizationmc3()

    M = MagnetizationMC((1, 2, 3), (3, 2, 1))
    Mcopy = copy(M)
    Mcopyto = MagnetizationMC(2)
    copyto!(Mcopyto, M)

    @test M ≈ Mcopy ≈ Mcopyto
    @test M == Mcopy == Mcopyto
    return signal(M) == 4 + 4im

end

function magnetizationmc4()

    M = MagnetizationMC(3)
    MInt = convert(MagnetizationMC{Int,3}, M)
    MF64 = convert(MagnetizationMC{Float64,3}, M)

    return M == MInt == MF64

end

function magnetizationmc5()

    src = [0, 0, 0, 0, 0, 0]
    dst = MagnetizationMC((3, 2, 1), (4, 5, 6))
    copyto!(dst, src)

    return Vector(dst) == src

end

function magnetizationmc6()

    src = MagnetizationMC((3, 2, 1), (4, 5, 6))
    dst = [0, 0, 0, 0, 0, 0]
    copyto!(dst, src)

    return dst == Vector(src)

end

@testset "Magnetization" begin

    @test magnetization1()
    @test magnetization2()
    @test magnetization3()
    @test magnetization4()
    @test magnetization5()
    @test magnetization6()

end

@testset "MagnetizationMC" begin

    @test magnetizationmc1()
    @test magnetizationmc2()
    @test magnetizationmc3()
    @test magnetizationmc4()
    @test magnetizationmc5()
    @test magnetizationmc6()

end


text:
# This file is a part of BAT.jl, licensed under the MIT License (MIT).

@recipe function f(prior::NamedTupleDist;
                params=collect(1:5),
                diagonal = Dict(),
                upper = Dict(),
                lower = Dict(),
                param_labels = [])

    parsel = params
    parsel = parsel[parsel .<= length(keys(prior))]

    if Base.size(param_labels, 1) == 0
        param_labels = [latexstring("\\theta_$i") for i in parsel]
        param_labels_y = [latexstring("p(\\theta_$i)") for i in parsel]
    else
        param_labels_y = [latexstring("p("*param_labels[i]*")") for i in 1:length(param_labels)]
        param_labels = [latexstring(param_labels[i]) for i in 1:length(param_labels)]
    end

    nparams = length(parsel)
    layout --> nparams^2
    size --> (1000, 600)

    for i in 1:nparams

        # diagonal
        @series begin
            subplot := i + (i-1)*nparams

            seriestype --> get(diagonal, "seriestype", :stephist)
            bins --> get(diagonal, "bins", 200)
            colors --> get(diagonal, "colors", standard_colors)
            intervals --> get(diagonal, "intervals", standard_confidence_vals)
            legend --> get(diagonal, "legend", false)
            xlims --> get(diagonal, "xlims", :auto)
            ylims --> get(diagonal, "ylims", :auto)
            linecolor --> get(diagonal, "linecolor", :black)
            xguide --> param_labels[i]
            yguide --> param_labels_y[i]

            prior, (parsel[i])
        end

        # upper right plots
        for j in i+1:nparams

            @series begin
                subplot := j + (i-1)*nparams

                seriestype --> get(upper, "seriestype", :smallest_intervals_contour)
                bins --> get(upper, "bins", 50)
                more_colors = []
                colors --> get(upper, "colors", colormap("Blues", 10))
                more_intervals = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 0.99]
                intervals --> get(upper, "intervals", more_intervals)
                legend --> get(upper, "legend", false)
                xlims --> get(upper, "xlims", :auto)
                ylims --> get(upper, "ylims", :auto)
                xguide --> param_labels[i]
                yguide --> param_labels[j]

                prior, (parsel[i], parsel[j])
            end

            # lower left plots
            @series begin
                subplot := i + (j-1)*nparams

                seriestype --> get(lower, "seriestype", :smallest_intervals_contour)
                bins --> get(lower, "bins", 50)
                colors --> get(lower, "colors", standard_colors)
                intervals --> get(lower, "intervals", standard_confidence_vals)
                legend --> get(lower, "legend", false)
                xlims --> get(lower, "xlims", :auto)
                ylims --> get(lower, "ylims", :auto)
                xguide --> param_labels[i]
                yguide --> param_labels[j]

                prior, (parsel[i], parsel[j])
            end

        end
    end

end


text:
mutable struct DynamicRegularization{T}
    ρ::T
    dρ::T
end

"""
    iLQRSolver

A fast solver for unconstrained trajectory optimization that uses a Riccati recursion
to solve for a local feedback controller around the current trajectory, and then
simulates the system forward using the derived feedback control law.

# Constructor
    Altro.iLQRSolver(prob, opts; kwarg_opts...)

With the following keyword arguments:
- `use_static` Pass `Val(true)` to use StaticArrays as the internal storage for the trajectories.
This will, by default, then use `RobotDynamics.StaticReturn()` methods for the dynamics.

All other keyword arguments are interpreted as solver options passed to [`set_options!`](@ref).
"""
struct iLQRSolver{L,O,Nx,Ne,Nu,T,V} <: UnconstrainedSolver{T}
    model::L
    obj::O

    x0::Vector{T}
    tf::T
    N::Int

    opts::SolverOptions{T}
    stats::SolverStats{T}

    Z::SampledTrajectory{Nx,Nu,T,KnotPoint{Nx,Nu,V,T}}
    Z̄::SampledTrajectory{Nx,Nu,T,KnotPoint{Nx,Nu,V,T}}
    dx::Vector{Vector{T}}
    du::Vector{Vector{T}}

    gains::Vector{Matrix{T}}  # N-1 × (Nu,Ne+1)
    K::Vector{SubArray{T,2,Matrix{T},Tuple{ColonSlice, UnitRange{Int}}, true}}  # N-1 × (Nu,Ne)
    d::Vector{SubArray{T,1,Matrix{T},Tuple{ColonSlice, Int}, true}}             # N-1 × (Nu,)

    D::Vector{DynamicsExpansion{T}}
    G::Vector{Matrix{T}}  # N × (Nx,Ne)

    Efull::CostExpansion{T}  # Cost expansion (full state)
    Eerr::CostExpansion{T}   # Cost expansion (error state)
    Q::Vector{StateControlExpansion{T}}      # Action-value expansion
    S::Vector{StateControlExpansion{T}}      # Action-value expansion
    ΔV::Vector{T}

    Qtmp::StateControlExpansion{T}
    Quu_reg::Matrix{T}
    Qux_reg::Matrix{T}
    reg::DynamicRegularization{T}

    grad::Vector{T}
    xdot::Vector{T}

    logger::SolverLogging.Logger
end

function iLQRSolver(
        prob::Problem{T},
        opts::SolverOptions=SolverOptions{T}(),
        stats::SolverStats=SolverStats{T}(parent=solvername(iLQRSolver));
        use_static::Val{USE_STATIC}=usestaticdefault(get_model(prob)[1]),
        dynamics_diffmethod::RD.DiffMethod=RD.default_diffmethod(TO.get_model(prob,1)),
        kwarg_opts...
    ) where {T, USE_STATIC}
    set_options!(opts; dynamics_diffmethod=dynamics_diffmethod, kwarg_opts...)

    nx,nu = dims(prob)
    N = TO.horizonlength(prob)
    ne = RD.errstate_dim.(prob.model)
    push!(ne,ne[end])

    x0 = copy(prob.x0)
    samestatedim = all(x->x == nx[1], nx)
    samecontroldim = all(u->u == nu[1], nu)
    if USE_STATIC
        @assert  samestatedim && samecontroldim "Cannot use StaticArrays with varying state or control dimension."
        Nx = nx[1]
        Nu = nu[1]
        Ne = ne[1]
    else
        Nx = samestatedim ? nx[1] : Any
        Ne = samestatedim ? ne[1] : Any
        Nu = samecontroldim ? nu[1] : Any
    end
    V = USE_STATIC ? SVector{nx[1]+nu[1],T} : Vector{T}
    Z = SampledTrajectory(map(enumerate(prob.Z)) do (k,z)
        RD.KnotPoint{Nx,Nu}(nx[k], nu[k], V(RD.getdata(z)), RD.getparams(z)...)
    end)
    Z̄ = copy(Z)

    # Rollout out dynamics if initial state contains NaN values
    if any(isnan, state(Z[1]))
        RD.rollout!(dynamics_signature(Z), prob.model[1], Z, prob.x0)
    end
    RD.setstate!(Z[1], prob.x0)  # set initial state

    dx = [zeros(T,ne[k]) for k = 1:N]
    du = [zeros(T,nu[k]) for k = 1:N-1]

    gains = [zeros(T,nu[k],ne[k]+1) for k = 1:N-1]
    K = [view(gain,:,1:ne[k]) for (k,gain) in enumerate(gains)]
    d = [view(gain,:,ne[k]+1) for (k,gain) in enumerate(gains)]

    D = [DynamicsExpansion{T}(nx[k],ne[k],nu[k]) for k = 1:N-1]
    G = [Matrix(one(T)*I,nx[k],ne[k]) for k = 1:N]
    push!(G, copy(G[end]))

    Eerr = CostExpansion{T}(ne,nu)
    Efull = FullStateExpansion(Eerr, prob.model[1])
    Q = [StateControlExpansion{T}(ne[k],nu[k]) for k = 1:N]
    S = [StateControlExpansion{T}(ne[k]) for k = 1:N]
    ΔV = zeros(T,2)

    # TODO: Change this to accept multiple sizes
    Qtmp = StateControlExpansion{T}(ne[1],nu[1])
    Quu_reg = zeros(T,nu[1],nu[1])
    Qux_reg = zeros(T,nu[1],ne[1])
    reg = DynamicRegularization{T}(opts.bp_reg_initial, 0)

    grad = zeros(T,N-1)
    xdot = zeros(T,nx[1])  # TODO: change this to support multiple sizes

    # logger = SolverLogging_v1.default_logger(opts.verbose >= 2)
    lg = SolverLogging.Logger()
    setentry(lg, "iter", Int, width=6)
    setentry(lg, "AL iter", Int, width=8, level=2)
    setentry(lg, "iLQR iter", Int, width=10, level=2)
    setentry(lg, "cost", fmt="%.3f")
    setentry(lg, "||v||")
    setentry(lg, "expected", level=2)
    setentry(lg, "dJ", level=2)
    setentry(lg, "grad", level=2)
    setentry(lg, "z", level=3, fmt="%.2f")
    setentry(lg, "α", level=3)
    setentry(lg, "ρ", level=3)
    setentry(lg, "μ_max", level=4)
    setentry(lg, "dJ_zero", Int, level=4)
    setentry(lg, "ls_iter", Int, width=8, level=5)
    setentry(lg, "info", String, width=40)

	L = typeof(prob.model)
	O = typeof(prob.obj)
    solver = iLQRSolver{L,O,Nx,Ne,Nu,T,V}(
        prob.model, prob.obj, x0, TO.get_final_time(prob), N, opts, stats, Z, Z̄, dx, du,
        gains, K, d, D, G, Efull, Eerr, Q, S, ΔV, Qtmp, Quu_reg, Qux_reg, reg, grad, xdot,
        lg,
    )
    reset!(solver)
end

# Getters
RD.dims(solver::iLQRSolver) = RD.dims(solver.Z)
RD.state_dim(::iLQRSolver{<:Any,<:Any,n}) where n = n
RD.errstate_dim(::iLQRSolver{<:Any,<:Any,<:Any,e}) where e = e
RD.control_dim(::iLQRSolver{<:Any,<:Any,<:Any,<:Any,m}) where m = m
RD.state_dim(solver::iLQRSolver, k::Integer) = RD.state_dim(solver.model[k])
RD.errstate_dim(solver::iLQRSolver, k::Integer) = RD.errstate_dim(solver.model[k])
RD.control_dim(solver::iLQRSolver, k::Integer) = RD.control_dim(solver.model[k])
@inline TO.get_trajectory(solver::iLQRSolver) = solver.Z
@inline TO.get_objective(solver::iLQRSolver) = solver.obj
@inline TO.get_model(solver::iLQRSolver) = solver.model
@inline get_initial_state(solver::iLQRSolver) = solver.x0
solvername(::Type{<:iLQRSolver}) = :iLQR
getlogger(solver::iLQRSolver) = solver.logger

get_feedbackgains(solver::iLQRSolver) = solver.K

RD.vectype(::iLQRSolver{<:Any,<:Any,<:Any,<:Any,<:Any,<:Any,V}) where V = V
usestatic(obj) = RD.vectype(obj) <: SVector
dynamics_signature(obj) = usestatic(obj) ? RD.StaticReturn() : RD.InPlace()
function_signature(obj) = usestatic(obj) ? RD.StaticReturn() : RD.InPlace()
usestaticdefault(model::RD.AbstractFunction) = Val(RD.default_signature(model) == RD.StaticReturn())

log_level(::iLQRSolver) = InnerLoop

function reset!(solver::iLQRSolver)
    opts = options(solver)::SolverOptions
    reset!(stats(solver), opts.iterations, solvername(solver))
    solver.reg.ρ = solver.opts.bp_reg_initial
    solver.reg.dρ = 0.0
    return solver
end

function dynamics_expansion!(solver::iLQRSolver, Z=solver.Z)
    diff = solver.opts.dynamics_diffmethod
    D = solver.D
    model = solver.model

    # Dynamics Jacobians
    for k in eachindex(D)
        RobotDynamics.jacobian!(dynamics_signature(solver), diff, model[k], D[k], Z[k])
    end

    # Calculate the expansion on the error state
    error_expansion!(solver.model, D, solver.G)
end

function cost_expansion!(solver::iLQRSolver)
    # Calculate normal cost expansion
    Altro.cost_expansion!(solver.obj, solver.Efull, solver.Z)

    # Calculate the error state expansion
    Altro.error_expansion!(solver.model, solver.Err, solver.Efull, solver.G, solver.G)
end

function increaseregularization!(solver::iLQRSolver)
    reg = solver.reg
    ρdot = solver.opts.bp_reg_increase_factor
    ρmin = solver.opts.bp_reg_min
    reg.dρ =  max(reg.dρ*ρdot, ρdot)
    reg.ρ = max(reg.ρ*reg.dρ, ρmin)
    return reg
end

function decreaseregularization!(solver::iLQRSolver)
    reg = solver.reg
    ρdot = solver.opts.bp_reg_increase_factor
    ρmin = solver.opts.bp_reg_min
    reg.dρ = min(reg.dρ/ρdot, 1/ρdot)
    reg.ρ = max(ρmin, reg.ρ*reg.dρ)
    return reg
end

"""
    reset_gains!(solver::iLQRSolver)

Resets the gains to zero. Useful if you want to call `solve!` multiple times,
without using information from the previous solve, since the feedback gains
are used to perform the initial rollout. See [`initialize!(::iLQRSolver)`](@ref).
"""
function reset_gains!(solver::iLQRSolver)
    for k in eachindex(solver.K)
        solver.K[k] .= 0
        solver.d[k] .= 0
    end
    return nothing
end

function state_diff_jacobian!(model, G, Z)
    for k = 1:length(Z)
        RD.errstate_jacobian!(model, G[k], Z[k])
    end
end