smcx 1.0.0__tar.gz

smcx-1.0.0/PKG-INFO

@@ -0,0 +1,181 @@
+Metadata-Version: 2.4
+Name: smcx
+Version: 1.0.0
+Summary: Sequential Monte Carlo for Apple silicon, built on MLX
+Keywords: apple-silicon,particle-filter,sequential-monte-carlo,bayesian-inference,state-space-models,hidden-markov-models
+Author: Michael Ellis
+Author-email: Michael Ellis <michaelellis003@gmail.com>
+License-Expression: Apache-2.0
+Classifier: Development Status :: 3 - Alpha
+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.10
+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)](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.
+
+smcx succeeds [smcjax](https://github.com/michaelellis003/smcjax); all of
+its functionality was ported forward and extended.
+
+## 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, 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)
+
+
+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.
+
+## 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.
+
+## Design notes
+
+Decision records live in `docs/adr/` and dated evidence in
+`docs/research/` — including the benchmark study that shaped the
+architecture: a native-MLX implementation of this library was built,
+measured against jax-mps on equal terms, and retired when the
+measurements showed the compatibility path could be made equally fast
+with fixes contributed upstream
+(`docs/research/2026-07-16-jax-mps-internals.md`). The MLX
+implementation remains in git history at the `mlx-final` tag.
+
+## Acknowledgments
+
+smcx's design draws on the SMC ecosystem:
+[smcjax](https://github.com/michaelellis003/smcjax) (the predecessor
+library), [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
+and `docs/research/licensing.md` for the full provenance record.
+
+## License
+
+Apache-2.0

smcx-1.0.0/README.md

@@ -0,0 +1,147 @@
+# 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)](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.
+
+smcx succeeds [smcjax](https://github.com/michaelellis003/smcjax); all of
+its functionality was ported forward and extended.
+
+## 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, 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)
+
+
+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.
+
+## 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.
+
+## Design notes
+
+Decision records live in `docs/adr/` and dated evidence in
+`docs/research/` — including the benchmark study that shaped the
+architecture: a native-MLX implementation of this library was built,
+measured against jax-mps on equal terms, and retired when the
+measurements showed the compatibility path could be made equally fast
+with fixes contributed upstream
+(`docs/research/2026-07-16-jax-mps-internals.md`). The MLX
+implementation remains in git history at the `mlx-final` tag.
+
+## Acknowledgments
+
+smcx's design draws on the SMC ecosystem:
+[smcjax](https://github.com/michaelellis003/smcjax) (the predecessor
+library), [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
+and `docs/research/licensing.md` for the full provenance record.
+
+## License
+
+Apache-2.0

smcx-1.0.0/pyproject.toml

@@ -0,0 +1,193 @@
+[project]
+name = "smcx"
+version = "1.0.0"
+description = "Sequential Monte Carlo for Apple silicon, built on MLX"
+readme = "README.md"
+authors = [
+    { name = "Michael Ellis", email = "michaelellis003@gmail.com" }
+]
+license = "Apache-2.0"
+requires-python = ">=3.11"
+keywords = [
+        "apple-silicon",
+    "particle-filter",
+    "sequential-monte-carlo",
+    "bayesian-inference",
+    "state-space-models",
+    "hidden-markov-models",
+]
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Intended Audience :: Science/Research",
+    "License :: OSI Approved :: Apache Software License",
+    "Operating System :: MacOS :: MacOS X",
+    "Operating System :: POSIX :: Linux",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
+    "Topic :: Scientific/Engineering :: Mathematics",
+    "Typing :: Typed",
+]
+dependencies = [
+    "jax>=0.10,<0.11",
+    "jaxlib>=0.10,<0.11",
+    "jaxtyping>=0.3.7",
+    "numpy>=2.2.6",
+]
+
+[project.urls]
+Homepage = "https://github.com/michaelellis003/smcx"
+Repository = "https://github.com/michaelellis003/smcx"
+Issues = "https://github.com/michaelellis003/smcx/issues"
+Documentation = "https://michaelellis003.github.io/smcx/"
+Changelog = "https://github.com/michaelellis003/smcx/blob/main/CHANGELOG.md"
+
+[build-system]
+requires = ["uv_build>=0.9.5,<0.12.0"]
+build-backend = "uv_build"
+
+[project.optional-dependencies]
+# Apple-GPU acceleration via the jax-mps PJRT plugin (ADR-0018).
+metal = [
+    # jax-mps only ships macOS-arm64 wheels; the marker keeps the
+    # extra resolvable (as a no-op) on other platforms, e.g. Linux CI.
+    "jax-mps>=0.10.9,<0.11; sys_platform == 'darwin' and platform_machine == 'arm64'",
+]
+
+[dependency-groups]
+dev = [
+    "beartype>=0.22.9",
+    "dynamax>=1.0.2",
+    "hypothesis>=6.156.6",
+    "mkdocs-gen-files>=0.6.1",
+    "mkdocs-literate-nav>=0.6.3",
+    "mkdocs-material>=9.7.6",
+    "mkdocstrings[python]>=1.0.3",
+    "nanobind>=2.13.0",
+    "pre-commit>=4.5.1",
+    "properdocs>=1.6.7",
+    "pytest>=9.0.3",
+    "pytest-cov>=7.1.0",
+    "ruff>=0.15.9",
+    "ty>=0.0.29",
+]
+
+[tool.ruff]
+# Target version enables modern syntax upgrades where possible
+target-version = "py310"
+preview = true
+line-length = 80
+indent-width = 4
+# Generated mlx stubs (ADR-0007) are vendored artifacts, not code.
+extend-exclude = ["typings"]
+
+[tool.ruff.lint]
+select = [
+    "E", "W", "F", "I", "D", # Core errors, warnings, isort, and docstrings
+    "N",   # pep8-naming (Google style includes naming conventions)
+    "UP",  # pyupgrade (modernizes syntax)
+    "B",   # flake8-bugbear (catches design flaws/bugs)
+    "SIM", # flake8-simplify (cleans up messy logic)
+    "RUF", # Ruff's custom advanced rules
+]
+ignore = [
+    # jaxtyping uses string-shape annotations (e.g. Float[mx.array, " n"])
+    # which ruff parses as forward references and rejects (ADR-0007).
+    "F722",
+]
+
+[tool.ruff.lint.pydocstyle]
+# Enforce Google-style docstrings
+convention = "google"
+
+[tool.ruff.lint.per-file-ignores]
+# Tests use class-based grouping; method docstrings are noisy. LGSSM
+# tests keep the literature's matrix notation (F, Q, H, R, P0).
+"tests/*" = ["D102", "D103", "D104", "N806", "N803"]
+# The version fallback try/except is standard packaging boilerplate.
+"src/smcx/__init__.py" = ["RUF067"]
+# Protocol stubs: the class docstring documents the contract.
+"src/smcx/types.py" = ["D102"]
+"src/smcx/containers.py" = ["D102"]
+# Preserved research/verification scripts (see docs/research/*),
+# not library code — exempt from library-grade lint.
+"benchmarks/exploratory/*" = ["ALL"]
+# Kill-test harness: maintained, protocol-critical code — full lint
+# except docstrings on model closures/helpers.
+"benchmarks/killtest/*" = ["D101", "D102", "D103"]
+"benchmarks/smc2/bench_smc2.py" = ["D101", "D102", "D103"]
+# particles_side follows the `particles` API: uppercase model
+# methods (PX0/PX/PY) and a mutable default_params dict are required
+# by StateSpaceModel, not smcx style.
+"benchmarks/smc2/particles_side.py" = ["D101", "D102", "D103", "N802", "RUF012"]
+# conftest installs the jaxtyping import hook before importing smcx,
+# which forces module-level imports below code (smcjax convention).
+"tests/conftest.py" = ["E402"]
+
+[tool.ruff.format]
+docstring-code-format = true
+
+[tool.semantic_release]
+version_toml = ["pyproject.toml:project.version"]
+version_variables = ["src/smcx/__init__.py:__version__"]
+commit_message = "chore(release): {version}"
+# Stay 0.x on BREAKING CHANGE commits until stability is declared.
+major_on_zero = false
+
+[tool.semantic_release.changelog.default_templates]
+changelog_file = "CHANGELOG.md"
+
+[tool.semantic_release.remote]
+type = "github"
+
+[tool.semantic_release.publish]
+upload_to_vcs_release = true
+
+[tool.ty.src]
+# Research scripts are lint/type-exempt artifacts; typings/ holds
+# generated stubs (checked as a search path, not as source);
+# killtest and the SMC² particles baseline span separate venvs
+# (jax_side / particles_side import packages unresolvable here by
+# construction). The native-vs-jax-mps JAX worker is the same case: it
+# runs only inside the pinned `uv run --no-project` env, so jax and the
+# jax_plugins.mps backend are unresolvable in the project environment.
+# The MLX-era benchmark workers reference the pre-ADR-0018 API and the
+# vendored MLX stubs; they are retained as historical harness code until
+# the harness is re-aimed at the jax-mps backend (pivot spec, slice 7).
+exclude = [
+    "benchmarks/exploratory/**",
+    "benchmarks/killtest/**",
+    "benchmarks/native_vs_jax_mps/jax_worker.py",
+    "benchmarks/native_vs_jax_mps/mlx_worker.py",
+    "benchmarks/smc2/**",
+    "typings/**",
+]
+
+[tool.ty.environment]
+# Vendored mlx.core stubs (ADR-0007): upstream ships py.typed but no
+# .pyi for the compiled extension. Regenerate on every mlx floor bump:
+#   uv run python -m nanobind.stubgen -m mlx.core -O typings/mlx
+extra-paths = ["typings"]
+
+[tool.pytest.ini_options]
+minversion = "7.0"
+addopts = "-ra -q"
+testpaths = [
+    "tests",
+]
+
+[tool.coverage.run]
+source = ["smcx"]
+branch = true
+
+[tool.coverage.report]
+fail_under = 80
+show_missing = true
+exclude_lines = [
+    "pragma: no cover",
+    "if __name__ == .__main__.",
+    "if TYPE_CHECKING:",
+    "@overload",
+]

smcx-1.0.0/src/smcx/__init__.py

@@ -0,0 +1,94 @@
+# Copyright 2026 Michael Ellis
+# SPDX-License-Identifier: Apache-2.0
+
+"""Sequential Monte Carlo and particle filtering in JAX."""
+
+from importlib.metadata import PackageNotFoundError as _PackageNotFoundError
+from importlib.metadata import version as _version
+
+from smcx.auxiliary import auxiliary_filter
+from smcx.bootstrap import bootstrap_filter
+from smcx.containers import (
+    LiuWestPosterior,
+    ParticleFilterPosterior,
+    ParticleFilterResult,
+    ParticleState,
+    SMC2Posterior,
+    TemperedPosterior,
+)
+from smcx.diagnostics import (
+    crps,
+    cumulative_log_score,
+    diagnose,
+    log_bayes_factor,
+    log_ml_increments,
+    param_weighted_mean,
+    param_weighted_quantile,
+    pareto_k_diagnostic,
+    particle_diversity,
+    posterior_predictive_sample,
+    replicated_log_ml,
+    tail_ess,
+    weighted_mean,
+    weighted_quantile,
+    weighted_variance,
+)
+from smcx.exceptions import DegenerateWeightsError
+from smcx.guided import guided_filter
+from smcx.liu_west import liu_west_filter
+from smcx.resampling import (
+    multinomial,
+    residual,
+    stratified,
+    systematic,
+)
+from smcx.simulate import simulate
+from smcx.smc2 import smc2
+from smcx.tempering import temper
+from smcx.weights import ess, log_ess, log_normalize, normalize
+
+try:
+    __version__ = _version("smcx")
+except _PackageNotFoundError:
+    __version__ = "1.0.0"
+
+__all__ = [
+    "DegenerateWeightsError",
+    "LiuWestPosterior",
+    "ParticleFilterPosterior",
+    "ParticleFilterResult",
+    "ParticleState",
+    "SMC2Posterior",
+    "TemperedPosterior",
+    "__version__",
+    "auxiliary_filter",
+    "bootstrap_filter",
+    "crps",
+    "cumulative_log_score",
+    "diagnose",
+    "ess",
+    "guided_filter",
+    "liu_west_filter",
+    "log_bayes_factor",
+    "log_ess",
+    "log_ml_increments",
+    "log_normalize",
+    "multinomial",
+    "normalize",
+    "param_weighted_mean",
+    "param_weighted_quantile",
+    "pareto_k_diagnostic",
+    "particle_diversity",
+    "posterior_predictive_sample",
+    "replicated_log_ml",
+    "residual",
+    "simulate",
+    "smc2",
+    "stratified",
+    "systematic",
+    "tail_ess",
+    "temper",
+    "weighted_mean",
+    "weighted_quantile",
+    "weighted_variance",
+]

smcx-1.0.0/src/smcx/__main__.py

@@ -0,0 +1,15 @@
+# Copyright 2026 Michael Ellis
+# SPDX-License-Identifier: Apache-2.0
+
+"""Allow running the package with ``python -m smcx``."""
+
+from smcx import __version__
+
+
+def main() -> None:
+    """Print package version."""
+    print(f"smcx {__version__}")
+
+
+if __name__ == "__main__":
+    main()