Remove the global dependence of functions related to phialpha in DeePKS.

ErjieWu · ErjieWu · commit 56194a93d69c · 2025-01-08T16:09:29.000+08:00
diff --git a/source/Makefile.Objects b/source/Makefile.Objects
@@ -202,8 +202,8 @@ OBJS_DEEPKS=LCAO_deepks.o\
         deepks_vdpre.o\
         deepks_hmat.o\
         deepks_pdm.o\
+        deepks_phialpha.o\
         LCAO_deepks_io.o\
-        LCAO_deepks_phialpha.o\
         LCAO_deepks_interface.o\
         
 
diff --git a/source/module_esolver/lcao_before_scf.cpp b/source/module_esolver/lcao_before_scf.cpp
@@ -211,13 +211,19 @@ void ESolver_KS_LCAO<TK, TR>::before_scf(UnitCell& ucell, const int istep)
     {
         const Parallel_Orbitals* pv = &this->pv;
         // allocate <phi(0)|alpha(R)>, phialpha is different every ion step, so it is allocated here
-        GlobalC::ld.allocate_phialpha(PARAM.inp.cal_force, ucell, orb_, this->gd);
+        DeePKS_domain::allocate_phialpha(PARAM.inp.cal_force, ucell, orb_, this->gd, pv, GlobalC::ld.phialpha);
         // build and save <phi(0)|alpha(R)> at beginning
-        GlobalC::ld.build_phialpha(PARAM.inp.cal_force, ucell, orb_, this->gd, *(two_center_bundle_.overlap_orb_alpha));
+        DeePKS_domain::build_phialpha(PARAM.inp.cal_force,
+                                      ucell,
+                                      orb_,
+                                      this->gd,
+                                      pv,
+                                      *(two_center_bundle_.overlap_orb_alpha),
+                                      GlobalC::ld.phialpha);
 
         if (PARAM.inp.deepks_out_unittest)
         {
-            GlobalC::ld.check_phialpha(PARAM.inp.cal_force, ucell, orb_, this->gd);
+            DeePKS_domain::check_phialpha(PARAM.inp.cal_force, ucell, orb_, this->gd, pv, GlobalC::ld.phialpha);
         }
     }
 #endif
diff --git a/source/module_esolver/lcao_others.cpp b/source/module_esolver/lcao_others.cpp
@@ -217,13 +217,19 @@ void ESolver_KS_LCAO<TK, TR>::others(UnitCell& ucell, const int istep)
     {
         const Parallel_Orbitals* pv = &this->pv;
         // allocate <phi(0)|alpha(R)>, phialpha is different every ion step, so it is allocated here
-        GlobalC::ld.allocate_phialpha(PARAM.inp.cal_force, ucell, orb_, this->gd);
+        DeePKS_domain::allocate_phialpha(PARAM.inp.cal_force, ucell, orb_, this->gd, pv, GlobalC::ld.phialpha);
         // build and save <phi(0)|alpha(R)> at beginning
-        GlobalC::ld.build_phialpha(PARAM.inp.cal_force, ucell, orb_, this->gd, *(two_center_bundle_.overlap_orb_alpha));
+        DeePKS_domain::build_phialpha(PARAM.inp.cal_force,
+                                      ucell,
+                                      orb_,
+                                      this->gd,
+                                      pv,
+                                      *(two_center_bundle_.overlap_orb_alpha),
+                                      GlobalC::ld.phialpha);
 
         if (PARAM.inp.deepks_out_unittest)
         {
-            GlobalC::ld.check_phialpha(PARAM.inp.cal_force, ucell, orb_, this->gd);
+            DeePKS_domain::check_phialpha(PARAM.inp.cal_force, ucell, orb_, this->gd, pv, GlobalC::ld.phialpha);
         }
     }
 #endif
diff --git a/source/module_hamilt_lcao/module_deepks/CMakeLists.txt b/source/module_hamilt_lcao/module_deepks/CMakeLists.txt
@@ -12,8 +12,8 @@ if(ENABLE_DEEPKS)
       deepks_vdpre.cpp
       deepks_hmat.cpp
       deepks_pdm.cpp
+      deepks_phialpha.cpp
       LCAO_deepks_io.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
@@ -11,6 +11,7 @@
 #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"
@@ -158,47 +159,14 @@ class LCAO_Deepks
     /// Allocate memory for correction to Hamiltonian
     void allocate_V_delta(const int nat, const int nks = 1);
 
