A JAX library for state-space model inference (filtering, smoothing, static parameter estimation).

Disclaimer: The name cuthbert was chosen as a playful nod to the well-known caterpillar cake rivalry between Aldi and M&S in the UK, as the classic state-space model diagram looks vaguely like a caterpillar. However, this software project has no formal connection to Aldi, M&S, or any food products (notwithstanding the coffee drunk during its writeup). cuthbert is simply a fun name for this state-space model library and should not be interpreted as an endorsement, association, or affiliation with any brand or animal themed baked goods.

Goals

• Simple, flexible and performant interface for state-space model inference.
• Decoupling of model specification and inference. cuthbert is built to swap between different inference methods without be tied to a specific model specification.
• Compose with the JAX ecosystem for extensive external tools.
• Functional API: The only classes in cuthbert are NamedTuples and Protocols. All functions are pure and work seamlessly with jax.grad, jax.jit, jax.vmap etc.
• Methods for filtering: $p(x_t \mid y_{0:t}, \theta)$.
• Methods for smoothing: $p(x_{0:T} \mid y_{0:T}, \theta)$ or $p(x_{t} \mid y_{0:T}, \theta)$.
• Methods for static parameter estimation: $p(\theta \mid y_{0:T})$ or $\text{argmax} p(y_{0:T} \mid \theta)$.
• This includes support for forward-backward/Baum-Welch, particle filtering/sequential Monte Carlo, Kalman filtering (+ extended/unscented/ensemble), expectation-maximization and more!

Non-goals

• Tools for defining models and distributions. cuthbert is not a probabilistic programming language (PPL). But can easily compose with dynamax, distrax, numpyro and pymc in a similar way to how blackjax does.
• "SMC Samplers" which sample from a posterior distribution which is not (necessarily) a state-space model - blackjax is great for this.

Codebase structure

The codebase is structured as follows:

• cuthbert: The main package with unified interface for filtering and smoothing.
• cuthbertlib: A collection of atomic, smaller-scoped tools useful for state-space model inference, that represent the building blocks that power the main cuthbert package.

• docs: Source code for the documentation for cuthbert and cuthbertlib.
• tests: Tests for the cuthbert and cuthbertlib packages.

Installation

cuthbert depends on JAX, so you'll need to install JAX for the available hardware (CPU, GPU, or TPU). For example, on computers with NVIDIA GPUs:

pip install -U "jax[cuda13]"

Now install cuthbert from PyPI:

pip install -U cuthbert

Installing cuthbert will also install cuthbertlib.

Ecosystem

• cuthbert is built on top of jax and composes easily with other JAX packages, e.g. optax for optimization, flax for neural networks, and blackjax for (SG)MCMC as well as the PPLs mentioned above.
• What about dynamax?
  • dynamax is a great library for state-space model specification and inference with discrete or Gaussian state-space models. cuthbert is focused on inference with arbitrary state-space models via e.g. SMC that is not supported in dynamax. However as they are both built on jax they can be used together! A dynamax model can be passed to cuthbert for inference.
• And particles?
  • particles and the accompanying book Sequential Monte Carlo Methods in Practice are wonderful learning materials for state-space models and SMC. cuthbert is more focused on performance and composability with the JAX ecosystem.
• Much of the code in cuthbert is built on work from sqrt-parallel-smoothers, mocat and abile.

Contributing

We're always looking for contributions! Check out the contributing guide for more information.