ENH(mutation): expose acceptance-based adaptive budgets in factories

docs/user_guide/sampling_options/mutation_configuration.rst

@@ -166,3 +166,33 @@ You can also pass a configured step as ``mutation``. Set its target, transform,
 and options directly: the helper does not apply ``mutation_n_steps`` or
 ``mutation_kwargs`` to it. See :doc:`../running_inference/smc_configuration` for custom pipelines
 and the API links above for method-specific options and outputs.
+
+
+Acceptance-based mutation budgets
+---------------------------------
+
+For standard and persistent SMC, set ``adaptive_mutation`` to adjust the number
+of transitions from the previous iteration's acceptance rate. Named
+``"minipcn"`` and ``"differential-evolution"`` mutations support this option:
+
+.. code-block:: python
+
+    mutation_n_steps=100,
+    adaptive_mutation={
+        "min_steps": 50,
+        "max_steps": 1000,
+        "target_accepted_moves": 5.0,
+    }
+
+``mutation_n_steps`` sets the initial budget and must lie within the bounds.
+The mapping defaults to ``min_steps=0``, ``max_steps=1000``, and
+``target_accepted_moves=5.0``. Omit it or set ``enabled=False`` to keep a fixed
+budget.
+
+Each mutation block uses a fixed budget; its acceptance rate determines the
+next block's budget, clipped to the bounds. Accepted moves need not be
+independent samples. Resume checkpoints with the same mutation configuration
+and bounds. See :func:`~sequince.steps.mutation.factory.get_mutation_step` for
+option details, or use
+:class:`~sequince.steps.mutation.adaptive.steps.AdaptMutationSteps` in a custom
+pipeline.

src/sequince/prebuilt_samplers.py

