add sample_weight parameter to dbscan.fit

cpp/bench/sg/dbscan.cu

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019-2022, NVIDIA CORPORATION.
+ * Copyright (c) 2019-2023, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -59,6 +59,7 @@ class Dbscan : public BlobsFixture<D, int> {
                       raft::distance::L2SqrtUnexpanded,
                       this->data.y.data(),
                       this->core_sample_indices,
+                      nullptr,
                       dParams.max_bytes_per_batch);
       state.SetItemsProcessed(this->params.nrows * this->params.ncols);
     });

cpp/examples/dbscan/dbscan_example.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019-2022, NVIDIA CORPORATION.
+ * Copyright (c) 2019-2023, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -205,6 +205,7 @@ int main(int argc, char* argv[])
                   raft::distance::L2SqrtUnexpanded,
                   d_labels,
                   nullptr,
+                  nullptr,
                   max_bytes_per_batch,
                   false);
   CUDA_RT_CALL(cudaMemcpyAsync(

cpp/include/cuml/cluster/dbscan.hpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2018-2022, NVIDIA CORPORATION.
+ * Copyright (c) 2018-2023, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -43,6 +43,10 @@ namespace Dbscan {
  *             indices of each core point. If the number of core points is less
  *             than n_rows, the right will be padded with -1. Setting this to
  *             NULL will prevent calculating the core sample indices
+ * @param[in] sample_weight (size n_rows) input array containing the
+ *            weight of each sample to be taken instead of a plain sum to
+ *            fulfill the min_pts criteria for core points.
+ *            NULL will default to weights of 1 for all samples
  * @param[in] max_bytes_per_batch the maximum number of megabytes to be used for
  *            each batch of the pairwise distance calculation. This enables the
  *            trade off between memory usage and algorithm execution time.
@@ -60,6 +64,7 @@ void fit(const raft::handle_t& handle,
          raft::distance::DistanceType metric,
          int* labels,
          int* core_sample_indices   = nullptr,
+         float* sample_weight       = nullptr,
          size_t max_bytes_per_batch = 0,
          int verbosity              = CUML_LEVEL_INFO,
          bool opg                   = false);
@@ -72,6 +77,7 @@ void fit(const raft::handle_t& handle,
          raft::distance::DistanceType metric,
          int* labels,
          int* core_sample_indices   = nullptr,
+         double* sample_weight      = nullptr,
          size_t max_bytes_per_batch = 0,
          int verbosity              = CUML_LEVEL_INFO,
          bool opg                   = false);
@@ -85,6 +91,7 @@ void fit(const raft::handle_t& handle,
          raft::distance::DistanceType metric,
          int64_t* labels,
          int64_t* core_sample_indices = nullptr,
+         float* sample_weight         = nullptr,
          size_t max_bytes_per_batch   = 0,
          int verbosity                = CUML_LEVEL_INFO,
          bool opg                     = false);
@@ -97,6 +104,7 @@ void fit(const raft::handle_t& handle,
          raft::distance::DistanceType metric,
          int64_t* labels,
          int64_t* core_sample_indices = nullptr,
+         double* sample_weight        = nullptr,
          size_t max_bytes_per_batch   = 0,
          int verbosity                = CUML_LEVEL_INFO,
          bool opg                     = false);

cpp/src/dbscan/corepoints/compute.cuh

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2021-2022, NVIDIA CORPORATION.
+ * Copyright (c) 2021-2023, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -28,16 +28,16 @@ namespace CorePoints {
 /**
  * Compute the core points from the vertex degrees and min_pts criterion
  * @param[in]  handle          cuML handle
- * @param[in]  vd              Vertex degrees
+ * @param[in]  vd              Vertex degrees (optionally weighted)
  * @param[out] mask            Boolean core point mask
  * @param[in]  min_pts         Core point criterion
  * @param[in]  start_vertex_id First point of the batch
  * @param[in]  batch_size      Batch size
  * @param[in]  stream          CUDA stream
  */
-template <typename Index_ = int>
+template <typename Values_ = int, typename Index_ = int>
 void compute(const raft::handle_t& handle,
-             const Index_* vd,
+             const Values_* vd,
              bool* mask,
              Index_ min_pts,
              Index_ start_vertex_id,
@@ -47,7 +47,7 @@ void compute(const raft::handle_t& handle,
   auto counting = thrust::make_counting_iterator<Index_>(0);
   thrust::for_each(
     handle.get_thrust_policy(), counting, counting + batch_size, [=] __device__(Index_ idx) {
-      mask[idx + start_vertex_id] = vd[idx] >= min_pts;
+      mask[idx + start_vertex_id] = (Index_)vd[idx] >= min_pts;
     });
 }
 

cpp/src/dbscan/dbscan.cu

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2018-2022, NVIDIA CORPORATION.
+ * Copyright (c) 2018-2023, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -31,6 +31,7 @@ void fit(const raft::handle_t& handle,
          raft::distance::DistanceType metric,
          int* labels,
          int* core_sample_indices,
+         float* sample_weight,
          size_t max_bytes_per_batch,
          int verbosity,
          bool opg)
@@ -45,6 +46,7 @@ void fit(const raft::handle_t& handle,
                                     metric,
                                     labels,
                                     core_sample_indices,
+                                    sample_weight,
                                     max_bytes_per_batch,
                                     handle.get_stream(),
                                     verbosity);
@@ -58,6 +60,7 @@ void fit(const raft::handle_t& handle,
                                      metric,
                                      labels,
                                      core_sample_indices,
+                                     sample_weight,
                                      max_bytes_per_batch,
                                      handle.get_stream(),
                                      verbosity);
@@ -72,6 +75,7 @@ void fit(const raft::handle_t& handle,
          raft::distance::DistanceType metric,
          int* labels,
          int* core_sample_indices,
+         double* sample_weight,
          size_t max_bytes_per_batch,
          int verbosity,
          bool opg)
@@ -86,6 +90,7 @@ void fit(const raft::handle_t& handle,
                                      metric,
                                      labels,
                                      core_sample_indices,
+                                     sample_weight,
                                      max_bytes_per_batch,
                                      handle.get_stream(),
                                      verbosity);
@@ -99,6 +104,7 @@ void fit(const raft::handle_t& handle,
                                       metric,
                                       labels,
                                       core_sample_indices,
+                                      sample_weight,
                                       max_bytes_per_batch,
                                       handle.get_stream(),
                                       verbosity);
@@ -113,6 +119,7 @@ void fit(const raft::handle_t& handle,
          raft::distance::DistanceType metric,
          int64_t* labels,
          int64_t* core_sample_indices,
+         float* sample_weight,
          size_t max_bytes_per_batch,
          int verbosity,
          bool opg)
@@ -127,6 +134,7 @@ void fit(const raft::handle_t& handle,
                                         metric,
                                         labels,
                                         core_sample_indices,
+                                        sample_weight,
                                         max_bytes_per_batch,
                                         handle.get_stream(),
                                         verbosity);
@@ -140,6 +148,7 @@ void fit(const raft::handle_t& handle,
                                          metric,
                                          labels,
                                          core_sample_indices,
+                                         sample_weight,
                                          max_bytes_per_batch,
                                          handle.get_stream(),
                                          verbosity);
@@ -154,6 +163,7 @@ void fit(const raft::handle_t& handle,
          raft::distance::DistanceType metric,
          int64_t* labels,
          int64_t* core_sample_indices,
+         double* sample_weight,
          size_t max_bytes_per_batch,
          int verbosity,
          bool opg)
@@ -168,6 +178,7 @@ void fit(const raft::handle_t& handle,
                                          metric,
                                          labels,
                                          core_sample_indices,
+                                         sample_weight,
                                          max_bytes_per_batch,
                                          handle.get_stream(),
                                          verbosity);
@@ -181,6 +192,7 @@ void fit(const raft::handle_t& handle,
                                           metric,
                                           labels,
                                           core_sample_indices,
+                                          sample_weight,
                                           max_bytes_per_batch,
                                           handle.get_stream(),
                                           verbosity);

cpp/src/dbscan/dbscan.cuh

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2018-2022, NVIDIA CORPORATION.
+ * Copyright (c) 2018-2023, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -104,6 +104,7 @@ void dbscanFitImpl(const raft::handle_t& handle,
                    raft::distance::DistanceType metric,
                    Index_* labels,
                    Index_* core_sample_indices,
+                   T* sample_weight,
                    size_t max_mbytes_per_batch,
                    cudaStream_t stream,
                    int verbosity)
@@ -177,6 +178,7 @@ void dbscanFitImpl(const raft::handle_t& handle,
                                                      min_pts,
                                                      labels,
                                                      core_sample_indices,
+                                                     sample_weight,
                                                      algo_vd,
                                                      algo_adj,
                                                      algo_ccl,
@@ -198,6 +200,7 @@ void dbscanFitImpl(const raft::handle_t& handle,
                               min_pts,
                               labels,
                               core_sample_indices,
+                              sample_weight,
                               algo_vd,
                               algo_adj,
                               algo_ccl,

cpp/src/dbscan/dbscan_api.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2018-2022, NVIDIA CORPORATION.
+ * Copyright (c) 2018-2023, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -47,6 +47,7 @@ cumlError_t cumlSpDbscanFit(cumlHandle_t handle,
                       raft::distance::L2SqrtUnexpanded,
                       labels,
                       core_sample_indices,
+                      NULL,
                       max_bytes_per_batch,
                       verbosity);
     }
@@ -88,6 +89,7 @@ cumlError_t cumlDpDbscanFit(cumlHandle_t handle,
                       raft::distance::L2SqrtUnexpanded,
                       labels,
                       core_sample_indices,
+                      NULL,
                       max_bytes_per_batch,
                       verbosity);
     }

cpp/src/dbscan/runner.cuh

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2018-2022, NVIDIA CORPORATION.
+ * Copyright (c) 2018-2023, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -110,6 +110,7 @@ std::size_t run(const raft::handle_t& handle,
                 Index_ min_pts,
                 Index_* labels,
                 Index_* core_indices,
+                const Type_f* sample_weight,
                 int algo_vd,
                 int algo_adj,
                 int algo_ccl,
@@ -146,6 +147,8 @@ std::size_t run(const raft::handle_t& handle,
   std::size_t ex_scan_size  = raft::alignTo<std::size_t>(sizeof(Index_) * batch_size, align);
   std::size_t row_cnt_size  = raft::alignTo<std::size_t>(sizeof(Index_) * batch_size, align);
   std::size_t labels_size   = raft::alignTo<std::size_t>(sizeof(Index_) * N, align);
+  std::size_t wght_sum_size =
+    sample_weight != nullptr ? raft::alignTo<std::size_t>(sizeof(Type_f) * batch_size, align) : 0;
 
   Index_ MAX_LABEL = std::numeric_limits<Index_>::max();
 
@@ -157,8 +160,8 @@ std::size_t run(const raft::handle_t& handle,
          (unsigned long)batch_size);
 
   if (workspace == NULL) {
-    auto size =
-      adj_size + core_pts_size + m_size + vd_size + ex_scan_size + row_cnt_size + 2 * labels_size;
+    auto size = adj_size + core_pts_size + m_size + vd_size + ex_scan_size + row_cnt_size +
+                2 * labels_size + wght_sum_size;
     return size;
   }
 
@@ -183,6 +186,11 @@ std::size_t run(const raft::handle_t& handle,
   temp += labels_size;
   Index_* work_buffer = (Index_*)temp;
   temp += labels_size;
+  Type_f* wght_sum = nullptr;
+  if (sample_weight != nullptr) {
+    wght_sum = (Type_f*)temp;
+    temp += wght_sum_size;
+  }
 
   // Compute the mask
   // 1. Compute the part owned by this worker (reversed order of batches to
@@ -196,13 +204,31 @@ std::size_t run(const raft::handle_t& handle,
 
     CUML_LOG_DEBUG("--> Computing vertex degrees");
     raft::common::nvtx::push_range("Trace::Dbscan::VertexDeg");
-    VertexDeg::run<Type_f, Index_>(
-      handle, adj, vd, x, eps, N, D, algo_vd, start_vertex_id, n_points, stream, metric);
+    VertexDeg::run<Type_f, Index_>(handle,
+                                   adj,
+                                   vd,
+                                   wght_sum,
+                                   x,
+                                   sample_weight,
+                                   eps,
+                                   N,
+                                   D,
+                                   algo_vd,
+                                   start_vertex_id,
+                                   n_points,
+                                   stream,
+                                   metric);
     raft::common::nvtx::pop_range();
 
     CUML_LOG_DEBUG("--> Computing core point mask");
     raft::common::nvtx::push_range("Trace::Dbscan::CorePoints");
-    CorePoints::compute<Index_>(handle, vd, core_pts, min_pts, start_vertex_id, n_points, stream);
+    if (wght_sum != nullptr) {
+      CorePoints::compute<Type_f, Index_>(
+        handle, wght_sum, core_pts, min_pts, start_vertex_id, n_points, stream);
+    } else {
+      CorePoints::compute<Index_, Index_>(
+        handle, vd, core_pts, min_pts, start_vertex_id, n_points, stream);
+    }
     raft::common::nvtx::pop_range();
   }
   // 2. Exchange with the other workers
@@ -224,8 +250,20 @@ std::size_t run(const raft::handle_t& handle,
     if (i > 0) {
       CUML_LOG_DEBUG("--> Computing vertex degrees");
       raft::common::nvtx::push_range("Trace::Dbscan::VertexDeg");
-      VertexDeg::run<Type_f, Index_>(
-        handle, adj, vd, x, eps, N, D, algo_vd, start_vertex_id, n_points, stream, metric);
+      VertexDeg::run<Type_f, Index_>(handle,
+                                     adj,
+                                     vd,
+                                     nullptr,
+                                     x,
+                                     nullptr,
+                                     eps,
+                                     N,
+                                     D,
+                                     algo_vd,
+                                     start_vertex_id,
+                                     n_points,
+                                     stream,
+                                     metric);
       raft::common::nvtx::pop_range();
     }
     raft::update_host(&curradjlen, vd + n_points, 1, stream);