Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions source/api_cc/include/common.h
Original file line number Diff line number Diff line change
Expand Up @@ -102,6 +102,9 @@ void select_real_atoms_coord(std::vector<VALUETYPE>& dcoord,
const int& nall,
const bool aparam_nall = false);

void select_real_atoms_sendlist(const deepmd::InputNlist& inlist,
const std::vector<int>& fwd_map);

/**
* @brief Apply the given map to a vector.
* @param[out] out The output vector.
Expand Down
1 change: 1 addition & 0 deletions source/api_cc/src/DeepPotPD.cc
Original file line number Diff line number Diff line change
Expand Up @@ -393,6 +393,7 @@ void DeepPotPD::compute(ENERGYVTYPE& ener,
auto sendlist_tensor = predictor_fl->GetInputHandle("send_list");

int nswap = lmp_list.nswap;
select_real_atoms_sendlist(lmp_list, fwd_map);
sendproc_tensor->Reshape({nswap});
sendproc_tensor->CopyFromCpu(lmp_list.sendproc);

Expand Down
1 change: 1 addition & 0 deletions source/api_cc/src/DeepPotPT.cc
Original file line number Diff line number Diff line change
Expand Up @@ -174,6 +174,7 @@ void DeepPotPT::compute(ENERGYVTYPE& ener,
nlist_data.padding();
if (do_message_passing) {
int nswap = lmp_list.nswap;
select_real_atoms_sendlist(lmp_list, fwd_map);
torch::Tensor sendproc_tensor =
torch::from_blob(lmp_list.sendproc, {nswap}, int32_option);
torch::Tensor recvproc_tensor =
Expand Down
1 change: 1 addition & 0 deletions source/api_cc/src/DeepSpinPT.cc
Original file line number Diff line number Diff line change
Expand Up @@ -182,6 +182,7 @@ void DeepSpinPT::compute(ENERGYVTYPE& ener,
nlist_data.padding();
if (do_message_passing) {
int nswap = lmp_list.nswap;
select_real_atoms_sendlist(lmp_list, fwd_map);
torch::Tensor sendproc_tensor =
torch::from_blob(lmp_list.sendproc, {nswap}, int32_option);
torch::Tensor recvproc_tensor =
Expand Down
25 changes: 25 additions & 0 deletions source/api_cc/src/common.cc
Original file line number Diff line number Diff line change
Expand Up @@ -232,6 +232,31 @@ template void deepmd::select_real_atoms_coord<float>(
const int& nall,
const bool aparam_nall);

void deepmd::select_real_atoms_sendlist(const deepmd::InputNlist& inlist,
const std::vector<int>& fwd_map) {
int nswap = inlist.nswap;
std::vector<std::vector<int>> sendlist_new;
sendlist_new.resize(nswap);
// select real atoms in sendlist
for (int s = 0; s < nswap; ++s) {
int cnt = 0;
sendlist_new[s].reserve(inlist.sendnum[s]);
for (int k = 0; k < inlist.sendnum[s]; ++k) {
const int old_idx = inlist.sendlist[s][k];
int mapped = (old_idx >= 0 && old_idx < (int)fwd_map.size())
? fwd_map[old_idx]
: -1;
if (mapped >= 0) {
sendlist_new[s].push_back(mapped);
++cnt;
}
}
std::memcpy(inlist.sendlist[s], sendlist_new[s].data(), cnt * sizeof(int));
inlist.sendnum[s] = cnt;
inlist.recvnum[s] = cnt;
}
}

void deepmd::NeighborListData::copy_from_nlist(const InputNlist& inlist,
const int natoms) {
int inum = natoms >= 0 ? natoms : inlist.inum;
Expand Down
6 changes: 6 additions & 0 deletions source/lmp/tests/test_lammps_dpa_pt.py
Original file line number Diff line number Diff line change
Expand Up @@ -475,6 +475,12 @@ def test_pair_deepmd_type_map(lammps_type_map) -> None:
lammps_type_map.run(1)


def test_pair_deepmd_type_map_with_null(lammps_type_map) -> None:
lammps_type_map.pair_style(f"deepmd {pb_file.resolve()}")
lammps_type_map.pair_coeff("* * H NULL")
lammps_type_map.run(0)


def test_pair_deepmd_real(lammps_real) -> None:
lammps_real.pair_style(f"deepmd {pb_file.resolve()}")
lammps_real.pair_coeff("* *")
Expand Down
Loading