diff --git a/CHANGELOG.md b/CHANGELOG.md index ba9ed33..49e84bd 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,12 @@ The format follows [Keep a Changelog](https://keepachangelog.com/en/1.1.0/). ## [Unreleased] +### Added + +- Allow adaptive SMC to set `beta_tolerance` independently of the minimum + beta increment, including through `make_standard_bayesian_smc`'s + `weighting_kwargs`. The default continues to use `min_delta` as the tolerance. + ## [0.1.0] - 2026-10-29 First stable release. Changes below are relative to 0.1.0b1. diff --git a/src/sequince/prebuilt_samplers.py b/src/sequince/prebuilt_samplers.py index 0946790..7bb4506 100644 --- a/src/sequince/prebuilt_samplers.py +++ b/src/sequince/prebuilt_samplers.py @@ -122,6 +122,9 @@ def make_standard_bayesian_smc( Criterion for adaptive tempering; one of ``"ess"`` (default) or ``"cess"``. weighting_kwargs : dict[str, Any] or None, optional Extra keyword arguments forwarded to the weighting step constructor. + For adaptive tempering, ``beta_tolerance`` sets the beta-search + precision independently of ``min_delta``; ``max_bisection_iter`` + limits the number of search iterations. resampling : {"systematic", "stratified", "multinomial"}, optional Resampling strategy to use; one of ``"systematic"``, ``"stratified"``, or ``"multinomial"``. diff --git a/src/sequince/steps/weighting.py b/src/sequince/steps/weighting.py index a6c31ab..6d71ace 100644 --- a/src/sequince/steps/weighting.py +++ b/src/sequince/steps/weighting.py @@ -783,6 +783,10 @@ class AdaptiveTemperedWeighting(BaseTemperedStep): max_bisection_iter: The maximum number of bisection iterations to perform when finding the beta increment. + beta_tolerance : float or None, optional + Stop bisection when the beta bracket is no wider than this finite, + positive value. Defaults to ``min_delta`` when None, preserving the + existing search behavior. Does not change the minimum beta increment. criterion: Criterion used when selecting the beta increment. ``"ess"`` targets the ESS ratio after weighting. ``"cess"`` targets the @@ -805,6 +809,7 @@ def __init__( max_delta: float | None = None, normalize_after: bool = True, max_bisection_iter: int = 30, + beta_tolerance: float | None = None, criterion: Literal["ess", "cess"] = "ess", refresh_components: Collection[str] = (), ): @@ -816,6 +821,10 @@ def __init__( raise ValueError("max_beta must be positive.") if min_delta < 0.0: raise ValueError("min_delta must be non-negative.") + if beta_tolerance is not None and ( + not math.isfinite(beta_tolerance) or beta_tolerance <= 0.0 + ): + raise ValueError("beta_tolerance must be finite and positive.") if max_delta is not None and max_delta <= 0.0: raise ValueError("max_delta must be positive.") if max_delta is not None and max_delta < min_delta: @@ -834,6 +843,9 @@ def __init__( self._target_ess_ratio = target_ess_ratio self._max_beta = max_beta self._min_delta = min_delta + self._beta_tolerance = ( + min_delta if beta_tolerance is None else beta_tolerance + ) self._max_delta = max_delta self._max_bisection_iter = max_bisection_iter self._criterion = criterion @@ -1126,7 +1138,9 @@ def _find_delta( low, high, target_ess_ratio, - tolerance=min_delta, + tolerance=asarray( + self._beta_tolerance, xp=xp, like=base_log_weights + ), max_iterations=self._max_bisection_iter, ) diff --git a/tests/test_functional_smc.py b/tests/test_functional_smc.py index f316f77..a01d4ef 100644 --- a/tests/test_functional_smc.py +++ b/tests/test_functional_smc.py @@ -788,6 +788,7 @@ def test_adaptive_tempering_functional_step_can_be_jitted( target_ess_ratio=0.8, max_beta=0.6, min_delta=1e-6, + beta_tolerance=1e-8, criterion=criterion, ) state = ParticleState(jnp.zeros((4, 1))) diff --git a/tests/test_prebuilt_samplers.py b/tests/test_prebuilt_samplers.py index 994b20e..8762655 100644 --- a/tests/test_prebuilt_samplers.py +++ b/tests/test_prebuilt_samplers.py @@ -57,6 +57,29 @@ def test_make_standard_bayesian_smc_builds_default_steps() -> None: assert isinstance(steps[3], LoggingStep) +def test_standard_smc_accepts_independent_beta_tolerance( + rng: RandomGenerator, +) -> None: + smc, _, metadata = make_standard_bayesian_smc( + lambda x: np.zeros(x.shape[0]), + lambda x: x[:, 0] * np.log(4.0), + target_ess_ratio=49 / 74, + min_delta=0.1, + weighting_kwargs={"beta_tolerance": 1e-12, "max_bisection_iter": 60}, + mutation="none", + ) + state = ParticleState( + np.array([[0.0], [1.0]]), log_weights=np.log([0.25, 0.75]) + ) + result = smc.run(state, num_steps=1, rng=rng, metadata=metadata) + assert float(result.context.metadata["beta"]) == pytest.approx( + 0.5, abs=1e-11 + ) + assert result.final_state.log_evidence == pytest.approx( + np.log(7 / 4), abs=1e-11 + ) + + def test_make_standard_bayesian_smc_reaches_max_beta(rng) -> None: max_beta = 0.5 smc, _, metadata = make_standard_bayesian_smc( diff --git a/tests/test_steps/test_steps.py b/tests/test_steps/test_steps.py index 926d0ed..df3cbd2 100644 --- a/tests/test_steps/test_steps.py +++ b/tests/test_steps/test_steps.py @@ -564,7 +564,8 @@ def test_adaptive_tempering_solves_known_two_particle_target( lambda x: np.array([0.0, np.log(4.0)]), criterion=criterion, target_ess_ratio=target_ratio, - min_delta=1e-10, + min_delta=0.1, + beta_tolerance=1e-10, max_bisection_iter=60, ) @@ -578,6 +579,41 @@ def test_adaptive_tempering_solves_known_two_particle_target( assert not context.stop_requested +@pytest.mark.parametrize("tolerance", [0.0, -1.0, np.nan, np.inf, -np.inf]) +def test_adaptive_tempering_rejects_invalid_beta_tolerance( + tolerance: float, +) -> None: + with pytest.raises( + ValueError, match="beta_tolerance.*finite and positive" + ): + AdaptiveTemperedWeighting(lambda x: x[:, 0], beta_tolerance=tolerance) + + +@pytest.mark.parametrize("min_delta", [0.1, 0.7]) +def test_beta_tolerance_preserves_minimum_increment_and_default( + min_delta: float, rng: RandomGenerator +) -> None: + state = ParticleState( + np.array([[0.0], [1.0]]), log_weights=np.log([0.25, 0.75]) + ) + + def run(tolerance: float | None) -> float: + step = AdaptiveTemperedWeighting( + lambda x: np.array([0.0, np.log(4.0)]), + target_ess_ratio=49 / 74, + min_delta=min_delta, + beta_tolerance=tolerance, + max_bisection_iter=60, + ) + _, _, output = step(state, StepContext(t=0), rng) + return float(output["delta_beta"]) + + assert run(None) == run(min_delta) + assert run(1e-12) == pytest.approx(max(0.5, min_delta), abs=1e-11) + if min_delta == 0.1: + assert abs(run(None) - 0.5) == pytest.approx(0.03125) + + def test_adaptive_endpoint_skips_bisection(monkeypatch) -> None: step = AdaptiveTemperedWeighting(lambda _: np.zeros(4), min_delta=1e-8) calls = 0