\section{The code for NMO}

Now I'll show how to implement the above algorithm in Python using the NumPy and SciPy libraries \citep{van_der_Walt_2011}.
We'll split the algorithm into three functions.
This is very important when programming any moderately complex code because it allows us to test each part of our code independently.
It also reduces the amount of code we have to search through to find that bug that is messing up our results.
Modular code is easier to understand and to reuse.

The first function I'll define performs the NMO correction on a given CMP gather.
We'll assume that the CMP gather is a 2D array of amplitudes and that the
velocities are a 1D array with $v_\mathrm{NMO}$ for each time sample.

\begin{lstlisting}[language=python]
    import numpy as np

    def nmo_correction(cmp, dt, offsets, velocities):
        nmo = np.zeros_like(cmp)
        nsamples = cmp.shape[0]
        times = np.arange(0, nsamples*dt, dt)
        for i, t0 in enumerate(times):
            for j, x in enumerate(offsets):
                t = reflection_time(t0, x, velocities[i])
                amplitude = sample_trace(cmp[:, j], t, dt)
                if amplitude is not None:
                    nmo[i, j] = amplitude
        return nmo
\end{lstlisting}

This function is essentially the algorithm above translated to Python with some
of the details pushed into the \texttt{reflection\_time} and
\texttt{sample\_trace} functions, which we will define below.

First, the function that calculates the reflection travel-time:

\begin{lstlisting}[language=python]
    def reflection_time(t0, x, vnmo):
        t = np.sqrt(t0**2 + x**2/vnmo**2)
        return t
\end{lstlisting}

For the \texttt{sample\_trace} function, we'll use cubic splines from the
\texttt{scipy.interpolate} package.
For more information on interpolation with scipy, see the tutorial by
\citet{Hall_2016}.

\begin{lstlisting}[language=python]
    from scipy.interpolate import CubicSpline

    def sample_trace(trace, time, dt):
        before = int(np.floor(time/dt))
        N = trace.size
        samples = np.arange(before - 1, before + 3)
        if any(samples < 0) or any(samples >= N):
            amplitude = None
        else:
            times = dt*samples
            amps = trace[samples]
            interpolator = CubicSpline(times, amps)
            amplitude = interpolator(time)
        return amplitude
\end{lstlisting}

The Jupyter notebook contains the full code for these functions, including
documentation through Python documentation strings or ``docstrings'' and code
that tests that the functions work as expected.
Also included is an application of our \texttt{nmo\_correction} function to a
synthetic CMP (Figure 2).
