plotting

py4dstem · Jan 30, 2025 · 7669fcd · 7669fcd
1 parent 8a2dc3b
commit 7669fcd
Showing 1 changed file with 17 additions and 5 deletions.
diff --git a/py4DSTEM/tomography/tomography.py b/py4DSTEM/tomography/tomography.py
@@ -1829,7 +1829,7 @@ def widget(
         vmax = kwargs.pop("vmax", None)
 
         _, vmin, vmax = return_scaled_histogram_ordering(
-            ((obj_6D) * diffraction_kernel[None, None, None, :, :]).mean((3, 4, 5)),
+            ((obj_6D) * diffraction_kernel[None, None, None, :, :, :]).mean((3, 4, 5)),
             vmin=vmin,
             vmax=vmax,
         )
@@ -1849,7 +1849,7 @@ def widget(
         min_distance = 1
 
         ax0.imshow(
-            (obj_6D[:, :, z] * diffraction_kernel[None, None, :, :]).mean((2, 3, 4)),
+            (obj_6D[:, :, z] * diffraction_kernel[None, None, :, :, :]).mean((2, 3, 4)),
             cmap="gray",
             vmin=vmin,
             vmax=vmax,
@@ -1863,7 +1863,7 @@ def widget(
                     y,
                     :,
                 ]
-                * diffraction_kernel[None, None, :, :]
+                * diffraction_kernel[None, None, :, :, :]
             )
             .mean((2, 3, 4))
             .T,
@@ -1902,6 +1902,12 @@ def widget(
         ax2.set_yticks([])
         ax2.set_zticks([])
 
+        ax0.set_xlabel('y')
+        ax0.set_ylabel('x')
+
+        ax1.set_xlabel('x')
+        ax1.set_ylabel('z')
+
         ax2.set_xlim([0, obj_6D.shape[3]])
         ax2.set_ylim([0, obj_6D.shape[4]])
         ax2.set_zlim([0, obj_6D.shape[5]])
@@ -1929,7 +1935,7 @@ def update_images(
                         :,
                         z,
                     ]
-                    * diffraction_kernel[None, None, :, :]
+                    * diffraction_kernel[None, None, :, :, :]
                 ).mean((2, 3, 4)),
                 cmap="gray",
                 vmin=vmin,
@@ -1945,7 +1951,7 @@ def update_images(
                         y,
                         :,
                     ]
-                    * diffraction_kernel[None, None, :, :]
+                    * diffraction_kernel[None, None, :, :, :]
                 )
                 .mean((2, 3, 4))
                 .T,
@@ -1985,6 +1991,12 @@ def update_images(
             ax2.set_yticks([])
             ax2.set_zticks([])
 
+            ax0.set_xlabel('y')
+            ax0.set_ylabel('x')
+
+            ax1.set_xlabel('x')
+            ax1.set_ylabel('z')
+
             ax2.set_xlim([0, obj_6D.shape[3]])
             ax2.set_ylim([0, obj_6D.shape[4]])
             ax2.set_zlim([0, obj_6D.shape[5]])