91 lines
3.4 KiB
Python
91 lines
3.4 KiB
Python
"""Кластеризация персонажей в стилометрическом пространстве."""
|
||
|
||
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)
|