From fc015209384795c50df942984e1a9d8259af5e78 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Fri, 10 Jul 2026 10:40:44 +0100 Subject: [PATCH] =?UTF-8?q?fix(graphical):=20EP=20statistics=20fix=20batch?= =?UTF-8?q?=20=E2=80=94=20F1/F2/F4/F8=20from=20#1332=20(#1350)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - F1/F7(a): MeanField.__truediv__/__pow__ passed log_norm positionally into the plates ctor slot — evidence silently dropped and a float stored in _plates. log_norm now passed by keyword; the per-variable- exponent __pow__ branch sets log_norm=0.0 (no meaningful scalar aggregate exists; the old self.log_norm * other.log_norm was meaningless even in the right slot). - F2 (+extended): GammaMessage.kl and BetaMessage.kl computed the reverse direction KL(dist||self); both now follow the family-wide contract self.kl(other) = KL(self||other) (matching Normal / TruncatedNormal), so EPHistory.kl_divergence no longer sums KLs measured in opposite directions on mixed graphs. Certified per family by a Monte-Carlo direction property test. - F4: AbstractMessage.update_invalid scalar branch empirically verified correct — the "fairly certain this would not work" TODO was a stale false alarm; replaced with an explanatory comment and pinned by unit tests on both branches. - F8: dead/suspect quasi-Newton variants deleted from laplace/newton.py (diag_sr1_update_ unused; diag_sr1_bfgs_update returned None; bfgs1_update sign-disputed vs the exported bfgs_update). Adds test_ep_statistics_fixes.py (9 tests). Full suite: 1465 passed / 14 skipped / 1 pre-existing flake (test_full_hierachical — fails on clean main with this batch fully stashed; filed as PyAutoMind/bug/autofit/hierarchical_ep_test_flaky.md). Co-Authored-By: Claude Opus 4.8 --- autofit/graphical/laplace/newton.py | 66 ++-------- autofit/graphical/mean_field.py | 15 ++- autofit/messages/abstract.py | 7 +- autofit/messages/beta.py | 10 +- autofit/messages/gamma.py | 17 ++- .../functionality/test_ep_statistics_fixes.py | 116 ++++++++++++++++++ 6 files changed, 160 insertions(+), 71 deletions(-) create mode 100644 test_autofit/graphical/functionality/test_ep_statistics_fixes.py diff --git a/autofit/graphical/laplace/newton.py b/autofit/graphical/laplace/newton.py index 347f32f5a..a4d3eaf6c 100644 --- a/autofit/graphical/laplace/newton.py +++ b/autofit/graphical/laplace/newton.py @@ -76,65 +76,13 @@ def diag_sr1_update( return state1 -def diag_sr1_update_( - state1: OptimisationState, state: OptimisationState, tol=1e-8, **kwargs -) -> OptimisationState: - yk = VariableData.sub(state1.gradient, state.gradient) - dk = VariableData.sub(state1.parameters, state.parameters) - Bk = state.hessian - zk = yk + Bk * dk - dzk = dk * zk - # alpha = -zk.dot(dk) / dzk.dot(dzk) - alpha = -(zk * dk).var_sum() - tols = tol * dk.var_norm() ** 2 * zk.var_norm() ** 2 - for v, d in (dzk ** 2).var_sum().items(): - if d > tols[v]: - alpha[v] /= d - else: - alpha[v] = 0.0 - - Bk = Bk.diagonalupdate((zk ** 2) * alpha) - - state1.hessian = Bk - return state1 - - -def diag_sr1_bfgs_update( - state1: OptimisationState, state: OptimisationState, **kwargs -) -> OptimisationState: - yk = VariableData.sub(state1.gradient, state.gradient) - dk = VariableData.sub(state1.parameters, state.parameters) - Bk = state.hessian - zk = yk + Bk * dk - dzk = dk * zk - - -def bfgs1_update( - state1: OptimisationState, state: OptimisationState, **kwargs, -) -> OptimisationState: - """ - y_k = g_{k+1} - g{k} - d_k = x_{k+1} - x{k} - B_{k+1} = B_{k} - + \frac - {y_{k}y_{k}^T} - {y_{k}^T d_{k}}} - - \frac - {B_{k} d_{k} (B_{k} d_{k})^T} - {d_{k}^T B_{k} d_{k}}}} - """ - yk = VariableData.sub(state.gradient, state1.gradient) - dk = VariableData.sub(state1.parameters, state.parameters) - Bk = state.hessian - - ykTdk = yk.dot(dk) - Bdk = Bk.dot(dk) - dkTBdk = -VariableData.dot(Bdk, dk) - - state1.hessian = Bk.update( - (yk, VariableData(yk).div(ykTdk)), (Bdk, VariableData(Bdk).div(dkTBdk)) - ) - return state1 +# Dead/suspect quasi-Newton variants removed (#1332 F8): +# - `diag_sr1_update_` — never referenced anywhere in the codebase. +# - `diag_sr1_bfgs_update` — computed locals then implicitly returned None +# (would have nuked the optimiser state if ever wired in). +# - `bfgs1_update` — disagreed with the exported `bfgs_update` on the sign of +# the y_k difference and carried a stray minus on d_k^T B d_k; the exported +# `bfgs_update` below is the verified-correct implementation. def bfgs_update( diff --git a/autofit/graphical/mean_field.py b/autofit/graphical/mean_field.py index d27bc91e8..99a4fd8f4 100755 --- a/autofit/graphical/mean_field.py +++ b/autofit/graphical/mean_field.py @@ -320,19 +320,28 @@ def prod(self, *approxs: "MeanField", default=None) -> "MeanField": __mul__ = prod def __truediv__(self, other: "MeanField") -> "MeanField": + # log_norm must be passed by keyword: the second positional ctor slot + # is `plates`, and passing it positionally silently dropped the + # evidence AND stored a float in `_plates` (#1332 F1). return type(self)( {k: m / other.get(k, 1.0) for k, m in self.items()}, - self.log_norm - other.log_norm, + log_norm=self.log_norm - other.log_norm, ) def __pow__(self, other: Union[float, "MeanField"]) -> "MeanField": if isinstance(other, Real): return type(self)( - {k: m ** other for k, m in self.items()}, self.log_norm * other + {k: m ** other for k, m in self.items()}, + log_norm=self.log_norm * other, ) + # Per-variable exponents: there is no meaningful scalar aggregate of + # the two mean-field log_norms (the per-message log_norms carry + # through each message's __pow__); the previous + # `self.log_norm * other.log_norm` was not a meaningful quantity and + # additionally landed in the plates slot (#1332 F1). return type(self)( {key: value ** other[key] for key, value in self.items()}, - self.log_norm * other.log_norm, + log_norm=0.0, ) def log_normalisation(self, other: "MeanField") -> float: diff --git a/autofit/messages/abstract.py b/autofit/messages/abstract.py index de2b097ed..a9779a34f 100644 --- a/autofit/messages/abstract.py +++ b/autofit/messages/abstract.py @@ -350,7 +350,12 @@ def update_invalid(self, other: "AbstractMessage") -> "AbstractMessage": np.where(valid, p, p_safe) for p, p_safe in zip(self, other) ) else: - # TODO: Fairly certain this would not work + # Scalar (ndim == 0) messages: `valid` is a scalar bool and + # messages iterate over their parameters, so this all-or-nothing + # swap is correct — keep self's parameters when valid, take + # other's when not. Empirically verified (#1332 F4); the previous + # "fairly certain this would not work" TODO was a stale false + # alarm, now pinned by unit tests on both branches. valid_parameters = iter(self if valid else other) cls = self._Base_class or type(self) new = cls( diff --git a/autofit/messages/beta.py b/autofit/messages/beta.py index b489d2251..8bb21bc22 100644 --- a/autofit/messages/beta.py +++ b/autofit/messages/beta.py @@ -343,13 +343,15 @@ def kl(self, dist: "BetaMessage") -> float: from scipy.special import betaln from scipy.special import psi - # TODO check this is correct - # https://arxiv.org/pdf/0911.4863.pdf + # Family-wide contract (#1332 F2): kl(self, dist) = KL(self || dist). + # Previously P was taken from `dist` and Q from `self`, computing the + # reverse direction KL(dist || self) despite this docstring. + # Reference: https://arxiv.org/pdf/0911.4863.pdf if self._support != dist._support: raise TypeError('Support does not match') - aP, bP = dist.parameters - aQ, bQ = self.parameters + aP, bP = self.parameters + aQ, bQ = dist.parameters return ( betaln(aQ, bQ) - betaln(aP, bP) - (aQ - aP) * psi(aP) diff --git a/autofit/messages/gamma.py b/autofit/messages/gamma.py index c744571ba..0982f1783 100644 --- a/autofit/messages/gamma.py +++ b/autofit/messages/gamma.py @@ -86,12 +86,21 @@ def from_mode(cls, mode, covariance, **kwargs): return cls(alpha, beta, **kwargs) def kl(self, dist): + """ + The Kullback-Leibler divergence KL(self || dist). + + Family-wide contract (#1332 F2): ``message.kl(other)`` returns + ``KL(message || other)``, matching ``NormalMessage.kl``. Previously + this method assigned ``P, Q = dist, self`` and returned the reverse + direction ``KL(dist || self)`` — a graph mixing Normal and Gamma + variables summed KLs measured in opposite directions inside + ``EPHistory.kl_divergence``. + + Reference: https://arxiv.org/pdf/0911.4863.pdf + """ from scipy import special - P, Q = dist, self - logP = np.log(P.alpha) - # TODO check this is correct - # https://arxiv.org/pdf/0911.4863.pdf + P, Q = self, dist return ( (P.alpha - Q.alpha) * special.psi(P.alpha) - special.gammaln(P.alpha) diff --git a/test_autofit/graphical/functionality/test_ep_statistics_fixes.py b/test_autofit/graphical/functionality/test_ep_statistics_fixes.py new file mode 100644 index 000000000..0bf20274d --- /dev/null +++ b/test_autofit/graphical/functionality/test_ep_statistics_fixes.py @@ -0,0 +1,116 @@ +"""Regression tests for the EP statistics fix batch (PyAutoFit #1350; +findings F1/F2/F4/F8 from the #1332 audit). Numpy-only. +""" +import numpy as np +import pytest + +from autofit.graphical.mean_field import MeanField +from autofit.mapper.variable import Variable +from autofit.messages.abstract import AbstractMessage +from autofit.messages.beta import BetaMessage +from autofit.messages.gamma import GammaMessage +from autofit.messages.normal import NormalMessage + + +# --- F1: MeanField __truediv__ / __pow__ keep log_norm in the log_norm slot --- + +def _mean_field(log_norm): + v = Variable("x") + return MeanField({v: NormalMessage(0.0, 1.0)}, log_norm=log_norm), v + + +def test__mean_field_truediv_preserves_log_norm(): + mf1, v = _mean_field(3.0) + mf2 = MeanField({v: NormalMessage(0.5, 2.0)}, log_norm=1.0) + out = mf1 / mf2 + # Previously: log_norm silently 0.0 and the float 2.0 stored in _plates. + assert out.log_norm == 3.0 - 1.0 + assert not isinstance(out.plates, float) + + +def test__mean_field_pow_scalar_preserves_log_norm(): + mf1, _ = _mean_field(3.0) + out = mf1 ** 2.0 + assert out.log_norm == 6.0 + assert not isinstance(out.plates, float) + + +def test__mean_field_pow_mean_field_exponent_constructs_cleanly(): + mf1, v = _mean_field(3.0) + exponents = MeanField({v: 0.5}, log_norm=2.0) + out = mf1 ** exponents + # No meaningful scalar aggregate exists for per-variable exponents; the + # previous self.log_norm * other.log_norm (= 6.0) was meaningless and + # landed in the plates slot. + assert out.log_norm == 0.0 + assert not isinstance(out.plates, float) + + +# --- F2: KL direction contract — message.kl(other) == KL(message || other), +# verified per family by Monte Carlo E_self[logpdf_self - logpdf_other] --- + +def _mc_kl(p, q, n=400_000, seed=1): + rng_state = np.random.get_state() + np.random.seed(seed) + try: + x = p.sample(n) + finally: + np.random.set_state(rng_state) + return float(np.mean(p.logpdf(x) - q.logpdf(x))) + + +@pytest.mark.parametrize( + "p, q", + [ + (NormalMessage(0.0, 1.0), NormalMessage(1.0, 2.0)), + (GammaMessage(2.0, 1.0), GammaMessage(5.0, 3.0)), + (BetaMessage(2.0, 5.0), BetaMessage(4.0, 2.0)), + ], + ids=["normal", "gamma", "beta"], +) +def test__kl_direction_contract(p, q): + analytic = float(p.kl(q)) + mc = _mc_kl(p, q) + # Direction reversal changes these asymmetric KLs by far more than the MC + # noise (~1e-2), so this test pins the KL(self || other) contract. + assert analytic == pytest.approx(mc, abs=0.05) + # And the reverse direction is genuinely different (asymmetric cases). + assert abs(float(q.kl(p)) - analytic) > 0.05 + + +# --- F4: update_invalid — both branches, scalar branch un-flagged --- + +def test__update_invalid_scalar_branch(): + good = NormalMessage(1.0, 2.0) + invalid = NormalMessage.from_natural_parameters(np.array([np.nan, -0.5])) + assert not invalid.check_valid() + repaired = invalid.update_invalid(good) + assert repaired.check_valid() + assert (repaired.mean, repaired.sigma) == (1.0, 2.0) + # valid scalar keeps itself + kept = good.update_invalid(NormalMessage(9.0, 9.0)) + assert (kept.mean, kept.sigma) == (1.0, 2.0) + + +def test__update_invalid_array_branch(): + good = NormalMessage(np.array([1.0, 1.0]), np.array([2.0, 2.0])) + mixed = NormalMessage.from_natural_parameters( + np.array([[0.5, np.nan], [-0.5, -0.5]]) + ) + valid = mixed.check_valid() + assert valid.tolist() == [True, False] + repaired = mixed.update_invalid(good) + assert repaired.check_valid().all() + # element 0 kept, element 1 replaced by good's parameters + assert repaired.mean[1] == 1.0 and repaired.sigma[1] == 2.0 + + +# --- F8: dead quasi-Newton variants are gone; exported ones remain --- + +def test__dead_newton_variants_removed(): + from autofit.graphical.laplace import newton + + for dead in ("diag_sr1_update_", "diag_sr1_bfgs_update", "bfgs1_update"): + assert not hasattr(newton, dead) + for alive in ("bfgs_update", "sr1_update", "full_bfgs_update"): + assert hasattr(newton, alive)