smcx
Sequential inference for state-space models in JAX: Kalman-family and DLM/DGLM filters, and SMC methods including particle filters, tempered SMC, 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 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
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
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 | Tempered SMC targets a fixed posterior through a temperature path. SMC² nests a particle filter inside parameter-space SMC. Liu-West is approximate online parameter learning through kernel shrinkage | 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.
1 maintainer
Release files for smcx 2.1.0
For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.
Source distribution (sdist)
| File | Size | Uploaded | |
|---|---|---|---|
| smcx-2.1.0.tar.gz | 118.9 kB | Details |
Built distribution (wheel)
| File | Interpreter | ABI | Platform | Reset |
|---|---|---|---|---|
| smcx-2.1.0-py3-none-any.whl | Python 3 | none | any | Details |
Total release size: 253.6 kB
Release files / smcx-2.1.0.tar.gz
| Download URL | smcx-2.1.0.tar.gz |
|---|---|
| Size | 118.9 kB |
| Tags | Source |
|
SHA-256 checksum How to use checksums |
da95faa1d50eecb21809766935b16a739b50e525974c3ff809d52a8394104ae4
|
|
BLAKE2b-256 checksum How to use checksums |
3522c5c4c076d393bd53401ac77772b84302a786bcc0defff4383e01aaec33bf
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
Yes |
| Uploaded via |
twine/6.1.0 CPython/3.13.14
|
Provenance
Provenance describes where a file came from. On PyPI, provenance is shared via attestations, which provide a verifiable record of the build or publishing details. View details, limitations and caveats.
PyPI Publish Attestation
PyPI verified that this artifact, at this checksum, originated from the publisher listed below.
Signed by GitHub Actions, verified by PyPI on Jul 30, 2026.
Transparency logRelease files / smcx-2.1.0-py3-none-any.whl
| Download URL | smcx-2.1.0-py3-none-any.whl |
|---|---|
| Size | 134.7 kB |
| Tags | Python 3 |
|
SHA-256 checksum How to use checksums |
b81a15c67b7e17442d23a31b313582108fd681091ccf21b60570030d9d532a2d
|
|
BLAKE2b-256 checksum How to use checksums |
6f03628e7e4358c8492c280c71b30054f609add6c385b3c301e8abc94c0cfc04
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
Yes |
| Uploaded via |
twine/6.1.0 CPython/3.13.14
|
Provenance
Provenance describes where a file came from. On PyPI, provenance is shared via attestations, which provide a verifiable record of the build or publishing details. View details, limitations and caveats.
PyPI Publish Attestation
PyPI verified that this artifact, at this checksum, originated from the publisher listed below.
Signed by GitHub Actions, verified by PyPI on Jul 30, 2026.
Transparency log
