Merge pull request #632 from PyAutoLabs/fix/potential-correction-warm-start-imaging

Add opt-in gauge_project_x0 to the imaging iterative warm start

Author: Jammy2211

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
+        <dpsi, 1> = <dpsi, x> = <dpsi, y> = 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
+            <dpsi, 1> = <dpsi, x> = <dpsi, y> = 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)

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
+    # <dpsi,1> 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)