diff --git a/.gitignore b/.gitignore index 11772985..36bd8286 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,5 @@ +*.DS_Store + # Byte-compiled / optimized / DLL files __pycache__/ *.py[cod] diff --git a/AI_POLICY.md b/AI_POLICY.md index 1e92d7b3..ee4771a2 100644 --- a/AI_POLICY.md +++ b/AI_POLICY.md @@ -1,12 +1,13 @@ # Usage of Generative AI -The sequince maintainers welcome contributions, including those developed with the assistance of generative AI tools. However, the use of such tools does not change our expectations for contributor responsibility, engagement, or code quality. The following apply to all contributions. +The Sequince maintainers welcome contributions, including those developed with the assistance of generative AI tools. However, the use of such tools does not change our expectations for contributor responsibility, engagement, or code quality. The following apply to all contributions. 1. **Human ownership and accountability** - i. Contributors are responsible for all submitted content, regardless of whether generative AI tools were used. - ii. Contributors must be able to explain and justify all changes during review. - iii. Contributors are responsible for ensuring all submitted code complies with Sequince's licensing requirements, regardless of whether generative AI tools were used. - iv. Contributors must remain the clear human owner of the contribution. Pull requests submitted primarily by automated systems or where no accountable human author is evident may be closed at the maintainers' discretion. + We require that: + - Contributors are responsible for all submitted content, regardless of whether generative AI tools were used. + - Contributors must be able to explain and justify all changes during review. + - Contributors are responsible for ensuring all submitted code complies with Sequince's licensing requirements, regardless of whether generative AI tools were used. + - Contributors must remain the clear human owner of the contribution. Pull requests submitted primarily by automated systems or where no accountable human author is evident may be closed at the maintainers' discretion. 2. **Engagement** The pull request process is collaborative and iterative. Contributors are expected to actively engage with reviewer feedback themselves. Contributors are expected to understand the changes they are proposing, not act as a message courier for a third party. @@ -20,5 +21,6 @@ The sequince maintainers welcome contributions, including those developed with t 5. **Maintainer discretion** Maintainers may close or decline pull requests that are low quality, excessively burdensome to review, appear to lack meaningful human oversight, or otherwise are not in the best interests of the project. -> [!NOTE] +> **Note:** +> > This policy is informed by similar guidance adopted in the Scientific Python ecosystem, including Astropy, Matplotlib, and scikit-learn. diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 02d3c8ab..9e092a49 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -8,10 +8,12 @@ workflow and conventions we use to keep changes easy to review and maintain. Treat other contributors with respect and follow the [Python Community Code of Conduct](https://www.python.org/psf/codeofconduct/). -Contributions must comply with the project's [license](LICENSE). Preserve required +Contributions must comply with the project's +[license](https://github.com/sequince-dev/sequince/blob/main/LICENSE). Preserve required attribution when incorporating third-party material. -AI-assisted contributions are welcome under our [AI policy](AI_POLICY.md). +AI-assisted contributions are welcome under our +[AI policy](https://github.com/sequince-dev/sequince/blob/main/AI_POLICY.md). Contributors are responsible for understanding, validating, and explaining their changes, including their numerical correctness and performance. @@ -157,6 +159,22 @@ def particle_mean(particles: Array) -> Array: return xp.mean(particles, axis=0) ``` +## Building the documentation + +The documentation uses Sphinx and lives in `docs/`. From the repository root, +build the HTML documentation with the dedicated Pixi environment: + +```bash +pixi run -e docs docs +``` + +Pixi installs the documentation dependencies defined in `pyproject.toml` when +preparing this environment. Open `docs/_build/html/index.html` in a browser to +review the result. Rebuild after editing documentation or API docstrings. + +The developer guide embeds this file and `AI_POLICY.md` using MyST. Edit the +root Markdown files to update their content on GitHub and in the documentation. + ## Checks and tests Run the following checks before requesting review: diff --git a/docs/.gitignore b/docs/.gitignore new file mode 100644 index 00000000..4ebf789e --- /dev/null +++ b/docs/.gitignore @@ -0,0 +1 @@ +autoapi/ diff --git a/docs/Makefile b/docs/Makefile new file mode 100644 index 00000000..bcd747f9 --- /dev/null +++ b/docs/Makefile @@ -0,0 +1,22 @@ +# Minimal makefile for Sphinx documentation +# + +# You can set these variables from the command line, and also +# from the environment for the first two. +SPHINXOPTS ?= +SPHINXBUILD ?= sphinx-build +SOURCEDIR = . +BUILDDIR = _build + +# Put it first so that "make" without argument is like "make help". +help: + @$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) + +.PHONY: help docs Makefile + +docs: html + +# Catch-all target: route all unknown targets to Sphinx using the new +# "make mode" option. $(O) is meant as a shortcut for $(SPHINXOPTS). +%: Makefile + @$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) diff --git a/docs/api/core.rst b/docs/api/core.rst deleted file mode 100644 index ca3c4d76..00000000 --- a/docs/api/core.rst +++ /dev/null @@ -1,60 +0,0 @@ -Core modules -============ - -Top-level API -------------- - -.. automodule:: sequince - :members: - :undoc-members: - :show-inheritance: - -Sequential Monte Carlo driver ------------------------------ - -Reconstruct a runner with ``SequentialMonteCarlo.from_config(smc.config)``. -This preserves runner settings and shares the referenced steps, callbacks, and -checkpoint writers. It does not clone mutable component state or resume a run; -use the checkpoint APIs to restore runtime state. - -.. automodule:: sequince.smc - :members: - :undoc-members: - :show-inheritance: - -Particle state and results --------------------------- - -.. automodule:: sequince.state - :members: - :undoc-members: - :show-inheritance: - -.. automodule:: sequince.result - :members: - :undoc-members: - :show-inheritance: - -Checkpointing -------------- - -.. automodule:: sequince.checkpoint - :members: - :undoc-members: - :show-inheritance: - -Posterior utilities -------------------- - -.. automodule:: sequince.posterior - :members: - :undoc-members: - :show-inheritance: - -Pre-built samplers ------------------- - -.. automodule:: sequince.prebuilt_samplers - :members: - :undoc-members: - :show-inheritance: diff --git a/docs/api/index.rst b/docs/api/index.rst deleted file mode 100644 index f42fa2d6..00000000 --- a/docs/api/index.rst +++ /dev/null @@ -1,9 +0,0 @@ -API reference -============= - -.. toctree:: - :maxdepth: 2 - - core - steps - utilities diff --git a/docs/api/steps.rst b/docs/api/steps.rst deleted file mode 100644 index b8e6c391..00000000 --- a/docs/api/steps.rst +++ /dev/null @@ -1,214 +0,0 @@ -SMC steps -========= - -Common interface ----------------- - -.. automodule:: sequince.steps.base - :members: - :undoc-members: - :show-inheritance: - -Before functional initialization, the runner calls each step's -``validate_functional_configuration(config, state)`` hook, passing its -``smc.config`` snapshot and the initial particle state. -``SMCConfig`` exposes ``steps``, ``callbacks``, ``checkpoint_writers``, and -``auto_normalize`` as read-only fields with tuple containers. Referenced step -and callback objects retain their identity and mutable state. The config is -host-side configuration, separate from runtime metadata and checkpoints. -Steps can reject incompatible pipeline settings here, before any target -evaluation or cache preparation. The default implementation does nothing. -Metadata and cache preparation remain in ``initialize_functional_metadata`` -and ``initialize_functional_state_and_context``. Algorithm-specific checks -belong to the relevant steps, rather than the runner. - -Functional conditional execution --------------------------------- - -``sequince.steps.apply_functional_if(predicate, apply, state, context, key)`` -is a reusable control-flow helper for functional steps, using ``acai.cond``. -It executes -``apply(state, context, key)`` only when the scalar predicate is true. Otherwise -it returns the unchanged state, context, and RNG. On NumPy, CuPy, and PyTorch, -only the selected branch executes and skipped output is None. PyTorch control -flow is eager, without structured ``torch.compile`` guarantees. On JAX, skipped -outputs contain zero array placeholders in the normal output structure, with -static fields retained. Both JAX branches are traced; the operation must be -pure and traceable even when skipped, and initialization must establish matching -carry structures, shapes, and dtypes. Placeholder outputs do not describe -executed transitions. - -After each functional step, the runner guards subsequent steps using the -runtime ``context.stop_immediately`` flag. ``StopIf(immediate=True)``, NS-SMC -shell steps, and custom steps returning ``context.request_stop_immediately()`` -all stop execution without an additional capability declaration. An ordinary -stop request still permits the rest of the iteration to execute. Use -``initialize_functional`` before directly invoking a compiled iteration so -all guarded steps have compatible carries; ``run_functional`` and -``compile_functional_iteration`` perform this initialization automatically. -As in eager execution, runner-level automatic normalization still runs after -the step sequence when enabled. - -Weighting strategies --------------------- - -.. automodule:: sequince.steps.weighting - :members: - :undoc-members: - :show-inheritance: - -Nested-sampling SMC -------------------- - -Use ``make_nested_sampling_smc`` to construct the composable pipeline and wrap its -result in ``NestedSamplingSMCResult`` for shell evidence and posterior samples. -``NestedSamplingLevel`` provides the per-shell diagnostics. The ordinary -``SMCResult.log_evidence()`` and ``SMCResult.posterior_samples()`` methods raise -an error for NS-SMC results and direct callers to this wrapper. -``SMCResult.log_normalizing_constant()`` remains available for the constrained -prior normalization. - -Custom threshold schedules must implement -``schedule(log_likelihoods, state, context) -> float | scalar array``. All three arguments are -passed positionally. Raise ``StopIteration`` to integrate the remaining live -tail and finish; other exceptions propagate to the caller. Fixed, weighted -quantile, and worst-particle schedules implement this same signature. - -During execution, ``context.metadata["nested_sampling"]`` holds the current -threshold and bookkeeping as a plain dictionary. Steps replace this dictionary -rather than mutate an earlier context. Shell expectations and evidence -increments are stored in ``NestedSamplingLevel`` records. -State metadata records snapshots after survivor -weighting or terminal-shell integration, preserving result histories. Resume -through the standard checkpoint API to restore both particles and context. - -Functional and compiled execution -~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - -``run_functional`` supports JAX arrays with unconditional resampling -(``always_resample=True``) and an unchanged population size. BlackJAX and -MiniPCN provide functional mutation kernels; host-only kernels such as emcee, -Eryn, and the current Metropolis-Hastings implementations are rejected. Set -``logging_frequency=0`` and use runner callbacks for host logging. Existing -functional restrictions on transforms and mutation options still apply. - -For example, using JAX-compatible density functions and particles:: - - import jax - from sequince import ParticleState, make_nested_sampling_smc - from sequince.jax import compile_functional_iteration - - smc, target, metadata = make_nested_sampling_smc( - log_prior, log_likelihood, - threshold_schedule=[-10.0, -5.0, -1.0], - evidence_tolerance=None, - mutation="blackjax", - mutation_kwargs={"algorithm": "random-walk", - "algorithm_kwargs": {"sigma": 0.1}}, - particle_transform=None, - logging_frequency=0, - ) - state = ParticleState(particles) - key = jax.random.key(42) - iteration = compile_functional_iteration(smc, state, key, metadata=metadata) - result, key = smc.run_functional( - state, random_state=key, metadata=metadata, iteration=iteration, - ) - nested = result.as_nested_sampling_result() - samples, weights = nested.posterior_samples() - -Omit ``iteration`` for uncompiled functional execution, or supply -``jax.jit(smc.functional_iteration)``. The runner initializes numeric metadata -and a complete target cache before the first iteration. Call -``initialize_functional`` yourself when invoking ``functional_iteration`` -directly. Context metadata remains authoritative; state metadata contains -independent historical snapshots. Stop reasons are numeric during functional -execution; ``nested.stop_reason`` and ``nested.summary()`` expose readable -strings. Shapes, dtypes, dictionary keys, and cache components must remain fixed. - -Compiled schedules use the same three-argument signature. They must use JAX -operations without host conversions, Python branching on traced values, mutable -iterator state, or data-dependent exceptions. Return positive infinity to -integrate the terminal tail. The built-in fixed schedule uses this sentinel on -exhaustion while retaining ``StopIteration`` in eager execution. An explicit -infinite threshold (including one returned by a custom schedule) is reported -as ``no_survivors``; exhausting a built-in fixed schedule is reported as -``schedule_exhausted``. Ordinary user exceptions propagate during execution or -tracing. Target functions and transformations must also be pure and traceable; -log densities must be finite or negative infinity. Shape and cache ownership -checks run during tracing; density values and threshold errors are checked at -initialization and after each host-controlled iteration. Direct compiled calls -must invoke the selection step's ``validate_functional_state_and_context`` at -the host boundary. Rejected internal mutation proposals are not inspected by -these host checks. - -After terminal-shell integration, later steps are conditionally skipped, -including weighting, resampling, and mutation. Their carry (particles, caches, -and RNG) is unchanged. Skipped outputs retain the usual structure with zero -array placeholders; they do not represent executed transitions. Functional -shell outputs use population-sized arrays and a particle mask. Optional -prior-mass tail outputs also carry a ``valid`` flag, allowing two shells in one -iteration. The result wrapper filters these flags and masks when reconstructing -levels and posterior samples. Eager shell storage remains compact, and -``store_shell_particles=False`` avoids retaining shell particle arrays in both -modes. - -Evidence remains separate from constrained-prior normalizing constants. A -``num_steps`` limit suspends an incomplete run without integrating its live -tail. Use ``resume_functional_from(checkpoint)`` and then -``run_functional(runtime=runtime, ...)`` to continue. With a newly constructed -target, initialization refreshes incompatible cached evaluations before -compilation/execution; RNG and historical shells are restored. Portable result -serialization retains shell masks, diagnostics, and posterior interpretation. - -.. automodule:: sequince.steps.constraints - :members: - :undoc-members: - :show-inheritance: - -Resampling ----------- - -.. automodule:: sequince.steps.resampling - :members: - :undoc-members: - :show-inheritance: - -Mutation kernels ----------------- - -.. automodule:: sequince.steps.mutation - :members: - :undoc-members: - :show-inheritance: - -Proposals and dimensionality control ------------------------------------- - -.. automodule:: sequince.steps.proposal - :members: - :undoc-members: - :show-inheritance: - -.. automodule:: sequince.steps.dimensionality - :members: - :undoc-members: - :show-inheritance: - -Control flow and utilities --------------------------- - -.. automodule:: sequince.steps.control - :members: - :undoc-members: - :show-inheritance: - -.. automodule:: sequince.steps.log - :members: - :undoc-members: - :show-inheritance: - -.. automodule:: sequince.steps.waste_free - :members: - :undoc-members: - :show-inheritance: diff --git a/docs/api/utilities.rst b/docs/api/utilities.rst deleted file mode 100644 index cb646f7c..00000000 --- a/docs/api/utilities.rst +++ /dev/null @@ -1,23 +0,0 @@ -Utilities -========= - -Array API ---------- - -.. automodule:: sequince.array_api - :members: - :undoc-members: - :show-inheritance: - -General utilities ------------------ - -.. automodule:: sequince.utils - :members: - :undoc-members: - :show-inheritance: - -.. automodule:: sequince.logging_utils - :members: - :undoc-members: - :show-inheritance: diff --git a/docs/citation.rst b/docs/citation.rst new file mode 100644 index 00000000..30bcefbd --- /dev/null +++ b/docs/citation.rst @@ -0,0 +1,8 @@ +Citing sequince +=============== + +If you use ``sequince`` in your work, please cite: + +.. TODO: add citation details (e.g. paper reference and/or Zenodo DOI, BibTeX). + +*Citation details coming soon.* diff --git a/docs/conf.py b/docs/conf.py index c439fb13..5fceeb6f 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -2,21 +2,13 @@ from __future__ import annotations -import os -import sys from datetime import datetime -# Ensure the package is importable without installation -PROJECT_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) -SRC_DIR = os.path.join(PROJECT_ROOT, "src") -if SRC_DIR not in sys.path: - sys.path.insert(0, SRC_DIR) - # -- Project information ----------------------------------------------------- project = "sequince" -copyright = f"{datetime.now():%Y}, Michael J. Williams" -author = "Michael J. Williams" +copyright = f"{datetime.now():%Y}, Sequince Developers" +author = "Sequince Developers" try: from importlib import metadata as importlib_metadata @@ -33,24 +25,28 @@ # -- General configuration --------------------------------------------------- extensions = [ - "sphinx.ext.autodoc", - "sphinx.ext.autosummary", + "autoapi.extension", + "myst_parser", "sphinx.ext.napoleon", "sphinx.ext.viewcode", "sphinx.ext.intersphinx", ] -autosummary_generate = True autodoc_typehints = "description" napoleon_google_docstring = False napoleon_numpy_docstring = True - -autodoc_mock_imports = [ - "numpy", - "scipy", - "matplotlib", - "emcee", - "minipcn", +# AutoAPI already emits attribute directives from source annotations. +napoleon_use_ivar = True + +# Generate module pages from source without importing optional backends. +autoapi_dirs = ["../src/sequince"] +autoapi_add_toctree_entry = False +# Document each object in its defining module to avoid ambiguous re-exports. +autoapi_options = [ + "members", + "undoc-members", + "show-inheritance", + "show-module-summary", ] templates_path = ["_templates"] diff --git a/docs/developer_guide/AI_POLICY.rst b/docs/developer_guide/AI_POLICY.rst new file mode 100644 index 00000000..27ca7c99 --- /dev/null +++ b/docs/developer_guide/AI_POLICY.rst @@ -0,0 +1,2 @@ +.. include:: ../../AI_POLICY.md + :parser: myst_parser.sphinx_ diff --git a/docs/developer_guide/CONTRIBUTING.rst b/docs/developer_guide/CONTRIBUTING.rst new file mode 100644 index 00000000..7a661a30 --- /dev/null +++ b/docs/developer_guide/CONTRIBUTING.rst @@ -0,0 +1,2 @@ +.. include:: ../../CONTRIBUTING.md + :parser: myst_parser.sphinx_ diff --git a/docs/developer_guide/index.rst b/docs/developer_guide/index.rst new file mode 100644 index 00000000..8b4f6a42 --- /dev/null +++ b/docs/developer_guide/index.rst @@ -0,0 +1,23 @@ +Developer guide +=============== + +These guides describe how to contribute to Sequince and the responsibilities +that apply to contributions. The pages include the maintained Markdown files +from the repository so updates appear here automatically. + +* The :doc:`contribution guidelines ` + cover development setup, building the documentation, code style, validation, + and the pull request process. +* The :doc:`AI policy ` + explains contributor ownership, review expectations, and validation of + AI-assisted work. + +For implementing new sampling components, see +:doc:`../user_guide/extending_sequince/index`. + +.. toctree:: + :maxdepth: 1 + :caption: Contents + + Contributing + AI policy diff --git a/docs/examples.rst b/docs/examples.rst index 66ee3b8c..2d407fe8 100644 --- a/docs/examples.rst +++ b/docs/examples.rst @@ -1,17 +1,5 @@ Examples ======== -The repository ships with a handful of executable scripts under ``examples/``. -They demonstrate how to assemble SMC pipelines for different scenarios. - -.. csv-table:: - :header: "Script", "Highlights" - :widths: 25, 75 - - ``adaptive_schedule.py``, "Adaptive tempering starting from the prior distribution." - ``adaptive_schedule_proposal.py``, "Tempering with an informative proposal distribution." - ``gmm_transdimensional.py``, "Adaptive annealing with on-the-fly component expansion for a Gaussian mixture model." - -Run any script with ``python examples/.py`` after installing the project -in a virtual environment. Each script exposes constants at the top that can be -tweaked to explore different behaviours. +An evolving set of code examples with ``sequince`` can be found in the +[sequince-examples](https://github.com/sequince-dev/sequince-examples) repository. diff --git a/docs/getting_started.rst b/docs/getting_started.rst deleted file mode 100644 index 4eb23b21..00000000 --- a/docs/getting_started.rst +++ /dev/null @@ -1,71 +0,0 @@ -Getting started -=============== - -Installation ------------- - -The project is published from the main branch and can be installed directly -from GitHub:: - - pip install git+https://github.com/sequince-dev/sequince.git - -``sequince`` targets Python 3.10+ and depends on ``numpy``, ``scipy``, -``matplotlib``, and ``emcee``. These will be installed automatically when the -package is installed via ``pip``. - -Quick start ------------ - -The library exposes a high-level helper, -``sequince.make_standard_bayesian_smc``, that wires together a sensible set of -steps for many Bayesian posterior inference workflows. The example below mirrors -the usage demonstrated in the README - -.. code-block:: python - - import numpy as np - from sequince import ParticleState, make_standard_bayesian_smc - - # User-supplied log-density functions - def log_prior(theta): - ... - - def log_likelihood(theta): - ... - - smc, posterior, metadata = make_standard_bayesian_smc( - log_prior, - log_likelihood, - ) - - shared_context = dict() - - # Build the initial particle cloud - n = 2_048 - initial_particles = np.random.normal(size=(n, 2)) - initial_state = ParticleState( - particles=initial_particles, - log_weights=np.full(n, -np.log(n)), - ) - - results = smc.run( - initial_state, - rng=np.random.default_rng(), - shared_context=shared_context, - metadata=metadata, - ) - samples = results.posterior_samples(1_000) - -Check the :doc:`user_guide` for a deeper look at how to customise the SMC loop -and swap in different building blocks. - - -Meta data & shared context --------------------------- - -The ``metadata`` dictionary is intended to hold information that needs to be -passed between steps in the SMC loop. It can be updated at each step. - -The ``shared`` dictionary is intended to hold information that needs to be -passed between steps in the SMC loop, but is not intended to be updated. For example, -it can be used to store static data that is needed by the likelihood function. diff --git a/docs/index.rst b/docs/index.rst index 6e719811..1e579966 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -5,15 +5,67 @@ sequince documentation Carlo (SMC) algorithms for Bayesian inference, particle filtering, and related sequential methods. +If you use ``sequince`` in your work, please follow our suggested :doc:`citation guidelines `. Thank you! + +Installation +------------ + +Install from PyPI:: + + pip install sequince + +``sequince`` targets Python 3.11+. Its core dependencies are installed +automatically. The example below uses the included MiniPCN mutation backend, +see :doc:`installation` for all optional integrations supported. + +.. _quick-start: + +Quick start +----------- + +Here, we fit a straight line, ``y = slope * x + intercept``, to synthetic data with +independent Gaussian measurement errors and uniform priors on both parameters. +The log-prior and log-likelihood callbacks accept particles shaped +``(n_particles, n_parameters)`` and return one log density per particle. + +.. literalinclude:: ../README.md + :language: python + :start-after: ```python + :end-before: ``` + +Next steps +---------- + +* :doc:`user_guide/running_inference/results` explains the returned result, + how to extract posterior samples, and how to inspect diagnostics; + :doc:`user_guide/running_inference/evidence_uncertainty` covers the evidence + estimate and its uncertainty. +* To customise the run, see :doc:`user_guide/running_inference/smc_configuration`, + :doc:`user_guide/sampling_options/resampling_configuration`, and + :doc:`user_guide/sampling_options/mutation_configuration`. +* The :doc:`user_guide/index` covers the full picture, starting with + :doc:`user_guide/core_concepts/index` (including the + :doc:`user_guide/core_concepts/step_context` distinction between shared + inputs and evolving metadata). +* :doc:`examples` links to worked examples in the + `sequince-examples `_ + repository. +* The :doc:`API reference ` lists classes, functions, + and their parameters. +* The :doc:`developer_guide/index` links to the contribution guidelines and + AI policy. + .. toctree:: :maxdepth: 2 :caption: Contents + :hidden: - getting_started - user_guide - checkpointing + installation + user_guide/index examples - api/index + API reference + developer_guide/index + citation Indices and tables ------------------ diff --git a/docs/installation.rst b/docs/installation.rst new file mode 100644 index 00000000..d758c080 --- /dev/null +++ b/docs/installation.rst @@ -0,0 +1,47 @@ +Installation +============ + +``sequince`` requires Python 3.11 or later. Install it from PyPI: + +.. code-block:: bash + + python -m pip install sequince + +The core installation includes NumPy, SciPy, MiniPCN, and the Array API +compatibility packages. It supports NumPy-based sampling with the default +MiniPCN mutation backend. + +Optional integrations +--------------------- + +Install extras for the integrations you need, for example: + +.. code-block:: bash + + python -m pip install "sequince[jax,blackjax,hdf5]" + +.. list-table:: + :header-rows: 1 + + * - Extra + - Provides + * - ``jax`` + - JAX arrays and functional execution + * - ``blackjax`` + - BlackJAX mutation algorithms + * - ``emcee`` + - emcee ensemble mutation + * - ``eryn`` + - Eryn ensemble mutation + * - ``hdf5`` + - HDF5 result storage and checkpointing + * - ``corner`` + - Corner plots + +For example, install ``sequince[emcee]`` to select emcee mutation instead of +the default MiniPCN backend. See +:doc:`user_guide/sampling_options/mutation_configuration` for mutation-specific +requirements. + +See :doc:`user_guide/core_concepts/array_backends` for backend selection, +accelerator setup, and SciPy Array API configuration. diff --git a/docs/requirements.txt b/docs/requirements.txt index 947b98c6..0d18d5b0 100644 --- a/docs/requirements.txt +++ b/docs/requirements.txt @@ -1,2 +1,3 @@ sphinx>=7.2 sphinx-book-theme>=1.1 +sphinx-autoapi>=3.8,<4 diff --git a/docs/user_guide.rst b/docs/user_guide.rst deleted file mode 100644 index 72818a7a..00000000 --- a/docs/user_guide.rst +++ /dev/null @@ -1,141 +0,0 @@ -User guide -========== - -The Sequential Monte Carlo (SMC) loop in ``sequince`` is built from a set of -composable *steps*. Each step inspects the current :class:`~sequince.state.ParticleState`, -optionally consumes observations and shared configuration, and returns an updated -state. At the end of each loop iteration, a set of *callbacks* are executed -that can perform operations like logging, diagnostics, and plotting. -The :class:`~sequince.smc.SequentialMonteCarlo` class drives the loop by -sequentially executing the configured steps and callbacks at every iteration. - -Configuring the runner ----------------------- - -The runner accepts a ``steps`` sequence containing callables that conform to the -``SMCStep`` protocol defined in :mod:`sequince.steps.base`. Similarly, the runner -accepts a sequence of optional ``callbacks`` that adhere to the -``SMCCallback`` protocol in :mod:`sequince.callbacks`. -The quickest way to build a useful pipeline is to reuse the pre-built helper: - -.. code-block:: python - - from sequince import make_standard_bayesian_smc - from sequince.callbackss import TracePlot - smc, posterior, shared_context = make_standard_bayesian_smc( - log_prior, - log_likelihood, - mutation="emcee", - callbacks = [TracePlot()], # Optional - ) - -``make_standard_bayesian_smc`` wires together -:class:`~sequince.steps.weighting.AdaptiveTemperedWeighting`, -:class:`~sequince.steps.resampling.SystematicResampling`, and a mutation step -chosen via the ``mutation`` argument. Pass ``mutation="minipcn"`` to opt into -the :class:`~sequince.steps.mutation.MiniPCNMutation` transition kernel or -``mutation="none"`` to skip rejuvenation. - -BlackJAX mutation ------------------ - -Install the optional BlackJAX integration with ``uv add "sequince[blackjax]"``. -BlackJAX mutation requires JAX-backed particles and JAX-compatible target -functions. NUTS is the default algorithm: - -.. code-block:: python - - import jax.numpy as jnp - - smc, posterior, metadata = make_standard_bayesian_smc( - log_prior, - log_likelihood, - mutation="blackjax", - mutation_n_steps=20, - mutation_kwargs={ - "algorithm": "nuts", - "algorithm_kwargs": { - "step_size": 0.1, - "inverse_mass_matrix": jnp.ones(n_dimensions), - }, - }, - ) - -``algorithm`` may also be ``"hmc"``, ``"random-walk"``, or a callable that -constructs a BlackJAX sampling algorithm from the current scalar log-density. -Acceptance and divergence diagnostics are returned in the mutation step output -under ``"acceptance"`` and ``"divergence"``. Per-iteration step outputs are -available through ``SMCResult.outputs``. - -If you need full control, construct ``SequentialMonteCarlo`` manually by -instantiating step and callback objects and supplying them in the desired order: - -.. code-block:: python - - from sequince import SequentialMonteCarlo - from sequince.steps import ( - AdaptiveTemperedWeighting, - SystematicResampling, - EmceeMutation, - ) - from sequince.callbacks import TracePlot - - smc = SequentialMonteCarlo( - steps=[ - AdaptiveTemperedWeighting(posterior, target_ess_ratio=0.8), - SystematicResampling(always_resample=True), - EmceeMutation(posterior, n_steps=10), - ], - callbacks=[TracePlot()], # Optional - ) - -Working with particle state ---------------------------- - -:class:`~sequince.state.ParticleState` stores the particle locations, -log-weights, accumulated log-evidence, and auxiliary metadata. A few convenience -methods are helpful when implementing new steps: - -``ParticleState.normalize()`` - Returns a fresh state with renormalised weights while keeping the original - untouched. - -``ParticleState.resample(indices)`` - Constructs a new state from ancestor indices, resetting the weights to a - uniform distribution. - -``ParticleState.with_updates(**kwargs)`` - Produces a copy with selected fields replaced. - -When possible, treat the state as immutable and return a new instance instead of -mutating in place. This keeps the iteration trace recorded in -:class:`~sequince.result.SMCResults` consistent and easier to debug. - -Incorporating observations --------------------------- - -The runner accepts an optional iterable of observations. Each iteration pulls -the next observation (if available) and exposes it through the -:class:`~sequince.steps.base.StepContext` supplied to steps. This allows for -particle filtering style updates where likelihood evaluations depend on time -indexed data. - -Custom steps and callbacks --------------------------- - -To implement bespoke step logic, derive from -:class:`~sequince.steps.base.SMCStep`. Override the ``apply`` method to -perform work and return ``(state, context, output)``. ``StepContext.shared`` -contains read-only user data, while current cross-step coordination values live -in ``StepContext.metadata``. Update metadata with -``context.with_metadata(...)`` and return the replacement context. Historical -diagnostics belong in the explicit step output; the runner stores ordered -per-iteration output tuples on :class:`~sequince.result.SMCResult`. - -Similarly, custom callbacks can be implemented by sub-classing -:class:`~sequince.callbacks.SMCCallback` and overriding the ``apply`` method -Callbacks are provided with the current state, context, and result history. An -observational callback returns ``None``; a callback that updates metadata or -stop controls returns a replacement context. - -See :mod:`sequince.steps` and :mod:`sequince.callbacks` for a catalogue of the built-in components. diff --git a/docs/checkpointing.rst b/docs/user_guide/advanced_workflows/checkpointing.rst similarity index 100% rename from docs/checkpointing.rst rename to docs/user_guide/advanced_workflows/checkpointing.rst diff --git a/docs/user_guide/advanced_workflows/index.rst b/docs/user_guide/advanced_workflows/index.rst new file mode 100644 index 00000000..4ba91404 --- /dev/null +++ b/docs/user_guide/advanced_workflows/index.rst @@ -0,0 +1,28 @@ +Advanced workflows +================== + +These guides cover running mutation algorithms independently, compiling SMC +iterations with JAX, and continuing a run from a saved checkpoint. Each builds +on the configuration and result interfaces introduced in +:doc:`../running_inference/index`. + +* :doc:`mcmc_sampler` shows how to use mutation steps for posterior sampling + with a standalone MCMC runner, then extract chains and diagnostics. +* :doc:`jax_support` explains functional execution with explicit random keys, + how to compile an iteration, and the requirements on steps, targets, and + runtime state. +* :doc:`checkpointing` covers checkpoint writers, saved state, and resuming + execution, including restoring the random-number-generator state. + +For JAX workflows, first configure your +:doc:`array backend <../core_concepts/array_backends>` and review +:doc:`../core_concepts/random_number_generation`. To select mutation and +resampling algorithms, see :doc:`../sampling_options/index`. + +.. toctree:: + :maxdepth: 1 + :caption: Contents + + mcmc_sampler + jax_support + checkpointing diff --git a/docs/user_guide/advanced_workflows/jax_support.rst b/docs/user_guide/advanced_workflows/jax_support.rst new file mode 100644 index 00000000..cc7aefc4 --- /dev/null +++ b/docs/user_guide/advanced_workflows/jax_support.rst @@ -0,0 +1,168 @@ +JAX and functional execution +============================ + +Sequince can run an SMC pipeline with JAX arrays and an explicitly threaded JAX +random key. The functional API separates pure iteration work from the host-side +loop, making one SMC iteration suitable for ``jax.jit`` while retaining +normal result collection, callbacks, and checkpoint handling between +iterations. + +Install JAX support with ``pip install "sequince[jax]"``. To use BlackJAX mutation, +install ``pip install "sequince[blackjax]"`` instead; BlackJAX brings its JAX +dependency with it. + +The functional API is not limited to BlackJAX. Any pipeline whose steps obey +the functional contract can use it; the pre-built MiniPCN pipeline is one other +supported option. + +Running a functional pipeline +----------------------------- + +Use JAX arrays for the particles and ensure that every target function called +inside the iteration uses JAX-compatible operations. Pass a JAX key as +``random_state`` rather than using the mutable ``rng`` argument accepted by +:meth:`~sequince.smc.SequentialMonteCarlo.run`: + +.. code-block:: python + + import jax + import jax.numpy as jnp + + from sequince import ParticleState, make_standard_bayesian_smc + from sequince.jax import compile_functional_iteration + + + def log_prior(particles, context): + inside = jnp.all(jnp.abs(particles) <= 5.0, axis=-1) + return jnp.where(inside, 0.0, -jnp.inf) + + + def log_likelihood(particles, context): + return -0.5 * jnp.sum(particles**2, axis=-1) + + + key = jax.random.key(42) + particle_key, key = jax.random.split(key) + initial_state = ParticleState( + particles=jax.random.uniform( + particle_key, + shape=(500, 2), + minval=-5.0, + maxval=5.0, + ) + ) + + smc, posterior, metadata = make_standard_bayesian_smc( + log_prior, + log_likelihood, + mutation="blackjax", + always_resample=True, + logging_frequency=0, + ) + + iteration = compile_functional_iteration( + smc, + initial_state, + key, + metadata=metadata, + ) + result, key = smc.run_functional( + initial_state, + random_state=key, + metadata=metadata, + iteration=iteration, + ) + +:meth:`~sequince.smc.SequentialMonteCarlo.run_functional` returns both the +:class:`~sequince.result.SMCResult` and the advanced random state. Reuse the +returned key, rather than the input key, when continuing random work. +See :doc:`../core_concepts/random_number_generation` for the corresponding eager-execution +contract and reproducibility guidance. + +The example selects BlackJAX's default NUTS mutation. Configure the number of +transitions with ``mutation_n_steps`` and pass algorithm options through +``mutation_kwargs``. Supported named algorithms are ``"nuts"``, ``"hmc"``, and +``"random-walk"``; a callable algorithm factory is also accepted. Mutation +diagnostics such as acceptance and divergence are recorded in +``result.outputs``. + +Compiling an iteration +---------------------- + +:func:`sequince.jax.compile_functional_iteration` initializes the state and +context fields required by the configured steps, lowers one iteration, verifies +that its input and output carry have matching tree structures, shapes, and +dtypes, and returns the compiled callable. Compilation does not execute or +discard a warm-up SMC transition. + +The ``observation``, ``shared_context``, and ``metadata`` values supplied while +compiling must have structures compatible with the subsequent run. Array values +may change between iterations, but their tree structure, shapes, and dtypes must +remain compatible with the compiled signature. + +For a simpler setup, JIT the bound iteration directly: + +.. code-block:: python + + iteration = jax.jit(smc.functional_iteration) + result, key = smc.run_functional( + initial_state, + random_state=key, + metadata=metadata, + iteration=iteration, + ) + +This compiles on the first call. Passing no ``iteration`` also works, but runs +``smc.functional_iteration`` without an external transformation. + +The functional step contract +---------------------------- + +An eager step returns ``(state, context, output)``. A functional step receives +an explicit random state and returns four values: + +.. code-block:: python + + def apply_functional(self, state, context, random_state): + # Perform array operations and advance random_state. + return state, context, next_random_state, output + +The default :meth:`~sequince.steps.base.SMCStep.apply_functional` adapter wraps +the backend random state in an ORNG generator and calls ``apply``. This is useful +when an eager implementation already consists of traceable array operations. +Override ``apply_functional`` when the functional implementation needs explicit +random-state control or JAX-specific control flow. + +Compiled carries must keep a stable pytree structure. A custom step that adds +metadata or state leaves during its first call should create shape-complete +placeholder leaves before compilation by overriding +``initialize_functional_metadata`` or +``initialize_functional_state_and_context``. The runner invokes these hooks +automatically at the beginning of :meth:`~sequince.smc.SequentialMonteCarlo.run_functional`, +and the compilation helper invokes them before lowering. + +Current functional limitations +------------------------------ + +* Functional resampling currently requires ``always_resample=True`` because a + traced ESS-triggered RNG transition is not yet implemented. +* A functional iteration always completes its configured step pipeline. An + immediate stop request is handled by the host loop after the iteration rather + than skipping later compiled steps. +* Python side effects and NumPy-only operations inside a compiled step or target + function are not JAX traceable. Keep logging in host-side callbacks or disable + the built-in logging step with ``logging_frequency=0``. +* Changing carry structure, array shape, or dtype triggers a new JAX signature; + the compilation helper rejects a pipeline whose input and output carry + signatures differ. + +The host loop in ``run_functional`` remains Python-driven. For a fixed-length +loop entirely inside a JAX transformation, call ``smc.functional_iteration`` +from ``jax.lax.scan`` and explicitly carry the state, context, and random +key. + +API reference +------------- + +See :func:`~sequince.jax.compile_functional_iteration` in the generated API +reference. diff --git a/docs/user_guide/advanced_workflows/mcmc_sampler.rst b/docs/user_guide/advanced_workflows/mcmc_sampler.rst new file mode 100644 index 00000000..ddb9dea7 --- /dev/null +++ b/docs/user_guide/advanced_workflows/mcmc_sampler.rst @@ -0,0 +1,120 @@ +Standalone MCMC +=============== + +Sequince's mutation implementations can be run as a standalone Markov chain +Monte Carlo sampler, without SMC weighting or resampling. This is useful for +ordinary posterior sampling and for testing a mutation configuration before +using it to rejuvenate an SMC population. + +Creating a sampler +------------------ + +Use :func:`~sequince.mcmc.make_mcmc` with a vectorised log-density. Each row of +the particle array is an independent chain: + +.. code-block:: python + + import numpy as np + + from sequince import ParticleState, make_mcmc + + + def log_prob(particles, context): + return -0.5 * np.sum(particles**2, axis=-1) + + + rng = np.random.default_rng(42) + initial_state = ParticleState( + particles=rng.normal(size=(32, 2)), + ) + + mcmc = make_mcmc( + log_prob, + mutation="minipcn", + n_steps=1_000, + ) + result = mcmc.run(initial_state, rng=rng) + +``n_steps`` is the number of MCMC transitions in the returned chain, not the +number of chains. The number of rows in ``initial_state.particles`` determines +the number of chains. + +Reading chains and diagnostics +------------------------------ + +:class:`~sequince.mcmc.MCMCResult.chain` has shape +``(n_steps, n_chains, n_dimensions)``. Use ``samples`` to discard burn-in, +thin, and optionally flatten the chain axis: + +.. code-block:: python + + chain = result.chain + samples = result.samples(burn_in=200, thin=5) + unflattened = result.samples(burn_in=200, thin=5, flat=False) + + final_positions = result.final_state.particles + diagnostics = result.diagnostics + +The flattened result has shape ``(retained_steps * n_chains, n_dimensions)``. +``diagnostics`` contains the final mutation block's backend-specific values, +such as acceptance information and the number of transitions. + +Choosing a mutation backend +--------------------------- + +``make_mcmc`` accepts the same named mutation backends used by the standard SMC +helper: ``"minipcn"``, ``"emcee"``, ``"blackjax"``, and +``"metropolis-hastings"``. Backend-specific settings are passed through +``mutation_kwargs``: + +.. code-block:: python + + mcmc = make_mcmc( + log_prob, + mutation="metropolis-hastings", + n_steps=1_000, + mutation_kwargs={ + "kernel": "gaussian-random-walk", + "kernel_kwargs": {"scale": 0.2}, + }, + ) + +See :doc:`../sampling_options/mutation_configuration` for available backends, optional dependencies, and +kernel configuration. A configured +:class:`~sequince.steps.mutation.base.MutationStep` can also be passed directly to +:class:`~sequince.mcmc.MarkovChainMonteCarlo`; the runner enables chain storage +on that step automatically. + +Context and particle transforms +------------------------------- + +A target may omit ``context`` when it does not need run-level inputs. When it +does, supply values using ``shared_context`` just as for SMC: + +.. code-block:: python + + def log_prob(particles, context): + scale = context.shared["scale"] + return -0.5 * np.sum((particles / scale) ** 2, axis=-1) + + + result = mcmc.run( + initial_state, + rng=42, + shared_context={"scale": 2.0}, + ) + +Pass a configured :class:`~sequince.transforms.ParticleTransform` as +``particle_transform`` when the chains should evolve in transformed +coordinates. Stored chains and final particles are returned in the original +particle coordinates. + +Checkpointing +------------- + +The standalone runner accepts callbacks and checkpoint writers. It executes +one mutation block per call to ``run``; a checkpoint therefore represents the +end of that block. Reconstruct the same mutation configuration, call +``resume_from(checkpoint)``, and pass the returned runtime to ``run`` to append +another block to the stored chain. See :doc:`checkpointing` for formats and RNG +restoration. diff --git a/docs/user_guide/core_concepts/array_backends.rst b/docs/user_guide/core_concepts/array_backends.rst new file mode 100644 index 00000000..04038bbd --- /dev/null +++ b/docs/user_guide/core_concepts/array_backends.rst @@ -0,0 +1,77 @@ +Array backends +============== + +Choosing an array backend +------------------------- + +Sequince uses the Python Array API through ``array-api-compat`` to work with +NumPy, JAX, PyTorch, and CuPy arrays. The initial particle array selects the +backend; there is no separate global backend setting. Install the array +library you intend to use: + +* NumPy is included in the core installation. +* For JAX, install ``sequince[jax]``. See the + `JAX installation guide `_ + for accelerator-specific packages. +* For PyTorch, follow the + `PyTorch installation selector `_ + for your platform and compute device. +* For CuPy, follow the + `CuPy installation guide `_ + to select a package compatible with your CUDA environment. + +PyTorch and CuPy are installed separately; Sequince does not provide extras +for them. Backend support also depends on the steps and mutation algorithms +in your pipeline. In particular, selecting JAX arrays alone does not compile +a run; see :doc:`../advanced_workflows/jax_support` for functional +execution and compilation requirements. + +Enabling SciPy Array API support +-------------------------------- + +SciPy's experimental Array API dispatch must be enabled before SciPy is +imported. Set the environment variable before starting Python: + +.. code-block:: bash + + export SCIPY_ARRAY_API=1 + python your_script.py + +Alternatively, set it at the beginning of your script, before importing +SciPy or packages that import SciPy: + +.. code-block:: python + + import os + + os.environ["SCIPY_ARRAY_API"] = "1" + + import sequince + +Sequince sets ``SCIPY_ARRAY_API=1`` during import if it is unset. It preserves +an existing value and warns if that value differs from ``"1"``. Setting it +explicitly avoids import-order problems. If SciPy has already been imported +in a notebook, restart the kernel and set the variable before importing it. +See `SciPy's Array API guide +`_ for details. + +Writing backend-compatible targets +---------------------------------- + +Create particles with your chosen library, and keep target evaluations in +the same array namespace. The following target works with each of the +backends above: + +.. code-block:: python + + from sequince.array_api import namespace + + + def log_likelihood(particles, context): + xp = namespace(particles) + return -0.5 * xp.sum(particles**2, axis=-1) + +Keep related arrays on compatible devices and use floating-point particles. +Avoid converting particles to NumPy inside a target when using another +backend. For JAX runs requiring 64-bit precision, call +``jax.config.update("jax_enable_x64", True)`` before creating arrays. diff --git a/docs/user_guide/core_concepts/index.rst b/docs/user_guide/core_concepts/index.rst new file mode 100644 index 00000000..aa1d485f --- /dev/null +++ b/docs/user_guide/core_concepts/index.rst @@ -0,0 +1,29 @@ +Core concepts +============= + +An SMC run passes a particle population, a context, and a source of randomness +through a sequence of steps. These guides explain how those objects fit +together and how the choice of array backend affects your target functions. + +* :doc:`array_backends` explains backend selection, SciPy Array API setup, + and writing target functions that use the same namespace as the particles. +* :doc:`particle_state` introduces particle locations, weights, and metadata, + how to return updated states, and how the result records iteration history. +* :doc:`step_context` describes shared inputs, evolving run metadata, + observations, and stop requests used to coordinate steps and callbacks. +* :doc:`random_number_generation` covers seeds, generator ownership, and + explicit random keys for reproducible eager and functional execution. + +Read these guides before building a custom pipeline or implementing new +steps. To configure a sampler and interpret its output, continue with +:doc:`../running_inference/index`; to implement your own components, see +:doc:`../extending_sequince/index`. + +.. toctree:: + :maxdepth: 1 + :caption: Contents + + array_backends + particle_state + step_context + random_number_generation diff --git a/docs/user_guide/core_concepts/particle_state.rst b/docs/user_guide/core_concepts/particle_state.rst new file mode 100644 index 00000000..bed3e9d3 --- /dev/null +++ b/docs/user_guide/core_concepts/particle_state.rst @@ -0,0 +1,70 @@ +Particle state and results +========================== + +:class:`~sequince.state.ParticleState` stores the particle locations, +log-weights, accumulated log-evidence, and auxiliary metadata. The state passed +to a step represents the current particle population; the step returns the +state that should be passed onward. + +When possible, treat a particle state as immutable and return a new instance. +This keeps the iteration trace recorded in +:class:`~sequince.result.SMCResult` consistent and easier to debug. + +Updating state +-------------- + +The main copy-returning operations are: + +``ParticleState.normalize()`` + Returns a state with normalized log-weights and the log-normalizer, while + leaving the original state untouched. + +``ParticleState.resample(indices)`` + Constructs a state from the selected ancestor indices and resets the + log-weights to a uniform distribution. + +``ParticleState.with_updates(**updates)`` + Returns a copy with the selected fields replaced. + +For example, a step can replace the particles without mutating its input: + +.. code-block:: python + + updated_state = state.with_updates(particles=proposed_particles) + +State and context serve different purposes. Particle-specific sampling data +belongs in ``ParticleState``. Run-level configuration and values shared between +steps belong in :doc:`step_context`. + +Result history +-------------- + +The runner records the state produced by each completed iteration in +:class:`~sequince.result.SMCResult`. ``result.final_state`` returns the final +population, while indexing the result accesses earlier populations. Step-level +diagnostics that do not describe particle state should be returned as the +step's explicit output instead; the runner stores these in ``result.outputs``. + +Summarising a result +-------------------- + +The final state remains weighted unless the last configured step resampled it. +Use the result helpers to preserve those weights when computing summaries: + +.. code-block:: python + + final_state = result.final_state + mean = result.weighted_mean() + covariance = result.weighted_covariance() + samples, weights = result.posterior_samples( + n_samples=1_000, method="multinomial", rng=42 + ) + +With ``method="multinomial"``, ``posterior_samples`` draws an equally weighted +sample with replacement from the final weighted population. The default +``method="weights"`` returns weighted particles. For diagnostics across the +complete run, use ``result.summary()``, +``result.effective_sample_sizes()``, ``result.log_evidence_history()``, and +``result.metadata_history(key)``. See :doc:`../running_inference/evidence_uncertainty` for +interpreting the final evidence estimate and its uncertainty, and +:doc:`../running_inference/results` for posterior extraction and nested-sampling results. diff --git a/docs/user_guide/core_concepts/random_number_generation.rst b/docs/user_guide/core_concepts/random_number_generation.rst new file mode 100644 index 00000000..697a89e2 --- /dev/null +++ b/docs/user_guide/core_concepts/random_number_generation.rst @@ -0,0 +1,117 @@ +Random number generation +======================== + +Sequince makes the source and ownership of randomness explicit so runs can be +reproduced, resumed, and executed with different array backends. Every +stochastic step receives its random-number generator from the runner; custom +steps should use that object rather than global functions such as +``numpy.random.random``. + +Eager execution +--------------- + +The ``rng`` argument accepted by +:meth:`~sequince.smc.SequentialMonteCarlo.run` may be: + +.. list-table:: + :header-rows: 1 + :widths: 35 65 + + * - Value + - Behaviour + * - An integer seed + - Creates a NumPy-backed ORNG generator. Reusing the seed for an otherwise + identical run reproduces the random stream. + * - A ``numpy.random.Generator`` + - Wraps and advances that mutable generator. + * - An ``orng.RandomGenerator`` + - Uses the supplied generator directly, including its configured backend. + * - ``None`` + - Creates a fresh NumPy-backed generator without a user-specified seed. + +For a simple reproducible run, pass an integer: + +.. code-block:: python + + result = smc.run(initial_state, rng=42, metadata=metadata) + +When initialisation is also random, create one generator and pass the same +object onward. The runner continues from its current position: + +.. code-block:: python + + import numpy as np + + rng = np.random.default_rng(42) + initial_particles = rng.normal(size=(500, 2)) + initial_state = ParticleState(initial_particles) + + result = smc.run(initial_state, rng=rng, metadata=metadata) + +The generator is mutable in eager execution. Supplying the same generator +object to another operation continues its stream; it does not restart from the +original seed. To reproduce a run, recreate both the initial state and runner +inputs from the same seed and configuration. + +RNG in custom steps +------------------- + +An eager step receives an :class:`orng.RandomGenerator` as its third argument: + +.. code-block:: python + + from sequince.steps.base import SMCStep + + + class GaussianJitter(SMCStep): + def apply(self, state, context, rng): + noise = rng.normal(size=state.particles.shape) + particles = state.particles + 0.01 * noise + return state.with_updates(particles=particles), context, None + +Do not create a new fixed-seed generator inside ``apply``. Doing so restarts the +same random sequence on every iteration and bypasses checkpointed RNG state. + +Functional and JAX execution +---------------------------- + +Functional execution does not mutate a generator. Pass backend-native state as +``random_state`` and retain the advanced state returned by the runner: + +.. code-block:: python + + key = jax.random.key(42) + result, key = smc.run_functional( + initial_state, + random_state=key, + metadata=metadata, + ) + +Each functional step accepts the current random state and must return its +replacement. Reusing the input key instead of the returned key repeats the same +random stream. See :doc:`../advanced_workflows/jax_support` for compiling the functional iteration +and splitting a key for random particle initialisation. + +Checkpointing and resuming +-------------------------- + +Checkpoint writers save RNG state by default. When that state is present, +``resume_from`` restores it and ignores any replacement RNG supplied to the +resume call. This makes a split run follow the same random stream as an +uninterrupted run. + +With ``save_rng=False``, pass ``rng`` to ``resume_from`` to deliberately start +the continuation from a new random state. Once a runtime has been prepared, +pass only ``runtime`` to ``run``; RNG and context arguments cannot be supplied +again. See :doc:`../advanced_workflows/checkpointing` for the complete resume workflow. + +Random posterior draws +---------------------- + +:meth:`~sequince.result.SMCResult.posterior_samples` also accepts a seed, +NumPy generator, or ORNG generator. Pass one explicitly when the resampled +posterior array must be reproducible: + +.. code-block:: python + + samples = result.posterior_samples(1_000, rng=42) diff --git a/docs/user_guide/core_concepts/step_context.rst b/docs/user_guide/core_concepts/step_context.rst new file mode 100644 index 00000000..98692545 --- /dev/null +++ b/docs/user_guide/core_concepts/step_context.rst @@ -0,0 +1,148 @@ +Step context +============ + +Every SMC step receives a :class:`~sequince.steps.context.StepContext` alongside +the particle state and random-number generator. The context describes the +current iteration and carries information that does not belong to an individual +particle. The same interface may also be passed to target functions and is +passed to callbacks. + +You normally do not construct a context yourself. The +:class:`~sequince.smc.SequentialMonteCarlo` runner creates one at the start of +each iteration from the arguments passed to +:meth:`~sequince.smc.SequentialMonteCarlo.run` or +:meth:`~sequince.smc.SequentialMonteCarlo.run_functional`. Steps may return +updated copies, and callbacks may return an updated copy or ``None``. + +Fields +------ + +``t`` + The zero-based index of the current SMC iteration. ``iteration`` is a + read-only alias for the same value. + +``observation`` + The observation selected for the current iteration, or ``None`` when no + observation is available. If ``num_steps`` is specified, a run can continue + after its observations are exhausted; those additional iterations receive + ``None``. See :doc:`../running_inference/observations` for complete iteration and exhaustion + semantics. + +``next_observation_index`` + The index at which the runner will look for the next observation. When the + current iteration consumed an observation at index ``i``, this value is + already ``i + 1``. It is primarily useful for checkpointing and resuming a + run; use ``observation`` to access the current value. + +``shared`` + User-supplied, run-wide configuration from ``shared_context``. Treat this + mapping as read-only. It is suitable for static data such as bounds, + constants, or model configuration. + +``metadata`` + Values that steps and callbacks communicate to later steps, callbacks, and + iterations. Use :meth:`~sequince.steps.context.StepContext.with_metadata` + to return a context containing additions or replacements. + +``stop_requested`` and ``stop_immediately`` + Control flags set by the stop-request methods described below. Most code + should use those methods rather than setting the fields directly. + +Shared data and metadata +------------------------ + +The distinction between ``shared`` and ``metadata`` is whether a value changes +during a run. Supply static dependencies once through ``shared_context``: + +.. code-block:: python + + result = smc.run( + initial_state, + rng=rng, + shared_context={"bounds": (-5.0, 5.0)}, + metadata={"accepted": 0}, + ) + +A custom step can read both mappings and publish updated metadata: + +.. code-block:: python + + from sequince.steps.base import SMCStep + + + class CountCalls(SMCStep): + def apply(self, state, context, rng): + bounds = context.shared["bounds"] + calls = context.metadata.get("calls", 0) + 1 + context = context.with_metadata(calls=calls) + return state, context, {"bounds_used": bounds} + +``StepContext`` is a frozen dataclass. Its helper methods return replacements +instead of changing the context object in place. ``with_metadata`` performs a +shallow copy and merge: keys supplied to the method replace existing keys, and +all other keys are retained. Values inside either mapping may still be mutable, +so avoid modifying them in place if iteration history or functional execution +must remain reproducible. + +The runner carries the final ``metadata`` mapping into the next iteration. +Within an iteration, each step receives the context returned by the preceding +step. After the steps finish, callbacks run in order; a callback that returns a +context passes it to the next callback, while returning ``None`` leaves the +context unchanged. The final context is available as ``result.context``. + +Stopping a run +-------------- + +:meth:`~sequince.steps.context.StepContext.request_stop` requests a deferred +stop. The remaining steps, automatic weight normalization, and callbacks for +the current iteration still run. The runner then finishes the run. + +:meth:`~sequince.steps.context.StepContext.request_stop_immediately` skips the +remaining steps in the current iteration. Automatic normalization and callbacks +still run before the runner finishes. Because callbacks execute after the step +pipeline, either stop method called from a callback stops after that iteration; +an immediate request does not skip later callbacks. + +Functional execution is intentionally different: a call to +:meth:`~sequince.smc.SequentialMonteCarlo.functional_iteration` always completes +the fixed step pipeline so that it remains suitable for compilation. In a +functional run, both kinds of stop request are handled by the host loop only +after the complete iteration. + +:meth:`~sequince.steps.context.StepContext.clear_stop_request` clears both stop +flags. This is useful in control-flow steps that conditionally override a stop +requested by an inner step. + +.. code-block:: python + + class StopAtIteration(SMCStep): + def __init__(self, final_iteration): + self.final_iteration = final_iteration + + def apply(self, state, context, rng): + if context.t >= self.final_iteration: + context = context.request_stop() + return state, context, None + +Context lifecycle +----------------- + +For each iteration, the runner performs the following operations: + +#. Selects the current observation and creates a fresh context with the current + ``t``, observation position, shared data, and metadata. +#. Passes the state and context through each configured step in order. Eager + execution can end this pipeline early for an immediate stop request. +#. Normalizes particle weights when ``auto_normalize=True``. +#. Records the state and step outputs, then invokes callbacks in order. +#. Writes any due checkpoints and either stops or advances to the next + iteration with the returned metadata. + +When resuming from a checkpoint, the iteration number, observation position, +and metadata are restored. ``shared_context`` is supplied again when creating +the resumed runtime because it is deliberately not stored in the checkpoint. + +API reference +------------- + +See :class:`~sequince.steps.context.StepContext` in the generated API reference. diff --git a/docs/user_guide/extending_sequince/custom_callbacks.rst b/docs/user_guide/extending_sequince/custom_callbacks.rst new file mode 100644 index 00000000..82cc9694 --- /dev/null +++ b/docs/user_guide/extending_sequince/custom_callbacks.rst @@ -0,0 +1,41 @@ +Custom callbacks +================ + +Callbacks run after the step pipeline and can inspect completed iterations, +report progress, or control the run. They receive a +:class:`~sequince.steps.context.StepContext`; see +:doc:`../core_concepts/step_context` for its lifecycle and update methods. +To perform work within the step pipeline, see :doc:`custom_steps`. + +For ready-made plotting callbacks and their configuration, see +:doc:`../running_inference/callbacks`. + +Derive a callback from :class:`~sequince.callbacks.SMCCallback` and implement +``apply``. A callback receives the state produced by the current iteration, the +current context, and the result accumulated so far: + +.. code-block:: python + + from sequince.callbacks import SMCCallback + + + class ReportProgress(SMCCallback): + def apply(self, state, context, result): + print(f"iteration={context.t}, states={len(result)}") + +An observational callback may return ``None``. To update metadata or request a +stop, return a replacement context: + +.. code-block:: python + + class StopAfter(SMCCallback): + def __init__(self, final_iteration): + self.final_iteration = final_iteration + + def apply(self, state, context, result): + if context.t >= self.final_iteration: + return context.request_stop() + return None + +Callbacks run in their configured order. A replacement context returned by one +callback is passed to the next callback and becomes ``result.context``. diff --git a/docs/user_guide/extending_sequince/custom_steps.rst b/docs/user_guide/extending_sequince/custom_steps.rst new file mode 100644 index 00000000..189dc7c9 --- /dev/null +++ b/docs/user_guide/extending_sequince/custom_steps.rst @@ -0,0 +1,30 @@ +Custom steps +============ + +Steps perform work within an SMC iteration, receiving the current particle +state and a :class:`~sequince.steps.context.StepContext`. For reporting and +run-level control after an iteration, see :doc:`custom_callbacks`. + +Derive a custom step from :class:`~sequince.steps.base.SMCStep` and implement +``apply``. It receives the current state, context, and random-number generator, +and must return ``(state, context, output)``: + +.. code-block:: python + + from sequince.steps.base import SMCStep + + + class RecordParticleMean(SMCStep): + def apply(self, state, context, rng): + mean = state.particles.mean(axis=0) + context = context.with_metadata(last_mean=mean) + return state, context, {"particle_mean": mean} + +Use ``context.shared`` for read-only configuration and +``context.with_metadata(...)`` for current coordination values. Return +historical diagnostics as the explicit output so that the runner records them +in the corresponding entry of :attr:`~sequince.result.SMCResult.outputs`. + +For compiled JAX execution, custom steps must also obey the functional +interface's fixed-structure requirements. See :doc:`../advanced_workflows/jax_support` for the +functional step contract and initialization hooks. diff --git a/docs/user_guide/extending_sequince/index.rst b/docs/user_guide/extending_sequince/index.rst new file mode 100644 index 00000000..86d83e10 --- /dev/null +++ b/docs/user_guide/extending_sequince/index.rst @@ -0,0 +1,31 @@ +Extending Sequince +================== + +Extend a pipeline with custom steps to update the particle population or +perform work within an iteration. Use callbacks to inspect completed +iterations, report progress, update run metadata, or request a stop. + +* :doc:`custom_steps` explains the step interface, returning updated state + and context, and recording diagnostics in the result. It includes an + example that records the particle mean. +* :doc:`custom_callbacks` explains the callback interface and returning + context updates. It includes examples that report progress or stop a run. + +Before implementing a component, review +:doc:`particle state <../core_concepts/particle_state>` and +:doc:`step context <../core_concepts/step_context>` to decide where its data +belongs. Stochastic steps should use the runner's generator as described in +:doc:`../core_concepts/random_number_generation`. + +For components that will run inside a compiled JAX iteration, also follow +the functional step contract and initialization requirements in +:doc:`../advanced_workflows/jax_support`. See +:doc:`../running_inference/smc_configuration` for assembling the components +into a pipeline. + +.. toctree:: + :maxdepth: 1 + :caption: Contents + + custom_steps + custom_callbacks diff --git a/docs/user_guide/index.rst b/docs/user_guide/index.rst new file mode 100644 index 00000000..c553ecae --- /dev/null +++ b/docs/user_guide/index.rst @@ -0,0 +1,24 @@ +User guide +========== + +The Sequential Monte Carlo (SMC) loop in ``sequince`` is assembled from +composable steps. Each step receives the current particle state and runtime +context, then returns an updated state, context, and optional output. +At the end of an iteration, callbacks can inspect the accumulated result and +update the context. This proceeds until a user-specified termination +criterion is satisfied, which returns a result object for further processing +and results visualisation. + +Start with the core concepts to understand the objects passed through the +SMC loop. The running-inference guides cover pre-built samplers, configuration, +and interpreting outputs. The remaining sections explain sampling choices and +more specialised execution and extension points. + +.. toctree:: + :maxdepth: 2 + + core_concepts/index + running_inference/index + sampling_options/index + advanced_workflows/index + extending_sequince/index diff --git a/docs/user_guide/running_inference/callbacks.rst b/docs/user_guide/running_inference/callbacks.rst new file mode 100644 index 00000000..c6f6c85f --- /dev/null +++ b/docs/user_guide/running_inference/callbacks.rst @@ -0,0 +1,68 @@ +Built-in callbacks +================== + +Callbacks inspect the state and accumulated result after each completed SMC +iteration. Pass a list through ``callbacks`` when constructing a pre-built +sampler or a ``SequentialMonteCarlo`` runner. + +Available callbacks +------------------- + +The following callbacks are available from :mod:`sequince.callbacks`: + +.. list-table:: + :header-rows: 1 + :widths: 35 65 + + * - Callback + - Purpose + * - :class:`~sequince.callbacks.TracePlot` + - Plot particle coordinates across the iteration history. + * - :class:`~sequince.callbacks.IntermediateParticlesPlot` + - Save a corner-style plot of the current particle population. + * - :class:`~sequince.callbacks.StatePlot` + - Plot selected metadata histories, ESS, or the normalizing constant. + * - :class:`~sequince.callbacks.AdaptiveMutationPlot` + - Plot completed mutation budgets, acceptance rates, and accepted moves. + * - :class:`~sequince.callbacks.MutationChainPlot` + - Plot stored chains from mutation step outputs. + * - :class:`~sequince.callbacks.TransformedMutationChainPlot` + - Plot stored mutation chains in the active transform's coordinates. + * - :class:`~sequince.callbacks.ParticleTransformPlot` + - Save a corner-style plot of the current particles in transformed + coordinates. + +Configuring plot callbacks +-------------------------- + +For example, configure a trace plot and histories of ESS and temperature: + +.. code-block:: python + + from sequince.callbacks import StatePlot, TracePlot + + callbacks = [ + TracePlot(outdir="plots", every=5), + StatePlot(metadata_keys=["ess", "beta"], outdir="plots", every=5), + ] + +Pass this list as ``callbacks=callbacks`` to the sampler constructor. +``outdir`` selects the output directory, which defaults to the current +directory. ``every`` controls the interval between callback invocations; +callbacks run on their first invocation and then at the requested interval. + +``TracePlot``, ``StatePlot``, and ``AdaptiveMutationPlot`` update their +configured ``filename``. The other plot callbacks use ``base_filename`` and +append the iteration number to save separate images. + +Chain plots require mutation steps configured with ``store_chains=True``. +``MutationChainPlot`` can select a step with ``step_index``; by default it +searches all step outputs. ``AdaptiveMutationPlot`` selects the mutation output +at ``step_index=2`` by default, so adjust this for your pipeline. +Transform plots require a fitted transform in context metadata, selected by +``metadata_key``. Corner-style plots require the ``corner`` extra; see +:doc:`../../installation` for optional dependencies. + +For your own reporting or run-control logic, see +:doc:`../extending_sequince/custom_callbacks`. Checkpoint writers are configured +separately; see :doc:`../advanced_workflows/checkpointing`. diff --git a/docs/user_guide/running_inference/evidence_uncertainty.rst b/docs/user_guide/running_inference/evidence_uncertainty.rst new file mode 100644 index 00000000..6a2c1cab --- /dev/null +++ b/docs/user_guide/running_inference/evidence_uncertainty.rst @@ -0,0 +1,126 @@ +Evidence and uncertainty +======================== + +For Bayesian inference, the evidence is the integral of the likelihood against +the normalized prior. It is the normalizing constant of the unnormalized +posterior density. Interpreting this value as Bayesian evidence requires +retaining the likelihood's normalization constants. + +For tempered SMC, the full posterior is reached at ``beta=1``. Evidence at an +earlier temperature is the normalizing constant of that intermediate target. +For nested-sampling SMC, use the shell-evidence interface described in +:doc:`results`; the ordinary state normalizing constant tracks constrained-prior +mass instead. The uncertainty estimators below apply to ordinary SMC weighting. + +:class:`~sequince.result.SMCResult` provides three estimates of log-evidence +uncertainty from a single SMC run: + +.. list-table:: + :header-rows: 1 + + * - Method + - Use + - Main limitation + * - ``"incremental"`` + - Default; uses each incremental weight update + - Ignores dependence between SMC iterations + * - ``"lee-whiteley"`` + - Uses the complete particle genealogy + - Requires multinomial resampling at every iteration + * - ``"fixed-lag-lee-whiteley"`` + - Uses only recent particle ancestry + - Trades genealogical stability for truncation bias + +Usage +----- + +``log_evidence_variance`` estimates +:math:`\operatorname{Var}(\log \widehat Z)`. ``log_evidence_error`` returns the +corresponding standard error: + +.. code-block:: python + + log_evidence = result.log_evidence() + variance = result.log_evidence_variance() + error = result.log_evidence_error() + + lw_variance = result.log_evidence_variance(method="lee-whiteley") + + fixed_lag_variance = result.log_evidence_variance( + method="fixed-lag-lee-whiteley", + lag=10, + ) + +These values describe Monte Carlo error in the SMC evidence estimate, not +uncertainty in the statistical model or posterior. + +Assumptions and limitations +--------------------------- + +``incremental`` + This estimator is inexpensive, requires no ancestry, and supports adaptive + temperatures and adaptive resampling. It ignores dependence between + successive SMC iterations and may therefore underestimate uncertainty. + + It requires ``log_evidence_increment_variance`` in every particle state's + metadata. This is recorded by + :class:`~sequince.steps.weighting.FixedTemperedWeighting` and + :class:`~sequince.steps.weighting.AdaptiveTemperedWeighting`; missing values + produce ``nan``. The correct particle-count scaling is already present in + the squared normalised weights, so no extra :math:`N` or :math:`N-1` factor + is needed. + +``lee-whiteley`` + This is the full-genealogy estimator described by `Lee and Whiteley (2018) + `_. It requires at least two + particles, a constant particle count, exactly one resampling step per + iteration, resampling at every iteration, and multinomial resampling. + + ``sequince`` records the scheme used by its built-in resampling steps. The + estimator warns when it finds systematic, stratified, missing, or unknown + scheme information. A custom resampling step should set its + ``resampling_scheme`` class attribute. The calculation still returns a value + after a warning, but it does not then have the stated theoretical + justification. The estimator also becomes unreliable as the time-zero + ancestor families collapse. + +``fixed-lag-lee-whiteley`` + This estimator uses only the most recent resampling events. Short lags retain + more ancestor families but omit older variance contributions; long lags + approach the full-genealogy estimator and its degeneracy. + + ``lag`` counts resampling events, not temperature updates or SMC iterations. + Skipped adaptive-resampling events do not consume lag. The constant particle + count and multinomial-resampling assumptions still apply, and each iteration + must contain exactly one resampling-step output. Mutation should not mix + particle paths. + + Recent-ancestor variance estimation is discussed by `Olsson and Douc (2019) + `_. Their estimator targets normalised + particle expectations rather than evidence. The method provided here is an + evidence-specific fixed-lag Lee--Whiteley estimator. + +Diagnostics +----------- + +Inspect the selected genealogy before relying on a fixed-lag estimate: + +.. code-block:: python + + diagnostics = result.fixed_lag_lee_whiteley_diagnostics(lag=10) + print(diagnostics.traced_resampling_events) + print(diagnostics.distinct_ancestors) + print(diagnostics.family_ess) + +The diagnostics report the requested and traced lag, total resampling events, +distinct selected-lag ancestors, and family ESS. A family ESS near one indicates +severe genealogical depletion. Comparing several plausible lags is generally +more informative than reporting one arbitrary choice. + +Finite-sample Lee--Whiteley estimates can be negative. The variance method +preserves such values because they signal an unreliable estimate; +``log_evidence_error`` raises ``ValueError`` because no real-valued standard +error can then be computed. + +Repeated independent SMC runs remain the most direct uncertainty check when +they are affordable. diff --git a/docs/user_guide/running_inference/index.rst b/docs/user_guide/running_inference/index.rst new file mode 100644 index 00000000..465c665e --- /dev/null +++ b/docs/user_guide/running_inference/index.rst @@ -0,0 +1,36 @@ +Running inference +================= + +This section explains how to choose a pre-built sampler, configure a run, +and interpret its results and evidence estimate (normalizing constant). +Custom pipelines can also be assembled from composable steps. + +Configuring inference +--------------------- + +Start with the available pre-built samplers, then configure the SMC runner. +For sequential data or particle filtering, see :doc:`observations` for +passing an observation stream to the runner. + +.. toctree:: + :maxdepth: 1 + :caption: Contents + + prebuilt_samplers + smc_configuration + callbacks + observations + +Understanding outputs +--------------------- + +After a run, inspect the particle population, extract posterior samples, and +check diagnostics. The evidence guide explains the normalizing constant and +the assumptions behind its uncertainty estimates. + +.. toctree:: + :maxdepth: 1 + :caption: Contents + + results + evidence_uncertainty diff --git a/docs/user_guide/running_inference/observations.rst b/docs/user_guide/running_inference/observations.rst new file mode 100644 index 00000000..e36131db --- /dev/null +++ b/docs/user_guide/running_inference/observations.rst @@ -0,0 +1,185 @@ +Using observations +================== + +A common use-case of SMC is to iteratively process a sequence of observations +constituting a data stream. To operate `sequince` in this mode, an iterable can +be passed to the SMC runner, which consumes at most one item at the beginning +of each iteration and exposes it as +:attr:`~sequince.steps.context.StepContext.observation` to each SMC step / callback. + +As the precise structure of an observation depends only on what the constituent SMC +steps expect, it can be any object: a scalar measurement, an array, a batch, or a +structured value containing several fields. + +A sequential likelihood update +------------------------------ + +For particle filtering or online Bayesian updates, an incremental weighting +step can evaluate the likelihood contribution from the current observation: + +.. code-block:: python + + import numpy as np + + from sequince import ParticleState, SequentialMonteCarlo + from sequince.steps import LogLikelihoodWeighting, SystematicResampling + + + def log_likelihood_increment(particles, context): + observation = context.observation + variance = context.shared["observation_variance"] + residual = observation - particles[:, 0] + return -0.5 * residual**2 / variance + + + smc = SequentialMonteCarlo( + steps=[ + LogLikelihoodWeighting( + log_likelihood_increment, + normalize_after=False, + ), + SystematicResampling(ess_threshold=0.5), + ] + ) + + initial_state = ParticleState( + particles=np.random.default_rng(42).normal(size=(500, 1)) + ) + result = smc.run( + initial_state, + observations=[0.2, -0.1, 0.4], + rng=42, + shared_context={"observation_variance": 1.0}, + ) + +This example run produces three states: ``result.states[i]`` is the population after +processing ``observations[i]``. The state before any observations were processed +is available as ``result.initial_state``. + +Interplay with ``num_steps`` +---------------------------- + +The interaction between ``observations`` and ``num_steps`` is important: + +.. list-table:: + :header-rows: 1 + :widths: 30 25 45 + + * - ``observations`` + - ``num_steps`` + - Behavior + * - Finite iterable + - ``None`` + - Run until the iterable is exhausted. + * - Finite iterable + - Positive integer ``N`` + - Run at most ``N`` iterations. If the iterable ends first, later + iterations receive ``observation=None``. + * - Infinite iterable + - Positive integer ``N`` + - Consume the first ``N`` items and stop. + * - ``None`` + - Positive integer ``N`` + - Run ``N`` iterations, each with ``observation=None``. + * - ``None`` or infinite iterable + - ``None`` + - Continue until a step or callback requests a stop. + +A stop request can always end a run before either limit. Conversely, passing +neither a finite observation stream nor ``num_steps`` requires the pipeline to +request its own stop; otherwise the run will continue indefinitely. + +A run must complete at least one iteration. Consequently, an empty observation +iterable with no ``num_steps`` produces no state and raises ``ValueError``. + +Iterables and generators +------------------------ + +The runner consumes a non-sequence iterable lazily, so a generator can represent +a stream without first loading it into memory: + +.. code-block:: python + + def measurements(source): + for record in source: + yield record.value + + + result = smc.run( + initial_state, + observations=measurements(source), + rng=rng, + ) + +If ``num_steps`` is omitted, normal generator exhaustion ends the run. If it is +specified, exhaustion instead supplies ``None`` through the remaining requested +iterations. A step that permits ``None`` observations should handle that case +explicitly. + +Observation position +-------------------- + +While ``context.t`` is the SMC iteration, it is not necessarily an +observation index because iterations may continue after the observations are exhausted. +``context.next_observation_index`` counts observations already consumed and is +the position from which a resumable sequence will continue. During the +iteration processing ``observations[i]``, it is already ``i + 1``. + +Use ``context.observation`` to read the current item. The next-index field is +primarily runtime and checkpoint state; it should not normally be used to index +the current observation. + +Checkpointing and resuming +-------------------------- + +Checkpoints store ``next_observation_index`` but do not store the full input +stream, as the flexible nature of this iterable (and its contents) means it is not +guaranteed to be compatible with inbuilt checkpointing schemes. +When resuming a run that has already consumed observations, pass the +original indexable sequence to ``run``: + +.. code-block:: python + + runtime = smc.resume_from( + checkpoint, + shared_context=shared_context, + ) + result = smc.run( + runtime=runtime, + observations=all_observations, + ) + +The runner resumes with ``all_observations[checkpoint.next_observation_index]``. +As generators cannot be repositioned, they are not compatible with this procedure. +See :doc:`../advanced_workflows/checkpointing` for the full checkpoint workflow. + +Observations with the JAX backend +--------------------------------- + +The jit-compilable :meth:`~sequince.smc.SequentialMonteCarlo.run_functional` follows the same +same behaviour as the standard runner regarding observations. To use a compiled iteration, +all observations must be compatible with tracing (e.g. static shapes and dtypes). +For easy compilation, supply a representative first observation to the compilation helper: + +.. code-block:: python + + iteration = compile_functional_iteration( + smc, + initial_state, + key, + observation=observations[0], + next_observation_index=1, + shared_context=shared_context, + metadata=metadata, + ) + result, key = smc.run_functional( + initial_state, + observations=observations, + random_state=key, + shared_context=shared_context, + metadata=metadata, + iteration=iteration, + ) + +See :doc:`../advanced_workflows/jax_support` for compilation, random-state threading, and fixed-carry +requirements. diff --git a/docs/user_guide/running_inference/prebuilt_samplers.rst b/docs/user_guide/running_inference/prebuilt_samplers.rst new file mode 100644 index 00000000..71935e4b --- /dev/null +++ b/docs/user_guide/running_inference/prebuilt_samplers.rst @@ -0,0 +1,55 @@ +Pre-built samplers +================== + +``sequince`` provides three helpers for assembling Bayesian SMC pipelines. +Each returns a runner, its target, and initial metadata. Supply your own +initial :class:`~sequince.state.ParticleState` and pass the returned metadata +to ``runner.run(..., metadata=metadata)``. + +Standard Bayesian SMC +--------------------- + +:func:`~sequince.prebuilt_samplers.make_standard_bayesian_smc` moves a +particle population through a sequence of tempered distributions, combining +weighting, resampling, and mutation. It uses adaptive temperatures by default; +``adaptive=False`` with ``num_steps`` selects a fixed temperature schedule. +The usual posterior endpoint is ``max_beta=1.0``. + +This is the helper used in the :ref:`quick start `. See +:doc:`smc_configuration` for configuration and manual pipeline assembly. + +Persistent Bayesian SMC +----------------------- + +:func:`~sequince.prebuilt_samplers.make_persistent_bayesian_smc` retains an +archive of previous particle batches and reweights that archive as the +temperature changes. ``n_particles`` sets the active batch size; +``target_ess`` and ``target_total_ess`` are absolute effective sample sizes +of the archive. + +The initial population must contain a positive multiple of ``n_particles`` +equally weighted draws from the prior, or from ``log_proposal`` when supplied. +By default, the sampler finishes with an equally weighted active population. +Set ``resample_at_end=False`` to return the complete weighted archive instead. + +Nested-sampling SMC +------------------- + +:func:`~sequince.prebuilt_samplers.make_nested_sampling_smc` raises a +likelihood threshold and explores the remaining constrained prior. It +accumulates evidence over the discarded likelihood shells and the final +live population. Initial particles must represent a normalized prior; +arbitrary importance proposals are not supported. + +Use ``threshold_quantile`` or ``threshold_schedule`` to control the thresholds. +``evidence_tolerance`` controls an estimated remaining-evidence stopping +criterion; it is not a guaranteed error bound. + +Read evidence and reconstruct posterior samples through +``result.as_nested_sampling_result()``. Posterior reconstruction requires a +completed run and ``store_shell_particles=True`` (the default). See +:doc:`results` for the corresponding result interface. + +All three helpers expose mutation, resampling, and particle-transform options; +see :doc:`../sampling_options/index` and their API references for supported settings. +For standalone MCMC, see :doc:`../advanced_workflows/mcmc_sampler`. diff --git a/docs/user_guide/running_inference/results.rst b/docs/user_guide/running_inference/results.rst new file mode 100644 index 00000000..1d3d856a --- /dev/null +++ b/docs/user_guide/running_inference/results.rst @@ -0,0 +1,71 @@ +Working with results +==================== + +The SMC runner returns an :class:`~sequince.result.SMCResult`. It contains +the state after each completed iteration in ``result.states``, the starting +population in ``result.initial_state``, and step outputs in ``result.outputs``. +``result.final_state`` is the final population; ``result[i]`` accesses an +earlier population. See :doc:`../core_concepts/particle_state` for the particle-state interface. + +Posterior samples and summaries +------------------------------- + +For standard and persistent Bayesian SMC runs that have reached the posterior +target, use the final population to compute weighted summaries: + +.. code-block:: python + + mean = result.weighted_mean() + covariance = result.weighted_covariance() + samples, weights = result.posterior_samples() + +The default ``method="weights"`` returns the final particles and their +normalized weights. To draw an equally weighted sample with replacement: + +.. code-block:: python + + samples, weights = result.posterior_samples( + n_samples=1_000, + method="multinomial", + rng=42, + ) + +Respect the weights when analysing the particles directly. Intermediate +tempered populations do not yet represent the full posterior, and +resampling does not make particles independent. + +Run diagnostics +--------------- + +Inspect the run summary and histories to understand how the population evolved: + +.. code-block:: python + + summary = result.summary() + ess = result.effective_sample_sizes() + beta = result.metadata_history("beta") + log_evidence = result.log_evidence() + log_evidence_history = result.log_evidence_history() + +The ``beta`` history is relevant to tempered runs. ESS describes weight +concentration; it does not measure the number of independent posterior draws. +Step-specific diagnostics are available in ``result.outputs``. +See :doc:`evidence_uncertainty` for interpreting evidence uncertainty. + +Nested-sampling results +----------------------- + +For nested-sampling SMC, use the shell-weighted result interface: + +.. code-block:: python + + nested = result.as_nested_sampling_result() + complete = nested.complete + log_evidence = nested.log_evidence + # Once complete, with store_shell_particles=True: + samples, weights = nested.posterior_samples() + +The final state represents the constrained prior. Bayesian evidence and +posterior weights come from the accumulated shells, so use this interface +for posterior analysis. A run suspended by ``num_steps`` must be resumed to +completion before extracting posterior samples. diff --git a/docs/user_guide/running_inference/smc_configuration.rst b/docs/user_guide/running_inference/smc_configuration.rst new file mode 100644 index 00000000..9cd543a7 --- /dev/null +++ b/docs/user_guide/running_inference/smc_configuration.rst @@ -0,0 +1,74 @@ +Configuring an SMC runner +========================= + +The :class:`~sequince.smc.SequentialMonteCarlo` runner applies a sequence of +steps and optional callbacks at every iteration. The quickest way to build a +Bayesian inference pipeline is the pre-built helper: + +.. code-block:: python + + from sequince import make_standard_bayesian_smc + from sequince.callbacks import TracePlot + + smc, posterior, metadata = make_standard_bayesian_smc( + log_prior, + log_likelihood, + mutation="emcee", + resampling="systematic", + callbacks=[TracePlot()], + ) + +``make_standard_bayesian_smc`` constructs an SMC runner that adaptively +steps through a series of simulated annealing temperatures +(:class:`~sequince.steps.weighting.AdaptiveTemperedWeighting`), applying a resampling step (:class:`~sequince.steps.resampling.SystematicResampling`) and a mutation step +(:class:`~sequince.steps.mutation.base.MutationStep`) at each iteration. +See :doc:`../sampling_options/resampling_configuration` and :doc:`../sampling_options/mutation_configuration` for the available +reasmpling and mutation options and configuration. + +See :doc:`callbacks` for the available plotting callbacks and their settings. + +When starting the run, pass the ``metadata`` returned by the convenience constructor: + +.. code-block:: python + + result = smc.run( + initial_state, + rng=rng, + shared_context=shared_context, + metadata=metadata, + ) + +See :doc:`../core_concepts/step_context` for the distinction between static ``shared_context`` +data and evolving metadata. + +Building a pipeline manually +---------------------------- + +For full control, instantiate the runner with steps in execution order. For example, +to reproduce the pre-built SMC defined above: + +.. code-block:: python + + from sequince import SequentialMonteCarlo + from sequince.callbacks import TracePlot + from sequince.steps import ( + AdaptiveTemperedWeighting, + EmceeMutation, + SystematicResampling, + ) + + smc = SequentialMonteCarlo( + steps=[ + AdaptiveTemperedWeighting(posterior), + SystematicResampling(always_resample=True), + EmceeMutation(posterior, n_steps=10), + ], + callbacks=[TracePlot()], + ) + +Each entry in ``steps`` must follow the +:class:`~sequince.steps.base.SMCStep` interface. Each callback must follow the +:class:`~sequince.callbacks.SMCCallback` interface. See +:doc:`../extending_sequince/custom_steps` and +:doc:`../extending_sequince/custom_callbacks` when implementing these +interfaces yourself. diff --git a/docs/user_guide/sampling_options/index.rst b/docs/user_guide/sampling_options/index.rst new file mode 100644 index 00000000..c1217851 --- /dev/null +++ b/docs/user_guide/sampling_options/index.rst @@ -0,0 +1,31 @@ +Sampling options +================ + +Resampling and mutation control how a particle population is carried between +SMC targets. Resampling draws from the weighted population and assigns equal +weights to the selected particles. Mutation applies MCMC transitions that +explore the current target while preserving the particle weights. + +The pre-built samplers expose these choices through their ``resampling`` and +``mutation`` arguments. For a custom pipeline, configure the corresponding +steps directly. See :doc:`../running_inference/prebuilt_samplers` for the +available sampler designs. + +* :doc:`resampling_configuration` explains systematic, stratified, and + multinomial resampling, when to trigger resampling, and how to inspect + ancestry and ESS diagnostics. +* :doc:`mutation_configuration` covers mutation algorithms, transition counts, + proposal kernels, particle transforms, and mutation diagnostics. + +Choose options compatible with your array backend and execution mode; see +:doc:`../core_concepts/array_backends` and +:doc:`../advanced_workflows/jax_support`. If you use genealogy-based evidence +uncertainty estimates, check their resampling assumptions in +:doc:`../running_inference/evidence_uncertainty` before choosing a scheme. + +.. toctree:: + :maxdepth: 1 + :caption: Contents + + resampling_configuration + mutation_configuration diff --git a/docs/user_guide/sampling_options/mutation_configuration.rst b/docs/user_guide/sampling_options/mutation_configuration.rst new file mode 100644 index 00000000..f03c5f7c --- /dev/null +++ b/docs/user_guide/sampling_options/mutation_configuration.rst @@ -0,0 +1,160 @@ +Mutation configuration +====================== + +Mutation steps apply MCMC transitions to a particle population without changing +its weights. Use them in an SMC pipeline or with the standalone +:class:`~sequince.mcmc.MarkovChainMonteCarlo` runner; see :doc:`../advanced_workflows/mcmc_sampler`. +Each step takes a target callable or :class:`~sequince.target.Target` object. + +.. list-table:: + :header-rows: 1 + :widths: 18 24 28 30 + + * - Method + - ``mutation`` name + - Implementation + - Typical use + * - MiniPCN + - ``"minipcn"`` + - :class:`~sequince.steps.mutation.minipcn.MiniPCNMutation` + - Supports eager and functional execution. + * - emcee + - ``"emcee"`` + - :class:`~sequince.steps.mutation.emcee.EmceeMutation` + - Ensemble MCMC using emcee moves. + * - BlackJAX + - ``"blackjax"`` + - :class:`~sequince.steps.mutation.blackjax.BlackJAXMutation` + - JAX-backed NUTS, HMC, or random-walk mutation. + * - Eryn + - ``"eryn"`` + - :class:`~sequince.steps.mutation.eryn.ErynMutation` + - Fixed-dimensional ensemble MCMC using Eryn moves. + * - Differential evolution + - ``"differential-evolution"`` + - :class:`~sequince.steps.mutation.metropolis.DifferentialEvolutionMutation` + - Red-blue ensemble updates using differential-evolution proposals. + * - Red-blue Metropolis--Hastings + - ``"red-blue-metropolis-hastings"``, ``"red-blue-mh"``, ``"red-blue"`` + - :class:`~sequince.steps.mutation.metropolis.RedBlueMetropolisHastingsMutation` + - Custom kernels using a frozen complementary ensemble. + * - Metropolis--Hastings + - ``"metropolis-hastings"`` + - :class:`~sequince.steps.mutation.metropolis.MetropolisHastingsMutation` + - User-supplied proposal kernel and transition correction. + +MH steps take a proposal through the ``kernel`` argument. See +:mod:`sequince.steps.mutation.kernels` for Gaussian random walks, masked +updates, mixtures, and differential-evolution kernels. Red-blue MH requires a +:class:`~sequince.steps.mutation.kernels.ComplementaryEnsembleKernel`. + +Common configuration +-------------------- + +Set ``n_steps`` on the mutation step to choose the number of transitions per +particle per call. For red-blue methods, one transition is a sweep over both +groups. For a target callable or ``Target`` object named ``target``: + +.. code-block:: python + + from sequince.steps.mutation import MiniPCNMutation + from sequince.transforms import AffineTransform + + mutation_step = MiniPCNMutation( + target, + n_steps=10, + store_chains=True, + transform=AffineTransform(), + ) + +The implementations in the table accept these constructor options: + +* ``store_chains``: store intermediate positions (default: ``False``). +* ``transform_refit``: fit the transform on every call (``"always"``, default), + once per iteration and observation position (``"per-context"``), or once + (``"once"``). Eager execution also accepts a callable predicate. +* ``metrics``: diagnostic objects such as ``ESJDMetric()`` or + ``RelativeJumpDistanceMetric()``. Results appear in the output's + ``metrics.values``; chains need not be stored. + +In a custom SMC pipeline, place mutation where it should run, usually after +resampling. Wrap it in +:class:`~sequince.steps.mutation.conditional.MutateIfResampled` to run it only +when resampling occurs. + +Install optional backends with ``pip install "sequince[emcee]"``, +``pip install "sequince[eryn]"``, or ``pip install "sequince[blackjax]"``. +See :doc:`../advanced_workflows/jax_support` for compiled execution. + +Particle transforms +------------------- + +Pass a :class:`~sequince.transforms.ParticleTransform` instance through the +step's ``transform`` argument. ``AffineTransform`` standardizes coordinates; +``WhiteningTransform`` also removes correlations. Bounds and periodic +transforms are available in :mod:`sequince.transforms`. + +Configure the transform's parameters when constructing it. Mutation maps +particles into transformed coordinates and back after the transitions. With +``transform=None`` (the constructor default), it uses the original coordinates. + +Auxiliary data +~~~~~~~~~~~~~~ + +A custom transform's ``forward`` method can return +``(values, log_det, auxiliary)``. Mutation passes ``auxiliary`` to the matching +``inverse`` call, which returns ``(particles, log_det)``. For example, a +transform taking absolute values can save each particle's signs for inversion. + +Auxiliary data can be an array or an Optree-compatible pytree, such as a tuple +or dictionary. Use :func:`~sequince.utils.pytrees.pytree_dataclass` to register +immutable dataclasses with Optree and JAX. + +When mutation selects a particle subset, it also selects auxiliary array leaves +with ``leaf.ndim > 0`` and ``leaf.shape[0] == n_particles``. Other leaves are +unchanged. This rule uses shape alone, so unrelated arrays of that leading size +are also selected. Particle-aligned auxiliary data works with batch-preserving +mutation and BlackJAX; emcee and Eryn reject it because their callbacks do not +expose the original walker indices. + +Blocked updates, repetition, and adaptation +------------------------------------------- + +Use these steps to compose mutation updates: + +* :class:`~sequince.steps.mutation.blocked.BlockedMetropolisHastingsMutation`: + sequential parameter-block updates, each with its own kernel and step count. + Block indices refer to transformed coordinates. +* :class:`~sequince.steps.mutation.repeated.RepeatedMutationStep`: split a total + transition count across fixed rounds, with an optional step after each round. +* :class:`~sequince.steps.mutation.adaptive.steps.AdaptMutationSteps`: choose + the next iteration's transition count from completed diagnostics. +* :class:`~sequince.steps.mutation.adaptive.minipcn.AdaptiveMiniPCNMutation`: + tune MiniPCN rho between iterations. Requires equal, finite weights and a + child with ``adaptive=False``. +* :class:`~sequince.steps.mutation.adaptive.de.AdaptDEScaleWeights`: tune DE + mixture weights between iterations. +* :class:`~sequince.steps.mutation.adaptive.proposal.AdaptiveGaussianProposal`: + update a scale in context metadata for a kernel factory to read. + +The adaptation wrappers support eager execution only. Place conditional steps +outside them, for example ``MutateIfResampled(mutation_step)``. + +Prebuilt samplers +----------------- + +:func:`~sequince.prebuilt_samplers.make_standard_bayesian_smc` accepts a method +name from the table in ``mutation`` (default: ``"minipcn"``), a transition count in +``mutation_n_steps`` (default: five), and constructor options in +``mutation_kwargs``. Use ``mutation="none"`` or ``None`` to disable mutation. + +The helper defaults to ``particle_transform="affine"``; configure it through +``particle_transform_kwargs`` or pass a transform object. With default ESS +tempering, mutation runs every iteration. With ``weighting_criterion="cess"`` +and default execution flags, it runs only after resampling unless +``always_mutate=True``. + +You can also pass a configured step as ``mutation``. Set its target, transform, +and options directly: the helper does not apply ``mutation_n_steps`` or +``mutation_kwargs`` to it. See :doc:`../running_inference/smc_configuration` for custom pipelines +and the API links above for method-specific options and outputs. diff --git a/docs/user_guide/sampling_options/resampling_configuration.rst b/docs/user_guide/sampling_options/resampling_configuration.rst new file mode 100644 index 00000000..392fcae7 --- /dev/null +++ b/docs/user_guide/sampling_options/resampling_configuration.rst @@ -0,0 +1,88 @@ +Resampling configuration +======================== + +Resampling draws particles from a weighted population and assigns equal +weights to the result. Some particles may be duplicated; mutation can then +move them to new positions. + +.. list-table:: + :header-rows: 1 + :widths: 18 20 30 32 + + * - Method + - ``resampling`` name + - Implementation + - Ancestor selection + * - Systematic + - ``"systematic"`` + - :class:`~sequince.steps.resampling.SystematicResampling` + - Equally spaced positions with one random offset. + * - Stratified + - ``"stratified"`` + - :class:`~sequince.steps.resampling.StratifiedResampling` + - One random position in each equal-width stratum. + * - Multinomial + - ``"multinomial"`` + - :class:`~sequince.steps.resampling.MultinomialResampling` + - Independent draws according to particle weights. + +Common configuration +-------------------- + +The method determines how ancestors are drawn. Constructor options control +when to resample and how many particles to draw: + +.. code-block:: python + + from sequince.steps.resampling import StratifiedResampling + + resampling_step = StratifiedResampling( + ess_threshold=0.4, + always_resample=False, + resample_at_end=True, + ) + +All three methods accept: + +* ``ess_threshold``: resample when ESS divided by the input particle count + falls below this value (default: ``0.5``). ``None`` disables the ESS trigger. +* ``always_resample``: resample on every call, regardless of ESS + (default: ``False``). +* ``resample_at_end``: resample when ``context.stop_requested`` is set, + regardless of ESS (default: ``True``). +* ``n_samples``: number of output particles when resampling occurs. ``None`` + (default) keeps the input count. Skipped resampling leaves the count unchanged. + +To disable resampling, set ``ess_threshold=None``, ``always_resample=False``, +and ``resample_at_end=False``. + +Functional execution currently requires ``always_resample=True``. Compiled +pipelines must also preserve the particle count; see :doc:`../advanced_workflows/jax_support`. + +Outputs and ancestry +-------------------- + +The step returns :class:`~sequince.steps.resampling.ResamplingOutput` with +``ess`` (before resampling), ``resampled``, ``ancestors``, and +``resampling_scheme``. ``ancestors`` indexes the input population; when +resampling is skipped, it contains the identity indices. The context metadata +also records ``ess`` and ``resampled``. + +The state tracks each particle's original ancestor in +``initial_ancestor_indices``. Resampling leaves the accumulated normalizing +constant unchanged. See :doc:`../running_inference/evidence_uncertainty` for the assumptions needed +by genealogy-based evidence estimators, including multinomial resampling. + +Prebuilt samplers +----------------- + +:func:`~sequince.prebuilt_samplers.make_standard_bayesian_smc` selects a method +with ``resampling`` (default: ``"systematic"``). Set the ESS threshold through +``resampling_ess_threshold`` (default: ``0.5``); ``always_resample`` and +``resample_at_end`` are also helper arguments. Pass extra constructor options, +such as ``n_samples``, through ``resampling_kwargs``. + +With ``weighting_criterion="ess"`` (default), the helper resamples every +iteration and rejects ``always_resample=False``. With +``weighting_criterion="cess"``, it defaults to ESS-triggered resampling. +Terminal resampling remains enabled unless ``resample_at_end=False``. diff --git a/pyproject.toml b/pyproject.toml index e5dfdf95..48ac2f87 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -164,10 +164,11 @@ typecheck = "mypy" [tool.pixi.feature.doc.dependencies] sphinx = ">=7.2" sphinx-book-theme = ">=1.1" +sphinx-autoapi = ">=3.8,<4" +myst-parser = ">=4,<5" [tool.pixi.feature.doc.tasks] -# Runtime dependencies are installed here; mocking NumPy breaks orng's types. -docs = "sphinx-build -b html -D autodoc_mock_imports= docs docs/_build/html" +docs = "sphinx-build -b html docs docs/_build/html" [tool.pixi.feature.jax-cpu.activation.env] JAX_PLATFORMS = "cpu" diff --git a/src/sequince/result.py b/src/sequince/result.py index 82d68bb5..8ae4189e 100644 --- a/src/sequince/result.py +++ b/src/sequince/result.py @@ -8,9 +8,10 @@ from typing import TYPE_CHECKING, Any, Callable, Literal, Self import numpy as np +from array_api_compat import device as get_device from numpy.random import Generator as NumPyGenerator -from .array_api import asarray, get_device, namespace, to_numpy +from .array_api import asarray, namespace, to_numpy from .state import ParticleState from .steps.constraints import _STOP_REASONS, NestedSamplingLevel from .steps.context import StepContext @@ -820,11 +821,12 @@ def posterior_samples( Method used to construct the returned posterior representation: - ``"weights"``: return weighted particles, optionally uniformly - down-selected without replacement. - - ``"rejection"``: independently retain each particle with probability - proportional to its weight, returning equal-weight unique samples. + down-selected without replacement. + - ``"rejection"``: independently retain each particle with + probability proportional to its weight, returning equal-weight + unique samples. - ``"multinomial"``: draw equal-weight samples with replacement - according to the particle weights. + according to the particle weights. Returns ------- diff --git a/src/sequince/smc.py b/src/sequince/smc.py index e50e2cf7..7388985d 100644 --- a/src/sequince/smc.py +++ b/src/sequince/smc.py @@ -373,8 +373,10 @@ def run_functional( Returns ------- - result, random_state: - The completed SMC result and final backend-native random state. + result : SMCResult + The completed SMC result. + random_state : object + The final backend-native random state. """ if runtime is None and random_state is None: raise ValueError( diff --git a/src/sequince/steps/constraints.py b/src/sequince/steps/constraints.py index c973c58e..d871e875 100644 --- a/src/sequince/steps/constraints.py +++ b/src/sequince/steps/constraints.py @@ -8,9 +8,10 @@ from typing import TYPE_CHECKING, Any, Callable, ClassVar import numpy as np +from array_api_compat import is_jax_namespace from orng import RandomGenerator -from ..array_api import asarray, is_jax_namespace, logsumexp, namespace +from ..array_api import asarray, logsumexp, namespace from ..array_api import copy as array_copy from ..state import ( ParticleState, diff --git a/src/sequince/steps/control.py b/src/sequince/steps/control.py index 3cf4b427..67e79917 100644 --- a/src/sequince/steps/control.py +++ b/src/sequince/steps/control.py @@ -6,8 +6,9 @@ from typing import TYPE_CHECKING, Any, Callable, Sequence from acai import cond +from array_api_compat import is_jax_namespace -from ..array_api import is_jax_namespace, namespace +from ..array_api import namespace from ..typing import Array from .base import SMCStep, StepContext diff --git a/src/sequince/steps/mutation/metrics/de.py b/src/sequince/steps/mutation/metrics/de.py index eebf7dd7..9dde32af 100644 --- a/src/sequince/steps/mutation/metrics/de.py +++ b/src/sequince/steps/mutation/metrics/de.py @@ -5,7 +5,9 @@ from dataclasses import dataclass from typing import ClassVar -from ....array_api import asarray, get_device, namespace +from array_api_compat import device as get_device + +from ....array_api import asarray, namespace from ....typing import Array from ...context import StepContext from .esjd import _squared_jump_distances diff --git a/src/sequince/steps/mutation/metrics/rjd.py b/src/sequince/steps/mutation/metrics/rjd.py index b5fe283e..5b9e6c0b 100644 --- a/src/sequince/steps/mutation/metrics/rjd.py +++ b/src/sequince/steps/mutation/metrics/rjd.py @@ -6,9 +6,10 @@ from dataclasses import dataclass from typing import Literal +from array_api_compat import device as get_device from orng import RandomGenerator -from ....array_api import asarray, get_device, namespace, to_numpy +from ....array_api import asarray, namespace, to_numpy from ....typing import Array from ....utils.pytrees import pytree_dataclass from ...base import validate_n_steps diff --git a/src/sequince/steps/mutation/metrics/state.py b/src/sequince/steps/mutation/metrics/state.py index 265f619c..1f48b7c5 100644 --- a/src/sequince/steps/mutation/metrics/state.py +++ b/src/sequince/steps/mutation/metrics/state.py @@ -71,8 +71,9 @@ class MutationTransition: Parameters ---------- - displacement : array, shape (n_particles, n_dimensions) - Accepted displacement in mutation coordinates, zero for rejected rows. + displacement : array + Accepted displacement in mutation coordinates with shape + ``(n_particles, n_dimensions)``, zero for rejected rows. kernel_auxiliary : pytree, optional Read-only auxiliary information from the kernel sample, or None. Red-blue sweeps restore particle-aligned leaves to population order. diff --git a/src/sequince/steps/mutation/metropolis.py b/src/sequince/steps/mutation/metropolis.py index 89f7d91e..4807f679 100644 --- a/src/sequince/steps/mutation/metropolis.py +++ b/src/sequince/steps/mutation/metropolis.py @@ -7,8 +7,9 @@ import array_api_extra as xpx import optree +from array_api_compat import device as get_device -from ...array_api import asarray, get_device, namespace, to_numpy +from ...array_api import asarray, namespace, to_numpy from ...target import EvaluationCounts from ...transforms import ParticleTransform, select_transform_auxiliary from ...typing import Array diff --git a/src/sequince/transforms.py b/src/sequince/transforms.py index a3ec2866..f3c8d3af 100644 --- a/src/sequince/transforms.py +++ b/src/sequince/transforms.py @@ -343,14 +343,9 @@ def _log_det_entries(value, arr, xp): class ParticleTransform: """Coordinate transform interface for particle mutation. - Methods - ------- - fit(particles, context=None, *, rng=None) - Optionally adapt internal statistics using the provided particle cloud. - forward(particles, context=None, *, rng=None) - Return transformed particles, log-Jacobian values, and auxiliary data. - inverse(values, context=None, *, auxiliary=None, rng=None) - Undo the transform applied by :meth:`forward`. + Use :meth:`fit` to optionally adapt internal statistics using the particle + cloud. :meth:`forward` returns transformed particles, log-Jacobian values, + and auxiliary data. :meth:`inverse` undoes the forward transform. """ def clone(