Merge pull request #298 from ComputationalCryoEM/replace_fft

includ centered fft functions in facade modules and replace FFT functions in other submodules

Author: junchaoxia

src/aspire/basis/ffb_2d.py

@@ -2,13 +2,13 @@
 
 import numpy as np
 from numpy import pi
-from scipy.fftpack import fft, ifft
 from scipy.special import jv
 
 from aspire.basis import FBBasis2D
 from aspire.basis.basis_utils import lgwt
 from aspire.image import Image
 from aspire.nufft import anufft, nufft
+from aspire.numeric import fft, xp
 from aspire.utils import complex_type
 from aspire.utils.matlab_compat import m_reshape
 
@@ -165,7 +165,7 @@ def evaluate(self, v):
             ind_pos = ind_pos + 2 * self.k_max[ell]
 
         # 1D inverse FFT in the degree of polar angle
-        pf = 2 * pi * ifft(pf, axis=1, overwrite_x=True)
+        pf = 2 * pi * xp.asnumpy(fft.ifft(xp.asarray(pf), axis=1))
 
         # Only need "positive" frequencies.
         hsize = int(np.size(pf, 1) / 2)
@@ -235,7 +235,7 @@ def evaluate_t(self, x):
             )
 
         #  1D FFT on the angular dimension for each concentric circle
-        pf = 2 * pi / (2 * n_theta) * fft(pf, 2 * n_theta, 2)
+        pf = 2 * pi / (2 * n_theta) * xp.asnumpy(fft.fft(xp.asarray(pf)))
 
         # This only makes it easier to slice the array later.
         v = np.zeros((n_images, self.count), dtype=x.dtype)

src/aspire/basis/fpswf_2d.py

@@ -3,13 +3,13 @@
 import numpy as np
 from numpy import pi
 from numpy.linalg import lstsq
-from scipy.fftpack import fft
 from scipy.optimize import least_squares
 from scipy.special import jn
 
 from aspire.basis.basis_utils import lgwt, t_x_mat, t_x_mat_dot
 from aspire.basis.pswf_2d import PSWFBasis2D
 from aspire.nufft import nufft
+from aspire.numeric import fft, xp
 from aspire.utils import complex_type
 
 logger = logging.getLogger(__name__)
@@ -400,7 +400,7 @@ def _pswf_integration(self, images_nufft):
                 :,
             ]
             curr_r_mat = np.concatenate((curr_r_mat, np.conj(curr_r_mat)))
