Wiggle plot for seismic

wiggle.py

import numpy as np
from scipy.interpolate import interp1d
import matplotlib.pyplot as plt


def wiggle_2d(data,
              time,
              ax=None,
              skip=1,
              perc=99.0,
              gain=1.0,
              oversampling=100,
              rgb=(0, 0, 0),
              alpha=1.0,
              lw=0.5,
              ):
        """
        Plots wiggle traces of seismic data.
        
        Args:
            data (ndarray): 2D seismic array.
            time (ndarray): 1D seismic time array.
            ax (Axes): matplotlib Axes object (optional).
            skip (int): Skip=1, every trace, skip=2, every second trace, etc.
            perc (float): Percentile to scale to, default 99%.
            gain (float): Gain value.
            oversampling (int): The approximate number of new samples per sample.
                Higher numbers result in smoother interpolated wavelets.
            rgb (tuple): 3-tuple of RGB for the trace colour.
            alpha (float): Opacity. Default 1.0.
            lw (float): Lineweight. Default 0.5.

        Returns:
            Axes: A matplotlib Axes object.
        """
        if ax is None:
            fig, ax = plt.subplots(figsize=(2, 6))

        ntraces, nt = data.shape
        rgba = list(rgb) + [alpha]
        sc = np.percentile(data, perc)  # Normalization factor
        wigdata = data[::skip, :]
        xpos = np.arange(ntraces)[::skip]

        for x, trace in zip(xpos, wigdata):
            # Compute high resolution trace.
            amp = gain * trace / sc + x
            t = 1000 * time
            hypertime = np.linspace(t[0], t[-1], (oversampling * t.size - 1) + 1)
            interp = interp1d(t, amp, kind='cubic')
            hyperamp = interp(hypertime)

            # Plot the line, then the fill.
            ax.plot(hyperamp, hypertime, 'k', lw=lw)
            ax.fill_betweenx(hypertime, hyperamp, x,
                             where=hyperamp > x,
                             facecolor=rgba,
                             interpolate=True,
                             lw=0,
                             )

        ax.invert_yaxis()
        plt.tight_layout()
        return ax

Call this function like:

from wiggle import wiggle_2d
t = np.arange(0, 1.004, 0.004)
wiggle_2d(vol[100], time=t, gain=2)

Where vol is a 3D array, with inlines, xlines, and time samples in dimensions 1, 2, and 3.