feat: added preallocate to SplitBregman

Author: mrava87

pylops/optimization/cls_sparsity.py

@@ -1687,9 +1687,7 @@ def step(self, x: NDArray, show: bool = False) -> Tuple[NDArray, float]:
         if not self.preallocate:
             res: NDArray = self.y - self.Opmatvec(x)
         else:
-            self.ncp.subtract(
-                self.y, self.Opmatvec(x), out=self.res if self.preallocate else res
-            )
+            self.ncp.subtract(self.y, self.Opmatvec(x), out=self.res)
 
         if self.monitorres:
             self.normres = np.linalg.norm(self.res if self.preallocate else res)
@@ -1707,9 +1705,9 @@ def step(self, x: NDArray, show: bool = False) -> Tuple[NDArray, float]:
             grad: NDArray = self.alpha * (self.Oprmatvec(res))
         else:
             self.ncp.multiply(
-                self.Oprmatvec(self.res if self.preallocate else res),
+                self.Oprmatvec(self.res),
                 self.alpha,
-                out=self.grad if self.preallocate else grad,
+                out=self.grad,
             )
 
         # update inverted model
@@ -1718,8 +1716,8 @@ def step(self, x: NDArray, show: bool = False) -> Tuple[NDArray, float]:
         else:
             self.ncp.add(
                 x,
-                self.grad if self.preallocate else grad,
-                out=self.x_unthesh if self.preallocate else x_unthesh,
+                self.grad,
+                out=self.x_unthesh,
             )
 
         # apply SOp.H to current x
@@ -2653,6 +2651,7 @@ def setup(
         tol: float = 1e-10,
         tau: float = 1.0,
         restart: bool = False,
+        preallocate: bool = False,
         show: bool = False,
     ) -> NDArray:
         r"""Setup solver
@@ -2698,6 +2697,10 @@ def setup(
             the initial guess (``True``) or with the last estimate (``False``).
             Note that when this is set to ``True``, the ``x0`` provided in the setup will
             be used in all iterations.
+        preallocate : :obj:`bool`, optional
+            .. versionadded:: 2.5.0
+
+            Pre-allocate all variables used by the solver
         show : :obj:`bool`, optional
             Display setup log
 
@@ -2719,14 +2722,19 @@ def setup(
         self.tol = tol
         self.tau = tau
         self.restart = restart
+
         self.ncp = get_array_module(y)
+        self.isjax = get_module_name(self.ncp) == "jax"
+        self._setpreallocate(preallocate)
 
         # L1 regularizations
         self.nregsL1 = len(RegsL1)
         self.b = [
             self.ncp.zeros(RegL1.shape[0], dtype=self.Op.dtype) for RegL1 in RegsL1
         ]
-        self.d = self.b.copy()
+        self.d = [
+            self.ncp.zeros(RegL1.shape[0], dtype=self.Op.dtype) for RegL1 in RegsL1
+        ]
 
         # L2 regularizations
         self.nregsL2 = 0 if RegsL2 is None else len(RegsL2)
@@ -2797,11 +2805,22 @@ def step(
             Updated model vector
 
         """
+        # add preallocate to keywords of solver
+        if self.preallocate and (engine == "pylops" or self.ncp != np):
+            kwargs_solver["preallocate"] = True
+
         for _ in range(self.niter_inner):
             # regularized problem
-            dataregs = self.dataregsL2 + [
-                self.d[ireg] - self.b[ireg] for ireg in range(self.nregsL1)
-            ]
+            if not self.preallocate:
+                dataregs = self.dataregsL2 + [
+                    self.d[ireg] - self.b[ireg] for ireg in range(self.nregsL1)
+                ]
+            else:
+                for ireg in range(self.nregsL1):
+                    self.ncp.subtract(self.d[ireg], self.b[ireg], out=self.d[ireg])
+                dataregs = self.dataregsL2 + [
+                    self.d[ireg] for ireg in range(self.nregsL1)
+                ]
             x = regularized_inversion(
                 self.Op,
                 self.y,
@@ -2813,11 +2832,18 @@ def step(
                 engine=engine,
                 **kwargs_solver,
             )[0]
-            # Shrinkage
-            for ireg in range(self.nregsL1):
-                self.d[ireg] = _softthreshold(
-                    self.RegsL1[ireg].matvec(x) + self.b[ireg], self.epsRL1s[ireg]
-                )
+            # shrinkage
+            if not self.preallocate:
+                for ireg in range(self.nregsL1):
+                    self.d[ireg] = _softthreshold(
+                        self.RegsL1[ireg].matvec(x) + self.b[ireg], self.epsRL1s[ireg]
+                    )
+            else:
+                for ireg in range(self.nregsL1):
+                    self.ncp.add(
+                        self.RegsL1[ireg].matvec(x), self.b[ireg], out=self.d[ireg]
+                    )
+                    self.d[ireg] = _softthreshold(self.d[ireg], self.epsRL1s[ireg])
 
         # Bregman update
         for ireg in range(self.nregsL1):

pylops/optimization/sparsity.py

@@ -636,6 +636,7 @@ def splitbregman(
     itershow: Tuple[int, int, int] = (10, 10, 10),
     show_inner: bool = False,
     callback: Optional[Callable] = None,
+    preallocate: bool = False,
     **kwargs_lsqr,
 ) -> Tuple[NDArray, int, NDArray]:
     r"""Split Bregman for mixed L2-L1 norms.
@@ -698,6 +699,10 @@ def splitbregman(
     callback : :obj:`callable`, optional
         Function with signature (``callback(x)``) to call after each iteration
         where ``x`` is the current model vector
+    preallocate : :obj:`bool`, optional
+            .. versionadded:: 2.5.0
+
+            Pre-allocate all variables used by the solver
     **kwargs_lsqr
         Arbitrary keyword arguments for
         :py:func:`scipy.sparse.linalg.lsqr` solver used to solve the first
@@ -735,6 +740,7 @@ def splitbregman(
         tau=tau,
         restart=restart,
         engine=engine,
+        preallocate=preallocate,
         show=show,
         itershow=itershow,
         show_inner=show_inner,