Implement MPIFFT2D and MPIFFTND

.github/workflows/build.yml

@@ -36,12 +36,24 @@ jobs:
         uses: actions/setup-python@v5
         with:
           python-version: ${{ matrix.python-version }}
+      - name: Install fftw libraries
+        run: |
+          case $(uname) in
+            Linux)
+              sudo apt update
+              sudo apt install -y -q libfftw3-dev
+              ;;
+            Darwin)
+              export HOMEBREW_NO_INSTALLED_DEPENDENTS_CHECK=1
+              brew install fftw
+              ;;
+          esac
       - name: Installing Dependencies
         run: |
           python -m pip install --upgrade pip setuptools
           if [ -f requirements.txt ]; then pip install -r requirements-dev.txt; fi
       - name: Install pylops-mpi
-        run: pip install .
+        run: pip install .[all]
       - name: Testing using pytest-mpi
         run: |
           if [ "${{ matrix.mpi }}" = "openmpi" ]; then

.github/workflows/deploy-docs.yml

@@ -22,12 +22,15 @@ jobs:
         uses: actions/setup-python@v4
         with:
           python-version: 3.11
-
+      - name: Install fftw libraries
+        run: |
+          sudo apt update
+          sudo apt install -y -q libfftw3-dev
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
           if [ -f requirements.txt ]; then pip install -r requirements-dev.txt; fi
-          pip install .
+          pip install .[all]
       - name: Build docs
         run: |
           sphinx-build -b html ./docs/source ./docs/build

Makefile

@@ -43,6 +43,10 @@ dev-install_conda:
 dev-install_conda_nccl:
 	conda env create -f environment-dev.yml && conda activate pylops_mpi && conda install -c conda-forge cupy nccl && pip install -e .
 
+dev-install_fft:
+	make pipcheck
+	$(PIP) install -r requirements-dev.txt && $(PIP) install -e ".[fft]"
+
 lint:
 	flake8 pylops_mpi/ tests/ examples/ tutorials/
 

docs/source/api/index.rst

@@ -71,7 +71,8 @@ Signal Processing
    :toctree: generated/
 
     MPIFredholm1
-
+    MPIFFT2D
+    MPIFFTND
 
 Wave-Equation processing
 ~~~~~~~~~~~~~~~~~~~~~~~~
@@ -107,7 +108,7 @@ Basic
     cgls
 
 Sparsity
-~~~~~
+~~~~~~~~
 
 .. currentmodule:: pylops_mpi.optimization.cls_sparsity
 

docs/source/conf.py

@@ -28,7 +28,8 @@
     "scipy": ("https://docs.scipy.org/doc/scipy/reference", None),
     "matplotlib": ("https://matplotlib.org/", None),
     "mpi4py": ("https://mpi4py.readthedocs.io/en/stable/", None),
-    "pylops": ("https://pylops.readthedocs.io/en/stable/", None)
+    "pylops": ("https://pylops.readthedocs.io/en/stable/", None),
+    "mpi4py_fft": ("https://mpi4py-fft.readthedocs.io/en/stable/", None)
 }
 
 # Generate autodoc stubs with summaries from code

docs/source/installation.rst

@@ -79,6 +79,13 @@ using `pip`:
 
    Replace `12x` with your CUDA version (e.g., `11x` for CUDA 11.x).
 
+Install fft dependencies (optional)
+===================================
+Similarly, to use the FFT classes with distributed arrays, install PyLops-MPI with the ``fft`` extra:
+
+.. code-block:: bash
+
+   >> pip install pylops-mpi[fft]
 
 .. _UserInstall:
 
@@ -94,6 +101,12 @@ command in your terminal to install the PyPI distribution:
 
 Note that when installing via `pip`, only *required* dependencies are installed.
 
+Alternatively, optional dependencies can be installed individually via:
+
+.. code-block:: bash
+
+   >> pip install pylops-mpi[all]
+
 
 .. _DevInstall:
 
@@ -169,6 +182,12 @@ Otherwise, you can change the command in `Makefile` to an appropriate CUDA versi
 i.e., If you use CUDA 11.x, change ``cupy-cuda12x`` and ``nvidia-nccl-cu12`` to 
 ``cupy-cuda11x`` and ``nvidia-nccl-cu11`` and run the command.
 
+If you want to be able to use FFT classes with distributed arrays, run:
+
+.. code-block:: bash
+
+   >> make dev-install_fft
+
 Run tests
 =========
 To ensure that everything has been setup correctly, run tests:

environment-dev.yml

@@ -21,5 +21,6 @@ dependencies:
   - nbsphinx
   - pydata-sphinx-theme
   - flake8
