diff --git a/docs/user_guide/sampling_options/mutation_configuration.rst b/docs/user_guide/sampling_options/mutation_configuration.rst index 38f295e9..1c35ba69 100644 --- a/docs/user_guide/sampling_options/mutation_configuration.rst +++ b/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. diff --git a/src/sequince/prebuilt_samplers.py b/src/sequince/prebuilt_samplers.py index 09467900..9f6bbe6f 100644 --- a/src/sequince/prebuilt_samplers.py +++ b/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, ) diff --git a/src/sequince/steps/mutation/conditional.py b/src/sequince/steps/mutation/conditional.py index 1e788d7d..35f250a7 100644 --- a/src/sequince/steps/mutation/conditional.py +++ b/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) diff --git a/src/sequince/steps/mutation/factory.py b/src/sequince/steps/mutation/factory.py index 4e8d21fb..065286f6 100644 --- a/src/sequince/steps/mutation/factory.py +++ b/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 diff --git a/tests/integration_tests/test_persistent_sampling.py b/tests/integration_tests/test_persistent_sampling.py index efc7f392..119ded95 100644 --- a/tests/integration_tests/test_persistent_sampling.py +++ b/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, ) diff --git a/tests/integration_tests/test_standard_sampling.py b/tests/integration_tests/test_standard_sampling.py index 42434a55..704dc825 100644 --- a/tests/integration_tests/test_standard_sampling.py +++ b/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( diff --git a/tests/test_mutation/test_factory.py b/tests/test_mutation/test_factory.py index 25e25d23..57782a23 100644 --- a/tests/test_mutation/test_factory.py +++ b/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) diff --git a/tests/test_prebuilt_samplers.py b/tests/test_prebuilt_samplers.py index 994b20ef..4d2fd1a0 100644 --- a/tests/test_prebuilt_samplers.py +++ b/tests/test_prebuilt_samplers.py @@ -650,3 +650,64 @@ def test_prebuilt_mutation_default_and_explicit_minipcn( DifferentialEvolutionMutation if mutation is None else MiniPCNMutation ) assert any(isinstance(step, expected) for step in sampler.config.steps) + + +@pytest.mark.parametrize( + "factory", [make_standard_bayesian_smc, make_persistent_bayesian_smc] +) +@pytest.mark.parametrize("name", ["minipcn", "differential-evolution"]) +@pytest.mark.parametrize("enabled", [True, False]) +def test_prebuilt_acceptance_adaptation(factory, name, enabled) -> None: + from sequince.steps.mutation.adaptive.steps import AdaptMutationSteps + + smc, _, _ = factory( + _quadratic_log_prior, + _zero_log_likelihood, + mutation=name, + mutation_n_steps=2, + adaptive_mutation={"enabled": enabled, "max_steps": 4}, + logging_frequency=0, + ) + mutation = smc._steps[-1] # noqa: SLF001 + assert isinstance(mutation, AdaptMutationSteps) == enabled + + +@pytest.mark.parametrize("mutation", ["minipcn", "differential-evolution"]) +def test_standard_adaptive_checkpoint_resume(tmp_path, mutation) -> None: + from sequince.checkpoint import PickleCheckpointWriter, SMCCheckpoint + + def make_sampler(writers=()): + return make_standard_bayesian_smc( + _quadratic_log_prior, + _zero_log_likelihood, + max_delta=0.25, + mutation=mutation, + mutation_n_steps=2, + weighting_criterion="cess", + always_mutate=False, + resampling_ess_threshold=1.1, + adaptive_mutation={"min_steps": 1, "max_steps": 8}, + checkpoint_writers=writers, + logging_frequency=0, + ) + + initial = ParticleState(np.linspace(-2, 2, 16)[:, None]) + reference_sampler, _, metadata = make_sampler() + reference = reference_sampler.run( + initial.copy(), rng=123, metadata=metadata + ) + path = tmp_path / "adaptive.pkl" + sampler, _, metadata = make_sampler( + [PickleCheckpointWriter(path, every=1)] + ) + sampler.run(initial.copy(), num_steps=2, rng=123, metadata=metadata) + restored, _, _ = make_sampler() + result = restored.run( + runtime=restored.resume_from(SMCCheckpoint.load_pickle(path)) + ) + np.testing.assert_array_equal( + result.final_state.particles, reference.final_state.particles + ) + np.testing.assert_array_equal( + result.final_state.log_weights, reference.final_state.log_weights + )