fix(dpmodel): preserve virtual spin atom types

[njzjz-bot]

Closes #5663.

Summary

• preserve negative placeholder types when creating dense and lower-interface spin partners
• use one array-API-safe lookup that maps virtual atom types to zero instead of applying Python negative indexing
• zero placeholder spin displacements, virial corrections, magnetic outputs, and magnetic masks
• add dense, lower, end-to-end invariance, and Array API strict regressions

Why existing tests missed this

Existing spin consistency tests use only nonnegative real atom types, including their lower-interface ghost atoms. Negative-type tests cover base atomic models and neighbor-list builders, where masking happens before type-dependent work, but not SpinModel preprocessing, which runs earlier. Cross-backend spin tests therefore never supplied a nonzero padded spin or asserted that both the real placeholder and its generated spin partner remained negative.

Validation

• pytest source/tests/common/dpmodel/test_spin_model_virtual_types.py source/tests/common/dpmodel/test_finetune_spin.py source/tests/common/dpmodel/test_spin_model_legacy_routing.py -q (12 passed)
• ruff format . (1664 files unchanged on final pass)
• ruff check . (passed)
• git diff --check (passed)

Coding agent: Codex
Codex version: codex-cli 0.144.4
Model: gpt-5.6-sol
Reasoning effort: xhigh

Summary by CodeRabbit

• Bug Fixes

  • Improved spin-model handling of padded and virtual atom placeholders.
  • Prevented placeholder atoms from receiving displacement, virial, scaling, or magnetic-output corrections.
  • Ensured consistent behavior across standard and lower-level spin processing paths.
  • Improved compatibility across supported array backends and processing devices.

• Tests

  • Added regression coverage for virtual placeholders, masking, output handling, and consistency between model backends.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 78.82%. Comparing base (cc908a8) to head (777927d).
⚠️ Report is 2 commits behind head on master.

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #5854      +/-   ##
==========================================
- Coverage   79.03%   78.82%   -0.22%     
==========================================
  Files        1055     1067      +12     
  Lines      122233   123813    +1580     
  Branches     4401     4515     +114     
==========================================
+ Hits        96607    97590     +983     
- Misses      24061    24613     +552     
- Partials     1565     1610      +45     

☔ View full report in Codecov by Harness.
📢 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.

[njzjz]

Possible reviewers based on changed lines, exact file history, and exact-file review history:

• @wanghan-iapcm — 7 commits on changed files; 42 reviews on exact changed files (deepmd/dpmodel/model/spin_model.py).
• @iProzd — 4 commits on changed files (deepmd/dpmodel/model/spin_model.py).

No review request was made automatically.

Coding agent: Codex
Codex version: codex-cli 0.144.4
Model: gpt-5.6-sol
Reasoning effort: xhigh

[wanghan-iapcm]

The dpmodel change is correct. I checked out the base (6c3b985cb) and ran only the new test file: all three fail, with atype_updated coming back [[0, -1, 2, 3, 2, 5]] against the expected [[0, -1, 2, 3, -1, 5]] - the placeholder's partner takes type -1 + ntypes_real = 2. All three pass at the head. Genuine regression tests, and putting the placeholder mid-array rather than trailing is the right choice. The all-real path is untouched: xp.where collapses to the original expression.

Requesting changes on one point: the fix is dpmodel-only and deepmd/pt has the identical defect. deepmd/pt_expt/model/spin_model.py declares class SpinModel(SpinModelDP) and overrides none of process_spin_input, process_spin_input_lower, process_spin_output* or call_common*, so pt_expt inherits the fix - pt does not. After this lands, dpmodel and pt_expt disagree with pt on identical weights and inputs for any frame containing a placeholder. Two inline comments cover the two halves of that. One further inline comment on the lower-interface test is non-blocking.

Other notes, none blocking:

1. source/tests/consistent/model/test_spin_ener.py hard-codes non-negative atype at :166 and :359. That harness already compares the dpmodel, pt and pt_expt classes, so adding a -1 entry is the natural gate for the divergence above - it would fail today on pt.

2. Separate pre-existing bug, worth its own issue rather than this PR: deepmd/pt/model/model/sezm_native_spin_model.py:120 and :427 call self.spin_mask.index_select(0, atype.reshape(-1)) with unclamped atype. torch.index_select rejects negative indices, so the DPA4 native spin model raises IndexError on any placeholder frame instead of corrupting silently.

3. process_spin_output and process_spin_output_lower still call _to_xp on the mask and then pass it into _lookup_type_values, which calls _to_xp again. The other three call sites pass the raw attribute. Harmless, but the local call can go.

[coderabbitai]

📝 Walkthrough

Walkthrough

SpinModel now preserves negative placeholder atom types during spin expansion and safely handles their scaling and masks in DP and PT implementations. New regression tests cover dense, lower-interface, strict array-namespace, and backend-equivalence behavior.

Changes

Spin placeholder handling

