Models

Continuous-discrete nonlinear SSM model definitions.

models

ContDiscreteNonlinearSSM

Bases: SSM

Continuous-discrete nonlinear SSM with generic (possibly non-Gaussian) initial and emission distributions.

We assume a model of the form

\[ dz=f(z,u_t,t)dt + L(z, u, t) db(t)\]

with diffusion covariance \(Q_c\).

We allow for arbitrary initial and emission distributions, $\(p(z_0) = p(z_0)\)$ $\(p(y_{t_k} | z_{t_k}) = p(y_{t_k} | z_{t_k})\)$

where the model parameters are * \(z_t\) = hidden variables of size state_dim,

• \(f\) = dynamics deterministic function (RHS), used to compute transition function
• \(L\) = dynamics coefficient multiplying brownian motion
• \(Q\) = dynamics brownian motion's covariance (system) noise

• \(u_t\) = input covariates of size input_dim (defaults to 0).

• \(y_t\) = observed variables of size emission_dim

filter(params: ParamsCDNLSSM, emissions: Array, t_emissions: Optional[Array] = None, inputs: Optional[Array] = None, filter_type: str = 'DPF', filter_state_order: str = 'first', filter_state_cov_rescaling: float = 1.0, filter_dt_average: float = 0.1, N_particles: int = 1000, diffeqsolve_max_steps: int = 100, diffeqsolve_dt0: float = 0.01, output_fields=None, key: PRNGKeyArray = jr.PRNGKey(0), diffeqsolve_kwargs: Optional[dict] = None, extra_filter_kwargs: Optional[dict] = None, warn: bool = True)

Filters a CD-NLSSM; by default, this runs a bootstrap differentiable particle filter (DPF).

Depending on the filter_type, certain arguments are ignored.

Parameters:

Name	Type	Description	Default
params	ParamsCDNLSSM	Parameters of the CDNLSSM.	required
emissions	Array	Emission sequence.	required
t_emissions	Optional[Array]	Time instants of observations.	None
inputs	Optional[Array]	Inputs.	None
filter_state_order	str	Order of Taylor expansion for dynamics used in the filter.	'first'
filter_emission_order		Order of Taylor expansion for emissions used in the filter.	required
filter_num_iter		Number of iterations for iterated filters.	required
filter_state_cov_rescaling	float	Rescale state covariance by this factor after each update (inflation delta is better for accurate likelihoods)	1.0
filter_dt_average	float	[Only for state_order="Discrete"] Average step size to determine constant state noise cov in filter.	0.1
N_particles	int	Number of particles (for DPF only).	1000
diffeqsolve_max_steps	int	Max steps for ODE solver between observations.	100
diffeqsolve_dt0	float	Initial step size for ODE/SDE solver (default is fixed step size).	0.01
output_fields		Which fields to return from the filter.	None
key	PRNGKeyArray	Random key.	PRNGKey(0)
diffeqsolve_kwargs	Optional[dict]	Extra kwargs for the ODE solver (e.g., {"solver": diffrax.Heun(), "dt0": 1e-2}).	None
filter_kwargs		Extra kwargs specific to the chosen filter (e.g., {"emission_order": "zeroth"} for EKF).	required
warn	bool	whether to issue warnings (e.g., about PSD issues)	True

Returns:

Name	Type	Description
PosteriorCDNLSSMFiltered		Posterior distribution of the CDNLSSM.

filter_and_forecast(params: ParamsCDNLSSM, emissions_filter: Array, t_emissions_filter: Array, t_emissions_forecast: Array, inputs_filter: Optional[Array] = None, inputs_forecast: Optional[Array] = None, filter_type: str = 'DPF', filter_state_order: str = 'first', filter_state_cov_rescaling: float = 1.0, filter_dt_average: float = 0.1, N_particles: int = 1000, diffeqsolve_max_steps: int = 100, diffeqsolve_dt0: float = 0.01, output_fields=None, key: PRNGKeyArray = jr.PRNGKey(0), diffeqsolve_kwargs: Optional[dict] = None, extra_filter_kwargs: Optional[dict] = None, warn: bool = True)

Filters a CD-NLSSM; by default, this runs a bootstrap differentiable particle filter (DPF).

Depending on the filter_type, certain arguments are ignored.

Parameters:

