Skip to content

Companion code in JAX for the paper Parallel Iterated Extended and Sigma-Point Kalman Smoothers.

License

Notifications You must be signed in to change notification settings

EEA-sensors/parallel-non-linear-gaussian-smoothers

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

73 Commits
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

Parallel Iterated Extended and Sigma-Point Kalman Smoothers

Companion code in JAX for the paper Parallel Iterated Extended and Sigma-Point Kalman Smoothers [2].

What is it?

This is an implementation of parallelized Extended and Sigma-Points Bayesian Filters and Smoothers with CPU/GPU/TPU support coded using JAX primitives, in particular associative scan.

Supported features

  • Extended Kalman Filtering and Smoothing
  • Cubature Kalman Filtering and Smoothing
  • Iterated versions of the above

Installation

Manually install the dependencies jax and jaxlib, and for examples only matplotlib, numba, tqdm

Example

from parsmooth.parallel import ieks
from parsmooth.utils import MVNormalParameters

initial_guess = MVNormalParameters(...)
data = ...
Q = ...  # transition noise covariance matrix
R = ...  # observation error covariance matrix

def transition_function(x):
  ...
  return next_x
  
def observation_function(x):
  ...
  return obs
  
iterated_smoothed_trajectories = ieks(initial_guess, 
                                      data, 
                                      transition_function, 
                                      Q, 
                                      observation_function, 
                                      R, 
                                      n_iter=100)  # runs the parallel IEKS 100 times.

For more examples, see the notebooks.

Acknowlegments

This JAX-based code was created by Adrien Corenflos to implement the original idea by Fatemeh Yaghoobi [2] who provided the initial code for the parallelized extended Kalman filter in pure Python. The sequential cubature filtering code was adapted from some original code by Zheng Zhao.

References

[1] S. Särkkä and A. F. García-Fernández. Temporal Parallelization of Bayesian Smoothers. In: IEEE Transactions on Automatic Control 2020.

[2] F. Yaghoobi and A. Corenflos and S. Hassan and S. Särkkä. Parallel Iterated Extended and Sigma-Points Kalman Smoothers. To appear in Proceedings of IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP). (arXiv)