diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index f05dfad14b..685862b6c9 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -427,7 +427,7 @@ if(BUILD_CUML_CPP_LIBRARY) endif() if(all_algo OR knn_algo) - target_sources(cuml_objs PRIVATE src/knn/knn.cu src/knn/knn_sparse.cu) + target_sources(cuml_objs PRIVATE src/knn/knn.cu src/knn/knn_sparse.cu src/kde/kde.cu) endif() if(all_algo OR hierarchicalclustering_algo) diff --git a/cpp/include/cuml/neighbors/kde.hpp b/cpp/include/cuml/neighbors/kde.hpp new file mode 100644 index 0000000000..8d6fd1c76f --- /dev/null +++ b/cpp/include/cuml/neighbors/kde.hpp @@ -0,0 +1,90 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ +#pragma once + +#include + +#include + +#include + +namespace ML::KDE { + +enum class DensityKernelType : int { + Gaussian = 0, + Tophat = 1, + Epanechnikov = 2, + Exponential = 3, + Linear = 4, + Cosine = 5 +}; + +/** + * @brief Compute normalized log-density scores for query samples. + * + * The query and training arrays must be dense row-major (C-contiguous) + * device arrays with shapes `(n_query, n_features)` and + * `(n_train, n_features)`, respectively. + * + * @tparam T floating point type, either float or double + * @param[in] handle raft resources used to launch work + * @param[in] query device pointer to query samples in row-major order + * @param[in] train device pointer to training samples in row-major order + * @param[in] weights optional device pointer to sample weights of length + * `n_train`, or nullptr for uniform weights + * @param[out] output device pointer to log-density scores of length `n_query` + * @param[in] n_query number of query samples + * @param[in] n_train number of training samples + * @param[in] n_features number of features per sample + * @param[in] bandwidth positive KDE bandwidth + * @param[in] sum_weights sum of `weights`, or `n_train` when weights is null + * @param[in] kernel density kernel to evaluate + * @param[in] metric distance metric used between query and training samples + * @param[in] metric_arg metric-specific argument, such as p for Minkowski + */ +template +void score_samples(raft::resources const& handle, + const T* query, + const T* train, + const T* weights, + T* output, + std::int64_t n_query, + std::int64_t n_train, + std::int64_t n_features, + T bandwidth, + T sum_weights, + DensityKernelType kernel, + ML::distance::DistanceType metric, + T metric_arg); + +extern template void score_samples(raft::resources const&, + const float*, + const float*, + const float*, + float*, + std::int64_t, + std::int64_t, + std::int64_t, + float, + float, + DensityKernelType, + ML::distance::DistanceType, + float); + +extern template void score_samples(raft::resources const&, + const double*, + const double*, + const double*, + double*, + std::int64_t, + std::int64_t, + std::int64_t, + double, + double, + DensityKernelType, + ML::distance::DistanceType, + double); + +} // namespace ML::KDE diff --git a/cpp/src/kde/kde.cu b/cpp/src/kde/kde.cu new file mode 100644 index 0000000000..d2251ed251 --- /dev/null +++ b/cpp/src/kde/kde.cu @@ -0,0 +1,83 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include + +#include + +#include + +#include + +namespace ML::KDE { + +template +void score_samples(raft::resources const& handle, + const T* query, + const T* train, + const T* weights, + T* output, + std::int64_t n_query, + std::int64_t n_train, + std::int64_t n_features, + T bandwidth, + T sum_weights, + DensityKernelType kernel, + ML::distance::DistanceType metric, + T metric_arg) +{ + auto query_view = + raft::make_device_matrix_view(query, n_query, n_features); + auto train_view = + raft::make_device_matrix_view(train, n_train, n_features); + auto output_view = raft::make_device_vector_view(output, n_query); + auto weights_view = + weights + ? std::make_optional(raft::make_device_vector_view(weights, n_train)) + : std::nullopt; + auto cuvs_kernel = static_cast(kernel); + auto cuvs_metric = static_cast(metric); + + cuvs::distance::kde(handle, + query_view, + train_view, + weights_view, + output_view, + bandwidth, + sum_weights, + cuvs_kernel, + cuvs_metric, + metric_arg); +} + +template void score_samples(raft::resources const&, + const float*, + const float*, + const float*, + float*, + std::int64_t, + std::int64_t, + std::int64_t, + float, + float, + DensityKernelType, + ML::distance::DistanceType, + float); + +template void score_samples(raft::resources const&, + const double*, + const double*, + const double*, + double*, + std::int64_t, + std::int64_t, + std::int64_t, + double, + double, + DensityKernelType, + ML::distance::DistanceType, + double); + +} // namespace ML::KDE diff --git a/python/cuml/cuml/neighbors/CMakeLists.txt b/python/cuml/cuml/neighbors/CMakeLists.txt index ad744c9669..9269f4a6ce 100644 --- a/python/cuml/cuml/neighbors/CMakeLists.txt +++ b/python/cuml/cuml/neighbors/CMakeLists.txt @@ -1,11 +1,12 @@ # ============================================================================= # cmake-format: off -# SPDX-FileCopyrightText: Copyright (c) 2022-2025, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION. # SPDX-License-Identifier: Apache-2.0 # cmake-format: on # ============================================================================= set(cython_sources "") +add_module_gpu_default("kernel_density.pyx" ${knn_algo} ${neighbors_algo}) add_module_gpu_default("kneighbors_classifier.pyx" ${kneighbors_classifier_algo} ${neighbors_algo}) add_module_gpu_default("kneighbors_regressor.pyx" ${kneighbors_regressor_algo} ${neighbors_algo}) add_module_gpu_default("nearest_neighbors.pyx" ${nearest_neighbors_algo} ${neighbors_algo}) diff --git a/python/cuml/cuml/neighbors/kernel_density.py b/python/cuml/cuml/neighbors/kernel_density.pyx similarity index 63% rename from python/cuml/cuml/neighbors/kernel_density.py rename to python/cuml/cuml/neighbors/kernel_density.pyx index 8acaedbf9b..0b56426248 100644 --- a/python/cuml/cuml/neighbors/kernel_density.py +++ b/python/cuml/cuml/neighbors/kernel_density.pyx @@ -3,14 +3,15 @@ # SPDX-License-Identifier: Apache-2.0 # -import math +import warnings import cupy as cp import numpy as np from cupyx.scipy.special import gammainc +from sklearn.exceptions import DataConversionWarning from cuml.internals.array import CumlArray -from cuml.internals.base import Base +from cuml.internals.base import Base, get_handle from cuml.internals.interop import InteropMixin, UnsupportedOnGPU from cuml.internals.outputs import reflect, run_in_internal_context from cuml.internals.validation import ( @@ -19,122 +20,84 @@ check_non_negative, check_random_seed, ) -from cuml.metrics import pairwise_distances from cuml.metrics.pairwise_distances import ( PAIRWISE_DISTANCE_METRICS as SUPPORTED_METRICS, ) -VALID_KERNELS = [ - "gaussian", - "tophat", - "epanechnikov", - "exponential", - "linear", - "cosine", -] +from libc.stdint cimport int64_t, uintptr_t +from libcpp cimport bool as cpp_bool +from pylibraft.common.handle cimport handle_t + +from cuml.metrics.distance_type cimport DistanceType + + +cdef extern from "cuml/neighbors/kde.hpp" nogil: + + ctypedef enum class DensityKernelType "ML::KDE::DensityKernelType": + Gaussian "ML::KDE::DensityKernelType::Gaussian" + Tophat "ML::KDE::DensityKernelType::Tophat" + Epanechnikov "ML::KDE::DensityKernelType::Epanechnikov" + Exponential "ML::KDE::DensityKernelType::Exponential" + Linear "ML::KDE::DensityKernelType::Linear" + Cosine "ML::KDE::DensityKernelType::Cosine" + + void _cuml_kde_score_samples \ + "ML::KDE::score_samples"(const handle_t &handle, + const float *query, + const float *train, + const float *weights, + float *output, + int64_t n_query, + int64_t n_train, + int64_t n_features, + float bandwidth, + float sum_weights, + DensityKernelType kernel, + DistanceType metric, + float metric_arg) except + + + void _cuml_kde_score_samples \ + "ML::KDE::score_samples"(const handle_t &handle, + const double *query, + const double *train, + const double *weights, + double *output, + int64_t n_query, + int64_t n_train, + int64_t n_features, + double bandwidth, + double sum_weights, + DensityKernelType kernel, + DistanceType metric, + double metric_arg) except + + + +KDE_KERNEL_TYPES = { + "gaussian": DensityKernelType.Gaussian, + "tophat": DensityKernelType.Tophat, + "epanechnikov": DensityKernelType.Epanechnikov, + "exponential": DensityKernelType.Exponential, + "linear": DensityKernelType.Linear, + "cosine": DensityKernelType.Cosine, +} +VALID_KERNELS = list(KDE_KERNEL_TYPES.keys()) -@cp.fuse() -def gaussian_log_kernel(x, h): - return -(x * x) / (2 * h * h) +def _coerce_russellrao_binary(arr, *, input_name): + """Coerce values to {0, 1} for the russellrao metric, warning on non-binary input. -@cp.fuse() -def tophat_log_kernel(x, h): - """ - if x < h: - return 0.0 - else: - return -FLOAT_MIN + The fused KDE kernel computes RussellRao assuming binary inputs, matching + the behavior of ``cuml.metrics.pairwise_distances`` for this metric. """ - y = (x >= h) * np.finfo(x.dtype).min - return y - - -@cp.fuse() -def epanechnikov_log_kernel(x, h): - # don't call log(0) otherwise we get NaNs - z = cp.maximum(1.0 - (x * x) / (h * h), 1e-30) - y = (x < h) * cp.log(z) - y += (x >= h) * np.finfo(y.dtype).min - return y - - -@cp.fuse() -def exponential_log_kernel(x, h): - return -x / h - - -@cp.fuse() -def linear_log_kernel(x, h): - # don't call log(0) otherwise we get NaNs - z = cp.maximum(1.0 - x / h, 1e-30) - y = (x < h) * cp.log(z) - y += (x >= h) * np.finfo(y.dtype).min - return y - - -@cp.fuse() -def cosine_log_kernel(x, h): - # don't call log(0) otherwise we get NaNs - z = cp.maximum(cp.cos(0.5 * np.pi * x / h), 1e-30) - y = (x < h) * cp.log(z) - y += (x >= h) * np.finfo(y.dtype).min - return y - - -log_probability_kernels_ = { - "gaussian": gaussian_log_kernel, - "tophat": tophat_log_kernel, - "epanechnikov": epanechnikov_log_kernel, - "exponential": exponential_log_kernel, - "linear": linear_log_kernel, - "cosine": cosine_log_kernel, -} - - -def logVn(n): - return 0.5 * n * np.log(np.pi) - math.lgamma(0.5 * n + 1) - - -def logSn(n): - return np.log(2 * np.pi) + logVn(n - 1) - - -def norm_factor(kernel, h, d): - if kernel == "gaussian": - factor = 0.5 * d * np.log(2 * np.pi) - elif kernel == "tophat": - factor = logVn(d) - elif kernel == "epanechnikov": - factor = logVn(d) + np.log(2.0 / (d + 2.0)) - elif kernel == "exponential": - factor = logSn(d - 1) + math.lgamma(d) - elif kernel == "linear": - factor = logVn(d) - np.log(d + 1.0) - elif kernel == "cosine": - factor = 0.0 - tmp = 2.0 / np.pi - for k in range(1, d + 1, 2): - factor += tmp - tmp *= -(d - k) * (d - k - 1) * (2.0 / np.pi) ** 2 - factor = np.log(factor) + logSn(d - 1) - else: - raise ValueError("Unsupported kernel.") - - return factor + d * np.log(h) - - -# Implements a faster (but simpler) version of `cupyx.scipy.special.logsumexp` -logsumexp = cp.ReductionKernel( - "T d", - "T out", - "exp(d)", - "a + b", - "out = log(a)", - "0", - "logsumexp", -) + if not bool(cp.logical_or(arr == 0, arr == 1).all()): + warnings.warn( + f"{input_name} was converted to boolean for metric 'russellrao'", + DataConversionWarning, + stacklevel=2, + ) + return cp.where(arr != 0, arr.dtype.type(1), arr.dtype.type(0)) + return arr class KernelDensity(Base, InteropMixin): @@ -224,7 +187,7 @@ def _attrs_from_cpu(self, model): sample_weight = ( None if model.tree_.sample_weight is None - else cp.asarray(model.tree_.sample_weight, dtype=cp.float32) + else cp.asarray(model.tree_.sample_weight, dtype=X.dtype) ) return { "bandwidth_": model.bandwidth_, @@ -291,6 +254,12 @@ def fit( if self.kernel not in VALID_KERNELS: raise ValueError(f"kernel={self.kernel!r} is not supported") + if self.metric == "nan_euclidean": + raise NotImplementedError( + "metric='nan_euclidean' is not supported by cuML's " + "KernelDensity; the fused kernel has no NaN-aware path." + ) + if isinstance(self.bandwidth, str): if self.bandwidth not in ("scott", "silverman"): raise ValueError( @@ -308,6 +277,8 @@ def fit( order="C", reset=True, ) + if self.metric == "russellrao": + self._X = _coerce_russellrao_binary(self._X, input_name="X") if self._sample_weight is not None: check_non_negative(self._sample_weight, input_name="sample_weight") @@ -350,67 +321,96 @@ def score_samples(self, X, *, convert_dtype=True) -> CumlArray: convert_dtype=convert_dtype, order="C", ) + if self.metric == "russellrao": + X = _coerce_russellrao_binary(X, input_name="X") + if self.metric_params: if len(self.metric_params) != 1: raise ValueError( "Cuml only supports metrics with a single arg." ) - metric_arg = list(self.metric_params.values())[0] - distances = pairwise_distances( - X, - self._X, - metric=self.metric, - metric_arg=metric_arg, - ) + metric_arg = float(next(iter(self.metric_params.values()))) else: - distances = pairwise_distances(X, self._X, metric=self.metric) + metric_arg = 2.0 - distances = cp.asarray(distances) + if self.metric not in SUPPORTED_METRICS: + raise ValueError(f"metric={self.metric!r} is not supported") - h = distances.dtype.type(self.bandwidth_) - if self.kernel in log_probability_kernels_: - # XXX: passing `h` as a 0-dim array works around dtype inference - # issues in cupy.fuse. See https://github.com/cupy/cupy/issues/9400 - distances = log_probability_kernels_[self.kernel]( - distances, cp.array(h, dtype=distances.dtype) - ) - else: - raise ValueError("Unsupported kernel.") - - if self._sample_weight is not None: - distances += cp.log(self._sample_weight) - - # To avoid overflow, we apply - # log(exp(x).sum()) -> log(exp(x - x.max())) + x.max() - # We subtract the max inplace to avoid an extra allocation, - # since `distances` is no longer needed after this point. - max_distances = distances.max(axis=1) - distances -= max_distances[:, None] - log_probabilities = logsumexp(distances, axis=1) - log_probabilities += max_distances - - # Note that sklearns user guide is wrong - # It says the (unnormalised) probability output for - # the kernel density is sum(K(x,h)). - # In fact what they implement is (1/n)*sum(K(x,h)) - # Here we divide by n in normal probability space - # Which becomes -log(n) in log probability space sum_weights = ( - cp.sum(self._sample_weight) + float(cp.sum(self._sample_weight)) if self._sample_weight is not None - else distances.shape[1] + else float(self._X.shape[0]) ) - log_probabilities -= np.log(sum_weights) - # norm - if len(X.shape) == 1: - # if X is one dimensional, we have 1 feature - dimension = 1 - else: - dimension = X.shape[1] - log_probabilities -= norm_factor(self.kernel, h, dimension) + cdef DensityKernelType kernel_enum = KDE_KERNEL_TYPES[self.kernel] + cdef DistanceType metric_enum = SUPPORTED_METRICS[self.metric] + + cdef cpp_bool is_float32 = X.dtype == np.float32 + cdef int64_t n_query = X.shape[0] + cdef int64_t n_train = self._X.shape[0] + cdef int64_t n_features = X.shape[1] if len(X.shape) > 1 else 1 + + output = cp.empty(n_query, dtype=X.dtype) + + cdef uintptr_t query_ptr = X.data.ptr + cdef uintptr_t train_ptr = self._X.data.ptr + cdef uintptr_t weight_ptr = 0 + if self._sample_weight is not None: + weight_ptr = self._sample_weight.data.ptr + cdef uintptr_t output_ptr = output.data.ptr + + cdef const float* weights_f = ( + weight_ptr if weight_ptr != 0 + else NULL + ) + cdef const double* weights_d = ( + weight_ptr if weight_ptr != 0 + else NULL + ) + + cdef double c_bandwidth = self.bandwidth_ + cdef double c_sum_weights = sum_weights + cdef double c_metric_arg = metric_arg + + handle = get_handle() + cdef handle_t* handle_ = handle.getHandle() + + with nogil: + if is_float32: + _cuml_kde_score_samples( + handle_[0], + query_ptr, + train_ptr, + weights_f, + output_ptr, + n_query, + n_train, + n_features, + c_bandwidth, + c_sum_weights, + kernel_enum, + metric_enum, + c_metric_arg, + ) + else: + _cuml_kde_score_samples( + handle_[0], + query_ptr, + train_ptr, + weights_d, + output_ptr, + n_query, + n_train, + n_features, + c_bandwidth, + c_sum_weights, + kernel_enum, + metric_enum, + c_metric_arg, + ) + handle.sync() - return log_probabilities + return CumlArray(data=output) @run_in_internal_context def score(self, X, y=None) -> float: diff --git a/python/cuml/tests/test_kernel_density.py b/python/cuml/tests/test_kernel_density.py index 93c48a849d..27fca89b25 100644 --- a/python/cuml/tests/test_kernel_density.py +++ b/python/cuml/tests/test_kernel_density.py @@ -2,30 +2,65 @@ # SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION. # SPDX-License-Identifier: Apache-2.0 # +import warnings import cupy as cp import numpy as np import pytest +import scipy.special import sklearn.neighbors from hypothesis import assume, example, given, settings from hypothesis import strategies as st from hypothesis.extra.numpy import arrays from sklearn.datasets import make_blobs -from sklearn.exceptions import NotFittedError +from sklearn.exceptions import DataConversionWarning, NotFittedError from sklearn.metrics import pairwise_distances as skl_pairwise_distances from sklearn.model_selection import GridSearchCV from sklearn.neighbors._ball_tree import kernel_norm import cuml from cuml.neighbors import VALID_KERNELS, KernelDensity -from cuml.neighbors.kernel_density import logsumexp from cuml.testing.utils import as_type +def _cosine_kernel_norm(h, d): + """Normalization constant for cosine kernel in d dimensions. + + Matches the recurrence in cuvs::distance for DensityKernelType::Cosine: + I_0 = 2/pi + I_1 = 2/pi - (2/pi)^2 + I_n = 2/pi - n*(n-1)*(2/pi)^2 * I_{n-2} for n >= 2 + norm = 1 / (S_{d-1} * I_{d-1} * h^d) + where S_{d-1} = 2*pi^(d/2) / Gamma(d/2). + + sklearn's kernel_norm returns NaN for cosine at d >= 4 due to a bug in + its integration-by-parts formula, so we use this custom implementation. + """ + two_over_pi = 2.0 / np.pi + two_over_pi_sq = two_over_pi**2 + I_prev = two_over_pi # I_0 + I_curr = two_over_pi - two_over_pi_sq # I_1 + n = d - 1 + if n == 0: + integral = I_prev + else: + for j in range(2, n + 1): + I_next = two_over_pi - j * (j - 1) * two_over_pi_sq * I_prev + I_prev, I_curr = I_curr, I_next + integral = I_curr + Sn = 2.0 * np.pi ** (d / 2.0) / scipy.special.gamma(d / 2.0) + return 1.0 / (Sn * integral * h**d) + + # not in log probability space def compute_kernel_naive(Y, X, kernel, metric, h, sample_weight): - d = skl_pairwise_distances(Y, X, metric) - norm = kernel_norm(h, X.shape[1], kernel) + with warnings.catch_warnings(): + warnings.simplefilter("ignore", DataConversionWarning) + d = skl_pairwise_distances(Y, X, metric) + if kernel == "cosine": + norm = _cosine_kernel_norm(h, X.shape[1]) + else: + norm = kernel_norm(h, X.shape[1], kernel) if kernel == "gaussian": k = np.exp(-0.5 * (d * d) / (h * h)) @@ -38,7 +73,7 @@ def compute_kernel_naive(Y, X, kernel, metric, h, sample_weight): elif kernel == "linear": k = (1 - d / h) * (d < h) elif kernel == "cosine": - k = np.cos(0.5 * np.pi * d / h) * (d < h) + k = np.maximum(np.cos(0.5 * np.pi * d / h), 1e-30) * (d < h) else: raise ValueError("kernel not recognized") return norm * np.average(k, -1, sample_weight) @@ -185,16 +220,6 @@ def test_score_samples_output_type_and_dtype(kernel, fit_dtype, score_dtype): assert isinstance(res, cp.ndarray) -def test_logsumexp(): - X = np.array([[0.0, 0.0], [0.0, 0.0]]) - out = logsumexp(cp.asarray(X), axis=1).get() - assert np.allclose(out, np.logaddexp.reduce(X, axis=1)) - - X = np.array([[3.0, 1.0], [0.2, 0.7]]) - out = logsumexp(cp.asarray(X), axis=1).get() - assert np.allclose(out, np.logaddexp.reduce(X, axis=1)) - - def test_metric_params(): X = np.array([[0.0, 1.0], [2.0, 0.5]]) kde = KernelDensity(metric="minkowski", metric_params={"p": 1.0}).fit(X) @@ -236,3 +261,232 @@ def test_bad_sample_weight_errors(): ValueError, match="Sample weights must be 1D array or scalar" ): kde.fit(X, sample_weight=np.array([[1, 2], [3, 4]])) + + +# ----------------------------------------------------------------------------- +# Reference pairwise distances for metrics absent from sklearn.pairwise +# (must match the corresponding DistOp accumulate/finalize in kde.cu exactly) +# ----------------------------------------------------------------------------- + + +def _hellinger_dist(X, Y): + """sqrt(max(0, 1 - sum sqrt(xi * yi))) - matches DistOp.""" + sx = np.sqrt(np.maximum(X, 0.0)) + sy = np.sqrt(np.maximum(Y, 0.0)) + return np.sqrt(np.maximum(1.0 - sx @ sy.T, 0.0)) + + +def _jensenshannon_dist(X, Y): + """sqrt(0.5 * sum(a * log(a/m) + b * log(b/m))) - matches DistOp.""" + out = np.zeros((len(X), len(Y))) + for i, a in enumerate(X): + for j, b in enumerate(Y): + m = 0.5 * (a + b) + # Mirror device guards: log(0) -> 0 + logM = np.where(m > 0, np.log(np.where(m > 0, m, 1.0)), 0.0) + logA = np.where(a > 0, np.log(np.where(a > 0, a, 1.0)), 0.0) + logB = np.where(b > 0, np.log(np.where(b > 0, b, 1.0)), 0.0) + acc = np.sum(-a * (logM - logA) + -b * (logM - logB)) + out[i, j] = np.sqrt(0.5 * max(float(acc), 0.0)) + return out + + +def _kldivergence_dist(X, Y): + """sum a * log(a/b) for a,b > 0 - matches DistOp.""" + out = np.zeros((len(X), len(Y))) + for i, a in enumerate(X): + for j, b in enumerate(Y): + mask = (a > 0) & (b > 0) + out[i, j] = float(np.sum(a[mask] * np.log(a[mask] / b[mask]))) + return out + + +def _kde_naive_custom(Y, X, kernel, dist_fn, h, sample_weight): + """Like compute_kernel_naive but accepts a callable pairwise distance.""" + d = dist_fn(Y, X) + if kernel == "cosine": + norm = _cosine_kernel_norm(h, X.shape[1]) + else: + norm = kernel_norm(h, X.shape[1], kernel) + if kernel == "gaussian": + k = np.exp(-0.5 * d * d / (h * h)) + elif kernel == "tophat": + k = (d < h).astype(float) + elif kernel == "epanechnikov": + k = np.maximum(1.0 - d * d / (h * h), 0.0) * (d < h) + elif kernel == "exponential": + k = np.exp(-d / h) + elif kernel == "linear": + k = np.maximum(1.0 - d / h, 0.0) * (d < h) + elif kernel == "cosine": + k = np.maximum(np.cos(0.5 * np.pi * d / h), 1e-30) * (d < h) + else: + raise ValueError(kernel) + return norm * np.average(k, axis=1, weights=sample_weight) + + +# Custom distance functions for metrics not in sklearn.pairwise_distances +_CUSTOM_DIST_FN = { + "hellinger": _hellinger_dist, + "jensenshannon": _jensenshannon_dist, + "kldivergence": _kldivergence_dist, +} + +# Metrics that require non-negative inputs +_NONNEG_METRICS = {"hellinger", "jensenshannon"} +# Metrics that require strictly positive inputs +_POSONLY_METRICS = {"kldivergence"} +# Metrics defined for binary {0,1} inputs (our DistOp matches sklearn only for binary) +_BINARY_METRICS = {"russellrao"} + + +def _make_metric_data(metric, n_train=40, n_query=8, d=4, seed=7): + """Generate float64 test data appropriate for the given metric.""" + rng = np.random.RandomState(seed) + if metric in _BINARY_METRICS: + X = rng.randint(0, 2, size=(n_train, d)).astype(np.float64) + Q = rng.randint(0, 2, size=(n_query, d)).astype(np.float64) + elif metric in _POSONLY_METRICS: + X = ( + rng.exponential(scale=1.0, size=(n_train, d)).astype(np.float64) + + 0.1 + ) + Q = ( + rng.exponential(scale=1.0, size=(n_query, d)).astype(np.float64) + + 0.1 + ) + elif metric in _NONNEG_METRICS: + X = np.abs(rng.randn(n_train, d)).astype(np.float64) + 0.05 + Q = np.abs(rng.randn(n_query, d)).astype(np.float64) + 0.05 + else: + X = rng.randn(n_train, d).astype(np.float64) + Q = rng.randn(n_query, d).astype(np.float64) + return X, Q + + +@pytest.mark.parametrize("kernel", VALID_KERNELS) +@pytest.mark.parametrize( + "metric", + [ + # sklearn-pairwise-compatible + "euclidean", + "manhattan", + "chebyshev", + "minkowski", + "sqeuclidean", + "canberra", + "hamming", + "cosine", + "correlation", + "russellrao", + # custom reference required + "hellinger", + "jensenshannon", + "kldivergence", + ], +) +def test_all_kernels_all_metrics(metric, kernel): + """Every metric x kernel combination produces output matching the reference. + + For metrics supported by sklearn.pairwise_distances the reference is + compute_kernel_naive; for metrics absent from sklearn a matching numpy + reference is used that mirrors the DistOp accumulate/finalize logic in + kde.cu exactly. + """ + X, Q = _make_metric_data(metric) + h = 1.0 + + kde = KernelDensity(kernel=kernel, metric=metric, bandwidth=h) + # fit with convert_dtype=False so float64 test data stays float64, + # matching the float64 Python reference distances. + kde.fit(X, convert_dtype=False) + cuml_log = as_type("numpy", kde.score_samples(Q)) + + # -inf is valid (zero density when all train points are beyond the bandwidth); + # only NaN indicates a real bug. + assert not np.any(np.isnan(cuml_log)), ( + f"NaN output for metric={metric}, kernel={kernel}" + ) + + dist_fn = _CUSTOM_DIST_FN.get(metric) + if dist_fn is not None: + ref = _kde_naive_custom(Q, X, kernel, dist_fn, h, None) + else: + ref = compute_kernel_naive(Q, X, kernel, metric, h, None) + + # exp(-inf) == 0 == reference density, so this naturally handles the + # all-zero-density case (compact-support kernels with small bandwidth). + cuml_prob = np.exp(cuml_log) + assert np.allclose(cuml_prob, ref, rtol=1e-3, atol=1e-3, equal_nan=True), ( + f"metric={metric}, kernel={kernel}: max err=" + f"{np.max(np.abs(cuml_prob - ref)):.4e}" + ) + + +def test_nan_euclidean_not_supported(): + """metric='nan_euclidean' is rejected with a clear error. + + The fused KDE kernel has no NaN-aware path, so we raise instead of + silently producing wrong results. + """ + rng = np.random.RandomState(0) + X = rng.randn(20, 3).astype(np.float64) + kde = KernelDensity(metric="nan_euclidean") + with pytest.raises(NotImplementedError, match="nan_euclidean"): + kde.fit(X) + + +def test_russellrao_coerces_non_binary_with_warning(): + """Non-binary inputs to metric='russellrao' are coerced to {0, 1} with a warning. + + Mirrors the long-standing behavior of cuml.metrics.pairwise_distances for + this metric. Output must match a reference computed on the coerced data. + """ + rng = np.random.RandomState(0) + X = rng.uniform(-1.0, 2.0, size=(30, 4)).astype(np.float64) + Q = rng.uniform(-1.0, 2.0, size=(5, 4)).astype(np.float64) + + kde = KernelDensity(kernel="gaussian", metric="russellrao", bandwidth=1.0) + with pytest.warns(DataConversionWarning, match="converted to boolean"): + kde.fit(X) + with pytest.warns(DataConversionWarning, match="converted to boolean"): + cuml_log = as_type("numpy", kde.score_samples(Q)) + + X_bin = np.where(X != 0.0, 1.0, 0.0) + Q_bin = np.where(Q != 0.0, 1.0, 0.0) + ref = compute_kernel_naive( + Q_bin, X_bin, "gaussian", "russellrao", 1.0, None + ) + assert np.allclose(np.exp(cuml_log), ref, rtol=1e-4, atol=1e-4) + + +def test_russellrao_binary_no_warning(): + """Already-binary inputs to metric='russellrao' do not trigger a warning.""" + rng = np.random.RandomState(0) + X = rng.randint(0, 2, size=(30, 4)).astype(np.float64) + Q = rng.randint(0, 2, size=(5, 4)).astype(np.float64) + kde = KernelDensity(kernel="gaussian", metric="russellrao", bandwidth=1.0) + with warnings.catch_warnings(): + warnings.simplefilter("error", DataConversionWarning) + kde.fit(X) + kde.score_samples(Q) + + +def test_tiling_multipass(): + """Multi-pass tiling path (small n_query, large n_train) matches reference. + + When n_query is small enough that the 2-D grid / multi-pass reduction + code path is taken the result must match the naive single-pass reference. + """ + rng = np.random.RandomState(0) + X_train = rng.randn(2000, 4).astype(np.float64) + X_query = rng.randn(2, 4).astype(np.float64) + + kde = KernelDensity(kernel="gaussian", metric="euclidean", bandwidth=0.5) + kde.fit(X_train) + cuml_scores = as_type("numpy", kde.score_samples(X_query)) + + ref = compute_kernel_naive( + X_query, X_train, "gaussian", "euclidean", 0.5, None + ) + assert np.allclose(np.exp(cuml_scores), ref, rtol=1e-3, atol=1e-3)