fix error

py4dstem · Jan 19, 2025 · 0af513c · 0af513c
1 parent a21f822
commit 0af513c
Showing 1 changed file with 58 additions and 55 deletions.
diff --git a/py4DSTEM/tomography/tomography.py b/py4DSTEM/tomography/tomography.py
@@ -1004,48 +1004,48 @@ def _solve_for_indicies(
 
         ind_real = np.ravel_multi_index((ind0, ind1), (s[1], s[2]), mode="clip")
 
-        # # solve for real space normalization
-        # num_points_norm = s[2]
-        # line_z_norm = np.linspace(0, 1, num_points_norm) * (s[2] - 1)
-        # line_y_norm = line_z_norm * np.tan(tilt)
-        # line_y_norm -= np.mean(line_y_norm)
-        # offset = np.arange(s[1], dtype="int")
-
-        # yF_norm = np.floor(line_y_norm).astype("int")
-        # zF_norm = np.floor(line_z_norm).astype("int")
-        # dy_norm = line_y_norm - yF_norm
-        # dz_norm = line_z_norm - zF_norm
-
-        # ind0_norm = np.hstack(
-        #     (
-        #         np.tile(yF_norm, (s[1], 1)) + offset[:, None],
-        #         np.tile(yF_norm + 1, (s[1], 1)) + offset[:, None],
-        #         np.tile(yF_norm, (s[1], 1)) + offset[:, None],
-        #         np.tile(yF_norm + 1, (s[1], 1)) + offset[:, None],
-        #     )
-        # )
+        # solve for real space normalization
+        num_points_norm = s[2]
+        line_z_norm = np.linspace(0, 1, num_points_norm) * (s[2] - 1)
+        line_y_norm = line_z_norm * np.tan(tilt)
+        line_y_norm -= np.mean(line_y_norm)
+        offset = np.arange(s[1], dtype="int")
 
-        # ind1_norm = np.hstack(
-        #     (
-        #         np.tile(zF_norm, (s[1], 1)),
-        #         np.tile(zF_norm, (s[1], 1)),
-        #         np.tile(zF_norm + 1, (s[1], 1)),
-        #         np.tile(zF_norm + 1, (s[1], 1)),
-        #     )
-        # )
+        yF_norm = np.floor(line_y_norm).astype("int")
+        zF_norm = np.floor(line_z_norm).astype("int")
+        dy_norm = line_y_norm - yF_norm
+        dz_norm = line_z_norm - zF_norm
 
-        # weights_real_norm = np.hstack(
-        #     (
-        #         np.tile(((1 - dy_norm) * (1 - dz_norm)), (s[1], 1)),
-        #         np.tile(((dy_norm) * (1 - dz_norm)), (s[1], 1)),
-        #         np.tile(((1 - dy_norm) * (dz_norm)), (s[1], 1)),
-        #         np.tile(((dy_norm) * (dz_norm)), (s[1], 1)),
-        #     )
-        # )
+        ind0_norm = np.hstack(
+            (
+                np.tile(yF_norm, (s[1], 1)) + offset[:, None],
+                np.tile(yF_norm + 1, (s[1], 1)) + offset[:, None],
+                np.tile(yF_norm, (s[1], 1)) + offset[:, None],
+                np.tile(yF_norm + 1, (s[1], 1)) + offset[:, None],
+            )
+        )
 
-        # ind_real_norm = np.ravel_multi_index(
-        #     (ind0_norm, ind1_norm), (s[1], s[2]), mode="clip"
-        # )
+        ind1_norm = np.hstack(
+            (
+                np.tile(zF_norm, (s[1], 1)),
+                np.tile(zF_norm, (s[1], 1)),
+                np.tile(zF_norm + 1, (s[1], 1)),
+                np.tile(zF_norm + 1, (s[1], 1)),
+            )
+        )
+
+        weights_real_norm = np.hstack(
+            (
+                np.tile(((1 - dy_norm) * (1 - dz_norm)), (s[1], 1)),
+                np.tile(((dy_norm) * (1 - dz_norm)), (s[1], 1)),
+                np.tile(((1 - dy_norm) * (dz_norm)), (s[1], 1)),
+                np.tile(((dy_norm) * (dz_norm)), (s[1], 1)),
+            )
+        )
+
+        ind_real_norm = np.ravel_multi_index(
+            (ind0_norm, ind1_norm), (s[1], s[2]), mode="clip"
+        )
 
         # solve for diffraction space coordinates
         length = s[-1] * np.cos(tilt)
@@ -1154,19 +1154,19 @@ def _solve_for_indicies(
         )
         ind_real_bincount = np.bincount(ind_real.ravel(), minlength=bincount_real_max)
 
-        # ind_real_bincount_weight_norm = np.bincount(
-        #     ind_real_norm.ravel(),
-        #     weights_real_norm.ravel(),
-        #     minlength=bincount_real_max,
-        # )
-        # ind_real_bincount_norm = np.bincount(
-        #     ind_real_norm.ravel(), minlength=bincount_real_max
-        # )
+        ind_real_bincount_weight_norm = np.bincount(
+            ind_real_norm.ravel(),
+            weights_real_norm.ravel(),
+            minlength=bincount_real_max,
+        )
+        ind_real_bincount_norm = np.bincount(
+            ind_real_norm.ravel(), minlength=bincount_real_max
+        )
 
-        # ind_real_bincount_weight_norm = ind_real_bincount_weight_norm[
-        #     ind_real_bincount > 0
-        # ]
-        # ind_real_bincount_norm = ind_real_bincount_norm[ind_real_bincount > 0]
+        ind_real_bincount_weight_norm = ind_real_bincount_weight_norm[
+            ind_real_bincount > 0
+        ]
+        ind_real_bincount_norm = ind_real_bincount_norm[ind_real_bincount > 0]
 
         ind_real_bincount_weight = ind_real_bincount_weight[ind_real_bincount > 0]
         ind_real_bincount = ind_real_bincount[ind_real_bincount > 0]
@@ -1185,7 +1185,8 @@ def _solve_for_indicies(
         #     / ind_real_bincount
         # )
 
-        correction_factor_real = 1 / ind_real_bincount_weight
+        correction_factor_real = ind_real_bincount_weight_norm / ind_real_bincount_weight
+        # correction_factor_real = 1 / ind_real_bincount_weight
 
         correction_factor_real = np.repeat(correction_factor_real, ind_real_bincount)
         sorted_indicies = np.argsort(np.argsort(ind_real.ravel()))
@@ -1194,6 +1195,8 @@ def _solve_for_indicies(
         )
         weights_real = weights_real * correction_factor_real
 
+
+
         # normalization reciprocal space
         bincount_diff_max = np.max((ind_diff.max(), ind_diff_norm.max())) + 1
 
@@ -1225,16 +1228,16 @@ def _solve_for_indicies(
         correction_factor_diff = (
             ind_diff_bincount_weight_norm
             / ind_diff_bincount_weight
-            / ind_diff_bincount_norm
-            * ind_diff_bincount
+            # / ind_diff_bincount_norm
+            # * ind_diff_bincount
         )
 
         correction_factor_diff = np.repeat(correction_factor_diff, ind_diff_bincount)
         sorted_indicies = np.argsort(np.argsort(ind_diff.ravel()))
         correction_factor_diff = correction_factor_diff[sorted_indicies].reshape(
             ind_diff.shape
         )
-        # weights_diff = weights_diff * correction_factor_diff
+        weights_diff = weights_diff * correction_factor_diff
 
         if datacube_number == 0:
             self._ind_real = []