diff --git a/deepmd/pt/model/model/make_model.py b/deepmd/pt/model/model/make_model.py index 167ad81923..9685766a9c 100644 --- a/deepmd/pt/model/model/make_model.py +++ b/deepmd/pt/model/model/make_model.py @@ -469,6 +469,11 @@ def get_nnei(self) -> int: """Returns the total number of selected neighboring atoms in the cut-off radius.""" return self.atomic_model.get_nnei() + @torch.jit.export + def get_ntypes(self) -> int: + """Returns the number of element types.""" + return self.atomic_model.get_ntypes() + def atomic_output_def(self) -> FittingOutputDef: """Get the output def of the atomic model.""" return self.atomic_model.atomic_output_def() diff --git a/source/api_c/include/c_api.h b/source/api_c/include/c_api.h index 911813e428..1c5c4a0bdc 100644 --- a/source/api_c/include/c_api.h +++ b/source/api_c/include/c_api.h @@ -24,7 +24,8 @@ extern DP_Nlist* DP_NewNlist(int inum_, int* ilist_, int* numneigh_, int** firstneigh_); - +extern DP_Nlist* DP_NewNlist_mapping( + int inum_, int* ilist_, int* numneigh_, int** firstneigh_, int* mapping); /** * @brief Delete a neighbor list. * diff --git a/source/api_c/include/deepmd.hpp b/source/api_c/include/deepmd.hpp index 16b8f08cad..b7def8bac3 100644 --- a/source/api_c/include/deepmd.hpp +++ b/source/api_c/include/deepmd.hpp @@ -522,6 +522,16 @@ struct InputNlist { nl(DP_NewNlist(inum_, ilist_, numneigh_, firstneigh_)) { DP_CHECK_OK(DP_NlistCheckOK, nl); }; + InputNlist( + int inum_, int *ilist_, int *numneigh_, int **firstneigh_, int *mapping) + : inum(inum_), + ilist(ilist_), + numneigh(numneigh_), + firstneigh(firstneigh_), + nl(DP_NewNlist_mapping( + inum_, ilist_, numneigh_, firstneigh_, mapping)) { + DP_CHECK_OK(DP_NlistCheckOK, nl); + }; ~InputNlist() { DP_DeleteNlist(nl); }; /// @brief C API neighbor list. DP_Nlist *nl; diff --git a/source/api_c/src/c_api.cc b/source/api_c/src/c_api.cc index 79dc486e0d..e8943ad265 100644 --- a/source/api_c/src/c_api.cc +++ b/source/api_c/src/c_api.cc @@ -24,7 +24,12 @@ DP_Nlist* DP_NewNlist(int inum_, deepmd::InputNlist nl(inum_, ilist_, numneigh_, firstneigh_); DP_Nlist* new_nl = new DP_Nlist(nl); return new_nl;) } - +DP_Nlist* DP_NewNlist_mapping( + int inum_, int* ilist_, int* numneigh_, int** firstneigh_, int* mapping) { + deepmd::InputNlist nl(inum_, ilist_, numneigh_, firstneigh_, mapping); + DP_Nlist* new_nl = new DP_Nlist(nl); + return new_nl; +} void DP_DeleteNlist(DP_Nlist* nl) { delete nl; } DP_DeepPot::DP_DeepPot() {} diff --git a/source/api_cc/src/DeepPotPT.cc b/source/api_cc/src/DeepPotPT.cc index 2c3fd1d865..59467139d2 100644 --- a/source/api_cc/src/DeepPotPT.cc +++ b/source/api_cc/src/DeepPotPT.cc @@ -131,7 +131,13 @@ void DeepPotPT::compute(ENERGYVTYPE& ener, at::Tensor firstneigh = createNlistTensor(nlist_data.jlist); firstneigh_tensor = firstneigh.to(torch::kInt64).to(device); bool do_atom_virial_tensor = true; - c10::optional optional_tensor; + c10::optional mapping_tensor; + if (lmp_list.mapping != nullptr) { + mapping_tensor = + torch::from_blob(lmp_list.mapping, {1, nall_real}, int32_options) + .to(torch::kInt64) + .to(device); + } c10::optional fparam_tensor; if (!fparam.empty()) { fparam_tensor = @@ -151,7 +157,7 @@ void DeepPotPT::compute(ENERGYVTYPE& ener, c10::Dict outputs = module .run_method("forward_lower", coord_wrapped_Tensor, atype_Tensor, - firstneigh_tensor, optional_tensor, fparam_tensor, + firstneigh_tensor, mapping_tensor, fparam_tensor, aparam_tensor, do_atom_virial_tensor) .toGenericDict(); c10::IValue energy_ = outputs.at("energy"); diff --git a/source/lib/include/neighbor_list.h b/source/lib/include/neighbor_list.h index eb510eb25b..6d4217986a 100644 --- a/source/lib/include/neighbor_list.h +++ b/source/lib/include/neighbor_list.h @@ -26,12 +26,23 @@ struct InputNlist { int* numneigh; /// Array stores the core region atom's neighbor index int** firstneigh; - InputNlist() : inum(0), ilist(NULL), numneigh(NULL), firstneigh(NULL){}; + + int* mapping; + InputNlist() + : inum(0), ilist(NULL), numneigh(NULL), firstneigh(NULL), mapping(NULL){}; InputNlist(int inum_, int* ilist_, int* numneigh_, int** firstneigh_) : inum(inum_), ilist(ilist_), numneigh(numneigh_), - firstneigh(firstneigh_){}; + firstneigh(firstneigh_), + mapping(NULL){}; + InputNlist( + int inum_, int* ilist_, int* numneigh_, int** firstneigh_, int* mapping) + : inum(inum_), + ilist(ilist_), + numneigh(numneigh_), + firstneigh(firstneigh_), + mapping(mapping){}; ~InputNlist(){}; }; diff --git a/source/lmp/pair_deepmd.cpp b/source/lmp/pair_deepmd.cpp index 90aa453143..b3073b73b1 100644 --- a/source/lmp/pair_deepmd.cpp +++ b/source/lmp/pair_deepmd.cpp @@ -462,6 +462,7 @@ void PairDeepMD::compute(int eflag, int vflag) { double **x = atom->x; double **f = atom->f; + int *tag_array = atom->tag; int *type = atom->type; int nlocal = atom->nlocal; int nghost = 0; @@ -484,6 +485,11 @@ void PairDeepMD::compute(int eflag, int vflag) { } } } + // make mapping array + int *mapping = new int[nall]; + for (int i = 0; i < nall; ++i) { + mapping[i] = atom->map(tag_array[i]); + } vector dtype(nall); for (int ii = 0; ii < nall; ++ii) { @@ -551,7 +557,7 @@ void PairDeepMD::compute(int eflag, int vflag) { (numb_models > 1 && (out_freq > 0 && update->ntimestep % out_freq == 0)); if (do_ghost) { deepmd_compat::InputNlist lmp_list(list->inum, list->ilist, list->numneigh, - list->firstneigh); + list->firstneigh, mapping); deepmd_compat::InputNlist extend_lmp_list; if (atom->sp_flag) { extend(extend_inum, extend_ilist, extend_numneigh, extend_neigh, @@ -1280,6 +1286,8 @@ void PairDeepMD::coeff(int narg, char **arg) { void PairDeepMD::init_style() { #if LAMMPS_VERSION_NUMBER >= 20220324 neighbor->add_request(this, NeighConst::REQ_FULL); + atom->map_user = 2; + atom->map_init(1); #else int irequest = neighbor->request(this, instance_me); neighbor->requests[irequest]->half = 0;