API Reference

LinearDynamicalSystem{T<:Real, S<:AbstractStateModel{T}, O<:AbstractObservationModel{T}}

Represents a unified Linear Dynamical System with customizable state and observation models.

Fields

• state_model::S: The state model (e.g., GaussianStateModel)
• obs_model::O: The observation model (e.g., GaussianObservationModel or PoissonObservationModel)
• latent_dim::Int: Dimension of the latent state
• obs_dim::Int: Dimension of the observations
• fit_bool::Vector{Bool}: Vector indicating which parameters to fit during optimization. Length 6 for the Gaussian path ([x0, P0, A&b&B, Q, C&d&D, R]); the M-step regression fits each row jointly. Length 5 for the Poisson path ([x0, P0, A&b, Q, C&d]).

SLDS{T,S,O,TM,ISV}

A Switching Linear Dynamical System (SLDS). A hierarchical time-series model of the form:

\[z_t | z_{t-1} ~ Categorical(A_{z_{t-1}, :}) x_t | x_{t-1}, z_t ~ N(A^{(z_t)} x_{t-1} + b^{(z_t)}, Q^{(z_t)}) y_t | x_t, z_t ~ N(C^{(z_t)} x_t + d^{(z_t)}, R^{(z_t)})\]

Fields

• A::TM: Transition matrix for the discrete states (K x K)
• πₖ::ISV: Initial state distribution for the discrete states (K-dimensional vector)
• LDSs::Vector{LinearDynamicalSystem{T,S,O}}: Vector of K Linear Dynamical Systems, one for each discrete state

GaussianStateModel{T<:Real, M<:AbstractMatrix{T}, V<:AbstractVector{T}}

Represents the state model of a Linear Dynamical System with Gaussian noise.

State evolution:

\[x_1 ~ N(x_0, P_0) x_{t+1} | x_t ~ N(A x_t + b + B ux_t, Q)\]

where B·ux_t is present only when B is supplied (i.e., has nonzero columns).

Fields

