"""Deep Ensemble Clustering — application-level realisation of EnFormer's idea. EnFormer (the backbone in EnFormer-main) builds feature extraction around two steps: (i) *Ensemble Generation*, where several differentiable base clustering methods capture diverse semantic structures, and (ii) *Consensus Aggregation*, which fuses those base clusterings into a refined result. This module mirrors that design at the level of a practical image-clustering pipeline: Ensemble Generation Run a diverse pool of base clusterers -- Partitional (K-Means), Fuzzy C-Means, Gaussian-Mixture (probabilistic) and Spectral / Agglomerative -- echoing EnFormer's PartitionalClustering, FuzzyClustering and ProbabilisticClustering base modules. Diversity is injected via varied cluster counts, random seeds and random feature subspaces. Consensus Aggregation Accumulate the base partitions into a co-association matrix (Evidence Accumulation Clustering, Fred & Jain 2005) and derive the final consensus partition from it. This is the differentiable-in-spirit fusion step. Everything runs on top of CPU-friendly scikit-learn so it deploys cleanly on a Hugging Face Space. """ from dataclasses import dataclass, field from typing import List, Optional, Sequence, Dict import numpy as np from sklearn.cluster import KMeans, AgglomerativeClustering, SpectralClustering from sklearn.mixture import GaussianMixture from sklearn.decomposition import PCA from sklearn.preprocessing import StandardScaler, normalize from sklearn.metrics import ( silhouette_score, davies_bouldin_score, calinski_harabasz_score, adjusted_rand_score, normalized_mutual_info_score, ) # --------------------------------------------------------------------------- FCM def fuzzy_cmeans(X: np.ndarray, k: int, m: float = 1.8, n_iter: int = 60, seed: int = 0, tol: float = 1e-4) -> np.ndarray: """Minimal Fuzzy C-Means; returns a hard partition (argmax membership). Mirrors EnFormer's FuzzyClustering base module (soft membership assignment). """ rng = np.random.default_rng(seed) n = X.shape[0] U = rng.random((n, k)) U /= U.sum(axis=1, keepdims=True) prev = None for _ in range(n_iter): Um = U ** m centers = (Um.T @ X) / (Um.sum(axis=0)[:, None] + 1e-12) d = np.linalg.norm(X[:, None, :] - centers[None, :, :], axis=2) + 1e-12 power = 2.0 / (m - 1.0) ratio = d[:, :, None] / d[:, None, :] U = 1.0 / (ratio ** power).sum(axis=2) if prev is not None and np.abs(U - prev).max() < tol: break prev = U.copy() return U.argmax(axis=1) # ----------------------------------------------------------------- preprocessing def preprocess_features(X: np.ndarray, pca_dim: Optional[int] = 50, whiten: bool = True, l2: bool = True) -> np.ndarray: """Standardise -> optional PCA(whiten) -> optional L2 normalise. PCA denoising markedly improves cluster separation on high-dim deep features; L2 makes Euclidean distance behave like cosine (right for deep embeddings). """ Z = StandardScaler().fit_transform(X) if pca_dim is not None and pca_dim < min(Z.shape): Z = PCA(n_components=pca_dim, whiten=whiten, random_state=0).fit_transform(Z) if l2: Z = normalize(Z) return Z.astype(np.float32) # --------------------------------------------------------------------- container @dataclass class ClusterResult: labels: np.ndarray n_clusters: int metrics: Dict[str, float] co_association: Optional[np.ndarray] = None base_partitions: List[np.ndarray] = field(default_factory=list) embedding_2d: Optional[np.ndarray] = None Z: Optional[np.ndarray] = None # ---------------------------------------------------------------------- ensemble class EnsembleClusterer: """Deep ensemble clustering with consensus aggregation. Parameters ---------- n_clusters : target number of consensus clusters. If None, chosen by silhouette over `k_range`. base_methods : subset of {"kmeans","fuzzy","gmm","spectral","agglomerative"}. n_runs : number of ensemble members per method. subspace_frac : fraction of feature dims sampled per member (random subspace). k_jitter : base members use K in [n_clusters - k_jitter, n_clusters + k_jitter]. """ def __init__(self, n_clusters: Optional[int] = None, base_methods: Sequence[str] = ("kmeans", "fuzzy", "gmm", "spectral"), n_runs: int = 6, subspace_frac: float = 0.8, k_jitter: int = 2, k_range: Sequence[int] = range(2, 9), random_state: int = 0): self.n_clusters = n_clusters self.base_methods = tuple(base_methods) self.n_runs = n_runs self.subspace_frac = subspace_frac self.k_jitter = k_jitter self.k_range = list(k_range) self.random_state = random_state # ---------------------------------------------------------------- base member def _base_partition(self, X: np.ndarray, method: str, k: int, seed: int) -> np.ndarray: rng = np.random.default_rng(seed) d = X.shape[1] n_sub = max(2, int(round(self.subspace_frac * d))) cols = rng.choice(d, size=n_sub, replace=False) Xs = X[:, cols] k = max(2, min(k, X.shape[0] - 1)) if method == "kmeans": return KMeans(k, n_init=5, random_state=seed).fit_predict(Xs) if method == "fuzzy": return fuzzy_cmeans(Xs, k, seed=seed) if method == "gmm": return GaussianMixture(k, covariance_type="diag", n_init=1, random_state=seed, reg_covar=1e-4).fit_predict(Xs) if method == "spectral": return SpectralClustering(k, affinity="nearest_neighbors", n_neighbors=12, assign_labels="kmeans", random_state=seed).fit_predict(Xs) if method == "agglomerative": return AgglomerativeClustering(k, linkage="ward").fit_predict(Xs) raise ValueError(f"unknown base method {method!r}") # -------------------------------------------------------------- ensemble gen. def _generate_ensemble(self, X: np.ndarray, k: int) -> List[np.ndarray]: parts: List[np.ndarray] = [] s = self.random_state for method in self.base_methods: for r in range(self.n_runs): jitter = 0 if self.k_jitter == 0 else ((r % (2 * self.k_jitter + 1)) - self.k_jitter) kk = max(2, k + jitter) try: parts.append(self._base_partition(X, method, kk, seed=s)) except Exception: pass # a single failed member must not sink the ensemble s += 1 return parts # ------------------------------------------------------ consensus aggregation @staticmethod def _co_association(parts: List[np.ndarray], n: int) -> np.ndarray: """Fraction of partitions in which each pair of points is co-clustered.""" M = np.zeros((n, n), dtype=np.float32) for p in parts: # one-hot co-membership without an n x n python loop same = (p[:, None] == p[None, :]) M += same M /= max(1, len(parts)) return M def _consensus(self, M: np.ndarray, k: int) -> np.ndarray: dist = 1.0 - M np.fill_diagonal(dist, 0.0) return AgglomerativeClustering( n_clusters=k, metric="precomputed", linkage="average" ).fit_predict(dist) # ------------------------------------------------------------------- fit API def _score(self, Z: np.ndarray, labels: np.ndarray) -> float: if len(set(labels.tolist())) < 2: return -1.0 return silhouette_score(Z, labels) def fit(self, X: np.ndarray, ground_truth: Optional[Dict[str, np.ndarray]] = None, compute_2d: bool = True) -> ClusterResult: Z = X n = Z.shape[0] # choose K by consensus silhouette if not given if self.n_clusters is None: best_k, best_s, best = None, -2, None for k in self.k_range: parts = self._generate_ensemble(Z, k) M = self._co_association(parts, n) lab = self._consensus(M, k) s = self._score(Z, lab) if s > best_s: best_k, best_s, best = k, s, (lab, M, parts) k = best_k labels, M, parts = best else: k = self.n_clusters parts = self._generate_ensemble(Z, k) M = self._co_association(parts, n) labels = self._consensus(M, k) metrics = self.compute_metrics(Z, labels, ground_truth) emb = None if compute_2d: emb = self._embed_2d(Z) return ClusterResult(labels=labels, n_clusters=k, metrics=metrics, co_association=M, base_partitions=parts, embedding_2d=emb, Z=Z) # -------------------------------------------------------------------- extras @staticmethod def compute_metrics(Z, labels, ground_truth=None) -> Dict[str, float]: out: Dict[str, float] = {} if len(set(labels.tolist())) > 1: out["silhouette"] = float(silhouette_score(Z, labels)) out["davies_bouldin"] = float(davies_bouldin_score(Z, labels)) out["calinski_harabasz"] = float(calinski_harabasz_score(Z, labels)) if ground_truth: for name, gt in ground_truth.items(): out[f"ARI_{name}"] = float(adjusted_rand_score(gt, labels)) out[f"NMI_{name}"] = float(normalized_mutual_info_score(gt, labels)) return out @staticmethod def _embed_2d(Z: np.ndarray) -> np.ndarray: try: from sklearn.manifold import TSNE perp = min(30, max(5, Z.shape[0] // 4)) return TSNE(n_components=2, perplexity=perp, init="pca", random_state=0).fit_transform(Z) except Exception: return PCA(n_components=2, random_state=0).fit_transform(Z)