STUB: Prototype regularize object scaling

src/tike/ptycho/object.py

@@ -56,6 +56,9 @@ class ObjectOptions:
     clip_magnitude: bool = True
     """Whether to force the object magnitude to remain <= 1."""
 
+    lasso_penalty: float = 0.0
+    """Weight of the penalty to keep object coefficients near 1 + 0j."""
+
     def copy_to_device(self, comm):
         """Copy to the current GPU memory."""
         if self.v is not None:

src/tike/ptycho/solvers/lstsq.py

@@ -380,6 +380,7 @@ def _update_nearplane(
             psi_update_denominator,
             probe_update_denominator,
             patches,
+            psi,
             op=op,
             m=m,
             recover_psi=recover_psi,
@@ -642,6 +643,7 @@ def _precondition_nearplane_gradients(
     psi_update_denominator,
     probe_update_denominator,
     patches,
+    psi,
     *,
     op,
     m,
@@ -656,11 +658,15 @@ def _precondition_nearplane_gradients(
     eps = 1e-9 / (diff.shape[-2] * diff.shape[-1])
 
     if recover_psi:
-        common_grad_psi /= ((1 - alpha) * psi_update_denominator +
-                            alpha * psi_update_denominator.max(
-                                axis=(-2, -1),
-                                keepdims=True,
-                            ))
+
+        b = cp.complex64(1.0 + 0.0j)
+
+        common_grad_psi = (common_grad_psi + b - psi) / (
+            (1 - alpha) * psi_update_denominator +
+            alpha * psi_update_denominator.max(
+                axis=(-2, -1),
+                keepdims=True,
+            ) + b)
         dOP = op.diffraction.patch.fwd(
             patches=cp.zeros(patches.shape, dtype='complex64')[..., 0, 0, :, :],
             images=common_grad_psi,

src/tike/ptycho/solvers/rpie.py

@@ -196,6 +196,7 @@ def _update_nearplane(
     position_options=None,
     *,
     probe_options=None,
+    object_options=None,
 ):
 
     patches = comm.pool.map(_get_patches, nearplane_, psi, scan_, op=op)
@@ -233,12 +234,14 @@ def _update_nearplane(
             psi_update_denominator = comm.reduce(psi_update_denominator,
                                                  'gpu')[0]
 
-        psi[0] += step_length * psi_update_numerator / (
+        b = cp.complex64(1.0 + 0.0j)
+
+        psi[0] += step_length * (psi_update_numerator + object_options.lasso_penalty * (b - psi[0])) / (
             (1 - alpha) * psi_update_denominator +
             alpha * psi_update_denominator.max(
                 axis=(-2, -1),
                 keepdims=True,
-            ))
+            ) + object_options.lasso_penalty)
 
         psi = comm.pool.bcast([psi[0]])