Models

Continuous-discrete linear Gaussian SSM model definitions.

models

ContDiscreteLinearGaussianSSM

Bases: SSM

Definition of a Continuous-Discrete Linear Gaussian State Space Model.

The CD-LGSSM model is defined in the following way, according to equation (3.134) in Sarkka (2019):

\[p(z_0) = \mathcal{N}(z_0 \mid m, S)\]

\[dz = F_t z_t dt + B_t u_t dt + b_t dt + L_t d\beta_t\]

with \(\beta_t\) a standard Brownian motion, implying

\[p(z_{t_k} \mid z_{t_{k-1}}, u_t) = \mathcal{N}(z_{t_k} \mid A_{t_k} z_{t_{k-1}} + B_{t_k} u_{t_k} + b_{t_k}, Q_{t_k})\]

where \(A_{t_k}\) and \(Q_{t_k}\) are computed as the solution to the SDE above over the interval \([t_{t_{k-1}}, t_{t_k}]\), and emissions defined as

\[p(y_{t_k} \mid z_{t_k}) = \mathcal{N}(y_{t_k} \mid H_{t_k} z_{t_k} + D_{t_k} u_{t_k} + d_{t_k}, R_{t_k})\]

where

• \(z_{t_k}\) is a latent state of size state_dim,
• \(m\) = mean of initial state

• \(S\) = covariance matrix of initial state

• \(F_t\) = is the linear dynamics of the state, as in the SDE definition above

• \(A_t\) = are the dynamics (transition) matrix of the state ---computed as the pushforward of the continuous dynamics--- A_t is the solution to the ODE in eq (3.135) in Sarkka (2019)
• \(L\) = diffusion coefficient of the dynamics SDE, used to compute Q_t
• \(Q\) = diffusion covariance matrix of dynamics (system) ---brownian motion--- noise, computed as the solution to the ODE in eq (3.135) in Sarkka (2019)
• \(u_t\) is an input of size input_dim (defaults to 0)
• \(B\) = (optional) input-to-state weight matrix

• \(b\) = (optional) state bias vector

• \(y_{t_k}\) is an observed variable (emissions) of size emission_dim

• \(H\) = emission (observation) matrix
• \(R\) = covariance function for emission (observation) noise
• \(D\) = (optional) input-to-emission weight matrix
• \(d\) = (optional) emission bias vector

The parameters of the model are stored in a ParamsCDLGSSM. You can create the parameters manually, or by calling initialize.

Parameters:

Name	Type	Description	Default
state_dim	int	Dimensionality of latent state.	required
emission_dim	int	Dimensionality of observation vector.	required
input_dim	int	Dimensionality of input vector. Defaults to 0.	0
has_dynamics_bias	bool	Whether model contains an offset term \(b\). Defaults to True.	True
has_emissions_bias	bool	Whether model contains an offset term \(d\). Defaults to True.	True
diffeqsolve_settings	dict	Settings to pass to the differential equation solver when computing the pushforward of the continuous-time dynamics. Defaults to {}.	{}

emission_distribution(params: ParamsCDLGSSM, state: Float[Array, ' state_dim'], inputs: Optional[Float[Array, 'ntime input_dim']] = None) -> tfd.Distribution

CD-LGSSM Emission distribution. Args: params: CD-LGSSM model parameters. state: current state. inputs: optional inputs.

Returns:

Type	Description
Distribution	Gaussian, state-conditional emission distribution

filter(params: ParamsCDLGSSM, emissions: Float[Array, 'ntime emission_dim'], t_emissions: Optional[Float[Array, 'ntime 1']] = None, filter_hyperparams: Optional[KFHyperParams] = KFHyperParams(), inputs: Optional[Float[Array, 'ntime input_dim']] = None, warn: bool = True) -> PosteriorGSSMFiltered

Run the CD-Kalman filter to compute the filtered posterior distribution over states. Args: params: CD-LGSSM model parameters. emissions: sequence of observations. t_emissions: continuous-time specific time instants of observations: if not None, it is an array filter_hyperparams: hyperparameters for the Kalman filter. inputs: optional sequence of inputs. warn: whether to warn about numerical issues. Returns: filtered posterior distribution over states.

filter_and_forecast(params, emissions_filter, t_emissions_filter, t_emissions_forecast, inputs_filter=None, inputs_forecast=None, warn: bool = True)

Run the CD-Kalman filter to compute the filtered posterior distributions over states, and then run forecasting from the last filtered state.

Parameters:

Name	Type	Description	Default
params		model parameters.	required
emissions_filter		sequence of observations for filtering.	required
t_emissions_filter		continuous-time specific time instants of observations for filtering: if not None, it is an array	required
t_emissions_forecast		continuous-time specific time instants for forecasting: if not None, it is an array	required
inputs_filter		optional sequence of inputs for filtering.	None
inputs_forecast		optional sequence of inputs for forecasting.	None
warn	bool	whether to warn about numerical issues.	True

Returns:

Type	Description
	filtered and forecasted posterior distributions over states.

