Implementation of rao-blackwellised particle filter (rbpf) and rbpf with optimal resampling.

[kostastsa]

• Addresses implement rao-blackwellised particle filtering using dynamax and blackjax #271.
• To implement this I have created a new model class SLDS based on SSM. I have put everything in a folder dynamax/slds.
• The filters are part of a inference.py file. Other inference methods for SLDS, (such as generalized pseudo Bayes etc) can be included in the same file.
• Also, methods for learning SLDS can be build on this model.
• I have included a file inference_test.py which compares the new implementation to an older one referenced in implement rao-blackwellised particle filtering using dynamax and blackjax #271 which has been included in dynamax/slds/mixture_kalman_filter.py. Since the algorithm is stochastic and there is no exact baseline I wasn't sure how to test it, so I ended up doing an allclose with a large tolerance between the two implementations.
• From the plots and MSE it is seen that new implementation does what it is supposed to do.

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[kostastsa]

Hi @murphyk @slinderman , I was wondering if one would kindly take a look at this PR implementing the RBPF. Would you let me know if there is something that needs to be changed? I understand that implementing a SLDS model class goes beyond what was raised in issue #271 and might not be necessary, in which case I'd be happy to fix this and make a more minimal implementation.

[murphyk]

Maybe use multinomial resampling?

[murphyk]

For the resampling step, maybe it would be worth making a jax version of https://github.com/nchopin/particles/blob/master/particles/resampling.py#L541?

[murphyk]

LGTM. We can modify resampling code later (if necessary).

[slinderman]

Thanks, @kostastsa. One of the big items on my wishlist is to have a solid implementation of SLDS model variants and inference algorithms in dynamax. This is a great start.

[kostastsa]

@murphyk Cool, I will check the multinomial resampling from Chopin and do a jax version.
@slinderman That's great to hear! I would be super interested in being part of this effort. I can start working on this soon, since I will also need it for my own research.