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 )