Merge pull request py4dstem#674 from py4dstem/phase_contrast

Phase contrast quality of life improvements

py4DSTEM/process/phase/magnetic_ptychographic_tomography.py

@@ -225,6 +225,7 @@ def preprocess(
         padded_diffraction_intensities_shape: Tuple[int, int] = None,
         region_of_interest_shape: Tuple[int, int] = None,
         dp_mask: np.ndarray = None,
+        in_place_datacube_modification: bool = False,
         fit_function: str = "plane",
         plot_probe_overlaps: bool = True,
         rotation_real_space_degrees: float = None,
@@ -266,6 +267,9 @@ def preprocess(
             at the diffraction plane to allow comparison with experimental data
         dp_mask: ndarray, optional
             Mask for datacube intensities (Qx,Qy)
+        in_place_datacube_modification: bool, optional
+            If True, the datacube will be preprocessed in-place. Note this is not possible
+            when either crop_patterns or positions_mask are used.
         fit_function: str, optional
             2D fitting function for CoM fitting. One of 'plane','parabola','bezier_two'
         plot_probe_overlaps: bool, optional
@@ -479,12 +483,14 @@ def preprocess(
                 amplitudes,
                 mean_diffraction_intensity_temp,
                 self._crop_mask,
+                self._crop_mask_shape,
             ) = self._normalize_diffraction_intensities(
                 intensities,
                 com_fitted_x,
                 com_fitted_y,
                 self._positions_mask[index],
                 crop_patterns,
+                in_place_datacube_modification,
             )
 
             self._mean_diffraction_intensity.append(mean_diffraction_intensity_temp)

py4DSTEM/process/phase/magnetic_ptychography.py

@@ -208,6 +208,7 @@ def preprocess(
         padded_diffraction_intensities_shape: Tuple[int, int] = None,
         region_of_interest_shape: Tuple[int, int] = None,
         dp_mask: np.ndarray = None,
+        in_place_datacube_modification: bool = False,
         fit_function: str = "plane",
         plot_rotation: bool = True,
         maximize_divergence: bool = False,
@@ -259,6 +260,9 @@ def preprocess(
             at the diffraction plane to allow comparison with experimental data
         dp_mask: ndarray, optional
             Mask for datacube intensities (Qx,Qy)
+        in_place_datacube_modification: bool, optional
+            If True, the datacube will be preprocessed in-place. Note this is not possible
+            when either crop_patterns or positions_mask are used.
         fit_function: str, optional
             2D fitting function for CoM fitting. One of 'plane','parabola','bezier_two'
         plot_rotation: bool, optional
@@ -373,10 +377,6 @@ def preprocess(
                 f"datacube must be the same length as magnetic_contribution_sign, not length {len(self._datacube)}."
             )
 
-        dc_shapes = [dc.shape for dc in self._datacube]
-        if dc_shapes.count(dc_shapes[0]) != self._num_measurements:
-            raise ValueError("datacube intensities must be the same size.")
-
         if self._positions_mask is not None:
             self._positions_mask = np.asarray(self._positions_mask, dtype="bool")
 
@@ -551,12 +551,14 @@ def preprocess(
                 amplitudes,
                 mean_diffraction_intensity_temp,
                 self._crop_mask,
+                self._crop_mask_shape,
             ) = self._normalize_diffraction_intensities(
                 intensities,
                 com_fitted_x,
                 com_fitted_y,
                 self._positions_mask[index],
                 crop_patterns,
+                in_place_datacube_modification,
             )
 
             self._mean_diffraction_intensity.append(mean_diffraction_intensity_temp)

py4DSTEM/process/phase/mixedstate_multislice_ptychography.py

@@ -267,6 +267,7 @@ def preprocess(
         padded_diffraction_intensities_shape: Tuple[int, int] = None,
         region_of_interest_shape: Tuple[int, int] = None,
         dp_mask: np.ndarray = None,
+        in_place_datacube_modification: bool = False,
         fit_function: str = "plane",
         plot_center_of_mass: str = "default",
         plot_rotation: bool = True,
@@ -318,6 +319,9 @@ def preprocess(
             at the diffraction plane to allow comparison with experimental data
         dp_mask: ndarray, optional
             Mask for datacube intensities (Qx,Qy)
+        in_place_datacube_modification: bool, optional
+            If True, the datacube will be preprocessed in-place. Note this is not possible
+            when either crop_patterns or positions_mask are used.
         fit_function: str, optional
             2D fitting function for CoM fitting. One of 'plane','parabola','bezier_two'
         plot_center_of_mass: str, optional
@@ -486,17 +490,21 @@ def preprocess(
             self._amplitudes,
             self._mean_diffraction_intensity,
             self._crop_mask,
+            self._crop_mask_shape,
         ) = self._normalize_diffraction_intensities(
             _intensities,
             self._com_fitted_x,
             self._com_fitted_y,
             self._positions_mask,
             crop_patterns,
+            in_place_datacube_modification,
         )
 
         # explicitly transfer arrays to storage
+        if not in_place_datacube_modification:
+            del _intensities
+
         self._amplitudes = copy_to_device(self._amplitudes, storage)
-        del _intensities
 
         self._num_diffraction_patterns = self._amplitudes.shape[0]
         self._amplitudes_shape = np.array(self._amplitudes.shape[-2:])

py4DSTEM/process/phase/mixedstate_ptychography.py

@@ -213,6 +213,7 @@ def preprocess(
         padded_diffraction_intensities_shape: Tuple[int, int] = None,
         region_of_interest_shape: Tuple[int, int] = None,
         dp_mask: np.ndarray = None,
+        in_place_datacube_modification: bool = False,
         fit_function: str = "plane",
         plot_center_of_mass: str = "default",
         plot_rotation: bool = True,
@@ -264,6 +265,9 @@ def preprocess(
             at the diffraction plane to allow comparison with experimental data
         dp_mask: ndarray, optional
             Mask for datacube intensities (Qx,Qy)
+        in_place_datacube_modification: bool, optional
+            If True, the datacube will be preprocessed in-place. Note this is not possible
+            when either crop_patterns or positions_mask are used.
         fit_function: str, optional
             2D fitting function for CoM fitting. One of 'plane','parabola','bezier_two'
         plot_center_of_mass: str, optional
@@ -432,17 +436,21 @@ def preprocess(
             self._amplitudes,
             self._mean_diffraction_intensity,
             self._crop_mask,
+            self._crop_mask_shape,
         ) = self._normalize_diffraction_intensities(
             _intensities,
             self._com_fitted_x,
             self._com_fitted_y,
             self._positions_mask,
             crop_patterns,
+            in_place_datacube_modification,
         )
 
         # explicitly transfer arrays to storage
+        if not in_place_datacube_modification:
+            del _intensities
+
         self._amplitudes = copy_to_device(self._amplitudes, storage)
-        del _intensities
 
         self._num_diffraction_patterns = self._amplitudes.shape[0]
         self._amplitudes_shape = np.array(self._amplitudes.shape[-2:])

py4DSTEM/process/phase/multislice_ptychography.py

@@ -241,6 +241,7 @@ def preprocess(
         padded_diffraction_intensities_shape: Tuple[int, int] = None,
         region_of_interest_shape: Tuple[int, int] = None,
         dp_mask: np.ndarray = None,
+        in_place_datacube_modification: bool = False,
         fit_function: str = "plane",
         plot_center_of_mass: str = "default",
         plot_rotation: bool = True,
@@ -292,6 +293,9 @@ def preprocess(
             at the diffraction plane to allow comparison with experimental data
         dp_mask: ndarray, optional
             Mask for datacube intensities (Qx,Qy)
+        in_place_datacube_modification: bool, optional
+            If True, the datacube will be preprocessed in-place. Note this is not possible
+            when either crop_patterns or positions_mask are used.
         fit_function: str, optional
             2D fitting function for CoM fitting. One of 'plane','parabola','bezier_two'
         plot_center_of_mass: str, optional
@@ -460,17 +464,21 @@ def preprocess(
             self._amplitudes,
             self._mean_diffraction_intensity,
             self._crop_mask,
+            self._crop_mask_shape,
         ) = self._normalize_diffraction_intensities(
             _intensities,
             self._com_fitted_x,
             self._com_fitted_y,
             self._positions_mask,
             crop_patterns,
+            in_place_datacube_modification,
         )
 
         # explicitly transfer arrays to storage
+        if not in_place_datacube_modification:
+            del _intensities
+
         self._amplitudes = copy_to_device(self._amplitudes, storage)
-        del _intensities
 
         self._num_diffraction_patterns = self._amplitudes.shape[0]
         self._amplitudes_shape = np.array(self._amplitudes.shape[-2:])