From cd3aadfc54696487810260a0f0b764510bbb9a63 Mon Sep 17 00:00:00 2001 From: mrava87 Date: Sat, 18 May 2024 21:30:18 +0300 Subject: [PATCH] minor: moved repeated code into private method --- pylops/signalprocessing/patch2d.py | 59 ++--- pylops/signalprocessing/patch3d.py | 325 ++++++--------------------- pylops/signalprocessing/sliding1d.py | 23 +- pylops/signalprocessing/sliding2d.py | 23 +- pylops/signalprocessing/sliding3d.py | 59 ++--- 5 files changed, 133 insertions(+), 356 deletions(-) diff --git a/pylops/signalprocessing/patch2d.py b/pylops/signalprocessing/patch2d.py index 1a352303..3aa157a5 100644 --- a/pylops/signalprocessing/patch2d.py +++ b/pylops/signalprocessing/patch2d.py @@ -297,6 +297,27 @@ def __init__( 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) @@ -397,48 +418,14 @@ def _rmatvec_nosavetaper(self, x: NDArray) -> NDArray: if self.tapertype is not None: for iwin0 in range(self.dims[0]): for iwin1 in range(self.dims[1]): - 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] + 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: - 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] + ywins = self._apply_taper(ywins, iwin0, iwin1) y[iwin0, iwin1] = self.Op.rmatvec( ywins[iwin0, iwin1].ravel() ).reshape(self.dims[2], self.dims[3]) diff --git a/pylops/signalprocessing/patch3d.py b/pylops/signalprocessing/patch3d.py index 7b969984..71ffe341 100644 --- a/pylops/signalprocessing/patch3d.py +++ b/pylops/signalprocessing/patch3d.py @@ -513,6 +513,69 @@ def __init__( 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) @@ -680,136 +743,7 @@ def _rmatvec_nosavetaper(self, x: NDArray) -> NDArray: for iwin0 in range(self.dims[0]): for iwin1 in range(self.dims[1]): for iwin2 in range(self.dims[2]): - 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] - ) + ywins = self._apply_taper(ywins, iwin0, iwin1, iwin2) y = self.Op.H @ ywins else: y = ncp.zeros(self.dims, dtype=self.dtype) @@ -817,136 +751,7 @@ def _rmatvec_nosavetaper(self, x: NDArray) -> NDArray: for iwin1 in range(self.dims[1]): for iwin2 in range(self.dims[2]): if self.tapertype is not None: - 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] - ) + 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]) diff --git a/pylops/signalprocessing/sliding1d.py b/pylops/signalprocessing/sliding1d.py index 767045d1..9bb331bb 100644 --- a/pylops/signalprocessing/sliding1d.py +++ b/pylops/signalprocessing/sliding1d.py @@ -212,6 +212,15 @@ def __init__( 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) @@ -282,23 +291,13 @@ def _rmatvec_nosavetaper(self, x: NDArray) -> NDArray: if self.simOp: if self.tapertype is not None: for iwin0 in range(self.dims[0]): - if iwin0 == 0: - ywins[0] = ywins[0] * self.taps[0] - elif iwin0 == self.dims[0] - 1: - ywins[-1] = ywins[-1] * self.taps[-1] - else: - ywins[iwin0] = ywins[iwin0] * self.taps[1] + 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: - 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] + ywins = self._apply_taper(ywins, iwin0) y[iwin0] = self.Op.rmatvec(ywins[iwin0]) return y diff --git a/pylops/signalprocessing/sliding2d.py b/pylops/signalprocessing/sliding2d.py index f3b3cfb3..4a76f5e6 100644 --- a/pylops/signalprocessing/sliding2d.py +++ b/pylops/signalprocessing/sliding2d.py @@ -255,6 +255,15 @@ def __init__( 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) @@ -335,23 +344,13 @@ def _rmatvec_nosavetaper(self, x: NDArray) -> NDArray: if self.simOp: if self.tapertype is not None: for iwin0 in range(self.dims[0]): - if iwin0 == 0: - ywins[0] = ywins[0] * self.taps[0] - elif iwin0 == self.dims[0] - 1: - ywins[-1] = ywins[-1] * self.taps[-1] - else: - ywins[iwin0] = ywins[iwin0] * self.taps[1] + 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: - 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] + ywins = self._apply_taper(ywins, iwin0) y[iwin0] = self.Op.rmatvec(ywins[iwin0].ravel()).reshape( self.dims[1], self.dims[2] ) diff --git a/pylops/signalprocessing/sliding3d.py b/pylops/signalprocessing/sliding3d.py index 9dc7e255..49324230 100644 --- a/pylops/signalprocessing/sliding3d.py +++ b/pylops/signalprocessing/sliding3d.py @@ -307,6 +307,27 @@ def __init__( 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) @@ -417,48 +438,14 @@ def _rmatvec_nosavetaper(self, x: NDArray) -> NDArray: if self.tapertype is not None: for iwin0 in range(self.dims[0]): for iwin1 in range(self.dims[1]): - 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] + 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: - 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] + 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])