Extended Kalman Filter

Extended Kalman filter inference for continuous-discrete nonlinear Gaussian SSMs.

inference_ekf

EKFHyperParams

Bases: NamedTuple

Lightweight container for EKF hyperparameters.

emissions_extended_kalman_filter(params: ParamsCDNLGSSM, t_states: Float[Array, 'num_timesteps 1'], state_means: Float[Array, 'num_timesteps state_dim'], state_covs: Optional[Float[Array, 'num_timesteps state_dim state_dim']] = None, inputs: Optional[Float[Array, 'num_timesteps input_dim']] = None, filter_hyperparams: EKFHyperParams = EKFHyperParams(), warn: bool = True) -> Tuple[Float[Array, 'num_timesteps emission_dim'], Optional[Float[Array, 'num_timesteps emission_dim emission_dim']]]

Compute the emissions corresponding to the EKF linearization of the CD-NLGSSM model.

Parameters:

Name	Type	Description	Default
params	ParamsCDNLGSSM	CD-NLGSSM parameters, containing - dynamics RHS drift function - diffusion coefficient matrix L - Brownian covariance matrix Q	required
t_states	Float[Array, 'num_timesteps 1']	continuous-time specific time instants of states	required
state_means	Float[Array, 'num_timesteps state_dim']	state means at time instants t_states, always required	required
state_covs	Optional[Float[Array, 'num_timesteps state_dim state_dim']]	state covariances at time instants t_states, optional - if None, then we assume that the states are point estimates, and simply push through emission function	None
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	EKFHyperParams	hyper-parameters of the filter	EKFHyperParams()
warn	bool	whether to issue warnings (e.g., about PSD issues)	True

Returns:

Name	Type	Description
emissions_mean	Float[Array, 'num_timesteps emission_dim']	mean of emissions
emissions_covariance	Optional[Float[Array, 'num_timesteps emission_dim emission_dim']]	covariance of emissions, if available

extended_kalman_filter(params: ParamsCDNLGSSM, emissions: Float[Array, 'ntime emission_dim'], t_emissions: Optional[Float[Array, 'num_timesteps 1']] = None, filter_hyperparams: EKFHyperParams = EKFHyperParams(), inputs: Optional[Float[Array, 'ntime input_dim']] = None, num_iter: int = 1, output_fields: Optional[List[str]] = ['filtered_means', 'filtered_covariances', 'predicted_means', 'predicted_covariances'], warn: bool = True) -> PosteriorGSSMFiltered

Run an (iterated) extended Kalman filter to produce the marginal likelihood and filtered state estimates.

Two implementations are available,
    based on first- and second-order approximations
        i.e. Algorithms 3.21 and 3.22 in Sarkka's thesis

Parameters:

Name	Type	Description	Default
params	ParamsCDNLGSSM	CD-NLGSSM parameters, containing - dynamics RHS drift function - diffusion coefficient matrix L - Brownian covariance matrix Q	required
emissions	Float[Array, 'ntime emission_dim']	observation sequence.	required
t_emissions	Optional[Float[Array, 'num_timesteps 1']]	continuous-time specific time instants of observations: if not None, it is an array	None
filter_hyperparams	EKFHyperParams	hyper-parameters of the EKF, related to the approximation order	EKFHyperParams()
inputs	Optional[Float[Array, 'ntime input_dim']]	optional array of inputs.	None
num_iter	int	number of linearizations around posterior for update step (default 1).	1
output_fields	Optional[List[str]]	list of fields to return in posterior object. These can take the values "filtered_means", "filtered_covariances", "predicted_means", "predicted_covariances", and "marginal_loglik".	['filtered_means', 'filtered_covariances', 'predicted_means', 'predicted_covariances']
warn	bool	whether to issue warnings during filtering (e.g., PSD issues).	True

Returns:

Name	Type	Description
filtered_posterior	PosteriorGSSMFiltered	posterior object.

extended_kalman_posterior_sample(key: PRNGKey, params: ParamsCDNLGSSM, emissions: Float[Array, 'ntime emission_dim'], filter_hyperparams: EKFHyperParams = EKFHyperParams(), t_emissions: Optional[Float[Array, 'num_timesteps 1']] = None, inputs: Optional[Float[Array, 'ntime input_dim']] = None, warn: bool = True) -> Float[Array, 'ntime state_dim']

Run forward-filtering, backward-sampling to draw samples.

Parameters:

Name	Type	Description	Default
key	PRNGKey	random number key.	required
params	ParamsCDNLGSSM	CD-NLGSSM parameters, containing - dynamics RHS drift function - diffusion coefficient matrix L - Brownian covariance matrix Q	required
emissions	Float[Array, 'ntime emission_dim']	observation sequence.	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, 'ntime input_dim']]	optional array of inputs	None
warn	bool	whether to issue warnings during filtering (e.g., PSD issues).	True

Returns:

Type	Description
Float[Array, 'ntime state_dim']	Float[Array, "ntime state_dim"]: one sample of \(z_{1:T}\) from the posterior distribution on latent states.

extended_kalman_smoother(params: ParamsCDNLGSSM, emissions: Float[Array, 'ntime emission_dim'], filter_hyperparams: EKFHyperParams = EKFHyperParams(), t_emissions: Optional[Float[Array, 'num_timesteps 1']] = None, filtered_posterior: Optional[PosteriorGSSMFiltered] = None, inputs: Optional[Float[Array, 'ntime input_dim']] = None, warn: bool = True) -> PosteriorGSSMSmoothed

Run forward-filtering, backward-smoother based on extended Kalman equations, to produce the smoothed state estimates. as described in Algorithm 3.23 in Sarkka's thesis

