[FEA] membership_vector for HDBSCAN

cpp/include/cuml/cluster/hdbscan.hpp

@@ -460,6 +460,16 @@ void compute_all_points_membership_vectors(
   raft::distance::DistanceType metric,
   float* membership_vec);
 
+void compute_membership_vector(const raft::handle_t& handle,
+                               HDBSCAN::Common::CondensedHierarchy<int, float>& condensed_tree,
+                               HDBSCAN::Common::PredictionData<int, float>& prediction_data,
+                               const float* X,
+                               const float* points_to_predict,
+                               size_t n_prediction_points,
+                               int min_samples,
+                               raft::distance::DistanceType metric,
+                               float* membership_vec);
+
 void out_of_sample_predict(const raft::handle_t& handle,
                            HDBSCAN::Common::CondensedHierarchy<int, float>& condensed_tree,
                            HDBSCAN::Common::PredictionData<int, float>& prediction_data,

cpp/src/hdbscan/detail/kernels/soft_clustering.cuh

@@ -13,9 +13,10 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-
 #pragma once
 
+#include <raft/util/fast_int_div.cuh>
+
 namespace ML {
 namespace HDBSCAN {
 namespace detail {
@@ -28,6 +29,7 @@ __global__ void merge_height_kernel(value_t* heights,
                                     value_idx* parents,
                                     size_t m,
                                     value_idx n_selected_clusters,
+                                    raft::util::FastIntDiv n,
                                     value_idx* selected_clusters)
 {
   value_idx idx = blockDim.x * blockIdx.x + threadIdx.x;
@@ -62,25 +64,47 @@ __global__ void merge_height_kernel(value_t* heights,
   }
 }
 
-template <typename value_idx, typename value_t>
-__global__ void prob_in_some_cluster_kernel(value_t* heights,
-                                            value_t* height_argmax,
-                                            value_t* deaths,
-                                            value_idx* index_into_children,
-                                            value_idx* selected_clusters,
-                                            value_t* lambdas,
-                                            value_t* prob_in_some_cluster,
-                                            value_idx n_selected_clusters,
-                                            value_idx n_leaves,
-                                            size_t m)
+template <typename value_idx, typename value_t, int tpb = 256>
+__global__ void merge_height_kernel(value_t* heights,
+                                    value_t* lambdas,
+                                    value_t* prediction_lambdas,
+                                    value_idx* min_mr_indices,
+                                    value_idx* index_into_children,
+                                    value_idx* parents,
+                                    size_t n_prediction_points,
+                                    value_idx n_selected_clusters,
+                                    raft::util::FastIntDiv n,
+                                    value_idx* selected_clusters)
 {
   value_idx idx = blockDim.x * blockIdx.x + threadIdx.x;
-  if (idx < (value_idx)m) {
-    value_t max_lambda = max(lambdas[index_into_children[idx]],
-                             deaths[selected_clusters[(int)height_argmax[idx]] - n_leaves]);
-    prob_in_some_cluster[idx] =
-      heights[idx * n_selected_clusters + (int)height_argmax[idx]] / max_lambda;
-    return;
+  if (idx < value_idx(n_prediction_points * n_selected_clusters)) {
+    value_idx row           = idx / n;
+    value_idx col           = idx % n;
+    value_idx right_cluster = selected_clusters[col];
+    value_idx left_cluster  = parents[index_into_children[min_mr_indices[row]]];
+    bool took_right_parent  = false;
+    bool took_left_parent   = false;
+    value_idx last_cluster;
+
+    while (left_cluster != right_cluster) {
+      if (left_cluster > right_cluster) {
+        took_left_parent = true;
+        last_cluster     = left_cluster;
+        left_cluster     = parents[index_into_children[left_cluster]];
+      } else {
+        took_right_parent = true;
+        last_cluster      = right_cluster;
+        right_cluster     = parents[index_into_children[right_cluster]];
+      }
+    }
+
+    if (took_left_parent && took_right_parent) {
+      heights[idx] = lambdas[index_into_children[last_cluster]];
+    }
+
+    else {
+      heights[idx] = prediction_lambdas[row];
+    }
   }
 }
 

cpp/src/hdbscan/detail/predict.cuh

@@ -36,7 +36,6 @@ namespace Predict {
 Find the nearest mutual reachability neighbor of a point, and  compute
 the associated lambda value for the point, given the mutual reachability
 distance to a nearest neighbor.
- *
  * @tparam value_idx
  * @tparam value_t
  * @tparam tpb
@@ -149,6 +148,62 @@ void _find_cluster_and_probability(const raft::handle_t& handle,
                                                            out_labels,
                                                            out_probabilities);
 }
+
+// Build the mutual reachability graph and obtain the nearest neighbors for the prediction points.
+// The KNN and core distances of prediction points are computed here.
+template <typename value_idx, typename value_t, int tpb = 256>
+void _compute_knn_and_nearest_neighbor(const raft::handle_t& handle,
+                                       Common::PredictionData<value_idx, value_t>& prediction_data,
+                                       const value_t* X,
+                                       const value_t* points_to_predict,
+                                       int min_samples,
+                                       size_t n_prediction_points,
+                                       value_idx* min_mr_inds,
+                                       value_t* prediction_lambdas,
+                                       raft::distance::DistanceType metric)
+{
+  auto stream               = handle.get_stream();
+  size_t m                  = prediction_data.n_rows;
+  size_t n                  = prediction_data.n_cols;
+  value_t* input_core_dists = prediction_data.get_core_dists();
+  int neighborhood          = (min_samples - 1) * 2;
+
+  rmm::device_uvector<value_idx> inds(neighborhood * n_prediction_points, stream);
+  rmm::device_uvector<value_t> dists(neighborhood * n_prediction_points, stream);
+  rmm::device_uvector<value_t> prediction_core_dists(n_prediction_points, stream);
+
+  // perform knn
+  Reachability::compute_knn(handle,
+                            X,
+                            inds.data(),
+                            dists.data(),
+                            m,
+                            n,
+                            points_to_predict,
+                            n_prediction_points,
+                            neighborhood,
+                            metric);
+
+  // Slice core distances (distances to kth nearest neighbor). The index of the neighbor is
+  // consistent with Scikit-learn Contrib
+  Reachability::core_distances<value_idx>(dists.data(),
+                                          min_samples,
+                                          neighborhood,
+                                          n_prediction_points,
+                                          prediction_core_dists.data(),
+                                          stream);
+
+  _find_neighbor_and_lambda(handle,
+                            input_core_dists,
+                            prediction_core_dists.data(),
+                            dists.data(),
+                            inds.data(),
+                            n_prediction_points,
+                            neighborhood,
+                            min_mr_inds,
+                            prediction_lambdas);
+}
+
 /**
  * Predict the cluster label and the probability of the label for new points.
  * The returned labels are those of the original clustering,
@@ -165,7 +220,7 @@ void _find_cluster_and_probability(const raft::handle_t& handle,
  * @param[in] labels converted monotonic labels of the input data points
  * @param[in] points_to_predict input prediction points (size n_prediction_points * n)
  * @param[in] n_prediction_points number of prediction points
- * @param[in] metric distance metric to use
+ * @param[in] metric distance metric
  * @param[in] min_samples neighborhood size during training (includes self-loop)
  * @param[out] out_labels output cluster labels
  * @param[out] out_probabilities output probabilities
@@ -189,50 +244,18 @@ void approximate_predict(const raft::handle_t& handle,
   auto stream      = handle.get_stream();
   auto exec_policy = handle.get_thrust_policy();
 
-  size_t m                  = prediction_data.n_rows;
-  size_t n                  = prediction_data.n_cols;
-  value_t* input_core_dists = prediction_data.get_core_dists();
-
-  // this is the neighborhood of prediction points for which MR distances are computed
-  int neighborhood = (min_samples - 1) * 2;
-
-  rmm::device_uvector<value_idx> inds(neighborhood * n_prediction_points, stream);
-  rmm::device_uvector<value_t> dists(neighborhood * n_prediction_points, stream);
-  rmm::device_uvector<value_t> prediction_core_dists(n_prediction_points, stream);
-
-  // perform knn
-  Reachability::compute_knn(handle,
-                            X,
-                            inds.data(),
-                            dists.data(),
-                            m,
-                            n,
-                            points_to_predict,
-                            n_prediction_points,
-                            neighborhood,
-                            metric);
-
-  // Slice core distances (distances to kth nearest neighbor). The index of the neighbor is
-  // consistent with Scikit-learn Contrib
-  Reachability::core_distances<value_idx>(dists.data(),
-                                          min_samples,
-                                          neighborhood,
-                                          n_prediction_points,
-                                          prediction_core_dists.data(),
-                                          stream);
-
   // Obtain lambdas for each prediction point using the closest point in mutual reachability space
   rmm::device_uvector<value_t> prediction_lambdas(n_prediction_points, stream);
   rmm::device_uvector<value_idx> min_mr_inds(n_prediction_points, stream);
-  _find_neighbor_and_lambda(handle,
-                            input_core_dists,
-                            prediction_core_dists.data(),
-                            dists.data(),
-                            inds.data(),
-                            n_prediction_points,
-                            neighborhood,
-                            min_mr_inds.data(),
-                            prediction_lambdas.data());
+  _compute_knn_and_nearest_neighbor(handle,
+                                    prediction_data,
+                                    X,
+                                    points_to_predict,
+                                    min_samples,
+                                    n_prediction_points,
+                                    min_mr_inds.data(),
+                                    prediction_lambdas.data(),
+                                    metric);
 
   // Using the nearest neighbor indices, find the assigned cluster label and probability
   _find_cluster_and_probability(handle,