Refactor hpsi_func in hsolver (#5202)

* Refactor hpsi_func of dav_subspace

* Modified the hpsi_func in pyabacus to maintain definition consistency

* Fix hpsi_func in pyabacus-dav_subspace

* [pre-commit.ci lite] apply automatic fixes

* Refactor hpsi_func of dav

* Change hpsi_func of hsolver_lrtd

* Modify hpsi_func definition in dav tests

* Modify the hpsi_func in pyabacus to maintain definition consistency

* [pre-commit.ci lite] apply automatic fixes

* Modify hsolver_pw_sup mock func signature

* Update docs for new hpsi_func

* Update docs

* Change indent to make it prettier

* Update docs

* Update spsi docs

* Update parameter name of spsi_func interface

* Rename leading dimension vars from camel to snake case

* Rename leading dimension in pyabacus to snake case

---------

Co-authored-by: pre-commit-ci-lite[bot] <117423508+pre-commit-ci-lite[bot]@users.noreply.github.com>

Author: Cstandardlib

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 ld_psi,
+            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({ld_psi, 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 * ld_psi, 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 * ld_psi, 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 ld_psi, 
+            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({ld_psi, 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 * ld_psi, 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 * ld_psi, 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

@@ -152,7 +152,7 @@ int DiagoDavid<T, Device>::diag_once(const HPsiFunc& hpsi_func,
                                      const SPsiFunc& spsi_func,
                                      const int dim,
                                      const int nband,
-                                     const int ldPsi,
+                                     const int ld_psi,
                                      T *psi_in,
                                      Real* eigenvalue_in,
                                      const Real david_diag_thr,
@@ -191,20 +191,20 @@ int DiagoDavid<T, Device>::diag_once(const HPsiFunc& hpsi_func,
         if(this->use_paw)
         {
 #ifdef USE_PAW
-            GlobalC::paw_cell.paw_nl_psi(1, reinterpret_cast<const std::complex<double>*> (psi_in + m*ldPsi),
+            GlobalC::paw_cell.paw_nl_psi(1, reinterpret_cast<const std::complex<double>*> (psi_in + m*ld_psi),
                 reinterpret_cast<std::complex<double>*>(&this->spsi[m * dim]));
 #endif
         }
         else
         {
-            // phm_in->sPsi(psi_in + m*ldPsi, &this->spsi[m * dim], dim, dim, 1);
-            spsi_func(psi_in + m*ldPsi,&this->spsi[m*dim],dim,dim,1);
+            // phm_in->sPsi(psi_in + m*ld_psi, &this->spsi[m * dim], dim, dim, 1);
+            spsi_func(psi_in + m*ld_psi,&this->spsi[m*dim],dim,dim,1);
         }
     }
     // begin SchmidtOrth
     for (int m = 0; m < nband; m++)
     {
-        syncmem_complex_op()(this->ctx, this->ctx, basis + dim*m, psi_in + m*ldPsi, dim);
+        syncmem_complex_op()(this->ctx, this->ctx, basis + dim*m, psi_in + m*ld_psi, dim);
 
         this->SchmidtOrth(dim,
                          nband,
@@ -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);
 
@@ -287,7 +289,7 @@ int DiagoDavid<T, Device>::diag_once(const HPsiFunc& hpsi_func,
 
             // update eigenvectors of Hamiltonian
 
-            setmem_complex_op()(this->ctx, psi_in, 0, nband * ldPsi);
+            setmem_complex_op()(this->ctx, psi_in, 0, nband * ld_psi);
             //<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
             gemm_op<T, Device>()(this->ctx,
                                       'N',
@@ -302,7 +304,7 @@ int DiagoDavid<T, Device>::diag_once(const HPsiFunc& hpsi_func,
                                       nbase_x,
                                       this->zero,
                                       psi_in,       // C dim * nband
-                                      ldPsi
+                                      ld_psi
             );
 
             if (!this->notconv || (dav_iter == david_maxiter))
@@ -322,7 +324,7 @@ int DiagoDavid<T, Device>::diag_once(const HPsiFunc& hpsi_func,
                               nbase_x,
                               eigenvalue_in,
                               psi_in,
-                              ldPsi,
+                              ld_psi,
                               this->hpsi,
                               this->spsi,
                               this->hcc,
@@ -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);
@@ -785,7 +788,7 @@ void DiagoDavid<T, Device>::diag_zhegvx(const int& nbase,
  * @param nbase_x The maximum dimension of the reduced basis set.
  * @param eigenvalue_in Pointer to the array of eigenvalues.
  * @param psi_in Pointer to the array of wavefunctions.
- * @param ldPsi The leading dimension of the wavefunction array.
+ * @param ld_psi The leading dimension of the wavefunction array.
  * @param hpsi Pointer to the output array for the updated basis set.
  * @param spsi Pointer to the output array for the updated basis set (nband-th column).
  * @param hcc Pointer to the output array for the updated reduced Hamiltonian.
@@ -800,7 +803,7 @@ void DiagoDavid<T, Device>::refresh(const int& dim,
                                          const int nbase_x, // maximum dimension of the reduced basis set
                                          const Real* eigenvalue_in,
                                          const T *psi_in,
-                                         const int ldPsi,
+                                         const int ld_psi,
                                          T* hpsi,
                                          T* spsi,
                                          T* hcc,
@@ -866,7 +869,7 @@ void DiagoDavid<T, Device>::refresh(const int& dim,
 
     for (int m = 0; m < nband; m++)
     {
-        syncmem_complex_op()(this->ctx, this->ctx, basis + dim*m,psi_in + m*ldPsi, dim);
+        syncmem_complex_op()(this->ctx, this->ctx, basis + dim*m,psi_in + m*ld_psi, dim);
         /*for (int ig = 0; ig < npw; ig++)
             basis(m, ig) = psi(m, ig);*/
     }
@@ -1149,15 +1152,13 @@ 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).
  * @param hpsi_func The function object that computes the matrix-blockvector product H * psi.
  * @param spsi_func The function object that computes the matrix-blockvector product overlap S * psi.
- * @param ldPsi The leading dimension of the psi_in array.
+ * @param ld_psi The leading dimension of the psi_in array.
  * @param psi_in The input wavefunction.
  * @param eigenvalue_in The array to store the eigenvalues.
  * @param david_diag_thr The convergence threshold for the diagonalization.
@@ -1172,7 +1173,7 @@ void DiagoDavid<T, Device>::planSchmidtOrth(const int nband, std::vector<int>& p
 template <typename T, typename Device>
 int DiagoDavid<T, Device>::diag(const HPsiFunc& hpsi_func,
                                 const SPsiFunc& spsi_func,
-                                const int ldPsi,
+                                const int ld_psi,
                                 T *psi_in,
                                 Real* eigenvalue_in,
                                 const Real david_diag_thr,
@@ -1187,7 +1188,7 @@ int DiagoDavid<T, Device>::diag(const HPsiFunc& hpsi_func,
     int sum_dav_iter = 0;
     do
     {
-        sum_dav_iter += this->diag_once(hpsi_func, spsi_func, dim, nband, ldPsi, psi_in, eigenvalue_in, david_diag_thr, david_maxiter);
+        sum_dav_iter += this->diag_once(hpsi_func, spsi_func, dim, nband, ld_psi, psi_in, eigenvalue_in, david_diag_thr, david_maxiter);
         ++ntry;
     } while (!check_block_conv(ntry, this->notconv, ntry_max, notconv_max));
 

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] ld_spsi   Leading dimension of spsi. Dimension of SX: nbands * nrow.
+     * @param[in] ld_psi    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 ld_psi,           // 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;
@@ -130,7 +128,7 @@ class DiagoDavid : public DiagH<T, Device>
                   const SPsiFunc& spsi_func,
                   const int dim,
                   const int nband,
-                  const int ldPsi,
+                  const int ld_psi,
                   T *psi_in,
                   Real* eigenvalue_in,
                   const Real david_diag_thr,
@@ -163,20 +161,20 @@ class DiagoDavid : public DiagH<T, Device>
                  const int nbase_x,
                  const Real* eigenvalue,
                  const T *psi_in,
-                 const int ldPsi,
+                 const int ld_psi,
                  T* hpsi,
                  T* spsi,
                  T* hcc,
                  T* scc,
                  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);