add OBF-like weights

py4DSTEM/process/phase/direct_ptychography.py

@@ -1229,6 +1229,7 @@ def unwrap_trotter_phase(complex_data, mask):
 
     def reconstruct(
         self,
+        use_OBF_weighting=False,
         wigner_deconvolution_wiener_epsilon=None,
         virtual_detector_masks: Sequence[np.ndarray] = None,
         polar_parameters: Mapping[str, float] = None,
@@ -1242,6 +1243,9 @@ def reconstruct(
 
         Parameters
         --------
+        use_OBF_weighting: bool
+            If True, the complex double probe overlap function is normalized using the
+            expresion in 10.1016/j.ultramic.2020.113133 instead of the amplitude.
         wigner_deconvolution_wiener_epsilon: float, optional
             If None (default), uses the direct inversion algorithm described in 10.1016/j.ultramic.2016.09.002.
             If not None, Wigner Distribution Deconvolution is performed, using the specified relative Wiener epsilon
@@ -1347,6 +1351,7 @@ def reconstruct(
 
         if wigner_deconvolution_wiener_epsilon is None:
             return self._reconstruct_SSB_gamma(
+                use_OBF_weighting=use_OBF_weighting,
                 virtual_detector_masks=virtual_detector_masks,
                 progress_bar=progress_bar,
             )
@@ -1447,6 +1452,7 @@ def _reconstruct_WDD(
 
     def _reconstruct_SSB_gamma(
         self,
+        use_OBF_weighting: bool = False,
         virtual_detector_masks: Sequence[np.ndarray] = None,
         progress_bar: bool = True,
     ):
@@ -1455,6 +1461,9 @@ def _reconstruct_SSB_gamma(
 
         Parameters
         --------
+        use_OBF_weighting: bool
+            If True, the complex double probe overlap function is normalized using the
+            expresion in 10.1016/j.ultramic.2020.113133 instead of the amplitude.
         virtual_detector_masks: np.ndarray
             List of corner-centered boolean masks for binning forward model trotters,
             to allow comparison with arbitrary geometry detector datasets. TO-DO
@@ -1480,6 +1489,14 @@ def _reconstruct_SSB_gamma(
         if virtual_detector_masks is not None:
             virtual_detector_masks = xp.asarray(virtual_detector_masks).astype(xp.bool_)
 
+        if use_OBF_weighting:
+            probe_normalization = xp.abs(self._fourier_probe) ** 2
+            probe_normalization /= probe_normalization.sum()
+            if virtual_detector_masks is not None:
+                probe_normalization = mask_array_using_virtual_detectors(
+                    probe_normalization, virtual_detector_masks, in_place=True
+                )
+
         # main loop
         for ind_x, ind_y in tqdmnd(
             sx,
@@ -1533,8 +1550,14 @@ def _reconstruct_SSB_gamma(
 
                 gamma_abs = np.abs(gamma)
                 gamma_ind = gamma_abs > threshold
+
+                normalization = gamma_abs[gamma_ind]
+                if use_OBF_weighting:
+                    d = probe_normalization[gamma_ind]
+                    normalization = d * xp.sqrt(xp.sum(normalization**2 / d))
+
                 psi[ind_x, ind_y] = (
-                    G[gamma_ind] * xp.conj(gamma[gamma_ind]) / gamma_abs[gamma_ind]
+                    G[gamma_ind] * xp.conj(gamma[gamma_ind]) / normalization
                 ).sum()
 
         self._object = xp.fft.ifft2(psi) / self._mean_diffraction_intensity