Merge pull request #676 from maclariz/maclariz_rdf_patch

Update polar_analysis.py

py4DSTEM/process/polar/polar_analysis.py

@@ -225,6 +225,8 @@ def plot_radial_var_norm(
 
 def calculate_pair_dist_function(
     self,
+    RxRy=None,
+    hxwy=None,
     k_min=0.05,
     k_max=None,
     k_width=0.25,
@@ -275,6 +277,13 @@ def calculate_pair_dist_function(
 
     Parameters
     ----------
+    RxRy : None or 2-tuple
+        None calculates on the whole radial average for dataset
+        A 2-tuple calculates on the radial profile for one scan pixel defined by Rx and Ry
+        or gives the top left corner of a box to average with widths defined by hxwy
+    hxwy : None or 2-tuple
+        A 2-tuple gives the height and width of a box to average anchored at Rx, Ry
+        hx and wy need to be larger than 1
     k_min : number
         Minimum scattering vector to include in the calculation
     k_max : number or None
@@ -322,13 +331,20 @@ def calculate_pair_dist_function(
     k = self.qq
     dk = k[1] - k[0]
     k2 = k**2
-    Ik = self.radial_mean
+    if RxRy == None:
+        Ik = self.radial_mean
+    elif RxRy != None and hxwy == None:
+        Ik = self.radial_all[RxRy[0], RxRy[1]]
+    elif RxRy != None and hxwy != None:
+        Ik = self.radial_all[
+            RxRy[0] : RxRy[0] + hxwy[0], RxRy[1] : RxRy[1] + hxwy[1]
+        ].mean(axis=(0, 1))
     int_mean = np.mean(Ik)
     sub_fit = k >= k_min
 
     # initial guesses for background coefs
-    const_bg = np.min(self.radial_mean) / int_mean
-    int0 = np.median(self.radial_mean) / int_mean - const_bg
+    const_bg = np.min(Ik) / int_mean
+    int0 = np.median(Ik) / int_mean - const_bg
     sigma0 = np.mean(k)
     coefs = [const_bg, int0, sigma0, int0, sigma0]
     lb = [0, 0, 0, 0, 0]
@@ -453,7 +469,7 @@ def calculate_pair_dist_function(
 
     # Plots
     if plot_background_fits:
-        fig, ax = self.plot_background_fits(figsize=figsize, returnfig=True)
+        fig, ax = self.plot_background_fits(Ik=Ik, figsize=figsize, returnfig=True)
         if returnfig:
             ans.append((fig, ax))
 
@@ -478,16 +494,24 @@ def calculate_pair_dist_function(
 
 def plot_background_fits(
     self,
+    Ik=None,
     figsize=(8, 4),
     returnfig=False,
 ):
     """
     TODO
+    Ik : numpy array
+        Ik calculated in calculate_pair_dist_function. Defaults to self.radial_mean for a pdf calculated
+        over whole dataset, but correctly calculated for sub areas, if defined
     """
+    if isinstance(Ik, np.ndarray):
+        pass
+    else:
+        Ik = self.radial_mean
     fig, ax = plt.subplots(figsize=figsize)
     ax.plot(
         self.qq,
-        self.radial_mean,
+        Ik,
         color="k",
     )
     ax.plot(
@@ -502,8 +526,8 @@ def plot_background_fits(
     ax.set_ylabel("I(k) and Background Fit Estimates")
     ax.set_ylim(
         (
-            np.min(self.radial_mean[self.radial_mean > 0]) * 0.8,
-            np.max(self.radial_mean * self.Sk_mask) * 1.25,
+            np.min(Ik[Ik > 0]) * 0.8,
+            np.max(Ik * self.Sk_mask) * 1.25,
         )
     )
     ax.set_yscale("log")