Bug: fixed Deghosting to work with either sparsity or restriction as well as both

pylops/waveeqprocessing/oneway.py

@@ -3,7 +3,7 @@
     "Deghosting",
 ]
 
-from collections.abc import Callable, Sequence
+from collections.abc import Callable
 
 import numpy as np
 from scipy.sparse.linalg import lsqr
@@ -13,7 +13,7 @@
 from pylops.utils import dottest as Dottest
 from pylops.utils.backend import to_cupy_conditional
 from pylops.utils.tapers import taper2d, taper3d
-from pylops.utils.typing import DTypeLike, NDArray, Tfftengine_nsf
+from pylops.utils.typing import DTypeLike, NDArray, SamplingLike, Tfftengine_nsf
 
 
 class _PhaseShift(LinearOperator):
@@ -204,7 +204,7 @@ def Deghosting(
     nt: int,
     nr: int | tuple[int, int],
     dt: float,
-    dr: Sequence[float],
+    dr: float | SamplingLike,
     vel: float,
     zrec: float,
     kind: str | None = "p",
@@ -238,7 +238,9 @@ def Deghosting(
     nt : :obj:`int`
         Number of samples along the time axis
     nr : :obj:`int` or :obj:`tuple`
-        Number of samples along the receiver axis (or axes)
+        Number of samples along the receiver axis (or axes) of the separated
+        pressure consituents. If ``restriction`` is not provided, this equals
+        to the number of samples along the receiver axis (or axes) of ``p``.
     dt : :obj:`float`
         Sampling along the time axis
     dr : :obj:`float` or :obj:`tuple`
@@ -267,9 +269,10 @@ def Deghosting(
     restriction : :obj:`pylops.LinearOperator`, optional
         Restriction operator
     sptransf : :obj:`pylops.LinearOperator`, optional
-        Sparsifying operator
+        Sparsifying operator used to represent the model in a
+        compressed domain.
     solver : :obj:`float`, optional
-        Function handle of solver to be used if ``kind='inverse'``
+        Function handle of solver
     dottest : :obj:`bool`, optional
         Apply dot-test
     dtype : :obj:`str`, optional
@@ -283,7 +286,11 @@ def Deghosting(
     pup : :obj:`numpy.ndarray`
         Up-going pressure (or particle velocity) wavefield
     pdown : :obj:`numpy.ndarray`
-        Down-going (or particle velocity) wavefield
+        Down-going (or particle velocity) wavefield. Note that when a ``restriction``
+        operator is provided, this is the down-going wavefield reconstructed at the
+        finely sampled spatial axes by applying the ghosting operator to the estimated
+        up-going wavefield (therefore it does not contain the direct arrival and
+        its source ghost).
 
     Raises
     ------
@@ -317,6 +324,14 @@ def Deghosting(
     used for the pressure data, whilst :math:`-` is used for the vertical velocity
     data.
 
+    In the case of densely sampled data, the above equation can be solved in a
+    least-squares sense via a linear solver (e.g., :func:`scipy.sparse.linalg.lsqr`).
+    Conversely, in the case of sparsely sampled data, a restriction operator can be
+    applied to the propagator to account for the missing data. In this case, the
+    solution can be regularized by applying a sparsifying transform to the
+    model (e.g., a patched FK transform) and solving the problem in a compressed
+    domain via a sparse solver (e.g., :func:`pylops.optimization.sparsity.ista`).
+
     .. [1] Amundsen, L., 1993, Wavenumber-based filtering of marine point-source
        data: GEOPHYSICS, 58, 1335–1348.
 
@@ -340,7 +355,7 @@ def Deghosting(
         nkx = nr[0] + 2 * npad[0]
         kx = np.fft.ifftshift(np.fft.fftfreq(nkx, dr[0]))
         nky = nr[1] + 2 * npad[1]
-        ky = np.fft.ifftshift(np.fft.fftfreq(nky, dr))
+        ky = np.fft.ifftshift(np.fft.fftfreq(nky, dr[1]))
     nf = nt
     freq = np.fft.rfftfreq(nf, dt)
 
@@ -370,8 +385,8 @@ def Deghosting(
         Dottest(Dupop, nt * nrs, nt * nrs, verb=True)
 
     # Add restriction
+    Dupop_norestr = Dupop
     if restriction is not None:
-        Dupop_norestr = Dupop
         Dupop = restriction * Dupop
 
     # Add sparsify transform
@@ -388,11 +403,14 @@ def Deghosting(
     # Inversion
     pup = solver(Dupop, d.ravel(), **kwargs_solver)[0]
 
-    # Apply sparse transform
+    # Reconstruct total wavefield at finely sampled spatial axes
+    if restriction is not None:
+        p = Dupop_norestr * pup
+        p = np.real(p).reshape(dims)
+
+    # Reconstruct up-going wavefield (apply sparsifying transform if provided)
     if sptransf is not None:
-        p = Dupop_norestr * pup  # reconstruct p at finely sampled spatial axes
         pup = sptransf * pup
-        p = np.real(p).reshape(dims)
 
     # Finalize estimates
     pup = np.real(pup).reshape(dims)

pytests/test_oneway.py

@@ -11,28 +11,18 @@
 import numpy as npp
 import pytest
 
-from pylops.basicoperators import Identity
+from pylops.basicoperators import Identity, Restriction
 from pylops.optimization.basic import lsqr
+from pylops.optimization.sparsity import fista
+from pylops.signalprocessing import FFT2D, FFTND
 from pylops.utils import dottest
-from pylops.utils.seismicevents import hyperbolic2d, makeaxis
+from pylops.utils.seismicevents import hyperbolic2d, hyperbolic3d, makeaxis
 from pylops.utils.wavelets import ricker
 from pylops.waveeqprocessing.oneway import Deghosting, PhaseShift
 
 np.random.seed(10)
 
-parmod = {
-    "ox": -400,
-    "dx": 10,
-    "nx": 81,
-    "oy": -50,
-    "dy": 10,
-    "ny": 11,
-    "ot": 0,
-    "dt": 0.004,
-    "nt": 50,
-    "f0": 40,
-}
-
+# phaseshift params
 par1 = {
     "ny": 8,
     "nx": 10,
@@ -79,6 +69,18 @@
 }  # odd, vz, numpy
 
 # deghosting params
+parmod = {
+    "ox": -400,
+    "dx": 10,
+    "nx": 81,
+    "oy": -50,
+    "dy": 10,
+    "ny": 11,
+    "ot": 0,
+    "dt": 0.004,
+    "nt": 100,
+    "f0": 40,
+}
 vel_sep = 1000.0  # velocity at separation level
 zrec = 20.0  # depth of receivers
 
@@ -109,6 +111,31 @@ def create_data2D(datakind):
     return np.asarray(p2d), np.asarray(p2d_minus)
 
 
+def create_data3D(datakind):
+    """Create 3d dataset"""
+    t0_plus = npp.array([0.02, 0.08])
+    t0_minus = t0_plus + 0.04
+    vrms = npp.array([1400.0, 1800.0])
+    amp = npp.array([1.0, -0.6])
+
+    p3d_minus = hyperbolic3d(x, y, t, t0_minus, vrms, vrms, amp, wav)[1].transpose(
+        2, 1, 0
+    )
+
+    kx = npp.fft.ifftshift(npp.fft.fftfreq(parmod["nx"], parmod["dx"]))
+    ky = npp.fft.ifftshift(npp.fft.fftfreq(parmod["ny"], parmod["dy"]))
+    freq = npp.fft.rfftfreq(parmod["nt"], parmod["dt"])
+
+    Pop = -PhaseShift(vel_sep, 2 * zrec, parmod["nt"], freq, kx, ky)
+
+    # Decomposition operator
+    Dupop = Identity(parmod["nt"] * parmod["nx"] * parmod["ny"]) + datakind * Pop
+
+    p3d = Dupop * p3d_minus.ravel()
+    p3d = p3d.reshape(parmod["nt"], parmod["nx"], parmod["ny"])
+    return np.asarray(p3d), np.asarray(p3d_minus)
+
+
 @pytest.mark.parametrize("par", [(par1), (par2), (par1s), (par1w)])
 @pytest.mark.parametrize("dtype", [np.float32, np.float64])
 def test_PhaseShift_2dsignal(par, dtype):
@@ -193,3 +220,180 @@ def test_Deghosting_2dsignal(par):
     )
 
     assert np.linalg.norm(p2d_minus_inv - p2d_minus) / np.linalg.norm(p2d_minus) < 3e-1
+
+
+@pytest.mark.parametrize("par", [(par1), (par2), (par1v), (par2v)])
+def test_Deghosting_2dsignal_sptransf(par):
+    """Deghosting of 2d data with FK sparsifying transform"""
+    p2d, p2d_minus = create_data2D(1 if par["kind"] == "p" else -1)
+
+    FOp = FFT2D(
+        dims=(parmod["nt"], parmod["nx"]),
+        sampling=(parmod["dt"], parmod["dx"]),
+        dtype=np.complex128,
+    )
+
+    p2d_minus_inv, _ = Deghosting(
+        p2d,
+        parmod["nt"],
+        parmod["nx"],
+        parmod["dt"],
+        parmod["dx"],
+        vel_sep,
+        zrec,
+        kind=par["kind"],
+        win=np.ones_like(p2d),
+        npad=0,
+        ntaper=0,
+        sptransf=FOp.H,
+        solver=fista,
+        dtype=np.float32,
+        **dict(eps=1e-4, niter=100),
+    )
+
+    assert np.linalg.norm(p2d_minus_inv - p2d_minus) / np.linalg.norm(p2d_minus) < 2e-1
+
+
+@pytest.mark.parametrize("par", [(par1), (par2), (par1v), (par2v)])
+def test_Deghosting_2dsignal_restr(par):
+    """Deghosting of 2d data with restriction and FK sparsifying transform"""
+    npp.random.seed(10)
+
+    p2d, p2d_minus = create_data2D(1 if par["kind"] == "p" else -1)
+
+    # Restriction operator (90% of available traces)
+    nsub = int(npp.round(0.9 * parmod["nx"]))
+    iava = np.asarray(npp.sort(npp.random.permutation(parmod["nx"])[:nsub]))
+    Rop = Restriction((parmod["nt"], parmod["nx"]), iava, axis=1, dtype=np.complex128)
+    p2d_sub = np.real(Rop * p2d.ravel()).reshape(parmod["nt"], nsub)
+
+    FOp = FFT2D(
+        dims=(parmod["nt"], parmod["nx"]),
+        sampling=(parmod["dt"], parmod["dx"]),
+        dtype=np.complex128,
+    )
+
+    p2d_minus_inv, p2d_plus_inv = Deghosting(
+        p2d_sub,
+        parmod["nt"],
+        parmod["nx"],
+        parmod["dt"],
+        parmod["dx"],
+        vel_sep,
+        zrec,
+        kind=par["kind"],
+        win=np.ones_like(p2d_sub),
+        npad=0,
+        ntaper=0,
+        restriction=Rop,
+        sptransf=FOp.H,
+        solver=fista,
+        dtype=np.float32,
+        **dict(eps=1e-4, niter=100),
+    )
+
+    assert p2d_minus_inv.shape == (parmod["nt"], parmod["nx"])
+    assert p2d_plus_inv.shape == (parmod["nt"], parmod["nx"])
+    assert np.linalg.norm(p2d_minus_inv - p2d_minus) / np.linalg.norm(p2d_minus) < 4e-1
+
+
+@pytest.mark.parametrize("par", [(par1), (par1v)])
+def test_Deghosting_3dsignal(par):
+    """Deghosting of 3d data"""
+    p3d, p3d_minus = create_data3D(1 if par["kind"] == "p" else -1)
+
+    p3d_minus_inv, _ = Deghosting(
+        p3d,
+        parmod["nt"],
+        (parmod["nx"], parmod["ny"]),
+        parmod["dt"],
+        (parmod["dx"], parmod["dy"]),
+        vel_sep,
+        zrec,
+        kind=par["kind"],
+        win=np.ones_like(p3d),
+        npad=(0, 0),
+        ntaper=(0, 0),
+        solver=lsqr,
+        dtype=np.float32,
+        **dict(damp=1e-10, niter=60),
+    )
+
+    assert np.linalg.norm(p3d_minus_inv - p3d_minus) / np.linalg.norm(p3d_minus) < 3e-1
+
+
+@pytest.mark.parametrize("par", [(par1), (par1v)])
+def test_Deghosting_3dsignal_sptransf(par):
+    """Deghosting of 3d data with FK sparsifying transform"""
+    p3d, p3d_minus = create_data3D(1 if par["kind"] == "p" else -1)
+
+    FOp = FFTND(
+        dims=(parmod["nt"], parmod["nx"], parmod["ny"]),
+        sampling=(parmod["dt"], parmod["dx"], parmod["dy"]),
+        dtype=np.complex128,
+    )
+
+    p3d_minus_inv, _ = Deghosting(
+        p3d,
+        parmod["nt"],
+        (parmod["nx"], parmod["ny"]),
+        parmod["dt"],
+        (parmod["dx"], parmod["dy"]),
+        vel_sep,
+        zrec,
+        kind=par["kind"],
+        win=np.ones_like(p3d),
+        npad=(0, 0),
+        ntaper=(0, 0),
+        sptransf=FOp.H,
+        solver=fista,
+        dtype=np.float32,
+        **dict(eps=1e-4, niter=50),
+    )
+
+    assert np.linalg.norm(p3d_minus_inv - p3d_minus) / np.linalg.norm(p3d_minus) < 2e-1
+
+
+@pytest.mark.parametrize("par", [(par1), (par1v)])
+def test_Deghosting_3dsignal_restr(par):
+    """Deghosting of 3d data with restriction and FK sparsifying transform"""
+    npp.random.seed(10)
+
+    p3d, p3d_minus = create_data3D(1 if par["kind"] == "p" else -1)
+
+    # Restriction operator (90% of available traces along x)
+    nsub = int(npp.round(0.9 * parmod["nx"]))
+    iava = np.asarray(npp.sort(npp.random.permutation(parmod["nx"])[:nsub]))
+    Rop = Restriction(
+        (parmod["nt"], parmod["nx"], parmod["ny"]), iava, axis=1, dtype=np.complex128
+    )
+    p3d_sub = np.real(Rop * p3d.ravel()).reshape(parmod["nt"], nsub, parmod["ny"])
+
+    FOp = FFTND(
+        dims=(parmod["nt"], parmod["nx"], parmod["ny"]),
+        sampling=(parmod["dt"], parmod["dx"], parmod["dy"]),
+        dtype=np.complex128,
+    )
+
+    p3d_minus_inv, p3d_plus_inv = Deghosting(
+        p3d_sub,
+        parmod["nt"],
+        (parmod["nx"], parmod["ny"]),
+        parmod["dt"],
+        (parmod["dx"], parmod["dy"]),
+        vel_sep,
+        zrec,
+        kind=par["kind"],
+        win=np.ones_like(p3d_sub),
+        npad=(0, 0),
+        ntaper=(0, 0),
+        restriction=Rop,
+        sptransf=FOp.H,
+        solver=fista,
+        dtype=np.float32,
+        **dict(eps=1e-4, niter=50),
+    )
+
+    assert p3d_minus_inv.shape == (parmod["nt"], parmod["nx"], parmod["ny"])
+    assert p3d_plus_inv.shape == (parmod["nt"], parmod["nx"], parmod["ny"])
+    assert np.linalg.norm(p3d_minus_inv - p3d_minus) / np.linalg.norm(p3d_minus) < 4e-1