feat: LaplaceOptimiser(projection="moments") — moment-matching projection for the hierarchical scatter (#1654)

[Jammy2211]

Summary

This PR addresses the Laplace-on-scatter caveat (autofit/graphical/README.md §3.5). EP's LaplaceOptimiser projects every factor's tilted distribution onto a Gaussian at its mode and curvature. For a hierarchical factor, the tilted density of the scatter σ often has no interior mode: it piles up against σ = 0, or is strongly skewed. The mode path then returns BAD_PROJECTION or a biased σ, and the autofit-EP column of analytic_gaussian / analytic_gaussian_priors misses the closed form. Those scripts are parked in no_run.yaml for this reason.

This PR adds an opt-in moment-matching projection, LaplaceOptimiser(projection="moments"). It is deterministic nested quadrature, ported from the analytic_ep_minimal referee in autofit_workspace_test:

• Outer: a Gauss–Legendre rule over each scale (or bounded-support) variable on its support, windowed to the cavity mean ± quadrature_half_width · std. Up to 4 re-windowing passes run until the tilted density is resolved.
• Inner: a conditional Laplace approximation over the remaining variables at each outer node. Its gradient and Hessian are value-based central differences, followed by a Newton polish.
• The weighted nodes are moment-matched back onto the mean field's message families.

The default projection="mode" is unchanged. Any factor the moments path does not apply to falls back to the mode path bit-for-bit.

Design departure from the prompt. The prompt proposed Gauss–Hermite quadrature centred on the Hessian (mode) approximation. That centring is exactly what breaks here: when the tilted σ density has no interior mode, a Hessian-centred rule is placed at the wrong location and scale. A support-aware outer Gauss–Legendre rule with re-windowing does not depend on a mode existing, and it is the scheme the referee already validates.

API Changes

• af.LaplaceOptimiser(..., projection="mode" | "moments", n_quadrature=64, quadrature_half_width=8.0, moment_max_size=4, moment_max_outer=2): new keyword arguments. The default "mode" keeps today's behaviour exactly. Any other projection string raises ValueError.
• New MeanField.from_weighted_nodes(nodes, log_weights, log_norm=0.0): a moment-matching projection of a weighted node set onto the mean field's message families.
• New _HierarchicalFactor.scale_variables (property): the priors that parameterise the parent distribution's scale.
• New module autofit.graphical.laplace.moments (numpy-only).
• A successful moments projection reports the status message moments: n_nodes=…, passes=…, f_count=…, g_count=….
• No removals and no default change, so there is no downstream ripple. autolens/autogalaxy do not construct LaplaceOptimiser.
  See full details below.

Test Plan

• test_autofit/graphical/functionality/test_moment_projection.py, 13 tests. The initial 9 failed red (TypeError: unexpected kwarg 'projection') before the change.
  • Known moments, N(x|0,s) with xN(1,0.5) and sTN(0.8,0.6;0,100), against scipy.integrate.quad: dE[x] 3.3e-9, dVar[x] 9.0e-10, dE[s] 7.7e-13, dVar[s] 2.9e-13, dlogZ 1.5e-13 (asserted at 1e-6).
  • Gaussian check, mode/cov errors ≤ 4.3e-10. Narrow tilted density (14x narrower than the cavity window): passes=3, errors ≤ 9e-10.
  • Determinism: bit-equal across seeds 0/1/12345 and across 1/3/6 throwaway Variables.
  • Fallbacks: a factor with no outer variable, moment_max_size=1 or moment_max_outer=0 gives the same bits and messages as "mode" and logs why. projection="median" raises ValueError.
  • Cavity TN(-100,1) on (0,100) gives BAD_PROJECTION, updated=False, and the model dist is returned unchanged.
  • End to end on two toys. On toy (45,52,58) every hierarchical update ends SUCCESS, and the scatter parent mean is 9.297 against exact 9.654±3.468. On the tight toy (50,50.5,49.5) the mode path has 0 SUCCESS (all BAD_PROJECTION) while moments gives 3/3 SUCCESS.
  • scale_variables == {hf.sigma}.
• Full suite after rebase on origin/main: python -m pytest test_autofit -x, 2947 passed, 2 skipped. Existing Laplace tests are unchanged, and so is the mode-path code.
• Referee cross-check (analytic_ep_minimal.ep_leg_b, moments, theta=sigma): toy (45,52,58) parent 9.2880±3.6524, identical to this PR at 4 dp.
• Witness: analytic_autofit.py in a scratch copy of autofit_workspace_test with LaplaceOptimiser(projection="moments"), 3 runs of 92–108 s:

leg / column	moments (this PR)	mode-path control
Leg A	18/18 PASS (a, b = 0.000)	—
Leg B autofit EP, sigma	6.3401±2.4609, a 0.079 / b 0.146 PASS	9.3722±3.5663, a 0.973 FAIL
Leg B autofit EP, mu	a 0.021 / b 0.061 PASS	—
Leg B autofit EP, x_0..x_4	a ≤ 0.018 / b ≤ 0.010 PASS	x_2, x_4 FAIL
HierarchicalFactor status	SUCCESS=25	BAD_PROJECTION=19
Overall PARITY	40/41 (41/41 in 1 of 3 runs)	37/41

The single moments miss is the "autofit graphical" DynestyStatic joint-fit sigma row (a 0.120–0.129 vs tol 0.10). That sampler is unseeded, the row passed in 1 of 3 runs, and it misses the same way under the mode path. It is not an EP column.

• Cost of one hierarchical projection (toy, numpy, uncontended): moments 0.64 s / 1234 factor calls, against mode 0.04 s / 84 calls. The witness leg B EP column takes 32–45 s, against about 2 s on the mode path. The moments path is opt-in for this reason.
• Phase 2 (workspace, after merge): autofit_workspace_test analytic_autofit.py gets the projection kwarg and labels, analytic_gaussian / analytic_gaussian_priors come out of no_run.yaml, and smoke_tests.txt is curated.

Full API Changes (for automation & release notes)

Added

• autofit.graphical.laplace.optimiser.LaplaceOptimiser.__init__(..., projection: str = "mode", n_quadrature: int = 64, quadrature_half_width: float = 8.0, moment_max_size: int = 4, moment_max_outer: int = 2)
• autofit.graphical.laplace.optimiser.LaplaceOptimiser._moment_projection(factor_approx, mean_field, params=None, **kwargs) -> Optional[Tuple[MeanField, Status]] (private). Returns None when the mode path should run.
• autofit.graphical.laplace.moments (new module). Public helpers are moment_projection(optimiser, factor_approx, mean_field, params=None, **kwargs), MomentProjection, split_variables, fallback_reason and gauss_legendre(n, lo, hi). Module constants are SCALE_SUPPORT = (0, inf), MAX_PASSES = 4, NARROW_FRACTION = 0.25, EDGE_MASS = 1e-6, WEIGHTED_NODE_MASS = 1e-10, NEWTON_DECREMENT_TOL = 1e-12 and MAX_NEWTON_POLISH = 5.
• autofit.graphical.mean_field.MeanField.from_weighted_nodes(nodes: Dict[Variable, np.ndarray], log_weights: np.ndarray, log_norm: float = 0.0) -> MeanField
• autofit.graphical.declarative.factor.hierarchical._HierarchicalFactor.scale_variables -> frozenset (property)
• Status message on a successful moments projection: moments: n_nodes=…, passes=…, f_count=…, g_count=…

Changed

• LaplaceOptimiser.__init__ accepts the five kwargs above (defaults preserve behaviour) and validates projection. LaplaceOptimiser.optimise tries the moments path first when projection="moments", then falls back to the unchanged mode path.
• autofit/graphical/expectation_propagation/diagnostics.py: a comment marks _SCALE_ARGUMENT_NAMES as shared with _HierarchicalFactor.scale_variables.
• autofit/graphical/README.md §3.2, §3.3, §3.5 document the moments path and the caveat it cures.

Migration

• None. Opt in with af.LaplaceOptimiser(projection="moments").

Deviations from the approved plan

1. The inner Hessian and polish gradient are value-based central differences, not newton.finite_difference_hessian on gradients. At σ≈0.017 the forward-difference gradient carried a 0.065 error: catastrophic cancellation of η₁x, η₂x² and A(η), which are around 1e6–1e7. That error made the Hessian indefinite (min eig -0.47) and every hierarchical update ended BAD_PROJECTION or FAILURE. The value differences are exact to about 4e-9. Quasi-Newton still runs first, warm-started.
2. Up to 4 re-windowing passes (the plan had 2). Two passes under-resolved a tilted density 14x narrower than the cavity (a 15% variance error); three passes give 1e-10.
3. _SCALE_ARGUMENT_NAMES stays in expectation_propagation/diagnostics.py, which imports nothing from autofit so no cycle arises, rather than moving.
4. The scale-variable support is (0, inf), applied by argument name rather than read per distribution class.
5. moments.py is numpy-only, with its own logsumexp and Cholesky solve. A module-level scipy.special import broke test_lazy_transformed_support.
6. The end-to-end test uses two toys, because the mode path succeeds on the first one.

Generated by the PyAutoLabs agent workflow.

Heart RED override (development only)

• Authorised by the live human in the Claude Code session at 2026-09-30 ~11:50 BST. Question: "Do you authorise the Heart RED development override for push + pending-release PR-open on PyAutoFit#1653, and, when its implementation lands green, on PyAutoFit#1654 (moment projection)? No merge, no release; recorded in the four sinks." Answer: "Yes, both fix: EP projection raises ProjectionException on non-finite stats #1653 and feat: LaplaceOptimiser(projection="moments") for hierarchical scatter #1654 (Recommended)".
• Exact RED reasons (readiness --json, ts 2026-09-30T09:51:49Z, re-read ~11:45 BST unchanged): "PyAutoArray: 2 commit(s) behind origin"; "PyAutoLens: 2 commit(s) behind origin"; "release validation FAILED (stage integrate)". Yellow: "workspace validation not passing (0 failed, 1 timeout, cloud#36404726969: autolens_test scripts/multi_dataset/rectangular.py)"; "manifest drift: public front-door organ tables (generated) — 1 mismatch(es) vs PyAutoMind/repos.yaml".
• Branch gates passed: rebase on origin/main clean; pytest test_autofit -x 2947 passed, 2 skipped; the new tests were confirmed red first; witness 40/41 with the one miss outside EP.
• This PR does not claim to fix Heart. Heart stays RED for release purposes. Merge needs its own human command (/prm) and green required checks.
• File sets are disjoint from PyAutoFit#1655 (fix: EP projection raises ProjectionException on non-finite stats #1653); the two PRs can merge in either order.

🤖 Generated with Claude Code

https://claude.ai/code/session_01Up1fQmjtGAh4RhLrWHfcfg

[Jammy2211]

Workspace PR: PyAutoLabs/autofit_workspace_test#105