probml / dynamax

State Space Models library in JAX
https://probml.github.io/dynamax/
MIT License
661 stars 75 forks source link

Implement UKF #63

Closed petergchang closed 2 years ago

petergchang commented 2 years ago

Implement UKF, UK smoother, unit tests, pendulum demo.

AdrienCorenflos commented 2 years ago

Hi,

I see that you have been implementing UKF, EKF etc as fully different methods. I have done some work of coding this in a more general fashion here: https://github.com/EEA-sensors/sqrt-parallel-smoothers/tree/main/parsmooth/linearization

It follows the "generalised statistical linear regression" formulation. See this for example: https://scholar.google.fi/citations?view_op=view_citation&hl=sv&user=q0rtB0EAAAAJ&citation_for_view=q0rtB0EAAAAJ:IjCSPb-OGe4C

Doing it in a similar fashion would IMO reduce boilerplate and allow for more precise testing than with fully independent implementations. WDYT?

The idea is fairly simple: if you have a model $xt \sim p(\cdot \mid x{t-1})$ for which you know the conditional first and second moments: $\mathbb{E}[xt \mid x{t-1}]$ and $\mathbb{C}[xt \mid x{t-1}]$, for example, if the transformation has additive noise, then you can form an approximate Gaussian model $xt \approx N(\cdot; F{t-1} x{t-1} + b{t-1}, Q_{t-1})$ by linearising the way you want, and then apply the Kalman equations.

petergchang commented 2 years ago

Hi Adrien! Thank you for your comment,

That's a great idea; we've an nlgssm folder that contains models and tests common to all non-linear Gaussian ssms, which has been ever-growing precisely because of the commonality that you mentioned.

Factoring out all but the linearization method is a great idea, @murphyk WDYT?

slinderman commented 2 years ago

I like the simplification, both in code and in concepts! Great suggestion!

murphyk commented 2 years ago

Hi @AdrienCorenflos . I totally agree. I have actually rewritten my book chapter on SSM inference to use this more abstract formulation, and I think we should refactor the code to follow suit, since it is much more modular and elegant.

Here is the description of Kalman filtering in my new notation:

kf

Here is the corresponding version for EKF

ekf

and UKF

ukf

and finally the method you mentioned, which works with exponential family likelihoods using generalized statistical linear regression:

giplf

I will upload a new version of my full book shortly, which has all the details ;)

murphyk commented 2 years ago

I just uploaded the latest version of the book to https://github.com/probml/pml2-book/releases/tag/2022-07-13. Please see sec 8.3-8.8 - I would love feedback.

slinderman commented 2 years ago

I think we can close this now that we've merged Peter's code.