Name	Type	Description	Default
params	ParamsCDNLSSM	Parameters of the CDNLSSM.	required
emissions_filter	Array	Emission sequence for filtering.	required
t_emissions_filter	Array	Time instants of observations for filtering.	required
t_emissions_forecast	Array	Time instants for forecasting.	required
inputs_filter	Optional[Array]	Inputs for filtering.	None
inputs_forecast	Optional[Array]	Inputs for forecasting.	None
filter_state_order	str	Order of Taylor expansion for dynamics used in the filter.	'first'
filter_state_cov_rescaling	float	Rescale state covariance by this factor after each update (inflation delta is better for accurate likelihoods)	1.0
filter_dt_average	float	[Only for state_order="Discrete"] Average step size to determine constant state noise cov in filter.	0.1
N_particles	int	Number of particles (for DPF only).	1000
diffeqsolve_max_steps	int	Max steps for ODE solver between observations.	100
diffeqsolve_dt0	float	Initial step size for ODE/SDE solver (default is fixed step size).	0.01
output_fields		Which fields to return from the filter.	None
key	PRNGKeyArray	Random key.	PRNGKey(0)
diffeqsolve_kwargs	Optional[dict]	Extra kwargs for the ODE solver (e.g., {"solver": diffrax.Heun(), "dt0": 1e-2}).	None
filter_kwargs		Extra kwargs specific to the chosen filter (e.g., {"emission_order": "zeroth"} for EKF).	required
warn	bool	whether to issue warnings (e.g., about PSD issues)	True

Returns:

Name	Type	Description
filtered_posterior		PosteriorCDNLSSMFiltered, posterior distribution of the CDNLSSM.
forecasted		Float[Array, "num_timesteps state_dim M"], forecasted states over time.

initialize(key: Float[Array, 2] = jr.PRNGKey(0), init_prior: Prior = None, initial_distribution: dict = None, dynamics_drift: dict = None, dynamics_diffusion_coefficient: dict = None, dynamics_diffusion_cov: dict = None, dynamics_approx_order: Optional[float] = 1.0, emission_distribution: dict = None) -> Tuple[ParamsCDNLSSM, ParamsCDNLSSM]

Initialize the model parameters.

Parameters:

Name	Type	Description	Default
key	Float[Array, 2]	Random key.	PRNGKey(0)
init_prior	Prior	Prior distribution.	None
initial_distribution	dict	Initial distribution.	None
dynamics_drift	dict	Dynamics drift.	None
dynamics_diffusion_coefficient	dict	Dynamics diffusion coefficient.	None
dynamics_diffusion_cov	dict	Dynamics diffusion covariance.	None
dynamics_approx_order	Optional[float]	Dynamics approximation order.	1.0
emission_distribution	dict	Emission distribution.	None

rns

Tuple[ParamsCDNLSSM, ParamsCDNLSSM]: Parameters and their properties.

