feat(pt/dpa4): Support property fitting in DPA4/SeZM models#5587
Conversation
OutisLi
requested review from
iProzd,
njzjz and
wanghan-iapcm
and removed request for
iProzd,
njzjz and
wanghan-iapcm
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Repository UI Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (3)
✅ Files skipped from review due to trivial changes (2)
📝 WalkthroughWalkthroughThe PR adds invariant property fitting support for DPA4/SeZM models, including a new property model class, factory and schema support, adjusted output-stat handling and non-conservative compute paths, plus docs, examples, and tests. ChangesSeZM property fitting
Sequence Diagram(s)sequenceDiagram
participant Factory as get_sezm_model
participant PropertyModel as SeZMPropertyModel
participant CoreCompute as SeZMModel.core_compute
participant OutputMap as fit_output_to_model_output
Factory->>PropertyModel: construct model when fitting_net.type="property"
PropertyModel->>CoreCompute: core_compute(..., conservative=false)
CoreCompute->>OutputMap: convert reduced outputs without force/virial graph
PropertyModel->>PropertyModel: translate atom_{var_name} and var_name
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Suggested labels
Suggested reviewers
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ 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}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
deepmd/pt/model/model/sezm_property_model.py (1)
121-153: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winKeep
core_compute()signature aligned withSeZMModel.Line 121 overrides the widened
SeZMModel.core_compute()API but drops the newconservativekeyword. That makesSeZMPropertyModelfail withTypeErrorfor any generic caller that uses the base-class interface.Proposed fix
def core_compute( self, coord: torch.Tensor, atype: torch.Tensor, edge_index: torch.Tensor, edge_vec: torch.Tensor, edge_scatter_index: torch.Tensor, edge_mask: torch.Tensor, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, charge_spin: torch.Tensor | None = None, comm_dict: dict[str, torch.Tensor] | None = None, extended_atype: torch.Tensor | None = None, extended_coord_corr: torch.Tensor | None = None, embedding_only: bool = False, + conservative: bool = False, ) -> dict[str, torch.Tensor]: """Compute property outputs through the SeZM forward-only graph.""" return super().core_compute( coord, atype, edge_index, edge_vec, edge_scatter_index, edge_mask, fparam=fparam, aparam=aparam, charge_spin=charge_spin, comm_dict=comm_dict, extended_atype=extended_atype, extended_coord_corr=extended_coord_corr, embedding_only=embedding_only, conservative=False, )🤖 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/pt/model/model/sezm_property_model.py` around lines 121 - 153, `SeZMPropertyModel.core_compute` is missing the new `conservative` keyword that exists on `SeZMModel.core_compute`, so the override no longer matches the base-class API and generic callers can hit a `TypeError`. Update `core_compute` in `SeZMPropertyModel` to accept the same signature as `SeZMModel.core_compute`, including `conservative`, and forward that argument through the `super().core_compute(...)` call unchanged.
🤖 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/utils/argcheck.py`:
- Around line 3263-3271: The DPA4/SeZM fitting schema is missing the supported
sezm_ener variant, causing validation to reject configs that the runtime still
accepts. Update the fitting type Variant in the schema builder around the
DPA4/SeZM arguments to include fitting_args_plugin.get_argument("sezm_ener")
alongside dpa4_ener and property, keeping the existing default_tag and doc
behavior intact so get_sezm_model() and get_sezm_spin_model() remain valid.
---
Nitpick comments:
In `@deepmd/pt/model/model/sezm_property_model.py`:
- Around line 121-153: `SeZMPropertyModel.core_compute` is missing the new
`conservative` keyword that exists on `SeZMModel.core_compute`, so the override
no longer matches the base-class API and generic callers can hit a `TypeError`.
Update `core_compute` in `SeZMPropertyModel` to accept the same signature as
`SeZMModel.core_compute`, including `conservative`, and forward that argument
through the `super().core_compute(...)` call unchanged.
🪄 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: 2e654586-7523-4e51-a585-094ed7f5851a
📒 Files selected for processing (8)
deepmd/pt/model/atomic_model/sezm_atomic_model.pydeepmd/pt/model/model/__init__.pydeepmd/pt/model/model/sezm_model.pydeepmd/pt/model/model/sezm_property_model.pydeepmd/utils/argcheck.pydoc/model/dpa4.mdexamples/water/dpa4/input_property.jsonsource/tests/pt/model/test_sezm_model.py
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## master #5587 +/- ##
==========================================
+ Coverage 82.27% 82.30% +0.02%
==========================================
Files 887 888 +1
Lines 100331 100533 +202
Branches 4060 4056 -4
==========================================
+ Hits 82550 82744 +194
- Misses 16320 16326 +6
- Partials 1461 1463 +2 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
|
Thanks for adding the SeZM/DPA4 invariant property path. The eager forward contract, public key translation ( One blocker before this can merge: the new compiled-property training test is currently failing in CI on both Python 3.10 and 3.13. The failing test is: The failure happens at Examples:
Since the PR exposes Reviewed by OpenClaw 2026.6.8 (844f405) (model: custom-chat-jinzhezeng-group/gpt-5.5) |
|
One more documentation/example issue I noticed: "property_name": "band_prop"with training systems Could we either add matching property-label arrays for the example data, or change the docs/example wording to make clear this is a template that requires user-supplied Reviewed by OpenClaw 2026.6.8 (844f405) (model: custom-chat-jinzhezeng-group/gpt-5.5) |
wanghan-iapcm
left a comment
wanghan-iapcm
left a comment
There was a problem hiding this comment.
please resolve njzjz's comment.
|
Thanks for the follow-up fixes. I re-reviewed the current head (
One remaining issue I would fix before merge: the DPA4/SeZM input schema still omits if fitting_net["type"] in ("dpa4_ener", "sezm_ener"):but the schema variant currently lists only fitting_args_plugin.get_argument("sezm_ener"),next to Reviewed by OpenClaw 2026.6.8 (844f405) (model: custom-chat-jinzhezeng-group/gpt-5.5) |
4 participants
Summary by CodeRabbit
New Features
Bug Fixes
Documentation