From ff8e29482c779fa5d6476e3b0a7e9789b5c2d66b Mon Sep 17 00:00:00 2001 From: "Michael J. Williams" Date: Fri, 21 Aug 2026 12:27:10 +0200 Subject: [PATCH 01/23] DOC: add missing docs --- docs/api/core.rst | 16 ++ docs/custom_steps_callbacks.rst | 67 ++++++++ docs/evidence_uncertainty.rst | 115 +++++++++++++ docs/getting_started.rst | 90 ++++++---- docs/index.rst | 4 +- docs/jax_support.rst | 167 +++++++++++++++++++ docs/mcmc_sampler.rst | 120 ++++++++++++++ docs/mutation_options.rst | 267 ++++++++++++++++++++++++++++++ docs/observations.rst | 189 +++++++++++++++++++++ docs/particle_state.rst | 65 ++++++++ docs/random_number_generation.rst | 118 +++++++++++++ docs/resampling_options.rst | 94 +++++++++++ docs/smc_configuration.rst | 70 ++++++++ docs/step_context.rst | 149 +++++++++++++++++ docs/user_guide.rst | 163 ++++-------------- 15 files changed, 1531 insertions(+), 163 deletions(-) create mode 100644 docs/custom_steps_callbacks.rst create mode 100644 docs/evidence_uncertainty.rst create mode 100644 docs/jax_support.rst create mode 100644 docs/mcmc_sampler.rst create mode 100644 docs/mutation_options.rst create mode 100644 docs/observations.rst create mode 100644 docs/particle_state.rst create mode 100644 docs/random_number_generation.rst create mode 100644 docs/resampling_options.rst create mode 100644 docs/smc_configuration.rst create mode 100644 docs/step_context.rst diff --git a/docs/api/core.rst b/docs/api/core.rst index ca3c4d76..cfe75ce6 100644 --- a/docs/api/core.rst +++ b/docs/api/core.rst @@ -22,6 +22,14 @@ use the checkpoint APIs to restore runtime state. :undoc-members: :show-inheritance: +Markov chain Monte Carlo driver +------------------------------- + +.. automodule:: sequince.mcmc + :members: + :undoc-members: + :show-inheritance: + Particle state and results -------------------------- @@ -35,6 +43,14 @@ Particle state and results :undoc-members: :show-inheritance: +Particle transforms +------------------- + +.. automodule:: sequince.transforms + :members: + :undoc-members: + :show-inheritance: + Checkpointing ------------- diff --git a/docs/custom_steps_callbacks.rst b/docs/custom_steps_callbacks.rst new file mode 100644 index 00000000..49c59ee1 --- /dev/null +++ b/docs/custom_steps_callbacks.rst @@ -0,0 +1,67 @@ +Custom steps and callbacks +========================== + +Steps transform particle state within an iteration. Callbacks run after the +step pipeline and are intended for observation, reporting, or run-level control. +Both receive a :class:`~sequince.steps.context.StepContext`; see +:doc:`step_context` for its lifecycle and update methods. + +Custom steps +------------ + +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`. + +Custom 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``. + +For compiled JAX execution, custom steps must also obey the functional +interface's fixed-structure requirements. See :doc:`jax_support` for the +functional step contract and initialization hooks. diff --git a/docs/evidence_uncertainty.rst b/docs/evidence_uncertainty.rst new file mode 100644 index 00000000..53e4cae7 --- /dev/null +++ b/docs/evidence_uncertainty.rst @@ -0,0 +1,115 @@ +Evidence and uncertainty +======================== + +: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/getting_started.rst b/docs/getting_started.rst index 4eb23b21..d8baab96 100644 --- a/docs/getting_started.rst +++ b/docs/getting_started.rst @@ -9,63 +9,85 @@ 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``. +``sequince`` targets Python 3.11+. Its core dependencies are installed +automatically; integrations such as emcee, BlackJAX, and HDF5 checkpointing are +available as optional extras. 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 +steps for many Bayesian posterior inference workflows. The basic usage is as +follows: .. 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): - ... - + # Set RNG + rng = np.random.default_rng(42) + + # User-supplied log-density functions, these must be vectorized + # Theta is a 2D array of shape (n_particles, n_parameters) + # Context is dataclass that contains additional information + def log_prior(theta, context): + """Uniform prior over the specified bounds""" + bounds = context.shared["bounds"] + return np.where( + np.all((theta >= bounds[0]) & (theta <= bounds[1]), axis=1), + 0.0, + -np.inf, + ) + + def log_likelihood(theta, context): + """Unnormalized Gaussian likelihood with mean 0 and std 1""" + return -0.5 * np.sum(theta**2, axis=1) + + + # Create a standard SMC runner for Bayesian inference + # smc is the main runner + # posterior is a convenience wrapper around the log-likelihood and log-prior + # metadata is dictionary that can be used to pass information between steps + # in the SMC loop 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)), + # Shared context is a dictionary that can be used to pass user-defined + # information to steps and the log-density functions. + shared_context = dict( + bounds=[-5.0, 5.0] ) - results = smc.run( + # Build the initial particle cloud with samples from the prior + n = 500 + initial_particles = rng.uniform(*shared_context["bounds"], size=(n, 2)) + # Create a particle state object to hold the initial particles + # Unless specified, the particles are assumed to be equally weighted + initial_state = ParticleState(particles=initial_particles) + + # Run the SMC sampler + result = smc.run( initial_state, - rng=np.random.default_rng(), + rng=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 --------------------------- + # Extract the posterior samples and log-evidence from the result object + samples = result.posterior_samples() + log_evidence = result.log_evidence() + log_evidence_error = result.log_evidence_error() -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. +Next steps +---------- -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. +See :doc:`evidence_uncertainty` for interpreting the evidence estimate and its +uncertainty. To customise the run, continue with :doc:`smc_configuration` and +the sampling choices in :doc:`resampling_options` and :doc:`mutation_options`. +The :doc:`step_context` guide explains the distinction between shared inputs +and evolving metadata. For sampling without SMC weighting and resampling, see +:doc:`mcmc_sampler`. diff --git a/docs/index.rst b/docs/index.rst index 6e719811..58cff83e 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -3,7 +3,8 @@ sequince documentation ``sequince`` provides modular building blocks for assembling Sequential Monte Carlo (SMC) algorithms for Bayesian inference, particle filtering, and related -sequential methods. +sequential methods. Its mutation backends can also be used through a standalone +Markov chain Monte Carlo runner. .. toctree:: :maxdepth: 2 @@ -11,7 +12,6 @@ sequential methods. getting_started user_guide - checkpointing examples api/index diff --git a/docs/jax_support.rst b/docs/jax_support.rst new file mode 100644 index 00000000..db0d75d3 --- /dev/null +++ b/docs/jax_support.rst @@ -0,0 +1,167 @@ +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 ``uv add "sequince[jax]"``. To use BlackJAX mutation, +install ``uv add "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:`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 +------------- + +.. autofunction:: sequince.jax.compile_functional_iteration diff --git a/docs/mcmc_sampler.rst b/docs/mcmc_sampler.rst new file mode 100644 index 00000000..9989cf85 --- /dev/null +++ b/docs/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:`mutation_options` for available backends, optional dependencies, and +kernel configuration. A configured +:class:`~sequince.steps.mutation.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/mutation_options.rst b/docs/mutation_options.rst new file mode 100644 index 00000000..1b798621 --- /dev/null +++ b/docs/mutation_options.rst @@ -0,0 +1,267 @@ +Mutation options +================ + +Mutation rejuvenates the particle population after weighting and resampling. +It applies MCMC transitions that preserve the current target distribution, +helping duplicated particles explore again without changing their weights. + +The :func:`~sequince.prebuilt_samplers.make_standard_bayesian_smc` helper +selects a mutation implementation with ``mutation``: + +.. list-table:: + :header-rows: 1 + :widths: 22 30 48 + + * - Value + - Implementation + - Typical use + * - ``"minipcn"`` + - :class:`~sequince.steps.mutation.MiniPCNMutation` + - General default with eager and functional execution support. + * - ``"emcee"`` + - :class:`~sequince.steps.mutation.EmceeMutation` + - Ensemble MCMC with emcee's collection of proposal moves. + * - ``"blackjax"`` + - :class:`~sequince.steps.mutation.BlackJAXMutation` + - JAX-backed NUTS, HMC, or random-walk mutation. + * - ``"metropolis-hastings"`` + - :class:`~sequince.steps.mutation.MetropolisHastingsMutation` + - Explicit control over a proposal kernel and transition ratio. + * - ``"none"`` or ``None`` + - No mutation step + - Reweighting and resampling only. + * - An ``SMCStep`` instance + - The supplied step + - A custom, already configured mutation implementation. + +Common configuration +-------------------- + +``mutation_n_steps`` controls how many MCMC transitions are applied to every +particle during one SMC iteration. The pre-built sampler defaults to five. +Strategy-specific options are passed through ``mutation_kwargs``: + +.. code-block:: python + + smc, posterior, metadata = make_standard_bayesian_smc( + log_prior, + log_likelihood, + mutation="minipcn", + mutation_n_steps=10, + mutation_kwargs={ + "store_chains": True, + "transform_refit": "per-context", + }, + particle_transform="affine", + ) + +All built-in mutation steps support these common options: + +``store_chains`` + Include the intermediate MCMC positions in the step output. It defaults to + ``False`` because full chains can be large. Acceptance diagnostics and the + number of transitions are returned even when chains are not stored. + +``transform_refit`` + Control how an adaptive particle transform is fitted. ``"always"`` fits on + every mutation call, ``"per-context"`` reuses it within the same iteration + and observation position, and ``"once"`` reuses the first fit. Eager + execution also accepts a callable predicate; compiled functional execution + requires one of the named policies. + +Particle transforms +------------------- + +The helper mutates particles in transformed coordinates by default using +``particle_transform="affine"``. This standardizes the particle cloud before +the MCMC transition and maps accepted positions back afterward. Configure it +through ``particle_transform_kwargs``: + +.. code-block:: python + + smc, posterior, metadata = make_standard_bayesian_smc( + log_prior, + log_likelihood, + particle_transform="affine", + particle_transform_kwargs={"eps": 1e-5}, + ) + +Other registered names are ``"identity"``, +``"bounded_to_unit_hypercube"``, ``"bounded_to_unbounded"``, and +``"periodic_boundary"``. The bounded and periodic transforms require bounds in +``particle_transform_kwargs``. Pass ``None`` to mutate directly in the original +coordinates, or pass a configured +:class:`~sequince.transforms.ParticleTransform` instance. + +Auxiliary information for custom transforms +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +A custom transform may return invocation-specific information that its inverse +needs without storing mutable state on the transform. For example, a +many-to-one transform can preserve the branch assigned to each particle:: + + from sequince.array_api import namespace + from sequince.transforms import ParticleTransform + + + class SignFoldingTransform(ParticleTransform): + def forward(self, values, context=None, *, rng=None): + xp = namespace(values) + signs = xp.where(values < 0, -xp.ones_like(values), xp.ones_like(values)) + transformed = xp.abs(values) + log_det = xp.zeros_like(values[..., 0]) + return transformed, log_det, signs + + def inverse( + self, + values, + context=None, + *, + auxiliary=None, + rng=None, + ): + if auxiliary is None: + raise ValueError("Sign information is required for inversion.") + xp = namespace(values) + return values * auxiliary, xp.zeros_like(values[..., 0]) + +Auxiliary information is ephemeral: mutation passes it from one ``forward`` +call through transformed-space operations to the corresponding ``inverse``. +It may be a single array or any Optree-compatible structure, including a tuple, +list, or dictionary such as ``(mode_ids, mask)`` or +``{"mode_ids": mode_ids, "mask": mask}``. Simple transforms do not need a +Sequince-specific wrapper. Libraries needing named fields may register their +own custom pytree types; immutable dataclasses can use +:func:`sequince.utils.pytree_dataclass` for Optree and JAX registration. + +For a batch of ``n_particles``, Sequince treats an auxiliary array leaf as +particle-aligned exactly when ``leaf.ndim > 0`` and +``leaf.shape[0] == n_particles``. A mutation operating on a particle subset +applies the same selection to those leaves before inversion, while preserving +other leaves unchanged. Consequently, an unrelated array whose leading size +happens to equal ``n_particles`` is also selected; this is an intentional +limitation of the current shape-based convention. + +Particle-aligned auxiliary works with batch-preserving mutation and BlackJAX. +It is currently rejected by emcee: emcee selects walker subsets internally, +but its standard log-probability callback does not expose their original walker +indices, so Sequince cannot guarantee the correct auxiliary entry. + +MiniPCN +------- + +``"minipcn"`` is the default and requires no extra installation. It uses the +``"tpcn"`` step function by default and targets an acceptance rate of ``0.234``. +Both can be changed through ``mutation_kwargs``: + +.. code-block:: python + + mutation="minipcn", + mutation_kwargs={ + "step_fn": "tpcn", + "target_acceptance_rate": 0.3, + }, + +MiniPCN implements explicit random-state threading and can be used in a compiled +functional pipeline. See :doc:`jax_support` for the functional runner. + +emcee +----- + +Install emcee support with ``uv add "sequince[emcee]"``. With no move specified, +emcee uses its default stretch move. Named choices are ``"de"``, +``"de-snooker"``, ``"gaussian"``, ``"kde"``, ``"stretch"``, and ``"walk"``. +The ``"bilby-de"`` preset supplies a Bilby-style weighted +differential-evolution schedule: + +.. code-block:: python + + mutation="emcee", + mutation_n_steps=20, + mutation_kwargs={ + "move": "bilby-de", + }, + +The preset is equivalent to using emcee's ``DEMove`` with its automatically +chosen differential-evolution scale 90% of the time and ``gamma0=1.0`` 10% of +the time. Both moves use ``sigma=1e-5``: + +.. code-block:: python + + mutation_kwargs={ + "move": [ + { + "move": "de", + "move_kwargs": {"gamma0": None, "sigma": 1e-5}, + "weight": 0.9, + }, + { + "move": "de", + "move_kwargs": {"gamma0": 1.0, "sigma": 1e-5}, + "weight": 0.1, + }, + ], + }, + +This preset does not require Bilby itself; it constructs the schedule from +emcee move objects. The expanded form can be adapted to create a custom +Bilby-like schedule with different moves, parameters, or weights. + +``move`` may also be an emcee move object, a mapping containing ``move`` and +``move_kwargs``, or a weighted sequence of move specifications. Other keyword +arguments are forwarded to ``emcee.EnsembleSampler``; ``run_kwargs`` is +forwarded to its ``run_mcmc`` method. + +Metropolis--Hastings +-------------------- + +The Metropolis--Hastings option requires a proposal ``kernel``. Sequince +provides ``"gaussian-random-walk"`` and ``"mixture"`` as named kernels: + +.. code-block:: python + + mutation="metropolis-hastings", + mutation_n_steps=10, + mutation_kwargs={ + "kernel": "gaussian-random-walk", + "kernel_kwargs": {"scale": 0.25}, + }, + +The mixture kernel accepts a ``kernels`` sequence and optional ``weights`` in +``kernel_kwargs``. A custom kernel object or factory can be passed directly; +it must provide proposals and one log transition ratio per particle. When a +kernel object is supplied, configure it directly rather than also passing +``kernel_kwargs``. + +BlackJAX +-------- + +``"blackjax"`` accepts ``"nuts"`` (the default), ``"hmc"``, +``"random-walk"``, or a callable algorithm factory through +``mutation_kwargs["algorithm"]``. Algorithm parameters belong in +``mutation_kwargs["algorithm_kwargs"]``. NUTS and HMC require differentiable, +JAX-compatible targets and JAX-backed particles. + +See :doc:`jax_support` for installation, a complete compiled example, random-key +threading, and current functional limitations. + +Disabling or replacing mutation +------------------------------- + +Use ``mutation="none"`` when mutation is deliberately unnecessary. Resampling +without mutation can rapidly reduce particle diversity, so monitor effective +sample size and duplicated particles when choosing this mode. + +You can also pass an initialized :class:`~sequince.steps.base.SMCStep` instance +as ``mutation``. The helper inserts that object unchanged; it does not inject +the generated posterior, particle transform, ``mutation_n_steps``, or +``mutation_kwargs``. Configure a custom step completely before passing it, or +construct the full pipeline manually as described in :doc:`smc_configuration`. + +Reading mutation outputs +------------------------ + +Each mutation step returns diagnostics in its position within the iteration's +output tuple. The runner collects those tuples in ``result.outputs``. Built-in +outputs include acceptance information, the configured number of transitions, +and ``chain``; the chain value is ``None`` unless ``store_chains=True``. diff --git a/docs/observations.rst b/docs/observations.rst new file mode 100644 index 00000000..45c7f0c6 --- /dev/null +++ b/docs/observations.rst @@ -0,0 +1,189 @@ +Using observations +================== + +Pass an iterable to ``observations`` when an SMC iteration should process one +item from a data stream. The runner consumes at most one item at the beginning +of each iteration and exposes it as +:attr:`~sequince.steps.context.StepContext.observation`. Every step and callback +in that iteration sees the same observation. + +An observation can be any object: a scalar measurement, an array, a batch, or a +structured value containing several fields. Sequince does not interpret it; the +configured steps or target functions decide how it affects the particle state. + +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 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``. + +Each item is one iteration +-------------------------- + +``observations`` defines iteration input, not an automatically vectorized data +axis. A list containing three arrays causes three iterations; each array is the +complete ``context.observation`` for its iteration. All configured steps run +once for each consumed item. + +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 has no natural bound. + +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 +-------------------- + +``context.t`` is the zero-based SMC iteration. It is not necessarily an +observation index because iterations may continue after the input is 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. 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]``. +A one-shot iterator cannot be repositioned and is rejected after any +observations have already been consumed. See :doc:`checkpointing` for the full +checkpoint workflow. + +JAX observations +---------------- + +``run_functional`` has the same host-side observation and termination semantics. +When using a compiled iteration, all observations must have a pytree structure, +array shapes, and dtypes compatible with the compiled signature. 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:`jax_support` for compilation, random-state threading, and fixed-carry +requirements. diff --git a/docs/particle_state.rst b/docs/particle_state.rst new file mode 100644 index 00000000..aa330a76 --- /dev/null +++ b/docs/particle_state.rst @@ -0,0 +1,65 @@ +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 = result.posterior_samples(1_000, rng=42) + +``posterior_samples`` draws an equally weighted sample from the final weighted +population. For diagnostics across the complete run, use ``result.summary()``, +``result.effective_sample_sizes()``, ``result.log_evidence_history()``, and +``result.metadata_history(key)``. See :doc:`evidence_uncertainty` for +interpreting the final evidence estimate and its uncertainty. diff --git a/docs/random_number_generation.rst b/docs/random_number_generation.rst new file mode 100644 index 00000000..4d9087da --- /dev/null +++ b/docs/random_number_generation.rst @@ -0,0 +1,118 @@ +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` and +:meth:`~sequince.mcmc.MarkovChainMonteCarlo.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. + +Randomness 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:`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:`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/resampling_options.rst b/docs/resampling_options.rst new file mode 100644 index 00000000..6ed18be2 --- /dev/null +++ b/docs/resampling_options.rst @@ -0,0 +1,94 @@ +Resampling options +================== + +Resampling replaces a weighted particle population with equally weighted +particles drawn from it. This prevents a small number of particles from +carrying almost all of the weight, at the cost of duplicating some particles. +Mutation is commonly applied afterwards to restore diversity. + +Choosing a strategy +------------------- + +The :func:`~sequince.prebuilt_samplers.make_standard_bayesian_smc` helper +selects the ancestor-sampling algorithm with ``resampling``: + +``"systematic"`` + Use :class:`~sequince.steps.resampling.SystematicResampling`. This is the + default and uses one random offset for the complete population. + +``"stratified"`` + Use :class:`~sequince.steps.resampling.StratifiedResampling`, with one + random position in each equal-width stratum. + +``"multinomial"`` + Use :class:`~sequince.steps.resampling.MultinomialResampling`, drawing each + ancestor independently according to the particle weights. + +For example: + +.. code-block:: python + + smc, posterior, metadata = make_standard_bayesian_smc( + log_prior, + log_likelihood, + resampling="stratified", + ) + +Choosing when to resample +------------------------- + +The resampling algorithm and its schedule are configured separately: + +.. code-block:: python + + smc, posterior, metadata = make_standard_bayesian_smc( + log_prior, + log_likelihood, + weighting_criterion="cess", + resampling="stratified", + resampling_ess_threshold=0.4, + always_resample=False, + ) + +With ``always_resample=False``, resampling occurs when the effective sample size +divided by the number of particles falls below ``resampling_ess_threshold``. +Set the threshold to ``None`` to disable threshold-triggered resampling. With +``always_resample=True``, the runner resamples every iteration and ignores the +threshold. + +When ``always_resample`` is omitted, its default depends on +``weighting_criterion``: ESS-based tempering resamples every iteration, while +CESS-based tempering uses the ESS threshold. Setting ``always_resample=False`` +with ``weighting_criterion="ess"`` is invalid because ESS-based tempering relies +on resampling to meet its target. + +Functional execution currently requires ``always_resample=True``. See +:doc:`jax_support` for the other functional-execution constraints. + +Configuring a resampling step directly +--------------------------------------- + +In a manually assembled pipeline, instantiate the selected step with the same +schedule arguments: + +.. code-block:: python + + from sequince.steps import StratifiedResampling + + resampling = StratifiedResampling( + ess_threshold=0.4, + always_resample=False, + ) + +The step records the pre-resampling ESS and whether resampling occurred in both +the context metadata and its output. The particle state records selected +ancestor indices when resampling occurs, which allows particle genealogies to +be reconstructed. + +Resampling and evidence estimates +--------------------------------- + +The resampling strategy affects Monte Carlo variability even though all three +built-in algorithms target the same weighted population. See +:doc:`evidence_uncertainty` for the uncertainty estimators and their resampling +requirements. diff --git a/docs/smc_configuration.rst b/docs/smc_configuration.rst new file mode 100644 index 00000000..06dbba8e --- /dev/null +++ b/docs/smc_configuration.rst @@ -0,0 +1,70 @@ +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`` wires together +:class:`~sequince.steps.weighting.AdaptiveTemperedWeighting`, +a configurable resampling step, and a configurable mutation step. See +:doc:`resampling_options` and :doc:`mutation_options` for the available +sampling strategies, and :doc:`jax_support` for JAX-backed execution. +For a mutation kernel without SMC weighting and resampling, see +:doc:`mcmc_sampler`. + +Pass the returned metadata when starting the run: + +.. code-block:: python + + result = smc.run( + initial_state, + rng=rng, + shared_context=shared_context, + metadata=metadata, + ) + +See :doc:`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: + +.. code-block:: python + + from sequince import SequentialMonteCarlo + from sequince.callbacks import TracePlot + from sequince.steps import ( + AdaptiveTemperedWeighting, + EmceeMutation, + SystematicResampling, + ) + + smc = SequentialMonteCarlo( + steps=[ + AdaptiveTemperedWeighting(posterior, target_ess_ratio=0.8), + 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:`custom_steps_callbacks` when implementing either interface yourself. diff --git a/docs/step_context.rst b/docs/step_context.rst new file mode 100644 index 00000000..376b5173 --- /dev/null +++ b/docs/step_context.rst @@ -0,0 +1,149 @@ +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:`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 +------------- + +.. autoclass:: sequince.steps.context.StepContext + :members: diff --git a/docs/user_guide.rst b/docs/user_guide.rst index 72818a7a..ad221cd8 100644 --- a/docs/user_guide.rst +++ b/docs/user_guide.rst @@ -1,141 +1,50 @@ 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. +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. The same +mutation implementations can also be used through the standalone MCMC runner. +At the end of an iteration, callbacks can inspect the accumulated result and +update the context. -Configuring the runner ----------------------- +Start with the workflow guides when configuring a run and interpreting its +output. The remaining sections explain sampling choices, the objects passed +through the SMC loop, and more specialised execution and extension points. -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: +.. toctree:: + :maxdepth: 1 + :caption: Running inference -.. code-block:: python + smc_configuration + mcmc_sampler + observations + evidence_uncertainty - 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 - ) +.. toctree:: + :maxdepth: 1 + :caption: Sampling options -``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. + resampling_options + mutation_options -BlackJAX mutation ------------------ +.. toctree:: + :maxdepth: 1 + :caption: Core concepts -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: + particle_state + step_context + random_number_generation -.. code-block:: python +.. toctree:: + :maxdepth: 1 + :caption: Advanced workflows - import jax.numpy as jnp + jax_support + checkpointing - 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), - }, - }, - ) +.. toctree:: + :maxdepth: 1 + :caption: Extending Sequince -``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. + custom_steps_callbacks From ed485301f0652338b5427519a2f7126e3640ddf9 Mon Sep 17 00:00:00 2001 From: "Michael J. Williams" Date: Tue, 8 Sep 2026 09:46:34 +0100 Subject: [PATCH 02/23] DOC: add makefile for docs --- docs/Makefile | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) create mode 100644 docs/Makefile diff --git a/docs/Makefile b/docs/Makefile new file mode 100644 index 00000000..d4bb2cbb --- /dev/null +++ b/docs/Makefile @@ -0,0 +1,20 @@ +# 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 Makefile + +# 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) From 71420d47924eece9808ec2c23735802bbbc83625 Mon Sep 17 00:00:00 2001 From: Christian Chapman-Bird Date: Tue, 8 Sep 2026 10:34:32 +0100 Subject: [PATCH 03/23] gitignore DS_Store files --- .gitignore | 2 ++ 1 file changed, 2 insertions(+) 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] From c6036f72bbefa8f2e58a22df9b5ecf9be117111b Mon Sep 17 00:00:00 2001 From: Christian Chapman-Bird Date: Tue, 8 Sep 2026 10:35:03 +0100 Subject: [PATCH 04/23] DOC: modify install instructions for PyPI release --- docs/getting_started.rst | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/docs/getting_started.rst b/docs/getting_started.rst index d8baab96..db35f87d 100644 --- a/docs/getting_started.rst +++ b/docs/getting_started.rst @@ -4,10 +4,9 @@ Getting started Installation ------------ -The project is published from the main branch and can be installed directly -from GitHub:: +Install from PyPI:: - pip install git+https://github.com/sequince-dev/sequince.git + pip install sequince ``sequince`` targets Python 3.11+. Its core dependencies are installed automatically; integrations such as emcee, BlackJAX, and HDF5 checkpointing are From 1abfb20d460b1fb116a022778af6adf27614dde6 Mon Sep 17 00:00:00 2001 From: Christian Chapman-Bird Date: Tue, 8 Sep 2026 12:26:09 +0100 Subject: [PATCH 05/23] small tweaks --- docs/smc_configuration.rst | 5 +++-- docs/user_guide.rst | 4 +++- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/docs/smc_configuration.rst b/docs/smc_configuration.rst index 06dbba8e..7bf5c2ce 100644 --- a/docs/smc_configuration.rst +++ b/docs/smc_configuration.rst @@ -43,7 +43,8 @@ data and evolving metadata. Building a pipeline manually ---------------------------- -For full control, instantiate the runner with steps in execution order: +For full control, instantiate the runner with steps in execution order. For example, +to reproduce the pre-built SMC defined above: .. code-block:: python @@ -57,7 +58,7 @@ For full control, instantiate the runner with steps in execution order: smc = SequentialMonteCarlo( steps=[ - AdaptiveTemperedWeighting(posterior, target_ess_ratio=0.8), + AdaptiveTemperedWeighting(posterior), SystematicResampling(always_resample=True), EmceeMutation(posterior, n_steps=10), ], diff --git a/docs/user_guide.rst b/docs/user_guide.rst index ad221cd8..5e550bc0 100644 --- a/docs/user_guide.rst +++ b/docs/user_guide.rst @@ -6,7 +6,9 @@ composable steps. Each step receives the current particle state and runtime context, then returns an updated state, context, and optional output. The same mutation implementations can also be used through the standalone MCMC runner. At the end of an iteration, callbacks can inspect the accumulated result and -update the context. +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 workflow guides when configuring a run and interpreting its output. The remaining sections explain sampling choices, the objects passed From c6e1601921b8eafcf602677f2d1281d078ca408e Mon Sep 17 00:00:00 2001 From: Christian Chapman-Bird Date: Thu, 8 Oct 2026 15:52:16 +0200 Subject: [PATCH 06/23] tidy observations.rst --- docs/observations.rst | 52 ++++++++++++++++++++----------------------- 1 file changed, 24 insertions(+), 28 deletions(-) diff --git a/docs/observations.rst b/docs/observations.rst index 45c7f0c6..cd7a4d0d 100644 --- a/docs/observations.rst +++ b/docs/observations.rst @@ -1,15 +1,15 @@ Using observations ================== -Pass an iterable to ``observations`` when an SMC iteration should process one -item from a data stream. The runner consumes at most one item at the beginning +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`. Every step and callback -in that iteration sees the same observation. +:attr:`~sequince.steps.context.StepContext.observation` to each SMC step / callback. -An observation can be any object: a scalar measurement, an array, a batch, or a -structured value containing several fields. Sequince does not interpret it; the -configured steps or target functions decide how it affects the particle state. +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 ------------------------------ @@ -52,17 +52,12 @@ step can evaluate the likelihood contribution from the current observation: shared_context={"observation_variance": 1.0}, ) -This run produces three states: ``result.states[i]`` is the population after +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``. -Each item is one iteration --------------------------- - -``observations`` defines iteration input, not an automatically vectorized data -axis. A list containing three arrays causes three iterations; each array is the -complete ``context.observation`` for its iteration. All configured steps run -once for each consumed item. +Interplay with ``num_steps`` +---------------------------- The interaction between ``observations`` and ``num_steps`` is important: @@ -92,7 +87,7 @@ The interaction between ``observations`` and ``num_steps`` is important: 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 has no natural bound. +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``. @@ -124,8 +119,8 @@ explicitly. Observation position -------------------- -``context.t`` is the zero-based SMC iteration. It is not necessarily an -observation index because iterations may continue after the input is exhausted. +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``. @@ -138,7 +133,9 @@ Checkpointing and resuming -------------------------- Checkpoints store ``next_observation_index`` but do not store the full input -stream. When resuming a run that has already consumed observations, pass the +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 @@ -153,17 +150,16 @@ original indexable sequence to ``run``: ) The runner resumes with ``all_observations[checkpoint.next_observation_index]``. -A one-shot iterator cannot be repositioned and is rejected after any -observations have already been consumed. See :doc:`checkpointing` for the full -checkpoint workflow. +As generators cannot be repositioned, they are not compatible with this procedure. +See :doc:`checkpointing` for the full checkpoint workflow. -JAX observations ----------------- +Observations with the JAX backend +--------------------------------- -``run_functional`` has the same host-side observation and termination semantics. -When using a compiled iteration, all observations must have a pytree structure, -array shapes, and dtypes compatible with the compiled signature. Supply a -representative first observation to the compilation helper: +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 From 5b457a92c25b2048524de798c35b202e9bc55c47 Mon Sep 17 00:00:00 2001 From: Christian Chapman-Bird Date: Thu, 8 Oct 2026 16:41:53 +0200 Subject: [PATCH 07/23] doc: repoint examples.rst to sequnice-examples --- docs/examples.rst | 16 ++-------------- 1 file changed, 2 insertions(+), 14 deletions(-) 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. From 5e19ae507bf8fda96ef12375dc0f0e32cd4a2e9f Mon Sep 17 00:00:00 2001 From: Christian Chapman-Bird Date: Fri, 9 Oct 2026 16:10:59 +0200 Subject: [PATCH 08/23] fix formatting of AI_POLICY --- AI_POLICY.md | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/AI_POLICY.md b/AI_POLICY.md index 1e92d7b3..7b94ca07 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. From 497bbcdac81d235146bd5c3d88b95a5c85c784c9 Mon Sep 17 00:00:00 2001 From: Christian Chapman-Bird Date: Fri, 16 Oct 2026 14:20:51 +0100 Subject: [PATCH 09/23] more docs improvements --- docs/getting_started.rst | 142 +++++++++++++++++++++---------------- docs/smc_configuration.rst | 16 ++--- docs/user_guide.rst | 6 +- 3 files changed, 91 insertions(+), 73 deletions(-) diff --git a/docs/getting_started.rst b/docs/getting_started.rst index db35f87d..0764f717 100644 --- a/docs/getting_started.rst +++ b/docs/getting_started.rst @@ -15,71 +15,89 @@ available as optional extras. 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 basic usage is as -follows: - -.. code-block:: python - - import numpy as np - from sequince import ParticleState, make_standard_bayesian_smc - - # Set RNG - rng = np.random.default_rng(42) - - # User-supplied log-density functions, these must be vectorized - # Theta is a 2D array of shape (n_particles, n_parameters) - # Context is dataclass that contains additional information - def log_prior(theta, context): - """Uniform prior over the specified bounds""" - bounds = context.shared["bounds"] - return np.where( - np.all((theta >= bounds[0]) & (theta <= bounds[1]), axis=1), - 0.0, - -np.inf, - ) - - def log_likelihood(theta, context): - """Unnormalized Gaussian likelihood with mean 0 and std 1""" - return -0.5 * np.sum(theta**2, axis=1) - - - # Create a standard SMC runner for Bayesian inference - # smc is the main runner - # posterior is a convenience wrapper around the log-likelihood and log-prior - # metadata is dictionary that can be used to pass information between steps - # in the SMC loop - smc, posterior, metadata = make_standard_bayesian_smc( - log_prior, - log_likelihood, +Below, we show a short example of how to use `sequince`. +For more details, see the [user guide](https://sequince.readthedocs.io/user_guide), and +see [`sequince-examples`](https://github.com/sequince-dev/sequince-examples) for more advanced usage. +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. + +```python +import numpy as np +from sequince import ParticleState, make_standard_bayesian_smc + + +def log_prior(theta, context): + """Uniform log-prior""" + # log-prior and -likelihood can optionally accept context to access shared data. + lower, upper = context.shared["bounds"] + inside = np.all((theta >= lower) & (theta <= upper), axis=-1) + log_density = -theta.shape[-1] * np.log(upper - lower) + return np.where(inside, log_density, -np.inf) + + +def log_likelihood(theta, context): + """Gaussian log-likelihood""" + x = context.shared["x"] + y = context.shared["y"] + sigma = context.shared["sigma"] + slope = theta[..., 0, None] + intercept = theta[..., 1, None] + residual = (y - (slope * x + intercept)) / sigma + return -0.5 * np.sum( + residual**2 + np.log(2 * np.pi * sigma**2), axis=-1 ) - # Shared context is a dictionary that can be used to pass user-defined - # information to steps and the log-density functions. - shared_context = dict( - bounds=[-5.0, 5.0] - ) - - # Build the initial particle cloud with samples from the prior - n = 500 - initial_particles = rng.uniform(*shared_context["bounds"], size=(n, 2)) - # Create a particle state object to hold the initial particles - # Unless specified, the particles are assumed to be equally weighted - initial_state = ParticleState(particles=initial_particles) - - # Run the SMC sampler - result = smc.run( - initial_state, - rng=rng, - shared_context=shared_context, - metadata=metadata, - ) - # Extract the posterior samples and log-evidence from the result object - samples = result.posterior_samples() - log_evidence = result.log_evidence() - log_evidence_error = result.log_evidence_error() +# Construct the runner +# `smc` is the main runner +# `posterior` contains definitions of the tempered posterior +# `metadata` contains metadata passed between the SMC steps +smc, posterior, metadata = make_standard_bayesian_smc( + log_prior, + log_likelihood, +) + +rng = np.random.default_rng(42) + +# Generate observations with true slope 2 and intercept 1. +x = np.linspace(-1, 1, 20) +sigma = 0.5 +y = 2 * x + 1 + rng.normal(scale=sigma, size=x.shape) + +bounds = [-5, 5] + +# Draw initial (slope, intercept) particles from the prior. +n = 1000 +initial_particles = rng.uniform(*bounds, size=(n, 2)) +# Construct the initial state +initial_state = ParticleState(particles=initial_particles) + +# Construct the shared context dictionary that will be exposed to the log_prior +# and log_likelihood +shared_context = dict(bounds=bounds, x=x, y=y, sigma=sigma) + +# Run the SMC +results = smc.run( + initial_state, + rng=rng, + shared_context=shared_context, + metadata=metadata, +) + +# Obtain equal-weight posterior samples and their weights +samples, sample_weights = results.posterior_samples( + rng=rng, method="rejection" +) + +# Compare the posterior means with the true slope 2 and intercept 1. +mean_slope, mean_intercept = np.average( + samples, weights=sample_weights, axis=0 +) +print(f"Posterior mean slope: {mean_slope:.2f}") +print(f"Posterior mean intercept: {mean_intercept:.2f}") +``` Next steps ---------- diff --git a/docs/smc_configuration.rst b/docs/smc_configuration.rst index 7bf5c2ce..2307eaee 100644 --- a/docs/smc_configuration.rst +++ b/docs/smc_configuration.rst @@ -18,15 +18,15 @@ Bayesian inference pipeline is the pre-built helper: callbacks=[TracePlot()], ) -``make_standard_bayesian_smc`` wires together -:class:`~sequince.steps.weighting.AdaptiveTemperedWeighting`, -a configurable resampling step, and a configurable mutation step. See -:doc:`resampling_options` and :doc:`mutation_options` for the available -sampling strategies, and :doc:`jax_support` for JAX-backed execution. -For a mutation kernel without SMC weighting and resampling, see -:doc:`mcmc_sampler`. +``make_standard_bayesian_smc`` constructs an SMC runner that adaptively +steps through a series of simulated annealing temperatures +(:class:`~sequince.steps.weighting.AdaptiveTemperedWeighting`), applying +See :doc:`resampling_options` and :doc:`mutation_options` for the available +sampling strategies. +See :doc:`jax_support` for JAX-backed execution, which needs some additional +function calls to enable full JAX functionality (``jit``, ``vmap``, ``grad``, ...). -Pass the returned metadata when starting the run: +When starting the run, pass the ``metadata`` returned by the convenience constructor: .. code-block:: python diff --git a/docs/user_guide.rst b/docs/user_guide.rst index 5e550bc0..badcbefd 100644 --- a/docs/user_guide.rst +++ b/docs/user_guide.rst @@ -3,8 +3,7 @@ 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. The same -mutation implementations can also be used through the standalone MCMC runner. +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 @@ -19,9 +18,10 @@ through the SMC loop, and more specialised execution and extension points. :caption: Running inference smc_configuration - mcmc_sampler observations evidence_uncertainty + jax_support + mcmc_sampler .. toctree:: :maxdepth: 1 From fb0c3808791911c34016782c30f5fcdf44f3caa1 Mon Sep 17 00:00:00 2001 From: "Michael J. Williams" Date: Mon, 19 Oct 2026 09:28:29 +0200 Subject: [PATCH 10/23] DOC: use autoapi to API docs --- docs/api/core.rst | 76 --------------- docs/api/index.rst | 9 -- docs/api/steps.rst | 214 ----------------------------------------- docs/api/utilities.rst | 23 ----- docs/conf.py | 27 ++---- docs/index.rst | 2 +- docs/jax_support.rst | 3 +- docs/requirements.txt | 1 + docs/step_context.rst | 3 +- pyproject.toml | 4 +- 10 files changed, 17 insertions(+), 345 deletions(-) delete mode 100644 docs/api/core.rst delete mode 100644 docs/api/index.rst delete mode 100644 docs/api/steps.rst delete mode 100644 docs/api/utilities.rst diff --git a/docs/api/core.rst b/docs/api/core.rst deleted file mode 100644 index cfe75ce6..00000000 --- a/docs/api/core.rst +++ /dev/null @@ -1,76 +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: - -Markov chain Monte Carlo driver -------------------------------- - -.. automodule:: sequince.mcmc - :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: - -Particle transforms -------------------- - -.. automodule:: sequince.transforms - :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/conf.py b/docs/conf.py index c439fb13..9ecb403b 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -2,16 +2,8 @@ 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" @@ -33,24 +25,25 @@ # -- General configuration --------------------------------------------------- extensions = [ - "sphinx.ext.autodoc", - "sphinx.ext.autosummary", + "autoapi.extension", "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", +# Generate module pages from source without importing optional backends. +autoapi_dirs = ["../src/sequince"] +autoapi_add_toctree_entry = False +autoapi_options = [ + "members", + "undoc-members", + "show-inheritance", + "show-module-summary", + "imported-members", ] templates_path = ["_templates"] diff --git a/docs/index.rst b/docs/index.rst index 58cff83e..74a65247 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -13,7 +13,7 @@ Markov chain Monte Carlo runner. getting_started user_guide examples - api/index + API reference Indices and tables ------------------ diff --git a/docs/jax_support.rst b/docs/jax_support.rst index db0d75d3..3590f78e 100644 --- a/docs/jax_support.rst +++ b/docs/jax_support.rst @@ -164,4 +164,5 @@ key. API reference ------------- -.. autofunction:: sequince.jax.compile_functional_iteration +See :func:`~sequince.jax.compile_functional_iteration` in the generated API +reference. 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/step_context.rst b/docs/step_context.rst index 376b5173..aca0ad93 100644 --- a/docs/step_context.rst +++ b/docs/step_context.rst @@ -145,5 +145,4 @@ the resumed runtime because it is deliberately not stored in the checkpoint. API reference ------------- -.. autoclass:: sequince.steps.context.StepContext - :members: +See :class:`~sequince.steps.context.StepContext` in the generated API reference. diff --git a/pyproject.toml b/pyproject.toml index 3b3bbe42..1cea954f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -164,10 +164,10 @@ typecheck = "mypy" [tool.pixi.feature.doc.dependencies] sphinx = ">=7.2" sphinx-book-theme = ">=1.1" +sphinx-autoapi = ">=3.8,<4" [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" From 4ec7a83029b0c6cea6fe62d85fc4f3fcea924765 Mon Sep 17 00:00:00 2001 From: "Michael J. Williams" Date: Mon, 19 Oct 2026 09:32:30 +0200 Subject: [PATCH 11/23] DOC: fix code in getting started --- docs/getting_started.rst | 90 +++++----------------------------------- 1 file changed, 10 insertions(+), 80 deletions(-) diff --git a/docs/getting_started.rst b/docs/getting_started.rst index 0764f717..c4df8770 100644 --- a/docs/getting_started.rst +++ b/docs/getting_started.rst @@ -15,89 +15,19 @@ available as optional extras. Quick start ----------- -Below, we show a short example of how to use `sequince`. -For more details, see the [user guide](https://sequince.readthedocs.io/user_guide), and -see [`sequince-examples`](https://github.com/sequince-dev/sequince-examples) for more advanced usage. -Here, we fit a straight line, `y = slope * x + intercept`, to synthetic data with +Below, we show a short example of how to use ``sequince``. +For more details, see the :doc:`user_guide`, and see +`sequince-examples `_ +for more advanced usage. +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. +``(n_particles, n_parameters)`` and return one log density per particle. -```python -import numpy as np -from sequince import ParticleState, make_standard_bayesian_smc - - -def log_prior(theta, context): - """Uniform log-prior""" - # log-prior and -likelihood can optionally accept context to access shared data. - lower, upper = context.shared["bounds"] - inside = np.all((theta >= lower) & (theta <= upper), axis=-1) - log_density = -theta.shape[-1] * np.log(upper - lower) - return np.where(inside, log_density, -np.inf) - - -def log_likelihood(theta, context): - """Gaussian log-likelihood""" - x = context.shared["x"] - y = context.shared["y"] - sigma = context.shared["sigma"] - slope = theta[..., 0, None] - intercept = theta[..., 1, None] - residual = (y - (slope * x + intercept)) / sigma - return -0.5 * np.sum( - residual**2 + np.log(2 * np.pi * sigma**2), axis=-1 - ) - - -# Construct the runner -# `smc` is the main runner -# `posterior` contains definitions of the tempered posterior -# `metadata` contains metadata passed between the SMC steps -smc, posterior, metadata = make_standard_bayesian_smc( - log_prior, - log_likelihood, -) - -rng = np.random.default_rng(42) - -# Generate observations with true slope 2 and intercept 1. -x = np.linspace(-1, 1, 20) -sigma = 0.5 -y = 2 * x + 1 + rng.normal(scale=sigma, size=x.shape) - -bounds = [-5, 5] - -# Draw initial (slope, intercept) particles from the prior. -n = 1000 -initial_particles = rng.uniform(*bounds, size=(n, 2)) -# Construct the initial state -initial_state = ParticleState(particles=initial_particles) - -# Construct the shared context dictionary that will be exposed to the log_prior -# and log_likelihood -shared_context = dict(bounds=bounds, x=x, y=y, sigma=sigma) - -# Run the SMC -results = smc.run( - initial_state, - rng=rng, - shared_context=shared_context, - metadata=metadata, -) - -# Obtain equal-weight posterior samples and their weights -samples, sample_weights = results.posterior_samples( - rng=rng, method="rejection" -) - -# Compare the posterior means with the true slope 2 and intercept 1. -mean_slope, mean_intercept = np.average( - samples, weights=sample_weights, axis=0 -) -print(f"Posterior mean slope: {mean_slope:.2f}") -print(f"Posterior mean intercept: {mean_intercept:.2f}") -``` +.. literalinclude:: ../README.md + :language: python + :start-after: ```python + :end-before: ``` Next steps ---------- From 3bd746287ed5a6e21065573ccccef30cafe5fc65 Mon Sep 17 00:00:00 2001 From: "Michael J. Williams" Date: Mon, 19 Oct 2026 09:39:57 +0200 Subject: [PATCH 12/23] DOC: fix sphinx reference warnings --- docs/conf.py | 4 +++- docs/mcmc_sampler.rst | 2 +- docs/mutation_options.rst | 8 ++++---- docs/user_guide.rst | 1 - src/sequince/result.py | 9 +++++---- src/sequince/smc.py | 6 ++++-- src/sequince/steps/mutation/metrics/state.py | 5 +++-- src/sequince/transforms.py | 11 +++-------- 8 files changed, 23 insertions(+), 23 deletions(-) diff --git a/docs/conf.py b/docs/conf.py index 9ecb403b..d5c45564 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -34,16 +34,18 @@ autodoc_typehints = "description" napoleon_google_docstring = False napoleon_numpy_docstring = True +# 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", - "imported-members", ] templates_path = ["_templates"] diff --git a/docs/mcmc_sampler.rst b/docs/mcmc_sampler.rst index 9989cf85..42a55fe5 100644 --- a/docs/mcmc_sampler.rst +++ b/docs/mcmc_sampler.rst @@ -81,7 +81,7 @@ helper: ``"minipcn"``, ``"emcee"``, ``"blackjax"``, and See :doc:`mutation_options` for available backends, optional dependencies, and kernel configuration. A configured -:class:`~sequince.steps.mutation.MutationStep` can also be passed directly to +: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. diff --git a/docs/mutation_options.rst b/docs/mutation_options.rst index 1b798621..1e2d6cc5 100644 --- a/docs/mutation_options.rst +++ b/docs/mutation_options.rst @@ -16,16 +16,16 @@ selects a mutation implementation with ``mutation``: - Implementation - Typical use * - ``"minipcn"`` - - :class:`~sequince.steps.mutation.MiniPCNMutation` + - :class:`~sequince.steps.mutation.minipcn.MiniPCNMutation` - General default with eager and functional execution support. * - ``"emcee"`` - - :class:`~sequince.steps.mutation.EmceeMutation` + - :class:`~sequince.steps.mutation.emcee.EmceeMutation` - Ensemble MCMC with emcee's collection of proposal moves. * - ``"blackjax"`` - - :class:`~sequince.steps.mutation.BlackJAXMutation` + - :class:`~sequince.steps.mutation.blackjax.BlackJAXMutation` - JAX-backed NUTS, HMC, or random-walk mutation. * - ``"metropolis-hastings"`` - - :class:`~sequince.steps.mutation.MetropolisHastingsMutation` + - :class:`~sequince.steps.mutation.metropolis.MetropolisHastingsMutation` - Explicit control over a proposal kernel and transition ratio. * - ``"none"`` or ``None`` - No mutation step diff --git a/docs/user_guide.rst b/docs/user_guide.rst index badcbefd..d2b7074d 100644 --- a/docs/user_guide.rst +++ b/docs/user_guide.rst @@ -20,7 +20,6 @@ through the SMC loop, and more specialised execution and extension points. smc_configuration observations evidence_uncertainty - jax_support mcmc_sampler .. toctree:: diff --git a/src/sequince/result.py b/src/sequince/result.py index 597d8e9d..14ede7e4 100644 --- a/src/sequince/result.py +++ b/src/sequince/result.py @@ -820,11 +820,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/mutation/metrics/state.py b/src/sequince/steps/mutation/metrics/state.py index 74ba11d6..b6735e19 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/transforms.py b/src/sequince/transforms.py index c41f5e1b..ec435749 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( From 58da6a38642b1f02f78c3b5f91c80e323695bb87 Mon Sep 17 00:00:00 2001 From: "Michael J. Williams" Date: Tue, 20 Oct 2026 12:04:17 +0200 Subject: [PATCH 13/23] DOC: simplify mutation and resampling docs --- docs/Makefile | 4 +- docs/jax_support.rst | 4 +- docs/mutation_options.rst | 373 +++++++++++++----------------------- docs/resampling_options.rst | 166 ++++++++-------- 4 files changed, 218 insertions(+), 329 deletions(-) diff --git a/docs/Makefile b/docs/Makefile index d4bb2cbb..bcd747f9 100644 --- a/docs/Makefile +++ b/docs/Makefile @@ -12,7 +12,9 @@ BUILDDIR = _build help: @$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) -.PHONY: help Makefile +.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). diff --git a/docs/jax_support.rst b/docs/jax_support.rst index 3590f78e..836a0800 100644 --- a/docs/jax_support.rst +++ b/docs/jax_support.rst @@ -7,8 +7,8 @@ loop, making one SMC iteration suitable for ``jax.jit`` while retaining normal result collection, callbacks, and checkpoint handling between iterations. -Install JAX support with ``uv add "sequince[jax]"``. To use BlackJAX mutation, -install ``uv add "sequince[blackjax]"`` instead; BlackJAX brings its JAX +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 diff --git a/docs/mutation_options.rst b/docs/mutation_options.rst index 1e2d6cc5..b04f433d 100644 --- a/docs/mutation_options.rst +++ b/docs/mutation_options.rst @@ -1,267 +1,160 @@ -Mutation options -================ +Mutation configuration +====================== -Mutation rejuvenates the particle population after weighting and resampling. -It applies MCMC transitions that preserve the current target distribution, -helping duplicated particles explore again without changing their weights. - -The :func:`~sequince.prebuilt_samplers.make_standard_bayesian_smc` helper -selects a mutation implementation with ``mutation``: +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:`mcmc_sampler`. +Each step takes a target callable or :class:`~sequince.target.Target` object. .. list-table:: :header-rows: 1 - :widths: 22 30 48 + :widths: 18 24 28 30 - * - Value + * - Method + - ``mutation`` name - Implementation - Typical use - * - ``"minipcn"`` + * - MiniPCN + - ``"minipcn"`` - :class:`~sequince.steps.mutation.minipcn.MiniPCNMutation` - - General default with eager and functional execution support. - * - ``"emcee"`` + - Supports eager and functional execution. + * - emcee + - ``"emcee"`` - :class:`~sequince.steps.mutation.emcee.EmceeMutation` - - Ensemble MCMC with emcee's collection of proposal moves. - * - ``"blackjax"`` + - Ensemble MCMC using emcee moves. + * - BlackJAX + - ``"blackjax"`` - :class:`~sequince.steps.mutation.blackjax.BlackJAXMutation` - JAX-backed NUTS, HMC, or random-walk mutation. - * - ``"metropolis-hastings"`` + * - 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` - - Explicit control over a proposal kernel and transition ratio. - * - ``"none"`` or ``None`` - - No mutation step - - Reweighting and resampling only. - * - An ``SMCStep`` instance - - The supplied step - - A custom, already configured mutation implementation. + - 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 -------------------- -``mutation_n_steps`` controls how many MCMC transitions are applied to every -particle during one SMC iteration. The pre-built sampler defaults to five. -Strategy-specific options are passed through ``mutation_kwargs``: +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 - smc, posterior, metadata = make_standard_bayesian_smc( - log_prior, - log_likelihood, - mutation="minipcn", - mutation_n_steps=10, - mutation_kwargs={ - "store_chains": True, - "transform_refit": "per-context", - }, - particle_transform="affine", - ) - -All built-in mutation steps support these common options: + from sequince.steps.mutation import MiniPCNMutation + from sequince.transforms import AffineTransform -``store_chains`` - Include the intermediate MCMC positions in the step output. It defaults to - ``False`` because full chains can be large. Acceptance diagnostics and the - number of transitions are returned even when chains are not stored. - -``transform_refit`` - Control how an adaptive particle transform is fitted. ``"always"`` fits on - every mutation call, ``"per-context"`` reuses it within the same iteration - and observation position, and ``"once"`` reuses the first fit. Eager - execution also accepts a callable predicate; compiled functional execution - requires one of the named policies. - -Particle transforms -------------------- - -The helper mutates particles in transformed coordinates by default using -``particle_transform="affine"``. This standardizes the particle cloud before -the MCMC transition and maps accepted positions back afterward. Configure it -through ``particle_transform_kwargs``: - -.. code-block:: python - - smc, posterior, metadata = make_standard_bayesian_smc( - log_prior, - log_likelihood, - particle_transform="affine", - particle_transform_kwargs={"eps": 1e-5}, + mutation_step = MiniPCNMutation( + target, + n_steps=10, + store_chains=True, + transform=AffineTransform(), ) -Other registered names are ``"identity"``, -``"bounded_to_unit_hypercube"``, ``"bounded_to_unbounded"``, and -``"periodic_boundary"``. The bounded and periodic transforms require bounds in -``particle_transform_kwargs``. Pass ``None`` to mutate directly in the original -coordinates, or pass a configured -:class:`~sequince.transforms.ParticleTransform` instance. - -Auxiliary information for custom transforms -~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - -A custom transform may return invocation-specific information that its inverse -needs without storing mutable state on the transform. For example, a -many-to-one transform can preserve the branch assigned to each particle:: - - from sequince.array_api import namespace - from sequince.transforms import ParticleTransform - - - class SignFoldingTransform(ParticleTransform): - def forward(self, values, context=None, *, rng=None): - xp = namespace(values) - signs = xp.where(values < 0, -xp.ones_like(values), xp.ones_like(values)) - transformed = xp.abs(values) - log_det = xp.zeros_like(values[..., 0]) - return transformed, log_det, signs - - def inverse( - self, - values, - context=None, - *, - auxiliary=None, - rng=None, - ): - if auxiliary is None: - raise ValueError("Sign information is required for inversion.") - xp = namespace(values) - return values * auxiliary, xp.zeros_like(values[..., 0]) - -Auxiliary information is ephemeral: mutation passes it from one ``forward`` -call through transformed-space operations to the corresponding ``inverse``. -It may be a single array or any Optree-compatible structure, including a tuple, -list, or dictionary such as ``(mode_ids, mask)`` or -``{"mode_ids": mode_ids, "mask": mask}``. Simple transforms do not need a -Sequince-specific wrapper. Libraries needing named fields may register their -own custom pytree types; immutable dataclasses can use -:func:`sequince.utils.pytree_dataclass` for Optree and JAX registration. - -For a batch of ``n_particles``, Sequince treats an auxiliary array leaf as -particle-aligned exactly when ``leaf.ndim > 0`` and -``leaf.shape[0] == n_particles``. A mutation operating on a particle subset -applies the same selection to those leaves before inversion, while preserving -other leaves unchanged. Consequently, an unrelated array whose leading size -happens to equal ``n_particles`` is also selected; this is an intentional -limitation of the current shape-based convention. - -Particle-aligned auxiliary works with batch-preserving mutation and BlackJAX. -It is currently rejected by emcee: emcee selects walker subsets internally, -but its standard log-probability callback does not expose their original walker -indices, so Sequince cannot guarantee the correct auxiliary entry. - -MiniPCN -------- - -``"minipcn"`` is the default and requires no extra installation. It uses the -``"tpcn"`` step function by default and targets an acceptance rate of ``0.234``. -Both can be changed through ``mutation_kwargs``: - -.. code-block:: python - - mutation="minipcn", - mutation_kwargs={ - "step_fn": "tpcn", - "target_acceptance_rate": 0.3, - }, - -MiniPCN implements explicit random-state threading and can be used in a compiled -functional pipeline. See :doc:`jax_support` for the functional runner. - -emcee ------ - -Install emcee support with ``uv add "sequince[emcee]"``. With no move specified, -emcee uses its default stretch move. Named choices are ``"de"``, -``"de-snooker"``, ``"gaussian"``, ``"kde"``, ``"stretch"``, and ``"walk"``. -The ``"bilby-de"`` preset supplies a Bilby-style weighted -differential-evolution schedule: - -.. code-block:: python - - mutation="emcee", - mutation_n_steps=20, - mutation_kwargs={ - "move": "bilby-de", - }, +The implementations in the table accept these constructor options: -The preset is equivalent to using emcee's ``DEMove`` with its automatically -chosen differential-evolution scale 90% of the time and ``gamma0=1.0`` 10% of -the time. Both moves use ``sigma=1e-5``: +* ``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. -.. code-block:: python - - mutation_kwargs={ - "move": [ - { - "move": "de", - "move_kwargs": {"gamma0": None, "sigma": 1e-5}, - "weight": 0.9, - }, - { - "move": "de", - "move_kwargs": {"gamma0": 1.0, "sigma": 1e-5}, - "weight": 0.1, - }, - ], - }, - -This preset does not require Bilby itself; it constructs the schedule from -emcee move objects. The expanded form can be adapted to create a custom -Bilby-like schedule with different moves, parameters, or weights. - -``move`` may also be an emcee move object, a mapping containing ``move`` and -``move_kwargs``, or a weighted sequence of move specifications. Other keyword -arguments are forwarded to ``emcee.EnsembleSampler``; ``run_kwargs`` is -forwarded to its ``run_mcmc`` method. - -Metropolis--Hastings --------------------- +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. -The Metropolis--Hastings option requires a proposal ``kernel``. Sequince -provides ``"gaussian-random-walk"`` and ``"mixture"`` as named kernels: +Install optional backends with ``pip install "sequince[emcee]"``, +``pip install "sequince[eryn]"``, or ``pip install "sequince[blackjax]"``. +See :doc:`jax_support` for compiled execution. -.. code-block:: python - - mutation="metropolis-hastings", - mutation_n_steps=10, - mutation_kwargs={ - "kernel": "gaussian-random-walk", - "kernel_kwargs": {"scale": 0.25}, - }, - -The mixture kernel accepts a ``kernels`` sequence and optional ``weights`` in -``kernel_kwargs``. A custom kernel object or factory can be passed directly; -it must provide proposals and one log transition ratio per particle. When a -kernel object is supplied, configure it directly rather than also passing -``kernel_kwargs``. - -BlackJAX --------- - -``"blackjax"`` accepts ``"nuts"`` (the default), ``"hmc"``, -``"random-walk"``, or a callable algorithm factory through -``mutation_kwargs["algorithm"]``. Algorithm parameters belong in -``mutation_kwargs["algorithm_kwargs"]``. NUTS and HMC require differentiable, -JAX-compatible targets and JAX-backed particles. - -See :doc:`jax_support` for installation, a complete compiled example, random-key -threading, and current functional limitations. - -Disabling or replacing mutation -------------------------------- - -Use ``mutation="none"`` when mutation is deliberately unnecessary. Resampling -without mutation can rapidly reduce particle diversity, so monitor effective -sample size and duplicated particles when choosing this mode. - -You can also pass an initialized :class:`~sequince.steps.base.SMCStep` instance -as ``mutation``. The helper inserts that object unchanged; it does not inject -the generated posterior, particle transform, ``mutation_n_steps``, or -``mutation_kwargs``. Configure a custom step completely before passing it, or -construct the full pipeline manually as described in :doc:`smc_configuration`. - -Reading mutation outputs ------------------------- +Particle transforms +------------------- -Each mutation step returns diagnostics in its position within the iteration's -output tuple. The runner collects those tuples in ``result.outputs``. Built-in -outputs include acceptance information, the configured number of transitions, -and ``chain``; the chain value is ``None`` unless ``store_chains=True``. +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:`smc_configuration` for custom pipelines +and the API links above for method-specific options and outputs. diff --git a/docs/resampling_options.rst b/docs/resampling_options.rst index 6ed18be2..df3dc2d1 100644 --- a/docs/resampling_options.rst +++ b/docs/resampling_options.rst @@ -1,94 +1,88 @@ -Resampling options -================== - -Resampling replaces a weighted particle population with equally weighted -particles drawn from it. This prevents a small number of particles from -carrying almost all of the weight, at the cost of duplicating some particles. -Mutation is commonly applied afterwards to restore diversity. - -Choosing a strategy -------------------- - -The :func:`~sequince.prebuilt_samplers.make_standard_bayesian_smc` helper -selects the ancestor-sampling algorithm with ``resampling``: - -``"systematic"`` - Use :class:`~sequince.steps.resampling.SystematicResampling`. This is the - default and uses one random offset for the complete population. - -``"stratified"`` - Use :class:`~sequince.steps.resampling.StratifiedResampling`, with one - random position in each equal-width stratum. - -``"multinomial"`` - Use :class:`~sequince.steps.resampling.MultinomialResampling`, drawing each - ancestor independently according to the particle weights. - -For example: +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 - smc, posterior, metadata = make_standard_bayesian_smc( - log_prior, - log_likelihood, - resampling="stratified", - ) - -Choosing when to resample -------------------------- - -The resampling algorithm and its schedule are configured separately: + from sequince.steps.resampling import StratifiedResampling -.. code-block:: python - - smc, posterior, metadata = make_standard_bayesian_smc( - log_prior, - log_likelihood, - weighting_criterion="cess", - resampling="stratified", - resampling_ess_threshold=0.4, - always_resample=False, - ) - -With ``always_resample=False``, resampling occurs when the effective sample size -divided by the number of particles falls below ``resampling_ess_threshold``. -Set the threshold to ``None`` to disable threshold-triggered resampling. With -``always_resample=True``, the runner resamples every iteration and ignores the -threshold. - -When ``always_resample`` is omitted, its default depends on -``weighting_criterion``: ESS-based tempering resamples every iteration, while -CESS-based tempering uses the ESS threshold. Setting ``always_resample=False`` -with ``weighting_criterion="ess"`` is invalid because ESS-based tempering relies -on resampling to meet its target. - -Functional execution currently requires ``always_resample=True``. See -:doc:`jax_support` for the other functional-execution constraints. - -Configuring a resampling step directly ---------------------------------------- - -In a manually assembled pipeline, instantiate the selected step with the same -schedule arguments: - -.. code-block:: python - - from sequince.steps import StratifiedResampling - - resampling = StratifiedResampling( + resampling_step = StratifiedResampling( ess_threshold=0.4, always_resample=False, + resample_at_end=True, ) -The step records the pre-resampling ESS and whether resampling occurred in both -the context metadata and its output. The particle state records selected -ancestor indices when resampling occurs, which allows particle genealogies to -be reconstructed. - -Resampling and evidence estimates ---------------------------------- - -The resampling strategy affects Monte Carlo variability even though all three -built-in algorithms target the same weighted population. See -:doc:`evidence_uncertainty` for the uncertainty estimators and their resampling -requirements. +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:`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:`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``. From b8ef7bd7ce64c77f811ece41061c37d927103cf8 Mon Sep 17 00:00:00 2001 From: "Michael J. Williams" Date: Tue, 20 Oct 2026 12:07:07 +0200 Subject: [PATCH 14/23] DOC: rename files --- docs/getting_started.rst | 2 +- docs/mcmc_sampler.rst | 2 +- docs/mutation_options.rst | 160 ------------------------------------ docs/resampling_options.rst | 88 -------------------- docs/smc_configuration.rst | 2 +- docs/user_guide.rst | 4 +- 6 files changed, 5 insertions(+), 253 deletions(-) delete mode 100644 docs/mutation_options.rst delete mode 100644 docs/resampling_options.rst diff --git a/docs/getting_started.rst b/docs/getting_started.rst index c4df8770..e5bfd638 100644 --- a/docs/getting_started.rst +++ b/docs/getting_started.rst @@ -34,7 +34,7 @@ Next steps See :doc:`evidence_uncertainty` for interpreting the evidence estimate and its uncertainty. To customise the run, continue with :doc:`smc_configuration` and -the sampling choices in :doc:`resampling_options` and :doc:`mutation_options`. +the sampling choices in :doc:`resampling_configuration` and :doc:`mutation_configuration`. The :doc:`step_context` guide explains the distinction between shared inputs and evolving metadata. For sampling without SMC weighting and resampling, see :doc:`mcmc_sampler`. diff --git a/docs/mcmc_sampler.rst b/docs/mcmc_sampler.rst index 42a55fe5..cd188b20 100644 --- a/docs/mcmc_sampler.rst +++ b/docs/mcmc_sampler.rst @@ -79,7 +79,7 @@ helper: ``"minipcn"``, ``"emcee"``, ``"blackjax"``, and }, ) -See :doc:`mutation_options` for available backends, optional dependencies, and +See :doc:`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 diff --git a/docs/mutation_options.rst b/docs/mutation_options.rst deleted file mode 100644 index b04f433d..00000000 --- a/docs/mutation_options.rst +++ /dev/null @@ -1,160 +0,0 @@ -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:`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:`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:`smc_configuration` for custom pipelines -and the API links above for method-specific options and outputs. diff --git a/docs/resampling_options.rst b/docs/resampling_options.rst deleted file mode 100644 index df3dc2d1..00000000 --- a/docs/resampling_options.rst +++ /dev/null @@ -1,88 +0,0 @@ -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:`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:`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/docs/smc_configuration.rst b/docs/smc_configuration.rst index 2307eaee..8399166a 100644 --- a/docs/smc_configuration.rst +++ b/docs/smc_configuration.rst @@ -21,7 +21,7 @@ Bayesian inference pipeline is the pre-built helper: ``make_standard_bayesian_smc`` constructs an SMC runner that adaptively steps through a series of simulated annealing temperatures (:class:`~sequince.steps.weighting.AdaptiveTemperedWeighting`), applying -See :doc:`resampling_options` and :doc:`mutation_options` for the available +See :doc:`resampling_configuration` and :doc:`mutation_configuration` for the available sampling strategies. See :doc:`jax_support` for JAX-backed execution, which needs some additional function calls to enable full JAX functionality (``jit``, ``vmap``, ``grad``, ...). diff --git a/docs/user_guide.rst b/docs/user_guide.rst index d2b7074d..a5e9519f 100644 --- a/docs/user_guide.rst +++ b/docs/user_guide.rst @@ -26,8 +26,8 @@ through the SMC loop, and more specialised execution and extension points. :maxdepth: 1 :caption: Sampling options - resampling_options - mutation_options + resampling_configuration + mutation_configuration .. toctree:: :maxdepth: 1 From 81b297b1a150552ebedeb3fc9dfa6409d5e4310b Mon Sep 17 00:00:00 2001 From: "Michael J. Williams" Date: Tue, 20 Oct 2026 12:11:20 +0200 Subject: [PATCH 15/23] DOC: add missing files --- docs/mutation_configuration.rst | 160 ++++++++++++++++++++++++++++++ docs/resampling_configuration.rst | 88 ++++++++++++++++ docs/smc_configuration.rst | 7 +- 3 files changed, 251 insertions(+), 4 deletions(-) create mode 100644 docs/mutation_configuration.rst create mode 100644 docs/resampling_configuration.rst diff --git a/docs/mutation_configuration.rst b/docs/mutation_configuration.rst new file mode 100644 index 00000000..b04f433d --- /dev/null +++ b/docs/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:`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:`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:`smc_configuration` for custom pipelines +and the API links above for method-specific options and outputs. diff --git a/docs/resampling_configuration.rst b/docs/resampling_configuration.rst new file mode 100644 index 00000000..df3dc2d1 --- /dev/null +++ b/docs/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:`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:`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/docs/smc_configuration.rst b/docs/smc_configuration.rst index 8399166a..4c638cce 100644 --- a/docs/smc_configuration.rst +++ b/docs/smc_configuration.rst @@ -20,11 +20,10 @@ Bayesian inference pipeline is the pre-built helper: ``make_standard_bayesian_smc`` constructs an SMC runner that adaptively steps through a series of simulated annealing temperatures -(:class:`~sequince.steps.weighting.AdaptiveTemperedWeighting`), applying +(: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:`resampling_configuration` and :doc:`mutation_configuration` for the available -sampling strategies. -See :doc:`jax_support` for JAX-backed execution, which needs some additional -function calls to enable full JAX functionality (``jit``, ``vmap``, ``grad``, ...). +reasmpling and mutation options and configuration. When starting the run, pass the ``metadata`` returned by the convenience constructor: From f7f3efee2b4b3bf183ab5f63f32c506292ceced8 Mon Sep 17 00:00:00 2001 From: "Michael J. Williams" Date: Tue, 20 Oct 2026 13:52:00 +0200 Subject: [PATCH 16/23] DOC: update index --- docs/index.rst | 16 ++++++++++++++-- docs/user_guide.rst | 40 ++++++---------------------------------- 2 files changed, 20 insertions(+), 36 deletions(-) diff --git a/docs/index.rst b/docs/index.rst index 74a65247..ea54306f 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -3,8 +3,20 @@ sequince documentation ``sequince`` provides modular building blocks for assembling Sequential Monte Carlo (SMC) algorithms for Bayesian inference, particle filtering, and related -sequential methods. Its mutation backends can also be used through a standalone -Markov chain Monte Carlo runner. +sequential methods. + +Overview +-------- + +* :doc:`getting_started` covers installation and a complete straight-line + fitting example. +* :doc:`user_guide` explains how to configure sampling, work with particle + states and results, and build custom pipelines. After the quickstart, + continue with :doc:`smc_configuration` and :doc:`evidence_uncertainty`. +* :doc:`examples` links to worked examples in the ``sequince-examples`` + repository. +* The :doc:`API reference ` lists classes, functions, + and their parameters. .. toctree:: :maxdepth: 2 diff --git a/docs/user_guide.rst b/docs/user_guide.rst index a5e9519f..be0eeb3a 100644 --- a/docs/user_guide.rst +++ b/docs/user_guide.rst @@ -14,38 +14,10 @@ output. The remaining sections explain sampling choices, the objects passed through the SMC loop, and more specialised execution and extension points. .. toctree:: - :maxdepth: 1 - :caption: Running inference + :maxdepth: 2 - smc_configuration - observations - evidence_uncertainty - mcmc_sampler - -.. toctree:: - :maxdepth: 1 - :caption: Sampling options - - resampling_configuration - mutation_configuration - -.. toctree:: - :maxdepth: 1 - :caption: Core concepts - - particle_state - step_context - random_number_generation - -.. toctree:: - :maxdepth: 1 - :caption: Advanced workflows - - jax_support - checkpointing - -.. toctree:: - :maxdepth: 1 - :caption: Extending Sequince - - custom_steps_callbacks + user_guide/running_inference + user_guide/sampling_options + user_guide/core_concepts + user_guide/advanced_workflows + user_guide/extending_sequince From a4b1c8c6731a26c590bd4150f86994ba87659c3d Mon Sep 17 00:00:00 2001 From: "Michael J. Williams" Date: Thu, 22 Oct 2026 11:44:13 +0200 Subject: [PATCH 17/23] DOC: restructure and add more details --- docs/.gitignore | 1 + docs/custom_steps_callbacks.rst | 67 ---------------- docs/getting_started.rst | 25 +++--- docs/index.rst | 10 ++- docs/installation.rst | 47 +++++++++++ .../advanced_workflows}/checkpointing.rst | 0 docs/user_guide/advanced_workflows/index.rst | 28 +++++++ .../advanced_workflows}/jax_support.rst | 2 +- .../advanced_workflows}/mcmc_sampler.rst | 2 +- .../core_concepts/array_backends.rst | 77 +++++++++++++++++++ docs/user_guide/core_concepts/index.rst | 29 +++++++ .../core_concepts}/particle_state.rst | 17 ++-- .../random_number_generation.rst | 11 ++- .../core_concepts}/step_context.rst | 2 +- .../extending_sequince/custom_callbacks.rst | 41 ++++++++++ .../extending_sequince/custom_steps.rst | 30 ++++++++ docs/user_guide/extending_sequince/index.rst | 30 ++++++++ docs/{user_guide.rst => user_guide/index.rst} | 17 ++-- .../running_inference/callbacks.rst | 68 ++++++++++++++++ .../evidence_uncertainty.rst | 11 +++ docs/user_guide/running_inference/index.rst | 36 +++++++++ .../running_inference}/observations.rst | 4 +- .../running_inference/prebuilt_samplers.rst | 55 +++++++++++++ docs/user_guide/running_inference/results.rst | 71 +++++++++++++++++ .../running_inference}/smc_configuration.rst | 10 ++- docs/user_guide/sampling_options/index.rst | 31 ++++++++ .../mutation_configuration.rst | 6 +- .../resampling_configuration.rst | 4 +- 28 files changed, 619 insertions(+), 113 deletions(-) create mode 100644 docs/.gitignore delete mode 100644 docs/custom_steps_callbacks.rst create mode 100644 docs/installation.rst rename docs/{ => user_guide/advanced_workflows}/checkpointing.rst (100%) create mode 100644 docs/user_guide/advanced_workflows/index.rst rename docs/{ => user_guide/advanced_workflows}/jax_support.rst (98%) rename docs/{ => user_guide/advanced_workflows}/mcmc_sampler.rst (97%) create mode 100644 docs/user_guide/core_concepts/array_backends.rst create mode 100644 docs/user_guide/core_concepts/index.rst rename docs/{ => user_guide/core_concepts}/particle_state.rst (77%) rename docs/{ => user_guide/core_concepts}/random_number_generation.rst (92%) rename docs/{ => user_guide/core_concepts}/step_context.rst (98%) create mode 100644 docs/user_guide/extending_sequince/custom_callbacks.rst create mode 100644 docs/user_guide/extending_sequince/custom_steps.rst create mode 100644 docs/user_guide/extending_sequince/index.rst rename docs/{user_guide.rst => user_guide/index.rst} (55%) create mode 100644 docs/user_guide/running_inference/callbacks.rst rename docs/{ => user_guide/running_inference}/evidence_uncertainty.rst (87%) create mode 100644 docs/user_guide/running_inference/index.rst rename docs/{ => user_guide/running_inference}/observations.rst (97%) create mode 100644 docs/user_guide/running_inference/prebuilt_samplers.rst create mode 100644 docs/user_guide/running_inference/results.rst rename docs/{ => user_guide/running_inference}/smc_configuration.rst (83%) create mode 100644 docs/user_guide/sampling_options/index.rst rename docs/{ => user_guide/sampling_options}/mutation_configuration.rst (97%) rename docs/{ => user_guide/sampling_options}/resampling_configuration.rst (94%) 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/custom_steps_callbacks.rst b/docs/custom_steps_callbacks.rst deleted file mode 100644 index 49c59ee1..00000000 --- a/docs/custom_steps_callbacks.rst +++ /dev/null @@ -1,67 +0,0 @@ -Custom steps and callbacks -========================== - -Steps transform particle state within an iteration. Callbacks run after the -step pipeline and are intended for observation, reporting, or run-level control. -Both receive a :class:`~sequince.steps.context.StepContext`; see -:doc:`step_context` for its lifecycle and update methods. - -Custom steps ------------- - -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`. - -Custom 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``. - -For compiled JAX execution, custom steps must also obey the functional -interface's fixed-structure requirements. See :doc:`jax_support` for the -functional step contract and initialization hooks. diff --git a/docs/getting_started.rst b/docs/getting_started.rst index e5bfd638..dfc858f2 100644 --- a/docs/getting_started.rst +++ b/docs/getting_started.rst @@ -9,14 +9,16 @@ Install from PyPI:: pip install sequince ``sequince`` targets Python 3.11+. Its core dependencies are installed -automatically; integrations such as emcee, BlackJAX, and HDF5 checkpointing are -available as optional extras. +automatically. The example below uses the included MiniPCN mutation backend. +See :doc:`installation` for optional integrations and +:doc:`user_guide/core_concepts/array_backends` for backend selection and the +``SCIPY_ARRAY_API`` setting for SciPy dispatch. Quick start ----------- Below, we show a short example of how to use ``sequince``. -For more details, see the :doc:`user_guide`, and see +For more details, see the :doc:`user_guide/index`, and see `sequince-examples `_ for more advanced usage. Here, we fit a straight line, ``y = slope * x + intercept``, to synthetic data with @@ -32,9 +34,14 @@ The log-prior and log-likelihood callbacks accept particles shaped Next steps ---------- -See :doc:`evidence_uncertainty` for interpreting the evidence estimate and its -uncertainty. To customise the run, continue with :doc:`smc_configuration` and -the sampling choices in :doc:`resampling_configuration` and :doc:`mutation_configuration`. -The :doc:`step_context` guide explains the distinction between shared inputs -and evolving metadata. For sampling without SMC weighting and resampling, see -:doc:`mcmc_sampler`. +Continue with :doc:`user_guide/running_inference/results` to understand the +returned result, extract posterior samples, and inspect diagnostics. See +:doc:`user_guide/running_inference/evidence_uncertainty` for interpreting the +evidence estimate and its uncertainty. + +To customise the run, see +:doc:`user_guide/running_inference/smc_configuration` and the sampling choices +in :doc:`user_guide/sampling_options/resampling_configuration` and +:doc:`user_guide/sampling_options/mutation_configuration`. +The :doc:`user_guide/core_concepts/step_context` guide explains the distinction +between shared inputs and evolving metadata. diff --git a/docs/index.rst b/docs/index.rst index ea54306f..25de8e37 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -8,11 +8,12 @@ sequential methods. Overview -------- -* :doc:`getting_started` covers installation and a complete straight-line +* :doc:`getting_started` gives a complete straight-line fitting example. -* :doc:`user_guide` explains how to configure sampling, work with particle +* :doc:`installation` covers dependencies and optional integrations. +* :doc:`user_guide/index` explains how to configure sampling, work with particle states and results, and build custom pipelines. After the quickstart, - continue with :doc:`smc_configuration` and :doc:`evidence_uncertainty`. + continue with :doc:`user_guide/core_concepts/index` and :doc:`user_guide/running_inference/index`. * :doc:`examples` links to worked examples in the ``sequince-examples`` repository. * The :doc:`API reference ` lists classes, functions, @@ -23,7 +24,8 @@ Overview :caption: Contents getting_started - user_guide + installation + user_guide/index examples API reference 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/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/jax_support.rst b/docs/user_guide/advanced_workflows/jax_support.rst similarity index 98% rename from docs/jax_support.rst rename to docs/user_guide/advanced_workflows/jax_support.rst index 836a0800..cc7aefc4 100644 --- a/docs/jax_support.rst +++ b/docs/user_guide/advanced_workflows/jax_support.rst @@ -76,7 +76,7 @@ inside the iteration uses JAX-compatible operations. Pass a JAX key as :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:`random_number_generation` for the corresponding eager-execution +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 diff --git a/docs/mcmc_sampler.rst b/docs/user_guide/advanced_workflows/mcmc_sampler.rst similarity index 97% rename from docs/mcmc_sampler.rst rename to docs/user_guide/advanced_workflows/mcmc_sampler.rst index cd188b20..ddb9dea7 100644 --- a/docs/mcmc_sampler.rst +++ b/docs/user_guide/advanced_workflows/mcmc_sampler.rst @@ -79,7 +79,7 @@ helper: ``"minipcn"``, ``"emcee"``, ``"blackjax"``, and }, ) -See :doc:`mutation_configuration` for available backends, optional dependencies, and +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 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/particle_state.rst b/docs/user_guide/core_concepts/particle_state.rst similarity index 77% rename from docs/particle_state.rst rename to docs/user_guide/core_concepts/particle_state.rst index aa330a76..bed3e9d3 100644 --- a/docs/particle_state.rst +++ b/docs/user_guide/core_concepts/particle_state.rst @@ -56,10 +56,15 @@ Use the result helpers to preserve those weights when computing summaries: final_state = result.final_state mean = result.weighted_mean() covariance = result.weighted_covariance() - samples = result.posterior_samples(1_000, rng=42) - -``posterior_samples`` draws an equally weighted sample from the final weighted -population. For diagnostics across the complete run, use ``result.summary()``, + 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:`evidence_uncertainty` for -interpreting the final evidence estimate and its uncertainty. +``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/random_number_generation.rst b/docs/user_guide/core_concepts/random_number_generation.rst similarity index 92% rename from docs/random_number_generation.rst rename to docs/user_guide/core_concepts/random_number_generation.rst index 4d9087da..697a89e2 100644 --- a/docs/random_number_generation.rst +++ b/docs/user_guide/core_concepts/random_number_generation.rst @@ -11,8 +11,7 @@ Eager execution --------------- The ``rng`` argument accepted by -:meth:`~sequince.smc.SequentialMonteCarlo.run` and -:meth:`~sequince.mcmc.MarkovChainMonteCarlo.run` may be: +:meth:`~sequince.smc.SequentialMonteCarlo.run` may be: .. list-table:: :header-rows: 1 @@ -54,8 +53,8 @@ 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. -Randomness in custom steps --------------------------- +RNG in custom steps +------------------- An eager step receives an :class:`orng.RandomGenerator` as its third argument: @@ -90,7 +89,7 @@ Functional execution does not mutate a generator. Pass backend-native state as 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:`jax_support` for compiling the functional iteration +random stream. See :doc:`../advanced_workflows/jax_support` for compiling the functional iteration and splitting a key for random particle initialisation. Checkpointing and resuming @@ -104,7 +103,7 @@ 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:`checkpointing` for the complete resume workflow. +again. See :doc:`../advanced_workflows/checkpointing` for the complete resume workflow. Random posterior draws ---------------------- diff --git a/docs/step_context.rst b/docs/user_guide/core_concepts/step_context.rst similarity index 98% rename from docs/step_context.rst rename to docs/user_guide/core_concepts/step_context.rst index aca0ad93..98692545 100644 --- a/docs/step_context.rst +++ b/docs/user_guide/core_concepts/step_context.rst @@ -25,7 +25,7 @@ Fields 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:`observations` for complete iteration and exhaustion + ``None``. See :doc:`../running_inference/observations` for complete iteration and exhaustion semantics. ``next_observation_index`` 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..579812a6 --- /dev/null +++ b/docs/user_guide/extending_sequince/index.rst @@ -0,0 +1,30 @@ +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 + + custom_steps + custom_callbacks diff --git a/docs/user_guide.rst b/docs/user_guide/index.rst similarity index 55% rename from docs/user_guide.rst rename to docs/user_guide/index.rst index be0eeb3a..c553ecae 100644 --- a/docs/user_guide.rst +++ b/docs/user_guide/index.rst @@ -9,15 +9,16 @@ 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 workflow guides when configuring a run and interpreting its -output. The remaining sections explain sampling choices, the objects passed -through the SMC loop, and more specialised execution and extension points. +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 - user_guide/running_inference - user_guide/sampling_options - user_guide/core_concepts - user_guide/advanced_workflows - user_guide/extending_sequince + 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/evidence_uncertainty.rst b/docs/user_guide/running_inference/evidence_uncertainty.rst similarity index 87% rename from docs/evidence_uncertainty.rst rename to docs/user_guide/running_inference/evidence_uncertainty.rst index 53e4cae7..6a2c1cab 100644 --- a/docs/evidence_uncertainty.rst +++ b/docs/user_guide/running_inference/evidence_uncertainty.rst @@ -1,6 +1,17 @@ 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: 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/observations.rst b/docs/user_guide/running_inference/observations.rst similarity index 97% rename from docs/observations.rst rename to docs/user_guide/running_inference/observations.rst index cd7a4d0d..e36131db 100644 --- a/docs/observations.rst +++ b/docs/user_guide/running_inference/observations.rst @@ -151,7 +151,7 @@ original indexable sequence to ``run``: The runner resumes with ``all_observations[checkpoint.next_observation_index]``. As generators cannot be repositioned, they are not compatible with this procedure. -See :doc:`checkpointing` for the full checkpoint workflow. +See :doc:`../advanced_workflows/checkpointing` for the full checkpoint workflow. Observations with the JAX backend --------------------------------- @@ -181,5 +181,5 @@ For easy compilation, supply a representative first observation to the compilati iteration=iteration, ) -See :doc:`jax_support` for compilation, random-state threading, and fixed-carry +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..e045b736 --- /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 :doc:`../../getting_started`. 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/smc_configuration.rst b/docs/user_guide/running_inference/smc_configuration.rst similarity index 83% rename from docs/smc_configuration.rst rename to docs/user_guide/running_inference/smc_configuration.rst index 4c638cce..9cd543a7 100644 --- a/docs/smc_configuration.rst +++ b/docs/user_guide/running_inference/smc_configuration.rst @@ -22,9 +22,11 @@ Bayesian inference pipeline is the pre-built helper: 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:`resampling_configuration` and :doc:`mutation_configuration` for the available +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 @@ -36,7 +38,7 @@ When starting the run, pass the ``metadata`` returned by the convenience constru metadata=metadata, ) -See :doc:`step_context` for the distinction between static ``shared_context`` +See :doc:`../core_concepts/step_context` for the distinction between static ``shared_context`` data and evolving metadata. Building a pipeline manually @@ -67,4 +69,6 @@ to reproduce the pre-built SMC defined above: 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:`custom_steps_callbacks` when implementing either interface yourself. +: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/mutation_configuration.rst b/docs/user_guide/sampling_options/mutation_configuration.rst similarity index 97% rename from docs/mutation_configuration.rst rename to docs/user_guide/sampling_options/mutation_configuration.rst index b04f433d..f03c5f7c 100644 --- a/docs/mutation_configuration.rst +++ b/docs/user_guide/sampling_options/mutation_configuration.rst @@ -3,7 +3,7 @@ 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:`mcmc_sampler`. +: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:: @@ -84,7 +84,7 @@ when resampling occurs. Install optional backends with ``pip install "sequince[emcee]"``, ``pip install "sequince[eryn]"``, or ``pip install "sequince[blackjax]"``. -See :doc:`jax_support` for compiled execution. +See :doc:`../advanced_workflows/jax_support` for compiled execution. Particle transforms ------------------- @@ -156,5 +156,5 @@ and default execution flags, it runs only after resampling unless 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:`smc_configuration` for custom pipelines +``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/resampling_configuration.rst b/docs/user_guide/sampling_options/resampling_configuration.rst similarity index 94% rename from docs/resampling_configuration.rst rename to docs/user_guide/sampling_options/resampling_configuration.rst index df3dc2d1..392fcae7 100644 --- a/docs/resampling_configuration.rst +++ b/docs/user_guide/sampling_options/resampling_configuration.rst @@ -57,7 +57,7 @@ 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:`jax_support`. +pipelines must also preserve the particle count; see :doc:`../advanced_workflows/jax_support`. Outputs and ancestry -------------------- @@ -70,7 +70,7 @@ 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:`evidence_uncertainty` for the assumptions needed +constant unchanged. See :doc:`../running_inference/evidence_uncertainty` for the assumptions needed by genealogy-based evidence estimators, including multinomial resampling. Prebuilt samplers From 2e0b9ee099f75bca8e8c8597a910a68f7e25befa Mon Sep 17 00:00:00 2001 From: "Michael J. Williams" Date: Thu, 22 Oct 2026 11:44:30 +0200 Subject: [PATCH 18/23] DOC: add section on documentation to contributing guidelines --- CONTRIBUTING.md | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index d2ac1ecf..1874188e 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -157,6 +157,19 @@ 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. + ## Checks and tests Run the following checks before requesting review: From 879704bf4a334afc1c0f7b23a15daced4cd44468 Mon Sep 17 00:00:00 2001 From: "Michael J. Williams" Date: Thu, 22 Oct 2026 11:53:26 +0200 Subject: [PATCH 19/23] DOC: link contributing guidelines and AI policy from docs --- AI_POLICY.md | 3 ++- CONTRIBUTING.md | 6 +++++- docs/conf.py | 1 + docs/developer_guide/index.rst | 22 ++++++++++++++++++++ docs/index.rst | 3 +++ docs/user_guide/extending_sequince/index.rst | 1 + pyproject.toml | 1 + 7 files changed, 35 insertions(+), 2 deletions(-) create mode 100644 docs/developer_guide/index.rst diff --git a/AI_POLICY.md b/AI_POLICY.md index 7b94ca07..ee4771a2 100644 --- a/AI_POLICY.md +++ b/AI_POLICY.md @@ -21,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 1874188e..615153a9 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -8,7 +8,8 @@ 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). @@ -170,6 +171,9 @@ 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/conf.py b/docs/conf.py index d5c45564..5cc91244 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -26,6 +26,7 @@ extensions = [ "autoapi.extension", + "myst_parser", "sphinx.ext.napoleon", "sphinx.ext.viewcode", "sphinx.ext.intersphinx", diff --git a/docs/developer_guide/index.rst b/docs/developer_guide/index.rst new file mode 100644 index 00000000..8e2dc90b --- /dev/null +++ b/docs/developer_guide/index.rst @@ -0,0 +1,22 @@ +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 + + Contributing + AI policy diff --git a/docs/index.rst b/docs/index.rst index 25de8e37..949ea916 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -18,6 +18,8 @@ Overview 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 @@ -28,6 +30,7 @@ Overview user_guide/index examples API reference + developer_guide/index Indices and tables ------------------ diff --git a/docs/user_guide/extending_sequince/index.rst b/docs/user_guide/extending_sequince/index.rst index 579812a6..86d83e10 100644 --- a/docs/user_guide/extending_sequince/index.rst +++ b/docs/user_guide/extending_sequince/index.rst @@ -25,6 +25,7 @@ into a pipeline. .. toctree:: :maxdepth: 1 + :caption: Contents custom_steps custom_callbacks diff --git a/pyproject.toml b/pyproject.toml index 1cea954f..5e5b40e4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -165,6 +165,7 @@ typecheck = "mypy" sphinx = ">=7.2" sphinx-book-theme = ">=1.1" sphinx-autoapi = ">=3.8,<4" +myst-parser = ">=4,<5" [tool.pixi.feature.doc.tasks] docs = "sphinx-build -b html docs docs/_build/html" From 673fa336eb219c7fed7f390d0a59dff976a68db6 Mon Sep 17 00:00:00 2001 From: "Michael J. Williams" Date: Thu, 22 Oct 2026 12:22:41 +0200 Subject: [PATCH 20/23] DOC: change copyright and authors --- docs/conf.py | 4 ++-- docs/developer_guide/index.rst | 1 + 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/docs/conf.py b/docs/conf.py index 5cc91244..5fceeb6f 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -7,8 +7,8 @@ # -- 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 diff --git a/docs/developer_guide/index.rst b/docs/developer_guide/index.rst index 8e2dc90b..8b4f6a42 100644 --- a/docs/developer_guide/index.rst +++ b/docs/developer_guide/index.rst @@ -17,6 +17,7 @@ For implementing new sampling components, see .. toctree:: :maxdepth: 1 + :caption: Contents Contributing AI policy From e04ae8885d663e4463b0f1aedcd799084e11acec Mon Sep 17 00:00:00 2001 From: "Michael J. Williams" Date: Thu, 22 Oct 2026 12:25:54 +0200 Subject: [PATCH 21/23] DOC: fix developer guide --- CONTRIBUTING.md | 3 ++- docs/developer_guide/AI_POLICY.rst | 2 ++ docs/developer_guide/CONTRIBUTING.rst | 2 ++ 3 files changed, 6 insertions(+), 1 deletion(-) create mode 100644 docs/developer_guide/AI_POLICY.rst create mode 100644 docs/developer_guide/CONTRIBUTING.rst diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 615153a9..7ea0ac26 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -12,7 +12,8 @@ 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. 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_ From 27ecedf984115fbddf022ee1738afdde9fca345d Mon Sep 17 00:00:00 2001 From: "Michael J. Williams" Date: Thu, 22 Oct 2026 12:26:12 +0200 Subject: [PATCH 22/23] MAINT: fix invalid imports --- src/sequince/result.py | 3 ++- src/sequince/steps/constraints.py | 3 ++- src/sequince/steps/control.py | 3 ++- src/sequince/steps/mutation/metrics/de.py | 4 +++- src/sequince/steps/mutation/metrics/rjd.py | 3 ++- src/sequince/steps/mutation/metropolis.py | 3 ++- 6 files changed, 13 insertions(+), 6 deletions(-) diff --git a/src/sequince/result.py b/src/sequince/result.py index 14ede7e4..679f058a 100644 --- a/src/sequince/result.py +++ b/src/sequince/result.py @@ -8,10 +8,11 @@ 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 ._typing import Array -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 diff --git a/src/sequince/steps/constraints.py b/src/sequince/steps/constraints.py index 93ff0183..78d7ce5e 100644 --- a/src/sequince/steps/constraints.py +++ b/src/sequince/steps/constraints.py @@ -8,10 +8,11 @@ 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 .._typing import Array -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 222b6aa2..a738d8d8 100644 --- a/src/sequince/steps/control.py +++ b/src/sequince/steps/control.py @@ -6,9 +6,10 @@ from typing import TYPE_CHECKING, Any, Callable, Sequence from acai import cond +from array_api_compat import is_jax_namespace from .._typing import Array -from ..array_api import is_jax_namespace, namespace +from ..array_api import namespace from .base import SMCStep, StepContext if TYPE_CHECKING: diff --git a/src/sequince/steps/mutation/metrics/de.py b/src/sequince/steps/mutation/metrics/de.py index d1be2681..0b55824d 100644 --- a/src/sequince/steps/mutation/metrics/de.py +++ b/src/sequince/steps/mutation/metrics/de.py @@ -5,8 +5,10 @@ from dataclasses import dataclass from typing import ClassVar +from array_api_compat import device as get_device + from ...._typing import Array -from ....array_api import asarray, get_device, namespace +from ....array_api import asarray, namespace from ...context import StepContext from .esjd import _squared_jump_distances from .state import MutationTransition diff --git a/src/sequince/steps/mutation/metrics/rjd.py b/src/sequince/steps/mutation/metrics/rjd.py index add5d323..46f7be77 100644 --- a/src/sequince/steps/mutation/metrics/rjd.py +++ b/src/sequince/steps/mutation/metrics/rjd.py @@ -6,10 +6,11 @@ from dataclasses import dataclass from typing import Literal +from array_api_compat import device as get_device from orng import RandomGenerator from ...._typing import Array -from ....array_api import asarray, get_device, namespace, to_numpy +from ....array_api import asarray, namespace, to_numpy from ....utils.pytrees import pytree_dataclass from ...base import validate_n_steps from ...context import StepContext diff --git a/src/sequince/steps/mutation/metropolis.py b/src/sequince/steps/mutation/metropolis.py index 6ea63383..905108db 100644 --- a/src/sequince/steps/mutation/metropolis.py +++ b/src/sequince/steps/mutation/metropolis.py @@ -7,9 +7,10 @@ import array_api_extra as xpx import optree +from array_api_compat import device as get_device from ..._typing import Array -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 ...utils.pytrees import OPTREE_NAMESPACE From 588860bbf2ca623af35b935420374a5373f53b55 Mon Sep 17 00:00:00 2001 From: Christian Chapman-Bird Date: Thu, 22 Oct 2026 17:30:24 +0100 Subject: [PATCH 23/23] add citation placeholder, merge getting started with index --- docs/citation.rst | 8 +++ docs/getting_started.rst | 47 --------------- docs/index.rst | 57 +++++++++++++++---- .../running_inference/prebuilt_samplers.rst | 2 +- 4 files changed, 55 insertions(+), 59 deletions(-) create mode 100644 docs/citation.rst delete mode 100644 docs/getting_started.rst 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/getting_started.rst b/docs/getting_started.rst deleted file mode 100644 index dfc858f2..00000000 --- a/docs/getting_started.rst +++ /dev/null @@ -1,47 +0,0 @@ -Getting started -=============== - -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 optional integrations and -:doc:`user_guide/core_concepts/array_backends` for backend selection and the -``SCIPY_ARRAY_API`` setting for SciPy dispatch. - -Quick start ------------ - -Below, we show a short example of how to use ``sequince``. -For more details, see the :doc:`user_guide/index`, and see -`sequince-examples `_ -for more advanced usage. -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 ----------- - -Continue with :doc:`user_guide/running_inference/results` to understand the -returned result, extract posterior samples, and inspect diagnostics. See -:doc:`user_guide/running_inference/evidence_uncertainty` for interpreting the -evidence estimate and its uncertainty. - -To customise the run, see -:doc:`user_guide/running_inference/smc_configuration` and the sampling choices -in :doc:`user_guide/sampling_options/resampling_configuration` and -:doc:`user_guide/sampling_options/mutation_configuration`. -The :doc:`user_guide/core_concepts/step_context` guide explains the distinction -between shared inputs and evolving metadata. diff --git a/docs/index.rst b/docs/index.rst index 949ea916..1e579966 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -5,16 +5,50 @@ sequince documentation Carlo (SMC) algorithms for Bayesian inference, particle filtering, and related sequential methods. -Overview --------- - -* :doc:`getting_started` gives a complete straight-line - fitting example. -* :doc:`installation` covers dependencies and optional integrations. -* :doc:`user_guide/index` explains how to configure sampling, work with particle - states and results, and build custom pipelines. After the quickstart, - continue with :doc:`user_guide/core_concepts/index` and :doc:`user_guide/running_inference/index`. -* :doc:`examples` links to worked examples in the ``sequince-examples`` +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. @@ -24,13 +58,14 @@ Overview .. toctree:: :maxdepth: 2 :caption: Contents + :hidden: - getting_started installation user_guide/index examples API reference developer_guide/index + citation Indices and tables ------------------ diff --git a/docs/user_guide/running_inference/prebuilt_samplers.rst b/docs/user_guide/running_inference/prebuilt_samplers.rst index e045b736..71935e4b 100644 --- a/docs/user_guide/running_inference/prebuilt_samplers.rst +++ b/docs/user_guide/running_inference/prebuilt_samplers.rst @@ -15,7 +15,7 @@ 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 :doc:`../../getting_started`. See +This is the helper used in the :ref:`quick start `. See :doc:`smc_configuration` for configuration and manual pipeline assembly. Persistent Bayesian SMC