diff --git a/py4DSTEM/process/polar/polar_analysis.py b/py4DSTEM/process/polar/polar_analysis.py index aee863e50..78a95c24a 100644 --- a/py4DSTEM/process/polar/polar_analysis.py +++ b/py4DSTEM/process/polar/polar_analysis.py @@ -12,6 +12,7 @@ def calculate_radial_statistics( self, + mask_realspace=None, plot_results_mean=False, plot_results_var=False, figsize=(8, 4), @@ -93,16 +94,30 @@ 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 + ) + ) - 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) - ) + 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) + ) + 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