Report a non-finite EP evidence term by name instead of warn-and-sum (#1713)

autofit/graphical/expectation_propagation/diagnostics.py

@@ -84,6 +84,9 @@ def __init__(self):
         self.scale_variables: Set[str] = set()
         self._previous_mean_field = None
         self._step = 0
+        # {term label: (first step, kind)} for every non-finite term of the
+        # log evidence seen during the run (PyAutoFit#1713).
+        self.evidence_nonfinite: Dict[str, Tuple[int, str]] = {}
 
     def register_hierarchical_scales(self, factor_graph) -> None:
         """
@@ -133,6 +136,18 @@ def snapshot(self, factor, model_approx, status) -> None:
             logger.warning(f"EPDiagnostics: log_evidence failed: {e}")
             log_evidence = float("nan")
 
+        # A non-finite evidence is returned as NaN by `EPMeanField.log_evidence`;
+        # record which terms caused it, and from which step, for
+        # `check_evidence_nonfinite` (PyAutoFit#1713).
+        if not np.isfinite(log_evidence):
+            try:
+                terms = model_approx.nonfinite_evidence_terms
+            except Exception as e:
+                logger.warning(f"EPDiagnostics: evidence term scan failed: {e}")
+                terms = {}
+            for label, kind in terms.items():
+                self.evidence_nonfinite.setdefault(label, (self._step, kind))
+
         if self._previous_mean_field is not None:
             try:
                 kl = float(np.sum(mean_field.kl(self._previous_mean_field)))
@@ -246,6 +261,39 @@ def plot(self, output_path: Path) -> None:
         plt.close(fig)
 
 
+def check_evidence_nonfinite(diagnostics: EPDiagnostics) -> List[str]:
+    """
+    Report the non-finite terms of the EP log evidence seen during a run.
+
+    ``EPMeanField.log_evidence`` returns NaN, and logs once, when a variable's
+    log evidence or a factor's ``log_norm`` is NaN or ±inf (PyAutoFit#1713).
+    This turns the terms ``EPDiagnostics.snapshot`` recorded into a single
+    ``evidence_nonfinite: <names>`` line for ``ep_diagnostics.results``, so a
+    batch run keeps a named trace of it. The posteriors are unaffected; the
+    line says that the ``log_evidence`` column of ``ep_history.csv`` is NaN
+    from the step named.
+
+    Returns
+    -------
+    A list holding that one line, or an empty list when every evidence was
+    finite.
+    """
+    terms = getattr(diagnostics, "evidence_nonfinite", {})
+    if not terms:
+        return []
+    names = ", ".join(
+        f"{label} ({kind}, first at step {step})"
+        for label, (step, kind) in sorted(terms.items())
+    )
+    return [
+        f"evidence_nonfinite: {names} — the EP log evidence is NaN from that "
+        f"step on (ep_history.csv log_evidence column). A message's "
+        f"log-partition is non-finite in floating point; the mean field and "
+        f"its posteriors are unaffected, only the evidence readout is "
+        f"(PyAutoFit#1713)."
+    ]
+
+
 def mean_field_summary(mean_field) -> str:
     """
     A human-readable summary table of a mean field.

autofit/graphical/expectation_propagation/ep_mean_field.py

@@ -1,4 +1,5 @@
 import logging
+import weakref
 from typing import Dict, Tuple, Optional, List
 
 import numpy as np
@@ -13,6 +14,14 @@
 
 logger = logging.getLogger(__name__)
 
+#: Non-finite evidence terms already reported by ``EPMeanField.log_evidence``,
+#: keyed weakly on the factor graph (one fit = one graph, shared by every mean
+#: field the EP updates build over it) with the set of reported
+#: ``(label, kind)`` sets as the value — so each is logged once per fit rather
+#: than on every EP update, and nothing is kept alive once the fit's graph is
+#: gone (PyAutoFit#1713).
+_reported_nonfinite = weakref.WeakKeyDictionary()
+
 
 class EPMeanField(FactorGraph):
     """
@@ -306,6 +315,61 @@ def factor_evidence(self) -> Dict[Factor, np.ndarray]:
             for factor, meanfield in self.factor_mean_field.items()
         }
 
+    def _evidence_terms(self) -> Tuple[Dict[Variable, float], Dict[Factor, float]]:
+        """
+        The per-variable log evidences ``log Zᵢ`` and per-factor ``log_norm``
+        sums that ``log_evidence`` combines.
+
+        Evaluated under ``np.errstate`` so a message whose log-partition is
+        non-finite does not emit numpy ``RuntimeWarning``s here; such a term
+        is reported by name by ``log_evidence`` instead.
+        """
+        with np.errstate(invalid="ignore", over="ignore", divide="ignore"):
+            variable_evidence = {
+                v: np.sum(logz) for v, logz in self.variable_evidence.items()
+            }
+            factor_log_norm = {
+                factor: np.sum(meanfield.log_norm)
+                for factor, meanfield in self.factor_mean_field.items()
+            }
+        return variable_evidence, factor_log_norm
+
+    @staticmethod
+    def _nonfinite_terms(variable_evidence, factor_log_norm) -> Dict:
+        """
+        ``{term: kind}`` for every non-finite evidence term, keyed by the
+        ``Variable`` or ``Factor`` itself; ``kind`` is ``"nan"``, ``"+inf"``
+        or ``"-inf"``.
+        """
+        nonfinite = {}
+        for term, value in list(variable_evidence.items()) + list(
+            factor_log_norm.items()
+        ):
+            if not np.isfinite(value):
+                nonfinite[term] = (
+                    "nan" if np.isnan(value) else ("+inf" if value > 0 else "-inf")
+                )
+        return nonfinite
+
+    @staticmethod
+    def _term_label(term) -> str:
+        if isinstance(term, Factor):
+            return f"factor '{term.name}' log_norm"
+        return f"variable '{term.name}'"
+
+    @property
+    def nonfinite_evidence_terms(self) -> Dict[str, str]:
+        """
+        The non-finite terms of ``log_evidence`` as ``{label: kind}``, e.g.
+        ``{"variable 'sigma'": "+inf"}``; empty when the evidence is finite.
+
+        A variable's term is its log evidence ``log Zᵢ`` (the log
+        normalisation of the product of its messages); a factor's is its mean
+        field's ``log_norm``.
+        """
+        nonfinite = self._nonfinite_terms(*self._evidence_terms())
+        return {self._term_label(term): kind for term, kind in nonfinite.items()}
+
     @property
     def log_evidence(self):
         """
