Merge pull request #593 from mrava87/feature-devitowav

feature: added updatesrc method to AcousticWave2D

Author: mrava87

examples/plot_twoway.py

@@ -0,0 +1,179 @@
+r"""
+Acoustic Wave Equation modelling
+================================
+
+This example shows how to perform acoustic wave equation modelling
+using the :class:`pylops.waveeqprocessing.AcousticWave2D` operator,
+which brings the power of finite-difference modelling via the Devito
+modelling engine to PyLops.
+"""
+import matplotlib.pyplot as plt
+import numpy as np
+from scipy.ndimage import gaussian_filter
+
+import pylops
+
+plt.close("all")
+np.random.seed(0)
+
+
+###############################################################################
+# To begin with, we will create a simple layered velocity model. We will also
+# define a background velocity model by smoothing the original velocity model
+# which will be responsible of the kinematic of the wavefield modelled via
+# Born modelling, and the perturbation velocity model which will lead to
+# scattering effects and therefore guide the dynamic of the modelled wavefield.
+
+# Velocity Model
+nx, nz = 61, 40
+dx, dz = 4, 4
+x, z = np.arange(nx) * dx, np.arange(nz) * dz
+vel = 1000 * np.ones((nx, nz))
+vel[:, 15:] = 1200
+vel[:, 35:] = 1600
+
+# Smooth velocity model
+v0 = gaussian_filter(vel, sigma=10)
+
+# Born perturbation from m - m0
+dv = vel ** (-2) - v0 ** (-2)
+
+# Receivers
+nr = 101
+rx = np.linspace(0, x[-1], nr)
+rz = 20 * np.ones(nr)
+recs = np.vstack((rx, rz))
+dr = recs[0, 1] - recs[0, 0]
+
+# Sources
+ns = 3
+sx = np.linspace(0, x[-1], ns)
+sz = 10 * np.ones(ns)
+sources = np.vstack((sx, sz))
+
+plt.figure(figsize=(10, 5))
+im = plt.imshow(vel.T, cmap="summer", extent=(x[0], x[-1], z[-1], z[0]))
+plt.scatter(recs[0], recs[1], marker="v", s=150, c="b", edgecolors="k")
+plt.scatter(sources[0], sources[1], marker="*", s=150, c="r", edgecolors="k")
+cb = plt.colorbar(im)
+cb.set_label("[m/s]")
+plt.axis("tight")
+plt.xlabel("x [m]"), plt.ylabel("z [m]")
+plt.title("Velocity")
+plt.xlim(x[0], x[-1])
+plt.tight_layout()
+
+plt.figure(figsize=(10, 5))
+im = plt.imshow(dv.T, cmap="seismic", extent=(x[0], x[-1], z[-1], z[0]))
+plt.scatter(recs[0], recs[1], marker="v", s=150, c="b", edgecolors="k")
+plt.scatter(sources[0], sources[1], marker="*", s=150, c="r", edgecolors="k")
+cb = plt.colorbar(im)
+cb.set_label("[m/s]")
+plt.axis("tight")
+plt.xlabel("x [m]"), plt.ylabel("z [m]")
+plt.title("Velocity perturbation")
+plt.xlim(x[0], x[-1])
+plt.tight_layout()
+
+###############################################################################
+# Let us now define the Born modelling operator
+
+Aop = pylops.waveeqprocessing.AcousticWave2D(
+    (nx, nz),
+    (0, 0),
+    (dx, dz),
+    v0,
+    sources[0],
+    sources[1],
+    recs[0],
+    recs[1],
+    0.0,
+    0.5 * 1e3,
+    "Ricker",
+    space_order=4,
+    nbl=100,
+    f0=15,
+    dtype="float32",
+)
+
+###############################################################################
+# And we use it to model our data
+
+dobs = Aop @ dv
+
+fig, axs = plt.subplots(1, 3, sharey=True, figsize=(10, 6))
+fig.suptitle("FD modelling with Ricker", y=0.99)
+
+for isrc in range(ns):
+    axs[isrc].imshow(
+        dobs[isrc].reshape(Aop.geometry.nrec, Aop.geometry.nt).T,
+        cmap="gray",
+        vmin=-1e-7,
+        vmax=1e-7,
+        extent=(
+            recs[0, 0],
+            recs[0, -1],
+            Aop.geometry.time_axis.time_values[-1] * 1e-3,
+            0,
+        ),
+    )
+    axs[isrc].axis("tight")
+    axs[isrc].set_xlabel("rec [m]")
+axs[0].set_ylabel("t [s]")
+fig.tight_layout()
+
+###############################################################################
+# Finally, we are going to show how despite the
+# :class:`pylops.waveeqprocessing.AcousticWave2D` operator allows a user to
+# specify a limited number of source wavelets (this is directly borrowed from
+# Devito), a simple modification can be applied to pass any user defined wavelet.
+# We are going to do that with a Ormsby wavelet
+
+# Extract Ricker wavelet
+wav = Aop.geometry.src.data[:, 0]
+wavc = np.argmax(wav)
+
+# Define Ormsby wavelet
+wavest = pylops.utils.wavelets.ormsby(
+    Aop.geometry.time_axis.time_values[:wavc] * 1e-3, f=[3, 20, 30, 45]
+)[0]
+
+# Update wavelet in operator and model new data
+Aop.updatesrc(wavest)
+
+dobs1 = Aop @ dv
+
+fig, axs = plt.subplots(1, 3, sharey=True, figsize=(10, 6))
+fig.suptitle("FD modelling with Ormsby", y=0.99)
+
+for isrc in range(ns):
+    axs[isrc].imshow(
+        dobs1[isrc].reshape(Aop.geometry.nrec, Aop.geometry.nt).T,
+        cmap="gray",
+        vmin=-1e-7,
+        vmax=1e-7,
+        extent=(
+            recs[0, 0],
+            recs[0, -1],
+            Aop.geometry.time_axis.time_values[-1] * 1e-3,
+            0,
+        ),
+    )
+    axs[isrc].axis("tight")
+    axs[isrc].set_xlabel("rec [m]")
+axs[0].set_ylabel("t [s]")
+fig.tight_layout()
+
+fig, axs = plt.subplots(1, 2, figsize=(10, 3))
+axs[0].plot(wav[: 2 * wavc], "k")
+axs[0].plot(wavest, "r")
+axs[1].plot(
+    dobs[isrc].reshape(Aop.geometry.nrec, Aop.geometry.nt)[nr // 2], "k", label="Ricker"
+)
+axs[1].plot(
+    dobs1[isrc].reshape(Aop.geometry.nrec, Aop.geometry.nt)[nr // 2],
+    "r",
+    label="Ormsby",
+)
+axs[1].legend()
+fig.tight_layout()

pylops/waveeqprocessing/twoway.py

@@ -14,6 +14,49 @@
 if devito_message is None:
     from examples.seismic import AcquisitionGeometry, Model
     from examples.seismic.acoustic import AcousticWaveSolver
+    from examples.seismic.utils import PointSource
+
+
+class _CustomSource(PointSource):
+    """Custom source
+
+    This class creates a Devito symbolic object that encapsulates a set of
+    sources with a user defined source signal wavelet ``wav``
+
+    Parameters
+    ----------
+    name : :obj:`str`
+        Name for the resulting symbol.
+    grid : :obj:`devito.types.grid.Grid`
+        The computational domain.
+    time_range : :obj:`examples.seismic.source.TimeAxis`
+        TimeAxis(start, step, num) object.
+    wav : :obj:`numpy.ndarray`
+        Wavelet of size
+
+    """
+
+    __rkwargs__ = PointSource.__rkwargs__ + ["wav"]
+
+    @classmethod
+    def __args_setup__(cls, *args, **kwargs):
+        kwargs.setdefault("npoint", 1)
+
+        return super().__args_setup__(*args, **kwargs)
+
+    def __init_finalize__(self, *args, **kwargs):
+        super().__init_finalize__(*args, **kwargs)
+
+        self.wav = kwargs.get("wav")
+
+        if not self.alias:
+            for p in range(kwargs["npoint"]):
+                self.data[:, p] = self.wavelet
+
+    @property
+    def wavelet(self):
+        """Return user-provided wavelet"""
+        return self.wav
 
 
 class AcousticWave2D(LinearOperator):
@@ -38,9 +81,9 @@ class AcousticWave2D(LinearOperator):
     rec_z : :obj:`numpy.ndarray` or :obj:`float`
         Receiver z-coordinates in m
     t0 : :obj:`float`
-        Initial time
+        Initial time in ms
     tn : :obj:`int`
-        Number of time samples
+        Final time in ms
     src_type : :obj:`str`
         Source type
     space_order : :obj:`int`, optional
@@ -79,7 +122,7 @@ def __init__(
         rec_x: NDArray,
         rec_z: NDArray,
         t0: float,
-        tn: int,
+        tn: float,
         src_type: str = "Ricker",
         space_order: int = 6,
         nbl: int = 20,
@@ -155,7 +198,7 @@ def _create_geometry(
         rec_x: NDArray,
         rec_z: NDArray,
         t0: float,
-        tn: int,
+        tn: float,
         src_type: str,
         f0: float = 20.0,
     ) -> None:
@@ -174,7 +217,7 @@ def _create_geometry(
         t0 : :obj:`float`
             Initial time
         tn : :obj:`int`
-            Number of time samples
+            Final time in ms
         src_type : :obj:`str`
             Source type
         f0 : :obj:`float`, optional
@@ -201,6 +244,28 @@ def _create_geometry(
             f0=None if f0 is None else f0 * 1e-3,
         )
 
+    def updatesrc(self, wav):
+        """Update source wavelet
+
+        This routines is used to allow users to pass a custom source
+        wavelet to replace the source wavelet generated when the
+        object is initialized
+
+        Parameters
+        ----------
+        wav : :obj:`numpy.ndarray`
+            Wavelet
+
+        """
+        wav_padded = np.pad(wav, (0, self.geometry.nt - len(wav)))
+
+        self.wav = _CustomSource(
+            name="src",
+            grid=self.model.grid,
+            wav=wav_padded,
+            time_range=self.geometry.time_axis,
+        )
+
     def _srcillumination_oneshot(self, isrc: int) -> Tuple[NDArray, NDArray]:
         """Source wavefield and illumination for one shot
 
@@ -229,8 +294,15 @@ def _srcillumination_oneshot(self, isrc: int) -> Tuple[NDArray, NDArray]:
         )
         solver = AcousticWaveSolver(self.model, geometry, space_order=self.space_order)
 
+        # assign source location to source object with custom wavelet
+        if hasattr(self, "wav"):
+            self.wav.coordinates.data[0, :] = self.geometry.src_positions[isrc, :]
+
         # source wavefield
-        u0 = solver.forward(save=True)[1]
+        u0 = solver.forward(
+            save=True, src=None if not hasattr(self, "wav") else self.wav
+        )[1]
+
         # source illumination
         src_ill = self._crop_model((u0.data**2).sum(axis=0), self.model.nbl)
         return u0, src_ill
@@ -286,9 +358,15 @@ def _born_oneshot(self, isrc: int, dm: NDArray) -> NDArray:
         dmext = np.zeros(self.model.grid.shape, dtype=np.float32)
         dmext[self.model.nbl : -self.model.nbl, self.model.nbl : -self.model.nbl] = dm
 
+        # assign source location to source object with custom wavelet
+        if hasattr(self, "wav"):
+            self.wav.coordinates.data[0, :] = self.geometry.src_positions[isrc, :]
+
         # solve
         solver = AcousticWaveSolver(self.model, geometry, space_order=self.space_order)
-        d = solver.jacobian(dmext)[0]
+        d = solver.jacobian(dmext, src=None if not hasattr(self, "wav") else self.wav)[
+            0
+        ]
         d = d.resample(geometry.dt).data[:][: geometry.nt].T
         return d
 
@@ -347,11 +425,18 @@ def _bornadj_oneshot(self, isrc, dobs):
 
         solver = AcousticWaveSolver(self.model, geometry, space_order=self.space_order)
 
+        # assign source location to source object with custom wavelet
+        if hasattr(self, "wav"):
+            self.wav.coordinates.data[0, :] = self.geometry.src_positions[isrc, :]
+
         # source wavefield
         if hasattr(self, "src_wavefield"):
             u0 = self.src_wavefield[isrc]
         else:
-            u0 = solver.forward(save=True)[1]
+            u0 = solver.forward(
+                save=True, src=None if not hasattr(self, "wav") else self.wav
+            )[1]
+
         # adjoint modelling (reverse wavefield plus imaging condition)
         model = solver.jacobian_adjoint(
             rec=recs, u=u0, checkpointing=self.checkpointing