Implement our own FFT functions for xr.DataArray

.github/environment.yml

@@ -13,7 +13,6 @@ dependencies:
   - numba>=0.57
   - xarray>=2022.03
   - verde>=1.9.0
-  - xrft>=1.0
   - choclo>=0.1
   - boule>=0.6.0
   # Build

README.md

@@ -52,9 +52,8 @@ Things that will *not* be covered in Harmonica:
 - Multi-physics partial differential equation solvers. Use
   [SimPEG](http://www.simpeg.xyz/) or [PyGIMLi](https://www.pygimli.org/)
   instead.
-- Generic grid processing methods (like FFT and standard interpolation).
-  We'll rely on [Verde](https://www.fatiando.org/verde),
-  [xrft](https://xrft.readthedocs.io/en/latest/) and
+- Generic grid processing methods (like standard interpolation).
+  We'll rely on [Verde](https://www.fatiando.org/verde) and
   [xarray](https://xarray.dev) for those.
 - Data visualization.
 - GUI applications.

doc/api/index.rst

@@ -51,8 +51,6 @@ Define filters in the frequency domain.
     filters.gaussian_highpass_kernel
     filters.reduction_to_pole_kernel
 
-Use :func:`xrft.xrft.fft` and :func:`xrft.xrft.ifft` to apply Fast-Fourier
-Transforms and its inverse on :class:`xarray.DataArray`.
 
 Equivalent Sources
 ------------------

doc/conf.py

@@ -49,7 +49,6 @@
     "scipy": ("https://docs.scipy.org/doc/scipy/reference", None),
     "pandas": ("http://pandas.pydata.org/pandas-docs/stable/", None),
     "xarray": ("http://xarray.pydata.org/en/stable/", None),
-    "xrft": ("https://xrft.readthedocs.io/en/stable/", None),
     "pooch": ("https://www.fatiando.org/pooch/latest/", None),
     "ensaio": ("https://www.fatiando.org/ensaio/latest/", None),
     "verde": ("https://www.fatiando.org/verde/latest/", None),

doc/install.rst

@@ -74,7 +74,6 @@ Required:
 * `scikit-learn <https://scikit-learn.org>`__
 * `pooch <http://www.fatiando.org/pooch/>`__
 * `verde <http://www.fatiando.org/verde/>`__
-* `xrft <https://xrft.readthedocs.io/>`__
 
 Optional:
 

doc/overview.rst

@@ -19,8 +19,7 @@ Harmonica *will not* provide:
   instead.
 - Generic processing methods like grid transformations (use `Verde
   <https://www.fatiando.org/verde>`__ or `Xarray <https://docs.xarray.dev>`__
-  instead) or multidimensional FFT calculations (use `xrft
-  <https://xrft.readthedocs.io>`__ instead).
+  instead).
 - Reference ellipsoid representations and computations like normal gravity. Use
   `Boule <https://www.fatiando.org/boule>`__ instead.
 - Data visualization functions. Use `matplotlib <https://matplotlib.org/>`__

environment.yml

@@ -15,7 +15,6 @@ dependencies:
   - scikit-learn
   - verde>=1.8.1
   - xarray
-  - xrft>=1.0
   - choclo>=0.1
   - boule>=0.6
   # Optional requirements

pyproject.toml

@@ -37,7 +37,6 @@ dependencies = [
     "numba >= 0.57",
     "xarray >= 2022.03",
     "verde >= 1.8.1",
-    "xrft >= 1.0",
     "choclo >= 0.1",
     "boule >= 0.6.0"
 ]

src/harmonica/filters/_fft.py

@@ -5,13 +5,18 @@
 # This code is part of the Fatiando a Terra project (https://www.fatiando.org)
 #
 """
-Wrap xrft functions to compute FFTs and inverse FFTs.
+Custom FFT and inverse FFT functions that work with :class:`xarray.DataArray`.
+
+These functions are inspired in the ``fft`` and ``ifft`` functions provided by ``xrft``,
+which are released under the MIT license.
 """
 
-import xrft
+import numpy as np
+import numpy.typing as npt
+import xarray as xr
 
 
-def fft(grid, true_phase=True, true_amplitude=True, **kwargs):
+def fft(grid, *, prefix="freq_"):
     """
     Compute Fast Fourier Transform of a 2D regular grid.
 
@@ -22,60 +27,213 @@ def fft(grid, true_phase=True, true_amplitude=True, **kwargs):
         evenly spaced (regular grid). Its dimensions should be in the following
         order: *northing*, *easting*. Its coordinates should be defined in the
         same units.
-    true_phase : bool (optional)
-        Take the coordinates into consideration, keeping the original phase of
-        the coordinates in the spatial domain (``direct_lag``) and multiplies
-        the FFT with an exponential function corresponding to this phase.
-        Defaults to True.
-    true_amplitude : bool (optional)
-        If True, the FFT is multiplied by the spacing of the transformed
-        variables to match theoretical FT amplitude.
-        Defaults to True.
-    **kwargs
-        Any extra keyword arguments will be passed the :func:`xrft.fft` function.
+    prefix : str, optional
+        Prefix used for the name of the frequency coordinates and dimensions.
+
 
     Returns
     -------
     fourier_transform : :class:`xarray.DataArray`
         Array with the Fourier transform of the original grid.
     """
-    return xrft.fft(
-        grid, true_phase=true_phase, true_amplitude=true_amplitude, **kwargs
+    if not isinstance(grid, xr.DataArray):
+        msg = (
+            f"Invalid 'grid' of type '{type(grid).__name__}'. "
+            "It must be an xarray.DataArray."
+        )
+        raise TypeError(msg)
+    if grid.ndim != 2:
+        msg = (
+            f"Invalid grid array with '{grid.ndim}' dimensions. It must be a 2D array."
+        )
+        raise ValueError(msg)
+
+    # Get dimensional coordinates and coordinates' shifts
+    dimensional_coords = tuple(
+        _get_dimensional_coordinate(grid, dim) for dim in grid.dims
     )
+    shifts = tuple(grid.coords[coord].values.min() for coord in dimensional_coords)
+
+    # Generate new coordinates
+    fft_dims = tuple(f"{prefix}{dim}" for dim in grid.dims)
+    fft_coords = {
+        f"{prefix}{coord}": (dim, _fftfreq(grid.coords[coord]))
+        for coord, dim in zip(dimensional_coords, fft_dims, strict=True)
+    }
+
+    # Compute FFT
+    fft = np.fft.fftshift(np.fft.fftn(grid.values))
+
+    # Build new xr.DataArray
+    da_fft = xr.DataArray(fft, dims=fft_dims, coords=fft_coords)
+
+    # Add shifts to frequency coordinates
+    for coord, shift in zip(fft_coords, shifts, strict=True):
+        da_fft.coords[coord].attrs.update({"shift": shift})
+    return da_fft
 
 
-def ifft(fourier_transform, true_phase=True, true_amplitude=True, **kwargs):
+def ifft(fft_grid, *, prefix="freq_"):
     """
-    Compute Inverse Fast Fourier Transform of a 2D regular grid.
+    Compute the inverse Fast Fourier Transform of a 2D regular grid.
+
+    If the frequency coordinates have a *shift* attribute, it will be used to shift the
+    coordinates in the spatial domain to such value.
+
+    .. important::
+
+        Assumes that the ``fft_grid`` is *shifted*: it was passed to
+        :func:`numpy.fft.fftshift`. The outputs of the ``fft`` function satisfy this
+        condition.
 
     Parameters
     ----------
-    fourier_transform : :class:`xarray.DataArray`
+    fft_grid : :class:`xarray.DataArray`
         Array with a regular grid defined in the frequency domain.
         Its dimensions should be in the following order:
         *freq_northing*, *freq_easting*.
-    true_phase : bool (optional)
-        Take the coordinates into consideration, recovering the original
-        coordinates in the spatial domain returning to the the original phase
-        (``direct_lag``), and multiplies the iFFT with an exponential function
-        corresponding to this phase.
-        Defaults to True.
-    true_amplitude : bool (optional)
-        If True, output is divided by the spacing of the transformed variables
-        to match theoretical IFT amplitude.
-        Defaults to True.
-    **kwargs
-        Any extra keyword arguments will be passed the :func:`xrft.ifft` function.
+    prefix : str, optional
+        Prefix used for the name of the frequency coordinates and dimensions.
 
     Returns
     -------
     grid : :class:`xarray.DataArray`
         Array with the inverse Fourier transform of the passed grid.
     """
-    return xrft.ifft(
-        fourier_transform,
-        true_phase=true_phase,
-        true_amplitude=true_amplitude,
-        lag=(None, None),  # Mutes an annoying FutureWarning from xrft
-        **kwargs,
+    if not isinstance(fft_grid, xr.DataArray):
+        msg = (
+            f"Invalid 'grid' of type '{type(fft_grid).__name__}'. "
+            "It must be an xarray.DataArray."
+        )
+        raise TypeError(msg)
+    if fft_grid.ndim != 2:
+        msg = (
+            f"Invalid grid array with '{fft_grid.ndim}' dimensions. "
+            "It must be a 2D array."
+        )
+        raise ValueError(msg)
+
+    for dim in fft_grid.dims:
+        if not dim.startswith(prefix):
+            msg = (
+                f"Invalid frequency dimension '{dim}'. "
+                f"It doesn't start with prefix '{prefix}'."
+            )
+            raise ValueError(msg)
+
+    # Get dimensional frequency coordinates and spacings
+    dimensional_fft_coords = tuple(
+        _get_dimensional_coordinate(fft_grid, dim) for dim in fft_grid.dims
     )
+    for coord in dimensional_fft_coords:
+        if not coord.startswith(prefix):
+            msg = (
+                f"Invalid dimensional coordinate '{coord}'. "
+                f"It doesn't start with prefix '{prefix}'."
+            )
+            raise ValueError(msg)
+
+    # Generate new coordinates
+    dims = tuple(dim.removeprefix(prefix) for dim in fft_grid.dims)
+
+    coords = {
+        coord.removeprefix(prefix): (dim, _ifftfreq(fft_grid.coords[coord]))
+        for coord, dim in zip(dimensional_fft_coords, dims, strict=True)
+    }
+
+    # Compute iFFT
+    ifft = np.fft.ifftn(np.fft.ifftshift(fft_grid.values))
+
+    # Build new xr.DataArray
+    da = xr.DataArray(ifft, dims=dims, coords=coords)
+
+    return da
+
+
+def _get_spacing(coordinate: xr.DataArray) -> float:
+    """
+    Return spacing of a grid coordinate.
+
+    Parameters
+    ----------
+    coordinate : xarray.DataArray
+        DataArray containing the coordinate.
+    coordinate : str
+        Coordinate name.
+
+    Returns
+    -------
+    spacing : float
+    """
+    spacing = coordinate.values[1] - coordinate.values[0]
+    if not np.allclose(spacing, coordinate.values[1:] - coordinate.values[:-1]):
+        msg = f"Invalid '{coordinate.name}' coordinates: they must be evenly spaced."
+        raise ValueError(msg)
+    if spacing <= 0:
+        msg = (
+            f"Invalid coordinate '{coordinate.name}': it must be increasingly ordered."
+        )
+        raise ValueError(msg)
+    return spacing
+
+
+def _get_dimensional_coordinate(grid: xr.DataArray, dim: str) -> str:
+    """
+    Get dimensional coordinate in the grid for a particular dimension.
+
+    Parameters
+    ----------
+    grid : xarray.DataArray
+        DataArray containing the coordinate.
+    dim : str
+        Dimension name.
+
+    Returns
+    -------
+    dimensional_coordinate : str
+    """
+    potential_coords = [
+        coord for coord in grid.coords if grid.coords[coord].dims == (dim,)
+    ]
+    if not potential_coords:
+        msg = f"Couldn't find dimensional coordinate for dimension '{dim}'."
+        raise ValueError(msg)
+    if len(potential_coords) > 1:
+        bad_coords = ", ".join(potential_coords)
+        msg = (
+            f"Multiple dimensional coordinates ({bad_coords}) found "
+            f"for the '{dim}' dimension. "
+            "Leave only one dimensional coordinate per dimension."
+        )
+        raise ValueError(msg)
+    (dimensional_coordinate,) = potential_coords
+    return dimensional_coordinate
+
+
+def _fftfreq(coordinate: xr.DataArray) -> npt.NDArray:
+    """
+    Get coordinate into the frequency domain.
+    """
+    if coordinate.ndim != 1:
+        raise ValueError()
+    spacing = _get_spacing(coordinate)
+    return np.fft.fftshift(np.fft.fftfreq(coordinate.size, spacing))
+
+
+def _ifftfreq(freq: xr.DataArray) -> npt.NDArray:
+    """
+    Recover coordinate in the space domain from the frequency domain.
+
+    Shifts the coordinates in the spatial domain if the ``freq`` has a *shift*
+    attribute.
+    """
+    if freq.ndim != 1:
+        raise ValueError()
+    spacing = _get_spacing(freq)
+    coordinate = np.fft.fftshift(np.fft.fftfreq(freq.size, spacing))
+
+    # Apply static shift if any
+    if "shift" in freq.attrs:
+        coordinate += freq.attrs["shift"] - coordinate.min()
+
+    return coordinate

test/test_filters.py

@@ -142,7 +142,7 @@ def fixture_invalid_grid_with_nans(sample_grid):
 
 def test_fft_round_trip(sample_grid):
     """
-    Test if the wrapped fft and ifft functions satisfy a round trip.
+    Test if the fft and ifft functions satisfy a round trip.
     """
     xrt.assert_allclose(sample_grid, ifft(fft(sample_grid)))