From 1e6f069fc9017e445eea01859a5d5bb35264f4bb Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Tue, 12 Nov 2024 18:29:49 -0500 Subject: [PATCH 1/2] fix(pt): pass mapping from LAMMPS to PT This will be useful for the external MACE model. Signed-off-by: Jinzhe Zeng --- source/api_cc/src/DeepPotPT.cc | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/source/api_cc/src/DeepPotPT.cc b/source/api_cc/src/DeepPotPT.cc index 780a8007f3..ce104b0f8e 100644 --- a/source/api_cc/src/DeepPotPT.cc +++ b/source/api_cc/src/DeepPotPT.cc @@ -203,6 +203,15 @@ void DeepPotPT::compute(ENERGYVTYPE& ener, comm_dict.insert("recv_num", recvnum_tensor); comm_dict.insert("communicator", communicator_tensor); } + if (lmp_list.mapping) { + std::vector mapping(nall_real); + for (size_t ii = 0; ii < nall_real; ii++) { + mapping[ii] = lmp_list.mapping[fwd_map[ii]]; + } + mapping_tensor = + torch::from_blob(mapping.data(), {1, nall_real}, int_option) + .to(device); + } } at::Tensor firstneigh = createNlistTensor(nlist_data.jlist); firstneigh_tensor = firstneigh.to(torch::kInt64).to(device); From 76cd66a79d9cc29bdf58db4d4a406643a67dcddf Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Wed, 13 Nov 2024 03:53:05 -0500 Subject: [PATCH 2/2] sync mapping interface with #4307 Signed-off-by: Jinzhe Zeng --- source/api_c/include/c_api.h | 16 +++++++++++++--- source/api_c/include/deepmd.hpp | 5 +++++ source/api_c/src/c_api.cc | 3 +++ source/lib/include/neighbor_list.h | 6 ++++++ source/lmp/fix_dplr.cpp | 11 +++++++++++ source/lmp/pair_deepmd.cpp | 11 +++++++++++ 6 files changed, 49 insertions(+), 3 deletions(-) diff --git a/source/api_c/include/c_api.h b/source/api_c/include/c_api.h index c34a6909a5..a8798ba24d 100644 --- a/source/api_c/include/c_api.h +++ b/source/api_c/include/c_api.h @@ -12,7 +12,7 @@ extern "C" { /** C API version. Bumped whenever the API is changed. * @since API version 22 */ -#define DP_C_API_VERSION 24 +#define DP_C_API_VERSION 25 /** * @brief Neighbor list. @@ -31,7 +31,7 @@ extern DP_Nlist* DP_NewNlist(int inum_, int* ilist_, int* numneigh_, int** firstneigh_); -/* +/** * @brief Create a new neighbor list with communication capabilities. * @details This function extends DP_NewNlist by adding support for parallel * communication, allowing the neighbor list to be used in distributed @@ -68,7 +68,7 @@ extern DP_Nlist* DP_NewNlist_comm(int inum_, int* recvproc, void* world); -/* +/** * @brief Set mask for a neighbor list. * * @param nl Neighbor list. @@ -78,6 +78,16 @@ extern DP_Nlist* DP_NewNlist_comm(int inum_, **/ extern void DP_NlistSetMask(DP_Nlist* nl, int mask); +/** + * @brief Set mapping for a neighbor list. + * + * @param nl Neighbor list. + * @param mapping mapping from all atoms to real atoms, in size nall. + * @since API version 25 + * + **/ +extern void DP_NlistSetMapping(DP_Nlist* nl, int* mapping); + /** * @brief Delete a neighbor list. * diff --git a/source/api_c/include/deepmd.hpp b/source/api_c/include/deepmd.hpp index dd212e9dec..8a3656bfc2 100644 --- a/source/api_c/include/deepmd.hpp +++ b/source/api_c/include/deepmd.hpp @@ -863,6 +863,11 @@ struct InputNlist { * @brief Set mask for this neighbor list. */ void set_mask(int mask) { DP_NlistSetMask(nl, mask); }; + /** + * @brief Set mapping for this neighbor list. + * @param mapping mapping from all atoms to real atoms, in size nall. + */ + void set_mapping(int *mapping) { DP_NlistSetMapping(nl, mapping); }; }; /** diff --git a/source/api_c/src/c_api.cc b/source/api_c/src/c_api.cc index eba9be3664..4a0cff1520 100644 --- a/source/api_c/src/c_api.cc +++ b/source/api_c/src/c_api.cc @@ -43,6 +43,9 @@ DP_Nlist* DP_NewNlist_comm(int inum_, return new_nl; } void DP_NlistSetMask(DP_Nlist* nl, int mask) { nl->nl.set_mask(mask); } +void DP_NlistSetMapping(DP_Nlist* nl, int* mapping) { + nl->nl.set_mapping(mapping); +} void DP_DeleteNlist(DP_Nlist* nl) { delete nl; } // DP Base Model diff --git a/source/lib/include/neighbor_list.h b/source/lib/include/neighbor_list.h index bb4b8cf13c..5b39ea7454 100644 --- a/source/lib/include/neighbor_list.h +++ b/source/lib/include/neighbor_list.h @@ -44,6 +44,8 @@ struct InputNlist { void* world; /// mask to the neighbor index int mask = 0xFFFFFFFF; + /// mapping from all atoms to real atoms, in the size of nall + int* mapping = nullptr; InputNlist() : inum(0), ilist(NULL), @@ -99,6 +101,10 @@ struct InputNlist { * @brief Set mask for this neighbor list. */ void set_mask(int mask_) { mask = mask_; }; + /** + * @brief Set mapping for this neighbor list. + */ + void set_mapping(int* mapping_) { mapping = mapping_; }; }; /** diff --git a/source/lmp/fix_dplr.cpp b/source/lmp/fix_dplr.cpp index 8e54410d0a..ac161730db 100644 --- a/source/lmp/fix_dplr.cpp +++ b/source/lmp/fix_dplr.cpp @@ -467,6 +467,14 @@ void FixDPLR::pre_force(int vflag) { int nghost = atom->nghost; int nall = nlocal + nghost; + // mapping (for DPA-2 JAX) + std::vector mapping_vec(nall, -1); + if (comm->nprocs == 1 && atom->map_style != Atom::MAP_NONE) { + for (size_t ii = 0; ii < nall; ++ii) { + mapping_vec[ii] = atom->map(atom->tag[ii]); + } + } + // if (eflag_atom) { // error->all(FLERR,"atomic energy calculation is not supported by this // fix\n"); @@ -499,6 +507,9 @@ void FixDPLR::pre_force(int vflag) { deepmd_compat::InputNlist lmp_list(list->inum, list->ilist, list->numneigh, list->firstneigh); lmp_list.set_mask(NEIGHMASK); + if (comm->nprocs == 1 && atom->map_style != Atom::MAP_NONE) { + lmp_list.set_mapping(mapping_vec.data()); + } // declear output vector tensor; // compute diff --git a/source/lmp/pair_deepmd.cpp b/source/lmp/pair_deepmd.cpp index 6d12fda20a..8127979cd1 100644 --- a/source/lmp/pair_deepmd.cpp +++ b/source/lmp/pair_deepmd.cpp @@ -155,6 +155,14 @@ void PairDeepMD::compute(int eflag, int vflag) { } } + // mapping (for DPA-2 JAX) + std::vector mapping_vec(nall, -1); + if (comm->nprocs == 1 && atom->map_style != Atom::MAP_NONE) { + for (size_t ii = 0; ii < nall; ++ii) { + mapping_vec[ii] = atom->map(atom->tag[ii]); + } + } + if (do_compute_aparam) { make_aparam_from_compute(daparam); } else if (aparam.size() > 0) { @@ -198,6 +206,9 @@ void PairDeepMD::compute(int eflag, int vflag) { commdata_->firstrecv, commdata_->sendlist, commdata_->sendproc, commdata_->recvproc, &world); lmp_list.set_mask(NEIGHMASK); + if (comm->nprocs == 1 && atom->map_style != Atom::MAP_NONE) { + lmp_list.set_mapping(mapping_vec.data()); + } deepmd_compat::InputNlist extend_lmp_list; if (single_model || multi_models_no_mod_devi) { // cvflag_atom is the right flag for the cvatom matrix