geossm.ssm.StateSpaceModel

class geossm.ssm.StateSpaceModel(H=None, R=None, F=None, Q=None, x0=None, Sigma0=None, Xbeta=None, beta=None, xbeta_names=None, backend='auto', dtype=<class 'jax.numpy.float32'>)

Bases: object

A linear-Gaussian state-space model with Kalman filtering, smoothing and simulation capabilities.

Implements

\[\begin{split}y_t &= H x_t + X_t \beta + e_t, \qquad e_t \sim N(0, R) \\ x_t &= F x_{t-1} + \eta_t, \qquad \eta_t \sim N(0, Q) \\ x_0 &\sim N(x_0, \Sigma_0)\end{split}\]

where y_t is the \(p\)-dimensional observation at time t, x_t is the \(q\)-dimensional latent state, and X_t \beta is a regression/fixed-effect term with \(b\) coefficients. The class stores the system matrices and dimensions, and exposes filter, smoother, estimate (filter + smoother), sim (simulation) and predict on top of the JIT-compiled JAX kernels defined at module level (_filter_kernelJAX, _smoother_kernelJAX, _sim_kernelJAX, …).

This class is meant to be reusable/subclassable: for example LRStateSpaceModel (in geossm.stmodel.lrssm) subclasses it to specialize the model for large-scale spatial data via a low-rank state covariance, while reusing the filter/smoother machinery defined here.

Notes

F and R are restricted to diagonal matrices in this implementation (stored internally as their 1D diagonal - see _prepare_diag_array); H and Q are general dense matrices.

Methods

computeExpectedValues	Compute predicted observations and EM sufficient statistics from a set of state means/covariances (filtered or smoothed).
estimate	Run the Kalman filter and smoother (filter + backward pass) and return the smoothed results, including the sufficient statistics (S11, S10, S00) needed by an EM M-step.
filter	Run the Kalman filter forward pass and return the filtered results (state estimates, log-likelihood, and one-step-ahead predicted observations/sufficient statistics computed from the filtered - not smoothed - states).
generate_summary	Build the two small key/value tables (model identity/config on the left, shape and mean system-matrix diagonals on the right) used by summary.
predict	Compute predicted observations and their state-driven predictive covariance from (typically smoothed) states.
set	Update any subset of the model's system matrices/metadata in place.
sim	Simulate a time series from the state-space model.
smoother	Run the Kalman filter followed by the RTS backward smoother, and return the smoothed results together with the EM sufficient statistics computed from the smoothed states.
summarize_ssm_variances	Compute (and optionally print) a compact summary comparing theoretical (model-implied) and empirical (simulation-based) variances of the latent effect, observation noise and response.
summary	Build a statsmodels-style structured summary of the model configuration (not of fitted results - see StateSpaceResults.summary for that).

Attributes

F	Diagonal of the state transition matrix F.
H	Observation/design matrix mapping the latent state to the mean of y_t.
Q	Process noise covariance matrix (covariance of eta_t in the state equation).
R	Diagonal of the observation noise covariance matrix R (variance of e_t per observed location/variable).
Sigma0	Covariance of the initial state \(x_0\).
T	Number of observed time steps.
Xbeta	Exogenous regressors entering the observation equation as X_t @ beta.
b	Number of regression coefficients (length of beta).
backend	The JAX compute device this model's arrays and kernels are pinned to (resolved from the backend constructor argument via geossm.utils._select_device).
beta	Regression coefficients for Xbeta.
order	Human-readable model order label (used in summaries), e.g. an ARMA-style "(p, q)" tuple string.
p	Dimension of the measurement/observation equation (number of observed locations/variables).
params	Free/estimated model parameters; currently an alias for beta (kept as a distinct attribute so subclasses, e.g. an EM/optimization-fitted model, can extend it with additional parameters).
params_dim	Dimension of params; currently an alias for b.
params_names	Names corresponding to params; currently an alias for xbeta_names.
q	Dimension of the latent state equation (state rank).
shape	(p, q, T) - observation dimension, state dimension and number of time steps.
type	Human-readable model type/family label (used in summaries), e.g. "Linear (Gaussian)".
x0	Mean of the initial state \(x_0\).
xbeta_names	Human-readable names for the b regression terms in Xbeta/beta.
y_t	Currently-set observed data (may contain NaN for missing entries).
yname	Name of the dependent variable y, used in summaries.

__call__(y_t)

Shorthand for smoother, i.e. run the Kalman filter followed by the RTS smoother (“estimation” of the state).

