Remove global dependence of pdm related functions in DeePKS.

ErjieWu · ErjieWu · commit a5e310999a29 · 2025-01-08T15:10:33.000+08:00
diff --git a/source/Makefile.Objects b/source/Makefile.Objects
@@ -201,8 +201,8 @@ OBJS_DEEPKS=LCAO_deepks.o\
         deepks_vdelta.o\
         deepks_vdpre.o\
         deepks_hmat.o\
+        deepks_pdm.o\
         LCAO_deepks_io.o\
-        LCAO_deepks_pdm.o\
         LCAO_deepks_phialpha.o\
         LCAO_deepks_interface.o\
         
diff --git a/source/module_esolver/esolver_ks_lcao.cpp b/source/module_esolver/esolver_ks_lcao.cpp
@@ -227,7 +227,14 @@ void ESolver_KS_LCAO<TK, TR>::before_all_runners(UnitCell& ucell, const Input_pa
         // load the DeePKS model from deep neural network
         DeePKS_domain::load_model(PARAM.inp.deepks_model, GlobalC::ld.model_deepks);
         // read pdm from file for NSCF or SCF-restart, do it only once in whole calculation
-        GlobalC::ld.read_projected_DM((PARAM.inp.init_chg == "file"), PARAM.inp.deepks_equiv, *orb_.Alpha);
+        DeePKS_domain::read_pdm((PARAM.inp.init_chg == "file"),
+                                PARAM.inp.deepks_equiv,
+                                GlobalC::ld.init_pdm,
+                                GlobalC::ld.inlmax,
+                                GlobalC::ld.lmaxd,
+                                GlobalC::ld.inl_l,
+                                *orb_.Alpha,
+                                GlobalC::ld.pdm);
     }
 #endif
 
