Refactor: Remove some redundant variables and global dependence in DeePKS. (#5791)

* Remove cal_orbital_precalc() and corresponding temporary variables from LCAO_Deepks; Partially change cal_gevdm() for future refactor; Remove the double pointer pdm and use tensor vector for replacement.

* clang-format adjustment.

Author: ErjieWu

source/Makefile.Objects

@@ -200,7 +200,7 @@ OBJS_DEEPKS=LCAO_deepks.o\
         LCAO_deepks_vdelta.o\
         deepks_hmat.o\
         LCAO_deepks_interface.o\
-        orbital_precalc.o\
+        deepks_orbpre.o\
         cal_gdmx.o\
         cal_gedm.o\
         cal_gvx.o\

source/module_hamilt_lcao/hamilt_lcaodft/FORCE_STRESS.cpp

@@ -7,12 +7,12 @@
 // new
 #include "module_base/timer.h"
 #include "module_cell/module_neighbor/sltk_grid_driver.h"
+#include "module_elecstate/elecstate_lcao.h"
 #include "module_elecstate/potentials/efield.h"           // liuyu add 2022-05-18
 #include "module_elecstate/potentials/gatefield.h"        // liuyu add 2022-09-13
 #include "module_hamilt_general/module_surchem/surchem.h" //sunml add 2022-08-10
 #include "module_hamilt_general/module_vdw/vdw.h"
 #include "module_parameter/parameter.h"
-#include "module_elecstate/elecstate_lcao.h"
 #ifdef __DEEPKS
 #include "module_hamilt_lcao/module_deepks/LCAO_deepks.h"    //caoyu add for deepks 2021-06-03
 #include "module_hamilt_lcao/module_deepks/LCAO_deepks_io.h" // mohan add 2024-07-22
@@ -540,7 +540,9 @@ void Force_Stress_LCAO<T>::getForceStress(UnitCell& ucell,
                     {
                         GlobalC::ld.check_gdmx(ucell.nat);
                     }
-                    GlobalC::ld.cal_gvx(ucell.nat);
+                    std::vector<torch::Tensor> gevdm;
+                    GlobalC::ld.cal_gevdm(ucell.nat, gevdm);
+                    GlobalC::ld.cal_gvx(ucell.nat, gevdm);
 
                     if (PARAM.inp.deepks_out_unittest)
                     {
@@ -758,7 +760,9 @@ void Force_Stress_LCAO<T>::getForceStress(UnitCell& ucell,
 
                 if (!PARAM.inp.deepks_equiv) // training with stress label not supported by equivariant version now
                 {
-                    GlobalC::ld.cal_gvepsl(ucell.nat);
+                    std::vector<torch::Tensor> gevdm;
+                    GlobalC::ld.cal_gevdm(ucell.nat, gevdm);
+                    GlobalC::ld.cal_gvepsl(ucell.nat, gevdm);
 
                     LCAO_deepks_io::save_npy_gvepsl(ucell.nat,
                                                     GlobalC::ld.des_per_atom,

source/module_hamilt_lcao/module_deepks/CMakeLists.txt

@@ -11,7 +11,7 @@ if(ENABLE_DEEPKS)
       LCAO_deepks_vdelta.cpp
       deepks_hmat.cpp
       LCAO_deepks_interface.cpp
-      orbital_precalc.cpp
+      deepks_orbpre.cpp
       cal_gdmx.cpp
       cal_gedm.cpp
       cal_gvx.cpp      

source/module_hamilt_lcao/module_deepks/LCAO_deepks.cpp

@@ -45,12 +45,7 @@ LCAO_Deepks::~LCAO_Deepks()
     delete[] inl_l;
 
     //=======1. to use deepks, pdm is required==========
-    // delete pdm**
-    for (int inl = 0; inl < this->inlmax; inl++)
-    {
-        delete[] pdm[inl];
-    }
-    delete[] pdm;
+    pdm.clear();
     //=======2. "deepks_scf" part==========
     // if (PARAM.inp.deepks_scf)
     if (gedm)
@@ -100,12 +95,33 @@ void LCAO_Deepks::init(const LCAO_Orbitals& orb,
 
     int pdm_size = 0;
     this->inlmax = tot_inl;
+    this->pdm.resize(this->inlmax);
+
+    // cal n(descriptor) per atom , related to Lmax, nchi(L) and m. (not total_nchi!)
+    if (!PARAM.inp.deepks_equiv)
+    {
+        this->des_per_atom = 0; // mohan add 2021-04-21
+        for (int l = 0; l <= this->lmaxd; l++)
+        {
+            this->des_per_atom += orb.Alpha[0].getNchi(l) * (2 * l + 1);
+        }
+        this->n_descriptor = nat * this->des_per_atom;
+
+        this->init_index(ntype, nat, na, tot_inl, orb);
+    }
+
     if (!PARAM.inp.deepks_equiv)
     {
         GlobalV::ofs_running << " total basis (all atoms) for descriptor = " << std::endl;
 
-        // init pdm**
-        pdm_size = (this->lmaxd * 2 + 1) * (this->lmaxd * 2 + 1);
+        // init pdm
+        for (int inl = 0; inl < this->inlmax; inl++)
+        {
+            int nm = 2 * inl_l[inl] + 1;
+            pdm_size += nm * nm;
+            this->pdm[inl] = torch::zeros({nm, nm}, torch::kFloat64);
+            // this->pdm[inl].requires_grad_(true);
+        }
     }
     else
     {
@@ -116,26 +132,10 @@ void LCAO_Deepks::init(const LCAO_Orbitals& orb,
         pdm_size = pdm_size * pdm_size;
         this->des_per_atom = pdm_size;
         GlobalV::ofs_running << " Equivariant version, size of pdm matrices : " << pdm_size << std::endl;
-    }
-
-    this->pdm = new double*[this->inlmax];
-    for (int inl = 0; inl < this->inlmax; inl++)
-    {
-        this->pdm[inl] = new double[pdm_size];
-        ModuleBase::GlobalFunc::ZEROS(this->pdm[inl], pdm_size);
-    }
-
-    // cal n(descriptor) per atom , related to Lmax, nchi(L) and m. (not total_nchi!)
-    if (!PARAM.inp.deepks_equiv)
-    {
-        this->des_per_atom = 0; // mohan add 2021-04-21
-        for (int l = 0; l <= this->lmaxd; l++)
+        for (int inl = 0; inl < this->inlmax; inl++)
         {
-            this->des_per_atom += orb.Alpha[0].getNchi(l) * (2 * l + 1);
+            this->pdm[inl] = torch::zeros({pdm_size}, torch::kFloat64);
         }
-        this->n_descriptor = nat * this->des_per_atom;
-
-        this->init_index(ntype, nat, na, tot_inl, orb);
     }
 
     this->pv = &pv_in;
@@ -343,41 +343,6 @@ void LCAO_Deepks::allocate_V_delta(const int nat, const int nks)
     return;
 }
 
-void LCAO_Deepks::init_orbital_pdm_shell(const int nks)
-{
-
-    this->orbital_pdm_shell = new double**[nks];
-
-    for (int iks = 0; iks < nks; iks++)
-    {
-        this->orbital_pdm_shell[iks] = new double*[this->inlmax];
-        for (int inl = 0; inl < this->inlmax; inl++)
-        {
-            this->orbital_pdm_shell[iks][inl] = new double[(2 * this->lmaxd + 1) * (2 * this->lmaxd + 1)];
-            ModuleBase::GlobalFunc::ZEROS(orbital_pdm_shell[iks][inl],
-                                            (2 * this->lmaxd + 1) * (2 * this->lmaxd + 1));
-        }
-        
-    }
-
-    return;
-}
-
-void LCAO_Deepks::del_orbital_pdm_shell(const int nks)
-{
-    for (int iks = 0; iks < nks; iks++)
-    {
-        for (int inl = 0; inl < this->inlmax; inl++)
-        {
-            delete[] this->orbital_pdm_shell[iks][inl];
-        }
-        delete[] this->orbital_pdm_shell[iks];
-    }
-    delete[] this->orbital_pdm_shell;
-
-    return;
-}
-
 void LCAO_Deepks::init_v_delta_pdm_shell(const int nks, const int nlocal)
 {
     const int mn_size = (2 * this->lmaxd + 1) * (2 * this->lmaxd + 1);

source/module_hamilt_lcao/module_deepks/LCAO_deepks.h

@@ -6,6 +6,7 @@
 #include "deepks_force.h"
 #include "deepks_hmat.h"
 #include "deepks_orbital.h"
+#include "deepks_orbpre.h"
 #include "module_base/complexmatrix.h"
 #include "module_base/intarray.h"
 #include "module_base/matrix.h"
@@ -57,12 +58,6 @@ class LCAO_Deepks
     /// Correction term to Hamiltonian, for multi-k
     std::vector<std::vector<std::complex<double>>> H_V_delta_k;
 
-    // k index of HOMO for multi-k bandgap label. QO added 2022-01-24
-    int h_ind = 0;
-
-    // k index of LUMO for multi-k bandgap label. QO added 2022-01-24
-    int l_ind = 0;
-
     // functions for hr status: 1. get value; 2. set value;
     int get_hr_cal()
     {
@@ -109,8 +104,9 @@ class LCAO_Deepks
     std::vector<hamilt::HContainer<double>*> phialpha;
 
     // projected density matrix
-    double** pdm; //[tot_Inl][2l+1][2l+1]	caoyu modified 2021-05-07; if equivariant version: [nat][nlm*nlm]
-    std::vector<torch::Tensor> pdm_tensor;
+    // [tot_Inl][2l+1][2l+1], here l is corresponding to inl;
+    // [nat][nlm*nlm] for equivariant version
+    std::vector<torch::Tensor> pdm;
 
     // descriptors
     std::vector<torch::Tensor> d_tensor;
@@ -138,17 +134,9 @@ class LCAO_Deepks
     // gvx:d(d)/dX, [natom][3][natom][des_per_atom]
     torch::Tensor gvx_tensor;
 
-    // d(d)/dD, autograd from torch::linalg::eigh
-    std::vector<torch::Tensor> gevdm_vector;
-
     // dD/dX, tensor form of gdmx
     std::vector<torch::Tensor> gdmr_vector;
 
-    // orbital_pdm_shell:[Inl,nm*nm]; \langle \phi_\mu|\alpha\rangle\langle\alpha|\phi_\nu\ranlge
-    double*** orbital_pdm_shell;
-    // orbital_precalc:[1,NAt,NDscrpt]; gvdm*orbital_pdm_shell
-    torch::Tensor orbital_precalc_tensor;
-
     // v_delta_pdm_shell[nks,nlocal,nlocal,Inl,nm*nm] = overlap * overlap
     double***** v_delta_pdm_shell;
     std::complex<double>***** v_delta_pdm_shell_complex; // for multi-k
@@ -223,12 +211,6 @@ class LCAO_Deepks
   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);
-    // data structure that saves <phi|alpha>
-    void allocate_nlm(const int nat);
-
-    // for bandgap label calculation; QO added on 2022-1-7
-    void init_orbital_pdm_shell(const int nks);
-    void del_orbital_pdm_shell(const int nks);
 
     // for v_delta label calculation; xinyuan added on 2023-2-22
     void init_v_delta_pdm_shell(const int nks, const int nlocal);
@@ -373,16 +355,16 @@ class LCAO_Deepks
     //       descriptors wrt strain tensor, calculated by
     //       d(des)/d\epsilon_{ab} = d(pdm)/d\epsilon_{ab} * d(des)/d(pdm) = gdm_epsl * gvdm
     //       using einsum
-    // 6. cal_gvdm : d(des)/d(pdm)
+    // 6. cal_gevdm : d(des)/d(pdm)
     //       calculated using torch::autograd::grad
     // 7. load_model : loads model for applying V_delta
     // 8. cal_gedm : calculates d(E_delta)/d(pdm)
     //       this is the term V(D) that enters the expression H_V_delta = |alpha>V(D)<alpha|
     //       caculated using torch::autograd::grad
     // 9. check_gedm : prints gedm for checking
     // 10. cal_orbital_precalc : orbital_precalc is usted for training with orbital label,
-    //                          which equals gvdm * orbital_pdm_shell,
-    //                          orbital_pdm_shell[Inl,nm*nm] = dm_hl * overlap * overlap
+    //                          which equals gvdm * orbital_pdm,
+    //                          orbital_pdm[nks,Inl,nm,nm] = dm_hl * overlap * overlap
     // 11. cal_v_delta_precalc : v_delta_precalc is used for training with v_delta label,
     //                         which equals gvdm * v_delta_pdm_shell,
     //                         v_delta_pdm_shell = overlap * overlap
@@ -408,11 +390,11 @@ class LCAO_Deepks
     ///  - b: the atoms whose force being calculated)
     /// gvdm*gdmx->gvx
     ///----------------------------------------------------
-    void cal_gvx(const int nat);
+    void cal_gvx(const int nat, const std::vector<torch::Tensor>& gevdm);
     void check_gvx(const int nat);
 
     // for stress
-    void cal_gvepsl(const int nat);
+    void cal_gvepsl(const int nat, const std::vector<torch::Tensor>& gevdm);
 
     // load the trained neural network model
     void load_model(const std::string& model_file);
@@ -423,20 +405,22 @@ class LCAO_Deepks
     void cal_gedm_equiv(const int nat);
 
     // calculates orbital_precalc
-    template <typename TK, typename TH>
-    void cal_orbital_precalc(const std::vector<TH>& dm_hl,
-                             const int lmaxd,
-                             const int inlmax,
-                             const int nat,
-                             const int nks,
-                             const int* inl_l,
-                             const std::vector<ModuleBase::Vector3<double>>& kvec_d,
-                             const std::vector<hamilt::HContainer<double>*> phialpha,
-                             const ModuleBase::IntArray* inl_index,
-                             const UnitCell& ucell,
-                             const LCAO_Orbitals& orb,
-                             const Parallel_Orbitals& pv,
-                             const Grid_Driver& GridD);
+    // template <typename TK, typename TH>
+    // void cal_orbital_precalc(const std::vector<TH>& dm_hl,
+    //                          const int lmaxd,
+    //                          const int inlmax,
+    //                          const int nat,
+    //                          const int nks,
+    //                          const int* inl_l,
+    //                          const std::vector<ModuleBase::Vector3<double>>& kvec_d,
+    //                          const std::vector<hamilt::HContainer<double>*> phialpha,
+    //                          const std::vector<torch::Tensor> gevdm,
+    //                          const ModuleBase::IntArray* inl_index,
+    //                          const UnitCell& ucell,
+    //                          const LCAO_Orbitals& orb,
+    //                          const Parallel_Orbitals& pv,
+    //                          const Grid_Driver& GridD,
+    //                          torch::Tensor& orbital_precalc);
 
     // calculates v_delta_precalc
     template <typename TK>
@@ -466,11 +450,11 @@ class LCAO_Deepks
 
     // prepare gevdm for outputting npy file
     void prepare_gevdm(const int nat, const LCAO_Orbitals& orb);
+    void cal_gevdm(const int nat, std::vector<torch::Tensor>& gevdm);
     void check_vdp_gevdm(const int nat);
 
   private:
     const Parallel_Orbitals* pv;
-    void cal_gvdm(const int nat);
 
 #ifdef __MPI
 

source/module_hamilt_lcao/module_deepks/LCAO_deepks_interface.cpp

@@ -113,14 +113,32 @@ void LCAO_Deepks_Interface<TK, TR>::out_deepks_labels(const double& etot,
 
                 std::vector<double> o_delta(nks, 0.0);
 
-                ld->cal_orbital_precalc<TK, TH>(dm_bandgap, ld->lmaxd, ld->inlmax, nat, nks, ld->inl_l, kvec_d, ld->phialpha, ld->inl_index, ucell, orb, *ParaV, GridD);
+                // calculate and save orbital_precalc: [nks,NAt,NDscrpt]
+                torch::Tensor orbital_precalc;
+                std::vector<torch::Tensor> gevdm;
+                ld->cal_gevdm(nat, gevdm);
+                DeePKS_domain::cal_orbital_precalc<TK, TH>(dm_bandgap,
+                                                           ld->lmaxd,
+                                                           ld->inlmax,
+                                                           nat,
+                                                           nks,
+                                                           ld->inl_l,
+                                                           kvec_d,
+                                                           ld->phialpha,
+                                                           gevdm,
+                                                           ld->inl_index,
+                                                           ucell,
+                                                           orb,
+                                                           *ParaV,
+                                                           GridD,
+                                                           orbital_precalc);
                 DeePKS_domain::cal_o_delta<TK, TH>(dm_bandgap, *h_delta, o_delta, *ParaV, nks);
 
                 // save obase and orbital_precalc
                 LCAO_deepks_io::save_npy_orbital_precalc(nat,
                                                          nks,
                                                          ld->des_per_atom,
-                                                         ld->orbital_precalc_tensor,
+                                                         orbital_precalc,
                                                          PARAM.globalv.global_out_dir,
                                                          my_rank);
                 const std::string file_obase = PARAM.globalv.global_out_dir + "deepks_obase.npy";
@@ -135,8 +153,8 @@ void LCAO_Deepks_Interface<TK, TR>::out_deepks_labels(const double& etot,
             {
                 const std::string file_obase = PARAM.globalv.global_out_dir + "deepks_obase.npy";
                 LCAO_deepks_io::save_npy_o(o_tot, file_obase, nks, my_rank); // no scf, o_tot=o_base
-            }                                                                       // end deepks_scf == 0
-        }                                                                           // end bandgap label
+            }                                                                // end deepks_scf == 0
+        }                                                                    // end bandgap label
 
         // save H(R) matrix
         if (true) // should be modified later!