Merge pull request #626 from PyAutoLabs/feature/potential-correction-interferometer

fix: Marquardt-scaled LM damping + zero-fill correction profiles

Author: Jammy2211

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)

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

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

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
     )