diff --git a/docs/release-notes/0.17.0.md b/docs/release-notes/0.17.0.md index 0a6d5acb6..79123e59b 100644 --- a/docs/release-notes/0.17.0.md +++ b/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 ``` diff --git a/src/rapids_singlecell/preprocessing/_neighbors/__init__.py b/src/rapids_singlecell/preprocessing/_neighbors/__init__.py index 6bda3a754..77359510e 100644 --- a/src/rapids_singlecell/preprocessing/_neighbors/__init__.py +++ b/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() diff --git a/src/rapids_singlecell/preprocessing/_neighbors/_algorithms/_all_neighbors.py b/src/rapids_singlecell/preprocessing/_neighbors/_algorithms/_all_neighbors.py index a5935f6c2..7a1b7bf1d 100644 --- a/src/rapids_singlecell/preprocessing/_neighbors/_algorithms/_all_neighbors.py +++ b/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,8 +65,8 @@ 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() @@ -41,23 +74,31 @@ def _all_neighbors_knn( 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, ) diff --git a/src/rapids_singlecell/preprocessing/_neighbors/_helper/__init__.py b/src/rapids_singlecell/preprocessing/_neighbors/_helper/__init__.py index e180a64b7..ecef3f590 100644 --- a/src/rapids_singlecell/preprocessing/_neighbors/_helper/__init__.py +++ b/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.") diff --git a/src/rapids_singlecell/preprocessing/_neighbors/_neighbors.py b/src/rapids_singlecell/preprocessing/_neighbors/_neighbors.py index 220893216..81e23d332 100644 --- a/src/rapids_singlecell/preprocessing/_neighbors/_neighbors.py +++ b/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, ) diff --git a/tests/test_mg_neighbors.py b/tests/test_mg_neighbors.py index 160b27be8..d60228f7f 100644 --- a/tests/test_mg_neighbors.py +++ b/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"):