diff --git a/docs/source/gpu.rst b/docs/source/gpu.rst index 864451c4..e861c567 100755 --- a/docs/source/gpu.rst +++ b/docs/source/gpu.rst @@ -5,7 +5,7 @@ GPU Support Overview -------- -PyLops supports computations on GPUs powered by `CuPy `_ (``cupy-cudaXX>=10.6.0``). +PyLops supports computations on GPUs powered by `CuPy `_ (``cupy-cudaXX>=v13.0.0``). This library must be installed *before* PyLops is installed. .. note:: diff --git a/pylops/signalprocessing/patch2d.py b/pylops/signalprocessing/patch2d.py index d95ddeb5..11b93fca 100644 --- a/pylops/signalprocessing/patch2d.py +++ b/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 @@ -22,6 +28,7 @@ def patch2d_design( nwin: Tuple[int, int], nover: Tuple[int, int], nop: Tuple[int, int], + verb: bool = True, ) -> Tuple[ Tuple[int, int], Tuple[int, int], @@ -45,6 +52,9 @@ def patch2d_design( Number of samples of overlapping part of window. nop : :obj:`tuple` Size of model in the transformed domain. + verb : :obj:`bool`, optional + Verbosity flag. If ``verb==True``, print the data + and model windows start-end indices Returns ------- @@ -73,35 +83,26 @@ def patch2d_design( mwins_inends = ((mwin0_ins, mwin0_ends), (mwin1_ins, mwin1_ends)) # print information about patching - logging.warning("%d-%d windows required...", nwins0, nwins1) - logging.warning( - "data wins - start:%s, end:%s / start:%s, end:%s", - dwin0_ins, - dwin0_ends, - dwin1_ins, - dwin1_ends, - ) - logging.warning( - "model wins - start:%s, end:%s / start:%s, end:%s", - mwin0_ins, - mwin0_ends, - mwin1_ins, - mwin1_ends, - ) + if verb: + logging.warning("%d-%d windows required...", nwins0, nwins1) + logging.warning( + "data wins - start:%s, end:%s / start:%s, end:%s", + dwin0_ins, + dwin0_ends, + dwin1_ins, + dwin1_ends, + ) + logging.warning( + "model wins - start:%s, end:%s / start:%s, end:%s", + mwin0_ins, + mwin0_ends, + mwin1_ins, + mwin1_ends, + ) 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 @@ -145,6 +146,11 @@ def Patch2D( Size of model in the transformed domain tapertype : :obj:`str`, optional Type of taper (``hanning``, ``cosine``, ``cosinesquare`` or ``None``) + savetaper: :obj:`bool`, optional + .. versionadded:: 2.3.0 + + 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. scalings : :obj:`tuple` or :obj:`list`, optional Set of scalings to apply to each patch. If ``None``, no scale will be applied @@ -172,104 +178,265 @@ 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", + savetaper: bool = True, + 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 + self.savetaper = savetaper + if self.tapertype is not None: + tap = taper2d(nwin[1], nwin[0], nover, tapertype=tapertype).astype(Op.dtype) + # topmost tapers + taptop = tap.copy() + taptop[: nover[0]] = tap[nwin[0] // 2] + # bottommost tapers + tapbottom = tap.copy() + tapbottom[-nover[0] :] = tap[nwin[0] // 2] + # leftmost tapers + tapleft = tap.copy() + tapleft[:, : nover[1]] = tap[:, nwin[1] // 2][:, np.newaxis] + # rightmost tapers + tapright = tap.copy() + tapright[:, -nover[1] :] = tap[:, nwin[1] // 2][:, np.newaxis] + # lefttopcorner taper + taplefttop = tap.copy() + taplefttop[:, : nover[1]] = tap[:, nwin[1] // 2][:, np.newaxis] + taplefttop[: nover[0]] = taplefttop[nwin[0] // 2] + # righttopcorner taper + taprighttop = tap.copy() + taprighttop[:, -nover[1] :] = tap[:, nwin[1] // 2][:, np.newaxis] + taprighttop[: nover[0]] = taprighttop[nwin[0] // 2] + # leftbottomcorner taper + tapleftbottom = tap.copy() + tapleftbottom[:, : nover[1]] = tap[:, nwin[1] // 2][:, np.newaxis] + tapleftbottom[-nover[0] :] = tapleftbottom[nwin[0] // 2] + # rightbottomcorner taper + taprightbottom = tap.copy() + taprightbottom[:, -nover[1] :] = tap[:, nwin[1] // 2][:, np.newaxis] + taprightbottom[-nover[0] :] = taprightbottom[nwin[0] // 2] + + if self.savetaper: + taps = [ + tap, + ] * nwins + for itap in range(0, nwins1): + taps[itap] = taptop + for itap in range(nwins - nwins1, nwins): + taps[itap] = tapbottom + for itap in range(0, nwins, nwins1): + taps[itap] = tapleft + for itap in range(nwins1 - 1, nwins, nwins1): + taps[itap] = tapright + taps[0] = taplefttop + taps[nwins1 - 1] = taprighttop + taps[nwins - nwins1] = tapleftbottom + taps[nwins - 1] = taprightbottom + self.taps = np.vstack(taps).reshape(nwins0, nwins1, nwin[0], nwin[1]) + else: + taps = [ + taplefttop, + taptop, + taprighttop, + tapleft, + tap, + tapright, + tapleftbottom, + tapbottom, + taprightbottom, + ] + self.taps = np.vstack(taps).reshape(3, 3, nwin[0], nwin[1]) + + # define scalings + self.scalings = [1.0] * nwins if scalings is None else 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) + self._register_multiplications(self.savetaper) - 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) + def _apply_taper(self, ywins, iwin0, iwin1): + if iwin0 == 0 and iwin1 == 0: + ywins[0, 0] = self.taps[0, 0] * ywins[0, 0] + elif iwin0 == 0 and iwin1 == self.dims[1] - 1: + ywins[0, -1] = self.taps[0, -1] * ywins[0, -1] + elif iwin0 == 0: + ywins[0, iwin1] = self.taps[0, 1] * ywins[0, iwin1] + elif iwin0 == self.dims[0] - 1 and iwin1 == 0: + ywins[-1, 0] = self.taps[-1, 0] * ywins[-1, 0] + elif iwin0 == self.dims[0] - 1 and iwin1 == self.dims[1] - 1: + ywins[-1, -1] = self.taps[-1, -1] * ywins[-1, -1] + elif iwin0 == self.dims[0] - 1: + ywins[-1, iwin1] = self.taps[-1, 1] * ywins[-1, iwin1] + elif iwin1 == 0: + ywins[iwin0, 0] = self.taps[1, 0] * ywins[iwin0, 0] + elif iwin1 == self.dims[1] - 1: + ywins[iwin0, -1] = self.taps[1, -1] * ywins[iwin0, -1] + else: + ywins[iwin0, iwin1] = self.taps[1, 1] * ywins[iwin0, iwin1] + return ywins + + @reshaped + 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 + 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 + + 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_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: + 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 + + @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]): + for iwin1 in range(self.dims[1]): + if self.simOp: + xxwin = x[iwin0, iwin1].reshape(self.nwin) + else: + xxwin = self.Op.matvec(x[iwin0, iwin1].ravel()).reshape(self.nwin) + if self.tapertype is not None: + if iwin0 == 0 and iwin1 == 0: + xxwin = self.taps[0, 0] * xxwin + elif iwin0 == 0 and iwin1 == self.dims[1] - 1: + xxwin = self.taps[0, -1] * xxwin + elif iwin0 == 0: + xxwin = self.taps[0, 1] * xxwin + elif iwin0 == self.dims[0] - 1 and iwin1 == 0: + xxwin = self.taps[-1, 0] * xxwin + elif iwin0 == self.dims[0] - 1 and iwin1 == self.dims[1] - 1: + xxwin = self.taps[-1, -1] * xxwin + elif iwin0 == self.dims[0] - 1: + xxwin = self.taps[-1, 1] * xxwin + elif iwin1 == 0: + xxwin = self.taps[1, 0] * xxwin + elif iwin1 == self.dims[1] - 1: + xxwin = self.taps[1, -1] * xxwin + else: + xxwin = self.taps[1, 1] * xxwin + + 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_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[0] - self.nover[0], :: self.nwin[1] - self.nover[1] + ].copy() + if self.simOp: + if self.tapertype is not None: + for iwin0 in range(self.dims[0]): + for iwin1 in range(self.dims[1]): + ywins = self._apply_taper(ywins, iwin0, iwin1) + 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]): + if self.tapertype is not None: + ywins = self._apply_taper(ywins, iwin0, iwin1) + y[iwin0, iwin1] = self.Op.rmatvec( + ywins[iwin0, iwin1].ravel() + ).reshape(self.dims[2], self.dims[3]) + 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 diff --git a/pylops/signalprocessing/patch3d.py b/pylops/signalprocessing/patch3d.py index 011963cc..26cd9086 100644 --- a/pylops/signalprocessing/patch3d.py +++ b/pylops/signalprocessing/patch3d.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 tapernd from pylops.utils.typing import InputDimsLike, NDArray @@ -22,6 +28,7 @@ def patch3d_design( nwin: Tuple[int, int, int], nover: Tuple[int, int, int], nop: Tuple[int, int, int], + verb: bool = True, ) -> Tuple[ Tuple[int, int, int], Tuple[int, int, int], @@ -45,6 +52,9 @@ def patch3d_design( Number of samples of overlapping part of window. nop : :obj:`tuple` Size of model in the transformed domain. + verb : :obj:`bool`, optional + Verbosity flag. If ``verb==True``, print the data + and model windows start-end indices Returns ------- @@ -84,39 +94,30 @@ def patch3d_design( ) # print information about patching - logging.warning("%d-%d-%d windows required...", nwins0, nwins1, nwins2) - logging.warning( - "data wins - start:%s, end:%s / start:%s, end:%s / start:%s, end:%s", - dwin0_ins, - dwin0_ends, - dwin1_ins, - dwin1_ends, - dwin2_ins, - dwin2_ends, - ) - logging.warning( - "model wins - start:%s, end:%s / start:%s, end:%s / start:%s, end:%s", - mwin0_ins, - mwin0_ends, - mwin1_ins, - mwin1_ends, - mwin2_ins, - mwin2_ends, - ) + if verb: + logging.warning("%d-%d-%d windows required...", nwins0, nwins1, nwins2) + logging.warning( + "data wins - start:%s, end:%s / start:%s, end:%s / start:%s, end:%s", + dwin0_ins, + dwin0_ends, + dwin1_ins, + dwin1_ends, + dwin2_ins, + dwin2_ends, + ) + logging.warning( + "model wins - start:%s, end:%s / start:%s, end:%s / start:%s, end:%s", + mwin0_ins, + mwin0_ends, + mwin1_ins, + mwin1_ends, + mwin2_ins, + mwin2_ends, + ) return nwins, dims, mwins_inends, dwins_inends -def Patch3D( - Op, - dims: InputDimsLike, - dimsd: InputDimsLike, - nwin: Tuple[int, int, int], - nover: Tuple[int, int, int], - nop: Tuple[int, int, int], - tapertype: str = "hanning", - scalings: Optional[Sequence[float]] = None, - name: str = "P", -) -> LinearOperator: +class Patch3D(LinearOperator): """3D Patch transform operator. Apply a transform operator ``Op`` repeatedly to patches of the model @@ -160,6 +161,11 @@ def Patch3D( Size of model in the transformed domain tapertype : :obj:`str`, optional Type of taper (``hanning``, ``cosine``, ``cosinesquare`` or ``None``) + savetaper: :obj:`bool`, optional + .. versionadded:: 2.3.0 + + 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. scalings : :obj:`tuple` or :obj:`list`, optional Set of scalings to apply to each patch. If ``None``, no scale will be applied @@ -185,272 +191,578 @@ def Patch3D( Patch2D: 2D 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]) - dwin2_ins, dwin2_ends = _slidingsteps(dimsd[2], nwin[2], nover[2]) - nwins0 = len(dwin0_ins) - nwins1 = len(dwin1_ins) - nwins2 = len(dwin2_ins) - nwins = nwins0 * nwins1 * nwins2 - - # check patching - if ( - nwins0 * nop[0] != dims[0] - or nwins1 * nop[1] != dims[1] - or nwins2 * nop[2] != dims[2] - ): - raise ValueError( - f"Model shape (dims={dims}) is not consistent with chosen " - f"number of windows. Run patch3d_design to identify the " - f"correct number of windows for the current " - "model size..." + + def __init__( + self, + Op: LinearOperator, + dims: InputDimsLike, + dimsd: InputDimsLike, + nwin: Tuple[int, int, int], + nover: Tuple[int, int, int], + nop: Tuple[int, int, int], + tapertype: str = "hanning", + savetaper: bool = True, + 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]) + dwin2_ins, dwin2_ends = _slidingsteps(dimsd[2], nwin[2], nover[2]) + self.dwins_inends = ( + (dwin0_ins, dwin0_ends), + (dwin1_ins, dwin1_ends), + (dwin2_ins, dwin2_ends), ) + nwins0 = len(dwin0_ins) + nwins1 = len(dwin1_ins) + nwins2 = len(dwin2_ins) + nwins = nwins0 * nwins1 * nwins2 + self.nwin = nwin + self.nover = nover - # create tapers - if tapertype is not None: - tap = tapernd(nwin, nover, tapertype=tapertype).astype(Op.dtype) - taps = {itap: tap for itap in range(nwins)} - # 1, sides - # topmost tapers - taptop = tap.copy() - taptop[: nover[0]] = tap[nwin[0] // 2] - for itap in range(0, nwins1 * nwins2): - taps[itap] = taptop - # bottommost tapers - tapbottom = tap.copy() - tapbottom[-nover[0] :] = tap[nwin[0] // 2] - for itap in range(nwins - nwins1 * nwins2, nwins): - taps[itap] = tapbottom - # frontmost tapers - tapfront = tap.copy() - tapfront[:, :, : nover[2]] = tap[:, :, nwin[2] // 2][:, :, np.newaxis] - for itap in range(0, nwins, nwins2): - taps[itap] = tapfront - # backmost tapers - tapback = tap.copy() - tapback[:, :, -nover[2] :] = tap[:, :, nwin[2] // 2][:, :, np.newaxis] - for itap in range(nwins2 - 1, nwins, nwins2): - taps[itap] = tapback - # leftmost tapers - tapleft = tap.copy() - tapleft[:, : nover[1]] = tap[:, nwin[1] // 2][:, np.newaxis, :] - for itap in range(0, nwins, nwins1 * nwins2): - for i in range(nwins2): - taps[itap + i] = tapleft - # rightmost tapers - tapright = tap.copy() - tapright[:, -nover[1] :] = tap[:, nwin[1] // 2][:, np.newaxis, :] - for itap in range(nwins2 * (nwins1 - 1), nwins, nwins2 * nwins1): - for i in range(nwins2): - taps[itap + i] = tapright - # 2. pillars - # topleftmost tapers - taplefttop = tap.copy() - taplefttop[:, : nover[1]] = tap[:, nwin[1] // 2][:, np.newaxis, :] - taplefttop[: nover[0]] = taplefttop[nwin[0] // 2] - for itap in range(nwins2): - taps[itap] = taplefttop - # toprightmost tapers - taprighttop = tap.copy() - taprighttop[:, -nover[1] :] = tap[:, nwin[1] // 2][:, np.newaxis, :] - taprighttop[: nover[0]] = taprighttop[nwin[0] // 2] - for itap in range(nwins2): - taps[nwins2 * (nwins1 - 1) + itap] = taprighttop - # topfrontmost tapers - tapfronttop = tap.copy() - tapfronttop[:, :, : nover[2]] = tap[:, :, nwin[2] // 2][:, :, np.newaxis] - tapfronttop[: nover[0]] = tapfronttop[nwin[0] // 2] - for itap in range(0, nwins1 * nwins2, nwins2): - taps[itap] = tapfronttop - # topbackmost tapers - tapbacktop = tap.copy() - tapbacktop[:, :, -nover[2] :] = tap[:, :, nwin[2] // 2][:, :, np.newaxis] - tapbacktop[: nover[0]] = tapbacktop[nwin[0] // 2] - for itap in range(nwins2 - 1, nwins1 * nwins2, nwins2): - taps[itap] = tapbacktop - # bottomleftmost tapers - tapleftbottom = tap.copy() - tapleftbottom[:, : nover[1]] = tap[:, nwin[1] // 2][:, np.newaxis, :] - tapleftbottom[-nover[0] :] = tapleftbottom[nwin[0] // 2] - for itap in range(nwins2): - taps[(nwins0 - 1) * nwins1 * nwins2 + itap] = tapleftbottom - # bottomrightmost tapers - taprightbottom = tap.copy() - taprightbottom[:, -nover[1] :] = tap[:, nwin[1] // 2][:, np.newaxis, :] - taprightbottom[-nover[0] :] = taprightbottom[nwin[0] // 2] - for itap in range(nwins2): - taps[ - (nwins0 - 1) * nwins1 * nwins2 + (nwins1 - 1) * nwins2 + itap - ] = taprightbottom - # bottomfrontmost tapers - tapfrontbottom = tap.copy() - tapfrontbottom[:, :, : nover[2]] = tap[:, :, nwin[2] // 2][:, :, np.newaxis] - tapfrontbottom[-nover[0] :] = tapfrontbottom[nwin[0] // 2] - for itap in range(0, nwins1 * nwins2, nwins2): - taps[(nwins0 - 1) * nwins1 * nwins2 + itap] = tapfrontbottom - # bottombackmost tapers - tapbackbottom = tap.copy() - tapbackbottom[:, :, -nover[2] :] = tap[:, :, nwin[2] // 2][:, :, np.newaxis] - tapbackbottom[-nover[0] :] = tapbackbottom[nwin[0] // 2] - for itap in range(0, nwins1 * nwins2, nwins2): - taps[(nwins0 - 1) * nwins1 * nwins2 + nwins2 + itap - 1] = tapbackbottom - # leftfrontmost tapers - tapleftfront = tap.copy() - tapleftfront[:, : nover[1]] = tap[:, nwin[1] // 2][:, np.newaxis, :] - tapleftfront[:, :, : nover[2]] = tapleftfront[:, :, nwin[2] // 2][ - :, :, np.newaxis - ] - for itap in range(0, nwins, nwins1 * nwins2): - taps[itap] = tapleftfront - # rightfrontmost tapers - taprightfront = tap.copy() - taprightfront[:, -nover[1] :] = tap[:, nwin[1] // 2][:, np.newaxis, :] - taprightfront[:, :, : nover[2]] = taprightfront[:, :, nwin[2] // 2][ - :, :, np.newaxis - ] - for itap in range(0, nwins, nwins1 * nwins2): - taps[(nwins1 - 1) * nwins2 + itap] = taprightfront - # leftbackmost tapers - tapleftback = tap.copy() - tapleftback[:, : nover[1]] = tap[:, nwin[1] // 2][:, np.newaxis, :] - tapleftback[:, :, -nover[2] :] = tapleftback[:, :, nwin[2] // 2][ - :, :, np.newaxis - ] - for itap in range(0, nwins, nwins1 * nwins2): - taps[nwins2 + itap - 1] = tapleftback - # rightbackmost tapers - taprightback = tap.copy() - taprightback[:, -nover[1] :] = tap[:, nwin[1] // 2][:, np.newaxis, :] - taprightback[:, :, -nover[2] :] = taprightback[:, :, nwin[2] // 2][ - :, :, np.newaxis - ] - for itap in range(0, nwins, nwins1 * nwins2): - taps[(nwins1 - 1) * nwins2 + nwins2 + itap - 1] = taprightback - # 3. corners - # lefttopfrontcorner taper - taplefttop = tap.copy() - taplefttop[: nover[0]] = tap[nwin[0] // 2] - taplefttop[:, : nover[1]] = taplefttop[:, nwin[1] // 2][:, np.newaxis, :] - taplefttop[:, :, : nover[2]] = taplefttop[:, :, nwin[2] // 2][:, :, np.newaxis] - taps[0] = taplefttop - # lefttopbackcorner taper - taplefttop = tap.copy() - taplefttop[: nover[0]] = tap[nwin[0] // 2] - taplefttop[:, : nover[1]] = taplefttop[:, nwin[1] // 2][:, np.newaxis, :] - taplefttop[:, :, -nover[2] :] = taplefttop[:, :, nwin[2] // 2][:, :, np.newaxis] - taps[nwins2 - 1] = taplefttop - # righttopfrontcorner taper - taprighttop = tap.copy() - taprighttop[: nover[0]] = tap[nwin[0] // 2] - taprighttop[:, -nover[1] :] = taprighttop[:, nwin[1] // 2][:, np.newaxis, :] - taprighttop[:, :, : nover[2]] = taprighttop[:, :, nwin[2] // 2][ - :, :, np.newaxis - ] - taps[(nwins1 - 1) * nwins2] = taprighttop - # righttopbackcorner taper - taprighttop = tap.copy() - taprighttop[: nover[0]] = tap[nwin[0] // 2] - taprighttop[:, -nover[1] :] = taprighttop[:, nwin[1] // 2][:, np.newaxis, :] - taprighttop[:, :, -nover[2] :] = taprighttop[:, :, nwin[2] // 2][ - :, :, np.newaxis - ] - taps[(nwins1 - 1) * nwins2 + nwins2 - 1] = taprighttop - # leftbottomfrontcorner taper - tapleftbottom = tap.copy() - tapleftbottom[-nover[0] :] = tap[nwin[0] // 2] - tapleftbottom[:, : nover[1]] = tapleftbottom[:, nwin[1] // 2][:, np.newaxis, :] - tapleftbottom[:, :, : nover[2]] = tapleftbottom[:, :, nwin[2] // 2][ - :, :, np.newaxis - ] - taps[(nwins0 - 1) * nwins1 * nwins2] = tapleftbottom - # leftbottombackcorner taper - tapleftbottom = tap.copy() - tapleftbottom[-nover[0] :] = tap[nwin[0] // 2] - tapleftbottom[:, : nover[1]] = tapleftbottom[:, nwin[1] // 2][:, np.newaxis, :] - tapleftbottom[:, :, -nover[2] :] = tapleftbottom[:, :, nwin[2] // 2][ - :, :, np.newaxis - ] - taps[(nwins0 - 1) * nwins1 * nwins2 + nwins2 - 1] = tapleftbottom - # rightbottomfrontcorner taper - taprightbottom = tap.copy() - taprightbottom[-nover[0] :] = tap[nwin[0] // 2] - taprightbottom[:, -nover[1] :] = taprightbottom[:, nwin[1] // 2][ - :, np.newaxis, : - ] - taprightbottom[:, :, : nover[2]] = taprightbottom[:, :, nwin[2] // 2][ - :, :, np.newaxis - ] - taps[(nwins0 - 1) * nwins1 * nwins2 + (nwins1 - 1) * nwins2] = taprightbottom - # rightbottombackcorner taper - taprightbottom = tap.copy() - taprightbottom[-nover[0] :] = tap[nwin[0] // 2] - taprightbottom[:, -nover[1] :] = taprightbottom[:, nwin[1] // 2][ - :, np.newaxis, : - ] - taprightbottom[:, :, -nover[2] :] = taprightbottom[:, :, nwin[2] // 2][ - :, :, np.newaxis - ] - taps[ - (nwins0 - 1) * nwins1 * nwins2 + (nwins1 - 1) * nwins2 + nwins2 - 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) + # check patching + if ( + nwins0 * nop[0] != dims[0] + or nwins1 * nop[1] != dims[1] + or nwins2 * nop[2] != dims[2] + ): + raise ValueError( + f"Model shape (dims={dims}) is not consistent with chosen " + f"number of windows. Run patch3d_design to identify the " + f"correct number of windows for the current " + "model size..." + ) + + # create tapers + self.tapertype = tapertype + self.savetaper = savetaper + if tapertype is not None: + tap = tapernd(nwin, nover, tapertype=tapertype).astype(Op.dtype) + # 1, sides + # topmost tapers + taptop = tap.copy() + taptop[: nover[0]] = tap[nwin[0] // 2] + # bottommost tapers + tapbottom = tap.copy() + tapbottom[-nover[0] :] = tap[nwin[0] // 2] + # frontmost tapers + tapfront = tap.copy() + tapfront[:, :, : nover[2]] = tap[:, :, nwin[2] // 2][:, :, np.newaxis] + # backmost tapers + tapback = tap.copy() + tapback[:, :, -nover[2] :] = tap[:, :, nwin[2] // 2][:, :, np.newaxis] + # leftmost tapers + tapleft = tap.copy() + tapleft[:, : nover[1]] = tap[:, nwin[1] // 2][:, np.newaxis, :] + # rightmost tapers + tapright = tap.copy() + tapright[:, -nover[1] :] = tap[:, nwin[1] // 2][:, np.newaxis, :] + + # 2. pillars + # topleftmost tapers + taplefttop = tap.copy() + taplefttop[:, : nover[1]] = tap[:, nwin[1] // 2][:, np.newaxis, :] + taplefttop[: nover[0]] = taplefttop[nwin[0] // 2] + # toprightmost tapers + taprighttop = tap.copy() + taprighttop[:, -nover[1] :] = tap[:, nwin[1] // 2][:, np.newaxis, :] + taprighttop[: nover[0]] = taprighttop[nwin[0] // 2] + # topfrontmost tapers + tapfronttop = tap.copy() + tapfronttop[:, :, : nover[2]] = tap[:, :, nwin[2] // 2][:, :, np.newaxis] + tapfronttop[: nover[0]] = tapfronttop[nwin[0] // 2] + # topbackmost tapers + tapbacktop = tap.copy() + tapbacktop[:, :, -nover[2] :] = tap[:, :, nwin[2] // 2][:, :, np.newaxis] + tapbacktop[: nover[0]] = tapbacktop[nwin[0] // 2] + # bottomleftmost tapers + tapleftbottom = tap.copy() + tapleftbottom[:, : nover[1]] = tap[:, nwin[1] // 2][:, np.newaxis, :] + tapleftbottom[-nover[0] :] = tapleftbottom[nwin[0] // 2] + # bottomrightmost tapers + taprightbottom = tap.copy() + taprightbottom[:, -nover[1] :] = tap[:, nwin[1] // 2][:, np.newaxis, :] + taprightbottom[-nover[0] :] = taprightbottom[nwin[0] // 2] + # bottomfrontmost tapers + tapfrontbottom = tap.copy() + tapfrontbottom[:, :, : nover[2]] = tap[:, :, nwin[2] // 2][:, :, np.newaxis] + tapfrontbottom[-nover[0] :] = tapfrontbottom[nwin[0] // 2] + # bottombackmost tapers + tapbackbottom = tap.copy() + tapbackbottom[:, :, -nover[2] :] = tap[:, :, nwin[2] // 2][:, :, np.newaxis] + tapbackbottom[-nover[0] :] = tapbackbottom[nwin[0] // 2] + # leftfrontmost tapers + tapleftfront = tap.copy() + tapleftfront[:, : nover[1]] = tap[:, nwin[1] // 2][:, np.newaxis, :] + tapleftfront[:, :, : nover[2]] = tapleftfront[:, :, nwin[2] // 2][ + :, :, np.newaxis + ] + # rightfrontmost tapers + taprightfront = tap.copy() + taprightfront[:, -nover[1] :] = tap[:, nwin[1] // 2][:, np.newaxis, :] + taprightfront[:, :, : nover[2]] = taprightfront[:, :, nwin[2] // 2][ + :, :, np.newaxis ] + # leftbackmost tapers + tapleftback = tap.copy() + tapleftback[:, : nover[1]] = tap[:, nwin[1] // 2][:, np.newaxis, :] + tapleftback[:, :, -nover[2] :] = tapleftback[:, :, nwin[2] // 2][ + :, :, np.newaxis + ] + # rightbackmost tapers + taprightback = tap.copy() + taprightback[:, -nover[1] :] = tap[:, nwin[1] // 2][:, np.newaxis, :] + taprightback[:, :, -nover[2] :] = taprightback[:, :, nwin[2] // 2][ + :, :, np.newaxis + ] + + # 3. corners + # lefttopfrontcorner taper + taplefttopfront = tap.copy() + taplefttopfront[: nover[0]] = tap[nwin[0] // 2] + taplefttopfront[:, : nover[1]] = taplefttopfront[:, nwin[1] // 2][ + :, np.newaxis, : + ] + taplefttopfront[:, :, : nover[2]] = taplefttopfront[:, :, nwin[2] // 2][ + :, :, np.newaxis + ] + # lefttopbackcorner taper + taplefttopback = tap.copy() + taplefttopback[: nover[0]] = tap[nwin[0] // 2] + taplefttopback[:, : nover[1]] = taplefttopback[:, nwin[1] // 2][ + :, np.newaxis, : + ] + taplefttopback[:, :, -nover[2] :] = taplefttopback[:, :, nwin[2] // 2][ + :, :, np.newaxis + ] + # righttopfrontcorner taper + taprighttopfront = tap.copy() + taprighttopfront[: nover[0]] = tap[nwin[0] // 2] + taprighttopfront[:, -nover[1] :] = taprighttopfront[:, nwin[1] // 2][ + :, np.newaxis, : + ] + taprighttopfront[:, :, : nover[2]] = taprighttopfront[:, :, nwin[2] // 2][ + :, :, np.newaxis + ] + # righttopbackcorner taper + taprighttopback = tap.copy() + taprighttopback[: nover[0]] = tap[nwin[0] // 2] + taprighttopback[:, -nover[1] :] = taprighttopback[:, nwin[1] // 2][ + :, np.newaxis, : + ] + taprighttopback[:, :, -nover[2] :] = taprighttopback[:, :, nwin[2] // 2][ + :, :, np.newaxis + ] + # leftbottomfrontcorner taper + tapleftbottomfront = tap.copy() + tapleftbottomfront[-nover[0] :] = tap[nwin[0] // 2] + tapleftbottomfront[:, : nover[1]] = tapleftbottomfront[:, nwin[1] // 2][ + :, np.newaxis, : + ] + tapleftbottomfront[:, :, : nover[2]] = tapleftbottomfront[ + :, :, nwin[2] // 2 + ][:, :, np.newaxis] + # leftbottombackcorner taper + tapleftbottomback = tap.copy() + tapleftbottomback[-nover[0] :] = tap[nwin[0] // 2] + tapleftbottomback[:, : nover[1]] = tapleftbottomback[:, nwin[1] // 2][ + :, np.newaxis, : + ] + tapleftbottomback[:, :, -nover[2] :] = tapleftbottomback[ + :, :, nwin[2] // 2 + ][:, :, np.newaxis] + # rightbottomfrontcorner taper + taprightbottomfront = tap.copy() + taprightbottomfront[-nover[0] :] = tap[nwin[0] // 2] + taprightbottomfront[:, -nover[1] :] = taprightbottomfront[:, nwin[1] // 2][ + :, np.newaxis, : + ] + taprightbottomfront[:, :, : nover[2]] = taprightbottomfront[ + :, :, nwin[2] // 2 + ][:, :, np.newaxis] + # rightbottombackcorner taper + taprightbottomback = tap.copy() + taprightbottomback[-nover[0] :] = tap[nwin[0] // 2] + taprightbottomback[:, -nover[1] :] = taprightbottomback[:, nwin[1] // 2][ + :, np.newaxis, : + ] + taprightbottomback[:, :, -nover[2] :] = taprightbottomback[ + :, :, nwin[2] // 2 + ][:, :, np.newaxis] + if self.savetaper: + taps = [ + tap, + ] * nwins + for itap in range(0, nwins1 * nwins2): + taps[itap] = taptop + for itap in range(nwins - nwins1 * nwins2, nwins): + taps[itap] = tapbottom + for itap in range(0, nwins, nwins2): + taps[itap] = tapfront + for itap in range(nwins2 - 1, nwins, nwins2): + taps[itap] = tapback + for itap in range(0, nwins, nwins1 * nwins2): + for i in range(nwins2): + taps[itap + i] = tapleft + for itap in range(nwins2 * (nwins1 - 1), nwins, nwins2 * nwins1): + for i in range(nwins2): + taps[itap + i] = tapright + for itap in range(nwins2): + taps[itap] = taplefttop + for itap in range(nwins2): + taps[nwins2 * (nwins1 - 1) + itap] = taprighttop + for itap in range(0, nwins1 * nwins2, nwins2): + taps[itap] = tapfronttop + for itap in range(nwins2 - 1, nwins1 * nwins2, nwins2): + taps[itap] = tapbacktop + for itap in range(nwins2): + taps[(nwins0 - 1) * nwins1 * nwins2 + itap] = tapleftbottom + for itap in range(nwins2): + taps[ + (nwins0 - 1) * nwins1 * nwins2 + (nwins1 - 1) * nwins2 + itap + ] = taprightbottom + for itap in range(0, nwins1 * nwins2, nwins2): + taps[(nwins0 - 1) * nwins1 * nwins2 + itap] = tapfrontbottom + for itap in range(0, nwins1 * nwins2, nwins2): + taps[ + (nwins0 - 1) * nwins1 * nwins2 + nwins2 + itap - 1 + ] = tapbackbottom + for itap in range(0, nwins, nwins1 * nwins2): + taps[itap] = tapleftfront + for itap in range(0, nwins, nwins1 * nwins2): + taps[(nwins1 - 1) * nwins2 + itap] = taprightfront + for itap in range(0, nwins, nwins1 * nwins2): + taps[nwins2 + itap - 1] = tapleftback + for itap in range(0, nwins, nwins1 * nwins2): + taps[(nwins1 - 1) * nwins2 + nwins2 + itap - 1] = taprightback + taps[0] = taplefttopfront + taps[nwins2 - 1] = taplefttopback + taps[(nwins1 - 1) * nwins2] = taprighttopfront + taps[(nwins1 - 1) * nwins2 + nwins2 - 1] = taprighttopback + taps[(nwins0 - 1) * nwins1 * nwins2] = tapleftbottomfront + taps[(nwins0 - 1) * nwins1 * nwins2 + nwins2 - 1] = tapleftbottomback + taps[ + (nwins0 - 1) * nwins1 * nwins2 + (nwins1 - 1) * nwins2 + ] = taprightbottomfront + taps[ + (nwins0 - 1) * nwins1 * nwins2 + (nwins1 - 1) * nwins2 + nwins2 - 1 + ] = taprightbottomback + self.taps = np.vstack(taps).reshape( + nwins0, nwins1, nwins2, nwin[0], nwin[1], nwin[2] + ) + else: + self.taps = np.zeros( + (3, 3, 3, nwin[0], nwin[1], nwin[2]), dtype=Op.dtype + ) + self.taps[0, 0, 0] = taplefttopfront + self.taps[0, 0, 1] = taplefttop + self.taps[0, 0, -1] = taplefttopback + self.taps[0, 1, 0] = tapfronttop + self.taps[0, 1, 1] = taptop + self.taps[0, 1, -1] = tapbacktop + self.taps[0, -1, 0] = taprighttopfront + self.taps[0, -1, 1] = taprighttop + self.taps[0, -1, -1] = taprighttopback + + self.taps[1, 0, 0] = tapleftfront + self.taps[1, 0, 1] = tapleft + self.taps[1, 0, -1] = tapleftback + self.taps[1, 1, 0] = tapfront + self.taps[1, 1, 1] = tap + self.taps[1, 1, -1] = tapback + self.taps[1, -1, 0] = taprightfront + self.taps[1, -1, 1] = tapright + self.taps[1, -1, -1] = taprightback + + self.taps[-1, 0, 0] = tapleftbottomfront + self.taps[-1, 0, 1] = tapleftbottom + self.taps[-1, 0, -1] = tapleftbottomback + self.taps[-1, 1, 0] = tapfrontbottom + self.taps[-1, 1, 1] = tapbottom + self.taps[-1, 1, -1] = tapbackbottom + self.taps[-1, -1, 0] = taprightbottomfront + self.taps[-1, -1, 1] = taprightbottom + self.taps[-1, -1, -1] = taprightbottomback + + # define scalings + self.scalings = [1.0] * nwins if scalings is None else 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, + nwins2, + int(dims[0] // nwins0), + int(dims[1] // nwins1), + int(dims[2] // nwins2), + ), + dimsd=dimsd, + clinear=False, + name=name, ) - hstack2 = HStack( - [ - Restriction( - (nwin[0], nwin[1], dimsd[2]), - range(win_in, win_end), - axis=2, - dtype=Op.dtype, - ).H - for win_in, win_end in zip(dwin2_ins, dwin2_ends) - ] - ) - combining2 = BlockDiag([hstack2] * (nwins1 * nwins0)) - - hstack1 = HStack( - [ - Restriction( - (nwin[0], dimsd[1], dimsd[2]), - range(win_in, win_end), - axis=1, - dtype=Op.dtype, - ).H - for win_in, win_end in zip(dwin1_ins, dwin1_ends) - ] - ) - combining1 = BlockDiag([hstack1] * nwins0) + self._register_multiplications(self.savetaper) + + def _apply_taper(self, ywins, iwin0, iwin1, iwin2): + if iwin0 == 0 and iwin1 == 0 and iwin2 == 0: + ywins[0, 0, 0] = self.taps[0, 0, 0] * ywins[0, 0, 0] + elif iwin0 == 0 and iwin1 == 0 and iwin2 == self.dims[2] - 1: + ywins[0, 0, -1] = self.taps[0, 0, -1] * ywins[0, 0, -1] + elif iwin0 == 0 and iwin1 == self.dims[1] - 1 and iwin2 == self.dims[2] - 1: + ywins[0, -1, -1] = self.taps[0, -1, -1] * ywins[0, -1, -1] + elif iwin0 == 0 and iwin1 == self.dims[1] - 1 and iwin2 == 0: + ywins[0, -1, 0] = self.taps[0, -1, 0] * ywins[0, -1, 0] + elif iwin0 == 0 and iwin1 == 0: + ywins[0, 0, iwin2] = self.taps[0, 0, 1] * ywins[0, 0, iwin2] + elif iwin0 == 0 and iwin1 == self.dims[1] - 1: + ywins[0, -1, iwin2] = self.taps[0, -1, 1] * ywins[0, -1, iwin2] + elif iwin0 == 0 and iwin2 == 0: + ywins[0, iwin1, 0] = self.taps[0, 1, 0] * ywins[0, iwin1, 0] + elif iwin0 == 0 and iwin2 == self.dims[2] - 1: + ywins[0, iwin1, -1] = self.taps[0, 1, -1] * ywins[0, iwin1, -1] + elif iwin0 == 0: + ywins[0, iwin1, iwin2] = self.taps[0, 1, 1] * ywins[0, iwin1, iwin2] + + elif iwin0 == self.dims[0] - 1 and iwin1 == 0 and iwin2 == 0: + ywins[-1, 0, 0] = self.taps[-1, 0, 0] * ywins[-1, 0, 0] + elif iwin0 == self.dims[0] - 1 and iwin1 == 0 and iwin2 == self.dims[2] - 1: + ywins[-1, 0, -1] = self.taps[-1, 0, -1] * ywins[-1, 0, -1] + elif ( + iwin0 == self.dims[0] - 1 + and iwin1 == self.dims[1] - 1 + and iwin2 == self.dims[2] - 1 + ): + ywins[-1, -1, -1] = self.taps[-1, -1, -1] * ywins[-1, -1, -1] + elif iwin0 == self.dims[0] - 1 and iwin1 == self.dims[1] - 1 and iwin2 == 0: + ywins[-1, -1, 0] = self.taps[-1, -1, 0] * ywins[-1, -1, 0] + elif iwin0 == self.dims[0] - 1 and iwin1 == 0: + ywins[-1, 0, iwin2] = self.taps[-1, 0, 1] * ywins[-1, 0, iwin2] + elif iwin0 == self.dims[0] - 1 and iwin1 == self.dims[1] - 1: + ywins[-1, -1, iwin2] = self.taps[-1, -1, 1] * ywins[-1, -1, iwin2] + elif iwin0 == self.dims[0] - 1 and iwin2 == 0: + ywins[-1, iwin1, 0] = self.taps[-1, 1, 0] * ywins[-1, iwin1, 0] + elif iwin0 == self.dims[0] - 1 and iwin2 == self.dims[2] - 1: + ywins[-1, iwin1, -1] = self.taps[-1, 1, -1] * ywins[-1, iwin1, -1] + elif iwin0 == self.dims[0] - 1: + ywins[-1, iwin1, iwin2] = self.taps[-1, 1, 1] * ywins[-1, iwin1, iwin2] - 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) + elif iwin1 == 0 and iwin2 == 0: + ywins[iwin0, 0, 0] = self.taps[1, 0, 0] * ywins[iwin0, 0, 0] + elif iwin1 == 0 and iwin2 == self.dims[2] - 1: + ywins[iwin0, 0, -1] = self.taps[1, 0, -1] * ywins[iwin0, 0, -1] + elif iwin1 == self.dims[1] - 1 and iwin2 == self.dims[2] - 1: + ywins[iwin0, -1, -1] = self.taps[1, -1, -1] * ywins[iwin0, -1, -1] + elif iwin1 == self.dims[1] - 1 and iwin2 == 0: + ywins[iwin0, -1, 0] = self.taps[1, -1, 0] * ywins[iwin0, -1, 0] + elif iwin1 == 0: + ywins[iwin0, 0, iwin2] = self.taps[1, 0, 1] * ywins[iwin0, 0, iwin2] + elif iwin1 == self.dims[1] - 1: + ywins[iwin0, -1, iwin2] = self.taps[1, -1, 1] * ywins[iwin0, -1, iwin2] + elif iwin2 == 0: + ywins[iwin0, iwin1, 0] = self.taps[1, 1, 0] * ywins[iwin0, iwin1, 0] + elif iwin2 == self.dims[2] - 1: + ywins[iwin0, iwin1, -1] = self.taps[1, 1, -1] * ywins[iwin0, iwin1, -1] + else: + ywins[iwin0, iwin1, iwin2] = self.taps[1, 1, 1] * ywins[iwin0, iwin1, iwin2] + return ywins + + @reshaped + 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 + for iwin0 in range(self.dims[0]): + for iwin1 in range(self.dims[1]): + for iwin2 in range(self.dims[2]): + if self.simOp: + xx = x[iwin0, iwin1, iwin2].reshape(self.nwin) + else: + xx = self.Op.matvec(x[iwin0, iwin1, iwin2].ravel()).reshape( + self.nwin + ) + if self.tapertype is not None: + xxwin = self.taps[iwin0, iwin1, iwin2] * xx + else: + xxwin = xx + + 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], + self.dwins_inends[2][0][iwin2] : self.dwins_inends[2][1][iwin2], + ] += xxwin + return y + + @reshaped + 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: + 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], + :: self.nwin[2] - self.nover[2], ] - ) + 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]): + for iwin2 in range(self.dims[2]): + y[iwin0, iwin1, iwin2] = self.Op.rmatvec( + ywins[iwin0, iwin1, iwin2].ravel() + ).reshape(self.dims[3], self.dims[4], self.dims[5]) + 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]): + for iwin1 in range(self.dims[1]): + for iwin2 in range(self.dims[2]): + if self.simOp: + xxwin = x[iwin0, iwin1, iwin2].reshape(self.nwin) + else: + xxwin = self.Op.matvec(x[iwin0, iwin1, iwin2].ravel()).reshape( + self.nwin + ) + if self.tapertype is not None: + if iwin0 == 0 and iwin1 == 0 and iwin2 == 0: + xxwin = self.taps[0, 0, 0] * xxwin + elif iwin0 == 0 and iwin1 == 0 and iwin2 == self.dims[2] - 1: + xxwin = self.taps[0, 0, -1] * xxwin + elif ( + iwin0 == 0 + and iwin1 == self.dims[1] - 1 + and iwin2 == self.dims[2] - 1 + ): + xxwin = self.taps[0, -1, -1] * xxwin + elif iwin0 == 0 and iwin1 == self.dims[1] - 1 and iwin2 == 0: + xxwin = self.taps[0, -1, 0] * xxwin + elif iwin0 == 0 and iwin1 == 0: + xxwin = self.taps[0, 0, 1] * xxwin + elif iwin0 == 0 and iwin1 == self.dims[1] - 1: + xxwin = self.taps[0, -1, 1] * xxwin + elif iwin0 == 0 and iwin2 == 0: + xxwin = self.taps[0, 1, 0] * xxwin + elif iwin0 == 0 and iwin2 == self.dims[2] - 1: + xxwin = self.taps[0, 1, -1] * xxwin + elif iwin0 == 0: + xxwin = self.taps[0, 1, 1] * xxwin + + elif iwin0 == self.dims[0] - 1 and iwin1 == 0 and iwin2 == 0: + xxwin = self.taps[-1, 0, 0] * xxwin + elif ( + iwin0 == self.dims[0] - 1 + and iwin1 == 0 + and iwin2 == self.dims[2] - 1 + ): + xxwin = self.taps[-1, 0, -1] * xxwin + elif ( + iwin0 == self.dims[0] - 1 + and iwin1 == self.dims[1] - 1 + and iwin2 == self.dims[2] - 1 + ): + xxwin = self.taps[-1, -1, -1] * xxwin + elif ( + iwin0 == self.dims[0] - 1 + and iwin1 == self.dims[1] - 1 + and iwin2 == 0 + ): + xxwin = self.taps[-1, -1, 0] * xxwin + elif iwin0 == self.dims[0] - 1 and iwin1 == 0: + xxwin = self.taps[-1, 0, 1] * xxwin + elif iwin0 == self.dims[0] - 1 and iwin1 == self.dims[1] - 1: + xxwin = self.taps[-1, -1, 1] * xxwin + elif iwin0 == self.dims[0] - 1 and iwin2 == 0: + xxwin = self.taps[-1, 1, 0] * xxwin + elif iwin0 == self.dims[0] - 1 and iwin2 == self.dims[2] - 1: + xxwin = self.taps[-1, 1, -1] * xxwin + elif iwin0 == self.dims[0] - 1: + xxwin = self.taps[-1, 1, 1] * xxwin + + elif iwin1 == 0 and iwin2 == 0: + xxwin = self.taps[1, 0, 0] * xxwin + elif iwin1 == 0 and iwin2 == self.dims[2] - 1: + xxwin = self.taps[1, 0, -1] * xxwin + elif iwin1 == self.dims[1] - 1 and iwin2 == self.dims[2] - 1: + xxwin = self.taps[1, -1, -1] * xxwin + elif iwin1 == self.dims[1] - 1 and iwin2 == 0: + xxwin = self.taps[1, -1, 0] * xxwin + elif iwin1 == 0: + xxwin = self.taps[1, 0, 1] * xxwin + elif iwin1 == self.dims[1] - 1: + xxwin = self.taps[1, -1, 1] * xxwin + elif iwin2 == 0: + xxwin = self.taps[1, 1, 0] * xxwin + elif iwin2 == self.dims[2] - 1: + xxwin = self.taps[1, 1, -1] * xxwin + else: + xxwin = self.taps[1, 1, 1] * xxwin + 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], + self.dwins_inends[2][0][iwin2] : self.dwins_inends[2][1][iwin2], + ] += 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[0] - self.nover[0], + :: self.nwin[1] - self.nover[1], + :: self.nwin[2] - self.nover[2], + ].copy() + if self.simOp: + if self.tapertype is not None: + for iwin0 in range(self.dims[0]): + for iwin1 in range(self.dims[1]): + for iwin2 in range(self.dims[2]): + ywins = self._apply_taper(ywins, iwin0, iwin1, iwin2) + 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]): + for iwin2 in range(self.dims[2]): + if self.tapertype is not None: + ywins = self._apply_taper(ywins, iwin0, iwin1, iwin2) + y[iwin0, iwin1, iwin2] = self.Op.rmatvec( + ywins[iwin0, iwin1, iwin2].ravel() + ).reshape(self.dims[3], self.dims[4], self.dims[5]) + return y - Pop = LinearOperator(combining0 * combining1 * combining2 * OOp) - Pop.dims, Pop.dimsd = ( - nwins0, - nwins1, - nwins2, - int(dims[0] // nwins0), - int(dims[1] // nwins1), - int(dims[2] // nwins2), - ), dimsd - Pop.name = name - return Pop + 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 diff --git a/pylops/signalprocessing/sliding1d.py b/pylops/signalprocessing/sliding1d.py index 5cb8c213..b203c1a1 100644 --- a/pylops/signalprocessing/sliding1d.py +++ b/pylops/signalprocessing/sliding1d.py @@ -6,10 +6,17 @@ import logging from typing import Tuple, Union +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 taper from pylops.utils.typing import InputDimsLike, NDArray @@ -21,6 +28,7 @@ def sliding1d_design( nwin: int, nover: int, nop: int, + verb: bool = True, ) -> Tuple[int, int, Tuple[NDArray, NDArray], Tuple[NDArray, NDArray]]: """Design Sliding1D operator @@ -39,6 +47,9 @@ def sliding1d_design( Number of samples of overlapping part of window. nop : :obj:`tuple` Size of model in the transformed domain. + verb : :obj:`bool`, optional + Verbosity flag. If ``verb==True``, print the data + and model windows start-end indices Returns ------- @@ -63,29 +74,22 @@ def sliding1d_design( mwins_inends = (mwin_ins, mwin_ends) # print information about patching - logging.warning("%d windows required...", nwins) - logging.warning( - "data wins - start:%s, end:%s", - dwin_ins, - dwin_ends, - ) - logging.warning( - "model wins - start:%s, end:%s", - mwin_ins, - mwin_ends, - ) + if verb: + logging.warning("%d windows required...", nwins) + logging.warning( + "data wins - start:%s, end:%s", + dwin_ins, + dwin_ends, + ) + logging.warning( + "model wins - start:%s, end:%s", + mwin_ins, + mwin_ends, + ) return nwins, dim, mwins_inends, dwins_inends -def Sliding1D( - Op: LinearOperator, - dim: Union[int, InputDimsLike], - dimd: Union[int, InputDimsLike], - nwin: int, - nover: int, - tapertype: str = "hanning", - name: str = "S", -) -> LinearOperator: +class Sliding1D(LinearOperator): r"""1D Sliding transform operator. Apply a transform operator ``Op`` repeatedly to slices of the model @@ -103,6 +107,12 @@ def Sliding1D( ``nover``, it is recommended to first run ``sliding1d_design`` to obtain the corresponding ``dims`` and number of windows. + .. note:: Two kind of operators ``Op`` can be provided: the first + applies a single transformation to each window separately; the second + applies the transformation to all of the windows at the same time. This + is directly inferred during initialization when the following condition + holds ``Op.shape[1] == dim[0]``. + .. warning:: Depending on the choice of `nwin` and `nover` as well as the size of the data, sliding windows may not cover the entire data. The start and end indices of each window will be displayed and returned @@ -122,16 +132,16 @@ def Sliding1D( Number of samples of overlapping part of window tapertype : :obj:`str`, optional Type of taper (``hanning``, ``cosine``, ``cosinesquare`` or ``None``) + savetaper: :obj:`bool`, optional + .. versionadded:: 2.3.0 + + 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 Name of operator (to be used by :func:`pylops.utils.describe.describe`) - Returns - ------- - Sop : :obj:`pylops.LinearOperator` - Sliding operator - Raises ------ ValueError @@ -139,50 +149,167 @@ def Sliding1D( shape (``dims``). """ - dim: Tuple[int, ...] = _value_or_sized_to_tuple(dim) - dimd: Tuple[int, ...] = _value_or_sized_to_tuple(dimd) - # data windows - dwin_ins, dwin_ends = _slidingsteps(dimd[0], nwin, nover) - nwins = len(dwin_ins) + def __init__( + self, + Op: LinearOperator, + dim: Union[int, InputDimsLike], + dimd: Union[int, InputDimsLike], + nwin: int, + nover: int, + tapertype: str = "hanning", + savetaper: bool = True, + name: str = "S", + ) -> None: - # check windows - if nwins * Op.shape[1] != dim[0]: - raise ValueError( - f"Model shape (dim={dim}) is not consistent with chosen " - f"number of windows. Run sliding1d_design to identify the " - f"correct number of windows for the current " - "model size..." - ) + dim: Tuple[int, ...] = _value_or_sized_to_tuple(dim) + dimd: Tuple[int, ...] = _value_or_sized_to_tuple(dimd) + + # data windows + dwin_ins, dwin_ends = _slidingsteps(dimd[0], nwin, nover) + self.dwin_inends = (dwin_ins, dwin_ends) + nwins = len(dwin_ins) + self.nwin = nwin + self.nover = nover + + # check windows + if nwins * Op.shape[1] != dim[0] and Op.shape[1] != dim[0]: + raise ValueError( + f"Model shape (dim={dim}) is not consistent with chosen " + f"number of windows. Run sliding1d_design to identify the " + f"correct number of windows for the current " + "model size..." + ) - # create tapers - if tapertype is not None: - tap = taper(nwin, nover, tapertype=tapertype).astype(Op.dtype) - tapin = tap.copy() - tapin[:nover] = 1 - tapend = tap.copy() - tapend[-nover:] = 1 - taps = {} - taps[0] = tapin - for i in range(1, nwins - 1): - taps[i] = tap - taps[nwins - 1] = tapend - - # transform to apply - if tapertype is None: - OOp = BlockDiag([Op for _ in range(nwins)]) - else: - OOp = BlockDiag( - [Diagonal(taps[itap].ravel(), dtype=Op.dtype) * Op for itap in range(nwins)] + # 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 + if self.savetaper: + self.taps = [ + tapin, + ] + for _ 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 + if Op.shape[1] == dim[0]: + self.simOp = True + self.Op = Op + + super().__init__( + dtype=Op.dtype, + dims=(nwins, int(dim[0] // nwins)), + dimsd=dimd, + clinear=False, + name=name, ) - combining = HStack( - [ - Restriction(dimd, range(win_in, win_end), dtype=Op.dtype).H - for win_in, win_end in zip(dwin_ins, dwin_ends) - ] - ) - Sop = LinearOperator(combining * OOp) - Sop.dims, Sop.dimsd = (nwins, int(dim[0] // nwins)), dimd - Sop.name = name - return Sop + self._register_multiplications(self.savetaper) + + def _apply_taper(self, ywins, iwin0): + 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] + return ywins + + @reshaped + 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: + xxwin = x[iwin0] + else: + 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_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: + self.taps = to_cupy_conditional(x, self.taps) + ywins = ncp_sliding_window_view(x, self.nwin)[:: self.nwin - self.nover] + 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]): + 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]): + ywins = self._apply_taper(ywins, iwin0) + y = self.Op.H @ ywins + else: + y = ncp.zeros(self.dims, dtype=self.dtype) + for iwin0 in range(self.dims[0]): + if self.tapertype is not None: + ywins = self._apply_taper(ywins, iwin0) + 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 diff --git a/pylops/signalprocessing/sliding2d.py b/pylops/signalprocessing/sliding2d.py index c97802c6..27c40089 100644 --- a/pylops/signalprocessing/sliding2d.py +++ b/pylops/signalprocessing/sliding2d.py @@ -9,7 +9,13 @@ import numpy as np from pylops import LinearOperator -from pylops.basicoperators import BlockDiag, Diagonal, HStack, Restriction +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 @@ -54,6 +60,7 @@ def sliding2d_design( nwin: int, nover: int, nop: Tuple[int, int], + verb: bool = True, ) -> Tuple[int, Tuple[int, int], Tuple[NDArray, NDArray], Tuple[NDArray, NDArray]]: """Design Sliding2D operator @@ -72,6 +79,9 @@ def sliding2d_design( Number of samples of overlapping part of window. nop : :obj:`tuple` Size of model in the transformed domain. + verb : :obj:`bool`, optional + Verbosity flag. If ``verb==True``, print the data + and model windows start-end indices Returns ------- @@ -96,29 +106,22 @@ def sliding2d_design( mwins_inends = (mwin_ins, mwin_ends) # print information about patching - logging.warning("%d windows required...", nwins) - logging.warning( - "data wins - start:%s, end:%s", - dwin_ins, - dwin_ends, - ) - logging.warning( - "model wins - start:%s, end:%s", - mwin_ins, - mwin_ends, - ) + if verb: + logging.warning("%d windows required...", nwins) + logging.warning( + "data wins - start:%s, end:%s", + dwin_ins, + dwin_ends, + ) + logging.warning( + "model wins - start:%s, end:%s", + mwin_ins, + mwin_ends, + ) return nwins, dims, mwins_inends, dwins_inends -def Sliding2D( - Op: LinearOperator, - dims: InputDimsLike, - dimsd: InputDimsLike, - nwin: int, - nover: int, - tapertype: str = "hanning", - name: str = "S", -) -> LinearOperator: +class Sliding2D(LinearOperator): """2D Sliding transform operator. Apply a transform operator ``Op`` repeatedly to slices of the model @@ -139,6 +142,12 @@ def Sliding2D( ``nover``, it is recommended to first run ``sliding2d_design`` to obtain the corresponding ``dims`` and number of windows. + .. note:: Two kind of operators ``Op`` can be provided: the first + applies a single transformation to each window separately; the second + applies the transformation to all of the windows at the same time. This + is directly inferred during initialization when the following condition + holds ``Op.shape[1] == np.prod(dims)``. + .. warning:: Depending on the choice of `nwin` and `nover` as well as the size of the data, sliding windows may not cover the entire data. The start and end indices of each window will be displayed and returned @@ -159,6 +168,11 @@ def Sliding2D( Number of samples of overlapping part of window tapertype : :obj:`str`, optional Type of taper (``hanning``, ``cosine``, ``cosinesquare`` or ``None``) + savetaper: :obj:`bool`, optional + .. versionadded:: 2.3.0 + + 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 @@ -176,47 +190,181 @@ def Sliding2D( shape (``dims``). """ - # data windows - dwin_ins, dwin_ends = _slidingsteps(dimsd[0], nwin, nover) - nwins = len(dwin_ins) - # check patching - if nwins * Op.shape[1] // dims[1] != dims[0]: - raise ValueError( - f"Model shape (dims={dims}) is not consistent with chosen " - f"number of windows. Run sliding2d_design to identify the " - f"correct number of windows for the current " - "model size..." + def __init__( + self, + Op: LinearOperator, + dims: InputDimsLike, + dimsd: InputDimsLike, + nwin: int, + nover: int, + tapertype: str = "hanning", + savetaper: bool = True, + name: str = "S", + ) -> None: + + dims: Tuple[int, ...] = _value_or_sized_to_tuple(dims) + dimsd: Tuple[int, ...] = _value_or_sized_to_tuple(dimsd) + + # data windows + dwin_ins, dwin_ends = _slidingsteps(dimsd[0], nwin, nover) + self.dwin_inends = (dwin_ins, dwin_ends) + nwins = len(dwin_ins) + self.nwin = nwin + self.nover = nover + + # check patching + if nwins * Op.shape[1] // dims[1] != dims[0] and Op.shape[1] != np.prod(dims): + raise ValueError( + f"Model shape (dims={dims}) is not consistent with chosen " + f"number of windows. Run sliding2d_design to identify the " + f"correct number of windows for the current " + "model size..." + ) + + # 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 + if self.savetaper: + self.taps = [ + tapin[np.newaxis, :], + ] + for _ 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 + if Op.shape[1] == np.prod(dims): + self.simOp = True + self.Op = Op + + super().__init__( + dtype=Op.dtype, + dims=(nwins, int(dims[0] // nwins), dims[1]), + dimsd=dimsd, + clinear=False, + name=name, ) - # create tapers - if tapertype is not None: - tap = taper2d(dimsd[1], nwin, nover, tapertype=tapertype).astype(Op.dtype) - tapin = tap.copy() - tapin[:nover] = 1 - tapend = tap.copy() - tapend[-nover:] = 1 - taps = {} - taps[0] = tapin - for i in range(1, nwins - 1): - taps[i] = tap - taps[nwins - 1] = tapend - - # transform to apply - if tapertype is None: - OOp = BlockDiag([Op for _ in range(nwins)]) - else: - OOp = BlockDiag( - [Diagonal(taps[itap].ravel(), dtype=Op.dtype) * Op for itap in range(nwins)] + self._register_multiplications(self.savetaper) + + def _apply_taper(self, ywins, iwin0): + 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] + return ywins + + @reshaped + 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 + for iwin0 in range(self.dims[0]): + if self.simOp: + xx = x[iwin0].reshape(self.nwin, self.dimsd[-1]) + else: + xx = self.Op.matvec(x[iwin0].ravel()).reshape(self.nwin, self.dimsd[-1]) + if self.tapertype is not None: + xxwin = self.taps[iwin0] * xx + else: + xxwin = xx + y[self.dwin_inends[0][iwin0] : self.dwin_inends[1][iwin0]] += xxwin + return y + + @reshaped + 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: + 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) + 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]): + y[iwin0] = self.Op.rmatvec(ywins[iwin0].ravel()).reshape( + 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() ) - - combining = HStack( - [ - Restriction(dimsd, range(win_in, win_end), axis=0, dtype=Op.dtype).H - for win_in, win_end in zip(dwin_ins, dwin_ends) - ] - ) - Sop = LinearOperator(combining * OOp) - Sop.dims, Sop.dimsd = (nwins, int(dims[0] // nwins), dims[1]), dimsd - Sop.name = name - return Sop + if self.simOp: + if self.tapertype is not None: + for iwin0 in range(self.dims[0]): + ywins = self._apply_taper(ywins, iwin0) + y = self.Op.H @ ywins + else: + y = ncp.zeros(self.dims, dtype=self.dtype) + for iwin0 in range(self.dims[0]): + if self.tapertype is not None: + ywins = self._apply_taper(ywins, iwin0) + 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 diff --git a/pylops/signalprocessing/sliding3d.py b/pylops/signalprocessing/sliding3d.py index 7d21018c..db052180 100644 --- a/pylops/signalprocessing/sliding3d.py +++ b/pylops/signalprocessing/sliding3d.py @@ -6,9 +6,17 @@ import logging from typing import Tuple +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 taper3d from pylops.utils.typing import InputDimsLike, NDArray @@ -20,6 +28,7 @@ def sliding3d_design( nwin: Tuple[int, int], nover: Tuple[int, int], nop: Tuple[int, int, int], + verb: bool = True, ) -> Tuple[ Tuple[int, int], Tuple[int, int, int], @@ -43,6 +52,9 @@ def sliding3d_design( Number of samples of overlapping part of window. nop : :obj:`tuple` Size of model in the transformed domain. + verb : :obj:`bool`, optional + Verbosity flag. If ``verb==True``, print the data + and model windows start-end indices Returns ------- @@ -71,35 +83,26 @@ def sliding3d_design( mwins_inends = ((mwin0_ins, mwin0_ends), (mwin1_ins, mwin1_ends)) # print information about patching - logging.warning("%d-%d windows required...", nwins0, nwins1) - logging.warning( - "data wins - start:%s, end:%s / start:%s, end:%s", - dwin0_ins, - dwin0_ends, - dwin1_ins, - dwin1_ends, - ) - logging.warning( - "model wins - start:%s, end:%s / start:%s, end:%s", - mwin0_ins, - mwin0_ends, - mwin1_ins, - mwin1_ends, - ) + if verb: + logging.warning("%d-%d windows required...", nwins0, nwins1) + logging.warning( + "data wins - start:%s, end:%s / start:%s, end:%s", + dwin0_ins, + dwin0_ends, + dwin1_ins, + dwin1_ends, + ) + logging.warning( + "model wins - start:%s, end:%s / start:%s, end:%s", + mwin0_ins, + mwin0_ends, + mwin1_ins, + mwin1_ends, + ) return nwins, dims, mwins_inends, dwins_inends -def Sliding3D( - Op: LinearOperator, - dims: InputDimsLike, - dimsd: InputDimsLike, - nwin: Tuple[int, int], - nover: Tuple[int, int], - nop: Tuple[int, int, int], - tapertype: str = "hanning", - nproc: int = 1, - name: str = "P", -) -> LinearOperator: +class Sliding3D(LinearOperator): """3D Sliding transform operator.w Apply a transform operator ``Op`` repeatedly to patches of the model @@ -121,6 +124,12 @@ def Sliding3D( ``nover``, it is recommended to first run ``sliding3d_design`` to obtain the corresponding ``dims`` and number of windows. + .. note:: Two kind of operators ``Op`` can be provided: the first + applies a single transformation to each window separately; the second + applies the transformation to all of the windows at the same time. This + is directly inferred during initialization when the following condition + holds ``Op.shape[1] == np.prod(dims)``. + .. warning:: Depending on the choice of `nwin` and `nover` as well as the size of the data, sliding windows may not cover the entire data. The start and end indices of each window will be displayed and returned @@ -145,9 +154,14 @@ def Sliding3D( to spatial axes in the data tapertype : :obj:`str`, optional Type of taper (``hanning``, ``cosine``, ``cosinesquare`` or ``None``) + savetaper: :obj:`bool`, optional + .. versionadded:: 2.3.0 + + 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. nproc : :obj:`int`, optional - Number of processes used to evaluate the N operators in parallel - using ``multiprocessing``. If ``nproc=1``, work in serial mode. + *Deprecated*, will be removed in v3.0.0. Simply kept for + back-compatibility with previous implementation name : :obj:`str`, optional .. versionadded:: 2.0.0 @@ -165,66 +179,287 @@ def Sliding3D( shape (``dims``). """ - # 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 windows - if nwins * Op.shape[1] // dims[2] != dims[0] * dims[1]: - raise ValueError( - f"Model shape (dims={dims}) is not consistent with chosen " - f"number of windows. Run sliding3d_design to identify the " - f"correct number of windows for the current " - "model size..." + + def __init__( + self, + Op: LinearOperator, + dims: InputDimsLike, + dimsd: InputDimsLike, + nwin: Tuple[int, int], + nover: Tuple[int, int], + nop: Tuple[int, int, int], + tapertype: str = "hanning", + savetaper: bool = True, + nproc: int = 1, + 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 + + # model windows + mwin0_ins, mwin0_ends = _slidingsteps(dims[0], nop[0], 0) + mwin1_ins, mwin1_ends = _slidingsteps(dims[1], nop[1], 0) + self.mwins_inends = ((mwin0_ins, mwin0_ends), (mwin1_ins, mwin1_ends)) + + # check windows + if nwins * Op.shape[1] // dims[2] != dims[0] * dims[1] and Op.shape[ + 1 + ] != np.prod(dims): + raise ValueError( + f"Model shape (dims={dims}) is not consistent with chosen " + f"number of windows. Run sliding3d_design to identify the " + f"correct number of windows for the current " + "model size..." + ) + + # create tapers + self.tapertype = tapertype + self.savetaper = savetaper + if self.tapertype is not None: + tap = taper3d(dimsd[2], nwin, nover, tapertype=tapertype).astype(Op.dtype) + # topmost tapers + taptop = tap.copy() + taptop[: nover[0]] = tap[nwin[0] // 2] + # bottommost tapers + tapbottom = tap.copy() + tapbottom[-nover[0] :] = tap[nwin[0] // 2] + # leftmost tapers + tapleft = tap.copy() + tapleft[:, : nover[1]] = tap[:, nwin[1] // 2][:, np.newaxis] + # rightmost tapers + tapright = tap.copy() + tapright[:, -nover[1] :] = tap[:, nwin[1] // 2][:, np.newaxis] + # lefttopcorner taper + taplefttop = tap.copy() + taplefttop[:, : nover[1]] = tap[:, nwin[1] // 2][:, np.newaxis] + taplefttop[: nover[0]] = taplefttop[nwin[0] // 2] + # righttopcorner taper + taprighttop = tap.copy() + taprighttop[:, -nover[1] :] = tap[:, nwin[1] // 2][:, np.newaxis] + taprighttop[: nover[0]] = taprighttop[nwin[0] // 2] + # leftbottomcorner taper + tapleftbottom = tap.copy() + tapleftbottom[:, : nover[1]] = tap[:, nwin[1] // 2][:, np.newaxis] + tapleftbottom[-nover[0] :] = tapleftbottom[nwin[0] // 2] + # rightbottomcorner taper + taprightbottom = tap.copy() + taprightbottom[:, -nover[1] :] = tap[:, nwin[1] // 2][:, np.newaxis] + taprightbottom[-nover[0] :] = taprightbottom[nwin[0] // 2] + + if self.savetaper: + taps = [ + tap, + ] * nwins + + for itap in range(0, nwins1): + taps[itap] = taptop + for itap in range(nwins - nwins1, nwins): + taps[itap] = tapbottom + for itap in range(0, nwins, nwins1): + taps[itap] = tapleft + for itap in range(nwins1 - 1, nwins, nwins1): + taps[itap] = tapright + taps[0] = taplefttop + taps[nwins1 - 1] = taprighttop + taps[nwins - nwins1] = tapleftbottom + taps[nwins - 1] = taprightbottom + self.taps = np.vstack(taps).reshape( + nwins0, nwins1, nwin[0], nwin[1], dimsd[2] + ) + else: + taps = [ + taplefttop, + taptop, + taprighttop, + tapleft, + tap, + tapright, + tapleftbottom, + tapbottom, + taprightbottom, + ] + self.taps = np.vstack(taps).reshape(3, 3, nwin[0], nwin[1], dimsd[2]) + # 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), + dims[2], + ), + dimsd=dimsd, + clinear=False, + name=name, ) - # create tapers - if tapertype is not None: - tap = taper3d(dimsd[2], nwin, nover, tapertype=tapertype).astype(Op.dtype) - - # transform to apply - if tapertype is None: - OOp = BlockDiag([Op for _ in range(nwins)], nproc=nproc) - else: - OOp = BlockDiag( - [Diagonal(tap.ravel(), dtype=Op.dtype) * Op for _ in range(nwins)], - nproc=nproc, + self._register_multiplications(self.savetaper) + + def _apply_taper(self, ywins, iwin0, iwin1): + if iwin0 == 0 and iwin1 == 0: + ywins[0, 0] = self.taps[0, 0] * ywins[0, 0] + elif iwin0 == 0 and iwin1 == self.dims[1] - 1: + ywins[0, -1] = self.taps[0, -1] * ywins[0, -1] + elif iwin0 == 0: + ywins[0, iwin1] = self.taps[0, 1] * ywins[0, iwin1] + elif iwin0 == self.dims[0] - 1 and iwin1 == 0: + ywins[-1, 0] = self.taps[-1, 0] * ywins[-1, 0] + elif iwin0 == self.dims[0] - 1 and iwin1 == self.dims[1] - 1: + ywins[-1, -1] = self.taps[-1, -1] * ywins[-1, -1] + elif iwin0 == self.dims[0] - 1: + ywins[-1, iwin1] = self.taps[-1, 1] * ywins[-1, iwin1] + elif iwin1 == 0: + ywins[iwin0, 0] = self.taps[1, 0] * ywins[iwin0, 0] + elif iwin1 == self.dims[1] - 1: + ywins[iwin0, -1] = self.taps[1, -1] * ywins[iwin0, -1] + else: + ywins[iwin0, iwin1] = self.taps[1, 1] * ywins[iwin0, iwin1] + return ywins + + @reshaped + 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 + for iwin0 in range(self.dims[0]): + for iwin1 in range(self.dims[1]): + if self.simOp: + xx = x[iwin0, iwin1].reshape( + self.nwin[0], self.nwin[1], self.dimsd[-1] + ) + else: + xx = self.Op.matvec(x[iwin0, iwin1].ravel()).reshape( + self.nwin[0], self.nwin[1], self.dimsd[-1] + ) + if self.tapertype is not None: + xxwin = self.taps[iwin0, iwin1] * xx + else: + xxwin = xx + 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_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: + self.taps = to_cupy_conditional(x, self.taps) + ywins = ncp_sliding_window_view(x, self.nwin, axis=(0, 1))[ + :: self.nwin[0] - self.nover[0], :: self.nwin[1] - self.nover[1] + ].transpose(0, 1, 3, 4, 2) + 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], self.dims[4]) + 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]): + for iwin1 in range(self.dims[1]): + if self.simOp: + xxwin = x[iwin0, iwin1].reshape( + self.nwin[0], self.nwin[1], self.dimsd[-1] + ) + else: + xxwin = self.Op.matvec(x[iwin0, iwin1].ravel()).reshape( + self.nwin[0], self.nwin[1], self.dimsd[-1] + ) + if self.tapertype is not None: + if iwin0 == 0 and iwin1 == 0: + xxwin = self.taps[0, 0] * xxwin + elif iwin0 == 0 and iwin1 == self.dims[1] - 1: + xxwin = self.taps[0, -1] * xxwin + elif iwin0 == 0: + xxwin = self.taps[0, 1] * xxwin + elif iwin0 == self.dims[0] - 1 and iwin1 == 0: + xxwin = self.taps[-1, 0] * xxwin + elif iwin0 == self.dims[0] - 1 and iwin1 == self.dims[1] - 1: + xxwin = self.taps[-1, -1] * xxwin + elif iwin0 == self.dims[0] - 1: + xxwin = self.taps[-1, 1] * xxwin + elif iwin1 == 0: + xxwin = self.taps[1, 0] * xxwin + elif iwin1 == self.dims[1] - 1: + xxwin = self.taps[1, -1] * xxwin + else: + xxwin = self.taps[1, 1] * xxwin + 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_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, 1))[ + :: self.nwin[0] - self.nover[0], :: self.nwin[1] - self.nover[1] + ] + .transpose(0, 1, 3, 4, 2) + .copy() ) + if self.simOp: + if self.tapertype is not None: + for iwin0 in range(self.dims[0]): + for iwin1 in range(self.dims[1]): + ywins = self._apply_taper(ywins, iwin0, iwin1) + 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]): + if self.tapertype is not None: + ywins = self._apply_taper(ywins, iwin0, iwin1) + y[iwin0, iwin1] = self.Op.rmatvec( + ywins[iwin0, iwin1].ravel() + ).reshape(self.dims[2], self.dims[3], self.dims[4]) + return y - hstack = HStack( - [ - Restriction( - (nwin[0], dimsd[1], dimsd[2]), - 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) - 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) - ] - ) - Sop = LinearOperator(combining0 * combining1 * OOp) - Sop.dims, Sop.dimsd = ( - nwins0, - nwins1, - int(dims[0] // nwins0), - int(dims[1] // nwins1), - dims[2], - ), dimsd - Sop.name = name - return Sop + 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 diff --git a/pylops/utils/backend.py b/pylops/utils/backend.py index 1cef8bdb..00bdb064 100644 --- a/pylops/utils/backend.py +++ b/pylops/utils/backend.py @@ -7,6 +7,7 @@ "get_oaconvolve", "get_correlate", "get_add_at", + "get_sliding_window_view", "get_block_diag", "get_toeplitz", "get_csc_matrix", @@ -228,6 +229,29 @@ def get_add_at(x: npt.ArrayLike) -> Callable: return cupyx.scatter_add +def get_sliding_window_view(x: npt.ArrayLike) -> Callable: + """Returns correct sliding_window_view module based on input + + Parameters + ---------- + x : :obj:`numpy.ndarray` + Array + + Returns + ------- + mod : :obj:`func` + Module to be used to process array (:mod:`numpy` or :mod:`cupy`) + + """ + if not deps.cupy_enabled: + return np.lib.stride_tricks.sliding_window_view + + if cp.get_array_module(x) == np: + return np.lib.stride_tricks.sliding_window_view + else: + return cp.lib.stride_tricks.sliding_window_view + + def get_block_diag(x: npt.ArrayLike) -> Callable: """Returns correct block_diag module based on input diff --git a/pytests/test_patching.py b/pytests/test_patching.py index 25656724..de61243e 100755 --- a/pytests/test_patching.py +++ b/pytests/test_patching.py @@ -25,6 +25,7 @@ "novert": 0, # "winst": 2, "tapertype": None, + "savetaper": True, } # no overlap, no taper par2 = { "ny": 6, @@ -43,6 +44,7 @@ "novert": 0, # "winst": 2, "tapertype": "hanning", + "savetaper": True, } # no overlap, with taper par3 = { "ny": 6, @@ -61,8 +63,28 @@ "novert": 2, # "winst": 4, "tapertype": None, + "savetaper": True, } # overlap, no taper par4 = { + "ny": 6, + "nx": 7, + "nt": 10, + "npy": 15, + "nwiny": 7, + "novery": 3, + # "winsy": 3, + "npx": 13, + "nwinx": 5, + "noverx": 2, + # "winsx": 3, + "npt": 10, + "nwint": 4, + "novert": 2, + # "winst": 4, + "tapertype": None, + "savetaper": False, +} # overlap, no taper (non saved +par5 = { "ny": 6, "nx": 7, "nt": 10, @@ -79,10 +101,30 @@ "novert": 2, # "winst": 4, "tapertype": "hanning", + "savetaper": True, } # overlap, with taper +par6 = { + "ny": 6, + "nx": 7, + "nt": 10, + "npy": 15, + "nwiny": 7, + "novery": 3, + # "winsy": 3, + "npx": 13, + "nwinx": 5, + "noverx": 2, + # "winsx": 3, + "npt": 10, + "nwint": 4, + "novert": 2, + # "winst": 4, + "tapertype": "hanning", + "savetaper": False, +} # overlap, with taper (non saved) -@pytest.mark.parametrize("par", [(par1), (par2), (par3), (par4)]) +@pytest.mark.parametrize("par", [(par1), (par2), (par3), (par4), (par5), (par6)]) def test_Patch2D(par): """Dot-test and inverse for Patch2D operator""" Op = MatrixMult(np.ones((par["nwiny"] * par["nwint"], par["ny"] * par["nt"]))) @@ -101,6 +143,7 @@ def test_Patch2D(par): nover=(par["novery"], par["novert"]), nop=(par["ny"], par["nt"]), tapertype=par["tapertype"], + savetaper=par["savetaper"], ) assert dottest( Pop, @@ -134,6 +177,7 @@ def test_Patch2D_scalings(par): nover=(par["novery"], par["novert"]), nop=(par["ny"], par["nt"]), tapertype=par["tapertype"], + savetaper=par["savetaper"], scalings=scalings, ) assert dottest( @@ -148,7 +192,7 @@ def test_Patch2D_scalings(par): assert_array_almost_equal(x.ravel(), xinv) -@pytest.mark.parametrize("par", [(par1), (par2), (par3), (par4)]) +@pytest.mark.parametrize("par", [(par1), (par2), (par3), (par4), (par5), (par6)]) def test_Patch3D(par): """Dot-test and inverse for Patch3D operator""" Op = MatrixMult( @@ -179,6 +223,7 @@ def test_Patch3D(par): nover=(par["novery"], par["noverx"], par["novert"]), nop=(par["ny"], par["nx"], par["nt"]), tapertype=par["tapertype"], + savetaper=par["savetaper"], ) assert dottest( Pop, diff --git a/pytests/test_sliding.py b/pytests/test_sliding.py index 6750bbdf..55a62199 100755 --- a/pytests/test_sliding.py +++ b/pytests/test_sliding.py @@ -22,6 +22,7 @@ "noverx": 0, # "winsx": 2, "tapertype": None, + "savetaper": True, } # no overlap, no taper par2 = { "ny": 6, @@ -36,6 +37,7 @@ "noverx": 0, # "winsx": 2, "tapertype": "hanning", + "savetaper": True, } # no overlap, with taper par3 = { "ny": 6, @@ -50,8 +52,24 @@ "noverx": 2, # "winsx": 4, "tapertype": None, + "savetaper": True, } # overlap, no taper par4 = { + "ny": 6, + "nx": 7, + "nt": 10, + "npy": 15, + "nwiny": 7, + "novery": 3, + # "winsy": 3, + "npx": 10, + "nwinx": 4, + "noverx": 2, + # "winsx": 4, + "tapertype": None, + "savetaper": False, +} # overlap, no taper (non saved) +par5 = { "ny": 6, "nx": 7, "nt": 10, @@ -64,10 +82,26 @@ "noverx": 2, # "winsx": 4, "tapertype": "hanning", + "savetaper": True, } # overlap, with taper +par6 = { + "ny": 6, + "nx": 7, + "nt": 10, + "npy": 15, + "nwiny": 7, + "novery": 3, + # "winsy": 3, + "npx": 10, + "nwinx": 4, + "noverx": 2, + # "winsx": 4, + "tapertype": "hanning", + "savetaper": False, +} # overlap, with taper (non saved) -@pytest.mark.parametrize("par", [(par1), (par2), (par3), (par4)]) +@pytest.mark.parametrize("par", [(par1), (par2), (par3), (par4), (par5), (par6)]) def test_Sliding1D(par): """Dot-test and inverse for Sliding1D operator""" Op = MatrixMult(np.ones((par["nwiny"], par["ny"]))) @@ -83,6 +117,7 @@ def test_Sliding1D(par): nwin=par["nwiny"], nover=par["novery"], tapertype=par["tapertype"], + savetaper=par["savetaper"], ) assert dottest(Slid, par["npy"], par["ny"] * nwins) x = np.ones(par["ny"] * nwins) @@ -92,7 +127,7 @@ def test_Sliding1D(par): assert_array_almost_equal(x.ravel(), xinv) -@pytest.mark.parametrize("par", [(par1), (par2), (par3), (par4)]) +@pytest.mark.parametrize("par", [(par1), (par2), (par3), (par4), (par5), (par6)]) def test_Sliding2D(par): """Dot-test and inverse for Sliding2D operator""" Op = MatrixMult(np.ones((par["nwiny"] * par["nt"], par["ny"] * par["nt"]))) @@ -107,6 +142,7 @@ def test_Sliding2D(par): nwin=par["nwiny"], nover=par["novery"], tapertype=par["tapertype"], + savetaper=par["savetaper"], ) assert dottest(Slid, par["npy"] * par["nt"], par["ny"] * par["nt"] * nwins) x = np.ones((par["ny"] * nwins, par["nt"])) @@ -116,7 +152,7 @@ def test_Sliding2D(par): assert_array_almost_equal(x.ravel(), xinv) -@pytest.mark.parametrize("par", [(par1), (par2), (par3), (par4)]) +@pytest.mark.parametrize("par", [(par1), (par2), (par3), (par4), (par5), (par6)]) def test_Sliding3D(par): """Dot-test and inverse for Sliding3D operator""" Op = MatrixMult( @@ -140,6 +176,7 @@ def test_Sliding3D(par): nover=(par["novery"], par["noverx"]), nop=(par["ny"], par["nx"]), tapertype=par["tapertype"], + savetaper=par["savetaper"], ) assert dottest( Slid,