initial_distribution(params: ParamsCDLGSSM, inputs: Optional[Float[Array, 'ntime input_dim']] = None) -> tfd.Distribution

Initial distribution. Args: params: CD-LGSSM model parameters. inputs: optional initial inputs (not used in state initialization). Returns: initial Gaussian state distribution.

initialize(key: PRNGKey = jr.PRNGKey(0), init_prior: Prior = None, initial_mean: dict = None, initial_cov: dict = None, dynamics_weights: dict = None, dynamics_bias: dict = None, dynamics_input_weights: dict = None, dynamics_diffusion_coefficient: dict = None, dynamics_diffusion_cov: dict = None, emission_weights: dict = None, emission_bias: dict = None, emission_input_weights: dict = None, emission_cov: dict = None) -> Tuple[ParamsCDLGSSM, ParamsCDLGSSM]

Initialize CD-LGSSM parameters, and their corresponding properties.

Parameters:

Name	Type	Description	Default
key	PRNGKey	Random number key. Defaults to jr.PRNGKey(0).	PRNGKey(0)
init_prior	Prior	prior distribution for the initialization. Defaults to None.	None
initial_mean	dict	parameter \(m\). Defaults to None.	None
initial_cov	dict	parameter \(S\). Defaults to None.	None
dynamics_weights	dict	parameter \(F\). Defaults to None.	None
dynamics_bias	dict	parameter \(b\). Defaults to None.	None
dynamics_input_weights	dict	parameter \(B\). Defaults to None.	None
dynamics_diffusion_coefficient	dict	parameter \(L\). Defaults to None.	None
dynamics_diffusion_cov	dict	parameter \(Q\). Defaults to None.	None
emission_weights	dict	parameter \(H\). Defaults to None.	None
emission_bias	dict	parameter \(d\). Defaults to None.	None
emission_input_weights	dict	parameter \(D\). Defaults to None.	None
emission_cov	dict	parameter \(R\). Defaults to None.	None

Returns:

Type	Description
Tuple[ParamsCDLGSSM, ParamsCDLGSSM]	Tuple[ParamsCDLGSSM, ParamsCDLGSSM]: parameters and their properties.

marginal_log_prob(params: ParamsCDLGSSM, emissions: Float[Array, 'ntime emission_dim'], t_emissions: Optional[Float[Array, 'ntime 1']] = None, filter_hyperparams: Optional[KFHyperParams] = KFHyperParams(), inputs: Optional[Float[Array, 'ntime input_dim']] = None, key: PRNGKey = jr.PRNGKey(0), warn: bool = True) -> Scalar

Compute the marginal log likelihood of a sequence of emissions under the CD-LGSSM model.

Parameters:

Name	Type	Description	Default
params	ParamsCDLGSSM	CD-LGSSM model parameters.	required
emissions	Float[Array, 'ntime emission_dim']	sequence of observations.	required
t_emissions	Optional[Float[Array, 'ntime 1']]	continuous-time specific time instants of observations: if not None, it is an array	None
filter_hyperparams	Optional[KFHyperParams]	hyperparameters for the Kalman filter.	KFHyperParams()
inputs	Optional[Float[Array, 'ntime input_dim']]	optional sequence of inputs.	None
key	PRNGKey	random number generator.	PRNGKey(0)
warn	bool	whether to warn about numerical issues.	True

Returns: Marginal log likelihood of the emissions, \(\log p(y_{1:T})\).

posterior_predictive(params: ParamsCDLGSSM, emissions: Float[Array, 'ntime emission_dim'], t_emissions: Optional[Float[Array, 'ntime 1']] = None, filter_hyperparams: Optional[KFHyperParams] = KFHyperParams(), inputs: Optional[Float[Array, 'ntime input_dim']] = None) -> Tuple[Float[Array, 'ntime emission_dim'], Float[Array, 'ntime emission_dim']]

Compute marginal posterior predictive smoothing distribution for each observation.

Parameters:

Name	Type	Description	Default
params	ParamsCDLGSSM	model parameters.	required
emissions	Float[Array, 'ntime emission_dim']	sequence of observations.	required
t_emissions	Optional[Float[Array, 'ntime 1']]	continuous-time specific time instants of observations: if not None, it is an array	None
filter_hyperparams	Optional[KFHyperParams]	hyperparameters for the Kalman filter.	KFHyperParams()
inputs	Optional[Float[Array, 'ntime input_dim']]	optional sequence of inputs.	None

Returns:

Type	Description
Tuple[Float[Array, 'ntime emission_dim'], Float[Array, 'ntime emission_dim']]	posterior predictive means \(\mathbb{E}[y_{t,d} \mid y_{1:T}]\) and standard deviations \(\mathrm{std}[y_{t,d} \mid y_{1:T}]\)

posterior_sample(key: PRNGKey, params: ParamsCDLGSSM, emissions: Float[Array, 'ntime emission_dim'], t_emissions: Optional[Float[Array, 'ntime 1']] = None, inputs: Optional[Float[Array, 'ntime input_dim']] = None) -> Float[Array, 'ntime state_dim']

