diff --git a/autofit/non_linear/search/nest/dynesty/search/abstract.py b/autofit/non_linear/search/nest/dynesty/search/abstract.py index 484316e79..7df9d239c 100644 --- a/autofit/non_linear/search/nest/dynesty/search/abstract.py +++ b/autofit/non_linear/search/nest/dynesty/search/abstract.py @@ -405,10 +405,7 @@ def iterations_from( # "unbounded" sentinel dynesty itself uses when ``maxcall=None``. return sys.maxsize, sys.maxsize - try: - total_iterations = np.sum(search_internal.results.ncall) - except AttributeError: - total_iterations = 0 + total_iterations = self.total_calls_from(search_internal=search_internal) if self.maxcall is not None: iterations = self.maxcall - total_iterations @@ -416,6 +413,52 @@ def iterations_from( return int(iterations), int(total_iterations) return self.iterations_per_full_update, int(total_iterations) + def total_calls_from(self, search_internal) -> int: + """ + The total number of likelihood calls the sampler has made so far, or 0 for a sampler that has not run. + + For the static sampler this is the sum of the per-sample call counts in `results.ncall`. + `DynestyDynamic` overrides this to use the sampler's own `ncall` counter, which is what dynesty + compares `maxcall` against: `results.ncall` omits the calls spent initialising each batch's live + points, so it undercounts and would compare a different quantity with the cumulative budget. + """ + try: + return int(np.sum(search_internal.results.ncall)) + except AttributeError: + return 0 + + def maxcall_from(self, iterations: int, total_iterations: int) -> int: + """ + The `maxcall` handed to `run_nested` for a chunk of `iterations` likelihood calls, given that + `total_iterations` calls have already been made. + + The static sampler counts calls per `run_nested` call, so the budget is the increment itself. + `DynestyDynamic` overrides this because the dynamic sampler counts cumulatively. + """ + return iterations + + def chunk_kwargs(self, total_iterations: int) -> Dict: + """ + Extra `run_nested` keyword arguments for a chunk, given the number of likelihood calls made before it. + + The static sampler continues a partially run sampler when `run_nested` is simply called again, so + nothing is needed. `DynestyDynamic` overrides this to run the baseline to completion in the first + chunk, because the dynamic sampler cannot continue a baseline run that was cut short. + """ + return {} + + def chunk_is_finished( + self, iterations: int, total_iterations: int, iterations_after_run: int + ) -> bool: + """ + Whether a `run_nested` chunk stopped on dynesty's own termination criterion rather than on its + budget, given the budget (`iterations`) and the total calls before and after the chunk. + + The static sampler stops on its budget only once its per-run counter is strictly above `maxcall`, so + a chunk that added no more calls than its budget converged. `DynestyDynamic` overrides this. + """ + return iterations_after_run - total_iterations <= iterations + def run_search_internal( self, search_internal: "Union[NestedSampler, DynamicNestedSampler]" ): @@ -429,17 +472,30 @@ def run_search_internal( 1. No output paths (`NullPaths`): there are no on-the-fly updates, so one pass is always final. 2. The global `maxcall` has been reached. - 3. Dynesty stopped itself inside the budget. Dynesty's per-run call counter resets on every - `run_nested` call and it stops on the budget only once that counter is *strictly above* `maxcall`, - while `sum(results.ncall)` grows by at least that counter. So a pass that added no more calls than - its budget stopped on dynesty's own termination criterion (`dlogz`, `maxiter`, ...) and is - converged. Under the default `iterations_per_full_update` (1e99) every converged run is therefore - finished in a single pass, with no intermediate `perform_update`, checkpoint restore or second - no-op `run_nested`. + 3. Dynesty stopped itself inside the budget (`chunk_is_finished`). The two samplers count `maxcall` + differently, so the budget handed to `run_nested` (`maxcall_from`) and this criterion are + per-sampler: + + - The static sampler's per-run call counter resets on every `run_nested` call and it stops on the + budget only once that counter is *strictly above* `maxcall`, while `sum(results.ncall)` grows by + at least that counter. So `maxcall` is the per-chunk increment, and a pass that added no more + calls than its budget stopped on dynesty's own termination criterion (`dlogz`, `maxiter`, ...). + - The dynamic sampler compares `maxcall` against its *cumulative* call count (`self.ncall` carried + over from previous `run_nested` calls). Handing it the per-chunk increment again on the second + chunk is a budget it has already spent, so it returns without sampling, criterion 4 fires and a + truncated baseline run is returned as the result. Nor can a baseline run that `maxcall` cut + short be continued (the next call restarts it, and `resume=True` refuses after `RUN_DONE`). + `DynestyDynamic` therefore runs the whole baseline in the first chunk (`maxbatch=0`, no budget), + then chunks the batch phase by cumulative budget `total_iterations + iterations`, and a batch + chunk is finished only when it stopped strictly inside that budget. See `DynestyDynamic.chunk_kwargs`. + + Under the default `iterations_per_full_update` (1e99) every converged run is finished in a single + pass for both samplers, with no intermediate `perform_update`, checkpoint restore or second no-op + `run_nested`. 4. Legacy criterion: the pass performed no new likelihood calls (e.g. re-running an already converged sampler restored from its checkpoint). - A pass that overran its budget (a finite `iterations_per_full_update` chunk) is not finished, so the + A pass that exhausted its budget (a finite `iterations_per_full_update` chunk) is not finished, so the search loops through `perform_update` and runs the next chunk. Parameters @@ -461,13 +517,16 @@ def run_search_internal( warnings.simplefilter("ignore") search_internal.run_nested( - maxcall=iterations, + maxcall=self.maxcall_from( + iterations=iterations, total_iterations=total_iterations + ), print_progress=not self.silence, checkpoint_file=self.checkpoint_file, **self.run_kwargs, + **self.chunk_kwargs(total_iterations=total_iterations), ) - iterations_after_run = np.sum(search_internal.results.ncall) + iterations_after_run = self.total_calls_from(search_internal=search_internal) if isinstance(self.paths, NullPaths): return True @@ -475,7 +534,11 @@ def run_search_internal( if self.maxcall is not None and iterations_after_run >= self.maxcall: return True - if iterations > 0 and iterations_after_run - total_iterations <= iterations: + if iterations > 0 and self.chunk_is_finished( + iterations=iterations, + total_iterations=total_iterations, + iterations_after_run=iterations_after_run, + ): return True return bool(total_iterations == iterations_after_run) diff --git a/autofit/non_linear/search/nest/dynesty/search/dynamic.py b/autofit/non_linear/search/nest/dynesty/search/dynamic.py index 705bb8a98..5e3df25bf 100644 --- a/autofit/non_linear/search/nest/dynesty/search/dynamic.py +++ b/autofit/non_linear/search/nest/dynesty/search/dynamic.py @@ -1,4 +1,6 @@ from __future__ import annotations + +import sys from typing import Dict, Optional from autofit.mapper.prior_model.abstract import AbstractPriorModel @@ -30,6 +32,7 @@ def __init__( maxcall_init: Optional[int] = None, maxiter: Optional[int] = None, maxiter_init: Optional[int] = None, + n_effective: Optional[int] = None, iterations_per_quick_update: int = None, iterations_per_full_update: int = None, number_of_cores: int = 1, @@ -57,6 +60,13 @@ def __init__( Number of live points used during the initial exploration phase. dlogz_init Stopping criterion for the initial baseline run. + n_effective + The minimum effective posterior sample size (ESS) of the whole run. After the baseline run, dynesty + keeps adding batches of live points until its default stopping function judges the estimated ESS to + have reached this value (`target_n_effective`). If `None` (the default), the argument is not passed + and dynesty's own default applies, which in dynesty 2.x is `max(10000, ndim ** 2)`. It has no effect + if a custom `stop_function` is used. It is not an identifier field, so changing it alone does not + give the search a new output folder. iterations_per_full_update The number of iterations performed between update (e.g. output latest model to hard-disk, visualization). number_of_cores @@ -78,6 +88,7 @@ def __init__( self.dlogz_init = dlogz_init self.logl_max_init = logl_max_init self.maxcall_init = maxcall_init + self.n_effective = n_effective self.maxiter = maxiter self.maxiter_init = maxiter_init @@ -89,7 +100,7 @@ def __init__( @property def run_kwargs(self) -> Dict: - return { + run_kwargs = { "dlogz_init": self.dlogz_init, "logl_max_init": self.logl_max_init, "maxcall_init": self.maxcall_init, @@ -97,6 +108,82 @@ def run_kwargs(self) -> Dict: "maxiter_init": self.maxiter_init, "nlive_init": self.nlive_init, } + if self.n_effective is not None: + run_kwargs["n_effective"] = self.n_effective + return run_kwargs + + @property + def chunked(self) -> bool: + """ + Whether a run is split into `run_nested` chunks for on-the-fly output, i.e. whether + `iterations_per_full_update` is finite rather than the inf-like default (1e99). + """ + return self.iterations_per_full_update is not None and self.iterations_per_full_update < sys.maxsize + + def total_calls_from(self, search_internal) -> int: + """ + The dynamic sampler's own cumulative call counter, the quantity dynesty compares the cumulative + `maxcall` with. `results.ncall` undercounts it (it omits each batch's live-point initialisation), + which would make a budget-limited batch chunk look converged. Falls back to the `results` count for + objects without the counter. + """ + try: + return int(search_internal.ncall) + except AttributeError: + return super().total_calls_from(search_internal=search_internal) + + def maxcall_from(self, iterations: int, total_iterations: int) -> int: + """ + The `maxcall` for a chunk of the dynamic sampler. + + dynesty's `DynamicNestedSampler.run_nested` compares `maxcall` with its *cumulative* call count + (`self.ncall`, carried over from earlier calls), unlike the static sampler whose counter resets per + call. The budget for a chunk is therefore the calls made so far plus the chunk size, capped at + `sys.maxsize` (dynesty's own "unbounded" value) so an inf-like value cannot overflow the C long + dynesty subtracts from. + + The first chunk of a chunked run is the exception: it is given the whole budget (`maxcall` if set, + otherwise unbounded) because the baseline run must never be cut short, see `chunk_kwargs`. + """ + if self.chunked and total_iterations == 0: + return self.maxcall if self.maxcall is not None else sys.maxsize + return int(min(total_iterations + iterations, sys.maxsize)) + + def chunk_kwargs(self, total_iterations: int) -> Dict: + """ + Extra `run_nested` keyword arguments for a chunk of the dynamic sampler. + + A dynamic run is a baseline nested sampling run followed by batches of live points. dynesty cannot + continue a baseline run that `maxcall` cut short: the next `run_nested` call restarts it from scratch + (`sample_initial` calls `reset()`), and `resume=True` refuses because the previous call ended in the + `RUN_DONE` state. Batches, however, can be added across `run_nested` calls: once `self.base` is set the + baseline is skipped and the batch loop continues from `self.batch`, evaluating the stopping criterion + (including `n_effective`) first. + + A chunked run therefore completes the baseline in its first chunk (`maxbatch=0`, no call budget) and + chunks only the batch phase, by cumulative `maxcall` (`maxcall_from`). Intermediate output starts + after the baseline run. + """ + if self.chunked and total_iterations == 0: + return {"maxbatch": 0} + return {} + + def chunk_is_finished( + self, iterations: int, total_iterations: int, iterations_after_run: int + ) -> bool: + """ + The baseline-only first chunk of a chunked run is never finished: the batch phase (and so the + `n_effective` stopping criterion) has not run yet. If it turns out to be satisfied already, the next + chunk adds no calls and `run_search_internal`'s no-new-calls criterion finishes the search. + + A batch chunk stops sampling once its cumulative call count reaches the cumulative `maxcall` (it may + land exactly on it), so it converged only if it stopped strictly inside that budget. + """ + if self.chunked and total_iterations == 0: + return False + return iterations_after_run < self.maxcall_from( + iterations=iterations, total_iterations=total_iterations + ) @property def search_internal(self): diff --git a/test_autofit/non_linear/search/conformance_roster.py b/test_autofit/non_linear/search/conformance_roster.py index 04819b35b..12e5d936e 100644 --- a/test_autofit/non_linear/search/conformance_roster.py +++ b/test_autofit/non_linear/search/conformance_roster.py @@ -404,6 +404,7 @@ class Golden: "maxcall_init", "maxiter", "maxiter_init", + "n_effective", "nlive_init", }, ), diff --git a/test_autofit/non_linear/search/nest/test_dynesty.py b/test_autofit/non_linear/search/nest/test_dynesty.py index 6eaaa2a8b..d78b67ca1 100644 --- a/test_autofit/non_linear/search/nest/test_dynesty.py +++ b/test_autofit/non_linear/search/nest/test_dynesty.py @@ -1,3 +1,5 @@ +import sys + import numpy as np import pytest @@ -62,6 +64,42 @@ def test__explicit_params(): assert search.number_of_cores == 1 +def test__dynamic_n_effective_passed_to_run_nested_only_when_set(): + + search = af.DynestyDynamic(n_effective=5000) + + assert search.n_effective == 5000 + assert search.run_kwargs["n_effective"] == 5000 + + search = af.DynestyDynamic() + + assert search.n_effective is None + assert "n_effective" not in search.run_kwargs + + +@pytest.mark.parametrize("search_cls", [af.DynestyStatic, af.DynestyDynamic]) +def test__null_paths_iteration_budget_fits_c_long(search_cls): + """ + Without an output directory (NullPaths) the per-chunk budget is "unbounded". + It must be ``sys.maxsize``, not ``int(1e99)``: dynesty's dynamic sampler + subtracts a numpy int64 ``ncall`` from it, which raised + ``OverflowError: Python int too large to convert to C long``. + """ + import sys + from autofit.non_linear.paths.null import NullPaths + + search = search_cls() + assert isinstance(search.paths, NullPaths) + + iterations, total = search.iterations_from(MockDynestySampler(results=None)) + + assert iterations == total == sys.maxsize + assert iterations - np.int64(40) == sys.maxsize - 40 + + search = search_cls(maxcall=100) + assert search.iterations_from(MockDynestySampler(results=None)) == (100, 100) + + @pytest.mark.parametrize("search_cls", [af.DynestyStatic, af.DynestyDynamic]) def test__null_paths_iteration_budget_fits_c_long(search_cls): """ @@ -141,6 +179,7 @@ def __init__(self, calls_per_run): samples=None, logl=None, logwt=None, ncall=[], logz=None, nlive=None ) self.maxcalls = [] + self.maxbatches = [] def run_nested(self, maxcall, **kwargs): self.maxcalls.append(maxcall) @@ -193,6 +232,115 @@ def test__run_search_internal__global_maxcall_is_finished(): assert sampler.maxcalls == [100] +class _FakeCumulativeRunSampler(_FakeRunSampler): + """ + Mirrors dynesty's *dynamic* sampler, whose `run_nested` compares `maxcall` against + the cumulative number of calls made so far: a run adds its scripted calls only up + to that cumulative budget, and adds nothing if the budget is already spent. + """ + + def run_nested(self, maxcall, **kwargs): + self.maxcalls.append(maxcall) + self.maxbatches.append(kwargs.get("maxbatch")) + total = sum(self.results.ncall) + wanted = self.calls_per_run.pop(0) + self.results.ncall.append(max(0, min(wanted, maxcall - total))) + + +def _dynamic_search_with_directory_paths(name, **kwargs): + search = af.DynestyDynamic(nlive_init=20, number_of_cores=1, silence=True, **kwargs) + search.paths = af.DirectoryPaths(name=name) + return search + + +def test__dynamic__first_chunk_runs_whole_baseline_then_batches_are_cumulative(): + """ + Under the default (inf-like) cadence nothing is chunked and the budget is simply + capped. With a finite cadence the first chunk is the whole baseline run (no call + budget, `maxbatch=0`) and later chunks are budgeted by cumulative calls. + """ + search = af.DynestyDynamic() + assert search.chunked is False + assert search.maxcall_from(iterations=1e99, total_iterations=0) == sys.maxsize + assert search.maxcall_from(iterations=1e99, total_iterations=120) == sys.maxsize + assert search.chunk_kwargs(total_iterations=0) == {} + + search = af.DynestyDynamic(iterations_per_full_update=50) + assert search.chunked is True + assert search.maxcall_from(iterations=50, total_iterations=0) == sys.maxsize + assert search.chunk_kwargs(total_iterations=0) == {"maxbatch": 0} + assert search.maxcall_from(iterations=50, total_iterations=120) == 170 + assert search.chunk_kwargs(total_iterations=120) == {} + + search = af.DynestyDynamic(iterations_per_full_update=50, maxcall=300) + assert search.maxcall_from(iterations=300, total_iterations=0) == 300 + + static = af.DynestyStatic(iterations_per_full_update=50) + assert static.maxcall_from(iterations=50, total_iterations=120) == 50 + assert static.chunk_kwargs(total_iterations=120) == {} + + +def test__dynamic__chunk_that_reaches_cumulative_budget_is_not_finished(): + """ + The baseline-only first chunk is never finished. A batch chunk that lands exactly + on its cumulative budget is budget-limited, not converged; one that stops short of + it converged. + """ + search = _dynamic_search_with_directory_paths( + "dynesty_dynamic_budget", iterations_per_full_update=50 + ) + + sampler = _FakeCumulativeRunSampler(calls_per_run=[30]) + assert search.run_search_internal(search_internal=sampler) is False + assert sampler.maxcalls == [sys.maxsize] + assert sampler.maxbatches == [0] + + sampler = _FakeCumulativeRunSampler(calls_per_run=[40, 50]) + search.run_search_internal(search_internal=sampler) + assert search.run_search_internal(search_internal=sampler) is False + assert sampler.maxcalls[-1] == 90 + + sampler = _FakeCumulativeRunSampler(calls_per_run=[40, 20]) + search.run_search_internal(search_internal=sampler) + assert search.run_search_internal(search_internal=sampler) is True + + +def test__fit__dynamic_finite_cadence_runs_to_convergence(monkeypatch): + """ + Regression test: with a finite `iterations_per_full_update` the dynamic sampler + used to be handed the same per-chunk `maxcall` on every pass. Because it compares + `maxcall` with its cumulative call count, the second pass added nothing and the + no-new-calls criterion returned a truncated baseline as finished. Now the first + pass is the whole baseline run, the batch passes get cumulative budgets, and the + run continues until dynesty converges (the third pass here, which stops short of + its budget), with an intermediate update after each unfinished pass. + """ + search = _dynamic_search_with_directory_paths( + "dynesty_dynamic_finite_cadence", iterations_per_full_update=50 + ) + sampler = _FakeCumulativeRunSampler(calls_per_run=[40, 50, 20]) + + monkeypatch.setattr( + search, "search_internal_from", lambda *args, **kwargs: sampler + ) + updates = [] + monkeypatch.setattr( + search, + "perform_update", + lambda *args, **kwargs: updates.append(kwargs.get("during_analysis")), + ) + + model = af.Model(af.ex.Gaussian) + analysis = af.ex.Analysis(data=np.full(10, 1.0), noise_map=np.full(10, 1.0)) + + search._fit(model=model, analysis=analysis) + + assert sampler.maxcalls == [sys.maxsize, 90, 140] + assert sampler.maxbatches == [0, None, None] + assert sampler.results.ncall == [40, 50, 20] + assert updates == [True, True] + + def test__fit__finite_cadence_loops_through_intermediate_updates(monkeypatch): """ A finite `iterations_per_full_update` still chunks the run: two budget-limited