fix(jax): handle DPA-2 pbc/nopbc without mapping #4363
Conversation
In the C++ API, generate the mapping for the no PBC and throw the error for PBC. Considering I forgot setting `atom_modify map yes` when testing it, others may also forget. Signed-off-by: Jinzhe Zeng <[email protected]>
📝 Walkthrough📝 WalkthroughWalkthroughThe pull request introduces 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: 3
🧹 Outside diff range and nitpick comments (3)
source/api_cc/include/DeepPotJAX.h (1)
192-193: Consider documenting the new member variable.The addition of
do_message_passingis well-placed and follows the class's conventions. Consider adding a brief comment explaining its purpose for better maintainability.+ // whether to perform message passing for periodic boundary conditions bool do_message_passing;deepmd/jax/jax2tf/serialization.py (1)
297-302: Add docstring to the new function.The implementation looks good and follows the established patterns in the file. However, consider adding a docstring to maintain consistency with other functions and improve maintainability.
@tf.function def has_message_passing() -> tf.Tensor: + """Return whether the model has message passing capabilities. + + Returns + ------- + tf.Tensor + Boolean tensor indicating if the model supports message passing. + """ return tf.constant(model.has_message_passing(), dtype=tf.bool)source/api_cc/src/DeepPotJAX.cc (1)
599-608: Provide clear user guidance when mapping is requiredThe code correctly throws an exception when
do_message_passingistruebut mapping is not provided. The exception message advises the user to setatom_modify map yesin LAMMPS, which is helpful. Ensure that this requirement is also documented in the user guides or documentation to prevent confusion.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
📒 Files selected for processing (3)
deepmd/jax/jax2tf/serialization.py(1 hunks)source/api_cc/include/DeepPotJAX.h(1 hunks)source/api_cc/src/DeepPotJAX.cc(3 hunks)
Signed-off-by: Jinzhe Zeng <[email protected]>
There was a problem hiding this comment.
Copilot reviewed 1 out of 3 changed files in this pull request and generated no suggestions.
Files not reviewed (2)
- source/api_cc/include/DeepPotJAX.h: Language not supported
- source/api_cc/src/DeepPotJAX.cc: Language not supported
Comments skipped due to low confidence (1)
deepmd/jax/jax2tf/serialization.py:298
- The new function
has_message_passingdoes not have explicit test coverage. Consider adding tests to ensure its correct behavior.
@tf.function
njzjz
changed the title
Signed-off-by: Jinzhe Zeng <[email protected]>
Codecov ReportAttention: Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## devel #4363 +/- ##
==========================================
+ Coverage 84.50% 84.54% +0.03%
==========================================
Files 596 597 +1
Lines 56665 56813 +148
Branches 3459 3487 +28
==========================================
+ Hits 47884 48031 +147
- Misses 7654 7655 +1
Partials 1127 1127 ☔ View full report in Codecov by Sentry. 🚨 Try these New Features:
|
Signed-off-by: Jinzhe Zeng <[email protected]>
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Outside diff range and nitpick comments (4)
source/api_cc/tests/test_deeppot_dpa_jax.cc (4)
16-18: Improve documentation for EPSILON definition.The comment "1e-10 cannot pass; unclear bug or not" lacks sufficient context. Please document:
- The specific test cases that fail with 1e-10
- The observed behavior vs expected behavior
- Any investigation results or hypotheses about the root cause
28-43: Consider moving Python code to a separate file.The Python code for generating test data could be moved to a separate script file and referenced in the comments. This would:
- Improve maintainability
- Make it easier to regenerate test data
- Allow version control of the data generation script
80-80: Consider making the model path configurable.The hardcoded path "../../tests/infer/deeppot_dpa.savedmodel" could make the tests brittle. Consider:
- Using an environment variable
- Making it a configurable parameter
- Using a relative path from the test binary location
329-331: Consider generating neighbor list data programmatically.The neighbor list data is currently hardcoded. Consider:
- Generating it programmatically based on distance criteria
- Making it more maintainable and less error-prone
- Adding comments explaining the neighbor selection criteria
Example:
std::vector<std::vector<int>> generate_nlist_data( const std::vector<VALUETYPE>& coord, int natoms, VALUETYPE cutoff) { std::vector<std::vector<int>> nlist_data(natoms); // Generate based on distance criteria return nlist_data; }Also applies to: 365-367
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
📒 Files selected for processing (1)
source/api_cc/tests/test_deeppot_dpa_jax.cc(1 hunks)
🧰 Additional context used
📓 Learnings (1)
source/api_cc/tests/test_deeppot_dpa_jax.cc (2)
Learnt from: njzjz
PR: deepmodeling/deepmd-kit#4144
File: source/api_cc/tests/test_deeppot_dpa_pt.cc:166-246
Timestamp: 2024-11-12T05:47:21.643Z
Learning: Refactoring between test classes `TestInferDeepPotDpaPt` and `TestInferDeepPotDpaPtNopbc` is addressed in PR #3905.
Learnt from: njzjz
PR: deepmodeling/deepmd-kit#4144
File: source/api_cc/tests/test_deeppot_dpa_pt.cc:166-246
Timestamp: 2024-11-12T05:47:21.643Z
Learning: Refactoring between test classes `TestInferDeepPotDpaPt` and `TestInferDeepPotDpaPtNopbc` is addressed in PR #3905.
In the C++ API, generate the mapping for the no PBC and throw the error for PBC. Considering I forgot setting
atom_modify map yeswhen testing it, others may also forget.Summary by CodeRabbit
New Features
DeepPotclass, validating its functionality under various conditions.Bug Fixes