diff --git a/CHANGELOG.md b/CHANGELOG.md index 09d6c06b..1f1c6780 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,7 +2,7 @@ Changelog ========= # 2.3.1 -* Fixed bug in :py:mod:`pylops.utils.backend` (see [Issue #606](https://github.com/PyLops/pylops/issues/606)) +* Fixed bug in `pylops.utils.backend` (see [Issue #606](https://github.com/PyLops/pylops/issues/606)) # 2.3.0 diff --git a/docs/source/api/index.rst b/docs/source/api/index.rst index 1d77dc15..59e40e51 100755 --- a/docs/source/api/index.rst +++ b/docs/source/api/index.rst @@ -109,6 +109,8 @@ Signal processing Seislet Radon2D Radon3D + FourierRadon2D + FourierRadon3D ChirpRadon2D ChirpRadon3D Sliding1D diff --git a/docs/source/api/others.rst b/docs/source/api/others.rst index a26d0cd7..82f5a7b1 100755 --- a/docs/source/api/others.rst +++ b/docs/source/api/others.rst @@ -118,6 +118,7 @@ Synthetics seismicevents.parabolic2d seismicevents.hyperbolic2d seismicevents.linear3d + seismicevents.parabolic3d seismicevents.hyperbolic3d .. currentmodule:: pylops.waveeqprocessing diff --git a/docs/source/gpu.rst b/docs/source/gpu.rst index a2d7d9f1..ecc87120 100755 --- a/docs/source/gpu.rst +++ b/docs/source/gpu.rst @@ -281,6 +281,14 @@ Signal processing: - |:white_check_mark:| - |:red_circle:| - |:red_circle:| + * - :class:`pylops.signalprocessing.FourierRadon2D` + - |:white_check_mark:| + - |:white_check_mark:| + - |:red_circle:| + * - :class:`pylops.signalprocessing.FourierRadon3D` + - |:white_check_mark:| + - |:white_check_mark:| + - |:red_circle:| * - :class:`pylops.signalprocessing.ChirpRadon2D` - |:white_check_mark:| - |:white_check_mark:| diff --git a/examples/plot_fourierradon.py b/examples/plot_fourierradon.py new file mode 100644 index 00000000..abe94819 --- /dev/null +++ b/examples/plot_fourierradon.py @@ -0,0 +1,325 @@ +r""" +Fourier Radon Transform +======================= +This example shows how to use the :py:class:`pylops.signalprocessing.FourierRadon2D` +and :py:class:`pylops.signalprocessing.FourierRadon3D` operators to apply the linear +and parabolic Radon Transform to 2-dimensional or 3-dimensional signals, respectively. + +These operators provides transformations equivalent to those of +:py:class:`pylops.signalprocessing.Radon2D` and :py:class:`pylops.signalprocessing.Radon3D`, +however since the shift-and-sum step is performed in the frequency domain, +this is analytically correct (compared to performing to shifting the data via +nearest or linear interpolation). + +""" + +import matplotlib.pyplot as plt +import numpy as np + +import pylops + +plt.close("all") + +############################################################################### +# Let's start by creating a empty 2d matrix of size :math:`n_x \times n_t` +# with a single linear event. + +par = { + "ot": 0, + "dt": 0.004, + "nt": 51, + "ox": -250, + "dx": 10, + "nx": 51, + "oy": -250, + "dy": 10, + "ny": 51, + "f0": 40, +} +theta = [10.0] +t0 = [0.1] +amp = [1.0] + +# Create axes +t, t2, x, y = pylops.utils.seismicevents.makeaxis(par) +dt, dx, dy = par["dt"], par["dx"], par["dy"] + +# Create wavelet +wav, _, wav_c = pylops.utils.wavelets.ricker(t[:41], f0=par["f0"]) + +# Generate data +_, d = pylops.utils.seismicevents.linear2d(x, t, 1500.0, t0, theta, amp, wav) + + +############################################################################### +# We can now define our operators and apply the forward and adjoint +# steps. +nfft = int(2 ** np.ceil(np.log2(par["nt"]))) +npx, pxmax = 2 * par["nx"], 5e-4 +px = np.linspace(-pxmax, pxmax, npx) + +R2Op = pylops.signalprocessing.FourierRadon2D( + t, x, px, nfft, kind="linear", engine="numpy", dtype="float64" +) +dL = R2Op.H * d +dadj = R2Op * dL + +fig, axs = plt.subplots(1, 3, figsize=(12, 4), sharey=True) +axs[0].imshow(d.T, vmin=-1, vmax=1, cmap="bwr_r", extent=(x[0], x[-1], t[-1], t[0])) +axs[0].set(xlabel=r"$x$ [m]", ylabel=r"$t$ [s]", title="Input linear") +axs[0].axis("tight") +axs[1].imshow( + dL.T, + cmap="bwr_r", + vmin=-dL.max(), + vmax=dL.max(), + extent=(1e3 * px[0], 1e3 * px[-1], t[-1], t[0]), +) +axs[1].scatter(1e3 * np.sin(np.deg2rad(theta[0])) / 1500.0, t0[0], s=50, color="k") +axs[1].set(xlabel=r"$p$ [s/km]", title="Radon") +axs[1].axis("tight") +axs[2].imshow( + dadj.T, + cmap="bwr_r", + vmin=-dadj.max(), + vmax=dadj.max(), + extent=(x[0], x[-1], t[-1], t[0]), +) +axs[2].set(xlabel=r"$x$ [m]", title="Adj Radon") +axs[2].axis("tight") +plt.tight_layout() + + +############################################################################### +# We repeat now the same with a parabolic event + +# Generate data +pxx = [1e-6] +_, d = pylops.utils.seismicevents.parabolic2d(x, t, t0, 0, np.array(pxx), amp, wav) + +# Radon transform +npx, pxmax = 2 * par["nx"], 5e-6 +px = np.linspace(-pxmax, pxmax, npx) + +R2Op = pylops.signalprocessing.FourierRadon2D( + t, x, px, nfft, kind="parabolic", engine="numpy", dtype="float64" +) +dL = R2Op.H * d +dadj = R2Op * dL + +fig, axs = plt.subplots(1, 3, figsize=(12, 4), sharey=True) +axs[0].imshow(d.T, vmin=-1, vmax=1, cmap="bwr_r", extent=(x[0], x[-1], t[-1], t[0])) +axs[0].set(xlabel=r"$x$ [m]", ylabel=r"$t$ [s]", title="Input parabolic") +axs[0].axis("tight") +axs[1].imshow( + dL.T, + cmap="bwr_r", + vmin=-dL.max(), + vmax=dL.max(), + extent=(1e3 * px[0], 1e3 * px[-1], t[-1], t[0]), +) +axs[1].scatter(1e3 * pxx[0], t0[0], s=50, color="k") +axs[1].set(xlabel=r"$p$ [s/km]", title="Radon") +axs[1].axis("tight") +axs[2].imshow( + dadj.T, + cmap="bwr_r", + vmin=-dadj.max(), + vmax=dadj.max(), + extent=(x[0], x[-1], t[-1], t[0]), +) +axs[2].set(xlabel=r"$x$ [m]", title="Adj Radon") +axs[2].axis("tight") +plt.tight_layout() + +############################################################################### +# Finally we repeat the same exercise with 3d data. + +par = { + "ot": 0, + "dt": 0.004, + "nt": 51, + "ox": -100, + "dx": 10, + "nx": 21, + "oy": -200, + "dy": 10, + "ny": 41, + "f0": 20, +} +theta = [30] +phi = [10] +t0 = [0.1] +amp = [1.0] + +# Create axes +t, t2, x, y = pylops.utils.seismicevents.makeaxis(par) +dt, dx, dy = par["dt"], par["dx"], par["dy"] + +# Generate linear data +pxx = np.sin(np.deg2rad(theta[0])) * np.cos(np.deg2rad(phi[0])) / 1500.0 +pyy = np.sin(np.deg2rad(theta[0])) * np.sin(np.deg2rad(phi[0])) / 1500.0 +_, d = pylops.utils.seismicevents.linear3d(x, y, t, 1500.0, t0, theta, phi, amp, wav) + +# Linear Radon +npy, pymax = par["ny"], 5e-4 +npx, pxmax = par["nx"], 5e-4 +py = np.linspace(-pymax, pymax, npy) +px = np.linspace(-pxmax, pxmax, npx) + +R3Op = pylops.signalprocessing.FourierRadon3D( + t, y, x, py, px, nfft, kind=("linear", "linear"), engine="numpy", dtype="float64" +) +dL = R3Op.H * d +dadj = R3Op * dL + +fig, axs = plt.subplots(1, 3, figsize=(12, 4), sharey=True) +axs[0].imshow( + d[par["ny"] // 2].T, + vmin=-1, + vmax=1, + cmap="bwr_r", + extent=(x[0], x[-1], t[-1], t[0]), +) +axs[0].set(xlabel=r"$x$ [m]", ylabel=r"$t$ [s]", title="Input linear 3d - y") +axs[0].axis("tight") +axs[1].imshow( + dL[np.argmin(np.abs(pyy - py))].T, + cmap="bwr_r", + vmin=-dL.max(), + vmax=dL.max(), + extent=(1e3 * px[0], 1e3 * px[-1], t[-1], t[0]), +) +axs[1].scatter(1e3 * pxx, t0[0], s=50, color="k") +axs[1].set(xlabel=r"$p_x$ [s/km]", title="Radon 3d - y") +axs[1].axis("tight") +axs[2].imshow( + dadj[par["ny"] // 2].T, + cmap="bwr_r", + vmin=-dadj.max(), + vmax=dadj.max(), + extent=(x[0], x[-1], t[-1], t[0]), +) +axs[2].set(xlabel=r"$x$ [m]", title="Adj Radon 3d - y") +axs[2].axis("tight") +plt.tight_layout() + +fig, axs = plt.subplots(1, 3, figsize=(12, 4), sharey=True) +axs[0].imshow( + d[:, par["nx"] // 2].T, + vmin=-1, + vmax=1, + cmap="bwr_r", + extent=(x[0], x[-1], t[-1], t[0]), +) +axs[0].set(xlabel=r"$y$ [m]", ylabel=r"$t$ [s]", title="Input linear 3d - x") +axs[0].axis("tight") +axs[1].imshow( + dL[:, np.argmin(np.abs(pxx - px))].T, + cmap="bwr_r", + vmin=-dL.max(), + vmax=dL.max(), + extent=(1e3 * py[0], 1e3 * py[-1], t[-1], t[0]), +) +axs[1].scatter(1e3 * pyy, t0[0], s=50, color="k") +axs[1].set(xlabel=r"$p_y$ [s/km]", title="Radon 3d - x") +axs[1].axis("tight") +axs[2].imshow( + dadj[:, par["nx"] // 2].T, + cmap="bwr_r", + vmin=-dadj.max(), + vmax=dadj.max(), + extent=(x[0], x[-1], t[-1], t[0]), +) +axs[2].set(xlabel=r"$y$ [m]", title="Adj Radon 3d - x") +axs[2].axis("tight") +plt.tight_layout() + +# Generate parabolic data +pxx = [1e-6] +pyy = [2e-6] +_, d = pylops.utils.seismicevents.parabolic3d( + x, y, t, t0, 0, 0, np.array(pxx), np.array(pyy), amp, wav +) + +# Parabolic Radon +npy, pymax = par["ny"], 5e-6 +npx, pxmax = par["nx"], 5e-6 +py = np.linspace(-pymax, pymax, npy) +px = np.linspace(-pxmax, pxmax, npx) + +R3Op = pylops.signalprocessing.FourierRadon3D( + t, + y, + x, + py, + px, + nfft, + kind=("parabolic", "parabolic"), + engine="numpy", + dtype="float64", +) +dL = R3Op.H * d +dadj = R3Op * dL + +fig, axs = plt.subplots(1, 3, figsize=(12, 4), sharey=True) +axs[0].imshow( + d[par["ny"] // 2].T, + vmin=-1, + vmax=1, + cmap="bwr_r", + extent=(x[0], x[-1], t[-1], t[0]), +) +axs[0].set(xlabel=r"$x$ [m]", ylabel=r"$t$ [s]", title="Input parabolic 3d - y") +axs[0].axis("tight") +axs[1].imshow( + dL[np.argmin(np.abs(pyy - py))].T, + cmap="bwr_r", + vmin=-dL.max(), + vmax=dL.max(), + extent=(1e3 * px[0], 1e3 * px[-1], t[-1], t[0]), +) +axs[1].scatter(1e3 * pxx[0], t0[0], s=50, color="k") +axs[1].set(xlabel=r"$p_x$ [s/km]", title="Radon 3d - y") +axs[1].axis("tight") +axs[2].imshow( + dadj[par["ny"] // 2].T, + cmap="bwr_r", + vmin=-dadj.max(), + vmax=dadj.max(), + extent=(x[0], x[-1], t[-1], t[0]), +) +axs[2].set(xlabel=r"$x$ [m]", title="Adj Radon 3d - y") +axs[2].axis("tight") +plt.tight_layout() + +fig, axs = plt.subplots(1, 3, figsize=(12, 4), sharey=True) +axs[0].imshow( + d[:, par["nx"] // 2].T, + vmin=-1, + vmax=1, + cmap="bwr_r", + extent=(x[0], x[-1], t[-1], t[0]), +) +axs[0].set(xlabel=r"$y$ [m]", ylabel=r"$t$ [s]", title="Input parabolic 3d - x") +axs[0].axis("tight") +axs[1].imshow( + dL[:, np.argmin(np.abs(pxx - px))].T, + cmap="bwr_r", + vmin=-dL.max(), + vmax=dL.max(), + extent=(1e3 * py[0], 1e3 * py[-1], t[-1], t[0]), +) +axs[1].scatter(1e3 * pyy[0], t0[0], s=50, color="k") +axs[1].set(xlabel=r"$p_y$ [s/km]", title="Radon 3d - x") +axs[1].axis("tight") +axs[2].imshow( + dadj[:, par["nx"] // 2].T, + cmap="bwr_r", + vmin=-dadj.max(), + vmax=dadj.max(), + extent=(x[0], x[-1], t[-1], t[0]), +) +axs[2].set(xlabel=r"$y$ [m]", title="Adj Radon 3d - x") +axs[2].axis("tight") +plt.tight_layout() diff --git a/examples/plot_seismicevents.py b/examples/plot_seismicevents.py index 27b6e9fc..03ebaaf5 100644 --- a/examples/plot_seismicevents.py +++ b/examples/plot_seismicevents.py @@ -106,11 +106,21 @@ ############################################ # Let's finally repeat the same exercise in 3d phi = [20, 0, -10] +py = [0, 0, 0] +pyy = [3e-5, 1e-5, 5e-6] mlin, mlinwav = pylops.utils.seismicevents.linear3d( x, y, t, v, t0, theta, phi, amp, wav ) +mpar, mparwav = pylops.utils.seismicevents.parabolic3d( + x, y, t, t0, px, py, pxx, pyy, amp, wav +) + +mhyp, mhypwav = pylops.utils.seismicevents.hyperbolic3d( + x, y, t, t0, vrms, vrms, amp, wav +) + fig, axs = plt.subplots(1, 2, figsize=(7, 5), sharey=True) fig.suptitle("Linear events in 3d", fontsize=12, fontweight="bold", y=0.95) axs[0].imshow( @@ -135,9 +145,30 @@ ) axs[1].set_xlabel(r"$y(m)$") -mhyp, mhypwav = pylops.utils.seismicevents.hyperbolic3d( - x, y, t, t0, vrms, vrms, amp, wav +fig, axs = plt.subplots(1, 2, figsize=(7, 5), sharey=True) +fig.suptitle("Parabolic events in 3d", fontsize=12, fontweight="bold", y=0.95) +axs[0].imshow( + mparwav[par["ny"] // 2].T, + aspect="auto", + interpolation="nearest", + vmin=-2, + vmax=2, + cmap="gray", + extent=(x.min(), x.max(), t.max(), t.min()), +) +axs[0].set_xlabel(r"$x(m)$") +axs[0].set_ylabel(r"$t(s)$") +axs[1].imshow( + mparwav[:, par["nx"] // 2].T, + aspect="auto", + interpolation="nearest", + vmin=-2, + vmax=2, + cmap="gray", + extent=(y.min(), y.max(), t.max(), t.min()), ) +axs[1].set_xlabel(r"$y(m)$") + fig, axs = plt.subplots(1, 2, figsize=(7, 5), sharey=True) fig.suptitle("Hyperbolic events in 3d", fontsize=12, fontweight="bold", y=0.95) diff --git a/pylops/signalprocessing/__init__.py b/pylops/signalprocessing/__init__.py index a8e5ed65..d8a5a4cc 100755 --- a/pylops/signalprocessing/__init__.py +++ b/pylops/signalprocessing/__init__.py @@ -28,6 +28,10 @@ DTCWT Dual-Tree Complex Wavelet Transform. Radon2D Two dimensional Radon transform. Radon3D Three dimensional Radon transform. + FourierRadon2D Two dimensional Fourier Radon transform. + FourierRadon3D Three dimensional Fourier Radon transform. + ChirpRadon2D Two dimensional Chirp Radon transform. + ChirpRadon3D Three dimensional Chirp Radon transform. Seislet Two dimensional Seislet operator. Sliding1D 1D Sliding transform operator. Sliding2D 2D Sliding transform operator. @@ -52,6 +56,8 @@ from .bilinear import * from .radon2d import * from .radon3d import * +from .fourierradon2d import * +from .fourierradon3d import * from .chirpradon2d import * from .chirpradon3d import * from .sliding1d import * @@ -67,7 +73,6 @@ from .dct import * from .dtcwt import * - __all__ = [ "FFT", "FFT2D", @@ -85,6 +90,8 @@ "Bilinear", "Radon2D", "Radon3D", + "FourierRadon2D", + "FourierRadon3D", "ChirpRadon2D", "ChirpRadon3D", "Sliding1D", diff --git a/pylops/signalprocessing/_fourierradon2d_cuda.py b/pylops/signalprocessing/_fourierradon2d_cuda.py new file mode 100755 index 00000000..2a0d35ba --- /dev/null +++ b/pylops/signalprocessing/_fourierradon2d_cuda.py @@ -0,0 +1,86 @@ +from cmath import exp +from math import pi + +from numba import cuda + + +@cuda.jit +def _radon_inner_2d_kernel(x, y, f, px, h, flim0, flim1, npx, nh): + """Cuda kernels for FourierRadon2D operator + + Cuda implementation of the on-the-fly kernel creation and application for the + FourierRadon2D operator. See :class:`pylops.signalprocessing.FourierRadon2D` + for details about input parameters. + + """ + ih, ifr = cuda.grid(2) + if ih < nh and ifr >= flim0 and ifr <= flim1: + for ipx in range(npx): + y[ih, ifr] += x[ipx, ifr] * exp(-1j * 2 * pi * f[ifr] * px[ipx] * h[ih]) + + +@cuda.jit +def _aradon_inner_2d_kernel(x, y, f, px, h, flim0, flim1, npx, nh): + """Cuda kernels for FourierRadon2D operator + + Cuda implementation of the on-the-fly kernel creation and application for the + FourierRadon2D operator. See :class:`pylops.signalprocessing.FourierRadon2D` + for details about input parameters. + + """ + ipx, ifr = cuda.grid(2) + if ipx < npx and ifr >= flim0 and ifr <= flim1: + for ih in range(nh): + x[ipx, ifr] += y[ih, ifr] * exp(1j * 2 * pi * f[ifr] * px[ipx] * h[ih]) + + +def _radon_inner_2d_cuda( + x, + y, + f, + px, + h, + flim0, + flim1, + npx, + nh, + num_blocks=(32, 32), + num_threads_per_blocks=(32, 32), +): + """Caller for FourierRadon2D operator + + Caller for cuda implementation of matvec kernel for FourierRadon2D operator. + See :class:`pylops.signalprocessing.FourierRadon2D` for details about + input parameters. + + """ + _radon_inner_2d_kernel[num_blocks, num_threads_per_blocks]( + x, y, f, px, h, flim0, flim1, npx, nh + ) + return y + + +def _aradon_inner_2d_cuda( + x, + y, + f, + px, + h, + flim0, + flim1, + npx, + nh, + num_blocks=(32, 32), + num_threads_per_blocks=(32, 32), +): + """Caller for FourierRadon2D operator + + Caller for cuda implementation of rmatvec kernel for FourierRadon2D operator. + See :class:`pylops.signalprocessing.FourierRadon2D` for details about + input parameters. + + """ + _aradon_inner_2d_kernel[num_blocks, num_threads_per_blocks]( + x, y, f, px, h, flim0, flim1, npx, nh + ) + return x diff --git a/pylops/signalprocessing/_fourierradon2d_numba.py b/pylops/signalprocessing/_fourierradon2d_numba.py new file mode 100755 index 00000000..8f09724e --- /dev/null +++ b/pylops/signalprocessing/_fourierradon2d_numba.py @@ -0,0 +1,25 @@ +import os +from cmath import exp +from math import pi + +from numba import jit, prange + +# detect whether to use parallel or not +numba_threads = int(os.getenv("NUMBA_NUM_THREADS", "1")) +parallel = True if numba_threads != 1 else False + + +@jit(nopython=True, parallel=parallel, nogil=True, cache=True, fastmath=True) +def _radon_inner_2d(X, Y, f, px, h, flim0, flim1, npx, nh): + for ih in prange(nh): + for ifr in range(flim0, flim1): + for ipx in range(npx): + Y[ih, ifr] += X[ipx, ifr] * exp(-1j * 2 * pi * f[ifr] * px[ipx] * h[ih]) + + +@jit(nopython=True, parallel=parallel, nogil=True, cache=True, fastmath=True) +def _aradon_inner_2d(X, Y, f, px, h, flim0, flim1, npx, nh): + for ipx in prange(npx): + for ifr in range(flim0, flim1): + for ih in range(nh): + X[ipx, ifr] += Y[ih, ifr] * exp(1j * 2 * pi * f[ifr] * px[ipx] * h[ih]) diff --git a/pylops/signalprocessing/_fourierradon3d_cuda.py b/pylops/signalprocessing/_fourierradon3d_cuda.py new file mode 100755 index 00000000..16f94f7e --- /dev/null +++ b/pylops/signalprocessing/_fourierradon3d_cuda.py @@ -0,0 +1,100 @@ +from cmath import exp +from math import pi + +from numba import cuda + + +@cuda.jit +def _radon_inner_3d_kernel(x, y, f, py, px, hy, hx, flim0, flim1, npy, npx, nhy, nhx): + """Cuda kernels for FourierRadon3D operator + + Cuda implementation of the on-the-fly kernel creation and application for the + FourierRadon3D operator. See :class:`pylops.signalprocessing.FourierRadon3D` + for details about input parameters. + + """ + ihy, ihx, ifr = cuda.grid(3) + if ihy < nhy and ihx < nhx and ifr >= flim0 and ifr <= flim1: + for ipy in range(npy): + for ipx in range(npx): + y[ihy, ihx, ifr] += x[ipy, ipx, ifr] * exp( + -1j * 2 * pi * f[ifr] * (py[ipy] * hy[ihy] + px[ipx] * hx[ihx]) + ) + + +@cuda.jit +def _aradon_inner_3d_kernel(x, y, f, py, px, hy, hx, flim0, flim1, npy, npx, nhy, nhx): + """Cuda kernels for FourierRadon3D operator + + Cuda implementation of the on-the-fly kernel creation and application for the + FourierRadon3D operator. See :class:`pylops.signalprocessing.FourierRadon3D` + for details about input parameters. + + """ + ipy, ipx, ifr = cuda.grid(3) + if ipy < npy and ipx < npx and ifr >= flim0 and ifr <= flim1: + for ihy in range(nhy): + for ihx in range(nhx): + x[ipy, ipx, ifr] += y[ihy, ihx, ifr] * exp( + 1j * 2 * pi * f[ifr] * (py[ipy] * hy[ihy] + px[ipx] * hx[ihx]) + ) + + +def _radon_inner_3d_cuda( + x, + y, + f, + py, + px, + hy, + hx, + flim0, + flim1, + npy, + npx, + nhy, + nhx, + num_blocks=(8, 8), + num_threads_per_blocks=(8, 8), +): + """Caller for FourierRadon2D operator + + Caller for cuda implementation of matvec kernel for FourierRadon3D operator. + See :class:`pylops.signalprocessing.FourierRadon3D` for details about + input parameters. + + """ + _radon_inner_3d_kernel[num_blocks, num_threads_per_blocks]( + x, y, f, py, px, hy, hx, flim0, flim1, npy, npx, nhy, nhx + ) + return y + + +def _aradon_inner_3d_cuda( + x, + y, + f, + py, + px, + hy, + hx, + flim0, + flim1, + npy, + npx, + nhy, + nhx, + num_blocks=(8, 8), + num_threads_per_blocks=(8, 8), +): + """Caller for FourierRadon3D operator + + Caller for cuda implementation of rmatvec kernel for FourierRadon3D operator. + See :class:`pylops.signalprocessing.FourierRadon3D` for details about + input parameters. + + """ + _aradon_inner_3d_kernel[num_blocks, num_threads_per_blocks]( + x, y, f, py, px, hy, hx, flim0, flim1, npy, npx, nhy, nhx + ) + return x diff --git a/pylops/signalprocessing/_fourierradon3d_numba.py b/pylops/signalprocessing/_fourierradon3d_numba.py new file mode 100755 index 00000000..e217fe95 --- /dev/null +++ b/pylops/signalprocessing/_fourierradon3d_numba.py @@ -0,0 +1,41 @@ +import os +from cmath import exp +from math import pi + +from numba import jit, prange + +# detect whether to use parallel or not +numba_threads = int(os.getenv("NUMBA_NUM_THREADS", "1")) +parallel = True if numba_threads != 1 else False + + +@jit(nopython=True, parallel=parallel, nogil=True, cache=True, fastmath=True) +def _radon_inner_3d(X, Y, f, py, px, hy, hx, flim0, flim1, npy, npx, nhy, nhx): + for ihy in prange(nhy): + for ihx in prange(nhx): + for ifr in range(flim0, flim1): + for ipy in range(npy): + for ipx in range(npx): + Y[ihy, ihx, ifr] += X[ipy, ipx, ifr] * exp( + -1j + * 2 + * pi + * f[ifr] + * (py[ipy] * hy[ihy] + px[ipx] * hx[ihx]) + ) + + +@jit(nopython=True, parallel=parallel, nogil=True, cache=True, fastmath=True) +def _aradon_inner_3d(X, Y, f, py, px, hy, hx, flim0, flim1, npy, npx, nhy, nhx): + for ipy in prange(npy): + for ipx in range(npx): + for ifr in range(flim0, flim1): + for ihy in range(nhy): + for ihx in range(nhx): + X[ipy, ipx, ifr] += Y[ihy, ihx, ifr] * exp( + 1j + * 2 + * pi + * f[ifr] + * (py[ipy] * hy[ihy] + px[ipx] * hx[ihx]) + ) diff --git a/pylops/signalprocessing/_nonstatconvolve2d_cuda.py b/pylops/signalprocessing/_nonstatconvolve2d_cuda.py index 5de400bb..feba8ed2 100644 --- a/pylops/signalprocessing/_nonstatconvolve2d_cuda.py +++ b/pylops/signalprocessing/_nonstatconvolve2d_cuda.py @@ -100,9 +100,9 @@ def _matvec_rmatvec_call( ): """Caller for NonStationaryConvolve2D operator - Caller for cuda implementation of matvec and rmatvec for NonStationaryConvolve2D operato, with same signature + Caller for cuda implementation of matvec and rmatvec for NonStationaryConvolve2D operator, with same signature as numpy/numba counterparts. See :class:`pylops.signalprocessing.NonStationaryConvolve2D` for details about - input parameters. + input parameters. """ _matvec_rmatvec[num_blocks, num_threads_per_blocks]( diff --git a/pylops/signalprocessing/fourierradon2d.py b/pylops/signalprocessing/fourierradon2d.py new file mode 100755 index 00000000..f8296d78 --- /dev/null +++ b/pylops/signalprocessing/fourierradon2d.py @@ -0,0 +1,296 @@ +__all__ = ["FourierRadon2D"] + +import logging +from typing import Optional, Tuple + +import numpy as np +import scipy as sp + +from pylops import LinearOperator +from pylops.utils import deps +from pylops.utils.backend import get_array_module, get_complex_dtype +from pylops.utils.decorators import reshaped +from pylops.utils.typing import DTypeLike, NDArray + +jit_message = deps.numba_import("the radon2d module") + +if jit_message is None: + from ._fourierradon2d_cuda import _aradon_inner_2d_cuda, _radon_inner_2d_cuda + from ._fourierradon2d_numba import _aradon_inner_2d, _radon_inner_2d + +logging.basicConfig(format="%(levelname)s: %(message)s", level=logging.WARNING) + + +class FourierRadon2D(LinearOperator): + r"""2D Fourier Radon transform + + Apply Radon forward (and adjoint) transform using Fast + Fourier Transform to a 2-dimensional array of size :math:`[n_{p_x} \times n_t]` + (and :math:`[n_x \times n_t]`). + + Note that forward and adjoint follow the same convention of the time-space + implementation in :class:`pylops.signalprocessing.Radon2D`. + + Parameters + ---------- + taxis : :obj:`np.ndarray` + Time axis + haxis : :obj:`np.ndarray` + Spatial axis + pxaxis : :obj:`np.ndarray` + Axis of scanning variable :math:`p_x` of parametric curve + nfft : :obj:`int` + Number of samples in Fourier transform + flims : :obj:`tuple`, optional + Indices of lower and upper limits of Fourier axis to be used in + the application of the Radon matrix (when ``None``, use entire axis) + kind : :obj:`str` + Curve to be used for stacking/spreading (``linear``, ``parabolic``) + engine : :obj:`str` + Engine used for computation (``numpy`` or ``numba`` or ``cuda``) + num_threads_per_blocks : :obj:`tuple` + Number of threads in each block (only when ``engine=cuda``) + dtype : :obj:`str` + Type of elements in input array. + name : :obj:`str` + 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 + ------ + NotImplementedError + If ``engine`` is neither ``numpy``, ``numba``, nor ``cuda``. + + Notes + ----- + The FourierRadon2D operator applies the Radon transform in the frequency domain. + After transforming a 2-dimensional array of size + :math:`[n_x \times n_t]` into the frequency domain, the following linear + transformation is applied to each frequency component in adjoint mode: + + .. math:: + \begin{bmatrix} + m(p_{x,1}, \omega_i) \\ + m(p_{x,2}, \omega_i) \\ + \vdots \\ + m(p_{x,N_p}, \omega_i) + \end{bmatrix} + = + \begin{bmatrix} + e^{-j \omega_i p_{x,1} x^l_1} & e^{-j \omega_i p_{x,1} x^l_2} & \ldots & e^{-j \omega_i p_{x,1} x^l_{N_x}} \\ + e^{-j \omega_i p_{x,2} x^l_1} & e^{-j \omega_i p_{x,2} x^l_2} & \ldots & e^{-j \omega_i p_{x,2} x^l_{N_x}} \\ + \vdots & \vdots & \ddots & \vdots \\ + e^{-j \omega_i p_{x,N_p} x^l_1} & e^{-j \omega_i p_{x,N_p} x^l_2} & \ldots & e^{-j \omega_i p_{x,N_p} x^l_{N_x}} \\ + \end{bmatrix} + \begin{bmatrix} + d(x_1, \omega_i) \\ + d(x_2, \omega_i) \\ + \vdots \\ + d(x_{N_x}, \omega_i) + \end{bmatrix} + + where :math:`l=1,2`. Similarly the forward mode is implemented by applying the + transpose and complex conjugate of the above matrix to the model transformed to + the Fourier domain. + + Refer to [1]_ for more theoretical and implementation details. + + .. [1] Sacchi, M. "Statistical and Transform Methods for + Geophysical Signal Processing", 2007. + + """ + + def __init__( + self, + taxis: NDArray, + haxis: NDArray, + pxaxis: NDArray, + nfft: int, + flims: Optional[Tuple[int, int]] = None, + kind: str = "linear", + engine: str = "numpy", + num_threads_per_blocks: Tuple[int, int] = (32, 32), + dtype: DTypeLike = "float64", + name: str = "R", + ) -> None: + # engine + if engine not in ["numpy", "numba", "cuda"]: + raise NotImplementedError("engine must be numpy or numba or cuda") + if engine == "numba" and jit_message is not None: + engine = "numpy" + + # dimensions and super + dims = len(pxaxis), len(taxis) + dimsd = len(haxis), len(taxis) + super().__init__(dtype=np.dtype(dtype), dims=dims, dimsd=dimsd, name=name) + + # other input params + self.taxis, self.haxis = taxis, haxis + self.nh, self.nt = self.dimsd + self.px = pxaxis + self.npx, self.nfft = self.dims[0], nfft + self.dt = taxis[1] - taxis[0] + self.dh = haxis[1] - haxis[0] + self.f = np.fft.rfftfreq(self.nfft, d=self.dt).astype(self.dtype) + self.nfft2 = len(self.f) + self.cdtype = get_complex_dtype(dtype) + self.flims = (0, self.nfft2) if flims is None else flims + + if kind == "parabolic": + self.haxis = self.haxis**2 + + # create additional input parameters for engine=cuda + if engine == "cuda": + self.num_threads_per_blocks = num_threads_per_blocks + ( + num_threads_per_blocks_hpx, + num_threads_per_blocks_f, + ) = num_threads_per_blocks + num_blocks_px = ( + self.dims[0] + num_threads_per_blocks_hpx - 1 + ) // num_threads_per_blocks_hpx + num_blocks_h = ( + self.dimsd[0] + num_threads_per_blocks_hpx - 1 + ) // num_threads_per_blocks_hpx + num_blocks_f = ( + self.dims[1] + num_threads_per_blocks_f - 1 + ) // num_threads_per_blocks_f + self.num_blocks_matvec = (num_blocks_h, num_blocks_f) + self.num_blocks_rmatvec = (num_blocks_px, num_blocks_f) + self._register_multiplications(engine) + + def _register_multiplications(self, engine: str) -> None: + if engine == "numba" and jit_message is None: + self._matvec = self._matvec_numba + self._rmatvec = self._rmatvec_numba + elif engine == "cuda": + self._matvec = self._matvec_cuda + self._rmatvec = self._rmatvec_cuda + else: + self._matvec = self._matvec_numpy + self._rmatvec = self._rmatvec_numpy + + @reshaped + def _matvec_numpy(self, x: NDArray) -> NDArray: + ncp = get_array_module(x) + self.f = ncp.asarray(self.f) + x = ncp.fft.rfft(x, n=self.nfft, axis=-1) + + H, PX, F = ncp.meshgrid( + self.haxis, self.px, self.f[self.flims[0] : self.flims[1]], indexing="ij" + ) + y = ncp.zeros((self.nh, self.nfft2), dtype=self.cdtype) + y[:, self.flims[0] : self.flims[1]] = ncp.einsum( + "ijk,jk->ik", + ncp.exp(-1j * 2 * ncp.pi * F * PX * H), + x[:, self.flims[0] : self.flims[1]], + ) + y = ncp.real(ncp.fft.irfft(y, n=self.nfft, axis=-1))[:, : self.nt] + return y + + @reshaped + def _rmatvec_numpy(self, y: NDArray) -> NDArray: + ncp = get_array_module(y) + self.f = ncp.asarray(self.f) + y = ncp.fft.rfft(y, n=self.nfft, axis=-1) + + PX, H, F = ncp.meshgrid( + self.px, self.haxis, self.f[self.flims[0] : self.flims[1]], indexing="ij" + ) + x = ncp.zeros((self.npx, self.nfft2), dtype=self.cdtype) + x[:, self.flims[0] : self.flims[1]] = ncp.einsum( + "ijk,jk->ik", + ncp.exp(1j * 2 * ncp.pi * F * PX * H), + y[:, self.flims[0] : self.flims[1]], + ) + x = ncp.real(ncp.fft.irfft(x, n=self.nfft, axis=-1))[:, : self.nt] + return x + + @reshaped + def _matvec_cuda(self, x: NDArray) -> NDArray: + ncp = get_array_module(x) + y = ncp.zeros((self.nh, self.nfft2), dtype=self.cdtype) + + x = ncp.fft.rfft(x, n=self.nfft, axis=-1) + y = _radon_inner_2d_cuda( + x, + y, + ncp.asarray(self.f), + self.px, + self.haxis, + self.flims[0], + self.flims[1], + self.npx, + self.nh, + num_blocks=self.num_blocks_matvec, + num_threads_per_blocks=self.num_threads_per_blocks, + ) + y = ncp.real(ncp.fft.irfft(y, n=self.nfft, axis=-1))[:, : self.nt] + return y + + @reshaped + def _rmatvec_cuda(self, y: NDArray) -> NDArray: + ncp = get_array_module(y) + x = ncp.zeros((self.npx, self.nfft2), dtype=self.cdtype) + + y = ncp.fft.rfft(y, n=self.nfft, axis=-1) + x = _aradon_inner_2d_cuda( + x, + y, + ncp.asarray(self.f), + self.px, + self.haxis, + self.flims[0], + self.flims[1], + self.npx, + self.nh, + num_blocks=self.num_blocks_rmatvec, + num_threads_per_blocks=self.num_threads_per_blocks, + ) + x = ncp.real(ncp.fft.irfft(x, n=self.nfft, axis=-1))[:, : self.nt] + return x + + @reshaped + def _matvec_numba(self, x: NDArray) -> NDArray: + y = np.zeros((self.nh, self.nfft2), dtype=self.cdtype) + + x = sp.fft.rfft(x, n=self.nfft, axis=-1) + _radon_inner_2d( + x, + y, + self.f, + self.px, + self.haxis, + self.flims[0], + self.flims[1], + self.npx, + self.nh, + ) + y = np.real(sp.fft.irfft(y, n=self.nfft, axis=-1))[:, : self.nt] + return y + + @reshaped + def _rmatvec_numba(self, y: NDArray) -> NDArray: + x = np.zeros((self.npx, self.nfft2), dtype=self.cdtype) + + y = sp.fft.rfft(y, n=self.nfft, axis=-1) + _aradon_inner_2d( + x, + y, + self.f, + self.px, + self.haxis, + self.flims[0], + self.flims[1], + self.npx, + self.nh, + ) + x = np.real(sp.fft.irfft(x, n=self.nfft, axis=-1))[:, : self.nt] + return x diff --git a/pylops/signalprocessing/fourierradon3d.py b/pylops/signalprocessing/fourierradon3d.py new file mode 100755 index 00000000..75622c85 --- /dev/null +++ b/pylops/signalprocessing/fourierradon3d.py @@ -0,0 +1,351 @@ +__all__ = ["FourierRadon3D"] + +import logging +from typing import Optional, Tuple + +import numpy as np +import scipy as sp + +from pylops import LinearOperator +from pylops.utils import deps +from pylops.utils.backend import get_array_module, get_complex_dtype +from pylops.utils.decorators import reshaped +from pylops.utils.typing import DTypeLike, NDArray + +jit_message = deps.numba_import("the radon2d module") + +if jit_message is None: + from ._fourierradon3d_cuda import _aradon_inner_3d_cuda, _radon_inner_3d_cuda + from ._fourierradon3d_numba import _aradon_inner_3d, _radon_inner_3d + +logging.basicConfig(format="%(levelname)s: %(message)s", level=logging.WARNING) + + +class FourierRadon3D(LinearOperator): + r"""3D Fourier Radon transform + + Apply Radon forward (and adjoint) transform using Fast + Fourier Transform to a 3-dimensional array of size :math:`[n_{p_y} \times n_{p_x} \times n_t]` + (and :math:`[n_y \times n_x \times n_t]`). + + Note that forward and adjoint follow the same convention of the time-space + implementation in :class:`pylops.signalprocessing.Radon3D`. + + Parameters + ---------- + taxis : :obj:`np.ndarray` + Time axis + hyaxis : :obj:`np.ndarray` + Slow spatial axis + hxaxis : :obj:`np.ndarray` + Fast spatial axis + pyaxis : :obj:`np.ndarray` + Axis of scanning variable :math:`p_y` of parametric curve + pxaxis : :obj:`np.ndarray` + Axis of scanning variable :math:`p_x` of parametric curve + nfft : :obj:`int` + Number of samples in Fourier transform + flims : :obj:`tuple`, optional + Indices of lower and upper limits of Fourier axis to be used in + the application of the Radon matrix (when ``None``, use entire axis) + kind : :obj:`tuple` + Curves to be used for stacking/spreading along the y- and x- axes + (``("linear", "linear")``, ``("linear", "parabolic")``, + ``("parabolic", "linear")``, or ``("parabolic", "parabolic")``) + engine : :obj:`str` + Engine used for computation (``numpy`` or ``numba`` or ``cuda``) + num_threads_per_blocks : :obj:`tuple` + Number of threads in each block (only when ``engine=cuda``) + dtype : :obj:`str` + Type of elements in input array. + name : :obj:`str` + 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 + ------ + NotImplementedError + If ``engine`` is neither ``numpy``, ``numba``, nor ``cuda``. + ValueError + If ``kind`` is not a tuple of two elements. + + Notes + ----- + The FourierRadon3D operator applies the Radon transform in the frequency domain. + After transforming a 3-dimensional array of size + :math:`[n_y \times n_x \times n_t]` into the frequency domain, the following linear + transformation is applied to each frequency component in adjoint mode: + + .. math:: + \begin{bmatrix} + \mathbf{m}(p_{y,1}, \mathbf{p}_{x}, \omega_i) \\ + \mathbf{m}(p_{y,2}, \mathbf{p}_{x}, \omega_i) \\ + \vdots \\ + \mathbf{m}(p_{y,N_{py}}, \mathbf{p}_{x}, \omega_i) + \end{bmatrix} + = + \begin{bmatrix} + e^{-j \omega_i (p_{y,1} y^l_1 + \mathbf{p}_x \cdot \mathbf{x}^l)} & e^{-j \omega_i (p_{y,1} y^l_2 + \mathbf{p}_x \cdot \mathbf{x}^l)} & \ldots & e^{-j \omega_i (p_{y,1} y^l_{N_y} + \mathbf{p}_x \cdot \mathbf{x}^l)} \\ + e^{-j \omega_i (p_{y,2} y^l_1 + \mathbf{p}_x \cdot \mathbf{x}^l)} & e^{-j \omega_i (p_{y,2} y^l_2 + \mathbf{p}_x \cdot \mathbf{x}^l)} & \ldots & e^{-j \omega_i (p_{y,2} y^l_{N_y} + \mathbf{p}_x \cdot \mathbf{x}^l)} \\ + \vdots & \vdots & \ddots & \vdots \\ + e^{-j \omega_i (p_{y,N_{py}} y^l_1 + \mathbf{p}_x \cdot \mathbf{x}^l)} & e^{-j \omega_i (p_{y,N_{py}} y^l_2 + \mathbf{p}_x \cdot \mathbf{x}^l)} & \ldots & e^{-j \omega_i (p_{y,N_{py}} y^l_{N_y} + \mathbf{p}_x \cdot \mathbf{x}^l)} \\ + \end{bmatrix} + \begin{bmatrix} + \mathbf{d}(y_1, \mathbf{x}, \omega_i) \\ + \mathbf{d}(y_2, \mathbf{x}, \omega_i) \\ + \vdots \\ + \mathbf{d}(y_{N_y}, \mathbf{x}, \omega_i) + \end{bmatrix} + + where :math:`\cdot` represents the element-wise multiplication of two vectors and + math:`l=1,2`. Similarly the forward mode is implemented by applying the + transpose and complex conjugate of the above matrix to the model transformed to + the Fourier domain. + + Refer to [1]_ for more theoretical and implementation details. + + .. [1] Sacchi, M. "Statistical and Transform Methods for + Geophysical Signal Processing", 2007. + """ + + def __init__( + self, + taxis: NDArray, + hyaxis: NDArray, + hxaxis: NDArray, + pyaxis: NDArray, + pxaxis: NDArray, + nfft: int, + flims: Optional[Tuple[int, int]] = None, + kind: Tuple[str, str] = ("linear", "linear"), + engine: str = "numpy", + num_threads_per_blocks: Tuple[int, int] = (32, 32), + dtype: DTypeLike = "float64", + name: str = "R", + ) -> None: + # engine + if engine not in ["numpy", "numba", "cuda"]: + raise NotImplementedError("engine must be numpy or numba or cuda") + if engine == "numba" and jit_message is not None: + engine = "numpy" + + # kind + if len(kind) != 2: + raise ValueError("kind must be a tuple of two elements") + + # dimensions and super + dims = len(pyaxis), len(pxaxis), len(taxis) + dimsd = len(hyaxis), len(hxaxis), len(taxis) + super().__init__(dtype=np.dtype(dtype), dims=dims, dimsd=dimsd, name=name) + + # other input params + self.taxis, self.hyaxis, self.hxaxis = taxis, hyaxis, hxaxis + self.nhy, self.nhx, self.nt = self.dimsd + self.py, self.px = pyaxis, pxaxis + self.npy, self.npx, self.nfft = self.dims[0], self.dims[1], nfft + self.dt = taxis[1] - taxis[0] + self.dhy = hyaxis[1] - hyaxis[0] + self.dhx = hxaxis[1] - hxaxis[0] + self.f = np.fft.rfftfreq(self.nfft, d=self.dt).astype(self.dtype) + self.nfft2 = len(self.f) + self.cdtype = get_complex_dtype(dtype) + self.flims = (0, self.nfft2) if flims is None else flims + + if kind[0] == "parabolic": + self.hyaxis = self.hyaxis**2 + if kind[1] == "parabolic": + self.hxaxis = self.hxaxis**2 + + # create additional input parameters for engine=cuda + if engine == "cuda": + self.num_threads_per_blocks = num_threads_per_blocks + ( + num_threads_per_blocks_hpy, + num_threads_per_blocks_hpx, + num_threads_per_blocks_f, + ) = num_threads_per_blocks + num_blocks_py = ( + self.dims[0] + num_threads_per_blocks_hpy - 1 + ) // num_threads_per_blocks_hpx + num_blocks_px = ( + self.dims[1] + num_threads_per_blocks_hpx - 1 + ) // num_threads_per_blocks_hpx + num_blocks_hy = ( + self.dimsd[0] + num_threads_per_blocks_hpy - 1 + ) // num_threads_per_blocks_hpx + num_blocks_hx = ( + self.dimsd[1] + num_threads_per_blocks_hpx - 1 + ) // num_threads_per_blocks_hpx + num_blocks_f = ( + self.dims[2] + num_threads_per_blocks_f - 1 + ) // num_threads_per_blocks_f + self.num_blocks_matvec = (num_blocks_hy, num_blocks_hx, num_blocks_f) + self.num_blocks_rmatvec = (num_blocks_py, num_blocks_px, num_blocks_f) + + self._register_multiplications(engine) + + def _register_multiplications(self, engine: str) -> None: + if engine == "numba" and jit_message is None: + self._matvec = self._matvec_numba + self._rmatvec = self._rmatvec_numba + elif engine == "cuda": + self._matvec = self._matvec_cuda + self._rmatvec = self._rmatvec_cuda + else: + self._matvec = self._matvec_numpy + self._rmatvec = self._rmatvec_numpy + + @reshaped + def _matvec_numpy(self, x: NDArray) -> NDArray: + ncp = get_array_module(x) + self.f = ncp.asarray(self.f) + x = ncp.fft.rfft(x.reshape(-1, self.dims[-1]), n=self.nfft, axis=-1) + + HY, HX = ncp.meshgrid(self.hyaxis, self.hxaxis, indexing="ij") + PY, PX = ncp.meshgrid(self.py, self.px, indexing="ij") + + HYY, PYY, F = ncp.meshgrid( + HY.ravel(), PY.ravel(), self.f[self.flims[0] : self.flims[1]], indexing="ij" + ) + HXX, PXX, _ = ncp.meshgrid( + HX.ravel(), PX.ravel(), self.f[self.flims[0] : self.flims[1]], indexing="ij" + ) + + y = ncp.zeros((self.nhy * self.nhx, self.nfft2), dtype=self.cdtype) + y[:, self.flims[0] : self.flims[1]] = ncp.einsum( + "ijk,jk->ik", + ncp.exp(-1j * 2 * ncp.pi * F * (PYY * HYY + PXX * HXX)), + x[:, self.flims[0] : self.flims[1]], + ) + y = ncp.real(ncp.fft.irfft(y, n=self.nfft, axis=-1))[:, : self.nt] + return y + + @reshaped + def _rmatvec_numpy(self, y: NDArray) -> NDArray: + ncp = get_array_module(y) + self.f = ncp.asarray(self.f) + y = ncp.fft.rfft(y.reshape(-1, self.dimsd[-1]), n=self.nfft, axis=-1) + + HY, HX = ncp.meshgrid(self.hyaxis, self.hxaxis, indexing="ij") + PY, PX = ncp.meshgrid(self.py, self.px, indexing="ij") + + PYY, HYY, F = ncp.meshgrid( + PY.ravel(), HY.ravel(), self.f[self.flims[0] : self.flims[1]], indexing="ij" + ) + PXX, HXX, _ = ncp.meshgrid( + PX.ravel(), HX.ravel(), self.f[self.flims[0] : self.flims[1]], indexing="ij" + ) + + x = ncp.zeros((self.npy * self.npx, self.nfft2), dtype=self.cdtype) + x[:, self.flims[0] : self.flims[1]] = ncp.einsum( + "ijk,jk->ik", + ncp.exp(1j * 2 * ncp.pi * F * (PYY * HYY + PXX * HXX)), + y[:, self.flims[0] : self.flims[1]], + ) + x = ncp.real(ncp.fft.irfft(x, n=self.nfft, axis=-1))[:, : self.nt] + return x + + @reshaped + def _matvec_cuda(self, x: NDArray) -> NDArray: + ncp = get_array_module(x) + y = ncp.zeros((self.nhy, self.nhx, self.nfft2), dtype=self.cdtype) + + x = ncp.fft.rfft(x, n=self.nfft, axis=-1) + y = _radon_inner_3d_cuda( + x, + y, + ncp.asarray(self.f), + self.py, + self.px, + self.hyaxis, + self.hxaxis, + self.flims[0], + self.flims[1], + self.npy, + self.npx, + self.nhy, + self.nhx, + num_blocks=self.num_blocks_matvec, + num_threads_per_blocks=self.num_threads_per_blocks, + ) + y = ncp.real(ncp.fft.irfft(y, n=self.nfft, axis=-1))[:, :, : self.nt] + return y + + @reshaped + def _rmatvec_cuda(self, y: NDArray) -> NDArray: + ncp = get_array_module(y) + x = ncp.zeros((self.npy, self.npx, self.nfft2), dtype=self.cdtype) + + y = ncp.fft.rfft(y, n=self.nfft, axis=-1) + x = _aradon_inner_3d_cuda( + x, + y, + ncp.asarray(self.f), + self.py, + self.px, + self.hyaxis, + self.hxaxis, + self.flims[0], + self.flims[1], + self.npy, + self.npx, + self.nhy, + self.nhx, + num_blocks=self.num_blocks_rmatvec, + num_threads_per_blocks=self.num_threads_per_blocks, + ) + x = ncp.real(ncp.fft.irfft(x, n=self.nfft, axis=-1))[:, :, : self.nt] + return x + + @reshaped + def _matvec_numba(self, x: NDArray) -> NDArray: + y = np.zeros((self.nhy, self.nhx, self.nfft2), dtype=self.cdtype) + + x = sp.fft.rfft(x, n=self.nfft, axis=-1) + _radon_inner_3d( + x, + y, + self.f, + self.py, + self.px, + self.hyaxis, + self.hxaxis, + self.flims[0], + self.flims[1], + self.npy, + self.npx, + self.nhy, + self.nhx, + ) + y = np.real(sp.fft.irfft(y, n=self.nfft, axis=-1))[:, :, : self.nt] + return y + + @reshaped + def _rmatvec_numba(self, y: NDArray) -> NDArray: + x = np.zeros((self.npy, self.npx, self.nfft2), dtype=self.cdtype) + + y = sp.fft.rfft(y, n=self.nfft, axis=-1) + _aradon_inner_3d( + x, + y, + self.f, + self.py, + self.px, + self.hyaxis, + self.hxaxis, + self.flims[0], + self.flims[1], + self.npy, + self.npx, + self.nhy, + self.nhx, + ) + x = np.real(sp.fft.irfft(x, n=self.nfft, axis=-1))[:, :, : self.nt] + return x diff --git a/pylops/signalprocessing/radon2d.py b/pylops/signalprocessing/radon2d.py index d045eb8a..a76696b4 100644 --- a/pylops/signalprocessing/radon2d.py +++ b/pylops/signalprocessing/radon2d.py @@ -218,7 +218,7 @@ def Radon2D( ----- The Radon2D operator applies the following linear transform in adjoint mode to the data after reshaping it into a 2-dimensional array of - size :math:`[n_x \times n_t]` in adjoint mode: + size :math:`[n_x \times n_t]`: .. math:: m(p_x, t_0) = \int{d(x, t = f(p_x, x, t))} \,\mathrm{d}x diff --git a/pylops/signalprocessing/radon3d.py b/pylops/signalprocessing/radon3d.py index 4357b983..939c14ba 100644 --- a/pylops/signalprocessing/radon3d.py +++ b/pylops/signalprocessing/radon3d.py @@ -189,10 +189,10 @@ def Radon3D( ---------- taxis : :obj:`np.ndarray` Time axis - hxaxis : :obj:`np.ndarray` - Fast patial axis hyaxis : :obj:`np.ndarray` Slow spatial axis + hxaxis : :obj:`np.ndarray` + False spatial axis pyaxis : :obj:`np.ndarray` Axis of scanning variable :math:`p_y` of parametric curve pxaxis : :obj:`np.ndarray` diff --git a/pylops/utils/estimators.py b/pylops/utils/estimators.py index 887fac3f..5e73f686 100644 --- a/pylops/utils/estimators.py +++ b/pylops/utils/estimators.py @@ -87,7 +87,7 @@ def trace_hutchinson( Operator trace. Raises - ------- + ------ ValueError If ``neval`` is smaller than 3. @@ -194,7 +194,7 @@ def trace_hutchpp( Operator trace. Raises - ------- + ------ ValueError If ``neval`` is smaller than 3. @@ -294,7 +294,7 @@ def trace_nahutchpp( Operator trace. Raises - ------- + ------ ValueError If ``neval`` not large enough to accomodate ``c1`` and ``c2``. diff --git a/pylops/utils/seismicevents.py b/pylops/utils/seismicevents.py index 1844c82d..b061734e 100755 --- a/pylops/utils/seismicevents.py +++ b/pylops/utils/seismicevents.py @@ -4,6 +4,7 @@ "parabolic2d", "hyperbolic2d", "linear3d", + "parabolic3d", "hyperbolic3d", ] @@ -13,6 +14,8 @@ import numpy.typing as npt import scipy.signal as filt +from pylops.utils._internal import _value_or_sized_to_array + def _filterdata( d: npt.NDArray, nt: int, wav: npt.ArrayLike, wcenter: int @@ -116,12 +119,9 @@ def linear2d( where :math:`p_{x,i}=\sin( \theta_i)/v` """ - if isinstance(t0, (float, int)): - t0 = (t0,) - if isinstance(theta, (float, int)): - theta = (theta,) - if isinstance(amp, (float, int)): - amp = (amp,) + t0 = _value_or_sized_to_array(t0) + theta = _value_or_sized_to_array(theta) + amp = _value_or_sized_to_array(amp) # identify dimensions dt = t[1] - t[0] @@ -160,7 +160,7 @@ def parabolic2d( r"""Parabolic 2D events Create 2d parabolic events given intercept time, - slowness, curvature, and amplitude of each event + slowness, curvature, and amplitude of each event. Parameters ---------- @@ -196,14 +196,10 @@ def parabolic2d( t_i(x) = t_{0,i} + p_{x,i} x + p_{xx,i} x^2 """ - if isinstance(t0, (float, int)): - t0 = (t0,) - if isinstance(px, (float, int)): - px = (px,) - if isinstance(pxx, (float, int)): - pxx = (pxx,) - if isinstance(amp, (float, int)): - amp = (amp,) + t0 = _value_or_sized_to_array(t0) + px = _value_or_sized_to_array(px) + pxx = _value_or_sized_to_array(pxx) + amp = _value_or_sized_to_array(amp) # identify dimensions dt = t[1] - t[0] @@ -272,12 +268,9 @@ def hyperbolic2d( t_i(x) = \sqrt{t_{0,i}^2 + \frac{x^2}{v_{\text{rms},i}^2}} """ - if isinstance(t0, (float, int)): - t0 = (t0,) - if isinstance(vrms, (float, int)): - vrms = (vrms,) - if isinstance(amp, (float, int)): - amp = (amp,) + t0 = _value_or_sized_to_array(t0) + vrms = _value_or_sized_to_array(vrms) + amp = _value_or_sized_to_array(amp) # identify dimensions dt = t[1] - t[0] @@ -330,7 +323,7 @@ def linear3d( v : :obj:`float` propagation velocity t0 : :obj:`tuple` or :obj:`float` - intercept time at :math:`x=0` of each linear event + intercept time at :math:`x=0` and :math:`y=0` of each linear event theta : :obj:`tuple` or :obj:`float` angle in x direction (in degrees) of each linear event phi : :obj:`tuple` or :obj:`float` @@ -360,14 +353,10 @@ def linear3d( and :math:`p_{x,i}=\frac{1}{v} \sin( \theta_i)\sin( \phi_i)`. """ - if isinstance(t0, (float, int)): - t0 = (t0,) - if isinstance(theta, (float, int)): - theta = (theta,) - if isinstance(phi, (float, int)): - phi = (phi,) - if isinstance(amp, (float, int)): - amp = (amp,) + t0 = _value_or_sized_to_array(t0) + theta = _value_or_sized_to_array(theta) + phi = _value_or_sized_to_array(phi) + amp = _value_or_sized_to_array(amp) # identify dimensions dt = t[1] - t[0] @@ -397,6 +386,100 @@ def linear3d( return d, dwav +def parabolic3d( + x: npt.NDArray, + y: npt.NDArray, + t: npt.NDArray, + t0: Union[float, Tuple[float]], + px: Union[float, Tuple[float]], + py: Union[float, Tuple[float]], + pxx: Union[float, Tuple[float]], + pyy: Union[float, Tuple[float]], + amp: Union[float, Tuple[float]], + wav: npt.NDArray, +) -> Tuple[npt.NDArray, npt.NDArray]: + r"""Parabolic 3D events + + Create 3d parabolic events given intercept time, + slowness, curvature, and amplitude of each event. + + Parameters + ---------- + x : :obj:`numpy.ndarray` + space axis in x direction + y : :obj:`numpy.ndarray` + space axis in y direction + t : :obj:`numpy.ndarray` + time axis + t0 : :obj:`tuple` or :obj:`float` + intercept time at :math:`x=0` and :math:`y=0` of each parabolic event + px : :obj:`tuple` or :obj:`float` + slowness of each parabolic event in x direction + py : :obj:`tuple` or :obj:`float` + slowness of each parabolic event in y direction + pxx : :obj:`tuple` or :obj:`float` + curvature of each parabolic event + amp : :obj:`tuple` or :obj:`float` + amplitude of each linear event + wav : :obj:`numpy.ndarray` + wavelet to be applied to data + + Returns + ------- + d : :obj:`numpy.ndarray` + data without wavelet of size + :math:`[n_y \times n_x \times n_t]` + dwav : :obj:`numpy.ndarray` + data with wavelet of size + :math:`[n_y \times n_x \times n_t]` + + Notes + ----- + Each event is created using the following relation: + + .. math:: + t_i(x, y) = t_{0,i} + p_{x,i} x + p_{y,i} x + p_{xx,i} x^2 + p_{yy,i} y^2 + + """ + t0 = _value_or_sized_to_array(t0) + px = _value_or_sized_to_array(px) + py = _value_or_sized_to_array(py) + pxx = _value_or_sized_to_array(pxx) + pyy = _value_or_sized_to_array(pyy) + amp = _value_or_sized_to_array(amp) + + # identify dimensions + dt = t[1] - t[0] + wcenter = int(len(wav) / 2) + nx = np.size(x) + ny = np.size(y) + nt = np.size(t) + wcenter + nevents = np.size(t0) + + # create events + d = np.zeros((ny, nx, nt)) + for ievent in range(nevents): + for iy in range(ny): + tevent = ( + t0[ievent] + + px[ievent] * x + + py[ievent] * y[iy] + + pxx[ievent] * x**2 + + pyy[ievent] * y[iy] ** 2 + ) + tevent = (tevent - t[0]) / dt + itevent = tevent.astype(int) + dtevent = tevent - itevent + for ix in range(nx): + if itevent[ix] < nt - 1 and itevent[ix] >= 0: + d[iy, ix, itevent[ix]] += amp[ievent] * (1 - dtevent[ix]) + d[iy, ix, itevent[ix] + 1] += amp[ievent] * dtevent[ix] + + # filter events with certain wavelet + d, dwav = _filterdata(d, nt, wav, wcenter) + return d, dwav + + def hyperbolic3d( x: npt.NDArray, y: npt.NDArray, @@ -451,14 +534,10 @@ def hyperbolic3d( simply control the curvature of the hyperboloid along the spatial axes. """ - if isinstance(t0, (float, int)): - t0 = (t0,) - if isinstance(vrms_x, (float, int)): - vrms_x = (vrms_x,) - if isinstance(vrms_y, (float, int)): - vrms_y = (vrms_y,) - if isinstance(amp, (float, int)): - amp = (amp,) + t0 = _value_or_sized_to_array(t0) + vrms_x = _value_or_sized_to_array(vrms_x) + vrms_y = _value_or_sized_to_array(vrms_y) + amp = _value_or_sized_to_array(amp) # identify dimensions dt = t[1] - t[0] diff --git a/pytests/test_fourierradon.py b/pytests/test_fourierradon.py new file mode 100755 index 00000000..accbf257 --- /dev/null +++ b/pytests/test_fourierradon.py @@ -0,0 +1,135 @@ +import numpy as np +import pytest +from numpy.testing import assert_array_almost_equal + +from pylops.optimization.sparsity import fista +from pylops.signalprocessing import FourierRadon2D, FourierRadon3D +from pylops.utils import dottest + +par1 = { + "nt": 11, + "nhx": 21, + "nhy": 10, + "npx": 21, + "npy": 17, + "pymax": 1e-2, + "pxmax": 2e-2, + "centeredh": True, + "kind": "linear", + "interp": True, + "engine": "numpy", +} # linear, numpy +par2 = { + "nt": 11, + "nhx": 21, + "nhy": 10, + "npx": 21, + "npy": 17, + "pymax": 1e-2, + "pxmax": 2e-2, + "kind": "linear", + "engine": "numba", +} # linear, numba +par3 = { + "nt": 11, + "nhx": 21, + "nhy": 10, + "npx": 21, + "npy": 17, + "pymax": 8e-3, + "pxmax": 7e-3, + "kind": "parabolic", + "engine": "numpy", +} # parabolic, numpy +par4 = { + "nt": 11, + "nhx": 21, + "nhy": 10, + "npx": 21, + "npy": 17, + "pymax": 8e-3, + "pxmax": 7e-3, + "kind": "parabolic", + "engine": "numba", +} # parabolic, numba + + +def test_unknown_engine2D(): + """Check error is raised if unknown engine is passed""" + with pytest.raises(NotImplementedError): + _ = FourierRadon2D(None, None, None, None, engine="foo") + + +def test_unknown_engine3D(): + """Check error is raised if unknown engine is passed""" + with pytest.raises(NotImplementedError): + _ = FourierRadon3D(None, None, None, None, None, None, engine="foo") + + +@pytest.mark.parametrize("par", [(par1), (par2), (par3), (par4)]) +def test_FourierRadon2D(par): + """Dot-test and sparse inverse for FourierRadon2D operator""" + dt, dh = 0.005, 1 + t = np.arange(par["nt"]) * dt + h = np.arange(par["nhx"]) * dh + px = np.linspace(0, par["pxmax"], par["npx"]) + nfft = int(2 ** np.ceil(np.log2(par["nt"]))) + + x = np.zeros((par["npx"], par["nt"])) + x[2, par["nt"] // 2] = 1 + + Rop = FourierRadon2D( + t, + h, + px, + nfft, + kind=par["kind"], + engine=par["engine"], + dtype="float64", + ) + assert dottest(Rop, par["nhx"] * par["nt"], par["npx"] * par["nt"], rtol=1e-3) + + y = Rop * x.ravel() + + if par["engine"] == "numba": # as numpy is too slow here... + xinv, _, _ = fista(Rop, y, niter=200, eps=3e0) + assert_array_almost_equal(x.ravel(), xinv, decimal=1) + + +@pytest.mark.parametrize("par", [(par1), (par2), (par3), (par4)]) +def test_FourierRadon3D(par): + """Dot-test and sparse inverse for FourierRadon3D operator""" + dt, dhy, dhx = 0.005, 1, 1 + t = np.arange(par["nt"]) * dt + hy = np.arange(par["nhy"]) * dhy + hx = np.arange(par["nhx"]) * dhx + py = np.linspace(0, par["pymax"], par["npy"]) + px = np.linspace(0, par["pxmax"], par["npx"]) + nfft = int(2 ** np.ceil(np.log2(par["nt"]))) + + x = np.zeros((par["npy"], par["npx"], par["nt"])) + x[3, 2, par["nt"] // 2] = 1 + + Rop = FourierRadon3D( + t, + hy, + hx, + py, + px, + nfft, + kind=(par["kind"], par["kind"]), + engine=par["engine"], + dtype="float64", + ) + assert dottest( + Rop, + par["nhy"] * par["nhx"] * par["nt"], + par["npy"] * par["npx"] * par["nt"], + rtol=1e-3, + ) + + y = Rop * x.ravel() + + if par["engine"] == "numba": # as numpy is too slow here... + xinv, _, _ = fista(Rop, y, niter=200, eps=1e1) + assert_array_almost_equal(x.ravel(), xinv, decimal=1) diff --git a/pytests/test_radon.py b/pytests/test_radon.py index cc4b8054..26d08ba1 100755 --- a/pytests/test_radon.py +++ b/pytests/test_radon.py @@ -1,5 +1,3 @@ -import multiprocessing - import numpy as np import pytest from numpy.testing import assert_array_almost_equal @@ -176,63 +174,59 @@ def test_Radon2D(par): "par", [(par1), (par2), (par3), (par4), (par5), (par6), (par7), (par8)] ) def test_Radon3D(par): - """Dot-test, forward and adjoint consistency check + """Dot-test, forward and adjoint consistency check (for onthefly parameter), and sparse inverse for Radon3D operator """ - if ( - par["engine"] == "numpy" or multiprocessing.cpu_count() >= 4 - ): # avoid timeout in travis for numba - - dt, dhy, dhx = 0.005, 1, 1 - t = np.arange(par["nt"]) * dt - hy = np.arange(par["nhy"]) * dhy - hx = np.arange(par["nhx"]) * dhx - py = np.linspace(0, par["pymax"], par["npy"]) - px = np.linspace(0, par["pxmax"], par["npx"]) - x = np.zeros((par["npy"], par["npx"], par["nt"])) - x[3, 2, par["nt"] // 2] = 1 - - Rop = Radon3D( - t, - hy, - hx, - py, - px, - centeredh=par["centeredh"], - interp=par["interp"], - kind=par["kind"], - onthefly=False, - engine=par["engine"], - dtype="float64", - ) - R1op = Radon3D( - t, - hy, - hx, - py, - px, - centeredh=par["centeredh"], - interp=par["interp"], - kind=par["kind"], - onthefly=True, - engine=par["engine"], - dtype="float64", - ) - - assert dottest( - Rop, - par["nhy"] * par["nhx"] * par["nt"], - par["npy"] * par["npx"] * par["nt"], - rtol=1e-3, - ) - y = Rop * x.ravel() - y1 = R1op * x.ravel() - assert_array_almost_equal(y, y1, decimal=4) - - xadj = Rop.H * y - xadj1 = R1op.H * y - assert_array_almost_equal(xadj, xadj1, decimal=4) - - if Rop.engine == "numba": # as numpy is too slow here... - xinv, _, _ = fista(Rop, y, niter=200, eps=3e0) - assert_array_almost_equal(x.ravel(), xinv, decimal=1) + dt, dhy, dhx = 0.005, 1, 1 + t = np.arange(par["nt"]) * dt + hy = np.arange(par["nhy"]) * dhy + hx = np.arange(par["nhx"]) * dhx + py = np.linspace(0, par["pymax"], par["npy"]) + px = np.linspace(0, par["pxmax"], par["npx"]) + x = np.zeros((par["npy"], par["npx"], par["nt"])) + x[3, 2, par["nt"] // 2] = 1 + + Rop = Radon3D( + t, + hy, + hx, + py, + px, + centeredh=par["centeredh"], + interp=par["interp"], + kind=par["kind"], + onthefly=False, + engine=par["engine"], + dtype="float64", + ) + R1op = Radon3D( + t, + hy, + hx, + py, + px, + centeredh=par["centeredh"], + interp=par["interp"], + kind=par["kind"], + onthefly=True, + engine=par["engine"], + dtype="float64", + ) + + assert dottest( + Rop, + par["nhy"] * par["nhx"] * par["nt"], + par["npy"] * par["npx"] * par["nt"], + rtol=1e-3, + ) + y = Rop * x.ravel() + y1 = R1op * x.ravel() + assert_array_almost_equal(y, y1, decimal=4) + + xadj = Rop.H * y + xadj1 = R1op.H * y + assert_array_almost_equal(xadj, xadj1, decimal=4) + + if Rop.engine == "numba": # as numpy is too slow here... + xinv, _, _ = fista(Rop, y, niter=200, eps=3e0) + assert_array_almost_equal(x.ravel(), xinv, decimal=1)