diff --git a/autolens/potential_correction/dense_util.py b/autolens/potential_correction/dense_util.py index 3b8b64e08..9531870b0 100644 --- a/autolens/potential_correction/dense_util.py +++ b/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) diff --git a/autolens/potential_correction/iterative.py b/autolens/potential_correction/iterative.py index 6c0816d9f..80eefd94b 100644 --- a/autolens/potential_correction/iterative.py +++ b/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.") diff --git a/autolens/potential_correction/iterative_interferometer.py b/autolens/potential_correction/iterative_interferometer.py index 564b0c30b..637ffc247 100644 --- a/autolens/potential_correction/iterative_interferometer.py +++ b/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,6 +635,24 @@ 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( @@ -629,6 +660,14 @@ def solve_joint_optimization(self, x0=None): ) 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.") diff --git a/test_autolens/potential_correction/test_dense_util.py b/test_autolens/potential_correction/test_dense_util.py index e96e4b068..70229eb7d 100644 --- a/test_autolens/potential_correction/test_dense_util.py +++ b/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] diff --git a/test_autolens/potential_correction/test_iterative.py b/test_autolens/potential_correction/test_iterative.py index 8212a89ad..7e88b82ea 100644 --- a/test_autolens/potential_correction/test_iterative.py +++ b/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 diff --git a/test_autolens/potential_correction/test_iterative_interferometer.py b/test_autolens/potential_correction/test_iterative_interferometer.py index 9d93d185f..983d37f53 100644 --- a/test_autolens/potential_correction/test_iterative_interferometer.py +++ b/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])