fix(dp/pt): support auto sel for dpa2 #4323
Conversation
📝 WalkthroughWalkthroughThis pull request introduces enhancements to the 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: 1
🧹 Outside diff range and nitpick comments (4)
source/tests/pt/test_update_sel.py (2)
177-194: Add docstring to explain test configuration.Consider adding a docstring to explain the dpa2 descriptor configuration, particularly the relationship between repinit and repformer parameters, and the significance of the different rcut values.
@patch("deepmd.pt.utils.update_sel.UpdateSel.get_nbor_stat") def test_update_sel_dpa2_auto(self, sel_mock): + """Test update_sel with dpa2 descriptor using auto selection. + + Tests the conversion of 'auto' selection parameters in both repinit + (including three-body interactions) and repformer configurations. + The repinit uses a larger rcut (6.0) than repformer (4.0) to capture + different interaction ranges. + """
213-214: Add specific assertions for individual parameters.While comparing the entire structure is good, consider adding specific assertions for individual parameters to:
- Make test failures more descriptive
- Explicitly verify each auto-selection conversion
jdata = update_sel(jdata) self.assertEqual(jdata, expected_out) +# Add specific assertions +descriptor = jdata["model"]["descriptor"] +self.assertEqual(descriptor["repinit"]["nsel"], 28, "repinit nsel not converted correctly") +self.assertEqual(descriptor["repinit"]["three_body_sel"], 28, "three_body_sel not converted correctly") +self.assertEqual(descriptor["repformer"]["nsel"], 28, "repformer nsel not converted correctly")deepmd/pt/model/descriptor/dpa2.py (1)
845-852: LGTM! Consider adding validation for three-body selection parameters.The implementation correctly adds support for three-body interaction selection parameters, following the same pattern as the existing two-body selection updates. The code is well-structured and maintains consistency with the existing implementation.
Consider adding validation to ensure that the three-body selection parameters are compatible with the two-body parameters. For example:
min_nbor_dist, repinit_three_body_sel = update_sel.update_one_sel( train_data, type_map, local_jdata_cpy["repinit"]["three_body_rcut"], local_jdata_cpy["repinit"]["three_body_sel"], True, ) +# Validate three-body selection parameters +if repinit_three_body_sel[0] > repinit_sel[0]: + raise ValueError("Three-body selection count cannot be larger than two-body selection count") local_jdata_cpy["repinit"]["three_body_sel"] = repinit_three_body_sel[0]deepmd/dpmodel/descriptor/dpa2.py (1)
1060-1067: Add error handling for missing parameters.Consider adding validation to ensure
three_body_rcutandthree_body_selexist in the input configuration when needed.+ if "three_body_rcut" not in local_jdata_cpy["repinit"] or "three_body_sel" not in local_jdata_cpy["repinit"]: + raise KeyError("Missing required three-body parameters in configuration") min_nbor_dist, repinit_three_body_sel = update_sel.update_one_sel( train_data, type_map, local_jdata_cpy["repinit"]["three_body_rcut"], local_jdata_cpy["repinit"]["three_body_sel"], True, )
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
📒 Files selected for processing (4)
deepmd/dpmodel/descriptor/dpa2.py(1 hunks)deepmd/pt/model/descriptor/dpa2.py(1 hunks)deepmd/utils/argcheck.py(5 hunks)source/tests/pt/test_update_sel.py(1 hunks)
🔇 Additional comments (4)
source/tests/pt/test_update_sel.py (1)
173-175: LGTM! Test setup follows established patterns.
The test method setup and mock configuration are consistent with other test methods in the file.
deepmd/dpmodel/descriptor/dpa2.py (1)
1060-1067: LGTM! Verify three-body parameters existence.
The implementation correctly updates the three-body selection parameters using the same pattern as other selections. However, we should verify that the required parameters are always present in the input configuration.
✅ Verification successful
Parameters are well-defined with proper defaults and documentation
The three-body parameters are properly integrated in the codebase:
- Default values are defined in
RepinitArgsclass andargcheck.py - Parameters are documented with clear descriptions
- Example configurations in
/examples/water/dpa2/demonstrate correct usage - Test files show both "auto" and fixed value configurations
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Description: Check if three-body parameters are properly documented in example configs
# and if there are any configs missing these parameters
# Test: Look for three-body parameter usage in configuration files
rg -A 5 'three_body_(rcut|sel)'
Length of output: 13153
deepmd/utils/argcheck.py (2)
943-945: Documentation looks good!
The documentation for selection parameters (nsel and three_body_sel) is clear, consistent, and provides good examples. The auto-selection feature is well explained.
Also applies to: 977-979, 1087-1089
987-987: Implementation is correct and consistent.
The implementation of selection parameters:
- Uses correct type hints ([int, str])
- Is properly integrated into the argument groups
- Maintains consistency with the existing codebase
Let's verify the usage of these parameters in the codebase:
Also applies to: 1073-1077, 1152-1152
✅ Verification successful
Based on the search results, I can now provide a final response:
Implementation is correct and consistent with codebase usage
The nsel and three_body_sel parameters are properly implemented and used consistently throughout the codebase:
- The parameters accept both
intandstrtypes as defined in argcheck.py - They are used extensively in descriptor implementations (dpa1.py, dpa2.py) and neighbor list utilities
- The parameters are properly documented with clear descriptions
- Test files verify their functionality
- Example configurations demonstrate proper usage with recommended values
The implementation maintains consistency with:
- Type checking
- Documentation standards
- Integration with the descriptor architecture
- Usage patterns across the codebase
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Description: Check the usage of nsel and three_body_sel parameters
# Test: Search for parameter usage in descriptor implementations
rg -A 5 "nsel|three_body_sel" --type python
Length of output: 76
Script:
#!/bin/bash
# Search for nsel parameter usage without file type restriction
rg -A 5 "nsel"
Length of output: 130805
Script:
#!/bin/bash
# Search for three_body_sel parameter usage without file type restriction
rg -A 5 "three_body_sel"
Length of output: 10019
Script:
#!/bin/bash
# Find Python files that might contain these parameters
fd ".py$" --exec grep -l "nsel\|three_body_sel" {}
Length of output: 2600
Fix #4314 .
Summary by CodeRabbit
New Features
Bug Fixes
Documentation
Tests
dpa2model.