import logging
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns
from sklearn.decomposition import PCA
import warnings
warnings.filterwarnings("ignore", category=RuntimeWarning)
logging.basicConfig(level=logging.INFO, format="%(levelname)s - %(message)s")
logger = logging.getLogger(__name__)
[docs]
def compute_pca(
df,
feature_cols=None,
n_components=2,
plot=True,
s=5,
color="C0",
edgecolors="black",
context="notebook",
figsize=(10, 5),
ax=None,
):
"""
Reduce the dimensionality of the data using PCA.
Parameters
----------
df : pandas.DataFrame
Data to reduce. If `feature_cols` is given, only those columns are
used to fit PCA (the returned dataframe keeps ALL of `df`'s
original columns, plus the PCA components) -- use this whenever
`df` also carries metadata columns (sim_name, replica, frame,
...) that must not be fed to PCA itself. If `feature_cols` is
None (default, backward-compatible), `df` is used as-is -- every
column must be numeric, or `PCA.fit_transform` will raise.
feature_cols : list of str, optional
Column names in `df` to actually reduce. Strongly recommended
whenever `df` has metadata columns beyond the raw numeric
features. Default None (use all of `df`).
n_components : int, optional
Number of components to keep. The default is 2.
plot : bool, optional
Plot the two first components. The default is True.
s : int, optional
Marker size for the scatter plot. The default is 5.
color : str, optional
Marker color for the scatter plot. The default is "C0".
context : str, optional
Seaborn plotting context. The default is "notebook".
figsize : tuple, optional
Figure size. The default is (10, 5).
ax : matplotlib.axes._subplots.AxesSubplot, optional
Axes on which to plot. The default is None.
Returns
-------
new_df : pandas.DataFrame
`df` (all original columns preserved) with the PCA components
added (columns "PC1", "PC2", ...).
fig : matplotlib.figure.Figure or None
Figure containing the plot (if plot=True) -- always retrievable
as an object for further tweaking, e.g. `fig.savefig(...)`.
pca : sklearn.decomposition.PCA
The fitted PCA object (e.g. for `pca.explained_variance_ratio_`,
`pca.components_`).
"""
X = df[feature_cols] if feature_cols is not None else df
logger.info("Running PCA with n_components=%d", n_components)
pca = PCA(n_components=n_components)
reduced = pca.fit_transform(X)
columns = [f"PC{i + 1}" for i in range(reduced.shape[1])]
explained = pca.explained_variance_ratio_
for i, ratio in enumerate(explained):
logger.info(
"PC%d explained variance ratio: %.2f%%",
i + 1,
ratio * 100,
)
logger.info(
"PCA completed. Total explained variance: %.2f%%",
np.sum(explained) * 100,
)
# Positional assignment: reduced preserves X's (and therefore df's)
# row order, so a plain array lines up correctly with new_df's rows.
new_df = df.copy()
new_df[columns] = reduced
fig = None
if plot:
logger.info("Plotting PCA results...")
sns.set_context(context)
if ax is None:
fig, ax = plt.subplots(figsize=figsize)
else:
fig = ax.figure
ax.scatter(
new_df[columns[0]],
new_df[columns[1]],
s=s,
color=color,
edgecolors=edgecolors,
)
ax.set_title("PCA")
ax.set_xlabel(f"{columns[0]} ({100 * explained[0]:.1f}%)")
ax.set_ylabel(f"{columns[1]} ({100 * explained[1]:.1f}%)")
fig.tight_layout()
return new_df, fig, pca