Metadata-Version: 2.4
Name: smcx
Version: 1.3.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 probability-space contract and float32-safe query grids; filters
  retain their weights in the log domain.
- **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.

Permanent SMC/PF tests use dependency-free mathematical oracles and fixed,
Monte-Carlo-calibrated gates. The outside SMC/PF implementations were run once
in isolated environments; their pinned sources, licenses, settings, and
numerical summaries are retained beside the corresponding unit tests, not as
test dependencies. Diagnostics retain their separately documented ArviZ
cross-validation dependency.

| Component | Permanent oracle | One-time independent comparison |
|---|---|---|
| Resampling | Exact offspring moments | particles, BlackJAX |
| Bootstrap / APF / guided | Kalman evidence and moments | particles; TFP where applicable |
| Liu–West | Conjugate posterior | nimbleSMC/NIMBLE under matched semantics |
| Tempering | Conjugate evidence and moments | particles, BlackJAX |
| SMC² | Grid-converged Kalman integral | particles; limited TFP diagnostic |

## 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  # log of an unbiased evidence estimate
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 bootstrap, auxiliary, and guided filters, plus
`simulate`, also accept a nonempty PyTree as the latent state. Every
particle-cloud leaf has leading axis `N`; every stored-history leaf has
leading axes `(T, N)`, and all leaves share the same resampling ancestry.
Liu–West, tempered SMC, and SMC² retain dense Euclidean parameter and
latent-state arrays.

All 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.
Trajectory reconstruction and posterior prediction preserve structured
states. Euclidean summaries such as `weighted_mean`, `tail_ess`, and
`diagnose` require a dense state history, so select or project a leaf
before calling them on a structured posterior.

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 example notebook uses
Dynamax models this way; unit tests use frozen dependency-free oracles).

## 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. Algorithm docstrings and validation tests carry
the formal references and immutable implementation pins.

## License

Apache-2.0
