Merge pull request py4dstem#652 from py4dstem/phase_contrast

Show me your bugs, and git will tell you who to blame

py4DSTEM/process/phase/magnetic_ptychographic_tomography.py

@@ -239,6 +239,7 @@ def preprocess(
         force_reciprocal_sampling: float = None,
         object_fov_mask: np.ndarray = True,
         crop_patterns: bool = False,
+        store_initial_arrays: bool = True,
         device: str = None,
         clear_fft_cache: bool = None,
         max_batch_size: int = None,
@@ -295,6 +296,8 @@ def preprocess(
             If None, probe_overlap intensity is thresholded
         crop_patterns: bool
             if True, crop patterns to avoid wrap around of patterns when centering
+        store_initial_arrays: bool
+            If True, preprocesed object and probe arrays are stored allowing reset=True in reconstruct.
         device: str, optional
             if not none, overwrites self._device to set device preprocess will be perfomed on.
         clear_fft_cache: bool, optional
@@ -534,8 +537,10 @@ def preprocess(
         else:
             self._object = obj
 
-        self._object_initial = self._object.copy()
-        self._object_type_initial = self._object_type
+        if store_initial_arrays:
+            self._object_initial = self._object.copy()
+            self._object_type_initial = self._object_type
+
         self._object_shape = self._object.shape[-2:]
         self._num_voxels = self._object.shape[1]
 
@@ -565,10 +570,14 @@ def preprocess(
 
         # initialize probe
         self._probes_all = []
-        self._probes_all_initial = []
-        self._probes_all_initial_aperture = []
         list_Q = isinstance(self._probe_init, (list, tuple))
 
+        if store_initial_arrays:
+            self._probes_all_initial = []
+            self._probes_all_initial_aperture = []
+        else:
+            self._probes_all_initial_aperture = [None] * self._num_measurements
+
         for index in range(self._num_measurements):
             _probe, self._semiangle_cutoff = self._initialize_probe(
                 self._probe_init[index] if list_Q else self._probe_init,
@@ -579,8 +588,9 @@ def preprocess(
             )
 
             self._probes_all.append(_probe)
-            self._probes_all_initial.append(_probe.copy())
-            self._probes_all_initial_aperture.append(xp.abs(xp.fft.fft2(_probe)))
+            if store_initial_arrays:
+                self._probes_all_initial.append(_probe.copy())
+                self._probes_all_initial_aperture.append(xp.abs(xp.fft.fft2(_probe)))
 
         del self._probe_init
 

py4DSTEM/process/phase/magnetic_ptychography.py

@@ -219,6 +219,7 @@ def preprocess(
         progress_bar: bool = True,
         object_fov_mask: np.ndarray = True,
         crop_patterns: bool = False,
+        store_initial_arrays: bool = True,
         device: str = None,
         clear_fft_cache: bool = None,
         max_batch_size: int = None,
@@ -285,6 +286,8 @@ def preprocess(
             If None, probe_overlap intensity is thresholded
         crop_patterns: bool
             if True, crop patterns to avoid wrap around of patterns when centering
+        store_initial_arrays: bool
+            If True, preprocesed object and probe arrays are stored allowing reset=True in reconstruct.
         device: str, optional
             if not none, overwrites self._device to set device preprocess will be perfomed on.
         clear_fft_cache: bool, optional
@@ -597,8 +600,9 @@ def preprocess(
         else:
             self._object = obj
 
-        self._object_initial = self._object.copy()
-        self._object_type_initial = self._object_type
+        if store_initial_arrays:
+            self._object_initial = self._object.copy()
+            self._object_type_initial = self._object_type
         self._object_shape = self._object.shape[-2:]
 
         # center probe positions
@@ -627,10 +631,14 @@ def preprocess(
 
         # initialize probe
         self._probes_all = []
-        self._probes_all_initial = []
-        self._probes_all_initial_aperture = []
         list_Q = isinstance(self._probe_init, (list, tuple))
 
+        if store_initial_arrays:
+            self._probes_all_initial = []
+            self._probes_all_initial_aperture = []
+        else:
+            self._probes_all_initial_aperture = [None] * self._num_measurements
+
         for index in range(self._num_measurements):
             _probe, self._semiangle_cutoff = self._initialize_probe(
                 self._probe_init[index] if list_Q else self._probe_init,
@@ -641,8 +649,9 @@ def preprocess(
             )
 
             self._probes_all.append(_probe)
-            self._probes_all_initial.append(_probe.copy())
-            self._probes_all_initial_aperture.append(xp.abs(xp.fft.fft2(_probe)))
+            if store_initial_arrays:
+                self._probes_all_initial.append(_probe.copy())
+                self._probes_all_initial_aperture.append(xp.abs(xp.fft.fft2(_probe)))
 
         del self._probe_init
 

py4DSTEM/process/phase/mixedstate_multislice_ptychography.py

@@ -278,6 +278,7 @@ def preprocess(
         force_reciprocal_sampling: float = None,
         object_fov_mask: np.ndarray = None,
         crop_patterns: bool = False,
+        store_initial_arrays: bool = True,
         device: str = None,
         clear_fft_cache: bool = None,
         max_batch_size: int = None,
@@ -347,6 +348,8 @@ def preprocess(
             If None, probe_overlap intensity is thresholded
         crop_patterns: bool
             if True, crop patterns to avoid wrap around of patterns when centering
+        store_initial_arrays: bool
+            If True, preprocesed object and probe arrays are stored allowing reset=True in reconstruct.
         device: str, optional
             If not None, overwrites self._device to set device preprocess will be perfomed on.
         clear_fft_cache: bool, optional
@@ -516,8 +519,9 @@ def preprocess(
             self._object_type,
         )
 
-        self._object_initial = self._object.copy()
-        self._object_type_initial = self._object_type
+        if store_initial_arrays:
+            self._object_initial = self._object.copy()
+            self._object_type_initial = self._object_type
         self._object_shape = self._object.shape[-2:]
 
         # center probe positions
@@ -553,8 +557,11 @@ def preprocess(
             device=self._device,
         )._evaluate_ctf()
 
-        self._probe_initial = self._probe.copy()
-        self._probe_initial_aperture = xp.abs(xp.fft.fft2(self._probe))
+        if store_initial_arrays:
+            self._probe_initial = self._probe.copy()
+            self._probe_initial_aperture = xp.abs(xp.fft.fft2(self._probe))
+        else:
+            self._probe_initial_aperture = None
 
         # precompute propagator arrays
         self._propagator_arrays = self._precompute_propagator_arrays(

py4DSTEM/process/phase/mixedstate_ptychography.py

@@ -224,6 +224,7 @@ def preprocess(
         force_reciprocal_sampling: float = None,
         object_fov_mask: np.ndarray = None,
         crop_patterns: bool = False,
+        store_initial_arrays: bool = True,
         device: str = None,
         clear_fft_cache: bool = None,
         max_batch_size: int = None,
@@ -293,6 +294,8 @@ def preprocess(
             If None, probe_overlap intensity is thresholded
         crop_patterns: bool
             if True, crop patterns to avoid wrap around of patterns when centering
+        store_initial_arrays: bool
+            If True, preprocesed object and probe arrays are stored allowing reset=True in reconstruct.
         device: str, optional
             if not none, overwrites self._device to set device preprocess will be perfomed on.
         clear_fft_cache: bool, optional
@@ -461,8 +464,9 @@ def preprocess(
             self._object_type,
         )
 
-        self._object_initial = self._object.copy()
-        self._object_type_initial = self._object_type
+        if store_initial_arrays:
+            self._object_initial = self._object.copy()
+            self._object_type_initial = self._object_type
         self._object_shape = self._object.shape
 
         # center probe positions
@@ -498,8 +502,11 @@ def preprocess(
             device=self._device,
         )._evaluate_ctf()
 
-        self._probe_initial = self._probe.copy()
-        self._probe_initial_aperture = xp.abs(xp.fft.fft2(self._probe))
+        if store_initial_arrays:
+            self._probe_initial = self._probe.copy()
+            self._probe_initial_aperture = xp.abs(xp.fft.fft2(self._probe))
+        else:
+            self._probe_initial_aperture = None
 
         if object_fov_mask is None or plot_probe_overlaps:
             # overlaps

py4DSTEM/process/phase/multislice_ptychography.py

@@ -252,6 +252,7 @@ def preprocess(
         force_reciprocal_sampling: float = None,
         object_fov_mask: np.ndarray = None,
         crop_patterns: bool = False,
+        store_initial_arrays: bool = True,
         device: str = None,
         clear_fft_cache: bool = None,
         max_batch_size: int = None,
@@ -321,6 +322,8 @@ def preprocess(
             If None, probe_overlap intensity is thresholded
         crop_patterns: bool
             if True, crop patterns to avoid wrap around of patterns when centering
+        store_initial_arrays: bool
+            If True, preprocesed object and probe arrays are stored allowing reset=True in reconstruct.
         device: str, optional
             If not None, overwrites self._device to set device preprocess will be perfomed on.
         clear_fft_cache: bool, optional
@@ -490,8 +493,9 @@ def preprocess(
             self._object_type,
         )
 
-        self._object_initial = self._object.copy()
-        self._object_type_initial = self._object_type
+        if store_initial_arrays:
+            self._object_initial = self._object.copy()
+            self._object_type_initial = self._object_type
         self._object_shape = self._object.shape[-2:]
 
         # center probe positions
@@ -527,8 +531,11 @@ def preprocess(
             device=self._device,
         )._evaluate_ctf()
 
-        self._probe_initial = self._probe.copy()
-        self._probe_initial_aperture = xp.abs(xp.fft.fft2(self._probe))
+        if store_initial_arrays:
+            self._probe_initial = self._probe.copy()
+            self._probe_initial_aperture = xp.abs(xp.fft.fft2(self._probe))
+        else:
+            self._probe_initial_aperture = None
 
         # precompute propagator arrays
         self._propagator_arrays = self._precompute_propagator_arrays(