@@ -77,6 +77,7 @@ def make_standard_bayesian_smc(
     mutation: MutationStrategy | SMCStep = "differential-evolution",
     mutation_n_steps: int = 5,
     mutation_kwargs: Dict[str, Any] | None = None,
+    adaptive_mutation: Mapping[str, object] | None = None,
     always_mutate: bool | None = None,
     particle_transform: ParticleTransform | str | None = "affine",
     particle_transform_kwargs: Dict[str, Any] | None = None,
@@ -152,6 +153,9 @@ def make_standard_bayesian_smc(
         weighted move configuration, or the ``"bilby-de"`` preset.
         For ``"blackjax"``, use ``algorithm`` and ``algorithm_kwargs`` to
         configure the JAX mutation kernel.
+    adaptive_mutation : Mapping[str, object] or None, optional
+        Acceptance-based budget configuration; see the user guide. Defaults
+        to disabled. Uses ``mutation_n_steps`` as the initial total budget.
     always_mutate : bool or None, optional
         Force mutation at every iteration regardless of resampling. If set to
         ``None``, defaults to ``True`` if ``always_resample=True`` and ``False``
@@ -276,6 +280,7 @@ def make_standard_bayesian_smc(
         target=posterior,
         transform=particle_transform,
         mutation_n_steps=mutation_n_steps,
+        adaptive_mutation=adaptive_mutation,
         mutation_kwargs=mutation_kwargs,
     )
 
@@ -331,6 +336,7 @@ def make_persistent_bayesian_smc(
     mutation: MutationStrategy | SMCStep = "differential-evolution",
     mutation_n_steps: int = 5,
     mutation_kwargs: Dict[str, Any] | None = None,
+    adaptive_mutation: Mapping[str, object] | None = None,
     particle_transform: ParticleTransform | str | None = "affine",
     particle_transform_kwargs: Dict[str, Any] | None = None,
     auto_normalize: bool = True,
@@ -380,6 +386,10 @@ def make_persistent_bayesian_smc(
     ``logging_frequency`` iterations and once more for the terminal
     state. Set ``logging_frequency=0`` to disable it.
 
+    ``adaptive_mutation`` optionally configures acceptance-based budgets for
+    named miniPCN or differential-evolution mutation, starting at
+    ``mutation_n_steps``. See the user guide for options and lagged behavior.
+
     For more details see https://arxiv.org/abs/2407.20722.
     """
     if isinstance(n_particles, bool) or not isinstance(n_particles, Integral):
@@ -476,6 +486,7 @@ def make_persistent_bayesian_smc(
         target=posterior,
         transform=particle_transform,
         mutation_n_steps=mutation_n_steps,
+        adaptive_mutation=adaptive_mutation,
         mutation_kwargs=mutation_kwargs,
     )
 

src/sequince/steps/mutation/conditional.py

@@ -1,7 +1,8 @@
 from __future__ import annotations
 
 import logging
-from typing import TYPE_CHECKING
+from collections.abc import Mapping
+from typing import TYPE_CHECKING, Any
 
 from ..base import FixedLengthMutation, SMCStep, validate_n_steps
 from ..context import StepContext
@@ -76,3 +77,18 @@ def apply_with_n_steps(
             state, context, rng, n_steps=n_steps
         )
         return updated, context, ConditionalStepOutput(True, output)
+
+    def state_dict(self) -> dict[str, Any]:
+        """Preserve the child mutation's adaptive and proposal state."""
+        return {"mutation": self.mutation.state_dict()}
+
+    def load_state_dict(self, state: Mapping[str, Any]) -> None:
+        """Restore nested state, accepting older empty wrapper checkpoints."""
+        if "mutation" not in state:
+            return
+        nested = state["mutation"]
+        if not isinstance(nested, Mapping):
+            raise TypeError(
+                "MutateIfResampled state 'mutation' must be a mapping."
+            )
+        self.mutation.load_state_dict(nested)

src/sequince/steps/mutation/factory.py

@@ -1,11 +1,15 @@
 from __future__ import annotations
 
-from collections.abc import Callable
+import math
+from collections.abc import Callable, Mapping
+from numbers import Integral, Real
 from typing import Any
 
 from ...target import Target
 from ...transforms import ParticleTransform
 from ..base import SMCStep
+from .adaptive.policies.acceptance import AcceptanceRatePolicy
+from .adaptive.steps import AdaptMutationSteps
 from .blackjax import BlackJAXMutation
 from .emcee import EmceeMutation
 from .eryn import ErynMutation
@@ -18,14 +22,100 @@
 from .minipcn import MiniPCNMutation
 
 
+def _acceptance_policy(
+    config: Mapping[str, object] | None, initial_steps: int
+) -> AcceptanceRatePolicy | None:
+    """Validate configuration even when adaptation is disabled."""
+    if config is None:
+        return None
+    if not isinstance(config, Mapping):
+        raise TypeError("adaptive_mutation must be a mapping or None.")
+    unknown = config.keys() - {
+        "enabled",
+        "min_steps",
+        "max_steps",
+        "target_accepted_moves",
+    }
+    if unknown:
+        raise ValueError(f"Unknown adaptive_mutation keys: {list(unknown)!r}")
+    enabled = config.get("enabled", True)
+    if not isinstance(enabled, bool):
+        raise TypeError("adaptive_mutation enabled must be a boolean.")
+    bounds = []
+    for name, default in (("min_steps", 0), ("max_steps", 1000)):
+        value = config.get(name, default)
+        if isinstance(value, bool) or not isinstance(value, Integral):
+            raise TypeError(f"adaptive_mutation {name} must be an integer.")
+        if value < 0:
+            raise ValueError(f"adaptive_mutation {name} must be >= 0.")
+        bounds.append(int(value))
+    minimum, maximum = bounds
+    if maximum < minimum:
+        raise ValueError("adaptive_mutation max_steps must be >= min_steps.")
+    target = config.get("target_accepted_moves", 5.0)
+    if isinstance(target, bool) or not isinstance(target, Real):
+        raise TypeError("target_accepted_moves must be a real number.")
+    if not math.isfinite(target) or target <= 0:
+        raise ValueError("target_accepted_moves must be finite and positive.")
+    if not enabled:
+        return None
+    if isinstance(initial_steps, bool) or not isinstance(
+        initial_steps, Integral
+    ):
+        raise TypeError("mutation_n_steps must be an integer.")
+    if not minimum <= initial_steps <= maximum:
+        raise ValueError(
+            "mutation_n_steps must be within [min_steps, max_steps]."
+        )
+    return AcceptanceRatePolicy(
+        float(target), min_steps=minimum, max_steps=maximum
+    )
+
+
 def get_mutation_step(
     mutation: str | SMCStep | None,
     target: "Target | Callable[..., Any]",
     transform: ParticleTransform | None = None,
     mutation_n_steps: int = 10,
     mutation_kwargs: dict[str, Any] | None = None,
+    adaptive_mutation: Mapping[str, object] | None = None,
 ) -> SMCStep | None:
-    """Construct a configured mutation step."""
+    """Construct a configured mutation step.
+
+    Parameters
+    ----------
+    mutation : str, SMCStep or None
+        Named mutation, custom step, or disabled mutation.
+    target : Target or callable
+        Target passed to the mutation constructor.
+    transform : ParticleTransform or None, optional
+        Transform passed to the mutation constructor.
+    mutation_n_steps : int, optional
+        Fixed total budget, or initial budget when adaptation is enabled.
+    mutation_kwargs : dict or None, optional
+        Additional mutation constructor options.
+    adaptive_mutation : Mapping[str, object] or None, optional
+        Named miniPCN and differential-evolution mutations support ``enabled``
+        (default True), ``min_steps`` (0), ``max_steps`` (1000), and
+        ``target_accepted_moves`` (5.0). Omission disables adaptation.
+        Completed acceptance metrics select the next iteration's bounded
+        budget; the current block is frozen before execution. Accepted moves
+        need not be independent samples. Proposal parameters are unchanged.
+
+    Returns
+    -------
+    SMCStep or None
+        Configured mutation, optionally wrapped in ``AdaptMutationSteps``.
+    """
+    policy = _acceptance_policy(adaptive_mutation, mutation_n_steps)
+    if policy is not None and (
+        not isinstance(mutation, str)
+        or mutation not in {"minipcn", "differential-evolution"}
+    ):
+        raise ValueError(
+            "adaptive_mutation requires named minipcn or "
+            "differential-evolution mutation."
+        )
     kwargs = dict(mutation_kwargs or {})
     if transform is not None:
         kwargs["transform"] = transform
@@ -117,4 +207,15 @@ def get_mutation_step(
             "'red-blue-metropolis-hastings', 'red-blue-mh', 'red-blue', "
             "'none', or an SMCStep instance."
         )
+    if policy is not None:
+        assert isinstance(
+            mutation_step, (MiniPCNMutation, DifferentialEvolutionMutation)
+        )
+        return AdaptMutationSteps(
+            mutation_step,
+            policy,
+            initial_steps=mutation_n_steps,
+            min_steps=policy.min_steps,
+            max_steps=policy.max_steps,
+        )
     return mutation_step

tests/integration_tests/test_persistent_sampling.py

@@ -161,9 +161,11 @@ def test_runner_accumulates_posterior_batches_until_target_total_ess(
 
 
 @pytest.mark.parametrize("resample_at_end", [False, True])
+@pytest.mark.parametrize("adaptive_mutation", [None, {"max_steps": 8}])
 def test_persistent_checkpoint_resume_matches_uninterrupted_run(
     tmp_path,
     resample_at_end: bool,
+    adaptive_mutation,
 ) -> None:
     n_particles = 32
     target_ess = 2 * n_particles
@@ -198,6 +200,7 @@ def make_sampler(
             target_total_ess=target_total_ess,
             resample_at_end=resample_at_end,
             mutation_n_steps=2,
+            adaptive_mutation=adaptive_mutation,
             checkpoint_writers=checkpoint_writers,
             callbacks=callbacks,
         )

tests/integration_tests/test_standard_sampling.py

@@ -12,7 +12,9 @@
 from sequince.steps.base import SMCStep
 
 
+@pytest.mark.parametrize("adaptive_mutation", [None, {"max_steps": 8}])
 def test_standard_sampler_reaches_gaussian_posterior(
+    adaptive_mutation,
     gaussian_log_prior: Callable[..., np.ndarray],
     gaussian_log_likelihood: Callable[..., np.ndarray],
 ) -> None:
@@ -21,6 +23,7 @@ def test_standard_sampler_reaches_gaussian_posterior(
         gaussian_log_prior,
         gaussian_log_likelihood,
         mutation_n_steps=3,
+        adaptive_mutation=adaptive_mutation,
         logging_frequency=0,
     )
     initial_particles = np.random.default_rng(1234).normal(

tests/test_mutation/test_factory.py

@@ -227,3 +227,75 @@ def test_factory_does_not_expose_proposal_rejuvenation(
             "proposal",
             standard_normal_posterior,
         )
+
+
+@pytest.mark.parametrize("name", ["minipcn", "differential-evolution"])
+def test_factory_adapts_completed_budget(
+    name, standard_normal_posterior, particle_state, context_factory, rng
+) -> None:
+    from sequince.steps.mutation.adaptive.steps import AdaptMutationSteps
+
+    step = get_mutation_step(
+        name,
+        standard_normal_posterior,
+        mutation_n_steps=2,
+        adaptive_mutation={
+            "min_steps": 1,
+            "max_steps": 4,
+            "target_accepted_moves": 100.0,
+        },
+    )
+    assert isinstance(step, AdaptMutationSteps)
+    state, _, output = step(particle_state, context_factory(t=0), rng)
+    assert output["metrics"].n_steps == 2
+    assert step.next_n_steps == 4
+    _, _, output = step(state, context_factory(t=1), rng)
+    assert output["metrics"].n_steps == 4
+
+
+@pytest.mark.parametrize(
+    "config,error,message",
+    [
+        (True, TypeError, "mapping"),
+        ({"rho": 0.5}, ValueError, "Unknown"),
+        ({"enabled": 1}, TypeError, "boolean"),
+        ({"min_steps": True}, TypeError, "integer"),
+        ({"max_steps": 2.5}, TypeError, "integer"),
+        ({"min_steps": -1}, ValueError, ">= 0"),
+        ({"min_steps": 5, "max_steps": 2}, ValueError, ">= min_steps"),
+        ({"target_accepted_moves": "5"}, TypeError, "real number"),
+        ({"target_accepted_moves": False}, TypeError, "real number"),
+        ({"target_accepted_moves": float("nan")}, ValueError, "finite"),
+        ({"target_accepted_moves": 0}, ValueError, "positive"),
+        ({"min_steps": 20}, ValueError, "mutation_n_steps"),
+        ({"enabled": False, "max_steps": -1}, ValueError, ">= 0"),
+    ],
+)
+def test_factory_invalid_adaptation(
+    config, error, message, standard_normal_posterior
+) -> None:
+    with pytest.raises(error, match=message):
+        get_mutation_step(
+            "minipcn", standard_normal_posterior, adaptive_mutation=config
+        )
+
+
+@pytest.mark.parametrize(
+    "mutation", [None, "none", "emcee", "red-blue", IdentityMutation()]
+)
+def test_factory_unsupported_adaptation(mutation, standard_normal_posterior):
+    with pytest.raises(ValueError, match="requires named"):
+        get_mutation_step(
+            mutation,
+            standard_normal_posterior,
+            adaptive_mutation={"enabled": True},
+        )
+
+
+@pytest.mark.parametrize("name", ["minipcn", "differential-evolution", "none"])
+def test_factory_disabled_adaptation(name, standard_normal_posterior) -> None:
+    plain = get_mutation_step(name, standard_normal_posterior)
+    disabled = get_mutation_step(
+        name, standard_normal_posterior, adaptive_mutation={"enabled": False}
+    )
+    assert type(disabled) is type(plain)