-  private:
-    // arrange index of descriptor in all atoms
-    void init_index(const int ntype, const int nat, std::vector<int> na, const int tot_inl, const LCAO_Orbitals& orb);
-
-    //-------------------
-    // LCAO_deepks_phialpha.cpp
-    //-------------------
-
-    // E.Wu 2024-12-24
-    // This file contains 3 subroutines:
-    // 1. allocate_phialpha, which allocates memory for phialpha
-    // 2. build_phialpha, which calculates the overlap
-    // between atomic basis and projector alpha : <phi_mu|alpha>
-    // which will be used in calculating pdm, gdmx, H_V_delta, F_delta;
-    // 3. check_phialpha, which prints the results into .dat files
-    // for checking
-
-  public:
-    // calculates <chi|alpha>
-    void allocate_phialpha(const bool& cal_deri,
-                           const UnitCell& ucell,
-                           const LCAO_Orbitals& orb,
-                           const Grid_Driver& GridD);
-
-    void build_phialpha(const bool& cal_deri /**< [in] 0 for 2-center intergration, 1 for its derivation*/,
-                        const UnitCell& ucell,
-                        const LCAO_Orbitals& orb,
-                        const Grid_Driver& GridD,
-                        const TwoCenterIntegrator& overlap_orb_alpha);
-
-    void check_phialpha(const bool& cal_deri /**< [in] 0 for 2-center intergration, 1 for its derivation*/,
-                        const UnitCell& ucell,
-                        const LCAO_Orbitals& orb,
-                        const Grid_Driver& GridD);
-
-  public:
     //! a temporary interface for cal_e_delta_band
     template <typename TK>
     void dpks_cal_e_delta_band(const std::vector<std::vector<TK>>& dm, const int nks);
 
   private:
+    // arrange index of descriptor in all atoms
+    void init_index(const int ntype, const int nat, std::vector<int> na, const int tot_inl, const LCAO_Orbitals& orb);
+
     const Parallel_Orbitals* pv;
 };
 
diff --git a/source/module_hamilt_lcao/module_deepks/deepks_phialpha.cpp b/source/module_hamilt_lcao/module_deepks/deepks_phialpha.cpp
@@ -8,21 +8,24 @@
 
 #ifdef __DEEPKS
 
-#include "LCAO_deepks.h"
+#include "deepks_phialpha.h"
+
 #include "module_base/timer.h"
 #include "module_base/vector3.h"
 #include "module_parameter/parameter.h"
 
-void LCAO_Deepks::allocate_phialpha(const bool& cal_deri,
-                                    const UnitCell& ucell,
-                                    const LCAO_Orbitals& orb,
-                                    const Grid_Driver& GridD)
+void DeePKS_domain::allocate_phialpha(const bool& cal_deri,
+                                      const UnitCell& ucell,
+                                      const LCAO_Orbitals& orb,
+                                      const Grid_Driver& GridD,
+                                      const Parallel_Orbitals* pv,
+                                      std::vector<hamilt::HContainer<double>*>& phialpha)
 {
-    ModuleBase::TITLE("LCAO_Deepks", "allocate_phialpha");
+    ModuleBase::TITLE("DeePKS_domain", "allocate_phialpha");
 
-    this->phialpha.resize(cal_deri ? 4 : 1);
+    phialpha.resize(cal_deri ? 4 : 1);
 
-    this->phialpha[0] = new hamilt::HContainer<double>(pv); // phialpha is always real
+    phialpha[0] = new hamilt::HContainer<double>(pv); // phialpha is always real
     // Do not use fix_gamma, since it may find wrong matrix for gamma-only case in DeePKS
 
     // cutoff for alpha is same for all types of atoms
@@ -63,31 +66,33 @@ void LCAO_Deepks::allocate_phialpha(const bool& cal_deri,
                 hamilt::AtomPair<double> pair(iat, ibt, R_index, pv);
                 // Notice: in AtomPair, the usage is set_size(ncol, nrow)
                 pair.set_size(nw_alpha, atom1->nw * PARAM.globalv.npol);
-                this->phialpha[0]->insert_pair(pair);
+                phialpha[0]->insert_pair(pair);
             }
         }
     }
 
-    this->phialpha[0]->allocate(nullptr, true);
+    phialpha[0]->allocate(nullptr, true);
     // whether to calculate the derivative of phialpha
     if (cal_deri)
     {
         for (int i = 1; i < 4; ++i)
         {
-            this->phialpha[i] = new hamilt::HContainer<double>(*this->phialpha[0], nullptr); // copy constructor
+            phialpha[i] = new hamilt::HContainer<double>(*phialpha[0], nullptr); // copy constructor
         }
     }
     return;
 }
 
