From d97164153c5b6e3a13d33a759a1ebc66ff65d0b1 Mon Sep 17 00:00:00 2001 From: ellisrae Date: Wed, 18 Oct 2023 14:19:14 -0700 Subject: [PATCH] Added realspace mask for radial statistics in polardatacube --- py4DSTEM/process/polar/polar_analysis.py | 29 +++++++++++++++++------- 1 file changed, 21 insertions(+), 8 deletions(-) diff --git a/py4DSTEM/process/polar/polar_analysis.py b/py4DSTEM/process/polar/polar_analysis.py index 8d2d92585..e5b604bcd 100644 --- a/py4DSTEM/process/polar/polar_analysis.py +++ b/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