• A::M: Transition matrix (size latent_dim × latent_dim).
• Q::M: Process noise covariance matrix.
• b::V: Bias vector (length latent_dim).
• x0::V: Initial state mean (length latent_dim).
• P0::M: Initial state covariance (size latent_dim × latent_dim).
• B::M: Optional dynamics input matrix (latent_dim × ux_dim). When supplied, inputs ux must be passed to fit!/smooth! via a keyword argument.
• Q_prior::Union{Nothing,IWPrior{T}} = nothing: Optional Inverse-Wishart prior on Q. If set, MAP updates use its mode.
• P0_prior::Union{Nothing,IWPrior{T}} = nothing: Optional Inverse-Wishart prior on P0. If set, MAP updates use its mode.
• AB_prior::Union{Nothing,MNPrior{T,Matrix{T}}} = nothing: Optional matrix-normal prior on the stacked dynamics matrix [A B]. Pair with Q_prior for a full MNIW prior on (AB, Q). Prior matrices are stored as plain Matrix{T} (decoupled from A's storage type M) so they match the internal KalmanWorkspace regardless of how A is stored.

GaussianObservationModel{T<:Real, M<:AbstractMatrix{T}, V<:AbstractVector{T}}

Represents the observation model of a Linear Dynamical System with Gaussian noise.

Fields

• C::M: Observation matrix of size (obs_dim × latent_dim). Maps latent states into observation space.
• R::M: Observation noise covariance of size (obs_dim × obs_dim).
• d::V: Bias vector of length (obs_dim).
• R_prior::Union{Nothing, IWPrior{T}} = nothing: Optional Inverse-Wishart prior for R.
• CD_prior::Union{Nothing,MNPrior{T,Matrix{T}}} = nothing: Optional matrix-normal prior on the stacked emission matrix [C D]. Pair with R_prior for a full MNIW prior on (CD, R). Prior matrices are stored as plain Matrix{T} (decoupled from C's storage type M) so they match the internal KalmanWorkspace regardless of how C is stored.

PoissonObservationModel{
    T<:Real,
    M<:AbstractMatrix{T},
    V<:AbstractVector{T}
} <: AbstractObservationModel{T}

Represents the observation model of a Linear Dynamical System with Poisson observations, with canonical log-link:

\[λ_t = exp(C x_t + d)\]

d is the standard Poisson-GLM intercept — the per-channel baseline log-rate, unconstrained in ℝ; positivity of the rate λ is provided by the exp.

Fields

• C::AbstractMatrix{T}: Observation matrix of size (obs_dim × latent_dim). Maps latent states into observation space.
• d::AbstractVector{T}: Per-neuron baseline log-rate (length obs_dim). Free in ℝ.
• CD_prior::Union{Nothing,MNPrior{T,Matrix{T}}} = nothing: Optional matrix-normal prior on the stacked emission matrix [C d] (treated as a single regression of log λ on [x; 1]). Prior matrices are stored as plain Matrix{T}, decoupled from C's storage type M. M₀ and Λ have shapes (obs_dim, latent_dim+1) and (latent_dim+1, latent_dim+1) respectively. Unlike the Gaussian path there is no IW counterpart since Poisson has no observation-noise covariance — this is an MN-only prior contributing ½ tr(([C d] - M₀) Λ ([C d] - M₀)') to the LBFGS objective.

ProbabilisticPCA{T<:Real, M<:AbstractMatrix{T}, V<:AbstractVector{T}}

Probabilistic Principal Component Analysis (PPCA) model. Observations x ∈ ℝ^D are generated from a k-dimensional latent factor z with isotropic Gaussian noise:

\[z ∼ N(0, I_k), \qquad x \mid z ∼ N(μ + W z, σ² I_D)\]

so that marginally x ∼ N(μ, W Wᵀ + σ² I). As σ² → 0, PPCA recovers classical PCA. Construct with ProbabilisticPCA(W, σ², μ); k and D are inferred from the size of W.

Fields

• W::M: Loading matrix (D × k). Identifiable only up to an orthogonal rotation of its columns.
• σ²::T: Isotropic observation noise variance.
• μ::V: Observation mean (length D).
• k::Int: Latent dimension.
• D::Int: Observation dimension.
• z::M: Posterior means of the latent factors from the most recent E-step (k × N; empty until fit! is called).

Data{T<:Real}

Container for the observed data passed to the Kalman path: observations y, dynamics (latent) inputs ux, and observation inputs uy, each (dim, tsteps, ntrials) (input fields may have zero rows when no controls are supplied).

AbstractStateModel{T<:Real}

Abstract supertype for latent-state models of a LinearDynamicalSystem (e.g. GaussianStateModel). A state model defines how the latent state evolves from one timestep to the next.

AbstractObservationModel{T<:Real}

Abstract supertype for observation (emission) models of a LinearDynamicalSystem (e.g. GaussianObservationModel or PoissonObservationModel). An observation model defines how observed data are generated from the latent state.

IWPrior{T<:Real, M<:AbstractMatrix}

Inverse-Wishart prior for a covariance matrix Σ ~ IW(Ψ, ν), with density p(Σ) ∝ |Σ|^{-(ν + d + 1)/2} exp(-½ tr(Ψ Σ^{-1})) for d = size(Σ,1).

Fields

• Ψ::M: Scale matrix (d×d, SPD).
• ν::T: Degrees of freedom (must satisfy ν > d + 1 for a proper mode).

Notes

• The MAP update for a posterior IW(Ψ + S, ν + n) is (Ψ + S) / (ν + n + d + 1).

MNPrior{T<:Real, M<:AbstractMatrix}

Matrix-normal prior on a regression coefficient matrix W (size k × p) appearing in a linear regression Y = W X + ε with row-noise covariance Σ:

\[W | Σ ~ MN(M₀, Σ ⊗ Λ⁻¹)\]

equivalently vec(W) ~ N(vec(M₀), Λ⁻¹ ⊗ Σ). Σ is the regression's row covariance (paired with IWPrior when both halves of an MNIW prior are desired); Λ is the column precision and is the only piece that enters the MAP update for W.

Fields

• M₀::M: prior mean (k × p, same shape as W). Use a zero matrix for plain ridge; an identity-like matrix for shrinkage toward a random walk on A.
• Λ::M: column precision (p × p, SPD).

Notes

• The MAP update for W given the regression sufficient statistics XX = X Xᵀ and XY = X Yᵀ is math W = (XYᵀ + M₀ Λ) (XX + Λ)⁻¹ Reduces to OLS when Λ = 0, and to ordinary ridge regression when M₀ = 0.
• Mathematically half of an MNIW prior; combine with an IWPrior on the same regression to recover the full conjugate prior on (W, Σ).

Random.rand([rng,] lds, tsteps::Integer; latent_inputs=nothing, obs_inputs=nothing)
Random.rand([rng,] lds, tsteps_per_trial::AbstractVector{<:Integer};
            latent_inputs=nothing, obs_inputs=nothing)

Sample from a Linear Dynamical System.

• With a scalar tsteps, returns one trial as (x::Matrix, y::Matrix) of sizes (latent_dim, tsteps) and (obs_dim, tsteps) respectively.
• With a vector of per-trial lengths, returns (x::Vector{Matrix}, y::Vector{Matrix}). Lengths may differ across trials.

Optional input sequences:

• latent_inputs: dynamics-input sequence consumed by B. Single-trial form is an (ux_dim, tsteps) matrix; multi-trial is a Vector{<:AbstractMatrix} of per-trial matrices. Required when size(state_model.B, 2) > 0.
• obs_inputs: same shape for the observation input D. Required when size(obs_model.D, 2) > 0. Gaussian observation model only.

rand([rng,] slds, tsteps::Integer)
rand([rng,] slds, tsteps_per_trial::AbstractVector{<:Integer})

Sample from a Switching Linear Dynamical System.

• Scalar tsteps: returns one trial as (z::Vector{Int}, x::Matrix, y::Matrix).
• Vector of per-trial lengths: returns (z::Vector{Vector{Int}}, x::Vector{Matrix}, y::Vector{Matrix}). Trial lengths may differ.

Random.rand([rng::AbstractRNG,] ppca::ProbabilisticPCA, n::Int)

Sample n observations from the PPCA generative model. Returns a tuple (X, z) where X is D × n and z holds the sampled latent factors (k × n).

smooth(lds, y::AbstractMatrix)

Direct smoothing for a single trial.

Arguments

• lds::LinearDynamicalSystem: The model.
• y::AbstractMatrix: Observations (obs_dim × tsteps).

Returns

• x_smooth::AbstractMatrix: Smoothed latent means (latent_dim × tsteps).
• p_smooth::Array{T,3}: Smoothed latent covariances (latentdim × latentdim × tsteps).

fit!(lds, y; max_iter=100, tol=1e-6, progress=true,
     latent_inputs=nothing, obs_inputs=nothing)

Fit a Gaussian Linear Dynamical System via Expectation-Maximization.

Arguments

• lds::LinearDynamicalSystem{T,S,O}: model to fit in place
• y: observations. Two shapes accepted:
  • AbstractMatrix{T} of size (obs_dim, T) — single trial
  • AbstractVector{<:AbstractMatrix{T}} — multi-trial, each (obs_dim, T_i), trial lengths may differ

Keywords

• max_iter::Int=100: maximum EM iterations
• tol::Float64=1e-6: convergence tolerance on ELBO change
• progress::Bool=true: show progress bar
• latent_inputs: optional dynamics-input sequence. Vector{<:AbstractMatrix} for multi-trial (each (ux_dim, T_i)); required when size(state_model.B, 2) > 0.
• obs_inputs: optional observation-input sequence (same shape) for the obs-side input matrix D. Required when size(obs_model.D, 2) > 0.

Returns a Vector{T} of ELBO values, one per iteration.

fit!(slds::SLDS, y::AbstractVector{<:AbstractMatrix}; max_iter=50, progress=true)
fit!(slds::SLDS, y::AbstractMatrix; max_iter=50, progress=true)

Fit SLDS using variational Laplace EM with stochastic ELBO estimates. Runs for exactly max_iter iterations.

y is either a single trial (obs_dim × T) matrix or a vector of per-trial matrices (ragged T_i allowed). Internally a single batched HMMs.ForwardBackwardStorage of length sum(T_i) is allocated, with seq_ends = cumsum(T_i) to demarcate trials.

fit!(plds, y; max_iter=100, tol=1e-6, progress=true, newton_max_iter=20, newton_tol=1e-6)

Fit a Poisson LDS via Laplace-EM.

Arguments

• plds: the Poisson LDS model (modified in place)
• y: observations. Two shapes accepted:
  • AbstractMatrix{T} of size (obs_dim, T) — single trial
  • AbstractVector{<:AbstractMatrix{T}} — multi-trial, each (obs_dim, T_i), ragged trial lengths allowed

Keywords

• max_iter: maximum EM iterations
• tol: convergence tolerance on ELBO change
• progress: show progress bar
• newton_max_iter: Newton iterations per E-step inner solve
• newton_tol: Newton convergence tolerance

fit!(ppca::ProbabilisticPCA, X::AbstractMatrix, max_iters::Int=100, tol::Float64=1e-6)

Fit the PPCA model to data X (D × N, one observation per column) with EM. μ is set directly to the sample mean (its exact ML estimate); W and σ² are then updated iteratively until the log-likelihood improves by less than tol or max_iters is reached.

Returns the vector of marginal log-likelihood values, one per iteration (non-decreasing).

loglikelihood(lds, y)

Marginal (observed-data) log-likelihood ∑_{t,n} log p(y_t^n | y_{1:t-1}^n) of a Gaussian LinearDynamicalSystem, computed by running the Kalman filter and summing the one-step-ahead predictive densities (latent states integrated out). Valid for any fitted model with a GaussianObservationModel.

This is the StatsAPI.loglikelihood method for the LDS; for the complete-data log-likelihood log p(x, y) given a trajectory x, see joint_loglikelihood.

Returns the total log-likelihood. Divide by obs_dim * tsteps * ntrials for a per-observation score that is comparable across configurations.

loglikelihood(plds, y)

Marginal (observed-data) log-likelihood for a Poisson LDS — not implemented.

The marginal log p(y) = ∫ p(x, y) dx is intractable for the Poisson observation model (non-conjugate; there is no closed-form Kalman filter as in the Gaussian case). Use joint_loglikelihood(x, plds, y) for the complete-data log-likelihood given a trajectory x, or the ELBO returned by fit! as a lower bound on log p(y).

loglikelihood(ppca::ProbabilisticPCA, X::AbstractMatrix)

Marginal log-likelihood ∑_n log p(x_n) of the data X (D × N, one observation per column) under the PPCA model, using the closed-form marginal x ∼ N(μ, W Wᵀ + σ² I).

elbo!(lds, suf, sws, total_entropy)

Total ELBO from aggregated sufficient statistics. Computes the same quantity as the legacy elbo(lds, tfs, y, sws_pool, ...) but in O(D³ + p²·D) instead of O(N·T·p²·D). The Gaussian-posterior entropy comes from each trial's fs.entropy (filled by the smoother) and is summed by the caller before this function is invoked.

elbo!(plds, suf, tfs, y, sws_pool)

Suf-based Poisson ELBO. Mirrors the Gaussian TD path's split:

• state-side Q-term via Q_state!(sws, plds, suf) from the aggregated sufficient statistics (O(D³) per E-step, not O(N·T·D²)),
• observation-side Q-term per-trial via the existing Poisson Q_obs!, which is irreducibly non-conjugate (no aggregator equivalent),
• posterior entropy from tfs[trial].entropy (filled by smooth!),
• IWPrior log-prior contributions on Q and P0, and the MN log-prior trace term on [C d] to match the LBFGS objective.

elbo!(slds, tfs, fb_storage, y, slds_ws; seq_ends)

Compute the stochastic ELBO for SLDS at the current variational posteriors.

• Computes E_q[log p(y, x | z)] weighted by discrete posteriors
• Computes log p(z1) and log p(zt | z_{t-1}) weighted by discrete posteriors
• Subtracts entropies: -H[q(x)] - H[q(z)]

Returns a scalar ELBO value.

validate_LDS(lds::LinearDynamicalSystem{T,S,O}) where {T,S,O}

Validate that all parameters in a LinearDynamicalSystem are dimensionally consistent and mathematically valid. Throws descriptive exceptions on validation failure.

Checks performed

• Matrix dimensions are consistent
• Covariance matrices are positive definite and symmetric
• fit_bool has correct length for the observation model type
• Stored dimensions match dimensions inferred from matrices

Throws

• DimensionMismatchError: If dimensions don't match
• NotPositiveDefiniteError: If covariance matrices aren't positive definite
• NotSymmetricError: If covariance matrices aren't symmetric
• NumericalStabilityError: If numerical issues are detected

Examples

# This will throw DimensionMismatchError if invalid
validate_LDS(my_lds)

# Can be caught for custom handling
try
    validate_LDS(my_lds)
    println("LDS is valid!")
catch e
    if e isa DimensionMismatchError
        println("Dimension error: ", e)
    end
end

validate_SLDS(slds::SLDS)

Validate SLDS structure. Throws descriptive exceptions on validation failure.

Checks performed

• Dimensions of A match the length of πₖ and the number of LDSs
• Rows of A and πₖ are valid probability vectors
• Each LDS has the same state dimension and observation dimension
• Each individual LDS is valid

Throws

• DimensionMismatchError: If dimensions are inconsistent
• InvalidProbabilityVectorError: If probability vectors are invalid
• Other exceptions from validate_LDS for individual LDS validation

Examples

# This will throw if invalid
validate_SLDS(my_slds)

# Can be caught for custom handling
try
    validate_SLDS(my_slds)
catch e
    if e isa InvalidProbabilityVectorError
        println("Probability vector error: ", e)
    end
end

validate_probvec(v::AbstractVector{T}; name::String="vector") where {T<:Real}

Validate that a vector is a valid probability vector (sums to 1, all non-negative, all ≤ 1). Throws InvalidProbabilityVectorError if validation fails.

Arguments

• v: The vector to validate
• name: Optional name for the vector (used in error messages)

Examples

v1 = [0.3, 0.5, 0.2]
validate_probvec(v1)  # No error

v2 = [0.3, 0.5, 0.3]
validate_probvec(v2)  # Throws InvalidProbabilityVectorError

DimensionMismatchError <: Exception

Custom exception for dimension mismatches in model parameters.

Fields

• parameter::String: Name of the parameter with incorrect dimensions
• expected::Union{Int,Tuple{Vararg{Int}}}: Expected dimension(s)
• got::Union{Int,Tuple{Vararg{Int}}}: Actual dimension(s)

NotPositiveDefiniteError <: Exception

Custom exception for matrices that should be positive definite but aren't.

Fields

• matrix_name::String: Name of the matrix
• min_eigenvalue::Float64: Minimum eigenvalue found

NotSymmetricError <: Exception

Custom exception for matrices that should be symmetric but aren't.

Fields

• matrix_name::String: Name of the matrix
• max_asymmetry::Float64: Maximum asymmetry measure

InvalidProbabilityVectorError <: Exception

Custom exception for invalid probability vectors.

Fields

• vector_name::String: Name of the probability vector
• sum_value::Float64: Sum of the vector
• has_negative::Bool: Whether the vector contains negative values
• has_greater_than_one::Bool: Whether the vector contains values > 1.0

NumericalStabilityError <: Exception

Custom exception for numerical stability issues.

Fields

• parameter::String: Name of the parameter
• issue::String: Description of the numerical issue

random_rotation_matrix(n)

Generate a random rotation matrix of size n x n.

gaussian_entropy(H::Symmetric{T}) where {T<:Real}

Entropy (in nats) of a Gaussian whose log-posterior Hessian at the MAP is H (i.e., H = ∇² log p so the precision is Λ = -H).

valid_Σ(Σ) -> Bool

Check whether Σ is a valid covariance matrix, i.e. Hermitian and positive definite.

block_tridgm(
    main_diag::Vector{Matrix{T}},
    upper_diag::Vector{Matrix{T}},
    lower_diag::Vector{Matrix{T}}
) where {T<:Real}

Construct a block tridiagonal matrix from three vectors of matrices.

Throws

• ErrorException if the lengths of upper_diag and lower_diag are not one less than the length of main_diag.

print_full([io::Union{IO, Base.TTY}, ] obj)

Prints full description of object obj, overriding both io-based limits as well as the limits set in the default pretty printing of StateSpaceDynamics objects.

info_update!(cache, P0, CiRC) -> PDMat

Return P = inv(inv(P0) + CiRC) as a PDMat, using the Cholesky cached inside P0 and the scratch buffers in cache. Both P0 and CiRC must be PDMat{T,Matrix{T}} of dimension n, matching cache.

The returned PDMat shares storage with cache.Pmat and cache.Pchol_factors. It is invalidated by the next call with the same cache — extract or accumulate what you need (e.g. add P.mat + μ*μ' into your running E[xxᵀ] sum) before the next invocation.

Cost: one potri + one cholesky! + one potri + one cholesky! on n × n matrices, ≈ 4n³/3 flops plus O(n²) for the add. Compare with the naive inv(inv(P0) + CiRC) which goes ≈ 8n³/3 via PDMat → Matrix → LU-based inv → PDMat, losing PD structure in the middle.

info_update!(P_dest, scratch_M, P0, CiRC) -> PDMat

In-place variant of info_update! that writes the result of inv(inv(P0) + CiRC) into the existing PDMat P_dest (overwriting both its mat and chol.factors fields). This is the form required when the result must persist across many calls — e.g. inside a loop where every step's output is read again later (Kalman forward pass → backward pass).

scratch_M is a single n × n workspace (re-used across calls). All inputs and the destination must be n × n.

The returned PDMat is P_dest itself — for caller convenience.

CovUpdateCache{T}(n)

Preallocated workspace for info_update!. Holds three n × n matrices: scratch for the intermediate sum and its Cholesky, plus storage for the output covariance matrix and its Cholesky factor.

StatsAPI.fit!(dl::SLDSDiscreteLayer, fb_storage, obs_inputs; seq_ends)

Update the discrete transition matrix dl.A and initial-state distribution dl.πₖ in place from forward-backward statistics. Mirrors HiddenMarkovModels.jl's fit!(::HMM, ...) pattern using the ξ[t2] scratch trick: for each sequence, ξ[t2] is zero by FB convention so it doubles as an accumulator for sum(ξ[t1:t2-1]).

Skips fitting observation distributions because the SLDS discrete layer doesn't have parametric obs distributions; per-state log-likelihoods are filled into dl.logL upstream by the SLDS E-step.