Fix mpi bug for DeePKS and add some timer ticks.

ErjieWu · ErjieWu · commit 188c3f189663 · 2025-02-07T17:16:55.000+08:00
diff --git a/source/module_esolver/lcao_before_scf.cpp b/source/module_esolver/lcao_before_scf.cpp
@@ -227,7 +227,7 @@ void ESolver_KS_LCAO<TK, TR>::before_scf(UnitCell& ucell, const int istep)
 
         if (PARAM.inp.deepks_out_unittest)
         {
-            DeePKS_domain::check_phialpha(PARAM.inp.cal_force, ucell, orb_, this->gd, pv, this->ld.phialpha);
+            DeePKS_domain::check_phialpha(PARAM.inp.cal_force, ucell, orb_, this->gd, pv, this->ld.phialpha, GlobalV::MY_RANK);
         }
     }
 #endif
diff --git a/source/module_esolver/lcao_others.cpp b/source/module_esolver/lcao_others.cpp
@@ -233,7 +233,7 @@ void ESolver_KS_LCAO<TK, TR>::others(UnitCell& ucell, const int istep)
 
         if (PARAM.inp.deepks_out_unittest)
         {
-            DeePKS_domain::check_phialpha(PARAM.inp.cal_force, ucell, orb_, this->gd, pv, this->ld.phialpha);
+            DeePKS_domain::check_phialpha(PARAM.inp.cal_force, ucell, orb_, this->gd, pv, this->ld.phialpha, GlobalV::MY_RANK);
         }
     }
 #endif
