Fix: a mpi bug in DeePKS. (#7485)

* Fix a mpi bug in DeePKS.

* Fix the mpi bug and update the reference value in tests.

* Update DeePKS UT.

Author: ErjieWu

source/source_lcao/module_deepks/LCAO_deepks_interface.cpp

@@ -629,7 +629,7 @@ void LCAO_Deepks_Interface<TK, TR>::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<double>(file_overlap, overlap_out, rank);
                             const std::string file_iRmat = PARAM.globalv.global_out_dir + "deepks_iRmat.npy";

source/source_lcao/module_deepks/deepks_vdrpre.cpp

@@ -19,6 +19,7 @@ void DeePKS_domain::prepare_phialpha_iRmat(const int nlocal,
                                             const std::vector<hamilt::HContainer<double>*> 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<double>(), overlap.numel());
+#endif
     iRmat = mapping_R(dRmat_tmp.unsqueeze(1) - dRmat_tmp.unsqueeze(2));
     ModuleBase::timer::end("DeePKS_domain", "prepare_phialpha_iRmat");
     return;

source/source_lcao/module_deepks/deepks_vdrpre.h

@@ -35,6 +35,7 @@ void prepare_phialpha_iRmat(const int nlocal,
                             const std::vector<hamilt::HContainer<double>*> phialpha,
                             const UnitCell& ucell,
                             const LCAO_Orbitals& orb,
+                            const Parallel_Orbitals& pv,
                             const Grid_Driver& GridD,
                             torch::Tensor& overlap,
                             torch::Tensor& iRmat);

source/source_lcao/module_deepks/test/LCAO_deepks_test.cpp

@@ -308,6 +308,7 @@ void test_deepks<T>::check_vdrpre()
                                         this->ld.phialpha,
                                         ucell,
                                         ORB,
+                                        ParaO,
                                         Test_Deepks::GridD,
                                         overlap_out,
                                         iRmat);

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

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