r"""
Normal Moveout (NMO) Correction
===============================
This example shows how to create your own operator for performing
normal moveout (NMO) correction to a seismic record.
We will perform classic NMO using an operator created from scratch,
as well as using the :py:class:`pylops.Spread` operator.
"""
from math import floor
from time import time
import matplotlib.pyplot as plt
import numpy as np
from mpl_toolkits.axes_grid1 import ImageGrid, make_axes_locatable
from numba import jit, prange
from scipy.interpolate import griddata
from scipy.ndimage import gaussian_filter
from pylops import LinearOperator, Spread
from pylops.utils import dottest
from pylops.utils.decorators import reshaped
from pylops.utils.seismicevents import hyperbolic2d, makeaxis
from pylops.utils.wavelets import ricker
def create_colorbar(im, ax):
divider = make_axes_locatable(ax)
cax = divider.append_axes("right", size="5%", pad=0.1)
cb = fig.colorbar(im, cax=cax, orientation="vertical")
return cax, cb
###############################################################################
# Given a common-shot or common-midpoint (CMP) record, the objective of NMO
# correction is to "flatten" events, that is, align events at later offsets
# to that of the zero offset. NMO has long been a staple of seismic data
# processing, used even today for initial velocity analysis and QC purposes.
# In addition, it can be the domain of choice for many useful processing
# steps, such as angle muting.
#
# To get started, let us create a 2D seismic dataset containing some hyperbolic
# events representing reflections from flat reflectors.
# Events are created with a true RMS velocity, which we will be using as if we
# picked them from, for example, a semblance panel.
par = dict(ox=0, dx=40, nx=80, ot=0, dt=0.004, nt=520)
t, _, x, _ = makeaxis(par)
t0s_true = np.array([0.5, 1.22, 1.65])
vrms_true = np.array([2000.0, 2400.0, 2500.0])
amps = np.array([1, 0.2, 0.5])
freq = 10 # Hz
wav, *_ = ricker(t[:41], f0=freq)
_, data = hyperbolic2d(x, t, t0s_true, vrms_true, amp=amps, wav=wav)
###############################################################################
# NMO correction plot
pclip = 0.5
dmax = np.max(np.abs(data))
opts = dict(
cmap="gray_r",
extent=[x[0], x[-1], t[-1], t[0]],
aspect="auto",
vmin=-pclip * dmax,
vmax=pclip * dmax,
)
# Offset-dependent traveltime of the first hyperbolic event
t_nmo_ev1 = np.sqrt(t0s_true[0] ** 2 + (x / vrms_true[0]) ** 2)
fig, ax = plt.subplots(figsize=(4, 5))
vmax = np.max(np.abs(data))
im = ax.imshow(data.T, **opts)
ax.plot(x, t_nmo_ev1, "C1--", label="Hyperbolic moveout")
ax.plot(x, t0s_true[0] + x * 0, "C1", label="NMO-corrected")
idx = 3 * par["nx"] // 4
ax.annotate(
"",
xy=(x[idx], t0s_true[0]),
xycoords="data",
xytext=(x[idx], t_nmo_ev1[idx]),
textcoords="data",
fontsize=7,
arrowprops=dict(edgecolor="w", arrowstyle="->", shrinkA=10),
)
ax.set(title="Data", xlabel="Offset [m]", ylabel="Time [s]")
cax, _ = create_colorbar(im, ax)
cax.set_ylabel("Amplitude")
ax.legend()
fig.tight_layout()
################################################################################
# NMO correction consists of applying an offset- and time-dependent shift to
# each sample of the trace in such a way that all events corresponding to the
# same reflection will be located at the same time intercept after correction.
#
# An arbitrary hyperbolic event at position :math:`(t, h)` is linked to its
# zero-offset traveltime :math:`t_0` by the following equation
#
# .. math::
# t(x) = \sqrt{t_0^2 + \frac{h^2}{v_\text{rms}^2(t_0)}}
#
# Our strategy in applying the correction is to loop over our time axis
# (which we will associate to :math:`t_0`) and respective RMS velocities
# and, for each offset, move the sample at :math:`t(x)` to location
# :math:`t_0(x) \equiv t_0`. In the figure above, we are considering a
# single :math:`t_0 = 0.5\mathrm{s}` which would have values along the dotted curve
# (i.e., :math:`t(x)`) moved to :math:`t_0` for every offset.
#
# Notice that we need NMO velocities for each sample of our time axis.
# In this example, we actually only have 3 samples, when we need ``nt`` samples.
# In practice, we would have many more samples, but probably not one for each
# ``nt``. To resolve this issue, we will interpolate these 3 samples to all samples
# of our time axis (or, more accurately, their slownesses to preserve traveltimes).
def interpolate_vrms(t0_picks, vrms_picks, taxis, smooth=None):
assert len(t0_picks) == len(vrms_picks)
# Sampled points in time axis
points = np.zeros((len(t0_picks) + 2,))
points[0] = taxis[0]
points[-1] = taxis[-1]
points[1:-1] = t0_picks
# Sampled values of slowness (in s/km)
values = np.zeros((len(vrms_picks) + 2,))
values[0] = 1000.0 / vrms_picks[0] # Use first slowness before t0_picks[0]
values[-1] = 1000.0 / vrms_picks[-1] # Use the last slowness after t0_picks[-1]
values[1:-1] = 1000.0 / vrms_picks
slowness = griddata(points, values, taxis, method="linear")
if smooth is not None:
slowness = gaussian_filter(slowness, sigma=smooth)
return 1000.0 / slowness
vel_t = interpolate_vrms(t0s_true, vrms_true, t, smooth=11)
###############################################################################
# Plot interpolated RMS velocities which will be used for NMO
fig, ax = plt.subplots(figsize=(4, 5))
ax.plot(vel_t, t, "k", lw=3, label="Interpolated", zorder=-1)
ax.plot(vrms_true, t0s_true, "C1o", markersize=10, label="Picks")
ax.invert_yaxis()
ax.set(xlabel="RMS Velocity [m/s]", ylabel="Time [s]", ylim=[t[-1], t[0]])
ax.legend()
fig.tight_layout()
###############################################################################
# NMO from scratch
# ----------------
# We are very close to building our NMO correction, we just need to take care of
# one final issue. When moving the sample from :math:`t(x)` to :math:`t_0`, we
# know that, by definition, :math:`t_0` is always on our time axis grid. In contrast,
# :math:`t(x)` may not fall exactly on a multiple of ``dt`` (our time axis
# sampling). Suppose its nearest sample smaller than itself (floor) is ``i``.
# Instead of moving only sample `i`, we will be moving samples both samples
# ``i`` and ``i+1`` with an appropriate weight to account for how far
# :math:`t(x)` is from ``i*dt`` and ``(i+1)*dt``.
@jit(nopython=True, fastmath=True, nogil=True, parallel=True)
def nmo_forward(data, taxis, haxis, vels_rms):
dt = taxis[1] - taxis[0]
ot = taxis[0]
nt = len(taxis)
nh = len(haxis)
dnmo = np.zeros_like(data)
# Parallel outer loop on slow axis
for ih in prange(nh):
h = haxis[ih]
for it0, (t0, vrms) in enumerate(zip(taxis, vels_rms)):
# Compute NMO traveltime
tx = np.sqrt(t0**2 + (h / vrms) ** 2)
it_frac = (tx - ot) / dt # Fractional index
it_floor = floor(it_frac)
it_ceil = it_floor + 1
w = it_frac - it_floor
if 0 <= it_floor and it_ceil < nt: # it_floor and it_ceil must be valid
# Linear interpolation
dnmo[ih, it0] += (1 - w) * data[ih, it_floor] + w * data[ih, it_ceil]
return dnmo
dnmo = nmo_forward(data, t, x, vel_t) # Compile and run
# Time execution
start = time()
nmo_forward(data, t, x, vel_t)
end = time()
print(f"Ran in {1e6*(end-start):.0f} μs")
###############################################################################
# Plot Data and NMO-corrected data
fig = plt.figure(figsize=(6.5, 5))
grid = ImageGrid(
fig,
111,
nrows_ncols=(1, 2),
axes_pad=0.15,
cbar_location="right",
cbar_mode="single",
cbar_size="7%",
cbar_pad=0.15,
aspect=False,
share_all=True,
)
im = grid[0].imshow(data.T, **opts)
grid[0].set(title="Data", xlabel="Offset [m]", ylabel="Time [s]")
grid[0].cax.colorbar(im)
grid[0].cax.set_ylabel("Amplitude")
grid[1].imshow(dnmo.T, **opts)
grid[1].set(title="NMO-corrected Data", xlabel="Offset [m]")
plt.show()
###############################################################################
# Now that we know how to compute the forward, we'll compute the adjoint pass.
# With these two functions, we can create a ``LinearOperator`` and ensure that
# it passes the dot-test.
@jit(nopython=True, fastmath=True, nogil=True, parallel=True)
def nmo_adjoint(dnmo, taxis, haxis, vels_rms):
dt = taxis[1] - taxis[0]
ot = taxis[0]
nt = len(taxis)
nh = len(haxis)
data = np.zeros_like(dnmo)
# Parallel outer loop on slow axis; use range if Numba is not installed
for ih in prange(nh):
h = haxis[ih]
for it0, (t0, vrms) in enumerate(zip(taxis, vels_rms)):
# Compute NMO traveltime
tx = np.sqrt(t0**2 + (h / vrms) ** 2)
it_frac = (tx - ot) / dt # Fractional index
it_floor = floor(it_frac)
it_ceil = it_floor + 1
w = it_frac - it_floor
if 0 <= it_floor and it_ceil < nt:
# Linear interpolation
# In the adjoint, we must spread the same it0 to both it_floor and
# it_ceil, since in the forward pass, both of these samples were
# pushed onto it0
data[ih, it_floor] += (1 - w) * dnmo[ih, it0]
data[ih, it_ceil] += w * dnmo[ih, it0]
return data
###############################################################################
# Finally, we can create our linear operator. To exemplify the
# class-based interface we will subclass :py:class:`pylops.LinearOperator` and
# implement the required methods: ``_matvec`` which will compute the forward and
# ``_rmatvec`` which will compute the adjoint. Note the use of the ``reshaped``
# decorator which allows us to pass ``x`` directly into our auxiliary function
# without having to do ``x.reshape(self.dims)`` and to output without having to
# call ``ravel()``.
class NMO(LinearOperator):
def __init__(self, taxis, haxis, vels_rms, dtype=None):
self.taxis = taxis
self.haxis = haxis
self.vels_rms = vels_rms
dims = (len(haxis), len(taxis))
if dtype is None:
dtype = np.result_type(taxis.dtype, haxis.dtype, vels_rms.dtype)
super().__init__(dims=dims, dimsd=dims, dtype=dtype)
@reshaped
def _matvec(self, x):
return nmo_forward(x, self.taxis, self.haxis, self.vels_rms)
@reshaped
def _rmatvec(self, y):
return nmo_adjoint(y, self.taxis, self.haxis, self.vels_rms)
###############################################################################
# With our new ``NMO`` linear operator, we can instantiate it with our current
# example and ensure that it passes the dot test which proves that our forward
# and adjoint transforms truly are adjoints of each other.
NMOOp = NMO(t, x, vel_t)
dottest(NMOOp, rtol=1e-4, verb=True)
###############################################################################
# NMO using :py:class:`pylops.Spread`
# -----------------------------------
# We learned how to implement an NMO correction and its adjoint from scratch.
# The adjoint has an interesting pattern, where energy taken from one domain
# is "spread" along a previously-defined parametric curve (the NMO hyperbola
# in this case). This pattern is very common in many algorithms, including
# Radon transform, Kirchhoff migration (also known as Total Focusing Method in
# ultrasonics) and many others.
#
# For these classes of operators, PyLops offers a :py:class:`pylops.Spread`
# constructor, which we will leverage to implement a version of the NMO correction.
# The :py:class:`pylops.Spread` operator will take a value in the "input" domain,
# and spread it along a parametric curve, defined in the "output" domain.
# In our case, the spreading operation is the *adjoint* of the NMO, so our
# "input" domain is the NMO domain, and the "output" domain is the original
# data domain.
#
# In order to use :py:class:`pylops.Spread`, we need to define the
# parametric curves. This can be done through the use of a table with shape
# :math:`(n_{x_i}, n_{t}, n_{x_o})`, where :math:`n_{x_i}` and :math:`n_{t}`
# represent the 2d dimensions of the "input" domain (NMO domain) and :math:`n_{x_o}`
# and :math:`n_{t}` the 2d dimensions of the "output" domain. In our NMO case,
# :math:`n_{x_i} = n_{x_o} = n_h` represents the number of offsets.
# Following the documentation of :py:class:`pylops.Spread`, the table will be
# used in the following manner:
#
# ``d_out[ix_o, table[ix_i, it, ix_o]] += d_in[ix_i, it]``
#
# In our case, ``ix_o = ix_i = ih``, and comparing with our NMO adjoint, ``it``
# refers to :math:`t_0` while ``table[ix, it, ix]`` should then provide the
# appropriate index for :math:`t(x)`. In our implementation we will also be
# constructing a second table containing the weights to be used for linear
# interpolation.
def create_tables(taxis, haxis, vels_rms):
dt = taxis[1] - taxis[0]
ot = taxis[0]
nt = len(taxis)
nh = len(haxis)
# NaN values will be not be spread.
# Using np.zeros has the same result but much slower.
table = np.full((nh, nt, nh), fill_value=np.nan)
dtable = np.full((nh, nt, nh), fill_value=np.nan)
for ih, h in enumerate(haxis):
for it0, (t0, vrms) in enumerate(zip(taxis, vels_rms)):
# Compute NMO traveltime
tx = np.sqrt(t0**2 + (h / vrms) ** 2)
it_frac = (tx - ot) / dt
it_floor = floor(it_frac)
w = it_frac - it_floor
# Both it_floor and it_floor + 1 must be valid indices for taxis
# when using two tables (interpolation).
if 0 <= it_floor and it_floor + 1 < nt:
table[ih, it0, ih] = it_floor
dtable[ih, it0, ih] = w
return table, dtable
nmo_table, nmo_dtable = create_tables(t, x, vel_t)
###############################################################################
SpreadNMO = Spread(
dims=data.shape, # "Input" shape: NMO-ed data shape
dimsd=data.shape, # "Output" shape: original data shape
table=nmo_table, # Table of time indices
dtable=nmo_dtable, # Table of weights for linear interpolation
engine="numba", # numba or numpy
).H # To perform NMO *correction*, we need the adjoint
dottest(SpreadNMO, rtol=1e-4)
###############################################################################
# We see it passes the dot test, but are the results right? Let's find out.
dnmo_spr = SpreadNMO @ data
start = time()
SpreadNMO @ data
end = time()
print(f"Ran in {1e6*(end-start):.0f} μs")
###############################################################################
# Note that since v2.0, we do not need to pass a flattened array. Consequently,
# the output will not be flattened, but will have ``SpreadNMO.dimsd`` as shape.
# Plot Data and NMO-corrected data
fig = plt.figure(figsize=(6.5, 5))
grid = ImageGrid(
fig,
111,
nrows_ncols=(1, 2),
axes_pad=0.15,
cbar_location="right",
cbar_mode="single",
cbar_size="7%",
cbar_pad=0.15,
aspect=False,
share_all=True,
)
im = grid[0].imshow(data.T, **opts)
grid[0].set(title="Data", xlabel="Offset [m]", ylabel="Time [s]")
grid[0].cax.colorbar(im)
grid[0].cax.set_ylabel("Amplitude")
grid[1].imshow(dnmo_spr.T, **opts)
grid[1].set(title="NMO correction using Spread", xlabel="Offset [m]")
plt.show()
###############################################################################
# Not as blazing fast as out original implementation, but pretty good (try the
# "numpy" backend for comparison!). In fact, using the ``Spread`` operator for
# NMO will always have a speed disadvantage. While iterating over the table, it must
# loop over the offsets twice: one for the "input" offsets and one for the "output"
# offsets. We know they are the same for NMO, but since ``Spread`` is a generic
# operator, it does not know that. So right off the bat we can expect an 80x
# slowdown (nh = 80). We diminished this cost to about 30x by setting values where
# ``ix_i != ix_o`` to NaN, but nothing beats the custom implementation. Despite this,
# we can still produce the same result to numerical accuracy:
np.allclose(dnmo, dnmo_spr)