Metadata-Version: 2.4
Name: smcx
Version: 1.14.7
Summary: State-space inference in JAX: Kalman and particle filters, tempered SMC, and SMC2 across CPU and accelerators
Keywords: jax,particle-filter,kalman-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
License-File: LICENSE
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: arviz>=0.23.4 ; extra == 'arviz'
Requires-Dist: jax-mps>=0.10.10,<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: arviz
Provides-Extra: metal
Description-Content-Type: text/markdown

# smcx

smcx is a [JAX](https://github.com/jax-ml/jax) library for state-space
inference: exact linear-Gaussian filtering and smoothing, extended and
unscented nonlinear Gaussian filtering, particle filters, adaptive
tempered SMC, and SMC² with a small, function-oriented API. It runs on
CPU, CUDA, and TPU through JAX, and on Apple-silicon GPUs through the optional
[jax-mps](https://github.com/tillahoffmann/jax-mps) backend.

Features include:

- exact linear-Gaussian Kalman filtering and RTS smoothing;
- extended and scaled unscented Kalman filtering with shared mean callbacks;
- bootstrap, auxiliary, guided, and Liu–West particle filters;
- a public runner for caller-owned particle-filter kernels;
- adaptive tempered SMC and nested SMC² parameter inference;
- systematic, stratified, multinomial, and residual resampling;
- filtering diagnostics, scoring rules, trajectory reconstruction, and
  ArviZ export; and
- structured latent-state PyTrees and explicit time-varying inputs.

smcx supplies inference algorithms, not model or distribution classes.
Its catalog covers core families rather than every named variant; additions
need concrete user or research demand and an independent validator.
Linear-Gaussian models use dense arrays. The EKF and UKF share mean callbacks;
the EKF alone adds explicit Jacobians. Particle models use JAX callbacks,
and `run_particle_filter` lets research code own the filter kernel. Typed
posteriors keep filtering and smoothing composable without a class hierarchy.

## Installation

smcx requires Python 3.11 or later.

```bash
pip install smcx
```

Install the optional extras for Apple-silicon GPU execution or ArviZ
reporting with:

```bash
pip install "smcx[metal]"
pip install "smcx[arviz]"
```

The `metal` extra uses jax-mps and is available on macOS 14 or later on
arm64. Metal is float32-only; releases are tested on a physical M-series
GPU as well as on CPU.

## Documentation

The [documentation](https://michaelellis003.github.io/smcx/) includes a
[quickstart](https://michaelellis003.github.io/smcx/guides/quickstart/),
guides for [custom models and custom particle filters][custom-models]
and [ArviZ reporting](https://michaelellis003.github.io/smcx/guides/arviz/),
and the complete
[API reference](https://michaelellis003.github.io/smcx/api/smcx/).

[custom-models]:
  https://michaelellis003.github.io/smcx/guides/custom-models/

## Quick example

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

import smcx

a, q, r = 0.9, 0.5, 0.3


def initial_sampler(key, num_particles):
    return jr.normal(key, (num_particles, 1))


def transition_sampler(key, state):
    return a * state + jnp.sqrt(q) * jr.normal(key, state.shape)


def log_observation(y, state):
    error = y[0] - state[0]
    return -0.5 * (jnp.log(2 * jnp.pi * r) + error**2 / r)


observations = jnp.array([0.2, -0.1, 0.4, 0.7, 0.3])[:, None]
posterior = smcx.bootstrap_filter(
    jr.key(0),
    initial_sampler,
    transition_sampler,
    log_observation,
    observations,
    num_particles=10_000,
)

posterior.marginal_loglik
smcx.weighted_mean(posterior)
smcx.diagnose(posterior)
```

Callbacks describe one particle; smcx vectorizes them over the cloud.
Every stochastic operation takes an explicit PRNG key, and posterior
containers are JAX PyTrees.

## Citation

If smcx contributes to academic work, cite the release used. The
repository's **Cite this repository** menu uses the base metadata in
[`CITATION.cff`](https://github.com/michaelellis003/smcx/blob/main/CITATION.cff)
to provide BibTeX and APA entries. Include the release tag or version and
release date in the final citation.

## Sources and attribution

The broader Feynman–Kac architecture follows Chopin and
Papaspiliopoulos's
[*An Introduction to Sequential Monte Carlo*](https://doi.org/10.1007/978-3-030-47845-2).
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](https://github.com/blackjax-devs/blackjax/blob/a9ef478c69d730a2caa13ca4b2d735c580e0feec/blackjax/smc/from_mcmc.py),
and by the separation of orchestration from history in
[particles 0.4](https://github.com/nchopin/particles/releases/tag/v0.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)](https://doi.org/10.1115/1.3662552) and
  [Rauch, Tung, and Striebel (1965)](https://doi.org/10.2514/3.3166).
- Nonlinear Gaussian filtering:
  [Schmidt (1966)](https://doi.org/10.1016/B978-1-4831-6716-9.50011-4) and
  [Julier (2002)](https://doi.org/10.1109/ACC.2002.1025369).
- Particle filters: [Gordon, Salmond, and Smith (1993)](https://doi.org/10.1049/ip-f-2.1993.0015),
  [Pitt and Shephard (1999)](https://doi.org/10.1080/01621459.1999.10474153),
  [Doucet, Godsill, and Andrieu (2000)](https://doi.org/10.1023/A:1008935410038),
  and [Liu and West (2001)](https://doi.org/10.1007/978-1-4757-3437-9_10).
- Static and parameter inference:
  [Del Moral, Doucet, and Jasra (2006)](https://doi.org/10.1111/j.1467-9868.2006.00553.x)
  and [Chopin, Jacob, and Papaspiliopoulos (2013)](https://doi.org/10.1111/j.1467-9868.2012.01046.x).
- Resampling and diagnostics:
  [Douc, Cappé, and Moulines (2005)](https://doi.org/10.1109/ISPA.2005.195385),
  [Lee and Whiteley (2018)](https://doi.org/10.1093/biomet/asy028),
  [Zhang and Stephens (2009)](https://doi.org/10.1198/TECH.2009.08017),
  and [Vehtari et al. (2024)](https://jmlr.org/papers/v25/19-556.html).
- Scoring rules:
  [Matheson and Winkler (1976)](https://doi.org/10.1287/mnsc.22.10.1087)
  and [Gneiting and Raftery (2007)](https://doi.org/10.1198/016214506000001437).
- Reporting: [ArviZ](https://doi.org/10.21105/joss.01143).

### Numerical validation references

The linear Kalman and RTS outputs are independently validated against
[Dynamax 1.0.2](https://github.com/probml/dynamax/releases/tag/1.0.2)
and
[statsmodels 0.14.6](https://github.com/statsmodels/statsmodels/releases/tag/v0.14.6);
the details are recorded with the
[frozen linear fixture](https://github.com/michaelellis003/smcx/blob/ac9572da46dad4c3829ccde99d1bc7fc05ead0dd/tests/_kalman_reference.py).

The extended and unscented Kalman outputs are independently validated against
[Stone Soup 1.9.1](https://github.com/dstl/Stone-Soup/releases/tag/v1.9.1),
cross-checked with
[Dynamax 1.0.2](https://github.com/probml/dynamax/releases/tag/1.0.2),
and checked against
[SciPy 1.18.0](https://github.com/scipy/scipy/releases/tag/v1.18.0)
innovation log densities. Exact commits, environments, licenses, and
observed differences are recorded with the frozen
[extended](https://github.com/michaelellis003/smcx/blob/ac9572da46dad4c3829ccde99d1bc7fc05ead0dd/tests/_extended_kalman_reference.py)
and
[unscented](https://github.com/michaelellis003/smcx/blob/ac9572da46dad4c3829ccde99d1bc7fc05ead0dd/tests/_unscented_kalman_reference.py)
fixtures.

The auxiliary-guided runner recipe is distributionally cross-checked against
particles 0.4's pinned
[`AuxiliaryPF` composition](https://github.com/nchopin/particles/blob/c5fcb0b6d34b3c8efea6f6dc21d73e0e91287d9f/particles/state_space_models.py#L352-L428)
and [SMC bookkeeping](https://github.com/nchopin/particles/blob/c5fcb0b6d34b3c8efea6f6dc21d73e0e91287d9f/particles/core.py#L299-L359).

These projects are numerical comparison implementations, not code
lineage; no implementation code was copied or translated.

## Contributing

Contributions are welcome. See
[`CONTRIBUTING.md`](https://github.com/michaelellis003/smcx/blob/main/CONTRIBUTING.md)
for the development setup and pull-request conventions.

## License

smcx is distributed under the
[Apache License 2.0](https://github.com/michaelellis003/smcx/blob/main/LICENSE).
