smcx 1.8.0

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

smcx

smcx is a JAX library for state-space inference: exact Gaussian filtering and smoothing, particle filters, adaptive tempered SMC, and SMC² with a small, function-oriented API. It runs on CPU, CUDA, and TPU through JAX, and on Apple-silicon GPUs through the optional jax-mps backend.

Features include:

• exact linear-Gaussian Kalman filtering and RTS smoothing;
• bootstrap, auxiliary, guided, and Liu–West particle filters;
• a public runner for caller-owned particle-filter kernels;
• adaptive tempered SMC and nested SMC² parameter inference;
• systematic, stratified, multinomial, and residual resampling;
• filtering diagnostics, scoring rules, trajectory reconstruction, and ArviZ export; and
• structured latent-state PyTrees and explicit time-varying inputs.

smcx supplies inference algorithms, not model or distribution classes. Linear-Gaussian models are dense arrays; particle models are ordinary JAX callables, so they can be written directly or adapted from libraries such as Dynamax. Filtering and smoothing remain separate functions joined by typed posterior containers, allowing research code to replace one stage without subclassing or rerunning the other.

Installation

smcx requires Python 3.11 or later.

pip install smcx

Install the optional extras for Apple-silicon GPU execution or ArviZ reporting with:

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

The metal extra uses jax-mps and is available on macOS arm64. Metal is float32-only; releases are tested on a physical M-series GPU as well as on CPU.

Documentation

The documentation includes a quickstart, guides for custom models and custom particle filters and ArviZ reporting, and the complete API reference.

Quick example

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

import smcx

a, q, r = 0.9, 0.5, 0.3


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


def transition_sampler(key, state):
    return a * state + jnp.sqrt(q) * jr.normal(key, state.shape)


def log_observation(y, state):
    error = y[0] - state[0]
    return -0.5 * (jnp.log(2 * jnp.pi * r) + error**2 / r)


observations = jnp.array([0.2, -0.1, 0.4, 0.7, 0.3])[:, None]
posterior = smcx.bootstrap_filter(
    jr.key(0),
    initial_sampler,
    transition_sampler,
    log_observation,
    observations,
    num_particles=10_000,
)

posterior.marginal_loglik
smcx.weighted_mean(posterior)
smcx.diagnose(posterior)

Callbacks describe one particle; smcx vectorizes them over the cloud. Every stochastic operation takes an explicit PRNG key, and posterior containers are JAX PyTrees.

Citation

If smcx contributes to academic work, cite the release used. The repository's Cite this repository menu is generated from CITATION.cff and provides BibTeX and APA entries.

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 was informed by the functional state/information protocol in BlackJAX 1.6.2 and 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 Gaussian state estimation: Kalman (1960) and Rauch, Tung, and Striebel (1965).
• 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.

Numerical validation references

The exact Gaussian outputs are independently validated against Dynamax 1.0.2 and statsmodels 0.14.6. They are comparison implementations, not code lineage: no code from either project is copied or translated. Exact commits, environments, licenses, and observed differences are recorded with the frozen multivariate fixture.

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.