Layer / File(s)	Summary
DP padding-safe processing deepmd/dpmodel/model/spin_model.py	Adds namespace-aware type lookup, preserves negative types during dense and lower expansion, and applies padding-safe scaling and magnetic masks.
PT placeholder preservation deepmd/pt/model/model/spin_model.py	Zeros lookup results for negative types and keeps placeholder entries negative when generating virtual types.
DP regression coverage source/tests/common/dpmodel/test_spin_model_virtual_types.py	Tests dense/lower expansion, strict array namespaces, lower-call behavior, and masks when virtual scaling is disabled.
PT regression coverage source/tests/pt/model/test_spin_model_virtual_types.py	Tests PT placeholder handling and equivalence with DP spin input/output processing.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Possibly related PRs

• deepmodeling/deepmd-kit#5391: Adds underlying atomic-model access used by related PT diagnostics.

Suggested reviewers: wanghan-iapcm

🚥 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 accurately states the main fix: preserving virtual spin atom types.
Linked Issues check	✅ Passed	The changes preserve negative placeholder types, keep virtual partners negative, avoid negative mask indexing, and add the requested dense/lower tests.
Out of Scope Changes check	✅ Passed	The diff stays focused on SpinModel virtual-placeholder handling and its regression tests.
Docstring Coverage	✅ Passed	Docstring coverage is 84.21% which is sufficient. The required threshold is 80.00%.

✨ 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

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

🧹 Nitpick comments (1)

source/tests/pt/model/test_spin_model_virtual_types.py (1)

61-140: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick win

Add PT coverage for process_spin_output_lower and the virtual_scale=False branch.

This file exercises process_spin_input, process_spin_input_lower, and process_spin_output, but not process_spin_output_lower (spin_model.py Lines 273-320) or the virtual_scale=False (spin_mask) path — both of which use the same _lookup_type_values fix and are equally susceptible to masking bugs. The dpmodel test file already has test_spin_mask_branch_zeroes_virtual_placeholders and lower-output masking coverage that could be mirrored here for backend parity.

🤖 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/pt/model/test_spin_model_virtual_types.py` around lines 61 -
140, Extend TestPtSpinModelVirtualTypes with coverage for
process_spin_output_lower and the virtual_scale=False spin_mask path, mirroring
the corresponding dpmodel tests. Verify virtual placeholders are masked and
produce zero magnetic outputs or mask values, including lower-output handling,
while preserving parity with dpmodel where applicable.

🤖 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/pt/model/test_spin_model_virtual_types.py`:
- Around line 61-140: Extend TestPtSpinModelVirtualTypes with coverage for
process_spin_output_lower and the virtual_scale=False spin_mask path, mirroring
the corresponding dpmodel tests. Verify virtual placeholders are masked and
produce zero magnetic outputs or mask values, including lower-output handling,
while preserving parity with dpmodel where applicable.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Pro Plus

Run ID: 247001c0-f4b3-4f80-bdbf-79276d0a2f73

📥 Commits

Reviewing files that changed from the base of the PR and between cc908a8 and 777927d.

📒 Files selected for processing (4)

• deepmd/dpmodel/model/spin_model.py
• deepmd/pt/model/model/spin_model.py
• source/tests/common/dpmodel/test_spin_model_virtual_types.py
• source/tests/pt/model/test_spin_model_virtual_types.py

[wanghan-iapcm]

All three blocking points are resolved, verified against 777927d68 rather than the replies.

Both pt sites now build the virtual types with torch.where(atype >= 0, atype + self.ntypes_real, ...) (process_spin_input, process_spin_input_lower), so pt, dpmodel and pt_expt agree and the divergence does not land.

Masking rather than clamping in _lookup_type_values is the right resolution: keeping index_select preserves the export device-stability reason the clamp existed for, while torch.where(real_atom, gathered, 0) makes the clamped index purely an index whose value is discarded. Replacing the old docstring matters as much as the code change -- it asserted that padding atoms carry zero spin and are dropped downstream, and nothing enforced either half.

The added coverage is what I was hoping for. test_call_lower_ignores_virtual_placeholders pins the leak I measured through the real lower interface, and test_matches_dpmodel_when_the_first_real_type_is_magnetic turns the use_spin=[True, False, False] scenario into a direct cross-backend gate, including the assertion that the placeholder's mask_mag is False -- the entry ener_spin.py would otherwise have counted in the magnetic-force loss. That supersedes my suggestion of adding a -1 to test_spin_ener.py; comparing the backends directly is better than doing it through the consistency harness.

Still open, and still out of scope for this PR: sezm_native_spin_model.py L120 and L427 call self.spin_mask.index_select(0, atype.reshape(-1)) with unclamped atype, and index_select rejects negative indices. The contrast is sharper after this PR: pt's ordinary spin model and pt_expt now handle a placeholder frame, while the DPA4 native spin model would raise IndexError on the same input. I will confirm whether that path can actually receive a -1 and open an issue if so -- no action needed here.