Add theta_map function and tests

harmonica/_transformations.py

@@ -341,6 +341,58 @@ def reduction_to_pole(
     )
 
 
+def total_horizontal_gradient(grid):
+    r"""
+    Calculate the total horizontal derivative of a potential field grid.
+
+    Compute the amplitude of the horizontal gradient of a regular gridded
+    potential field `M`. This is a measure of the rate of change in the x and y
+    (horizontal) directions. . The horizontal derivatives are calculated though 
+    finite-differences.
+
+    Parameters
+    ----------
+    grid : :class:`xarray.DataArray`
+        A two-dimensional :class:`xarray.DataArray` whose coordinates are
+        evenly spaced (regular grid). Its dimensions should be in the following
+        order: *northing*, *easting*. The coordinates must be defined in the same units.
+
+    Returns
+    -------
+    horizontal_derivative_grid : :class:`xarray.DataArray`
+        A :class:`xarray.DataArray` containing the total horizontal derivative of
+        the input ``grid``.
+
+    Notes
+    -----
+    The total horizontal derivative is calculated as:
+
+    .. math::
+
+        A(x, y) = \sqrt{
+            \left( \frac{\partial M}{\partial x} \right)^2
+            + \left( \frac{\partial M}{\partial y} \right)^2
+        }
+
+    where :math:`M` is the regularly gridded potential field.
+
+    References
+    ----------
+    [Blakely1995]_
+    [CordellGrauch1985]_
+    """
+
+    # Run sanity checks on the grid
+    grid_sanity_checks(grid)
+    # Calculate the horizontal gradients of the grid
+    horizontal_gradient = (
+        derivative_easting(grid, order=1),
+        derivative_northing(grid, order=1)
+    )
+    # return the total horizontal gradient
+    return np.sqrt(horizontal_gradient[0] ** 2 + horizontal_gradient[1] ** 2)
+
+
 def total_gradient_amplitude(grid):
     r"""
     Calculates the total gradient amplitude of a potential field grid
@@ -455,6 +507,72 @@ def tilt_angle(grid):
     return tilt
 
 
+def theta_map(grid):
+    r"""
+    Calculate the theta map of a potential field grid
+
+    Compute the theta map of a regularly gridded potential field
+    :math:`M`, used to enhance edges in gravity and magnetic data.
+    The horizontal and vertical derivatives are calculated from the
+    input grid, and the theta angle is defined as the arctangent
+    of the ratio between the total horizontal derivative and the
+    total gradient amplitude.
+
+    Parameters
+    ----------
+    grid : :class:`xarray.DataArray`
+       A two-dimensional :class:`xarray.DataArray` whose coordinates are
+        evenly spaced (regular grid). Its dimensions should be in the following
+        order: *northing*, *easting*. Its coordinates should be defined in the
+        same units.
+
+    Returns
+    -------
+    theta_grid : :class:`xarray.DataArray`
+       A :class:`xarray.DataArray` with the calculated theta angle
+        in radians.
+
+    Notes
+    -----
+    The theta angle is calculated as:
+
+    .. math::
+
+        \theta(M) = \tan^{-1} \left(
+            \frac{
+                \sqrt{
+                    \left( \frac{\partial M}{\partial x} \right)^2 +
+                    \left( \frac{\partial M}{\partial y} \right)^2
+                }
+            }{
+                \sqrt{
+                    \left( \frac{\partial M}{\partial x} \right)^2 +
+                    \left( \frac{\partial M}{\partial y} \right)^2 +
+                    \left( \frac{\partial M}{\partial z} \right)^2
+                }
+            }
+        \right)
+
+    where :math:`M` is the regularly gridded potential field.
+
+    References
+    ----------
+    [Wijns et al., 2005]_
+    """
+    # Run sanity checks on the grid
+    grid_sanity_checks(grid)
+    # Calculate the gradients of the grid
+    gradient = (
+        derivative_easting(grid, order=1),
+        derivative_northing(grid, order=1),
+        derivative_upward(grid, order=1),
+    )
+    # Calculate and return the theta map
+    total_gradient = np.sqrt(gradient[0] ** 2 + gradient[1] ** 2 + gradient[2] ** 2)
+    horiz_deriv = np.sqrt(gradient[0] ** 2 + gradient[1] ** 2)
+    return np.arctan2(horiz_deriv, total_gradient)
+
+
 def _get_dataarray_coordinate(grid, dimension_index):
     """
     Return the name of the easting or northing coordinate in the grid

harmonica/tests/test_transformations.py

@@ -26,7 +26,9 @@
     gaussian_lowpass,
     reduction_to_pole,
     tilt_angle,
+    theta_map,
     total_gradient_amplitude,
+    total_horizontal_gradient,
     upward_continuation,
 )
 from .utils import root_mean_square_error
@@ -602,6 +604,70 @@ def test_invalid_grid_with_nans(self, sample_potential):
             total_gradient_amplitude(sample_potential)
 
 
