feat: added savetaper to Sliding1d and Sliding2d

Author: mrava87

pylops/signalprocessing/sliding1d.py

@@ -127,6 +127,9 @@ class Sliding1D(LinearOperator):
         Number of samples of overlapping part of window
     tapertype : :obj:`str`, optional
         Type of taper (``hanning``, ``cosine``, ``cosinesquare`` or ``None``)
+    savetaper: :obj:`bool`, optional
+        Save all tapers and apply them in one go (``True``) or save unique tapers and apply them one by one (``False``).
+        The first option is more computationally efficient, whilst the second is more memory efficient.
     name : :obj:`str`, optional
         .. versionadded:: 2.0.0
 
@@ -148,6 +151,7 @@ def __init__(
         nwin: int,
         nover: int,
         tapertype: str = "hanning",
+        savetaper: bool = True,
         name: str = "S",
     ) -> None:
 
@@ -172,19 +176,23 @@ def __init__(
 
         # create tapers
         self.tapertype = tapertype
+        self.savetaper = savetaper
         if self.tapertype is not None:
             tap = taper(nwin, nover, tapertype=self.tapertype)
             tapin = tap.copy()
             tapin[:nover] = 1
             tapend = tap.copy()
             tapend[-nover:] = 1
-            self.taps = [
-                tapin,
-            ]
-            for i in range(1, nwins - 1):
-                self.taps.append(tap)
-            self.taps.append(tapend)
-            self.taps = np.vstack(self.taps)
+            if self.savetaper:
+                self.taps = [
+                    tapin,
+                ]
+                for i in range(1, nwins - 1):
+                    self.taps.append(tap)
+                self.taps.append(tapend)
+                self.taps = np.vstack(self.taps)
+            else:
+                self.taps = np.vstack([tapin, tap, tapend])
 
         # check if operator is applied to all windows simultaneously
         self.simOp = False
@@ -200,28 +208,30 @@ def __init__(
             name=name,
         )
 
+        self._register_multiplications(self.savetaper)
+
     @reshaped
-    def _matvec(self, x: NDArray) -> NDArray:
+    def _matvec_savetaper(self, x: NDArray) -> NDArray:
         ncp = get_array_module(x)
         if self.tapertype is not None:
             self.taps = to_cupy_conditional(x, self.taps)
         y = ncp.zeros(self.dimsd, dtype=self.dtype)
         if self.simOp:
             x = self.Op @ x
+            if self.tapertype is not None:
+                x = self.taps * x
         for iwin0 in range(self.dims[0]):
             if self.simOp:
-                xx = x[iwin0]
-            else:
-                xx = self.Op.matvec(x[iwin0])
-            if self.tapertype is not None:
-                xxwin = self.taps[iwin0] * xx
+                xxwin = x[iwin0]
             else:
-                xxwin = xx
+                xxwin = self.Op.matvec(x[iwin0])
+                if self.tapertype is not None:
+                    xxwin = self.taps[iwin0] * xxwin
             y[self.dwin_inends[0][iwin0] : self.dwin_inends[1][iwin0]] += xxwin
         return y
 
     @reshaped
-    def _rmatvec(self, x: NDArray) -> NDArray:
+    def _rmatvec_savetaper(self, x: NDArray) -> NDArray:
         ncp = get_array_module(x)
         ncp_sliding_window_view = get_sliding_window_view(x)
         if self.tapertype is not None:
@@ -236,3 +246,68 @@ def _rmatvec(self, x: NDArray) -> NDArray:
             for iwin0 in range(self.dims[0]):
                 y[iwin0] = self.Op.rmatvec(ywins[iwin0])
         return y
+
+    @reshaped
+    def _matvec_nosavetaper(self, x: NDArray) -> NDArray:
+        ncp = get_array_module(x)
+        if self.tapertype is not None:
+            self.taps = to_cupy_conditional(x, self.taps)
+        y = ncp.zeros(self.dimsd, dtype=self.dtype)
+        if self.simOp:
+            x = self.Op @ x
+        for iwin0 in range(self.dims[0]):
+            if self.simOp:
+                xxwin = x[iwin0]
+            else:
+                xxwin = self.Op.matvec(x[iwin0])
+            if self.tapertype is not None:
+                if iwin0 == 0:
+                    xxwin = self.taps[0] * xxwin
+                elif iwin0 == self.dims[0] - 1:
+                    xxwin = self.taps[-1] * xxwin
+                else:
+                    xxwin = self.taps[1] * xxwin
+            y[self.dwin_inends[0][iwin0] : self.dwin_inends[1][iwin0]] += xxwin
+        return y
+
+    @reshaped
+    def _rmatvec_nosavetaper(self, x: NDArray) -> NDArray:
+        ncp = get_array_module(x)
+        ncp_sliding_window_view = get_sliding_window_view(x)
+        if self.tapertype is not None:
+            self.taps = to_cupy_conditional(x, self.taps)
+        ywins = ncp_sliding_window_view(x, self.nwin)[:: self.nwin - self.nover].copy()
+        if self.simOp:
+            if self.tapertype is not None:
+                for iwin0 in range(self.dims[0]):
+                    if iwin0 == 0:
+                        ywins[0] = ywins[0] * self.taps[0]
+                    elif iwin0 == self.dims[0] - 1:
+                        ywins[-1] = ywins[-1] * self.taps[-1]
+                    else:
+                        ywins[iwin0] = ywins[iwin0] * self.taps[1]
+            y = self.Op.H @ ywins
+        else:
+            y = ncp.zeros(self.dims, dtype=self.dtype)
+            for iwin0 in range(self.dims[0]):
+                if iwin0 == 0:
+                    if self.tapertype is not None:
+                        ywins[0] = ywins[0] * self.taps[0]
+                    y[0] = self.Op.rmatvec(ywins[0])
+                elif iwin0 == self.dims[0] - 1:
+                    if self.tapertype is not None:
+                        ywins[-1] = ywins[-1] * self.taps[-1]
+                    y[-1] = self.Op.rmatvec(ywins[-1])
+                else:
+                    if self.tapertype is not None:
+                        ywins[iwin0] = ywins[iwin0] * self.taps[1]
+                    y[iwin0] = self.Op.rmatvec(ywins[iwin0])
+        return y
+
+    def _register_multiplications(self, savetaper: bool) -> None:
+        if savetaper:
+            self._matvec = self._matvec_savetaper
+            self._rmatvec = self._rmatvec_savetaper
+        else:
+            self._matvec = self._matvec_nosavetaper
+            self._rmatvec = self._rmatvec_nosavetaper

pylops/signalprocessing/sliding2d.py

@@ -163,6 +163,9 @@ class Sliding2D(LinearOperator):
         Number of samples of overlapping part of window
     tapertype : :obj:`str`, optional
         Type of taper (``hanning``, ``cosine``, ``cosinesquare`` or ``None``)
+    savetaper: :obj:`bool`, optional
+        Save all tapers and apply them in one go (``True``) or save unique tapers and apply them one by one (``False``).
+        The first option is more computationally efficient, whilst the second is more memory efficient.
     name : :obj:`str`, optional
         .. versionadded:: 2.0.0
 
@@ -189,6 +192,7 @@ def __init__(
         nwin: int,
         nover: int,
         tapertype: str = "hanning",
+        savetaper: bool = True,
         name: str = "S",
     ) -> None:
 
@@ -213,19 +217,25 @@ def __init__(
 
         # create tapers
         self.tapertype = tapertype
+        self.savetaper = savetaper
         if self.tapertype is not None:
             tap = taper2d(dimsd[1], nwin, nover, tapertype=self.tapertype)
             tapin = tap.copy()
             tapin[:nover] = 1
             tapend = tap.copy()
             tapend[-nover:] = 1
-            self.taps = [
-                tapin[np.newaxis, :],
-            ]
-            for i in range(1, nwins - 1):
-                self.taps.append(tap[np.newaxis, :])
-            self.taps.append(tapend[np.newaxis, :])
-            self.taps = np.concatenate(self.taps, axis=0)
+            if self.savetaper:
+                self.taps = [
+                    tapin[np.newaxis, :],
+                ]
+                for i in range(1, nwins - 1):
+                    self.taps.append(tap[np.newaxis, :])
+                self.taps.append(tapend[np.newaxis, :])
+                self.taps = np.concatenate(self.taps, axis=0)
+            else:
+                self.taps = np.vstack(
+                    [tapin[np.newaxis, :], tap[np.newaxis, :], tapend[np.newaxis, :]]
+                )
 
         # check if operator is applied to all windows simultaneously
         self.simOp = False
@@ -241,8 +251,10 @@ def __init__(
             name=name,
         )
 
+        self._register_multiplications(self.savetaper)
+
     @reshaped
-    def _matvec(self, x: NDArray) -> NDArray:
+    def _matvec_savetaper(self, x: NDArray) -> NDArray:
         ncp = get_array_module(x)
         if self.tapertype is not None:
             self.taps = to_cupy_conditional(x, self.taps)
@@ -262,7 +274,7 @@ def _matvec(self, x: NDArray) -> NDArray:
         return y
 
     @reshaped
-    def _rmatvec(self, x: NDArray) -> NDArray:
+    def _rmatvec_savetaper(self, x: NDArray) -> NDArray:
         ncp = get_array_module(x)
         ncp_sliding_window_view = get_sliding_window_view(x)
         if self.tapertype is not None:
@@ -281,3 +293,80 @@ def _rmatvec(self, x: NDArray) -> NDArray:
                     self.dims[1], self.dims[2]
                 )
         return y
+
+    @reshaped
+    def _matvec_nosavetaper(self, x: NDArray) -> NDArray:
+        ncp = get_array_module(x)
+        if self.tapertype is not None:
+            self.taps = to_cupy_conditional(x, self.taps)
+        y = ncp.zeros(self.dimsd, dtype=self.dtype)
+        if self.simOp:
+            x = self.Op @ x
+        for iwin0 in range(self.dims[0]):
+            if self.simOp:
+                xxwin = x[iwin0].reshape(self.nwin, self.dimsd[-1])
+            else:
+                xxwin = self.Op.matvec(x[iwin0].ravel()).reshape(
+                    self.nwin, self.dimsd[-1]
+                )
+            if self.tapertype is not None:
+                if iwin0 == 0:
+                    xxwin = self.taps[0] * xxwin
+                elif iwin0 == self.dims[0] - 1:
+                    xxwin = self.taps[-1] * xxwin
+                else:
+                    xxwin = self.taps[1] * xxwin
+            y[self.dwin_inends[0][iwin0] : self.dwin_inends[1][iwin0]] += xxwin
+        return y
+
+    @reshaped
+    def _rmatvec_nosavetaper(self, x: NDArray) -> NDArray:
+        ncp = get_array_module(x)
+        ncp_sliding_window_view = get_sliding_window_view(x)
+        if self.tapertype is not None:
+            self.taps = to_cupy_conditional(x, self.taps)
+        ywins = (
+            ncp_sliding_window_view(x, self.nwin, axis=0)[:: self.nwin - self.nover]
+            .transpose(0, 2, 1)
+            .copy()
+        )
+        if self.simOp:
+            if self.tapertype is not None:
+                for iwin0 in range(self.dims[0]):
+                    if iwin0 == 0:
+                        ywins[0] = ywins[0] * self.taps[0]
+                    elif iwin0 == self.dims[0] - 1:
+                        ywins[-1] = ywins[-1] * self.taps[-1]
+                    else:
+                        ywins[iwin0] = ywins[iwin0] * self.taps[1]
+            y = self.Op.H @ ywins
+        else:
+            y = ncp.zeros(self.dims, dtype=self.dtype)
+            for iwin0 in range(self.dims[0]):
+                if iwin0 == 0:
+                    if self.tapertype is not None:
+                        ywins[0] = ywins[0] * self.taps[0]
+                    y[0] = self.Op.rmatvec(ywins[0].ravel()).reshape(
+                        self.dims[1], self.dims[2]
+                    )
+                elif iwin0 == self.dims[0] - 1:
+                    if self.tapertype is not None:
+                        ywins[-1] = ywins[-1] * self.taps[-1]
+                    y[-1] = self.Op.rmatvec(ywins[-1].ravel()).reshape(
+                        self.dims[1], self.dims[2]
+                    )
+                else:
+                    if self.tapertype is not None:
+                        ywins[iwin0] = ywins[iwin0] * self.taps[1]
+                    y[iwin0] = self.Op.rmatvec(ywins[iwin0].ravel()).reshape(
+                        self.dims[1], self.dims[2]
+                    )
+        return y
+
+    def _register_multiplications(self, savetaper: bool) -> None:
+        if savetaper:
+            self._matvec = self._matvec_savetaper
+            self._rmatvec = self._rmatvec_savetaper
+        else:
+            self._matvec = self._matvec_nosavetaper
+            self._rmatvec = self._rmatvec_nosavetaper