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:
@@ -39,7 +39,8 @@ jobs:
         pip freeze
     - name: Test with pytest
       run: |
-        coverage run --source=qpretrieve -m pytest tests
+        # ignore the cupy imports, as we don't have a gpu-enabled pipeline setup
+        coverage run --source=qpretrieve -m pytest tests --ignore=tests/test_cupy_gpu
     - name: Lint with flake8
       run: |
         flake8 .

docs/README.md

@@ -6,6 +6,7 @@ To install the requirements for building the documentation, run
 
 To compile the documentation, run
 
+    cd docs
     sphinx-build . _build
 
 

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,7 @@ Examples
 .. fancy_include:: filter_visualization.py
 
 .. fancy_include:: fourier_scale.py
+
+.. fancy_include:: fft_options.py
+
+.. fancy_include:: fft_cupy3d_speed.py

examples/fft_cupy3d_speed.py

@@ -0,0 +1,66 @@
+"""Fourier Transform speeds for the Cupy 3D interface
+
+This example visualizes the speed for different batch sizes for
+the `FFTFilterCupy3D` FFT Filter. The y-axis shows the speed of a single
+FFT for the corresponding batch size.
+
+- Optimum batch size is between 64 and 256 for 256x256pix imgs (incl padding).
+- Here, batch size is the size of the 3D stack in z.
+
+"""
+import time
+import matplotlib.pylab as plt
+import numpy as np
+import qpretrieve
+
+# load the experimental data
+edata = np.load("./data/hologram_cell.npz")
+
+n_transforms_list = [8, 16, 32, 64, 128, 256, 512]
+subtract_mean = True
+padding = True
+fft_interface = qpretrieve.fourier.FFTFilterCupy3D
+
+results = {}
+for n_transforms in n_transforms_list:
+    print(f"Running {n_transforms} transforms...")
+
+    data_2d = edata["data"].copy()
+    data_2d_bg = edata["bg_data"].copy()
+    data_3d = np.repeat(
+        edata["data"].copy()[np.newaxis, ...],
+        repeats=n_transforms, axis=0)
+    data_3d_bg = np.repeat(
+        edata["bg_data"].copy()[np.newaxis, ...],
+        repeats=n_transforms, axis=0)
+    assert data_3d.shape == data_3d_bg.shape == (n_transforms,
+                                                 edata["data"].shape[0],
+                                                 edata["data"].shape[1])
+
+    t0 = time.time()
+
+    holo = qpretrieve.OffAxisHologram(data=data_3d,
+                                      fft_interface=fft_interface,
+                                      subtract_mean=subtract_mean,
+                                      padding=padding)
+    holo.run_pipeline(filter_name="disk", filter_size=1 / 2)
+    bg = qpretrieve.OffAxisHologram(data=data_3d_bg)
+    bg.process_like(holo)
+
+    t1 = time.time()
+    results[n_transforms] = t1 - t0
+
+speed_norm = [v / k for k, v in results.items()]
+
+fig, axes = plt.subplots(1, 1, figsize=(8, 5))
+ax1 = axes
+
+ax1.bar(range(len(n_transforms_list)), height=speed_norm, color='darkmagenta')
+ax1.set_xticks(range(len(n_transforms_list)), labels=n_transforms_list)
+ax1.set_xlabel("Fourier transform batch size")
+ax1.set_ylabel("Speed / batch size (s)")
+
+plt.suptitle("Speed of FFT Interface CuPy3D")
+plt.tight_layout()
+plt.show()
+# plt.savefig("fft_cupy3d_speed.png", dpi=150)

examples/fft_options.py

@@ -0,0 +1,109 @@
+"""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 3D should not be used for direct comparison here,
+  as it is not doing post-processing padding.
+
+"""
+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 = 200
+subtract_mean = True
+padding = True
+
+print("Running single transform...")
+# one transform
+results_1 = {}
+for fft_interface in interfaces_available:
+    t0 = time.time()
+    holo = qpretrieve.OffAxisHologram(data=edata["data"],
+                                      fft_interface=fft_interface,
+                                      subtract_mean=subtract_mean,
+                                      padding=padding)
+    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)
+print(f"Running {n_transforms} transforms...")
+results = {}
+for fft_interface in interfaces_available:
+
+    data_2d = edata["data"].copy()
+    data_2d_bg = edata["bg_data"].copy()
+
+    data_3d = np.repeat(
+        edata["data"].copy()[np.newaxis, ...],
+        repeats=n_transforms, axis=0)
+    data_3d_bg = np.repeat(
+        edata["bg_data"].copy()[np.newaxis, ...],
+        repeats=n_transforms, axis=0)
+    assert data_3d.shape == data_3d_bg.shape == (n_transforms,
+                                                 edata["data"].shape[0],
+                                                 edata["data"].shape[1])
+
+    t0 = time.time()
+
+    if fft_interface.__name__ == "FFTFilterCupy3D":
+        holo = qpretrieve.OffAxisHologram(data=data_3d,
+                                          fft_interface=fft_interface,
+                                          subtract_mean=subtract_mean,
+                                          padding=padding)
+        holo.run_pipeline(filter_name="disk", filter_size=1 / 2)
+        bg = qpretrieve.OffAxisHologram(data=data_3d_bg)
+        bg.process_like(holo)
+    else:
+        # 2d
+        for _ in range(n_transforms):
+            holo = qpretrieve.OffAxisHologram(
+                data=data_2d,
+                fft_interface=fft_interface,
+                subtract_mean=subtract_mean, padding=padding)
+            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)")
+# todo: fix code, then this title
+ax2.set_title(f"{n_transforms} Transforms\n(**Cupy comparison not valid)")
+
+plt.suptitle("Speed of FFT Interfaces")
+plt.tight_layout()
+# plt.show()
+plt.savefig("fft_options.png", dpi=150)

examples/requirements.txt

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

qpretrieve/filter.py

@@ -38,9 +38,9 @@ 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
+        The shape of the Fourier transformed image (2d) for which the
         filter will be applied. The shape must be squared (two
         identical integers).
 
@@ -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,42 @@
 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
+
+try:
+    from .ff_cupy3D import FFTFilterCupy3D
+except ImportError:
+    FFTFilterCupy3D = None
+
 PREFERRED_INTERFACE = None
 
 
+def get_available_interfaces():
+    """Return a list of available FFT algorithms"""
+    interfaces = [
+        FFTFilterPyFFTW,
+        FFTFilterNumpy,
+        FFTFilterScipy,
+        FFTFilterCupy,
+        FFTFilterCupy3D,
+    ]
+    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