marginal_log_prob(params: ParamsCDNLSSM, emissions: Float[Array, 'ntime emission_dim'], t_emissions: Optional[Float[Array, 'ntime 1']] = None, filter_hyperparams: Optional[DPFHyperParams] = DPFHyperParams(), 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-NLSSM model, Args: params: CD-NLSSM 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 Differentiable Particle Filter (DPF) algorithm. inputs: Optional input sequence. key: Random number generator key. Returns: Marginal log-likelihood of the emissions, \(\log p(y_{1:T})\).

sample_dist(params: ParamsCDNLSSM, key: PRNGKeyArray, num_timesteps: int, t_emissions: Optional[Array] = None, inputs: Optional[Array] = None)

Sample from the joint distribution to produce state and emission trajectories.

Parameters:

Name	Type	Description	Default
params	ParamsCDNLSSM	Parameters of the CDNLSSM.	required
key	PRNGKeyArray	Random key.	required
num_timesteps	int	Number of timesteps.	required
t_emissions	Optional[Array]	Time instants of observations.	None
inputs	Optional[Array]	Inputs.	None

Returns:

Type	Description
	Tuple[Array, Array]: States and emissions.

sample_path(params: ParamsCDNLSSM, key: PRNGKeyArray, num_timesteps: int, t_emissions: Optional[Array] = None, inputs: Optional[Array] = None)

Sample states and emissions by integrating the SDE and drawing from the emission distribution.

Parameters:

Name	Type	Description	Default
params	ParamsCDNLSSM	Parameters of the CDNLSSM.	required
key	PRNGKeyArray	Random key.	required
num_timesteps	int	Number of timesteps.	required
t_emissions	Optional[Array]	Time instants of observations.	None
inputs	Optional[Array]	Inputs.	None

Returns:

Type	Description
	Tuple[Array, Array]: States and emissions.

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

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

:param prior: prior distribution. :param M: number of samples to draw. :param init_params: dictionary of parameters to use as initialization if not provided, default parameters are used :param key: random number generator key

Returns:

Type	Description
Tuple[ParamsCDNLSSM, ParamsCDNLSSM]	Tuple[ParamsCDNLSSM, ParamsCDNLSSM]: Parameters and their properties.

cdnlssm_emissions(params: ParamsCDNLSSM, t_states: Float[Array, 'num_timesteps 1'], states: Float[Array, 'num_timesteps state_dim M'], inputs: Optional[Float[Array, 'num_timesteps input_dim']] = None, filter_hyperparams: Optional[DPFHyperParams] = DPFHyperParams(), key: PRNGKey = jr.PRNGKey(0), warn: bool = True) -> Float[Array, 'num_timesteps emission_dim M']

Compute the emissions corresponding to - a continuous-discrete nonlinear model, as specified by params

Parameters:

Name	Type	Description	Default
params	ParamsCDNLSSM	model parameters.	required
t_states	Float[Array, 'num_timesteps 1']	continuous-time specific time instants of states	required
states	Float[Array, 'num_timesteps state_dim M']	states at time instants t_states, always required it may handle multiple particles, in which case states is of shape num_timesteps \times state_dim \times M	required
inputs	Optional[Float[Array, 'num_timesteps input_dim']]	optional array of inputs, of shape (1 + num_timesteps) \times input_dim - The extra input is needed for the initial emission, i.e., it should be at time t_init	None
filter_hyperparams	Optional[DPFHyperParams]	hyper-parameters of the filter, optional, actually ignored	DPFHyperParams()
key	PRNGKey	random key for sampling	PRNGKey(0)

Returns:

Name	Type	Description
emissions	Float[Array, 'num_timesteps emission_dim M']	emissions at time instants t_states, of shape num_timesteps \times emission_dim \times M

cdnlssm_filter(params: ParamsCDNLSSM, emissions: Array, t_emissions: Optional[Array] = None, filter_hyperparams: Optional[DPFHyperParams] = None, inputs: Optional[Array] = None, output_fields: Optional[List[str]] = None, key: PRNGKeyArray = jr.PRNGKey(0), warn: bool = True)

Run particle filtering (configurable DPF) for a CD-NLSSM and return particles, log-weights, and log-evidence.

Parameters:

Name	Type	Description	Default
params	ParamsCDNLSSM	Parameters of the CDNLSSM.	required
emissions	Array	Emission sequence.	required
t_emissions	Optional[Array]	Time instants of observations.	None
filter_hyperparams	Optional[DPFHyperParams]	Hyperparameters of the filter.	None
inputs	Optional[Array]	Inputs.	None
output_fields	Optional[List[str]]	Fields to return.	None
key	PRNGKeyArray	Random key.	PRNGKey(0)
warn	bool	Whether to warn.	True

Returns:

Name	Type	Description
PosteriorCDNLSSMFiltered		Posterior distribution of the CDNLSSM.

cdnlssm_forecast(params: ParamsCDNLSSM, init_forecast: Float[Array, 'state_dim M'], t_init: Float[Array, '1 1'], t_forecast: Optional[Float[Array, 'num_timesteps 1']] = None, filter_hyperparams: Optional[DPFHyperParams] = DPFHyperParams(), inputs: Optional[Float[Array, 'ntime input_dim']] = None, key: PRNGKey = jr.PRNGKey(0), diffeqsolve_settings: dict = {}, warn: bool = True) -> Float[Array, 'num_timesteps state_dim M']

Run an continuous-discrete nonlinear model to produce the forecasted state estimates.

It supports two modes of forecasting:
1) Forecasting through nonlinear distributions, based on DPF: in this case, the initial condition of the forecast is a distribution (e.g., the filtering distribution at the last observation), and we forecast the evolution of such distribution based on DPF with no resampling.
2) Forecasting paths, based on solving the SDE: in this case, the initial condition of the forecast is a point estimate of state, and we

Parameters:

Name	Type	Description	Default
params	ParamsCDNLSSM	CD-NLSSM parameters.	required
init_forecast	Float[Array, 'state_dim M']	initial condition to start forecasting with, which we push forward starting at that state	required
t_init	Float[Array, '1 1']	time-instant of the initial condition of forecast	required
t_forecast	Optional[Float[Array, 'num_timesteps 1']]	continuous-time specific time instants of observations: if not None, it is an array	None
filter_hyperparams	Optional[DPFHyperParams]	hyper-parameters of the filter	DPFHyperParams()
inputs	Optional[Float[Array, 'ntime input_dim']]	optional array of inputs, of shape (1 + num_timesteps) imes input_dim - The extra input is needed for the initial emission, i.e., it should be at time t_init	None
key	PRNGKey	random key (e.g., for sampling).	PRNGKey(0)
diffeqsolve_settings	dict	settings for the SDE solver	{}
warn	bool	whether to issue warnings during forecasting (e.g., PSD issues).	True

Returns:

Name	Type	Description
post	Float[Array, 'num_timesteps state_dim M']	forecasted states over time of shape num_timesteps state_dim M.

cdnlssm_joint_sample(params: ParamsCDNLSSM, key: PRNGKeyArray, num_timesteps: int, t_emissions: Optional[Array] = None, inputs: Optional[Array] = None, diffeqsolve_settings: Optional[dict] = None)

Sample states and emissions jointly by integrating the SDE and drawing emissions.

compute_pushforward(x0: Array, P0: Array, params: ParamsCDNLSSM, t0: Float, t1: Float, inputs: Optional[Array] = None, diffeqsolve_settings: Optional[dict] = None) -> Tuple[Array, Array]

Compute the pushforward of particles through the CD-NLSSM dynamics.

Currently, as only Brownian motion-driven SDEs are supported, this simply calls the CDNLGSSM pushforward.

Returns:

Type	Description
Tuple[Array, Array]	Tuple[Array, Array]: Mean and covariance of the pushforward.