another day another bug: make tilting reshape work

py4dstem · Jun 15, 2024 · 1d089bf · 1d089bf
1 parent de38241
commit 1d089bf
Showing 1 changed file with 8 additions and 10 deletions.
diff --git a/py4DSTEM/tomography/tomography.py b/py4DSTEM/tomography/tomography.py
@@ -1276,7 +1276,7 @@ def _back(
             bincount_diff,
             update_reshaped.ravel(),
             minlength=((diff_max) * real_shape),
-        )[xp.tile(diff_bincount > 0, real_shape)].reshape((real_shape, -1))
+        ).reshape((real_shape, -1))[:, diff_bincount > 0]
 
         diff_shape_bin = update_q_summed.shape[-1]
 
@@ -1289,14 +1289,12 @@ def _back(
         )
 
         update_r_summed = (
-            (
-                xp.bincount(
-                    bincount_real,
-                    (update_q_summed * self._weights_real.ravel()[:, None]).ravel(),
-                    minlength=((real_max) * diff_shape_bin),
-                )
-            )[xp.tile(real_bincount > 0, diff_shape_bin)]
-        ).reshape((-1, diff_shape_bin))
+            xp.bincount(
+                bincount_real,
+                (update_q_summed * self._weights_real.ravel()[:, None]).ravel(),
+                minlength=((real_max) * diff_shape_bin),
+            )
+        ).reshape((-1, diff_shape_bin))[real_bincount > 0]
 
         yy, zz = xp.meshgrid(
             xp.unique(real_index), xp.unique(diff_index), indexing="ij"
@@ -1343,7 +1341,7 @@ def _make_test_object(
             6D test object
         """
         xp_storage = self._xp_storage
-        storage = storage
+        storage = self._storage
 
         test_object = xp_storage.zeros((sx, sy, sz, sq, sq, sq))