diff --git a/source/source_lcao/module_deepks/LCAO_deepks_interface.cpp b/source/source_lcao/module_deepks/LCAO_deepks_interface.cpp index e7c41587595..b18b448459a 100644 --- a/source/source_lcao/module_deepks/LCAO_deepks_interface.cpp +++ b/source/source_lcao/module_deepks/LCAO_deepks_interface.cpp @@ -629,7 +629,7 @@ void LCAO_Deepks_Interface::out_deepks_labels(const double& etot, int R_size = DeePKS_domain::get_R_size(*h_deltaR); torch::Tensor overlap_out; torch::Tensor iRmat; - DeePKS_domain::prepare_phialpha_iRmat(nlocal, R_size, deepks_param, phialpha, ucell, orb, GridD, overlap_out, iRmat); + DeePKS_domain::prepare_phialpha_iRmat(nlocal, R_size, deepks_param, phialpha, ucell, orb, *ParaV, GridD, overlap_out, iRmat); const std::string file_overlap = PARAM.globalv.global_out_dir + "deepks_phialpha_r.npy"; LCAO_deepks_io::save_tensor2npy(file_overlap, overlap_out, rank); const std::string file_iRmat = PARAM.globalv.global_out_dir + "deepks_iRmat.npy"; diff --git a/source/source_lcao/module_deepks/deepks_vdrpre.cpp b/source/source_lcao/module_deepks/deepks_vdrpre.cpp index 0bb0325bc2a..d640b48a953 100644 --- a/source/source_lcao/module_deepks/deepks_vdrpre.cpp +++ b/source/source_lcao/module_deepks/deepks_vdrpre.cpp @@ -19,6 +19,7 @@ void DeePKS_domain::prepare_phialpha_iRmat(const int nlocal, const std::vector*> phialpha, const UnitCell& ucell, const LCAO_Orbitals& orb, + const Parallel_Orbitals& pv, const Grid_Driver& GridD, torch::Tensor& overlap, torch::Tensor& iRmat) @@ -81,6 +82,10 @@ void DeePKS_domain::prepare_phialpha_iRmat(const int nlocal, for (int ix = 0; ix < nw_tot; ix++) { + if (pv.global2local_row(start + ix) < 0 || pv.global2local_col(start + ix) < 0) + { + continue; + } for (int iy = 0; iy < deepks_param.des_per_atom; iy++) { overlap_accessor[iat][nnmax_vec[iat]][start + ix][iy] = overlap_mat->get_value(ix, iy); @@ -89,6 +94,9 @@ void DeePKS_domain::prepare_phialpha_iRmat(const int nlocal, nnmax_vec[iat]++; } ); +#ifdef __MPI + Parallel_Reduce::reduce_all(overlap.data_ptr(), overlap.numel()); +#endif iRmat = mapping_R(dRmat_tmp.unsqueeze(1) - dRmat_tmp.unsqueeze(2)); ModuleBase::timer::end("DeePKS_domain", "prepare_phialpha_iRmat"); return; diff --git a/source/source_lcao/module_deepks/deepks_vdrpre.h b/source/source_lcao/module_deepks/deepks_vdrpre.h index c22eb408fa7..62edad9a894 100644 --- a/source/source_lcao/module_deepks/deepks_vdrpre.h +++ b/source/source_lcao/module_deepks/deepks_vdrpre.h @@ -35,6 +35,7 @@ void prepare_phialpha_iRmat(const int nlocal, const std::vector*> phialpha, const UnitCell& ucell, const LCAO_Orbitals& orb, + const Parallel_Orbitals& pv, const Grid_Driver& GridD, torch::Tensor& overlap, torch::Tensor& iRmat); diff --git a/source/source_lcao/module_deepks/test/LCAO_deepks_test.cpp b/source/source_lcao/module_deepks/test/LCAO_deepks_test.cpp index 38a0cac574d..6f53a757d6c 100644 --- a/source/source_lcao/module_deepks/test/LCAO_deepks_test.cpp +++ b/source/source_lcao/module_deepks/test/LCAO_deepks_test.cpp @@ -308,6 +308,7 @@ void test_deepks::check_vdrpre() this->ld.phialpha, ucell, ORB, + ParaO, Test_Deepks::GridD, overlap_out, iRmat); diff --git a/tests/09_DeePKS/22_NO_GO_deepks_vdelta_r_2/result.ref b/tests/09_DeePKS/22_NO_GO_deepks_vdelta_r_2/result.ref index 5678831f58a..3de8b5a4c8e 100644 --- a/tests/09_DeePKS/22_NO_GO_deepks_vdelta_r_2/result.ref +++ b/tests/09_DeePKS/22_NO_GO_deepks_vdelta_r_2/result.ref @@ -6,6 +6,6 @@ deepks_e_label 17.126764505645642 deepks_edelta 0.09815855485768665 deepks_hr_label_pass 0 deepks_vdelta_r_pass 0 -deepks_phialpha_r 40.10770566607966 +deepks_phialpha_r 73.40481582340394 deepks_gevdm 54.0 totaltimeref 1.81 diff --git a/tests/09_DeePKS/24_NO_KP_deepks_vdelta_r_2/result.ref b/tests/09_DeePKS/24_NO_KP_deepks_vdelta_r_2/result.ref index b7e5a997dd0..3fce0adcd52 100644 --- a/tests/09_DeePKS/24_NO_KP_deepks_vdelta_r_2/result.ref +++ b/tests/09_DeePKS/24_NO_KP_deepks_vdelta_r_2/result.ref @@ -6,6 +6,6 @@ deepks_e_label 17.12676450564565 deepks_edelta 0.09815855485768665 deepks_hr_label_pass 0 deepks_vdelta_r_pass 0 -deepks_phialpha_r 40.10770566607966 +deepks_phialpha_r 73.40481582340394 deepks_gevdm 54.0 totaltimeref 1.81