Merge pull request #1 from py4dstem/pr_683

Extending functionality to other ptycho classes

py4DSTEM/process/phase/magnetic_ptychographic_tomography.py

@@ -869,6 +869,7 @@ def reconstruct(
         shrinkage_rad: float = 0.0,
         fix_potential_baseline: bool = True,
         detector_fourier_mask: np.ndarray = None,
+        virtual_detector_masks: Sequence[np.ndarray] = None,
         tv_denoise: bool = True,
         tv_denoise_weights=None,
         tv_denoise_inner_iter=40,
@@ -984,6 +985,9 @@ def reconstruct(
         detector_fourier_mask: np.ndarray
             Corner-centered mask to multiply the detector-plane gradients with (a value of zero supresses those pixels).
             Useful when detector has artifacts such as dead-pixels. Usually binary.
+        virtual_detector_masks: np.ndarray
+            List of corner-centered boolean masks for binning forward model exit waves,
+            to allow comparison with arbitrary geometry detector datasets.
         store_iterations: bool, optional
             If True, reconstructed objects and probes are stored at each iteration
         progress_bar: bool, optional
@@ -1073,6 +1077,9 @@ def reconstruct(
         if detector_fourier_mask is not None:
             detector_fourier_mask = xp.asarray(detector_fourier_mask)
 
+        if virtual_detector_masks is not None:
+            virtual_detector_masks = xp.asarray(virtual_detector_masks)
+
         if gaussian_filter_sigma_m is None:
             gaussian_filter_sigma_m = gaussian_filter_sigma_e
 
@@ -1180,6 +1187,7 @@ def reconstruct(
                         amplitudes_device,
                         self._exit_waves,
                         detector_fourier_mask,
+                        virtual_detector_masks,
                         use_projection_scheme,
                         projection_a,
                         projection_b,

py4DSTEM/process/phase/magnetic_ptychography.py

@@ -1193,6 +1193,7 @@ def reconstruct(
         shrinkage_rad: float = 0.0,
         fix_potential_baseline: bool = True,
         detector_fourier_mask: np.ndarray = None,
+        virtual_detector_masks: Sequence[np.ndarray] = None,
         store_iterations: bool = False,
         collective_measurement_updates: bool = True,
         progress_bar: bool = True,
@@ -1309,6 +1310,9 @@ def reconstruct(
         detector_fourier_mask: np.ndarray
             Corner-centered mask to multiply the detector-plane gradients with (a value of zero supresses those pixels).
             Useful when detector has artifacts such as dead-pixels. Usually binary.
+        virtual_detector_masks: np.ndarray
+            List of corner-centered boolean masks for binning forward model exit waves,
+            to allow comparison with arbitrary geometry detector datasets.
         store_iterations: bool, optional
             If True, reconstructed objects and probes are stored at each iteration
         collective_measurement_updates: bool
@@ -1404,6 +1408,9 @@ def reconstruct(
         if detector_fourier_mask is not None:
             detector_fourier_mask = xp.asarray(detector_fourier_mask)
 
+        if virtual_detector_masks is not None:
+            virtual_detector_masks = xp.asarray(virtual_detector_masks)
+
         if gaussian_filter_sigma_m is None:
             gaussian_filter_sigma_m = gaussian_filter_sigma_e
 
@@ -1486,6 +1493,7 @@ def reconstruct(
                         amplitudes_device,
                         self._exit_waves,
                         detector_fourier_mask,
+                        virtual_detector_masks,
                         use_projection_scheme=use_projection_scheme,
                         projection_a=projection_a,
                         projection_b=projection_b,

py4DSTEM/process/phase/mixedstate_multislice_ptychography.py

@@ -767,6 +767,7 @@ def reconstruct(
         shrinkage_rad: float = 0.0,
         fix_potential_baseline: bool = True,
         detector_fourier_mask: np.ndarray = None,
+        virtual_detector_masks: Sequence[np.ndarray] = None,
         pure_phase_object: bool = False,
         tv_denoise_chambolle: bool = True,
         tv_denoise_weight_chambolle=None,
@@ -881,6 +882,9 @@ def reconstruct(
         detector_fourier_mask: np.ndarray
             Corner-centered mask to multiply the detector-plane gradients with (a value of zero supresses those pixels).
             Useful when detector has artifacts such as dead-pixels. Usually binary.
+        virtual_detector_masks: np.ndarray
+            List of corner-centered boolean masks for binning forward model exit waves,
+            to allow comparison with arbitrary geometry detector datasets.
         pure_phase_object: bool, optional
             If True, object amplitude is set to unity
         tv_denoise_chambolle: bool
@@ -978,6 +982,9 @@ def reconstruct(
         if detector_fourier_mask is not None:
             detector_fourier_mask = xp.asarray(detector_fourier_mask)
 
+        if virtual_detector_masks is not None:
+            virtual_detector_masks = xp.asarray(virtual_detector_masks)
+
         # main loop
         for a0 in tqdmnd(
             num_iter,
@@ -1025,6 +1032,7 @@ def reconstruct(
                     amplitudes_device,
                     self._exit_waves,
                     detector_fourier_mask,
+                    virtual_detector_masks,
                     use_projection_scheme,
                     projection_a,
                     projection_b,

py4DSTEM/process/phase/mixedstate_ptychography.py

@@ -676,6 +676,7 @@ def reconstruct(
         shrinkage_rad: float = 0.0,
         fix_potential_baseline: bool = True,
         detector_fourier_mask: np.ndarray = None,
+        virtual_detector_masks: Sequence[np.ndarray] = None,
         store_iterations: bool = False,
         progress_bar: bool = True,
         reset: bool = None,
@@ -787,6 +788,9 @@ def reconstruct(
         detector_fourier_mask: np.ndarray
             Corner-centered mask to multiply the detector-plane gradients with (a value of zero supresses those pixels).
             Useful when detector has artifacts such as dead-pixels. Usually binary.
+        virtual_detector_masks: np.ndarray
+            List of corner-centered boolean masks for binning forward model exit waves,
+            to allow comparison with arbitrary geometry detector datasets.
         store_iterations: bool, optional
             If True, reconstructed objects and probes are stored at each iteration
         progress_bar: bool, optional
@@ -872,6 +876,9 @@ def reconstruct(
         if detector_fourier_mask is not None:
             detector_fourier_mask = xp.asarray(detector_fourier_mask)
 
+        if virtual_detector_masks is not None:
+            virtual_detector_masks = xp.asarray(virtual_detector_masks)
+
         # main loop
         for a0 in tqdmnd(
             num_iter,
@@ -919,6 +926,7 @@ def reconstruct(
                     amplitudes_device,
                     self._exit_waves,
                     detector_fourier_mask,
+                    virtual_detector_masks,
                     use_projection_scheme,
                     projection_a,
                     projection_b,

py4DSTEM/process/phase/multislice_ptychography.py

@@ -739,6 +739,7 @@ def reconstruct(
         shrinkage_rad: float = 0.0,
         fix_potential_baseline: bool = True,
         detector_fourier_mask: np.ndarray = None,
+        virtual_detector_masks: Sequence[np.ndarray] = None,
         pure_phase_object: bool = False,
         tv_denoise_chambolle: bool = True,
         tv_denoise_weight_chambolle=None,
@@ -853,6 +854,9 @@ def reconstruct(
         detector_fourier_mask: np.ndarray
             Corner-centered mask to multiply the detector-plane gradients with (a value of zero supresses those pixels).
             Useful when detector has artifacts such as dead-pixels. Usually binary.
+        virtual_detector_masks: np.ndarray
+            List of corner-centered boolean masks for binning forward model exit waves,
+            to allow comparison with arbitrary geometry detector datasets.
         pure_phase_object: bool, optional
             If True, object amplitude is set to unity
         tv_denoise_chambolle: bool
@@ -954,6 +958,9 @@ def reconstruct(
         if detector_fourier_mask is not None:
             detector_fourier_mask = xp.asarray(detector_fourier_mask)
 
+        if virtual_detector_masks is not None:
+            virtual_detector_masks = xp.asarray(virtual_detector_masks)
+
         # main loop
         for a0 in tqdmnd(
             num_iter,
@@ -1001,6 +1008,7 @@ def reconstruct(
                     amplitudes_device,
                     self._exit_waves,
                     detector_fourier_mask,
+                    virtual_detector_masks,
                     use_projection_scheme,
                     projection_a,
                     projection_b,

py4DSTEM/process/phase/ptychographic_methods.py

@@ -2766,7 +2766,9 @@ def _return_farfield_amplitudes(self, fourier_overlap):
         xp = self._xp
         return xp.sqrt(xp.sum(xp.abs(fourier_overlap) ** 2, axis=1))
 
-    def _gradient_descent_fourier_projection(self, amplitudes, overlap, fourier_mask):
+    def _gradient_descent_fourier_projection(
+        self, amplitudes, overlap, fourier_mask, virtual_detector_masks
+    ):
         """
         Ptychographic fourier projection method for GD method.
 
@@ -2779,6 +2781,9 @@ def _gradient_descent_fourier_projection(self, amplitudes, overlap, fourier_mask
         fourier_mask: np.ndarray
             Mask to apply at the detector-plane for zeroing-out unreliable gradients
             Useful when detector has artifacts such as dead-pixels
+        virtual_detector_masks: np.ndarray
+            List of corner-centered boolean masks for binning forward model exit waves,
+            to allow comparison with arbitrary geometry detector datasets.
 
         Returns
         --------
@@ -2802,8 +2807,19 @@ def _gradient_descent_fourier_projection(self, amplitudes, overlap, fourier_mask
                 xp=xp,
             )
 
+        if virtual_detector_masks is not None:
+            masked_values = xp.sum(
+                fourier_overlap[:, :, None, :, :]
+                * virtual_detector_masks[None, None, :, :, :],
+                axis=(-1, -2),
+            ).transpose(2, 0, 1)
+            fourier_overlap = xp.zeros_like(fourier_overlap)
+            for mask, value in zip(virtual_detector_masks, masked_values):
+                fourier_overlap[..., mask] = value[:, :, None] / xp.sum(mask)
+
         if fourier_mask is not None:
             fourier_overlap *= fourier_mask
+
         farfield_amplitudes = self._return_farfield_amplitudes(fourier_overlap)
         error = xp.sum(xp.abs(amplitudes - farfield_amplitudes) ** 2)
 
@@ -2813,6 +2829,7 @@ def _gradient_descent_fourier_projection(self, amplitudes, overlap, fourier_mask
         fourier_modified_overlap = amplitude_modification[:, None] * fourier_overlap
 
         fourier_modified_overlap = fourier_modified_overlap - fourier_overlap
+
         if fourier_mask is not None:
             fourier_modified_overlap *= fourier_mask
 
@@ -2836,6 +2853,7 @@ def _projection_sets_fourier_projection(
         overlap,
         exit_waves,
         fourier_mask,
+        virtual_detector_masks,
         projection_a,
         projection_b,
         projection_c,
@@ -2867,6 +2885,9 @@ def _projection_sets_fourier_projection(
             Mask to apply at the detector-plane for zeroing-out unreliable gradients
             Useful when detector has artifacts such as dead-pixels
             Currently not implemented for projection sets
+        virtual_detector_masks: np.ndarray
+            List of corner-centered boolean masks for binning forward model exit waves,
+            to allow comparison with arbitrary geometry detector datasets.
         projection_a: float
         projection_b: float
         projection_c: float
@@ -2882,6 +2903,9 @@ def _projection_sets_fourier_projection(
         if fourier_mask is not None:
             raise NotImplementedError()
 
+        if virtual_detector_masks is not None:
+            raise NotImplementedError()
+
         xp = self._xp
         projection_x = 1 - projection_a - projection_b
         projection_y = 1 - projection_c

py4DSTEM/process/phase/ptychographic_tomography.py

@@ -788,6 +788,7 @@ def reconstruct(
         shrinkage_rad: float = 0.0,
         fix_potential_baseline: bool = True,
         detector_fourier_mask: np.ndarray = None,
+        virtual_detector_masks: Sequence[np.ndarray] = None,
         tv_denoise: bool = True,
         tv_denoise_weights: float = None,
         tv_denoise_inner_iter=40,
@@ -901,6 +902,9 @@ def reconstruct(
         detector_fourier_mask: np.ndarray
             Corner-centered mask to apply at the detector-plane for zeroing-out unreliable gradients.
             Useful when detector has artifacts such as dead-pixels. Usually binary.
+        virtual_detector_masks: np.ndarray
+            List of corner-centered boolean masks for binning forward model exit waves,
+            to allow comparison with arbitrary geometry detector datasets.
         store_iterations: bool, optional
             If True, reconstructed objects and probes are stored at each iteration
         progress_bar: bool, optional
@@ -979,6 +983,9 @@ def reconstruct(
         if detector_fourier_mask is not None:
             detector_fourier_mask = xp.asarray(detector_fourier_mask)
 
+        if virtual_detector_masks is not None:
+            virtual_detector_masks = xp.asarray(virtual_detector_masks)
+
         # main loop
         for a0 in tqdmnd(
             num_iter,
@@ -1071,6 +1078,7 @@ def reconstruct(
                         amplitudes_device,
                         self._exit_waves,
                         detector_fourier_mask,
+                        virtual_detector_masks,
                         use_projection_scheme,
                         projection_a,
                         projection_b,

py4DSTEM/process/phase/xray_magnetic_ptychography.py

@@ -1165,6 +1165,7 @@ def reconstruct(
         tv_denoise_weight: float = None,
         tv_denoise_inner_iter: float = 40,
         detector_fourier_mask: np.ndarray = None,
+        virtual_detector_masks: Sequence[np.ndarray] = None,
         store_iterations: bool = False,
         collective_measurement_updates: bool = True,
         progress_bar: bool = True,
@@ -1281,6 +1282,9 @@ def reconstruct(
         detector_fourier_mask: np.ndarray
             Corner-centered mask to multiply the detector-plane gradients with (a value of zero supresses those pixels).
             Useful when detector has artifacts such as dead-pixels. Usually binary.
+        virtual_detector_masks: np.ndarray
+            List of corner-centered boolean masks for binning forward model exit waves,
+            to allow comparison with arbitrary geometry detector datasets.
         store_iterations: bool, optional
             If True, reconstructed objects and probes are stored at each iteration
         collective_measurement_updates: bool
@@ -1376,6 +1380,9 @@ def reconstruct(
         if detector_fourier_mask is not None:
             detector_fourier_mask = xp.asarray(detector_fourier_mask)
 
+        if virtual_detector_masks is not None:
+            virtual_detector_masks = xp.asarray(virtual_detector_masks)
+
         if gaussian_filter_sigma_m is None:
             gaussian_filter_sigma_m = gaussian_filter_sigma_e
 
@@ -1458,6 +1465,7 @@ def reconstruct(
                         amplitudes_device,
                         self._exit_waves,
                         detector_fourier_mask,
+                        virtual_detector_masks,
                         use_projection_scheme=use_projection_scheme,
                         projection_a=projection_a,
                         projection_b=projection_b,