diff --git a/autofit/non_linear/search/nest/nautilus/search.py b/autofit/non_linear/search/nest/nautilus/search.py index 9b42fd183..8dd1054ba 100644 --- a/autofit/non_linear/search/nest/nautilus/search.py +++ b/autofit/non_linear/search/nest/nautilus/search.py @@ -603,7 +603,7 @@ def call_search(self, search_internal, model, analysis, fitness): search_internal=search_internal ) - search_internal.run( + converged = search_internal.run( f_live=self.f_live, n_shell=self.n_shell, n_eff=self.n_eff, @@ -615,8 +615,12 @@ def call_search(self, search_internal, model, analysis, fitness): iterations_after_run = self.iterations_from(search_internal=search_internal)[1] if ( - total_iterations == iterations_after_run - or iterations_after_run == self.n_like_max + converged + or total_iterations == iterations_after_run + or ( + self.n_like_max is not None + and iterations_after_run >= self.n_like_max + ) ): finished = True @@ -636,33 +640,27 @@ def iterations_from( self, search_internal ) -> Tuple[int, int]: """ - Returns the next number of iterations that a dynesty call will use and the total number of iterations - that have been performed so far. - - This is used so that the `iterations_per_full_update` input leads to on-the-fly output of dynesty results. + Return the next cumulative likelihood-call budget and calls performed. - It also ensures dynesty does not perform more samples than the `n_like_max` input variable. + Nautilus budgets likelihood evaluations, not posterior samples. Batches + can overshoot a budget, so advance from the sampler's actual call count. Parameters ---------- search_internal - The Dynesty sampler (static or dynamic) which is run and performs nested sampling. + The Nautilus sampler. Returns ------- - The next number of iterations that a dynesty run sampling will perform and the total number of iterations - it has performed so far. + The next cumulative likelihood-call limit and current likelihood count. """ + total_iterations = search_internal.n_like + if isinstance(self.paths, NullPaths): if self.n_like_max is not None and self.n_like_max != float("inf"): - return int(self.n_like_max), int(self.n_like_max) - return int(1e99), int(1e99) - - try: - total_iterations = len(search_internal.posterior()[1]) - except ValueError: - total_iterations = 0 + return int(self.n_like_max), total_iterations + return int(1e99), total_iterations iterations = total_iterations + self.iterations_per_full_update @@ -760,4 +758,4 @@ def samples_via_internal_from( @property def batch_size(self): - return self.n_batch \ No newline at end of file + return self.n_batch diff --git a/test_autofit/non_linear/search/nest/test_nautilus.py b/test_autofit/non_linear/search/nest/test_nautilus.py index 58b85ce2b..84829ca96 100644 --- a/test_autofit/non_linear/search/nest/test_nautilus.py +++ b/test_autofit/non_linear/search/nest/test_nautilus.py @@ -17,6 +17,60 @@ ) +@pytest.mark.parametrize( + "steps, limit, expected_budgets, expected_updates", + [ + ([(40, True)], None, [300], 0), + ([(320, False), (360, True)], None, [300, 620], 1), + ([(320, False), (640, False)], 600, [300, 600], 1), + ([(300, False)], 300, [300], 0), + ([(0, False)], None, [300], 0), + ], +) +def test__completion_and_likelihood_budgets( + monkeypatch, steps, limit, expected_budgets, expected_updates +): + """Convergence ends immediately; incomplete chunks advance by real calls.""" + search = af.Nautilus( + name="completion_budget", n_live=10, + iterations_per_full_update=300, n_like_max=limit, + ) + updates = [] + monkeypatch.setattr(search, "perform_update", lambda **kw: updates.append(kw)) + + class Sampler: + n_like = 0 + + def __init__(self): + self.budgets = [] + + def run(self, **kwargs): + self.budgets.append(kwargs["n_like_max"]) + self.n_like, converged = steps[len(self.budgets) - 1] + return converged + + def posterior(self): + raise AssertionError("Posterior sample count is not a call budget") + + sampler = Sampler() + assert search.call_search(sampler, None, None, None) is sampler + assert sampler.budgets == expected_budgets + assert len(updates) == expected_updates + assert all(update["during_analysis"] for update in updates) + + +@pytest.mark.parametrize("limit", [None, float("inf"), 100]) +def test__null_paths_likelihood_budget(limit): + from types import SimpleNamespace + from autofit.non_linear.paths.null import NullPaths + + search = af.Nautilus(n_like_max=limit) + assert isinstance(search.paths, NullPaths) + budget, calls = search.iterations_from(SimpleNamespace(n_like=40)) + assert calls == 40 + assert budget == (100 if limit == 100 else int(1e99)) + + def test__explicit_params(): search = af.Nautilus( n_live=500,