Metadata-Version: 2.4
Name: smcx
Version: 1.2.0
Summary: Sequential Monte Carlo in JAX: particle filters, tempered SMC, and SMC2 on CPU, CUDA, TPU, and Apple-silicon GPUs
Keywords: jax,particle-filter,sequential-monte-carlo,bayesian-inference,apple-silicon,state-space-models,hidden-markov-models
Author: Michael Ellis
Author-email: Michael Ellis <michaelellis003@gmail.com>
License-Expression: Apache-2.0
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: Apache Software License
Classifier: Operating System :: MacOS :: MacOS X
Classifier: Operating System :: POSIX :: Linux
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Topic :: Scientific/Engineering :: Mathematics
Classifier: Typing :: Typed
Requires-Dist: jax>=0.10,<0.11
Requires-Dist: jaxlib>=0.10,<0.11
Requires-Dist: jaxtyping>=0.3.7
Requires-Dist: numpy>=2.2.6
Requires-Dist: jax-mps>=0.10.9,<0.11 ; platform_machine == 'arm64' and sys_platform == 'darwin' and extra == 'metal'
Requires-Python: >=3.11
Project-URL: Homepage, https://github.com/michaelellis003/smcx
Project-URL: Repository, https://github.com/michaelellis003/smcx
Project-URL: Issues, https://github.com/michaelellis003/smcx/issues
Project-URL: Documentation, https://michaelellis003.github.io/smcx/
Project-URL: Changelog, https://github.com/michaelellis003/smcx/blob/main/CHANGELOG.md
Provides-Extra: metal
Description-Content-Type: text/markdown

# smcx

[![CI](https://github.com/michaelellis003/smcx/actions/workflows/ci.yml/badge.svg)](https://github.com/michaelellis003/smcx/actions/workflows/ci.yml)
[![PyPI](https://img.shields.io/pypi/v/smcx?cacheSeconds=3600)](https://pypi.org/project/smcx/)
[![License](https://img.shields.io/github/license/michaelellis003/smcx)](LICENSE)

Sequential Monte Carlo in [JAX](https://github.com/jax-ml/jax): particle
filters, adaptive tempered SMC, and SMC² with a small, flat API. Runs on
CPU, CUDA, and TPU through stock JAX, and on Apple-silicon GPUs through
the optional [jax-mps](https://github.com/tillahoffmann/jax-mps) backend.

## Install

```bash
pip install smcx            # CPU / CUDA / TPU via your jax install
pip install "smcx[metal]"   # + jax-mps for Apple-silicon GPUs
```

## What's in the box

- **Filters**: `bootstrap_filter`, `guided_filter` (general g·f/q
  proposal weights), `auxiliary_filter` (twisted potentials), and
  `liu_west_filter` (joint state–parameter, labeled approximate).
- **Static targets**: `temper` — adaptive tempered SMC with an
  ESS-bisection schedule and covariance-adapted random-walk moves.
- **Parameter inference**: `smc2` — nested SMC² with vmapped inner
  filters and PMMH rejuvenation.
- **Resampling**: systematic, stratified, multinomial, and residual —
  one contract, log-domain weights throughout, float32-safe query
  grids.
- **Diagnostics**: ESS traces, quantile tail-ESS, Pareto-k
  reliability, single-run log-evidence variance from the genealogy
  (Lee & Whiteley 2018), trajectory reconstruction, CRPS,
  cumulative log score, Bayes factors, posterior-predictive
  sampling, and a one-call `diagnose` summary.
- `store_history=False` on every filter drops memory from O(T·N) to
  O(N) with a bit-identical evidence estimate.

Every sampler is validated against exact references — Kalman oracles
for the filters, conjugate evidence for tempering, grid-integrated
posteriors for SMC² — with Monte-Carlo-calibrated gates, not loose
tolerances.

## Quick start

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

import smcx

# A 1-D linear-Gaussian state-space model.
A, Q, R = 0.9, 0.5, 0.3


def init(key, n):
    return jr.normal(key, (n, 1))


def transition(key, z):
    return A * z + jnp.sqrt(Q) * jr.normal(key, z.shape)


def log_observation(y, z):
    return -0.5 * (jnp.log(2 * jnp.pi * R) + (y[0] - z[0]) ** 2 / R)


def emission(key, z):
    return z + jnp.sqrt(R) * jr.normal(key, z.shape)


_, emissions = smcx.simulate(
    jr.key(1),
    lambda key: init(key, 1)[0],
    transition,
    emission,
    num_timesteps=100,
)

post = smcx.bootstrap_filter(
    jr.key(0),
    init,
    transition,
    log_observation,
    emissions,
    num_particles=10_000,
)
post.marginal_loglik  # unbiased evidence estimate (log-domain)
smcx.diagnose(post)  # ESS / diversity / Pareto-k health summary
```

Callbacks are per-particle; smcx vmaps them internally. Everything
takes an explicit PRNG key, and posteriors are NamedTuples — ordinary
JAX pytrees. The four filters and `simulate` accept a keyword-only
`inputs` sequence for controlled dynamics and covariate-driven
observations; input-aware callbacks receive the aligned `input_t` as
their final argument.

smcx is deliberately just the inference engine: it defines no model
classes and no distributions. Models enter as JAX callables — your
own closures, or thin wrappers around a model library such as
[Dynamax](https://github.com/probml/dynamax) (the test suite itself
uses Dynamax models this way).

## Apple silicon

The `[metal]` extra runs the same code on M-series GPUs via jax-mps.
Filter correctness on Metal is gate-verified in this repository's test
suite (`SMCX_TEST_PLATFORM=mps` runs it on the GPU), and several of the
performance fixes that make the backend fast for SMC-shaped workloads
were contributed upstream from this project (jax-mps #215, #216, #220).
Metal is float32-only; the suite runs float64 on CPU and float32 on
Metal automatically.

## Development

Requires Python 3.11+ and [uv](https://docs.astral.sh/uv/).

```bash
git clone https://github.com/michaelellis003/smcx.git
cd smcx
uv sync
uv run pre-commit install
uv run pre-commit install --hook-type commit-msg
uv run pre-commit install --hook-type pre-push
```

A `Makefile` covers common tasks:

```bash
make test        # lint + pytest
make lint        # ruff check, format check, license headers, ty
make format      # add license headers, ruff format, ruff fix
make docs        # build docs
```

Releases are automated: `python-semantic-release` reads conventional
commits on merge to main, bumps the version, tags, and publishes.

## Acknowledgments

smcx's design draws on the SMC ecosystem:
[particles](https://github.com/nchopin/particles) and Chopin &
Papaspiliopoulos's *An Introduction to Sequential Monte Carlo* (the
Feynman-Kac architecture),
[BlackJAX](https://github.com/blackjax-devs/blackjax) (the resampling
contract), [Dynamax](https://github.com/probml/dynamax) (container
conventions), TensorFlow Probability (criterion/trace hooks), and
design lessons from PyMC, FilterPy, pfilter, pyfilter, Stone Soup,
pomp, nimbleSMC, and ArviZ. See `CITATION.cff` for formal references.

## License

Apache-2.0
