feat(autofit): multi-start gradient convergence results contract (phase 2)

autofit/config/visualize/plots_search.yaml

@@ -4,4 +4,5 @@ mcmc:
   corner_cornerpy: true     # Output corner figure (using corner.py) during a non-linear search fit?
 mle:
   subplot_parameters: true   # Output a subplot of the best-fit parameters of the model?
-  log_likelihood_vs_iteration: true  # Output a plot of the log likelihood versus iteration number?
+  log_likelihood_vs_iteration: true  # Output a plot of the log likelihood versus iteration number?
+  figure_of_merit_vs_iteration: true  # Output the global-best figure-of-merit trace (auto-convergence gradient searches)?

autofit/non_linear/plot/__init__.py

@@ -1,4 +1,8 @@
 from autofit.non_linear.plot.samples_plotters import corner_cornerpy
 from autofit.non_linear.plot.nest_plotters import corner_anesthetic
-from autofit.non_linear.plot.mle_plotters import subplot_parameters, log_likelihood_vs_iteration
+from autofit.non_linear.plot.mle_plotters import (
+    subplot_parameters,
+    log_likelihood_vs_iteration,
+    figure_of_merit_vs_iteration,
+)
 from autofit.non_linear.plot.plot_util import output_figure

autofit/non_linear/plot/mle_plotters.py