@@ -324,16 +388,67 @@ def log_evidence(self):
         Evidence for model
 
         Z = ∏ᵢ Zᵢ ∏ₐ Zₐ
+
+        Non-finite terms
+        ----------------
+        If any variable's ``log Zᵢ`` or any factor's ``log_norm`` is NaN or
+        ±inf the sum is undefined (a ``+inf`` variable term and the ``-inf``
+        it induces in every adjacent factor term add to NaN). The evidence is
+        then returned as an explicit ``float("nan")`` without numpy
+        ``RuntimeWarning``s, and the offending terms are named in one
+        ``logger.warning`` per distinct set of terms (so once per fit, not
+        once per EP update). The usual cause is a message whose
+        log-partition is non-finite in floating point — e.g. a truncated
+        message whose mass inside its support underflows to zero
+        (PyAutoFit#1713). Only the evidence readout is affected: the mean
+        field and its posteriors are computed without it.
         """
-        variable_evidence = {
-            v: np.sum(logz) for v, logz in self.variable_evidence.items()
-        }
-        factor_evidence = sum(
-            np.sum(meanfield.log_norm)
-            - sum(variable_evidence[v] for v in factor.all_variables)
-            for factor, meanfield in self.factor_mean_field.items()
+        variable_evidence, factor_log_norm = self._evidence_terms()
+
+        nonfinite = self._nonfinite_terms(variable_evidence, factor_log_norm)
+        if nonfinite:
+            self._warn_nonfinite(nonfinite)
+            return float("nan")
+
+        with np.errstate(invalid="ignore", over="ignore"):
+            factor_evidence = sum(
+                factor_log_norm[factor]
+                - sum(variable_evidence[v] for v in factor.all_variables)
+                for factor in self.factor_mean_field
+            )
+            return factor_evidence + sum(variable_evidence.values())
+
+    def _warn_nonfinite(self, nonfinite: Dict) -> None:
+        """
+        Log the non-finite evidence terms once per distinct set of terms.
+
+        Every EP update builds a new ``EPMeanField`` over the same graph, and
+        the history re-evaluates past evidences, so the key is the set of
+        offending (term, kind) pairs rather than the instance.
+        """
+        key = frozenset(
+            (self._term_label(term), kind) for term, kind in nonfinite.items()
+        )
+        try:
+            reported = _reported_nonfinite.setdefault(self.factor_graph, set())
+        except TypeError:  # a graph that cannot be weakly referenced
+            reported = set()
+        if key in reported:
+            return
+        reported.add(key)
+        terms = ", ".join(
+            sorted(
+                f"{self._term_label(term)} ({kind})"
+                for term, kind in nonfinite.items()
+            )
+        )
+        logger.warning(
+            f"EP log evidence is non-finite and is reported as NaN: {terms}. "
+            f"A message's log-partition is non-finite in floating point (e.g. a "
+            f"truncated message whose mass inside its support underflows). The "
+            f"mean field and its posteriors are unaffected; only the evidence "
+            f"readout is. Reported once per fit (PyAutoFit#1713)."
         )
-        return factor_evidence + sum(variable_evidence.values())
 
     def __repr__(self) -> str:
         clsname = type(self).__name__

autofit/graphical/expectation_propagation/history.py

@@ -146,11 +146,17 @@ def evidence_divergence(self) -> Union[float, np.ndarray]:
         successful optimisations.
 
         If there are less than two successful optimisations then this is
-        infinite.
+        infinite. It is also infinite when either evidence is NaN (a
+        non-finite evidence term, see ``EPMeanField.log_evidence``), so an
+        evidence-based stopping rule can never claim convergence on a NaN
+        (PyAutoFit#1713).
         """
-        return (
+        divergence = (
                 self.latest_successful.log_evidence - self.previous_successful.log_evidence
         )
+        if np.any(np.isnan(divergence)):
+            return float("inf")
+        return divergence
 
     @property
     def evidences(self) -> List[float]:

autofit/graphical/expectation_propagation/optimiser.py

@@ -15,7 +15,12 @@
 from autofit.non_linear.paths import DirectoryPaths
 from autofit.non_linear.paths.abstract import AbstractPaths
 from autofit.tools.util import IntervalCounter
-from .diagnostics import EPDiagnostics, check_sigma_collapse, mean_field_summary
+from .diagnostics import (
+    EPDiagnostics,
+    check_evidence_nonfinite,
+    check_sigma_collapse,
+    mean_field_summary,
+)
 from .factor_optimiser import AbstractFactorOptimiser, ExactFactorFit
 from .visualise import Visualise
 
@@ -653,7 +658,8 @@ def _output_diagnostics(
         Write the diagnostics CSVs (`ep_history.csv`, `mean_field_history.csv`)
         and the mean-field evolution plot to the output folder. At the end of
         a run (``final=True``) also write ``ep_diagnostics.results``: the
-        mean-field summary table plus any sigma-collapse warnings.
+        mean-field summary table plus any stale-factor, sigma-collapse and
+        non-finite-evidence warnings.
         """
         if not self.output_path:
             return
@@ -662,7 +668,9 @@ def _output_diagnostics(
 
         if final and model_approx is not None:
             warnings_list = (
-                self._stale_factor_warnings() + check_sigma_collapse(self.diagnostics)
+                self._stale_factor_warnings()
+                + check_sigma_collapse(self.diagnostics)
+                + check_evidence_nonfinite(self.diagnostics)
             )
             with open(self.output_path / "ep_diagnostics.results", "w+") as f:
                 f.write(mean_field_summary(model_approx.mean_field))