geossm.stmodel.LRStateSpaceModel

class geossm.stmodel.LRStateSpaceModel(df, formulas, domain, verbose=True, backend='auto', dtype=<class 'jax.numpy.float32'>)

Bases: StateSpaceModel

Low-Rank State-Space Model (LRSSM) for large-scale spatio-temporal data.

Subclass of geossm.ssm.StateSpaceModel specialised for geostatistical data observed at many sites over time. Rather than a full-rank spatial covariance over all observation sites (which does not scale), the latent spatial field of each response variable (“block”, one per formula) is written as a linear combination of a low-rank set of latent factors living on a finite-element (SPDE/GMRF) mesh – one already-setup()-ed geossm.covmodel.spdeAppoxCov per latent factor (see setup). The state-space measurement matrix H is built from a loading matrix A (one row per response variable, one column per latent factor) and each factor’s finite-element basis evaluated at the observation sites, and the model is fit with an EM algorithm built on the parent class’s Kalman filter/smoother (_E_step/_M_step).

Typical workflow

1. Build the model from a geopandas.GeoDataFrame, one Patsy/R-style formula per response variable, and one domain polygon per formula:

   model = LRStateSpaceModel(df, formulas=["y ~ 1 + x"], domain=[poly])
   

2. Build a mesh over (a buffer of) the domain and attach one already setup()-ed covariance function per latent factor:

   cov_fun = spdeAppoxCov(latlon=True).setup(mesh_io, domain=[poly])
   model = model.setup(cov_fun=[cov_fun])
   

3. Fit the model with the EM algorithm:

   results = model.fit()  # a LRStateSpaceResults
   

4. Optionally simulate (sim), predict at new locations/times (predict), and inspect/report the fit (results.summary()).

See the examples/example_LRSSM_*.py scripts and the README Quick Start section for complete, runnable versions of this workflow.

Methods

