diff --git a/autofit/graphical/README.md b/autofit/graphical/README.md index b40b2c454..df804d236 100644 --- a/autofit/graphical/README.md +++ b/autofit/graphical/README.md @@ -154,6 +154,34 @@ subtraction of natural parameters is not closed in the family), previous message per-parameter (`update_invalid`) and flags `StatusFlag.BAD_PROJECTION`. +**Failed factor update**: a factor's own optimiser may *raise* rather +than return — most commonly `InitializerException`, when EP has driven +the factor to a state where every drawn start point has the same figure +of merit. `factor_step` catches this, degrades to the factor's previous +message, and flags `StatusFlag.EXCEPTION`, so one bad factor costs one +sweep's update rather than the whole graph fit. This is distinct from a +*returned* `StatusFlag.FAILURE` (e.g. the Laplace optimiser's "line +search failed"), which EP absorbs routinely. A factor that raises on +every sweep is not going to start working, so after +`max_consecutive_failures` (default 3) consecutive raises on one factor +`run` stops sweeping early; only raises are counted, and the count +resets on any sweep that does not raise. Every raise is recorded in +`ep_history.csv` as an `EXCEPTION` row and logged as a warning. + +The result is still returned in that state — a partly-failed graph may +still hold converged messages worth having — but never quietly. If +enough factors raise, *nothing* in the mean field changes, so the KL +step of Eq. (12) is zero and `EPHistory` declares convergence — in +practice within two sweeps, before any per-factor count reaches its +threshold — and the mean field holds the starting priors for those +factors. `run` therefore checks, once the sweeps are over, whether any +factor both raised and never once updated, and emits a **STALE FACTORS** +warning naming them: logged, and written into `ep_diagnostics.results` +beside the sigma-collapse warnings. Read that file before trusting a +mean field from a run that logged failures. A factor that failed +intermittently but landed at least one update is not stale and is not +reported. + ## 4. Convergence — `EPHistory` (`expectation_propagation/history.py`) After each factor update the history records the new `EPMeanField`. diff --git a/autofit/graphical/expectation_propagation/optimiser.py b/autofit/graphical/expectation_propagation/optimiser.py index 3370283a1..d3eb9fee1 100644 --- a/autofit/graphical/expectation_propagation/optimiser.py +++ b/autofit/graphical/expectation_propagation/optimiser.py @@ -3,7 +3,7 @@ import os from abc import ABC, abstractmethod from pathlib import Path -from typing import Dict, Optional, List, Tuple +from typing import Dict, Optional, List, Set, Tuple from autofit import exc from autofit.graphical.expectation_propagation.ep_mean_field import EPMeanField @@ -136,12 +136,37 @@ def factor_step(factor_approx, optimiser, model_approx=None): messages = status.messages + tuple(caught_warnings.messages) - status = Status(status.success, messages, status.flag, result=status.result) + # Keyword arguments matter here: `Status`'s third positional parameter is + # `updated`, not `flag`. Passing the flag positionally silently dropped it + # and left `flag` at its `SUCCESS` default, so a failed factor step was + # recorded in `ep_history.csv` as a success. + status = Status( + success=status.success, + messages=messages, + updated=status.updated, + flag=status.flag, + result=status.result, + ) - except (ValueError, ArithmeticError, RuntimeError) as e: + except ( + ValueError, + ArithmeticError, + RuntimeError, + exc.InitializerException, + ) as e: + # `InitializerException` is raised when a factor's own optimiser cannot + # find a start point — most commonly because EP has driven the factor to + # a state where every drawn point has the same figure of merit. That is a + # failure of this sweep's update for this factor, not of the graph fit: + # degrade to the factor's previous message and let the sweep continue, + # with the failure recorded. `EPOptimiser` aborts if one factor keeps + # failing (see `max_consecutive_failures`). logger.exception(e) status = Status( - False, (f"Factor: {factor} experienced error {e}",), StatusFlag.FAILURE, + success=False, + messages=(f"Factor: {factor} experienced error {e}",), + updated=False, + flag=StatusFlag.EXCEPTION, ) new_model_dist = factor_approx.model_dist @@ -210,6 +235,15 @@ def __init__( self.ep_history = ep_history or EPHistory() self.diagnostics = EPDiagnostics() + # Per-factor count of consecutive failed updates; see + # `_check_consecutive_failures`. Reset at the start of every `run`. + self._consecutive_failures: Dict[Factor, int] = {} + # Factors that raised at least once, and factors that landed at least + # one successful update; together these identify a factor whose message + # is still the one it started with. See `_stale_factor_warnings`. + self._factors_raised: Set[Factor] = set() + self._factors_updated: Set[Factor] = set() + self.visualiser = None if paths is None: try: @@ -294,6 +328,115 @@ def _log_factor(self, factor: Factor): def factor_step(self, factor_approx, optimiser, model_approx=None): return factor_step(factor_approx, optimiser, model_approx=model_approx) + def _check_consecutive_failures( + self, + factor: Factor, + status: Status, + max_consecutive_failures: int, + raised: bool, + ) -> bool: + """ + Track how many sweeps in a row a given factor's optimiser has *raised*. + + A single raise is survivable — the sweep continues on that factor's + previous message. A factor that raises on *every* sweep is not going to + start working, so once it has raised `max_consecutive_failures` times in + a row there is nothing to gain by sweeping further: stop, warn, and let + `run` return what it has. The result is still reported, but loudly + qualified — see `_stale_factor_warnings`. + + Only raises are counted. A returned `StatusFlag.FAILURE` is an ordinary, + recoverable outcome that EP absorbs by design — the Laplace optimiser + returns one whenever its line search fails — and counting those would + cut healthy fits short. + + Counting is per-factor and resets on any sweep that does not raise, so + an intermittent failure (the observed case) never trips it. + + Parameters + ---------- + raised + Whether this factor's optimiser raised on this sweep. Read from the + status `factor_step` returned, *before* the mean-field projection, + which may legitimately overwrite the flag with `BAD_PROJECTION`. + + Returns + ------- + True if this factor has now failed enough consecutive sweeps that the + run should stop early. + """ + if raised: + self._factors_raised.add(factor) + count = self._consecutive_failures.get(factor, 0) + 1 + self._consecutive_failures[factor] = count + + logger.warning( + "Factor %s raised on %d consecutive step(s) " + "(giving up on it at %d); continuing with its previous message. " + "Latest messages: %s", + factor.name, + count, + max_consecutive_failures, + "; ".join(status.messages) or "(none)", + ) + + if max_consecutive_failures and count >= max_consecutive_failures: + logger.warning( + "Factor %s has raised on %d consecutive steps; abandoning " + "further sweeps. Its message is whatever it last held, so " + "the returned mean field is not a posterior for this factor.", + factor.name, + count, + ) + return True + else: + self._factors_updated.add(factor) + self._consecutive_failures.pop(factor, None) + + return False + + def _stale_factor_warnings(self) -> List[str]: + """ + Warn about any factor whose message is still the one it started with. + + A per-factor failure count is not enough to detect this. When several + factors raise, *nothing* in the mean field changes, so the KL step + between sweeps is zero and `EPHistory` declares convergence — often + within two sweeps, before any count reaches its threshold. The run then + terminates "successfully" and the returned mean field holds the starting + priors for those factors, dressed up as a posterior (PyAutoFit#1405). + + The result is still returned — callers with a partly-failed graph may + well want the factors that did converge — but never quietly: these + strings are logged as warnings and written into `ep_diagnostics.results` + alongside the sigma-collapse warnings. + + The condition is deliberately narrow: a factor that raised at least once + and *never once* updated. A factor that failed intermittently but landed + at least one update has a real message and is not reported. + """ + stale = self._factors_raised - self._factors_updated + if not stale: + return [] + + names = ", ".join(sorted(factor.name for factor in stale)) + return [ + f"STALE FACTORS: {names} never completed a single update — their " + f"optimisers raised on every sweep. The mean field returned for " + f"them is the prior the fit started with, not a posterior. Do not " + f"read those values as a result. Note that EP may also report " + f"convergence in this state: with no factor updating, the KL step " + f"between sweeps is zero, which is indistinguishable from having " + f"converged." + ] + + def _warn_stale_factors(self): + """ + Log the stale-factor warnings, whether or not output paths are enabled. + """ + for warning in self._stale_factor_warnings(): + logger.warning(warning) + def run( self, model_approx: EPMeanField, @@ -301,6 +444,7 @@ def run( log_interval: int = 10, visualise_interval: int = 100, output_interval: int = 10, + max_consecutive_failures: int = 3, ) -> EPMeanField: """ Run the optimisation on an approximation of the model. @@ -322,6 +466,13 @@ def run( How steps should we wait before outputting information? This includes the model.results file which describes the current mean values of each message. + max_consecutive_failures + How many consecutive sweeps a single factor's optimiser may *raise* + on before the fit is aborted. One raise is not fatal — the sweep + continues on that factor's previous message — but a factor that + raises every sweep would leave EP converging on a stale message and + reporting success. A returned failure status (e.g. a failed line + search) is not counted. Set to 0 to never abort. Returns ------- @@ -331,6 +482,10 @@ def run( should_visualise = IntervalCounter(visualise_interval) should_output = IntervalCounter(output_interval) + self._consecutive_failures = {} + self._factors_raised = set() + self._factors_updated = set() + for _ in range(max_steps): _should_log = should_log() _should_visualise = should_visualise() @@ -340,10 +495,15 @@ def run( new_model_dist, status = self.factor_step( factor_approx, optimiser, model_approx=model_approx, ) + raised = status.flag is StatusFlag.EXCEPTION model_approx, status = self.updater.update_model_approx( new_model_dist, factor_approx, model_approx, status ) self.diagnostics.snapshot(factor, model_approx, status) + if self._check_consecutive_failures( + factor, status, max_consecutive_failures, raised=raised + ): + break if status and _should_log: self._log_factor(factor) @@ -365,6 +525,7 @@ def run( self._output_results(model_approx) self._output_diagnostics(final=True, model_approx=model_approx) self._warn_sigma_collapse() + self._warn_stale_factors() return model_approx @@ -391,7 +552,9 @@ def _output_diagnostics( self.diagnostics.plot(self.output_path) if final and model_approx is not None: - warnings_list = check_sigma_collapse(self.diagnostics) + warnings_list = ( + self._stale_factor_warnings() + check_sigma_collapse(self.diagnostics) + ) with open(self.output_path / "ep_diagnostics.results", "w+") as f: f.write(mean_field_summary(model_approx.mean_field)) f.write("\n") @@ -470,6 +633,7 @@ def run( log_interval: int = 10, visualise_interval: int = 100, output_interval: int = 10, + max_consecutive_failures: int = 3, ) -> EPMeanField: """ Run the optimisation on an approximation of the model. @@ -491,6 +655,9 @@ def run( How steps should we wait before outputting information? This includes the model.results file which describes the current mean values of each message. + max_consecutive_failures + How many consecutive sweeps a single factor's optimiser may raise on + before the fit is aborted. See `EPOptimiser.run`. Returns ------- @@ -500,6 +667,10 @@ def run( should_visualise = IntervalCounter(visualise_interval) should_output = IntervalCounter(output_interval) + self._consecutive_failures = {} + self._factors_raised = set() + self._factors_updated = set() + for _ in range(max_steps): _should_log = should_log() _should_visualise = should_visualise() @@ -515,11 +686,16 @@ def run( for (factor_approx, _), (new_model_dist, status) in zip( factor_approx_optimisers, new_dist_statuses ): + raised = status.flag is StatusFlag.EXCEPTION model_approx, status = self.updater.update_model_approx( new_model_dist, factor_approx, model_approx, status ) factor = factor_approx.factor self.diagnostics.snapshot(factor, model_approx, status) + if self._check_consecutive_failures( + factor, status, max_consecutive_failures, raised=raised + ): + break if status and _should_log: self._log_factor(factor) @@ -542,5 +718,6 @@ def run( self._output_results(model_approx) self._output_diagnostics(final=True, model_approx=model_approx) self._warn_sigma_collapse() + self._warn_stale_factors() return model_approx diff --git a/autofit/graphical/expectation_propagation/stochastic.py b/autofit/graphical/expectation_propagation/stochastic.py index 0d6aeb0e8..ceb275bb0 100644 --- a/autofit/graphical/expectation_propagation/stochastic.py +++ b/autofit/graphical/expectation_propagation/stochastic.py @@ -1,6 +1,7 @@ import logging from typing import Dict, List, Generator +from autofit import exc from autofit.graphical.expectation_propagation.ep_mean_field import EPMeanField from autofit.graphical.mean_field import Status from autofit.graphical.utils import StatusFlag, LogWarnings @@ -28,13 +29,27 @@ def factor_step(self, factor, subset_approx, optimiser): ) messages = status.messages + tuple(caught_warnings.messages) - status = Status(status.success, messages, status.flag) - except (ValueError, ArithmeticError, RuntimeError) as e: + # Keyword arguments: `Status`'s third positional parameter is + # `updated`, not `flag` — see the same fix in `optimiser.factor_step`. + status = Status( + success=status.success, + messages=messages, + updated=status.updated, + flag=status.flag, + ) + except ( + ValueError, + ArithmeticError, + RuntimeError, + exc.InitializerException, + ) as e: logger.exception(e) status = Status( - False, - status.messages + (f"Factor: {factor} experienced error {e}",), - StatusFlag.FAILURE, + success=False, + messages=status.messages + + (f"Factor: {factor} experienced error {e}",), + updated=False, + flag=StatusFlag.EXCEPTION, ) factor_logger.debug(status) diff --git a/autofit/graphical/utils.py b/autofit/graphical/utils.py index b45437ca4..cd1675d0d 100644 --- a/autofit/graphical/utils.py +++ b/autofit/graphical/utils.py @@ -259,6 +259,12 @@ class StatusFlag(Enum): SUCCESS = 1 NO_CHANGE = 2 BAD_PROJECTION = 3 + # The factor's optimiser *raised* rather than returning a failed status. + # Distinct from FAILURE, which an optimiser returns routinely and which EP + # is designed to absorb (e.g. "Line search failed" from the Laplace + # optimiser). Only EXCEPTION counts toward the consecutive-failure abort in + # `EPOptimiser`. + EXCEPTION = 4 @classmethod def get_flag(cls, success, n_iter): diff --git a/autofit/non_linear/initializer.py b/autofit/non_linear/initializer.py index e768cdf76..5b5b57eed 100644 --- a/autofit/non_linear/initializer.py +++ b/autofit/non_linear/initializer.py @@ -17,6 +17,27 @@ logger = logging.getLogger(__name__) +IDENTICAL_FIGURES_OF_MERIT_MESSAGE = """ + The initial samples all have the same figure of merit (e.g. log likelihood values). + + The non-linear search will therefore not progress correctly. + + Possible causes for this behaviour are: + + - The `log_likelihood_function` of the analysis class is defined incorrectly. + - The model parameterization creates numerically inaccurate log likelihoods. + - The model is so tightly constrained that every drawn point is effectively the + same point. This is the usual cause when the search is a factor optimiser + inside an outer loop that updates its priors, e.g. expectation propagation. + + Note that this is a check for *identical* figures of merit, made with + `np.allclose`, which is `False` for `nan`. `nan` draws are also discarded + by `figure_of_metric` before they ever reach the check. An all-`nan` + `log_likelihood_function` therefore cannot raise this exception and is not + a possible cause of it. + """ + + class AbstractInitializer(ABC): @abstractmethod @@ -117,19 +138,7 @@ def samples_from_model( if total_points > 1 and np.allclose( a=figures_of_merit_list[0], b=figures_of_merit_list[1:] ): - raise exc.InitializerException( - """ - The initial samples all have the same figure of merit (e.g. log likelihood values). - - The non-linear search will therefore not progress correctly. - - Possible causes for this behaviour are: - - - The `log_likelihood_function` of the analysis class is defined incorrectly. - - The model parameterization creates numerically inaccurate log likelihoods. - - The`log_likelihood_function` is always returning `nan` values. - """ - ) + raise exc.InitializerException(IDENTICAL_FIGURES_OF_MERIT_MESSAGE) logger.info(f"Initial samples generated, starting non-linear search") @@ -182,19 +191,7 @@ def samples_jax( if total_points > 1 and np.allclose( a=figures_of_merit_list[0], b=figures_of_merit_list[1:] ): - raise exc.InitializerException( - """ - The initial samples all have the same figure of merit (e.g. log likelihood values). - - The non-linear search will therefore not progress correctly. - - Possible causes for this behaviour are: - - - The `log_likelihood_function` of the analysis class is defined incorrectly. - - The model parameterization creates numerically inaccurate log likelihoods. - - The`log_likelihood_function` is always returning `nan` values. - """ - ) + raise exc.InitializerException(IDENTICAL_FIGURES_OF_MERIT_MESSAGE) logger.info(f"Initial samples generated, starting non-linear search") diff --git a/test_autofit/graphical/functionality/test_factor_failure_recovery.py b/test_autofit/graphical/functionality/test_factor_failure_recovery.py new file mode 100644 index 000000000..5acf2ac9d --- /dev/null +++ b/test_autofit/graphical/functionality/test_factor_failure_recovery.py @@ -0,0 +1,303 @@ +""" +An `InitializerException` in one factor should not kill the whole EP fit. + +A factor's own optimiser can fail to find a start point — most often because EP +has driven that factor to a state where every drawn point has the same figure of +merit. That is a failure of one sweep's update for one factor, not of the graph +fit, so the sweep should continue on that factor's previous message with the +failure recorded. + +These tests use a **shared-variable, non-hierarchical** graph, which is the shape +that reproduced this on the release leg (PyAutoFit#1405): two factors connected by +one shared variable, no `HierarchicalFactor` involved. +""" + +import logging + +import numpy as np +import pytest + +from autofit import exc +from autofit import graphical as graph +from autofit.graphical.expectation_propagation.factor_optimiser import ( + AbstractFactorOptimiser, + ExactFactorFit, +) +from autofit.graphical.expectation_propagation.history import EPHistory +from autofit.graphical.utils import StatusFlag +from autofit.mapper.variable import Variable +from autofit.messages.normal import NormalMessage +from autofit.non_linear.paths.directory import DirectoryPaths + + +def make_shared_variable_approx(): + """ + Two factors joined by one shared variable `x` — the minimal form of the + graph that failed on the release leg (a shared prior across several + `AnalysisFactor`s), small enough to converge in a few sweeps. + """ + x = Variable("x") + prior = NormalMessage(1.0, 2.0).as_factor(x, name="prior_x") + likelihood = NormalMessage(3.0, 0.5).as_factor(x, name="like_x") + factor_graph = graph.FactorGraph([prior, likelihood]) + model_approx = graph.EPMeanField.from_approx_dists( + factor_graph, {x: NormalMessage(0.0, 10.0)} + ) + return model_approx, factor_graph, prior, likelihood + + +class InitializerFailingOptimiser(AbstractFactorOptimiser): + """ + Stands in for a per-factor search whose initializer cannot find a start + point. Fails its first `n_failures` calls, then defers to an exact fit — so + a test can model either an intermittent failure (the observed case, ~23% of + runs) or a factor that never initialises. + """ + + def __init__(self, n_failures=1): + super().__init__() + self.n_failures = n_failures + self.call_count = 0 + + def optimise(self, factor_approx, status=graph.Status()): + self.call_count += 1 + if self.call_count <= self.n_failures: + raise exc.InitializerException( + "The initial samples all have the same figure of merit" + ) + return self.exact_fit(factor_approx, status) + + +def test_initializer_exception_does_not_abort_the_fit(): + """ + The headline behaviour: one factor failing to initialise on one sweep leaves + the graph fit running, and it still returns a usable mean field. + """ + model_approx, factor_graph, prior, likelihood = make_shared_variable_approx() + + failing = InitializerFailingOptimiser(n_failures=1) + optimiser = graph.EPOptimiser( + factor_graph, + factor_optimisers={prior: failing, likelihood: ExactFactorFit()}, + paths=False, + ) + + result = optimiser.run(model_approx, max_steps=4) + + assert failing.call_count > 1, "the failing factor was never retried" + (x,) = [v for v in result.mean_field if v.name == "x"] + assert np.isfinite(result.mean_field[x].mean) + + +def test_failure_is_recorded_as_a_failure_not_a_success(): + """ + The failure must stay loud. Degrading the crash to a skipped update is only + acceptable because it is still visible — a failed step recorded as a success + is the silent-failure mode this fix exists to avoid. + """ + model_approx, factor_graph, prior, likelihood = make_shared_variable_approx() + + optimiser = graph.EPOptimiser( + factor_graph, + factor_optimisers={ + prior: InitializerFailingOptimiser(n_failures=1), + likelihood: ExactFactorFit(), + }, + paths=False, + ) + optimiser.run(model_approx, max_steps=4) + + flags = [ + row["flag"] + for row in optimiser.diagnostics.factor_rows + if row["factor"] == prior.name + ] + assert StatusFlag.EXCEPTION.name in flags, ( + "the failed factor update was not recorded as a raise in the " + f"diagnostics rows: {flags}" + ) + + +def test_persistent_failure_stops_sweeping_early(): + """ + A factor that fails *every* sweep is not going to start working, so EP stops + rather than burning the full `max_steps` on it — but it still returns. + """ + model_approx, factor_graph, prior, likelihood = make_shared_variable_approx() + + failing = InitializerFailingOptimiser(n_failures=1000) + optimiser = graph.EPOptimiser( + factor_graph, + factor_optimisers={prior: failing, likelihood: ExactFactorFit()}, + # `kl_tol=None` disables the convergence check: this graph is exact and + # would otherwise be declared converged after one sweep, before the + # failure count could build up. + ep_history=EPHistory(kl_tol=None), + paths=False, + ) + + optimiser.run(model_approx, max_steps=20, max_consecutive_failures=3) + + assert failing.call_count == 3, ( + "expected the run to stop after 3 consecutive raises, not sweep on to " + f"max_steps; got {failing.call_count} attempts" + ) + + +def test_consecutive_failure_count_resets_on_success(): + """ + Counting is per-factor and consecutive, so an intermittent failure — the + common case — never trips the abort even over many sweeps. + """ + model_approx, factor_graph, prior, likelihood = make_shared_variable_approx() + + class IntermittentOptimiser(InitializerFailingOptimiser): + def optimise(self, factor_approx, status=graph.Status()): + self.call_count += 1 + if self.call_count % 2 == 1: + raise exc.InitializerException("degenerate start point") + return self.exact_fit(factor_approx, status) + + intermittent = IntermittentOptimiser() + optimiser = graph.EPOptimiser( + factor_graph, + factor_optimisers={prior: intermittent, likelihood: ExactFactorFit()}, + ep_history=EPHistory(kl_tol=None), + paths=False, + ) + + # Alternating failure/success over many sweeps: never two failures in a row, + # so this must not abort even with a threshold of 2. + optimiser.run(model_approx, max_steps=8, max_consecutive_failures=2) + + assert intermittent.call_count > 2 + + +def test_never_updating_factor_is_warned_about_loudly(caplog): + """ + The result is returned even when no factor ever updated — but it must not be + returned quietly. + + When every factor raises, nothing in the mean field changes, so the KL step + between sweeps is zero and `EPHistory` declares convergence — in practice + within two sweeps, before any per-factor count reaches its threshold. The + mean field then holds the starting priors, and a caller reading it without + the warning would take priors for a posterior. + + The threshold here is deliberately higher than the number of sweeps that + will run, so this can only pass via the end-of-run check. + """ + model_approx, factor_graph, prior, likelihood = make_shared_variable_approx() + + optimiser = graph.EPOptimiser( + factor_graph, + factor_optimisers={ + prior: InitializerFailingOptimiser(n_failures=1000), + likelihood: InitializerFailingOptimiser(n_failures=1000), + }, + paths=False, + ) + + with caplog.at_level(logging.WARNING): + result = optimiser.run(model_approx, max_steps=2, max_consecutive_failures=100) + + assert result is not None, "the result should still be returned" + + warnings = optimiser._stale_factor_warnings() + assert len(warnings) == 1 + assert "never completed a single update" in warnings[0] + assert prior.name in warnings[0] and likelihood.name in warnings[0] + + logged = caplog.text + assert "STALE FACTORS" in logged, "the stale-factor warning was not logged" + + +def test_stale_factor_warning_is_written_to_the_diagnostics_file(tmp_path): + """ + The warning has to survive the run, not just scroll past in a log — it goes + into `ep_diagnostics.results` beside the sigma-collapse warnings. + """ + model_approx, factor_graph, prior, likelihood = make_shared_variable_approx() + + optimiser = graph.EPOptimiser( + factor_graph, + factor_optimisers={ + prior: InitializerFailingOptimiser(n_failures=1000), + likelihood: InitializerFailingOptimiser(n_failures=1000), + }, + paths=DirectoryPaths(name="stale_factors", path_prefix=str(tmp_path)), + ) + + optimiser.run(model_approx, max_steps=2, max_consecutive_failures=100) + + written = (optimiser.output_path / "ep_diagnostics.results").read_text() + assert "STALE FACTORS" in written + assert prior.name in written and likelihood.name in written + + +def test_partially_updating_factor_is_not_treated_as_stale(): + """ + The end-of-run check must stay narrow: a factor that failed at some point + but landed at least one update has a real message, and its fit is returned. + """ + model_approx, factor_graph, prior, likelihood = make_shared_variable_approx() + + optimiser = graph.EPOptimiser( + factor_graph, + factor_optimisers={ + prior: InitializerFailingOptimiser(n_failures=1), + likelihood: ExactFactorFit(), + }, + paths=False, + ) + + result = optimiser.run(model_approx, max_steps=4) + + (x,) = [v for v in result.mean_field if v.name == "x"] + assert np.isfinite(result.mean_field[x].mean) + + +def test_returned_failure_status_does_not_trip_the_abort(): + """ + Only a *raise* counts toward the abort. Optimisers return + `StatusFlag.FAILURE` routinely — the Laplace optimiser does so every time + its line search fails — and EP is designed to absorb that. Counting returned + failures aborts healthy fits, which is exactly what an earlier revision of + this guard did to `test_full_hierachical`. + """ + model_approx, factor_graph, prior, likelihood = make_shared_variable_approx() + + class AlwaysReturnsFailure(AbstractFactorOptimiser): + def optimise(self, factor_approx, status=graph.Status()): + return ( + factor_approx.model_dist, + graph.Status( + success=False, + messages=("Line search failed",), + updated=False, + flag=StatusFlag.FAILURE, + ), + ) + + optimiser = graph.EPOptimiser( + factor_graph, + factor_optimisers={prior: AlwaysReturnsFailure(), likelihood: ExactFactorFit()}, + ep_history=EPHistory(kl_tol=None), + paths=False, + ) + + # Would raise if returned failures were counted: 10 sweeps, threshold 2. + optimiser.run(model_approx, max_steps=10, max_consecutive_failures=2) + + +def test_nan_likelihoods_cannot_raise_the_identical_merit_exception(): + """ + Guards the diagnostic wording. The exception's message used to offer "always + returning `nan`" as a possible cause, which sent one investigation chasing a + nan that cannot occur: the check is `np.allclose`, which is False for `nan`, + and nan draws are discarded before it anyway. + """ + from autofit.non_linear.initializer import IDENTICAL_FIGURES_OF_MERIT_MESSAGE + + assert not np.allclose(np.nan, [np.nan, np.nan]) + assert "always returning `nan`" not in IDENTICAL_FIGURES_OF_MERIT_MESSAGE