Skip to content

fix: LM damping parity + stall guards for potential-correction iterative engines - #676

Merged
Jammy2211 merged 1 commit into
mainfrom
feature/potential-correction-validation
Jul 31, 2026
Merged

fix: LM damping parity + stall guards for potential-correction iterative engines#676
Jammy2211 merged 1 commit into
mainfrom
feature/potential-correction-validation

Conversation

@Jammy2211

Copy link
Copy Markdown
Collaborator

Summary

Closes the potential-correction JAX-vs-Python discrepancy investigation of #672 (upstream author report: "for the same set of regularization parameters ... the recovered mass-perturbation signal differs between my original pure Python code and the new JAX-ported version").

Side-by-side on the reference 200x200 demo dataset (identical inputs, cold start, the author's demo-2 hyper-params, n_iter=5) localized the divergence to the LM damping matrix: the reference implementation damps H + mu*I (in both its numpy and JAX paths), while the port used the scale-invariant Marquardt form H + mu*diag(H) everywhere. The two are not interchangeable at a fixed iteration budget — identical initial costs (5.0276e6, cost functions agree exactly), but after 5 iterations the reference engine reaches cost 6.2694e3 while the Marquardt-damped port sits at 1.0367e4 with a noise dkappa (corr 0.024). With mu*I patched in, the port matches the reference trajectory to 4 significant figures at every iteration (6.2696e3) — the implementations are equivalent; the damping form was the entire discrepancy. One-shot (single-step) inversion needed no fix: machine-precision parity (evidence diff 6e-7, dkappa corr 1.000000).

The investigation also exposed an end-of-budget failure mode in both implementations: at a cost minimum no decreasing step exists, so the inner LM loop rejects trial steps (each a full Jacobian rebuild, ~minutes at 200x200) until mu exceeds 1e15 — 47 wasted minutes and, for the reference implementation on a 15GB machine, an OOM death mid-storm.

Changes:

  • dense_util.solve_lm_step_from gains damping="identity"|"marquardt" (function default "marquardt", unchanged — the JAX-parity smoke scripts that call it directly are unaffected).
  • IterFitDpsiSrcImaging gains damping (default "identity" — restores the reference implementation's behavior: cold-start converges in a few iterations at the author's published hyper-params) and max_consecutive_rejections (default 10).
  • IterFitDpsiSrcInterferometer gains the same two options with damping="marquardt" retained as its default (visibility-weighted curvatures ~1e11 are why the scale-invariant form exists; the research: realistic-uv iterative potential-correction recovery campaign #627-certified uv recipe is unchanged).
  • Stall guards in both engines' inner loops: a rejected step below tol now returns as converged (growing mu only shrinks the step further — previously the tol check fired only on accepted steps), and max_consecutive_rejections bounds the rejection storm.

API Changes

Changed Behaviour

  • al.pc.IterFitDpsiSrcImaging now defaults to identity LM damping (damping="identity"), matching the reference implementation of Cao et al. 2025. Cold-start recovery at published hyper-params works again without the feat: x0 warm start for the imaging iterative engine (parity) #630 warm start (the warm start remains supported and is still the certified recipe under damping="marquardt"). Pass damping="marquardt" for the previous behavior.
  • Both iterative engines now terminate when a rejected step falls below tol or after max_consecutive_rejections consecutive rejections, instead of rejecting until mu > 1e15. States returned at these exits are the same current-best states the mu-ceiling exit returned; runs that previously burned a rejection storm at convergence now return up to tens of minutes sooner.

Added

  • damping and max_consecutive_rejections parameters on al.pc.IterFitDpsiSrcImaging and al.pc.IterFitDpsiSrcInterferometer; damping parameter on al.pc.dense_util.solve_lm_step_from.

Test Plan

  • test_autolens/ full suite: 492 passed (includes new identity-damping solve tests, damping plumb-through tests, and a warm-start-at-optimum stall-guard test bounding Jacobian rebuilds).
  • Side-by-side validation vs the author's original code on the 200x200 demo (issue fix: potential-correction end-to-end validation (parity + evidence sampling) #672 comments: parity table + trajectory evidence).
  • autolens_workspace_test imaging/subhalo_recovery.py run against this branch under the exact CI smoke env — passes with identical one-shot evidence.

Generated by the PyAutoLabs agent workflow.

🤖 Generated with Claude Code

…ive engines

The iterative engines' Marquardt mu*diag(H) damping diverged from the
reference implementation's mu*I, changing recovered signals at fixed
iteration budgets (PyAutoLens#672). damping= option added to both engines
(imaging default restored to identity; interferometer keeps marquardt);
rejected steps below tol now return as converged and consecutive
rejections are capped, ending the end-of-budget rejection storm.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@Jammy2211

Copy link
Copy Markdown
Collaborator Author

Workspace PR (phase 3 acceptance test): PyAutoLabs/autolens_workspace_test#243

@Jammy2211
Jammy2211 merged commit 277b938 into main Jul 31, 2026
3 checks passed
@Jammy2211
Jammy2211 deleted the feature/potential-correction-validation branch July 31, 2026 22:11
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

pending-release PR queued for the next release build

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant