🚚 move PCA plot from acore normalization example (#8)

Author: enryH

src/vuecore/decomposition.py

@@ -1,8 +1,10 @@
 """Decompositon plots like pca, umap, tsne, etc."""
 
+import itertools
 from typing import Optional
 
 import matplotlib
+import matplotlib.pyplot as plt
 import pandas as pd
 import sklearn.decomposition
 
@@ -17,3 +19,53 @@ def plot_explained_variance(
     exp_var.index.name = "PC"
     ax = exp_var.plot(ax=ax)
     return ax
+
+
+def pca_grid(
+    PCs: pd.DataFrame,
+    meta_column: pd.Series,
+    n_components: int = 4,
+    meta_col_name: Optional[str] = None,
+    figsize=(6, 8),
+) -> plt.Figure:
+    """Plot a grid of scatter plots for the first n_components of PCA,  per default 4.
+
+    Parameters
+    ----------
+    PCs : pd.DataFrame
+        DataFrame with the principal components as columns.
+    meta_column : pd.Series
+        Series with categorical data to color the scatter plots.
+    n_components : int, optional
+        Number of first n components to plot, by default 4
+    meta_col_name : Optional[str], optional
+        If another name than the default series name shoudl be used, by default None
+
+    Returns
+    -------
+    plt.Figure
+        Matplotlib figure with the scatter plots.
+    """
+    if meta_col_name is None:
+        meta_col_name = meta_column.name
+    else:
+        meta_column = meta_column.rename(meta_col_name)
+    up_to = min(PCs.shape[-1], n_components)
+    fig, axes = plt.subplots(up_to - 1, 2, figsize=figsize, layout="constrained")
+    PCs = PCs.join(
+        meta_column.astype("category")
+    )  # ! maybe add a check that it's not continous
+    for k, (pos, ax) in enumerate(
+        zip(itertools.combinations(range(up_to), 2), axes.flatten())
+    ):
+        i, j = pos
+        plot_heatmap = bool(k % 2)
+        PCs.plot.scatter(
+            i,
+            j,
+            c=meta_col_name,
+            cmap="Paired",
+            ax=ax,
+            colorbar=plot_heatmap,
+        )
+    return fig