linear gaussian ssm

When the dynamics model and the observation model of a state space model are both Gaussian, we have a linear gaussian ssm,

The latent dynamics model is written as

with the observation model

variable	variable description	shape
	state vector	(N_states, 1)
	observation vector at time	(N_obs, 1)
	dynamics (transition) matrix	(N_states, N_states)
	covariance matrix of dynamics (system) noise	(N_states, N_states)
	emission (observation) matrix	(N_obs, N_states)
	covariance function for emission (observation) noise	(N_obs, N_obs)

extracts the relevant parts of the state vector . are exogenous inputs.

Inference can be performed efficiently using kalman filtering and smoothing. These models have applications in object tracking and structural time series models.

jax implementation (cuthbert)

cuthbert provides an exact Kalman filter and smoother in square-root form via cuthbert.gaussian.kalman.

from cuthbert import filter, smoother
from cuthbert.gaussian import kalman
 
filter_obj = kalman.build_filter(
    get_init_params=lambda inputs: (m0, chol_P0),          # p(x_0)
    get_dynamics_params=lambda inputs: (F, c, chol_Q),     # p(x_t | x_{t-1})
    get_observation_params=lambda inputs: (H, d, chol_R, y),  # p(y_t | x_t)
)
 
states = filter(filter_obj, model_inputs, parallel=True)

The filter supports temporal parallelisation via jax.lax.associative_scan (associative Kalman filter, cuthbert parallelization example).

See also ssm in dynamax and ssm resources.