Skip to content

smcx

Sequential inference in JAX: Kalman filters, particle filters, and sequential Monte Carlo, with model code decoupled from inference code.

When a state-space model is linear and Gaussian with known covariances, the Kalman filter computes the exact posterior by a two-step recursion. Relaxing linearity keeps a Gaussian approximation alive (the extended and unscented filters). Outside special conjugate cases, unknown noise parameters or non-Gaussian observations break the Gaussian closed form. smcx implements an exact variance-scaled DLM and approximate conjugate/linear-Bayes DGLMs for specific observation families. Particle filters carry the posterior as a weighted sample cloud and cover general nonlinear or non-Gaussian models. Broader sequential Monte Carlo methods target other distribution sequences, including tempered paths for static parameters and nested state-parameter inference. smcx implements all of these methods. The introduction develops them one assumption at a time, with the equations and references.

Models are supplied as plain JAX callbacks. StateSpaceModel groups the particle-model callbacks for reuse across the bootstrap, guided, and auxiliary Feynman–Kac derivations over one shared loop, while the named filters provide the shortest one-call interface; Liu–West, tempered SMC, SMC², and the caller-owned runner have their own drivers. Resamplers, criteria, proposals, potentials, look-ahead twists, mutation kernels, temperature schedules, and Gaussian linearization strategies are all caller-replaceable, with parameters as explicitly threaded PyTrees on the record path.

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]"

Start here

  • Quickstart establishes an exact Kalman baseline, then builds, diagnoses, and improves a particle filter.
  • Filtering tutorial runs a complete example and plots its filtering intervals and effective sample size.
  • Custom models covers the model record, Feynman–Kac derivations, structured latent states, time-varying inputs, and an optional Equinox representation.
  • Stochastic volatility learns a static parameter online with the Liu–West filter.
  • ArviZ reporting exports weighted particle output for downstream analysis.
  • The API reference documents every public function and posterior container from its source docstring.

Backends and precision

CPU, CUDA, and TPU use stock JAX. The optional metal extra uses jax-mps on macOS 14 or later on arm64; Metal is float32-only. Every stochastic operation takes an explicit PRNG key.

Citation and license

See the repository citation record and CITATION.cff for citation metadata. smcx is distributed under the Apache License 2.0.