Refactor: Remove some files in DeePKS and add some features. (#6065)

* Add HR precalc functions for DeePKS and fix some bugs.

* Remove and combine some files in DeePKS.

* Fix wrong function position.

Author: ErjieWu

source/Makefile.Objects

@@ -208,7 +208,7 @@ OBJS_DEEPKS=LCAO_deepks.o\
         deepks_orbpre.o\
         deepks_vdelta.o\
         deepks_vdpre.o\
-        deepks_hmat.o\
+        deepks_vdrpre.o\
         deepks_pdm.o\
         deepks_phialpha.o\
         LCAO_deepks_io.o\

source/module_hamilt_lcao/module_deepks/CMakeLists.txt

@@ -11,7 +11,7 @@ if(ENABLE_DEEPKS)
       deepks_orbpre.cpp
       deepks_vdelta.cpp
       deepks_vdpre.cpp
-      deepks_hmat.cpp
+      deepks_vdrpre.cpp
       deepks_pdm.cpp
       deepks_phialpha.cpp
       LCAO_deepks_io.cpp

source/module_hamilt_lcao/module_deepks/LCAO_deepks.h

@@ -7,14 +7,14 @@
 #include "deepks_descriptor.h"
 #include "deepks_force.h"
 #include "deepks_fpre.h"
-#include "deepks_hmat.h"
 #include "deepks_orbital.h"
 #include "deepks_orbpre.h"
 #include "deepks_pdm.h"
 #include "deepks_phialpha.h"
 #include "deepks_spre.h"
 #include "deepks_vdelta.h"
 #include "deepks_vdpre.h"
+#include "deepks_vdrpre.h"
 #include "module_base/complexmatrix.h"
 #include "module_base/intarray.h"
 #include "module_base/matrix.h"
@@ -56,10 +56,10 @@ class LCAO_Deepks
   public:
     ///(Unit: Ry) Correction energy provided by NN
     double E_delta = 0.0;
-    ///(Unit: Ry)  \f$tr(\rho H_\delta), \rho = \sum_i{c_{i, \mu}c_{i,\nu}} \f$ (for gamma_only)
+    ///(Unit: Ry)  \f$tr(\rho H_\delta), \rho = \sum_i{c_{i, \mu}c_{i,\nu}} \f$
     double e_delta_band = 0.0;
 
-    /// Correction term to the Hamiltonian matrix: \f$\langle\phi|V_\delta|\phi\rangle\f$ (for gamma only)
+    /// Correction term to the Hamiltonian matrix: \f$\langle\phi|V_\delta|\phi\rangle\f$
     /// The first dimension is for k-points V_delta(k)
     std::vector<std::vector<T>> V_delta;
 

source/module_hamilt_lcao/module_deepks/LCAO_deepks_interface.cpp

@@ -314,34 +314,18 @@ void LCAO_Deepks_Interface<TK, TR>::out_deepks_labels(const double& etot,
                     ofs_hr.close();
                 }
 
-                const std::string file_vdrpre = PARAM.globalv.global_out_dir + "deepks_vdrpre.csr";
-                std::vector<hamilt::HContainer<TR>*> h_deltaR_pre(inlmax);
-                for (int i = 0; i < inlmax; i++)
-                {
-                    h_deltaR_pre[i] = new hamilt::HContainer<TR>(*hR_tot);
-                    h_deltaR_pre[i]->set_zero();
-                }
-                // DeePKS_domain::cal_vdr_precalc<TR>();
-                if (rank == 0)
-                {
-                    std::ofstream ofs_hrp(file_vdrpre, std::ios::out);
-                    for (int iat = 0; iat < nat; iat++)
-                    {
-                        ofs_hrp << "- Index of atom: " << iat << std::endl;
-                        for (int nl = 0; nl < nlmax; nl++)
-                        {
-                            int inl = iat * nlmax + nl;
-                            ofs_hrp << "-- Index of nl: " << nl << std::endl;
-                            ofs_hrp << "Matrix Dimension of H_delta(R): " << h_deltaR_pre[inl]->get_nbasis() << std::endl;
-                            ofs_hrp << "Matrix number of H_delta(R): " << h_deltaR_pre[inl]->size_R_loop() << std::endl;
-                            hamilt::Output_HContainer<TR> out_hrp(h_deltaR_pre[inl], ofs_hrp, sparse_threshold, precision);
-                            out_hrp.write();
-                            ofs_hrp << std::endl;
-                        }
-                        ofs_hrp << std::endl;
-                    }
-                    ofs_hrp.close();
-                }
+                torch::Tensor phialpha_r_out;
+                torch::Tensor R_query;
+                DeePKS_domain::prepare_phialpha_r(nlocal, lmaxd, inlmax, nat, phialpha, ucell, orb, *ParaV, GridD, phialpha_r_out, R_query);
+                const std::string file_phialpha_r = PARAM.globalv.global_out_dir + "deepks_phialpha_r.npy";
+                const std::string file_R_query = PARAM.globalv.global_out_dir + "deepks_R_query.npy";
+                LCAO_deepks_io::save_tensor2npy<double>(file_phialpha_r, phialpha_r_out, rank);
+                LCAO_deepks_io::save_tensor2npy<int>(file_R_query, R_query, rank);
+
+                torch::Tensor gevdm_out;
+                DeePKS_domain::prepare_gevdm(nat, lmaxd, inlmax, orb, gevdm, gevdm_out);
+                const std::string file_gevdm = PARAM.globalv.global_out_dir + "deepks_gevdm.npy";
+                LCAO_deepks_io::save_tensor2npy<double>(file_gevdm, gevdm_out, rank);
             }
         }
 

