Add module for dispersion image calculation (#275)

* Add module for dispersion calculation, by Ariel Lellouch

* Updated documentation

* Fixed docstrings for testing module

* Speedup, style, and bugfixes

* remove deleted parameter from docstring, add exception tests

* Added dispersion example, and improved functionality

* Fixed dosctring

* Fixed dosctring

* Fixed testing

---------

Co-authored-by: derrick chambers <[email protected]>

dascore/core/patch.py

@@ -245,6 +245,7 @@ def iresample(self, *args, **kwargs):
     integrate = transform.integrate
     spectrogram = transform.spectrogram
     velocity_to_strain_rate = transform.velocity_to_strain_rate
+    dispersion_phase_shift = transform.dispersion_phase_shift
 
     # --- Method Namespaces
     # Note: these can't be cached_property (from functools) or references

dascore/data_registry.txt

@@ -13,3 +13,4 @@ prodml_2.1.h5 47875ae1fa17d99849cdb595e9fa6853d162605c1556be90225f37aaa123dded h
 h5_simple_1.h5 52803b26a5738da541cc9b32b867434cad0a845686c47dd93551b8eb431f8bc0 https://github.com/DASDAE/test_data/raw/master/das/h5_simple_1.h5
 h5_simple_2.h5 8a70b873c5c2c172871ecd63760d24fa2c305c015e2ca1c84018c6864d2fb304 https://github.com/dasdae/test_data/raw/master/das/h5_simple_2.h5
 conoco_segy_1.sgy 3944297d7c27dd265b40d56d4797f1a14caa5c2bed9f0af020b0f6ea193d4dfd https://github.com/dasdae/test_data/raw/master/das/conoco_segy_1.sgy
+dispersion_event.h5 598c8baa2a5610c930e1c003f2ba02da13f8d8686e3ccf2a034e94bfc5e1990c https://github.com/dasdae/test_data/raw/master/das/dispersion_event.h5

dascore/examples.py

@@ -247,6 +247,13 @@ def _diverse_spool():
     return dc.spool([y for x in all_spools for y in x])
 
 
+@register_func(EXAMPLE_PATCHES, key="dispersion_event")
+def _dispersion_event():
+    """Returns an example of a synthetic shot record that exhibits dispersion."""
+    path = fetch("dispersion_event.h5")
+    return dc.spool(path)[0]
+
+
 def spool_to_directory(spool, path=None, file_format="DASDAE", extention="hdf5"):
     """
     Write out each patch in a spool to a directory.

dascore/transform/__init__.py

@@ -11,3 +11,4 @@
 from .integrate import integrate
 from .spectro import spectrogram
 from .strain import velocity_to_strain_rate
+from .dispersion import dispersion_phase_shift

dascore/transform/dispersion.py

@@ -0,0 +1,150 @@
+"""Dispersion computation using the phase-shift (Park et al., 1999) method."""
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+import numpy as np
+import numpy.fft as nft
+
+from dascore.constants import PatchType
+from dascore.exceptions import ParameterError
+from dascore.utils.patch import patch_function
+
+
+@patch_function(required_dims=("time", "distance"))
+def dispersion_phase_shift(
+    patch: PatchType,
+    phase_velocities: Sequence[float],
+    approx_resolution: None | float = None,
+    approx_freq: [None, None] | float = None,
+) -> PatchType:
+    """
+    Compute dispersion images using the phase-shift method.
+
+    Parameters
+    ----------
+    patch
+        Patch to transform. Has to have dimensions of time and distance.
+    phase_velocities
+        NumPY array of positive velocities, monotonically increasing, for
+        which the dispersion will be computed.
+    approx_resolution
+        Approximated frequency (Hz) resolution for the output. If left empty,
+        the frequency resolution is dictated by the number of samples.
+    approx_min_freq
+        Minimum frequency to compute dispersion for. If left empty, 0 Hz
+    approx_max_freq
+        Maximum frequency to compute dispersion for. If left empty,
+        Nyquist frequency will be used.
+
+    Notes
+    -----
+    - See also @park1998imaging
+
+    - Inspired by https://geophydog.cool/post/masw_phase_shift/.
+
+    - Dims/Units of the output are forced to be 'frequency' ('Hz')
+      and 'velocity' ('m/s').
+
+    - The patch's distance coordinates are assumed to be ordered by
+    distance from the source, and not "fiber distance". In other
+    words, data are effectively mapped along a 2-D line.
+
+    Example
+    --------
+    import dascore as dc
+    import numpy as np
+
+    patch = (
+        dc.get_example_patch('dispersion_event')
+    )
+
+    disp_patch = patch.dispersion_phase_shift(np.arange(100,1500,1),
+                approx_resolution=0.1,approx_freq=[5,70])
+    ax = disp_patch.viz.waterfall(show=False,cmap=None)
+    ax.set_xlim(5, 70)
+    ax.set_ylim(1500, 100)
+    disp_patch.viz.waterfall(show=True, ax=ax)
+    """
+    patch_cop = patch.convert_units(distance="m").transpose("distance", "time")
+    dist = patch_cop.coords.get_array("distance")
+    time = patch_cop.coords.get_array("time")
+
+    dt = (time[1] - time[0]) / np.timedelta64(1, "s")
+
+    if not np.all(np.diff(phase_velocities) > 0):
+        raise ParameterError(
+            "Velocities for dispersion must be monotonically increasing"
+        )
+
+    if np.amin(phase_velocities) <= 0:
+        raise ParameterError("Velocities must be positive.")
+
+    if approx_resolution is not None and approx_resolution <= 0:
+        raise ParameterError("Frequency resolution has to be positive")
+
+    if not approx_freq:
+        approx_min_freq = 0
+        approx_max_freq = 0.5 / dt
+    else:
+        approx_min_freq = approx_freq[0]
+        approx_max_freq = approx_freq[1]
+        if approx_min_freq <= 0 or approx_max_freq <= 0:
+            msg = "Minimal and maximal frequencies have to be positive"
+            raise ParameterError(msg)
+
+        if approx_min_freq >= approx_max_freq:
+            msg = "Maximal frequency needs to be larger than minimal frequency"
+            raise ParameterError(msg)
+
+        if approx_min_freq >= 0.5 / dt or approx_max_freq >= 0.5 / dt:
+            msg = "Frequency range cannot exceed Nyquist"
+            raise ParameterError(msg)
+
+    nchan = dist.size
+    nt = time.size
+    assert (nchan, nt) == patch_cop.data.shape
+
+    fs = 1 / dt
+    if approx_resolution is not None:
+        approxnf = int(nt * (fs / (nt)) / approx_resolution)
+        f = np.arange(approxnf) * fs / (approxnf - 1)
+    else:
+        f = np.arange(nt) * fs / (nt - 1)
+
+    nf = np.size(f)
+
+    nv = np.size(phase_velocities)
+    w = 2 * np.pi * f
+    fft_d = np.zeros((nchan, nf), dtype=complex)
+    for i in range(nchan):
+        fft_d[i] = nft.fft(patch_cop.data[i, :], n=nf)
+
+    fft_d = np.divide(
+        fft_d, abs(fft_d), out=np.zeros_like(fft_d), where=abs(fft_d) != 0
+    )
+
+    fft_d[np.isnan(fft_d)] = 0
+
+    first_live_f = np.argmax(w >= 2 * np.pi * approx_min_freq)
+    last_live_f = np.argmax(w >= 2 * np.pi * approx_max_freq)
+    w = w[first_live_f:last_live_f]
+    fft_d = fft_d[:, first_live_f:last_live_f]
+    nlivef = last_live_f - first_live_f
+
+    if nlivef < 1:
+        msg = "Combination of frequency resolution and range is not an array"
+        raise ParameterError(msg)
+
+    fc = np.zeros(shape=(nv, nlivef))
+    preamb = 1j * np.outer(dist, w)
+    for ci in range(nv):
+        fc[ci, :] = abs(sum(np.exp(preamb / phase_velocities[ci]) * fft_d))
+
+    attrs = patch.attrs.update(category="dispersion")
+    coords = dict(velocity=phase_velocities, frequency=w / (2 * np.pi))
+
+    disp_patch = patch.new(
+        data=fc / nchan, coords=coords, attrs=attrs, dims=["velocity", "frequency"]
+    )
+    return disp_patch.set_units(velocity="m/s", frequency="Hz")

docs/contributors.md

@@ -19,6 +19,8 @@ A Huge thanks to all of our contributors:
 
 [Aaron Girard](https://github.com/aaronjgirard)
 
+[Ariel Lellouch](https://github.com/ariellellouch)
+
 [Manuel M. Mendoza](https:/github.com/SeisMatt)
 
 You can find more contributor information

docs/references.bib

@@ -42,3 +42,12 @@ @article{stanvek2022fracture
   volume={10},
   year={2022}
 }
+
+@incollection{park1998imaging,
+  title={Imaging dispersion curves of surface waves on multi-channel record},
+  author={Park, Choon Byong and Miller, Richard D and Xia, Jianghai},
+  booktitle={SEG technical program expanded abstracts 1998},
+  pages={1377--1380},
+  year={1998},
+  publisher={Society of Exploration Geophysicists}
+}

tests/test_transform/test_dispersion.py

@@ -0,0 +1,119 @@
+"""Tests for Dispersion transforms."""
+import numpy as np
+import pytest
+
+import dascore as dc
+from dascore import get_example_patch
+from dascore.exceptions import ParameterError
+from dascore.transform import dispersion_phase_shift
+from dascore.utils.misc import suppress_warnings
+
+
+class TestDispersion:
+    """Tests for the dispersion module."""
+
+    @pytest.fixture(scope="class")
+    def dispersion_patch(self, random_patch):
+        """Return the random patched transformed to frequency-velocity."""
+        test_vels = np.linspace(1500, 5000, 351)
+        with suppress_warnings(DeprecationWarning):
+            out = dispersion_phase_shift(random_patch, test_vels, approx_resolution=2.0)
+        return out
+
+    def test_dispersion(self, dispersion_patch):
+        """Check consistency of test_dispersion module."""
+        # assert velocity dimension
+        assert "velocity" in dispersion_patch.dims
+        # assert frequency dimension
+        assert "frequency" in dispersion_patch.dims
+        vels = dispersion_patch.coords.get_array("velocity")
+        freqs = dispersion_patch.coords.get_array("frequency")
+        assert np.array_equal(vels, np.linspace(1500, 5000, 351))
+        # Check that the velocity output is correct
+        assert freqs[1] - freqs[0] > 1.9 and freqs[1] - freqs[0] < 2.1
+        # check that the approximate frequency resolution is obtained
+
+    def test_dispersion_no_resolution(self, random_patch):
+        """Ensure dispersion calc works when no resolution nor limits are provided."""
+        test_vels = np.linspace(1500, 5000, 50)
+        # create a smaller patch so this runs quicker.
+        patch = random_patch.select(
+            time=(0, 50), distance=(0, 10), relative=True, samples=True
+        )
+        dispersive_patch = dispersion_phase_shift(patch, test_vels)
+        assert isinstance(dispersive_patch, dc.Patch)
+        assert "velocity" in dispersive_patch.dims
+        assert "frequency" in dispersive_patch.dims
+
+    def test_non_monotonic_velocities(self, random_patch):
+        """Ensure non-monotonic velocities raise Parameter Error."""
+        msg = "must be monotonically increasing"
+        velocities = np.array([10, -2, 100, 42])
+        with pytest.raises(ParameterError, match=msg):
+            random_patch.dispersion_phase_shift(phase_velocities=velocities)
+
+    def test_velocity_lt_0_raises(self, random_patch):
+        """Ensure velocity values < 0 raise ParameterError."""
+        msg = "Velocities must be positive"
+        velocities = np.array([-1, 0, 1])
+        with pytest.raises(ParameterError, match=msg):
+            random_patch.dispersion_phase_shift(phase_velocities=velocities)
+
+    def test_approx_resolution_gt_0(self, random_patch):
+        """Ensure velocity values < 0 raise ParameterError."""
+        msg = "Frequency resolution has to be positive"
+        test_vels = np.linspace(1500, 5000, 10)
+        with pytest.raises(ParameterError, match=msg):
+            random_patch.dispersion_phase_shift(
+                phase_velocities=test_vels,
+                approx_resolution=-1,
+            )
+
+    def test_freq_range_gt_0(self):
+        """Ensure negative Parameter Error."""
+        msg = "Minimal and maximal frequencies have to be positive"
+        velocities = np.linspace(1500, 5000, 50)
+        patch = get_example_patch("dispersion_event")
+        with pytest.raises(ParameterError, match=msg):
+            patch.dispersion_phase_shift(
+                phase_velocities=velocities,
+                approx_resolution=1.0,
+                approx_freq=[-10, 50],
+            )
+        with pytest.raises(ParameterError, match=msg):
+            patch.dispersion_phase_shift(
+                phase_velocities=velocities,
+                approx_resolution=1.0,
+                approx_freq=[10, -50],
+            )
+
+    def test_freq_range_raises(self, random_patch):
+        """Ensure negative Parameter Error."""
+        msg = "Maximal frequency needs to be larger than minimal frequency"
+        velocities = np.linspace(1500, 5000, 50)
+        with pytest.raises(ParameterError, match=msg):
+            random_patch.dispersion_phase_shift(
+                phase_velocities=velocities, approx_resolution=1.0, approx_freq=[23, 22]
+            )
+
+    def test_freq_range_nyquist(self, random_patch):
+        """Ensure negative Parameter Error."""
+        msg = "Frequency range cannot exceed Nyquist"
+        velocities = np.linspace(1500, 5000, 50)
+        with pytest.raises(ParameterError, match=msg):
+            random_patch.dispersion_phase_shift(
+                phase_velocities=velocities,
+                approx_resolution=1.0,
+                approx_freq=[5000, 8000],
+            )
+
+    def test_freq_range_yields_empty(self, random_patch):
+        """Ensure negative Parameter Error."""
+        msg = "Combination of frequency resolution and range is not an array"
+        velocities = np.linspace(1500, 5000, 50)
+        with pytest.raises(ParameterError, match=msg):
+            random_patch.dispersion_phase_shift(
+                phase_velocities=velocities,
+                approx_resolution=2.0,
+                approx_freq=[22.0, 22.1],
+            )