Skip to content

Commit

Permalink
feat(pt): pass mapping from LAMMPS to PT (#4351)
Browse files Browse the repository at this point in the history
As discussed, this PR passes mapping from LAMMPS to the PT C++
interface, which is helpful for the external GNN models.

The mapping interface is synced from #4307.

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

## Release Notes

- **New Features**
- Introduced a new function `DP_NlistSetMapping` for setting mappings in
neighbor lists.
- Added `set_mapping` method to `InputNlist` for mapping atoms to real
atoms.
- Enhanced `compute` methods in `DeepPotPT`, `PairDeepMD`, and `FixDPLR`
classes to support new mapping functionalities.

- **Bug Fixes**
- Improved error handling in various classes to ensure robustness during
execution.

- **Documentation**
- Updated and added comments for clarity and consistency in new and
existing functions.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Jinzhe Zeng <[email protected]>
  • Loading branch information
njzjz authored Nov 13, 2024
1 parent 9779454 commit d5295d5
Showing 1 changed file with 9 additions and 0 deletions.
9 changes: 9 additions & 0 deletions source/api_cc/src/DeepPotPT.cc
Original file line number Diff line number Diff line change
Expand Up @@ -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<std::int64_t> 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);
Expand Down

0 comments on commit d5295d5

Please sign in to comment.