diff --git a/source/module_hamilt_lcao/module_deepks/LCAO_deepks.cpp b/source/module_hamilt_lcao/module_deepks/LCAO_deepks.cpp
@@ -53,6 +53,7 @@ void LCAO_Deepks::init(const LCAO_Orbitals& orb,
                        std::vector<int> na)
 {
     ModuleBase::TITLE("LCAO_Deepks", "init");
+    ModuleBase::timer::tick("LCAO_Deepks", "init");
 
     GlobalV::ofs_running << " Initialize the descriptor index for DeePKS (lcao line)" << std::endl;
 
@@ -123,6 +124,7 @@ void LCAO_Deepks::init(const LCAO_Orbitals& orb,
 
     this->pv = &pv_in;
 
+    ModuleBase::timer::tick("LCAO_Deepks", "init");
     return;
 }
 
@@ -169,6 +171,7 @@ void LCAO_Deepks::init_index(const int ntype,
 void LCAO_Deepks::allocate_V_delta(const int nat, const int nks)
 {
     ModuleBase::TITLE("LCAO_Deepks", "allocate_V_delta");
+    ModuleBase::timer::tick("LCAO_Deepks", "allocate_V_delta");
 
     // initialize the H matrix H_V_delta
     if (PARAM.globalv.gamma_only_local)
@@ -205,6 +208,7 @@ void LCAO_Deepks::allocate_V_delta(const int nat, const int nks)
         ModuleBase::GlobalFunc::ZEROS(this->gedm[inl], pdm_size);
     }
 
+    ModuleBase::timer::tick("LCAO_Deepks", "allocate_V_delta");
     return;
 }
 
diff --git a/source/module_hamilt_lcao/module_deepks/LCAO_deepks_io.cpp b/source/module_hamilt_lcao/module_deepks/LCAO_deepks_io.cpp
@@ -261,11 +261,11 @@ void LCAO_deepks_io::save_matrix2npy(const std::string& file_name,
 template <typename T>
 void LCAO_deepks_io::save_tensor2npy(const std::string& file_name, const torch::Tensor& tensor, const int rank)
 {
+    ModuleBase::TITLE("LCAO_deepks_io", "save_tensor2npy");
     if (rank != 0)
     {
         return;
     }
-    ModuleBase::TITLE("LCAO_deepks_io", "save_tensor2npy");
     const int dim = tensor.dim();
     std::vector<long unsigned> shape(dim);
     for (int i = 0; i < dim; i++)
diff --git a/source/module_hamilt_lcao/module_deepks/deepks_basic.cpp b/source/module_hamilt_lcao/module_deepks/deepks_basic.cpp
@@ -15,6 +15,7 @@ void DeePKS_domain::cal_gevdm(const int nat,
                               std::vector<torch::Tensor>& gevdm)
 {
     ModuleBase::TITLE("DeePKS_domain", "cal_gevdm");
+    ModuleBase::timer::tick("DeePKS_domain", "cal_gevdm");
     // cal gevdm(d(EigenValue(D))/dD)
     int nlmax = inlmax / nat;
     for (int nl = 0; nl < nlmax; ++nl)
@@ -48,12 +49,14 @@ void DeePKS_domain::cal_gevdm(const int nat,
         gevdm.push_back(avmm);
     }
     assert(gevdm.size() == nlmax);
+    ModuleBase::timer::tick("DeePKS_domain", "cal_gevdm");
     return;
 }
 
 void DeePKS_domain::load_model(const std::string& model_file, torch::jit::script::Module& model)
 {
     ModuleBase::TITLE("DeePKS_domain", "load_model");
+    ModuleBase::timer::tick("DeePKS_domain", "load_model");
 
     try
     {
@@ -64,6 +67,7 @@ void DeePKS_domain::load_model(const std::string& model_file, torch::jit::script
         std::cerr << "error loading the model" << std::endl;
         return;
     }
+    ModuleBase::timer::tick("DeePKS_domain", "load_model");
     return;
 }
 
@@ -132,6 +136,7 @@ void DeePKS_domain::cal_edelta_gedm_equiv(const int nat,
                                    double& E_delta)
 {
     ModuleBase::TITLE("DeePKS_domain", "cal_edelta_gedm_equiv");
+    ModuleBase::timer::tick("DeePKS_domain", "cal_edelta_gedm_equiv");
 
     LCAO_deepks_io::save_npy_d(nat,
                                des_per_atom,
@@ -157,6 +162,9 @@ void DeePKS_domain::cal_edelta_gedm_equiv(const int nat,
 
     std::string cmd = "rm -f cal_edelta_gedm.py basis.yaml ec.npy gedm.npy";
     std::system(cmd.c_str());
+
+    ModuleBase::timer::tick("DeePKS_domain", "cal_edelta_gedm_equiv");
+    return;
 }
 
 // obtain from the machine learning model dE_delta/dDescriptor
@@ -179,6 +187,7 @@ void DeePKS_domain::cal_edelta_gedm(const int nat,
         return;
     }
     ModuleBase::TITLE("DeePKS_domain", "cal_edelta_gedm");
+    ModuleBase::timer::tick("DeePKS_domain", "cal_edelta_gedm");
 
     // forward
     std::vector<torch::jit::IValue> inputs;
@@ -213,6 +222,7 @@ void DeePKS_domain::cal_edelta_gedm(const int nat,
             }
         }
     }
+    ModuleBase::timer::tick("DeePKS_domain", "cal_edelta_gedm");
     return;
 }
 
diff --git a/source/module_hamilt_lcao/module_deepks/deepks_force.cpp b/source/module_hamilt_lcao/module_deepks/deepks_force.cpp
@@ -25,12 +25,11 @@ void DeePKS_domain::cal_f_delta(const std::vector<std::vector<TK>>& dm,
                                 ModuleBase::matrix& svnl_dalpha)
 {
     ModuleBase::TITLE("DeePKS_domain", "cal_f_delta");
-
+    ModuleBase::timer::tick("DeePKS_domain", "cal_f_delta");
     f_delta.zero_out();
 
     const double Rcut_Alpha = orb.Alpha[0].getRcut();
     const int lmaxd = orb.get_lmax_d();
-    const int nrow = pv.nrow;
 
     for (int T0 = 0; T0 < ucell.ntype; T0++)
     {
@@ -146,6 +145,22 @@ void DeePKS_domain::cal_f_delta(const std::vector<std::vector<TK>>& dm,
                         }
                     }
 
+                    hamilt::BaseMatrix<double>* overlap_1 = phialpha[0]->find_matrix(iat, ibt1, dR1);
+                    hamilt::BaseMatrix<double>* overlap_2 = phialpha[0]->find_matrix(iat, ibt2, dR2);
+                    if (overlap_1 == nullptr || overlap_2 == nullptr)
+                    {
+                        continue;
+                    }
+                    std::vector<hamilt::BaseMatrix<double>*> grad_overlap_1(3);
+                    std::vector<hamilt::BaseMatrix<double>*> grad_overlap_2(3);
+                    for (int i = 0; i < 3; ++i)
+                    {
+                        grad_overlap_1[i] = phialpha[i + 1]->find_matrix(iat, ibt1, dR1);
+                        grad_overlap_2[i] = phialpha[i + 1]->find_matrix(iat, ibt2, dR2);
+                    }
+
+                    assert(overlap_1->get_col_size() == overlap_2->get_col_size());
+
                     const double* dm_current = dm_pair.get_pointer();
 
                     for (int iw1 = 0; iw1 < row_indexes.size(); ++iw1)
@@ -155,18 +170,6 @@ void DeePKS_domain::cal_f_delta(const std::vector<std::vector<TK>>& dm,
                             double nlm[3] = {0, 0, 0};
                             double nlm_t[3] = {0, 0, 0}; // for stress
 
-                            hamilt::BaseMatrix<double>* overlap_1 = phialpha[0]->find_matrix(iat, ibt1, dR1);
-                            hamilt::BaseMatrix<double>* overlap_2 = phialpha[0]->find_matrix(iat, ibt2, dR2);
-                            std::vector<hamilt::BaseMatrix<double>*> grad_overlap_1(3);
-                            std::vector<hamilt::BaseMatrix<double>*> grad_overlap_2(3);
-                            for (int i = 0; i < 3; ++i)
-                            {
-                                grad_overlap_1[i] = phialpha[i + 1]->find_matrix(iat, ibt1, dR1);
-                                grad_overlap_2[i] = phialpha[i + 1]->find_matrix(iat, ibt2, dR2);
-                            }
-
-                            assert(overlap_1->get_col_size() == overlap_2->get_col_size());
-
                             if (!PARAM.inp.deepks_equiv)
                             {
                                 int ib = 0;
@@ -310,7 +313,7 @@ void DeePKS_domain::cal_f_delta(const std::vector<std::vector<TK>>& dm,
             }
         }
     }
-
+    ModuleBase::timer::tick("DeePKS_domain", "cal_f_delta");
     return;
 }
 
diff --git a/source/module_hamilt_lcao/module_deepks/deepks_fpre.cpp b/source/module_hamilt_lcao/module_deepks/deepks_fpre.cpp
@@ -136,23 +136,27 @@ void DeePKS_domain::cal_gdmx(const int lmaxd,
 
                     dm_current = dm_pair.get_pointer();
 
-                    for (int iw1 = 0; iw1 < row_indexes.size(); ++iw1)
+                    hamilt::BaseMatrix<double>* overlap_1 = phialpha[0]->find_matrix(iat, ibt1, dR1);
+                    hamilt::BaseMatrix<double>* overlap_2 = phialpha[0]->find_matrix(iat, ibt2, dR2);
+                    if (overlap_1 == nullptr || overlap_2 == nullptr)
                     {
-                        for (int iw2 = 0; iw2 < col_indexes.size(); ++iw2)
-                        {
-                            hamilt::BaseMatrix<double>* overlap_1 = phialpha[0]->find_matrix(iat, ibt1, dR1);
-                            hamilt::BaseMatrix<double>* overlap_2 = phialpha[0]->find_matrix(iat, ibt2, dR2);
-                            std::vector<hamilt::BaseMatrix<double>*> grad_overlap_1(3);
-                            std::vector<hamilt::BaseMatrix<double>*> grad_overlap_2(3);
+                        continue;
+                    }
+                    std::vector<hamilt::BaseMatrix<double>*> grad_overlap_1(3);
+                    std::vector<hamilt::BaseMatrix<double>*> grad_overlap_2(3);
 
-                            assert(overlap_1->get_col_size() == overlap_2->get_col_size());
+                    assert(overlap_1->get_col_size() == overlap_2->get_col_size());
 
-                            for (int i = 0; i < 3; ++i)
-                            {
-                                grad_overlap_1[i] = phialpha[i + 1]->find_matrix(iat, ibt1, dR1);
-                                grad_overlap_2[i] = phialpha[i + 1]->find_matrix(iat, ibt2, dR2);
-                            }
+                    for (int i = 0; i < 3; ++i)
+                    {
+                        grad_overlap_1[i] = phialpha[i + 1]->find_matrix(iat, ibt1, dR1);
+                        grad_overlap_2[i] = phialpha[i + 1]->find_matrix(iat, ibt2, dR2);
+                    }
 
+                    for (int iw1 = 0; iw1 < row_indexes.size(); ++iw1)
+                    {
+                        for (int iw2 = 0; iw2 < col_indexes.size(); ++iw2)
+                        {
                             int ib = 0;
                             for (int L0 = 0; L0 <= orb.Alpha[0].getLmax(); ++L0)
                             {
@@ -266,7 +270,7 @@ void DeePKS_domain::cal_gvx(const int nat,
                             torch::Tensor& gvx)
 {
     ModuleBase::TITLE("DeePKS_domain", "cal_gvx");
-
+    ModuleBase::timer::tick("DeePKS_domain", "cal_gvx");
     // gdmr : nat(derivative) * 3 * inl(projector) * nm * nm
     std::vector<torch::Tensor> gdmr;
     auto accessor = gdmx.accessor<double, 5>();
@@ -332,7 +336,7 @@ void DeePKS_domain::cal_gvx(const int nat,
         assert(gvx.size(2) == nat);
         assert(gvx.size(3) == des_per_atom);
     }
-
+    ModuleBase::timer::tick("DeePKS_domain", "cal_gvx");
     return;
 }
 
diff --git a/source/module_hamilt_lcao/module_deepks/deepks_orbital.cpp b/source/module_hamilt_lcao/module_deepks/deepks_orbital.cpp
@@ -15,6 +15,7 @@ void DeePKS_domain::cal_o_delta(const std::vector<TH>& dm_hl,
                                 const int nks)
 {
     ModuleBase::TITLE("DeePKS_domain", "cal_o_delta");
+    ModuleBase::timer::tick("DeePKS_domain", "cal_o_delta");
 
     for (int ik = 0; ik < nks; ik++)
     {
@@ -64,6 +65,7 @@ void DeePKS_domain::cal_o_delta(const std::vector<TH>& dm_hl,
             o_delta(ik, 0) = o_delta_tmp.real();
         }
     }
+    ModuleBase::timer::tick("DeePKS_domain", "cal_o_delta");
     return;
 }
 
diff --git a/source/module_hamilt_lcao/module_deepks/deepks_orbpre.cpp b/source/module_hamilt_lcao/module_deepks/deepks_orbpre.cpp
@@ -94,12 +94,9 @@ void DeePKS_domain::cal_orbital_precalc(const std::vector<TH>& dm_hl,
 
                 ModuleBase::Vector3<int> dR1(GridD.getBox(ad1).x, GridD.getBox(ad1).y, GridD.getBox(ad1).z);
 
-                if constexpr (std::is_same<TK, std::complex<double>>::value)
+                if (phialpha[0]->find_matrix(iat, ibt1, dR1.x, dR1.y, dR1.z) == nullptr)
                 {
-                    if (phialpha[0]->find_matrix(iat, ibt1, dR1.x, dR1.y, dR1.z) == nullptr)
-                    {
-                        continue;
-                    }
+                    continue;
                 }
 
                 auto row_indexes = pv.get_indexes_row(ibt1);
@@ -150,12 +147,9 @@ void DeePKS_domain::cal_orbital_precalc(const std::vector<TH>& dm_hl,
 
                     ModuleBase::Vector3<int> dR2(GridD.getBox(ad2).x, GridD.getBox(ad2).y, GridD.getBox(ad2).z);
 
-                    if constexpr (std::is_same<TK, std::complex<double>>::value)
+                    if (phialpha[0]->find_matrix(iat, ibt2, dR2.x, dR2.y, dR2.z) == nullptr)
                     {
-                        if (phialpha[0]->find_matrix(iat, ibt2, dR2.x, dR2.y, dR2.z) == nullptr)
-                        {
-                            continue;
-                        }
+                        continue;
                     }
 
                     auto col_indexes = pv.get_indexes_col(ibt2);
diff --git a/source/module_hamilt_lcao/module_deepks/deepks_pdm.cpp b/source/module_hamilt_lcao/module_deepks/deepks_pdm.cpp
@@ -102,6 +102,7 @@ void DeePKS_domain::cal_pdm(bool& init_pdm,
 
 {
     ModuleBase::TITLE("DeePKS_domain", "cal_pdm");
+    ModuleBase::timer::tick("DeePKS_domain", "cal_pdm");
 
     // if pdm has been initialized, skip the calculation
     if (init_pdm)
@@ -133,8 +134,6 @@ void DeePKS_domain::cal_pdm(bool& init_pdm,
         }
     }
 
-    ModuleBase::timer::tick("DeePKS_domain", "cal_pdm");
-
     const double Rcut_Alpha = orb.Alpha[0].getRcut();
     for (int T0 = 0; T0 < ucell.ntype; T0++)
     {
@@ -204,12 +203,9 @@ void DeePKS_domain::cal_pdm(bool& init_pdm,
                 }
 
                 ModuleBase::Vector3<int> dR1(GridD.getBox(ad1).x, GridD.getBox(ad1).y, GridD.getBox(ad1).z);
-                if constexpr (std::is_same<TK, std::complex<double>>::value)
+                if (phialpha[0]->find_matrix(iat, ibt1, dR1.x, dR1.y, dR1.z) == nullptr)
                 {
-                    if (phialpha[0]->find_matrix(iat, ibt1, dR1.x, dR1.y, dR1.z) == nullptr)
-                    {
-                        continue;
-                    }
+                    continue;
                 }
 
                 auto row_indexes = pv.get_indexes_row(ibt1);
@@ -242,12 +238,9 @@ void DeePKS_domain::cal_pdm(bool& init_pdm,
                     const int nw2_tot = atom2->nw * PARAM.globalv.npol;
 
                     ModuleBase::Vector3<int> dR2(GridD.getBox(ad2).x, GridD.getBox(ad2).y, GridD.getBox(ad2).z);
-                    if constexpr (std::is_same<TK, std::complex<double>>::value)
+                    if (phialpha[0]->find_matrix(iat, ibt2, dR2.x, dR2.y, dR2.z) == nullptr)
                     {
-                        if (phialpha[0]->find_matrix(iat, ibt2, dR2.x, dR2.y, dR2.z) == nullptr)
-                        {
-                            continue;
-                        }
+                        continue;
                     }
 
                     const double Rcut_AO2 = orb.Phi[T2].getRcut();
diff --git a/source/module_hamilt_lcao/module_deepks/deepks_phialpha.cpp b/source/module_hamilt_lcao/module_deepks/deepks_phialpha.cpp
diff --git a/source/module_hamilt_lcao/module_deepks/deepks_phialpha.h b/source/module_hamilt_lcao/module_deepks/deepks_phialpha.h
diff --git a/source/module_hamilt_lcao/module_deepks/deepks_spre.cpp b/source/module_hamilt_lcao/module_deepks/deepks_spre.cpp
diff --git a/source/module_hamilt_lcao/module_deepks/deepks_vdpre.cpp b/source/module_hamilt_lcao/module_deepks/deepks_vdpre.cpp

Original file line number	Diff line number	Diff line change
`@@ -227,7 +227,7 @@ void ESolver_KS_LCAO<TK, TR>::before_scf(UnitCell& ucell, const int istep)`
`227`	`227`
`228`	`228`	`if (PARAM.inp.deepks_out_unittest)`
`229`	`229`	`{`
`230`		`- DeePKS_domain::check_phialpha(PARAM.inp.cal_force, ucell, orb_, this->gd, pv, this->ld.phialpha);`
	`230`	`+ DeePKS_domain::check_phialpha(PARAM.inp.cal_force, ucell, orb_, this->gd, pv, this->ld.phialpha, GlobalV::MY_RANK);`
`231`	`231`	`}`
`232`	`232`	`}`
`233`	`233`	`#endif`
Original file line number	Diff line number	Diff line change
`@@ -233,7 +233,7 @@ void ESolver_KS_LCAO<TK, TR>::others(UnitCell& ucell, const int istep)`
`233`	`233`
`234`	`234`	`if (PARAM.inp.deepks_out_unittest)`
`235`	`235`	`{`
`236`		`- DeePKS_domain::check_phialpha(PARAM.inp.cal_force, ucell, orb_, this->gd, pv, this->ld.phialpha);`
	`236`	`+ DeePKS_domain::check_phialpha(PARAM.inp.cal_force, ucell, orb_, this->gd, pv, this->ld.phialpha, GlobalV::MY_RANK);`
`237`	`237`	`}`
`238`	`238`	`}`
`239`	`239`	`#endif`
Original file line number	Diff line number	Diff line change
`@@ -53,6 +53,7 @@ void LCAO_Deepks::init(const LCAO_Orbitals& orb,`
`53`	`53`	`std::vector<int> na)`
`54`	`54`	`{`
`55`	`55`	`ModuleBase::TITLE("LCAO_Deepks", "init");`
	`56`	`+ ModuleBase::timer::tick("LCAO_Deepks", "init");`
`56`	`57`
`57`	`58`	`GlobalV::ofs_running << " Initialize the descriptor index for DeePKS (lcao line)" << std::endl;`
`58`	`59`
`@@ -123,6 +124,7 @@ void LCAO_Deepks::init(const LCAO_Orbitals& orb,`
`123`	`124`
`124`	`125`	`this->pv = &pv_in;`
`125`	`126`
	`127`	`+ ModuleBase::timer::tick("LCAO_Deepks", "init");`
`126`	`128`	`return;`
`127`	`129`	`}`
`128`	`130`
`@@ -169,6 +171,7 @@ void LCAO_Deepks::init_index(const int ntype,`
`169`	`171`	`void LCAO_Deepks::allocate_V_delta(const int nat, const int nks)`
`170`	`172`	`{`
`171`	`173`	`ModuleBase::TITLE("LCAO_Deepks", "allocate_V_delta");`
	`174`	`+ ModuleBase::timer::tick("LCAO_Deepks", "allocate_V_delta");`
`172`	`175`
`173`	`176`	`// initialize the H matrix H_V_delta`
`174`	`177`	`if (PARAM.globalv.gamma_only_local)`
`@@ -205,6 +208,7 @@ void LCAO_Deepks::allocate_V_delta(const int nat, const int nks)`
`205`	`208`	`ModuleBase::GlobalFunc::ZEROS(this->gedm[inl], pdm_size);`
`206`	`209`	`}`
`207`	`210`
	`211`	`+ ModuleBase::timer::tick("LCAO_Deepks", "allocate_V_delta");`
`208`	`212`	`return;`
`209`	`213`	`}`
`210`	`214`
Original file line number	Diff line number	Diff line change
`@@ -261,11 +261,11 @@ void LCAO_deepks_io::save_matrix2npy(const std::string& file_name,`
`261`	`261`	`template <typename T>`
`262`	`262`	`void LCAO_deepks_io::save_tensor2npy(const std::string& file_name, const torch::Tensor& tensor, const int rank)`
`263`	`263`	`{`
	`264`	`+ ModuleBase::TITLE("LCAO_deepks_io", "save_tensor2npy");`
`264`	`265`	`if (rank != 0)`
`265`	`266`	`{`
`266`	`267`	`return;`
`267`	`268`	`}`
`268`		`- ModuleBase::TITLE("LCAO_deepks_io", "save_tensor2npy");`
`269`	`269`	`const int dim = tensor.dim();`
`270`	`270`	`std::vector<long unsigned> shape(dim);`
`271`	`271`	`for (int i = 0; i < dim; i++)`
Original file line number	Diff line number	Diff line change
`@@ -15,6 +15,7 @@ void DeePKS_domain::cal_gevdm(const int nat,`
`15`	`15`	`std::vector<torch::Tensor>& gevdm)`
`16`	`16`	`{`
`17`	`17`	`ModuleBase::TITLE("DeePKS_domain", "cal_gevdm");`
	`18`	`+ ModuleBase::timer::tick("DeePKS_domain", "cal_gevdm");`
`18`	`19`	`// cal gevdm(d(EigenValue(D))/dD)`
`19`	`20`	`int nlmax = inlmax / nat;`
`20`	`21`	`for (int nl = 0; nl < nlmax; ++nl)`
`@@ -48,12 +49,14 @@ void DeePKS_domain::cal_gevdm(const int nat,`
`48`	`49`	`gevdm.push_back(avmm);`
`49`	`50`	`}`
`50`	`51`	`assert(gevdm.size() == nlmax);`
	`52`	`+ ModuleBase::timer::tick("DeePKS_domain", "cal_gevdm");`
`51`	`53`	`return;`
`52`	`54`	`}`
`53`	`55`
`54`	`56`	`void DeePKS_domain::load_model(const std::string& model_file, torch::jit::script::Module& model)`
`55`	`57`	`{`
`56`	`58`	`ModuleBase::TITLE("DeePKS_domain", "load_model");`
	`59`	`+ ModuleBase::timer::tick("DeePKS_domain", "load_model");`
`57`	`60`
`58`	`61`	`try`
`59`	`62`	`{`
`@@ -64,6 +67,7 @@ void DeePKS_domain::load_model(const std::string& model_file, torch::jit::script`
`64`	`67`	`std::cerr << "error loading the model" << std::endl;`
`65`	`68`	`return;`
`66`	`69`	`}`
	`70`	`+ ModuleBase::timer::tick("DeePKS_domain", "load_model");`
`67`	`71`	`return;`
`68`	`72`	`}`
`69`	`73`
`@@ -132,6 +136,7 @@ void DeePKS_domain::cal_edelta_gedm_equiv(const int nat,`
`132`	`136`	`double& E_delta)`
`133`	`137`	`{`
`134`	`138`	`ModuleBase::TITLE("DeePKS_domain", "cal_edelta_gedm_equiv");`
	`139`	`+ ModuleBase::timer::tick("DeePKS_domain", "cal_edelta_gedm_equiv");`
`135`	`140`
`136`	`141`	`LCAO_deepks_io::save_npy_d(nat,`
`137`	`142`	`des_per_atom,`
`@@ -157,6 +162,9 @@ void DeePKS_domain::cal_edelta_gedm_equiv(const int nat,`
`157`	`162`
`158`	`163`	`std::string cmd = "rm -f cal_edelta_gedm.py basis.yaml ec.npy gedm.npy";`
`159`	`164`	`std::system(cmd.c_str());`
	`165`	`+`
	`166`	`+ ModuleBase::timer::tick("DeePKS_domain", "cal_edelta_gedm_equiv");`
	`167`	`+ return;`
`160`	`168`	`}`
`161`	`169`
`162`	`170`	`// obtain from the machine learning model dE_delta/dDescriptor`
`@@ -179,6 +187,7 @@ void DeePKS_domain::cal_edelta_gedm(const int nat,`
`179`	`187`	`return;`
`180`	`188`	`}`
`181`	`189`	`ModuleBase::TITLE("DeePKS_domain", "cal_edelta_gedm");`
	`190`	`+ ModuleBase::timer::tick("DeePKS_domain", "cal_edelta_gedm");`
`182`	`191`
`183`	`192`	`// forward`
`184`	`193`	`std::vector<torch::jit::IValue> inputs;`
`@@ -213,6 +222,7 @@ void DeePKS_domain::cal_edelta_gedm(const int nat,`
`213`	`222`	`}`
`214`	`223`	`}`
`215`	`224`	`}`
	`225`	`+ ModuleBase::timer::tick("DeePKS_domain", "cal_edelta_gedm");`
`216`	`226`	`return;`
`217`	`227`	`}`
`218`	`228`
Original file line number	Diff line number	Diff line change
`@@ -94,12 +94,9 @@ void DeePKS_domain::cal_orbital_precalc(const std::vector<TH>& dm_hl,`
`94`	`94`
`95`	`95`	`ModuleBase::Vector3<int> dR1(GridD.getBox(ad1).x, GridD.getBox(ad1).y, GridD.getBox(ad1).z);`
`96`	`96`
`97`		`- if constexpr (std::is_same<TK, std::complex<double>>::value)`
	`97`	`+ if (phialpha[0]->find_matrix(iat, ibt1, dR1.x, dR1.y, dR1.z) == nullptr)`
`98`	`98`	`{`
`99`		`- if (phialpha[0]->find_matrix(iat, ibt1, dR1.x, dR1.y, dR1.z) == nullptr)`
`100`		`- {`
`101`		`- continue;`
`102`		`- }`
	`99`	`+ continue;`
`103`	`100`	`}`
`104`	`101`
`105`	`102`	`auto row_indexes = pv.get_indexes_row(ibt1);`
`@@ -150,12 +147,9 @@ void DeePKS_domain::cal_orbital_precalc(const std::vector<TH>& dm_hl,`
`150`	`147`
`151`	`148`	`ModuleBase::Vector3<int> dR2(GridD.getBox(ad2).x, GridD.getBox(ad2).y, GridD.getBox(ad2).z);`
`152`	`149`
`153`		`- if constexpr (std::is_same<TK, std::complex<double>>::value)`
	`150`	`+ if (phialpha[0]->find_matrix(iat, ibt2, dR2.x, dR2.y, dR2.z) == nullptr)`
`154`	`151`	`{`
`155`		`- if (phialpha[0]->find_matrix(iat, ibt2, dR2.x, dR2.y, dR2.z) == nullptr)`
`156`		`- {`
`157`		`- continue;`
`158`		`- }`
	`152`	`+ continue;`
`159`	`153`	`}`
`160`	`154`
`161`	`155`	`auto col_indexes = pv.get_indexes_col(ibt2);`
Original file line number	Diff line number	Diff line change
`@@ -102,6 +102,7 @@ void DeePKS_domain::cal_pdm(bool& init_pdm,`
`102`	`102`
`103`	`103`	`{`
`104`	`104`	`ModuleBase::TITLE("DeePKS_domain", "cal_pdm");`
	`105`	`+ ModuleBase::timer::tick("DeePKS_domain", "cal_pdm");`
`105`	`106`
`106`	`107`	`// if pdm has been initialized, skip the calculation`
`107`	`108`	`if (init_pdm)`
`@@ -133,8 +134,6 @@ void DeePKS_domain::cal_pdm(bool& init_pdm,`
`133`	`134`	`}`
`134`	`135`	`}`
`135`	`136`
`136`		`- ModuleBase::timer::tick("DeePKS_domain", "cal_pdm");`
`137`		`-`
`138`	`137`	`const double Rcut_Alpha = orb.Alpha[0].getRcut();`
`139`	`138`	`for (int T0 = 0; T0 < ucell.ntype; T0++)`
`140`	`139`	`{`
`@@ -204,12 +203,9 @@ void DeePKS_domain::cal_pdm(bool& init_pdm,`
`204`	`203`	`}`
`205`	`204`
`206`	`205`	`ModuleBase::Vector3<int> dR1(GridD.getBox(ad1).x, GridD.getBox(ad1).y, GridD.getBox(ad1).z);`
`207`		`- if constexpr (std::is_same<TK, std::complex<double>>::value)`
	`206`	`+ if (phialpha[0]->find_matrix(iat, ibt1, dR1.x, dR1.y, dR1.z) == nullptr)`
`208`	`207`	`{`
`209`		`- if (phialpha[0]->find_matrix(iat, ibt1, dR1.x, dR1.y, dR1.z) == nullptr)`
`210`		`- {`
`211`		`- continue;`
`212`		`- }`
	`208`	`+ continue;`
`213`	`209`	`}`
`214`	`210`
`215`	`211`	`auto row_indexes = pv.get_indexes_row(ibt1);`
`@@ -242,12 +238,9 @@ void DeePKS_domain::cal_pdm(bool& init_pdm,`
`242`	`238`	`const int nw2_tot = atom2->nw * PARAM.globalv.npol;`
`243`	`239`
`244`	`240`	`ModuleBase::Vector3<int> dR2(GridD.getBox(ad2).x, GridD.getBox(ad2).y, GridD.getBox(ad2).z);`
`245`		`- if constexpr (std::is_same<TK, std::complex<double>>::value)`
	`241`	`+ if (phialpha[0]->find_matrix(iat, ibt2, dR2.x, dR2.y, dR2.z) == nullptr)`
`246`	`242`	`{`
`247`		`- if (phialpha[0]->find_matrix(iat, ibt2, dR2.x, dR2.y, dR2.z) == nullptr)`
`248`		`- {`
`249`		`- continue;`
`250`		`- }`
	`243`	`+ continue;`
`251`	`244`	`}`
`252`	`245`
`253`	`246`	`const double Rcut_AO2 = orb.Phi[T2].getRcut();`