Make CD filters return predicted observation distributions

[mattlevine22]

Summary

This PR adds predictive emission information to the continuous-discrete Gaussian filters so downstream code can evaluate both latent signal uncertainty and observation-level uncertainty from filter outputs.

Why? We want to evaluate filters using scoring rules that compare predicted observation distributions versus the observed data; this will allow us to not only evaluate, but also tune filters w.r.t. various metrics.

Specifically:

• KF, EKF, UKF, and EnKF now return predictive emission fields in posterior_extras
• filter docstrings were cleaned up around output_fields and predictive output semantics
• regression tests were added for the new outputs

External dependency

There is an open Dynestyx PR for filter-scoring that relies on this PR here.

What’s included

KF / EKF / UKF

These filters now expose both of the following in posterior_extras:

• y_pred_mean
• y_pred_cov
• y_obs_pred_mean
• y_obs_pred_cov

The intended distinction is:

• y_pred_*: predictive moments in observation space before adding observation noise
• y_obs_pred_*: predictive moments for the actual observed data

In other words:

• y_pred_mean = E[h(x_t) | y_{1:t-1}]
• y_pred_cov = Cov[h(x_t) | y_{1:t-1}]
• y_obs_pred_mean = y_pred_mean
• y_obs_pred_cov = y_pred_cov + R

This is useful because different downstream tasks want different objects:

• latent-signal analysis or uncertainty decomposition: use y_pred_* [OPTIONAL]
• scoring rules against observed data: use y_obs_pred_* [ESSENTIAL TO THIS PR]

EnKF

For EnKF, this PR newly exposes in posterior_extras:

• y_ens_pred: predicted observation ensemble

Unlike KF/EKF/UKF, EnKF does not return explicit y_pred_mean / y_pred_cov fields. Instead:

• noiseless predictive observation moments can be computed downstream from y_ens_pred
• the observation-level predictive covariance is already available as S

We now document S more clearly as the observation-predictive covariance used for assimilation and data scoring, i.e. the ensemble predictive covariance plus observation noise R.

This keeps EnKF lightweight while still exposing the important predictive information.

Why y_pred_* vs y_obs_pred_* matters

There’s an important modeling distinction here:

• y_pred_cov reflects uncertainty from the latent state propagated into observation space
• y_obs_pred_cov reflects uncertainty in the actual measured observation, so it includes R

If we want to evaluate the predictive distribution of the observed data, the correct object is y_obs_pred_cov, not y_pred_cov.

That distinction is now explicit in the API instead of being left implicit in downstream code.

Why DPF is not included

This PR does not extend DPF, as the predicted observations are not computed explicitly (hence, a bit more thought and work is needed for how to do this appropriately).

Testing

Added regression coverage for:

• presence and shapes of the new predictive fields
• consistency between y_pred_* and state predictive moments for KF/EKF/UKF
• consistency of y_obs_pred_cov = y_pred_cov + R
• output_fields=["posterior_extras"] behavior
• EnKF behavior via newly exposed y_ens_pred and documented S

[Copilot]

Pull request overview

This PR extends the continuous-discrete nonlinear Gaussian filtering APIs to expose predictive emission information via posterior_extras, enabling downstream evaluation of predicted observation distributions (e.g., via scoring rules).

Changes:

• Add per-step predictive emission outputs to EKF/UKF (y_pred_mean, y_pred_cov) and expose predicted observation ensembles in EnKF (y_ens_pred) via posterior_extras.
• Clarify output_fields / posterior_extras semantics in docstrings.
• Add regression tests covering posterior_extras and output_fields=["posterior_extras"] behavior.

Reviewed changes

Copilot reviewed 9 out of 9 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
tests/test_nonlinear_gaussian_filter_posterior_extras.py	Adds regression tests for new posterior_extras outputs and minimal-output mode.
cd_dynamax/src/continuous_discrete_nonlinear_gaussian_ssm/models.py	Updates docstrings to describe output_fields and posterior_extras usage.
cd_dynamax/src/continuous_discrete_nonlinear_gaussian_ssm/inference_ekf.py	Computes and returns EKF predictive emission moments in posterior_extras.
cd_dynamax/src/continuous_discrete_nonlinear_gaussian_ssm/inference_ukf.py	Refactors UKF conditioning to surface predictive emission moments in posterior_extras.
cd_dynamax/src/continuous_discrete_nonlinear_gaussian_ssm/inference_enkf.py	Exposes predicted observation ensemble (y_ens_pred) in EnKF posterior_extras.
cd_dynamax/dynamax/linear_gaussian_ssm/inference.py	Documents optional predicted-state fields and posterior_extras on PosteriorGSSMFiltered.

Comments suppressed due to low confidence (1)

cd_dynamax/src/continuous_discrete_nonlinear_gaussian_ssm/inference_ukf.py:466

