Skip to content

fix(pd): add default chg spin parameter#5469

Open
HydrogenSulfate wants to merge 3 commits into
deepmodeling:masterfrom
HydrogenSulfate:copilot/add-default-chg-spin-parameter
Open

fix(pd): add default chg spin parameter#5469
HydrogenSulfate wants to merge 3 commits into
deepmodeling:masterfrom
HydrogenSulfate:copilot/add-default-chg-spin-parameter

Conversation

@HydrogenSulfate
Copy link
Copy Markdown
Collaborator

@HydrogenSulfate HydrogenSulfate commented May 27, 2026

This pull request introduces support for handling charge_spin information in atomic and descriptor models. The main changes add new methods to query and provide default charge_spin values, update the forward methods to accept a charge_spin argument, and refactor code to use this argument consistently instead of the previous fparam field. Additionally, the DPA3 descriptor now supports default charge_spin values and serializes them.

Support for charge_spin in descriptor and atomic models:

  • All descriptor classes (dpa1.py, dpa2.py, dpa3.py, se_a.py, se_t_tebd.py) now implement get_dim_chg_spin, has_default_chg_spin, and get_default_chg_spin methods to standardize how charge and spin information is queried. [1] [2] [3] [4] [5]
  • The forward methods in all descriptor classes, as well as in the atomic model (dp_atomic_model.py), now accept a charge_spin argument instead of (or in addition to) fparam, and all internal logic is updated to use charge_spin. [1] [2] [3] [4] [5] [6]

DPA3 descriptor enhancements:

  • The DPA3 descriptor now supports a default_chg_spin parameter, validates its shape, and serializes it as part of its configuration. [1] [2] [3]
  • The logic for embedding charge and spin in DPA3 is updated to use the new charge_spin argument.

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

  • New Features
    • Extended charge and spin embedding support across all model descriptors with a standardized interface for consistency
    • Added configurable default charge/spin values option for enhanced descriptor flexibility

Review Change Stack

Copilot AI and others added 3 commits May 27, 2026 03:30
… 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>
Copilot AI review requested due to automatic review settings May 27, 2026 09:21
@dosubot dosubot Bot added the enhancement label May 27, 2026
@coderabbitai
Copy link
Copy Markdown
Contributor

coderabbitai Bot commented May 27, 2026

📝 Walkthrough

Walkthrough

This 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.

Changes

Charge-spin descriptor interface and DPA3 implementation

Layer / File(s) Summary
Non-supporting descriptors expose charge-spin interface
deepmd/pd/model/descriptor/dpa1.py, deepmd/pd/model/descriptor/dpa2.py, deepmd/pd/model/descriptor/se_a.py, deepmd/pd/model/descriptor/se_t_tebd.py
DPA1, DPA2, SeA, and SeTTebd each add get_dim_chg_spin(), has_default_chg_spin(), and get_default_chg_spin() methods indicating charge-spin is not supported (dimension 0, no defaults). All four update their forward() signatures to accept an optional charge_spin parameter for API uniformity.
DPA3 descriptor implements charge-spin support
deepmd/pd/model/descriptor/dpa3.py
DPA3 adds an optional default_chg_spin constructor parameter with validation (must be two floats if provided), stores and persists it via serialization, and exposes it through the query methods. When add_chg_spin_ebd is enabled, forward() accepts charge_spin input, asserts it is present, and computes embedded charge and spin values before adding them to node embeddings.
DPAtomicModel uses new charge-spin interface
deepmd/pd/model/atomic_model/dp_atomic_model.py
DPAtomicModel.forward_atomic() now routes charge/spin data through the descriptor's charge_spin keyword argument (when add_chg_spin_ebd is enabled) instead of the previous fparam mechanism, aligning with the new interface contract.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~12 minutes

Possibly related PRs

  • deepmodeling/deepmd-kit#5333: Both PRs update the DP atomic model/descriptor wiring for charge-spin embedding—specifically dp_atomic_model.forward_atomic conditionally passes the charge-spin-related parameter into the descriptor (retrieved PR via fparam, main PR via charge_spin) and this corresponds to the DescrptDPA3 add_chg_spin_ebd feature.
  • deepmodeling/deepmd-kit#5431: Both PRs implement the same DPA3 API decoupling: DPAtomicModel.forward_atomic/forward stop deriving/forwarding charge-spin via fparam and instead pass a charge_spin kwarg into the DPA3 descriptor (with PD doing this in deepmd/pd/... counterparts).

Suggested labels

Python

Suggested reviewers

  • njzjz
  • wanghan-iapcm
  • iProzd
🚥 Pre-merge checks | ✅ 5
✅ Passed checks (5 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The title 'fix(pd): add default chg spin parameter' accurately describes the main change: adding support for charge/spin parameters across descriptor and atomic model classes.
Docstring Coverage ✅ Passed Docstring coverage is 87.50% which is sufficient. The required threshold is 80.00%.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.

✏️ Tip: You can configure your own custom pre-merge checks in the settings.

✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests

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.

❤️ Share

Comment @coderabbitai help to get the list of available commands and usage tips.

Copy link
Copy Markdown
Contributor

@coderabbitai coderabbitai Bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

📥 Commits

Reviewing files that changed from the base of the PR and between 4e64f8b and db26f5e.

📒 Files selected for processing (6)
  • deepmd/pd/model/atomic_model/dp_atomic_model.py
  • deepmd/pd/model/descriptor/dpa1.py
  • deepmd/pd/model/descriptor/dpa2.py
  • deepmd/pd/model/descriptor/dpa3.py
  • deepmd/pd/model/descriptor/se_a.py
  • deepmd/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,
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major | ⚡ Quick win

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
Copy link
Copy Markdown

codecov Bot commented May 27, 2026

Codecov Report

❌ Patch coverage is 47.22222% with 19 lines in your changes missing coverage. Please review.
✅ Project coverage is 82.34%. Comparing base (f39a081) to head (db26f5e).
⚠️ Report is 1 commits behind head on master.

Files with missing lines Patch % Lines
deepmd/pd/model/descriptor/dpa3.py 41.66% 7 Missing ⚠️
deepmd/pd/model/descriptor/dpa1.py 50.00% 3 Missing ⚠️
deepmd/pd/model/descriptor/dpa2.py 50.00% 3 Missing ⚠️
deepmd/pd/model/descriptor/se_a.py 50.00% 3 Missing ⚠️
deepmd/pd/model/descriptor/se_t_tebd.py 50.00% 3 Missing ⚠️
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.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
  • 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants