LRCG (#5399)

docs/advanced/input_files/input-main.md

@@ -3925,7 +3925,7 @@ Currently supported: `RPA`, `LDA`, `PBE`, `HSE`, `HF`.
 
 - **Type**: String
 - **Description**: The method to solve the Casida equation $AX=\Omega X$ in LR-TDDFT under Tamm-Dancoff approximation (TDA), where $A_{ai,bj}=(\epsilon_a-\epsilon_i)\delta_{ij}\delta_{ab}+(ai|f_{Hxc}|bj)+\alpha_{EX}(ab|ij)$ is the particle-hole excitation matrix and $X$ is the transition amplitude.
-  - `dav`: Construct $AX$ and diagonalize the Hamiltonian matrix iteratively with Davidson algorithm.
+  - `dav`/`dav_subspace`/ `cg`: Construct $AX$ and diagonalize the Hamiltonian matrix iteratively with Davidson/Non-ortho-Davidson/CG algorithm.
   - `lapack`: Construct the full $A$ matrix and directly diagonalize with LAPACK.
   - `spectrum`: Calculate absorption spectrum only without solving Casida equation. The `OUT.${suffix}/` directory should contain the
   files for LR-TDDFT eigenstates and eigenvalues, i.e. `Excitation_Energy.dat` and `Excitation_Amplitude_${processor_rank}.dat`

source/module_hsolver/diago_iter_assist.cpp

@@ -20,7 +20,7 @@ namespace hsolver
 // Produces on output n_band eigenvectors (n_band <= nstart) in evc.
 //----------------------------------------------------------------------
 template <typename T, typename Device>
-void DiagoIterAssist<T, Device>::diagH_subspace(hamilt::Hamilt<T, Device>* pHamilt, // hamiltonian operator carrier
+void DiagoIterAssist<T, Device>::diagH_subspace(const hamilt::Hamilt<T, Device>* const pHamilt, // hamiltonian operator carrier
                                                 const psi::Psi<T, Device>& psi,     // [in] wavefunction
                                                 psi::Psi<T, Device>& evc,           // [out] wavefunction
                                                 Real* en,                           // [out] eigenvalues

source/module_hsolver/diago_iter_assist.h

@@ -29,7 +29,7 @@ class DiagoIterAssist
     static int SCF_ITER;
 
     // for psi::Psi structure
-    static void diagH_subspace(hamilt::Hamilt<T, Device>* pHamilt,
+    static void diagH_subspace(const hamilt::Hamilt<T, Device>* const pHamilt,
                                const psi::Psi<T, Device>& psi,
                                psi::Psi<T, Device>& evc,
                                Real* en,

source/module_lr/esolver_lrtd_lcao.cpp

@@ -86,7 +86,7 @@ inline int cal_nupdown_form_occ(const ModuleBase::matrix& wg)
 template<typename T, typename TR>
 void LR::ESolver_LR<T, TR>::parameter_check()const
 {
-    std::set<std::string> lr_solvers = { "dav", "lapack" , "spectrum", "dav_subspace" };
+    std::set<std::string> lr_solvers = { "dav", "lapack" , "spectrum", "dav_subspace", "cg" };
     std::set<std::string> xc_kernels = { "rpa", "lda", "pbe", "hf" , "hse" };
     if (lr_solvers.find(this->input.lr_solver) == lr_solvers.end()) {
         throw std::invalid_argument("ESolver_LR: unknown type of lr_solver");

source/module_lr/hsolver_lrtd.hpp

@@ -4,8 +4,10 @@
 #include "module_hsolver/diago_dav_subspace.h"
 #include "module_hsolver/diago_cg.h"
 #include "module_hsolver/diago_iter_assist.h"
+#include "module_hsolver/diago_cg.h"
 #include "module_lr/utils/lr_util.h"
 #include "module_lr/utils/lr_util_print.h"
+#include "module_base/module_container/ATen/core/tensor_map.h"
 
 namespace LR
 {
@@ -68,7 +70,7 @@ namespace LR
                 // 3. set precondition and diagethr
                 for (int i = 0; i < dim; ++i) { precondition[i] = static_cast<Real<T>>(1.0); }
 
-                const int david_maxiter = hsolver::DiagoIterAssist<T>::PW_DIAG_NMAX;
+                const int maxiter = hsolver::DiagoIterAssist<T>::PW_DIAG_NMAX;
 
                 auto hpsi_func = [&hm](T* psi_in, T* hpsi, const int ld_psi, const int nvec) {hm.hPsi(psi_in, hpsi, ld_psi, nvec);};
                 auto spsi_func = [&hm](const T* psi_in, T* spsi, const int ld_psi, const int nvec)
@@ -84,7 +86,7 @@ namespace LR
                     // do diag and add davidson iteration counts up to avg_iter
                     hsolver::DiagoDavid<T> david(precondition.data(), nband, dim, PARAM.inp.pw_diag_ndim, PARAM.inp.use_paw, comm_info);
                     hsolver::DiagoIterAssist<T>::avg_iter += static_cast<double>(david.diag(hpsi_func, spsi_func,
-                        dim, psi, eigenvalue.data(), diag_ethr, david_maxiter, ntry_max, 0));
+                        dim, psi, eigenvalue.data(), diag_ethr, maxiter, ntry_max, 0));
                 }
                 else if (method == "dav_subspace") //need refactor
                 {
@@ -93,7 +95,7 @@ namespace LR
                         dim,
                         PARAM.inp.pw_diag_ndim,
                         diag_ethr,
-                        david_maxiter,
+                        maxiter,
                         false, //always do the subspace diag (check the implementation)
                         comm_info);
                     std::vector<double> ethr_band(nband, diag_ethr);
@@ -105,6 +107,44 @@ namespace LR
                             ethr_band.data(),
                             false /*scf*/));
                 }
+                else if (method == "cg")
+                {
+                    ////// `diagH_subspace` needs refactor: 
+                    ////// replace `Hamilt*` with `hpsi_func`
+                    ////// or I cannot use `is_subspace=true` as my `HamiltLR` does not inherit `Hamilt`.
+
+                    // auto subspace_func = [&hm](const ct::Tensor& psi_in, ct::Tensor& psi_out) {
+                    //     const auto ndim = psi_in.shape().ndim();
+                    //     REQUIRES_OK(ndim == 2, "dims of psi_in should be less than or equal to 2");
+                    //     // Convert a Tensor object to a psi::Psi object
+                    //     auto psi_in_wrapper = psi::Psi<T>(psi_in.data<T>(),
+                    //         1,
+                    //         psi_in.shape().dim_size(0),
+                    //         psi_in.shape().dim_size(1));
+                    //     auto psi_out_wrapper = psi::Psi<T>(psi_out.data<T>(),
+                    //         1,
+                    //         psi_out.shape().dim_size(0),
+                    //         psi_out.shape().dim_size(1));
+                    //     auto eigen = ct::Tensor(ct::DataTypeToEnum<Real<T>>::value,
+                    //         ct::DeviceType::CpuDevice,
+                    //         ct::TensorShape({ psi_in.shape().dim_size(0) }));
+                    //     hsolver::DiagoIterAssist<T>::diagH_subspace(hm, psi_in_wrapper, psi_out_wrapper, eigen.data<Real<T>>());
+                    //     };
+
+                    ////// why diago_cg depends on basis_type?
+                    // hsolver::DiagoCG<T> cg("lcao", "nscf", true, subspace_func, diag_ethr, maxiter, GlobalV::NPROC_IN_POOL);
+
+                    auto subspace_func = [](const ct::Tensor& psi_in, ct::Tensor& psi_out) {};
+                    hsolver::DiagoCG<T> cg("lcao", "nscf", false, subspace_func, diag_ethr, maxiter, GlobalV::NPROC_IN_POOL);
+
+                    auto psi_tensor = ct::TensorMap(psi, ct::DataTypeToEnum<T>::value, ct::DeviceType::CpuDevice, ct::TensorShape({ nband, dim }));
+                    auto eigen_tensor = ct::TensorMap(eigenvalue.data(), ct::DataTypeToEnum<Real<T>>::value, ct::DeviceType::CpuDevice, ct::TensorShape({ nband }));
+                    auto precon_tensor = ct::TensorMap(precondition.data(), ct::DataTypeToEnum<Real<T>>::value, ct::DeviceType::CpuDevice, ct::TensorShape({ dim }));
+                    auto hpsi_func = [&hm](const ct::Tensor& psi_in, ct::Tensor& hpsi) {hm.hPsi(psi_in.data<T>(), hpsi.data<T>(), psi_in.shape().dim_size(0) /*nbasis_local*/, 1/*band-by-band*/);};
+                    auto spsi_func = [&hm](const ct::Tensor& psi_in, ct::Tensor& spsi)
+                        { std::memcpy(spsi.data<T>(), psi_in.data<T>(), sizeof(T) * psi_in.NumElements()); };
+                    cg.diag(hpsi_func, spsi_func, psi_tensor, eigen_tensor, precon_tensor);
+                }
                 else { throw std::runtime_error("HSolverLR::solve: method not implemented"); }
             }