Add mkl_fft in fft

Author: rohanbabbar04

environment-dev-arm.yml

@@ -23,6 +23,7 @@ dependencies:
   - autopep8
   - isort
   - black
+  - mkl_fft
   - pip:
     - torch
     - devito

environment-dev-gpu.yml

@@ -21,6 +21,7 @@ dependencies:
   - autopep8
   - isort
   - black
+  - mkl_fft
   - pip:
     - torch
     - pytest-runner

environment-dev.yml

@@ -24,6 +24,7 @@ dependencies:
   - autopep8
   - isort
   - black
+  - mkl_fft
   - pip:
     - torch
     - devito

examples/plot_fft.py

@@ -66,6 +66,30 @@
 axs[1].set_xlim([0, 3 * f0])
 plt.tight_layout()
 
+###############################################################################
+# PyLops also has a third FFT engine (engine='mkl_fft') that uses the well-known
+# `Intel MKL FFT <https://github.com/IntelPython/mkl_fft>`_. This is a Python wrapper around
+# the `Intel® oneAPI Math Kernel Library (oneMKL) <https://www.intel.com/content/www/us/en/docs/onemkl/developer-reference-c/2025-2/fourier-transform-functions.html>`_
+# Fourier Transform functions. It lets PyLops run discrete Fourier transforms faster
+# by using Intel’s highly optimized math routines.
+
+FFTop = pylops.signalprocessing.FFT(dims=nt, nfft=nfft, sampling=dt, engine="mkl_fft")
+D = FFTop * d
+
+# Inverse for FFT
+dinv = FFTop.H * D
+dinv = FFTop / D
+
+fig, axs = plt.subplots(1, 2, figsize=(10, 4))
+axs[0].plot(t, d, "k", lw=2, label="True")
+axs[0].plot(t, dinv.real, "--r", lw=2, label="Inverted")
+axs[0].legend()
+axs[0].set_title("Signal")
+axs[1].plot(FFTop.f[: int(FFTop.nfft / 2)], np.abs(D[: int(FFTop.nfft / 2)]), "k", lw=2)
+axs[1].set_title("Fourier Transform with MKL FFT")
+axs[1].set_xlim([0, 3 * f0])
+plt.tight_layout()
+
 ###############################################################################
 # We can also apply the one dimensional FFT to to a two-dimensional
 # signal (along one of the first axis)

pylops/signalprocessing/fft.py

@@ -16,10 +16,20 @@
 from pylops.utils.typing import DTypeLike, InputDimsLike, NDArray
 
 pyfftw_message = deps.pyfftw_import("the fft module")
+mkl_fft_message = deps.mkl_fft_import("the mkl fft module")
 
 if pyfftw_message is None:
     import pyfftw
 
+if mkl_fft_message is None:
+    import mkl_fft.interfaces.numpy_fft as mkl_backend
+
+    try:
+        import scipy.fft  # noqa: F401
+        import mkl_fft.interfaces.scipy_fft as mkl_backend
+    except ImportError:
+        pass
+
 logger = logging.getLogger(__name__)
 
 
@@ -394,6 +404,94 @@ def __truediv__(self, y: npt.ArrayLike) -> npt.ArrayLike:
         return self._rmatvec(y) / self._scale
 
 