+  - mpi4py-fft
   - pip:
       - sphinx-gallery

examples/plot_ffts.py

@@ -0,0 +1,109 @@
+"""
+Fourier Transform
+=================
+This example shows how to use the :py:class:`pylops_mpi.signalprocessing.MPIFFT2D`
+and :py:class:`pylops_mpi.signalprocessing.MPIFFTND` operators to apply the Fourier
+Transform to the model and the inverse Fourier Transform to the data.
+"""
+
+import matplotlib.pyplot as plt
+import numpy as np
+
+import pylops_mpi
+
+plt.close("all")
+
+###############################################################################
+# We start by applying the two dimensional MPI-distributed FFT to a
+# two-dimensional signal using :py:class:`pylops_mpi.signalprocessing.MPIFFT2D`.
+# The input signal is a :py:class:`pylops_mpi.DistributedArray` which is
+# distributed across MPI ranks before applying the transform.
+
+dt, dx = 0.005, 5
+nt, nx = 2**7, 2**8
+t = np.arange(nt) * dt
+x = np.arange(nx) * dx
+f0 = 10
+
+d = np.outer(np.sin(2 * np.pi * f0 * t), np.arange(nx) + 1)
+dist = pylops_mpi.DistributedArray.to_dist(x=d.ravel())
+
+FFTop = pylops_mpi.signalprocessing.MPIFFT2D(
+    dims=(nt, nx), sampling=(dt, dx)
+)
+
+D = FFTop * dist
+
+dinv = FFTop.H * D
+dinv = np.real(dinv.asarray()).reshape(nt, nx)
+
+D_2d = D.asarray().reshape(nt, nx)
+
+fig, axs = plt.subplots(2, 2, figsize=(10, 6))
+
+axs[0][0].imshow(d, vmin=-100, vmax=100, cmap="bwr")
+axs[0][0].set_title("Signal")
+axs[0][0].axis("tight")
+
+axs[0][1].imshow(
+    np.abs(np.fft.fftshift(D_2d, axes=1)[:nt // 2, :]), cmap="bwr"
+)
+axs[0][1].set_title("Fourier Transform")
+axs[0][1].axis("tight")
+
+axs[1][0].imshow(dinv, vmin=-100, vmax=100, cmap="bwr")
+axs[1][0].set_title("Inverted")
+axs[1][0].axis("tight")
+
+axs[1][1].imshow(d - dinv, vmin=-100, vmax=100, cmap="bwr")
+axs[1][1].set_title("Error")
+axs[1][1].axis("tight")
+
+fig.tight_layout()
+
+###############################################################################
+# We can also apply the three dimensional MPI-distributed FFT to a
+# three-dimensional signal using :py:class:`pylops_mpi.signalprocessing.MPIFFTND`.
+
+dt, dx, dy = 0.005, 5, 3
+nt, nx, ny = 2**7, 2**6, 13
+t = np.arange(nt) * dt
+x = np.arange(nx) * dx
+y = np.arange(ny) * dy
+f0 = 10
+
+d = np.outer(np.sin(2 * np.pi * f0 * t), np.arange(nx) + 1)
+d = np.tile(d[:, :, np.newaxis], [1, 1, ny])
+dist = pylops_mpi.DistributedArray.to_dist(x=d.ravel())
+
+FFTop = pylops_mpi.signalprocessing.MPIFFTND(
+    dims=(nt, nx, ny),
+    sampling=(dt, dx, dy)
+)
+
+D = FFTop * dist
+dinv = FFTop.H * D
+dinv = np.real(dinv.asarray()).reshape(nt, nx, ny)
+D_3d = D.asarray().reshape(nt, nx, ny)  # shape matches dims now
+
+fig, axs = plt.subplots(2, 2, figsize=(10, 6))
+
+axs[0][0].imshow(d[:, :, ny // 2], vmin=-20, vmax=20, cmap="bwr")
+axs[0][0].set_title("Signal")
+axs[0][0].axis("tight")
+axs[0][1].imshow(
+    np.abs(np.fft.fftshift(D_3d, axes=1)[:nx // 2, :, ny // 2]),
+    cmap="bwr"
+)
+axs[0][1].set_title("Fourier Transform")
+axs[0][1].axis("tight")
+
+axs[1][0].imshow(dinv[:, :, ny // 2], vmin=-20, vmax=20, cmap="bwr")
+axs[1][0].set_title("Inverted")
+axs[1][0].axis("tight")
+
+axs[1][1].imshow(d[:, :, ny // 2] - dinv[:, :, ny // 2], vmin=-20, vmax=20, cmap="bwr")
+axs[1][1].set_title("Error")
+axs[1][1].axis("tight")
+
+fig.tight_layout()