From ad97e5fd558d60511899afb9c5fb6fefde44787e Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Fri, 17 Jul 2026 14:34:35 +0100 Subject: [PATCH] Marquardt-scale LM damping + zero-fill correction profiles MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit solve_lm_step_from now damps with mu*diag(diag(H)) (clipped) instead of mu*I — scale-invariant damping; visibility-weighted curvatures (~1e11) stalled the interferometer LM under mu*I's 4 usable decades. Both iterative engines build their correction InputPotential with extrapolate=zero (companion PyAutoGalaxy PR) so arc-restricted corrections vanish outside their mesh. Part of PyAutoLabs/PyAutoLens#623. Co-Authored-By: Claude Fable 5 --- autolens/potential_correction/dense_util.py | 14 ++++++++++---- autolens/potential_correction/iterative.py | 4 ++++ .../iterative_interferometer.py | 4 ++++ .../potential_correction/test_dense_util.py | 3 ++- 4 files changed, 20 insertions(+), 5 deletions(-) diff --git a/autolens/potential_correction/dense_util.py b/autolens/potential_correction/dense_util.py index 6a26273ff..3b8b64e08 100644 --- a/autolens/potential_correction/dense_util.py +++ b/autolens/potential_correction/dense_util.py @@ -376,14 +376,20 @@ def lm_hessian_and_gradient_from( def solve_lm_step_from(H, minus_gradient, mu, constraint_matrix=None, x=None, xp=np): """ - The damped LM step delta_x solving (H + mu I) dx = -g, or — when a - ``constraint_matrix`` C is given — the equality-constrained step via the - KKT system enforcing C (x + dx) = 0. + The damped LM step delta_x solving (H + mu D) dx = -g with Marquardt + scaling D = diag(diag(H)) (clipped below at the mean diagonal times + 1e-12 so zero diagonal entries stay damped) — scale-invariant damping, + required when H's magnitude varies over many orders between datasets + (e.g. visibility-weighted interferometer curvatures ~1e11 vs imaging + ~1e4). When a ``constraint_matrix`` C is given, the equality-constrained + step solves the KKT system enforcing C (x + dx) = 0. """ H_d = as_dense(H, xp=xp) g = xp.asarray(minus_gradient) n_x = H_d.shape[0] - H_lm = H_d + mu * xp.eye(n_x, dtype=H_d.dtype) + diag = xp.diag(H_d) + diag = xp.clip(diag, 1e-12 * xp.mean(xp.abs(diag)), None) + H_lm = H_d + mu * xp.diag(diag) if constraint_matrix is None: return xp.linalg.solve(H_lm, g) diff --git a/autolens/potential_correction/iterative.py b/autolens/potential_correction/iterative.py index 0129a7540..5666039d5 100644 --- a/autolens/potential_correction/iterative.py +++ b/autolens/potential_correction/iterative.py @@ -212,10 +212,14 @@ def _updated_lens_galaxies_from_dpsi(self, dpsi_vec: np.ndarray): mask_dpsi_aa = self.pair_dpsi_data_obj.mask_dpsi_aa grid_dpsi = aa.Grid2D.from_mask(mask=mask_dpsi_aa) + # zero-fill extrapolation: the correction vanishes outside its mesh + # (nearest extrapolation would smear constant deflections across the + # full re-trace grid when the dpsi mesh is a sub-region of it) pix_mass_profile = InputPotential( lensing_potential=dpsi_vec, image_plane_grid=np.asarray(grid_dpsi), mask=mask_dpsi_aa, + extrapolate="zero", ) lens_macro = self.lens_start diff --git a/autolens/potential_correction/iterative_interferometer.py b/autolens/potential_correction/iterative_interferometer.py index 27322767a..32b210c18 100644 --- a/autolens/potential_correction/iterative_interferometer.py +++ b/autolens/potential_correction/iterative_interferometer.py @@ -230,10 +230,14 @@ def _updated_lens_galaxies_from_dpsi(self, dpsi_vec: np.ndarray): mask_dpsi_aa = self.pair_dpsi_data_obj.mask_dpsi_aa grid_dpsi = aa.Grid2D.from_mask(mask=mask_dpsi_aa) + # zero-fill extrapolation: the correction vanishes outside its mesh + # (nearest extrapolation would smear constant deflections across the + # full re-trace grid when the dpsi mesh is a sub-region of it) pix_mass_profile = InputPotential( lensing_potential=dpsi_vec, image_plane_grid=np.asarray(grid_dpsi), mask=mask_dpsi_aa, + extrapolate="zero", ) lens_macro = self.lens_start diff --git a/test_autolens/potential_correction/test_dense_util.py b/test_autolens/potential_correction/test_dense_util.py index 135164dc8..e96e4b068 100644 --- a/test_autolens/potential_correction/test_dense_util.py +++ b/test_autolens/potential_correction/test_dense_util.py @@ -219,7 +219,8 @@ def test__solve_lm_step_from__unconstrained_and_constrained(): mu = 0.7 step = dense_util.solve_lm_step_from(H, minus_gradient, mu) - assert (H + mu * np.eye(H.shape[0])) @ step == pytest.approx( + diag = np.clip(np.diag(H), 1e-12 * np.mean(np.abs(np.diag(H))), None) + assert (H + mu * np.diag(diag)) @ step == pytest.approx( minus_gradient, rel=1.0e-8 )