Source code for pylops.signalprocessing.patch3d

__all__ = [
    "patch3d_design",
    "Patch3D",
]

import logging
from typing import Optional, Sequence, Tuple

import numpy as np

from pylops import LinearOperator
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

logging.basicConfig(format="%(levelname)s: %(message)s", level=logging.WARNING)


[docs]def patch3d_design( dimsd: InputDimsLike, 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], Tuple[Tuple[NDArray, NDArray], Tuple[NDArray, NDArray], Tuple[NDArray, NDArray]], Tuple[Tuple[NDArray, NDArray], Tuple[NDArray, NDArray], Tuple[NDArray, NDArray]], ]: """Design Patch3D operator This routine can be used prior to creating the :class:`pylops.signalprocessing.Patch3D` operator to identify the correct number of windows to be used based on the dimension of the data (``dimsd``), dimension of the window (``nwin``), overlap (``nover``),a and dimension of the operator acting in the model space. Parameters ---------- dimsd : :obj:`tuple` Shape of 3-dimensional data. nwin : :obj:`tuple` Number of samples of window. nover : :obj:`tuple` 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 ------- nwins : :obj:`tuple` Number of windows. dims : :obj:`tuple` Shape of 3-dimensional model. mwins_inends : :obj:`tuple` Start and end indices for model patches (stored as tuple of tuples). dwins_inends : :obj:`tuple` Start and end indices for data patches (stored as tuple of tuples). """ # 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]) 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) # model windows dims = (nwins0 * nop[0], nwins1 * nop[1], nwins2 * nop[2]) mwin0_ins, mwin0_ends = _slidingsteps(dims[0], nop[0], 0) mwin1_ins, mwin1_ends = _slidingsteps(dims[1], nop[1], 0) mwin2_ins, mwin2_ends = _slidingsteps(dims[2], nop[2], 0) mwins_inends = ( (mwin0_ins, mwin0_ends), (mwin1_ins, mwin1_ends), (mwin2_ins, mwin2_ends), ) # print information about patching 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
[docs]class Patch3D(LinearOperator): """3D Patch transform operator. Apply a transform operator ``Op`` repeatedly to patches of the model vector in forward mode and patches of the data vector in adjoint mode. More specifically, in forward mode the model vector is divided into patches, each patch is transformed, and patches are then recombined together. Both model and data are internally reshaped and interpreted as 3-dimensional arrays: each patch contains a portion of the array in every axis. This operator can be used to perform local, overlapping transforms (e.g., :obj:`pylops.signalprocessing.FFTND` or :obj:`pylops.signalprocessing.Radon3D`) on 3-dimensional arrays. .. note:: The shape of the model has to be consistent with the number of windows for this operator not to return an error. As the number of windows depends directly on the choice of ``nwin`` and ``nover``, it is recommended to first run ``patch3d_design`` to obtain the corresponding ``dims`` and number of windows. .. 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 with running ``patch3d_design``. Parameters ---------- Op : :obj:`pylops.LinearOperator` Transform operator dims : :obj:`tuple` Shape of 3-dimensional model. Note that ``dims[0]``, ``dims[1]`` and ``dims[2]`` should be multiple of the model size of the transform in their respective dimensions dimsd : :obj:`tuple` Shape of 3-dimensional data nwin : :obj:`tuple` Number of samples of window nover : :obj:`tuple` Number of samples of overlapping part of window nop : :obj:`tuple` 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 name : :obj:`str`, optional Name of operator (to be used by :func:`pylops.utils.describe.describe`) Returns ------- Sop : :obj:`pylops.LinearOperator` Sliding operator Raises ------ ValueError Identified number of windows is not consistent with provided model shape (``dims``). See Also -------- Sliding1D: 1D Sliding transform operator. Sliding2D: 2D Sliding transform operator. Sliding3D: 3D Sliding transform operator. Patch2D: 2D Patching transform operator. """ 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 # 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, ) 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] 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 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