Add subspace diagonalization for orthogonal input (#6559)

* Update subspace_func where needed

* Refactor hsolver function to do subspace diagonalization, standard & generalized, used in cg

* Rename heevx params

* Update corresponding tests

* update pyabacus contributing md

* Update corresponding pyabacus code and remove need_subspace in dav_subspace code

* Update pyabacus api

* Fix typo

Author: Cstandardlib

python/pyabacus/CONTRIBUTING.md

@@ -189,7 +189,7 @@ list(APPEND _diago
     ${HSOLVER_PATH}/diago_david.cpp
     ${HSOLVER_PATH}/diag_const_nums.cpp
     ${HSOLVER_PATH}/diago_iter_assist.cpp
-    ${HSOLVER_PATH}/kernels/dngvd_op.cpp
+    ${HSOLVER_PATH}/kernels/hegvd_op.cpp
     ${HSOLVER_PATH}/kernels/bpcg_kernel_op.cpp
     ${BASE_PATH}/kernels/math_kernel_op.cpp
     ${BASE_PATH}/kernels/math_kernel_op_vec.cpp

python/pyabacus/src/hsolver/py_diago_cg.hpp

@@ -145,7 +145,7 @@ class PyDiagoCG
             std::copy(hpsi_ptr, hpsi_ptr + nvec * ld_psi, hpsi_out.data<std::complex<double>>());
         };
 
-        auto subspace_func = [] (const ct::Tensor& psi_in, ct::Tensor& psi_out) { /*do nothing*/ };
+        auto subspace_func = [](const ct::Tensor& psi_in, ct::Tensor& psi_out, const bool S_orth) { /*do nothing*/ };
 
         auto spsi_func = [this] (const ct::Tensor& psi_in, ct::Tensor& spsi_out) {
             const auto ndim = psi_in.shape().ndim();

python/pyabacus/src/hsolver/py_diago_dav_subspace.hpp

@@ -144,7 +144,6 @@ class PyDiagoDavSubspace
             dav_ndim, 
             tol, 
             max_iter, 
-            need_subspace, 
             comm_info,
             diag_subspace,
             nb2d

source/source_hsolver/diago_cg.cpp

@@ -28,7 +28,7 @@ template <typename T, typename Device>
 DiagoCG<T, Device>::DiagoCG(const std::string& basis_type,
                             const std::string& calculation,
                             const bool& need_subspace,
-                            const Func& subspace_func,
+                            const SubspaceFunc& subspace_func,
                             const Real& pw_diag_thr,
                             const int& pw_diag_nmax,
                             const int& nproc_in_pool)
@@ -54,7 +54,7 @@ DiagoCG<T, Device>::~DiagoCG()
 }
 
 template <typename T, typename Device>
-void DiagoCG<T, Device>::diag_mock(const ct::Tensor& prec_in,
+void DiagoCG<T, Device>::diag_once(const ct::Tensor& prec_in,
                                    ct::Tensor& psi,
                                    ct::Tensor& eigen,
                                    const std::vector<double>& ethr_band)
@@ -569,7 +569,7 @@ bool DiagoCG<T, Device>::test_exit_cond(const int& ntry, const int& notconv) con
     // In non-self consistent calculation, do until totally converged.
     const bool f2 = !scf && notconv > 0;
     // if self consistent calculation, if not converged > 5,
-    // using diagH_subspace and cg method again. ntry++
+    // using diag_subspace and cg method again. ntry++
     const bool f3 = scf && notconv > 5;
     return f1 && (f2 || f3);
 }
@@ -592,16 +592,29 @@ void DiagoCG<T, Device>::diag(const Func& hpsi_func,
     ct::Tensor psi_temp = psi.slice({0, 0}, {int(psi.shape().dim_size(0)), int(prec.shape().dim_size(0))});
     do
     {
-        if (need_subspace_ || ntry > 0)
+        // subspace diagonalization to get a better starting guess
+        // for cg diagonalization, restart from current psi approximation
+        // Note: if not the first try, then psi is already S-orthogonalized by CG iterations!
+        // Otherwise, if the first try, then psi is not assumed to be S-orthogonalized
+        if (ntry > 0)
         {
             ct::TensorMap psi_map = ct::TensorMap(psi.data(), psi_temp);
-            this->subspace_func_(psi_temp, psi_map);
+            const bool assume_S_orthogonal = true;
+            this->subspace_func_(psi_temp, psi_map, assume_S_orthogonal);
             psi_temp.sync(psi_map);
         }
+        else if (need_subspace_)
+        {
+            ct::TensorMap psi_map = ct::TensorMap(psi.data(), psi_temp);
+            const bool assume_S_orthogonal = false;
+            this->subspace_func_(psi_temp, psi_map, assume_S_orthogonal);
+            psi_temp.sync(psi_map);
+        }
+
 
         ++ntry;
         avg_iter_ += 1.0;
-        this->diag_mock(prec, psi_temp, eigen, ethr_band);
+        this->diag_once(prec, psi_temp, eigen, ethr_band);
     } while (this->test_exit_cond(ntry, this->notconv_));
 
     if (this->notconv_ > std::max(5, this->n_band_ / 4))

source/source_hsolver/diago_cg.h

@@ -22,6 +22,7 @@ class DiagoCG final
     using ct_Device = typename ct::PsiToContainer<Device>::type;
   public:
     using Func = std::function<void(const ct::Tensor&, ct::Tensor&)>;
+    using SubspaceFunc = std::function<void(const ct::Tensor&, ct::Tensor&, const bool)>;
     // Constructor need:
     // 1. temporary mock of Hamiltonian "Hamilt_PW"
     // 2. precondition pointer should point to place of precondition array.
@@ -30,7 +31,7 @@ class DiagoCG final
         const std::string& basis_type,
         const std::string& calculation,
         const bool& need_subspace,
-        const Func& subspace_func,
+        const SubspaceFunc& subspace_func,
         const Real& pw_diag_thr,
         const int& pw_diag_nmax,
         const int& nproc_in_pool);
@@ -72,11 +73,11 @@ class DiagoCG final
 
     bool need_subspace_ = false;
     /// A function object that performs the hPsi calculation.
-    std::function<void(const ct::Tensor&, ct::Tensor&)> hpsi_func_ = nullptr;
+    Func hpsi_func_ = nullptr;
     /// A function object that performs the sPsi calculation.
-    std::function<void(const ct::Tensor&, ct::Tensor&)> spsi_func_ = nullptr;
+    Func spsi_func_ = nullptr;
     /// A function object that performs the subspace calculation.
-    std::function<void(const ct::Tensor&, ct::Tensor&)> subspace_func_ = nullptr;
+    SubspaceFunc subspace_func_ = nullptr;
 
     void calc_grad(
         const ct::Tensor& prec,
@@ -119,7 +120,7 @@ class DiagoCG final
     void schmit_orth(const int& m, const ct::Tensor& psi, const ct::Tensor& sphi, ct::Tensor& phi_m);
 
     // used in diag() for template replace Hamilt with Hamilt_PW
-    void diag_mock(const ct::Tensor& prec,
+    void diag_once(const ct::Tensor& prec,
                    ct::Tensor& psi,
                    ct::Tensor& eigen,
                    const std::vector<double>& ethr_band);

source/source_hsolver/diago_dav_subspace.cpp

@@ -1,7 +1,7 @@
 #include "diago_dav_subspace.h"
 
 #include "diago_iter_assist.h"
-#include "source_base/memory.h"
+
 #include "source_base/module_device/device.h"
 #include "source_base/timer.h"
 #include "source_hsolver/kernels/hegvd_op.h"
@@ -26,12 +26,11 @@ Diago_DavSubspace<T, Device>::Diago_DavSubspace(const std::vector<Real>& precond
                                                 const int& david_ndim_in,
                                                 const double& diag_thr_in,
                                                 const int& diag_nmax_in,
-                                                const bool& need_subspace_in,
                                                 const diag_comm_info& diag_comm_in,
                                                 const int diag_subspace_in,
                                                 const int diago_subspace_bs_in)
     : precondition(precondition_in), n_band(nband_in), dim(nbasis_in), nbase_x(nband_in * david_ndim_in),
-      diag_thr(diag_thr_in), iter_nmax(diag_nmax_in), is_subspace(need_subspace_in), diag_comm(diag_comm_in),
+      diag_thr(diag_thr_in), iter_nmax(diag_nmax_in), diag_comm(diag_comm_in),
         diag_subspace(diag_subspace_in), diago_subspace_bs(diago_subspace_bs_in)
 {
     this->device = base_device::get_device_type<Device>(this->ctx);
@@ -76,7 +75,7 @@ Diago_DavSubspace<T, Device>::Diago_DavSubspace(const std::vector<Real>& precond
     {
         resmem_real_op()(this->d_precondition, nbasis_in);
         // syncmem_var_h2d_op()(this->ctx, this->cpu_ctx, this->d_precondition, this->precondition.data(), nbasis_in);
-        base_device::memory::resize_memory_op<T, Device>()(this->d_scc, this->nbase_x * this->nbase_x);
+        resmem_complex_op()(this->d_scc, this->nbase_x * this->nbase_x);
         resmem_real_op()(this->d_eigenvalue, this->nbase_x);
     }
 #endif
@@ -370,7 +369,6 @@ void Diago_DavSubspace<T, Device>::cal_grad(const HPsiFunc& hpsi_func,
     {
         Real* psi_norm = nullptr;
         resmem_real_op()(psi_norm, notconv);
-        using setmem_real_op = base_device::memory::set_memory_op<Real, Device>;
         setmem_real_op()(psi_norm, 0.0, notconv);
 
         normalize_op<T, Device>()(this->dim,
@@ -541,7 +539,7 @@ void Diago_DavSubspace<T, Device>::diag_zhegvx(const int& nbase,
 #if defined(__CUDA) || defined(__ROCM)
         if (this->diag_comm.rank == 0)
         {
-            base_device::memory::synchronize_memory_op<T, Device, Device>()(this->d_scc, scc, nbase * this->nbase_x);
+            syncmem_complex_op()(this->d_scc, scc, nbase * this->nbase_x);
             hegvd_op<T, Device>()(this->ctx, nbase, this->nbase_x, this->hcc, this->d_scc, this->d_eigenvalue, this->vcc);
             syncmem_var_d2h_op()((*eigenvalue_iter).data(), this->d_eigenvalue, this->nbase_x);
         }

source/source_hsolver/diago_dav_subspace.h

@@ -30,7 +30,6 @@ class Diago_DavSubspace
                       const int& david_ndim_in,
                       const double& diag_thr_in,
                       const int& diag_nmax_in,
-                      const bool& need_subspace_in,
                       const diag_comm_info& diag_comm_in,
                       const int diago_dav_method_in,
                       const int block_size_in);
@@ -58,9 +57,6 @@ class Diago_DavSubspace
     /// maximal iteration number
     const int iter_nmax;
 
-    /// is diagH_subspace needed?
-    const bool is_subspace;
-
     /// the first dimension of the matrix to be diagonalized
     const int n_band = 0;