[FEA] membership_vector for HDBSCAN

cpp/include/cuml/cluster/hdbscan.hpp

@@ -460,6 +460,17 @@ 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

@@ -54,6 +54,7 @@ __global__ void merge_height_kernel(value_t* heights,
 
     if (took_left_parent && took_right_parent) {
       heights[idx] = lambdas[index_into_children[last_cluster]];
+      // printf("%d %d %d %f\n", row, col, last_cluster, heights[idx]);
     }
 
     else {
@@ -62,6 +63,49 @@ __global__ void merge_height_kernel(value_t* heights,
   }
 }
 
+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,
+                                    value_idx* selected_clusters)
+{
+  value_idx idx = blockDim.x * blockIdx.x + threadIdx.x;
+  if (idx < value_idx(n_prediction_points * n_selected_clusters)) {
+    value_idx row           = idx / n_selected_clusters;
+    value_idx col           = idx % n_selected_clusters;
+    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];
+    }
+  }
+}
+
 template <typename value_idx, typename value_t>
 __global__ void prob_in_some_cluster_kernel(value_t* heights,
                                             value_t* height_argmax,
@@ -84,6 +128,30 @@ __global__ void prob_in_some_cluster_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* prediction_lambdas,
+                                            value_t* deaths,
+                                            value_idx* index_into_children,
+                                            value_idx* min_mr_indices,
+                                            value_idx* selected_clusters,
+                                            value_t* lambdas,
+                                            value_t* prob_in_some_cluster,
+                                            value_idx n_selected_clusters,
+                                            value_idx n_leaves,
+                                            size_t n_prediction_points)
+{
+  value_idx idx = blockDim.x * blockIdx.x + threadIdx.x;
+  if (idx < (value_idx)n_prediction_points) {
+    value_t max_lambda = max(prediction_lambdas[idx], 
+                             deaths[selected_clusters[(int)height_argmax[idx]] - n_leaves]) + 1e-8;
+    prob_in_some_cluster[idx] =
+      heights[idx * n_selected_clusters + (int)height_argmax[idx]] / max_lambda;
+    return;
+  }
+}
+
 };  // namespace Predict
 };  // namespace detail
 };  // namespace HDBSCAN

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,64 @@ 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,
@@ -189,50 +246,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,