Feature: update new version of dav_subspace with higher performance (#5199)

* Feature: update new version of dav_subspace

* Fix: UnitTest error

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

* Fix: UnitTest HSolver_PW

* Fix: comments in PR

* Fix: error in UT

---------

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

Author: 3 people

docs/advanced/input_files/input-main.md

@@ -39,7 +39,6 @@
     - [pw\_diag\_thr](#pw_diag_thr)
     - [pw\_diag\_nmax](#pw_diag_nmax)
     - [pw\_diag\_ndim](#pw_diag_ndim)
-    - [diago\_full\_acc](#diago_full_acc)
     - [erf\_ecut](#erf_ecut)
     - [fft\_mode](#fft_mode)
     - [erf\_height](#erf_height)
@@ -779,12 +778,6 @@ These variables are used to control the plane wave related parameters.
 - **Description**: Only useful when you use `ks_solver = dav` or `ks_solver = dav_subspace`. It indicates dimension of workspace(number of wavefunction packets, at least 2 needed) for the Davidson method. A larger value may yield a smaller number of iterations in the algorithm but uses more memory and more CPU time in subspace diagonalization.
 - **Default**: 4
 
-### diago_full_acc
-
-- **Type**: bool
-- **Description**: Only useful when you use `ks_solver = dav_subspace`. If `TRUE`, all the empty states are diagonalized at the same level of accuracy of the occupied ones. Otherwise the empty states are diagonalized using a larger threshold (10-5) (this should not affect total energy, forces, and other ground-state properties).
-- **Default**: false
-
 ### erf_ecut
 
 - **Type**: Real
@@ -925,7 +918,7 @@ calculations.
   - **cg**: cg method.
   - **bpcg**: bpcg method, which is a block-parallel Conjugate Gradient (CG) method, typically exhibits higher acceleration in a GPU environment.
   - **dav**: the Davidson algorithm.
-  - **dav_subspace**: subspace Davidson algorithm
+  - **dav_subspace**: Davidson algorithm without orthogonalization operation, this method is the most recommended for efficiency. `pw_diag_ndim` can be set to 2 for this method.
 
   For atomic orbitals basis,
 

examples/lr-tddft/lcao_Si2/INPUT

@@ -37,4 +37,3 @@ out_alllog	1
 
 nvirt 19
 abs_wavelen_range  100 175
-#diago_full_acc 1

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
     ) {
@@ -141,7 +141,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

@@ -252,31 +252,6 @@ void ElecState::init_ks(Charge* chg_in, // pointer for class Charge
     this->wg.create(nk_in, PARAM.inp.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

@@ -29,6 +29,12 @@ void PotLocal::cal_fixed_v(double *vl_pseudo // store the local pseudopotential
         }
     }
 
+    /// save the V_local at G=0
+    if(this->rho_basis_->npw > 0)
+    {
+        *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

@@ -12,8 +12,9 @@ class PotLocal : public PotBase
   public:
     PotLocal(const ModuleBase::matrix* vloc_in, // local pseduopotentials
              const ModuleBase::ComplexMatrix* sf_in,
-             const ModulePW::PW_Basis* rho_basis_in)
-        : vloc_(vloc_in), sf_(sf_in)
+             const ModulePW::PW_Basis* rho_basis_in,
+             double& vl_of_0)
+        : vloc_(vloc_in), sf_(sf_in), vl_of_0_(&vl_of_0)
     {
         assert(this->vloc_->nr == this->sf_->nr);
         this->rho_basis_ = rho_basis_in;
@@ -24,6 +25,11 @@ class PotLocal : public PotBase
 
     void cal_fixed_v(double* vl_pseudo) override;
 
+    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
+    double* vl_of_0_ = nullptr;
+
     // std::vector<double> vltot;
 
     const ModuleBase::matrix* vloc_ = nullptr; // local pseduopotentials

source/module_elecstate/potentials/potential_new.h

@@ -169,6 +169,14 @@ class Potential : public PotBase
         return this->v_effective_fixed.data();
     }
 
+
+    /// @brief get the value of vloc at G=0;
+    /// @return vl(0)
+    double get_vl_of_0() const
+    {
+        return this->vl_of_0;
+    }
+
   private:
     void cal_v_eff(const Charge*const chg, const UnitCell*const ucell, ModuleBase::matrix& v_eff) override;
     void cal_fixed_v(double* vl_pseudo) override;
@@ -196,6 +204,8 @@ class Potential : public PotBase
     double* etxc_ = nullptr;
     double* vtxc_ = nullptr;
 
+    double vl_of_0 = 0.0;
+
     std::vector<PotBase*> components;
 
     const UnitCell* ucell_ = nullptr;