Lanczos Solver which=SA,SM,LA,LM argument (#2628)

Resolves #2624
Resolves #2483

Authors:
  - Anupam (https://github.com/aamijar)
  - Corey J. Nolet (https://github.com/cjnolet)

Approvers:
  - Micka (https://github.com/lowener)
  - Divye Gala (https://github.com/divyegala)

URL: #2628

Author: aamijar

cpp/include/raft/sparse/solver/detail/lanczos.cuh

@@ -50,6 +50,7 @@
 #include <raft/linalg/transpose.cuh>
 #include <raft/linalg/unary_op.cuh>
 #include <raft/matrix/diagonal.cuh>
+#include <raft/matrix/gather.cuh>
 #include <raft/matrix/matrix.cuh>
 #include <raft/matrix/slice.cuh>
 #include <raft/matrix/triangular.cuh>
@@ -63,6 +64,7 @@
 #include <raft/util/cudart_utils.hpp>
 
 #include <cuda.h>
+#include <thrust/sort.h>
 
 #include <cublasLt.h>
 #include <curand.h>
@@ -1507,10 +1509,15 @@ void lanczos_solve_ritz(
   raft::device_matrix_view<ValueTypeT, uint32_t, raft::row_major> beta,
   std::optional<raft::device_vector_view<ValueTypeT, uint32_t>> beta_k,
   IndexTypeT k,
-  int which,
+  LANCZOS_WHICH which,
   int ncv,
   raft::device_matrix_view<ValueTypeT, uint32_t, raft::col_major> eigenvectors,
-  raft::device_vector_view<ValueTypeT> eigenvalues)
+  raft::device_vector_view<ValueTypeT> eigenvalues,
+  raft::device_matrix_view<ValueTypeT, uint32_t, raft::col_major>& eigenvectors_k,
+  raft::device_vector_view<ValueTypeT, uint32_t>& eigenvalues_k,
+  raft::device_matrix_view<ValueTypeT, IndexTypeT, raft::col_major>& eigenvectors_k_slice,
+  raft::device_vector_view<ValueTypeT> sm_eigenvalues,
+  raft::device_matrix_view<ValueTypeT, uint32_t, raft::col_major> sm_eigenvectors)
 {
   auto stream = resource::get_cuda_stream(handle);
 
@@ -1543,6 +1550,75 @@ void lanczos_solve_ritz(
       triangular_matrix.data_handle(), ncv, ncv);
 
   raft::linalg::eig_dc(handle, triangular_matrix_view, eigenvectors, eigenvalues);
+
+  IndexTypeT nEigVecs = k;
+
+  auto indices          = raft::make_device_vector<int>(handle, ncv);
+  auto selected_indices = raft::make_device_vector<int>(handle, nEigVecs);
+
+  if (which == LANCZOS_WHICH::SA) {
+    eigenvectors_k = raft::make_device_matrix_view<ValueTypeT, uint32_t, raft::col_major>(
+      eigenvectors.data_handle(), ncv, nEigVecs);
+    eigenvalues_k =
+      raft::make_device_vector_view<ValueTypeT, uint32_t>(eigenvalues.data_handle(), nEigVecs);
+    eigenvectors_k_slice = raft::make_device_matrix_view<ValueTypeT, IndexTypeT, raft::col_major>(
+      eigenvectors.data_handle(), ncv, nEigVecs);
+  } else if (which == LANCZOS_WHICH::LA) {
+    eigenvectors_k = raft::make_device_matrix_view<ValueTypeT, uint32_t, raft::col_major>(
+      eigenvectors.data_handle() + (ncv - nEigVecs) * ncv, ncv, nEigVecs);
+    eigenvalues_k = raft::make_device_vector_view<ValueTypeT, uint32_t>(
+      eigenvalues.data_handle() + (ncv - nEigVecs), nEigVecs);
+    eigenvectors_k_slice = raft::make_device_matrix_view<ValueTypeT, IndexTypeT, raft::col_major>(
+      eigenvectors.data_handle() + (ncv - nEigVecs) * ncv, ncv, nEigVecs);
+  } else if (which == LANCZOS_WHICH::SM || which == LANCZOS_WHICH::LM) {
+    thrust::sequence(thrust::device, indices.data_handle(), indices.data_handle() + ncv, 0);
+
+    // Sort indices by absolute eigenvalues (magnitude) using a custom comparator
+    thrust::sort(thrust::device,
+                 indices.data_handle(),
+                 indices.data_handle() + ncv,
+                 [eigenvalues = eigenvalues.data_handle()] __device__(int a, int b) {
+                   return fabsf(eigenvalues[a]) < fabsf(eigenvalues[b]);
+                 });
+
+    if (which == LANCZOS_WHICH::SM) {
+      // Take the first nEigVecs indices (smallest magnitude)
+      raft::copy(selected_indices.data_handle(), indices.data_handle(), nEigVecs, stream);
+    } else if (which == LANCZOS_WHICH::LM) {
+      // Take the last nEigVecs indices (largest magnitude)
+      raft::copy(
+        selected_indices.data_handle(), indices.data_handle() + (ncv - nEigVecs), nEigVecs, stream);
+    }
+
+    // Re-sort these indices by algebraic value to maintain algebraic ordering
+    thrust::sort(thrust::device,
+                 selected_indices.data_handle(),
+                 selected_indices.data_handle() + nEigVecs,
+                 [eigenvalues = eigenvalues.data_handle()] __device__(int a, int b) {
+                   return eigenvalues[a] < eigenvalues[b];
+                 });
+    raft::matrix::gather(
+      handle,
+      raft::make_device_matrix_view<const ValueTypeT, uint32_t, raft::row_major>(
+        eigenvalues.data_handle(), ncv, 1),
+      raft::make_device_vector_view<const int, uint32_t>(selected_indices.data_handle(), nEigVecs),
+      raft::make_device_matrix_view<ValueTypeT, uint32_t, raft::row_major>(
+        sm_eigenvalues.data_handle(), nEigVecs, 1));
+    raft::matrix::gather(
+      handle,
+      raft::make_device_matrix_view<const ValueTypeT, uint32_t, raft::row_major>(
+        eigenvectors.data_handle(), ncv, ncv),
+      raft::make_device_vector_view<const int, uint32_t>(selected_indices.data_handle(), nEigVecs),
+      raft::make_device_matrix_view<ValueTypeT, uint32_t, raft::row_major>(
+        sm_eigenvectors.data_handle(), nEigVecs, ncv));
+
+    eigenvectors_k = raft::make_device_matrix_view<ValueTypeT, uint32_t, raft::col_major>(
+      sm_eigenvectors.data_handle(), ncv, nEigVecs);
+    eigenvalues_k =
+      raft::make_device_vector_view<ValueTypeT, uint32_t>(sm_eigenvalues.data_handle(), nEigVecs);
+    eigenvectors_k_slice = raft::make_device_matrix_view<ValueTypeT, IndexTypeT, raft::col_major>(
+      sm_eigenvectors.data_handle(), ncv, nEigVecs);
+  }
 }
 
 template <typename IndexTypeT, typename ValueTypeT>
@@ -1695,6 +1771,7 @@ auto lanczos_smallest(
   int maxIter,
   int restartIter,
   ValueTypeT tol,
+  LANCZOS_WHICH which,
   ValueTypeT* eigVals_dev,
   ValueTypeT* eigVecs_dev,
   ValueTypeT* v0,
@@ -1752,20 +1829,28 @@ auto lanczos_smallest(
     raft::make_device_matrix<ValueTypeT, uint32_t, raft::col_major>(handle, ncv, ncv);
   auto eigenvalues = raft::make_device_vector<ValueTypeT, uint32_t>(handle, ncv);
 
+  raft::device_matrix_view<ValueTypeT, uint32_t, raft::col_major> eigenvectors_k;
+  raft::device_vector_view<ValueTypeT, uint32_t> eigenvalues_k;
+  raft::device_matrix_view<ValueTypeT, IndexTypeT, raft::col_major> eigenvectors_k_slice;
+
+  auto sm_eigenvalues = raft::make_device_vector<ValueTypeT>(handle, nEigVecs);
+  auto sm_eigenvectors =
+    raft::make_device_matrix<ValueTypeT, uint32_t, raft::col_major>(handle, ncv, nEigVecs);
+
   lanczos_solve_ritz<IndexTypeT, ValueTypeT>(handle,
                                              alpha.view(),
                                              beta.view(),
                                              std::nullopt,
                                              nEigVecs,
-                                             0,
+                                             which,
                                              ncv,
                                              eigenvectors.view(),
-                                             eigenvalues.view());
-
-  auto eigenvectors_k = raft::make_device_matrix_view<ValueTypeT, uint32_t, raft::col_major>(
-    eigenvectors.data_handle(), ncv, nEigVecs);
-  auto eigenvalues_k =
-    raft::make_device_vector_view<ValueTypeT, uint32_t>(eigenvalues.data_handle(), nEigVecs);
+                                             eigenvalues.view(),
+                                             eigenvectors_k,
+                                             eigenvalues_k,
+                                             eigenvectors_k_slice,
+                                             sm_eigenvalues.view(),
+                                             sm_eigenvectors.view());
 
   auto ritz_eigenvectors =
     raft::make_device_matrix_view<ValueTypeT, uint32_t, raft::col_major>(eigVecs_dev, n, nEigVecs);
@@ -1777,9 +1862,6 @@ auto lanczos_smallest(
 
   auto s = raft::make_device_vector<ValueTypeT>(handle, nEigVecs);
 
-  auto eigenvectors_k_slice =
-    raft::make_device_matrix_view<ValueTypeT, IndexTypeT, raft::col_major>(
-      eigenvectors.data_handle(), ncv, nEigVecs);
   auto S_matrix = raft::make_device_matrix_view<ValueTypeT, IndexTypeT, raft::col_major>(
     s.data_handle(), 1, nEigVecs);
 
@@ -2002,12 +2084,15 @@ auto lanczos_smallest(
                                                beta.view(),
                                                beta_k.view(),
                                                nEigVecs,
-                                               0,
+                                               which,
                                                ncv,
                                                eigenvectors.view(),
-                                               eigenvalues.view());
-    auto eigenvectors_k = raft::make_device_matrix_view<ValueTypeT, uint32_t, raft::col_major>(
-      eigenvectors.data_handle(), ncv, nEigVecs);
+                                               eigenvalues.view(),
+                                               eigenvectors_k,
+                                               eigenvalues_k,
+                                               eigenvectors_k_slice,
+                                               sm_eigenvalues.view(),
+                                               sm_eigenvectors.view());
 
     auto ritz_eigenvectors = raft::make_device_matrix_view<ValueTypeT, uint32_t, raft::col_major>(
       eigVecs_dev, n, nEigVecs);
@@ -2017,9 +2102,6 @@ auto lanczos_smallest(
     raft::linalg::gemm<ValueTypeT, uint32_t, raft::col_major, raft::col_major, raft::col_major>(
       handle, V_T, eigenvectors_k, ritz_eigenvectors);
 
-    auto eigenvectors_k_slice =
-      raft::make_device_matrix_view<ValueTypeT, IndexTypeT, raft::col_major>(
-        eigenvectors.data_handle(), ncv, nEigVecs);
     auto S_matrix = raft::make_device_matrix_view<ValueTypeT, IndexTypeT, raft::col_major>(
       s.data_handle(), 1, nEigVecs);
 
@@ -2066,6 +2148,7 @@ auto lanczos_compute_smallest_eigenvectors(
                             config.max_iterations,
                             config.ncv,
                             config.tolerance,
+                            config.which,
                             eigenvalues.data_handle(),
                             eigenvectors.data_handle(),
                             v0->data_handle(),
@@ -2082,6 +2165,7 @@ auto lanczos_compute_smallest_eigenvectors(
                             config.max_iterations,
                             config.ncv,
                             config.tolerance,
+                            config.which,
                             eigenvalues.data_handle(),
                             eigenvectors.data_handle(),
                             temp_v0.data_handle(),

cpp/include/raft/sparse/solver/lanczos_types.hpp

@@ -20,17 +20,57 @@
 
 namespace raft::sparse::solver {
 
+/**
+ * @enum LANCZOS_WHICH
+ * @brief Enumeration specifying which eigenvalues to compute in the Lanczos algorithm
+ */
+enum LANCZOS_WHICH {
+  /** @brief LA: Largest (algebraic) eigenvalues */
+  LA,
+  /** @brief LM: Largest (in magnitude) eigenvalues */
+  LM,
+  /** @brief SA: Smallest (algebraic) eigenvalues */
+  SA,
+  /** @brief SM: Smallest (in magnitude) eigenvalues */
+  SM
+};
+
+/**
+ * @brief Configuration parameters for the Lanczos eigensolver
+ *
+ * This structure encapsulates all configuration parameters needed to run the
+ * Lanczos algorithm for computing eigenvalues and eigenvectors of large sparse matrices.
+ *
+ * @tparam ValueTypeT Data type for values (float or double)
+ */
 template <typename ValueTypeT>
 struct lanczos_solver_config {
-  /** The number of eigenvalues and eigenvectors to compute. Must be 1 <= k < n.*/
+  /** @brief The number of eigenvalues and eigenvectors to compute
+   *  @note Must be 1 <= n_components < n, where n is the matrix dimension
+   */
   int n_components;
-  /** Maximum number of iteration. */
+
+  /** @brief Maximum number of iterations allowed for the algorithm to converge */
   int max_iterations;
-  /** The number of Lanczos vectors generated. Must be k + 1 < ncv < n. */
+
+  /** @brief The number of Lanczos vectors to generate
+   *  @note Must satisfy n_components + 1 < ncv < n, where n is the matrix dimension
+   */
   int ncv;
-  /** Tolerance for residuals ``||Ax - wx||`` */
+
+  /** @brief Convergence tolerance for residuals
+   *  @note Used to determine when to stop iteration based on ||Ax - wx|| < tolerance
+   */
   ValueTypeT tolerance;
-  /** random seed */
+
+  /** @brief Specifies which eigenvalues to compute in the Lanczos algorithm
+   *  @see LANCZOS_WHICH for possible values (SA, LA, SM, LM)
+   */
+  LANCZOS_WHICH which;
+
+  /** @brief Random seed for initialization of the algorithm
+   *  @note Controls reproducibility of results
+   */
   uint64_t seed;
 };
 

cpp/include/raft/spectral/eigen_solvers.cuh

@@ -40,6 +40,8 @@ struct eigen_solver_config_t {
     1234567};  // CAVEAT: this default value is now common to all instances of using seed in
                // Lanczos; was not the case before: there were places where a default seed = 123456
                // was used; this may trigger slightly different # solver iterations
+
+  raft::sparse::solver::LANCZOS_WHICH which{raft::sparse::solver::LANCZOS_WHICH::SA};
 };
 
 template <typename index_type_t, typename value_type_t, typename size_type_t = index_type_t>
@@ -79,8 +81,13 @@ struct lanczos_solver_t {
     RAFT_EXPECTS(eigVals != nullptr, "Null eigVals buffer.");
     RAFT_EXPECTS(eigVecs != nullptr, "Null eigVecs buffer.");
 
-    auto lanczos_config = raft::sparse::solver::lanczos_solver_config<value_type_t>{
-      config_.n_eigVecs, config_.maxIter, config_.restartIter, config_.tol, config_.seed};
+    auto lanczos_config =
+      raft::sparse::solver::lanczos_solver_config<value_type_t>{config_.n_eigVecs,
+                                                                config_.maxIter,
+                                                                config_.restartIter,
+                                                                config_.tol,
+                                                                config_.which,
+                                                                config_.seed};
     auto v0_opt = std::optional<raft::device_vector_view<value_type_t, uint32_t, raft::row_major>>{
       std::nullopt};
     auto input_structure = input.structure_view();