Refactor: add smooth threshold support for david method (deepmodeling#5697)

* change raw pointer to std::vector

* add ethr_band for dav method

* change unit test

* fix build bug

* fix pyabacus

* fix pyabacus dav-subspace

* fix pyabacus build

* fix pyabacus

* add & for vector

Author: haozhihan

python/pyabacus/src/hsolver/py_diago_dav_subspace.hpp

@@ -101,12 +101,12 @@ class PyDiagoDavSubspace
 
     int diag(
         std::function<py::array_t<std::complex<double>>(py::array_t<std::complex<double>>)> mm_op,
-        std::vector<double> precond_vec,
+        std::vector<double>& precond_vec,
         int dav_ndim,
         double tol,
         int max_iter,
         bool need_subspace,
-        std::vector<double> diag_ethr,
+        std::vector<double>& diag_ethr,
         bool scf_type,
         hsolver::diag_comm_info comm_info
     ) {
@@ -141,7 +141,7 @@ class PyDiagoDavSubspace
             comm_info
         );
 
-        return obj->diag(hpsi_func, psi, nbasis, eigenvalue, diag_ethr.data(), scf_type);
+        return obj->diag(hpsi_func, psi, nbasis, eigenvalue, diag_ethr, scf_type);
     }
 
 private:

python/pyabacus/src/hsolver/py_diago_david.hpp

@@ -101,9 +101,10 @@ class PyDiagoDavid
 
     int diag(
         std::function<py::array_t<std::complex<double>>(py::array_t<std::complex<double>>)> mm_op,
-        std::vector<double> precond_vec,
+        std::vector<double>& precond_vec,
         int dav_ndim,
         double tol,
+        std::vector<double>& diag_ethr,
         int max_iter,
         bool use_paw,
         hsolver::diag_comm_info comm_info
@@ -146,7 +147,7 @@ class PyDiagoDavid
             comm_info
         );
 
-        return obj->diag(hpsi_func, spsi_func, nbasis, psi, eigenvalue, tol, max_iter);
+        return obj->diag(hpsi_func, spsi_func, nbasis, psi, eigenvalue, diag_ethr, max_iter);
     }
 
 private:

python/pyabacus/src/hsolver/py_hsolver.cpp

@@ -121,6 +121,8 @@ void bind_hsolver(py::module& m)
                 eigenvectors to be calculated.
             tol : double
                 The tolerance for the convergence.
+            diag_ethr: np.ndarray
+                The tolerance vector.
             max_iter : int
                 The maximum number of iterations.
             use_paw : bool
@@ -130,6 +132,7 @@ void bind_hsolver(py::module& m)
         "precond_vec"_a, 
         "dav_ndim"_a, 
         "tol"_a, 
+        "diag_ethr"_a,
         "max_iter"_a, 
         "use_paw"_a, 
         "comm_info"_a)

python/pyabacus/src/pyabacus/hsolver/_hsolver.py

@@ -118,6 +118,7 @@ def davidson(
     dav_ndim: int = 2,
     tol: float = 1e-2,
     max_iter: int = 1000,
+    diag_ethr: Union[List[float], None] = None,
     use_paw: bool = False,
     # scf_type: bool = False
 ) -> Tuple[NDArray[np.float64], NDArray[np.complex128]]:
@@ -143,6 +144,8 @@ def davidson(
         The tolerance for the convergence, by default 1e-2.
     max_iter : int, optional    
         The maximum number of iterations, by default 1000.
+    diag_ethr : List[float] | None, optional
+        The list of thresholds of bands, by default None.    
     use_paw : bool, optional
         Whether to use projector augmented wave (PAW) method, by default False.
     
@@ -164,12 +167,16 @@ def davidson(
     _diago_obj_david.init_eigenvalue()
 
     comm_info = diag_comm_info(0, 1)
+
+    if diag_ethr is None:
+        diag_ethr = [tol] * num_eigs
 
     _ = _diago_obj_david.diag(
         mvv_op,
         precondition,
         dav_ndim,
         tol,
+        diag_ethr,
         max_iter,
         use_paw,
         comm_info

source/module_hsolver/diago_dav_subspace.cpp

@@ -88,7 +88,7 @@ int Diago_DavSubspace<T, Device>::diag_once(const HPsiFunc& hpsi_func,
                                             T* psi_in,
                                             const int psi_in_dmax,
                                             Real* eigenvalue_in_hsolver,
-                                            const double* ethr_band)
+                                            const std::vector<double>& ethr_band)
 {
     ModuleBase::timer::tick("Diago_DavSubspace", "diag_once");
 
@@ -726,7 +726,7 @@ int Diago_DavSubspace<T, Device>::diag(const HPsiFunc& hpsi_func,
                                        T* psi_in,
                                        const int psi_in_dmax,
                                        Real* eigenvalue_in_hsolver,
-                                       const double* ethr_band,
+                                       const std::vector<double>& ethr_band,
                                        const bool& scf_type)
 {
     /// record the times of trying iterative diagonalization

source/module_hsolver/diago_dav_subspace.h

@@ -42,7 +42,7 @@ class Diago_DavSubspace
              T* psi_in,
              const int psi_in_dmax,
              Real* eigenvalue_in,
-             const double* ethr_band,
+             const std::vector<double>& ethr_band,
              const bool& scf_type);
 
   private:
@@ -135,7 +135,7 @@ class Diago_DavSubspace
                   T* psi_in,
                   const int psi_in_dmax,
                   Real* eigenvalue_in,
-                  const double* ethr_band);
+                  const std::vector<double>& ethr_band);
 
     bool test_exit_cond(const int& ntry, const int& notconv, const bool& scf);
 

source/module_hsolver/diago_david.cpp

@@ -156,7 +156,7 @@ int DiagoDavid<T, Device>::diag_once(const HPsiFunc& hpsi_func,
                                      const int ld_psi,
                                      T *psi_in,
                                      Real* eigenvalue_in,
-                                     const Real david_diag_thr,
+                                     const std::vector<double>& ethr_band,
                                      const int david_maxiter)
 {
     if (test_david == 1)
@@ -273,7 +273,7 @@ int DiagoDavid<T, Device>::diag_once(const HPsiFunc& hpsi_func,
         this->notconv = 0;
         for (int m = 0; m < nband; m++)
         {
-            convflag[m] = (std::abs(this->eigenvalue[m] - eigenvalue_in[m]) < david_diag_thr);
+            convflag[m] = (std::abs(this->eigenvalue[m] - eigenvalue_in[m]) < ethr_band[m]);
             if (!convflag[m])
             {
                 unconv[this->notconv] = m;
@@ -1177,7 +1177,7 @@ int DiagoDavid<T, Device>::diag(const HPsiFunc& hpsi_func,
                                 const int ld_psi,
                                 T *psi_in,
                                 Real* eigenvalue_in,
-                                const Real david_diag_thr,
+                                const std::vector<double>& ethr_band,
                                 const int david_maxiter,
                                 const int ntry_max,
                                 const int notconv_max)
@@ -1189,7 +1189,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, ld_psi, 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, ethr_band, david_maxiter);
         ++ntry;
     } while (!check_block_conv(ntry, this->notconv, ntry_max, notconv_max));
 

source/module_hsolver/diago_david.h

@@ -79,7 +79,7 @@ class DiagoDavid
       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 std::vector<double>& ethr_band, // 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
@@ -134,7 +134,7 @@ class DiagoDavid
                   const int ld_psi,
                   T *psi_in,
                   Real* eigenvalue_in,
-                  const Real david_diag_thr,
+                  const std::vector<double>& ethr_band,
                   const int david_maxiter);
 
     void cal_grad(const HPsiFunc& hpsi_func,

source/module_hsolver/hsolver_pw.cpp

@@ -526,7 +526,7 @@ void HSolverPW<T, Device>::hamiltSolvePsiK(hamilt::Hamilt<T, Device>* hm,
                                                   comm_info);
 
         DiagoIterAssist<T, Device>::avg_iter += static_cast<double>(
-            dav_subspace.diag(hpsi_func, psi.get_pointer(), psi.get_nbasis(), eigenvalue, this->ethr_band.data(), scf));
+            dav_subspace.diag(hpsi_func, psi.get_pointer(), psi.get_nbasis(), eigenvalue, this->ethr_band, scf));
     }
     else if (this->method == "dav")
     {
@@ -589,7 +589,7 @@ void HSolverPW<T, Device>::hamiltSolvePsiK(hamilt::Hamilt<T, Device>* hm,
                                                                                ld_psi,
                                                                                psi.get_pointer(),
                                                                                eigenvalue,
-                                                                               david_diag_thr,
+                                                                               this->ethr_band,
                                                                                david_maxiter,
                                                                                ntry_max,
                                                                                notconv_max));

source/module_hsolver/test/diago_david_float_test.cpp

@@ -119,10 +119,12 @@ class DiagoDavPrepare
                         hpsi_info info(&psi_iter_wrapper, bands_range, hpsi_out);
                         phm->ops->hPsi(info);
                     };
-		auto spsi_func = [phm](const std::complex<float>* psi_in, std::complex<float>* spsi_out,const int ld_psi, const int nbands){
-			phm->sPsi(psi_in, spsi_out, ld_psi, ld_psi, nbands);
-		};
-		dav.diag(hpsi_func,spsi_func, ld_psi, phi.get_pointer(), en, eps, maxiter);
+        auto spsi_func = [phm](const std::complex<float>* psi_in,
+                               std::complex<float>* spsi_out,
+                               const int ld_psi,
+                               const int nbands) { phm->sPsi(psi_in, spsi_out, ld_psi, ld_psi, nbands); };
+        std::vector<double> ethr_band(phi.get_nbands(), eps);
+		dav.diag(hpsi_func,spsi_func, ld_psi, phi.get_pointer(), en, ethr_band, maxiter);
 
 #ifdef __MPI		
 		end = MPI_Wtime();