From 6830c2ee1b1df53dcce96784186c1cd04f31d535 Mon Sep 17 00:00:00 2001 From: James Nightingale Date: Sun, 25 Oct 2026 08:11:39 -0700 Subject: [PATCH] fix: guard MCMC thinning and short-chain autocorrelation --- .../non_linear/search/mcmc/emcee/search.py | 18 +++--- autofit/non_linear/search/mcmc/zeus/search.py | 2 +- .../non_linear/search/mcmc/test_emcee.py | 7 +-- .../search/mcmc/test_short_chain.py | 61 +++++++++++++++++++ .../non_linear/search/mcmc/test_zeus.py | 6 +- 5 files changed, 77 insertions(+), 17 deletions(-) create mode 100644 test_autofit/non_linear/search/mcmc/test_short_chain.py diff --git a/autofit/non_linear/search/mcmc/emcee/search.py b/autofit/non_linear/search/mcmc/emcee/search.py index ea4f9444d..6dcf23476 100644 --- a/autofit/non_linear/search/mcmc/emcee/search.py +++ b/autofit/non_linear/search/mcmc/emcee/search.py @@ -284,7 +284,7 @@ def raw_samples_from(self, model, search_internal): else: discard = int(3.0 * np.max(auto_correlations.times)) - thin = int(np.max(auto_correlations.times) / 2.0) + thin = max(1, int(np.max(auto_correlations.times) / 2.0)) parameters, log_posterior = ChainPosterior.from_sampler( search_internal, label="Emcee" @@ -311,12 +311,16 @@ def auto_correlations_from(self, search_internal=None): times = search_internal.get_autocorr_time(tol=0) - previous_auto_correlation_times = emcee.autocorr.integrated_time( - x=search_internal.get_chain()[ - : -self.auto_correlation_settings.check_size, :, : - ], - tol=0, - ) + previous_chain = search_internal.get_chain()[ + : -self.auto_correlation_settings.check_size, :, : + ] + if len(previous_chain) == 0: + self.logger.debug("Unable to compute previous auto correlation times.") + previous_auto_correlation_times = None + else: + previous_auto_correlation_times = emcee.autocorr.integrated_time( + x=previous_chain, tol=0, + ) return AutoCorrelations( check_size=self.auto_correlation_settings.check_size, diff --git a/autofit/non_linear/search/mcmc/zeus/search.py b/autofit/non_linear/search/mcmc/zeus/search.py index 44c6988d2..d3b7a39bb 100644 --- a/autofit/non_linear/search/mcmc/zeus/search.py +++ b/autofit/non_linear/search/mcmc/zeus/search.py @@ -318,7 +318,7 @@ def raw_samples_from(self, model, search_internal): else: discard = int(3.0 * np.max(auto_correlations.times)) - thin = int(np.max(auto_correlations.times) / 2.0) + thin = max(1, int(np.max(auto_correlations.times) / 2.0)) parameters, log_posterior = ChainPosterior.from_sampler( search_internal, label="Zeus" diff --git a/test_autofit/non_linear/search/mcmc/test_emcee.py b/test_autofit/non_linear/search/mcmc/test_emcee.py index 2ab719373..e74683a81 100644 --- a/test_autofit/non_linear/search/mcmc/test_emcee.py +++ b/test_autofit/non_linear/search/mcmc/test_emcee.py @@ -70,15 +70,13 @@ def test__log_posteriors_are_aligned_with_the_thinned_chain(monkeypatch): noise_map=np.full(100, 1.0), ) - # Test mode 1 runs the real sampler for 10 steps with 20 walkers, so the - # auto-correlation `check_size` must be smaller than the chain for the - # auto-correlation calculation to have samples to work with. + # Test mode runs only 10 steps; keep the default check_size=100 to + # exercise unavailable previous autocorrelations on a real short chain. search = af.Emcee( name="emcee_log_prob_alignment", unique_tag="log_prob_alignment_test", auto_correlation_settings=af.AutoCorrelationsSettings( check_for_convergence=False, - check_size=5, required_length=2, ), number_of_cores=1, @@ -119,7 +117,6 @@ def test__update_computes_auto_correlations_once(monkeypatch): unique_tag="auto_correlations_once_test", auto_correlation_settings=af.AutoCorrelationsSettings( check_for_convergence=False, - check_size=5, required_length=2, ), number_of_cores=1, diff --git a/test_autofit/non_linear/search/mcmc/test_short_chain.py b/test_autofit/non_linear/search/mcmc/test_short_chain.py new file mode 100644 index 000000000..2c0151239 --- /dev/null +++ b/test_autofit/non_linear/search/mcmc/test_short_chain.py @@ -0,0 +1,61 @@ +"""Short chains keep aligned draws and never request a zero thinning step.""" +from types import SimpleNamespace + +import numpy as np +import pytest + +import autofit as af +from autofit.non_linear.search.mcmc.auto_correlations import AutoCorrelations + + +@pytest.mark.parametrize("search_cls", [af.Emcee, af.Zeus]) +@pytest.mark.parametrize("tau, thin", [(0.5, 1), (1.9, 1), (4.0, 2)]) +def test_thinning_preserves_matching_draws(search_cls, tau, thin, monkeypatch): + monkeypatch.delenv("PYAUTO_TEST_MODE", raising=False) + chain = np.arange(60.0).reshape(30, 2, 1) + log_prob = -chain[..., 0] ** 2 + + def select(values, discard=0, thin=1, flat=False): + selected = values[discard + thin - 1::thin] + return selected.reshape((-1,) + values.shape[2:]) if flat else selected + + sampler = SimpleNamespace( + get_chain=lambda **kw: select(chain, **kw), + get_log_prob=lambda **kw: select(log_prob, **kw), + ncall_total=60, + ) + search = search_cls() + monkeypatch.setattr(search, "auto_correlations_from", lambda **_: AutoCorrelations( + times=np.array([tau]), previous_times=None, + check_size=100, required_length=50, change_threshold=0.01, + )) + raw = search.raw_samples_from(af.Collection(x=af.GaussianPrior(mean=0, sigma=1)), sampler) + expected = chain[int(3 * tau) + thin - 1::thin].reshape(-1, 1) + np.testing.assert_array_equal(raw.parameters, expected) + np.testing.assert_array_equal(raw.log_posterior, -expected[:, 0] ** 2) + + +@pytest.mark.parametrize("length", [10, 100, 101]) +def test_previous_autocorrelation_requires_a_nonempty_chain(length, monkeypatch): + import emcee + + calls = [] + + def integrated_time(x, tol): + calls.append(x) + assert len(x) > 0 + return np.array([2.0]) + + monkeypatch.setattr(emcee.autocorr, "integrated_time", integrated_time) + sampler = SimpleNamespace( + get_chain=lambda: np.ones((length, 4, 1)), + get_autocorr_time=lambda **_: np.array([3.0]), + ) + result = af.Emcee().auto_correlations_from(sampler) + if length <= 100: + assert result.previous_times is None + assert calls == [] + assert not result.check_if_converged(total_samples=length) + else: + np.testing.assert_array_equal(result.previous_times, [2.0]) + assert calls[0].shape == (1, 4, 1) diff --git a/test_autofit/non_linear/search/mcmc/test_zeus.py b/test_autofit/non_linear/search/mcmc/test_zeus.py index 8f509e0c5..8b4dac8dd 100644 --- a/test_autofit/non_linear/search/mcmc/test_zeus.py +++ b/test_autofit/non_linear/search/mcmc/test_zeus.py @@ -81,15 +81,13 @@ def test__log_posteriors_are_aligned_with_the_thinned_chain(monkeypatch): noise_map=np.full(100, 1.0), ) - # Test mode 1 runs the real sampler for 10 steps with 20 walkers, so the - # auto-correlation `check_size` must be smaller than the chain for the - # auto-correlation calculation to have samples to work with. + # Test mode runs only 10 steps; keep the default check_size=100 to + # exercise unavailable previous autocorrelations on a real short chain. search = af.Zeus( name="zeus_log_prob_alignment", unique_tag="log_prob_alignment_test", auto_correlation_settings=af.AutoCorrelationsSettings( check_for_convergence=False, - check_size=5, required_length=2, ), number_of_cores=1,