+class TestTotalHorizontalGradient:
+    """
+    Test total_horizontal_gradient function
+    """
+
+    def test_against_synthetic(
+        self, sample_potential, sample_g_n, sample_g_e
+    ):
+        """
+        Test total_horizontal_gradient function against the synthetic model
+        """
+        pad_width = {
+            "easting": sample_potential.easting.size // 3,
+            "northing": sample_potential.northing.size // 3,
+        }
+        potential_padded = xrft.pad(
+            sample_potential.drop_vars("upward"),
+            pad_width=pad_width,
+        )
+        thg = total_horizontal_gradient(potential_padded)
+        thg = xrft.unpad(thg, pad_width)
+
+        trim = 6
+        thg = thg[trim:-trim, trim:-trim]
+        g_e = sample_g_e[trim:-trim, trim:-trim] * 1e-5  # convert to SI
+        g_n = sample_g_n[trim:-trim, trim:-trim] * 1e-5
+        g_thg = np.sqrt(g_e**2 + g_n**2)
+        rms = root_mean_square_error(thg, g_thg)
+        assert rms / np.abs(g_thg).max() < 0.1
+
+    def test_invalid_grid_single_dimension(self):
+        """
+        Check if total_horizontal_gradient raises error on grid with single
+        dimension
+        """
+        x = np.linspace(0, 10, 11)
+        y = x**2
+        grid = xr.DataArray(y, coords={"x": x}, dims=("x",))
+        with pytest.raises(ValueError, match="Invalid grid with 1 dimensions."):
+            total_horizontal_gradient(grid)
+
+    def test_invalid_grid_three_dimensions(self):
+        """
+        Check if total_horizontal_gradient raises error on grid with three
+        dimensions
+        """
+        x = np.linspace(0, 10, 11)
+        y = np.linspace(-4, 4, 9)
+        z = np.linspace(20, 30, 5)
+        xx, yy, zz = np.meshgrid(x, y, z)
+        data = xx + yy + zz
+        grid = xr.DataArray(data, coords={"x": x, "y": y, "z": z}, dims=("y", "x", "z"))
+        with pytest.raises(ValueError, match="Invalid grid with 3 dimensions."):
+            total_horizontal_gradient(grid)
+
+    def test_invalid_grid_with_nans(self, sample_potential):
+        """
+        Check if total_horizontal_gradient raises error if grid contains nans
+        """
+        sample_potential.values[0, 0] = np.nan
+        with pytest.raises(ValueError, match="Found nan"):
+            total_horizontal_gradient(sample_potential)
+
+
 class TestTilt:
     """
     Test tilt function
@@ -672,6 +738,70 @@ def test_invalid_grid_with_nans(self, sample_potential):
         with pytest.raises(ValueError, match="Found nan"):
             tilt_angle(sample_potential)
 
+class TestThetaMap:
+    """
+    Test theta_map function
+    """
+
+    def test_against_synthetic(
+        self, sample_potential, sample_g_n, sample_g_e, sample_g_z
+    ):
+        """
+        Test theta_map function against the synthetic model
+        """
+        pad_width = {
+            "easting": sample_potential.easting.size // 3,
+            "northing": sample_potential.northing.size // 3,
+        }
+        potential_padded = xrft.pad(
+            sample_potential.drop_vars("upward"),
+            pad_width=pad_width,
+        )
+        theta_grid = theta_map(potential_padded)
+        theta_grid = xrft.unpad(theta_grid, pad_width)
+
+        trim = 6
+        theta_grid = theta_grid[trim:-trim, trim:-trim]
+        g_e = sample_g_e[trim:-trim, trim:-trim] * 1e-5  # SI units
+        g_n = sample_g_n[trim:-trim, trim:-trim] * 1e-5
+        g_z = sample_g_z[trim:-trim, trim:-trim] * 1e-5
+        g_thdr = np.sqrt(g_e**2 + g_n**2)
+        g_total = np.sqrt(g_thdr**2 + g_z**2)
+        g_theta = np.arctan2(g_thdr, g_total)
+        rms = root_mean_square_error(theta_grid, g_theta)
+        assert rms / np.abs(g_theta).max() < 0.1
+
+    def test_invalid_grid_single_dimension(self):
+        """
+        Check if theta_map raises error on grid with single dimension
+        """
+        x = np.linspace(0, 10, 11)
+        y = x**2
+        grid = xr.DataArray(y, coords={"x": x}, dims=("x",))
+        with pytest.raises(ValueError, match="Invalid grid with 1 dimensions."):
+            theta_map(grid)
+
+    def test_invalid_grid_three_dimensions(self):
+        """
+        Check if theta_map raises error on grid with three dimensions
+        """
+        x = np.linspace(0, 10, 11)
+        y = np.linspace(-4, 4, 9)
+        z = np.linspace(20, 30, 5)
+        xx, yy, zz = np.meshgrid(x, y, z)
+        data = xx + yy + zz
+        grid = xr.DataArray(data, coords={"x": x, "y": y, "z": z}, dims=("y", "x", "z"))
+        with pytest.raises(ValueError, match="Invalid grid with 3 dimensions."):
+            theta_map(grid)
+
+    def test_invalid_grid_with_nans(self, sample_potential):
+        """
+        Check if theta_map raises error if grid contains nans
+        """
+        sample_potential.values[0, 0] = np.nan
+        with pytest.raises(ValueError, match="Found nan"):
+            theta_map(sample_potential)
+
 
 class Testfilter:
     """