diff --git a/src/sequince/prebuilt_samplers.py b/src/sequince/prebuilt_samplers.py index c7f09ded..482a0fc7 100644 --- a/src/sequince/prebuilt_samplers.py +++ b/src/sequince/prebuilt_samplers.py @@ -365,12 +365,14 @@ def make_persistent_bayesian_smc( A one-batch initial state preserves the existing initialization behavior. The first persistence step prepares evaluations and archives the state at beta - zero. Once the terminal archive ESS condition is met, - ``resample_at_end=True`` (default) finishes resampling and the configured - mutation, returning an equally weighted active population of ``n_particles``. + zero. With ``resample_at_end=True`` (default), selecting beta one always + completes resampling and the configured mutation. The next iteration archives + that batch and reweights the complete archive before checking termination. + If ``target_total_ess`` is set, beta-one batches accumulate until the returned + weighted archive meets that threshold; no extra batch is generated afterward. + The final archive weighting also updates the evidence estimate. With ``resample_at_end=False``, stop immediately after weighting and return - the complete weighted archive. The final active population is not appended - to the archive, and its mutation does not change the evidence estimate. + the complete weighted archive without generating a final active batch. Progress is logged every ``logging_frequency`` iterations and once more for the terminal state. Set ``logging_frequency=0`` to disable it. diff --git a/src/sequince/steps/persistent.py b/src/sequince/steps/persistent.py index ae3d7046..c33dc00f 100644 --- a/src/sequince/steps/persistent.py +++ b/src/sequince/steps/persistent.py @@ -511,7 +511,9 @@ class PersistentAdaptiveTemperedWeighting(BasePersistentTemperedWeighting): multinomial unique entries. Defaults to False. The fixed reference is set by the initial threshold; the terminal target never adapts. resample_at_end : bool, optional - If True, let downstream resampling and mutation finish before stopping. + If True, let downstream resampling and mutation generate a beta-one + batch before checking termination on the next archived, weighted state. + The terminal ESS requirement is checked after that batch is included. Defaults to False, returning the terminal weighted archive immediately. beta_tolerance: Stop bisection once the beta bracket is no wider than this. @@ -711,15 +713,19 @@ def effective_sample_size(beta: Array) -> Array: context = context.with_metadata(target_total_ess=target_total_ess) output["target_total_ess"] = target_total_ess - if selected_terminal and ( - self._target_total_ess is None - or float(output["effective_sample_size"]) >= self._target_total_ess - ): - context = ( - context.request_stop() - if self._resample_at_end - else context.request_stop_immediately() + if ( + selected_terminal + # A newly selected beta of one still needs its resample/mutate + # cycle. On the next iteration, persistence includes that batch, + # and the stopping test uses the complete reweighted archive. + and (not self._resample_at_end or current_beta == 1.0) + and ( + self._target_total_ess is None + or float(output["effective_sample_size"]) + >= self._target_total_ess ) + ): + context = context.request_stop_immediately() return weighted, context, output def state_dict(self) -> dict[str, Any]: diff --git a/tests/integration_tests/test_persistent_sampling.py b/tests/integration_tests/test_persistent_sampling.py index 9bd632a5..d569869c 100644 --- a/tests/integration_tests/test_persistent_sampling.py +++ b/tests/integration_tests/test_persistent_sampling.py @@ -59,22 +59,18 @@ def test_make_persistent_bayesian_smc_reaches_gaussian_posterior( assert delta_betas[:2] == pytest.approx([0.0, 0.0]) assert betas[-1] == pytest.approx(1.0) assert result.context.stop_requested - assert result.context.stop_immediately is not resample_at_end - assert ( - isinstance(result.final_state, PersistentParticleState) - is not resample_at_end - ) - assert len(result.outputs[-1]) == (4 if resample_at_end else 2) - if resample_at_end: - assert result.final_state.particles.shape == (n_particles, 1) - np.testing.assert_allclose(result.final_state.weights, 1 / n_particles) + assert result.context.stop_immediately + assert isinstance(result.final_state, PersistentParticleState) + assert len(result.outputs[-1]) == 2 np.testing.assert_allclose(result.final_state.weights.sum(), 1.0) persist = smc._steps[0] # noqa: SLF001 assert isinstance(persist, PersistParticles) assert persist.persistent_state is not None - assert persist.persistent_state.n_batches == len(result) - assert persist.persistent_state.n_particles == len(result) * n_particles + n_batches = len(result) + assert persist.persistent_state.n_batches == n_batches + assert persist.persistent_state.n_particles == n_batches * n_particles + assert result.final_state.particles.shape == (n_batches * n_particles, 1) final_particles = np.asarray(result.final_state.particles)[:, 0] final_weights = np.asarray(result.final_state.weights) @@ -166,8 +162,10 @@ def test_runner_accumulates_posterior_batches_until_target_total_ess( ) +@pytest.mark.parametrize("resample_at_end", [False, True]) def test_persistent_checkpoint_resume_matches_uninterrupted_run( tmp_path, + resample_at_end: bool, ) -> None: n_particles = 32 target_ess = 2 * n_particles @@ -200,7 +198,7 @@ def make_sampler( n_particles=n_particles, target_ess=target_ess, target_total_ess=target_total_ess, - resample_at_end=False, + resample_at_end=resample_at_end, mutation_n_steps=2, mutation_kwargs={"step_fn": "tpcn"}, checkpoint_writers=checkpoint_writers, diff --git a/tests/test_prebuilt_samplers.py b/tests/test_prebuilt_samplers.py index d6eeaf93..67112e5c 100644 --- a/tests/test_prebuilt_samplers.py +++ b/tests/test_prebuilt_samplers.py @@ -365,12 +365,12 @@ def test_make_persistent_bayesian_smc_logs_archive_progress( assert "Persistent SMC iteration 0" in caplog.messages[0] terminal_message = caplog.messages[-1] - assert "Persistent SMC iteration 1" in terminal_message + assert "Persistent SMC iteration 2" in terminal_message assert "Beta=1" in terminal_message - assert "Batches=2" in terminal_message - assert f"Archive size={2 * n_particles}" in terminal_message + assert "Batches=3" in terminal_message + assert f"Archive size={3 * n_particles}" in terminal_message assert ( - f"Persistent ESS={2 * n_particles:.1f}/{2 * n_particles:.1f}" + f"Persistent ESS={3 * n_particles:.1f}/{2 * n_particles:.1f}" in terminal_message ) @@ -421,7 +421,7 @@ def apply( context: StepContext, rng: RandomGenerator, ) -> tuple[ParticleState, StepContext, None]: - if context.stop_requested: + if float(context.metadata["beta"]) == 1.0: terminal_calls.append(state) return state, context, None @@ -437,25 +437,135 @@ def apply( result = smc.run( ParticleState(np.arange(4.0)[:, None]), metadata=metadata, rng=123 ) - assert len(terminal_calls) == 1 + assert len(terminal_calls) >= 1 assert result.context.stop_requested - assert not result.context.stop_immediately - assert type(result.final_state) is ParticleState - assert result.final_state.particles.shape == (4, 1) - np.testing.assert_allclose(result.final_state.weights, 0.25) + assert result.context.stop_immediately + assert isinstance(result.final_state, PersistentParticleState) + n_batches = len(result) + assert result.final_state.particles.shape == (4 * n_batches, 1) + np.testing.assert_allclose(result.final_state.weights, 1 / (4 * n_batches)) + np.testing.assert_allclose( + result.final_state.particles[-4:], terminal_calls[-1].particles + ) + assert result.final_state.betas[-1] == pytest.approx(1.0) assert result.final_state.log_evidence == pytest.approx(0.0) required_ess = target_total_ess or 8 assert result.outputs[-1][1]["effective_sample_size"] >= required_ess -def test_persistent_zero_support_preserves_evidence(rng) -> None: +@pytest.mark.parametrize("auto_normalize", [False, True]) +def test_persistent_final_batch_reweights_complete_archive( + auto_normalize: bool, +) -> None: + terminal_batches: list[ParticleState] = [] + + class ShiftMutation(SMCStep): + def apply( + self, + state: ParticleState, + context: StepContext, + rng: RandomGenerator, + ) -> tuple[ParticleState, StepContext, None]: + # A deterministic marker verifies that the post-mutation batch is + # included, rather than its resampled parents. This is not MCMC. + state = state.replace_particles(state.particles + 0.5) + terminal_batches.append(state) + return state, context, None + + def log_likelihood(particles: np.ndarray) -> np.ndarray: + return -0.5 * particles[:, 0] ** 2 + + sampler, posterior, _ = make_persistent_bayesian_smc( + _quadratic_log_prior, + log_likelihood, + n_particles=4, + target_ess=1, + mutation=ShiftMutation(), + auto_normalize=auto_normalize, + logging_frequency=0, + ) + initial = np.linspace(-1, 2, 4)[:, None] + result = sampler.run(ParticleState(initial), rng=123) + final = result.final_state + assert isinstance(final, PersistentParticleState) + assert len(terminal_batches) == 1 + batch = terminal_batches[0] + np.testing.assert_array_equal(final.particles[:4], initial) + np.testing.assert_array_equal(final.particles[4:], batch.particles) + np.testing.assert_array_equal(final.betas, [0, 1]) + np.testing.assert_array_equal( + final.initial_ancestor_indices[4:], batch.initial_ancestor_indices + ) + log_z = np.log(np.mean(np.exp(log_likelihood(initial)))) + np.testing.assert_allclose(final.log_normalizing_constants, [0, log_z]) + likelihood = np.exp(log_likelihood(final.particles)) + mixture = (1 + likelihood / np.exp(log_z)) / 2 + raw_weights = likelihood / mixture + np.testing.assert_allclose(final.weights, raw_weights / raw_weights.sum()) + assert final.log_evidence == pytest.approx(np.log(raw_weights.mean())) + assert final.evaluation is not None + expected_evaluation = posterior.prepare_evaluation( + final.particles, result.context + ) + for name in expected_evaluation.components: + np.testing.assert_allclose( + final.evaluation[name], expected_evaluation[name] + ) + + +@pytest.mark.parametrize( + ("target_total_ess", "expected_batches"), [(3, 2), (9, 2), (14, 3)] +) +def test_persistent_checks_terminal_ess_after_mutation( + target_total_ess: float, expected_batches: int +) -> None: + mutated_batches: list[ParticleState] = [] + + class RecordMutation(SMCStep): + def apply( + self, + state: ParticleState, + context: StepContext, + rng: RandomGenerator, + ) -> tuple[ParticleState, StepContext, None]: + assert not context.stop_requested + assert float(context.metadata["beta"]) == 1.0 + mutated_batches.append(state) + return state, context, None + + sampler, _, _ = make_persistent_bayesian_smc( + _quadratic_log_prior, + _zero_log_likelihood, + n_particles=5, + target_ess=2, + target_total_ess=target_total_ess, + mutation=RecordMutation(), + logging_frequency=0, + ) + result = sampler.run(ParticleState(np.arange(5.0)[:, None]), rng=123) + final = result.final_state + assert isinstance(final, PersistentParticleState) + assert final.n_batches == expected_batches + assert len(mutated_batches) == expected_batches - 1 + assert final.effective_sample_size() >= target_total_ess + assert result.context.stop_immediately + np.testing.assert_array_equal( + final.particles[-5:], mutated_batches[-1].particles + ) + np.testing.assert_array_equal(final.betas[1:], 1.0) + + +@pytest.mark.parametrize("resample_at_end", [False, True]) +def test_persistent_zero_support_preserves_evidence( + rng, resample_at_end: bool +) -> None: smc, _, metadata = make_persistent_bayesian_smc( _zero_log_likelihood, lambda p: np.where(p[:, 0] > 0, 0.0, -np.inf), n_particles=4, target_ess=2, mutation="none", - resample_at_end=False, + resample_at_end=resample_at_end, logging_frequency=0, ) result = smc.run( @@ -466,8 +576,9 @@ def test_persistent_zero_support_preserves_evidence(rng) -> None: ) assert result.context.stop_requested assert result.context.metadata["beta"] == 1.0 + n_supported = 7 if resample_at_end else 3 np.testing.assert_allclose( - result.final_state.weights, [0, 1 / 3, 1 / 3, 1 / 3] + result.final_state.weights, [0] + [1 / n_supported] * n_supported ) assert result.final_state.log_evidence == pytest.approx(np.log(0.75))