diff --git a/cpp/include/cuml/neighbors/knn.hpp b/cpp/include/cuml/neighbors/knn.hpp index bed2e666d4..e6a5219465 100644 --- a/cpp/include/cuml/neighbors/knn.hpp +++ b/cpp/include/cuml/neighbors/knn.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2019-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ @@ -71,6 +71,38 @@ void rbc_knn_query(const raft::handle_t& handle, int64_t* out_inds, float* out_dists); +/** + * @brief Perform a radius neighbors query on the fit index. + * + * A single query requires two calls to this API: + * - The first call should pass adj_indices=nullptr and nnz=0. This will + * fill in adj_indptr, letting you get the size needed for adj_indices. + * - The second call should pass adj_indices and nnz, an array of size + * nnz=adj_indptr[-1]. This will fill in adj_indices. + * + * @param[in] handle: RAFT handle + * @param[in] rbc_index: the fit RBC index + * @param[in] query: the query points as a C-contiguous array + * @param[in] n_query: number of rows in the query + * @param[in] dim: number of columns in the query + * @param[in] radius: the neighborhood radius + * @param[out] adj_indptr: the indptr array in output CSR adjacency matrix, + * of shape n_query + 1. + * @param[out] adj_indices: the indices array in the output CSR adjacency + * matrix, of shape adj_indptr[-1]. Should be NULL on the first + * call. + * @param[in] nnz: the number of elements in adj_indices, or 0 on the first call. + */ +void rbc_radius_neighbors_graph(const raft::handle_t& handle, + const std::uintptr_t& rbc_index, + const float* query, + int64_t n_query, + int64_t dim, + float radius, + int64_t* adj_indptr, + int64_t* adj_indices = nullptr, + int64_t nnz = 0); + /** * @brief Free the RBC index * diff --git a/cpp/src/knn/knn.cu b/cpp/src/knn/knn.cu index 39c26e5ba8..4f946fde54 100644 --- a/cpp/src/knn/knn.cu +++ b/cpp/src/knn/knn.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2019-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ @@ -200,6 +200,28 @@ void rbc_knn_query(const raft::handle_t& handle, handle, *rbc_index_ptr, query_view, indices_view, distances_view, k); } +void rbc_radius_neighbors_graph(const raft::handle_t& handle, + const std::uintptr_t& rbc_index, + const float* query, + int64_t n_query, + int64_t dim, + float radius, + int64_t* adj_indptr, + int64_t* adj_indices, + int64_t nnz) +{ + auto index_ptr = reinterpret_cast*>(rbc_index); + + cuvs::neighbors::ball_cover::eps_nn( + handle, + *index_ptr, + raft::make_device_vector_view(adj_indptr, n_query + 1), + raft::make_device_vector_view(adj_indices, nnz), + raft::make_device_vector_view(nullptr, 0), + raft::make_device_matrix_view(query, n_query, dim), + radius); +} + void rbc_free_index(std::uintptr_t rbc_index) { // Cast back to the original type and delete diff --git a/python/cuml/cuml/accel/_wrappers/sklearn/neighbors.py b/python/cuml/cuml/accel/_wrappers/sklearn/neighbors.py index d5d34ea315..23733dc32c 100644 --- a/python/cuml/cuml/accel/_wrappers/sklearn/neighbors.py +++ b/python/cuml/cuml/accel/_wrappers/sklearn/neighbors.py @@ -1,10 +1,12 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. # SPDX-License-Identifier: Apache-2.0 # import cuml.neighbors from cuml.accel.estimator_proxy import ProxyBase +from cuml.common.sparse_utils import is_sparse +from cuml.internals.interop import UnsupportedOnGPU __all__ = ( "NearestNeighbors", @@ -18,6 +20,21 @@ class NearestNeighbors(ProxyBase): _gpu_class = cuml.neighbors.NearestNeighbors _other_attributes = frozenset(("_fit_method", "_tree", "_fit_X")) + def _gpu_radius_neighbors_graph( + self, X=None, radius=None, mode="connectivity", sort_results=False + ): + if mode != "connectivity": + raise UnsupportedOnGPU(f"`mode={mode!r}` is not supported") + if sort_results: + raise UnsupportedOnGPU("`sort_results=True` is not supported") + if is_sparse(X): + raise UnsupportedOnGPU("Sparse inputs are not supported") + if self.effective_metric_ not in ["l2", "euclidean"]: + raise UnsupportedOnGPU( + f"metric={self.effective_metric_!r} is not supported" + ) + return self._gpu.radius_neighbors_graph(X=X, radius=radius) + class KNeighborsClassifier(ProxyBase): _gpu_class = cuml.neighbors.KNeighborsClassifier diff --git a/python/cuml/cuml/neighbors/kneighbors_classifier.pyx b/python/cuml/cuml/neighbors/kneighbors_classifier.pyx index 7c7e220bbc..be032c2910 100644 --- a/python/cuml/cuml/neighbors/kneighbors_classifier.pyx +++ b/python/cuml/cuml/neighbors/kneighbors_classifier.pyx @@ -16,7 +16,7 @@ from cuml.internals.array import CumlArray from cuml.internals.interop import UnsupportedOnGPU from cuml.internals.mixins import ClassifierMixin, FMajorInputTagMixin from cuml.internals.outputs import reflect, run_in_internal_context -from cuml.neighbors.nearest_neighbors import NearestNeighbors +from cuml.neighbors.nearest_neighbors import NeighborsBase from cuml.neighbors.weights import compute_weights from libc.stdint cimport int64_t, uintptr_t @@ -49,9 +49,7 @@ cdef extern from "cuml/neighbors/knn.hpp" namespace "ML" nogil: ) except + -class KNeighborsClassifier(ClassifierMixin, - FMajorInputTagMixin, - NearestNeighbors): +class KNeighborsClassifier(ClassifierMixin, FMajorInputTagMixin, NeighborsBase): """ K-Nearest Neighbors Classifier is an instance-based learning technique, that keeps training samples around for prediction, rather than trying diff --git a/python/cuml/cuml/neighbors/kneighbors_regressor.pyx b/python/cuml/cuml/neighbors/kneighbors_regressor.pyx index c92c2da6a7..3326bb76da 100644 --- a/python/cuml/cuml/neighbors/kneighbors_regressor.pyx +++ b/python/cuml/cuml/neighbors/kneighbors_regressor.pyx @@ -10,7 +10,7 @@ from cuml.internals import get_handle, reflect from cuml.internals.array import CumlArray from cuml.internals.interop import UnsupportedOnGPU, to_cpu, to_gpu from cuml.internals.mixins import FMajorInputTagMixin, RegressorMixin -from cuml.neighbors.nearest_neighbors import NearestNeighbors +from cuml.neighbors.nearest_neighbors import NeighborsBase from cuml.neighbors.weights import compute_weights from libc.stdint cimport int64_t, uintptr_t @@ -32,7 +32,7 @@ cdef extern from "cuml/neighbors/knn.hpp" namespace "ML" nogil: ) except + -class KNeighborsRegressor(RegressorMixin, FMajorInputTagMixin, NearestNeighbors): +class KNeighborsRegressor(RegressorMixin, FMajorInputTagMixin, NeighborsBase): """ K-Nearest Neighbors Regressor is an instance-based learning technique, that keeps training samples around for prediction, rather than trying diff --git a/python/cuml/cuml/neighbors/nearest_neighbors.pyx b/python/cuml/cuml/neighbors/nearest_neighbors.pyx index a965090b33..84632954cc 100644 --- a/python/cuml/cuml/neighbors/nearest_neighbors.pyx +++ b/python/cuml/cuml/neighbors/nearest_neighbors.pyx @@ -67,6 +67,18 @@ cdef extern from "cuml/neighbors/knn.hpp" namespace "ML" nogil: float *out_dists ) except + + void rbc_radius_neighbors_graph( + const handle_t& handle, + const uintptr_t& rbc_index, + const float* query, + int64_t n_query, + int64_t dim, + float radius, + int64_t *adj_rows, + int64_t *adj_cols, + int64_t nnz + ) except + + void rbc_free_index( uintptr_t rbc_index ) except + @@ -254,6 +266,7 @@ cdef DistanceType _metric_to_distance_type(str metric): cdef class RBCIndex: """An RBC index.""" cdef uintptr_t index + cdef int64_t n_samples def __dealloc__(self): if self.index != 0: @@ -266,10 +279,6 @@ cdef class RBCIndex: @staticmethod def build(X, metric): """Build a new RBC index.""" - if X.shape[1] > 3: - raise ValueError( - "The rbc algorithm is not supported for >3 dimensions currently." - ) cdef RBCIndex self = RBCIndex.__new__(RBCIndex) handle = get_handle() @@ -289,10 +298,67 @@ cdef class RBCIndex: distance_type, ) handle.sync() + self.n_samples = n_rows return self + def radius_neighbors_graph( + RBCIndex self, + X, + float radius, + ): + """Query the index for neighbors within a radius""" + handle = get_handle() + cdef handle_t* handle_ = handle.getHandle() + cdef int64_t n_query = X.shape[0] + + indptr = cp.empty(n_query + 1, dtype=np.int64) + cdef float* X_ptr = X.ptr + cdef int64_t n_rows = X.shape[0] + cdef int64_t n_cols = X.shape[1] + cdef int64_t* indptr_ptr = indptr.data.ptr + + with nogil: + rbc_radius_neighbors_graph( + handle_[0], + self.index, + X_ptr, + n_rows, + n_cols, + radius, + indptr_ptr, + NULL, + 0, + ) + + cdef int64_t nnz = indptr[-1].item() + indices = cp.empty(nnz, dtype=np.int64) + cdef int64_t* indices_ptr = indices.data.ptr + + with nogil: + rbc_radius_neighbors_graph( + handle_[0], + self.index, + X_ptr, + n_rows, + n_cols, + radius, + indptr_ptr, + indices_ptr, + nnz, + ) + + data = cp.ones(nnz) + return cupyx.scipy.sparse.csr_matrix( + (data, indices, indptr), + shape=(n_rows, self.n_samples), + ) + def kneighbors(RBCIndex self, X, uint32_t n_neighbors): """Query the index for the k nearest neighbors.""" + if X.shape[1] > 3: + raise ValueError( + "The rbc algorithm is not supported for >3 dimensions currently." + ) distances = CumlArray.zeros( (X.shape[0], n_neighbors), dtype=np.float32, @@ -435,147 +501,10 @@ cdef class ApproxIndex: return distances, indices -class NearestNeighbors(Base, - InteropMixin, - CMajorInputTagMixin, - SparseInputTagMixin): - """ - NearestNeighbors is an queries neighborhoods from a given set of - datapoints. Currently, cuML supports k-NN queries, which define - the neighborhood as the closest `k` neighbors to each query point. - - Parameters - ---------- - n_neighbors : int (default=5) - Default number of neighbors to query - verbose : int or boolean, default=False - Sets logging level. It must be one of `cuml.common.logger.level_*`. - See :ref:`verbosity-levels` for more info. - algorithm : string (default='auto') - The query algorithm to use. Valid options are: - - - ``'auto'``: to automatically select brute-force or - random ball cover based on data shape and metric - - ``'rbc'``: for the random ball algorithm, which partitions - the data space and uses the triangle inequality to lower the - number of potential distances. Currently, this algorithm - supports Haversine (2d) and Euclidean in 2d and 3d. - - ``'brute'``: for brute-force, slow but produces exact results - - ``'ivfflat'``: for inverted file, divide the dataset in partitions - and perform search on relevant partitions only - - ``'ivfpq'``: for inverted file and product quantization, - same as inverted list, in addition the vectors are broken - in n_features/M sub-vectors that will be encoded thanks - to intermediary k-means clusterings. This encoding provide - partial information allowing faster distances calculations - - metric : string (default='euclidean'). - Distance metric to use. Supported metrics include: 'l1', 'cityblock', - 'taxicab', 'manhattan', 'euclidean', 'l2', 'sqeuclidean', 'canberra', - 'minkowski', 'lp', 'chebyshev', 'linf', 'jensenshannon', 'cosine', - 'braycurtis', 'jaccard', 'hellinger', 'correlation', 'inner_product'. - The ``'ivfflat'`` and ``'ivfpq'`` - algorithms only support: 'euclidean', 'l2', 'sqeuclidean', 'cosine', - 'correlation', 'inner_product', whereas the ``'rbc'`` algorithm only - supports 'euclidean', 'l2', and 'haversine' (≤3 dimensions only). - For sparse inputs, only the ``'brute'`` algorithm is supported, with - metrics: 'l1', 'cityblock', 'taxicab', 'manhattan', 'euclidean', 'l2', - 'canberra', 'minkowski', 'lp', 'chebyshev', 'linf', 'cosine', - 'inner_product', 'jaccard', 'hellinger'. - p : float (default=2) - Parameter for the Minkowski metric. When p = 1, this is equivalent to - manhattan distance (l1), and euclidean distance (l2) for p = 2. For - arbitrary p, minkowski distance (lp) is used. - algo_params : dict, optional (default=None) - Used to configure the nearest neighbor algorithm to be used. - If set to None, parameters will be generated automatically. - Parameters for algorithm ``'brute'`` when inputs are sparse: - - - batch_size_index : (int) number of rows in each batch of \ - index array - - batch_size_query : (int) number of rows in each batch of \ - query array - - Parameters for algorithm ``'ivfflat'``: - - - nlist: (int) number of cells to partition dataset into - - nprobe: (int) at query time, number of cells used for search - - Parameters for algorithm ``'ivfpq'``: - - - nlist: (int) number of cells to partition dataset into - - nprobe: (int) at query time, number of cells used for search - - M: (int) number of subquantizers - - n_bits: (int) bits allocated per subquantizer - - usePrecomputedTables : (bool) whether to use precomputed tables - metric_params : dict, optional (default = None) - This is currently ignored. - n_jobs : int (default = None) - Ignored, here for scikit-learn API compatibility. - output_type : {'input', 'array', 'dataframe', 'series', 'df_obj', \ - 'numba', 'cupy', 'numpy', 'cudf', 'pandas'}, default=None - Return results and set estimator attributes to the indicated output - type. If None, the output type set at the module level - (`cuml.global_settings.output_type`) will be used. See - :ref:`output-data-type-configuration` for more info. - - Examples - -------- - - .. code-block:: python - - >>> import cudf - >>> from cuml.neighbors import NearestNeighbors - >>> from cuml.datasets import make_blobs - - >>> X, _ = make_blobs(n_samples=5, centers=5, - ... n_features=10, random_state=42) - - >>> # build a cudf Dataframe - >>> X_cudf = cudf.DataFrame(X) - - >>> # fit model - >>> model = NearestNeighbors(n_neighbors=3) - >>> model.fit(X) - NearestNeighbors() - - >>> # get 3 nearest neighbors - >>> distances, indices = model.kneighbors(X_cudf) - - >>> # print results - >>> print(indices) # doctest: +SKIP - 0 1 2 - 0 0 3 1 - 1 1 3 0 - 2 2 4 0 - 3 3 0 1 - 4 4 2 0 - >>> print(distances) # doctest: +SKIP - 0 1 2 - 0 0.007812 24.786566 26.399996 - 1 0.000000 24.786566 30.045017 - 2 0.007812 5.458400 27.051241 - 3 0.000000 26.399996 27.543869 - 4 0.000000 5.458400 29.583437 - - Notes - ----- - For an additional example see `the NearestNeighbors notebook - `_. - - For additional docs, see `scikit-learn's NearestNeighbors - `_. - - Pickling ``NearestNeighbors`` instances is supported for all algorithms. - However, for RBC, IVFPQ or IVFFlat the index will currently be rebuilt upon - load rather than serialized as part of the pickled binary. For approximate - indices like IVFPQ or IVFFlat this may result in small differences between - the original and reloaded models, as the generated indices may differ. - """ +class NeighborsBase(Base, InteropMixin, CMajorInputTagMixin, SparseInputTagMixin): + """Base class for `cuml.neighbors` models""" _fit_X = CumlArrayDescriptor(order='C') - _cpu_class_path = "sklearn.neighbors.NearestNeighbors" - @classmethod def _get_param_names(cls): return [ @@ -1101,6 +1030,277 @@ class NearestNeighbors(Base, return self.metric_params or {} +class NearestNeighbors(NeighborsBase): + """ + NearestNeighbors is an queries neighborhoods from a given set of + datapoints. Currently, cuML supports k-NN queries, which define + the neighborhood as the closest `k` neighbors to each query point. + + Parameters + ---------- + n_neighbors : int (default=5) + Default number of neighbors to query + radius : float (default=1.0) + Range of parameter space to use by default for ``radius_neighbors`` + queries. + verbose : int or boolean, default=False + Sets logging level. It must be one of `cuml.common.logger.level_*`. + See :ref:`verbosity-levels` for more info. + algorithm : string (default='auto') + The query algorithm to use. Valid options are: + + - ``'auto'``: to automatically select brute-force or + random ball cover based on data shape and metric + - ``'rbc'``: for the random ball algorithm, which partitions + the data space and uses the triangle inequality to lower the + number of potential distances. Currently, this algorithm + supports Haversine (2d) and Euclidean in 2d and 3d. + - ``'brute'``: for brute-force, slow but produces exact results + - ``'ivfflat'``: for inverted file, divide the dataset in partitions + and perform search on relevant partitions only + - ``'ivfpq'``: for inverted file and product quantization, + same as inverted list, in addition the vectors are broken + in n_features/M sub-vectors that will be encoded thanks + to intermediary k-means clusterings. This encoding provide + partial information allowing faster distances calculations + + metric : string (default='euclidean'). + Distance metric to use. Supported metrics include: 'l1', 'cityblock', + 'taxicab', 'manhattan', 'euclidean', 'l2', 'sqeuclidean', 'canberra', + 'minkowski', 'lp', 'chebyshev', 'linf', 'jensenshannon', 'cosine', + 'braycurtis', 'jaccard', 'hellinger', 'correlation', 'inner_product'. + The ``'ivfflat'`` and ``'ivfpq'`` + algorithms only support: 'euclidean', 'l2', 'sqeuclidean', 'cosine', + 'correlation', 'inner_product', whereas the ``'rbc'`` algorithm only + supports 'euclidean', 'l2', and 'haversine' (≤3 dimensions only). + For sparse inputs, only the ``'brute'`` algorithm is supported, with + metrics: 'l1', 'cityblock', 'taxicab', 'manhattan', 'euclidean', 'l2', + 'canberra', 'minkowski', 'lp', 'chebyshev', 'linf', 'cosine', + 'inner_product', 'jaccard', 'hellinger'. + p : float (default=2) + Parameter for the Minkowski metric. When p = 1, this is equivalent to + manhattan distance (l1), and euclidean distance (l2) for p = 2. For + arbitrary p, minkowski distance (lp) is used. + algo_params : dict, optional (default=None) + Used to configure the nearest neighbor algorithm to be used. + If set to None, parameters will be generated automatically. + Parameters for algorithm ``'brute'`` when inputs are sparse: + + - batch_size_index : (int) number of rows in each batch of \ + index array + - batch_size_query : (int) number of rows in each batch of \ + query array + + Parameters for algorithm ``'ivfflat'``: + + - nlist: (int) number of cells to partition dataset into + - nprobe: (int) at query time, number of cells used for search + + Parameters for algorithm ``'ivfpq'``: + + - nlist: (int) number of cells to partition dataset into + - nprobe: (int) at query time, number of cells used for search + - M: (int) number of subquantizers + - n_bits: (int) bits allocated per subquantizer + - usePrecomputedTables : (bool) whether to use precomputed tables + metric_params : dict, optional (default = None) + This is currently ignored. + n_jobs : int (default = None) + Ignored, here for scikit-learn API compatibility. + output_type : {'input', 'array', 'dataframe', 'series', 'df_obj', \ + 'numba', 'cupy', 'numpy', 'cudf', 'pandas'}, default=None + Return results and set estimator attributes to the indicated output + type. If None, the output type set at the module level + (`cuml.global_settings.output_type`) will be used. See + :ref:`output-data-type-configuration` for more info. + + Examples + -------- + + .. code-block:: python + + >>> import cudf + >>> from cuml.neighbors import NearestNeighbors + >>> from cuml.datasets import make_blobs + + >>> X, _ = make_blobs(n_samples=5, centers=5, + ... n_features=10, random_state=42) + + >>> # build a cudf Dataframe + >>> X_cudf = cudf.DataFrame(X) + + >>> # fit model + >>> model = NearestNeighbors(n_neighbors=3) + >>> model.fit(X) + NearestNeighbors() + + >>> # get 3 nearest neighbors + >>> distances, indices = model.kneighbors(X_cudf) + + >>> # print results + >>> print(indices) # doctest: +SKIP + 0 1 2 + 0 0 3 1 + 1 1 3 0 + 2 2 4 0 + 3 3 0 1 + 4 4 2 0 + >>> print(distances) # doctest: +SKIP + 0 1 2 + 0 0.007812 24.786566 26.399996 + 1 0.000000 24.786566 30.045017 + 2 0.007812 5.458400 27.051241 + 3 0.000000 26.399996 27.543869 + 4 0.000000 5.458400 29.583437 + + Notes + ----- + For an additional example see `the NearestNeighbors notebook + `_. + + For additional docs, see `scikit-learn's NearestNeighbors + `_. + + Pickling ``NearestNeighbors`` instances is supported for all algorithms. + However, for RBC, IVFPQ or IVFFlat the index will currently be rebuilt upon + load rather than serialized as part of the pickled binary. For approximate + indices like IVFPQ or IVFFlat this may result in small differences between + the original and reloaded models, as the generated indices may differ. + """ + _cpu_class_path = "sklearn.neighbors.NearestNeighbors" + + def __init__( + self, + *, + n_neighbors=5, + radius=1.0, + algorithm="auto", + metric="euclidean", + p=2, + algo_params=None, + metric_params=None, + n_jobs=None, # Ignored, here for sklearn API compatibility + verbose=False, + output_type=None, + ): + self.radius = radius + super().__init__( + n_neighbors=n_neighbors, + algorithm=algorithm, + metric=metric, + p=p, + algo_params=algo_params, + metric_params=metric_params, + n_jobs=n_jobs, + verbose=verbose, + output_type=output_type + ) + + @classmethod + def _get_param_names(cls): + return ["radius", *super()._get_param_names()] + + @classmethod + def _params_from_cpu(cls, model): + return { + "radius": model.radius, + **super()._params_from_cpu(model), + } + + def _params_to_cpu(self): + return { + "radius": self.radius, + **super()._params_to_cpu(), + } + + @insert_into_docstring(parameters=[('dense', '(n_samples, n_features)')]) + @reflect + def radius_neighbors_graph(self, X=None, radius=None) -> SparseCumlArray: + """Compute the (weighted) graph of neighbors within a radius. + + Parameters + ---------- + X : array-like, default=None + The query point or points. If not provided, neighbors of each indexed + point are returned. In this case, the query point is not considered its + own neighbor. + + radius : float, default=None + Radius of neighborhoods. The default is the value passed to the + constructor. + + Returns + ------- + A : sparse-matrix of shape (n_queries, n_samples_fit) + The neighborhood graph, in CSR format. + + Notes + ----- + This method is most efficient when the instance is fit with + `algorithm="rbc"`. Other algorithms will build a temporary RBC index + per-call, which adds a small overhead. + + Only euclidean/l2 metrics and dense inputs are currently supported. + + Examples + -------- + >>> import cupy as cp + >>> from cuml.neighbors import NearestNeighbors + >>> X = cp.array([[0], [3], [1]]) + >>> nn = NearestNeighbors().fit(X) + >>> A = nn.radius_neighbors_graph(X, radius=1.5) + >>> A.toarray() + array([[1., 0., 1.], + [0., 1., 0.], + [1., 0., 1.]]) + """ + if not hasattr(self, "_fit_X"): + raise ValueError("This NearestNeighbors instance has not been " + "fitted yet, call 'fit' before using this " + "estimator") + + if isinstance(self._fit_X, SparseCumlArray) or is_sparse(X): + raise TypeError("`radius_neighbors_graph` doesn't support sparse inputs") + + if self.effective_metric_ not in ["l2", "euclidean"]: + raise ValueError( + f"`radius_neighbors_graph` doesn't support " + f"metric={self.effective_metric_!r}" + ) + + if radius is None: + radius = self.radius + + if radius <= 0: + raise ValueError(f"Expected `radius > 0`, got {radius}") + + if (using_fit_X := (X is None)): + X = self._fit_X + + X_m = input_to_cuml_array( + X, + order="C", + check_dtype=np.float32, + check_cols=self.n_features_in_, + convert_to_dtype=np.float32, + ).array + + if hasattr(self, "_index") and isinstance(self._index, RBCIndex): + # Already fit with RBC, reuse the index + index = self._index + else: + # Fit with another method, build a temporary index + index = RBCIndex.build(self._fit_X, self.effective_metric_) + + out = index.radius_neighbors_graph(X_m, radius) + if using_fit_X: + # When using the training data, the diagonal elements aren't included + out.setdiag(np.int64(0)) + out.eliminate_zeros() + + return out + + @reflect def kneighbors_graph( X=None, diff --git a/python/cuml/tests/test_nearest_neighbors.py b/python/cuml/tests/test_nearest_neighbors.py index 5c6e0fb5d3..93e6704a1a 100644 --- a/python/cuml/tests/test_nearest_neighbors.py +++ b/python/cuml/tests/test_nearest_neighbors.py @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2019-2025, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION. # SPDX-License-Identifier: Apache-2.0 # @@ -557,6 +557,56 @@ def test_knn_graph_algorithm(algorithm): assert ((sk_graph - cu_graph.get()) < 0).sum() == 0 +@pytest.mark.parametrize("n_features", [2, 3, 5]) +@pytest.mark.parametrize("radius", [2.5, None]) +@pytest.mark.parametrize("self_query", [True, False]) +def test_radius_neighbors_graph(n_features, radius, self_query): + X, _ = make_blobs(n_samples=500, n_features=n_features, random_state=42) + X_train, X_query = X[:400], X[400:] + + sk_model = skKNN(radius=1.5).fit(X_train.get()) + cu_model = cuKNN(radius=1.5).fit(X_train) + + sk_graph = sk_model.radius_neighbors_graph( + None if self_query else X_query.get(), radius=radius + ) + cu_graph = cu_model.radius_neighbors_graph( + None if self_query else X_query, radius=radius + ) + + assert cupyx.scipy.sparse.isspmatrix_csr(cu_graph) + np.testing.assert_array_equal( + sk_graph.toarray(), + cu_graph.toarray().get(), + ) + + +def test_radius_neighbors_graph_errors(): + X, _ = make_blobs(n_samples=100, random_state=42) + X_sparse = cupyx.scipy.sparse.random( + 100, 5, format="csr", density=0.5, random_state=42 + ) + + # Unsupported sparse inputs + model = cuKNN().fit(X) + model_sparse = cuKNN().fit(X_sparse) + + with pytest.raises(TypeError, match="doesn't support sparse inputs"): + model.radius_neighbors_graph(X_sparse) + + with pytest.raises(TypeError, match="doesn't support sparse inputs"): + model_sparse.radius_neighbors_graph(X) + + # Invalid radius + with pytest.raises(ValueError, match="Expected `radius > 0`, got -2"): + model.radius_neighbors_graph(radius=-2) + + # Unsupported metric + model = cuKNN(metric="linf").fit(X) + with pytest.raises(ValueError, match="doesn't support metric='linf'"): + model.radius_neighbors_graph() + + @pytest.mark.parametrize( "distance_dims", [("euclidean", 2), ("euclidean", 3), ("haversine", 2)] )