Update pdm before outputting DeePKS labels (#5857)

* rename cal_gedm as cal_edelta_gedm

* Update pdm and other properties before outputting deeepks labels

* fix bug of include

* change test ref because of updating pdm

source/module_hamilt_lcao/hamilt_lcaodft/FORCE_gamma.cpp

@@ -251,20 +251,21 @@ void Force_LCAO<double>::ftable(const bool isforce,
     if (PARAM.inp.deepks_scf)
     {
         const std::vector<std::vector<double>>& dm_gamma = dm->get_DMK_vector();
-        std::vector<torch::Tensor> descriptor;
-        // when deepks_scf is on, the init pdm should be same as the out pdm, so we should not recalculate the pdm
-        DeePKS_domain::cal_descriptor(ucell.nat, ld.inlmax, ld.inl_l, ld.pdm, descriptor, ld.des_per_atom);
-        DeePKS_domain::cal_gedm(ucell.nat,
-                                ld.lmaxd,
-                                ld.nmaxd,
-                                ld.inlmax,
-                                ld.des_per_atom,
-                                ld.inl_l,
-                                descriptor,
-                                ld.pdm,
-                                ld.model_deepks,
-                                ld.gedm,
-                                ld.E_delta);
+
+        // These calculations have been done in LCAO_Deepks_Interface in after_scf
+        // std::vector<torch::Tensor> descriptor;
+        // DeePKS_domain::cal_descriptor(ucell.nat, ld.inlmax, ld.inl_l, ld.pdm, descriptor, ld.des_per_atom);
+        // DeePKS_domain::cal_edelta_gedm(ucell.nat,
+        //                         ld.lmaxd,
+        //                         ld.nmaxd,
+        //                         ld.inlmax,
+        //                         ld.des_per_atom,
+        //                         ld.inl_l,
+        //                         descriptor,
+        //                         ld.pdm,
+        //                         ld.model_deepks,
+        //                         ld.gedm,
+        //                         ld.E_delta);
 
         const int nks = 1;
         DeePKS_domain::cal_f_delta<double>(dm_gamma,

source/module_hamilt_lcao/hamilt_lcaodft/FORCE_k.cpp

@@ -346,20 +346,21 @@ void Force_LCAO<std::complex<double>>::ftable(const bool isforce,
     if (PARAM.inp.deepks_scf)
     {
         const std::vector<std::vector<std::complex<double>>>& dm_k = dm->get_DMK_vector();
-        std::vector<torch::Tensor> descriptor;
-        // when deepks_scf is on, the init pdm should be same as the out pdm, so we should not recalculate the pdm
-        DeePKS_domain::cal_descriptor(ucell.nat, ld.inlmax, ld.inl_l, ld.pdm, descriptor, ld.des_per_atom);
-        DeePKS_domain::cal_gedm(ucell.nat,
-                                ld.lmaxd,
-                                ld.nmaxd,
-                                ld.inlmax,
-                                ld.des_per_atom,
-                                ld.inl_l,
-                                descriptor,
-                                ld.pdm,
-                                ld.model_deepks,
-                                ld.gedm,
-                                ld.E_delta);
+
+        // These calculations have been done in LCAO_Deepks_Interface in after_scf
+        // std::vector<torch::Tensor> descriptor;
+        // DeePKS_domain::cal_descriptor(ucell.nat, ld.inlmax, ld.inl_l, ld.pdm, descriptor, ld.des_per_atom);
+        // DeePKS_domain::cal_edelta_gedm(ucell.nat,
+        //                         ld.lmaxd,
+        //                         ld.nmaxd,
+        //                         ld.inlmax,
+        //                         ld.des_per_atom,
+        //                         ld.inl_l,
+        //                         descriptor,
+        //                         ld.pdm,
+        //                         ld.model_deepks,
+        //                         ld.gedm,
+        //                         ld.E_delta);
 
         DeePKS_domain::cal_f_delta<std::complex<double>>(dm_k,
                                                          ucell,

source/module_hamilt_lcao/hamilt_lcaodft/hamilt_lcao.h

@@ -13,6 +13,11 @@
 #include "module_hamilt_lcao/module_hcontainer/hcontainer.h"
 
 #include <vector>
+
+#ifdef __DEEPKS
+#include "module_hamilt_lcao/module_deepks/LCAO_deepks.h"
+#endif
+
 #ifdef __EXX
 #include "module_ri/Exx_LRI.h"
 #endif

source/module_hamilt_lcao/hamilt_lcaodft/operator_lcao/deepks_lcao.cpp

@@ -186,7 +186,7 @@ void hamilt::DeePKS<hamilt::OperatorLCAO<TK, TR>>::contributeHR()
                                       this->ld->pdm,
                                       descriptor,
                                       this->ld->des_per_atom);
-        DeePKS_domain::cal_gedm(this->ucell->nat,
+        DeePKS_domain::cal_edelta_gedm(this->ucell->nat,
                                 this->ld->lmaxd,
                                 this->ld->nmaxd,
                                 inlmax,

source/module_hamilt_lcao/module_deepks/LCAO_deepks.h

@@ -80,7 +80,7 @@ class LCAO_Deepks
     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.
+    // related derivatives. Used in cal_edelta_gedm.
     torch::jit::script::Module model_deepks;
 
     // saves <phi(0)|alpha(R)> and its derivatives

source/module_hamilt_lcao/module_deepks/LCAO_deepks_interface.cpp

@@ -39,6 +39,7 @@ void LCAO_Deepks_Interface<TK, TR>::out_deepks_labels(const double& etot,
     // These variables are frequently used in the following code
     const int inlmax = orb.Alpha[0].getTotal_nchi() * nat;
     const int lmaxd = orb.get_lmax_d();
+    const int nmaxd = ld->nmaxd;
 
     const int des_per_atom = ld->des_per_atom;
     const int* inl_l = ld->inl_l;
@@ -49,11 +50,56 @@ void LCAO_Deepks_Interface<TK, TR>::out_deepks_labels(const double& etot,
     bool init_pdm = ld->init_pdm;
     double E_delta = ld->E_delta;
     double e_delta_band = ld->e_delta_band;
-    double** gedm = ld->gedm;
 
     const int my_rank = GlobalV::MY_RANK;
     const int nspin = PARAM.inp.nspin;
 
+    // Note : update PDM and all other quantities with the current dm
+    // DeePKS PDM and descriptor
+    if (PARAM.inp.deepks_out_labels || PARAM.inp.deepks_scf)
+    {
+        // this part is for integrated test of deepks
+        // so it is printed no matter even if deepks_out_labels is not used
+        DeePKS_domain::cal_pdm<TK>
+                (init_pdm, inlmax, lmaxd, inl_l, inl_index, dm, phialpha, ucell, orb, GridD, *ParaV, pdm);
+
+        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,
+                                      des_per_atom); // final descriptor
+        DeePKS_domain::check_descriptor(inlmax, des_per_atom, inl_l, ucell, PARAM.globalv.global_out_dir, descriptor);
+
+        if (PARAM.inp.deepks_out_labels)
+        {
+            LCAO_deepks_io::save_npy_d(nat,
+                                       des_per_atom,
+                                       inlmax,
+                                       inl_l,
+                                       PARAM.inp.deepks_equiv,
+                                       descriptor,
+                                       PARAM.globalv.global_out_dir,
+                                       GlobalV::MY_RANK); // libnpy needed
+        }
+
+        if (PARAM.inp.deepks_scf)
+        {
+            // update E_delta and gedm
+            // new gedm is also useful in cal_f_delta, so it should be ld->gedm
+            DeePKS_domain::cal_edelta_gedm(nat,
+                        lmaxd,
+                        nmaxd,
+                        inlmax,
+                        des_per_atom,
+                        inl_l,
+                        descriptor,
+                        pdm,
+                        ld->model_deepks,
+                        ld->gedm,
+                        E_delta);
+        }
+    }
+
     // Used for deepks_bandgap == 1 and deepks_v_delta > 0
     std::vector<std::vector<TK>>* h_delta = nullptr;
     if constexpr (std::is_same<TK, double>::value)
@@ -65,7 +111,7 @@ void LCAO_Deepks_Interface<TK, TR>::out_deepks_labels(const double& etot,
         h_delta = &ld->H_V_delta_k;
     }
 
-    // calculating deepks correction to bandgap and save the results
+    // calculating deepks correction and save the results
     if (PARAM.inp.deepks_out_labels)
     {
         // Used for deepks_scf == 1
@@ -317,38 +363,6 @@ void LCAO_Deepks_Interface<TK, TR>::out_deepks_labels(const double& etot,
 
     } // end deepks_out_labels
 
-    // DeePKS PDM and descriptor
-    if (PARAM.inp.deepks_out_labels || PARAM.inp.deepks_scf)
-    {
-        // this part is for integrated test of deepks
-        // so it is printed no matter even if deepks_out_labels is not used
-        // 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)
-        {
-            DeePKS_domain::cal_pdm<
-                TK>(init_pdm, inlmax, lmaxd, inl_l, inl_index, dm, phialpha, ucell, orb, GridD, *ParaV, pdm);
-        }
-
-        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,
-                                      des_per_atom); // final descriptor
-        DeePKS_domain::check_descriptor(inlmax, des_per_atom, inl_l, ucell, PARAM.globalv.global_out_dir, descriptor);
-
-        if (PARAM.inp.deepks_out_labels)
-        {
-            LCAO_deepks_io::save_npy_d(nat,
-                                       des_per_atom,
-                                       inlmax,
-                                       inl_l,
-                                       PARAM.inp.deepks_equiv,
-                                       descriptor,
-                                       PARAM.globalv.global_out_dir,
-                                       GlobalV::MY_RANK); // libnpy needed
-        }
-    }
-
     /// print out deepks information to the screen
     if (PARAM.inp.deepks_scf)
     {
@@ -361,7 +375,7 @@ void LCAO_Deepks_Interface<TK, TR>::out_deepks_labels(const double& etot,
         {
             LCAO_deepks_io::print_dm(nks, PARAM.globalv.nlocal, ParaV->nrow, dm->get_DMK_vector());
 
-            DeePKS_domain::check_gedm(inlmax, inl_l, gedm);
+            DeePKS_domain::check_gedm(inlmax, inl_l, ld->gedm);
 
             std::ofstream ofs("E_delta_bands.dat");
             ofs << std::setprecision(10) << e_delta_band;

source/module_hamilt_lcao/module_deepks/deepks_basic.cpp

@@ -75,7 +75,7 @@ inline void generate_py_files(const int lmaxd, const int nmaxd, const std::strin
         return;
     }
 
-    std::ofstream ofs("cal_gedm.py");
+    std::ofstream ofs("cal_edelta_gedm.py");
     ofs << "import torch" << std::endl;
     ofs << "import numpy as np" << std::endl << std::endl;
     ofs << "import sys" << std::endl;
@@ -121,7 +121,7 @@ inline void generate_py_files(const int lmaxd, const int nmaxd, const std::strin
     }
 }
 
-void DeePKS_domain::cal_gedm_equiv(const int nat,
+void DeePKS_domain::cal_edelta_gedm_equiv(const int nat,
                                    const int lmaxd,
                                    const int nmaxd,
                                    const int inlmax,
@@ -131,7 +131,7 @@ void DeePKS_domain::cal_gedm_equiv(const int nat,
                                    double** gedm,
                                    double& E_delta)
 {
-    ModuleBase::TITLE("DeePKS_domain", "cal_gedm_equiv");
+    ModuleBase::TITLE("DeePKS_domain", "cal_edelta_gedm_equiv");
 
     LCAO_deepks_io::save_npy_d(nat,
                                des_per_atom,
@@ -146,7 +146,7 @@ void DeePKS_domain::cal_gedm_equiv(const int nat,
 
     if (GlobalV::MY_RANK == 0)
     {
-        std::string cmd = "python cal_gedm.py " + PARAM.inp.deepks_model;
+        std::string cmd = "python cal_edelta_gedm.py " + PARAM.inp.deepks_model;
         int stat = std::system(cmd.c_str());
         assert(stat == 0);
     }
@@ -155,13 +155,13 @@ void DeePKS_domain::cal_gedm_equiv(const int nat,
 
     LCAO_deepks_io::load_npy_gedm(nat, des_per_atom, gedm, E_delta, GlobalV::MY_RANK);
 
-    std::string cmd = "rm -f cal_gedm.py basis.yaml ec.npy gedm.npy";
+    std::string cmd = "rm -f cal_edelta_gedm.py basis.yaml ec.npy gedm.npy";
     std::system(cmd.c_str());
 }
 
 // obtain from the machine learning model dE_delta/dDescriptor
 // E_delta is also calculated here
-void DeePKS_domain::cal_gedm(const int nat,
+void DeePKS_domain::cal_edelta_gedm(const int nat,
                              const int lmaxd,
                              const int nmaxd,
                              const int inlmax,
@@ -175,10 +175,10 @@ void DeePKS_domain::cal_gedm(const int nat,
 {
     if (PARAM.inp.deepks_equiv)
     {
-        DeePKS_domain::cal_gedm_equiv(nat, lmaxd, nmaxd, inlmax, des_per_atom, inl_l, descriptor, gedm, E_delta);
+        DeePKS_domain::cal_edelta_gedm_equiv(nat, lmaxd, nmaxd, inlmax, des_per_atom, inl_l, descriptor, gedm, E_delta);
         return;
     }
-    ModuleBase::TITLE("DeePKS_domain", "cal_gedm");
+    ModuleBase::TITLE("DeePKS_domain", "cal_edelta_gedm");
 
     // forward
     std::vector<torch::jit::IValue> inputs;

source/module_hamilt_lcao/module_deepks/deepks_basic.h

@@ -18,11 +18,11 @@ namespace DeePKS_domain
 // The file contains 2 subroutines:
 // 1. load_model : loads model for applying V_delta
 // 2. cal_gevdm : d(des)/d(pdm), calculated using torch::autograd::grad
-// 3. cal_gedm : calculates d(E_delta)/d(pdm)
+// 3. cal_edelta_gedm : calculates E_delta and 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
 // 4. check_gedm : prints gedm for checking
-// 5. cal_gedm_equiv : calculates d(E_delta)/d(pdm) for equivariant version
+// 5. cal_edelta_gedm_equiv : calculates E_delta and d(E_delta)/d(pdm) for equivariant version
 
 // load the trained neural network models
 void load_model(const std::string& model_file, torch::jit::script::Module& model);
@@ -35,7 +35,7 @@ void cal_gevdm(const int nat,
                std::vector<torch::Tensor>& gevdm);
 
 /// calculate partial of energy correction to descriptors
-void cal_gedm(const int nat,
+void cal_edelta_gedm(const int nat,
               const int lmaxd,
               const int nmaxd,
               const int inlmax,
@@ -47,7 +47,7 @@ void cal_gedm(const int nat,
               double** gedm,
               double& E_delta);
 void check_gedm(const int inlmax, const int* inl_l, double** gedm);
-void cal_gedm_equiv(const int nat,
+void cal_edelta_gedm_equiv(const int nat,
                     const int lmaxd,
                     const int nmaxd,
                     const int inlmax,

source/module_hamilt_lcao/module_deepks/test/LCAO_deepks_test.cpp

@@ -372,7 +372,7 @@ void test_deepks::check_edelta(std::vector<torch::Tensor>& descriptor)
     {
         ld.allocate_V_delta(ucell.nat, kv.nkstot);
     }
-    DeePKS_domain::cal_gedm(ucell.nat,
+    DeePKS_domain::cal_edelta_gedm(ucell.nat,
                             this->ld.lmaxd,
                             this->ld.nmaxd,
                             this->ld.inlmax,

tests/deepks/602_NO_deepks_d_H2O_md_lda2pbe/result.ref

@@ -1,9 +1,9 @@
-etotref -465.9986234576679
-etotperatomref -155.3328744859
-totalforceref 5.535106
-totalstressref 1.522353
-totaldes 2.163703
-deepks_e_dm -57.8857180593137
-deepks_f_label 19.09559844689178
-deepks_s_label 19.250590727951906
-totaltimeref 12.58
+etotref -465.9986234579913
+etotperatomref -155.3328744860
+totalforceref 5.535112
+totalstressref 1.522354
+totaldes 2.163682
+deepks_e_dm -57.88576364957592
+deepks_f_label 19.095631983991726
+deepks_s_label 19.250613228828858
+totaltimeref 22.06

tests/deepks/603_NO_deepks_SiO2_bandgap_multik/result.ref

@@ -1,11 +1,11 @@
-etotref -1946.6192712234942519
-etotperatomref -324.4365452039
-totalforceref 32.200542
-totalstressref 111.359237
+etotref -1946.6192712237452724
+etotperatomref -324.4365452040
+totalforceref 32.200544
+totalstressref 111.359240
 totaldes 12.208233
-deepks_e_dm -233.1079723669145
-deepks_f_label 57.62761995350644
-deepks_s_label 106.76696458677478
-odelta -0.001860574012843015
-oprec -0.7420199384115197
-totaltimeref 23.80
+deepks_e_dm -233.10797375382373
+deepks_f_label 57.62761915264378
+deepks_s_label 106.76695772125963
+odelta -0.0018605740130539488
+oprec -0.7420199382995606
+totaltimeref 30.60