Refactor hpsi_func in hsolver

python/pyabacus/src/py_diago_dav_subspace.hpp

@@ -113,23 +113,21 @@ class PyDiagoDavSubspace
         auto hpsi_func = [mm_op] (
             std::complex<double> *psi_in,
             std::complex<double> *hpsi_out, 
-            const int nband_in,
-            const int nbasis_in, 
-            const int band_index1,
-            const int band_index2
+            const int ldPsi,
+            const int nvec
         ) {
             // Note: numpy's py::array_t is row-major, but
             //       our raw pointer-array is column-major
-            py::array_t<std::complex<double>, py::array::f_style> psi({nbasis_in, band_index2 - band_index1 + 1});
+            py::array_t<std::complex<double>, py::array::f_style> psi({ldPsi, nvec});
             py::buffer_info psi_buf = psi.request();
             std::complex<double>* psi_ptr = static_cast<std::complex<double>*>(psi_buf.ptr);
-            std::copy(psi_in + band_index1 * nbasis_in, psi_in + (band_index2 + 1) * nbasis_in, psi_ptr);
+            std::copy(psi_in, psi_in + nvec * ldPsi, psi_ptr);
 
             py::array_t<std::complex<double>, py::array::f_style> hpsi = mm_op(psi);
 
             py::buffer_info hpsi_buf = hpsi.request();
             std::complex<double>* hpsi_ptr = static_cast<std::complex<double>*>(hpsi_buf.ptr);
-            std::copy(hpsi_ptr, hpsi_ptr + (band_index2 - band_index1 + 1) * nbasis_in, hpsi_out);
+            std::copy(hpsi_ptr, hpsi_ptr + nvec * ldPsi, hpsi_out);
         };
 
         obj = std::make_unique<hsolver::Diago_DavSubspace<std::complex<double>, base_device::DEVICE_CPU>>(

python/pyabacus/src/py_diago_david.hpp

@@ -111,23 +111,21 @@ class PyDiagoDavid
         auto hpsi_func = [mm_op] (
             std::complex<double> *psi_in,
             std::complex<double> *hpsi_out, 
-            const int nband_in, 
-            const int nbasis_in, 
-            const int band_index1, 
-            const int band_index2
+            const int ldPsi, 
+            const int nvec
         ) {
             // Note: numpy's py::array_t is row-major, but
             //       our raw pointer-array is column-major
-            py::array_t<std::complex<double>, py::array::f_style> psi({nbasis_in, band_index2 - band_index1 + 1});
+            py::array_t<std::complex<double>, py::array::f_style> psi({ldPsi, nvec});
             py::buffer_info psi_buf = psi.request();
             std::complex<double>* psi_ptr = static_cast<std::complex<double>*>(psi_buf.ptr);
-            std::copy(psi_in + band_index1 * nbasis_in, psi_in + (band_index2 + 1) * nbasis_in, psi_ptr);
+            std::copy(psi_in, psi_in + nvec * ldPsi, psi_ptr);
 
             py::array_t<std::complex<double>, py::array::f_style> hpsi = mm_op(psi);
 
             py::buffer_info hpsi_buf = hpsi.request();
             std::complex<double>* hpsi_ptr = static_cast<std::complex<double>*>(hpsi_buf.ptr);
-            std::copy(hpsi_ptr, hpsi_ptr + (band_index2 - band_index1 + 1) * nbasis_in, hpsi_out);
+            std::copy(hpsi_ptr, hpsi_ptr + nvec * ldPsi, hpsi_out);
         };
 
         auto spsi_func = [this] (

source/module_hsolver/diago_dav_subspace.cpp

@@ -124,7 +124,8 @@ int Diago_DavSubspace<T, Device>::diag_once(const HPsiFunc& hpsi_func,
 
     // compute h*psi_in_iter
     // NOTE: bands after the first n_band should yield zero
-    hpsi_func(this->psi_in_iter, this->hphi, this->nbase_x, this->dim, 0, this->nbase_x - 1);
+    // hphi[:, 0:nbase_x] = H * psi_in_iter[:, 0:nbase_x]
+    hpsi_func(this->psi_in_iter, this->hphi, this->dim, this->nbase_x);
 
     // at this stage, notconv = n_band and nbase = 0
     // note that nbase of cal_elem is an inout parameter: nbase := nbase + notconv
@@ -421,7 +422,8 @@ void Diago_DavSubspace<T, Device>::cal_grad(const HPsiFunc& hpsi_func,
     }
 
     // update hpsi[:, nbase:nbase+notconv]
-    hpsi_func(psi_iter, &hphi[nbase * this->dim], this->nbase_x, this->dim, nbase, nbase + notconv - 1);
+    // hpsi[:, nbase:nbase+notconv] = H * psi_iter[:, nbase:nbase+notconv]
+    hpsi_func(psi_iter + nbase * dim, hphi + nbase * this->dim, this->dim, notconv);
 
     ModuleBase::timer::tick("Diago_DavSubspace", "cal_grad");
     return;
@@ -886,7 +888,8 @@ void Diago_DavSubspace<T, Device>::diagH_subspace(T* psi_pointer, // [in] & [out
 
     {
         // do hPsi for all bands
-        hpsi_func(psi_pointer, hphi, n_band, dmax, 0, nstart - 1);
+        // hphi[:, 0:nstart] = H * psi_pointer[:, 0:nstart]
+        hpsi_func(psi_pointer, hphi, dmax, nstart);
 
         gemm_op<T, Device>()(ctx,
                              'C',

source/module_hsolver/diago_dav_subspace.h

@@ -31,7 +31,8 @@ class Diago_DavSubspace : public DiagH<T, Device>
 
     virtual ~Diago_DavSubspace() override;
 
-    using HPsiFunc = std::function<void(T*, T*, const int, const int, const int, const int)>;
+    // See diago_david.h for information on the HPsiFunc function type
+    using HPsiFunc = std::function<void(T*, T*, const int, const int)>;
 
     int diag(const HPsiFunc& hpsi_func,
              T* psi_in,

source/module_hsolver/diago_david.cpp

@@ -230,7 +230,9 @@ int DiagoDavid<T, Device>::diag_once(const HPsiFunc& hpsi_func,
     // end of SchmidtOrth and calculate H|psi>
     // hpsi_info dav_hpsi_in(&basis, psi::Range(true, 0, 0, nband - 1), this->hpsi);
     // phm_in->ops->hPsi(dav_hpsi_in);
-    hpsi_func(basis, hpsi, nbase_x, dim, 0, nband - 1);
+    // hpsi[:, 0:nband] = H basis[:, 0:nband]
+    // slice index in this piece of code is in C manner. i.e. 0:id stands for [0,id)
+    hpsi_func(basis, hpsi, dim, nband);
 
     this->cal_elem(dim, nbase, nbase_x, this->notconv, this->hpsi, this->spsi, this->hcc, this->scc);
 
@@ -601,7 +603,8 @@ void DiagoDavid<T, Device>::cal_grad(const HPsiFunc& hpsi_func,
     //                       psi::Range(true, 0, nbase, nbase + notconv - 1),
     //                       &hpsi[nbase * dim]); // &hp(nbase, 0)
     // phm_in->ops->hPsi(dav_hpsi_in);
-    hpsi_func(basis, &hpsi[nbase * dim], nbase_x, dim, nbase, nbase + notconv - 1);
+    // hpsi[:, nbase:nbase+notcnv] = H basis[:, nbase:nbase+notcnv]
+    hpsi_func(basis + nbase * dim, hpsi + nbase * dim, dim, notconv);
 
     delmem_complex_op()(this->ctx, lagrange);
     delmem_complex_op()(this->ctx, vc_ev_vector);
@@ -1149,9 +1152,7 @@ void DiagoDavid<T, Device>::planSchmidtOrth(const int nband, std::vector<int>& p
 /**
  * @brief Performs iterative diagonalization using the David algorithm.
  * 
- * @warning Please see docs of `HPsiFunc` for more information.
- * @warning Please adhere strictly to the requirements of the function pointer
- * @warning for the hpsi mat-vec interface; it may seem counterintuitive.
+ * @warning Please see docs of `HPsiFunc` for more information about the hpsi mat-vec interface.
  * 
  * @tparam T The type of the elements in the matrix.
  * @tparam Device The device type (CPU or GPU).

source/module_hsolver/diago_david.h

@@ -38,50 +38,48 @@ class DiagoDavid : public DiagH<T, Device>
      * this function computes the product of the Hamiltonian matrix H and a blockvector X.
      * 
      * Called as follows:
-     * hpsi(X, HX, nvec, dim, id_start, id_end)
-     * Result is stored in HX.
-     * HX = H * X[id_start:id_end]
+     * hpsi(X, HX, ld, nvec) where X and HX are (ld, nvec)-shaped blockvectors.
+     * Result HX = H * X is stored in HX.
      *
      * @param[out] X      Head address of input blockvector of type `T*`.
-     * @param[in]  HX     Where to write output blockvector of type `T*`.
-     * @param[in]  nvec   Number of eigebpairs, i.e. number of vectors in a block.
-     * @param[in]  dim    Dimension of matrix.
-     * @param[in]  id_start Start index of blockvector.
-     * @param[in]  id_end   End index of blockvector.
+     * @param[in]  HX     Head address of output blockvector of type `T*`.
+     * @param[in]  ld     Leading dimension of blockvector.
+     * @param[in]  nvec   Number of vectors in a block.
      * 
-     * @warning HX is the exact address to store output H*X[id_start:id_end];
-     * @warning while X is the head address of input blockvector, \b without offset.
-     * @warning Calling function should pass X and HX[offset] as arguments,
-     * @warning where offset is usually id_start * leading dimension.
+     * @warning X and HX are the exact address to read input X and store output H*X,
+     * @warning both of size ld * nvec.
      */
-    using HPsiFunc = std::function<void(T*, T*, const int, const int, const int, const int)>;
+    using HPsiFunc = std::function<void(T*, T*, const int, const int)>;
 
     /**
      * @brief A function type representing the SX function.
+     * 
+     * nrow is leading dimension of spsi, npw is leading dimension of psi, nbands is number of vecs
      *
      * This function type is used to define a matrix-blockvector operator S.
      * For generalized eigenvalue problem HX = λSX,
      * this function computes the product of the overlap matrix S and a blockvector X.
      *
-     * @param[in]   X     Pointer to the input array.
-     * @param[out] SX     Pointer to the output array.
-     * @param[in] nrow    Dimension of SX: nbands * nrow.
-     * @param[in] npw     Number of plane waves.
-     * @param[in] nbands  Number of bands.
+     * @param[in]   X     Pointer to the input blockvector.
+     * @param[out] SX     Pointer to the output blockvector.
+     * @param[in] nrow    Leading dimension of spsi. Dimension of SX: nbands * nrow.
+     * @param[in] npw     Leading dimension of psi. Number of plane waves.
+     * @param[in] nbands  Number of vectors.
      * 
-     * @note called as spsi(in, out, dim, dim, 1)
+     * @note called like spsi(in, out, dim, dim, 1)
      */
     using SPsiFunc = std::function<void(T*, T*, const int, const int, const int)>;
 
-    int diag(const HPsiFunc& hpsi_func,           // function void hpsi(T*, T*, const int, const int, const int, const int) 
-             const SPsiFunc& spsi_func,           // function void spsi(T*, T*, const int, const int, const int) 
-                      const int ldPsi,            // Leading dimension of the psi input
-                      T *psi_in,                  // Pointer to eigenvectors
-                      Real* eigenvalue_in,        // Pointer to store the resulting eigenvalues
-                      const Real david_diag_thr,  // Convergence threshold for the Davidson iteration
-                      const int david_maxiter,    // Maximum allowed iterations for the Davidson method
-                      const int ntry_max = 5,     // Maximum number of diagonalization attempts (default is 5)
-                      const int notconv_max = 0); // Maximum number of allowed non-converged eigenvectors
+    int diag(
+      const HPsiFunc& hpsi_func,  // function void hpsi(T*, T*, const int, const int) 
+      const SPsiFunc& spsi_func,  // function void spsi(T*, T*, const int, const int, const int) 
+      const int ldPsi,            // Leading dimension of the psi input
+      T *psi_in,                  // Pointer to eigenvectors
+      Real* eigenvalue_in,        // Pointer to store the resulting eigenvalues
+      const Real david_diag_thr,  // Convergence threshold for the Davidson iteration
+      const int david_maxiter,    // Maximum allowed iterations for the Davidson method
+      const int ntry_max = 5,     // Maximum number of diagonalization attempts (5 by default)
+      const int notconv_max = 0); // Maximum number of allowed non-converged eigenvectors
 
   private:
     bool use_paw = false;
@@ -171,12 +169,12 @@ class DiagoDavid : public DiagH<T, Device>
                  T* vcc);
 
     void SchmidtOrth(const int& dim,
-                    const int nband,
-                    const int m,
-                    const T* spsi,
-                    T* lagrange_m,
-                    const int mm_size,
-                    const int mv_size);
+                     const int nband,
+                     const int m,
+                     const T* spsi,
+                     T* lagrange_m,
+                     const int mm_size,
+                     const int mv_size);
 
     void planSchmidtOrth(const int nband, std::vector<int>& pre_matrix_mm_m, std::vector<int>& pre_matrix_mv_m);
 

source/module_hsolver/hsolver_pw.cpp

@@ -434,18 +434,17 @@ void HSolverPW<T, Device>::hamiltSolvePsiK(hamilt::Hamilt<T, Device>* hm,
     else if (this->method == "dav_subspace")
     {
         auto ngk_pointer = psi.get_ngk_pointer();
-        auto hpsi_func = [hm, ngk_pointer](T* psi_in,
-                                           T* hpsi_out,
-                                           const int nband_in,
-                                           const int nbasis_in,
-                                           const int band_index1,
-                                           const int band_index2) {
+        // hpsi_func (X, HX, ld, nvec) -> HX = H(X), X and HX blockvectors of size ld x nvec
+        auto hpsi_func = [hm, ngk_pointer](T *psi_in,
+                                           T *hpsi_out,
+                                           const int ldPsi,
+                                           const int nvec) {
             ModuleBase::timer::tick("DavSubspace", "hpsi_func");
 
             // Convert "pointer data stucture" to a psi::Psi object
-            auto psi_iter_wrapper = psi::Psi<T, Device>(psi_in, 1, nband_in, nbasis_in, ngk_pointer);
+            auto psi_iter_wrapper = psi::Psi<T, Device>(psi_in, 1, nvec, ldPsi, ngk_pointer);
 
-            psi::Range bands_range(true, 0, band_index1, band_index2);
+            psi::Range bands_range(true, 0, 0, nvec-1);
 
             using hpsi_info = typename hamilt::Operator<T, Device>::hpsi_info;
             hpsi_info info(&psi_iter_wrapper, bands_range, hpsi_out);
@@ -492,19 +491,17 @@ void HSolverPW<T, Device>::hamiltSolvePsiK(hamilt::Hamilt<T, Device>* hm,
 
         auto ngk_pointer = psi.get_ngk_pointer();
         /// wrap hpsi into lambda function, Matrix \times blockvector
-        /// hpsi(X, HX, nband, dim, band_index1, band_index2)
-        auto hpsi_func = [hm, ngk_pointer](T* psi_in,
-                                           T* hpsi_out,
-                                           const int nband_in,
-                                           const int nbasis_in,
-                                           const int band_index1,
-                                           const int band_index2) {
+        // hpsi_func (X, HX, ld, nvec) -> HX = H(X), X and HX blockvectors of size ld x nvec
+        auto hpsi_func = [hm, ngk_pointer](T *psi_in,
+                                           T *hpsi_out,
+                                           const int ldPsi,
+                                           const int nvec) {
             ModuleBase::timer::tick("David", "hpsi_func");
 
-            // Convert "pointer data stucture" to a psi::Psi object
-            auto psi_iter_wrapper = psi::Psi<T, Device>(psi_in, 1, nband_in, nbasis_in, ngk_pointer);
+            // Convert pointer of psi_in to a psi::Psi object
+            auto psi_iter_wrapper = psi::Psi<T, Device>(psi_in, 1, nvec, ldPsi, ngk_pointer);
 
-            psi::Range bands_range(true, 0, band_index1, band_index2);
+            psi::Range bands_range(true, 0, 0, nvec-1);
 
             using hpsi_info = typename hamilt::Operator<T, Device>::hpsi_info;
             hpsi_info info(&psi_iter_wrapper, bands_range, hpsi_out);
@@ -515,14 +512,15 @@ void HSolverPW<T, Device>::hamiltSolvePsiK(hamilt::Hamilt<T, Device>* hm,
 
         /// wrap spsi into lambda function, Matrix \times blockvector
         /// spsi(X, SX, nrow, npw, nbands)
+        /// nrow is leading dimension of spsi, npw is leading dimension of psi, nbands is number of vecs
         auto spsi_func = [hm](const T* psi_in, T* spsi_out,
-                               const int nrow,  // dimension of spsi: nbands * nrow
-                               const int npw,   // number of plane waves
-                               const int nbands // number of bands
+                               const int ldSpsi,  // dimension of spsi: nbands * nrow
+                               const int ldPsi,   // number of plane waves
+                               const int nvec     // number of bands
                             ){
             ModuleBase::timer::tick("David", "spsi_func");
             // sPsi determines S=I or not by GlobalV::use_uspp inside
-            hm->sPsi(psi_in, spsi_out, nrow, npw, nbands);
+            hm->sPsi(psi_in, spsi_out, ldSpsi, ldPsi, nvec);
             ModuleBase::timer::tick("David", "spsi_func");
         };
 

source/module_hsolver/test/diago_david_float_test.cpp

@@ -108,11 +108,10 @@ class DiagoDavPrepare
 
 
 		auto hpsi_func = [phm](std::complex<float>* psi_in,std::complex<float>* hpsi_out,
-					const int nband_in, const int nbasis_in,
-                    const int band_index1, const int band_index2)
+					const int ldPsi, const int nvec)
                     {
-                        auto psi_iter_wrapper = psi::Psi<std::complex<float>>(psi_in, 1, nband_in, nbasis_in, nullptr);
-                        psi::Range bands_range(1, 0, band_index1, band_index2);
+                        auto psi_iter_wrapper = psi::Psi<std::complex<float>>(psi_in, 1, nvec, ldPsi, nullptr);
+                        psi::Range bands_range(1, 0, 0, nvec-1);
                         using hpsi_info = typename hamilt::Operator<std::complex<float>>::hpsi_info;
                         hpsi_info info(&psi_iter_wrapper, bands_range, hpsi_out);
                         phm->ops->hPsi(info);