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..828fc72 100644 --- a/src/scib_rapids/utils/_diffusion_nn.py +++ b/src/scib_rapids/utils/_diffusion_nn.py @@ -48,12 +48,21 @@ 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] 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