From 90fc327600b323bc8e5851572c046c949e6cf936 Mon Sep 17 00:00:00 2001 From: Matteo Ravasi Date: Fri, 17 Feb 2023 09:35:22 +0300 Subject: [PATCH] NonStationaryFilters1D and NonStationaryFilters2D operators (#478) * feat: Added NonStationaryFilters1D and NonStationaryFilters2D First implemenation of non-stationary filter estimation operators (point 4 in https://github.com/PyLops/pylops/issues/466). * minor: added tests for nonstatfilters operators * minor: added checks that filter locations are within dims * minor: reintroduced parallel in rmatvec --- docs/source/api/index.rst | 2 + examples/plot_nonstatfilter.py | 135 +++++++++ pylops/signalprocessing/__init__.py | 4 + pylops/signalprocessing/nonstatconvolve1d.py | 185 +++++++++++- pylops/signalprocessing/nonstatconvolve2d.py | 296 ++++++++++++++++++- pytests/test_nonstatconvolve.py | 81 ++++- 6 files changed, 687 insertions(+), 16 deletions(-) create mode 100644 examples/plot_nonstatfilter.py diff --git a/docs/source/api/index.rst b/docs/source/api/index.rst index 76ca7344..111c026e 100755 --- a/docs/source/api/index.rst +++ b/docs/source/api/index.rst @@ -91,6 +91,8 @@ Signal processing ConvolveND NonStationaryConvolve1D NonStationaryConvolve2D + NonStationaryFilters1D + NonStationaryFilters2D Interp Bilinear FFT diff --git a/examples/plot_nonstatfilter.py b/examples/plot_nonstatfilter.py new file mode 100644 index 00000000..3d4d58a7 --- /dev/null +++ b/examples/plot_nonstatfilter.py @@ -0,0 +1,135 @@ +""" +Non-stationary Filter Estimation +================================ +This example shows how to use the :py:class:`pylops.signalprocessing.NonStationaryFilters1D` +and :py:class:`pylops.signalprocessing.NonStationaryFilters2D` operators to perform non-stationary +filter estimation that when convolved with an input signal produces a desired output signal. + +""" +import matplotlib.pyplot as plt +import numpy as np +from scipy.signal.windows import gaussian + +import pylops +from pylops.utils.wavelets import ricker + +plt.close("all") + + +############################################################################### +# We will start by creating a zero signal of length `nt` and we will +# place a comb of unitary spikes. We also create a non-stationary filter defined by +# 5 equally spaced `Ricker wavelets `_ +# with dominant frequencies of :math:`f = 10, 15, 20, 25` and :math:`30` Hz. +nt = 601 +dt = 0.004 +t = np.arange(nt) * dt +tw = ricker(t[:51], f0=5)[1] + +fs = [10, 15, 20, 25, 30] +wavs = np.stack([ricker(t[:51], f0=f)[0] for f in fs]) + +x = np.zeros(nt) +x[64 : nt - 64 : 64] = 1.0 + +Cop = pylops.signalprocessing.NonStationaryFilters1D( + x, hsize=101, ih=(101, 201, 301, 401, 501) +) + +y = Cop @ wavs + +wavsinv = Cop.div(y, niter=20) +wavsinv = wavsinv.reshape(wavs.shape) + +fig, axs = plt.subplots(1, len(fs), sharey=True, figsize=(14, 3)) +fig.suptitle("1D Non-stationary filter estimation") +for i, ax in enumerate(axs): + ax.plot(wavs[i], "k", lw=4, label="True") + ax.plot(wavsinv[i], "r", lw=2, label="Estimate") + ax.set_xlabel("Time [sec]") +axs[0].legend() +plt.tight_layout() + +############################################################################### +# Finally, we repeat the same exercise with a 2-dimensional non-stationary +# filter +nx, nz = 201, 101 + +wav1a, _, _ = ricker(t[:9], f0=12) +wav1b, _, _ = ricker(t[:9], f0=30) +wav2a = gaussian(15, 2.0) +wav2b = gaussian(15, 4.0) + +wav11 = np.outer(wav1a, wav2a[np.newaxis]).T +wav12 = np.outer(wav1b, wav2a[np.newaxis]).T +wav21 = np.outer(wav1b, wav2b[np.newaxis]).T +wav22 = np.outer(wav1a, wav2b[np.newaxis]).T +wavsize = wav11.shape + +hs = np.zeros((2, 2, *wavsize)) +hs[0, 0] = wav11 +hs[0, 1] = wav12 +hs[1, 0] = wav21 +hs[1, 1] = wav22 + +x = np.zeros((nx, nz)) +x[:, 21] = 1.0 +x[:, 41] = -1.0 + +Cop = pylops.signalprocessing.NonStationaryFilters2D( + inp=x, hshape=wavsize, ihx=(21, 41), ihz=(21, 41), engine="numba" +) + +y = Cop @ hs + +hsinv = Cop.div(y.ravel(), niter=50) +hsinv = hsinv.reshape(hs.shape) + +fig, axs = plt.subplots(2, 2, figsize=(10, 5)) +fig.suptitle("True filters") +axs[0, 0].imshow(hs[0, 0], cmap="gray", vmin=-1, vmax=1) +axs[0, 0].axis("tight") +axs[0, 0].set_title(r"$H_{1,1}$") +axs[0, 1].imshow(hs[0, 1], cmap="gray", vmin=-1, vmax=1) +axs[0, 1].axis("tight") +axs[0, 1].set_title(r"$H_{1,2}$") +axs[1, 0].imshow(hs[1, 0], cmap="gray", vmin=-1, vmax=1) +axs[1, 0].axis("tight") +axs[1, 0].set_title(r"$H_{2,1}$") +axs[1, 1].imshow(hs[1, 1], cmap="gray", vmin=-1, vmax=1) +axs[1, 1].axis("tight") +axs[1, 1].set_title(r"$H_{2,2}$") +plt.tight_layout() + +fig, axs = plt.subplots(2, 2, figsize=(10, 5)) +fig.suptitle("Estimated filters") +axs[0, 0].imshow(hsinv[0, 0], cmap="gray", vmin=-1, vmax=1) +axs[0, 0].axis("tight") +axs[0, 0].set_title(r"$H_{1,1}$") +axs[0, 1].imshow(hsinv[0, 1], cmap="gray", vmin=-1, vmax=1) +axs[0, 1].axis("tight") +axs[0, 1].set_title(r"$H_{1,2}$") +axs[1, 0].imshow(hsinv[1, 0], cmap="gray", vmin=-1, vmax=1) +axs[1, 0].axis("tight") +axs[1, 0].set_title(r"$H_{2,1}$") +axs[1, 1].imshow(hsinv[1, 1], cmap="gray", vmin=-1, vmax=1) +axs[1, 1].axis("tight") +axs[1, 1].set_title(r"$H_{2,2}$") +plt.tight_layout() + + +fig, axs = plt.subplots(2, 2, figsize=(10, 5)) +fig.suptitle("Estimation error") +axs[0, 0].imshow(hs[0, 0] - hsinv[0, 0], cmap="gray", vmin=-1, vmax=1) +axs[0, 0].axis("tight") +axs[0, 0].set_title(r"$H_{1,1}$") +axs[0, 1].imshow(hs[0, 1] - hsinv[0, 1], cmap="gray", vmin=-1, vmax=1) +axs[0, 1].axis("tight") +axs[0, 1].set_title(r"$H_{1,2}$") +axs[1, 0].imshow(hs[1, 0] - hsinv[1, 0], cmap="gray", vmin=-1, vmax=1) +axs[1, 0].axis("tight") +axs[1, 0].set_title(r"$H_{2,1}$") +axs[1, 1].imshow(hs[1, 1] - hsinv[1, 1], cmap="gray", vmin=-1, vmax=1) +axs[1, 1].axis("tight") +axs[1, 1].set_title(r"$H_{2,2}$") +plt.tight_layout() diff --git a/pylops/signalprocessing/__init__.py b/pylops/signalprocessing/__init__.py index fcf6f61a..b4af9884 100755 --- a/pylops/signalprocessing/__init__.py +++ b/pylops/signalprocessing/__init__.py @@ -12,6 +12,8 @@ ConvolveND ND convolution operator. NonStationaryConvolve1D 1D nonstationary convolution operator. NonStationaryConvolve2D 2D nonstationary convolution operator. + NonStationaryFilters1D 1D nonstationary filter estimation operator. + NonStationaryFilters2D 2D nonstationary filter estimation operator. Interp Interpolation operator. Bilinear Bilinear interpolation operator. FFT One dimensional Fast-Fourier Transform. @@ -67,6 +69,8 @@ "Convolve2D", "NonStationaryConvolve1D", "NonStationaryConvolve2D", + "NonStationaryFilters1D", + "NonStationaryFilters2D", "Interp", "Bilinear", "Radon2D", diff --git a/pylops/signalprocessing/nonstatconvolve1d.py b/pylops/signalprocessing/nonstatconvolve1d.py index fce77238..45daeed5 100644 --- a/pylops/signalprocessing/nonstatconvolve1d.py +++ b/pylops/signalprocessing/nonstatconvolve1d.py @@ -1,4 +1,7 @@ -__all__ = ["NonStationaryConvolve1D"] +__all__ = [ + "NonStationaryConvolve1D", + "NonStationaryFilters1D", +] from typing import Union @@ -50,13 +53,15 @@ class NonStationaryConvolve1D(LinearOperator): If filters ``hs`` have even size ValueError If ``ih`` is not regularly sampled + ValueError + If ``ih`` is outside the bounds defined by ``dims[axis]`` Notes ----- The NonStationaryConvolve1D operator applies non-stationary one-dimensional convolution between the input signal :math:`d(t)` and a bank of compact filter kernels :math:`h(t; t_i)`. Assuming - an input signal composed of :math:`N=4` samples, and filters at locations + an input signal composed of :math:`N=5` samples, and filters at locations :math:`t_1` and :math:`t_3`, the forward operator can be represented as follows: .. math:: @@ -100,12 +105,16 @@ def __init__( raise ValueError( "the indices of filters 'ih' are must be regularly sampled" ) + dims = _value_or_sized_to_tuple(dims) + if min(ih) < 0 or max(ih) >= dims[axis]: + raise ValueError( + "the indices of filters 'ih' must be larger than 0 and smaller than `dims`" + ) self.hs = hs self.hsize = hs.shape[1] self.oh, self.dh, self.nh, self.eh = ih[0], ih[1] - ih[0], len(ih), ih[-1] self.axis = axis - dims = _value_or_sized_to_tuple(dims) super().__init__(dtype=np.dtype(dtype), dims=dims, dimsd=dims, name=name) @property @@ -183,3 +192,173 @@ def todense(self): ) H = H[:, int(self.hsize // 2) : -int(self.hsize // 2) - 1] return H + + +class NonStationaryFilters1D(LinearOperator): + r"""1D non-stationary filter estimation operator. + + Estimate a non-stationary one-dimensional filter by non-stationary convolution. + In forward mode, a varying compact filter on a coarser grid is on-the-fly linearly interpolated prior + to being convolved with a fixed input signal. In adjoint mode, the output signal is first weighted by the + fixed input signal and then spread across multiple filters (i.e., adjoint of linear interpolation). + + Parameters + ---------- + inp : :obj:`numpy.ndarray` + Fixed input signal of size :math:`n_x`. + hsize : :obj:`int` + Size of the 1d compact filters (filters must have odd number of samples and are assumed to be + centered in the middle of the filter support). + ih : :obj:`tuple` + Indices of the locations of the filters ``hs`` in the model (and data). Note + that the filters must be regularly sampled, i.e. :math:`dh=\text{diff}(ih)=\text{const.}` + dtype : :obj:`str`, optional + Type of elements in input array. + name : :obj:`str`, optional + Name of operator (to be used by :func:`pylops.utils.describe.describe`) + + Attributes + ---------- + shape : :obj:`tuple` + Operator shape + explicit : :obj:`bool` + Operator contains a matrix that can be solved explicitly (``True``) or + not (``False``) + + Raises + ------ + ValueError + If filters ``hsize`` is a even number + ValueError + If ``ih`` is not regularly sampled + ValueError + If ``ih`` is outside the bounds defined by ``dims[axis]`` + + See Also + -------- + NonStationaryConvolve1D : 1D non-stationary convolution operator. + NonStationaryFilters2D : 2D non-stationary filter estimation operator. + + Notes + ----- + The NonStationaryFilters1D operates in a similar fashion to the + :class:`pylops.signalprocessing.NonStationaryConvolve1D` operator. In practical applications, + this operator shall be used when interested to estimate a 1-dimensional non-stationary filter + given an input and output signal. + + In forward mode, this operator uses the same implementation of the + :class:`pylops.signalprocessing.NonStationaryConvolve1D`, with the main difference that + the role of the filters and the input signal is swapped. Nevertheless, to understand how + to implement adjoint, mathematically we arrange the forward operator in a slightly different way. + Assuming once again an input signal composed of :math:`N=5` samples, and filters at locations + :math:`t_1` and :math:`t_3`, the forward operator can be represented as follows: + + .. math:: + \mathbf{y} = + \begin{bmatrix} + \mathbf{X}_0 & \mathbf{X}_1 & \vdots & \mathbf{X}_4 + \end{bmatrix} \mathbf{L} + \begin{bmatrix} + \mathbf{h}_1 \\ \mathbf{h}_3 + \end{bmatrix} + + where :math:`\mathbf{L}` is an operator that linearly interpolates the filters from the available locations to + the entire input grid -- i.e., :math:`[\hat{\mathbf{h}}_0 \quad \mathbf{h}_1 \quad \hat{\mathbf{h}}_2 \quad + \mathbf{h}_3 \quad \hat{\mathbf{h}}_4]^T = \mathbf{L} [ \mathbf{h}_1 \quad \mathbf{h}_3]`. Finally, + :math:`\mathbf{X}_i` is a diagonal matrix containing the value :math:`x_i` along the + main diagonal. Note that in practice the filter may be shorter than the input and output signals and + the :math:`x_i` values are placed only at the effective positions of the filter along the diagonal matrices + :math:`\mathbf{X}_i`. + + In adjoint mode, the output signal is first weighted by the fixed input signal and then spread across + multiple filters (i.e., adjoint of linear interpolation) as follows + + .. math:: + \mathbf{h} = + \mathbf{L}^H + \begin{bmatrix} + \mathbf{X}_0 \\ \mathbf{X}_1 \\ \vdots \\ \mathbf{X}_4 + \end{bmatrix} + \mathbf{y} + + """ + + def __init__( + self, + inp: NDArray, + hsize: int, + ih: InputDimsLike, + dtype: DTypeLike = "float64", + name: str = "C", + ) -> None: + if hsize % 2 == 0: + raise ValueError("filters hs must have odd length") + if len(np.unique(np.diff(ih))) > 1: + raise ValueError( + "the indices of filters 'ih' are must be regularly sampled" + ) + if min(ih) < 0 or max(ih) >= inp.size: + raise ValueError( + "the indices of filters 'ih' must be larger than 0 and smaller than `dims`" + ) + self.inp = inp + self.hsize = hsize + self.oh, self.dh, self.nh, self.eh = ih[0], ih[1] - ih[0], len(ih), ih[-1] + + super().__init__( + dtype=np.dtype(dtype), dims=(len(ih), hsize), dimsd=inp.shape, name=name + ) + + # use same interpolation method as inNonStationaryConvolve1D + _interpolate_h = staticmethod(NonStationaryConvolve1D._interpolate_h) + + @staticmethod + def _interpolate_hadj(htmp, hs, hextremes, ix, oh, dh, nh): + """find closest filters and spread weighted psf""" + ih_closest = int(np.floor((ix - oh) / dh)) + if ih_closest < 0: + hs[0, hextremes[0] : hextremes[1]] += htmp + elif ih_closest >= nh - 1: + hs[nh - 1, hextremes[0] : hextremes[1]] += htmp + else: + dh_closest = (ix - oh) / dh - ih_closest + hs[ih_closest, hextremes[0] : hextremes[1]] += (1 - dh_closest) * htmp + hs[ih_closest + 1, hextremes[0] : hextremes[1]] += dh_closest * htmp + return hs + + @reshaped + def _matvec(self, x: NDArray) -> NDArray: + y = np.zeros(self.dimsd, dtype=self.dtype) + for ix in range(self.dimsd[0]): + h = self._interpolate_h(x, ix, self.oh, self.dh, self.nh) + xextremes = ( + max(0, ix - self.hsize // 2), + min(ix + self.hsize // 2 + 1, self.dimsd[0]), + ) + hextremes = ( + max(0, -ix + self.hsize // 2), + min(self.hsize, self.hsize // 2 + (self.dimsd[0] - ix)), + ) + y[..., xextremes[0] : xextremes[1]] += ( + self.inp[..., ix : ix + 1] * h[hextremes[0] : hextremes[1]] + ) + return y + + @reshaped + def _rmatvec(self, x: NDArray) -> NDArray: + hs = np.zeros(self.dims, dtype=self.dtype) + for ix in range(self.dimsd[0]): + xextremes = ( + max(0, ix - self.hsize // 2), + min(ix + self.hsize // 2 + 1, self.dimsd[0]), + ) + hextremes = ( + max(0, -ix + self.hsize // 2), + min(self.hsize, self.hsize // 2 + (self.dimsd[0] - ix)), + ) + + htmp = self.inp[ix] * x[..., xextremes[0] : xextremes[1]] + hs = self._interpolate_hadj( + htmp, hs, hextremes, ix, self.oh, self.dh, self.nh + ) + return hs diff --git a/pylops/signalprocessing/nonstatconvolve2d.py b/pylops/signalprocessing/nonstatconvolve2d.py index cc4bb138..569736d6 100644 --- a/pylops/signalprocessing/nonstatconvolve2d.py +++ b/pylops/signalprocessing/nonstatconvolve2d.py @@ -1,4 +1,7 @@ -__all__ = ["NonStationaryConvolve2D"] +__all__ = [ + "NonStationaryConvolve2D", + "NonStationaryFilters2D", +] import os from typing import Tuple, Union @@ -31,14 +34,14 @@ class NonStationaryConvolve2D(LinearOperator): r"""2D non-stationary convolution operator. - Apply non-stationary two-dimensional convolution. A varying compact filter - is provided on a coarser grid and on-the-fly interpolation is applied - in forward and adjoint modes. + Apply non-stationary two-dimensional convolution. A varying compact filter is provided on a + coarser grid and on-the-fly interpolation is applied in forward and adjoint modes. + Both input and output have size :math`n_x \time n_z`. Parameters ---------- dims : :obj:`list` or :obj:`int` - Number of samples for each dimension + Number of samples for each dimension (which we refer to as :math`n_x \time n_z`). hs : :obj:`numpy.ndarray` Bank of 2d compact filters of size :math:`n_{\text{filts},x} \times n_{\text{filts},z} \times n_h \times n_{h,x} \times n_{h,z}`. @@ -73,6 +76,8 @@ class NonStationaryConvolve2D(LinearOperator): If filters ``hs`` have even size ValueError If ``ihx`` or ``ihz`` is not regularly sampled + ValueError + If ``ihx`` or ``ihz`` are outside the bounds defined by ``dims`` NotImplementedError If ``engine`` is neither ``numpy``, ``fftw``, nor ``scipy``. @@ -132,7 +137,10 @@ def __init__( raise ValueError( "the indices of filters 'ih' are must be regularly sampled" ) - + if min(ihx) < 0 or min(ihz) < 0 or max(ihx) >= dims[0] or max(ihz) >= dims[1]: + raise ValueError( + "the indices of filters 'ih' must be larger than 0 and smaller than `dims`" + ) self.hs = hs self.hshape = hs.shape[2:] self.ohx, self.dhx, self.nhx = ihx[0], ihx[1] - ihx[0], len(ihx) @@ -292,3 +300,279 @@ def _rmatvec(self, x: NDArray) -> NDArray: **self.kwargs_cuda ) return y + + +class NonStationaryFilters2D(LinearOperator): + r"""2D non-stationary filter estimation operator. + + Estimate a non-stationary two-dimensional filter by non-stationary convolution. + + Parameters + ---------- + inp : :obj:`numpy.ndarray` + Fixed input signal of size :math:`n_x \times n_z`. + hshape : :obj:`list` or :obj:`tuple` + Shape of the 2d compact filters (filters must have odd number of samples and are assumed to be + centered in the middle of the filter support). + ihx : :obj:`tuple` + Indices of the x locations of the filters ``hs`` in the model (and data). Note + that the filters must be regularly sampled, i.e. :math:`dh_x=\text{diff}(ihx)=\text{const.}` + ihz : :obj:`tuple` + Indices of the z locations of the filters ``hs`` in the model (and data). Note + that the filters must be regularly sampled, i.e. :math:`dh_z=\text{diff}(ihz)=\text{const.}` + engine : :obj:`str`, optional + Engine used for spread computation (``numpy``, ``numba``, or ``cuda``) + num_threads_per_blocks : :obj:`tuple`, optional + Number of threads in each block (only when ``engine=cuda``) + dtype : :obj:`str`, optional + Type of elements in input array. + name : :obj:`str`, optional + Name of operator (to be used by :func:`pylops.utils.describe.describe`) + + Attributes + ---------- + shape : :obj:`tuple` + Operator shape + explicit : :obj:`bool` + Operator contains a matrix that can be solved explicitly (``True``) or + not (``False``) + + Raises + ------ + ValueError + If filters ``hs`` have even size + ValueError + If ``ihx`` or ``ihz`` is not regularly sampled + NotImplementedError + If ``engine`` is neither ``numpy``, ``fftw``, nor ``scipy``. + + Notes + ----- + The NonStationaryConvolve2D operator is used to estimate a non-stationary + two-dimensional filter betwen two signals, an input signal (provided directly + to the operator) and the desired output signal. + + For more details on the numerical implementation of the forward and adjoint, + see :class:`pylops.signalprocessing.NonStationaryFilters1D`. + + """ + + def __init__( + self, + inp: NDArray, + hshape: InputDimsLike, + ihx: InputDimsLike, + ihz: InputDimsLike, + engine: str = "numpy", + num_threads_per_blocks: Tuple[int, int] = (32, 32), + dtype: DTypeLike = "float64", + name: str = "C", + ) -> None: + if engine not in ["numpy", "numba", "cuda"]: + raise NotImplementedError("engine must be numpy or numba or cuda") + if hshape[0] % 2 == 0 or hshape[1] % 2 == 0: + raise ValueError("filters hs must have odd length") + if len(np.unique(np.diff(ihx))) > 1 or len(np.unique(np.diff(ihz))) > 1: + raise ValueError( + "the indices of filters 'ih' are must be regularly sampled" + ) + if ( + min(ihx) < 0 + or min(ihz) < 0 + or max(ihx) >= inp.shape[0] + or max(ihz) >= inp.shape[1] + ): + raise ValueError( + "the indices of filters 'ih' must be larger than 0 and smaller than `dims`" + ) + self.inp = inp + self.inpdims = inp.shape + self.hshape = hshape + self.ohx, self.dhx, self.nhx = ihx[0], ihx[1] - ihx[0], len(ihx) + self.ohz, self.dhz, self.nhz = ihz[0], ihz[1] - ihz[0], len(ihz) + self.ehx, self.ehz = ihx[-1], ihz[-1] + self.engine = engine + super().__init__( + dtype=np.dtype(dtype), + dims=(self.nhx, self.nhz, *hshape), + dimsd=self.inpdims, + name=name, + ) + + # create additional input parameters for engine=cuda + self.kwargs_cuda = {} + if engine == "cuda": + self.kwargs_cuda["num_threads_per_blocks"] = num_threads_per_blocks + num_threads_per_blocks_x, num_threads_per_blocks_z = num_threads_per_blocks + num_blocks_x = ( + self.dims[0] + num_threads_per_blocks_x - 1 + ) // num_threads_per_blocks_x + num_blocks_z = ( + self.dims[1] + num_threads_per_blocks_z - 1 + ) // num_threads_per_blocks_z + self.kwargs_cuda["num_blocks"] = (num_blocks_x, num_blocks_z) + self._register_multiplications(engine) + + def _register_multiplications(self, engine: str) -> None: + if engine == "numba": + numba_opts = dict(nopython=True, fastmath=True, nogil=True, parallel=True) + self._mv = jit(**numba_opts)(self.__matvec) + self._rmv = jit(**numba_opts)(self.__rmatvec) + elif engine == "cuda": + raise NotImplementedError("engine=cuda is currently not available") + else: + self._mv = self.__matvec + self._rmv = self.__rmatvec + + # use same matvec method as inNonStationaryConvolve2D + __matvec = staticmethod(NonStationaryConvolve2D._matvec_rmatvec) + + @staticmethod + def __rmatvec( + x: NDArray, + y: NDArray, + hs: NDArray, + hshape: Tuple[int, int], + xdims: Tuple[int, int], + ohx: float, + ohz: float, + dhx: float, + dhz: float, + nhx: int, + nhz: int, + ) -> NDArray: + # Currently a race condition seem to occur when updating parts of hs multiple times within + # the same loop (see https://numba.pydata.org/numba-doc/latest/user/parallel.html?highlight=njit). + # Until atomic operations are provided we create a temporary filter array and store intermediate + # results from each ix and reduce them at the end. + hstmp = np.zeros((xdims[0], *hs.shape)) + for ix in prange(xdims[0]): + for iz in range(xdims[1]): + + # find extremes of model where to apply h (in case h is going out of model) + xextremes = ( + max(0, ix - hshape[0] // 2), + min(ix + hshape[0] // 2 + 1, xdims[0]), + ) + zextremes = ( + max(0, iz - hshape[1] // 2), + min(iz + hshape[1] // 2 + 1, xdims[1]), + ) + # find extremes of h (in case h is going out of model) + hxextremes = ( + max(0, -ix + hshape[0] // 2), + min(hshape[0], hshape[0] // 2 + (xdims[0] - ix)), + ) + hzextremes = ( + max(0, -iz + hshape[1] // 2), + min(hshape[1], hshape[1] // 2 + (xdims[1] - iz)), + ) + + htmp = ( + x[ix, iz] + * y[xextremes[0] : xextremes[1], zextremes[0] : zextremes[1]] + ) + + # find closest filters and interpolate h + ihx_l = int(np.floor((ix - ohx) / dhx)) + ihz_t = int(np.floor((iz - ohz) / dhz)) + dhx_r = (ix - ohx) / dhx - ihx_l + dhz_b = (iz - ohz) / dhz - ihz_t + if ihx_l < 0: + ihx_l = ihx_r = 0 + dhx_l = dhx_r = 0.5 + elif ihx_l >= nhx - 1: + ihx_l = ihx_r = nhx - 1 + dhx_l = dhx_r = 0.5 + else: + ihx_r = ihx_l + 1 + dhx_l = 1.0 - dhx_r + + if ihz_t < 0: + ihz_t = ihz_b = 0 + dhz_t = dhz_b = 0.5 + elif ihz_t >= nhz - 1: + ihz_t = ihz_b = nhz - 1 + dhz_t = dhz_b = 0.5 + else: + ihz_b = ihz_t + 1 + dhz_t = 1.0 - dhz_b + + hstmp[ + ix, + ihx_l, + ihz_t, + hxextremes[0] : hxextremes[1], + hzextremes[0] : hzextremes[1], + ] += ( + dhz_t * dhx_l * htmp + ) + hstmp[ + ix, + ihx_l, + ihz_b, + hxextremes[0] : hxextremes[1], + hzextremes[0] : hzextremes[1], + ] += ( + dhz_b * dhx_l * htmp + ) + hstmp[ + ix, + ihx_r, + ihz_t, + hxextremes[0] : hxextremes[1], + hzextremes[0] : hzextremes[1], + ] += ( + dhz_t * dhx_r * htmp + ) + hstmp[ + ix, + ihx_r, + ihz_b, + hxextremes[0] : hxextremes[1], + hzextremes[0] : hzextremes[1], + ] += ( + dhz_b * dhx_r * htmp + ) + hs = hstmp.sum(axis=0) + return hs + + @reshaped + def _matvec(self, x: NDArray) -> NDArray: + ncp = get_array_module(x) + y = ncp.zeros(self.dimsd, dtype=self.dtype) + y = self._mv( + self.inp, + y, + x, + self.hshape, + self.dimsd, + self.ohx, + self.ohz, + self.dhx, + self.dhz, + self.nhx, + self.nhz, + **self.kwargs_cuda + ) + return y + + @reshaped + def _rmatvec(self, x: NDArray) -> NDArray: + ncp = get_array_module(x) + hs = ncp.zeros(self.dims, dtype=self.dtype) + hs = self._rmv( + self.inp, + x, + hs, + self.hshape, + self.dimsd, + self.ohx, + self.ohz, + self.dhx, + self.dhz, + self.nhx, + self.nhz, + **self.kwargs_cuda + ) + return hs diff --git a/pytests/test_nonstatconvolve.py b/pytests/test_nonstatconvolve.py index abb49c92..83a355c0 100755 --- a/pytests/test_nonstatconvolve.py +++ b/pytests/test_nonstatconvolve.py @@ -9,13 +9,15 @@ Convolve2D, NonStationaryConvolve1D, NonStationaryConvolve2D, + NonStationaryFilters1D, + NonStationaryFilters2D, ) from pylops.utils import dottest # filters -nfilt = (5, 7) -h1 = triang(nfilt[0], sym=True) -h2 = np.outer(triang(nfilt[0], sym=True), triang(nfilt[1], sym=True)) +nfilts = (5, 7) +h1 = triang(nfilts[0], sym=True) +h2 = np.outer(triang(nfilts[0], sym=True), triang(nfilts[1], sym=True)) h1stat = np.vstack([h1, h1, h1]) h1ns = np.vstack([h1, -h1, 2 * h1]) h2stat = np.vstack( @@ -27,7 +29,7 @@ h2.ravel(), h2.ravel(), ] -).reshape(3, 2, nfilt[0], nfilt[1]) +).reshape(3, 2, nfilts[0], nfilts[1]) h2ns = np.vstack( [ 2 * h2.ravel(), @@ -37,7 +39,7 @@ -h2.ravel(), 2 * h2.ravel(), ] -).reshape(3, 2, nfilt[0], nfilt[1]) +).reshape(3, 2, nfilts[0], nfilts[1]) par1_1d = { "nz": 21, @@ -73,6 +75,20 @@ def test_even_filter(par): ihz=(int(par["nz"] // 4), int(3 * par["nz"] // 4)), ) + with pytest.raises(ValueError): + _ = NonStationaryFilters1D( + inp=np.arange(par["nx"]), + hsize=nfilts[0] - 1, + ih=(int(par["nx"] // 4), int(2 * par["nx"] // 4), int(3 * par["nx"] // 4)), + ) + with pytest.raises(ValueError): + _ = NonStationaryFilters2D( + inp=np.ones((par["nx"], par["nz"])), + hshape=(nfilts[0] - 1, nfilts[1] - 1), + ihx=(int(par["nx"] // 4), int(2 * par["nx"] // 4), int(3 * par["nx"] // 4)), + ihz=(int(par["nz"] // 4), int(3 * par["nz"] // 4)), + ) + @pytest.mark.parametrize("par", [(par_2d)]) def test_ih_irregular(par): @@ -103,6 +119,14 @@ def test_unknown_engine_2d(par): ihz=(int(par["nz"] // 3), int(2 * par["nz"] // 3)), engine="foo", ) + with pytest.raises(NotImplementedError): + _ = NonStationaryFilters2D( + inp=np.ones((par["nx"], par["nz"])), + hshape=(nfilts[0] - 1, nfilts[1] - 1), + ihx=(int(par["nx"] // 3), int(2 * par["nx"] // 3)), + ihz=(int(par["nz"] // 3), int(2 * par["nz"] // 3)), + engine="foo", + ) @pytest.mark.parametrize("par", [(par1_1d), (par2_1d)]) @@ -158,7 +182,7 @@ def test_StationaryConvolve1D(par): Cop_stat = Convolve1D( dims=par["nx"], h=h1, - offset=nfilt[0] // 2, + offset=nfilts[0] // 2, dtype="float64", ) @@ -202,9 +226,52 @@ def test_StationaryConvolve2D(par): Cop_stat = Convolve2D( dims=(par["nx"], par["nz"]), h=h2, - offset=(nfilt[0] // 2, nfilt[1] // 2), + offset=(nfilts[0] // 2, nfilts[1] // 2), dtype="float64", ) x = np.random.normal(0, 1, (par["nx"], par["nz"])) assert_array_almost_equal(Cop_stat * x, Cop * x, decimal=10) + + +@pytest.mark.parametrize( + "par", + [ + (par1_1d), + ], +) +def test_NonStationaryFilters2D(par): + """Dot-test and inversion for NonStationaryFilters2D operator""" + x = np.zeros((par["nx"])) + x[par["nx"] // 4], x[par["nx"] // 2], x[3 * par["nx"] // 4] = 1.0, 1.0, 1.0 + Cop = NonStationaryFilters1D( + inp=x, + hsize=nfilts[0], + ih=(int(par["nx"] // 4), int(2 * par["nx"] // 4), int(3 * par["nx"] // 4)), + dtype="float64", + ) + assert dottest(Cop, par["nx"], 3 * nfilts[0]) + + h1lsqr = lsqr(Cop, Cop * h1ns, damp=1e-20, iter_lim=200, show=0)[0] + assert_array_almost_equal(h1ns.ravel(), h1lsqr, decimal=1) + + +@pytest.mark.parametrize("par", [(par_2d)]) +def test_NonStationaryFilters2D(par): + """Dot-test and inversion for NonStationaryFilters2D operator""" + x = np.zeros((par["nx"], par["nz"])) + x[int(par["nx"] // 4)] = 1.0 + x[int(par["nx"] // 2)] = 1.0 + x[int(3 * par["nx"] // 4)] = 1.0 + + Cop = NonStationaryFilters2D( + inp=x, + hshape=nfilts, + ihx=(int(par["nx"] // 4), int(2 * par["nx"] // 4), int(3 * par["nx"] // 4)), + ihz=(int(par["nz"] // 4), int(3 * par["nz"] // 4)), + dtype="float64", + ) + assert dottest(Cop, par["nx"] * par["nz"], 6 * nfilts[0] * nfilts[1]) + + h2lsqr = lsqr(Cop, Cop * h2ns.ravel(), damp=1e-20, iter_lim=400, show=0)[0] + assert_array_almost_equal(h2ns.ravel(), h2lsqr, decimal=1)