Parameters:

Name	Type	Description	Default
params	ParamsCDNLGSSM	CD-NLGSSM parameters, containing - dynamics RHS drift function - diffusion coefficient matrix L - Brownian covariance matrix Q	required
emissions	Float[Array, 'ntime emission_dim']	observation sequence.	required
filter_hyperparams	EKFHyperParams	hyper-parameters of the EKF, related to the approximation order	EKFHyperParams()
t_emissions	Optional[Float[Array, 'num_timesteps 1']]	continuous-time specific time instants of observations: if not None, it is an array	None
filtered_posterior	Optional[PosteriorGSSMFiltered]	optional output from filtering step.	None
inputs	Optional[Float[Array, 'ntime input_dim']]	optional array of inputs.	None
warn	bool	whether to issue warnings during smoothing (e.g., PSD issues).	True

Returns:

Name	Type	Description
post	PosteriorGSSMSmoothed	posterior object.

forecast_extended_kalman_filter(params: ParamsCDNLGSSM, init_forecast: tfd.Distribution, t_init: Float[Array, '1 1'], t_forecast: Float[Array, 'num_timesteps 1'], filter_hyperparams: EKFHyperParams = EKFHyperParams(), inputs: Optional[Float[Array, 'ntime input_dim']] = None, output_fields: Optional[List[str]] = ['forecasted_state_means', 'forecasted_state_covariances'], warn: bool = True) -> GSSMForecast

Run an extended Kalman filter to forecast states Two implementations are available, based on first- and second-order approximations i.e. Algorithms 3.21 and 3.22 in Sarkka's thesis

Parameters:

Name	Type	Description	Default
params	ParamsCDNLGSSM	CD-NLGSSM parameters, containing - dynamics RHS drift function - diffusion coefficient matrix L - Brownian covariance matrix Q	required
init_forecast	Distribution	initial distribution to forecast with.	required
t_init	Float[Array, '1 1']	time-instant of the initial condition of forecast	required
t_forecast	Float[Array, 'num_timesteps 1']	continuous-time specific time instants to forecast	required
filter_hyperparams	EKFHyperParams	hyper-parameters of the EKF, related to the approximation order	EKFHyperParams()
inputs	Optional[Float[Array, 'ntime input_dim']]	optional array of inputs.	None
output_fields	Optional[List[str]]	list of fields to return	['forecasted_state_means', 'forecasted_state_covariances']
warn	bool	whether to issue warnings (e.g., about PSD issues)	True

Returns:

Name	Type	Description
forecast	GSSMForecast	forecast object.

iterated_extended_kalman_filter(params: ParamsCDNLGSSM, emissions: Float[Array, 'ntime emission_dim'], t_emissions: Optional[Float[Array, 'num_timesteps 1']] = None, filter_hyperparams: EKFHyperParams = EKFHyperParams(), inputs: Optional[Float[Array, 'ntime input_dim']] = None, num_iter: int = 2, output_fields: Optional[List[str]] = ['filtered_means', 'filtered_covariances', 'predicted_means', 'predicted_covariances'], warn: bool = True) -> PosteriorGSSMFiltered

Run an iterated extended Kalman filter to produce the marginal likelihood and filtered state estimates.

It simply calls the extended_kalman_filter multiple times,
re-linearizing around the posterior from the previous iteration.

Parameters:

Name	Type	Description	Default
params	ParamsCDNLGSSM	CD-NLGSSM parameters, containing - dynamics RHS drift function - diffusion coefficient matrix L - Brownian covariance matrix Q	required
emissions	Float[Array, 'ntime emission_dim']	observation sequence.	required
t_emissions	Optional[Float[Array, 'num_timesteps 1']]	continuous-time specific time instants of observations: if not None, it is an array	None
filter_hyperparams	EKFHyperParams	hyper-parameters of the EKF, related to the approximation order	EKFHyperParams()
inputs	Optional[Float[Array, 'ntime input_dim']]	optional array of inputs.	None
num_iter	int	number of linearizations around posterior for update step (default 2).	2
output_fields	Optional[List[str]]	list of fields to return in posterior object.	['filtered_means', 'filtered_covariances', 'predicted_means', 'predicted_covariances']
warn	bool	whether to issue warnings during filtering (e.g., PSD issues).	True

Returns: post: posterior object.

iterated_extended_kalman_smoother(params: ParamsCDNLGSSM, emissions: Float[Array, 'ntime emission_dim'], filter_hyperparams: EKFHyperParams = EKFHyperParams(), t_emissions: Optional[Float[Array, 'num_timesteps 1']] = None, num_iter: int = 2, inputs: Optional[Float[Array, 'ntime input_dim']] = None, warn: bool = True) -> PosteriorGSSMSmoothed

Run an iterated extended Kalman smoother (IEKS).

Parameters:

Name	Type	Description	Default
params	ParamsCDNLGSSM	CD-NLGSSM parameters, containing - dynamics RHS drift function - diffusion coefficient matrix L - Brownian covariance matrix Q	required
emissions	Float[Array, 'ntime emission_dim']	observation sequence.	required
filter_hyperparams	EKFHyperParams	EKFHyperParams = EKFHyperParams(),	EKFHyperParams()
t_emissions	Optional[Float[Array, 'num_timesteps 1']]	continuous-time specific time instants of observations: if not None, it is an array	None
num_iter	int	number of linearizations around posterior for update step (default 2).	2
inputs	Optional[Float[Array, 'ntime input_dim']]	optional array of inputs.	None
warn	bool	whether to issue warnings during smoothing (e.g., PSD issues).	True

Returns:

Name	Type	Description
post	PosteriorGSSMSmoothed	posterior object.