Skip to content

Add Differentiable Particle Filter Support#1

Merged
mattlevine22 merged 9 commits intohd-UQ:dev-numpyro-apifrom
DanWaxman:dw-dpf
Jan 23, 2026
Merged

Add Differentiable Particle Filter Support#1
mattlevine22 merged 9 commits intohd-UQ:dev-numpyro-apifrom
DanWaxman:dw-dpf

Conversation

@DanWaxman
Copy link
Contributor

@DanWaxman DanWaxman commented Jan 11, 2026

This PR adds differentiable particle filter support. These entail a general ContDiscreteNonlinearSSM class, which allows for non-Gaussian, non-linear observation processes. We include a differentiable particle filter using either "soft" or "stop gradient" resampling methods.

These methods are integrated into the Lorenz-63 example notebook (the new API one), as well as a new dpf_examples.ipynb notebook, which shows a few examples of non-Gaussian observations (an OU-driven count process with Poisson observations, and a tracking example with bearing (S^1) and signal power (R) observations.

Remaining to-dos:

  • Add more documentation, including some citations to the literature
  • Clean up code after documentation

Two things that could be done in a follow-up PR, IMO:

Unit tests will fail otherwise because they compare sample paths, and each sampler will use a different key for sampling the initial state. This enforces they derive keys from the same pattern, instead.
@DanWaxman DanWaxman marked this pull request as draft January 11, 2026 14:51
inputs: Optional[Array] = None,
diffeqsolve_settings: Optional[dict] = None,
) -> Tuple[Array, Array]:
"""Propagate mean/covariance using chosen approximation order."""
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Dispatch to CDNLGSSM case if using Brownian motions, give errors otherwise.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We don't really have a nice setup for this right now. I would propose a refactor that takes a lot of common utilities (like compute_pushforward) to SDEs.py, or something similar, though this seems out-of-scope for the current PR.

@DanWaxman DanWaxman marked this pull request as ready for review January 16, 2026 18:27
Copy link
Contributor

@mattlevine22 mattlevine22 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks great, thanks Dan!

@mattlevine22 mattlevine22 merged commit 31e558b into hd-UQ:dev-numpyro-api Jan 23, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants

Comments