fix: correct internal behaviour of ISTA and FISTA with SOp

Author: mrava87

pylops/optimization/cls_sparsity.py

@@ -1618,6 +1618,10 @@ def setup(
             self.grad = self.ncp.empty_like(x)
             self.x_unthesh = self.ncp.empty_like(x)
             self.xold = self.ncp.empty_like(x)
+            if self.SOp is not None:
+                self.SOpx_unthesh: NDArray = self.ncp.zeros(
+                    self.SOp.shape[1], dtype=self.SOp.dtype
+                )
 
         # create variable to track residual
         if monitorres:
@@ -1657,6 +1661,10 @@ def step(self, x: NDArray, show: bool = False) -> Tuple[NDArray, float]:
             res: NDArray = self.ncp.zeros_like(self.y)
             grad: NDArray = self.ncp.zeros_like(x)
             x_unthesh: NDArray = self.ncp.zeros_like(x)
+            if self.SOp is not None:
+                SOpx_unthesh: NDArray = self.ncp.zeros(
+                    self.SOp.shape[1], dtype=self.SOp.dtype
+                )
 
         # store old vector
         if self.preallocate:
@@ -1703,20 +1711,29 @@ def step(self, x: NDArray, show: bool = False) -> Tuple[NDArray, float]:
                 out=self.x_unthesh if self.preallocate else x_unthesh,
             )
 
+        # apply SOp.H to current x
         if self.SOp is not None:
             if self.preallocate:
-                self.x_unthesh[:] = self.SOprmatvec(self.x_unthesh)
+                self.SOpx_unthesh[:] = self.SOprmatvec(self.x_unthesh)
             else:
-                x_unthesh[:] = self.SOprmatvec(x_unthesh)
-        if self.perc is None and self.decay is not None:
-            x = self.threshf(
-                self.x_unthesh if self.preallocate else x_unthesh,
-                self.decay[self.iiter] * self.thresh,
+                SOpx_unthesh[:] = self.SOprmatvec(x_unthesh)
+        # threshold current solution or current solution projected onto SOp.H space
+        if self.SOp is None:
+            x_unthesh_or_SOpx_unthesh = (
+                self.x_unthesh if self.preallocate else x_unthesh
+            )
+        else:
+            x_unthesh_or_SOpx_unthesh = (
+                self.SOpx_unthesh if self.preallocate else SOpx_unthesh
             )
-        elif self.perc is not None:
+        if self.perc is None:
             x = self.threshf(
-                self.x_unthesh if self.preallocate else x_unthesh, 100 - self.perc
+                x_unthesh_or_SOpx_unthesh,
+                self.decay[self.iiter] * self.thresh,
             )
+        else:
+            x = self.threshf(x_unthesh_or_SOpx_unthesh, 100 - self.perc)
+        # apply SOp to thresholded x
         if self.SOp is not None:
             x = self.SOpmatvec(x)
 
@@ -2022,6 +2039,10 @@ def step(self, x: NDArray, z: NDArray, show: bool = False) -> NDArray:
             res: NDArray = self.ncp.zeros_like(self.y)
             grad: NDArray = self.ncp.zeros_like(x)
             x_unthesh: NDArray = self.ncp.zeros_like(x)
+            if self.SOp is not None:
+                SOpx_unthesh: NDArray = self.ncp.zeros(
+                    self.SOp.shape[1], dtype=self.SOp.dtype
+                )
 
         # store old vector
         if self.preallocate:
@@ -2064,21 +2085,29 @@ def step(self, x: NDArray, z: NDArray, show: bool = False) -> NDArray:
                 out=self.x_unthesh if self.preallocate else x_unthesh,
             )
 
-        # apply regularization operator
+        # apply SOp.H to current x
         if self.SOp is not None:
             if self.preallocate:
-                self.x_unthesh[:] = self.SOprmatvec(self.x_unthesh)
+                self.SOpx_unthesh[:] = self.SOprmatvec(self.x_unthesh)
             else:
-                x_unthesh[:] = self.SOprmatvec(x_unthesh)
-        if self.perc is None and self.decay is not None:
-            x = self.threshf(
-                self.x_unthesh if self.preallocate else x_unthesh,
-                self.decay[self.iiter] * self.thresh,
+                SOpx_unthesh[:] = self.SOprmatvec(x_unthesh)
+        # threshold current solution or current solution projected onto SOp.H space
+        if self.SOp is None:
+            x_unthesh_or_SOpx_unthesh = (
+                self.x_unthesh if self.preallocate else x_unthesh
             )
-        elif self.perc is not None:
+        else:
+            x_unthesh_or_SOpx_unthesh = (
+                self.SOpx_unthesh if self.preallocate else SOpx_unthesh
+            )
+        if self.perc is None:
             x = self.threshf(
-                self.x_unthesh if self.preallocate else x_unthesh, 100 - self.perc
+                x_unthesh_or_SOpx_unthesh,
+                self.decay[self.iiter] * self.thresh,
             )
+        else:
+            x = self.threshf(x_unthesh_or_SOpx_unthesh, 100 - self.perc)
+        # apply SOp to thresholded x
         if self.SOp is not None:
             x = self.SOpmatvec(x)