ModelTrainEval/tests/test_cluster.py

"""无监督聚类的原语。"""

import pytest

torch = pytest.importorskip("torch")

from engine.cluster import (  # noqa: E402
    ClusterResult,
    centroid_distances,
    cluster_features,
    cluster_stability,
    kmeans,
    normalized_mutual_info,
    pca_reduce,
    project_2d,
    random_fit_index,
    silhouette_score,
)


@pytest.fixture
def blobs():
    """4 团分得很开的高斯，外加 200 个纯噪声维度。"""
    torch.manual_seed(0)
    centers = torch.tensor([[6.0, 0, 0], [0, 6, 0], [0, 0, 6], [-6, -6, -6]])
    signal = torch.cat([center + 0.4 * torch.randn(50, 3) for center in centers])
    noise = 0.01 * torch.randn(200, 200)
    return signal, torch.cat([signal, noise], dim=1)


def test_kmeans_recovers_well_separated_blobs(blobs):
    signal, _ = blobs
    labels, centroids, inertia = kmeans(signal, 4, n_init=5, seed=0)

    assert sorted(torch.bincount(labels).tolist()) == [50, 50, 50, 50]
    assert centroids.shape == (4, 3)
    assert inertia > 0
    assert silhouette_score(signal, labels) > 0.8


def test_kmeans_rejects_bad_arguments(blobs):
    signal, _ = blobs
    with pytest.raises(ValueError, match="至少为 2"):
        kmeans(signal, 1)
    with pytest.raises(ValueError, match="少于簇数"):
        kmeans(signal[:3], 4)


def test_nmi_is_invariant_to_label_renaming():
    labels = [0, 0, 1, 1, 2, 2]
    assert normalized_mutual_info(labels, labels) == pytest.approx(1.0)
    assert normalized_mutual_info(labels, [2, 2, 0, 0, 1, 1]) == pytest.approx(1.0)

    torch.manual_seed(0)
    unrelated = torch.randint(0, 2, (200,)).tolist()
    assert normalized_mutual_info([0] * 100 + [1] * 100, unrelated) < 0.1


def test_whitening_amplifies_noise_directions(blobs):
    """回归测试：pca 白化 + 较大的 dim 会把噪声方向抬到与信号同权，簇随即塌掉。

    这正是默认 `whiten=False` 的原因——早期版本默认开白化，真实数据上
    最大的簇吃掉了 96% 的样本。
    """
    _, noisy = blobs

    clean = cluster_features(noisy, 4, pca_dim=32, whiten=False)
    whitened = cluster_features(noisy, 4, pca_dim=32, whiten=True)

    assert clean.silhouette > 0.8
    assert sorted(clean.sizes) == [50, 50, 50, 50]
    assert whitened.silhouette < clean.silhouette
    assert max(whitened.sizes) > max(clean.sizes)


def test_cluster_features_is_insensitive_to_pca_dim_without_whitening(blobs):
    _, noisy = blobs
    scores = [cluster_features(noisy, 4, pca_dim=dim).silhouette for dim in (8, 32, 64)]
    assert all(score > 0.8 for score in scores)


def test_pca_reduce_shapes_and_whitening_effect(blobs):
    _, noisy = blobs
    projected = pca_reduce(noisy, 16)
    assert projected.shape == (200, 16)
    # 不白化时各主成分方差递减；白化后大致拉平
    variances = projected.var(dim=0)
    assert variances[0] > variances[-1]
    whitened = pca_reduce(noisy, 16, whiten=True).var(dim=0)
    assert float(whitened.max() / whitened.min()) < float(variances.max() / variances.min())


def test_stability_is_high_on_real_structure_and_low_on_noise(blobs):
    signal, _ = blobs
    torch.manual_seed(1)
    assert cluster_stability(signal, 4, seeds=(0, 7), pca_dim=0) > 0.9
    assert cluster_stability(torch.randn(300, 16), 4, seeds=(0, 7), pca_dim=0) < 0.5


def test_cluster_features_returns_full_result(blobs):
    signal, _ = blobs
    result = cluster_features(signal, 4, pca_dim=0)

    assert isinstance(result, ClusterResult)
    assert result.labels.shape == (200,)
    assert result.centroids.shape[0] == 4
    assert result.distances.shape == (200,)
    assert sum(result.sizes) == 200


def test_fit_index_only_fits_a_subset_but_assigns_everything(blobs):
    signal, _ = blobs
    index = random_fit_index(200, 60, seed=0)

    assert index is not None and len(index) == 60
    assert random_fit_index(200, 500) is None  # 样本本来就少于上限
    assert random_fit_index(200, 0) is None  # 0 表示不采样

    result = cluster_features(signal, 4, pca_dim=0, fit_index=index)
    assert result.labels.shape == (200,)  # 拟合用 60 个，分配仍是全部
    assert sorted(result.sizes) == [50, 50, 50, 50]


def test_silhouette_subsamples_large_inputs(blobs):
    signal, _ = blobs
    labels = kmeans(signal, 4, n_init=3, seed=0)[0]
    assert silhouette_score(signal, labels, max_samples=50) > 0.5
    assert torch.isnan(torch.tensor(silhouette_score(signal, torch.zeros_like(labels))))


def test_cluster_result_keeps_the_space_kmeans_ran_in(blobs):
    """可视化要复用 KMeans 实际使用的空间，否则画出来的距离和聚类依据对不上。"""
    _, noisy = blobs
    result = cluster_features(noisy, 4, pca_dim=16)

    assert result.points is not None
    assert result.points.shape == (200, 16)
    assert result.centroids.shape[1] == result.points.shape[1]


def test_centroid_distances_is_a_proper_metric_matrix(blobs):
    signal, _ = blobs
    result = cluster_features(signal, 4, pca_dim=0)
    matrix = centroid_distances(result.centroids)

    assert matrix.shape == (4, 4)
    assert torch.allclose(matrix, matrix.T)                      # 对称
    assert torch.allclose(matrix.diagonal(), torch.zeros(4), atol=1e-5)  # 对角为 0
    assert matrix[matrix > 0].min() > 0


def test_project_2d_shapes_and_shared_basis(blobs):
    signal, _ = blobs
    result = cluster_features(signal, 4, pca_dim=0)
    points, centroids, ratio = project_2d(result.points, result.centroids)

    assert points.shape == (200, 2)
    assert centroids.shape == (4, 2)
    assert len(ratio) == 2 and 0 <= sum(ratio) <= 1.0001
    # 分离良好的团：前两个主成分应当解释大部分方差
    assert sum(ratio) > 0.5

    # 质心必须用与散点相同的变换，否则图上质心会飘到簇外
    for cluster in range(4):
        members = points[result.labels == cluster]
        assert torch.allclose(members.mean(dim=0), centroids[cluster], atol=1e-4)