feature: improved Patch2D

Changed patch2d from using other operators to being
directly implemented (including also an option to apply
Op simultaneously to all windows).

Author: mrava87

pylops/signalprocessing/patch2d.py

@@ -9,8 +9,14 @@
 import numpy as np
 
 from pylops import LinearOperator
-from pylops.basicoperators import BlockDiag, Diagonal, HStack, Restriction
 from pylops.signalprocessing.sliding2d import _slidingsteps
+from pylops.utils._internal import _value_or_sized_to_tuple
+from pylops.utils.backend import (
+    get_array_module,
+    get_sliding_window_view,
+    to_cupy_conditional,
+)
+from pylops.utils.decorators import reshaped
 from pylops.utils.tapers import taper2d
 from pylops.utils.typing import InputDimsLike, NDArray
 
@@ -91,17 +97,7 @@ def patch2d_design(
     return nwins, dims, mwins_inends, dwins_inends
 
 
-def Patch2D(
-    Op: LinearOperator,
-    dims: InputDimsLike,
-    dimsd: InputDimsLike,
-    nwin: Tuple[int, int],
-    nover: Tuple[int, int],
-    nop: Tuple[int, int],
-    tapertype: str = "hanning",
-    scalings: Optional[Sequence[float]] = None,
-    name: str = "P",
-) -> LinearOperator:
+class Patch2D(LinearOperator):
     """2D Patch transform operator.
 
     Apply a transform operator ``Op`` repeatedly to patches of the model
@@ -172,104 +168,154 @@ def Patch2D(
     Patch3D: 3D Patching transform operator.
 
     """
-    # data windows
-    dwin0_ins, dwin0_ends = _slidingsteps(dimsd[0], nwin[0], nover[0])
-    dwin1_ins, dwin1_ends = _slidingsteps(dimsd[1], nwin[1], nover[1])
-    nwins0 = len(dwin0_ins)
-    nwins1 = len(dwin1_ins)
-    nwins = nwins0 * nwins1
-
-    # check patching
-    if nwins0 * nop[0] != dims[0] or nwins1 * nop[1] != dims[1]:
-        raise ValueError(
-            f"Model shape (dims={dims}) is not consistent with chosen "
-            f"number of windows. Run patch2d_design to identify the "
-            f"correct number of windows for the current "
-            "model size..."
-        )
 
-    # create tapers
-    if tapertype is not None:
-        tap = taper2d(nwin[1], nwin[0], nover, tapertype=tapertype).astype(Op.dtype)
-        taps = {itap: tap for itap in range(nwins)}
-        # topmost tapers
-        taptop = tap.copy()
-        taptop[: nover[0]] = tap[nwin[0] // 2]
-        for itap in range(0, nwins1):
-            taps[itap] = taptop
-        # bottommost tapers
-        tapbottom = tap.copy()
-        tapbottom[-nover[0] :] = tap[nwin[0] // 2]
-        for itap in range(nwins - nwins1, nwins):
-            taps[itap] = tapbottom
-        # leftmost tapers
-        tapleft = tap.copy()
-        tapleft[:, : nover[1]] = tap[:, nwin[1] // 2][:, np.newaxis]
-        for itap in range(0, nwins, nwins1):
-            taps[itap] = tapleft
-        # rightmost tapers
-        tapright = tap.copy()
-        tapright[:, -nover[1] :] = tap[:, nwin[1] // 2][:, np.newaxis]
-        for itap in range(nwins1 - 1, nwins, nwins1):
-            taps[itap] = tapright
-        # lefttopcorner taper
-        taplefttop = tap.copy()
-        taplefttop[:, : nover[1]] = tap[:, nwin[1] // 2][:, np.newaxis]
-        taplefttop[: nover[0]] = taplefttop[nwin[0] // 2]
-        taps[0] = taplefttop
-        # righttopcorner taper
-        taprighttop = tap.copy()
-        taprighttop[:, -nover[1] :] = tap[:, nwin[1] // 2][:, np.newaxis]
-        taprighttop[: nover[0]] = taprighttop[nwin[0] // 2]
-        taps[nwins1 - 1] = taprighttop
-        # leftbottomcorner taper
-        tapleftbottom = tap.copy()
-        tapleftbottom[:, : nover[1]] = tap[:, nwin[1] // 2][:, np.newaxis]
-        tapleftbottom[-nover[0] :] = tapleftbottom[nwin[0] // 2]
-        taps[nwins - nwins1] = tapleftbottom
-        # rightbottomcorner taper
-        taprightbottom = tap.copy()
-        taprightbottom[:, -nover[1] :] = tap[:, nwin[1] // 2][:, np.newaxis]
-        taprightbottom[-nover[0] :] = taprightbottom[nwin[0] // 2]
-        taps[nwins - 1] = taprightbottom
-
-    # define scalings
-    if scalings is None:
-        scalings = [1.0] * nwins
-
-    # transform to apply
-    if tapertype is None:
-        OOp = BlockDiag([scalings[itap] * Op for itap in range(nwins)])
-    else:
-        OOp = BlockDiag(
-            [
-                scalings[itap] * Diagonal(taps[itap].ravel(), dtype=Op.dtype) * Op
-                for itap in range(nwins)
-            ]
+    def __init__(
+        self,
+        Op: LinearOperator,
+        dims: InputDimsLike,
+        dimsd: InputDimsLike,
+        nwin: Tuple[int, int],
+        nover: Tuple[int, int],
+        nop: Tuple[int, int],
+        tapertype: str = "hanning",
+        scalings: Optional[Sequence[float]] = None,
+        name: str = "P",
+    ) -> None:
+
+        dims: Tuple[int, ...] = _value_or_sized_to_tuple(dims)
+        dimsd: Tuple[int, ...] = _value_or_sized_to_tuple(dimsd)
+
+        # data windows
+        dwin0_ins, dwin0_ends = _slidingsteps(dimsd[0], nwin[0], nover[0])
+        dwin1_ins, dwin1_ends = _slidingsteps(dimsd[1], nwin[1], nover[1])
+        self.dwins_inends = ((dwin0_ins, dwin0_ends), (dwin1_ins, dwin1_ends))
+        nwins0 = len(dwin0_ins)
+        nwins1 = len(dwin1_ins)
+        nwins = nwins0 * nwins1
+        self.nwin = nwin
+        self.nover = nover
+
+        # check patching
+        if nwins0 * nop[0] != dims[0] or nwins1 * nop[1] != dims[1]:
+            raise ValueError(
+                f"Model shape (dims={dims}) is not consistent with chosen "
+                f"number of windows. Run patch2d_design to identify the "
+                f"correct number of windows for the current "
+                "model size..."
+            )
+
+        # create tapers
+        self.tapertype = tapertype
+        if self.tapertype is not None:
+            tap = taper2d(nwin[1], nwin[0], nover, tapertype=tapertype).astype(Op.dtype)
+            taps = [
+                tap,
+            ] * nwins
+            # topmost tapers
+            taptop = tap.copy()
+            taptop[: nover[0]] = tap[nwin[0] // 2]
+            for itap in range(0, nwins1):
+                taps[itap] = taptop
+            # bottommost tapers
+            tapbottom = tap.copy()
+            tapbottom[-nover[0] :] = tap[nwin[0] // 2]
+            for itap in range(nwins - nwins1, nwins):
+                taps[itap] = tapbottom
+            # leftmost tapers
+            tapleft = tap.copy()
+            tapleft[:, : nover[1]] = tap[:, nwin[1] // 2][:, np.newaxis]
+            for itap in range(0, nwins, nwins1):
+                taps[itap] = tapleft
+            # rightmost tapers
+            tapright = tap.copy()
+            tapright[:, -nover[1] :] = tap[:, nwin[1] // 2][:, np.newaxis]
+            for itap in range(nwins1 - 1, nwins, nwins1):
+                taps[itap] = tapright
+            # lefttopcorner taper
+            taplefttop = tap.copy()
+            taplefttop[:, : nover[1]] = tap[:, nwin[1] // 2][:, np.newaxis]
+            taplefttop[: nover[0]] = taplefttop[nwin[0] // 2]
+            taps[0] = taplefttop
+            # righttopcorner taper
+            taprighttop = tap.copy()
+            taprighttop[:, -nover[1] :] = tap[:, nwin[1] // 2][:, np.newaxis]
+            taprighttop[: nover[0]] = taprighttop[nwin[0] // 2]
+            taps[nwins1 - 1] = taprighttop
+            # leftbottomcorner taper
+            tapleftbottom = tap.copy()
+            tapleftbottom[:, : nover[1]] = tap[:, nwin[1] // 2][:, np.newaxis]
+            tapleftbottom[-nover[0] :] = tapleftbottom[nwin[0] // 2]
+            taps[nwins - nwins1] = tapleftbottom
+            # rightbottomcorner taper
+            taprightbottom = tap.copy()
+            taprightbottom[:, -nover[1] :] = tap[:, nwin[1] // 2][:, np.newaxis]
+            taprightbottom[-nover[0] :] = taprightbottom[nwin[0] // 2]
+            taps[nwins - 1] = taprightbottom
+            self.taps = np.vstack(taps).reshape(nwins0, nwins1, nwin[0], nwin[1])
+
+        # define scalings
+        if scalings is None:
+            self.scalings = [1.0] * nwins
+        else:
+            self.scalings = scalings
+
+        # check if operator is applied to all windows simultaneously
+        self.simOp = False
+        if Op.shape[1] == np.prod(dims):
+            self.simOp = True
+        self.Op = Op
+
+        super().__init__(
+            dtype=Op.dtype,
+            dims=(nwins0, nwins1, int(dims[0] // nwins0), int(dims[1] // nwins1)),
+            dimsd=dimsd,
+            clinear=False,
+            name=name,
         )
 
-    hstack = HStack(
-        [
-            Restriction(
-                (nwin[0], dimsd[1]), range(win_in, win_end), axis=1, dtype=Op.dtype
-            ).H
-            for win_in, win_end in zip(dwin1_ins, dwin1_ends)
-        ]
-    )
-    combining1 = BlockDiag([hstack] * nwins0)
+    @reshaped()
+    def _matvec(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]):
+            for iwin1 in range(self.dims[1]):
+                if self.simOp:
+                    xx = x[iwin0, iwin1].reshape(self.nwin)
+                else:
+                    xx = self.Op.matvec(x[iwin0, iwin1].ravel()).reshape(self.nwin)
+                if self.tapertype is not None:
+                    xxwin = self.taps[iwin0, iwin1] * xx
+                else:
+                    xxwin = xx
 
-    combining0 = HStack(
-        [
-            Restriction(dimsd, range(win_in, win_end), axis=0, dtype=Op.dtype).H
-            for win_in, win_end in zip(dwin0_ins, dwin0_ends)
+                y[
+                    self.dwins_inends[0][0][iwin0] : self.dwins_inends[0][1][iwin0],
+                    self.dwins_inends[1][0][iwin1] : self.dwins_inends[1][1][iwin1],
+                ] += xxwin
+        return y
+
+    @reshaped
+    def _rmatvec(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[0] - self.nover[0], :: self.nwin[1] - self.nover[1]
         ]
-    )
-    Pop = LinearOperator(combining0 * combining1 * OOp)
-    Pop.dims, Pop.dimsd = (
-        nwins0,
-        nwins1,
-        int(dims[0] // nwins0),
-        int(dims[1] // nwins1),
-    ), dimsd
-    Pop.name = name
-    return Pop
+        if self.tapertype is not None:
+            ywins = ywins * self.taps
+        if self.simOp:
+            y = self.Op.H @ ywins
+        else:
+            y = ncp.zeros(self.dims, dtype=self.dtype)
+            for iwin0 in range(self.dims[0]):
+                for iwin1 in range(self.dims[1]):
+                    y[iwin0, iwin1] = self.Op.rmatvec(
+                        ywins[iwin0, iwin1].ravel()
+                    ).reshape(self.dims[2], self.dims[3])
+        return y

pylops/signalprocessing/sliding3d.py

@@ -218,7 +218,7 @@ def __init__(
             tap = taper3d(dimsd[2], nwin, nover, tapertype=tapertype).astype(Op.dtype)
             taps = [
                 tap,
-            ] * nwins  # {itap: tap for itap in range(nwins)}
+            ] * nwins
 
             # topmost tapers
             taptop = tap.copy()