Speedup of autocenter

CHANGELOG.rst

@@ -1,6 +1,13 @@
 Changelog
 =========
 
+.. currentmodule:: skued
+
+Release 2.1.2
+-------------
+
+* Speedup of :func:`autocenter`, :func:`align`, :func:`ialign`, and :func:`itrack_peak` by 50%.
+
 Release 2.1.1
 -------------
 

requirements.txt

@@ -3,6 +3,6 @@ npstreams    >= 1.6.4, < 2
 numpy        >= 1.17, < 2
 pywavelets   >= 1.0.0
 scikit-image >= 0.17, <1
-scipy        >= 1.0.0
+scipy        >= 1.5.0
 pyyaml       >= 3.1
 matplotlib   >= 3.0.0, <4

skued/image/alignment.py

@@ -5,10 +5,12 @@
 """
 from functools import partial
 
+from os import cpu_count
 import numpy as np
 from warnings import warn
 from npstreams import array_stream
 from scipy import ndimage as ndi
+from scipy.fft import set_workers
 from skimage.registration import phase_cross_correlation
 
 
@@ -50,15 +52,20 @@ def itrack_peak(images, row_slice=None, col_slice=None, precision=1 / 10):
     ref = np.array(first[row_slice, col_slice], copy=True)
     sub = np.empty_like(ref)
 
-    for image in images:
-        sub[:] = image[row_slice, col_slice]
-        shift = phase_cross_correlation(
-            reference_image=ref,
-            moving_image=sub,
-            return_error=False,
-            upsample_factor=int(1 / precision),
-        )
-        yield np.asarray(shift)
+    # scikit-image will use scipy.fft module
+    # so we can increase the number of FFT workers
+    # to get a performance speedup (50% in my tests)
+    with set_workers(cpu_count()):
+        for image in images:
+            sub[:] = image[row_slice, col_slice]
+
+            shift = phase_cross_correlation(
+                reference_image=ref,
+                moving_image=sub,
+                return_error=False,
+                upsample_factor=int(1 / precision),
+            )
+            yield np.asarray(shift)
 
 
 def align(image, reference, mask=None, fill_value=0.0):
@@ -85,12 +92,16 @@ def align(image, reference, mask=None, fill_value=0.0):
     --------
     ialign : generator of aligned images
     """
-    shift = phase_cross_correlation(
-        reference_image=reference,
-        moving_image=image,
-        reference_mask=mask,
-        return_error=False,
-    )
+    # scikit-image will use scipy.fft module
+    # so we can increase the number of FFT workers
+    # to get a performance speedup (50% in my tests)
+    with set_workers(cpu_count()):
+        shift = phase_cross_correlation(
+            reference_image=reference,
+            moving_image=image,
+            reference_mask=mask,
+            return_error=False,
+        )
     return ndi.shift(image, shift=shift, order=2, mode="constant", cval=fill_value)
 
 

skued/image/center.py

@@ -3,10 +3,11 @@
 Determine the center of diffraction images
 ==========================================
 """
-
+from os import cpu_count
 import numpy as np
 from skimage.registration import phase_cross_correlation
 from scipy.ndimage import shift
+from scipy.fft import set_workers
 
 
 def autocenter(im, mask=None):
@@ -53,12 +54,16 @@ def autocenter(im, mask=None):
     im_i = _fast_radial_inversion(im, center=(r_rough, c_rough), cval=0.0)
     mask_i = _fast_radial_inversion(mask, center=(r_rough, c_rough), cval=False)
 
-    shift = phase_cross_correlation(
-        reference_image=im,
-        moving_image=im_i,
-        reference_mask=mask,
-        moving_mask=mask_i,
-    )
+    # scikit-image will use scipy.fft module
+    # so we can increase the number of FFT workers
+    # to get a performance speedup (50% in my tests)
+    with set_workers(cpu_count()):
+        shift = phase_cross_correlation(
+            reference_image=im,
+            moving_image=im_i,
+            reference_mask=mask,
+            moving_mask=mask_i,
+        )
 
     return np.array([r_rough, c_rough]) + shift / 2 - np.array([1 / 2, 1 / 2])