smcx 2.0.0

State-space inference in JAX: Kalman and particle filters, tempered SMC, and SMC2 across CPU and accelerators

smcx

Sequential inference for state-space models (SSMs) in JAX: Kalman-family filters, particle filters, and sequential Monte Carlo (SMC) samplers. Following a theme of some modern probabilistic programming languages like NumPyro and other projects in the JAX ecosystem like dynestyx, I aimed to decouple the inference code from the model code. smcx goes one step further and has no modeling language of its own. Users define their models as plain JAX functions. This also lets them wrap models built in other SSM libraries like dynamax.

An introduction to the Kalman and SMC methods is developed in the documentation. Below is a quick start and a map of the methods.

pip install smcx

Quick start

A simple model to start with is a linear-Gaussian model

$$ \begin{aligned} y_t &\sim \mathcal{N}(\theta_t,\ 0.3), \ \theta_t &\sim \mathcal{N}(0.8,\theta_{t-1},\ 0.2), \qquad \theta_0 \sim \mathcal{N}(0, 1). \end{aligned} $$

In this case we can calculate the exact filtering distribution in closed form using the Kalman filter. The Kalman filter assumes the model is linear and Gaussian, so all you need to provide are the model parameters.

import jax.numpy as jnp
import jax.random as jr

import smcx

# fmt: off
y = jnp.array([
    -0.54, -1.09, -0.77, -0.03, 0.92, -0.45, 1.19, 0.24, 1.13,
    -0.42, 0.63, 1.18, 1.13, 0.64, 1.35, 2.25, 1.98, 1.65, 2.01,
    1.63, 0.80, 0.39, -0.68, -0.87, -0.96,
])[:, None]
# fmt: on

m0 = jnp.zeros(1)
C0 = jnp.eye(1)
G = 0.8 * jnp.eye(1)
W = 0.2 * jnp.eye(1)
F = jnp.eye(1)
V = 0.3 * jnp.eye(1)

kalman = smcx.kalman_filter(m0, C0, G, W, F, V, y)
print(kalman.marginal_loglik)  # -29.26, exact

If you want to define your own model, of any form, you instead write functions for the initial distribution, the transition distribution, and the observation distribution, and pass them to a particle filter. Here is the same model through the bootstrap particle filter:

def sample_initial(key, num_particles):
    return jr.normal(key, (num_particles, 1))


def sample_transition(key, state):
    return 0.8 * state + jnp.sqrt(0.2) * jr.normal(key, state.shape)


def log_observation(obs, state):
    residual = obs[0] - state[0]
    return -0.5 * (jnp.log(2 * jnp.pi * 0.3) + residual**2 / 0.3)


particle = smcx.bootstrap_filter(
    jr.key(0),
    sample_initial,
    sample_transition,
    log_observation,
    y,
    num_particles=10_000,
)
print(particle.marginal_loglik)  # -29.16, N = 10,000

The two estimates agree because the model is the same. If our model leaves the linear-Gaussian family, we can no longer use the Kalman filter. We only change the three functions of the bootstrap call to the new densities. The table below maps each model class to its methods, and the introduction in the documentation develops the theory with four worked examples, relaxing one assumption at a time.

Methods

smcx implements the standard sequential inference methods:

Setting	Methods	Functions
Linear-Gaussian, fully known	Kalman filter and RTS smoother, exact	kalman_filter, rts_smoother
Known nonlinear functions	Extended and unscented Kalman filters, approximate; the linearization strategy is an argument	extended_kalman_filter, unscented_kalman_filter, gaussian_filter
Observation variance unknown, variance-scaled	Conjugate DLM, exact	dlm_filter
Count and binary observations	Conjugate/linear-Bayes DGLM, approximate; the observation family is an argument	dglm_filter with poisson(), bernoulli(), or binomial(trials=n)
General densities	Bootstrap, auxiliary, and guided particle filters	bootstrap_filter, auxiliary_filter, guided_filter
Custom particle algorithms	Feynman–Kac derivations over one generic loop	StateSpaceModel, FeynmanKac, run_smc, run_particle_filter
Static parameters	Adaptive tempered SMC, SMC², and the joint Liu-West filter	temper, smc2, liu_west_filter
Simulation and prediction	Model simulation and posterior predictive draws	simulate, posterior_predictive_sample
Resampling	Systematic, stratified, multinomial, residual	systematic, stratified, multinomial, residual
Diagnostics and reporting	ESS, scoring rules, trajectory reconstruction, ArviZ export	diagnose, crps, reconstruct_trajectories, to_arviz

smcx runs on CPU, CUDA, and TPU through JAX, and on Apple-silicon GPUs through the optional jax-mps backend.

Installation

smcx requires Python 3.11 or later.

pip install smcx

Optional extras add Apple-silicon GPU execution or ArviZ reporting:

pip install "smcx[metal]"
pip install "smcx[arviz]"

Documentation

Available at michaelellis003.github.io/smcx.

Citation

If smcx contributes to academic work, please cite the release used. The repository's Cite this repository menu uses CITATION.cff to provide BibTeX and APA entries; include the version and release date in the final citation.

See also

State-space models and SMC

• dynamax: probabilistic state-space models with learning via EM and SGD.
• dynestyx: NumPyro-based inference for dynamical systems.
• particles: the reference Python companion to Chopin and Papaspiliopoulos (2020).
• BlackJAX: MCMC and SMC samplers for JAX.

The JAX ecosystem

• Equinox: neural networks and PyTree modules.
• Diffrax: numerical differential equation solvers.
• jaxtyping: shape and dtype annotations for arrays.
• ArviZ: exploratory analysis of Bayesian models.

Sources and attribution

The broader Feynman–Kac architecture follows Chopin and Papaspiliopoulos's An Introduction to Sequential Monte Carlo. The caller-owned particle-filter runner and dependency-free tempering mutation boundary were informed by BlackJAX's functional state/information protocol and pinned SMC-from-MCMC split, and by the separation of orchestration from history in particles 0.4. These are design credits; no code was copied or translated. The implemented methods draw on these primary sources:

• Exact linear-Gaussian state estimation: Kalman (1960) and Rauch, Tung, and Striebel (1965).
• Conjugate dynamic models: West and Harrison (1997) and West, Harrison, and Migon (1985).
• Nonlinear Gaussian filtering: Schmidt (1966) and Julier (2002).
• Particle filters: Gordon, Salmond, and Smith (1993), Pitt and Shephard (1999), Doucet, Godsill, and Andrieu (2000), and Liu and West (2001).
• Static and parameter inference: Del Moral, Doucet, and Jasra (2006) and Chopin, Jacob, and Papaspiliopoulos (2013).
• Resampling and diagnostics: Douc, Cappé, and Moulines (2005), Lee and Whiteley (2018), Zhang and Stephens (2009), and Vehtari et al. (2024).
• Scoring rules: Matheson and Winkler (1976) and Gneiting and Raftery (2007).
• Reporting: ArviZ.

Contributing

Contributions are welcome. See CONTRIBUTING.md for the development setup and pull-request conventions.

License

smcx is distributed under the Apache License 2.0.