feat(autofit): multi-start gradient auto-convergence (phase 1)

autofit/__init__.py

@@ -92,6 +92,9 @@
 from .non_linear.search.mle.multi_start_gradient.search import MultiStartADABelief
 from .non_linear.search.mle.multi_start_gradient.search import MultiStartLion
 from .non_linear.search.mle.multi_start_gradient.search import MultiStartProdigy
+from .non_linear.search.mle.multi_start_gradient.convergence import (
+    MultiStartGradientConvergence,
+)
 from .non_linear.paths.abstract import AbstractPaths
 from .non_linear.paths import DirectoryPaths
 from .non_linear.paths import DatabasePaths

autofit/non_linear/search/mle/multi_start_gradient/convergence.py

@@ -0,0 +1,105 @@
+import numpy as np
+
+from autofit.non_linear.test_mode import is_test_mode
+
+
+class MultiStartGradientConvergence:
+    def __init__(
+        self,
+        check_for_convergence: bool = True,
+        window: int = 50,
+        rtol: float = 1.0e-4,
+        atol: float = 1.0e-3,
+        min_steps: int = 100,
+    ):
+        """
+        Settings for the auto-convergence (early-stopping) check of the multi-start
+        gradient searches (``MultiStartAdam`` / ``MultiStartADABelief`` /
+        ``MultiStartLion`` / ``MultiStartProdigy``).
+
+        This mirrors the ``AutoCorrelationsSettings`` precedent that lets the
+        ensemble MCMC samplers (Emcee / Zeus) terminate before their full step
+        budget: the search still has a hard ``n_steps`` ceiling (it never runs
+        forever), but when ``check_for_convergence`` is ``True`` it stops early once
+        the global-best figure-of-merit has plateaued.
+
+        The search minimises ``-2 * log_posterior`` (a chi-squared-like quantity),
+        so the tracked global best figure-of-merit ``best_fom`` is monotonically
+        non-increasing. Convergence is declared when the improvement over the
+        trailing ``window`` of the best-fom history is within tolerance::
+
+            improvement = fom_history[-window] - fom_history[-1]   # >= 0
+            converged   = improvement <= atol + rtol * abs(fom_history[-window])
+
+        evaluated only once at least ``max(min_steps, window)`` steps have been
+        taken, so early transient plateaus (before the population has descended)
+        cannot trigger a false stop.
+
+        Scope: this is the parametric-source (MGE / Sersic) regime, where a
+        global-best plateau genuinely means converged. It is deliberately **not**
+        applied when ``resurrect=True`` (the pixelized regime), whose best-fom
+        climbs in long plateaus punctuated by breakthrough jumps that plateau
+        detection would false-stop on; there the search leans on the ``n_steps``
+        ceiling instead.
+
+        Parameters
+        ----------
+        check_for_convergence
+            Whether the global-best figure-of-merit is checked to terminate the
+            search early. If ``True`` the search may stop before ``n_steps`` has
+            been reached; if ``False`` all ``n_steps`` steps are taken.
+        window
+            The number of trailing steps of the best-fom history over which the
+            plateau (improvement within tolerance) is measured. A longer window
+            requires the best-fom to be flat for longer before stopping, trading
+            a few extra steps for robustness against noise.
+        rtol
+            The relative tolerance on the best-fom improvement over the window.
+        atol
+            The absolute tolerance on the best-fom improvement over the window.
+        min_steps
+            The minimum number of steps that must be taken before the convergence
+            check is allowed to terminate the search, so the population has time to
+            descend out of its broad starts before a plateau can stop it.
+        """
+        self.check_for_convergence = check_for_convergence
+        self.window = window
+        self.rtol = rtol
+        self.atol = atol
+        self.min_steps = min_steps
+
+        if is_test_mode():
+            self.window = 1
+            self.min_steps = 1
+
+    def check_if_converged(self, fom_history) -> bool:
+        """
+        Whether the multi-start gradient search has converged: the global-best
+        figure-of-merit has plateaued over the trailing ``window``.
+
+        Parameters
+        ----------
+        fom_history
+            The history of the global best figure-of-merit, one entry per step
+            (monotonically non-increasing). Convergence is only assessed once at
+            least ``max(min_steps, window)`` entries are present.
+        """
+        if not self.check_for_convergence:
+            return False
+
+        fom_history = np.asarray(fom_history)
+
+        if fom_history.size < max(self.min_steps, self.window):
+            return False
+
+        latest = fom_history[-1]
+        past = fom_history[-self.window]
+
+        # Both non-finite (no finite basin found yet) -> not converged; the plateau
+        # of the actual descent is what we want to detect, not a flat +inf run.
+        if not np.isfinite(latest) or not np.isfinite(past):
+            return False
+
+        improvement = past - latest
+
+        return bool(improvement <= self.atol + self.rtol * abs(past))

