Skip to content

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

Merged
Jammy2211 merged 4 commits into
mainfrom
feature/ep-moment-projection
Sep 30, 2026
Merged

Jammy2211 merged 4 commits into
mainfrom
feature/ep-moment-projection

Conversation

@Jammy2211

Copy link
Copy Markdown
Collaborator

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 and others added 4 commits September 30, 2026 11:56
autofit/graphical/laplace/moments.py: outer Gauss-Legendre over the
scale / bounded-support variables in base coordinates, windowed to the
support and the cavity (re-windowed while under-resolved, <= 4 passes);
inner conditional Laplace at each node (warm-started quasi-Newton,
Newton polish on value central differences); log w = log w_GL + log|J|
+ l(m,s) + d/2 log 2pi - 1/2 log det(-H); moments via an order-3
Gauss-Hermite expansion of each conditional. Deterministic, numpy only.

Tests: known two-variable moments vs scipy quad (1e-6), Gaussian
mode/cov, narrow re-window, RNG / Variable-id determinism, mode
fallbacks, BAD_PROJECTION on an empty window, hierarchical toy vs the
exact scatter posterior, near-zero scatter mode-vs-moments.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Up1fQmjtGAh4RhLrWHfcfg
New kwargs projection ("mode" default | "moments"), n_quadrature=64,
quadrature_half_width=8.0, moment_max_size=4, moment_max_outer=2;
optimise_approx tries the moment projection first and falls back to the
unchanged mode path when it does not apply.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Up1fQmjtGAh4RhLrWHfcfg
…iables (#1654)

from_weighted_nodes projects weighted nodes through each message's own
project (log_norm 0 per message, projection log_norm = log Z).
scale_variables maps _SCALE_ARGUMENT_NAMES (kept in the import-free
diagnostics module, now shared) to the parent distribution's priors.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Up1fQmjtGAh4RhLrWHfcfg
@Jammy2211 Jammy2211 added the pending-release PR queued for the next release build label Sep 30, 2026
@Jammy2211
Jammy2211 merged commit b13169e into main Sep 30, 2026
4 checks passed
@Jammy2211
Jammy2211 deleted the feature/ep-moment-projection branch September 30, 2026 13:52
@Jammy2211

Copy link
Copy Markdown
Collaborator Author

Workspace PR: PyAutoLabs/autofit_workspace_test#105

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant