feat: add plot_feature_clustering and plot_cluster_populations

Add clustering visualization helpers to mdpp.plots:
- plot_feature_clustering: scatter of PCA/TICA projections colored by
  cluster labels with cluster centers overlaid
- plot_cluster_populations: bar chart of frame counts per cluster,
  ranked largest-first, with noise frame exclusion

Update examples/gromacs/clustering notebook to use the new functions
instead of inline plotting code.

Author: FridrichMethod

examples/gromacs/clustering.ipynb

@@ -12,30 +12,7 @@
    "id": "3c411c15",
    "metadata": {},
    "outputs": [],
-   "source": [
-    "from __future__ import annotations\n",
-    "\n",
-    "from pathlib import Path\n",
-    "\n",
-    "import matplotlib.pyplot as plt\n",
-    "import numpy as np\n",
-    "from mplplots.utils import auto_ticks\n",
-    "\n",
-    "from mdpp.analysis.clustering import (\n",
-    "    DBSCAN,\n",
-    "    HDBSCAN,\n",
-    "    Gromos,\n",
-    "    Hierarchical,\n",
-    "    KMeans,\n",
-    "    MiniBatchKMeans,\n",
-    "    RegularSpace,\n",
-    "    compute_rmsd_matrix,\n",
-    ")\n",
-    "from mdpp.analysis.decomposition import compute_pca, featurize_backbone_torsions\n",
-    "from mdpp.core.trajectory import align_trajectory, load_trajectory\n",
-    "\n",
-    "plt.style.use(\"mplplots.styles.GraphPadPrism\")"
-   ]
+   "source": "from __future__ import annotations\n\nfrom pathlib import Path\n\nimport matplotlib.pyplot as plt\nimport numpy as np\nfrom mplplots.utils import auto_ticks\n\nfrom mdpp.analysis.clustering import (\n    DBSCAN,\n    HDBSCAN,\n    Gromos,\n    Hierarchical,\n    KMeans,\n    MiniBatchKMeans,\n    RegularSpace,\n    compute_rmsd_matrix,\n)\nfrom mdpp.analysis.decomposition import compute_pca, featurize_backbone_torsions\nfrom mdpp.core.trajectory import align_trajectory, load_trajectory\nfrom mdpp.plots import plot_cluster_populations, plot_feature_clustering\n\nplt.style.use(\"mplplots.styles.GraphPadPrism\")"
   },
   {
    "cell_type": "code",
@@ -213,24 +190,7 @@
    "id": "bde64aa1",
    "metadata": {},
    "outputs": [],
-   "source": [
-    "fig, axes = plt.subplots(2, 2, figsize=(12, 8), dpi=120, sharey=True)\n",
-    "\n",
-    "for ax, (name, r) in zip(axes.ravel(), results.items()):\n",
-    "    valid = r.labels[r.labels >= 0]\n",
-    "    if len(valid) > 0:\n",
-    "        counts = np.bincount(valid)\n",
-    "        top_k = min(20, len(counts))\n",
-    "        ax.bar(range(top_k), counts[:top_k])\n",
-    "    ax.set_xlabel(\"Cluster\")\n",
-    "    ax.set_title(f\"{name} ({r.n_clusters} clusters)\")\n",
-    "    auto_ticks(ax)\n",
-    "\n",
-    "axes[0, 0].set_ylabel(\"Frames\")\n",
-    "axes[1, 0].set_ylabel(\"Frames\")\n",
-    "fig.suptitle(f\"Cluster Populations (cutoff = {CUTOFF_NM} nm)\", y=1.02)\n",
-    "fig.tight_layout()"
-   ]
+   "source": "fig, axes = plt.subplots(2, 2, figsize=(12, 8), dpi=120, sharey=True)\n\nfor ax, (name, r) in zip(axes.ravel(), results.items()):\n    plot_cluster_populations(r, top_k=20, ax=ax)\n    ax.set_title(f\"{name} ({r.n_clusters} clusters)\")\n    auto_ticks(ax)\n\naxes[0, 0].set_ylabel(\"Frames\")\naxes[1, 0].set_ylabel(\"Frames\")\nfig.suptitle(f\"Cluster Populations (cutoff = {CUTOFF_NM} nm)\", y=1.02)\nfig.tight_layout()"
   },
   {
    "cell_type": "markdown",
@@ -279,34 +239,7 @@
    "id": "c0bc4780",
    "metadata": {},
    "outputs": [],
-   "source": [
-    "fig, axes = plt.subplots(1, 3, figsize=(16, 4.5), dpi=120)\n",
-    "\n",
-    "for ax, (name, r) in zip(axes, [(\"KMeans\", km), (\"MiniBatch\", mb), (\"RegularSpace\", rs)]):\n",
-    "    sc = ax.scatter(\n",
-    "        pca.projections[:, 0],\n",
-    "        pca.projections[:, 1],\n",
-    "        c=r.labels,\n",
-    "        cmap=\"tab10\",\n",
-    "        s=2,\n",
-    "        alpha=0.4,\n",
-    "        rasterized=True,\n",
-    "    )\n",
-    "    ax.scatter(\n",
-    "        r.cluster_centers[:, 0],\n",
-    "        r.cluster_centers[:, 1],\n",
-    "        c=\"black\",\n",
-    "        marker=\"x\",\n",
-    "        s=100,\n",
-    "        linewidths=2,\n",
-    "        zorder=5,\n",
-    "    )\n",
-    "    ax.set_xlabel(\"PC1\")\n",
-    "    ax.set_ylabel(\"PC2\")\n",
-    "    ax.set_title(f\"{name} ({r.n_clusters} clusters)\")\n",
-    "\n",
-    "fig.tight_layout()"
-   ]
+   "source": "fig, axes = plt.subplots(1, 3, figsize=(16, 4.5), dpi=120)\n\nfor ax, (name, r) in zip(axes, [(\"KMeans\", km), (\"MiniBatch\", mb), (\"RegularSpace\", rs)]):\n    plot_feature_clustering(r, pca, s=2, alpha=0.4, ax=ax)\n    ax.set_title(f\"{name} ({r.n_clusters} clusters)\")\n\nfig.tight_layout()"
   },
   {
    "cell_type": "markdown",

src/mdpp/plots/__init__.py

@@ -1,5 +1,6 @@
 """Plotting helpers for MD post-analysis outputs."""
 
+from mdpp.plots.clustering import plot_cluster_populations, plot_feature_clustering
 from mdpp.plots.contacts import contact_frequency_to_matrix, plot_contact_map
 from mdpp.plots.fes import plot_fes
 from mdpp.plots.matrix import plot_dccm
@@ -31,11 +32,13 @@
     "draw_mols",
     "get_highlight_bonds",
     "make_atom_labels_3d",
+    "plot_cluster_populations",
     "plot_contact_map",
     "plot_dccm",
     "plot_delta_rmsf",
     "plot_distances",
     "plot_energy",
+    "plot_feature_clustering",
     "plot_fes",
     "plot_hbond_counts",
     "plot_hbond_occupancy",

src/mdpp/plots/clustering.py

@@ -0,0 +1,134 @@
+"""Clustering visualization helpers."""
+
+from __future__ import annotations
+
+import numpy as np
+from matplotlib.axes import Axes
+from numpy.typing import NDArray
+
+from mdpp.analysis.clustering import ClusteringResult, FeatureClusteringResult
+from mdpp.analysis.decomposition import PCAResult, TICAResult
+from mdpp.plots.utils import get_axis
+
+
+def plot_feature_clustering(
+    result: FeatureClusteringResult,
+    projections: PCAResult | TICAResult | NDArray[np.floating],
+    *,
+    x_index: int = 0,
+    y_index: int = 1,
+    cmap: str = "tab10",
+    center_color: str = "black",
+    center_marker: str = "x",
+    center_size: float = 80.0,
+    center_linewidths: float = 2.0,
+    s: float = 2.0,
+    alpha: float = 0.4,
+    rasterized: bool = True,
+    show_centers: bool = True,
+    ax: Axes | None = None,
+) -> Axes:
+    """Scatter plot of projections colored by cluster labels with centers.
+
+    Args:
+        result: Feature clustering result (KMeans, MiniBatchKMeans, or
+            RegularSpace).
+        projections: 2-D projection coordinates. Accepts a ``PCAResult``,
+            ``TICAResult``, or a raw ``(n_samples, n_features)`` array.
+        x_index: Column index for the x-axis.
+        y_index: Column index for the y-axis.
+        cmap: Colormap for cluster labels.
+        center_color: Color for cluster center markers.
+        center_marker: Marker style for cluster centers.
+        center_size: Marker size for cluster centers.
+        center_linewidths: Line width for cluster center markers.
+        s: Marker size for data points.
+        alpha: Marker transparency for data points.
+        rasterized: Whether to rasterize the scatter layer (recommended
+            for large trajectories to keep file sizes small).
+        show_centers: Whether to overlay cluster center markers.
+        ax: Optional matplotlib axis.
+
+    Returns:
+        The matplotlib axis with the scatter plot.
+    """
+    axis = get_axis(ax)
+
+    if isinstance(projections, PCAResult | TICAResult):
+        coords = projections.projections
+    else:
+        coords = np.asarray(projections)
+
+    x = coords[:, x_index]
+    y = coords[:, y_index]
+
+    axis.scatter(
+        x,
+        y,
+        c=result.labels,
+        cmap=cmap,
+        s=s,
+        alpha=alpha,
+        edgecolors="none",
+        rasterized=rasterized,
+    )
+
+    if show_centers:
+        axis.scatter(
+            result.cluster_centers[:, x_index],
+            result.cluster_centers[:, y_index],
+            c=center_color,
+            marker=center_marker,
+            s=center_size,
+            linewidths=center_linewidths,
+            zorder=5,
+        )
+
+    if isinstance(projections, PCAResult):
+        axis.set_xlabel(f"PC{x_index + 1}")
+        axis.set_ylabel(f"PC{y_index + 1}")
+    elif isinstance(projections, TICAResult):
+        axis.set_xlabel(f"IC{x_index + 1}")
+        axis.set_ylabel(f"IC{y_index + 1}")
+    else:
+        axis.set_xlabel(f"Component {x_index + 1}")
+        axis.set_ylabel(f"Component {y_index + 1}")
+
+    return axis
+
+
+def plot_cluster_populations(
+    result: ClusteringResult | FeatureClusteringResult,
+    *,
+    top_k: int = 20,
+    color: str | None = None,
+    ax: Axes | None = None,
+) -> Axes:
+    """Bar chart of cluster populations (frame counts per cluster).
+
+    Args:
+        result: Clustering result from any method.
+        top_k: Maximum number of clusters to show (largest first).
+        color: Bar color. If ``None``, uses the default color cycle.
+        ax: Optional matplotlib axis.
+
+    Returns:
+        The matplotlib axis with the bar chart.
+    """
+    axis = get_axis(ax)
+
+    valid = result.labels[result.labels >= 0]
+    if len(valid) > 0:
+        counts = np.bincount(valid)
+        ranked = np.argsort(counts)[::-1]
+        n_show = min(top_k, len(counts))
+        ranked_counts = counts[ranked[:n_show]]
+        if color is not None:
+            axis.bar(range(n_show), ranked_counts, color=color)
+        else:
+            axis.bar(range(n_show), ranked_counts)
+
+    axis.set_xlabel("Cluster (ranked)")
+    axis.set_ylabel("Frames")
+
+    return axis