Utilities

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

cdnlgssm_utils

LearnableDiagonalMatrix

Bases: NamedTuple

A learnable diagonal matrix i.e., f(x,u,t) = diag_params = diagonal matrix

LearnableFunction

Bases: NamedTuple

A definiton of a function that takes as input x, u and t

All Learnable functions should have params properties

f(x, u=None, t=None) abstractmethod

A function to be defined by each specific class With inputs x: state u: inputs t: time

LearnableLinear

Bases: NamedTuple

A linear function with learnable parameters i.e., f(x,u,t) = weights @ x + bias

weights: weights of the linear function bias: bias of the linear function

LearnableMatrix

Bases: NamedTuple

A learnable matrix i.e., f(x,u,t) = params = matrix

LearnableScaledMatrix

Bases: NamedTuple

A learnable scaled matrix i.e., f(x,u,t) = scale * matrix

LearnableVector

Bases: NamedTuple

A learnable vector i.e., f(x,u,t) = params = vector

ParamsCDNLGSSM

Bases: NamedTuple

Parameters of a nonlinear Gaussian CD-NLGSSM.

Parameters:

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

ParamsCDNLGSSMDynamics

Bases: NamedTuple

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

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

The resulting transition distribution is \(p(z_{t1}| z_{t0}, u_{t1}) = N(z_{t1} | m(z_{t0}, u_{t1}), P)\) where the mean m(z_{t0}, u_{t1}) and covariance Q are computed based on numerically solving the SDE defined by f, L_t and Q.

For the solution of the SDE, we use an approximation of order defined by dynamics_approx:

This model does not obey the solution to the SDE as in Särkkä's equation (3.151): the true solution is not necessarily a Gaussian Process (note there are cases where that is indeed the case) we here approximate the solution at each time step with a Gaussian distribution.

Parameters:

Name	Type	Description	Default
drift_function		Drift \(f\).	required
diffusion_coefficient		Diffusion coefficient \(L_t\).	required
diffusion_cov		Diffusion covariance \(Q\).	required
dynamics_approx		One of 'zeroth', 'first', or 'second'.	required

ParamsCDNLGSSMEmissions

Bases: NamedTuple

Parameters of the CD-NLGSSM emission model. The tuple doubles as a container for the ParameterProperties.

We assume a Gaussian observation model \(p(y_k | z_k) = N(y_k | h(z_k), R)\) where h is the emission function and R the observation noise covariance.

Parameters:

Name	Type	Description	Default
emission_function		Emission function h.	required
emission_cov		Observation noise covariance R.	required

create_cdnlgssm_params_and_props(params: dict) -> Tuple[ParamsCDNLGSSM, ParameterProperties]

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

Parameters:

Name	Type	Description	Default
params	dict	Dictionary of parameters.	required

Returns:

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

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

Initialize CD-NLGSSM 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. Defaults to jr.PRNGKey(0).	PRNGKey(0)

Returns:

Type	Description
Tuple[ParamsCDNLGSSM, ParamsCDNLGSSM]	Tuple of CD-NLGSSM parameters and properties objects.

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

Sample CD-NLGSSM 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[ParamsCDNLGSSM, ParamsCDNLGSSM]	Tuple of CD-NLGSSM parameters and properties objects.

update_params(params, updates: dict)

Update parameters of a CD-NLGSSM. Returns a copy of params with all updates applied.

updates: dict with keys like "initial.mean.params" or "dynamics.drift.sigma"

Example usage

updates = { "initial.mean.params": jnp.ones(3), # Vector of ones "dynamics.drift.sigma": 11.3 # Scalar value } new_params = update_params(params, updates)