From e10b1101394b053fbfa76fab9eb79f02490f866b Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Fri, 10 Jul 2026 10:06:10 +0100 Subject: [PATCH] feat(graphical): EP diagnostics and monitoring tooling (#1335) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit EP framework review Phase 4 — built-in tools for analysing a finished EP fit and monitoring a running one: - expectation_propagation/diagnostics.py (new): EPDiagnostics collector (per-factor-update snapshot: status, global log-evidence, KL step, every variable's mean/std); live-overwritten ep_history.csv + mean_field_history.csv; mean_field_evolution.png; post-fit mean_field_summary() table; check_sigma_collapse() F10 guard (floor + monotone-shrink detection) - optimiser.py: EPOptimiser/ParallelEPOptimiser snapshot each update, emit artifacts at output_interval + run end; F3 fixed (parallel end-of-run visualiser()/_output_results guarded for paths=None) - visualise.py: graph_factors.png per-factor evidence/KL curves (replaces dead commented-out code) - exports: EPDiagnostics, mean_field_summary, check_sigma_collapse from autofit.graphical (additive) - test_diagnostics.py: 7 tests (snapshot schema, artifacts, summary, collapse detector, F3 guards) Co-Authored-By: Claude Opus 4.8 --- autofit/graphical/__init__.py | 1 + .../expectation_propagation/__init__.py | 1 + .../expectation_propagation/diagnostics.py | 264 ++++++++++++++++++ .../expectation_propagation/optimiser.py | 52 +++- .../expectation_propagation/visualise.py | 35 ++- .../functionality/test_diagnostics.py | 202 ++++++++++++++ 6 files changed, 544 insertions(+), 11 deletions(-) create mode 100644 autofit/graphical/expectation_propagation/diagnostics.py create mode 100644 test_autofit/graphical/functionality/test_diagnostics.py diff --git a/autofit/graphical/__init__.py b/autofit/graphical/__init__.py index a3e6ae235..f369c77c8 100644 --- a/autofit/graphical/__init__.py +++ b/autofit/graphical/__init__.py @@ -3,6 +3,7 @@ from .declarative.collection import FactorGraphModel from .declarative.factor.analysis import AnalysisFactor, EPAnalysisFactor from .declarative.factor.hierarchical import _HierarchicalFactor, HierarchicalFactor +from .expectation_propagation.diagnostics import EPDiagnostics, check_sigma_collapse, mean_field_summary from .expectation_propagation.ep_mean_field import EPMeanField from .expectation_propagation.optimiser import EPOptimiser from .expectation_propagation import StochasticEPOptimiser diff --git a/autofit/graphical/expectation_propagation/__init__.py b/autofit/graphical/expectation_propagation/__init__.py index 83faa5902..07b62ea09 100644 --- a/autofit/graphical/expectation_propagation/__init__.py +++ b/autofit/graphical/expectation_propagation/__init__.py @@ -1,3 +1,4 @@ +from .diagnostics import EPDiagnostics, check_sigma_collapse, mean_field_summary from .ep_mean_field import EPMeanField from .history import FactorHistory, EPHistory from .optimiser import AbstractFactorOptimiser diff --git a/autofit/graphical/expectation_propagation/diagnostics.py b/autofit/graphical/expectation_propagation/diagnostics.py new file mode 100644 index 000000000..6bd3d3a18 --- /dev/null +++ b/autofit/graphical/expectation_propagation/diagnostics.py @@ -0,0 +1,264 @@ +""" +Diagnostics and monitoring for expectation-propagation fits. + +An EP fit is a sequence of per-factor updates to a global mean-field +approximation. This module records that sequence so it can be inspected +while the fit runs and analysed after it finishes: + +- ``EPDiagnostics`` — collects one snapshot per factor update (status, + factor log-evidence, KL step size, and every variable's mean/std) and + writes them as machine-readable CSVs and a matplotlib evolution plot. +- ``mean_field_summary`` — a human-readable table of a mean field, + suitable for printing at the end of any example or script. +- ``check_sigma_collapse`` — guards against the known pathology where + repeated undamped EP updates over-count shared information and every + sigma collapses towards zero around the starting point (rather than + the data); see PyAutoFit issue #1332 (F10). + +Outputs written to the EP output folder by ``EPOptimiser`` when paths +are enabled: + +- ``ep_history.csv`` — one row per factor update: + ``step, factor, success, updated, flag, log_evidence, kl_divergence`` +- ``mean_field_history.csv`` — one row per (factor update, variable): + ``step, factor, variable, mean, std`` +- ``mean_field_evolution.png`` — per-variable mean ± std vs update. +""" +import csv +import logging +from pathlib import Path +from typing import Dict, List, Optional, Tuple + +import numpy as np + +logger = logging.getLogger(__name__) + + +def _scalar_mean_std(message) -> Tuple[float, float]: + """ + Reduce a message's mean/std to scalars for logging. + + Scalar variables pass through; array (plated) variables are + summarised as the mean of their means and the maximum of their + stds — the max is chosen so a single collapsing element is not + hidden by averaging. + + The std is derived from ``variance``, which every message family + exposes (``TransformedMessage`` — e.g. any ``UniformPrior``-backed + variable — has no ``std`` attribute). + """ + mean = np.asarray(message.mean, dtype=float) + std = np.sqrt(np.asarray(message.variance, dtype=float)) + return float(mean.mean()), float(std.max()) + + +class EPDiagnostics: + def __init__(self): + """ + Collects a per-factor-update record of an EP fit. + + ``EPOptimiser`` calls ``snapshot`` after every factor update. + Each snapshot stores the update's status, the current global + log-evidence, the KL divergence of the global mean field from + the previous snapshot (the "step size" EP convergence watches), + and every variable's (mean, std). + """ + self.factor_rows: List[dict] = [] + self.variable_rows: List[dict] = [] + self._previous_mean_field = None + self._step = 0 + + def snapshot(self, factor, model_approx, status) -> None: + """ + Record the state of the approximation after one factor update. + + Parameters + ---------- + factor + The factor that was just updated. + model_approx + The ``EPMeanField`` after the update. + status + The update's ``Status``. + """ + mean_field = model_approx.mean_field + + try: + log_evidence = float(model_approx.log_evidence) + except Exception as e: # diagnostics must never kill the fit + logger.warning(f"EPDiagnostics: log_evidence failed: {e}") + log_evidence = float("nan") + + if self._previous_mean_field is not None: + try: + kl = float(np.sum(mean_field.kl(self._previous_mean_field))) + except Exception as e: + logger.warning(f"EPDiagnostics: kl computation failed: {e}") + kl = float("nan") + else: + kl = float("nan") + + self.factor_rows.append( + { + "step": self._step, + "factor": factor.name, + "success": status.success, + "updated": status.updated, + "flag": status.flag.name, + "log_evidence": log_evidence, + "kl_divergence": kl, + } + ) + + for variable, message in mean_field.items(): + mean, std = _scalar_mean_std(message) + self.variable_rows.append( + { + "step": self._step, + "factor": factor.name, + "variable": variable.name, + "mean": mean, + "std": std, + } + ) + + self._previous_mean_field = mean_field + self._step += 1 + + @property + def variable_history(self) -> Dict[str, List[Tuple[int, float, float]]]: + """ + The recorded history keyed by variable name: a list of + (step, mean, std) tuples in update order. + """ + history: Dict[str, List[Tuple[int, float, float]]] = {} + for row in self.variable_rows: + history.setdefault(row["variable"], []).append( + (row["step"], row["mean"], row["std"]) + ) + return history + + def write(self, output_path: Path) -> None: + """ + Write ``ep_history.csv`` and ``mean_field_history.csv`` to the + output folder (overwriting — the CSVs always reflect the full + history so far, so they can be watched during a run). + """ + output_path = Path(output_path) + output_path.mkdir(parents=True, exist_ok=True) + + for filename, rows in ( + ("ep_history.csv", self.factor_rows), + ("mean_field_history.csv", self.variable_rows), + ): + if not rows: + continue + with open(output_path / filename, "w", newline="") as f: + writer = csv.DictWriter(f, fieldnames=list(rows[0].keys())) + writer.writeheader() + writer.writerows(rows) + + def plot(self, output_path: Path) -> None: + """ + Write ``mean_field_evolution.png``: each variable's mean with a + ± std band against factor-update number. + """ + if not self.variable_rows: + return + + import matplotlib.pyplot as plt + + history = self.variable_history + + fig, ax = plt.subplots(figsize=(10, 6)) + for name, rows in history.items(): + steps, means, stds = map(np.array, zip(*rows)) + (line,) = ax.plot(steps, means, label=name) + ax.fill_between( + steps, means - stds, means + stds, alpha=0.2, color=line.get_color() + ) + ax.set_xlabel("factor update") + ax.set_ylabel("posterior mean ± std") + ax.set_title("Mean-field evolution") + ax.legend(fontsize="small") + fig.savefig(str(Path(output_path) / "mean_field_evolution.png")) + plt.close(fig) + + +def mean_field_summary(mean_field) -> str: + """ + A human-readable summary table of a mean field. + + One row per variable: name, posterior mean and std (arrays + summarised as mean-of-means / max-std). Call this at the end of an + EP example to display the approximate posterior: + + print(mean_field_summary(result.updated_ep_mean_field.mean_field)) + """ + rows = [] + for variable, message in mean_field.items(): + mean, std = _scalar_mean_std(message) + rows.append((variable.name, mean, std)) + + name_width = max([len(name) for name, _, _ in rows] + [len("variable")]) + lines = [ + f"{'variable':<{name_width}} {'mean':>14} {'std':>14}", + "-" * (name_width + 32), + ] + for name, mean, std in rows: + lines.append(f"{name:<{name_width}} {mean:>14.6g} {std:>14.6g}") + return "\n".join(lines) + + +def check_sigma_collapse( + diagnostics: EPDiagnostics, + std_floor: float = 1e-8, + monotone_steps: int = 5, + shrink_factor: float = 1e-3, +) -> List[str]: + """ + Detect the EP sigma-collapse pathology (PyAutoFit #1332, F10). + + Repeated undamped EP updates can over-count shared-variable + information: every std shrinks monotonically towards zero around + the *starting* means, while the KL convergence criterion never + triggers. This check flags a variable when either: + + - its latest std is below ``std_floor``, or + - its std has shrunk monotonically for the last ``monotone_steps`` + updates *and* by more than a factor ``1 / shrink_factor`` overall. + + Returns + ------- + A list of warning strings, one per flagged variable (empty when + healthy). ``EPOptimiser`` logs these and appends them to the + results text at the end of a run. + """ + warnings_list = [] + + for name, rows in diagnostics.variable_history.items(): + stds = np.array([std for _, _, std in rows], dtype=float) + if len(stds) == 0: + continue + + if stds[-1] < std_floor: + warnings_list.append( + f"sigma-collapse: variable '{name}' has std {stds[-1]:.3g} " + f"below the floor {std_floor:.1g} — the fit has likely " + f"collapsed to a point (see PyAutoFit #1332 F10; consider " + f"damping, e.g. delta < 1, or per-factor sampler optimisers)." + ) + continue + + if len(stds) > monotone_steps: + tail = stds[-(monotone_steps + 1):] + if np.all(np.diff(tail) < 0) and stds[-1] < shrink_factor * stds[0]: + warnings_list.append( + f"sigma-collapse: variable '{name}' std has shrunk " + f"monotonically over the last {monotone_steps} updates to " + f"{stds[-1]:.3g} ({stds[-1] / stds[0]:.1e} of its initial " + f"value) — possible information over-counting (PyAutoFit " + f"#1332 F10)." + ) + + return warnings_list diff --git a/autofit/graphical/expectation_propagation/optimiser.py b/autofit/graphical/expectation_propagation/optimiser.py index 97ccc048c..bfbe292d9 100644 --- a/autofit/graphical/expectation_propagation/optimiser.py +++ b/autofit/graphical/expectation_propagation/optimiser.py @@ -16,6 +16,7 @@ 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 .factor_optimiser import AbstractFactorOptimiser, ExactFactorFit from .visualise import Visualise @@ -206,6 +207,7 @@ def __init__( self.ep_history = ep_history or EPHistory() + self.diagnostics = EPDiagnostics() self.visualiser = None if paths is None: @@ -340,6 +342,7 @@ def run( model_approx, status = self.updater.update_model_approx( new_model_dist, factor_approx, model_approx, status ) + self.diagnostics.snapshot(factor, model_approx, status) if status and _should_log: self._log_factor(factor) @@ -352,12 +355,15 @@ def run( self.visualiser() if self.output_path and _should_output: self._output_results(model_approx) + self._output_diagnostics() continue break # stop iterations if self.paths: self.visualiser() self._output_results(model_approx) + self._output_diagnostics(final=True, model_approx=model_approx) + self._warn_sigma_collapse() return model_approx @@ -369,6 +375,38 @@ def _output_results(self, model_approx: EPMeanField): with open(self.output_path / "graph.results", "w+") as f: f.write(self.factor_graph.make_results_text(model_approx)) + def _output_diagnostics( + self, final: bool = False, model_approx: Optional[EPMeanField] = None + ): + """ + 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. + """ + if not self.output_path: + return + self.diagnostics.write(self.output_path) + self.diagnostics.plot(self.output_path) + + if final and model_approx is not None: + warnings_list = 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") + if warnings_list: + f.write("\nWARNINGS\n\n") + f.write("\n".join(warnings_list)) + f.write("\n") + + def _warn_sigma_collapse(self): + """ + Log any sigma-collapse warnings (PyAutoFit #1332 F10) at the end of a + run — emitted regardless of whether output paths are enabled. + """ + for warning in check_sigma_collapse(self.diagnostics): + logger.warning(warning) + class ParallelEPOptimiser(EPOptimiser): def __init__( @@ -480,6 +518,7 @@ def run( new_model_dist, factor_approx, model_approx, status ) factor = factor_approx.factor + self.diagnostics.snapshot(factor, model_approx, status) if status and _should_log: self._log_factor(factor) @@ -488,14 +527,19 @@ def run( break # callback controls convergence else: # If no break do next iteration - if _should_visualise: + if self.visualiser and _should_visualise: self.visualiser() - if _should_output: + if self.output_path and _should_output: self._output_results(model_approx) + self._output_diagnostics() continue break # stop iterations - self.visualiser() - self._output_results(model_approx) + if self.visualiser: + self.visualiser() + if self.output_path: + self._output_results(model_approx) + self._output_diagnostics(final=True, model_approx=model_approx) + self._warn_sigma_collapse() return model_approx diff --git a/autofit/graphical/expectation_propagation/visualise.py b/autofit/graphical/expectation_propagation/visualise.py index 443ee72c3..8332e08d7 100644 --- a/autofit/graphical/expectation_propagation/visualise.py +++ b/autofit/graphical/expectation_propagation/visualise.py @@ -38,16 +38,37 @@ def __call__(self): fig.suptitle("Evidence and KL Divergence") evidence_plot.plot(self.ep_history.evidences(), label="evidence") kl_plot.semilogy(self.ep_history.kl_divergences(), label="KL divergence") - # for factor, factor_history in self.ep_history.items(): - # evidence_plot.plot( - # factor_history.evidences, label=f"{factor.name} evidence" - # ) - # kl_plot.plot( - # factor_history.kl_divergences, label=f"{factor.name} divergence" - # ) evidence_plot.legend() kl_plot.legend() with warnings.catch_warnings(): warnings.simplefilter("ignore") plt.savefig(str(self.output_path / "graph.png")) plt.close() + + self.plot_factors() + + def plot_factors(self): + """ + Save `graph_factors.png`: each factor's evidence and KL-divergence + history on its own curve, so a single misbehaving factor (failing + fits, oscillating KL) is visible instead of being averaged into the + global curves of `graph.png`. + """ + import matplotlib.pyplot as plt + + fig, (evidence_plot, kl_plot) = plt.subplots(2) + fig.suptitle("Per-factor Evidence and KL Divergence") + for factor, factor_history in self.ep_history.items(): + evidence_plot.plot( + factor_history.evidences, label=f"{factor.name}" + ) + kl_plot.plot( + factor_history.kl_divergences, label=f"{factor.name}" + ) + kl_plot.set_yscale("log") + evidence_plot.legend(fontsize="small") + kl_plot.legend(fontsize="small") + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + plt.savefig(str(self.output_path / "graph_factors.png")) + plt.close() diff --git a/test_autofit/graphical/functionality/test_diagnostics.py b/test_autofit/graphical/functionality/test_diagnostics.py new file mode 100644 index 000000000..b69d02c77 --- /dev/null +++ b/test_autofit/graphical/functionality/test_diagnostics.py @@ -0,0 +1,202 @@ +import csv + +import numpy as np +import pytest + +from autofit import graphical as graph +from autofit.graphical.expectation_propagation.diagnostics import EPDiagnostics +from autofit.mapper.variable import Variable +from autofit.messages.normal import NormalMessage +from autofit.non_linear.paths.directory import DirectoryPaths + + +def make_model_approx(): + """ + A tiny exact-conjugate normal/normal EP graph (prior_x * like_x on a + single scalar variable x). Cheap enough to run to convergence in a + handful of factor updates, which is all these diagnostics tests need. + """ + 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") + fg = graph.FactorGraph([prior, likelihood]) + model_approx = graph.EPMeanField.from_approx_dists(fg, {x: NormalMessage(0.0, 10.0)}) + return model_approx, x + + +def test_snapshot_records_rows(): + model_approx, x = make_model_approx() + + opt = graph.EPOptimiser.from_meanfield(model_approx, paths=False) + opt.run(model_approx, max_steps=4) + + factor_rows = opt.diagnostics.factor_rows + assert factor_rows + + expected_columns = { + "step", + "factor", + "success", + "updated", + "flag", + "log_evidence", + "kl_divergence", + } + for row in factor_rows: + assert set(row.keys()) == expected_columns + + variable_rows = opt.diagnostics.variable_rows + assert variable_rows + + # one variable_row per (step, variable) -- here there is a single + # variable "x", so the count matches the number of factor updates + steps = [row["step"] for row in variable_rows] + assert len(steps) == len(set(steps)) + assert len(variable_rows) == len(factor_rows) + + # kl_divergence lives on factor_rows, not variable_rows -- first step + # has no previous mean field to diverge from, so it is NaN. + assert np.isnan(factor_rows[0]["kl_divergence"]) + assert all(np.isfinite(row["kl_divergence"]) for row in factor_rows[1:]) + + +def test_csv_outputs_written(tmp_path): + model_approx, x = make_model_approx() + + paths = DirectoryPaths(name="ep_diag_test", path_prefix=str(tmp_path)) + opt = graph.EPOptimiser.from_meanfield(model_approx, paths=paths) + opt.run(model_approx, max_steps=4) + + output_path = opt.output_path + assert output_path is not None + + ep_history_path = output_path / "ep_history.csv" + mean_field_history_path = output_path / "mean_field_history.csv" + assert ep_history_path.exists() + assert mean_field_history_path.exists() + + with open(ep_history_path, newline="") as f: + reader = csv.DictReader(f) + assert reader.fieldnames == [ + "step", + "factor", + "success", + "updated", + "flag", + "log_evidence", + "kl_divergence", + ] + ep_rows = list(reader) + assert len(ep_rows) == len(opt.diagnostics.factor_rows) + + with open(mean_field_history_path, newline="") as f: + reader = csv.DictReader(f) + assert reader.fieldnames == ["step", "factor", "variable", "mean", "std"] + mf_rows = list(reader) + assert len(mf_rows) == len(opt.diagnostics.variable_rows) + + assert (output_path / "mean_field_evolution.png").exists() + assert (output_path / "graph_factors.png").exists() + + results_path = output_path / "ep_diagnostics.results" + assert results_path.exists() + results_text = results_path.read_text() + assert "x" in results_text + assert "WARNINGS" not in results_text + + +def test_mean_field_summary(): + model_approx, x = make_model_approx() + + opt = graph.EPOptimiser.from_meanfield(model_approx, paths=False) + result = opt.run(model_approx, max_steps=4) + + summary = graph.mean_field_summary(result.mean_field) + assert "variable" in summary + assert "x" in summary + + message = result.mean_field[x] + assert f"{message.mean:.6g}" in summary + assert f"{message.std:.6g}" in summary + + +def test_sigma_collapse_floor(): + diagnostics = EPDiagnostics() + diagnostics.variable_rows = [ + {"step": 0, "factor": "f", "variable": "collapsed", "mean": 0.0, "std": 1e-12}, + ] + + warnings_list = graph.check_sigma_collapse(diagnostics) + + assert len(warnings_list) == 1 + assert "collapsed" in warnings_list[0] + assert "floor" in warnings_list[0] + + +def test_sigma_collapse_monotone(): + diagnostics = EPDiagnostics() + + shrinking_stds = np.geomspace(1.0, 1e-5, num=8) + healthy_stds = [1.0, 1.05, 0.95, 1.02, 0.98, 1.01, 0.99, 1.0] + + rows = [] + for step, std in enumerate(shrinking_stds): + rows.append( + {"step": step, "factor": "f", "variable": "shrinking", "mean": 0.0, "std": float(std)} + ) + for step, std in enumerate(healthy_stds): + rows.append( + {"step": step, "factor": "f", "variable": "healthy", "mean": 0.0, "std": float(std)} + ) + diagnostics.variable_rows = rows + + warnings_list = graph.check_sigma_collapse(diagnostics) + + shrinking_warnings = [w for w in warnings_list if "shrinking" in w] + healthy_warnings = [w for w in warnings_list if "healthy" in w] + + assert len(shrinking_warnings) == 1 + assert "monotonically" in shrinking_warnings[0] + assert healthy_warnings == [] + + +def test_no_collapse_on_healthy_run(): + model_approx, x = make_model_approx() + + opt = graph.EPOptimiser.from_meanfield(model_approx, paths=False) + opt.run(model_approx, max_steps=4) + + assert graph.check_sigma_collapse(opt.diagnostics) == [] + + +def test_parallel_end_of_run_guards(): + """ + F3: a real ParallelEPOptimiser run needs n_cores >= 3 and spins up a + multiprocessing.Pool. The only existing ParallelEPOptimiser test in + this repo, `_test_parallel_laplace` in + test_autofit/graphical/regression/test_linear_regression.py, is + prefixed with an underscore so pytest never collects it -- i.e. a + real parallel run is already established as impractical to exercise + routinely under this test suite (pool startup cost / environment + flakiness), so we don't add a second one here. + + Instead we verify directly the end-of-run diagnostics guards that + ParallelEPOptimiser.run shares with EPOptimiser.run (both call + self._output_diagnostics(...) and self._warn_sigma_collapse() after + the main loop): with paths=False, self.visualiser stays None and + self.output_path stays None, and both calls must be no-ops rather + than raising. + """ + model_approx, x = make_model_approx() + + opt = graph.EPOptimiser.from_meanfield(model_approx, paths=False) + + assert opt.visualiser is None + assert opt.output_path is None + + # no snapshots taken yet -- factor_rows/variable_rows are empty, which + # is the same state ParallelEPOptimiser.run's guards must tolerate + # before touching self.output_path / self.visualiser. + opt._output_diagnostics() + opt._output_diagnostics(final=True, model_approx=model_approx) + opt._warn_sigma_collapse()