Parameters:

y_t (array_like, shape (p, T)) – Observed data (may contain NaN for missing entries).

Returns:

StateSpaceResults – Smoothed results; see smoother.

set(H=None, R=None, F=None, Q=None, x0=None, Sigma0=None, y_t=None, Xbeta=None, beta=None, xbeta_names=None, yname=None)

Update any subset of the model’s system matrices/metadata in place.

Only arguments that are not None are updated; omitted ones keep their current value. This is the method the constructor and every public entry point (filter, smoother, sim, …) use internally to merge caller-supplied overrides into the model before running a kernel. No shape/consistency validation is performed here - see _check_parameters for that.

Parameters:

• H (array_like, shape (p, q), optional) – Observation/design matrix.

• R (array_like, shape (p, p) or (p,), optional) – Observation noise covariance (diagonal).

• F (array_like, shape (q, q) or (q,), optional) – State transition matrix (diagonal).

• Q (array_like, shape (q, q), optional) – Process noise covariance.

• x0 (array_like, shape (q,), optional) – Initial state mean.

• Sigma0 (array_like, shape (q, q), optional) – Initial state covariance.

• y_t (array_like, shape (p, T), optional) – Observed data.

• Xbeta (array_like, shape (p, b, T), optional) – Exogenous regressors.

• beta (array_like, shape (b,), optional) – Regression coefficients for Xbeta.

• xbeta_names (list of str, optional) – Names for the exogenous regression terms.

• yname (str, optional) – Name for the dependent variable (used in summaries).

Returns:

None

property backend

The JAX compute device this model’s arrays and kernels are pinned to (resolved from the backend constructor argument via geossm.utils._select_device).

Type:

jax.Device

estimate(y_t, yname=None, H=None, R=None, F=None, Q=None, x0=None, Sigma0=None, Xbeta=None, beta=None, xbeta_names=None, light=False)

Run the Kalman filter and smoother (filter + backward pass) and return the smoothed results, including the sufficient statistics (S11, S10, S00) needed by an EM M-step.

This is equivalent to smoother - kept as a distinct, named entry point for API compatibility - and does not recompute S11/S10/S00/y_hat a second time on top of what smoother already computes.

Parameters:

• y_t (array_like, shape (p, T)) – Observed data (may contain NaN for missing entries).

• yname (str, optional) – Name for the dependent variable (used in summaries).

• H (optional) – Overrides for the corresponding model parameters; see set. If omitted, the model’s currently-set values are used.

• R (optional) – Overrides for the corresponding model parameters; see set. If omitted, the model’s currently-set values are used.

• F (optional) – Overrides for the corresponding model parameters; see set. If omitted, the model’s currently-set values are used.

• Q (optional) – Overrides for the corresponding model parameters; see set. If omitted, the model’s currently-set values are used.

• x0 (optional) – Overrides for the corresponding model parameters; see set. If omitted, the model’s currently-set values are used.

• Sigma0 (optional) – Overrides for the corresponding model parameters; see set. If omitted, the model’s currently-set values are used.

• Xbeta (optional) – Overrides for the corresponding model parameters; see set. If omitted, the model’s currently-set values are used.

• beta (optional) – Overrides for the corresponding model parameters; see set. If omitted, the model’s currently-set values are used.

• xbeta_names (optional) – Overrides for the corresponding model parameters; see set. If omitted, the model’s currently-set values are used.

• light (bool, default False) – Forwarded to smoother; if True, the returned StateSpaceResults omits the filter-stage arrays (x_filtered, P_filtered, K, x_pred, P_pred).

Returns:

StateSpaceResults – Smoothed results, including filtered arrays (unless light=True), smoothed arrays, the log-likelihood and the EM sufficient statistics.

predict(H, x_T, P_T, Xbeta=None, beta=None)

Compute predicted observations and their state-driven predictive covariance from (typically smoothed) states.

Thin wrapper around _compute_predict_kernel_JAX; see its docstring for the exact equations. Unlike filter/smoother, this does not update the model’s own state (H, Xbeta, beta are used as passed, not read from self).

Parameters:

• H (array_like, shape (p, q)) – Observation matrix.

• x_T (array_like, shape (q, T+1)) – State means (index 0 is the prior/initial state and is excluded from the output).

• P_T (array_like, shape (q, q, T+1)) – State covariances (index 0 excluded from the output).

• Xbeta (array_like, shape (p, b, T), optional) – Exogenous regressors.

• beta (array_like, shape (b,), optional) – Regression coefficients for Xbeta.

