Repository: fatiando/harmonica (python)

Issue:
Total gradient amplitude transformation
**Description of the desired feature:**
<!--
Add as much detail as you can in your description below. If possible, include an
example of how you would like to use this feature (even better if it's a code
example).
-->

A nice transformation to add to `harmonica` is the _total gradient amplitude_ (aka _analytic signal_), defined in Roest et al (1992) as 
$$|A(x, y)| = \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}$$

Roest, W. & Verhoef, Joyce & Pilkington, Mark. (1992). Magnetic interpretation using 3-D analytic signal. GEOPHYSICS. 57. 116-125. 10.1190/1.1443174. 



**Are you willing to help implement and maintain this feature?** 
<!--
Every feature we add is code that we will have to maintain and keep updated.
Please let us know if you're willing to help maintain this feature in the future.
-->
I will try to implement it but no promises about maintenance 😅

Pull request:
Add total gradient amplitude transformation
<!--
Thank you for contributing a pull request to Fatiando! 💖

👆🏽 ABOVE: Describe the changes proposed and WHY you made them.

👇🏽 BELOW: Link to any relevant issue or pull request.

Please ensure you have taken a look at the CONTRIBUTING.md file 
in this repository (if available) and the general guidelines at 
https://github.com/fatiando/community/blob/main/CONTRIBUTING.md
-->

Add a new `total_gradient_amplitude()` function that computes the total gradient amplitude for a given 2D grid. It makes use of the derivative functions we already have in Harmonica, computing the horizontal ones through finite differences and the upward one through FFT. Add new gallery and tutorial examples, tests against a synthetic model, and list the new function in the API reference.

**Relevant issues/PRs:**
<!--
Example: "Fixes #1234" / "See also #345" / "Relevant to #111"
Use keywords (e.g., Fixes, Closes) to create the links and automatically
close issues when this PR is merged. 
See https://github.com/blog/1506-closing-issues-via-pull-requests
-->
This PR resolves #470. I

- [x] implemented the _total gradient amplitude_ function in `_transformations.py`
- [x] created the `tga.py` example
- [x] added the section _Total gradient amplitude (aka Analytic Signal)_ to `transformations.rst`
- [x] added `total_gradient_amplitude_test` test to `test_transformations.py`

Notes:
* I had difficulty checking if the mathematical expression in the `transformations.rst` compiled correctly. 
* I also don't know if I commented the function properly since I added a mathematical expression to it.

Patch:
diff --git a/doc/api/index.rst b/doc/api/index.rst
index 8a522a2..2b85813 100644
--- a/doc/api/index.rst
+++ b/doc/api/index.rst
@@ -32,6 +32,7 @@ Apply well known transformations regular gridded potential fields data.
     gaussian_lowpass
     gaussian_highpass
     reduction_to_pole
+    total_gradient_amplitude
 
 Frequency domain filters
 ------------------------
diff --git a/doc/user_guide/transformations.rst b/doc/user_guide/transformations.rst
index 6491c9f..ce9ac97 100644
--- a/doc/user_guide/transformations.rst
+++ b/doc/user_guide/transformations.rst
@@ -396,6 +396,52 @@ Let's plot the results side by side:
     )
     plt.show()
 
+
+Total gradient amplitude
+------------------------
+
+.. hint::
+
+    Total gradient amplitude is also known as *analytic signal*.
+
+We can also calculate the total gradient amplitude of any magnetic anomaly grid.
+This transformation consists in obtaining the amplitude of the gradient of the
+magnetic field in all the three spatial directions by applying
+
+.. math::
+
+   A(x, y) = \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
+   }.
+
+
+We can apply it through the :func:`harmonica.total_gradient_amplitude` function.
+
+.. jupyter-execute::
+
+    tga_grid = hm.total_gradient_amplitude(
+        magnetic_grid_padded
+    )
+
+    # Unpad the total gradient amplitude grid
+    tga_grid = xrft.unpad(tga_grid, pad_width)
+    tga_grid
+
+And plot it:
+
+.. jupyter-execute::
+
+    import verde as vd
+
+    tmp = tga_grid.plot(cmap="viridis", add_colorbar=False)
+    plt.gca().set_aspect("equal")
+    plt.title("Total gradient amplitude of the magnetic anomaly")
+    plt.gca().ticklabel_format(style="sci", scilimits=(0, 0))
+    plt.colorbar(tmp, label="nT/m")
+    plt.show()
+
 ----
 
 .. grid:: 2
