Added realspace mask for radial statistics in polardatacube

py4DSTEM/process/polar/polar_analysis.py

@@ -10,6 +10,7 @@
 
 def calculate_radial_statistics(
     self,
+    mask_realspace=None,
     plot_results_mean=False,
     plot_results_var=False,
     figsize=(8, 4),
@@ -89,16 +90,28 @@ def calculate_radial_statistics(
         unit=" probe positions",
         disable=not progress_bar,
     ):
-        self.radial_all[rx, ry] = np.mean(self.data[rx, ry], axis=0)
-        self.radial_all_std[rx, ry] = np.sqrt(
-            np.mean((self.data[rx, ry] - self.radial_all[rx, ry][None]) ** 2, axis=0)
+        if mask_realspace is None or mask_realspace[rx,ry]:
+            self.radial_all[rx, ry] = np.mean(self.data[rx, ry], axis=0)
+            self.radial_all_std[rx, ry] = np.sqrt(
+                np.mean((self.data[rx, ry] - self.radial_all[rx, ry][None]) ** 2, axis=0)
+            )
+
+    if mask_realspace is None:
+        self.radial_mean = np.mean(self.radial_all, axis=(0, 1))
+        self.radial_var = np.mean(
+            (self.radial_all - self.radial_mean[None, None]) ** 2, axis=(0, 1)
         )
 
-    self.radial_mean = np.mean(self.radial_all, axis=(0, 1))
-    self.radial_var = np.mean(
-        (self.radial_all - self.radial_mean[None, None]) ** 2, axis=(0, 1)
-    )
-
+    else:
+        self.radial_mean = np.sum(self.radial_all, axis=(0, 1)) / np.sum(mask_realspace)
+        self.radial_var = np.zeros_like(self.radial_mean)
+        for rx in range(self._datacube.shape[0]):
+            for ry in  range(self._datacube.shape[1]):
+                if mask_realspace[rx,ry]:
+                    self.radial_var += (self.radial_all[rx,ry] - self.radial_mean) ** 2
+        self.radial_var /= np.sum(mask_realspace)
+
+    # Compute normalized variance
     self.radial_var_norm = np.copy(self.radial_var)
     sub = self.radial_mean > 0.0
     self.radial_var_norm[sub] /= self.radial_mean[sub] ** 2