autofit/non_linear/search/mle/multi_start_gradient/search.py

@@ -12,6 +12,9 @@
 from autofit.non_linear.initializer import AbstractInitializer
 from autofit.non_linear.samples.sample import Sample
 from autofit.non_linear.samples.samples import Samples
+from autofit.non_linear.search.mle.multi_start_gradient.convergence import (
+    MultiStartGradientConvergence,
+)
 
 
 class AbstractMultiStartGradient(AbstractMLE):
@@ -38,6 +41,7 @@ def __init__(
         start_lower_limit: float = 0.15,
         start_upper_limit: float = 0.85,
         resurrect: bool = False,
+        convergence: Optional[MultiStartGradientConvergence] = None,
         initializer: Optional[AbstractInitializer] = None,
         iterations_per_full_update: int = None,
         iterations_per_quick_update: int = None,
@@ -112,6 +116,16 @@ def __init__(
             searchable at all. (Even so, on such landscapes a nested sampler
             still wins decisively — resurrection makes gradient MAP *viable*
             there, not competitive.)
+        convergence
+            Auto-convergence (early-stopping) settings. When
+            ``check_for_convergence`` is ``True`` (the default) the search stops
+            early once the global-best figure-of-merit has plateaued, so users do
+            not have to hand-tune ``n_steps`` — which remains a hard ceiling / max
+            budget (the search never runs forever). ``None`` builds the default
+            ``MultiStartGradientConvergence()``. The check is deliberately skipped
+            when ``resurrect=True`` (the pixelized regime, whose best-fom climbs in
+            breakthrough jumps that a plateau check would false-stop on); there the
+            search leans on the ``n_steps`` ceiling.
         """
 
         super().__init__(
@@ -136,6 +150,9 @@ def __init__(
         self.start_lower_limit = start_lower_limit
         self.start_upper_limit = start_upper_limit
         self.resurrect = resurrect
+        self.convergence = (
+            convergence if convergence is not None else MultiStartGradientConvergence()
+        )
 
         self.logger.debug(f"Creating {self.optax_method} MultiStartGradient Search")
 
@@ -223,6 +240,7 @@ def batched_value_and_grad(params):
             fom_history = list(search_internal["fom_history"])
             total_steps = int(search_internal["total_steps"])
             n_resurrections = int(search_internal.get("n_resurrections", 0))
+            stop_reason = search_internal.get("stop_reason", None)
 
             self.logger.info(
                 "Resuming MultiStartGradient search (previous samples found)."
@@ -246,6 +264,7 @@ def batched_value_and_grad(params):
             fom_history = []
             total_steps = 0
             n_resurrections = 0
+            stop_reason = None
 
             self.logger.info(
                 f"Starting new {self.optax_method} MultiStartGradient search "
@@ -256,13 +275,26 @@ def batched_value_and_grad(params):
         # ``resurrect`` is on); seeded independently of the broad-start draw.
         resurrect_rng = np.random.default_rng(1)
 
-        while total_steps < self.n_steps:
+        # ``n_steps`` is the hard ceiling / max budget; ``stop_reason`` becomes
+        # ``"converged"`` if the auto-convergence check stops the search early
+        # (this also short-circuits the loop on a resumed, already-converged run).
+        while total_steps < self.n_steps and stop_reason != "converged":
 
             steps_remaining = self.n_steps - total_steps
             iterations = min(
                 self.iterations_per_full_update or self.n_steps, steps_remaining
             )
 
+            # Convergence is assessed every step (``fom_history`` updates every
+            # step) rather than only at the ``iterations_per_full_update`` boundary,
+            # so early-stopping works even with the default single-chunk update
+            # cadence (``iterations_per_full_update=None``). Checkpointing and the
+            # expensive ``perform_update`` still happen only at the boundary (or on
+            # convergence), so the extra work per step is a cheap NumPy plateau
+            # check. The check is skipped when ``resurrect`` is on (pixelized
+            # regime), where a plateau does not mean converged.
+            converged = False
+
             for _ in range(iterations):
                 foms, grads = batched_value_and_grad(params)
 
@@ -299,7 +331,18 @@ def batched_value_and_grad(params):
                 updates, opt_state = step_update(grads, opt_state, params, foms)
                 params = optax.apply_updates(params, updates)
 
-            total_steps += iterations
+                total_steps += 1
+
+                if not self.resurrect and self.convergence.check_if_converged(
+                    fom_history
+                ):
+                    converged = True
+                    break
+
+            if converged:
+                stop_reason = "converged"
+            elif total_steps >= self.n_steps:
+                stop_reason = "max_steps"
 
             search_internal = {
                 "params": np.asarray(params),
@@ -309,17 +352,27 @@ def batched_value_and_grad(params):
                 "fom_history": np.asarray(fom_history),
                 "total_steps": total_steps,
                 "n_resurrections": n_resurrections,
+                "stop_reason": stop_reason,
             }
             self.paths.save_search_internal(obj=search_internal)
 
+            # A converged (or ceiling-reached) boundary is the final update, so it
+            # runs with ``during_analysis=False``; intermediate boundaries do not.
+            is_final = converged or total_steps >= self.n_steps
             self.perform_update(
                 model=model,
                 analysis=analysis,
-                during_analysis=total_steps < self.n_steps,
+                during_analysis=not is_final,
                 fitness=fitness,
                 search_internal=search_internal,
             )
 
+            if converged:
+                self.logger.info(
+                    f"{self.optax_method} MultiStartGradient converged early at "
+                    f"step {total_steps} (best figure-of-merit plateaued)."
+                )
+
         self.logger.info(f"{self.optax_method} MultiStartGradient sampling complete.")
 
         return search_internal, fitness
@@ -374,9 +427,9 @@ def _build_optimizer(self, optax, jax):
 
             @jax.jit
             def step_update(grads, opt_state, params, values):
-                return jax.vmap(
-                    lambda g, s, p, v: optimizer.update(g, s, p, value=v)
-                )(grads, opt_state, params, values)
+                return jax.vmap(lambda g, s, p, v: optimizer.update(g, s, p, value=v))(
+                    grads, opt_state, params, values
+                )
 
         else:
 
@@ -490,9 +543,7 @@ def samples_via_internal_from(
         # Fitness.call returns -2 * log_posterior, so log_posterior = -0.5 * fom.
         best_log_posterior = -0.5 * float(search_internal["best_fom"])
         log_likelihood_list = [best_log_posterior - log_prior_list[0]]
-        log_likelihood_list += [
-            np.nan for _ in range(len(parameter_lists) - 1)
-        ]
+        log_likelihood_list += [np.nan for _ in range(len(parameter_lists) - 1)]
 
         weight_list = [1.0] + [0.0] * (len(parameter_lists) - 1)
 

test_autofit/non_linear/search/mle/test_multi_start_gradient.py

@@ -114,6 +114,96 @@ def test__batch_size_is_carried_to_every_rule():
         assert cls(batch_size=8).batch_size == 8
 
 
+def test__convergence_default_is_on_and_carried():
+    """Auto-convergence is a shared base-class setting, on by default (so users do
+    not hand-tune ``n_steps``), and a custom settings object is carried through."""
+    for cls in (
+        af.MultiStartAdam,
+        af.MultiStartADABelief,
+        af.MultiStartLion,
+        af.MultiStartProdigy,
+    ):
+        default = cls().convergence
+        assert isinstance(default, af.MultiStartGradientConvergence)
+        assert default.check_for_convergence is True
+        # documented defaults for the plateau check.
+        assert default.window == 50
+        assert default.min_steps == 100
+        assert default.rtol == pytest.approx(1.0e-4)
+        assert default.atol == pytest.approx(1.0e-3)
+
+        custom = af.MultiStartGradientConvergence(check_for_convergence=False)
+        assert cls(convergence=custom).convergence is custom
+
+
+def test__check_if_converged__plateau_stops_climbing_does_not():
+    # rtol/atol both zero: converges only on an exactly-flat window.
+    convergence = af.MultiStartGradientConvergence(
+        window=3, min_steps=3, rtol=0.0, atol=0.5
+    )
+
+    # Flat global-best over the window -> converged.
+    assert convergence.check_if_converged([10.0, 10.0, 10.0]) is True
+
+    # Still descending over the window -> not converged.
+    assert convergence.check_if_converged([10.0, 5.0, 1.0]) is False
+
+
+def test__check_if_converged__respects_min_steps_and_window_length():
+    convergence = af.MultiStartGradientConvergence(
+        window=3, min_steps=5, rtol=0.0, atol=1.0
+    )
+
+    # Fewer than max(min_steps, window) entries -> never terminates early.
+    assert convergence.check_if_converged([10.0, 10.0, 10.0]) is False
+    assert convergence.check_if_converged([10.0, 10.0, 10.0, 10.0]) is False
+
+    # Once min_steps entries exist and the window is flat -> converged.
+    assert convergence.check_if_converged([10.0] * 5) is True
+
+
+def test__check_if_converged__relative_tolerance():
+    convergence = af.MultiStartGradientConvergence(
+        window=3, min_steps=3, rtol=1.0e-4, atol=0.0
+    )
+
+    # Improvement 0.1 over a best-fom ~1000 is within rtol * 1000 = 0.1 -> converged.
+    assert convergence.check_if_converged([1000.1, 1000.05, 1000.0]) is True
+
+    # A large improvement over the same window is not within tolerance.
+    assert convergence.check_if_converged([1002.0, 1001.0, 1000.0]) is False
+
+
+def test__check_if_converged__disabled_and_non_finite():
+    # check_for_convergence=False never terminates, even on a flat window.
+    disabled = af.MultiStartGradientConvergence(
+        check_for_convergence=False, window=3, min_steps=3
+    )
+    assert disabled.check_if_converged([10.0, 10.0, 10.0]) is False
+
+    # A non-finite window (no finite basin found yet) is not a real plateau.
+    convergence = af.MultiStartGradientConvergence(window=3, min_steps=3, atol=1.0)
+    assert convergence.check_if_converged([np.inf, np.inf, np.inf]) is False
+
+
+def test__dict_round_trip__convergence():
+    # The convergence settings must survive serialisation so a resumed search
+    # keeps the same early-stopping behaviour.
+    search = af.MultiStartAdam(
+        n_starts=4,
+        convergence=af.MultiStartGradientConvergence(
+            check_for_convergence=False, window=7, min_steps=9
+        ),
+    )
+    restored = from_dict(to_dict(search))
+
+    assert isinstance(restored, af.MultiStartAdam)
+    assert isinstance(restored.convergence, af.MultiStartGradientConvergence)
+    assert restored.convergence.check_for_convergence is False
+    assert restored.convergence.window == 7
+    assert restored.convergence.min_steps == 9
+
+
 def test__dict_round_trip():
     dictionary = to_dict(af.MultiStartLion(n_starts=7, n_steps=33))
     restored = from_dict(dictionary)