From e98717af69645fafe7f29724b72b757e325f689b Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Fri, 10 Jul 2026 11:18:40 +0100 Subject: [PATCH] =?UTF-8?q?fix(graphical):=20EP=20statistics=20completion?= =?UTF-8?q?=20=E2=80=94=20F6=20truncated=20KL=20+=20F7(b)=20evidence=20rec?= =?UTF-8?q?ording=20(#1353)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Closes out the #1332 audit's statistical findings: - F6: TruncatedNormalMessage.kl now computes the exact same-support KL via truncated moments (scipy truncnorm.stats); reduces to the untruncated Gaussian formula as bounds recede. The old untruncated approximation distorted the EPHistory convergence metric as posterior mass approached the bounds (live in HowToFit ch3 via af.TruncatedGaussianPrior). MC-verified at interior, near-bounds and half-bounded geometries. - F7(b): the search-driven factor projection now records the sampler's log-evidence of the tilted fit on the projected mean field's log_norm (AbstractSearch.optimise -> MeanField.from_priors(log_norm=...)), implementing the design README §5 already documented. Both levels guarded (StaticResult has no samples; MCMC/MLE have no evidence -> 0.0 default). With F7(a) (#1351) and F7(c) (#1345) this makes EPMeanField.log_evidence trustworthy for model comparison when factor searches are nested samplers. - README §5: stale audit warning replaced with the F7-fixed statement + the nested-sampler requirement; §6: #1336 Recommendation-A decision recorded (keep all three deterministic-composition mechanisms, trade-off documented, compound pattern is the declarative surface). Adds test_ep_statistics_completion.py (7 tests: MC KL parity x3 geometries, wide-bounds reduction, near-bounds divergence from the old formula, support mismatch raise, from_priors log_norm plumbing). Full suite: 1474 passed / 14 skipped / 1 pre-existing order-dependent failure (test_full_hierachical, triaged on #1352 — deterministic test-order dependence, not a flake; unrelated to this diff). Co-Authored-By: Claude Opus 4.8 --- autofit/graphical/README.md | 27 +++-- autofit/graphical/mean_field.py | 16 ++- autofit/messages/truncated_normal.py | 41 ++++++-- autofit/non_linear/search/abstract_search.py | 17 +++- .../test_ep_statistics_completion.py | 99 +++++++++++++++++++ 5 files changed, 185 insertions(+), 15 deletions(-) create mode 100644 test_autofit/graphical/functionality/test_ep_statistics_completion.py diff --git a/autofit/graphical/README.md b/autofit/graphical/README.md index 57ced1a88..b40b2c454 100644 --- a/autofit/graphical/README.md +++ b/autofit/graphical/README.md @@ -185,15 +185,30 @@ The EP approximation to the model evidence factorises as: log ∫ ∏ₖ qₖ = Σₖ (log hₖ − A(ηₖ)) − ( log h − A(Σₖ ηₖ) ) (16) and the per-factor `Ẑₐ` is carried on `MeanField.log_norm` by the -projection. (Audit note: as of 2026-07 the `log_norm` bookkeeping is -broken in three places — issue #1332 finding F7 — so Eq. (15) should -not be used for model comparison until those fixes land; the -decomposition itself is correct.) +projection: a search-driven factor update records the sampler's +log-evidence of the tilted fit there (`AbstractSearch.optimise` → +`MeanField.from_priors(..., log_norm=...)`). + +All three legs of the 2026-07 audit's finding F7 (#1332) are now fixed: +(a) `MeanField.__truediv__`/`__pow__` propagate `log_norm` (#1351), +(b) the search-driven projection records `Ẑₐ` (this section), +(c) the truncated-normal log-partition is complete (#1345). Evidence- +correct model comparison additionally requires **nested-sampling factor +searches** — MCMC/MLE searches carry no evidence estimate and contribute +`log_norm = 0`. ## 6. Deterministic variables -Three composition mechanisms exist; they are **not** interchangeable -and reconciling them is an open design item (EP review Phase 5): +Three composition mechanisms exist; they are **not** interchangeable. +**Decision (EP review Phase 5, #1336, 2026-07-10): keep all three — +no unification, no deprecation.** The trade-off users should choose by: +compound priors / shared variables are statistically *tighter* (the +relation holds exactly inside each factor), while `factor_out` trades +that exactness for modularity (the deterministic variable receives its +own messages and `q(z)` factorises from its parents). The declarative +surface for deterministic quantities is the explicit compound pattern +(e.g. `model.sigma * 2.355`), pinned by the seam tests (§8); the +`model.` sugar from #1153 was deliberately retired. 1. **Graph-level deterministic variables**: `Factor(..., factor_out=v)` declares outputs computed by the factor diff --git a/autofit/graphical/mean_field.py b/autofit/graphical/mean_field.py index 99a4fd8f4..1bd83ade6 100755 --- a/autofit/graphical/mean_field.py +++ b/autofit/graphical/mean_field.py @@ -123,7 +123,9 @@ def fixed_values(self): ) @classmethod - def from_priors(cls, priors: Iterable[Prior]) -> "MeanField": + def from_priors( + cls, priors: Iterable[Prior], log_norm: float = 0.0 + ) -> "MeanField": """ Create a MeanField from a list of priors. @@ -134,12 +136,22 @@ def from_priors(cls, priors: Iterable[Prior]) -> "MeanField": ---------- priors A list of priors + log_norm + The log-normalisation carried by this mean field. For a + search-driven factor update this is the sampler's log-evidence of + the tilted-distribution fit — the per-factor ``Ẑₐ`` of README §5 + Eq. (14), which the projection is documented to carry on + ``MeanField.log_norm`` (#1332 F7(b)). Defaults to 0.0 for callers + with no evidence to record. Returns ------- A mean field """ - return MeanField({prior: prior.message for prior in priors}) + return MeanField( + {prior: prior.message for prior in priors}, + log_norm=log_norm, + ) pop = dict.pop values = dict.values diff --git a/autofit/messages/truncated_normal.py b/autofit/messages/truncated_normal.py index 10ac3bd1d..3d7c6e646 100644 --- a/autofit/messages/truncated_normal.py +++ b/autofit/messages/truncated_normal.py @@ -304,11 +304,21 @@ def sample(self, n_samples: Optional[int] = None) -> np.ndarray: def kl(self, dist : "TruncatedNormalMessage") -> float: """ - Compute the Kullback-Leibler (KL) divergence between two truncated Gaussian distributions. + The exact Kullback-Leibler divergence KL(self || dist) between two + truncated Gaussians sharing the same support (#1332 F6). - This is an approximate KL divergence that assumes both distributions are truncated Gaussians - with the same support (i.e. the same lower and upper limits). If the supports differ, this - expression is invalid and should raise an error or be corrected for normalization. + With p, q truncated to the same ``[a, b]``, standardised bounds + ``α = (a − μ)/σ``, ``β = (b − μ)/σ``, truncation mass + ``Z = Φ(β) − Φ(α)`` and the *truncated* moments ``m_p``, ``V_p`` of p: + + KL(p‖q) = log(σ_q/σ_p) + log(Z_q/Z_p) + + ½·[ (V_p + (m_p − μ_q)²)/σ_q² − (V_p + (m_p − μ_p)²)/σ_p² ] + + As the bounds recede (Z → 1, m_p → μ_p, V_p → σ_p²) this reduces to the + untruncated Gaussian KL. Previously this method *used* the untruncated + formula, which degrades as posterior mass approaches the bounds — + directly distorting the ``EPHistory`` convergence metric for models + built on ``af.TruncatedGaussianPrior`` (e.g. HowToFit chapter 3). Parameters ---------- @@ -323,10 +333,29 @@ def kl(self, dist : "TruncatedNormalMessage") -> float: if (self.lower_limit != dist.lower_limit) or (self.upper_limit != dist.upper_limit): raise ValueError("KL divergence between truncated Gaussians with different support is not implemented.") + from scipy.stats import norm, truncnorm + + a_p = (self.lower_limit - self.mean) / self.sigma + b_p = (self.upper_limit - self.mean) / self.sigma + a_q = (self.lower_limit - dist.mean) / dist.sigma + b_q = (self.upper_limit - dist.mean) / dist.sigma + + log_Z_p = np.log(norm.cdf(b_p) - norm.cdf(a_p)) + log_Z_q = np.log(norm.cdf(b_q) - norm.cdf(a_q)) + + # Truncated mean and variance of p (scipy returns them for the + # standardised bounds with loc/scale applied). + m_p, V_p = truncnorm.stats( + a_p, b_p, loc=self.mean, scale=self.sigma, moments="mv" + ) + + e_zp2 = (V_p + (m_p - self.mean) ** 2) / self.sigma**2 + e_zq2 = (V_p + (m_p - dist.mean) ** 2) / dist.sigma**2 + return ( np.log(dist.sigma / self.sigma) - + (self.sigma**2 + (self.mean - dist.mean) ** 2) / 2 / dist.sigma**2 - - 1 / 2 + + (log_Z_q - log_Z_p) + + 0.5 * (e_zq2 - e_zp2) ) @classmethod diff --git a/autofit/non_linear/search/abstract_search.py b/autofit/non_linear/search/abstract_search.py index 0ce101302..6701ffdb9 100644 --- a/autofit/non_linear/search/abstract_search.py +++ b/autofit/non_linear/search/abstract_search.py @@ -367,7 +367,22 @@ def optimise( result = self.fit(model=model, analysis=analysis) - new_model_dist = MeanField.from_priors(result.projected_model.priors) + # Record the sampler's log-evidence of this tilted-distribution fit on + # the projected mean field — the per-factor Ẑₐ that README §5 documents + # `MeanField.log_norm` as carrying (#1332 F7(b)). Previously always 0, + # so `EPMeanField.log_evidence` could not be trusted for model + # comparison in sampler-driven EP fits. Searches with no evidence + # estimate (MCMC / MLE) yield None and keep the 0.0 default — evidence- + # correct model comparison requires nested-sampling factor searches. + # (Both levels guarded: e.g. StaticResult carries no samples at all.) + log_evidence = getattr( + getattr(result, "samples", None), "log_evidence", None + ) + + new_model_dist = MeanField.from_priors( + result.projected_model.priors, + log_norm=log_evidence if log_evidence is not None else 0.0, + ) status.result = result diff --git a/test_autofit/graphical/functionality/test_ep_statistics_completion.py b/test_autofit/graphical/functionality/test_ep_statistics_completion.py new file mode 100644 index 000000000..afb19f242 --- /dev/null +++ b/test_autofit/graphical/functionality/test_ep_statistics_completion.py @@ -0,0 +1,99 @@ +"""Regression tests for the EP statistics completion (PyAutoFit #1353; +findings F6 and F7(b) from the #1332 audit). Numpy-only. +""" +import numpy as np +import pytest +from scipy.stats import truncnorm + +import autofit as af +from autofit.graphical.mean_field import MeanField +from autofit.messages.truncated_normal import TruncatedNormalMessage + + +# --- F6: exact same-support truncated KL via truncated moments --- + +def _mc_kl_truncnorm(p, q, n=400_000, seed=3): + a_p = (p.lower_limit - p.mean) / p.sigma + b_p = (p.upper_limit - p.mean) / p.sigma + a_q = (q.lower_limit - q.mean) / q.sigma + b_q = (q.upper_limit - q.mean) / q.sigma + rng = np.random.default_rng(seed) + x = truncnorm.rvs( + a_p, b_p, loc=p.mean, scale=p.sigma, size=n, random_state=rng + ) + return float( + np.mean( + truncnorm.logpdf(x, a_p, b_p, loc=p.mean, scale=p.sigma) + - truncnorm.logpdf(x, a_q, b_q, loc=q.mean, scale=q.sigma) + ) + ) + + +def _old_untruncated_kl(p, q): + return ( + np.log(q.sigma / p.sigma) + + (p.sigma**2 + (p.mean - q.mean) ** 2) / 2 / q.sigma**2 + - 1 / 2 + ) + + +@pytest.mark.parametrize( + "p, q", + [ + # comfortable interior — old approximation was nearly right here + ( + TruncatedNormalMessage(0.0, 1.0, -10.0, 10.0), + TruncatedNormalMessage(0.5, 2.0, -10.0, 10.0), + ), + # mass pressed against the bounds — old approximation degrades badly + ( + TruncatedNormalMessage(0.9, 1.5, -1.0, 1.0), + TruncatedNormalMessage(-0.5, 0.7, -1.0, 1.0), + ), + # half-bounded, mean below the bound (prior-passing shape) + ( + TruncatedNormalMessage(-0.5, 1.0, 0.0, np.inf), + TruncatedNormalMessage(1.0, 2.0, 0.0, np.inf), + ), + ], + ids=["interior", "near-bounds", "half-bounded"], +) +def test__truncated_kl_matches_monte_carlo(p, q): + analytic = float(p.kl(q)) + mc = _mc_kl_truncnorm(p, q) + assert analytic == pytest.approx(mc, abs=0.02) + + +def test__truncated_kl_reduces_to_gaussian_for_wide_bounds(): + p = TruncatedNormalMessage(0.0, 1.0, -50.0, 50.0) + q = TruncatedNormalMessage(1.0, 2.0, -50.0, 50.0) + assert float(p.kl(q)) == pytest.approx(float(_old_untruncated_kl(p, q)), rel=1e-9) + + +def test__truncated_kl_near_bounds_differs_from_old_formula(): + # The case the audit quantified (errors 1.5% -> 140% as mass reaches the + # bounds): the exact value must both match MC and differ measurably from + # the old untruncated formula. + p = TruncatedNormalMessage(0.9, 1.5, -1.0, 1.0) + q = TruncatedNormalMessage(-0.5, 0.7, -1.0, 1.0) + exact = float(p.kl(q)) + old = float(_old_untruncated_kl(p, q)) + assert abs(exact - old) / abs(exact) > 0.05 + + +def test__truncated_kl_different_support_still_raises(): + p = TruncatedNormalMessage(0.0, 1.0, -1.0, 1.0) + q = TruncatedNormalMessage(0.0, 1.0, -2.0, 2.0) + with pytest.raises(ValueError): + p.kl(q) + + +# --- F7(b): from_priors records the per-factor evidence on log_norm --- + +def test__from_priors_log_norm_default_and_explicit(): + priors = [af.GaussianPrior(0.0, 1.0), af.GaussianPrior(1.0, 2.0)] + assert MeanField.from_priors(priors).log_norm == 0.0 + mf = MeanField.from_priors(priors, log_norm=-42.5) + assert mf.log_norm == -42.5 + # and it survives the (now-fixed, #1351) operator algebra + assert (mf / MeanField.from_priors(priors, log_norm=-2.5)).log_norm == -40.0