Updating RDF calcs

py4DSTEM/process/polar/polar_analysis.py

@@ -5,6 +5,7 @@
 import matplotlib.pyplot as plt
 from scipy.optimize import curve_fit
 from scipy.ndimage import gaussian_filter
+from scipy.signal import argrelextrema
 from sklearn.decomposition import PCA
 
 from emdfile import tqdmnd
@@ -319,13 +320,61 @@ def calculate_pair_dist_function(
         which `plot_*` is True are returned.
     """
 
+    k_width = np.array(k_width)
+    if k_width.size == 1:
+        k_width = k_width*np.ones(2)
+
     # set up coordinates and scaling
     k = self.qq
     dk = k[1] - k[0]
     k2 = k**2
     Ik = self.radial_mean
     int_mean = np.mean(Ik)
-    sub_fit = k >= k_min
+    # sub_fit = k >= k_min
+    # sub_fit = np.logical_and(
+    #     k >= k_min
+
+    # Calculate structure factor mask 
+    if k_max is None:
+        k_max = np.max(k)
+    mask_low = np.sin(
+        np.clip(
+            (k - k_min) / k_width[0],
+            0,
+            1,
+        ) * np.pi/2.0,
+    )**2
+    mask_high = np.sin(
+        np.clip(
+            (k_max - k) / k_width[1],
+            0,
+            1,
+        ) * np.pi/2.0,
+    )**2
+    mask = mask_low * mask_high
+
+    # weighting function for fitting atomic scattering factors
+    weights_fit = np.divide(
+        1,
+        mask_low,
+        where = mask_low > 1e-4,
+    )
+    weights_fit[mask_low <= 1e-4] = np.inf
+    # Scale weighting to favour high k values
+    weights_fit *= ((k[-1] - 0.9*k + dk))
+
+
+    # fig,ax = plt.subplots()
+    # ax.plot(
+    #     k,
+    #     weights_fit,
+    # )
+    # ax.set_ylim([0,4])
+    # ax.plot(
+    #     k,
+    #     mask_high,
+    # )
+
 
     # initial guesses for background coefs
     const_bg = np.min(self.radial_mean) / int_mean
@@ -335,25 +384,31 @@ def calculate_pair_dist_function(
     lb = [0, 0, 0, 0, 0]
     ub = [np.inf, np.inf, np.inf, np.inf, np.inf]
     # Weight the fit towards high k values
-    noise_est = k[-1] - k + dk
+    # noise_est = k[-1] - k + dk
 
     # Estimate the mean atomic form factor + background
     if maxfev is None:
         coefs = curve_fit(
             scattering_model,
-            k2[sub_fit],
-            Ik[sub_fit] / int_mean,
-            sigma=noise_est[sub_fit],
+            # k2[sub_fit],
+            # Ik[sub_fit] / int_mean,
+            # sigma=noise_est[sub_fit],
+            k2,
+            Ik / int_mean,
+            sigma=weights_fit,
             p0=coefs,
             xtol=1e-8,
             bounds=(lb, ub),
         )[0]
     else:
         coefs = curve_fit(
             scattering_model,
-            k2[sub_fit],
-            Ik[sub_fit] / int_mean,
-            sigma=noise_est[sub_fit],
+            # k2[sub_fit],
+            # Ik[sub_fit] / int_mean,
+            # sigma=noise_est[sub_fit],
+            k2,
+            Ik / int_mean,
+            sigma=weights_fit,
             p0=coefs,
             xtol=1e-8,
             bounds=(lb, ub),
@@ -369,20 +424,7 @@ def calculate_pair_dist_function(
     # fk = scattering_model(k2, coefs_fk)
     bg = scattering_model(k2, coefs)
     fk = bg - coefs[0]
-
-    # mask for structure factor estimate
-    if k_max is None:
-        k_max = np.max(k)
-    mask = np.clip(
-        np.minimum(
-            (k - 0.0) / k_width,
-            (k_max - k) / k_width,
-        ),
-        0,
-        1,
-    )
-    mask = np.sin(mask * (np.pi / 2))
-
+
     # Estimate the reduced structure factor S(k)
     Sk = (Ik - bg) * k / fk
 
@@ -408,14 +450,40 @@ def calculate_pair_dist_function(
             axis=0,
         )
     )
+    pdf_reduced[0] = 0.0
+
+    # # Damp the unphysical fluctuations at the PDF origin
+    # if damp_origin_fluctuations:
+    #     # Find radial value of primary peak
+    #     ind_max = np.argmax(pdf_reduced)
+    #     r_max = r[ind_max]
+
+    #     # find adjacent local minimum
+    #     inds = argrelextrema(pdf_reduced, np.less)
+    #     r_local_min = r[inds]
+    #     r_thresh = np.max(r_local_min[r_local_min < r_max])
+    #     ind_thresh = np.argmin(np.abs(r_thresh - r))
+    #     pdf_thresh = pdf_reduced[ind_thresh]
+    #     # r_thresh = np.max(r_local_min[r_local_min < r_max])
+
+    #     # replace low r values with linear fit
+    #     m = pdf_thresh / r_thresh
+    #     pdf_reduced[r < r_thresh] = r[r < r_thresh] * m
+
+        # pdf_reduced[r < r[ind_thresh]] = m * r[ind_thresh]
+
+        # r_ind_max = r[ind_max]
+        # r_mask = np.minimum(r / r_ind_max, 1.0)
+        # r_mask = np.sin(r_mask * np.pi / 2) ** 2
+        # pdf_reduced *= r_mask
+
 
-    # Damp the unphysical fluctuations at the PDF origin
-    if damp_origin_fluctuations:
-        ind_max = np.argmax(pdf_reduced)
-        r_ind_max = r[ind_max]
-        r_mask = np.minimum(r / r_ind_max, 1.0)
-        r_mask = np.sin(r_mask * np.pi / 2) ** 2
-        pdf_reduced *= r_mask
+        # original version        
+        # ind_max = np.argmax(pdf_reduced)
+        # r_ind_max = r[ind_max]
+        # r_mask = np.minimum(r / r_ind_max, 1.0)
+        # r_mask = np.sin(r_mask * np.pi / 2) ** 2
+        # pdf_reduced *= r_mask
 
     # Store results
     self.pdf_r = r
@@ -433,9 +501,29 @@ def calculate_pair_dist_function(
         pdf[1:] /= 4 * np.pi * density * r[1:] * (r[1] - r[0])
         pdf += 1
 
-        # damp and clip values below zero
+
+        # Damp the unphysical fluctuations at the PDF origin
         if damp_origin_fluctuations:
-            pdf *= r_mask
+            # Find radial value of primary peak
+            ind_max = np.argmax(pdf)
+            r_max = r[ind_max]
+
+            # find adjacent local minimum
+            inds = argrelextrema(pdf, np.less)
+            r_local_min = r[inds]
+            r_thresh = np.max(r_local_min[r_local_min < r_max])
+            ind_thresh = np.argmin(np.abs(r_thresh - r))
+            pdf_thresh = pdf[ind_thresh]
+            # r_thresh = np.max(r_local_min[r_local_min < r_max])
+
+            # replace low r values with linear fit
+            m = pdf_thresh / r_thresh
+            pdf[r < r_thresh] = r[r < r_thresh] * m
+
+
+        # damp and clip values below zero
+        # if damp_origin_fluctuations:
+        #     pdf *= r_mask
         if enforce_positivity:
             pdf = np.maximum(pdf, 0.0)