add self

scverse · Feb 11, 2025 · d304459 · d304459
1 parent caeb583
commit d304459
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/src/rapids_singlecell/preprocessing/_neighbors.py b/src/rapids_singlecell/preprocessing/_neighbors.py
@@ -239,6 +239,14 @@ def _nn_descent_knn(
     )
     if metric == "euclidean":
         distances = cp.sqrt(distances)
+    if metric in ("cosine", "euclidean", "sqeuclidean"):
+        add_self_neighbors = cp.arange(X.shape[0], dtype=cp.uint32)
+        neighbors = cp.concatenate(
+            (add_self_neighbors[:, None], neighbors[:, :-1]), axis=1
+        )
+        add_self_distances = cp.zeros((X.shape[0], 1), dtype=cp.float32)
+        distances = cp.concatenate((add_self_distances, distances[:, :-1]), axis=1)
+    print(neighbors.shape, distances.shape, k)
     return neighbors, distances