From d9989ad5867bf44786cb4b33ec38f5b42814cf8a Mon Sep 17 00:00:00 2001 From: maarten-devries Date: Tue, 24 Mar 2026 00:33:14 +0000 Subject: [PATCH 1/2] Fix pynndescent to match scib-metrics params and add deterministic eigsh - pynndescent now accepts random_state and n_jobs, and uses the same UMAP/scanpy-matching NNDescent parameters as scib-metrics (n_trees, n_iters, max_candidates, low_memory, compressed). - Use deterministic v0 in eigsh for reproducible diffusion embeddings. Co-Authored-By: Claude Opus 4.6 (1M context) --- .../nearest_neighbors/_pynndescent.py | 37 ++++++++++++++----- src/scib_rapids/utils/_diffusion_nn.py | 4 +- 2 files changed, 30 insertions(+), 11 deletions(-) diff --git a/src/scib_rapids/nearest_neighbors/_pynndescent.py b/src/scib_rapids/nearest_neighbors/_pynndescent.py index 6a4cf90..975aaf2 100644 --- a/src/scib_rapids/nearest_neighbors/_pynndescent.py +++ b/src/scib_rapids/nearest_neighbors/_pynndescent.py @@ -4,20 +4,37 @@ from ._dataclass import NeighborsResults -def pynndescent(X: np.ndarray, n_neighbors: int = 30) -> NeighborsResults: - """Compute nearest neighbors using PyNNDescent. +def pynndescent(X: np.ndarray, n_neighbors: int, random_state: int = 0, n_jobs: int = 1) -> NeighborsResults: + """Run pynndescent approximate nearest neighbor search. Parameters ---------- X - Array of shape (n_cells, n_features). + Data matrix. n_neighbors - Number of nearest neighbors. - - Returns - ------- - NeighborsResults + Number of neighbors to search for. + random_state + Random state. + n_jobs + Number of jobs to use. """ - index = NNDescent(X, n_neighbors=n_neighbors) - indices, distances = index.neighbor_graph + # Variables from umap (https://github.com/lmcinnes/umap/blob/3f19ce19584de4cf99e3d0ae779ba13a57472cd9/umap/umap_.py#LL326-L327) + # which is used by scanpy under the hood + n_trees = min(64, 5 + int(round((X.shape[0]) ** 0.5 / 20.0))) + n_iters = max(5, int(round(np.log2(X.shape[0])))) + max_candidates = 60 + + knn_search_index = NNDescent( + X, + n_neighbors=n_neighbors, + random_state=random_state, + low_memory=True, + n_jobs=n_jobs, + compressed=False, + n_trees=n_trees, + n_iters=n_iters, + max_candidates=max_candidates, + ) + indices, distances = knn_search_index.neighbor_graph + return NeighborsResults(indices=indices, distances=distances) diff --git a/src/scib_rapids/utils/_diffusion_nn.py b/src/scib_rapids/utils/_diffusion_nn.py index 36fd9f6..da34148 100644 --- a/src/scib_rapids/utils/_diffusion_nn.py +++ b/src/scib_rapids/utils/_diffusion_nn.py @@ -48,7 +48,9 @@ def _compute_eigen( ncv = None which = "LM" if sort == "decrease" else "SM" matrix = matrix.astype(np.float64) - evals, evecs = scipy.sparse.linalg.eigsh(matrix, k=n_comps, which=which, ncv=ncv) + # Use deterministic starting vector to ensure reproducibility + v0 = np.ones(matrix.shape[0]) / np.sqrt(matrix.shape[0]) + evals, evecs = scipy.sparse.linalg.eigsh(matrix, k=n_comps, which=which, ncv=ncv, v0=v0) evals, evecs = evals.astype(np.float32), evecs.astype(np.float32) if sort == "decrease": evals = evals[::-1] From af7afcb542a05b3e9a387ebfab050d8d92808552 Mon Sep 17 00:00:00 2001 From: maarten-devries Date: Tue, 24 Mar 2026 01:02:19 +0000 Subject: [PATCH 2/2] Add eigenvector sign canonicalization in diffusion maps Eigsh can return eigenvectors with arbitrary signs, causing non-deterministic diffusion embeddings across processes. Canonicalize by making the largest absolute element of each eigenvector positive. Note: this does not fully resolve kbet_per_label diffs for small clusters with degenerate eigenvalues (repeated eigenvalues produce a non-unique eigenspace that sign-flipping alone cannot canonicalize). Co-Authored-By: Claude Opus 4.6 (1M context) --- src/scib_rapids/utils/_diffusion_nn.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/src/scib_rapids/utils/_diffusion_nn.py b/src/scib_rapids/utils/_diffusion_nn.py index da34148..828fc72 100644 --- a/src/scib_rapids/utils/_diffusion_nn.py +++ b/src/scib_rapids/utils/_diffusion_nn.py @@ -56,6 +56,13 @@ def _compute_eigen( evals = evals[::-1] evecs = evecs[:, ::-1] + # Canonicalize eigenvector signs: make the largest absolute element positive + # This resolves sign ambiguity that can differ across runs/processes + max_abs_idx = np.argmax(np.abs(evecs), axis=0) + signs = np.sign(evecs[max_abs_idx, np.arange(evecs.shape[1])]) + signs[signs == 0] = 1 + evecs *= signs + return evals, evecs