Skip to content

Utilities

Utility classes and functions for continuous-discrete nonlinear SSMs.

cdnlssm_utils

LearnableDistribution

Bases: NamedTuple

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

Compute the log probability of the distribution

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

Sample from the distribution

LearnablePoissonEmission

Bases: NamedTuple

Learnable Poisson emission distribution for count data.

The Poisson distribution is defined by the following probability mass function

\(P(Y=y) = \frac{\lambda^y e^{-\lambda}}{y!}\)

where \(\lambda\) is the rate parameter of the distribution.

ParamsCDNLSSM

Bases: NamedTuple

Parameters of a continuous-discrete nonlinear SSM.

:param initial: initial distribution parameters :param dynamics: dynamics distribution parameters :param emissions: emission distribution parameters

The assumed transition and emission distributions are $\(p(z_0) = p_initial(z_0)\)$ $\(p(z_{t_k} | z_{t_{k-1}}, u_{t_k}) = solve_sde(z_{t_{k-1}}, u_{t_k}, t_{k-1}, t_k, f_dynamics, L_dynamics, Q_dynamics)\)$ $\(p(y_{t_k} | z_{t_k}) = p_emissions(y_{t_k} | z_{t_k})\)$

create_cdnlssm_params_and_props(params: dict) -> Tuple[ParamsCDNLSSM, ParameterProperties]

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

Parameters:

Name Type Description Default

param params: dictionary of parameters

required

Returns:

Type Description
Tuple[ParamsCDNLSSM, ParameterProperties]

return: Tuple of parameters and properties objects

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

Create CD-NLSSM parameters and properties using provided prior, init_values or defaults.

Parameters:

Name Type Description Default
default_params

dictionary of default parameters: we at least need some default values

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[ParamsCDNLSSM, ParamsCDNLSSM]

Tuple of CD-NLSSM parameters and properties objects

sample_cdnlssm_params(prior, M, init_params, key=jr.PRNGKey(0)) -> ParamsCDNLSSM

Sample CD-NLSSM parameters from the provided prior. Args: :param prior: prior distribution for the initialization. :param M: number of samples to draw :param init_params: dictionary of all parameters :param key: random key for sampling from the prior. Defaults to jr.PRNGKey(0).

Returns:

Type Description
ParamsCDNLSSM

return: Tuple of CD-NLSSM parameters and properties objects

update_params(params, updates: dict)

Returns a copy of params with all updates applied.