Adding background PCA function to polar datacubes

py4DSTEM/process/polar/polar_analysis.py

@@ -4,6 +4,7 @@
 import matplotlib.pyplot as plt
 from scipy.optimize import curve_fit
 from scipy.ndimage import gaussian_filter
+from sklearn.decomposition import PCA
 
 from emdfile import tqdmnd
 
@@ -924,3 +925,66 @@ def scattering_model(k2, *coefs):
     # int1*np.exp(k2/(-2*sigma1**2))
 
     return int_model
+
+
+def background_pca(
+    self,
+    pca_index=0,
+    intensity_range=(0, 1),
+    normalize_mean=True,
+    normalize_std=True,
+    plot_result=True,
+    plot_coef=False,
+):
+    """
+    Generate PCA decompositions of the background signal
+
+    Parameters
+    --------
+    progress_bar: bool
+        Turns on the progress bar for the polar transformation
+
+    Returns
+    --------
+    im_pca: np,array
+        rgb image array
+    coef_pca: np.array
+        radial PCA component selected
+
+    """
+
+    # PCA decomposition
+    shape = self.radial_all.shape
+    A = np.reshape(self.radial_all, (shape[0] * shape[1], shape[2]))
+    if normalize_mean:
+        A -= np.mean(A, axis=0)
+    if normalize_std:
+        A /= np.std(A, axis=0)
+    pca = PCA(n_components=np.maximum(pca_index + 1, 2))
+    pca.fit(A)
+
+    components = pca.components_
+    loadings = pca.transform(A)
+
+    # output image data
+    sig_pca = np.reshape(loadings[:, pca_index], shape[0:2])
+    sig_pca -= intensity_range[0]
+    sig_pca /= intensity_range[1] - intensity_range[0]
+    sig_pca = np.clip(sig_pca, 0, 1)
+    im_pca = np.tile(sig_pca[:, :, None], (1, 1, 3))
+
+    # output PCA coefficient
+    coef_pca = np.vstack((self.radial_bins, components[pca_index, :])).T
+
+    if plot_result:
+        fig, ax = plt.subplots(figsize=(8, 8))
+        ax.imshow(
+            im_pca,
+            vmin=0,
+            vmax=5,
+        )
+    if plot_coef:
+        fig, ax = plt.subplots(figsize=(8, 4))
+        ax.plot(coef_pca[:, 0], coef_pca[:, 1])
+
+    return im_pca, coef_pca

py4DSTEM/process/polar/polar_datacube.py

@@ -104,6 +104,7 @@ def __init__(
         plot_sf_estimate,
         plot_reduced_pdf,
         plot_pdf,
+        background_pca,
     )
     from py4DSTEM.process.polar.polar_peaks import (
         find_peaks_single_pattern,