diff --git a/autolens/potential_correction/iterative.py b/autolens/potential_correction/iterative.py index 5666039d5..7cecb5ab5 100644 --- a/autolens/potential_correction/iterative.py +++ b/autolens/potential_correction/iterative.py @@ -354,7 +354,7 @@ def _init_joint_optimization(self): # The Levenberg-Marquardt loop # ----------------------------- - def solve_joint_optimization(self, xp=np): + def solve_joint_optimization(self, xp=np, x0=None): """ Runs the Levenberg-Marquardt loop on the combined state x = [s | dpsi], accepting steps that decrease the penalized cost and @@ -368,10 +368,22 @@ def solve_joint_optimization(self, xp=np): ``jax.numpy`` for accelerator-ready execution. The autolens ray-tracing / source-mesh updates between steps always run on the host. + x0 + An optional warm-start state [s | dpsi] (e.g. the solution of + 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). """ n_s, n_dpsi = self._init_joint_optimization() - x = xp.zeros(n_s + n_dpsi) + 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,): + raise ValueError( + f"x0 has shape {x.shape}; expected ({n_s + n_dpsi},)" + ) data_slim = xp.asarray(np.asarray(self.masked_imaging.data.slim)) inv_var = xp.asarray(self.inverse_noise_variance)