import numpy as np import pandas as pd import matplotlib.pyplot as plt from collections import Counter from scipy.stats import rankdata from tqdm import tqdm from sklearn.base import BaseEstimator, ClusterMixin, clone as sklearn_clone from sklearn.metrics import (silhouette_samples, silhouette_score, calinski_harabasz_score, davies_bouldin_score, confusion_matrix) from sklearn.ensemble import RandomForestClassifier from sklearn.cluster import SpectralClustering, KMeans from umap import UMAP def consensus_cluster_labels(cluster_labels_df): """Compute consensus cluster labels from multiple clustering runs. Given a DataFrame where each column contains the cluster assignments from one run, this function finds the most common cross-run label patterns and remaps them to canonical cluster IDs. The final label for each sample is the mode (most frequent value) across runs after remapping. Parameters ---------- cluster_labels_df : pandas.DataFrame of shape (n_samples, n_runs) Each column holds the integer cluster labels from one clustering run. All runs must use the same label set (e.g., {0, 1}). Returns ------- cluster_labels : numpy.ndarray of shape (n_samples,) Consensus cluster labels (0-indexed, dense-ranked). probability : numpy.ndarray of shape (n_samples,) For each sample, the fraction of runs that agree with the assigned consensus label. Values range from ``1/n_runs`` (no agreement) to ``1.0`` (perfect agreement across all runs). Notes ----- The number of canonical clusters ``K`` is determined by the number of unique labels present in ``cluster_labels_df``. The top-K most frequent cross-run patterns are each assigned a canonical ID; remaining patterns are resolved by the mode operation. Examples -------- >>> import pandas as pd >>> runs = pd.DataFrame({0: [0, 0, 1, 1], 1: [0, 1, 1, 1], 2: [0, 0, 1, 1]}) >>> labels, probs = consensus_cluster_labels(runs) >>> labels array([0, 0, 1, 1]) >>> probs array([1. , 0.67, 1. , 1. ]) """ n_clusters = len(np.unique(cluster_labels_df.values)) # Build a string key per sample from its labels across all runs total_labels = ['-'.join(r) for r in cluster_labels_df.astype(str).values] counts = Counter(total_labels) sorted_keys = sorted(counts, key=counts.get, reverse=True) # Remap the top-K patterns to canonical cluster IDs (offset by n_clusters # to avoid collisions during the in-place update). df = cluster_labels_df.copy() for new_id, key in enumerate(sorted_keys[:n_clusters]): run_labels = np.array([int(i) for i in key.split('-')]) for run_idx, label in enumerate(run_labels): df.iloc[df.values[:, run_idx] == label, run_idx] = new_id + n_clusters df = df - n_clusters # Mode across runs → consensus label; fraction agreeing → probability cluster_mode = df.mode(axis=1)[0] probability = df.apply(lambda x: x == cluster_mode).sum(axis=1) / df.shape[1] cluster_labels = (rankdata(cluster_mode, method='dense') - 1).astype(int) return cluster_labels, probability.values def evaluate_cluster_stability( X, cluster_labels, n_iterations=100, train_fraction=0.2, n_estimators=100, max_depth=10, cluster_names=None, figsize=(12, 9), ): """Assess cluster stability via repeated random sub-sampling cross-validation. A Random Forest is trained on a random subset of samples and used to predict labels on the held-out set. After many iterations the prediction-frequency matrix reveals how consistently each cluster is recovered. Parameters ---------- X : ndarray of shape (n_samples, n_features) Feature matrix (original space or UMAP-embedded). cluster_labels : ndarray of shape (n_samples,) Ground-truth (or reference) cluster labels for each sample. n_iterations : int, default=100 Number of random train/test splits. train_fraction : float, default=0.2 Fraction of samples used for training in each iteration. n_estimators : int, default=100 Number of trees in the Random Forest. max_depth : int, default=10 Maximum tree depth. cluster_names : list of str or None, default=None Human-readable names for each cluster. If ``None``, integer labels are used. figsize : tuple of (float, float), default=(12, 9) Size of *each* subplot. Returns ------- fraction_matrix : ndarray of shape (n_samples, n_clusters) For each sample, the fraction of iterations in which it was predicted as each cluster. predicted_labels : ndarray of shape (n_samples,) The most-frequently predicted cluster for each sample. Notes ----- Two diagnostic plots are produced side by side: * **Normalized Confusion Matrix** — compares reference labels to the majority-vote predicted labels. Near-diagonal = stable clustering. * **Mean Prediction Fraction (%)** — for every true cluster, the average fraction of iterations each sample was assigned to each predicted cluster. High diagonal values indicate that the classifier reliably recovers the cluster. """ cluster_labels = np.asarray(cluster_labels) unique_labels = np.sort(np.unique(cluster_labels)) n_clusters = len(unique_labels) n_samples = len(X) if cluster_names is None: cluster_names = [str(c) for c in unique_labels] # Accumulator: for each sample, count how many times it was predicted as # each cluster across iterations. pred_counts = np.zeros((n_samples, n_clusters), dtype=float) label_to_col = {lab: i for i, lab in enumerate(unique_labels)} rf = RandomForestClassifier(n_estimators=n_estimators, max_depth=max_depth) for _ in tqdm(range(n_iterations), desc="Cross-validation iterations"): idx_train = np.random.choice( n_samples, size=int(train_fraction * n_samples), replace=False ) mask_test = np.ones(n_samples, dtype=bool) mask_test[idx_train] = False rf.fit(X[idx_train], cluster_labels[idx_train]) preds = rf.predict(X[mask_test]) for lab in unique_labels: pred_counts[mask_test, label_to_col[lab]] += (preds == lab).astype(float) # Normalize to fractions (avoid /0 for samples that were never in a test set) totals = pred_counts.sum(axis=1, keepdims=True) totals[totals == 0] = 1.0 fraction_matrix = pred_counts / totals # Most-frequently predicted label for each sample predicted_labels = unique_labels[np.argmax(fraction_matrix, axis=1)] # ---- Diagnostic Plots ------------------------------------------------ cm = confusion_matrix(cluster_labels, predicted_labels, labels=unique_labels) cm_norm = cm / cm.sum(axis=1, keepdims=True) fig, axes = plt.subplots(1, 2, figsize=(figsize[0] * 1.8, figsize[1])) # Plot 1 — Normalized confusion matrix ax = axes[0] im = ax.imshow(cm_norm, interpolation="nearest", cmap="Blues", vmin=0, vmax=1) ax.set_title("Normalized Confusion Matrix") fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04) ax.set_xlabel("Predicted Cluster") ax.set_ylabel("True Cluster") ax.set_xticks(range(n_clusters)) ax.set_xticklabels(cluster_names[:n_clusters], rotation=90) ax.set_yticks(range(n_clusters)) ax.set_yticklabels(cluster_names[:n_clusters]) for i in range(cm.shape[0]): for j in range(cm.shape[1]): ax.text( j, i, format(cm[i, j], "d"), ha="center", va="center", color="white" if cm_norm[i, j] > 0.6 else "black", fontsize=12, ) # Plot 2 — Mean prediction-fraction per true cluster frac_per_cluster = 100.0 * np.vstack( [fraction_matrix[cluster_labels == c].mean(axis=0) for c in unique_labels] ) ax2 = axes[1] im2 = ax2.imshow(frac_per_cluster, cmap="Blues", aspect="auto", vmin=0, vmax=100) ax2.set_title("Mean Predicted-Cluster Fraction (%)") fig.colorbar(im2, ax=ax2, fraction=0.046, pad=0.04, label="Percentage") ax2.set_xlabel("Predicted Cluster") ax2.set_ylabel("True Cluster") ax2.set_xticks(range(n_clusters)) ax2.set_xticklabels(cluster_names[:n_clusters], rotation=90) ax2.set_yticks(range(n_clusters)) ax2.set_yticklabels(cluster_names[:n_clusters]) for i in range(frac_per_cluster.shape[0]): for j in range(frac_per_cluster.shape[1]): ax2.text( j, i, f"{frac_per_cluster[i, j]:.1f}", ha="center", va="center", color="white" if frac_per_cluster[i, j] > 60 else "black", fontsize=12, ) plt.tight_layout() plt.show() return fraction_matrix, predicted_labels class DivisiveClustering(BaseEstimator, ClusterMixin): """Divisive (top-down) clustering using UMAP + a pluggable inner clusterer + consensus voting. The algorithm starts with all samples in a single cluster and iteratively bisects the cluster whose split yields the best improvement in the chosen quality criterion. Each bisection is performed by: 1. Embedding the cluster's samples with UMAP. 2. Running the inner clusterer (K=2) multiple times on random subsamples. 3. Propagating each run's sub-labels to the full cluster via a Random Forest. 4. Taking a consensus vote across runs to produce stable bisection labels. Splitting stops when no bisection improves the chosen quality criterion. The API mirrors scikit-learn estimators (``fit``, ``predict``, ``fit_predict``), so it can be used as a drop-in replacement for ``KMeans`` or ``SpectralClustering`` in pipelines and evaluation code. Parameters ---------- criterion : str, default='silhouette' Cluster-quality metric used to decide whether a split is accepted. Supported values: * ``'silhouette'`` — Silhouette score (higher is better, range [-1, 1]). Measures how similar each sample is to its own cluster vs. the nearest other cluster. * ``'calinski_harabasz'`` — Calinski-Harabasz index (higher is better). Ratio of between-cluster dispersion to within-cluster dispersion. Also known as the Variance Ratio Criterion. * ``'davies_bouldin'`` — Davies-Bouldin index (**lower** is better). Average similarity of each cluster with its most similar cluster. inner_clusterer : sklearn estimator or None, default=None Any scikit-learn-compatible clustering estimator that implements ``fit_predict(X)``. Must produce exactly 2 clusters per bisection. The estimator is **cloned** for each run, so the original is never mutated. If ``None``, defaults to:: SpectralClustering(n_clusters=2, affinity='rbf', eigen_solver='arpack') Common alternatives: * ``KMeans(n_clusters=2)`` * ``AgglomerativeClustering(n_clusters=2)`` * ``MiniBatchKMeans(n_clusters=2)`` n_components : int, default=3 Number of UMAP embedding dimensions. n_neighbors : int, default=60 Number of neighbors for the UMAP graph. metric : str, default='correlation' Distance metric used by UMAP. output_dens : bool, default=True If True, UMAP also estimates local radii (useful for density-aware embeddings). n_consensus_runs : int, default=10 Number of independent clustering runs used to form the consensus. subsampling : bool, default=True Whether to subsample before each inner-clustering run. n_subsample : int, default=5000 Maximum number of samples drawn (without replacement) for each inner-clustering run. If a cluster has fewer samples, all are used. Only effective when ``subsampling=True``. rf_n_estimators : int, default=100 Number of trees in the Random Forest used to propagate sub-labels and for the final ``predict`` classifier. rf_max_depth : int, default=10 Maximum depth of each tree in the Random Forest. random_state : int or None, default=None Seed for reproducibility. Controls numpy and the internal estimators. Attributes ---------- labels_ : ndarray of shape (n_samples,) Cluster label for each training sample (0-indexed). probabilities_ : ndarray of shape (n_samples,) Consensus probability for each sample's assigned label, i.e. the fraction of consensus runs that agreed on the final label. Ranges from near 0 (low confidence) to 1.0 (perfect agreement). n_clusters_ : int Total number of clusters discovered. criterion_score_ : float Final value of the chosen quality criterion. criterion_name_ : str Name of the criterion used (mirrors the ``criterion`` parameter). inner_clusterer_ : sklearn estimator The resolved inner clusterer (after applying the default if ``None`` was passed). Examples -------- >>> X = np.random.rand(500, 50) >>> model = DivisiveClustering(n_neighbors=30, n_components=3) >>> labels = model.fit_predict(X) >>> print(model.n_clusters_, model.criterion_score_) 3 0.42 >>> >>> # With KMeans + Calinski-Harabasz criterion >>> model_km = DivisiveClustering( ... inner_clusterer=KMeans(n_clusters=2), ... criterion='calinski_harabasz', ... ) >>> labels_km = model_km.fit_predict(X) """ # Map criterion names → (scoring_function, whether higher is better) _CRITERIA = { "silhouette": (silhouette_score, True), "calinski_harabasz": (calinski_harabasz_score, True), "davies_bouldin": (davies_bouldin_score, False), } def __init__( self, criterion="silhouette", inner_clusterer=None, n_components=3, n_neighbors=60, metric="correlation", output_dens=True, n_consensus_runs=10, subsampling=True, n_subsample=5000, rf_n_estimators=100, rf_max_depth=10, random_state=None, ): self.criterion = criterion self.inner_clusterer = inner_clusterer self.n_components = n_components self.n_neighbors = n_neighbors self.metric = metric self.output_dens = output_dens self.n_consensus_runs = n_consensus_runs self.subsampling = subsampling self.n_subsample = n_subsample self.rf_n_estimators = rf_n_estimators self.rf_max_depth = rf_max_depth self.random_state = random_state # --------------------------------------------------------------------- # # Internal helpers # # --------------------------------------------------------------------- # def _resolve_inner_clusterer(self, random_state=None): """Return a fresh clone of the inner clusterer. If ``self.inner_clusterer`` is None, a default SpectralClustering is created. The estimator is always cloned so the original is never mutated. Parameters ---------- random_state : int or None Passed to the clone if the estimator accepts ``random_state``. Returns ------- estimator : sklearn estimator A fresh, un-fitted clone. """ if self.inner_clusterer is None: # default: Spectral Clustering return SpectralClustering( n_clusters=2, affinity="rbf", eigen_solver="arpack", random_state=random_state, ) est = sklearn_clone(self.inner_clusterer) # Inject random_state if the estimator supports it if hasattr(est, "random_state"): est.random_state = random_state return est def _evaluate_criterion(self, X, labels): """Compute the quality criterion for a given partition. Parameters ---------- X : ndarray of shape (n_samples, n_features) labels : ndarray of shape (n_samples,) Returns ------- score : float The criterion value. For 'silhouette' and 'calinski_harabasz' higher is better; for 'davies_bouldin' lower is better. """ score_fn, _ = self._CRITERIA[self.criterion] return score_fn(X, labels) def _is_improvement(self, new_score, old_score): """Return True if *new_score* is better than *old_score* under the current criterion's direction (higher-is-better vs lower-is-better).""" _, higher_is_better = self._CRITERIA[self.criterion] if higher_is_better: return new_score > old_score return new_score < old_score # --------------------------------------------------------------------- # # Public API # # --------------------------------------------------------------------- # def fit(self, X, max_clusters=None): """Run the divisive clustering algorithm on *X*. Parameters ---------- X : array-like of shape (n_samples, n_features) Training data. max_clusters : int or None, optional The maximum number of clusters to find. If None, the algorithm will continue until no further improvement is possible. Returns ------- self : DivisiveClustering The fitted estimator (allows method-chaining). """ X = np.asarray(X, dtype=np.float64) rng = np.random.RandomState(self.random_state) if self.criterion not in self._CRITERIA: raise ValueError( f"Unknown criterion '{self.criterion}'. " f"Choose from {list(self._CRITERIA.keys())}." ) # Store the resolved clusterer for inspection self.inner_clusterer_ = self._resolve_inner_clusterer() _, higher_is_better = self._CRITERIA[self.criterion] labels = np.zeros(X.shape[0], dtype=int) # Track per-sample consensus probabilities across successive splits probabilities = np.ones(X.shape[0], dtype=float) # Initial score: worst possible for the chosen criterion direction current_score = -np.inf if higher_is_better else np.inf stop = False while not stop: best_score = current_score best_candidate = None best_probs = None for cluster_id in np.unique(labels): mask = labels == cluster_id # A cluster must have enough samples for a meaningful bisection if mask.sum() < max(4, self.n_neighbors + 1): continue # 1. UMAP embedding of the current cluster umap_model = UMAP( n_components=self.n_components, n_neighbors=min(self.n_neighbors, mask.sum() - 1), metric=self.metric, output_dens=self.output_dens, random_state=self.random_state, ) umap_model.fit(X[mask]) X_embedded = umap_model.transform(X[mask]) # 2. Consensus bisection using the inner clusterer runs_df = pd.DataFrame() for run_idx in range(self.n_consensus_runs): if self.subsampling: n_sub = min(self.n_subsample, np.shape(X_embedded)[0]) idx_sub = rng.choice(np.shape(X_embedded)[0], n_sub, replace=False) else: idx_sub = np.arange(np.shape(X_embedded)[0]) # Clone a fresh inner clusterer for this run inner = self._resolve_inner_clusterer(random_state=run_idx) sub_labels = inner.fit_predict(X_embedded[idx_sub]) # Propagate labels to full cluster via Random Forest rf = RandomForestClassifier( n_estimators=self.rf_n_estimators, max_depth=self.rf_max_depth, random_state=run_idx, ) rf.fit(X_embedded[idx_sub], sub_labels) runs_df[run_idx] = rf.predict(X_embedded) consensus_labels, cons_probs = consensus_cluster_labels(runs_df) # 3. Build candidate global labels with this bisection applied new_id = labels.max() + 1 candidate = labels.copy() candidate[mask] = np.where(consensus_labels == 0, cluster_id, new_id) # 4. Evaluate the candidate partition with the chosen criterion score = self._evaluate_criterion(X, candidate) if self._is_improvement(score, best_score): best_score = score best_candidate = candidate.copy() best_probs = probabilities.copy() best_probs[mask] = cons_probs # Accept the best-improving split, or stop if best_candidate is not None and (max_clusters is None or len(np.unique(best_candidate)) <= max_clusters): labels = best_candidate probabilities = best_probs current_score = best_score else: stop = True # Dense-rank labels to 0 .. K-1 self.labels_ = (rankdata(labels, method="dense") - 1).astype(int) self.probabilities_ = probabilities self.n_clusters_ = len(np.unique(self.labels_)) self.criterion_score_ = current_score self.criterion_name_ = self.criterion # Train a final Random Forest so .predict() works on new data self._classifier = RandomForestClassifier( n_estimators=self.rf_n_estimators, max_depth=self.rf_max_depth, random_state=self.random_state, ) self._classifier.fit(X, self.labels_) self._is_fitted = True return self def predict(self, X): """Predict cluster labels for unseen data. Uses the Random Forest trained on the original data and the discovered cluster labels to classify new samples. Parameters ---------- X : array-like of shape (n_samples, n_features) New data to classify. Returns ------- labels : ndarray of shape (n_samples,) Predicted cluster labels (same label space as ``labels_``). Raises ------ RuntimeError If called before ``fit``. """ if not getattr(self, "_is_fitted", False): raise RuntimeError( "This DivisiveClustering instance is not fitted yet. " "Call .fit(X) first." ) return self._classifier.predict(np.asarray(X, dtype=np.float64)) def fit_predict(self, X, max_clusters=None): """Fit the model and return cluster labels (convenience method). Equivalent to ``self.fit(X, max_clusters=max_clusters).labels_``. Parameters ---------- X : array-like of shape (n_samples, n_features) Training data. max_clusters : int or None, optional The maximum number of clusters to find. If None, the algorithm will continue until no further improvement is possible. Returns ------- labels : ndarray of shape (n_samples,) Cluster labels for the training data. """ return self.fit(X, max_clusters=max_clusters).labels_