From c67bf681a872731c0083a371f73dec98a7955ab0 Mon Sep 17 00:00:00 2001 From: ErjieWu Date: Thu, 18 Jun 2026 17:53:52 +0800 Subject: [PATCH 1/3] Fix a mpi bug in DeePKS. --- .../module_deepks/deepks_vdrpre.cpp | 588 +++++++++--------- 1 file changed, 297 insertions(+), 291 deletions(-) diff --git a/source/source_lcao/module_deepks/deepks_vdrpre.cpp b/source/source_lcao/module_deepks/deepks_vdrpre.cpp index 0bb0325bc2a..a00937d27bd 100644 --- a/source/source_lcao/module_deepks/deepks_vdrpre.cpp +++ b/source/source_lcao/module_deepks/deepks_vdrpre.cpp @@ -1,291 +1,297 @@ -// prepare_phialpha_iRmat : prepare phialpha_r and iR_mat for outputting npy file - -#ifdef __MLALGO - -#include "deepks_vdrpre.h" - -#include "LCAO_deepks_io.h" // mohan add 2024-07-22 -#include "deepks_iterate.h" -#include "source_base/constants.h" -#include "source_base/libm/libm.h" -#include "source_base/module_external/blas_connector.h" -#include "source_base/parallel_reduce.h" -#include "source_io/module_parameter/parameter.h" -#include "source_lcao/module_hcontainer/atom_pair.h" - -void DeePKS_domain::prepare_phialpha_iRmat(const int nlocal, - const int R_size, - const DeePKS_Param& deepks_param, - const std::vector*> phialpha, - const UnitCell& ucell, - const LCAO_Orbitals& orb, - const Grid_Driver& GridD, - torch::Tensor& overlap, - torch::Tensor& iRmat) -{ - ModuleBase::TITLE("DeePKS_domain", "prepare_phialpha_iRmat"); - ModuleBase::timer::start("DeePKS_domain", "prepare_phialpha_iRmat"); - constexpr torch::Dtype dtype = torch::kFloat64; - - // get the maximum nnmax - std::vector nnmax_vec(ucell.nat, 0); - DeePKS_domain::iterate_ad1( - ucell, - GridD, - orb, - false, // no trace_alpha - [&](const int iat, - const ModuleBase::Vector3& tau0, - const int ibt, - const ModuleBase::Vector3& tau1, - const int start, - const int nw_tot, - ModuleBase::Vector3 dR) - { - if (phialpha[0]->find_matrix(iat, ibt, dR.x, dR.y, dR.z) == nullptr) - { - return; // to next loop - } - nnmax_vec[iat]++; - } - ); - - int nnmax = *std::max_element(nnmax_vec.begin(), nnmax_vec.end()); - overlap = torch::zeros({ucell.nat, nnmax, nlocal, deepks_param.des_per_atom}, dtype); - torch::Tensor dRmat_tmp = torch::zeros({ucell.nat, nnmax, 3}, torch::kInt32); - auto overlap_accessor = overlap.accessor(); - auto dRmat_accessor = dRmat_tmp.accessor(); - - std::fill(nnmax_vec.begin(), nnmax_vec.end(), 0); - DeePKS_domain::iterate_ad1( - ucell, - GridD, - orb, - false, // no trace_alpha - [&](const int iat, - const ModuleBase::Vector3& tau0, - const int ibt, - const ModuleBase::Vector3& tau1, - const int start, - const int nw_tot, - ModuleBase::Vector3 dR) - { - hamilt::BaseMatrix* overlap_mat = phialpha[0]->find_matrix(iat, ibt, dR); - if (overlap_mat == nullptr) - { - return; // to next loop - } - dRmat_accessor[iat][nnmax_vec[iat]][0] = dR.x; - dRmat_accessor[iat][nnmax_vec[iat]][1] = dR.y; - dRmat_accessor[iat][nnmax_vec[iat]][2] = dR.z; - - for (int ix = 0; ix < nw_tot; ix++) - { - 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); - } - } - nnmax_vec[iat]++; - } - ); - iRmat = mapping_R(dRmat_tmp.unsqueeze(1) - dRmat_tmp.unsqueeze(2)); - ModuleBase::timer::end("DeePKS_domain", "prepare_phialpha_iRmat"); - return; -} - -void DeePKS_domain::cal_vdr_precalc(const int nlocal, - const int nat, - const int nks, - const int R_size, - const DeePKS_Param& deepks_param, - const std::vector>& kvec_d, - const std::vector*> phialpha, - const std::vector gevdm, - const UnitCell& ucell, - const LCAO_Orbitals& orb, - const Parallel_Orbitals& pv, - const Grid_Driver& GridD, - torch::Tensor& vdr_precalc) -{ - ModuleBase::TITLE("DeePKS_domain", "calc_vdr_precalc"); - ModuleBase::timer::start("DeePKS_domain", "calc_vdr_precalc"); - - torch::Tensor vdr_pdm = torch::zeros({R_size, - R_size, - R_size, - nlocal, - nlocal, - deepks_param.inlmax, - (2 * deepks_param.lmaxd + 1), - (2 * deepks_param.lmaxd + 1)}, - torch::TensorOptions().dtype(torch::kFloat64)); - auto accessor = vdr_pdm.accessor(); - - DeePKS_domain::iterate_ad2(ucell, - GridD, - orb, - false, // no trace_alpha - [&](const int iat, - const ModuleBase::Vector3& tau0, - const int ibt1, - const ModuleBase::Vector3& tau1, - const int start1, - const int nw1_tot, - ModuleBase::Vector3 dR1, - const int ibt2, - const ModuleBase::Vector3& tau2, - const int start2, - const int nw2_tot, - ModuleBase::Vector3 dR2) { - const int T0 = ucell.iat2it[iat]; - const int I0 = ucell.iat2ia[iat]; - if (phialpha[0]->find_matrix(iat, ibt1, dR1.x, dR1.y, dR1.z) == nullptr - || phialpha[0]->find_matrix(iat, ibt2, dR2.x, dR2.y, dR2.z) == nullptr) - { - return; // to next loop - } - - hamilt::BaseMatrix* overlap_1 = phialpha[0]->find_matrix(iat, ibt1, dR1); - hamilt::BaseMatrix* overlap_2 = phialpha[0]->find_matrix(iat, ibt2, dR2); - assert(overlap_1->get_col_size() == overlap_2->get_col_size()); - ModuleBase::Vector3 dR = dR2 - dR1; - int iRx = DeePKS_domain::mapping_R(dR.x); - int iRy = DeePKS_domain::mapping_R(dR.y); - int iRz = DeePKS_domain::mapping_R(dR.z); - // Make sure the index is in range we need to save - if (iRx >= R_size || iRy >= R_size || iRz >= R_size) - { - return; // to next loop - } - - for (int iw1 = 0; iw1 < nw1_tot; ++iw1) - { - const int iw1_all = start1 + iw1; // this is \mu - const int iw1_local = pv.global2local_row(iw1_all); - if (iw1_local < 0) - { - continue; - } - for (int iw2 = 0; iw2 < nw2_tot; ++iw2) - { - const int iw2_all = start2 + iw2; // this is \nu - const int iw2_local = pv.global2local_col(iw2_all); - if (iw2_local < 0) - { - continue; - } - - int ib = 0; - for (int L0 = 0; L0 <= orb.Alpha[0].getLmax(); ++L0) - { - for (int N0 = 0; N0 < orb.Alpha[0].getNchi(L0); ++N0) - { - const int inl = deepks_param.inl_index[T0](I0, L0, N0); - const int nm = 2 * L0 + 1; - - for (int m1 = 0; m1 < nm; ++m1) // nm = 1 for s, 3 for p, 5 for d - { - for (int m2 = 0; m2 < nm; ++m2) // nm = 1 for s, 3 for p, 5 for d - { - double tmp = overlap_1->get_value(iw1, ib + m1) - * overlap_2->get_value(iw2, ib + m2); - accessor[iRx][iRy][iRz][iw1_all][iw2_all][inl][m1][m2] - += tmp; - } - } - ib += nm; - } - } - } // iw2 - } // iw1 - }); - -#ifdef __MPI - const int size = R_size * R_size * R_size * nlocal * nlocal * deepks_param.inlmax * (2 * deepks_param.lmaxd + 1) - * (2 * deepks_param.lmaxd + 1); - double* data_ptr = vdr_pdm.data_ptr(); - Parallel_Reduce::reduce_all(data_ptr, size); -#endif - - // transfer v_delta_pdm to v_delta_pdm_vector - int nlmax = deepks_param.inlmax / nat; - std::vector vdr_pdm_vector; - for (int nl = 0; nl < nlmax; ++nl) - { - int nm = 2 * deepks_param.inl2l[nl] + 1; - torch::Tensor vdr_pdm_sliced - = vdr_pdm.slice(5, nl, deepks_param.inlmax, nlmax).slice(6, 0, nm, 1).slice(7, 0, nm, 1); - vdr_pdm_vector.push_back(vdr_pdm_sliced); - } - - assert(vdr_pdm_vector.size() == nlmax); - - // einsum for each nl: - std::vector vdr_vector; - for (int nl = 0; nl < nlmax; ++nl) - { - vdr_vector.push_back(at::einsum("pqrxyamn, avmn->pqrxyav", {vdr_pdm_vector[nl], gevdm[nl]})); - } - - vdr_precalc = torch::cat(vdr_vector, -1); - - ModuleBase::timer::end("DeePKS_domain", "calc_vdr_precalc"); - return; -} - -int DeePKS_domain::mapping_R(int R) -{ - // R_index mapping: index(R) = 2R-1 if R > 0, index(R) = -2R if R <= 0 - // after mapping, the new index [0,1,2,3,4,...] -> old index [0,1,-1,2,-2,...] - // This manipulation makes sure that the new index is natural number - // which makes it available to be used as index in torch::Tensor - int R_index = 0; - if (R > 0) - { - R_index = 2 * R - 1; - } - else - { - R_index = -2 * R; - } - return R_index; -} - -torch::Tensor DeePKS_domain::mapping_R(const torch::Tensor& R_tensor) -{ - auto R = R_tensor.to(torch::kInt32); - auto pos = R > 0; - auto twoR_minus1 = R * 2 - 1; - auto neg_minus2R = -2 * R; - return at::where(pos, twoR_minus1, neg_minus2R); -} - -template -int DeePKS_domain::get_R_size(const hamilt::HContainer& hcontainer) -{ - // get R_size from hcontainer - int R_size = 0; - if (hcontainer.size_R_loop() > 0) - { - for (int iR = 0; iR < hcontainer.size_R_loop(); ++iR) - { - ModuleBase::Vector3 R_vec; - hcontainer.loop_R(iR, R_vec.x, R_vec.y, R_vec.z); - int R_min = std::min({R_vec.x, R_vec.y, R_vec.z}); - int R_max = std::max({R_vec.x, R_vec.y, R_vec.z}); - int tmp_R_size = std::max(DeePKS_domain::mapping_R(R_min), DeePKS_domain::mapping_R(R_max)) + 1; - if (tmp_R_size > R_size) - { - R_size = tmp_R_size; - } - } - } - assert(R_size > 0); - return R_size; -} - -template int DeePKS_domain::get_R_size(const hamilt::HContainer& hcontainer); -template int DeePKS_domain::get_R_size>( - const hamilt::HContainer>& hcontainer); -#endif +// prepare_phialpha_iRmat : prepare phialpha_r and iR_mat for outputting npy file + +#ifdef __MLALGO + +#include "deepks_vdrpre.h" + +#include "LCAO_deepks_io.h" // mohan add 2024-07-22 +#include "deepks_iterate.h" +#include "source_base/constants.h" +#include "source_base/libm/libm.h" +#include "source_base/module_external/blas_connector.h" +#include "source_base/parallel_reduce.h" +#include "source_io/module_parameter/parameter.h" +#include "source_lcao/module_hcontainer/atom_pair.h" + +void DeePKS_domain::prepare_phialpha_iRmat(const int nlocal, + const int R_size, + const DeePKS_Param& deepks_param, + const std::vector*> phialpha, + const UnitCell& ucell, + const LCAO_Orbitals& orb, + const Grid_Driver& GridD, + torch::Tensor& overlap, + torch::Tensor& iRmat) +{ + ModuleBase::TITLE("DeePKS_domain", "prepare_phialpha_iRmat"); + ModuleBase::timer::start("DeePKS_domain", "prepare_phialpha_iRmat"); + constexpr torch::Dtype dtype = torch::kFloat64; + + // get the maximum nnmax + std::vector nnmax_vec(ucell.nat, 0); + DeePKS_domain::iterate_ad1( + ucell, + GridD, + orb, + false, // no trace_alpha + [&](const int iat, + const ModuleBase::Vector3& tau0, + const int ibt, + const ModuleBase::Vector3& tau1, + const int start, + const int nw_tot, + ModuleBase::Vector3 dR) + { + if (phialpha[0]->find_matrix(iat, ibt, dR.x, dR.y, dR.z) == nullptr) + { + return; // to next loop + } + nnmax_vec[iat]++; + } + ); + + int nnmax = *std::max_element(nnmax_vec.begin(), nnmax_vec.end()); + overlap = torch::zeros({ucell.nat, nnmax, nlocal, deepks_param.des_per_atom}, dtype); + torch::Tensor dRmat_tmp = torch::zeros({ucell.nat, nnmax, 3}, torch::kInt32); + auto overlap_accessor = overlap.accessor(); + auto dRmat_accessor = dRmat_tmp.accessor(); + + std::fill(nnmax_vec.begin(), nnmax_vec.end(), 0); + DeePKS_domain::iterate_ad1( + ucell, + GridD, + orb, + false, // no trace_alpha + [&](const int iat, + const ModuleBase::Vector3& tau0, + const int ibt, + const ModuleBase::Vector3& tau1, + const int start, + const int nw_tot, + ModuleBase::Vector3 dR) + { + hamilt::BaseMatrix* overlap_mat = phialpha[0]->find_matrix(iat, ibt, dR); + if (overlap_mat == nullptr) + { + return; // to next loop + } + dRmat_accessor[iat][nnmax_vec[iat]][0] = dR.x; + dRmat_accessor[iat][nnmax_vec[iat]][1] = dR.y; + dRmat_accessor[iat][nnmax_vec[iat]][2] = dR.z; + + for (int ix = 0; ix < nw_tot; ix++) + { + 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); + } + } + nnmax_vec[iat]++; + } + ); +#ifdef __MPI + // Each MPI process only holds the rows (orbitals) assigned to it in the + // distributed matrix pv. Sum contributions from all processes so the + // saved phialpha_r.npy contains the full nlocal × des_per_atom overlap. + 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; +} + +void DeePKS_domain::cal_vdr_precalc(const int nlocal, + const int nat, + const int nks, + const int R_size, + const DeePKS_Param& deepks_param, + const std::vector>& kvec_d, + const std::vector*> phialpha, + const std::vector gevdm, + const UnitCell& ucell, + const LCAO_Orbitals& orb, + const Parallel_Orbitals& pv, + const Grid_Driver& GridD, + torch::Tensor& vdr_precalc) +{ + ModuleBase::TITLE("DeePKS_domain", "calc_vdr_precalc"); + ModuleBase::timer::start("DeePKS_domain", "calc_vdr_precalc"); + + torch::Tensor vdr_pdm = torch::zeros({R_size, + R_size, + R_size, + nlocal, + nlocal, + deepks_param.inlmax, + (2 * deepks_param.lmaxd + 1), + (2 * deepks_param.lmaxd + 1)}, + torch::TensorOptions().dtype(torch::kFloat64)); + auto accessor = vdr_pdm.accessor(); + + DeePKS_domain::iterate_ad2(ucell, + GridD, + orb, + false, // no trace_alpha + [&](const int iat, + const ModuleBase::Vector3& tau0, + const int ibt1, + const ModuleBase::Vector3& tau1, + const int start1, + const int nw1_tot, + ModuleBase::Vector3 dR1, + const int ibt2, + const ModuleBase::Vector3& tau2, + const int start2, + const int nw2_tot, + ModuleBase::Vector3 dR2) { + const int T0 = ucell.iat2it[iat]; + const int I0 = ucell.iat2ia[iat]; + if (phialpha[0]->find_matrix(iat, ibt1, dR1.x, dR1.y, dR1.z) == nullptr + || phialpha[0]->find_matrix(iat, ibt2, dR2.x, dR2.y, dR2.z) == nullptr) + { + return; // to next loop + } + + hamilt::BaseMatrix* overlap_1 = phialpha[0]->find_matrix(iat, ibt1, dR1); + hamilt::BaseMatrix* overlap_2 = phialpha[0]->find_matrix(iat, ibt2, dR2); + assert(overlap_1->get_col_size() == overlap_2->get_col_size()); + ModuleBase::Vector3 dR = dR2 - dR1; + int iRx = DeePKS_domain::mapping_R(dR.x); + int iRy = DeePKS_domain::mapping_R(dR.y); + int iRz = DeePKS_domain::mapping_R(dR.z); + // Make sure the index is in range we need to save + if (iRx >= R_size || iRy >= R_size || iRz >= R_size) + { + return; // to next loop + } + + for (int iw1 = 0; iw1 < nw1_tot; ++iw1) + { + const int iw1_all = start1 + iw1; // this is \mu + const int iw1_local = pv.global2local_row(iw1_all); + if (iw1_local < 0) + { + continue; + } + for (int iw2 = 0; iw2 < nw2_tot; ++iw2) + { + const int iw2_all = start2 + iw2; // this is \nu + const int iw2_local = pv.global2local_col(iw2_all); + if (iw2_local < 0) + { + continue; + } + + int ib = 0; + for (int L0 = 0; L0 <= orb.Alpha[0].getLmax(); ++L0) + { + for (int N0 = 0; N0 < orb.Alpha[0].getNchi(L0); ++N0) + { + const int inl = deepks_param.inl_index[T0](I0, L0, N0); + const int nm = 2 * L0 + 1; + + for (int m1 = 0; m1 < nm; ++m1) // nm = 1 for s, 3 for p, 5 for d + { + for (int m2 = 0; m2 < nm; ++m2) // nm = 1 for s, 3 for p, 5 for d + { + double tmp = overlap_1->get_value(iw1, ib + m1) + * overlap_2->get_value(iw2, ib + m2); + accessor[iRx][iRy][iRz][iw1_all][iw2_all][inl][m1][m2] + += tmp; + } + } + ib += nm; + } + } + } // iw2 + } // iw1 + }); + +#ifdef __MPI + const int size = R_size * R_size * R_size * nlocal * nlocal * deepks_param.inlmax * (2 * deepks_param.lmaxd + 1) + * (2 * deepks_param.lmaxd + 1); + double* data_ptr = vdr_pdm.data_ptr(); + Parallel_Reduce::reduce_all(data_ptr, size); +#endif + + // transfer v_delta_pdm to v_delta_pdm_vector + int nlmax = deepks_param.inlmax / nat; + std::vector vdr_pdm_vector; + for (int nl = 0; nl < nlmax; ++nl) + { + int nm = 2 * deepks_param.inl2l[nl] + 1; + torch::Tensor vdr_pdm_sliced + = vdr_pdm.slice(5, nl, deepks_param.inlmax, nlmax).slice(6, 0, nm, 1).slice(7, 0, nm, 1); + vdr_pdm_vector.push_back(vdr_pdm_sliced); + } + + assert(vdr_pdm_vector.size() == nlmax); + + // einsum for each nl: + std::vector vdr_vector; + for (int nl = 0; nl < nlmax; ++nl) + { + vdr_vector.push_back(at::einsum("pqrxyamn, avmn->pqrxyav", {vdr_pdm_vector[nl], gevdm[nl]})); + } + + vdr_precalc = torch::cat(vdr_vector, -1); + + ModuleBase::timer::end("DeePKS_domain", "calc_vdr_precalc"); + return; +} + +int DeePKS_domain::mapping_R(int R) +{ + // R_index mapping: index(R) = 2R-1 if R > 0, index(R) = -2R if R <= 0 + // after mapping, the new index [0,1,2,3,4,...] -> old index [0,1,-1,2,-2,...] + // This manipulation makes sure that the new index is natural number + // which makes it available to be used as index in torch::Tensor + int R_index = 0; + if (R > 0) + { + R_index = 2 * R - 1; + } + else + { + R_index = -2 * R; + } + return R_index; +} + +torch::Tensor DeePKS_domain::mapping_R(const torch::Tensor& R_tensor) +{ + auto R = R_tensor.to(torch::kInt32); + auto pos = R > 0; + auto twoR_minus1 = R * 2 - 1; + auto neg_minus2R = -2 * R; + return at::where(pos, twoR_minus1, neg_minus2R); +} + +template +int DeePKS_domain::get_R_size(const hamilt::HContainer& hcontainer) +{ + // get R_size from hcontainer + int R_size = 0; + if (hcontainer.size_R_loop() > 0) + { + for (int iR = 0; iR < hcontainer.size_R_loop(); ++iR) + { + ModuleBase::Vector3 R_vec; + hcontainer.loop_R(iR, R_vec.x, R_vec.y, R_vec.z); + int R_min = std::min({R_vec.x, R_vec.y, R_vec.z}); + int R_max = std::max({R_vec.x, R_vec.y, R_vec.z}); + int tmp_R_size = std::max(DeePKS_domain::mapping_R(R_min), DeePKS_domain::mapping_R(R_max)) + 1; + if (tmp_R_size > R_size) + { + R_size = tmp_R_size; + } + } + } + assert(R_size > 0); + return R_size; +} + +template int DeePKS_domain::get_R_size(const hamilt::HContainer& hcontainer); +template int DeePKS_domain::get_R_size>( + const hamilt::HContainer>& hcontainer); +#endif From 2c843fbeba98657d2c67c85dbf2610f12b2cdda5 Mon Sep 17 00:00:00 2001 From: ErjieWu Date: Fri, 19 Jun 2026 00:08:47 +0800 Subject: [PATCH 2/3] Fix the mpi bug and update the reference value in tests. --- .../module_deepks/LCAO_deepks_interface.cpp | 2 +- .../module_deepks/deepks_vdrpre.cpp | 8 +- .../source_lcao/module_deepks/deepks_vdrpre.h | 125 +++++++++--------- .../22_NO_GO_deepks_vdelta_r_2/result.ref | 2 +- .../24_NO_KP_deepks_vdelta_r_2/result.ref | 2 +- 5 files changed, 71 insertions(+), 68 deletions(-) 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 a00937d27bd..a3d0b8ab781 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); @@ -90,9 +95,6 @@ void DeePKS_domain::prepare_phialpha_iRmat(const int nlocal, } ); #ifdef __MPI - // Each MPI process only holds the rows (orbitals) assigned to it in the - // distributed matrix pv. Sum contributions from all processes so the - // saved phialpha_r.npy contains the full nlocal × des_per_atom overlap. Parallel_Reduce::reduce_all(overlap.data_ptr(), overlap.numel()); #endif iRmat = mapping_R(dRmat_tmp.unsqueeze(1) - dRmat_tmp.unsqueeze(2)); diff --git a/source/source_lcao/module_deepks/deepks_vdrpre.h b/source/source_lcao/module_deepks/deepks_vdrpre.h index c22eb408fa7..ce2bd52f29a 100644 --- a/source/source_lcao/module_deepks/deepks_vdrpre.h +++ b/source/source_lcao/module_deepks/deepks_vdrpre.h @@ -1,63 +1,64 @@ -#ifndef DEEPKS_VDRPRE_H -#define DEEPKS_VDRPRE_H - -#ifdef __MLALGO - -#include "deepks_param.h" -#include "source_base/complexmatrix.h" -#include "source_base/intarray.h" -#include "source_base/matrix.h" -#include "source_base/timer.h" -#include "source_basis/module_ao/parallel_orbitals.h" -#include "source_cell/module_neighbor/sltk_grid_driver.h" -#include "source_lcao/module_hcontainer/hcontainer.h" - -#include -#include - -namespace DeePKS_domain -{ -//------------------------ -// deepks_vdrpre.cpp -//------------------------ - -// This file contains 1 subroutine for calculating v_delta, -// cal_vdr_precalc : v_delta_r_precalc is used for training with v_delta_r label, -// which equals gevdm * v_delta_pdm, -// v_delta_pdm = overlap * overlap - -// for deepks_v_delta = -1 -// calculates v_delta_r_precalc - -void prepare_phialpha_iRmat(const int nlocal, - const int R_size, - const DeePKS_Param& deepks_param, - const std::vector*> phialpha, - const UnitCell& ucell, - const LCAO_Orbitals& orb, - const Grid_Driver& GridD, - torch::Tensor& overlap, - torch::Tensor& iRmat); - -void cal_vdr_precalc(const int nlocal, - const int nat, - const int nks, - const int R_size, - const DeePKS_Param& deepks_param, - const std::vector>& kvec_d, - const std::vector*> phialpha, - const std::vector gevdm, - const UnitCell& ucell, - const LCAO_Orbitals& orb, - const Parallel_Orbitals& pv, - const Grid_Driver& GridD, - torch::Tensor& vdr_precalc); - -int mapping_R(int R); -torch::Tensor mapping_R(const torch::Tensor& R_tensor); - -template -int get_R_size(const hamilt::HContainer& hcontainer); -} // namespace DeePKS_domain -#endif +#ifndef DEEPKS_VDRPRE_H +#define DEEPKS_VDRPRE_H + +#ifdef __MLALGO + +#include "deepks_param.h" +#include "source_base/complexmatrix.h" +#include "source_base/intarray.h" +#include "source_base/matrix.h" +#include "source_base/timer.h" +#include "source_basis/module_ao/parallel_orbitals.h" +#include "source_cell/module_neighbor/sltk_grid_driver.h" +#include "source_lcao/module_hcontainer/hcontainer.h" + +#include +#include + +namespace DeePKS_domain +{ +//------------------------ +// deepks_vdrpre.cpp +//------------------------ + +// This file contains 1 subroutine for calculating v_delta, +// cal_vdr_precalc : v_delta_r_precalc is used for training with v_delta_r label, +// which equals gevdm * v_delta_pdm, +// v_delta_pdm = overlap * overlap + +// for deepks_v_delta = -1 +// calculates v_delta_r_precalc + +void prepare_phialpha_iRmat(const int nlocal, + const int R_size, + const DeePKS_Param& deepks_param, + 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); + +void cal_vdr_precalc(const int nlocal, + const int nat, + const int nks, + const int R_size, + const DeePKS_Param& deepks_param, + const std::vector>& kvec_d, + const std::vector*> phialpha, + const std::vector gevdm, + const UnitCell& ucell, + const LCAO_Orbitals& orb, + const Parallel_Orbitals& pv, + const Grid_Driver& GridD, + torch::Tensor& vdr_precalc); + +int mapping_R(int R); +torch::Tensor mapping_R(const torch::Tensor& R_tensor); + +template +int get_R_size(const hamilt::HContainer& hcontainer); +} // namespace DeePKS_domain +#endif #endif \ No newline at end of file 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 From 18cd3d7772fffd343b29b8da7638c8aef371913c Mon Sep 17 00:00:00 2001 From: ErjieWu Date: Fri, 19 Jun 2026 00:16:02 +0800 Subject: [PATCH 3/3] Update DeePKS UT. --- .../module_deepks/deepks_vdrpre.cpp | 598 +++++++++--------- .../source_lcao/module_deepks/deepks_vdrpre.h | 126 ++-- .../module_deepks/test/LCAO_deepks_test.cpp | 1 + 3 files changed, 363 insertions(+), 362 deletions(-) diff --git a/source/source_lcao/module_deepks/deepks_vdrpre.cpp b/source/source_lcao/module_deepks/deepks_vdrpre.cpp index a3d0b8ab781..d640b48a953 100644 --- a/source/source_lcao/module_deepks/deepks_vdrpre.cpp +++ b/source/source_lcao/module_deepks/deepks_vdrpre.cpp @@ -1,299 +1,299 @@ -// prepare_phialpha_iRmat : prepare phialpha_r and iR_mat for outputting npy file - -#ifdef __MLALGO - -#include "deepks_vdrpre.h" - -#include "LCAO_deepks_io.h" // mohan add 2024-07-22 -#include "deepks_iterate.h" -#include "source_base/constants.h" -#include "source_base/libm/libm.h" -#include "source_base/module_external/blas_connector.h" -#include "source_base/parallel_reduce.h" -#include "source_io/module_parameter/parameter.h" -#include "source_lcao/module_hcontainer/atom_pair.h" - -void DeePKS_domain::prepare_phialpha_iRmat(const int nlocal, - const int R_size, - const DeePKS_Param& deepks_param, - 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) -{ - ModuleBase::TITLE("DeePKS_domain", "prepare_phialpha_iRmat"); - ModuleBase::timer::start("DeePKS_domain", "prepare_phialpha_iRmat"); - constexpr torch::Dtype dtype = torch::kFloat64; - - // get the maximum nnmax - std::vector nnmax_vec(ucell.nat, 0); - DeePKS_domain::iterate_ad1( - ucell, - GridD, - orb, - false, // no trace_alpha - [&](const int iat, - const ModuleBase::Vector3& tau0, - const int ibt, - const ModuleBase::Vector3& tau1, - const int start, - const int nw_tot, - ModuleBase::Vector3 dR) - { - if (phialpha[0]->find_matrix(iat, ibt, dR.x, dR.y, dR.z) == nullptr) - { - return; // to next loop - } - nnmax_vec[iat]++; - } - ); - - int nnmax = *std::max_element(nnmax_vec.begin(), nnmax_vec.end()); - overlap = torch::zeros({ucell.nat, nnmax, nlocal, deepks_param.des_per_atom}, dtype); - torch::Tensor dRmat_tmp = torch::zeros({ucell.nat, nnmax, 3}, torch::kInt32); - auto overlap_accessor = overlap.accessor(); - auto dRmat_accessor = dRmat_tmp.accessor(); - - std::fill(nnmax_vec.begin(), nnmax_vec.end(), 0); - DeePKS_domain::iterate_ad1( - ucell, - GridD, - orb, - false, // no trace_alpha - [&](const int iat, - const ModuleBase::Vector3& tau0, - const int ibt, - const ModuleBase::Vector3& tau1, - const int start, - const int nw_tot, - ModuleBase::Vector3 dR) - { - hamilt::BaseMatrix* overlap_mat = phialpha[0]->find_matrix(iat, ibt, dR); - if (overlap_mat == nullptr) - { - return; // to next loop - } - dRmat_accessor[iat][nnmax_vec[iat]][0] = dR.x; - dRmat_accessor[iat][nnmax_vec[iat]][1] = dR.y; - dRmat_accessor[iat][nnmax_vec[iat]][2] = dR.z; - - 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); - } - } - 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; -} - -void DeePKS_domain::cal_vdr_precalc(const int nlocal, - const int nat, - const int nks, - const int R_size, - const DeePKS_Param& deepks_param, - const std::vector>& kvec_d, - const std::vector*> phialpha, - const std::vector gevdm, - const UnitCell& ucell, - const LCAO_Orbitals& orb, - const Parallel_Orbitals& pv, - const Grid_Driver& GridD, - torch::Tensor& vdr_precalc) -{ - ModuleBase::TITLE("DeePKS_domain", "calc_vdr_precalc"); - ModuleBase::timer::start("DeePKS_domain", "calc_vdr_precalc"); - - torch::Tensor vdr_pdm = torch::zeros({R_size, - R_size, - R_size, - nlocal, - nlocal, - deepks_param.inlmax, - (2 * deepks_param.lmaxd + 1), - (2 * deepks_param.lmaxd + 1)}, - torch::TensorOptions().dtype(torch::kFloat64)); - auto accessor = vdr_pdm.accessor(); - - DeePKS_domain::iterate_ad2(ucell, - GridD, - orb, - false, // no trace_alpha - [&](const int iat, - const ModuleBase::Vector3& tau0, - const int ibt1, - const ModuleBase::Vector3& tau1, - const int start1, - const int nw1_tot, - ModuleBase::Vector3 dR1, - const int ibt2, - const ModuleBase::Vector3& tau2, - const int start2, - const int nw2_tot, - ModuleBase::Vector3 dR2) { - const int T0 = ucell.iat2it[iat]; - const int I0 = ucell.iat2ia[iat]; - if (phialpha[0]->find_matrix(iat, ibt1, dR1.x, dR1.y, dR1.z) == nullptr - || phialpha[0]->find_matrix(iat, ibt2, dR2.x, dR2.y, dR2.z) == nullptr) - { - return; // to next loop - } - - hamilt::BaseMatrix* overlap_1 = phialpha[0]->find_matrix(iat, ibt1, dR1); - hamilt::BaseMatrix* overlap_2 = phialpha[0]->find_matrix(iat, ibt2, dR2); - assert(overlap_1->get_col_size() == overlap_2->get_col_size()); - ModuleBase::Vector3 dR = dR2 - dR1; - int iRx = DeePKS_domain::mapping_R(dR.x); - int iRy = DeePKS_domain::mapping_R(dR.y); - int iRz = DeePKS_domain::mapping_R(dR.z); - // Make sure the index is in range we need to save - if (iRx >= R_size || iRy >= R_size || iRz >= R_size) - { - return; // to next loop - } - - for (int iw1 = 0; iw1 < nw1_tot; ++iw1) - { - const int iw1_all = start1 + iw1; // this is \mu - const int iw1_local = pv.global2local_row(iw1_all); - if (iw1_local < 0) - { - continue; - } - for (int iw2 = 0; iw2 < nw2_tot; ++iw2) - { - const int iw2_all = start2 + iw2; // this is \nu - const int iw2_local = pv.global2local_col(iw2_all); - if (iw2_local < 0) - { - continue; - } - - int ib = 0; - for (int L0 = 0; L0 <= orb.Alpha[0].getLmax(); ++L0) - { - for (int N0 = 0; N0 < orb.Alpha[0].getNchi(L0); ++N0) - { - const int inl = deepks_param.inl_index[T0](I0, L0, N0); - const int nm = 2 * L0 + 1; - - for (int m1 = 0; m1 < nm; ++m1) // nm = 1 for s, 3 for p, 5 for d - { - for (int m2 = 0; m2 < nm; ++m2) // nm = 1 for s, 3 for p, 5 for d - { - double tmp = overlap_1->get_value(iw1, ib + m1) - * overlap_2->get_value(iw2, ib + m2); - accessor[iRx][iRy][iRz][iw1_all][iw2_all][inl][m1][m2] - += tmp; - } - } - ib += nm; - } - } - } // iw2 - } // iw1 - }); - -#ifdef __MPI - const int size = R_size * R_size * R_size * nlocal * nlocal * deepks_param.inlmax * (2 * deepks_param.lmaxd + 1) - * (2 * deepks_param.lmaxd + 1); - double* data_ptr = vdr_pdm.data_ptr(); - Parallel_Reduce::reduce_all(data_ptr, size); -#endif - - // transfer v_delta_pdm to v_delta_pdm_vector - int nlmax = deepks_param.inlmax / nat; - std::vector vdr_pdm_vector; - for (int nl = 0; nl < nlmax; ++nl) - { - int nm = 2 * deepks_param.inl2l[nl] + 1; - torch::Tensor vdr_pdm_sliced - = vdr_pdm.slice(5, nl, deepks_param.inlmax, nlmax).slice(6, 0, nm, 1).slice(7, 0, nm, 1); - vdr_pdm_vector.push_back(vdr_pdm_sliced); - } - - assert(vdr_pdm_vector.size() == nlmax); - - // einsum for each nl: - std::vector vdr_vector; - for (int nl = 0; nl < nlmax; ++nl) - { - vdr_vector.push_back(at::einsum("pqrxyamn, avmn->pqrxyav", {vdr_pdm_vector[nl], gevdm[nl]})); - } - - vdr_precalc = torch::cat(vdr_vector, -1); - - ModuleBase::timer::end("DeePKS_domain", "calc_vdr_precalc"); - return; -} - -int DeePKS_domain::mapping_R(int R) -{ - // R_index mapping: index(R) = 2R-1 if R > 0, index(R) = -2R if R <= 0 - // after mapping, the new index [0,1,2,3,4,...] -> old index [0,1,-1,2,-2,...] - // This manipulation makes sure that the new index is natural number - // which makes it available to be used as index in torch::Tensor - int R_index = 0; - if (R > 0) - { - R_index = 2 * R - 1; - } - else - { - R_index = -2 * R; - } - return R_index; -} - -torch::Tensor DeePKS_domain::mapping_R(const torch::Tensor& R_tensor) -{ - auto R = R_tensor.to(torch::kInt32); - auto pos = R > 0; - auto twoR_minus1 = R * 2 - 1; - auto neg_minus2R = -2 * R; - return at::where(pos, twoR_minus1, neg_minus2R); -} - -template -int DeePKS_domain::get_R_size(const hamilt::HContainer& hcontainer) -{ - // get R_size from hcontainer - int R_size = 0; - if (hcontainer.size_R_loop() > 0) - { - for (int iR = 0; iR < hcontainer.size_R_loop(); ++iR) - { - ModuleBase::Vector3 R_vec; - hcontainer.loop_R(iR, R_vec.x, R_vec.y, R_vec.z); - int R_min = std::min({R_vec.x, R_vec.y, R_vec.z}); - int R_max = std::max({R_vec.x, R_vec.y, R_vec.z}); - int tmp_R_size = std::max(DeePKS_domain::mapping_R(R_min), DeePKS_domain::mapping_R(R_max)) + 1; - if (tmp_R_size > R_size) - { - R_size = tmp_R_size; - } - } - } - assert(R_size > 0); - return R_size; -} - -template int DeePKS_domain::get_R_size(const hamilt::HContainer& hcontainer); -template int DeePKS_domain::get_R_size>( - const hamilt::HContainer>& hcontainer); -#endif +// prepare_phialpha_iRmat : prepare phialpha_r and iR_mat for outputting npy file + +#ifdef __MLALGO + +#include "deepks_vdrpre.h" + +#include "LCAO_deepks_io.h" // mohan add 2024-07-22 +#include "deepks_iterate.h" +#include "source_base/constants.h" +#include "source_base/libm/libm.h" +#include "source_base/module_external/blas_connector.h" +#include "source_base/parallel_reduce.h" +#include "source_io/module_parameter/parameter.h" +#include "source_lcao/module_hcontainer/atom_pair.h" + +void DeePKS_domain::prepare_phialpha_iRmat(const int nlocal, + const int R_size, + const DeePKS_Param& deepks_param, + 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) +{ + ModuleBase::TITLE("DeePKS_domain", "prepare_phialpha_iRmat"); + ModuleBase::timer::start("DeePKS_domain", "prepare_phialpha_iRmat"); + constexpr torch::Dtype dtype = torch::kFloat64; + + // get the maximum nnmax + std::vector nnmax_vec(ucell.nat, 0); + DeePKS_domain::iterate_ad1( + ucell, + GridD, + orb, + false, // no trace_alpha + [&](const int iat, + const ModuleBase::Vector3& tau0, + const int ibt, + const ModuleBase::Vector3& tau1, + const int start, + const int nw_tot, + ModuleBase::Vector3 dR) + { + if (phialpha[0]->find_matrix(iat, ibt, dR.x, dR.y, dR.z) == nullptr) + { + return; // to next loop + } + nnmax_vec[iat]++; + } + ); + + int nnmax = *std::max_element(nnmax_vec.begin(), nnmax_vec.end()); + overlap = torch::zeros({ucell.nat, nnmax, nlocal, deepks_param.des_per_atom}, dtype); + torch::Tensor dRmat_tmp = torch::zeros({ucell.nat, nnmax, 3}, torch::kInt32); + auto overlap_accessor = overlap.accessor(); + auto dRmat_accessor = dRmat_tmp.accessor(); + + std::fill(nnmax_vec.begin(), nnmax_vec.end(), 0); + DeePKS_domain::iterate_ad1( + ucell, + GridD, + orb, + false, // no trace_alpha + [&](const int iat, + const ModuleBase::Vector3& tau0, + const int ibt, + const ModuleBase::Vector3& tau1, + const int start, + const int nw_tot, + ModuleBase::Vector3 dR) + { + hamilt::BaseMatrix* overlap_mat = phialpha[0]->find_matrix(iat, ibt, dR); + if (overlap_mat == nullptr) + { + return; // to next loop + } + dRmat_accessor[iat][nnmax_vec[iat]][0] = dR.x; + dRmat_accessor[iat][nnmax_vec[iat]][1] = dR.y; + dRmat_accessor[iat][nnmax_vec[iat]][2] = dR.z; + + 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); + } + } + 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; +} + +void DeePKS_domain::cal_vdr_precalc(const int nlocal, + const int nat, + const int nks, + const int R_size, + const DeePKS_Param& deepks_param, + const std::vector>& kvec_d, + const std::vector*> phialpha, + const std::vector gevdm, + const UnitCell& ucell, + const LCAO_Orbitals& orb, + const Parallel_Orbitals& pv, + const Grid_Driver& GridD, + torch::Tensor& vdr_precalc) +{ + ModuleBase::TITLE("DeePKS_domain", "calc_vdr_precalc"); + ModuleBase::timer::start("DeePKS_domain", "calc_vdr_precalc"); + + torch::Tensor vdr_pdm = torch::zeros({R_size, + R_size, + R_size, + nlocal, + nlocal, + deepks_param.inlmax, + (2 * deepks_param.lmaxd + 1), + (2 * deepks_param.lmaxd + 1)}, + torch::TensorOptions().dtype(torch::kFloat64)); + auto accessor = vdr_pdm.accessor(); + + DeePKS_domain::iterate_ad2(ucell, + GridD, + orb, + false, // no trace_alpha + [&](const int iat, + const ModuleBase::Vector3& tau0, + const int ibt1, + const ModuleBase::Vector3& tau1, + const int start1, + const int nw1_tot, + ModuleBase::Vector3 dR1, + const int ibt2, + const ModuleBase::Vector3& tau2, + const int start2, + const int nw2_tot, + ModuleBase::Vector3 dR2) { + const int T0 = ucell.iat2it[iat]; + const int I0 = ucell.iat2ia[iat]; + if (phialpha[0]->find_matrix(iat, ibt1, dR1.x, dR1.y, dR1.z) == nullptr + || phialpha[0]->find_matrix(iat, ibt2, dR2.x, dR2.y, dR2.z) == nullptr) + { + return; // to next loop + } + + hamilt::BaseMatrix* overlap_1 = phialpha[0]->find_matrix(iat, ibt1, dR1); + hamilt::BaseMatrix* overlap_2 = phialpha[0]->find_matrix(iat, ibt2, dR2); + assert(overlap_1->get_col_size() == overlap_2->get_col_size()); + ModuleBase::Vector3 dR = dR2 - dR1; + int iRx = DeePKS_domain::mapping_R(dR.x); + int iRy = DeePKS_domain::mapping_R(dR.y); + int iRz = DeePKS_domain::mapping_R(dR.z); + // Make sure the index is in range we need to save + if (iRx >= R_size || iRy >= R_size || iRz >= R_size) + { + return; // to next loop + } + + for (int iw1 = 0; iw1 < nw1_tot; ++iw1) + { + const int iw1_all = start1 + iw1; // this is \mu + const int iw1_local = pv.global2local_row(iw1_all); + if (iw1_local < 0) + { + continue; + } + for (int iw2 = 0; iw2 < nw2_tot; ++iw2) + { + const int iw2_all = start2 + iw2; // this is \nu + const int iw2_local = pv.global2local_col(iw2_all); + if (iw2_local < 0) + { + continue; + } + + int ib = 0; + for (int L0 = 0; L0 <= orb.Alpha[0].getLmax(); ++L0) + { + for (int N0 = 0; N0 < orb.Alpha[0].getNchi(L0); ++N0) + { + const int inl = deepks_param.inl_index[T0](I0, L0, N0); + const int nm = 2 * L0 + 1; + + for (int m1 = 0; m1 < nm; ++m1) // nm = 1 for s, 3 for p, 5 for d + { + for (int m2 = 0; m2 < nm; ++m2) // nm = 1 for s, 3 for p, 5 for d + { + double tmp = overlap_1->get_value(iw1, ib + m1) + * overlap_2->get_value(iw2, ib + m2); + accessor[iRx][iRy][iRz][iw1_all][iw2_all][inl][m1][m2] + += tmp; + } + } + ib += nm; + } + } + } // iw2 + } // iw1 + }); + +#ifdef __MPI + const int size = R_size * R_size * R_size * nlocal * nlocal * deepks_param.inlmax * (2 * deepks_param.lmaxd + 1) + * (2 * deepks_param.lmaxd + 1); + double* data_ptr = vdr_pdm.data_ptr(); + Parallel_Reduce::reduce_all(data_ptr, size); +#endif + + // transfer v_delta_pdm to v_delta_pdm_vector + int nlmax = deepks_param.inlmax / nat; + std::vector vdr_pdm_vector; + for (int nl = 0; nl < nlmax; ++nl) + { + int nm = 2 * deepks_param.inl2l[nl] + 1; + torch::Tensor vdr_pdm_sliced + = vdr_pdm.slice(5, nl, deepks_param.inlmax, nlmax).slice(6, 0, nm, 1).slice(7, 0, nm, 1); + vdr_pdm_vector.push_back(vdr_pdm_sliced); + } + + assert(vdr_pdm_vector.size() == nlmax); + + // einsum for each nl: + std::vector vdr_vector; + for (int nl = 0; nl < nlmax; ++nl) + { + vdr_vector.push_back(at::einsum("pqrxyamn, avmn->pqrxyav", {vdr_pdm_vector[nl], gevdm[nl]})); + } + + vdr_precalc = torch::cat(vdr_vector, -1); + + ModuleBase::timer::end("DeePKS_domain", "calc_vdr_precalc"); + return; +} + +int DeePKS_domain::mapping_R(int R) +{ + // R_index mapping: index(R) = 2R-1 if R > 0, index(R) = -2R if R <= 0 + // after mapping, the new index [0,1,2,3,4,...] -> old index [0,1,-1,2,-2,...] + // This manipulation makes sure that the new index is natural number + // which makes it available to be used as index in torch::Tensor + int R_index = 0; + if (R > 0) + { + R_index = 2 * R - 1; + } + else + { + R_index = -2 * R; + } + return R_index; +} + +torch::Tensor DeePKS_domain::mapping_R(const torch::Tensor& R_tensor) +{ + auto R = R_tensor.to(torch::kInt32); + auto pos = R > 0; + auto twoR_minus1 = R * 2 - 1; + auto neg_minus2R = -2 * R; + return at::where(pos, twoR_minus1, neg_minus2R); +} + +template +int DeePKS_domain::get_R_size(const hamilt::HContainer& hcontainer) +{ + // get R_size from hcontainer + int R_size = 0; + if (hcontainer.size_R_loop() > 0) + { + for (int iR = 0; iR < hcontainer.size_R_loop(); ++iR) + { + ModuleBase::Vector3 R_vec; + hcontainer.loop_R(iR, R_vec.x, R_vec.y, R_vec.z); + int R_min = std::min({R_vec.x, R_vec.y, R_vec.z}); + int R_max = std::max({R_vec.x, R_vec.y, R_vec.z}); + int tmp_R_size = std::max(DeePKS_domain::mapping_R(R_min), DeePKS_domain::mapping_R(R_max)) + 1; + if (tmp_R_size > R_size) + { + R_size = tmp_R_size; + } + } + } + assert(R_size > 0); + return R_size; +} + +template int DeePKS_domain::get_R_size(const hamilt::HContainer& hcontainer); +template int DeePKS_domain::get_R_size>( + const hamilt::HContainer>& hcontainer); +#endif diff --git a/source/source_lcao/module_deepks/deepks_vdrpre.h b/source/source_lcao/module_deepks/deepks_vdrpre.h index ce2bd52f29a..62edad9a894 100644 --- a/source/source_lcao/module_deepks/deepks_vdrpre.h +++ b/source/source_lcao/module_deepks/deepks_vdrpre.h @@ -1,64 +1,64 @@ -#ifndef DEEPKS_VDRPRE_H -#define DEEPKS_VDRPRE_H - -#ifdef __MLALGO - -#include "deepks_param.h" -#include "source_base/complexmatrix.h" -#include "source_base/intarray.h" -#include "source_base/matrix.h" -#include "source_base/timer.h" -#include "source_basis/module_ao/parallel_orbitals.h" -#include "source_cell/module_neighbor/sltk_grid_driver.h" -#include "source_lcao/module_hcontainer/hcontainer.h" - -#include -#include - -namespace DeePKS_domain -{ -//------------------------ -// deepks_vdrpre.cpp -//------------------------ - -// This file contains 1 subroutine for calculating v_delta, -// cal_vdr_precalc : v_delta_r_precalc is used for training with v_delta_r label, -// which equals gevdm * v_delta_pdm, -// v_delta_pdm = overlap * overlap - -// for deepks_v_delta = -1 -// calculates v_delta_r_precalc - -void prepare_phialpha_iRmat(const int nlocal, - const int R_size, - const DeePKS_Param& deepks_param, - 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); - -void cal_vdr_precalc(const int nlocal, - const int nat, - const int nks, - const int R_size, - const DeePKS_Param& deepks_param, - const std::vector>& kvec_d, - const std::vector*> phialpha, - const std::vector gevdm, - const UnitCell& ucell, - const LCAO_Orbitals& orb, - const Parallel_Orbitals& pv, - const Grid_Driver& GridD, - torch::Tensor& vdr_precalc); - -int mapping_R(int R); -torch::Tensor mapping_R(const torch::Tensor& R_tensor); - -template -int get_R_size(const hamilt::HContainer& hcontainer); -} // namespace DeePKS_domain -#endif +#ifndef DEEPKS_VDRPRE_H +#define DEEPKS_VDRPRE_H + +#ifdef __MLALGO + +#include "deepks_param.h" +#include "source_base/complexmatrix.h" +#include "source_base/intarray.h" +#include "source_base/matrix.h" +#include "source_base/timer.h" +#include "source_basis/module_ao/parallel_orbitals.h" +#include "source_cell/module_neighbor/sltk_grid_driver.h" +#include "source_lcao/module_hcontainer/hcontainer.h" + +#include +#include + +namespace DeePKS_domain +{ +//------------------------ +// deepks_vdrpre.cpp +//------------------------ + +// This file contains 1 subroutine for calculating v_delta, +// cal_vdr_precalc : v_delta_r_precalc is used for training with v_delta_r label, +// which equals gevdm * v_delta_pdm, +// v_delta_pdm = overlap * overlap + +// for deepks_v_delta = -1 +// calculates v_delta_r_precalc + +void prepare_phialpha_iRmat(const int nlocal, + const int R_size, + const DeePKS_Param& deepks_param, + 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); + +void cal_vdr_precalc(const int nlocal, + const int nat, + const int nks, + const int R_size, + const DeePKS_Param& deepks_param, + const std::vector>& kvec_d, + const std::vector*> phialpha, + const std::vector gevdm, + const UnitCell& ucell, + const LCAO_Orbitals& orb, + const Parallel_Orbitals& pv, + const Grid_Driver& GridD, + torch::Tensor& vdr_precalc); + +int mapping_R(int R); +torch::Tensor mapping_R(const torch::Tensor& R_tensor); + +template +int get_R_size(const hamilt::HContainer& hcontainer); +} // namespace DeePKS_domain +#endif #endif \ No newline at end of file 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);