Returns:

• y_hat (jax.numpy.ndarray, shape (p, T)) – Predicted observation means.

• Sigma_y_hat (jax.numpy.ndarray, shape (p, p, T)) – State-driven predictive covariance \(H P_{t} H^T\) (excludes the observation noise covariance R).

filter(y_t, yname=None, H=None, R=None, F=None, Q=None, x0=None, Sigma0=None, Xbeta=None, beta=None, xbeta_names=None)

Run the Kalman filter forward pass and return the filtered results (state estimates, log-likelihood, and one-step-ahead predicted observations/sufficient statistics computed from the filtered - not smoothed - states).

See _filter_kernelJAX for the exact filter recursion (predict/update, Woodbury identity, missing-data handling).

Parameters:

• y_t (array_like, shape (p, T)) – Observed data (may contain NaN for missing entries).

• yname (str, optional) – Name for the dependent variable (used in summaries).

• H (optional) – Overrides for the corresponding model parameters; see set. If omitted, the model’s currently-set values are used.

• R (optional) – Overrides for the corresponding model parameters; see set. If omitted, the model’s currently-set values are used.

• F (optional) – Overrides for the corresponding model parameters; see set. If omitted, the model’s currently-set values are used.

• Q (optional) – Overrides for the corresponding model parameters; see set. If omitted, the model’s currently-set values are used.

• x0 (optional) – Overrides for the corresponding model parameters; see set. If omitted, the model’s currently-set values are used.

• Sigma0 (optional) – Overrides for the corresponding model parameters; see set. If omitted, the model’s currently-set values are used.

• Xbeta (optional) – Overrides for the corresponding model parameters; see set. If omitted, the model’s currently-set values are used.

• beta (optional) – Overrides for the corresponding model parameters; see set. If omitted, the model’s currently-set values are used.

• xbeta_names (optional) – Overrides for the corresponding model parameters; see set. If omitted, the model’s currently-set values are used.

Returns:

StateSpaceResults – Filtered results: filtered/predicted state means and covariances, the Kalman gain, the log-likelihood, and the EM sufficient statistics computed from the filtered states.

Return type:

tuple

References

Durbin, J., & Koopman, S. J. (2012). Time Series Analysis by State Space Methods. Oxford University Press.

smoother(y_t, yname=None, H=None, R=None, F=None, Q=None, x0=None, Sigma0=None, Xbeta=None, beta=None, xbeta_names=None, light=False)

Run the Kalman filter followed by the RTS backward smoother, and return the smoothed results together with the EM sufficient statistics computed from the smoothed states.

This is the main “estimation” entry point (filtering and smoothing algorithm for linear-Gaussian state-space models); see _filter_kernelJAX and _smoother_kernelJAX for the exact recursions.

Parameters:

• y_t (array_like, shape (p, T)) – Observed data (may contain NaN for missing entries).

• yname (str, optional) – Name for the dependent variable (used in summaries).

• H (optional) – Overrides for the corresponding model parameters; see set. If omitted, the model’s currently-set values are used.

• R (optional) – Overrides for the corresponding model parameters; see set. If omitted, the model’s currently-set values are used.

• F (optional) – Overrides for the corresponding model parameters; see set. If omitted, the model’s currently-set values are used.

• Q (optional) – Overrides for the corresponding model parameters; see set. If omitted, the model’s currently-set values are used.

• x0 (optional) – Overrides for the corresponding model parameters; see set. If omitted, the model’s currently-set values are used.

• Sigma0 (optional) – Overrides for the corresponding model parameters; see set. If omitted, the model’s currently-set values are used.

• Xbeta (optional) – Overrides for the corresponding model parameters; see set. If omitted, the model’s currently-set values are used.

• beta (optional) – Overrides for the corresponding model parameters; see set. If omitted, the model’s currently-set values are used.

• xbeta_names (optional) – Overrides for the corresponding model parameters; see set. If omitted, the model’s currently-set values are used.

• light (bool, default False) – If True, the returned StateSpaceResults omits the filter-stage arrays (x_filtered, P_filtered, K, x_pred, P_pred) instead of keeping numpy copies of them alive for the whole lifetime of the results object. They are only used internally by this method (as smoother-kernel inputs); once the smoothed states are computed they serve no further purpose for callers - such as an EM loop - that never read them back. The default False preserves the previous behaviour for callers (e.g. model.smoother(y) used directly) that do want to inspect filtered vs. smoothed state.

Returns:

StateSpaceResults – Smoothed results: smoothed state means/covariances and lag-one covariances, the log-likelihood, predicted observations and the EM sufficient statistics (S11, S10, S00) computed from the smoothed states, plus (unless light=True) the filter-stage arrays.

Return type:

tuple

References

Durbin, J., & Koopman, S. J. (2012). Time Series Analysis by State Space Methods. Oxford University Press.

computeExpectedValues(x_T, P_T, P_T_1)

Compute predicted observations and EM sufficient statistics from a set of state means/covariances (filtered or smoothed).

Thin, timed wrapper around _compute_expected_values_kernelJAX; see its docstring for the exact equations for y_hat, S11, S10 and S00. Called internally by both filter (on the filtered states) and smoother (on the smoothed states).

Parameters:

• x_T (array_like, shape (q, T+1)) – State means (index 0 is the initial prior).

• P_T (array_like, shape (q, q, T+1)) – State covariances.

• P_T_1 (array_like, shape (q, q, T+1)) – Lag-one covariances \(\mathrm{Cov}(x_t, x_{t-1})\) (as produced by _smoother_kernelJAX; when called on filtered-only states this argument is typically not meaningful/available in the same sense).

Returns:

• y_hat (jax.numpy.ndarray, shape (p, T)) – Predicted observation means.

• S11, S10, S00 (jax.numpy.ndarray, shape (q, q)) – EM sufficient statistics.

• tDelta (float) – Wall-clock time (seconds) spent computing these quantities.

Return type:

tuple

sim(seed=1234, R=None, F=None, H=None, Q=None, x0=None, Sigma0=None, Xbeta=None, beta=None, block_p=None, block_q=None, stats=True, verbose=False)

Simulate a time series from the state-space model.

Draws \(x_0 \sim N(x_0, \Sigma_0)\) and then propagates the observation/state equations forward in time via _sim_kernelJAX (see its docstring for the exact equations). Uses a Python for-loop under the hood, so this is not JIT-compiled and is slower than the filter/smoother kernels.

Parameters:

• seed (int or geossm.utils.KeyStream, default 1234) – Either an integer seed used to build a new KeyStream, or an existing KeyStream to continue drawing from.

• R (optional) – Overrides for the corresponding model parameters; see set. If omitted, the model’s currently-set values are used.

• F (optional) – Overrides for the corresponding model parameters; see set. If omitted, the model’s currently-set values are used.

• H (optional) – Overrides for the corresponding model parameters; see set. If omitted, the model’s currently-set values are used.

• Q (optional) – Overrides for the corresponding model parameters; see set. If omitted, the model’s currently-set values are used.

• x0 (optional) – Overrides for the corresponding model parameters; see set. If omitted, the model’s currently-set values are used.

• Sigma0 (optional) – Overrides for the corresponding model parameters; see set. If omitted, the model’s currently-set values are used.

• Xbeta (optional) – Overrides for the corresponding model parameters; see set. If omitted, the model’s currently-set values are used.

• beta (optional) – Overrides for the corresponding model parameters; see set. If omitted, the model’s currently-set values are used.

• block_p (list of int, optional) – Block boundaries along the observation dimension p, forwarded to summarize_ssm_variances when stats=True (e.g. to separate multiple co-simulated processes). Defaults to a single block [0, p].

• block_q (list of int, optional) – Block boundaries along the latent dimension q, forwarded to summarize_ssm_variances. Defaults to a single block [0, q].

• stats (bool, default True) – If True, additionally compute and return a theoretical-vs- empirical variance summary via summarize_ssm_variances.

• verbose (bool, default False) – If True (and stats=True), print the variance summary.

Returns:

• y_t (jax.numpy.ndarray, shape (p, T)) – Simulated observations.

• x_t (jax.numpy.ndarray, shape (q, T+1)) – Simulated state vectors \([x_0, \dots, x_T]\).

• stats (dict or None) – Output of summarize_ssm_variances if stats=True, otherwise None.

• tdelta (float) – Wall-clock time (seconds) spent in the simulation kernel.

Return type:

Array

summarize_ssm_variances(x_sim, y_sim, block_p, block_q, decimals=4, verbose=True)

Compute (and optionally print) a compact summary comparing theoretical (model-implied) and empirical (simulation-based) variances of the latent effect, observation noise and response.

