fix(pd): add default chg spin parameter#5469
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>
📝 WalkthroughWalkthroughThis PR updates the charge/spin input contract across multiple PaddlePaddle descriptor classes by introducing a uniform interface. Four descriptors (DPA1, DPA2, SeA, SeTTebd) declare this interface as unsupported, DPA3 implements it with default value support and embedding logic, and DPAtomicModel is updated to use it. ChangesCharge-spin descriptor interface and DPA3 implementation
Estimated code review effort🎯 2 (Simple) | ⏱️ ~12 minutes Possibly related PRs
Suggested labels
Suggested reviewers
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 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 |
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, |
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% 82.34% -0.13%
==========================================
Files 829 829
Lines 88763 88796 +33
Branches 4225 4225
==========================================
- Hits 73197 73116 -81
- Misses 14274 14387 +113
- Partials 1292 1293 +1 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
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
Release Notes