diff --git a/autolens/potential_correction/iterative.py b/autolens/potential_correction/iterative.py index 7cecb5ab5..6c0816d9f 100644 --- a/autolens/potential_correction/iterative.py +++ b/autolens/potential_correction/iterative.py @@ -354,7 +354,33 @@ def _init_joint_optimization(self): # The Levenberg-Marquardt loop # ----------------------------- - def solve_joint_optimization(self, xp=np, x0=None): + def _gauge_project_dpsi(self, dpsi_vec: np.ndarray) -> np.ndarray: + """ + Project a dpsi vector onto the gauge-constrained subspace + = = = 0 by removing its constant and + linear (deflection-drift) modes. + + These modes are the null space of the Hamiltonian dpsi -> dkappa + map, so the projection leaves the physical convergence correction + unchanged. It is needed when warm-starting a ``gauge_constraints`` + solve from an un-gauged ``x0`` (e.g. the one-shot + ``FitDpsiSrcImaging`` solution, which is not gauge-fixed). Each + accepted LM step re-imposes the gauge on the updated state + (C (x + dx) = 0), but from a near-optimal un-gauged start the + gauge-fixing move slides along the model-degenerate constant / + deflection-drift modes and so is not cost-decreasing: it is rejected + and the solve stalls off the constraint surface. Projecting ``x0`` + up front places the warm start on the surface, where a cold start + from zero already begins. + """ + dpsi_vec = np.asarray(dpsi_vec, dtype=float) + x = self.dpsi_points[:, 1] + y = self.dpsi_points[:, 0] + basis = np.column_stack([np.ones_like(x), x, y]) + coeffs, *_ = np.linalg.lstsq(basis, dpsi_vec, rcond=None) + return dpsi_vec - basis @ coeffs + + def solve_joint_optimization(self, xp=np, x0=None, gauge_project_x0=False): """ Runs the Levenberg-Marquardt loop on the combined state x = [s | dpsi], accepting steps that decrease the penalized cost and @@ -373,17 +399,32 @@ def solve_joint_optimization(self, xp=np, x0=None): the one-shot ``FitDpsiSrcImaging``): the LM then refines inside that basin rather than searching from zero — the recipe certified by the interferometer uv campaign (PyAutoLens#627). + gauge_project_x0 + When warm-starting a ``gauge_constraints`` solve, project the + dpsi component of ``x0`` onto the gauge subspace + = = = 0 before iterating (see + :meth:`_gauge_project_dpsi`). Defaults to ``False`` (``x0`` used + verbatim); set ``True`` when ``x0`` is not itself gauge-fixed — + a near-optimal un-gauged start cannot reach the constraint + surface without a non-cost-decreasing move, which the LM + rejects, so it would otherwise stall with the gauge violated. + The dkappa recovery is unaffected by the projection. """ n_s, n_dpsi = self._init_joint_optimization() if x0 is None: x = xp.zeros(n_s + n_dpsi) else: - x = xp.asarray(np.asarray(x0, dtype=float)) - if x.shape != (n_s + n_dpsi,): + x0 = np.asarray(x0, dtype=float) + if x0.shape != (n_s + n_dpsi,): raise ValueError( - f"x0 has shape {x.shape}; expected ({n_s + n_dpsi},)" + f"x0 has shape {x0.shape}; expected ({n_s + n_dpsi},)" + ) + if gauge_project_x0: + x0 = np.concatenate( + [x0[:n_s], self._gauge_project_dpsi(x0[n_s:])] ) + x = xp.asarray(x0) data_slim = xp.asarray(np.asarray(self.masked_imaging.data.slim)) inv_var = xp.asarray(self.inverse_noise_variance) diff --git a/test_autolens/potential_correction/test_iterative.py b/test_autolens/potential_correction/test_iterative.py index ffaaecf8b..8212a89ad 100644 --- a/test_autolens/potential_correction/test_iterative.py +++ b/test_autolens/potential_correction/test_iterative.py @@ -71,6 +71,53 @@ def test__solve_joint_optimization__gauge_constraints_are_satisfied( assert G @ dpsi_opt == pytest.approx(np.zeros(3), abs=1.0e-6) +def test__gauge_project_dpsi__zeroes_the_three_functionals(masked_imaging_7x7): + fit = iter_fit_from(masked_imaging_7x7, gauge_constraints=True) + fit._init_joint_optimization() + + x = fit.dpsi_points[:, 1] + y = fit.dpsi_points[:, 0] + n_dpsi = fit.dpsi_points.shape[0] + + # an arbitrary field carrying a deliberate constant + linear (gauge) part + rng = np.random.default_rng(0) + dpsi = rng.standard_normal(n_dpsi) + 4.0 + 2.0 * x - 3.0 * y + + projected = fit._gauge_project_dpsi(dpsi) + + functionals = np.array( + [np.sum(projected), np.sum(x * projected), np.sum(y * projected)] + ) + assert functionals == pytest.approx(np.zeros(3), abs=1.0e-8) + # removing only constant + linear modes: a already-gauged field is unchanged + assert fit._gauge_project_dpsi(projected) == pytest.approx(projected, abs=1.0e-8) + + +def test__solve_joint_optimization__gauge_project_x0__gauges_an_ungauged_warm_start( + masked_imaging_7x7, +): + # a cold solve gives a gauge-satisfied optimum; offset its dpsi by a pure + # constant to build a warm start that sits on the optimum but violates the + # constraint (an un-gauged x0, as the one-shot solution is). + fit0 = iter_fit_from(masked_imaging_7x7, gauge_constraints=True) + s_opt, dpsi_opt = fit0.solve_joint_optimization() + x0 = np.concatenate([s_opt, dpsi_opt + 0.7]) + + def gauge_of(fit, dpsi): + n_dpsi = dpsi.shape[0] + G = np.zeros((3, n_dpsi)) + G[0, :] = 1.0 / n_dpsi + G[1, :] = fit.dpsi_points[:, 1] / n_dpsi + G[2, :] = fit.dpsi_points[:, 0] / n_dpsi + return G @ dpsi + + assert np.abs(gauge_of(fit0, dpsi_opt + 0.7)[0]) > 1.0e-2 # x0 is un-gauged + + fit_on = iter_fit_from(masked_imaging_7x7, gauge_constraints=True) + _, dpsi_on = fit_on.solve_joint_optimization(x0=x0, gauge_project_x0=True) + assert gauge_of(fit_on, dpsi_on) == pytest.approx(np.zeros(3), abs=1.0e-6) + + def test__log_evidence__finite_at_optimum(masked_imaging_7x7): fit = iter_fit_from(masked_imaging_7x7)