-            fft_plan = fft(curr_r_mat, curr_r_mat.shape[0], axis=0)
+            fft_plan = xp.asnumpy(fft.fft(xp.asarray(curr_r_mat), axis=0))
             angular_eval = fft_plan * self.quad_rule_radial_wts[i]
 
             r_n_eval_mat[i, :, :] = np.tile(

src/aspire/image/image.py

@@ -2,16 +2,16 @@
 
 import mrcfile
 import numpy as np
-from scipy.fftpack import fft2, ifft2, ifftshift
 from scipy.interpolate import RegularGridInterpolator
 from scipy.linalg import lstsq
 
 import aspire.volume
 from aspire.nufft import anufft
-from aspire.utils import anorm, ensure
+from aspire.numeric import fft, xp
+from aspire.utils import ensure
 from aspire.utils.coor_trans import grid_2d
-from aspire.utils.fft import centered_fft2, centered_ifft2
 from aspire.utils.matlab_compat import m_reshape
+from aspire.utils.matrix import anorm
 
 logger = logging.getLogger(__name__)
 
@@ -46,7 +46,9 @@ def _im_translate2(im, shifts):
         raise ValueError("The number of shifts must be 1 or match the number of images")
 
     resolution = im.res
-    grid = np.fft.ifftshift(np.ceil(np.arange(-resolution / 2, resolution / 2)))
+    grid = xp.asnumpy(
+        fft.ifftshift(xp.asarray(np.ceil(np.arange(-resolution / 2, resolution / 2))))
+    )
     om_y, om_x = np.meshgrid(grid, grid)
     phase_shifts = np.einsum("ij, k -> ijk", om_x, shifts[:, 0]) + np.einsum(
         "ij, k -> ijk", om_y, shifts[:, 1]
@@ -56,9 +58,9 @@ def _im_translate2(im, shifts):
     phase_shifts /= resolution
 
     mult_f = np.exp(-2 * np.pi * 1j * phase_shifts)
-    im_f = np.fft.fft2(im.asnumpy())
+    im_f = xp.asnumpy(fft.fft2(xp.asarray(im.asnumpy())))
     im_translated_f = im_f * mult_f
-    im_translated = np.real(np.fft.ifft2(im_translated_f))
+    im_translated = np.real(xp.asnumpy(fft.ifft2(xp.asarray(im_translated_f))))
 
     return Image(im_translated)
 
@@ -200,7 +202,10 @@ def downsample(self, ds_res):
         mask = (np.abs(grid["x"]) < ds_res / self.res) & (
             np.abs(grid["y"]) < ds_res / self.res
         )
-        im = np.real(centered_ifft2(centered_fft2(self.data) * mask))
+        im_shifted = fft.centered_ifft2(
+            fft.centered_fft2(xp.asarray(self.data)) * xp.asarray(mask)
+        )
+        im = np.real(xp.asnumpy(im_shifted))
 
         for s in range(im_ds.shape[0]):
             interpolator = RegularGridInterpolator(
@@ -219,12 +224,13 @@ def filter(self, filter):
         """
         filter_values = filter.evaluate_grid(self.res)
 
-        im_f = centered_fft2(self.data)
+        im_f = xp.asnumpy(fft.centered_fft2(xp.asarray(self.data)))
+
         if im_f.ndim > filter_values.ndim:
             im_f *= filter_values
         else:
             im_f = filter_values * im_f
-        im = centered_ifft2(im_f)
+        im = xp.asnumpy(fft.centered_ifft2(xp.asarray(im_f)))
         im = np.real(im)
 
         return Image(im)
@@ -263,13 +269,11 @@ def _im_translate(self, shifts):
         shifts = shifts.astype(self.dtype)
 
         L = self.res
-        im_f = fft2(im, axes=(1, 2))
-        grid_1d = (
-            ifftshift(np.ceil(np.arange(-L / 2, L / 2, dtype=self.dtype)))
-            * 2
-            * np.pi
-            / L
+        im_f = xp.asnumpy(fft.fft2(xp.asarray(im)))
+        grid_shifted = fft.ifftshift(
+            xp.asarray(np.ceil(np.arange(-L / 2, L / 2, dtype=self.dtype)))
         )
+        grid_1d = xp.asnumpy(grid_shifted) * 2 * np.pi / L
         om_x, om_y = np.meshgrid(grid_1d, grid_1d, indexing="ij")
 
         phase_shifts_x = -shifts[:, 0].reshape((n_shifts, 1, 1))
@@ -281,7 +285,7 @@ def _im_translate(self, shifts):
         )
         mult_f = np.exp(-1j * phase_shifts)
         im_translated_f = im_f * mult_f
-        im_translated = ifft2(im_translated_f, axes=(1, 2))
+        im_translated = xp.asnumpy(fft.ifft2(xp.asarray(im_translated_f)))
         im_translated = np.real(im_translated)
 
         return Image(im_translated)
@@ -316,7 +320,7 @@ def backproject(self, rot_matrices):
         pts_rot = np.moveaxis(pts_rot, 1, 2)
         pts_rot = m_reshape(pts_rot, (3, -1))
 
-        im_f = centered_fft2(self.data) / (L ** 2)
+        im_f = xp.asnumpy(fft.centered_fft2(xp.asarray(self.data))) / (L ** 2)
         if L % 2 == 0:
             im_f[:, 0, :] = 0
             im_f[:, :, 0] = 0

src/aspire/image/preprocess.py

@@ -2,9 +2,9 @@
 import math
 
 import numpy as np
-from scipy.fftpack import fft, fft2, fftn, fftshift, ifft, ifft2, ifftn, ifftshift
 from scipy.special import erf
 
+from aspire.numeric import fft, xp
 from aspire.utils import ensure
 
 logger = logging.getLogger(__name__)
@@ -148,23 +148,36 @@ def downsample(insamples, szout, mask=None):
         # stack of one dimension objects
 
         for idata in range(ndata):
-            insamples_fft = crop_pad(fftshift(fft(insamples[idata])), L_out) * mask
-            outsamples[idata] = np.real(ifft(ifftshift(insamples_fft)) * (L_out / L_in))
+            insamples_shifted = fft.fftshift(fft.fft(xp.asarray(insamples[idata])))
+            insamples_fft = crop_pad(insamples_shifted, L_out) * mask
+
+            outsamples_shifted = fft.ifft(fft.ifftshift(xp.asarray(insamples_fft)))
+            outsamples[idata] = np.real(xp.asnumpy(outsamples_shifted) * (L_out / L_in))
 
     elif insamples.ndim == 3:
         # stack of two dimension objects
         for idata in range(ndata):
-            insamples_fft = crop_pad(fftshift(fft2(insamples[idata])), L_out) * mask
+            insamples_shifted = fft.fftshift(fft.fft2(xp.asarray(insamples[idata])))
+            insamples_fft = crop_pad(insamples_shifted, L_out) * mask
+
+            outsamples_shifted = fft.ifft2(fft.ifftshift(xp.asarray(insamples_fft)))
             outsamples[idata] = np.real(
-                ifft2(ifftshift(insamples_fft)) * (L_out ** 2 / L_in ** 2)
+                xp.asnumpy(outsamples_shifted) * (L_out ** 2 / L_in ** 2)
             )
 
     elif insamples.ndim == 4:
         # stack of three dimension objects
         for idata in range(ndata):
-            insamples_fft = crop_pad(fftshift(fftn(insamples[idata])), L_out) * mask
+            insamples_shifted = fft.fftshift(
+                fft.fftn(xp.asarray(insamples[idata]), axes=(0, 1, 2))
+            )
+            insamples_fft = crop_pad(insamples_shifted, L_out) * mask
+
+            outsamples_shifted = fft.ifftn(
+                fft.ifftshift(xp.asarray(insamples_fft)), axes=(0, 1, 2)
+            )
             outsamples[idata] = np.real(
-                ifftn(ifftshift(insamples_fft)) * (L_out ** 3 / L_in ** 3)
+                xp.asnumpy(outsamples_shifted) * (L_out ** 3 / L_in ** 3)
             )
 
     else:

src/aspire/noise/noise.py

@@ -2,9 +2,9 @@
 
 import numpy as np
 
+from aspire.numeric import fft, xp
 from aspire.operators import ArrayFilter, ScalarFilter
 from aspire.utils.coor_trans import grid_2d
-from aspire.utils.fft import centered_fft2
 
 logger = logging.getLogger(__name__)
 
@@ -122,7 +122,7 @@ def estimate_noise_psd(self):
 
             _denominator = self.n * np.sum(mask)
             mean_est += np.sum(images_masked) / _denominator
-            im_masked_f = centered_fft2(images_masked)
+            im_masked_f = xp.asnumpy(fft.centered_fft2(xp.asarray(images_masked)))
             noise_psd_est += np.sum(np.abs(im_masked_f ** 2), axis=0) / _denominator
 
         mid = self.L // 2

src/aspire/numeric/__init__.py

@@ -12,7 +12,7 @@
 xp = NumericClass()
 
 
-def fft_class(which):
+def fft_object(which):
     if which == "pyfftw":
         from .pyfftw_fft import PyfftwFFT as FFTClass
     elif which == "cupy":
@@ -21,7 +21,7 @@ def fft_class(which):
         from .scipy_fft import ScipyFFT as FFTClass
     else:
         raise RuntimeError(f"Invalid selection for fft class: {which}")
-    return FFTClass
+    return FFTClass()
 
 
-fft = fft_class(config.common.fft)
+fft = fft_object(config.common.fft)

src/aspire/numeric/base_fft.py

@@ -0,0 +1,67 @@
+class FFT:
+    """
+    Define a customized interface for FFT functions
+
+    To make consistent among Pyfftw, Scipyfft and cupy fft,
+    not all arguments are included.
+    """
+
+    def fft(self, x, axis=-1, workers=-1):
+        raise NotImplementedError("subclasses must implement this")
+
+    def ifft(self, x, axis=-1, workers=-1):
+        raise NotImplementedError("subclasses must implement this")
+
+    def fft2(self, x, axes=(-2, -1), workers=-1):
+        raise NotImplementedError("subclasses must implement this")
+
+    def ifft2(self, x, axes=(-2, -1), workers=-1):
+        raise NotImplementedError("subclasses must implement this")
+
+    def fftn(self, x, axes=None, workers=-1):
+        raise NotImplementedError("subclasses must implement this")
+
+    def ifftn(self, x, axes=None, workers=-1):
+        raise NotImplementedError("subclasses must implement this")
+
+    def fftshift(self, x, axes=None):
+        raise NotImplementedError("subclasses must implement this")
+
+    def ifftshift(self, x, axes=None):
+        raise NotImplementedError("subclasses must implement this")
+
+    def centered_ifft(self, x, axis=-1, workers=-1):
+        x = self.ifftshift(x, axes=axis)
+        x = self.ifft(x, axis=axis, workers=workers)
+        x = self.fftshift(x, axes=axis)
+        return x
+
+    def centered_fft(self, x, axis=-1, workers=-1):
+        x = self.ifftshift(x, axes=axis)
+        x = self.fft(x, axis=axis, workers=workers)
+        x = self.fftshift(x, axes=axis)
+        return x
+
+    def centered_ifft2(self, x, axes=(-2, -1), workers=-1):
+        x = self.ifftshift(x, axes=axes)
+        x = self.ifft2(x, axes=axes, workers=workers)
+        x = self.fftshift(x, axes=axes)
+        return x
+
+    def centered_fft2(self, x, axes=(-2, -1), workers=-1):
+        x = self.ifftshift(x, axes=axes)
+        x = self.fft2(x, axes=axes, workers=workers)
+        x = self.fftshift(x, axes=axes)
+        return x
+
+    def centered_ifftn(self, x, axes=None, workers=-1):
+        x = self.ifftshift(x, axes=axes)
+        x = self.ifftn(x, axes=axes, workers=workers)
+        x = self.fftshift(x, axes=axes)
+        return x
+
+    def centered_fftn(self, x, axes=None, workers=-1):
+        x = self.ifftshift(x, axes=axes)
+        x = self.fftn(x, axes=axes, workers=workers)
+        x = self.fftshift(x, axes=axes)
+        return x

src/aspire/numeric/cupy_fft.py

@@ -1,41 +1,35 @@
 import cupy as cp
 
+from aspire.numeric.base_fft import FFT
 
-class CupyFFT:
+
+class CupyFFT(FFT):
     """
     Define a unified wrapper class for Cupy FFT functions
 
     To be consistent with Scipy and Pyfftw, not all arguments are included.
     """
 
-    @staticmethod
-    def fft(x, axis=-1, workers=-1):
+    def fft(self, x, axis=-1, workers=-1):
         return cp.fft.fft(x, axis=axis)
 
-    @staticmethod
-    def ifft(x, axis=-1, workers=-1):
+    def ifft(self, x, axis=-1, workers=-1):
         return cp.fft.ifft(x, axis=axis)
 
-    @staticmethod
-    def fft2(x, axes=(-2, -1), workers=-1):
+    def fft2(self, x, axes=(-2, -1), workers=-1):
         return cp.fft.fft2(x, axes=axes)
 
-    @staticmethod
-    def ifft2(x, axes=(-2, -1), workers=-1):
+    def ifft2(self, x, axes=(-2, -1), workers=-1):
         return cp.fft.ifft2(x, axes=axes)
 
-    @staticmethod
-    def fftn(x, axes=None, workers=-1):
+    def fftn(self, x, axes=None, workers=-1):
         return cp.fft.fftn(x, axes=axes)
 
-    @staticmethod
-    def ifftn(x, axes=None, workers=-1):
-        return cp.fft.ifft2(x, axes=axes)
+    def ifftn(self, x, axes=None, workers=-1):
+        return cp.fft.ifftn(x, axes=axes)
 
-    @staticmethod
-    def fftshift(x, axes=None):
+    def fftshift(self, x, axes=None):
         return cp.fft.fftshift(x, axes=axes)
 
-    @staticmethod
-    def ifftshift(x, axes=None):
+    def ifftshift(self, x, axes=None):
         return cp.fft.ifftshift(x, axes=axes)