diff --git a/py4DSTEM/process/phase/magnetic_ptychographic_tomography.py b/py4DSTEM/process/phase/magnetic_ptychographic_tomography.py index 597ae3d7f..d763c8886 100644 --- a/py4DSTEM/process/phase/magnetic_ptychographic_tomography.py +++ b/py4DSTEM/process/phase/magnetic_ptychographic_tomography.py @@ -874,6 +874,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, probe_real_space_support_mask: np.ndarray = None, tv_denoise: bool = True, tv_denoise_weights=None, @@ -990,6 +991,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. probe_real_space_support_mask: np.ndarray Corner-centered boolean mask, outside of which the probe amplitude will be set to zero. store_iterations: bool, optional @@ -1081,6 +1085,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 @@ -1188,6 +1195,7 @@ def reconstruct( amplitudes_device, self._exit_waves, detector_fourier_mask, + virtual_detector_masks, use_projection_scheme, projection_a, projection_b, diff --git a/py4DSTEM/process/phase/magnetic_ptychography.py b/py4DSTEM/process/phase/magnetic_ptychography.py index 4872193d3..b5fb1dc46 100644 --- a/py4DSTEM/process/phase/magnetic_ptychography.py +++ 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, probe_real_space_support_mask: np.ndarray = None, store_iterations: bool = False, collective_measurement_updates: bool = True, @@ -1310,6 +1311,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. probe_real_space_support_mask: np.ndarray Corner-centered boolean mask, outside of which the probe amplitude will be set to zero. store_iterations: bool, optional @@ -1407,6 +1411,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 @@ -1489,6 +1496,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, diff --git a/py4DSTEM/process/phase/mixedstate_multislice_ptychography.py b/py4DSTEM/process/phase/mixedstate_multislice_ptychography.py index f56245e64..f0cffac3d 100644 --- a/py4DSTEM/process/phase/mixedstate_multislice_ptychography.py +++ b/py4DSTEM/process/phase/mixedstate_multislice_ptychography.py @@ -768,6 +768,7 @@ def reconstruct( fix_potential_baseline: bool = True, vacuum_mask: np.ndarray = None, detector_fourier_mask: np.ndarray = None, + virtual_detector_masks: Sequence[np.ndarray] = None, probe_real_space_support_mask: np.ndarray = None, pure_phase_object: bool = False, tv_denoise_chambolle: bool = True, @@ -885,6 +886,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. probe_real_space_support_mask: np.ndarray Corner-centered boolean mask, outside of which the probe amplitude will be set to zero. pure_phase_object: bool, optional @@ -984,6 +988,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, @@ -1031,6 +1038,7 @@ def reconstruct( amplitudes_device, self._exit_waves, detector_fourier_mask, + virtual_detector_masks, use_projection_scheme, projection_a, projection_b, diff --git a/py4DSTEM/process/phase/mixedstate_ptychography.py b/py4DSTEM/process/phase/mixedstate_ptychography.py index b0b222539..02ed99708 100644 --- a/py4DSTEM/process/phase/mixedstate_ptychography.py +++ b/py4DSTEM/process/phase/mixedstate_ptychography.py @@ -677,6 +677,7 @@ def reconstruct( fix_potential_baseline: bool = True, vacuum_mask: np.ndarray = None, detector_fourier_mask: np.ndarray = None, + virtual_detector_masks: Sequence[np.ndarray] = None, probe_real_space_support_mask: np.ndarray = None, store_iterations: bool = False, progress_bar: bool = True, @@ -791,6 +792,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. probe_real_space_support_mask: np.ndarray Corner-centered boolean mask, outside of which the probe amplitude will be set to zero. store_iterations: bool, optional @@ -878,6 +882,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, @@ -925,6 +932,7 @@ def reconstruct( amplitudes_device, self._exit_waves, detector_fourier_mask, + virtual_detector_masks, use_projection_scheme, projection_a, projection_b, diff --git a/py4DSTEM/process/phase/multislice_ptychography.py b/py4DSTEM/process/phase/multislice_ptychography.py index e371b4e51..570f1c796 100644 --- a/py4DSTEM/process/phase/multislice_ptychography.py +++ b/py4DSTEM/process/phase/multislice_ptychography.py @@ -740,6 +740,7 @@ def reconstruct( fix_potential_baseline: bool = True, vacuum_mask: np.ndarray = None, detector_fourier_mask: np.ndarray = None, + virtual_detector_masks: Sequence[np.ndarray] = None, probe_real_space_support_mask: np.ndarray = None, pure_phase_object: bool = False, tv_denoise_chambolle: bool = True, @@ -857,6 +858,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. probe_real_space_support_mask: np.ndarray Corner-centered boolean mask, outside of which the probe amplitude will be set to zero. pure_phase_object: bool, optional @@ -960,6 +964,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, @@ -1007,6 +1014,7 @@ def reconstruct( amplitudes_device, self._exit_waves, detector_fourier_mask, + virtual_detector_masks, use_projection_scheme, projection_a, projection_b, diff --git a/py4DSTEM/process/phase/ptychographic_methods.py b/py4DSTEM/process/phase/ptychographic_methods.py index 97ddb4433..b468e9833 100644 --- a/py4DSTEM/process/phase/ptychographic_methods.py +++ b/py4DSTEM/process/phase/ptychographic_methods.py @@ -1679,7 +1679,9 @@ def cross_correlate_amplitudes_to_probe_aperture( return self - 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. @@ -1692,6 +1694,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 -------- @@ -1715,6 +1720,15 @@ 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, :, :, :], + axis=(-1, -2), + ).transpose() + 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 @@ -1746,6 +1760,7 @@ def _projection_sets_fourier_projection( overlap, exit_waves, fourier_mask, + virtual_detector_masks, projection_a, projection_b, projection_c, @@ -1777,6 +1792,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 @@ -1792,6 +1810,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 @@ -1849,6 +1870,7 @@ def _forward( amplitudes, exit_waves, fourier_mask, + virtual_detector_masks, use_projection_scheme, projection_a, projection_b, @@ -1871,6 +1893,9 @@ def _forward( 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. use_projection_scheme: bool, If True, use generalized projection update projection_a: float @@ -1907,6 +1932,7 @@ def _forward( overlap, exit_waves, fourier_mask, + virtual_detector_masks, projection_a, projection_b, projection_c, @@ -1917,6 +1943,7 @@ def _forward( amplitudes, overlap, fourier_mask, + virtual_detector_masks, ) return shifted_probes, object_patches, overlap, exit_waves, error @@ -2904,7 +2931,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. @@ -2917,6 +2946,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 -------- @@ -2940,8 +2972,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) @@ -2951,6 +2994,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 @@ -2974,6 +3018,7 @@ def _projection_sets_fourier_projection( overlap, exit_waves, fourier_mask, + virtual_detector_masks, projection_a, projection_b, projection_c, @@ -3005,6 +3050,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 @@ -3020,6 +3068,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 diff --git a/py4DSTEM/process/phase/ptychographic_tomography.py b/py4DSTEM/process/phase/ptychographic_tomography.py index 25de7e7ed..2d3ef8c4b 100644 --- a/py4DSTEM/process/phase/ptychographic_tomography.py +++ b/py4DSTEM/process/phase/ptychographic_tomography.py @@ -796,6 +796,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, probe_real_space_support_mask: np.ndarray = None, tv_denoise: bool = True, tv_denoise_weights: float = None, @@ -910,6 +911,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. probe_real_space_support_mask: np.ndarray Corner-centered boolean mask, outside of which the probe amplitude will be set to zero. store_iterations: bool, optional @@ -990,6 +994,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, @@ -1082,6 +1089,7 @@ def reconstruct( amplitudes_device, self._exit_waves, detector_fourier_mask, + virtual_detector_masks, use_projection_scheme, projection_a, projection_b, diff --git a/py4DSTEM/process/phase/singleslice_ptychography.py b/py4DSTEM/process/phase/singleslice_ptychography.py index 92f0f2df3..f459385eb 100644 --- a/py4DSTEM/process/phase/singleslice_ptychography.py +++ b/py4DSTEM/process/phase/singleslice_ptychography.py @@ -650,6 +650,7 @@ def reconstruct( fix_potential_baseline: bool = True, vacuum_mask: np.ndarray = None, detector_fourier_mask: np.ndarray = None, + virtual_detector_masks: Sequence[np.ndarray] = None, probe_real_space_support_mask: np.ndarray = None, store_iterations: bool = False, progress_bar: bool = True, @@ -762,6 +763,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. probe_real_space_support_mask: np.ndarray Corner-centered boolean mask, outside of which the probe amplitude will be set to zero. store_iterations: bool, optional @@ -849,6 +853,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, @@ -896,6 +903,7 @@ def reconstruct( amplitudes_device, self._exit_waves, detector_fourier_mask, + virtual_detector_masks, use_projection_scheme, projection_a, projection_b, diff --git a/py4DSTEM/process/phase/xray_magnetic_ptychography.py b/py4DSTEM/process/phase/xray_magnetic_ptychography.py index 86e4367b5..2dc74cfe4 100644 --- a/py4DSTEM/process/phase/xray_magnetic_ptychography.py +++ 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, probe_real_space_support_mask: np.ndarray = None, store_iterations: bool = False, collective_measurement_updates: bool = True, @@ -1282,6 +1283,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. probe_real_space_support_mask: np.ndarray Corner-centered boolean mask, outside of which the probe amplitude will be set to zero. store_iterations: bool, optional @@ -1379,6 +1383,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 @@ -1461,6 +1468,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,