Skip to content

smcx

smcx provides state-space inference algorithms for JAX: exact linear-Gaussian filtering and smoothing, extended and unscented nonlinear Gaussian filtering, particle filters, adaptive tempered SMC, and SMC². It is a function-oriented inference library, not a modeling framework.

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 explains nonlinear Gaussian and particle callback boundaries, 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.

Model boundary

Linear-Gaussian models enter as dense arrays. The extended and unscented Kalman filters share nonlinear mean callbacks; the extended filter adds explicit state-Jacobian callbacks. Particle-model callbacks describe one particle; smcx vectorizes them over the cloud. Every stochastic operation takes an explicit PRNG key. Bootstrap, auxiliary, and guided filters can carry nonempty latent-state PyTrees, and every filter accepts an explicit sequence of time-varying inputs.

The linear, extended, and unscented Kalman filters and RTS smoother return typed Gaussian posteriors. Particle algorithms similarly expose model, proposal, and resampling callbacks. These boundaries allow research code to replace supported pieces without adopting a class hierarchy.

CPU, CUDA, and TPU use stock JAX. The optional metal extra uses jax-mps on macOS arm64; Metal is float32-only.

Citation and license

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