The theoretical stationary state covariance is obtained by solving the discrete Lyapunov equation \(P = F P F^T + Q\) for \(P\); since F is diagonal here, this has the closed form \(P_{ij} = Q_{ij} / (1 - f_i f_j)\) elementwise, avoiding a dense-matrix Lyapunov solve. The theoretical variance of the latent effect H @ x_t is then diag(H @ P @ H.T), and the theoretical response variance is the latent variance plus the observation noise variance (R); these are compared against the corresponding sample variances of x_sim/y_sim, averaged within each block of block_p/block_q (useful when several processes are stacked/simulated together).

Parameters:

• x_sim (array_like, shape (q, T+1) or (q, T)) – Simulated latent states, as returned by sim.

• y_sim (array_like, shape (p, T)) – Simulated observations with the fixed-effect term Xbeta @ beta already removed (i.e. the H @ x_t + e_t part).

• block_p (list of int) – Block boundaries along the observation dimension p (e.g. [0, p1, p1+p2, ...]) delimiting separate processes for the purpose of averaging.

• block_q (list of int) – Block boundaries along the latent dimension q.

• decimals (int, default 4) – Currently unused by this implementation.

• verbose (bool, default True) – If True, print a formatted summary table.

Returns:

dict – Dictionary with keys var_latent_theoretical, var_latent_empirical, var_noise, var_y_theoretical, var_y_empirical, ratios (a dict of empirical/theoretical ratios) and P (the theoretical stationary state covariance).

Raises:

AttributeError – If the model is missing any of F, Q, H, R.

Return type:

dict

property T

Number of observed time steps.

Type:

int

property p

Dimension of the measurement/observation equation (number of observed locations/variables).

Type:

int

property q

Dimension of the latent state equation (state rank).

Type:

int

property b

Number of regression coefficients (length of beta).

Type:

int

property y_t

Currently-set observed data (may contain NaN for missing entries).

Type:

jax.numpy.ndarray, shape (p, T)

property params

Free/estimated model parameters; currently an alias for beta (kept as a distinct attribute so subclasses, e.g. an EM/optimization-fitted model, can extend it with additional parameters).

Type:

array_like

property params_names

Names corresponding to params; currently an alias for xbeta_names.

Type:

list of str

property params_dim

Dimension of params; currently an alias for b.

Type:

int

property xbeta_names

Human-readable names for the b regression terms in Xbeta/beta.

Type:

list of str

property yname

Name of the dependent variable y, used in summaries.

Type:

str

property type

Human-readable model type/family label (used in summaries), e.g. "Linear (Gaussian)".

Type:

str

property order

Human-readable model order label (used in summaries), e.g. an ARMA-style "(p, q)" tuple string.

Type:

str

property shape

(p, q, T) - observation dimension, state dimension and number of time steps.

Type:

tuple of int

property Xbeta

Exogenous regressors entering the observation equation as X_t @ beta.

Type:

jax.numpy.ndarray, shape (p, b, T)

property beta

Regression coefficients for Xbeta.

Type:

jax.numpy.ndarray, shape (b,)

property H

Observation/design matrix mapping the latent state to the mean of y_t.

Type:

jax.numpy.ndarray, shape (p, q)

property R

Diagonal of the observation noise covariance matrix R (variance of e_t per observed location/variable).

Type:

jax.numpy.ndarray, shape (p,)

property F

Diagonal of the state transition matrix F.

Type:

jax.numpy.ndarray, shape (q,)

property Q

Process noise covariance matrix (covariance of eta_t in the state equation).

Type:

jax.numpy.ndarray, shape (q, q)

property x0

Mean of the initial state \(x_0\).

Type:

jax.numpy.ndarray, shape (q,)

property Sigma0

Covariance of the initial state \(x_0\).

Type:

jax.numpy.ndarray, shape (q, q)

__getstate__()

Return a serializable state for pickling.

Convert JAX arrays to NumPy arrays and store basic metadata.

__setstate__(state)

Restore object state from pickled state.

Arrays are converted back to JAX arrays with the original dtype.

generate_summary()

Build the two small key/value tables (model identity/config on the left, shape and mean system-matrix diagonals on the right) used by summary.

Returns:

tuple of list – (gen_top_left, gen_top_right), each a list of (label, [value]) pairs suitable for statsmodels.iolib.summary.Summary.add_table_2cols.

summary()

Build a statsmodels-style structured summary of the model configuration (not of fitted results - see StateSpaceResults.summary for that).

Returns:

statsmodels.iolib.summary.Summary – Printable summary table with model identity/type/order, shape (p, q, T) and mean system-matrix diagonals.

Return type:

Summary

__str__()

String representation of the model.