JSL is a JAX library for Bayesian inference in state space models. We support discrete state spaces (HMM) and continuous state spaces. For HMMs, we support exact inference (forwards-backwards and Viterbi) and MLE learning (using EM or SGD). For continuous SSMs, we assume that the model parameters are known, and just focus on state estimation. We support linear and non-linear dynamics/ observations; we assume all noise is Gaussian. For linear dynamical systems (LDS), we support Kalman filtering and RTS smoothing. For nonlinear dynamical systems (NLDS), we support extended Kalman filtering (EKF) with full and diagonal covariance, unscented Kalman filtering (UKF) and particle filtering (PF).

We assume you have already installed JAX and Tensorflow, since the details on how to do this depend on whether you have a CPU, GPU, etc. (This step is not necessary in Colab.)

Now install these packages:

!pip install --upgrade git+https://github.com/google/flax.git
!pip install --upgrade tensorflow-probability
!pip install git+https://github.com/blackjax-devs/blackjax.git
!pip install git+https://github.com/deepmind/distrax.git

Then install JSL:

!pip install git+https://github.com/probml/jsl

Alternatively, you can clone the repo locally, into say ~/github/JSL, and then install it as a package, as follows:

!git clone https://github.com/probml/JSL.git
cd JSL
!pip install -e .

You can see how to use the library by looking at some of the demos. You can run the demos from inside a notebook like this

%run JSL/jsl/demos/kf_tracking.py
%run JSL/jsl/demos/hmm_casino_em_train.py

Or from inside an ipython shell like this

from jsl.demos import kf_tracking
figdict = kf_tracking.main()

Most of the demos create figures. If you want to save them (in both png and pdf format), you need to specify the FIGDIR environment variable, like this:

import os
os.environ["FIGDIR"]='/Users/kpmurphy/figures'

from jsl.demos.plot_utils import savefig
savefig(figdict)

Gerardo Durán-Martín (@gerdm), Aleyna Kara(@karalleyna), Kevin Murphy (@murphyk).