Refactor hpsi_func of dav

Author: Cstandardlib

source/module_hsolver/diago_david.cpp

@@ -230,7 +230,7 @@ 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_func(basis, hpsi, dim, nband);
 
     this->cal_elem(dim, nbase, nbase_x, this->notconv, this->hpsi, this->spsi, this->hcc, this->scc);
 
@@ -601,7 +601,7 @@ 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_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 +1149,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

@@ -44,17 +44,13 @@ class DiagoDavid : public DiagH<T, Device>
      *
      * @param[out] X      Head address of input blockvector of type `T*`.
      * @param[in]  HX     Where to write output blockvector of type `T*`.
+     * @param[in]  ld    Leading dimension of matrix.
      * @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.
      * 
-     * @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.

source/module_hsolver/hsolver_pw.cpp

@@ -492,18 +492,16 @@ 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) {
+        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);
+            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);