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

autolens/potential_correction/dense_util.py

@@ -374,22 +374,38 @@ def lm_hessian_and_gradient_from(
     return H, minus_gradient, residual, chi2, reg_s, reg_dpsi, cost
 
 
-def solve_lm_step_from(H, minus_gradient, mu, constraint_matrix=None, x=None, xp=np):
+def solve_lm_step_from(
+    H, minus_gradient, mu, constraint_matrix=None, x=None, xp=np, damping="marquardt"
+):
     """
-    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.
+    The damped LM step delta_x solving (H + mu D) dx = -g.
+
+    ``damping="identity"`` uses D = I — the damping of the reference
+    implementation (Cao et al. 2025): at moderate mu it barely perturbs
+    high-curvature directions, so early steps are near full Gauss-Newton and
+    the imaging problem converges in a few iterations from a cold start.
+    ``damping="marquardt"`` uses the scale-invariant D = diag(diag(H))
+    (clipped below at the mean diagonal times 1e-12 so zero diagonal entries
+    stay damped), required when H's magnitude varies over many orders between
+    datasets (e.g. visibility-weighted interferometer curvatures ~1e11 vs
+    imaging ~1e4); its steps are far more conservative at the same mu, so a
+    small iteration budget under-converges relative to identity damping.
+    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]
-    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 damping == "identity":
+        H_lm = H_d + mu * xp.eye(n_x, dtype=H_d.dtype)
+    elif damping == "marquardt":
+        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)
+    else:
+        raise ValueError(
+            f"damping must be 'identity' or 'marquardt', got {damping!r}"
+        )
 
     if constraint_matrix is None:
         return xp.linalg.solve(H_lm, g)

autolens/potential_correction/iterative.py

@@ -57,6 +57,8 @@ def __init__(
         preloads: Optional[dict] = None,
         n_iter: int = 20,
         tol: float = 1e-6,
+        damping: str = "identity",
+        max_consecutive_rejections: int = 10,
         verbose: bool = False,
         visualize_output_dir: Optional[str] = None,
         visualize_every_n: int = 1000000,
@@ -98,7 +100,22 @@ def __init__(
         n_iter
             The maximum number of outer LM iterations.
         tol
-            The step-norm convergence tolerance.
+            The step-norm convergence tolerance. Also applied to rejected
+            steps: once a proposed step is smaller than ``tol``, growing the
+            damping can only shrink it further, so the solve returns rather
+            than rejecting its way to the mu ceiling.
+        damping
+            The LM damping matrix (``dense_util.solve_lm_step_from``):
+            ``"identity"`` (default) is the reference implementation's
+            ``H + mu I`` — near Gauss-Newton early steps, converging the
+            imaging problem in a few iterations from a cold start;
+            ``"marquardt"`` is the scale-invariant ``H + mu diag(H)``, whose
+            conservative steps need a much larger iteration budget.
+        max_consecutive_rejections
+            Stop after this many consecutive rejected trial steps (each costs
+            a full Jacobian rebuild); at a cost minimum no decreasing step
+            exists and unbounded rejection wastes the runtime driving mu to
+            its ceiling.
         verbose
             Whether to log per-iteration costs.
         visualize_output_dir
@@ -115,6 +132,8 @@ def __init__(
         self.src_image_mesh = src_image_mesh
         self.n_iter = int(n_iter)
         self.tol = float(tol)
+        self.damping = str(damping)
+        self.max_consecutive_rejections = int(max_consecutive_rejections)
         self.verbose = bool(verbose)
         self.visualize_output_dir = visualize_output_dir
         self.visualize_every_n = int(visualize_every_n)
@@ -486,12 +505,14 @@ def solve_joint_optimization(self, xp=np, x0=None, gauge_project_x0=False):
                 )
 
             step_accepted = False
+            consecutive_rejections = 0
             while not step_accepted:
                 delta_x = None
                 try:
                     delta_x = dense_util.solve_lm_step_from(
                         H, minus_gradient, mu,
                         constraint_matrix=constraint_matrix, x=x, xp=xp,
+                        damping=self.damping,
                     )
                     if np.any(np.isnan(np.asarray(delta_x))):
                         delta_x = None
@@ -530,6 +551,30 @@ def solve_joint_optimization(self, xp=np, x0=None, gauge_project_x0=False):
                             self.dpsi_opt = np.asarray(x[n_s:])
                             return self.s_opt, self.dpsi_opt
                     else:
+                        # rejected step below the step tolerance: growing mu
+                        # only shrinks it further — the state is converged
+                        # (at a cost minimum no decreasing step exists), so
+                        # return instead of rejecting to the mu ceiling.
+                        if float(xp.linalg.norm(delta_x)) < self.tol:
+                            if self.verbose:
+                                logger.info(
+                                    "Converged at iteration %d (rejected step "
+                                    "below tolerance).",
+                                    i,
+                                )
+                            self.s_opt = np.asarray(x[:n_s])
+                            self.dpsi_opt = np.asarray(x[n_s:])
+                            return self.s_opt, self.dpsi_opt
+                        consecutive_rejections += 1
+                        if consecutive_rejections >= self.max_consecutive_rejections:
+                            logger.warning(
+                                "%d consecutive rejected LM steps (each a full "
+                                "Jacobian rebuild); stopping at the current state.",
+                                consecutive_rejections,
+                            )
+                            self.s_opt = np.asarray(x[:n_s])
+                            self.dpsi_opt = np.asarray(x[n_s:])
+                            return self.s_opt, self.dpsi_opt
                         mu *= 5.0
                         if mu > 1e15:
                             logger.warning(
@@ -539,6 +584,16 @@ def solve_joint_optimization(self, xp=np, x0=None, gauge_project_x0=False):
                             self.dpsi_opt = np.asarray(x[n_s:])
                             return self.s_opt, self.dpsi_opt
                 else:
+                    consecutive_rejections += 1
+                    if consecutive_rejections >= self.max_consecutive_rejections:
+                        logger.warning(
+                            "%d consecutive failed LM solves; stopping at the "
+                            "current state.",
+                            consecutive_rejections,
+                        )
+                        self.s_opt = np.asarray(x[:n_s])
+                        self.dpsi_opt = np.asarray(x[n_s:])
+                        return self.s_opt, self.dpsi_opt
                     mu *= 5.0
                     if mu > 1e15:
                         logger.warning("LM solver failed repeatedly; stopping.")

autolens/potential_correction/iterative_interferometer.py

@@ -70,6 +70,8 @@ def __init__(
         preloads: Optional[dict] = None,
         n_iter: int = 20,
         tol: float = 1e-6,
+        damping: str = "marquardt",
+        max_consecutive_rejections: int = 10,
         reg_optimize_every: Optional[int] = None,
         reg_optimize_grid: int = 5,
         verbose: bool = False,
@@ -119,7 +121,14 @@ def __init__(
         n_iter
             The maximum number of outer LM iterations.
         tol
-            The step-norm convergence tolerance.
+            The step-norm convergence tolerance; also applied to rejected
+            steps (a sub-tolerance rejected step means converged).
+        damping
+            LM damping matrix (see ``IterFitDpsiSrcImaging``): default
+            ``"marquardt"`` here — visibility-weighted curvatures (~1e11)
+            need the scale-invariant form.
+        max_consecutive_rejections
+            Stop after this many consecutive rejected trial steps.
         reg_optimize_every
             When set, every N accepted outer iterations (and once at the
             start) the regularization strength multipliers (a_src, a_dpsi)
@@ -144,6 +153,8 @@ def __init__(
         self.dpsi_mask = dpsi_mask
         self.n_iter = int(n_iter)
         self.tol = float(tol)
+        self.damping = str(damping)
+        self.max_consecutive_rejections = int(max_consecutive_rejections)
         self.reg_optimize_every = reg_optimize_every
         self.reg_optimize_grid = int(reg_optimize_grid)
         self.reg_scales = (1.0, 1.0)
@@ -565,12 +576,14 @@ def solve_joint_optimization(self, x0=None):
                 )
 
             step_accepted = False
+            consecutive_rejections = 0
             while not step_accepted:
                 delta_x = None
                 try:
                     delta_x = dense_util.solve_lm_step_from(
                         H, minus_gradient, mu,
                         constraint_matrix=constraint_matrix, x=x,
+                        damping=self.damping,
                     )
                     if np.any(np.isnan(np.asarray(delta_x))):
                         delta_x = None
@@ -622,13 +635,39 @@ def solve_joint_optimization(self, x0=None):
                             self._final_state = (F, D, A, R)
                             return self.s_opt, self.dpsi_opt
                     else:
+                        # rejected step below the step tolerance: growing mu
+                        # only shrinks it further — the state is converged.
+                        if float(np.linalg.norm(delta_x)) < self.tol:
+                            if self.verbose:
+                                logger.info(
+                                    "Converged at iteration %d (rejected step "
+                                    "below tolerance).",
+                                    i,
+                                )
+                            break
+                        consecutive_rejections += 1
+                        if consecutive_rejections >= self.max_consecutive_rejections:
+                            logger.warning(
+                                "%d consecutive rejected LM steps; stopping at "
+                                "the current state.",
+                                consecutive_rejections,
+                            )
+                            break
                         mu *= 5.0
                         if mu > 1e15:
                             logger.warning(
                                 "LM damping parameter exceeded 1e15; stopping."
                             )
                             break
                 else:
+                    consecutive_rejections += 1
+                    if consecutive_rejections >= self.max_consecutive_rejections:
+                        logger.warning(
+                            "%d consecutive failed LM solves; stopping at the "
+                            "current state.",
+                            consecutive_rejections,
+                        )
+                        break
                     mu *= 5.0
                     if mu > 1e15:
                         logger.warning("LM solver failed repeatedly; stopping.")

test_autolens/potential_correction/test_dense_util.py

@@ -233,6 +233,38 @@ def test__solve_lm_step_from__unconstrained_and_constrained():
     assert float((C @ (x + step_c))[0]) == pytest.approx(0.0, abs=1.0e-8)
 
 
+def test__solve_lm_step_from__identity_damping():
+    data, noise, mapping, src_reg, dpsi_reg = joint_problem()
+    n_src = src_reg.shape[0]
+    inv_var = 1.0 / noise**2
+    x = np.zeros(mapping.shape[1])
+
+    H, minus_gradient, *_ = dense_util.lm_hessian_and_gradient_from(
+        data, inv_var, x, mapping[:, :n_src], mapping[:, n_src:], src_reg, dpsi_reg
+    )
+
+    mu = 0.7
+    step = dense_util.solve_lm_step_from(H, minus_gradient, mu, damping="identity")
+    assert (H + mu * np.eye(H.shape[0])) @ step == pytest.approx(
+        minus_gradient, rel=1.0e-8
+    )
+
+    # the two damping forms genuinely differ once diag(H) is not ~1
+    step_m = dense_util.solve_lm_step_from(H, minus_gradient, mu, damping="marquardt")
+    assert not np.allclose(step, step_m)
+
+    # constrained identity step stays on the constraint surface
+    C = np.zeros((1, mapping.shape[1]))
+    C[0, n_src:] = 1.0
+    step_c = dense_util.solve_lm_step_from(
+        H, minus_gradient, mu, constraint_matrix=C, x=x, damping="identity"
+    )
+    assert float((C @ (x + step_c))[0]) == pytest.approx(0.0, abs=1.0e-8)
+
+    with pytest.raises(ValueError):
+        dense_util.solve_lm_step_from(H, minus_gradient, mu, damping="not-a-mode")
+
+
 def test__log_evidence_lm_from__matches_hand_computed():
     data, noise, mapping, src_reg, dpsi_reg = joint_problem()
     n_src = src_reg.shape[0]

test_autolens/potential_correction/test_iterative.py

@@ -5,7 +5,7 @@
 import autolens as al
 
 
-def iter_fit_from(masked_imaging, gauge_constraints=False, n_iter=2):
+def iter_fit_from(masked_imaging, gauge_constraints=False, n_iter=2, **kwargs):
     lens = al.Galaxy(
         redshift=0.5,
         mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=1.0),
@@ -25,6 +25,7 @@ def iter_fit_from(masked_imaging, gauge_constraints=False, n_iter=2):
         src_pixelization=src_pixelization,
         gauge_constraints=gauge_constraints,
         n_iter=n_iter,
+        **kwargs,
     )
 
 
@@ -132,3 +133,40 @@ def test__log_evidence__requires_state_or_solve(masked_imaging_7x7):
 
     with pytest.raises(ValueError):
         fit.log_evidence()
+
+
+def test__damping_marquardt__solves_finite(masked_imaging_7x7):
+    fit = iter_fit_from(masked_imaging_7x7, damping="marquardt")
+
+    s_opt, dpsi_opt = fit.solve_joint_optimization()
+
+    assert np.isfinite(s_opt).all()
+    assert np.isfinite(dpsi_opt).all()
+
+
+def test__warm_start_at_optimum__stall_guards_bound_the_rebuilds(
+    masked_imaging_7x7,
+):
+    # solving again from the converged optimum admits no cost-decreasing step;
+    # the tol-on-rejected-step / consecutive-rejection guards must return after
+    # a bounded number of Jacobian rebuilds instead of rejecting mu to 1e15
+    # (mu grows x5 per rejection: reaching 1e15 from 1.0 takes ~22 rejections).
+    fit0 = iter_fit_from(masked_imaging_7x7)
+    s_opt, dpsi_opt = fit0.solve_joint_optimization()
+    x0 = np.concatenate([s_opt, dpsi_opt])
+
+    fit = iter_fit_from(masked_imaging_7x7, n_iter=5, max_consecutive_rejections=3)
+    original = fit.get_L_Js_Jdpsi
+    calls = {"n": 0}
+
+    def counting(*args, **kwargs):
+        calls["n"] += 1
+        return original(*args, **kwargs)
+
+    fit.get_L_Js_Jdpsi = counting
+    s_new, dpsi_new = fit.solve_joint_optimization(x0=x0)
+
+    assert np.isfinite(s_new).all()
+    assert np.isfinite(dpsi_new).all()
+    # init + at most (accepts + rejections-per-iteration capped at 3) trials
+    assert calls["n"] <= 1 + 5 * 4

test_autolens/potential_correction/test_iterative_interferometer.py

@@ -16,7 +16,7 @@
 )
 
 
-def iter_fit_from(dataset, gauge_constraints=False, n_iter=2):
+def iter_fit_from(dataset, gauge_constraints=False, n_iter=2, **kwargs):
     lens = al.Galaxy(
         redshift=0.5,
         mass=al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=1.0),
@@ -36,6 +36,7 @@ def iter_fit_from(dataset, gauge_constraints=False, n_iter=2):
         src_pixelization=src_pixelization,
         gauge_constraints=gauge_constraints,
         n_iter=n_iter,
+        **kwargs,
     )
 
 
@@ -59,6 +60,16 @@ def test__solve_joint_optimization__finite_state_and_decreasing_cost(
     assert np.isfinite(s_opt).all()
     assert np.isfinite(dpsi_opt).all()
 
+
+def test__solve_joint_optimization__identity_damping_finite(interferometer_7):
+    dataset = interferometer_7.apply_sparse_operator()
+    fit = iter_fit_from(dataset, damping="identity", max_consecutive_rejections=3)
+
+    s_opt, dpsi_opt = fit.solve_joint_optimization()
+
+    assert np.isfinite(s_opt).all()
+    assert np.isfinite(dpsi_opt).all()
+
     # the optimized state must beat the zero starting state, whose penalized
     # cost is 0.5 d^H C^-1 d
     x = np.concatenate([s_opt, dpsi_opt])