ENH: Add Scipy and Cupy as fft interfaces

.github/workflows/check.yml

@@ -10,7 +10,7 @@ jobs:
     runs-on: ${{ matrix.os }}
     strategy:
       matrix:
-        python-version: ['3.10']
+        python-version: ['3.10', '3.11']
         os: [macos-latest, ubuntu-latest, windows-latest]
 
     steps:

docs/requirements.txt

@@ -1,4 +1,3 @@
-sphinx==4.3.0
-sphinxcontrib.bibtex>=2.0
-sphinx_rtd_theme==1.0
-
+sphinx
+sphinxcontrib.bibtex
+sphinx_rtd_theme

docs/sec_code_reference.rst

@@ -23,21 +23,34 @@ Fourier transform methods
 =========================
 
 .. _sec_code_fourier_numpy:
-
 Numpy
 -----
 .. automodule:: qpretrieve.fourier.ff_numpy
    :members:
    :inherited-members:
 
 .. _sec_code_fourier_pyfftw:
-
 PyFFTW
 ------
 .. automodule:: qpretrieve.fourier.ff_pyfftw
    :members:
    :inherited-members:
 
+.. _sec_code_fourier_scipy:
+Scipy
+------
+.. automodule:: qpretrieve.fourier.ff_scipy
+   :members:
+   :inherited-members:
+
+.. _sec_code_fourier_cupy:
+Cupy
+----
+.. automodule:: qpretrieve.fourier.ff_cupy
+   :members:
+   :inherited-members:
+
+
 .. _sec_code_ifer:
 
 Interference image analysis

docs/sec_examples.rst

@@ -12,3 +12,5 @@ Examples
 .. fancy_include:: filter_visualization.py
 
 .. fancy_include:: fourier_scale.py
+
+.. fancy_include:: fft_options.py

examples/fft_options.py

@@ -0,0 +1,73 @@
+"""Fourier Transform interfaces available
+
+This example visualizes the different backends and packages available to the
+user for performing Fourier transforms.
+
+- PyFFTW is initially slow, but over many FFTs is very quick.
+- CuPy using CUDA can be very fast, but is currently limited because we are
+  transferring one image at a time to the GPU.
+
+"""
+import time
+import matplotlib.pylab as plt
+import numpy as np
+import qpretrieve
+
+# load the experimental data
+edata = np.load("./data/hologram_cell.npz")
+
+# get the available fft interfaces
+interfaces_available = qpretrieve.fourier.get_available_interfaces()
+
+n_transforms = 100
+
+# one transform
+results_1 = {}
+for fft_interface in interfaces_available:
+    t0 = time.time()
+    holo = qpretrieve.OffAxisHologram(data=edata["data"],
+                                      fft_interface=fft_interface)
+    holo.run_pipeline(filter_name="disk", filter_size=1 / 2)
+    bg = qpretrieve.OffAxisHologram(data=edata["bg_data"])
+    bg.process_like(holo)
+    t1 = time.time()
+    results_1[fft_interface.__name__] = t1 - t0
+num_interfaces = len(results_1)
+
+# multiple transforms (should see speed increase for PyFFTW)
+results = {}
+for fft_interface in interfaces_available:
+    t0 = time.time()
+    for _ in range(n_transforms):
+        holo = qpretrieve.OffAxisHologram(data=edata["data"],
+                                          fft_interface=fft_interface)
+        holo.run_pipeline(filter_name="disk", filter_size=1 / 2)
+        bg = qpretrieve.OffAxisHologram(data=edata["bg_data"])
+        bg.process_like(holo)
+    t1 = time.time()
+    results[fft_interface.__name__] = t1 - t0
+num_interfaces = len(results)
+
+fft_interfaces = list(results.keys())
+speed_1 = list(results_1.values())
+speed = list(results.values())
+
+fig, axes = plt.subplots(1, 2, figsize=(8, 5))
+ax1, ax2 = axes
+labels = [fftstr[9:] for fftstr in fft_interfaces]
+
+ax1.bar(range(num_interfaces), height=speed_1, color='lightseagreen')
+ax1.set_xticks(range(num_interfaces), labels=labels,
+               rotation=45)
+ax1.set_ylabel("Speed (s)")
+ax1.set_title("1 Transform")
+
+ax2.bar(range(num_interfaces), height=speed, color='lightseagreen')
+ax2.set_xticks(range(num_interfaces), labels=labels,
+               rotation=45)
+ax2.set_ylabel("Speed (s)")
+ax2.set_title(f"{n_transforms} Transforms")
+
+plt.suptitle("Speed of FFT Interfaces")
+plt.tight_layout()
+plt.show()

examples/requirements.txt

@@ -0,0 +1 @@
+matplotlib

qpretrieve/filter.py

