A generic library for linear and non-linear Gaussian smoothing problems. The code leverages JAX and implements several linearization algorithms, both in a sequential and parallel fashion, as well as low-memory cost algorithms computing gradients of required quantities (such as the pseudo-loglikelihood of the system).
This code was written by Adrien Corenflos and Fatemeh Yaghoobi as a companion code for the article "Parallel square-root statistical linear regression for inference in nonlinear state space models" by Fatemeh Yaghoobi, Adrien Corenflos, Sakira Hassan, and Simo Särkkä, ArXiv link: https://arxiv.org/abs/2207.00426
- Create a virtual environment and clone this repository
- Install JAX (preferably with GPU support) following https://github.com/google/jax#installation
- Run
pip install . - (optional) If you want to run the examples, run
pip install -r requirements-examples.txt
Example uses (reproducing the experiments of our paper) can be found in the examples folder. More low-level examples can be found in the test folder.
If you find this work useful, please cite us in the following way:
@misc{yaghoobi2022sqrt,
Author = {Fatemeh Yaghoobi, Adrien Corenflos, Sakira Hassan, Simo S\"arkk\"a},
Title = {Parallel square-root statistical linear regression for inference in nonlinear state space models},
Year = {2022},
Eprint = {arXiv:2207.00426},
}