make all neighbors more robust

docs/release-notes/0.17.0.md

@@ -6,6 +6,9 @@
 * Fix {func}`~rapids_singlecell.pp.harmony_integrate` with RMM managed memory by making multi-key clustering arrays optional and allocating them only when needed. {pr}`766` {smaller}`A Holly & S Dicks`
 * Fix {func}`~rapids_singlecell.tl.umap` on neighbor graphs whose ``metric`` ``cuml`` does not know {pr}`768` {smaller}`S Dicks`
 * Make {func}`~rapids_singlecell.pp.regress_out` numerically robust {pr}`767` {smaller}`S Dicks`
+* Fix {func}`~rapids_singlecell.pp.neighbors` with ``algorithm="all_neighbors"`` losing the neighbors that cross cluster boundaries on multiple GPUs; ``overlap_factor`` now scales with the cluster count {pr}`769` {smaller}`S Dicks`
+* Fix {func}`~rapids_singlecell.pp.neighbors` with ``metric="inner_product"`` failing in the connectivity step for every ``algorithm``. The metric is no longer forwarded to ``cuml``'s ``fuzzy_simplicial_set``, which rejects the metrics it does not know and ignores it anyway when the neighbors are precomputed {pr}`769` {smaller}`S Dicks`
+* Fix {func}`~rapids_singlecell.pp.neighbors` with ``algorithm="all_neighbors"`` and ``algorithm_kwds={"n_clusters": 2}`` always raising; the derived ``overlap_factor`` is now capped at ``n_clusters - 1`` {pr}`769` {smaller}`S Dicks`
 
 ```{rubric} Features
 ```
@@ -17,16 +20,19 @@
 * Add ``key_added`` to {func}`~rapids_singlecell.tl.diffmap` and {func}`~rapids_singlecell.tl.draw_graph` {pr}`751` {smaller}`S Dicks`
 * Derive {func}`~rapids_singlecell.pp.harmony_integrate`'s ``adjusted_basis`` from ``basis`` instead of hard-coding ``"X_pca_harmony"``, so it follows ``preset`` (``"pca"`` gives ``"pca_harmony"``) and a custom ``basis`` is suffixed rather than ignored {pr}`751` {smaller}`S Dicks`
 * Add ``mean_in_log_space`` to {func}`~rapids_singlecell.tl.rank_genes_groups` {pr}`751` {smaller}`S Dicks`
+* Add ``cosine`` and ``inner_product`` to the metrics supported by ``algorithm="all_neighbors"`` in {func}`~rapids_singlecell.pp.neighbors`. ``algorithm_kwds={"algo": "ivf_pq"}`` still requires ``euclidean``/``sqeuclidean`` {pr}`769` {smaller}`S Dicks`
 
 ```{rubric} Performance
 ```
 * Speed up {func}`~rapids_singlecell.tl.umap` and cut its peak memory by no longer moving the representation to the GPU. {pr}`768` {smaller}`S Dicks`
+* Keep the dataset on device for unbatched ``all_neighbors`` builds in {func}`~rapids_singlecell.pp.neighbors` {pr}`769` {smaller}`S Dicks`
 
 ```{rubric} Misc
 ```
 * Build the local Docker images against CUDA 12.9.1 on CUDA 12, matching the published images {pr}`757` {smaller}`S Dicks`
 * Require ``scverse-misc[settings]>=0.1.3`` {pr}`751` {smaller}`S Dicks`
 * Drop the remaining ``librmm``/``rapids_logger`` build- and link-time dependencies. No RAPIDS C++ package is needed to build, so isolated source builds now work on both CUDA 12 and CUDA 13 without ``--no-build-isolation`` {pr}`762` {smaller}`S Dicks`
+* Default ``all_neighbors`` to nn-descent's ``graph_degree`` of 64/128 like `cuml.UMAP` instead of ``n_neighbors`` {pr}`769` {smaller}`S Dicks`
 
 ```{rubric} Deprecations
 ```

src/rapids_singlecell/preprocessing/_neighbors/__init__.py

