feat(pt_expt): implement DeepSpin model in pt_expt backend #5293
Open
wanghan-iapcm wants to merge 9 commits intodeepmodeling:masterfrom
Open
feat(pt_expt): implement DeepSpin model in pt_expt backend #5293wanghan-iapcm wants to merge 9 commits intodeepmodeling:masterfrom
wanghan-iapcm wants to merge 9 commits intodeepmodeling:masterfrom
Conversation
added 2 commits
March 5, 2026 16:32
…and load-from-file support - Add coord_corr_for_virial support to dpmodel and pt_expt spin model, matching the pt backend's virial correction for virtual atom displacement. The correction flows through process_spin_input -> call_common -> call_common_lower -> forward_common_atomic -> fit_output_to_model_output. - Add torch.export test (TestSpinEnerModelExportable) verifying make_fx tracing and torch.export.export work for the spin energy model. - Fix get_spin_model in dpmodel to deepcopy input data dict, preventing in-place mutation of type_map/sel/exclude_types (pt backend already did this). - Un-skip test_load_stat_from_file for spin model; the load-from-file mechanism already worked correctly once the data mutation bug was fixed. - Add virial output comparison to spin consistency tests (TestSpinEner and TestSpinEnerLower) across pt and pt_expt backends.
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 0abbf38d24
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
Contributor
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughThreads coordinate-correction tensors (coord_corr_for_virial / extended_coord_corr) through DP, JAX, and PT_EXPT model call paths for virial computations; migrates DP spin processing to array‑API xp usage; adds PT_EXPT SpinEnergyModel with an exportable lower-forward; and expands cross‑backend tests and stat sampling support. Changes
Sequence Diagram(s)sequenceDiagram
participant Caller as Caller
participant CM as CM / public model
participant Backbone as BackboneModel
participant Atomic as AtomicModel
participant Transform as fit_output_to_model_output
Note over CM,Backbone: coord_corr_for_virial supplied at top-level
Caller->>CM: call_common(..., coord_corr_for_virial)
CM->>Backbone: model_call_from_call_lower(..., coord_corr_for_virial)
Backbone->>Atomic: forward_common_atomic(..., extended_coord_corr)
Atomic-->>Backbone: atomic outputs (+ extended_virial if requested)
Backbone->>Transform: fit_output_to_model_output(outputs, extended_coord_corr)
Transform-->>Caller: final outputs (energy/forces/virial)
Estimated code review effort🎯 5 (Critical) | ⏱️ ~120 minutes Possibly related PRs
Suggested reviewers
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 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: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
deepmd/dpmodel/model/make_model.py (1)
374-399:⚠️ Potential issue | 🔴 Critical
extended_coord_corris dropped before output transformation.
extended_coord_corrreachesforward_common_atomic(Line 383) but is never consumed, so virial coordinate correction is effectively not applied on this path.Proposed fix
def forward_common_atomic( self, extended_coord: Array, extended_atype: Array, nlist: Array, mapping: Array | None = None, fparam: Array | None = None, aparam: Array | None = None, do_atomic_virial: bool = False, extended_coord_corr: Array | None = None, ) -> dict[str, Array]: atomic_ret = self.atomic_model.forward_common_atomic( extended_coord, extended_atype, nlist, mapping=mapping, fparam=fparam, aparam=aparam, ) return fit_output_to_model_output( atomic_ret, self.atomic_output_def(), extended_coord, + extended_coord_corr=extended_coord_corr, do_atomic_virial=do_atomic_virial, mask=atomic_ret["mask"] if "mask" in atomic_ret else None, )🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@deepmd/dpmodel/model/make_model.py` around lines 374 - 399, The function forward_common_atomic drops the extended_coord_corr parameter instead of using it for virial corrections; forward extended_coord_corr into the downstream calls — pass extended_coord_corr to atomic_model.forward_common_atomic (so the atomic model can apply coordinate corrections) and also supply extended_coord_corr into fit_output_to_model_output (so virial/correction logic runs on the transformed output) while leaving all other args unchanged.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@deepmd/dpmodel/model/spin_model.py`:
- Line 248: The tuple unpacking "nframes, nloc, nnei = nlist.shape" introduces
an unused variable `nnei` causing a linter warning; change the unpacking in
spin_model.py (the line using `nlist.shape`) to discard the third value (e.g.,
"nframes, nloc, _ = nlist.shape" or remove the third target) so only used
variables remain, then run ruff check/format.
In `@source/tests/pt_expt/model/test_spin_ener_model.py`:
- Line 403: Remove the unused local variable nall assigned from
ext_coord.shape[1] in test_spin_ener_model.py (around the test that references
ext_coord) to satisfy lint (Ruff F841); simply delete the line "nall =
ext_coord.shape[1]" or replace it with a used expression if the count is needed
elsewhere (search for ext_coord and the surrounding test function to locate the
assignment).
---
Outside diff comments:
In `@deepmd/dpmodel/model/make_model.py`:
- Around line 374-399: The function forward_common_atomic drops the
extended_coord_corr parameter instead of using it for virial corrections;
forward extended_coord_corr into the downstream calls — pass extended_coord_corr
to atomic_model.forward_common_atomic (so the atomic model can apply coordinate
corrections) and also supply extended_coord_corr into fit_output_to_model_output
(so virial/correction logic runs on the transformed output) while leaving all
other args unchanged.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 302ec291-5d78-47e1-9bb8-0fb738497162
📒 Files selected for processing (10)
deepmd/dpmodel/model/make_model.pydeepmd/dpmodel/model/model.pydeepmd/dpmodel/model/spin_model.pydeepmd/pt_expt/model/__init__.pydeepmd/pt_expt/model/make_model.pydeepmd/pt_expt/model/spin_ener_model.pydeepmd/pt_expt/model/spin_model.pydeepmd/pt_expt/model/transform_output.pysource/tests/consistent/model/test_spin_ener.pysource/tests/pt_expt/model/test_spin_ener_model.py
Signed-off-by: Han Wang <92130845+wanghan-iapcm@users.noreply.github.com>
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
Comment on lines
+43
to
+52
| def forward( | ||
| self, | ||
| coord: torch.Tensor, | ||
| atype: torch.Tensor, | ||
| spin: torch.Tensor, | ||
| box: torch.Tensor | None = None, | ||
| fparam: torch.Tensor | None = None, | ||
| aparam: torch.Tensor | None = None, | ||
| do_atomic_virial: bool = False, | ||
| ) -> dict[str, torch.Tensor]: |
Check warning
Code scanning / CodeQL
Signature mismatch in overriding method Warning
Comment on lines
+75
to
+85
| def forward_lower( | ||
| self, | ||
| extended_coord: torch.Tensor, | ||
| extended_atype: torch.Tensor, | ||
| extended_spin: torch.Tensor, | ||
| nlist: torch.Tensor, | ||
| mapping: torch.Tensor | None = None, | ||
| fparam: torch.Tensor | None = None, | ||
| aparam: torch.Tensor | None = None, | ||
| do_atomic_virial: bool = False, | ||
| ) -> dict[str, torch.Tensor]: |
Check warning
Code scanning / CodeQL
Signature mismatch in overriding method Warning
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@deepmd/pt_expt/model/__init__.py`:
- Around line 26-28: SpinEnergyModel is not registered with the BaseModel
registry and conflicts with EnergyModel via model_type, so add the decorator
`@BaseModel.register`("spin_ener") above the SpinEnergyModel class definition in
spin_ener_model.py and update SpinEnergyModel.model_type = "spin_ener" (and any
related registry key uses) so get_model() can discover and instantiate it via
config without colliding with EnergyModel.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 8004a2eb-6b10-4f6b-8ae5-da2a9b68f557
📒 Files selected for processing (3)
deepmd/pt_expt/model/__init__.pydeepmd/pt_expt/model/make_model.pydeepmd/pt_expt/model/transform_output.py
🚧 Files skipped from review as they are similar to previous changes (1)
- deepmd/pt_expt/model/make_model.py
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (1)
deepmd/dpmodel/model/spin_model.py (1)
248-248:⚠️ Potential issue | 🟡 Minor
nneiis still unused inextend_nlist.Please rename it to
_nnei(or_) to clear the Ruff warning.Proposed fix
- nframes, nloc, nnei = nlist.shape + nframes, nloc, _nnei = nlist.shapeAs per coding guidelines
**/*.py: Always runruff check .andruff format .before committing changes or CI will fail.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@deepmd/dpmodel/model/spin_model.py` at line 248, The tuple unpacking in extend_nlist assigns nframes, nloc, nnei from nlist but nnei is unused and triggers Ruff; change the third variable name to _nnei (or _) in the unpacking within the extend_nlist function (where nframes, nloc, nnei = nlist.shape is declared) to silence the unused-variable warning and run ruff check/format before committing.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@deepmd/dpmodel/model/spin_model.py`:
- Around line 630-638: translated_output_def currently omits the virial keys
that call()/call_lower() now emit; update the translated_output_def
implementation in the SpinModel (or related class) to include "virial" and
"atom_virial" in the output schema whenever the model can compute gradients
(e.g., when backbone_model.do_grad_c(var_name) is true and do_atomic_virial is
possible). Specifically, add entries for "virial" and conditional "atom_virial"
matching the runtime names produced in model_predict (those set from
model_ret[f"{var_name}_derv_c_redu"] and model_ret[f"{var_name}_derv_c"]) so
schema and runtime keys align; apply the same change to the other
translated_output_def occurrence referenced in the file.
---
Duplicate comments:
In `@deepmd/dpmodel/model/spin_model.py`:
- Line 248: The tuple unpacking in extend_nlist assigns nframes, nloc, nnei from
nlist but nnei is unused and triggers Ruff; change the third variable name to
_nnei (or _) in the unpacking within the extend_nlist function (where nframes,
nloc, nnei = nlist.shape is declared) to silence the unused-variable warning and
run ruff check/format before committing.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 0ea24918-0bd0-4f9f-84b5-b3e833befeb9
📒 Files selected for processing (6)
deepmd/dpmodel/model/spin_model.pydeepmd/jax/model/base_model.pydeepmd/jax/model/dp_model.pydeepmd/jax/model/dp_zbl_model.pydeepmd/pt/model/model/spin_model.pydeepmd/pt_expt/model/spin_ener_model.py
🚧 Files skipped from review as they are similar to previous changes (1)
- deepmd/pt_expt/model/spin_ener_model.py
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (1)
deepmd/dpmodel/model/spin_model.py (1)
630-638:⚠️ Potential issue | 🟠 MajorAdd virial keys to
translated_output_def()to match emitted outputs.
call()/call_lower()now emit virial fields, buttranslated_output_def()still omits virial schema keys. This leaves a runtime/schema mismatch for output-definition consumers.Proposed fix
if self.backbone_model.do_grad_r(var_name): output_def["force"] = deepcopy(out_def_data[f"{var_name}_derv_r"]) output_def["force"].squeeze(-2) output_def["force_mag"] = deepcopy(out_def_data[f"{var_name}_derv_r_mag"]) output_def["force_mag"].squeeze(-2) + if self.backbone_model.do_grad_c(var_name): + output_def["virial"] = deepcopy(out_def_data[f"{var_name}_derv_c_redu"]) + output_def["virial"].squeeze(-2) + output_def["atom_virial"] = deepcopy(out_def_data[f"{var_name}_derv_c"]) + output_def["atom_virial"].squeeze(-2) return output_defAlso applies to: 813-821, 824-846
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@deepmd/dpmodel/model/spin_model.py` around lines 630 - 638, translated_output_def() is missing schema entries for the newly emitted virial fields causing a runtime/schema mismatch; update translated_output_def() to include keys "virial" and "atom_virial" (matching the outputs set in call() / call_lower()) with the appropriate shape/type descriptors, and ensure any conditional inclusion logic mirrors the conditions used when populating model_predict (e.g., respect do_grad_c and do_atomic_virial); also add the same keys in the other translated_output_def() occurrences mentioned so the schema aligns with the outputs at emission sites.
🧹 Nitpick comments (3)
source/tests/pt_expt/model/test_spin_ener_model.py (3)
170-177: Includevirialin serialize/deserialize parity checks.This test validates round-trip correctness but currently skips
virial, which is a core path touched in this PR.Proposed patch
- for key in ["energy", "atom_energy", "force", "force_mag", "mask_mag"]: + for key in ["energy", "atom_energy", "force", "force_mag", "mask_mag", "virial"]: np.testing.assert_allclose( ret1[key].detach().cpu().numpy(), ret2[key].detach().cpu().numpy(), rtol=1e-10, atol=1e-10, err_msg=f"Mismatch in {key} after round-trip", )🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@source/tests/pt_expt/model/test_spin_ener_model.py` around lines 170 - 177, The parity test loop in test_spin_ener_model.py currently checks keys ["energy","atom_energy","force","force_mag","mask_mag"] but omits "virial"; update the keys list used in the loop (the one iterating over keys and calling np.testing.assert_allclose on ret1[...] vs ret2[...]) to include "virial" so the serialize/deserialize round-trip asserts also verify virial equality for ret1 and ret2.
323-332: Avoid repeated CPU→device transfers in virial finite-difference evaluations.
np_inferbuilds CPU tensors each step andeval_modelmoves them toenv.DEVICE; constructing directly onenv.DEVICEreduces overhead in this hot loop.Proposed patch
def np_infer(new_cell): + stretched_coord = torch.tensor( + stretch_box(coord, cell, new_cell), dtype=dtype, device=env.DEVICE + ).unsqueeze(0) + new_cell_t = torch.tensor( + new_cell, dtype=dtype, device=env.DEVICE + ).unsqueeze(0) result = eval_model( self.model, - torch.tensor( - stretch_box(coord, cell, new_cell), device="cpu" - ).unsqueeze(0), - torch.tensor(new_cell, device="cpu").unsqueeze(0), + stretched_coord, + new_cell_t, atype, spin.unsqueeze(0), )🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@source/tests/pt_expt/model/test_spin_ener_model.py` around lines 323 - 332, np_infer is creating tensors on CPU and relying on eval_model to move them to env.DEVICE each finite-difference step; change the tensor constructions inside np_infer to use device=env.DEVICE (e.g., torch.tensor(stretch_box(...), device=env.DEVICE) and torch.tensor(new_cell, device=env.DEVICE)) and ensure spin.unsqueeze(0) is also on env.DEVICE so eval_model doesn't perform repeated CPU→device transfers. Keep the same shapes/unsqueeze calls and use the existing stretch_box/new_cell inputs and atype argument names so callers remain unchanged.
95-102: Consider extracting shared random-system setup into one helper.Cell/coord/atype/spin initialization is repeated in several tests; a common helper would reduce duplication and drift.
Also applies to: 121-130, 155-162, 271-280, 310-319, 372-381
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@source/tests/pt_expt/model/test_spin_ener_model.py` around lines 95 - 102, Extract the repeated random-system setup (torch.Generator(device="cpu").manual_seed(GLOBAL_SEED), creation of cell, coord, atype, and spin using the test-local dtype and generator) into a single helper function (e.g., make_random_system or build_test_system) and call that helper from each test; ensure the helper accepts dtype and optional seed or generator and returns the tuple (generator, cell, coord, atype, spin) so existing tests that reference generator, cell, coord, atype, and spin keep using the same symbols without duplicating initialization code.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@deepmd/dpmodel/model/spin_model.py`:
- Around line 347-364: The per-atom field "aparam" is copied unchanged when
doubling virtual spin atoms causing shape mismatches; in the block inside
compute_or_load_stat()/the spin-sampling loop where tmp_dict is built (same area
that expands "natoms"), detect if "aparam" is in sys and expand it the same way
as other per-atom arrays: get the array namespace via
array_api_compat.array_namespace(sys["aparam"]) or using the existing
expand_aparam helper, and concat the slices to duplicate the spin-local atoms
(e.g. xp.concat([2 * aparam[:, :2], aparam[:, 2:], aparam[:, 2:]], axis=-1) or
call expand_aparam(aparam, nloc * 2)); assign the result to tmp_dict["aparam"]
so per-atom shapes stay synchronized with coord/atype.
---
Duplicate comments:
In `@deepmd/dpmodel/model/spin_model.py`:
- Around line 630-638: translated_output_def() is missing schema entries for the
newly emitted virial fields causing a runtime/schema mismatch; update
translated_output_def() to include keys "virial" and "atom_virial" (matching the
outputs set in call() / call_lower()) with the appropriate shape/type
descriptors, and ensure any conditional inclusion logic mirrors the conditions
used when populating model_predict (e.g., respect do_grad_c and
do_atomic_virial); also add the same keys in the other translated_output_def()
occurrences mentioned so the schema aligns with the outputs at emission sites.
---
Nitpick comments:
In `@source/tests/pt_expt/model/test_spin_ener_model.py`:
- Around line 170-177: The parity test loop in test_spin_ener_model.py currently
checks keys ["energy","atom_energy","force","force_mag","mask_mag"] but omits
"virial"; update the keys list used in the loop (the one iterating over keys and
calling np.testing.assert_allclose on ret1[...] vs ret2[...]) to include
"virial" so the serialize/deserialize round-trip asserts also verify virial
equality for ret1 and ret2.
- Around line 323-332: np_infer is creating tensors on CPU and relying on
eval_model to move them to env.DEVICE each finite-difference step; change the
tensor constructions inside np_infer to use device=env.DEVICE (e.g.,
torch.tensor(stretch_box(...), device=env.DEVICE) and torch.tensor(new_cell,
device=env.DEVICE)) and ensure spin.unsqueeze(0) is also on env.DEVICE so
eval_model doesn't perform repeated CPU→device transfers. Keep the same
shapes/unsqueeze calls and use the existing stretch_box/new_cell inputs and
atype argument names so callers remain unchanged.
- Around line 95-102: Extract the repeated random-system setup
(torch.Generator(device="cpu").manual_seed(GLOBAL_SEED), creation of cell,
coord, atype, and spin using the test-local dtype and generator) into a single
helper function (e.g., make_random_system or build_test_system) and call that
helper from each test; ensure the helper accepts dtype and optional seed or
generator and returns the tuple (generator, cell, coord, atype, spin) so
existing tests that reference generator, cell, coord, atype, and spin keep using
the same symbols without duplicating initialization code.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 0c984edb-03ca-4bc1-b82e-7822457e9aca
📒 Files selected for processing (2)
deepmd/dpmodel/model/spin_model.pysource/tests/pt_expt/model/test_spin_ener_model.py
…ing all api. fix tranform_output_def issue
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
deepmd/pt_expt/model/spin_model.py (1)
126-139: Keepdeserialize()model-type driven to avoid future backbone lock-in.
SpinModel.deserialize()currently hardcodesDPEnergyAtomicModel; consider dispatching from serialized backbone metadata so the baseSpinModelstays extensible.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@deepmd/pt_expt/model/spin_model.py` around lines 126 - 139, SpinModel.deserialize currently hardcodes DPEnergyAtomicModel; change it to pick the backbone class from the serialized metadata instead. In SpinModel.deserialize use the model identifier stored in data["backbone_model"] (e.g. a "type" or "model_name" field) to resolve the correct class via make_model or your model registry, then call .deserialize on that instance (replace the direct DPEnergyAtomicModel reference), keeping Spin.deserialize(data["spin"]) and the cls(...) return path unchanged so the SpinModel remains extensible.deepmd/pt_expt/model/spin_ener_model.py (1)
96-111: Extract shared lower-output translation logic to reduce drift risk.
forward_lower()andforward_lower_exportable()duplicate nearly the same key translation/squeeze steps; a shared helper would make future output-key changes safer.Also applies to: 181-196
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@deepmd/pt_expt/model/spin_ener_model.py` around lines 96 - 111, Forward_lower() and forward_lower_exportable() duplicate the key mapping/squeeze logic that builds model_predict (keys like "atom_energy", "energy", "extended_mask_mag", "extended_force", "extended_force_mag", "virial", "extended_virial") and the conditional squeezes based on self.backbone_model.do_grad_r("energy") and do_grad_c("energy"); extract that block into a single helper method (e.g., _translate_model_output or build_model_predict) that takes model_ret and do_atomic_virial and returns the normalized model_predict dict, then call this helper from both forward_lower and forward_lower_exportable (and the other duplicated location around lines 181-196) so all key translations and squeeze operations are centralized.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@deepmd/pt_expt/model/spin_ener_model.py`:
- Around line 96-111: Forward_lower() and forward_lower_exportable() duplicate
the key mapping/squeeze logic that builds model_predict (keys like
"atom_energy", "energy", "extended_mask_mag", "extended_force",
"extended_force_mag", "virial", "extended_virial") and the conditional squeezes
based on self.backbone_model.do_grad_r("energy") and do_grad_c("energy");
extract that block into a single helper method (e.g., _translate_model_output or
build_model_predict) that takes model_ret and do_atomic_virial and returns the
normalized model_predict dict, then call this helper from both forward_lower and
forward_lower_exportable (and the other duplicated location around lines
181-196) so all key translations and squeeze operations are centralized.
In `@deepmd/pt_expt/model/spin_model.py`:
- Around line 126-139: SpinModel.deserialize currently hardcodes
DPEnergyAtomicModel; change it to pick the backbone class from the serialized
metadata instead. In SpinModel.deserialize use the model identifier stored in
data["backbone_model"] (e.g. a "type" or "model_name" field) to resolve the
correct class via make_model or your model registry, then call .deserialize on
that instance (replace the direct DPEnergyAtomicModel reference), keeping
Spin.deserialize(data["spin"]) and the cls(...) return path unchanged so the
SpinModel remains extensible.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: d186edf4-c54b-4127-a381-283b3502a87a
📒 Files selected for processing (4)
deepmd/dpmodel/model/spin_model.pydeepmd/pt_expt/model/spin_ener_model.pydeepmd/pt_expt/model/spin_model.pysource/tests/consistent/model/test_spin_ener.py
…load_stat When computing statistics for spin models with aparam, the aparam array must be expanded to include virtual atoms (zeros appended) before being passed to the backbone model, matching how aparam is expanded during forward inference. Without this, the fitting net computes incorrect aparam statistics (mean diluted by factor of 2). Also make dpmodel general_fitting.compute_input_stats array-API compatible (np.concatenate/np.sum → xp.concat/xp.sum) so it works when called with torch tensors from the pt_expt backend. Add fparam and aparam (numb_fparam=2, numb_aparam=3) to the spin model consistency tests to exercise these code paths.
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## master #5293 +/- ##
==========================================
+ Coverage 82.30% 82.36% +0.06%
==========================================
Files 767 769 +2
Lines 76984 77151 +167
Branches 3659 3660 +1
==========================================
+ Hits 63359 63547 +188
+ Misses 12454 12432 -22
- Partials 1171 1172 +1 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
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
1 participant
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.
Summary
SpinModelandSpinEnergyModelin the pt_expt backend, supporting spin degrees of freedom for magnetic systemsSpinModelarray-API compatible so the same code works across numpy/torch/jax backendscoord_corr_for_virial) to dpmodel and pt_expt, matching the pt backendget_spin_modelin dpmodel to not mutate the caller's input data dict (pt backend already useddeepcopy)Changes
dpmodel (
deepmd/dpmodel/model/)spin_model.py: Replace allnp.*operations witharray_api_compatequivalents (xp.concat,xp.where,xp.zeroswithdevice=, slicing instead ofxp.split). Addcompute_or_load_statand virial correction support viacoord_corr_for_virial/extended_coord_corr.make_model.py: Threadcoord_corr_for_virialthroughcall_common→model_call_from_call_lower(extends to ghost atoms via mapping) →call_common_lower→forward_common_atomic.model.py: Addcopy.deepcopy(data)inget_spin_modelto prevent in-place mutation of input dict.pt_expt (
deepmd/pt_expt/model/)spin_model.py(new):@torch_modulewrapper inheriting from dpmodelSpinModel.spin_ener_model.py(new):SpinEnergyModelwithforward()/forward_lower()/forward_lower_exportable()providing user-facing output translation.make_model.py,transform_output.py: Acceptextended_coord_corrfor virial correction.Tests
test_spin_ener_model.py(new): Unit tests for output keys/shapes, serialize/deserialize round-trip, dpmodel consistency, force finite-difference, virial finite-difference, andtorch.exportexportability.test_spin_ener.py: Cross-backend consistency tests forcall/call_lower,compute_or_load_stat, and load-from-file. Virial output now compared across pt and pt_expt.Test plan
python -m pytest source/tests/pt_expt/model/ -v— all 28 tests passpython -m pytest source/tests/consistent/model/test_spin_ener.py -v— all 12 tests pass (18 skipped for uninstalled backends)torch.export.exportverified onforward_lower_exportablecompute_or_load_statload-from-file verified across dp/pt/pt_exptSummary by CodeRabbit
New Features
Bug Fixes
Tests