+class _FFT_mklfft(_BaseFFT):
+    """One-dimensional Fast-Fourier Transform using mkl_fft"""
+
+    def __init__(
+        self,
+        dims: Union[int, InputDimsLike],
+        axis: int = -1,
+        nfft: Optional[int] = None,
+        sampling: float = 1.0,
+        norm: str = "ortho",
+        real: bool = False,
+        ifftshift_before: bool = False,
+        fftshift_after: bool = False,
+        dtype: DTypeLike = "complex128",
+        **kwargs_fft,
+    ) -> None:
+        super().__init__(
+            dims=dims,
+            axis=axis,
+            nfft=nfft,
+            sampling=sampling,
+            norm=norm,
+            real=real,
+            ifftshift_before=ifftshift_before,
+            fftshift_after=fftshift_after,
+            dtype=dtype,
+        )
+        self._kwargs_fft = kwargs_fft
+        self._norm_kwargs = {"norm": None}
+        if self.norm is _FFTNorms.ORTHO:
+            self._norm_kwargs["norm"] = "ortho"
+            self._scale = np.sqrt(1 / self.nfft)
+        elif self.norm is _FFTNorms.NONE:
+            self._scale = self.nfft
+        elif self.norm is _FFTNorms.ONE_OVER_N:
+            self._scale = 1.0 / self.nfft
+
+    @reshaped
+    def _matvec(self, x: NDArray) -> NDArray:
+        if self.ifftshift_before:
+            x = mkl_backend.ifftshift(x, axes=self.axis)
+        if not self.clinear:
+            x = np.real(x)
+        if self.real:
+            y = mkl_backend.rfft(x, n=self.nfft, axis=self.axis, **self._norm_kwargs)
+            y = np.swapaxes(y, -1, self.axis)
+            y[..., 1 : 1 + (self.nfft - 1) // 2] *= np.sqrt(2)
+            y = np.swapaxes(y, self.axis, -1)
+        else:
+            y = mkl_backend.fft(x, n=self.nfft, axis=self.axis, **self._norm_kwargs)
+        if self.norm is _FFTNorms.ONE_OVER_N:
+            y *= self._scale
+        if self.fftshift_after:
+            y = mkl_backend.fftshift(y, axes=self.axis)
+        return y
+
+    @reshaped
+    def _rmatvec(self, x: NDArray) -> NDArray:
+        if self.fftshift_after:
+            x = mkl_backend.ifftshift(x, axes=self.axis)
+        if self.real:
+            x = x.copy()
+            x = np.swapaxes(x, -1, self.axis)
+            x[..., 1 : 1 + (self.nfft - 1) // 2] /= np.sqrt(2)
+            x = np.swapaxes(x, self.axis, -1)
+            y = mkl_backend.irfft(x, n=self.nfft, axis=self.axis, **self._norm_kwargs)
+        else:
+            y = mkl_backend.ifft(x, n=self.nfft, axis=self.axis, **self._norm_kwargs)
+        if self.norm is _FFTNorms.NONE:
+            y *= self._scale
+
+        if self.nfft > self.dims[self.axis]:
+            y = np.take(y, range(0, self.dims[self.axis]), axis=self.axis)
+        elif self.nfft < self.dims[self.axis]:
+            y = np.pad(y, self.ifftpad)
+
+        if not self.clinear:
+            y = np.real(y)
+        if self.ifftshift_before:
+            y = mkl_backend.fftshift(y, axes=self.axis)
+        return y
+
+    def __truediv__(self, y):
+        if self.norm is not _FFTNorms.ORTHO:
+            return self._rmatvec(y) / self._scale
+        return self._rmatvec(y)
+
+
 def FFT(
     dims: Union[int, InputDimsLike],
     axis: int = -1,
@@ -481,7 +579,7 @@ def FFT(
         frequencies are arranged from zero to largest positive, and then from negative
         Nyquist to the frequency bin before zero.
     engine : :obj:`str`, optional
-        Engine used for fft computation (``numpy``, ``fftw``, or ``scipy``). Choose
+        Engine used for fft computation (``numpy``, ``fftw``, ``scipy`` or ``mkl_fft``). Choose
         ``numpy`` when working with cupy and jax arrays.
 
         .. note:: Since version 1.17.0, accepts "scipy".
@@ -534,7 +632,7 @@ def FFT(
         - If ``dims`` is provided and ``axis`` is bigger than ``len(dims)``.
         - If ``norm`` is not one of "ortho", "none", or "1/n".
     NotImplementedError
-        If ``engine`` is neither ``numpy``, ``fftw``, nor ``scipy``.
+        If ``engine`` is neither ``numpy``, ``fftw``, ``scipy`` nor ``mkl_fft``.
 
     See Also
     --------
@@ -579,7 +677,24 @@ def FFT(
             dtype=dtype,
             **kwargs_fft,
         )
-    elif engine == "numpy" or (engine == "fftw" and pyfftw_message is not None):
+    elif engine == "mkl_fft" and mkl_fft_message is None:
+        f = _FFT_mklfft(
+            dims,
+            axis=axis,
+            nfft=nfft,
+            sampling=sampling,
+            norm=norm,
+            real=real,
+            ifftshift_before=ifftshift_before,
+            fftshift_after=fftshift_after,
+            dtype=dtype,
+            **kwargs_fft,
+        )
+    elif (
+        engine == "numpy"
+        or (engine == "fftw" and pyfftw_message is not None)
+        or (engine == "mkl_fft" and mkl_fft_message is not None)
+    ):
         if engine == "fftw" and pyfftw_message is not None:
             logger.warning(pyfftw_message)
         f = _FFT_numpy(
@@ -608,6 +723,6 @@ def FFT(
             **kwargs_fft,
         )
     else:
-        raise NotImplementedError("engine must be numpy, fftw or scipy")
+        raise NotImplementedError("engine must be numpy, fftw, scipy or mkl_fft")
     f.name = name
     return f