Sample from the posterior distribution over states given emissions using the CD-Kalman filter/smoother. Args: key: random number generator. params: model parameters. emissions: sequence of observations. t_emissions: continuous-time specific time instants of observations: if not None, it is an array inputs: optional sequence of inputs. Returns: sampled latent states.

sample_dist(params: ParamsCDLGSSM, key: PRNGKey, num_timesteps: int, t_emissions: Optional[Float[Array, 'ntime 1']] = None, inputs: Optional[Float[Array, 'ntime input_dim']] = None) -> Tuple[Float[Array, 'num_timesteps state_dim'], Float[Array, 'num_timesteps emission_dim']]

Sample from the joint distribution of the CD-LGSSM to produce states and emission trajectories.

Parameters:

Name	Type	Description	Default
params	ParamsCDLGSSM	CD-LGSSM model parameters	required
key	PRNGKey	random number generator key	required
num_timesteps	int	number of time steps to sample	required
t_emissions	Optional[Float[Array, 'ntime 1']]	continuous-time specific time instants of observations: if not None, it is an array	None
inputs	Optional[Float[Array, 'ntime input_dim']]	optional array of inputs.	None

Returns: latent states and observed emissions

sample_path(params: ParamsCDLGSSM, key: PRNGKey, num_timesteps: int, t_emissions: Optional[Float[Array, 'num_timesteps 1']] = None, inputs: Optional[Float[Array, 'num_timesteps input_dim']] = None) -> Tuple[Float[Array, 'num_timesteps state_dim'], Float[Array, 'num_timesteps emission_dim']]

Sample from a forward path of the CD-LGSSM to produce state and emission trajectories.

Parameters:

Name	Type	Description	Default
params	ParamsCDLGSSM	model parameters	required
key	PRNGKey	random number generator key	required
num_timesteps	int	number of time steps to sample	required
t_emissions	Optional[Float[Array, 'num_timesteps 1']]	continuous-time specific time instants of observations: if not None, it is an array	None
inputs	Optional[Float[Array, 'num_timesteps input_dim']]	optional array of inputs.	None

Returns:

Type	Description
Tuple[Float[Array, 'num_timesteps state_dim'], Float[Array, 'num_timesteps emission_dim']]	latent states and observed emissions

sample_prior(prior: Prior, M: int, init_params: Optional[ParamsCDLGSSM] = None, key: Optional[PRNGKey] = jr.PRNGKey(0)) -> Tuple[ParamsCDLGSSM, ParamsCDLGSSM]

Sample from the prior distribution over CD-LGSSM model parameters.

Parameters:

Name	Type	Description	Default
prior	Prior	Prior distribution.	required
M	int	Number of samples to draw.	required
init_params	Optional[ParamsCDLGSSM]	Dictionary of parameters to use as initialization; if not provided, default parameters are used.	None
key	Optional[PRNGKey]	Random number generator key.	PRNGKey(0)

Returns:

Type	Description
Tuple[ParamsCDLGSSM, ParamsCDLGSSM]	Tuple of sampled CD-LGSSM parameters and properties objects.

smoother(params: ParamsCDLGSSM, emissions: Float[Array, 'ntime emission_dim'], t_emissions: Optional[Float[Array, 'ntime 1']] = None, filter_hyperparams: Optional[KFHyperParams] = KFHyperParams(), inputs: Optional[Float[Array, 'ntime input_dim']] = None, warn: bool = True) -> PosteriorGSSMSmoothed

Compute the smoothing distribution over states using the CD-Kalman smoother.

Parameters:

Name	Type	Description	Default
params	ParamsCDLGSSM	model parameters.	required
emissions	Float[Array, 'ntime emission_dim']	sequence of observations.	required
t_emissions	Optional[Float[Array, 'ntime 1']]	continuous-time specific time instants of observations: if not None, it is an array	None
filter_hyperparams	Optional[KFHyperParams]	hyperparameters for the Kalman filter.	KFHyperParams()
inputs	Optional[Float[Array, 'ntime input_dim']]	optional sequence of inputs.	None

Returns:

Type	Description
PosteriorGSSMSmoothed	smoothed posterior distributions over states.

transition_distribution(params: ParamsCDLGSSM, state: Float[Array, ' state_dim'], t0: Optional[Float] = None, t1: Optional[Float] = None, inputs: Optional[Float[Array, 'ntime input_dim']] = None) -> tfd.Distribution

CD-LGSSM Transition distribution, defined as the pushforward of the continuous-time, linear dynamics and discrete inputs (controls).

Parameters:

Name	Type	Description	Default
params	ParamsCDLGSSM	CD-LGSSM model parameters.	required
state	Float[Array, ' state_dim']	current state.	required
t0	Optional[Float]	initial time.	None
t1	Optional[Float]	final time.	None
inputs	Optional[Float[Array, 'ntime input_dim']]	optional inputs.	None

Returns:

Type	Description
Distribution	Gaussian, state transition distribution.

SuffStatsCDLGSSM

Bases: Protocol

A NamedTuple with sufficient statistics for CDLGSSM parameter estimation.