-void LCAO_Deepks::build_phialpha(const bool& cal_deri,
-                                 const UnitCell& ucell,
-                                 const LCAO_Orbitals& orb,
-                                 const Grid_Driver& GridD,
-                                 const TwoCenterIntegrator& overlap_orb_alpha)
+void DeePKS_domain::build_phialpha(const bool& cal_deri,
+                                   const UnitCell& ucell,
+                                   const LCAO_Orbitals& orb,
+                                   const Grid_Driver& GridD,
+                                   const Parallel_Orbitals* pv,
+                                   const TwoCenterIntegrator& overlap_orb_alpha,
+                                   std::vector<hamilt::HContainer<double>*>& phialpha)
 {
-    ModuleBase::TITLE("LCAO_Deepks", "build_phialpha");
-    ModuleBase::timer::tick("LCAO_Deepks", "build_phialpha");
+    ModuleBase::TITLE("DeePKS_domain", "build_phialpha");
+    ModuleBase::timer::tick("DeePKS_domain", "build_phialpha");
 
     // cutoff for alpha is same for all types of atoms
     const double Rcut_Alpha = orb.Alpha[0].getRcut();
@@ -126,13 +131,13 @@ void LCAO_Deepks::build_phialpha(const bool& cal_deri,
                     continue;
                 }
 
-                double* data_pointer = this->phialpha[0]->data(iat, ibt, R);
+                double* data_pointer = phialpha[0]->data(iat, ibt, R);
                 std::vector<double*> grad_pointer(3);
                 if (cal_deri)
                 {
                     for (int i = 0; i < 3; ++i)
                     {
-                        grad_pointer[i] = this->phialpha[i + 1]->data(iat, ibt, R);
+                        grad_pointer[i] = phialpha[i + 1]->data(iat, ibt, R);
                     }
                 }
 
@@ -192,17 +197,19 @@ void LCAO_Deepks::build_phialpha(const bool& cal_deri,
         }
     }
 
-    ModuleBase::timer::tick("LCAO_Deepks", "build_phialpha");
+    ModuleBase::timer::tick("DeePKS_domain", "build_phialpha");
     return;
 }
 
-void LCAO_Deepks::check_phialpha(const bool& cal_deri,
-                                 const UnitCell& ucell,
-                                 const LCAO_Orbitals& orb,
-                                 const Grid_Driver& GridD)
+void DeePKS_domain::check_phialpha(const bool& cal_deri,
+                                   const UnitCell& ucell,
+                                   const LCAO_Orbitals& orb,
+                                   const Grid_Driver& GridD,
+                                   const Parallel_Orbitals* pv,
+                                   std::vector<hamilt::HContainer<double>*>& phialpha)
 {
-    ModuleBase::TITLE("LCAO_Deepks", "check_phialpha");
-    ModuleBase::timer::tick("LCAO_Deepks", "check_phialpha");
+    ModuleBase::TITLE("DeePKS_domain", "check_phialpha");
+    ModuleBase::timer::tick("DeePKS_domain", "check_phialpha");
 
     const double Rcut_Alpha = orb.Alpha[0].getRcut();
     // same for all types of atoms
@@ -280,13 +287,13 @@ void LCAO_Deepks::check_phialpha(const bool& cal_deri,
                     ofs_z << "R : " << R[0] << " " << R[1] << " " << R[2] << std::endl;
                 }
 
-                const double* data_pointer = this->phialpha[0]->data(iat, ibt, R);
+                const double* data_pointer = phialpha[0]->data(iat, ibt, R);
                 std::vector<double*> grad_pointer(3, nullptr);
                 if (cal_deri)
                 {
-                    grad_pointer[0] = this->phialpha[1]->data(iat, ibt, R);
-                    grad_pointer[1] = this->phialpha[2]->data(iat, ibt, R);
-                    grad_pointer[2] = this->phialpha[3]->data(iat, ibt, R);
+                    grad_pointer[0] = phialpha[1]->data(iat, ibt, R);
+                    grad_pointer[1] = phialpha[2]->data(iat, ibt, R);
+                    grad_pointer[2] = phialpha[3]->data(iat, ibt, R);
                 }
 
                 for (int iw1 = 0; iw1 < nw1_tot; ++iw1)
@@ -334,7 +341,7 @@ void LCAO_Deepks::check_phialpha(const bool& cal_deri,
         }         // end I0
     }             // end T0
 
