Fix numerical instability in silhouette/distance computation

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)

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)