Первичная версия: стилометрия реплик, LSA и протокол экспериментов

lab1/stylometry/clustering.py

@@ -0,0 +1,90 @@
+"""Кластеризация персонажей в стилометрическом пространстве."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Sequence
+
+import numpy as np
+import pandas as pd
+from scipy.sparse import issparse
+from sklearn.cluster import AgglomerativeClustering, KMeans
+from sklearn.metrics import (
+    adjusted_rand_score,
+    calinski_harabasz_score,
+    davies_bouldin_score,
+    normalized_mutual_info_score,
+    silhouette_score,
+)
+from sklearn.preprocessing import normalize
+
+from .pipelines import FeatureConfig, build_feature_union
+
+
+@dataclass
+class ClusteringResult:
+    metrics: pd.DataFrame
+    centroids: pd.DataFrame
+
+
+def _centroid_matrix(matrix: np.ndarray, labels: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
+    classes = np.unique(labels)
+    centroids = np.vstack([matrix[labels == label].mean(axis=0) for label in classes])
+    return classes, centroids
+
+
+def run_clustering(
+    X: pd.DataFrame,
+    y: Sequence,
+    feature_config: FeatureConfig,
+    n_clusters: int | None = None,
+    seed: int = 42,
+) -> ClusteringResult:
+    """Кластеризует центроиды персонажей и считает внешние метрики."""
+    y = np.asarray(y)
+    features = build_feature_union(feature_config)
+    matrix = features.fit_transform(X)
+    matrix = matrix.toarray() if issparse(matrix) else np.asarray(matrix)
+    matrix = normalize(matrix, norm="l2", copy=False)
+    classes, centroids = _centroid_matrix(matrix, y)
+    # Если число кластеров равно числу персонажей, кластеризация тривиальна,
+    # а силуэт не определён. По умолчанию берём примерно корень из числа
+    # классов: для 34 персонажей это 6 групп.
+    n_clusters = int(n_clusters or max(2, round(len(classes) ** 0.5)))
+    n_clusters = min(n_clusters, len(classes) - 1)
+
+    rows: list[dict[str, Any]] = []
+    for name, model in (
+        (
+            "agglomerative_cosine",
+            AgglomerativeClustering(
+                n_clusters=n_clusters, metric="cosine", linkage="average"
+            ),
+        ),
+        (
+            "kmeans",
+            KMeans(n_clusters=n_clusters, random_state=seed, n_init=20),
+        ),
+    ):
+        predicted = model.fit_predict(centroids)
+        n_predicted = len(set(predicted))
+        valid_partition = 2 <= n_predicted < len(centroids)
+        rows.append(
+            {
+                "method": name,
+                "n_clusters": n_clusters,
+                "adjusted_rand_index": float(adjusted_rand_score(classes, predicted)),
+                "normalized_mutual_info": float(normalized_mutual_info_score(classes, predicted)),
+                "silhouette_cosine": float(
+                    silhouette_score(centroids, predicted, metric="cosine")
+                ) if valid_partition else 0.0,
+                "calinski_harabasz": float(calinski_harabasz_score(centroids, predicted))
+                if valid_partition else 0.0,
+                "davies_bouldin": float(davies_bouldin_score(centroids, predicted))
+                if valid_partition else 0.0,
+            }
+        )
+
+    centroid_frame = pd.DataFrame(centroids, index=classes)
+    centroid_frame.index.name = "character"
+    return ClusteringResult(metrics=pd.DataFrame(rows), centroids=centroid_frame)