diff --git a/autofit/config/visualize/plots_search.yaml b/autofit/config/visualize/plots_search.yaml index 7552865d3..f99dc58f2 100644 --- a/autofit/config/visualize/plots_search.yaml +++ b/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? \ No newline at end of file + 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)? \ No newline at end of file diff --git a/autofit/non_linear/plot/__init__.py b/autofit/non_linear/plot/__init__.py index d581e9dd7..4b88c0ec0 100644 --- a/autofit/non_linear/plot/__init__.py +++ b/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 diff --git a/autofit/non_linear/plot/mle_plotters.py b/autofit/non_linear/plot/mle_plotters.py index c118d90a1..e7df552c4 100644 --- a/autofit/non_linear/plot/mle_plotters.py +++ b/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) diff --git a/autofit/non_linear/samples/samples.py b/autofit/non_linear/samples/samples.py index f44db636b..0de283e63 100644 --- a/autofit/non_linear/samples/samples.py +++ b/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: diff --git a/autofit/non_linear/search/mle/abstract_mle.py b/autofit/non_linear/search/mle/abstract_mle.py index 32f61e133..90956c947 100644 --- a/autofit/non_linear/search/mle/abstract_mle.py +++ b/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") diff --git a/autofit/non_linear/search/mle/multi_start_gradient/search.py b/autofit/non_linear/search/mle/multi_start_gradient/search.py index f94c7785d..6c8ea48b4 100644 --- a/autofit/non_linear/search/mle/multi_start_gradient/search.py +++ b/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, } diff --git a/test_autofit/non_linear/search/mle/test_multi_start_gradient.py b/test_autofit/non_linear/search/mle/test_multi_start_gradient.py index c105c4299..59ea78630 100644 --- a/test_autofit/non_linear/search/mle/test_multi_start_gradient.py +++ b/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