more careful normalizations

py4dstem · May 30, 2024 · 936df0e · 936df0e
1 parent 7ed3d2e
commit 936df0e
Showing 1 changed file with 77 additions and 37 deletions.
diff --git a/py4DSTEM/process/phase/xray_magnetic_ptychography.py b/py4DSTEM/process/phase/xray_magnetic_ptychography.py
@@ -887,49 +887,22 @@ def _gradient_descent_adjoint(
             + (normalization_min * xp.max(probe_electrostatic_normalization)) ** 2
         )
 
-        probe_magnetic_abs = xp.abs(shifted_probes * xp.exp(1.0j * object_patches[1]))
-        probe_magnetic_normalization = self._sum_overlapping_patches_bincounts(
-            probe_magnetic_abs**2,
-            positions_px,
-        )
-        probe_magnetic_normalization = 1 / xp.sqrt(
-            1e-16
-            + ((1 - normalization_min) * probe_magnetic_normalization) ** 2
-            + (normalization_min * xp.max(probe_magnetic_normalization)) ** 2
-        )
+        match (self._recon_mode, self._active_measurement_index):
+            case (0, 0) | (1, 0):  # reverse
 
-        if not fix_probe:
-            electrostatic_magnetic_abs = xp.abs(
-                object_patches[0] * xp.exp(1.0j * object_patches[1])
-            )
-            electrostatic_magnetic_normalization = xp.sum(
-                electrostatic_magnetic_abs**2,
-                axis=0,
-            )
-            electrostatic_magnetic_normalization = 1 / xp.sqrt(
-                1e-16
-                + ((1 - normalization_min) * electrostatic_magnetic_normalization) ** 2
-                + (normalization_min * xp.max(electrostatic_magnetic_normalization))
-                ** 2
-            )
+                magnetic_exp = xp.exp(1.0j * xp.conj(object_patches[1]))
 
-            if self._recon_mode > 0:
-                electrostatic_abs = xp.abs(object_patches[0])
-                electrostatic_normalization = xp.sum(
-                    electrostatic_abs**2,
-                    axis=0,
+                probe_magnetic_abs = xp.abs(shifted_probes * magnetic_exp)
+                probe_magnetic_normalization = self._sum_overlapping_patches_bincounts(
+                    probe_magnetic_abs**2,
+                    positions_px,
                 )
-                electrostatic_normalization = 1 / xp.sqrt(
+                probe_magnetic_normalization = 1 / xp.sqrt(
                     1e-16
-                    + ((1 - normalization_min) * electrostatic_normalization) ** 2
-                    + (normalization_min * xp.max(electrostatic_normalization)) ** 2
+                    + ((1 - normalization_min) * probe_magnetic_normalization) ** 2
+                    + (normalization_min * xp.max(probe_magnetic_normalization)) ** 2
                 )
 
-        match (self._recon_mode, self._active_measurement_index):
-            case (0, 0) | (1, 0):  # reverse
-
-                magnetic_exp = xp.exp(1.0j * xp.conj(object_patches[1]))
-
                 # P* exp(i M*)
                 electrostatic_update = self._sum_overlapping_patches_bincounts(
                     probe_conj * magnetic_exp * exit_waves,
@@ -950,6 +923,28 @@ def _gradient_descent_adjoint(
                 )
 
                 if not fix_probe:
+
+                    electrostatic_magnetic_abs = xp.abs(
+                        object_patches[0] * magnetic_exp
+                    )
+                    electrostatic_magnetic_normalization = xp.sum(
+                        electrostatic_magnetic_abs**2,
+                        axis=0,
+                    )
+                    electrostatic_magnetic_normalization = 1 / xp.sqrt(
+                        1e-16
+                        + (
+                            (1 - normalization_min)
+                            * electrostatic_magnetic_normalization
+                        )
+                        ** 2
+                        + (
+                            normalization_min
+                            * xp.max(electrostatic_magnetic_normalization)
+                        )
+                        ** 2
+                    )
+
                     # exp(i M*) C*
                     current_probe += step_size * (
                         xp.sum(
@@ -963,6 +958,17 @@ def _gradient_descent_adjoint(
 
                 magnetic_exp = xp.exp(-1.0j * xp.conj(object_patches[1]))
 
+                probe_magnetic_abs = xp.abs(shifted_probes * magnetic_exp)
+                probe_magnetic_normalization = self._sum_overlapping_patches_bincounts(
+                    probe_magnetic_abs**2,
+                    positions_px,
+                )
+                probe_magnetic_normalization = 1 / xp.sqrt(
+                    1e-16
+                    + ((1 - normalization_min) * probe_magnetic_normalization) ** 2
+                    + (normalization_min * xp.max(probe_magnetic_normalization)) ** 2
+                )
+
                 # P* exp(-i M*)
                 electrostatic_update = self._sum_overlapping_patches_bincounts(
                     probe_conj * magnetic_exp * exit_waves,
@@ -983,6 +989,28 @@ def _gradient_descent_adjoint(
                 )
 
                 if not fix_probe:
+
+                    electrostatic_magnetic_abs = xp.abs(
+                        object_patches[0] * magnetic_exp
+                    )
+                    electrostatic_magnetic_normalization = xp.sum(
+                        electrostatic_magnetic_abs**2,
+                        axis=0,
+                    )
+                    electrostatic_magnetic_normalization = 1 / xp.sqrt(
+                        1e-16
+                        + (
+                            (1 - normalization_min)
+                            * electrostatic_magnetic_normalization
+                        )
+                        ** 2
+                        + (
+                            normalization_min
+                            * xp.max(electrostatic_magnetic_normalization)
+                        )
+                        ** 2
+                    )
+
                     # exp(-i M*) C*
                     current_probe += step_size * (
                         xp.sum(
@@ -1015,6 +1043,18 @@ def _gradient_descent_adjoint(
                 )
 
                 if not fix_probe:
+
+                    electrostatic_abs = xp.abs(object_patches[0])
+                    electrostatic_normalization = xp.sum(
+                        electrostatic_abs**2,
+                        axis=0,
+                    )
+                    electrostatic_normalization = 1 / xp.sqrt(
+                        1e-16
+                        + ((1 - normalization_min) * electrostatic_normalization) ** 2
+                        + (normalization_min * xp.max(electrostatic_normalization)) ** 2
+                    )
+
                     # V*
                     current_probe += step_size * (
                         xp.sum(