diff --git a/examples/transformations/tga.py b/examples/transformations/tga.py
new file mode 100644
index 0000000..665b4f6
--- /dev/null
+++ b/examples/transformations/tga.py
@@ -0,0 +1,71 @@
+# Copyright (c) 2018 The Harmonica Developers.
+# Distributed under the terms of the BSD 3-Clause License.
+# SPDX-License-Identifier: BSD-3-Clause
+#
+# This code is part of the Fatiando a Terra project (https://www.fatiando.org)
+#
+"""
+Total gradient amplitude of a regular grid
+==========================================
+"""
+import ensaio
+import pygmt
+import verde as vd
+import xarray as xr
+import xrft
+
+import harmonica as hm
+
+# Fetch magnetic grid over the Lightning Creek Sill Complex, Australia using
+# Ensaio and load it with Xarray
+fname = ensaio.fetch_lightning_creek_magnetic(version=1)
+magnetic_grid = xr.load_dataarray(fname)
+
+# Pad the grid to increase accuracy of the FFT filter
+pad_width = {
+    "easting": magnetic_grid.easting.size // 3,
+    "northing": magnetic_grid.northing.size // 3,
+}
+# drop the extra height coordinate
+magnetic_grid_no_height = magnetic_grid.drop_vars("height")
+magnetic_grid_padded = xrft.pad(magnetic_grid_no_height, pad_width)
+
+# Compute the total gradient amplitude of the grid
+tga = hm.total_gradient_amplitude(magnetic_grid_padded)
+
+# Unpad the total gradient amplitude grid
+tga = xrft.unpad(tga, pad_width)
+
+# Show the total gradient amplitude
+print("\nTotal Gradient Amplitude:\n", tga)
+
+# Plot original magnetic anomaly and the total gradient amplitude
+fig = pygmt.Figure()
+with fig.subplot(nrows=1, ncols=2, figsize=("28c", "15c"), sharey="l"):
+    with fig.set_panel(panel=0):
+        # Make colormap of data
+        scale = 2500
+        pygmt.makecpt(cmap="polar+h", series=[-scale, scale], background=True)
+        # Plot magnetic anomaly grid
+        fig.grdimage(
+            grid=magnetic_grid,
+            projection="X?",
+            cmap=True,
+        )
+        # Add colorbar
+        fig.colorbar(
+            frame='af+l"Magnetic anomaly [nT]"',
+            position="JBC+h+o0/1c+e",
+        )
+    with fig.set_panel(panel=1):
+        # Make colormap for total gradient amplitude (saturate it a little bit)
+        scale = 0.6 * vd.maxabs(tga)
+        pygmt.makecpt(cmap="polar+h", series=[0, scale], background=True)
+        # Plot total gradient amplitude
+        fig.grdimage(grid=tga, projection="X?", cmap=True)
+        # Add colorbar
+        fig.colorbar(
+            frame='af+l"Total Gradient Amplitude [nT/m]"',
+            position="JBC+h+o0/1c+e",
+        )
+fig.show()
diff --git a/harmonica/__init__.py b/harmonica/__init__.py
index c61fe3f..73132f7 100644
--- a/harmonica/__init__.py
+++ b/harmonica/__init__.py
@@ -27,6 +27,7 @@ from ._transformations import (
     gaussian_highpass,
     gaussian_lowpass,
     reduction_to_pole,
+    total_gradient_amplitude,
     upward_continuation,
 )
 from ._utils import magnetic_angles_to_vec, magnetic_vec_to_angles
diff --git a/harmonica/_transformations.py b/harmonica/_transformations.py
index 4396a42..3957e73 100644
--- a/harmonica/_transformations.py
+++ b/harmonica/_transformations.py
@@ -7,6 +7,8 @@
 """
 Apply transformations to regular grids of potential fields
 """
+import numpy as np
+
 from .filters._filters import (
     derivative_easting_kernel,
     derivative_northing_kernel,
@@ -16,7 +18,7 @@ from .filters._filters import (
     reduction_to_pole_kernel,
     upward_continuation_kernel,
 )
-from .filters._utils import apply_filter
+from .filters._utils import apply_filter, grid_sanity_checks
 
 
 def derivative_upward(grid, order=1):
@@ -339,6 +341,58 @@ def reduction_to_pole(
     )
 
 
+def total_gradient_amplitude(grid):
+    r"""
+    Calculate the total gradient amplitude a magnetic field grid
+
+    Compute the total gradient amplitude of a regular gridded potential field
+    `M`. The horizontal derivatives are calculated though finite-differences
+    while the upward derivative is calculated using FFT.
+
+    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
+    -------
+    total_gradient_amplitude_grid : :class:`xarray.DataArray`
+        A :class:`xarray.DataArray` after calculating the
+        total gradient amplitude of the passed ``grid``.
+
+    Notes
+    -----
+    The total gradient amplitude is calculated as:
+
+    .. math::
+
+        A(x, y) = \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
+        }
+
+    where :math:`M` is the regularly gridded potential field.
+
+    References
+    ----------
+    [Blakely1995]_
+    """
+    # 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),
+    )
+    # return the total gradient amplitude
+    return np.sqrt(gradient[0] ** 2 + gradient[1] ** 2 + gradient[2] ** 2)
+
+
 def _get_dataarray_coordinate(grid, dimension_index):
     """
     Return the name of the easting or northing coordinate in the grid
diff --git a/harmonica/filters/_utils.py b/harmonica/filters/_utils.py
index 0325386..a3680f4 100644
--- a/harmonica/filters/_utils.py
+++ b/harmonica/filters/_utils.py
@@ -38,21 +38,10 @@ def apply_filter(grid, fft_filter, **kwargs):
         A :class:`xarray.DataArray` with the filtered version of the passed
         ``grid``. Defined are in the spatial domain.
     """
