Total gradient amplitude of a regular grid

Total Gradient Amplitude:
 <xarray.DataArray (northing: 370, easting: 346)> Size: 1MB
array([[1.08738592, 0.72940807, 0.75509218, ..., 2.0132328 , 1.95999329,
        3.07313939],
       [0.73942233, 0.22047172, 0.23189001, ..., 0.49121787, 0.46025515,
        1.96645552],
       [0.76571121, 0.23589359, 0.24513169, ..., 0.51108555, 0.49303789,
        2.03290899],
       ...,
       [3.80734867, 0.95574565, 0.87764784, ..., 0.59290377, 0.57163822,
        1.15583825],
       [3.71454961, 0.92058867, 0.87376757, ..., 0.63073088, 0.60064385,
        1.14353417],
       [5.63063027, 3.39067041, 3.24439853, ..., 0.61133653, 0.64716741,
        1.33503332]])
Coordinates:
  * easting   (easting) float64 3kB 4.655e+05 4.656e+05 ... 4.827e+05 4.828e+05
  * northing  (northing) float64 3kB 7.576e+06 7.576e+06 ... 7.595e+06 7.595e+06

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()