-    ModuleBase::timer::tick("LCAO_Deepks", "check_phialpha");
+    ModuleBase::timer::tick("DeePKS_domain", "check_phialpha");
     return;
 }
 
diff --git a/source/module_hamilt_lcao/module_deepks/deepks_phialpha.h b/source/module_hamilt_lcao/module_deepks/deepks_phialpha.h
@@ -0,0 +1,52 @@
+#ifndef DEEPKS_PHIALPHA_H
+#define DEEPKS_PHIALPHA_H
+
+#ifdef __DEEPKS
+
+#include "module_base/complexmatrix.h"
+#include "module_base/matrix.h"
+#include "module_base/timer.h"
+#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_hamilt_lcao/module_hcontainer/hcontainer.h"
+
+#include <torch/script.h>
+#include <torch/torch.h>
+
+namespace DeePKS_domain
+{
+// This file contains 3 subroutines:
+// 1. allocate_phialpha, which allocates memory for phialpha
+// 2. build_phialpha, which calculates the overlap
+// between atomic basis and projector alpha : <phi_mu|alpha>
+// which will be used in calculating pdm, gdmx, H_V_delta, F_delta;
+// 3. check_phialpha, which prints the results into .dat files
+// for checking
+
+// calculates <chi|alpha>
+void allocate_phialpha(const bool& cal_deri,
+                       const UnitCell& ucell,
+                       const LCAO_Orbitals& orb,
+                       const Grid_Driver& GridD,
+                       const Parallel_Orbitals* pv,
+                       std::vector<hamilt::HContainer<double>*>& phialpha);
+
+void build_phialpha(const bool& cal_deri /**< [in] 0 for 2-center intergration, 1 for its derivation*/,
+                    const UnitCell& ucell,
+                    const LCAO_Orbitals& orb,
+                    const Grid_Driver& GridD,
+                    const Parallel_Orbitals* pv,
+                    const TwoCenterIntegrator& overlap_orb_alpha,
+                    std::vector<hamilt::HContainer<double>*>& phialpha);
+
+void check_phialpha(const bool& cal_deri /**< [in] 0 for 2-center intergration, 1 for its derivation*/,
+                    const UnitCell& ucell,
+                    const LCAO_Orbitals& orb,
+                    const Grid_Driver& GridD,
+                    const Parallel_Orbitals* pv,
+                    std::vector<hamilt::HContainer<double>*>& phialpha);
+} // namespace DeePKS_domain
+
+#endif
+#endif
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
@@ -36,11 +36,22 @@ void test_deepks::check_phialpha()
     }
     GlobalC::ld.init(ORB, ucell.nat, ucell.ntype, kv.nkstot, ParaO, na);
 
-    GlobalC::ld.allocate_phialpha(PARAM.input.cal_force, ucell, ORB, Test_Deepks::GridD);
-
-    GlobalC::ld.build_phialpha(PARAM.input.cal_force, ucell, ORB, Test_Deepks::GridD, overlap_orb_alpha_);
-
-    GlobalC::ld.check_phialpha(PARAM.input.cal_force, ucell, ORB, Test_Deepks::GridD);
+    DeePKS_domain::allocate_phialpha(PARAM.input.cal_force,
+                                     ucell,
+                                     ORB,
+                                     Test_Deepks::GridD,
+                                     &ParaO,
+                                     GlobalC::ld.phialpha);
+
+    DeePKS_domain::build_phialpha(PARAM.input.cal_force,
+                                  ucell,
+                                  ORB,
+                                  Test_Deepks::GridD,
+                                  &ParaO,
+                                  overlap_orb_alpha_,
+                                  GlobalC::ld.phialpha);
+
+    DeePKS_domain::check_phialpha(PARAM.input.cal_force, ucell, ORB, Test_Deepks::GridD, &ParaO, GlobalC::ld.phialpha);
 
     this->compare_with_ref("phialpha.dat", "phialpha_ref.dat");
     this->compare_with_ref("dphialpha_x.dat", "dphialpha_x_ref.dat");
diff --git a/source/module_hamilt_lcao/module_deepks/test/Makefile.Objects b/source/module_hamilt_lcao/module_deepks/test/Makefile.Objects
@@ -11,9 +11,9 @@ deepks_basic.o\
 deepks_force.o\
 deepks_vdelta.o\
 deepks_pdm.o\
+deepks_phialpha.o\
 LCAO_deepks.o\
 LCAO_deepks_io.o\
