diff --git a/c/src/cluster/kmeans.cpp b/c/src/cluster/kmeans.cpp index 698f387616..57b6282c20 100644 --- a/c/src/cluster/kmeans.cpp +++ b/c/src/cluster/kmeans.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ @@ -63,12 +63,10 @@ void _fit(cuvsResources_t res, RAFT_FAIL("float64 is an unsupported dtype for hierarchical kmeans"); } else { auto kmeans_params = convert_balanced_params(params); - cuvs::cluster::kmeans::fit(*res_ptr, - kmeans_params, - cuvs::core::from_dlpack(X_tensor), - cuvs::core::from_dlpack(centroids_tensor)); - - *inertia = 0; + T inertia_temp; + auto inertia_view = raft::make_host_scalar_view(&inertia_temp); + cuvs::cluster::kmeans::fit(*res_ptr, kmeans_params, cuvs::core::from_dlpack(X_tensor), cuvs::core::from_dlpack(centroids_tensor), std::make_optional(inertia_view)); + *inertia = inertia_temp; *n_iter = params.hierarchical_n_iters; } } else { diff --git a/cpp/include/cuvs/cluster/kmeans.hpp b/cpp/include/cuvs/cluster/kmeans.hpp index 22c8e056ec..a839cecf56 100644 --- a/cpp/include/cuvs/cluster/kmeans.hpp +++ b/cpp/include/cuvs/cluster/kmeans.hpp @@ -11,6 +11,8 @@ #include #include +#include + namespace cuvs::cluster::kmeans { /** Base structure for parameters that are common to all k-means algorithms */ @@ -420,11 +422,14 @@ void fit(raft::resources const& handle, * kmeans algorithm are stored at the address * pointed by 'centroids'. * [dim = n_clusters x n_features] + * @param[out] inertia Sum of squared distances of samples to their + * closest cluster center. */ void fit(const raft::resources& handle, cuvs::cluster::kmeans::balanced_params const& params, raft::device_matrix_view X, - raft::device_matrix_view centroids); + raft::device_matrix_view centroids, + std::optional> inertia = std::nullopt); /** * @brief Find balanced clusters with k-means algorithm. @@ -454,11 +459,14 @@ void fit(const raft::resources& handle, * kmeans algorithm are stored at the address * pointed by 'centroids'. * [dim = n_clusters x n_features] + * @param[out] inertia Sum of squared distances of samples to their + * closest cluster center. */ void fit(const raft::resources& handle, cuvs::cluster::kmeans::balanced_params const& params, raft::device_matrix_view X, - raft::device_matrix_view centroids); + raft::device_matrix_view centroids, + std::optional> inertia = std::nullopt); /** * @brief Find balanced clusters with k-means algorithm. @@ -488,11 +496,14 @@ void fit(const raft::resources& handle, * kmeans algorithm are stored at the address * pointed by 'centroids'. * [dim = n_clusters x n_features] + * @param[out] inertia Sum of squared distances of samples to their + * closest cluster center. */ void fit(const raft::resources& handle, cuvs::cluster::kmeans::balanced_params const& params, raft::device_matrix_view X, - raft::device_matrix_view centroids); + raft::device_matrix_view centroids, + std::optional> inertia = std::nullopt); /** * @brief Find balanced clusters with k-means algorithm. @@ -522,11 +533,14 @@ void fit(const raft::resources& handle, * kmeans algorithm are stored at the address * pointed by 'centroids'. * [dim = n_clusters x n_features] + * @param[out] inertia Sum of squared distances of samples to their + * closest cluster center. */ void fit(const raft::resources& handle, cuvs::cluster::kmeans::balanced_params const& params, raft::device_matrix_view X, - raft::device_matrix_view centroids); + raft::device_matrix_view centroids, + std::optional> inertia = std::nullopt); /** * @brief Predict the closest cluster each sample in X belongs to. @@ -1380,7 +1394,7 @@ void transform(raft::resources const& handle, raft::device_matrix_view X_new); /** - * @brief Compute cluster cost + * @brief Compute (optionally weighted) cluster cost * * @param[in] handle The raft handle * @param[in] X Training instances to cluster. The data must @@ -1390,12 +1404,16 @@ void transform(raft::resources const& handle, * row-major format. * [dim = n_clusters x n_features] * @param[out] cost Resulting cluster cost + * @param[in] sample_weight Optional per-sample weights. + * [len = n_samples] * */ -void cluster_cost(const raft::resources& handle, - raft::device_matrix_view X, - raft::device_matrix_view centroids, - raft::host_scalar_view cost); +void cluster_cost( + const raft::resources& handle, + raft::device_matrix_view X, + raft::device_matrix_view centroids, + raft::host_scalar_view cost, + std::optional> sample_weight = std::nullopt); /** * @brief Compute cluster cost @@ -1408,13 +1426,57 @@ void cluster_cost(const raft::resources& handle, * row-major format. * [dim = n_clusters x n_features] * @param[out] cost Resulting cluster cost + * @param[in] sample_weight Optional per-sample weights. + * [len = n_samples] + */ +void cluster_cost( + const raft::resources& handle, + raft::device_matrix_view X, + raft::device_matrix_view centroids, + raft::host_scalar_view cost, + std::optional> sample_weight = std::nullopt); + +/** + * @brief Compute (optionally weighted) cluster cost * + * @param[in] handle The raft handle + * @param[in] X Training instances to cluster. The data must + * be in row-major format. + * [dim = n_samples x n_features] + * @param[in] centroids Cluster centroids. The data must be in + * row-major format. + * [dim = n_clusters x n_features] + * @param[out] cost Resulting cluster cost + * @param[in] sample_weight Optional per-sample weights. + * [len = n_samples] */ -void cluster_cost(const raft::resources& handle, - raft::device_matrix_view X, - raft::device_matrix_view centroids, - raft::host_scalar_view cost); +void cluster_cost( + const raft::resources& handle, + raft::device_matrix_view X, + raft::device_matrix_view centroids, + raft::host_scalar_view cost, + std::optional> sample_weight = std::nullopt); +/** + * @brief Compute (optionally weighted) cluster cost + * + * @param[in] handle The raft handle + * @param[in] X Training instances to cluster. The data must + * be in row-major format. + * [dim = n_samples x n_features] + * @param[in] centroids Cluster centroids. The data must be in + * row-major format. + * [dim = n_clusters x n_features] + * @param[out] cost Resulting cluster cost + * @param[in] sample_weight Optional per-sample weights. + * [len = n_samples] + */ +void cluster_cost( + const raft::resources& handle, + raft::device_matrix_view X, + raft::device_matrix_view centroids, + raft::host_scalar_view cost, + std::optional> sample_weight = std::nullopt); /** * @} */ diff --git a/cpp/src/cluster/detail/kmeans_balanced.cuh b/cpp/src/cluster/detail/kmeans_balanced.cuh index 7582ec900e..f5dc759725 100644 --- a/cpp/src/cluster/detail/kmeans_balanced.cuh +++ b/cpp/src/cluster/detail/kmeans_balanced.cuh @@ -999,22 +999,22 @@ auto build_fine_clusters(const raft::resources& handle, /** * @brief Hierarchical balanced k-means * - * @tparam T element type - * @tparam MathT type of the centroids and mapped data - * @tparam IdxT index type - * @tparam LabelT label type + * @tparam T element type + * @tparam MathT type of the centroids and mapped data + * @tparam IdxT index type * @tparam MappingOpT type of the mapping operation * - * @param[in] handle The raft handle. - * @param[in] params Structure containing the hyper-parameters - * @param dim number of columns in `centers` and `dataset` - * @param[in] dataset a device pointer to the source dataset [n_rows, dim] - * @param n_rows number of rows in the input - * @param[out] cluster_centers a device pointer to the found cluster centers [n_cluster, dim] - * @param n_cluster - * @param metric the distance type - * @param mapping_op Mapping operation from T to MathT - * @param stream + * @param[in] handle The raft handle. + * @param[in] params Structure containing the hyper-parameters + * @param[in] dim Number of columns in `cluster_centers` and `dataset` + * @param[in] dataset A device pointer to the source dataset [n_rows, dim] + * @param[in] n_rows Number of rows in the input + * @param[out] cluster_centers A device pointer to the found cluster centers [n_clusters, dim] + * @param[in] n_clusters Requested number of clusters + * @param[in] mapping_op Mapping operation from T to MathT + * @param[out] inertia (optional) If non-null, the sum of squared distances of samples to + * their closest cluster center is written here. + * Only supported when T == MathT (float/double). */ template void build_hierarchical(const raft::resources& handle, @@ -1025,7 +1025,7 @@ void build_hierarchical(const raft::resources& handle, MathT* cluster_centers, IdxT n_clusters, MappingOpT mapping_op, - const MathT* dataset_norm = nullptr) + MathT* inertia = nullptr) { auto stream = raft::resource::get_cuda_stream(handle); using LabelT = uint32_t; @@ -1044,9 +1044,10 @@ void build_hierarchical(const raft::resources& handle, // Precompute the L2 norm of the dataset if relevant and not yet computed. rmm::device_uvector dataset_norm_buf(0, stream, device_memory); - if (dataset_norm == nullptr && (params.metric == cuvs::distance::DistanceType::L2Expanded || - params.metric == cuvs::distance::DistanceType::L2SqrtExpanded || - params.metric == cuvs::distance::DistanceType::CosineExpanded)) { + const MathT* dataset_norm = nullptr; + if ((params.metric == cuvs::distance::DistanceType::L2Expanded || + params.metric == cuvs::distance::DistanceType::L2SqrtExpanded || + params.metric == cuvs::distance::DistanceType::CosineExpanded)) { dataset_norm_buf.resize(n_rows, stream); for (IdxT offset = 0; offset < n_rows; offset += max_minibatch_size) { IdxT minibatch_size = std::min(max_minibatch_size, n_rows - offset); @@ -1164,6 +1165,20 @@ void build_hierarchical(const raft::resources& handle, MathT{0.2}, mapping_op, device_memory); + + // Compute inertia if requested (only supported when T == MathT) + if (inertia != nullptr) { + if constexpr (std::is_same_v) { + auto X_view = raft::make_device_matrix_view( + reinterpret_cast(dataset), n_rows, dim); + auto centroids_view = + raft::make_device_matrix_view(cluster_centers, n_clusters, dim); + cuvs::cluster::kmeans::cluster_cost( + handle, X_view, centroids_view, raft::make_host_scalar_view(inertia)); + } else { + RAFT_LOG_WARN("Inertia is not computed for non float/double types"); + } + } } } // namespace cuvs::cluster::kmeans::detail diff --git a/cpp/src/cluster/kmeans.cuh b/cpp/src/cluster/kmeans.cuh index 7d37b9cf80..e4f9821990 100644 --- a/cpp/src/cluster/kmeans.cuh +++ b/cpp/src/cluster/kmeans.cuh @@ -404,14 +404,27 @@ void min_cluster_distance(raft::resources const& handle, workspace); } +/** + * @brief Compute (optionally weighted) cluster cost (inertia). + * + * @tparam DataT float or double + * @tparam IndexT Index type + * + * @param[in] handle The raft handle + * @param[in] X Input data [n_samples x n_features] + * @param[in] centroids Cluster centroids [n_clusters x n_features] + * @param[out] cost Sum of squared distances to nearest centroid + * @param[in] sample_weight Optional per-sample weights [n_samples] + */ template -void cluster_cost(raft::resources const& handle, - raft::device_matrix_view X, - raft::device_matrix_view centroids, - raft::host_scalar_view cost) +void cluster_cost( + raft::resources const& handle, + raft::device_matrix_view X, + raft::device_matrix_view centroids, + raft::host_scalar_view cost, + std::optional> sample_weight = std::nullopt) { - auto stream = raft::resource::get_cuda_stream(handle); - + auto stream = raft::resource::get_cuda_stream(handle); auto n_clusters = centroids.extent(0); auto n_samples = X.extent(0); auto n_features = X.extent(1); @@ -440,6 +453,15 @@ void cluster_cost(raft::resources const& handle, n_clusters, workspace); + // Apply sample weights if provided + if (sample_weight.has_value()) { + raft::linalg::map(handle, + min_cluster_distance.view(), + raft::mul_op{}, + raft::make_const_mdspan(min_cluster_distance.view()), + sample_weight.value()); + } + auto device_cost = raft::make_device_scalar(handle, DataT(0)); cuvs::cluster::kmeans::cluster_cost( diff --git a/cpp/src/cluster/kmeans_balanced.cuh b/cpp/src/cluster/kmeans_balanced.cuh index d3a85b1d94..0c0df03397 100644 --- a/cpp/src/cluster/kmeans_balanced.cuh +++ b/cpp/src/cluster/kmeans_balanced.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2022-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ @@ -61,13 +61,16 @@ namespace cuvs::cluster::kmeans_balanced { * @param[out] centroids The generated centroids [dim = n_clusters x n_features] * @param[in] mapping_op (optional) Functor to convert from the input datatype to the arithmetic * datatype. If DataT == MathT, this must be the identity. + * @param[out] inertia (optional) Sum of squared distances of samples to their + * closest cluster center. */ template void fit(const raft::resources& handle, cuvs::cluster::kmeans::balanced_params const& params, raft::device_matrix_view X, raft::device_matrix_view centroids, - MappingOpT mapping_op = raft::identity_op()) + MappingOpT mapping_op = raft::identity_op(), + std::optional> inertia = std::nullopt) { RAFT_EXPECTS(X.extent(1) == centroids.extent(1), "Number of features in dataset and centroids are different"); @@ -78,6 +81,8 @@ void fit(const raft::resources& handle, "The number of centroids must be strictly positive and cannot exceed the number of " "points in the training dataset."); + MathT* inertia_ptr = inertia.has_value() ? inertia.value().data_handle() : nullptr; + cuvs::cluster::kmeans::detail::build_hierarchical(handle, params, X.extent(1), @@ -85,7 +90,8 @@ void fit(const raft::resources& handle, X.extent(0), centroids.data_handle(), centroids.extent(0), - mapping_op); + mapping_op, + inertia_ptr); } /** diff --git a/cpp/src/cluster/kmeans_balanced_fit_float.cu b/cpp/src/cluster/kmeans_balanced_fit_float.cu index 37a6bb127c..f3ef94b7be 100644 --- a/cpp/src/cluster/kmeans_balanced_fit_float.cu +++ b/cpp/src/cluster/kmeans_balanced_fit_float.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2022-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ @@ -14,9 +14,10 @@ namespace cuvs::cluster::kmeans { void fit(const raft::resources& handle, cuvs::cluster::kmeans::balanced_params const& params, raft::device_matrix_view X, - raft::device_matrix_view centroids) + raft::device_matrix_view centroids, + std::optional> inertia) { cuvs::cluster::kmeans_balanced::fit( - handle, params, X, centroids, cuvs::spatial::knn::detail::utils::mapping{}); + handle, params, X, centroids, cuvs::spatial::knn::detail::utils::mapping{}, inertia); } } // namespace cuvs::cluster::kmeans diff --git a/cpp/src/cluster/kmeans_balanced_fit_half.cu b/cpp/src/cluster/kmeans_balanced_fit_half.cu index e554930293..7272e6087a 100644 --- a/cpp/src/cluster/kmeans_balanced_fit_half.cu +++ b/cpp/src/cluster/kmeans_balanced_fit_half.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ @@ -14,9 +14,10 @@ namespace cuvs::cluster::kmeans { void fit(const raft::resources& handle, cuvs::cluster::kmeans::balanced_params const& params, raft::device_matrix_view X, - raft::device_matrix_view centroids) + raft::device_matrix_view centroids, + std::optional> inertia) { cuvs::cluster::kmeans_balanced::fit( - handle, params, X, centroids, cuvs::spatial::knn::detail::utils::mapping{}); + handle, params, X, centroids, cuvs::spatial::knn::detail::utils::mapping{}, inertia); } } // namespace cuvs::cluster::kmeans diff --git a/cpp/src/cluster/kmeans_balanced_fit_int8.cu b/cpp/src/cluster/kmeans_balanced_fit_int8.cu index a46f7bcc94..3615c4675b 100644 --- a/cpp/src/cluster/kmeans_balanced_fit_int8.cu +++ b/cpp/src/cluster/kmeans_balanced_fit_int8.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2022-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ @@ -14,9 +14,10 @@ namespace cuvs::cluster::kmeans { void fit(const raft::resources& handle, cuvs::cluster::kmeans::balanced_params const& params, raft::device_matrix_view X, - raft::device_matrix_view centroids) + raft::device_matrix_view centroids, + std::optional> inertia) { cuvs::cluster::kmeans_balanced::fit( - handle, params, X, centroids, cuvs::spatial::knn::detail::utils::mapping{}); + handle, params, X, centroids, cuvs::spatial::knn::detail::utils::mapping{}, inertia); } } // namespace cuvs::cluster::kmeans diff --git a/cpp/src/cluster/kmeans_balanced_fit_uint8.cu b/cpp/src/cluster/kmeans_balanced_fit_uint8.cu index 8395dd107f..2a7211e48e 100644 --- a/cpp/src/cluster/kmeans_balanced_fit_uint8.cu +++ b/cpp/src/cluster/kmeans_balanced_fit_uint8.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ @@ -14,9 +14,10 @@ namespace cuvs::cluster::kmeans { void fit(const raft::resources& handle, cuvs::cluster::kmeans::balanced_params const& params, raft::device_matrix_view X, - raft::device_matrix_view centroids) + raft::device_matrix_view centroids, + std::optional> inertia) { cuvs::cluster::kmeans_balanced::fit( - handle, params, X, centroids, cuvs::spatial::knn::detail::utils::mapping{}); + handle, params, X, centroids, cuvs::spatial::knn::detail::utils::mapping{}, inertia); } } // namespace cuvs::cluster::kmeans diff --git a/cpp/src/cluster/kmeans_cluster_cost.cu b/cpp/src/cluster/kmeans_cluster_cost.cu index a806cba4a4..0cdc182fb9 100644 --- a/cpp/src/cluster/kmeans_cluster_cost.cu +++ b/cpp/src/cluster/kmeans_cluster_cost.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ @@ -11,16 +11,36 @@ namespace cuvs::cluster::kmeans { void cluster_cost(const raft::resources& handle, raft::device_matrix_view X, raft::device_matrix_view centroids, - raft::host_scalar_view cost) + raft::host_scalar_view cost, + std::optional> sample_weight) { - cuvs::cluster::kmeans::cluster_cost(handle, X, centroids, cost); + cuvs::cluster::kmeans::cluster_cost(handle, X, centroids, cost, sample_weight); } void cluster_cost(const raft::resources& handle, raft::device_matrix_view X, raft::device_matrix_view centroids, - raft::host_scalar_view cost) + raft::host_scalar_view cost, + std::optional> sample_weight) { - cuvs::cluster::kmeans::cluster_cost(handle, X, centroids, cost); + cuvs::cluster::kmeans::cluster_cost(handle, X, centroids, cost, sample_weight); +} + +void cluster_cost(const raft::resources& handle, + raft::device_matrix_view X, + raft::device_matrix_view centroids, + raft::host_scalar_view cost, + std::optional> sample_weight) +{ + cuvs::cluster::kmeans::cluster_cost(handle, X, centroids, cost, sample_weight); +} + +void cluster_cost(const raft::resources& handle, + raft::device_matrix_view X, + raft::device_matrix_view centroids, + raft::host_scalar_view cost, + std::optional> sample_weight) +{ + cuvs::cluster::kmeans::cluster_cost(handle, X, centroids, cost, sample_weight); } } // namespace cuvs::cluster::kmeans