From b4151162ae503bbdf7744210e1cba90f79bbbda6 Mon Sep 17 00:00:00 2001 From: maarten-devries Date: Mon, 23 Mar 2026 22:18:54 +0000 Subject: [PATCH] Fix numerical instability in silhouette/distance computation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit cdist used the expansion formula (||x||²+||y||²-2xy) which suffers from catastrophic cancellation in float32 for nearby points. Switch to the direct formula sqrt(sum((x-y)²)) to match scib-metrics (JAX) behavior. Also fix silhouette denominator handling and mean_other mode to exactly match scib-metrics. Verified: all metrics now agree to <0.0002 on both 1k and 20k cell Tabula Sapiens benchmarks. Co-Authored-By: Claude Opus 4.6 (1M context) --- src/scib_rapids/utils/_dist.py | 33 +++++++++++++++++++--------- src/scib_rapids/utils/_silhouette.py | 7 +++--- 2 files changed, 26 insertions(+), 14 deletions(-) diff --git a/src/scib_rapids/utils/_dist.py b/src/scib_rapids/utils/_dist.py index 32f51b1..c180085 100644 --- a/src/scib_rapids/utils/_dist.py +++ b/src/scib_rapids/utils/_dist.py @@ -5,7 +5,10 @@ def cdist_sq(x: np.ndarray, y: np.ndarray) -> cp.ndarray: - """Squared euclidean pairwise distances without materialising (n, m, d). + """Squared euclidean pairwise distances using expansion formula. + + Fast but less numerically stable — use only where exact distances + are not required (e.g., kmeans cluster assignment). Parameters ---------- @@ -27,9 +30,26 @@ def cdist_sq(x: np.ndarray, y: np.ndarray) -> cp.ndarray: return cp.maximum(dist_sq, 0.0) +def _euclidean_cdist_direct(x: cp.ndarray, y: cp.ndarray, sub_chunk: int = 2048) -> cp.ndarray: + """Euclidean pairwise distances using the direct formula sqrt(sum((x-y)^2)). + + Sub-chunks over y to limit memory of the (n_a, sub_chunk, d) intermediate. + """ + n_a = x.shape[0] + n_b = y.shape[0] + result = cp.empty((n_a, n_b), dtype=x.dtype) + for j in range(0, n_b, sub_chunk): + j_end = min(j + sub_chunk, n_b) + diff = x[:, None, :] - y[None, j:j_end, :] # (n_a, sub_chunk, d) + result[:, j:j_end] = cp.sqrt(cp.sum(diff * diff, axis=2)) + return result + + def cdist(x: np.ndarray, y: np.ndarray, metric: Literal["euclidean", "cosine"] = "euclidean") -> cp.ndarray: """CuPy implementation of pairwise distance computation. + Uses the direct formula for numerical stability (matches JAX/scib-metrics). + Parameters ---------- x @@ -51,7 +71,7 @@ def cdist(x: np.ndarray, y: np.ndarray, metric: Literal["euclidean", "cosine"] = y = cp.asarray(y, dtype=cp.float32) if metric == "euclidean": - return cp.sqrt(cdist_sq(x, y)) + return _euclidean_cdist_direct(x, y) else: # cosine distance = 1 - (x @ y^T) / (||x|| * ||y||) x_norm = cp.linalg.norm(x, axis=1, keepdims=True) @@ -75,11 +95,4 @@ def pdist_squareform(X: np.ndarray) -> cp.ndarray: Array of shape (n_cells, n_cells) """ X = cp.asarray(X, dtype=cp.float32) - x_sq = cp.sum(X**2, axis=1, keepdims=True) - dist_sq = x_sq + x_sq.T - 2.0 * X @ X.T - dist_sq = cp.maximum(dist_sq, 0.0) - dist = cp.sqrt(dist_sq) - # Ensure symmetry and zero diagonal - dist = (dist + dist.T) / 2.0 - cp.fill_diagonal(dist, 0.0) - return dist + return _euclidean_cdist_direct(X, X) diff --git a/src/scib_rapids/utils/_silhouette.py b/src/scib_rapids/utils/_silhouette.py index a4aeec6..5be13bc 100644 --- a/src/scib_rapids/utils/_silhouette.py +++ b/src/scib_rapids/utils/_silhouette.py @@ -36,8 +36,8 @@ def _silhouette_reduce( clust_dists = clust_dists / label_freqs inter_clust_dists = cp.max(clust_dists, axis=1) elif between_cluster_distances == "mean_other": - clust_dists[row_idx, chunk_labels] = 0.0 - total_other_dists = cp.sum(clust_dists, axis=1) + clust_dists[row_idx, chunk_labels] = cp.nan + total_other_dists = cp.nansum(clust_dists, axis=1) total_other_count = cp.sum(label_freqs) - label_freqs[chunk_labels] inter_clust_dists = total_other_dists / total_other_count elif between_cluster_distances == "nearest": @@ -114,8 +114,7 @@ def silhouette_samples( X_gpu, chunk_size=chunk_size, reduce_fn=reduce_fn, metric=metric ) - denom = label_freqs[labels_gpu] - 1 - denom = cp.maximum(denom, 1) # avoid division by zero + denom = (label_freqs - 1)[labels_gpu] intra_clust_dists = intra_clust_dists / denom sil_samples = inter_clust_dists - intra_clust_dists sil_samples = sil_samples / cp.maximum(intra_clust_dists, inter_clust_dists)