From 175ec7d85e0da44798769c5ab4d33d484dca7993 Mon Sep 17 00:00:00 2001 From: Georgios Varnavides Date: Tue, 8 Oct 2024 16:32:01 -0700 Subject: [PATCH] minor shape bug --- py4DSTEM/process/phase/direct_ptychography.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/py4DSTEM/process/phase/direct_ptychography.py b/py4DSTEM/process/phase/direct_ptychography.py index 56fa11a25..0eab30df9 100644 --- a/py4DSTEM/process/phase/direct_ptychography.py +++ b/py4DSTEM/process/phase/direct_ptychography.py @@ -499,14 +499,15 @@ def preprocess( del _intensities self._intensities = copy_to_device(self._intensities, storage) - self._intensities_shape = np.array(self._intensities.shape[-2:]) + self._intensities_shape = self._intensities.shape[-2:] # take FFT wrt real-space if vectorized_com_calculation: self._intensities_FFT = xp_storage.fft.fft2(self._intensities, axes=(0, 1)) else: - self._intensities_FFT = xp_storage.empty_like( - self._intensities, dtype=xp_storage.complex64 + self._intensities_FFT = xp_storage.empty( + self._intensities_shape + self._grid_scan_shape, + dtype=xp_storage.complex64, ) # transpose to loop over cols faster intensities = self._intensities.transpose((2, 3, 0, 1))