feat(pt): pass mapping from LAMMPS to PT #4351
Conversation
This will be useful for the external MACE model. Signed-off-by: Jinzhe Zeng <[email protected]>
Signed-off-by: Jinzhe Zeng <[email protected]>
📝 Walkthrough📝 WalkthroughWalkthroughThe pull request introduces several changes across multiple files to enhance the functionality and documentation of the API related to neighbor list management. The API version is incremented from 24 to 25, and a new function, Changes
Possibly related PRs
Suggested labels
Suggested reviewers
Thank you for using CodeRabbit. We offer it for free to the OSS community and would appreciate your support in helping us grow. If you find it useful, would you consider giving us a shout-out on your favorite social media? 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Outside diff range and nitpick comments (5)
source/lib/include/neighbor_list.h (1)
Line range hint
47-107: Consider architectural improvements for the mapping feature.While the implementation adds the required functionality, consider these architectural improvements:
Memory Safety:
- Consider using
std::shared_ptrorstd::unique_ptrfor better memory management- Add RAII principles to prevent memory leaks
API Design:
- Consider making this a breaking change with a new version number
- Add validation methods to ensure mapping integrity
Thread Safety:
- Document thread safety guarantees
- Consider adding synchronization if concurrent access is expected
Example smart pointer implementation:
private: std::shared_ptr<int[]> mapping; public: void set_mapping(std::shared_ptr<int[]> mapping_) { if (!mapping_) { throw std::invalid_argument("mapping cannot be null"); } mapping = std::move(mapping_); }source/api_cc/src/DeepPotPT.cc (1)
211-213: Consider adding error handling for tensor creationThe tensor creation could potentially fail if memory allocation fails. Consider wrapping it in a try-catch block.
- mapping_tensor = - torch::from_blob(mapping.data(), {1, nall_real}, int_option) - .to(device); + try { + mapping_tensor = + torch::from_blob(mapping.data(), {1, nall_real}, int_option) + .to(device); + } catch (const c10::Error& e) { + throw deepmd::deepmd_exception("Failed to create mapping tensor: " + + std::string(e.what())); + }source/lmp/pair_deepmd.cpp (1)
209-211: LGTM! Consider adding documentation.The mapping is correctly set in the InputNlist object with consistent conditions. Consider adding a comment explaining the purpose of this mapping for DPA-2 JAX and why it's only needed for single process.
Add documentation above the condition:
+ // Set mapping for DPA-2 JAX interface + // Only needed for single process to avoid complexity of distributed memory if (comm->nprocs == 1 && atom->map_style != Atom::MAP_NONE) { lmp_list.set_mapping(mapping_vec.data()); }source/api_c/include/c_api.h (1)
81-89: Documentation could be more descriptive.While the documentation follows the established style, consider enhancing the description of the
mappingparameter to:
- Explain the purpose of mapping from all atoms to real atoms
- Clarify the expected values/format of the mapping array
- Add an @return tag to document the behavior on success/failure
Apply this diff to improve the documentation:
/** * @brief Set mapping for a neighbor list. * * @param nl Neighbor list. - * @param mapping mapping from all atoms to real atoms, in size nall. + * @param mapping Array that maps indices from all atoms to real atoms. The array + * should be of size nall, where nall is the total number of atoms + * including ghost atoms. Each element maps an atom index to its + * corresponding real atom index. * @since API version 25 + * @return void * **/source/api_c/include/deepmd.hpp (1)
868-868: Consider enhancing the documentation with validation details.While the parameter documentation is clear, it would be helpful to add:
- Expected validation behavior for invalid mappings
- Whether null mappings are allowed
- Any assumptions about the mapping array's lifetime
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
📒 Files selected for processing (7)
source/api_c/include/c_api.h(4 hunks)source/api_c/include/deepmd.hpp(1 hunks)source/api_c/src/c_api.cc(1 hunks)source/api_cc/src/DeepPotPT.cc(1 hunks)source/lib/include/neighbor_list.h(2 hunks)source/lmp/fix_dplr.cpp(2 hunks)source/lmp/pair_deepmd.cpp(2 hunks)
🔇 Additional comments (9)
source/lib/include/neighbor_list.h (1)
47-48: 🛠️ Refactor suggestion
Add size validation and improve documentation for the mapping member.
The comment indicates the mapping size should be nall, but there's no member variable storing nall to validate this at runtime. Consider:
- Adding
nallas a member variable - Adding size validation in
set_mapping - Documenting the ownership semantics of the pointer
Let's verify if nall is available elsewhere in the codebase:
/// mapping from all atoms to real atoms, in the size of nall
-int* mapping = nullptr;
+/// @note Caller retains ownership of the mapping pointer
+/// @param nall Total number of atoms (including ghost atoms)
+int nall = 0;
+int* mapping = nullptr; // Size must match nallsource/api_cc/src/DeepPotPT.cc (2)
206-214: LGTM: Mapping implementation looks correct
The implementation properly handles the mapping from LAMMPS to PT by:
- Creating a mapping vector with the correct size
- Applying the forward mapping indices
- Converting to a tensor and moving it to the appropriate device
206-214: Verify error handling for null mapping
Consider adding a null check for lmp_list.mapping to prevent potential segmentation faults.
Consider adding a null check:
- if (lmp_list.mapping) {
+ if (lmp_list.mapping != nullptr) {source/lmp/pair_deepmd.cpp (1)
158-164: LGTM! Verify atom mapping configuration.
The mapping initialization is correctly implemented with proper boundary checks. The -1 initialization provides a safe default for invalid indices.
Let's verify the atom mapping configuration:
✅ Verification successful
The mapping implementation is correct and safe
The code's handling of atom mapping is properly implemented with appropriate checks:
- The mapping is only used when
map_style != Atom::MAP_NONE, preventing invalid access - The single process check (
comm->nprocs == 1) is consistent with other similar checks in the codebase (fix_dplr.cpp) - The mapping is properly used in the InputNlist with the same safety checks
- The example input scripts show proper atom mapping configuration with
atom_modify map array
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Description: Check atom mapping style configuration in input scripts
# Expect: atom_style and atom_modify commands that set up mapping
# Search for atom_style commands
echo "Checking atom_style configuration:"
rg "atom_style"
# Search for atom_modify map commands
echo "Checking atom_modify map configuration:"
rg "atom_modify.*map"
Length of output: 1958
Script:
#!/bin/bash
# Let's check the LAMMPS documentation or implementation for default mapping behavior
# and verify how atom mapping is handled in pair_deepmd.cpp
# Check for any documentation about atom mapping in the codebase
echo "Checking documentation about atom mapping:"
rg -i "atom.*map" "doc/" -A 3
# Check how atom mapping is used in pair_deepmd.cpp
echo "Checking atom mapping usage in pair_deepmd.cpp:"
rg "map.*style|map\(" "source/lmp/pair_deepmd.cpp" -A 2
# Check for any other relevant atom map handling in the codebase
echo "Checking other atom map handling:"
rg "map.*style|MAP_NONE" "source/lmp/"
Length of output: 37849
source/api_c/include/c_api.h (2)
15-15: API version increment is appropriate.
The increment from 24 to 25 correctly reflects the addition of the new DP_NlistSetMapping function to the API.
Line range hint 34-70: Documentation improvements are comprehensive.
The enhanced documentation for DP_NewNlist_comm is well-structured and thorough:
- Clear @details section explaining the function's purpose
- Well-documented parameters with descriptive explanations
- Proper formatting and consistent style
source/api_c/src/c_api.cc (1)
46-48: Verify the mapping implementation across the codebase.
The new mapping functionality needs to be consistently implemented across the codebase.
Run the following script to verify the mapping implementation:
✅ Verification successful
Based on the search results, I can now provide a final response:
The mapping implementation is consistent across the codebase
The mapping functionality is properly implemented throughout the codebase:
- In C API (
source/api_c/src/c_api.cc),DP_NlistSetMappingsets the mapping for neighbor lists - The mapping is consistently used in neighbor list handling across CPU/GPU implementations in
source/lib/src/coord.ccandsource/lib/src/gpu/coord.cu - The mapping is properly propagated through model implementations in PyTorch, JAX and TensorFlow
- Test coverage exists for mapping functionality across different model types and architectures
- The mapping is used consistently for:
- Neighbor list construction
- Force/energy/virial calculations
- Handling periodic boundary conditions
- Converting between extended and local indices
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Description: Verify the mapping implementation across the codebase.
# Test 1: Check for mapping-related functions and methods
echo "Checking for mapping-related functions and methods..."
rg -A 5 "set_mapping|mapping"
# Test 2: Check for InputNlist usage with mapping
echo "Checking for InputNlist usage with mapping..."
ast-grep --pattern $'InputNlist\s*\w+\s*\([^)]*\)'
Length of output: 463817
source/api_c/include/deepmd.hpp (2)
866-870: LGTM! The new mapping functionality is well integrated.
The new set_mapping method is correctly implemented and properly documented. The implementation appropriately delegates to the C API function DP_NlistSetMapping.
866-870: Verify the mapping functionality across the codebase.
Let's verify that this new mapping capability is properly handled throughout the codebase.
✅ Verification successful
Based on the search results, I can now provide a final response:
The mapping functionality is well-integrated and thoroughly tested across the codebase
The mapping capability is properly implemented and extensively tested:
- Core implementation is in
source/api_c/src/c_api.ccwithDP_NlistSetMapping - Comprehensive test coverage exists across multiple test files:
- Unit tests in
source/api_c/tests/covering various scenarios - Integration tests in
source/tests/validating mapping with different models - Cross-framework tests ensuring consistent behavior in TensorFlow, PyTorch and JAX
- Unit tests in
The mapping functionality is used consistently for:
- Neighbor list management
- Coordinate transformations
- Force/energy calculations
- Model evaluations
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Search for other uses of DP_NlistSetMapping to ensure consistent usage
rg "DP_NlistSetMapping"
# Search for any existing mapping-related functionality
ast-grep --pattern 'mapping'
# Look for tests covering this functionality
fd -e cc -e cpp -e py "test.*mapping"
Length of output: 85195
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.
Summary by CodeRabbit
Release Notes
New Features
DP_NlistSetMappingfor setting mappings in neighbor lists.set_mappingmethod toInputNlistfor mapping atoms to real atoms.computemethods inDeepPotPT,PairDeepMD, andFixDPLRclasses to support new mapping functionalities.Bug Fixes
Documentation