Feature: update new version of dav_subspace with higher performance

python/pyabacus/src/py_diago_dav_subspace.hpp

@@ -106,7 +106,7 @@ class PyDiagoDavSubspace
         double tol,
         int max_iter,
         bool need_subspace,
-        std::vector<bool> is_occupied,
+        std::vector<double> diag_ethr,
         bool scf_type,
         hsolver::diag_comm_info comm_info
     ) {
@@ -143,7 +143,7 @@ class PyDiagoDavSubspace
             comm_info
         );
 
-        return obj->diag(hpsi_func, psi, nbasis, eigenvalue, is_occupied, scf_type);
+        return obj->diag(hpsi_func, psi, nbasis, eigenvalue, diag_ethr.data(), scf_type);
     }
 
 private:

python/pyabacus/src/py_hsolver.cpp

@@ -59,8 +59,8 @@ void bind_hsolver(py::module& m)
                 The maximum number of iterations.
             need_subspace : bool
                 Whether to use the subspace function.
-            is_occupied : list[bool]
-                A list of boolean values indicating whether the band is occupied,
+            diag_ethr : list[float]
+                A list of float values indicating the thresholds of each band for the diagonalization,
                 meaning that the corresponding eigenvalue is to be calculated.
             scf_type : bool
                 Whether to use the SCF type, which is used to determine the
@@ -76,7 +76,7 @@ void bind_hsolver(py::module& m)
         "tol"_a, 
         "max_iter"_a, 
         "need_subspace"_a, 
-        "is_occupied"_a, 
+        "diag_ethr"_a, 
         "scf_type"_a, 
         "comm_info"_a)
         .def("set_psi", &py_hsolver::PyDiagoDavSubspace::set_psi, R"pbdoc(

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

@@ -25,7 +25,7 @@ def dav_subspace(
     tol: float = 1e-2,
     max_iter: int = 1000,
     need_subspace: bool = False,
-    is_occupied: Union[List[bool], None] = None,
+    diag_ethr: Union[List[float], None] = None,
     scf_type: bool = False
 ) -> Tuple[NDArray[np.float64], NDArray[np.complex128]]:
     """ A function to diagonalize a matrix using the Davidson-Subspace method.
@@ -52,10 +52,8 @@ def dav_subspace(
         The maximum number of iterations, by default 1000.
     need_subspace : bool, optional
         Whether to use subspace function, by default False.
-    is_occupied : List[bool] | None, optional
-        The list of occupied bands, by default None. This indicates how many eigenvalues 
-        need to be calculated, starting from the smallest eigenvalue. Only the energy levels 
-        occupied by electrons (occupied) need to be calculated.
+    diag_ethr : List[float] | None, optional
+        The list of thresholds of bands, by default None.
     scf_type : bool, optional
         Indicates whether the calculation is a self-consistent field (SCF) calculation. 
         If True, the initial precision of eigenvalue calculation can be coarse. 
@@ -72,8 +70,8 @@ def dav_subspace(
     if not callable(mvv_op):
         raise TypeError("mvv_op must be a callable object.")
 
-    if is_occupied is None:
-        is_occupied = [True] * num_eigs
+    if diag_ethr is None:
+        diag_ethr = [tol] * num_eigs
 
     if init_v.ndim != 1 or init_v.dtype != np.complex128:
         init_v = init_v.flatten().astype(np.complex128, order='C')
@@ -93,7 +91,7 @@ def dav_subspace(
         tol,
         max_iter,
         need_subspace,
-        is_occupied,
+        diag_ethr,
         scf_type,
         comm_info
     )
@@ -113,7 +111,6 @@ def davidson(
     tol: float = 1e-2,
     max_iter: int = 1000,
     use_paw: bool = False,
-    # is_occupied: Union[List[bool], None] = None,
     # scf_type: bool = False
 ) -> Tuple[NDArray[np.float64], NDArray[np.complex128]]:
     """ A function to diagonalize a matrix using the Davidson-Subspace method.

source/module_elecstate/elecstate.cpp

@@ -253,31 +253,6 @@ void ElecState::init_ks(Charge* chg_in, // pointer for class Charge
     this->wg.create(nk_in, GlobalV::NBANDS);
 }
 
-void set_is_occupied(std::vector<bool>& is_occupied,
-                     elecstate::ElecState* pes,
-                     const int i_scf,
-                     const int nk,
-                     const int nband,
-                     const bool diago_full_acc)
-{
-    if (i_scf != 0 && diago_full_acc == false)
-    {
-        for (int i = 0; i < nk; i++)
-        {
-            if (pes->klist->wk[i] > 0.0)
-            {
-                for (int j = 0; j < nband; j++)
-                {
-                    if (pes->wg(i, j) / pes->klist->wk[i] < 0.01)
-                    {
-                        is_occupied[i * nband + j] = false;
-                    }
-                }
-            }
-        }
-    }
-};
-
 
 
 } // namespace elecstate