@@ -92,3 +92,47 @@ def log_likelihood_vs_iteration(
         actual_filename += "_last_50_percent"
 
     output_figure(path=path, filename=actual_filename, format=format)
+
+
+@skip_in_test_mode
+def figure_of_merit_vs_iteration(
+    samples,
+    path=None,
+    filename="figure_of_merit_vs_iteration",
+    format="show",
+    **kwargs,
+):
+    """
+    Plot the global-best figure-of-merit trace of an auto-convergence gradient
+    search versus step, so the plateau the search stopped on can be inspected.
+
+    The trace is read from ``samples.samples_info["fom_history"]`` (the per-step
+    global best figure-of-merit, ``-2 * log_posterior``, that the multi-start
+    gradient searches record). Searches that do not record it (e.g. ``LBFGS`` /
+    ``Drawer``) leave it absent and this plot is skipped.
+    """
+    fom_history = samples.samples_info.get("fom_history")
+
+    if not fom_history:
+        return
+
+    import matplotlib.pyplot as plt
+
+    iteration_list = range(len(fom_history))
+
+    plt.figure(figsize=(12, 12))
+    plt.plot(iteration_list, fom_history, c="k")
+
+    plt.xlabel("Step", fontsize=16)
+    plt.ylabel("Global-Best Figure of Merit (-2 ln posterior)", fontsize=16)
+    plt.xticks(fontsize=16)
+    plt.yticks(fontsize=16)
+
+    stop_reason = samples.samples_info.get("stop_reason")
+    title = "Global-Best Figure of Merit vs Step (auto-convergence trace)"
+    if stop_reason is not None:
+        title += f"\nstopped: {stop_reason}"
+
+    plt.title(title, fontsize=20)
+
+    output_figure(path=path, filename=filename, format=format)

autofit/non_linear/samples/samples.py

@@ -309,8 +309,18 @@ def max_log_likelihood_sample(self) -> Sample:
     def max_log_likelihood_index(self) -> int:
         """
         The index of the sample with the highest log likelihood.
+
+        ``np.nanargmax`` (not ``np.argmax``) so ``NaN`` samples are never
+        selected: searches such as the multi-start gradient MAP optimizers record
+        their non-best starts as zero-weight diagnostic rows with a ``NaN`` log
+        likelihood, and plain ``argmax`` treats ``NaN`` as the maximum. This keeps
+        the index consistent with ``max_log_likelihood_sample``. If every sample
+        is ``NaN`` (a fully failed fit) the first index is returned.
         """
-        return int(np.argmax(self.log_likelihood_list))
+        log_likelihood_list = np.asarray(self.log_likelihood_list, dtype=float)
+        if np.all(np.isnan(log_likelihood_list)):
+            return 0
+        return int(np.nanargmax(log_likelihood_list))
 
     @to_instance
     def max_log_likelihood(self) -> List[float]:
@@ -333,8 +343,16 @@ def max_log_posterior_sample(self) -> Sample:
     def max_log_posterior_index(self) -> int:
         """
         The index of the sample with the highest log posterior.
+
+        ``np.nanargmax`` (not ``np.argmax``) so ``NaN`` samples are never
+        selected — see ``max_log_likelihood_index``; the zero-weight diagnostic
+        rows written by the multi-start gradient searches carry a ``NaN`` log
+        posterior. If every sample is ``NaN`` the first index is returned.
         """
-        return int(np.argmax(self.log_posterior_list))
+        log_posterior_list = np.asarray(self.log_posterior_list, dtype=float)
+        if np.all(np.isnan(log_posterior_list)):
+            return 0
+        return int(np.nanargmax(log_posterior_list))
 
     @to_instance
     def max_log_posterior(self) -> ModelInstance:

autofit/non_linear/search/mle/abstract_mle.py

@@ -3,16 +3,19 @@
 from autonerves import conf
 from autofit.non_linear.search.abstract_search import NonLinearSearch
 from autofit.non_linear.initializer import InitializerBall
-from autofit.non_linear.plot import subplot_parameters, log_likelihood_vs_iteration
+from autofit.non_linear.plot import (
+    subplot_parameters,
+    log_likelihood_vs_iteration,
+    figure_of_merit_vs_iteration,
+)
 
 
 class AbstractMLE(NonLinearSearch, ABC):
 
     def __init__(self, initializer=None, **kwargs):
         super().__init__(
-            initializer=initializer or InitializerBall(
-                lower_limit=0.49, upper_limit=0.51
-            ),
+            initializer=initializer
+            or InitializerBall(lower_limit=0.49, upper_limit=0.51),
             **kwargs,
         )
 
@@ -25,10 +28,29 @@ def should_plot(name):
 
         if should_plot("subplot_parameters"):
             subplot_parameters(samples=samples, path=plot_path, format="png")
-            subplot_parameters(samples=samples, use_log_y=True, path=plot_path, format="png")
-            subplot_parameters(samples=samples, use_last_50_percent=True, path=plot_path, format="png")
+            subplot_parameters(
+                samples=samples, use_log_y=True, path=plot_path, format="png"
+            )
+            subplot_parameters(
+                samples=samples, use_last_50_percent=True, path=plot_path, format="png"
+            )
 
         if should_plot("log_likelihood_vs_iteration"):
             log_likelihood_vs_iteration(samples=samples, path=plot_path, format="png")
-            log_likelihood_vs_iteration(samples=samples, use_log_y=True, path=plot_path, format="png")
-            log_likelihood_vs_iteration(samples=samples, use_last_50_percent=True, path=plot_path, format="png")
+            log_likelihood_vs_iteration(
+                samples=samples, use_log_y=True, path=plot_path, format="png"
+            )
+            log_likelihood_vs_iteration(
+                samples=samples, use_last_50_percent=True, path=plot_path, format="png"
+            )
+
+        # No-op unless the samples carry a ``fom_history`` (the auto-convergence
+        # gradient searches); LBFGS / Drawer leave it absent. Default-on and
+        # tolerant of a config predating this key (older workspaces that shadow
+        # the ``mle`` plots section) so it never KeyErrors on an MLE fit.
+        try:
+            plot_figure_of_merit = should_plot("figure_of_merit_vs_iteration")
+        except KeyError:
+            plot_figure_of_merit = True
+        if plot_figure_of_merit:
+            figure_of_merit_vs_iteration(samples=samples, path=plot_path, format="png")

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

@@ -535,6 +535,11 @@ def samples_via_internal_from(
         best_params = np.asarray(search_internal["best_params"])
         per_start_params = np.asarray(search_internal["params"])
         total_steps = int(search_internal["total_steps"])
+        # ``stop_reason`` is persisted by the fit loop ("converged" | "max_steps");
+        # older search_internal files (pre-auto-convergence) have neither it nor a
+        # ``fom_history``, so both are read defensively.
+        stop_reason = search_internal.get("stop_reason")
+        fom_history = search_internal.get("fom_history")
 
         parameter_lists = [list(best_params)] + [list(p) for p in per_start_params]
 
@@ -565,6 +570,24 @@ def samples_via_internal_from(
             "max_consecutive_nan": self.max_consecutive_nan,
             "resurrect": self.resurrect,
             "n_resurrections": int(search_internal.get("n_resurrections", 0)),
+            # Auto-convergence outcome: whether the run stopped on the plateau
+            # check ("converged") or exhausted the ``n_steps`` ceiling
+            # ("max_steps"), the settings that produced it, and the global-best
+            # figure-of-merit trace so the plateau can be inspected downstream.
+            "stop_reason": stop_reason,
+            "converged": stop_reason == "converged",
+            "convergence": {
+                "check_for_convergence": self.convergence.check_for_convergence,
+                "window": self.convergence.window,
+                "rtol": self.convergence.rtol,
+                "atol": self.convergence.atol,
+                "min_steps": self.convergence.min_steps,
+            },
+            "fom_history": (
+                [float(x) for x in np.asarray(fom_history)]
+                if fom_history is not None
+                else None
+            ),
             "time": self.timer.time if self.timer else None,
         }
 

test_autofit/non_linear/search/mle/test_multi_start_gradient.py

@@ -258,6 +258,7 @@ def test__samples_via_internal_from():
         "fom_history": np.asarray([-4.0, -8.0, best_fom]),
         "total_steps": 42,
         "n_resurrections": 7,
+        "stop_reason": "converged",
     }
 
     search = af.MultiStartAdam(n_starts=2, n_steps=42, resurrect=True)
@@ -289,3 +290,95 @@ def test__samples_via_internal_from():
     # restart-on-death diagnostics flow through to samples_info.
     assert samples.samples_info["resurrect"] is True
     assert samples.samples_info["n_resurrections"] == 7
+
+    # Auto-convergence outcome (phase-2 results contract) flows through to
+    # samples_info: the stop reason, the converged flag, the settings, and the
+    # global-best figure-of-merit trace (plain floats, JSON-round-trippable).
+    assert samples.samples_info["stop_reason"] == "converged"
+    assert samples.samples_info["converged"] is True
+    assert samples.samples_info["convergence"]["check_for_convergence"] is True
+    assert samples.samples_info["convergence"]["window"] == search.convergence.window
+    assert (
+        samples.samples_info["convergence"]["min_steps"] == search.convergence.min_steps
+    )
+    assert samples.samples_info["fom_history"] == pytest.approx([-4.0, -8.0, best_fom])
+    assert all(isinstance(x, float) for x in samples.samples_info["fom_history"])
+
+
+def test__samples_info__stop_reason_max_steps_and_legacy_search_internal():
+    model = af.Model(example.Gaussian)
+
+    best_params = np.asarray(model.vector_from_unit_vector([0.5] * model.prior_count))
+    per_start_params = np.stack([best_params])
+
+    base_internal = {
+        "params": per_start_params,
+        "best_params": best_params,
+        "best_fom": -2.0,
+        "total_steps": 300,
+        "n_resurrections": 0,
+    }
+
+    search = af.MultiStartAdam(n_starts=1, n_steps=300)
+
+    # Ceiling reached: converged is False, stop_reason is "max_steps".
+    samples = search.samples_via_internal_from(
+        model=model,
+        search_internal={
+            **base_internal,
+            "fom_history": np.asarray([-2.0]),
+            "stop_reason": "max_steps",
+        },
+    )
+    assert samples.samples_info["stop_reason"] == "max_steps"
+    assert samples.samples_info["converged"] is False
+
+    # Legacy search_internal (pre-auto-convergence: no stop_reason / fom_history)
+    # degrades gracefully rather than KeyError-ing.
+    legacy = search.samples_via_internal_from(
+        model=model, search_internal=base_internal
+    )
+    assert legacy.samples_info["stop_reason"] is None
+    assert legacy.samples_info["converged"] is False
+    assert legacy.samples_info["fom_history"] is None
+
+
+def test__variable_length_zero_weight_nan_rows_are_robust():
+    """The multi-start searches write zero-weight, NaN-log-likelihood diagnostic
+    rows for the non-best starts, and auto-convergence makes runs variable-length.
+    The max-likelihood/posterior accessors must never pick a NaN diagnostic row
+    (plain ``np.argmax`` would), and the aggregator summary must not raise for
+    runs of differing length."""
+    model = af.Model(example.Gaussian)
+    best_params = np.asarray(model.vector_from_unit_vector([0.5] * model.prior_count))
+
+    def samples_for(n_starts, total_steps):
+        per_start = np.stack(
+            [np.asarray(model.vector_from_unit_vector([0.4] * model.prior_count))]
+            * n_starts
+        )
+        search_internal = {
+            "params": per_start,
+            "best_params": best_params,
+            "best_fom": -10.0,
+            "fom_history": np.asarray([-4.0] * total_steps),
+            "total_steps": total_steps,
+            "n_resurrections": 0,
+            "stop_reason": "converged",
+        }
+        return af.MultiStartAdam(
+            n_starts=n_starts, n_steps=300
+        ).samples_via_internal_from(model=model, search_internal=search_internal)
+
+    for n_starts, total_steps in [(2, 40), (16, 175)]:
+        samples = samples_for(n_starts, total_steps)
+
+        # The NaN diagnostic rows are never selected as the best (argmax would).
+        assert samples.max_log_likelihood_index == 0
+        assert samples.max_log_posterior_index == 0
+        assert samples.max_log_likelihood(as_instance=False) == pytest.approx(
+            list(best_params)
+        )
+
+        # The aggregator summary is computed without an IndexError.
+        assert samples.summary() is not None