@@ -141,15 +141,17 @@ def neighbors(
 
         For `all_neighbors` algorithm, the following parameters can be specified:
 
-        * 'algo': The algorithm to use. Valid options are: 'ivf_pq' and 'nn_descent'. Default is 'nn_descent'.
+        * 'algo': The algorithm to use. Valid options are: 'ivf_pq' and 'nn_descent'. Default is 'nn_descent'. `ivf_pq` is restricted to the `euclidean` and `sqeuclidean` metrics; use `nn_descent` for `cosine` and `inner_product`.
 
-        * 'n_clusters': Number of clusters/batches to partition the dataset into (> overlap_factor). Default is number of GPUs.
+        * 'n_clusters': Number of clusters/batches to partition the dataset into (> overlap_factor). Default is 1 on a single GPU and the smallest multiple of the device count greater than `overlap_factor` otherwise.
 
-        * 'overlap_factor': Number of clusters each point is assigned to (must be < n_clusters). Default is 1.
+        * 'overlap_factor': Number of clusters each point is assigned to. Must be < n_clusters when the build is batched (`n_clusters > 1`). Default is 1 for an unbatched build and `min(max(2, ceil(log2(n_clusters))), n_clusters - 1)` otherwise. Lower values are faster but lose neighbors at cluster boundaries.
 
         * 'n_lists': Number of inverted lists for IVF indexing. Default is 2 * next_power_of_2(sqrt(n_samples)). Only available for `ivf_pq` algorithm.
 
-        * 'intermediate_graph_degree': The degree of the intermediate graph. Default is None. It is recommended to set it to `>= 1.5 * n_neighbors`. Only available for `nn_descent` algorithm.
+        * 'graph_degree': The degree of the graph nn-descent builds before selecting the final `n_neighbors`. Default is 64, raised to `n_neighbors` if larger. Only available for `nn_descent` algorithm.
+
+        * 'intermediate_graph_degree': The degree of the intermediate graph. Default is `max(128, int(1.5 * graph_degree))`, following the recommended `>= 1.5 * graph_degree`. A smaller user-supplied value is raised to `graph_degree`. Only available for `nn_descent` algorithm.
 
         For `mg_ivfflat` and `mg_ivfpq` algorithms, the following parameters can be specified:
 
@@ -208,7 +210,7 @@ def neighbors(
         )
 
     X = _choose_representation(adata, use_rep=use_rep, n_pcs=n_pcs)
-    X_contiguous = _check_neighbors_X(X, algorithm)
+    X_contiguous = _check_neighbors_X(X, algorithm, algorithm_kwds)
     _check_metrics(algorithm, metric)
 
     knn_indices, knn_dist = KNN_ALGORITHMS[algorithm](
@@ -240,7 +242,6 @@ def neighbors(
         n_obs=n_obs,
         n_neighbors=n_neighbors,
         rng=rng,
-        metric=metric,
         method=method,
     )
     if connectivities.nnz >= np.iinfo(np.int32).max:
@@ -402,7 +403,7 @@ def bbknn(
         adata._init_as_actual(adata.copy())
 
     X = _choose_representation(adata, use_rep=use_rep, n_pcs=n_pcs)
-    X_contiguous = _check_neighbors_X(X, algorithm)
+    X_contiguous = _check_neighbors_X(X, algorithm, algorithm_kwds)
     _check_metrics(algorithm, metric)
 
     n_obs = adata.shape[0]
@@ -464,7 +465,6 @@ def bbknn(
         n_obs=n_obs,
         n_neighbors=total_neighbors,
         rng=rng,
-        metric=metric,
     )
     if connectivities.nnz >= np.iinfo(np.int32).max:
         connectivities = connectivities.get().tocsr()

src/rapids_singlecell/preprocessing/_neighbors/_algorithms/_all_neighbors.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import math
 from typing import TYPE_CHECKING
 
 import cupy as cp
@@ -17,6 +18,38 @@
     from rapids_singlecell.preprocessing._neighbors import _Metrics
 
 
+def _default_overlap_factor(n_clusters: int) -> int:
+    """Overlap needed to hold recall as the dataset is split into more clusters."""
+    if n_clusters <= 1:
+        return 1
+    return max(2, math.ceil(math.log2(n_clusters)))
+
+
+def _all_neighbors_batching(algorithm_kwds: Mapping) -> tuple[int, int]:
+    """Resolve ``(n_clusters, overlap_factor)`` for the cuVS all-neighbors build."""
+    n_devices = cp.cuda.runtime.getDeviceCount()
+    n_clusters = algorithm_kwds.get("n_clusters")
+    overlap_factor = algorithm_kwds.get("overlap_factor")
+    if n_clusters is None:
+        n_clusters = 1 if n_devices == 1 else n_devices
+        while n_clusters > 1 and n_clusters <= (
+            _default_overlap_factor(n_clusters)
+            if overlap_factor is None
+            else overlap_factor
+        ):
+            n_clusters += n_devices
+    if overlap_factor is None:
+        overlap_factor = _default_overlap_factor(n_clusters)
+        if n_clusters > 1:
+            overlap_factor = min(overlap_factor, n_clusters - 1)
+    if n_clusters > 1 and overlap_factor >= n_clusters:
+        raise ValueError(
+            f"'n_clusters' ({n_clusters}) must be greater than 'overlap_factor' "
+            f"({overlap_factor}) when batching the all_neighbors build."
+        )
+    return n_clusters, overlap_factor
+
+
 def _all_neighbors_knn(
     X: np.ndarray,
     Y: np.ndarray,
@@ -32,32 +65,40 @@ def _all_neighbors_knn(
             "Please update your cuvs installation."
         )
     algo = algorithm_kwds.get("algo", "nn_descent")
-    n_devices = cp.cuda.runtime.getDeviceCount()
-    if n_devices == 1:
+    n_clusters, overlap_factor = _all_neighbors_batching(algorithm_kwds)
+    if n_clusters == 1:
         from cuvs.common import Resources
 
         res = Resources()
     else:
         from cuvs.common import MultiGpuResources
 
         res = MultiGpuResources()
-    n_clusters = algorithm_kwds.get("n_clusters", n_devices)
-    overlap_factor = algorithm_kwds.get("overlap_factor", 1)
+    cuvs_metric = "sqeuclidean" if metric == "euclidean" else metric
     if algo == "ivf_pq" or algo == "ivfpq":
         from cuvs.neighbors import ivf_pq
 
         algo = "ivf_pq"
+        if cuvs_metric != "sqeuclidean":
+            raise ValueError(
+                f"all_neighbors with algo='ivf_pq' only supports 'euclidean' and "
+                f"'sqeuclidean' metrics, got {metric!r}. Use algo='nn_descent' instead."
+            )
         n_lists = algorithm_kwds.get("n_lists", _compute_nlist(X.shape[0]))
-        ivf_pq_params = ivf_pq.IndexParams(n_lists=n_lists)
+        ivf_pq_params = ivf_pq.IndexParams(n_lists=n_lists, metric=cuvs_metric)
         nn_descent_params = None
     elif algo == "nn_descent":
         from cuvs.neighbors import nn_descent
 
+        graph_degree = max(algorithm_kwds.get("graph_degree", 64), k)
         intermediate_graph_degree = algorithm_kwds.get(
-            "intermediate_graph_degree", None
+            "intermediate_graph_degree", max(128, int(1.5 * graph_degree))
         )
+        intermediate_graph_degree = max(intermediate_graph_degree, graph_degree)
         nn_descent_params = nn_descent.IndexParams(
-            graph_degree=k, intermediate_graph_degree=intermediate_graph_degree
+            graph_degree=graph_degree,
+            intermediate_graph_degree=intermediate_graph_degree,
+            metric=cuvs_metric,
         )
         ivf_pq_params = None
     else:
@@ -66,7 +107,7 @@ def _all_neighbors_knn(
         algo=algo,
         overlap_factor=overlap_factor,
         n_clusters=n_clusters,
-        metric="sqeuclidean",
+        metric=cuvs_metric,
         ivf_pq_params=ivf_pq_params,
         nn_descent_params=nn_descent_params,
     )

src/rapids_singlecell/preprocessing/_neighbors/_helper/__init__.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import math
+from types import MappingProxyType
 from typing import TYPE_CHECKING
 
 import cupy as cp
@@ -11,6 +12,8 @@
 from scipy import sparse as sc_sparse
 
 if TYPE_CHECKING:
+    from collections.abc import Mapping
+
     from rapids_singlecell.preprocessing._neighbors import _Algorithms, _Metrics
 
 
@@ -27,6 +30,7 @@ def _cuvs_switch():
 def _check_neighbors_X(
     X: cp_sparse.spmatrix | sc_sparse.spmatrix | np.ndarray | cp.ndarray,
     algorithm: _Algorithms,
+    algorithm_kwds: Mapping = MappingProxyType({}),
 ) -> cp_sparse.spmatrix | cp.ndarray | np.ndarray:
     """Check and convert input X to the expected format based on algorithm.
 
@@ -41,13 +45,19 @@ def _check_neighbors_X(
     X_contiguous (cupy.ndarray or sparse.csr_matrix): Contiguous array or CSR matrix.
 
     """
+    from rapids_singlecell.preprocessing._neighbors._algorithms._all_neighbors import (
+        _all_neighbors_batching,
+    )
+
     if cp_sparse.issparse(X) or sc_sparse.issparse(X):
         if algorithm != "brute":
             raise ValueError(
                 f"Sparse input is not supported for {algorithm} algorithm. Use 'brute' instead."
             )
         X_contiguous = X.tocsr()
-    elif algorithm in ["all_neighbors", "mg_ivfflat", "mg_ivfpq"]:
+    elif algorithm in ["mg_ivfflat", "mg_ivfpq"] or (
+        algorithm == "all_neighbors" and _all_neighbors_batching(algorithm_kwds)[0] > 1
+    ):
         if isinstance(X, np.ndarray):
             X_contiguous = np.asarray(X, order="C", dtype=np.float32)
         elif isinstance(X, cp.ndarray):
@@ -101,9 +111,10 @@ def _check_metrics(algorithm: _Algorithms, metric: _Metrics) -> bool:
                 "nn_descent only supports 'euclidean', 'sqeuclidean', 'inner_product' and 'cosine' metrics."
             )
     elif algorithm == "all_neighbors":
-        if metric not in ["euclidean", "sqeuclidean"]:
+        if metric not in ["euclidean", "sqeuclidean", "cosine", "inner_product"]:
             raise ValueError(
-                "all_neighbors only supports 'euclidean' and 'sqeuclidean' metrics."
+                "all_neighbors only supports 'euclidean', 'sqeuclidean', "
+                "'inner_product' and 'cosine' metrics."
             )
     else:
         raise NotImplementedError(f"The {algorithm} algorithm is not implemented yet.")

src/rapids_singlecell/preprocessing/_neighbors/_neighbors.py

@@ -138,9 +138,13 @@ def _get_connectivities_umap(
     n_obs: int,
     n_neighbors: int,
     rng: np.random.Generator,
-    metric: str,
 ) -> cp_sparse.coo_matrix:
-    """UMAP fuzzy simplicial set connectivities."""
+    """UMAP fuzzy simplicial set connectivities.
+
+    The graph is built from the precomputed ``knn_indices``/``knn_dist``, so the
+    metric is never recomputed here. Forwarding it would only make cuML reject
+    the metrics it does not know itself, such as ``inner_product``.
+    """
     set_op_mix_ratio = 1.0
     local_connectivity = 1.0
 
@@ -151,7 +155,6 @@ def _get_connectivities_umap(
         n_neighbors,
         # cuML seeds its fuzzy simplicial set, so draw the seed right here
         _seed_from_rng(rng),
-        metric=metric,
         knn_indices=knn_indices,
         knn_dists=knn_dist,
         set_op_mix_ratio=set_op_mix_ratio,
@@ -268,7 +271,6 @@ def _calc_connectivities(
     n_obs: int,
     n_neighbors: int,
     rng: np.random.Generator,
-    metric: str,
     method: Literal["umap", "gauss", "jaccard"] = "umap",
 ) -> cp_sparse.spmatrix:
     """Compute connectivities from KNN arrays.
@@ -285,8 +287,6 @@ def _calc_connectivities(
         Number of nearest neighbors.
     rng
         Random generator (a seed is drawn for the UMAP fuzzy simplicial set).
-    metric
-        Distance metric name.
     method
         Method for computing connectivities.
 
@@ -312,5 +312,4 @@ def _calc_connectivities(
         n_obs=n_obs,
         n_neighbors=n_neighbors,
         rng=rng,
-        metric=metric,
     )

tests/test_mg_neighbors.py

@@ -83,6 +83,100 @@ def test_all_neighbors(algo):
     _calc_recall(distances, adata.obsp["distances"], tolerance=tolerance)
 
 
+@pytest.mark.parametrize(
+    ("n_devices", "expected"),
+    [(1, (1, 1)), (2, (4, 2)), (3, (3, 2)), (4, (4, 2)), (8, (8, 3)), (16, (16, 4))],
+)
+def test_all_neighbors_batching_defaults(monkeypatch, n_devices, expected):
+    import cupy as cp
+
+    from rapids_singlecell.preprocessing._neighbors._algorithms._all_neighbors import (
+        _all_neighbors_batching,
+    )
+
+    monkeypatch.setattr(cp.cuda.runtime, "getDeviceCount", lambda: n_devices)
+    n_clusters, overlap_factor = _all_neighbors_batching({})
+    assert (n_clusters, overlap_factor) == expected
+    assert n_clusters == 1 or overlap_factor < n_clusters
+
+
+def test_all_neighbors_batching_overrides():
+    from rapids_singlecell.preprocessing._neighbors._algorithms._all_neighbors import (
+        _all_neighbors_batching,
+    )
+
+    assert _all_neighbors_batching({"n_clusters": 16}) == (16, 4)
+    assert _all_neighbors_batching({"n_clusters": 8, "overlap_factor": 2}) == (8, 2)
+    # The default overlap is capped at n_clusters - 1, so small explicit cluster
+    # counts stay usable instead of tripping the guard below.
+    assert _all_neighbors_batching({"n_clusters": 2}) == (2, 1)
+    assert _all_neighbors_batching({"n_clusters": 3}) == (3, 2)
+    with pytest.raises(ValueError, match="must be greater than"):
+        _all_neighbors_batching({"n_clusters": 3, "overlap_factor": 3})
+
+
+@pytest.mark.parametrize("n_clusters", [4, 8])
+def test_all_neighbors_batched(n_clusters):
+    """These recall 0.91 and 0.47 with the previous ``overlap_factor=1``."""
+    if parse_version(cuvs.__version__) <= parse_version("25.08"):
+        pytest.skip("Skipping All-Neighbors")
+    adata = pbmc68k_reduced()
+    rsc.pp.neighbors(
+        adata,
+        n_pcs=50,
+        n_neighbors=15,
+        algorithm="all_neighbors",
+        algorithm_kwds={"n_clusters": n_clusters},
+    )
+    distances = adata.obsp["distances"].copy()
+    rsc.pp.neighbors(adata, n_pcs=50, n_neighbors=15, algorithm="brute")
+    _calc_recall(distances, adata.obsp["distances"], tolerance=0.95)
+
+
+@pytest.mark.parametrize("n_clusters", [1, 4])
+@pytest.mark.parametrize("metric", ["cosine", "sqeuclidean"])
+def test_all_neighbors_metrics(metric, n_clusters):
+    if parse_version(cuvs.__version__) <= parse_version("25.08"):
+        pytest.skip("Skipping All-Neighbors")
+    adata = pbmc68k_reduced()
+    rsc.pp.neighbors(
+        adata,
+        n_pcs=50,
+        n_neighbors=15,
+        algorithm="all_neighbors",
+        metric=metric,
+        algorithm_kwds={"n_clusters": n_clusters},
+    )
+    distances = adata.obsp["distances"].copy()
+    rsc.pp.neighbors(adata, n_pcs=50, n_neighbors=15, algorithm="brute", metric=metric)
+    _calc_recall(distances, adata.obsp["distances"], tolerance=0.95)
+
+
+@pytest.mark.parametrize("n_clusters", [1, 4])
+def test_all_neighbors_inner_product(n_clusters):
+    """``inner_product`` used to fail in the connectivity step, not the build."""
+    if parse_version(cuvs.__version__) <= parse_version("25.08"):
+        pytest.skip("Skipping All-Neighbors")
+    adata = pbmc68k_reduced()
+    rsc.pp.neighbors(
+        adata,
+        n_pcs=50,
+        n_neighbors=15,
+        algorithm="all_neighbors",
+        metric="inner_product",
+        algorithm_kwds={"n_clusters": n_clusters},
+    )
+    distances = adata.obsp["distances"].copy()
+    rsc.pp.neighbors(
+        adata, n_pcs=50, n_neighbors=15, algorithm="brute", metric="inner_product"
+    )
+    # inner_product is a similarity, so the top-k are the *largest* dot products and
+    # a point is not its own nearest neighbor. nn-descent refines a similarity graph
+    # less well than a distance one: recall measures 0.933 here, against 0.995 for
+    # the metrics in ``test_all_neighbors_metrics``.
+    _calc_recall(distances, adata.obsp["distances"], tolerance=0.9)
+
+
 @pytest.mark.parametrize("algo", ["mg_ivfflat", "mg_ivfpq"])
 def test_mg_bbknn(algo):
     if parse_version(cuvs.__version__) <= parse_version("25.08"):