source/module_elecstate/elecstate.h

@@ -194,13 +194,5 @@ class ElecState
     bool skip_weights = false;
 };
 
-// This is an independent function under the elecstate namespace and does not depend on any class.
-void set_is_occupied(std::vector<bool>& is_occupied,
-                     elecstate::ElecState* pes,
-                     const int i_scf,
-                     const int nk,
-                     const int nband,
-                     const bool diago_full_acc);
-
 } // namespace elecstate
 #endif

source/module_elecstate/potentials/pot_local.cpp

@@ -8,6 +8,8 @@
 namespace elecstate
 {
 
+double PotLocal::vl_of_0 = 0.0;
+
 //==========================================================
 // This routine computes the local potential in real space
 //==========================================================
@@ -29,6 +31,8 @@ void PotLocal::cal_fixed_v(double *vl_pseudo // store the local pseudopotential
         }
     }
 
+    PotLocal::vl_of_0 = vg[0].real();
+
     // recip2real should be a const function, but now it isn't
     // a dangerous usage appears here, which should be fix in the future.
     const_cast<ModulePW::PW_Basis *>(this->rho_basis_)->recip2real(vg, vl_pseudo);

source/module_elecstate/potentials/pot_local.h

@@ -24,6 +24,15 @@ class PotLocal : public PotBase
 
     void cal_fixed_v(double* vl_pseudo) override;
 
+    /// @brief get the value of vloc at G=0;
+    /// @return vl(0)
+    static double get_vl_of_0() { return vl_of_0; }
+
+    private:
+
+    /// @brief save the value of vloc at G=0; this is a static member because there is only one vl(0) for all instances
+    static double vl_of_0;
+
     // std::vector<double> vltot;
 
     const ModuleBase::matrix* vloc_ = nullptr; // local pseduopotentials

source/module_esolver/esolver_ks_pw.cpp

@@ -352,15 +352,6 @@ void ESolver_KS_PW<T, Device>::hamilt2density(const int istep, const int iter, c
         hsolver::DiagoIterAssist<T, Device>::SCF_ITER = iter;
         hsolver::DiagoIterAssist<T, Device>::PW_DIAG_THR = ethr;
         hsolver::DiagoIterAssist<T, Device>::PW_DIAG_NMAX = PARAM.inp.pw_diag_nmax;
-
-        std::vector<bool> is_occupied(this->kspw_psi->get_nk() * this->kspw_psi->get_nbands(), true);
-
-        elecstate::set_is_occupied(is_occupied,
-                                   this->pelec,
-                                   hsolver::DiagoIterAssist<T, Device>::SCF_ITER,
-                                   this->kspw_psi->get_nk(),
-                                   this->kspw_psi->get_nbands(),
-                                   PARAM.inp.diago_full_acc);
 
         hsolver::HSolverPW<T, Device> hsolver_pw_obj(this->pw_wfc, 
                                                      &this->wf, 
@@ -383,7 +374,6 @@ void ESolver_KS_PW<T, Device>::hamilt2density(const int istep, const int iter, c
                            this->kspw_psi[0],
                            this->pelec,
                            this->pelec->ekb.c,
-                           is_occupied,
                            GlobalV::RANK_IN_POOL,
                            GlobalV::NPROC_IN_POOL,
                            false);

source/module_esolver/pw_fun.cpp

@@ -71,15 +71,6 @@ void ESolver_KS_PW<T, Device>::hamilt2estates(const double ethr)
     hsolver::DiagoIterAssist<T, Device>::need_subspace = false;
     hsolver::DiagoIterAssist<T, Device>::PW_DIAG_THR = ethr;
 
-    std::vector<bool> is_occupied(this->kspw_psi->get_nk() * this->kspw_psi->get_nbands(), true);
-
-    elecstate::set_is_occupied(is_occupied,
-                                this->pelec,
-                                hsolver::DiagoIterAssist<T, Device>::SCF_ITER,
-                                this->kspw_psi->get_nk(),
-                                this->kspw_psi->get_nbands(),
-                                PARAM.inp.diago_full_acc);
-
     hsolver::HSolverPW<T, Device> hsolver_pw_obj(this->pw_wfc, 
                                                  &this->wf, 
 
@@ -101,7 +92,6 @@ void ESolver_KS_PW<T, Device>::hamilt2estates(const double ethr)
                          this->kspw_psi[0],
                          this->pelec,
                          this->pelec->ekb.c,
-                         is_occupied,
                          GlobalV::RANK_IN_POOL,
                          GlobalV::NPROC_IN_POOL,
                          true);