BuildDesignMatrix
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).
fit	Estimate the model parameters with an EM algorithm.
generate_summary	Build the left/right key-value rows used by summary()'s header table: model name/type/shape, then (if print_full=True) one "Grid ..." section per response variable's DesignMatrices and one "Latent.
generate_summary_header	Cheap identity/shape rows (name, type/order, dep.
getDesignMatrix	Build the design matrix for new (out-of-sample) locations/times -- e.g. a held-out test set -- exactly the way predict() builds it internally, plus the response y_obs, so the result can be compared directly against predict()'s output (e.g. y_obs - y_pred_list).
logger	Format one EM iteration's statistics (as built by _log_iteration) into a human-readable, multi-line block for console logging during fit().
predict	Out-of-sample prediction of the response variable(s) at new locations/times, from an already-fitted model.
print_info	Print msg prefixed with a UTC timestamp.
set	Update any subset of the model's system matrices/metadata in place.
setup	Build the model's latent covariance functions, one per latent factor, from already-built (and already setup()-ed) covariance functions.
sim	Simulate data from the low-rank 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	Return a statsmodels-style structured summary of the model (before fitting -- for the fitted results' own summary, see LRStateSpaceResults.summary).

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.
cov_function	The latent covariance functions set by setup, one per latent factor (self._cov_matern); None (or an empty list) until setup has been called.
domain	a list with one shapely Polygon or MultiPolygon per measurement equation (len(formulas)) -- see _setDomain.
domain_hull	Convex hull of the union of domain, i.e. of all measurement equations combined -- a single Polygon summarising the model's overall geographic footprint, e.g. to pass as domain to buildMesh2d.
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	Dimensions of the model as (p, q, T).
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.

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

setup(cov_fun)

Build the model’s latent covariance functions, one per latent factor, from already-built (and already setup()-ed) covariance functions.

Parameters:

cov_fun (list of spdeAppoxCov) – Already-built (and already setup()-ed) covariance functions, one per latent factor, used as-is. Each factor’s latent domain is already fixed as covi.domain, set when covi was itself setup()-ed (validated by spdeAppoxCov, via _validate_domain) – generally a different, usually smaller list than self.domain (the measurement-equation domain, one entry per formula); it is not defaulted from self.domain.

BuildDesignMatrix(df, formulas, verbose=None, domain=None)

getDesignMatrix(df, verbose=None)

Build the design matrix for new (out-of-sample) locations/times – e.g. a held-out test set – exactly the way predict() builds it internally, plus the response y_obs, so the result can be compared directly against predict()’s output (e.g. y_obs - y_pred_list).

Thin public wrapper around _getPredictionDesignMatrix (the single source of truth also used by predict()) with include_response=True: reuses each formula’s already-fitted DesignMatricesBuilder (self.builders) so stateful transforms (e.g. standardize()) are evaluated with the training mean/std rather than refit on df, and the model’s own measurement-equation domain (self.domain, fixed at __init__ time) – there is nothing left to check here that __init__ didn’t already validate once. Unlike predict()’s own use of it, df must contain the response column(s) referenced by self.formulas.

Parameters:

• df (geopandas.GeoDataFrame or list of geopandas.GeoDataFrame) – New locations/times, with the same columns required by the model’s formulas (including the response) and the same CRS as the training data. A single GeoDataFrame is used for every formula (e.g. a shared held-out set); a list with one GeoDataFrame per formula (see _asPredictionDataFrames) lets each response variable be compared at its own, different set of locations – e.g. model.getDesignMatrix([gdf_no2_test, gdf_pm10_test]) for a bivariate model, one test set per pollutant, instead of a single combined df that would give every formula the union of both pollutants’ locations.

• verbose (bool, optional) – Overrides self.verbose for this call’s logging.

Returns:

• points (list of ndarray) – One entry per response variable/formula, aligned with Xbeta’s/y_obs’s rows for that variable (see _buildPredictionGrid).

• gridList (list of DesignMatrices) – The per-formula DesignMatrices Xbeta/y_obs were built from (timestamps, CRS, …).

• ndim (int) – Total number of rows stacked across all response variables.

• pdim (list of int) – Per-variable row counts (block_p[i+1] - block_p[i]).

• block_p (ndarray) – Cumulative row-index boundaries, one more entry than pdim.

• T (list of int) – Per-variable number of time steps.

• Xbeta (ndarray) – Stacked, block-diagonal-by-variable fixed-effect design matrix – the same array predict() feeds through H/beta to build y_pred.

• y_obs (ndarray) – Stacked response, one row per observation (same row order as Xbeta), evaluated from df with the training y_design_info.

property shape

Dimensions of the model as (p, q, T).

Returns:

tuple of (int or None, int or None, int or None) – p : total number of observed sites stacked across all response variables (from y_train/Xbeta, whichever is set). q : total latent dimension (rank), i.e. the sum of cov.fem_solver.n_inner_points over every latent factor set in setup. T : number of time steps. Any entry is None if the corresponding data (y_train/Xbeta, _cov_matern) has not been set yet.

sim(formulas=None, seed=1234, params=None, verbose=None, stats=False)

Simulate data from the low-rank state-space model.

Builds the model matrices H, R, F, Q from params (or from default initial values if a field is left unset, see _parseParams) and the latent covariance functions set in setup, then delegates to the parent class’s StateSpaceModel.sim.

Parameters:

• formulas (list of str, optional) – Patsy/R-style formulas (right-hand side only is used to build Xbeta, e.g. “1 + temperature”), one per response variable, used to build a fresh observation grid/design matrix for the simulation. If None, reuses the design matrix/grid the model was built with at __init__ (self.Xbeta, self.points, …). self.formulas must not be None in that case.

• seed (int, default 1234) – Random seed for the simulation, forwarded to StateSpaceModel.sim.

• params (ModelParams, optional) – Model parameters to simulate with (beta, A, s2e, f, ks, one Param each). Fields left as None fall back to the model’s default initial values (see _parseParams). ks, if given, is also written into each covariance function’s rescale attribute (self.cov_function[i].rescale).

• verbose (bool, optional) – Overrides self.verbose for this call’s logging.

• stats (bool, default False) – If True, asks StateSpaceModel.sim to also compute and return variance diagnostics of the simulated state/observations (see StateSpaceModel.summarize_ssm_variances).

Returns:

• y_sim (ndarray) – Simulated observations, stacked across all response variables (shape (p, T), p = sum(pdim)).

• x_sim (ndarray) – Simulated latent state trajectory (shape (q, T)).

• info (dict) – Metadata about the simulation: formulas, y_name, xbeta_names, Xbeta, params, points, T, stats (variance diagnostics, if stats=True), qdim, nvar, nlat, pdim, block_p, block_q, and the ad-hoc sim_model (a plain StateSpaceModel built with the matrices used for the simulation).

• tdelta (float) – Wall-clock time (seconds) spent in StateSpaceModel.sim.

Raises:

ValueError – If formulas and self.formulas are both None; if the covariance functions have not been set via setup; or if the shapes of beta, A, or ks in params are inconsistent with the model (Xbeta’s width, (nvar, nlat), and the number of covariance functions, respectively).

predict(df, modelresults, verbose=True, std=True)

Out-of-sample prediction of the response variable(s) at new locations/times, from an already-fitted model.

For each formula, builds a new design matrix/grid over df via _getPredictionDesignMatrix – which reuses the formula’s fitted DesignMatricesBuilder (self.builders) so stateful transforms in the formula (e.g. standardize()) are evaluated with the training mean/std rather than recomputed on df – then maps the fitted model’s smoothed states (modelresults.x_smoothed, modelresults.P_smoothed) through the new observation matrix H and design matrix to get the predictive mean and covariance (see _predict). The public getDesignMatrix builds this same design matrix (plus the response y_obs) for inspection/comparison outside of predict().

Parameters:

• df (geopandas.GeoDataFrame or list of geopandas.GeoDataFrame) – New locations/times to predict at, with the same columns required by the model’s formulas (minus the response, which is not needed for prediction) and the same CRS as the training data. A single GeoDataFrame is used for every formula (e.g. a shared regular grid); a list with one GeoDataFrame per formula (see _asPredictionDataFrames) predicts each response variable at its own, different set of locations instead.

• modelresults (LRStateSpaceResults) – Results of a previous fit() call, providing the estimated parameters (modelresults.params) and smoothed states used to build the prediction. Updated and returned in place.

• verbose (bool, default True) – Verbosity for the observation-grid/design-matrix construction logging.

• std (bool, default True) – If False, skip the predictive covariance Sigma_y_pred (left as None) – useful when the (p, p, T) array does not fit in device memory.

Returns:

LRStateSpaceResults – The same modelresults object, with points_pred, y_pred/Sigma_y_pred (and their per-variable views y_pred_list/Sigma_y_pred_list; Sigma_y_pred is None if std=False), tdelta_pred, timestamps_pred, and crs_pred populated (one entry per response variable, except tdelta_pred/crs_pred).

Raises:

ValueError – If modelresults is None.

fit(params0=None, options=None)

Estimate the model parameters with an EM algorithm.

Each EM iteration alternates: an E-step (_E_step) that runs the parent class’s Kalman filter/smoother given the current parameters to get the smoothed states/covariances and the sufficient statistics S11/S10/S00, plus the observed-data log-likelihood; and an M-step (_M_step) that updates beta (regression coefficients), s2e (measurement-error variances), f (latent factors’ AR(1) coefficients), A (loading matrix), and ks (Matern rescale/range parameters, one per covariance function, optimised with L-BFGS-B against exact JAX gradients – see _build_ks_value_and_grad) from those statistics. Iteration stops when the relative change in log-likelihood drops to tol_relat or max_iter iterations are reached (see FitOptions).

Parameters:

• params0 (ModelParams, optional) – Initial/fixed values for the parameters. Any field left as None is filled in with data-driven initial values (OLS beta, moment-based s2e/A, a mesh-based initial ks, … – see _getInitialValues); a field with fixed=True on its Param is held at its given value and excluded from the M-step updates for the whole run (see _updateParams). If None, every parameter is initialised and freely estimated.

• options (FitOptions, optional) – EM stopping-rule and logging options: max_iter caps the number of EM iterations, tol_relat is the relative log-likelihood-improvement tolerance below which EM is considered converged, and verbose overrides self.verbose for the whole call (including the per-iteration log printed via logger). Defaults to max_iter=100, tol_relat=1e-3 if not given (note this differs from FitOptions’s own dataclass defaults, which are only used if options itself is None – passing an explicit FitOptions() uses its max_iter=20).

Returns:

LRStateSpaceResults – The fitted results: estimated parameters, per-iteration EM statistics (nstats/llf_path), in-sample fitted values (y_hat/y_hat_list and their covariances), the smoothed states, and the sufficient statistics S11/S10/S00.

property cov_function

The latent covariance functions set by setup, one per latent factor (self._cov_matern); None (or an empty list) until setup has been called.

Type:

list of spdeAppoxCov

property domain

a list with one shapely Polygon or MultiPolygon per measurement equation (len(formulas)) – see _setDomain. Distinct from the latent domain (one entry per latent factor, generally a different, usually smaller count): each factor’s latent domain is fixed on its own spdeAppoxCov when that object is setup()-ed, before it is passed to setup(cov_fun=…), and is available via cov_function[i].domain. See domain_hull for a single Polygon summarising every measurement equation’s domain combined.

Type:

The measurement-equation domain

property domain_hull

Convex hull of the union of domain, i.e. of all measurement equations combined – a single Polygon summarising the model’s overall geographic footprint, e.g. to pass as domain to buildMesh2d. See _domain_hull.

logger(stats, beta_decimals=2, scalar_decimals=2, relat_decimals=5)

Format one EM iteration’s statistics (as built by _log_iteration) into a human-readable, multi-line block for console logging during fit().

Parameters:

• stats (dict) – One entry of the nstats history produced by _log_iteration (keys niter, logL, deltaL, relatL, beta, s2e, f, ks, opt_success, A, x0, Sigma0, time_tot, tdelta_E, tdelta_M).

• beta_decimals (int, default 2) – Decimal places used to format stats[“beta”].

• scalar_decimals (int, default 2) – Decimal places used to format the other numeric fields (log-likelihood, s2e, f, ks, A, x0, Sigma0, timings).

• relat_decimals (int, default 5) – Decimal places used to format the relative log-likelihood change stats[“relatL”].

Returns:

str – The formatted, multi-line iteration summary.

generate_summary_header()

Cheap identity/shape rows (name, type/order, dep. variable, shape, backend/dtype) - the subset of generate_summary() that’s pure attribute lookups, with no per-matrix computation. Used by __repr__ (via format_info_table) so that a bare console/Jupyter display never touches R/Q/F/x0/Sigma0.

Returns:

gen_top_left, gen_top_right (list of (str, list)) – Two lists of (label, [value]) rows, matching the other generate_summary* methods’ return shape.

generate_summary(print_short=False)

Build the left/right key-value rows used by summary()’s header table: model name/type/shape, then (if print_full=True) one “Grid …” section per response variable’s DesignMatrices and one “Latent. …” section per latent covariance function, each via that object’s own generate_summary.

Parameters:

print_full (bool, default True) – If False, only the top model-level rows are returned (used when this table is embedded elsewhere, e.g. by LRStateSpaceResults.generate_summary, without repeating the per-grid/per-covariance detail).

Returns:

gen_top_left, gen_top_right (list of (str, list)) – Two lists of (label, [value]) rows, of equal length.

summary(print_short=False)

Return a statsmodels-style structured summary of the model (before fitting – for the fitted results’ own summary, see LRStateSpaceResults.summary).

Parameters:

print_full (bool, default True) – If True, includes one section per response variable’s observation grid and one per latent covariance function (see generate_summary); if False, only the top model-level rows (name, type, shape) are included – used by fit() to print a compact header before the EM iterations start.

Returns:

statsmodels.iolib.summary.Summary – Printable summary object (str(…)/print(…)).

Return type:

Summary

__str__()

Return a human-readable summary of the model (see summary).

__repr__()

Cheap, side-effect-free representation for interactive use (model name/type/shape/backend only, from generate_summary_header) – see summary()/str() for the full, more expensive report (per-grid/per-covariance detail). Does not touch self.model or any other state, unlike summary().

property F

Diagonal of the state transition matrix F.

Type:

jax.numpy.ndarray, shape (q,)

property H

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

Type:

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

property Q

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

Type:

jax.numpy.ndarray, shape (q, 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 Sigma0

Covariance of the initial state \(x_0\).

Type:

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

property T

Number of observed time steps.

Type:

int

property Xbeta

Exogenous regressors entering the observation equation as X_t @ beta.

Type:

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

__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.

__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.

property b

Number of regression coefficients (length of beta).

Type:

int

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

property beta

Regression coefficients for Xbeta.

Type:

jax.numpy.ndarray, shape (b,)

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

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.

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.

property p

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

Type:

int

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_dim

Dimension of params; currently an alias for b.

Type:

int

property params_names

Names corresponding to params; currently an alias for xbeta_names.

Type:

list of str

property q

Dimension of the latent state equation (state rank).

Type:

int

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

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.

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 x0

Mean of the initial state \(x_0\).

Type:

jax.numpy.ndarray, shape (q,)

property xbeta_names

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

Type:

list of str

property y_t

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

Type:

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

property yname

Name of the dependent variable y, used in summaries.

Type:

str

print_info(msg)

Print msg prefixed with a UTC timestamp.