fix(pd): add default chg spin parameter#5469
Open
HydrogenSulfate wants to merge 4 commits into
Open
Conversation
… charge_spin Co-authored-by: HydrogenSulfate <23737287+HydrogenSulfate@users.noreply.github.com>
…t_tebd descriptors Co-authored-by: HydrogenSulfate <23737287+HydrogenSulfate@users.noreply.github.com>
…e_spin keyword Co-authored-by: HydrogenSulfate <23737287+HydrogenSulfate@users.noreply.github.com>
Contributor
📝 WalkthroughWalkthroughUniform charge-spin interface methods are added to multiple descriptors; DPA3 gains default value support, validation, serialization, and uses ChangesCharge-spin descriptor interface and DPA3 implementation
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Possibly related PRs
Suggested labels
Suggested reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@deepmd/pd/model/atomic_model/dp_atomic_model.py`:
- Line 335: When constructing the call that sets charge_spin (the line currently
using "charge_spin=fparam if self.add_chg_spin_ebd else None"), ensure that when
self.add_chg_spin_ebd is True but fparam is None you fall back to the configured
default_chg_spin instead of passing None; change the expression to use fparam if
fparam is not None else self.default_chg_spin (only when self.add_chg_spin_ebd
is True), so charge_spin becomes (fparam if fparam is not None else
self.default_chg_spin) when self.add_chg_spin_ebd else None and avoids
triggering the DPA3 assertion path.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 6949ec6c-cc26-4f0f-a800-8b4381111ee6
📒 Files selected for processing (6)
deepmd/pd/model/atomic_model/dp_atomic_model.pydeepmd/pd/model/descriptor/dpa1.pydeepmd/pd/model/descriptor/dpa2.pydeepmd/pd/model/descriptor/dpa3.pydeepmd/pd/model/descriptor/se_a.pydeepmd/pd/model/descriptor/se_t_tebd.py
| mapping=mapping, | ||
| comm_dict=comm_dict, | ||
| fparam=fparam if self.add_chg_spin_ebd else None, | ||
| charge_spin=fparam if self.add_chg_spin_ebd else None, |
Contributor
There was a problem hiding this comment.
Missing default charge_spin fallback breaks the new default contract.
When self.add_chg_spin_ebd is enabled and fparam is None, this passes charge_spin=None, which trips the DPA3 assertion path and ignores configured default_chg_spin.
💡 Suggested fix
nframes, nloc, nnei = nlist.shape
atype = extended_atype[:, :nloc]
if self.do_grad_r() or self.do_grad_c():
extended_coord.stop_gradient = False
+ charge_spin = fparam if self.add_chg_spin_ebd else None
+ if self.add_chg_spin_ebd and charge_spin is None:
+ default_cs = self.descriptor.get_default_chg_spin()
+ if default_cs is not None:
+ cs = paddle.to_tensor(default_cs, dtype=extended_coord.dtype).to(
+ device=extended_coord.place
+ )
+ charge_spin = paddle.tile(cs.reshape([1, -1]), [nframes, 1])
descriptor, rot_mat, g2, h2, sw = self.descriptor(
extended_coord,
extended_atype,
nlist,
mapping=mapping,
comm_dict=comm_dict,
- charge_spin=fparam if self.add_chg_spin_ebd else None,
+ charge_spin=charge_spin,
)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@deepmd/pd/model/atomic_model/dp_atomic_model.py` at line 335, When
constructing the call that sets charge_spin (the line currently using
"charge_spin=fparam if self.add_chg_spin_ebd else None"), ensure that when
self.add_chg_spin_ebd is True but fparam is None you fall back to the configured
default_chg_spin instead of passing None; change the expression to use fparam if
fparam is not None else self.default_chg_spin (only when self.add_chg_spin_ebd
is True), so charge_spin becomes (fparam if fparam is not None else
self.default_chg_spin) when self.add_chg_spin_ebd else None and avoids
triggering the DPA3 assertion path.
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## master #5469 +/- ##
==========================================
- Coverage 82.46% 81.35% -1.11%
==========================================
Files 829 868 +39
Lines 88763 96406 +7643
Branches 4225 4233 +8
==========================================
+ Hits 73197 78434 +5237
- Misses 14274 16672 +2398
- Partials 1292 1300 +8 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
Co-authored-by: HydrogenSulfate <23737287+HydrogenSulfate@users.noreply.github.com>
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (1)
source/tests/consistent/descriptor/common.py (1)
272-282: ⚡ Quick winRemove dead
fparam_pdassignment in PD descriptor helper
fparam_pdis computed fromfparambut never used/added tokwargson the PD path; it’s safe to remove for clarity (Ruff ignoresF841in this repo, so it won’t fail CI).♻️ Proposed cleanup
- fparam_pd = ( - paddle.to_tensor(fparam).to(PD_DEVICE) if fparam is not None else None - ) charge_spin_pd = ( paddle.to_tensor(charge_spin).to(PD_DEVICE) if charge_spin is not None else None ) kwargs = {"nlist": nlist, "mapping": mapping} if charge_spin_pd is not None: kwargs["charge_spin"] = charge_spin_pd🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@source/tests/consistent/descriptor/common.py` around lines 272 - 282, The temporary paddle tensor fparam_pd is created but never used; remove the unused assignment "fparam_pd = (paddle.to_tensor(fparam).to(PD_DEVICE) if fparam is not None else None)" from the PD descriptor helper so we don't allocate an unused tensor, and ensure only charge_spin_pd is computed/added to kwargs (leave PD_DEVICE, paddle usage and the kwargs mapping/nlist logic intact).
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@source/tests/consistent/descriptor/common.py`:
- Around line 272-282: The temporary paddle tensor fparam_pd is created but
never used; remove the unused assignment "fparam_pd =
(paddle.to_tensor(fparam).to(PD_DEVICE) if fparam is not None else None)" from
the PD descriptor helper so we don't allocate an unused tensor, and ensure only
charge_spin_pd is computed/added to kwargs (leave PD_DEVICE, paddle usage and
the kwargs mapping/nlist logic intact).
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 7a52faf4-1efb-473c-9cc2-1c3928995a53
📒 Files selected for processing (3)
source/tests/consistent/descriptor/common.pysource/tests/consistent/descriptor/test_dpa3.pysource/tests/pd/model/test_dpa3.py
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
| pd_obj.get_sel(), | ||
| distinguish_types=(not mixed_types), | ||
| ) | ||
| fparam_pd = ( |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This pull request introduces support for handling
charge_spininformation in atomic and descriptor models. The main changes add new methods to query and provide defaultcharge_spinvalues, update theforwardmethods to accept acharge_spinargument, and refactor code to use this argument consistently instead of the previousfparamfield. Additionally, the DPA3 descriptor now supports defaultcharge_spinvalues and serializes them.Support for charge_spin in descriptor and atomic models:
dpa1.py,dpa2.py,dpa3.py,se_a.py,se_t_tebd.py) now implementget_dim_chg_spin,has_default_chg_spin, andget_default_chg_spinmethods to standardize how charge and spin information is queried. [1] [2] [3] [4] [5]forwardmethods in all descriptor classes, as well as in the atomic model (dp_atomic_model.py), now accept acharge_spinargument instead of (or in addition to)fparam, and all internal logic is updated to usecharge_spin. [1] [2] [3] [4] [5] [6]DPA3 descriptor enhancements:
default_chg_spinparameter, validates its shape, and serializes it as part of its configuration. [1] [2] [3]charge_spinargument.These changes standardize how charge and spin information is handled across descriptors and atomic models, making the codebase more extensible and robust for future features involving charge and spin.
Summary by CodeRabbit