Utilities

Utility classes and functions for continuous-discrete linear Gaussian SSMs.

cdlgssm_utils

GSSMForecast

Bases: NamedTuple

Object definition used when forecasting.

Parameters:

Name	Type	Description	Default
forecasted_state_means		Array of forecasted state means \(\mathbb{E}[z_{t+1:t+t_f} \mid y_{1:t}, u_{1:t}, u_{t+1:t+f}]\).	required
filtered_covariances		Array of forecasted state covariances \(\mathrm{Cov}[z_{t+1:t+t_f} \mid y_{1:t}, u_{1:t}, u_{t+1:t+f}]\).	required
forecasted_emission_means		Array of forecasted emission means \(\mathbb{E}[y_{t+1:t+t_f} \mid y_{1:t}, u_{1:t}, u_{t+1:t+f}]\).	required
forecasted_emission_covariances		Array of forecasted emission covariances \(\mathrm{Cov}[y_{t+1:t+t_f} \mid y_{1:t}, u_{1:t}, u_{t+1:t+f}]\).	required
forecasted_state_path		Array of forecasted state path \(z_{t+1:t+t_f}\).	required
forecasted_emission_path		Array of forecasted emission path \(y_{t+1:t+t_f}\).	required

ParamsCDLGSSM

Bases: NamedTuple

Parameters of a linear Gaussian CD-LGSSM.

Parameters:

Name	Type	Description	Default
initial		Initial distribution parameters, same as in LGSSM.	required
dynamics		Dynamics distribution parameters.	required
emissions		Emission distribution parameters.	required

ParamsCDLGSSMDynamics

Bases: NamedTuple

Parameters of the CD-LGSSM state dynamics The tuple doubles as a container for the ParameterProperties.

We assume a model of the form \(dz_t = F_t z_t dt + L_t dB_t\)

The resulting transition distribution is \(p(z_{t1}} | z_{t0}, u_t1) = N(z_{t1} | A z_{t0} + B u_{t1} + b, P)\) where A, B, b, P are computed based on the SDE defined by F_t, L_t and Q.

Parameters:

Name	Type	Description	Default
weights		Dynamics weights \(F_t\).	required
bias		Dynamics bias \(b\).	required
input_weights		Dynamics input weights \(B\).	required
diffusion_coefficient		Dynamics diffusion coefficient \(L_t\).	required
diffusion_cov		Dynamics covariance \(Q\).	required

create_cdlgssm_params_and_props(params: dict) -> Tuple[ParamsCDLGSSM, ParameterProperties]

Create CD-LGSSM parameters and properties, based on provided dictionaries.

Parameters:

Name	Type	Description	Default
params	dict	Dictionary of parameters.	required

Returns:

Type	Description
Tuple[ParamsCDLGSSM, ParameterProperties]	Tuple of parameters and properties objects.

init_cdlgssm_params(default_params, init_params=None, init_prior=None, key=jr.PRNGKey(0)) -> Tuple[ParamsCDLGSSM, ParamsCDLGSSM]

Initialize CD-LGSSM parameters and properties from prior, init_values, or defaults.

Parameters:

Name	Type	Description	Default
default_params		Dictionary of default parameters; at least some default values are required.	required
init_params		Dictionary of all parameters.	None
init_prior		Prior distribution for the initialization. Defaults to None.	None
key		Random key for sampling from the prior. Defaults to jr.PRNGKey(0).	PRNGKey(0)

Returns:

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

sample_cdlgssm_params(prior, M, init_params, key=jr.PRNGKey(0)) -> Tuple[ParamsCDLGSSM, ParamsCDLGSSM]

Sample CD-LGSSM parameters from the provided prior; init_params used for non-sampled parameters.

Parameters:

Name	Type	Description	Default
prior		Prior distribution for the initialization.	required
M		Number of samples to draw.	required
init_params		Dictionary of all parameters.	required
key		Random key for sampling from the prior. Defaults to jr.PRNGKey(0).	PRNGKey(0)

Returns:

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