Match pynndescent params and fix diffusion map reproducibility

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)

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