UMAP with new Spectral Embedding initialization

cpp/include/cuml/manifold/spectral_embedding.hpp

@@ -14,6 +14,7 @@
  * limitations under the License.
  */
 
+#include <raft/core/device_coo_matrix.hpp>
 #include <raft/core/device_mdspan.hpp>
 #include <raft/core/resources.hpp>
 
@@ -52,4 +53,9 @@ void transform(raft::resources const& handle,
                raft::device_matrix_view<float, int, raft::row_major> dataset,
                raft::device_matrix_view<float, int, raft::col_major> embedding);
 
+void transform(raft::resources const& handle,
+               ML::SpectralEmbedding::params config,
+               raft::device_coo_matrix_view<float, int, int, int> connectivity_graph,
+               raft::device_matrix_view<float, int, raft::col_major> embedding);
+
 }  // namespace ML::SpectralEmbedding

cpp/src/spectral/spectral_embedding.cu

@@ -44,4 +44,13 @@ void transform(raft::resources const& handle,
   cuvs::preprocessing::spectral_embedding::transform(handle, to_cuvs(config), dataset, embedding);
 }
 
+void transform(raft::resources const& handle,
+               ML::SpectralEmbedding::params config,
+               raft::device_coo_matrix_view<float, int, int, int> connectivity_graph,
+               raft::device_matrix_view<float, int, raft::col_major> embedding)
+{
+  cuvs::preprocessing::spectral_embedding::transform(
+    handle, to_cuvs(config), connectivity_graph, embedding);
+}
+
 }  // namespace ML::SpectralEmbedding

cpp/src/umap/init_embed/spectral_algo.cuh

@@ -16,9 +16,10 @@
 
 #pragma once
 
-#include <cuml/cluster/spectral.hpp>
+#include <cuml/manifold/spectral_embedding.hpp>
 #include <cuml/manifold/umapparams.h>
 
+#include <raft/core/device_coo_matrix.hpp>
 #include <raft/core/handle.hpp>
 #include <raft/linalg/add.cuh>
 #include <raft/linalg/transpose.cuh>
@@ -31,8 +32,6 @@
 
 #include <stdint.h>
 
-#include <iostream>
-
 namespace UMAPAlgo {
 
 namespace InitEmbed {
@@ -57,36 +56,44 @@ void launcher(const raft::handle_t& handle,
   ASSERT(n > static_cast<nnz_t>(params->n_components),
          "Spectral layout requires n_samples > n_components");
 
-  rmm::device_uvector<T> tmp_storage(n * params->n_components, stream);
+  auto tmp_embedding =
+    raft::make_device_matrix<float, int, raft::col_major>(handle, n, params->n_components);
+
+  auto connectivity_graph_view = raft::make_device_coo_matrix_view<float, int, int, int>(
+    coo->vals(),
+    raft::make_device_coordinate_structure_view<int, int, int>(
+      coo->rows(), coo->cols(), n, n, coo->nnz));
+
+  ML::SpectralEmbedding::params spectral_params;
+  spectral_params.n_neighbors    = params->n_neighbors;
+  spectral_params.norm_laplacian = true;
+  spectral_params.drop_first     = true;
+  spectral_params.seed           = params->random_state;
+  spectral_params.n_components =
+    spectral_params.drop_first ? params->n_components + 1 : params->n_components;
 
   uint64_t seed = params->random_state;
 
-  Spectral::fit_embedding(handle,
-                          coo->rows(),
-                          coo->cols(),
-                          coo->vals(),
-                          coo->nnz,
-                          n,
-                          params->n_components,
-                          tmp_storage.data(),
-                          seed);
+  ML::SpectralEmbedding::transform(
+    handle, spectral_params, connectivity_graph_view, tmp_embedding.view());
 
-  raft::linalg::transpose(handle, tmp_storage.data(), embedding, n, params->n_components, stream);
+  raft::linalg::transpose(
+    handle, tmp_embedding.data_handle(), embedding, n, params->n_components, stream);
 
   raft::linalg::unaryOp<T>(
-    tmp_storage.data(),
-    tmp_storage.data(),
+    tmp_embedding.data_handle(),
+    tmp_embedding.data_handle(),
     n * params->n_components,
     [=] __device__(T input) { return fabsf(input); },
     stream);
 
-  thrust::device_ptr<T> d_ptr = thrust::device_pointer_cast(tmp_storage.data());
+  thrust::device_ptr<T> d_ptr = thrust::device_pointer_cast(tmp_embedding.data_handle());
   T max =
     *(thrust::max_element(thrust::cuda::par.on(stream), d_ptr, d_ptr + (n * params->n_components)));
 
   // Reuse tmp_storage to add random noise
   raft::random::Rng r(seed);
-  r.normal(tmp_storage.data(), n * params->n_components, 0.0f, 0.0001f, stream);
+  r.normal(tmp_embedding.data_handle(), n * params->n_components, 0.0f, 0.0001f, stream);
 
   raft::linalg::unaryOp<T>(
     embedding,
@@ -95,7 +102,8 @@ void launcher(const raft::handle_t& handle,
     [=] __device__(T input) { return (10.0f / max) * input; },
     stream);
 
-  raft::linalg::add(embedding, embedding, tmp_storage.data(), n * params->n_components, stream);
+  raft::linalg::add(
+    embedding, embedding, tmp_embedding.data_handle(), n * params->n_components, stream);
 
   RAFT_CUDA_TRY(cudaPeekAtLastError());
 }