smcx

Sequential inference for state-space models in JAX: Kalman-family and DLM/DGLM filters, smoothers, and SMC methods including particle filtering and smoothing, tempered SMC, IBIS, and SMC². Algorithms consume plain JAX callables and small typed records, keeping model definitions separate from inference. smcx defines no probabilistic programming language. Models defined elsewhere can be used when the caller maps their components to these callables or records.

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, assuming the only unknown is the latent state, 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(
    initial_mean=m0,
    initial_covariance=C0,
    transition_matrix=G,
    transition_covariance=W,
    observation_matrix=F,
    observation_covariance=V,
    emissions=y,
)
print(kalman.marginal_loglik)  # -29.26, exact

A particle filter instead needs a Markov state-space model given by three functions: a sampler for the initial law, a sampler for the transition, and an evaluable observation log density. 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 particle estimate approximates the exact Kalman value. At this key and N = 10,000 the two log-likelihoods differ by 0.10. This single key does not characterize Monte Carlo error. Repeated keys are needed to estimate the bias and spread of the log-likelihood error at N = 10,000. 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

Filtering conditions each state on the observations up to its own time and runs online. Smoothing revisits every state once the complete record is in hand: a backward pass consumes the stored filter output, so rts_smoother takes the result of kalman_filter rather than rerunning it. smcx implements the standard sequential inference methods:

Setting	Methods	Functions
Linear-Gaussian, fully known	Kalman filter, RTS smoother, lag-one covariances, and joint posterior draws, exact	kalman_filter, rts_smoother, smoothed_cross_covariances, posterior_sample
Known nonlinear functions	Extended and unscented Kalman filters and RTS smoothers, approximate; the linearization strategy is an argument	extended_kalman_filter, unscented_kalman_filter, gaussian_filter, gaussian_smoother
Observation variance unknown, variance-scaled	Conjugate DLM filter and retrospective smoother, exact	dlm_filter, dlm_smoother
Count and binary observations	Conjugate/linear-Bayes DGLM filtering and retrospective state-moment smoothing, approximate	dglm_filter with poisson(), bernoulli(), or binomial(trials=n); dglm_smoother
General densities	Bootstrap, auxiliary, and guided particle filters	bootstrap_filter, auxiliary_filter, guided_filter
General densities, retrospective	Genealogy paths and particle FFBS trajectory draws, approximate	reconstruct_trajectories, backward_simulation
Custom particle algorithms	Feynman–Kac derivations over one generic loop	StateSpaceModel, FeynmanKac, run_smc, run_particle_filter
Static parameters	Tempered SMC targets a fixed posterior through a temperature path. IBIS updates an exact-likelihood posterior in data order. SMC² nests a particle filter inside parameter-space SMC. Liu-West is approximate online parameter learning through kernel shrinkage	temper, ibis, 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, ArviZ export	diagnose, crps, 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), with joint draws from Carter and Kohn (1994) and Frühwirth-Schnatter (1994).
• Conjugate dynamic models: West and Harrison (1997) and West, Harrison, and Migon (1985), with retrospective state moments from Alves et al. (2025).
• Nonlinear Gaussian filtering: Schmidt (1966) and Julier (2002).
• Particle filters and smoothing: Gordon, Salmond, and Smith (1993), Pitt and Shephard (1999), Doucet, Godsill, and Andrieu (2000), Godsill, Doucet, and West (2004), and Liu and West (2001).
• Static and parameter inference: Chopin (2002), 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.