• posterior_extras only returns y_pred_mean / y_pred_cov (noiseless predictive emission moments). For scoring rules against observed data, the observation-level predictive covariance (including R) is needed; consider also returning y_obs_pred_mean and y_obs_pred_cov = y_pred_cov + R so callers don’t have to recompute time-varying R downstream.

            filter_hyperparams,
            warn=warn,
        )

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[mattlevine22]

@copilot review again (new commit posted)

[Copilot]

@copilot review again (new commit posted)

Addressed all four review points (now in commits c0a800c and c0a800c's follow-up):

1. EKF/UKF y_obs_pred_* – both y_obs_pred_mean and y_obs_pred_cov = y_pred_cov + R are now included in posterior_extras.
2. Linear KF – posterior_extras is now populated with the same four fields (y_pred_mean, y_pred_cov, y_obs_pred_mean, y_obs_pred_cov).
3. Tests – regression tests now assert both y_pred_* and y_obs_pred_* fields, including the y_obs_pred_cov = y_pred_cov + R relation.
4. EnKF y_ens_pred vs S consistency – fixed in the latest commit (Fix EnKF y_ens_pred to use inflated ensemble, consistent with S): y_ens_pred now returns y_ensemble_infl instead of y_ensemble, making it consistent with S = y_pred_cov_infl + R. When inflation_delta == 0 (the default) the two are identical, so all 30 tests continue to pass.

[iurteaga]

param posterior_extras: stores per-step predictive emission, based on one-step-ahead state predictions z_{t+1}:
moments y_pred_mean & y_pred_cov: $\mathbb{E}[h(z_{t+1}) \mid y_{1:t}, u_{1:t}]$ & $\mathrm{Cov}[h(z_{t+1}) \mid y_{1:t}, u_{1:t}]$
moments y_obs_pred_mean & y_obs_pred_cov, where y_obs_pred_cov = y_pred_cov + R.

[iurteaga]

I would add the tuple "y_pred_mean", "y_pred_cov", "y_obs_pred_mean", "y_obs_pred_cov"

[iurteaga]

I left some comments and edit suggestions, mostly on having a consistent API for posterior extras.

Also, maybe worth revising the docstrings to specifically indicate that this predicted emissions are conditioned on predicted state z_{t+1}

[mattlevine22]

Thanks! Do you think these should all just become (optional) parts of the filtered posterior itself? This way a user can decide whether to return some/all of them. Currently, you have to take or leave the entire posterior_extras field, which I don't love.

[mattlevine22]

@iurteaga I updated things quite a bit to expand the PosteriorGSSMFiltered class to optionally include these predicted observation objects, and refreshed the documentation (and some error messages) to support these changes and a bit more clarity about the mathematical objects these variables represent (ensuring the user understands t vs t-1 in these structs)

[DanWaxman]

This all looks reasonable to me. I have two main comments:

1. Should we be also updated filtered dataclasses? E.g., we are often actually returning something like a PosteriorGSSMFiltered:
   

   cd_dynamax/cd_dynamax/dynamax/linear_gaussian_ssm/inference.py

   Lines 112 to 127 in df24b56

   	class PosteriorGSSMFiltered(NamedTuple):
   	r"""Marginals of the Gaussian filtering posterior.
   	:param marginal_loglik: marginal log likelihood, $p(y_{1:T} \mid u_{1:T})$
   	:param filtered_means: array of filtered means $\mathbb{E}[z_t \mid y_{1:t}, u_{1:t}]$
   	:param filtered_covariances: array of filtered covariances $\mathrm{Cov}[z_t \mid y_{1:t}, u_{1:t}]$
   	"""
   	# Default attributes
   	marginal_loglik: Union[Scalar, Float[Array, "ntime"]]
   	filtered_means: Optional[Float[Array, "ntime state_dim"]] = None
   	filtered_covariances: Optional[Float[Array, "ntime state_dim state_dim"]] = None
   	predicted_means: Optional[Float[Array, "ntime state_dim"]] = None
   	predicted_covariances: Optional[Float[Array, "ntime state_dim state_dim"]] = None
   	# Additional extras
   	posterior_extras: Optional[dict] = None

   
   This does not currently have all possible fields listed.

2. From a design perspective, I find it slightly unfortunate we're passing around a bunch of dictionaries with possibly missing keys or empty values. I think it would be beneficial to have class FilterOutput (or even pass around partial initializations of PosteriorGSSMFiltered). Though I acknowledge this may be better served as part of The Great Refactor.

[mattlevine22]

@DanWaxman

1. I agree---I think this is updated in the most recent commit? Looks like you link back to main

2. Yeah I think this is better to leave for The Great Refactor

[DanWaxman]

1. I agree---I think this is updated in the most recent commit? Looks like you link back to main

Wow, TIL that if you use the GitHub magic links feature it will just randomly throw you back to main when you try to follow a code reference... GitHub never fails to amaze me :D

2. Yeah I think this is better to leave for The Great Refactor

Fair enough -- I think this looks good me then!

[mattlevine22]

Perfect, made those changes and will merge!