diff --git a/source/module_hamilt_lcao/hamilt_lcaodft/FORCE_gamma.cpp b/source/module_hamilt_lcao/hamilt_lcaodft/FORCE_gamma.cpp
@@ -252,8 +252,6 @@ void Force_LCAO<double>::ftable(const bool isforce,
     if (PARAM.inp.deepks_scf)
     {
         // when deepks_scf is on, the init pdm should be same as the out pdm, so we should not recalculate the pdm
-        // GlobalC::ld.cal_projected_DM(dm, ucell, orb, gd);
-
         DeePKS_domain::cal_descriptor(ucell.nat,
                                       GlobalC::ld.inlmax,
                                       GlobalC::ld.inl_l,
diff --git a/source/module_hamilt_lcao/hamilt_lcaodft/FORCE_k.cpp b/source/module_hamilt_lcao/hamilt_lcaodft/FORCE_k.cpp
@@ -347,8 +347,6 @@ void Force_LCAO<std::complex<double>>::ftable(const bool isforce,
         const std::vector<std::vector<std::complex<double>>>& dm_k = dm->get_DMK_vector();
 
         // when deepks_scf is on, the init pdm should be same as the out pdm, so we should not recalculate the pdm
-        // GlobalC::ld.cal_projected_DM(dm, ucell, orb, gd);
-
         std::vector<torch::Tensor> descriptor;
         DeePKS_domain::cal_descriptor(ucell.nat,
                                       GlobalC::ld.inlmax,
diff --git a/source/module_hamilt_lcao/hamilt_lcaodft/operator_lcao/deepks_lcao.cpp b/source/module_hamilt_lcao/hamilt_lcaodft/operator_lcao/deepks_lcao.cpp
@@ -160,7 +160,18 @@ void hamilt::DeePKS<hamilt::OperatorLCAO<TK, TR>>::contributeHR()
     {
         ModuleBase::timer::tick("DeePKS", "contributeHR");
 
-        GlobalC::ld.cal_projected_DM<TK>(this->DM, *this->ucell, *ptr_orb_, *(this->gd));
+        DeePKS_domain::cal_pdm<TK>(GlobalC::ld.init_pdm,
+                                   GlobalC::ld.inlmax,
+                                   GlobalC::ld.lmaxd,
+                                   GlobalC::ld.inl_l,
+                                   GlobalC::ld.inl_index,
+                                   this->DM,
+                                   GlobalC::ld.phialpha,
+                                   *this->ucell,
+                                   *ptr_orb_,
+                                   *(this->gd),
+                                   *(this->hR->get_paraV()),
+                                   GlobalC::ld.pdm);
 
         std::vector<torch::Tensor> descriptor;
         DeePKS_domain::cal_descriptor(this->ucell->nat,
diff --git a/source/module_hamilt_lcao/module_deepks/CMakeLists.txt b/source/module_hamilt_lcao/module_deepks/CMakeLists.txt
@@ -11,8 +11,8 @@ if(ENABLE_DEEPKS)
       deepks_vdelta.cpp
       deepks_vdpre.cpp
       deepks_hmat.cpp
+      deepks_pdm.cpp
       LCAO_deepks_io.cpp
-      LCAO_deepks_pdm.cpp
       LCAO_deepks_phialpha.cpp
       LCAO_deepks_interface.cpp
   )
diff --git a/source/module_hamilt_lcao/module_deepks/LCAO_deepks.h b/source/module_hamilt_lcao/module_deepks/LCAO_deepks.h
@@ -10,6 +10,7 @@
 #include "deepks_hmat.h"
 #include "deepks_orbital.h"
 #include "deepks_orbpre.h"
+#include "deepks_pdm.h"
 #include "deepks_spre.h"
 #include "deepks_vdelta.h"
 #include "deepks_vdpre.h"
@@ -20,13 +21,11 @@
 #include "module_basis/module_ao/parallel_orbitals.h"
 #include "module_basis/module_nao/two_center_integrator.h"
 #include "module_cell/module_neighbor/sltk_grid_driver.h"
-#include "module_elecstate/module_dm/density_matrix.h"
 #include "module_hamilt_lcao/module_hcontainer/hcontainer.h"
 #include "module_io/winput.h"
 
 #include <torch/script.h>
 #include <torch/torch.h>
-#include <unordered_map>
 
 ///
 /// The LCAO_Deepks contains subroutines for implementation of the DeePKS method in atomic basis.
@@ -102,7 +101,7 @@ class LCAO_Deepks
     ModuleBase::IntArray* alpha_index; // seems not used in the code
     ModuleBase::IntArray* inl_index;   // caoyu add 2021-05-07
 
-    bool init_pdm = false; // for DeePKS NSCF calculation
+    bool init_pdm = false; // for DeePKS NSCF calculation, set init_pdm to skip the calculation of pdm in SCF iteration
 
     // deep neural network module that provides corrected Hamiltonian term and
     // related derivatives. Used in cal_gedm.
@@ -194,52 +193,6 @@ class LCAO_Deepks
                         const LCAO_Orbitals& orb,
                         const Grid_Driver& GridD);
 
-    //-------------------
-    // LCAO_deepks_pdm.cpp
-    //-------------------
-
-    // This file contains subroutines for calculating pdm,
-    // which is defind as sum_mu,nu rho_mu,nu (<chi_mu|alpha><alpha|chi_nu>);
-    // as well as gdmx, which is the gradient of pdm, defined as
-    // sum_mu,nu rho_mu,nu d/dX(<chi_mu|alpha><alpha|chi_nu>)
-
-    // It also contains subroutines for printing pdm and gdmx
-    // for checking purpose
-
-    // There are 2 subroutines in this file:
-    // 1. cal_projected_DM, which is used for calculating pdm
-    // 2. check_projected_dm, which prints pdm to descriptor.dat
-
-  public:
-    /**
-     * @brief calculate projected density matrix:
-     * pdm = sum_i,occ <phi_i|alpha1><alpha2|phi_k>
-     * 3 cases to skip calculation of pdm:
-     *    1. NSCF calculation of DeePKS, init_chg = file and pdm has been read
-     *    2. SCF calculation of DeePKS with init_chg = file and pdm has been read for restarting SCF
-     *    3. Relax/Cell-Relax/MD calculation, non-first step will use the convergence pdm from the last step as initial
-     * pdm
-     */
-    template <typename TK>
-    void cal_projected_DM(const elecstate::DensityMatrix<TK, double>* dm,
-                          const UnitCell& ucell,
-                          const LCAO_Orbitals& orb,
-                          const Grid_Driver& GridD);
-
-    void check_projected_dm();
-
-    /**
-     * @brief set init_pdm to skip the calculation of pdm in SCF iteration
-     */
-    void set_init_pdm(bool ipdm)
-    {
-        this->init_pdm = ipdm;
-    }
-    /**
-     * @brief read pdm from file, do it only once in whole calculation
-     */
-    void read_projected_DM(bool read_pdm_file, bool is_equiv, const Numerical_Orbital& alpha);
-
   public:
     //! a temporary interface for cal_e_delta_band
     template <typename TK>
diff --git a/source/module_hamilt_lcao/module_deepks/LCAO_deepks_interface.cpp b/source/module_hamilt_lcao/module_deepks/LCAO_deepks_interface.cpp
@@ -40,8 +40,8 @@ void LCAO_Deepks_Interface<TK, TR>::out_deepks_labels(const double& etot,
     const int des_per_atom = ld->des_per_atom;
     const int* inl_l = ld->inl_l;
     const ModuleBase::IntArray* inl_index = ld->inl_index;
-    const std::vector<torch::Tensor> pdm = ld->pdm;
     const std::vector<hamilt::HContainer<double>*> phialpha = ld->phialpha;
+    std::vector<torch::Tensor> pdm = ld->pdm;
 
     const int my_rank = GlobalV::MY_RANK;
     const int nspin = PARAM.inp.nspin;
@@ -348,10 +348,11 @@ void LCAO_Deepks_Interface<TK, TR>::out_deepks_labels(const double& etot,
         // when deepks_scf is on, the init pdm should be same as the out pdm, so we should not recalculate the pdm
         if (!PARAM.inp.deepks_scf)
         {
-            ld->cal_projected_DM<TK>(dm, ucell, orb, GridD);
+            DeePKS_domain::cal_pdm<
+                TK>(ld->init_pdm, inlmax, lmaxd, inl_l, inl_index, dm, phialpha, ucell, orb, GridD, *ParaV, pdm);
         }
 
-        ld->check_projected_dm(); // print out the projected dm for NSCF calculaiton
+        DeePKS_domain::check_pdm(inlmax, inl_l, pdm); // print out the projected dm for NSCF calculaiton
 
         std::vector<torch::Tensor> descriptor;
         DeePKS_domain::cal_descriptor(nat, inlmax, inl_l, pdm, descriptor,
diff --git a/source/module_hamilt_lcao/module_deepks/LCAO_deepks_interface.h b/source/module_hamilt_lcao/module_deepks/LCAO_deepks_interface.h
@@ -6,6 +6,7 @@
 #include "module_base/complexmatrix.h"
 #include "module_base/matrix.h"
 #include "module_hamilt_lcao/hamilt_lcaodft/hamilt_lcao.h"
+
 #include <memory>
 
 template <typename TK, typename TR>
@@ -26,11 +27,9 @@ class LCAO_Deepks_Interface
     /// @param[in] orb
     /// @param[in] GridD
     /// @param[in] ParaV
-    /// @param[in] psi
     /// @param[in] psid
-    /// @param[in] dm_gamma
-    /// @param[in] dm_k
-    // for Gamma-only
+    /// @param[in] dm
+    /// @param[in] p_ham
     void out_deepks_labels(const double& etot,
                            const int& nks,
                            const int& nat,
diff --git a/source/module_hamilt_lcao/module_deepks/deepks_pdm.cpp b/source/module_hamilt_lcao/module_deepks/deepks_pdm.cpp
diff --git a/source/module_hamilt_lcao/module_deepks/deepks_pdm.h b/source/module_hamilt_lcao/module_deepks/deepks_pdm.h
diff --git a/source/module_hamilt_lcao/module_deepks/test/LCAO_deepks_test.cpp b/source/module_hamilt_lcao/module_deepks/test/LCAO_deepks_test.cpp
diff --git a/source/module_hamilt_lcao/module_deepks/test/Makefile.Objects b/source/module_hamilt_lcao/module_deepks/test/Makefile.Objects

Original file line number	Diff line number	Diff line change
`@@ -252,8 +252,6 @@ void Force_LCAO<double>::ftable(const bool isforce,`
`252`	`252`	`if (PARAM.inp.deepks_scf)`
`253`	`253`	`{`
`254`	`254`	`// when deepks_scf is on, the init pdm should be same as the out pdm, so we should not recalculate the pdm`
`255`		`- // GlobalC::ld.cal_projected_DM(dm, ucell, orb, gd);`
`256`		`-`
`257`	`255`	`DeePKS_domain::cal_descriptor(ucell.nat,`
`258`	`256`	`GlobalC::ld.inlmax,`
`259`	`257`	`GlobalC::ld.inl_l,`