Adjust to multiple lines

Author: junchaoxia

src/aspire/image/__init__.py

@@ -204,13 +204,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(
-            xp.asnumpy(
-                fft.centered_ifft2(
-                    fft.centered_fft2(xp.asarray(self.data)) * xp.asarray(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(
@@ -274,16 +271,10 @@ def _im_translate(self, shifts):
 
         L = self.res
         im_f = xp.asnumpy(fft.fft2(xp.asarray(im)))
-        grid_1d = (
-            xp.asnumpy(
-                fft.ifftshift(
-                    xp.asarray(np.ceil(np.arange(-L / 2, L / 2, dtype=self.dtype)))
-                )
-            )
-            * 2
-            * np.pi
-            / L
+        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")
 

src/aspire/utils/preprocess.py

@@ -148,13 +148,11 @@ def downsample(insamples, szout, mask=None):
         # stack of one dimension objects
 
         for idata in range(ndata):
-            imsample_shifted = fft.fftshift(fft.fft(xp.asarray(insamples[idata])))
-            insamples_fft = crop_pad(imsample_shifted, L_out) * mask
+            insamples_shifted = fft.fftshift(fft.fft(xp.asarray(insamples[idata])))
+            insamples_fft = crop_pad(insamples_shifted, L_out) * mask
 
-            outsammples_shifted = fft.ifft(fft.ifftshift(xp.asarray(insamples_fft)))
-            outsamples[idata] = np.real(
-                xp.asnumpy(outsammples_shifted) * (L_out / L_in)
-            )
+            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