source/module_hamilt_lcao/module_deepks/LCAO_deepks_io.cpp

@@ -275,18 +275,18 @@ void LCAO_deepks_io::save_tensor2npy(const std::string& file_name, const torch::
 
     std::vector<T> data(tensor.numel());
 
-    if constexpr (std::is_same<T, double>::value)
-    {
-        std::memcpy(data.data(), tensor.data_ptr<double>(), tensor.numel() * sizeof(double));
-    }
-    else
+    if constexpr (std::is_same<T, std::complex<double>>::value)
     {
         auto tensor_data = tensor.data_ptr<c10::complex<double>>();
         for (size_t i = 0; i < tensor.numel(); ++i)
         {
             data[i] = std::complex<double>(tensor_data[i].real(), tensor_data[i].imag());
         }
     }
+    else
+    {
+        std::memcpy(data.data(), tensor.data_ptr<T>(), tensor.numel() * sizeof(T));
+    }
 
     npy::SaveArrayAsNumpy(file_name, false, shape.size(), shape.data(), data);
 }
@@ -313,6 +313,10 @@ template void LCAO_deepks_io::save_npy_h<std::complex<double>>(const std::vector
                                                                const int nks,
                                                                const int rank);
 
+template void LCAO_deepks_io::save_tensor2npy<int>(const std::string& file_name,
+                                                   const torch::Tensor& tensor,
+                                                   const int rank);
+
 template void LCAO_deepks_io::save_tensor2npy<double>(const std::string& file_name,
                                                       const torch::Tensor& tensor,
                                                       const int rank);

source/module_hamilt_lcao/module_deepks/deepks_hmat.cpp

@@ -77,6 +77,73 @@ void DeePKS_domain::cal_e_delta_band(const std::vector<std::vector<TK>>& dm,
     return;
 }
 
+template <typename TK, typename TH>
+void DeePKS_domain::collect_h_mat(const Parallel_Orbitals& pv,
+                                  const std::vector<std::vector<TK>>& h_in,
+                                  std::vector<TH>& h_out,
+                                  const int nlocal,
+                                  const int nks)
+{
+    ModuleBase::TITLE("DeePKS_domain", "collect_h_tot");
+
+    // construct the total H matrix
+    for (int k = 0; k < nks; k++)
+    {
+#ifdef __MPI
+        int ir = 0;
+        int ic = 0;
+        for (int i = 0; i < nlocal; i++)
+        {
+            std::vector<TK> lineH(nlocal - i, TK(0.0));
+
+            ir = pv.global2local_row(i);
+            if (ir >= 0)
+            {
+                // data collection
+                for (int j = i; j < nlocal; j++)
+                {
+                    ic = pv.global2local_col(j);
+                    if (ic >= 0)
+                    {
+                        int iic = 0;
+                        if (ModuleBase::GlobalFunc::IS_COLUMN_MAJOR_KS_SOLVER(PARAM.inp.ks_solver))
+                        {
+                            iic = ir + ic * pv.nrow;
+                        }
+                        else
+                        {
+                            iic = ir * pv.ncol + ic;
+                        }
+                        lineH[j - i] = h_in[k][iic];
+                    }
+                }
+            }
+            else
+            {
+                // do nothing
+            }
+
+            Parallel_Reduce::reduce_all(lineH.data(), nlocal - i);
+
+            for (int j = i; j < nlocal; j++)
+            {
+                h_out[k](i, j) = lineH[j - i];
+                h_out[k](j, i) = h_out[k](i, j); // H is a symmetric matrix
+            }
+        }
+#else
+        for (int i = 0; i < nlocal; i++)
+        {
+            for (int j = i; j < nlocal; j++)
+            {
+                h_out[k](i, j) = h_in[k][i * nlocal + j];
+                h_out[k](j, i) = h_out[k](i, j); // H is a symmetric matrix
+            }
+        }
+#endif
+    }
+}
+
 template void DeePKS_domain::cal_e_delta_band<double>(const std::vector<std::vector<double>>& dm,
                                                       const std::vector<std::vector<double>>& V_delta,
                                                       const int nks,
@@ -89,4 +156,18 @@ template void DeePKS_domain::cal_e_delta_band<std::complex<double>>(
     const Parallel_Orbitals* pv,
     double& e_delta_band);
 
+template void DeePKS_domain::collect_h_mat<double, ModuleBase::matrix>(
+    const Parallel_Orbitals& pv,
+    const std::vector<std::vector<double>>& h_in,
+    std::vector<ModuleBase::matrix>& h_out,
+    const int nlocal,
+    const int nks);
+
+template void DeePKS_domain::collect_h_mat<std::complex<double>, ModuleBase::ComplexMatrix>(
+    const Parallel_Orbitals& pv,
+    const std::vector<std::vector<std::complex<double>>>& h_in,
+    std::vector<ModuleBase::ComplexMatrix>& h_out,
+    const int nlocal,
+    const int nks);
+
 #endif