-LCAO_deepks_phialpha.o\
 
 OBJS_IO=output.o\
 

Original file line number	Diff line number	Diff line change
`@@ -211,13 +211,19 @@ void ESolver_KS_LCAO<TK, TR>::before_scf(UnitCell& ucell, const int istep)`
`211`	`211`	`{`
`212`	`212`	`const Parallel_Orbitals* pv = &this->pv;`
`213`	`213`	`// allocate <phi(0)\|alpha(R)>, phialpha is different every ion step, so it is allocated here`
`214`		`- GlobalC::ld.allocate_phialpha(PARAM.inp.cal_force, ucell, orb_, this->gd);`
	`214`	`+ DeePKS_domain::allocate_phialpha(PARAM.inp.cal_force, ucell, orb_, this->gd, pv, GlobalC::ld.phialpha);`
`215`	`215`	`// build and save <phi(0)\|alpha(R)> at beginning`
`216`		`- GlobalC::ld.build_phialpha(PARAM.inp.cal_force, ucell, orb_, this->gd, *(two_center_bundle_.overlap_orb_alpha));`
	`216`	`+ DeePKS_domain::build_phialpha(PARAM.inp.cal_force,`
	`217`	`+ ucell,`
	`218`	`+ orb_,`
	`219`	`+ this->gd,`
	`220`	`+ pv,`
	`221`	`+ *(two_center_bundle_.overlap_orb_alpha),`
	`222`	`+ GlobalC::ld.phialpha);`
`217`	`223`
`218`	`224`	`if (PARAM.inp.deepks_out_unittest)`
`219`	`225`	`{`
`220`		`- GlobalC::ld.check_phialpha(PARAM.inp.cal_force, ucell, orb_, this->gd);`
	`226`	`+ DeePKS_domain::check_phialpha(PARAM.inp.cal_force, ucell, orb_, this->gd, pv, GlobalC::ld.phialpha);`
`221`	`227`	`}`
`222`	`228`	`}`
`223`	`229`	`#endif`
Original file line number	Diff line number	Diff line change
`@@ -217,13 +217,19 @@ void ESolver_KS_LCAO<TK, TR>::others(UnitCell& ucell, const int istep)`
`217`	`217`	`{`
`218`	`218`	`const Parallel_Orbitals* pv = &this->pv;`
`219`	`219`	`// allocate <phi(0)\|alpha(R)>, phialpha is different every ion step, so it is allocated here`
`220`		`- GlobalC::ld.allocate_phialpha(PARAM.inp.cal_force, ucell, orb_, this->gd);`
	`220`	`+ DeePKS_domain::allocate_phialpha(PARAM.inp.cal_force, ucell, orb_, this->gd, pv, GlobalC::ld.phialpha);`
`221`	`221`	`// build and save <phi(0)\|alpha(R)> at beginning`
`222`		`- GlobalC::ld.build_phialpha(PARAM.inp.cal_force, ucell, orb_, this->gd, *(two_center_bundle_.overlap_orb_alpha));`
	`222`	`+ DeePKS_domain::build_phialpha(PARAM.inp.cal_force,`
	`223`	`+ ucell,`
	`224`	`+ orb_,`
	`225`	`+ this->gd,`
	`226`	`+ pv,`
	`227`	`+ *(two_center_bundle_.overlap_orb_alpha),`
	`228`	`+ GlobalC::ld.phialpha);`
`223`	`229`
`224`	`230`	`if (PARAM.inp.deepks_out_unittest)`
`225`	`231`	`{`
`226`		`- GlobalC::ld.check_phialpha(PARAM.inp.cal_force, ucell, orb_, this->gd);`
	`232`	`+ DeePKS_domain::check_phialpha(PARAM.inp.cal_force, ucell, orb_, this->gd, pv, GlobalC::ld.phialpha);`
`227`	`233`	`}`
`228`	`234`	`}`
`229`	`235`	`#endif`
Original file line number	Diff line number	Diff line change
`@@ -12,8 +12,8 @@ if(ENABLE_DEEPKS)`
`12`	`12`	`deepks_vdpre.cpp`
`13`	`13`	`deepks_hmat.cpp`
`14`	`14`	`deepks_pdm.cpp`
	`15`	`+ deepks_phialpha.cpp`
`15`	`16`	`LCAO_deepks_io.cpp`
`16`		`- LCAO_deepks_phialpha.cpp`
`17`	`17`	`LCAO_deepks_interface.cpp`
`18`	`18`	`)`
`19`	`19`