+    # Run sanity checks on the grid
+    grid_sanity_checks(grid)
     # Catch the dims of the grid
     dims = grid.dims
-    # Check if the array has two dimensions
-    if len(dims) != 2:
-        raise ValueError(
-            f"Invalid grid with {len(dims)} dimensions. "
-            + "The passed grid must be a 2 dimensional array."
-        )
-    # Check if the grid has nans
-    if np.isnan(grid).any():
-        raise ValueError(
-            "Found nan(s) on the passed grid. "
-            + "The grid must not have missing values before computing the "
-            + "Fast Fourier Transform."
-        )
     # Compute Fourier Transform of the grid
     fft_grid = fft(grid)
     # Build the filter
@@ -68,3 +57,35 @@ def apply_filter(grid, fft_filter, **kwargs):
     )
 
     return filtered_grid
+
+
+def grid_sanity_checks(grid):
+    """
+    Run sanity checks on the grid
+
+    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.
+
+    Raises
+    ------
+    ValueError
+        If the passed grid is not 2D or if it contains nan values.
+    """
+    # Check if the array has two dimensions
+    if (n_dims := len(grid.dims)) != 2:
+        raise ValueError(
+            f"Invalid grid with {n_dims} dimensions. "
+            + "The passed grid must be a 2 dimensional array."
+        )
+    # Check if the grid has nans
+    if np.isnan(grid).any():
+        raise ValueError(
+            "Found nan(s) on the passed grid. "
+            + "The grid must not have missing values before computing the "
+            + "Fast Fourier Transform."
+        )


Test patch:
diff --git a/harmonica/tests/test_transformations.py b/harmonica/tests/test_transformations.py
index 78d4592..d93b365 100644
--- a/harmonica/tests/test_transformations.py
+++ b/harmonica/tests/test_transformations.py
@@ -25,6 +25,7 @@ from .._transformations import (
     gaussian_highpass,
     gaussian_lowpass,
     reduction_to_pole,
+    total_gradient_amplitude,
     upward_continuation,
 )
 from .utils import root_mean_square_error
@@ -522,6 +523,74 @@ def test_upward_continuation(sample_g_z, sample_g_z_upward):
     xrt.assert_allclose(continuation, g_z_upward, atol=1e-8)
 
 
+class TestTotalGradientAmplitude:
+    """
+    Test total_gradient_amplitude function
+    """
+
+    def test_against_synthetic(
+        self, sample_potential, sample_g_n, sample_g_e, sample_g_z
+    ):
+        """
+        Test total_gradient_amplitude function against the synthetic model
+        """
+        # Pad the potential field grid to improve accuracy
+        pad_width = {
+            "easting": sample_potential.easting.size // 3,
+            "northing": sample_potential.northing.size // 3,
+        }
+        # need to drop upward coordinate (bug in xrft)
+        potential_padded = xrft.pad(
+            sample_potential.drop_vars("upward"),
+            pad_width=pad_width,
+        )
+        # Calculate total gradient amplitude and unpad it
+        tga = total_gradient_amplitude(potential_padded)
+        tga = xrft.unpad(tga, pad_width)
+        # Compare against g_tga (trim the borders to ignore boundary effects)
+        trim = 6
+        tga = tga[trim:-trim, trim:-trim]
+        g_e = sample_g_e[trim:-trim, trim:-trim] * 1e-5  # convert to SI units
+        g_n = sample_g_n[trim:-trim, trim:-trim] * 1e-5  # convert to SI units
+        g_z = sample_g_z[trim:-trim, trim:-trim] * 1e-5  # convert to SI units
+        g_tga = np.sqrt(g_e**2 + g_n**2 + g_z**2)
+        rms = root_mean_square_error(tga, g_tga)
+        assert rms / np.abs(g_tga).max() < 0.1
+
+    def test_invalid_grid_single_dimension(self):
+        """
+        Check if total_gradient_amplitude 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_gradient_amplitude(grid)
+
+    def test_invalid_grid_three_dimensions(self):
+        """
+        Check if total_gradient_amplitude 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_gradient_amplitude(grid)
+
+    def test_invalid_grid_with_nans(self, sample_potential):
+        """
+        Check if total_gradient_amplitude raises error if grid contains nans
+        """
+        sample_potential.values[0, 0] = np.nan
+        with pytest.raises(ValueError, match="Found nan"):
+            total_gradient_amplitude(sample_potential)
+
+
 class Testfilter:
     """
     Test filter result against the output from oasis montaj