@@ -38,7 +38,7 @@ def get_filter_array(filter_name, filter_size, freq_pos, fft_shape):
         and must be between 0 and `max(fft_shape)/2`
     freq_pos: tuple of floats
         The position of the filter in frequency coordinates as
-        returned by :func:`nunpy.fft.fftfreq`.
+        returned by :func:`numpy.fft.fftfreq`.
     fft_shape: tuple of int
         The shape of the Fourier transformed image for which the
         filter will be applied. The shape must be squared (two
@@ -104,8 +104,10 @@ def get_filter_array(filter_name, filter_size, freq_pos, fft_shape):
         # TODO: avoid the np.roll, instead use the indices directly
         alpha = 0.1
         rsize = int(min(fx.size, fy.size) * filter_size) * 2
-        tukey_window_x = signal.tukey(rsize, alpha=alpha).reshape(-1, 1)
-        tukey_window_y = signal.tukey(rsize, alpha=alpha).reshape(1, -1)
+        tukey_window_x = signal.windows.tukey(
+            rsize, alpha=alpha).reshape(-1, 1)
+        tukey_window_y = signal.windows.tukey(
+            rsize, alpha=alpha).reshape(1, -1)
         tukey = tukey_window_x * tukey_window_y
         base = np.zeros(fft_shape)
         s1 = (np.array(fft_shape) - rsize) // 2

qpretrieve/fourier/__init__.py

@@ -2,15 +2,36 @@
 import warnings
 
 from .ff_numpy import FFTFilterNumpy
+from .ff_scipy import FFTFilterScipy
 
 try:
     from .ff_pyfftw import FFTFilterPyFFTW
 except ImportError:
     FFTFilterPyFFTW = None
 
+try:
+    from .ff_cupy import FFTFilterCupy
+except ImportError:
+    FFTFilterCupy = None
+
 PREFERRED_INTERFACE = None
 
 
+def get_available_interfaces():
+    """Return a list of available FFT algorithms"""
+    interfaces = [
+        FFTFilterPyFFTW,
+        FFTFilterNumpy,
+        FFTFilterScipy,
+        FFTFilterCupy,
+    ]
+    interfaces_available = []
+    for interface in interfaces:
+        if interface is not None and interface.is_available:
+            interfaces_available.append(interface)
+    return interfaces_available
+
+
 def get_best_interface():
     """Return the fastest refocusing interface available
 

qpretrieve/fourier/base.py

@@ -70,8 +70,11 @@ def __init__(self, data, subtract_mean=True, padding=2, copy=True):
         else:
             # convert integer-arrays to floating point arrays
             dtype = float
+        if not copy:
+            # numpy v2.x behaviour requires asarray with copy=False
+            copy = None
         data_ed = np.array(data, dtype=dtype, copy=copy)
-        #: original data (with subtracted mean)
+#: original data (with subtracted mean)
         self.origin = data_ed
         #: whether padding is enabled
         self.padding = padding
@@ -175,7 +178,7 @@ def filter(self, filter_name: str, filter_size: float,
             and must be between 0 and `max(fft_shape)/2`
         freq_pos: tuple of floats
             The position of the filter in frequency coordinates as
-            returned by :func:`nunpy.fft.fftfreq`.
+            returned by :func:`numpy.fft.fftfreq`.
         scale_to_filter: bool or float
             Crop the image in Fourier space after applying the filter,
             effectively removing surplus (zero-padding) data and

qpretrieve/fourier/ff_cupy.py

@@ -0,0 +1,40 @@
+import scipy as sp
+import cupy as cp
+import cupyx.scipy.fft as cufft
+
+from .base import FFTFilter
+
+
+class FFTFilterCupy(FFTFilter):
+    """Wraps the cupy Fourier transform and uses it via the scipy backend
+    """
+    is_available = True
+    # sp.fft.set_backend(cufft)
+
+    def _init_fft(self, data):
+        """Perform initial Fourier transform of the input data
+
+        Parameters
+        ----------
+        data: 2d real-valued np.ndarray
+            Input field to be refocused
+
+        Returns
+        -------
+        fft_fdata: 2d complex-valued ndarray
+            Fourier transform `data`
+        """
+        data_gpu = cp.asarray(data)
+        # likely an inefficiency here, could use `set_global_backend`
+        with sp.fft.set_backend(cufft):
+            fft_gpu = sp.fft.fft2(data_gpu)
+        fft_cpu = fft_gpu.get()
+        return fft_cpu
+
+    def _ifft(self, data):
+        """Perform inverse Fourier transform"""
+        data_gpu = cp.asarray(data)
+        with sp.fft.set_backend(cufft):
+            ifft_gpu = sp.fft.ifft2(data_gpu)
+        ifft_cpu = ifft_gpu.get()
+        return ifft_cpu

qpretrieve/fourier/ff_scipy.py

@@ -0,0 +1,30 @@
+import scipy as sp
+
+
+from .base import FFTFilter
+
+
+class FFTFilterScipy(FFTFilter):
+    """Wraps the scipy Fourier transform
+    """
+    # always available, because scipy is a dependency
+    is_available = True
+
+    def _init_fft(self, data):
+        """Perform initial Fourier transform of the input data
+
+        Parameters
+        ----------
+        data: 2d real-valued np.ndarray
+            Input field to be refocused
+
+        Returns
+        -------
+        fft_fdata: 2d complex-valued ndarray
+            Fourier transform `data`
+        """
+        return sp.fft.fft2(data)
+
+    def _ifft(self, data):
+        """Perform inverse Fourier transform"""
+        return sp.fft.ifft2(data)

qpretrieve/interfere/base.py

@@ -2,7 +2,8 @@
 
 import numpy as np
 
-from ..fourier import get_best_interface
+from ..fourier import get_best_interface, get_available_interfaces
+from ..fourier.base import FFTFilter
 
 
 class BaseInterferogram(ABC):
@@ -15,11 +16,19 @@ class BaseInterferogram(ABC):
         "invert_phase": False,
     }
 
-    def __init__(self, data, subtract_mean=True, padding=2, copy=True,
+    def __init__(self, data, fft_interface: FFTFilter = None,
+                 subtract_mean=True, padding=2, copy=True,
                  **pipeline_kws):
         """
         Parameters
         ----------
+        fft_interface: FFTFilter
+            A Fourier transform interface.
+            See :func:`qpretrieve.fourier.get_available_interfaces`
+            to get a list of implemented interfaces.
+            Default is None, which will use
+            :func:`qpretrieve.fourier.get_best_interface`. This is in line
+            with old behaviour.
         subtract_mean: bool
             If True, remove the mean of the hologram before performing
             the Fourier transform. This setting is recommended as it
@@ -38,15 +47,26 @@ def __init__(self, data, subtract_mean=True, padding=2, copy=True,
             Any additional keyword arguments for :func:`run_pipeline`
             as defined in :const:`default_pipeline_kws`.
         """
-        ff_iface = get_best_interface()
+        if fft_interface == 'auto' or fft_interface is None:
+            self.ff_iface = get_best_interface()
+        else:
+            if fft_interface in get_available_interfaces():
+                self.ff_iface = fft_interface
+            else:
+                raise ValueError(
+                    f"User-chosen FFT Interface '{fft_interface}' is not "
+                    f"available. The available interfaces are: "
+                    f"{get_available_interfaces()}.\n"
+                    f"You can use `fft_interface='auto'` to get the best "
+                    f"available interface.")
         if len(data.shape) == 3:
             # take the first slice (we have alpha or RGB information)
             data = data[:, :, 0]
         #: qpretrieve Fourier transform interface class
-        self.fft = ff_iface(data=data,
-                            subtract_mean=subtract_mean,
-                            padding=padding,
-                            copy=copy)
+        self.fft = self.ff_iface(data=data,
+                                 subtract_mean=subtract_mean,
+                                 padding=padding,
+                                 copy=copy)
         #: originally computed Fourier transform
         self.fft_origin = self.fft.fft_origin
         #: filtered Fourier data from last run of `run_pipeline`
@@ -94,7 +114,7 @@ def compute_filter_size(self, filter_size, filter_size_interpretation,
                 raise ValueError("For sideband distance interpretation, "
                                  "`filter_size` must be between 0 and 1; "
                                  f"got '{filter_size}'!")
-            fsize = np.sqrt(np.sum(np.array(sideband_freq)**2)) * filter_size
+            fsize = np.sqrt(np.sum(np.array(sideband_freq) ** 2)) * filter_size
         elif filter_size_interpretation == "frequency index":
             # filter size given in Fourier index (number of Fourier pixels)
             # The user probably does not know that we are padding in

setup.py

@@ -2,7 +2,6 @@
 from setuptools import setup, find_packages
 import sys
 
-
 author = u"Paul Müller"
 authors = [author]
 description = 'library for phase retrieval from holograms'
@@ -27,8 +26,13 @@
         "numpy>=1.9.0",
         "scikit-image>=0.11.0",
         "scipy>=0.18.0",
-        ],
-    extras_require={"FFTW": "pyfftw>=0.12.0"},
+    ],
+    extras_require={
+        "FFTW": "pyfftw>=0.12.0",
+        # manually install 'cupy-cuda11x' if you have older CUDA.
+        # See https://cupy.dev/
+        "CUPY": "cupy-cuda12x",
+    },
     python_requires='>=3.10, <4',
     keywords=["digital holographic microscopy",
               "optics",
@@ -41,6 +45,6 @@
         'Operating System :: OS Independent',
         'Programming Language :: Python :: 3',
         'Intended Audience :: Science/Research'
-                 ],
+    ],
     platforms=['ALL'],
-    )
+)