fix(dpmodel): mask virtual EnvMat centers

[njzjz-bot]

Summary

• replace negative virtual-center types with a safe index before gathering normalization tables
• use neutral average and standard-deviation values for masked centers, then explicitly zero their descriptors
• add radial, angular, NumPy, and strict Array API regression coverage

Why existing tests missed this

The direct EnvMat unit test only used real center atoms. A separate model-level virtual-atom test did not expose the bug because the model masks virtual outputs after fitting and initializes descriptor averages to zero, so the invalid negative normalization lookup was not observable. The new tests use nonzero averages and a zero placeholder scale, which make both the negative-index behavior and masking order observable. The broader compiled DPA2 varying-natoms test had previously never been connected to virtual-center normalization; it is retained as a cross-cutting compile/autograd guard and was also reproduced after an isolated CI mismatch.

Validation

• ruff format .
• ruff check .
• focused dpmodel, Array API, PyTorch EnvMat, and virtual atomic-model tests: 7 passed with 2 subtests
• TensorFlow C++ core regression test: 1 passed
• TestCompiledVaryingNatoms::test_compiled_matches_uncompiled_varying_natoms_dpa2: passed locally in 205.82 s, covering four compiled/eager training steps, changing frame/atom counts, force and virial outputs, and second-order force-loss gradients
• manual NumPy, Array API Strict, PyTorch, and JAX virtual-center scenario

Closes #5628

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

Summary by CodeRabbit

• Bug Fixes

  • Improved environment-matrix normalization for virtual center atoms.
  • Virtual centers now consistently produce zero environment, derivative, and switching outputs across supported computation modes.

• Tests

  • Added coverage for all-virtual and mixed real/virtual center scenarios, including normalization with provided statistics.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Pro Plus

Run ID: abcd9e4d-1747-430e-906d-f442181424c7

📥 Commits

Reviewing files that changed from the base of the PR and between 27f2a5a and 8fe00ba.

📒 Files selected for processing (1)

• deepmd/dpmodel/utils/env_mat.py

🚧 Files skipped from review as they are similar to previous changes (1)

• deepmd/dpmodel/utils/env_mat.py

---

📝 Walkthrough

Walkthrough

EnvMat.call() safely handles negative virtual-center type sentinels during normalization and preserves zero outputs. Tests cover array APIs, mixed centers, and both radial modes.

Changes

EnvMat virtual center handling

Layer / File(s)	Summary
Safe normalization and virtual center masking deepmd/dpmodel/utils/env_mat.py	EnvMat.call() documents virtual-center neighbor rows, uses safe indices for normalization lookups, and neutralizes normalization terms for virtual centers.
Virtual center regression coverage source/tests/common/dpmodel/array_api/test_env_mat.py, source/tests/common/dpmodel/test_env_mat.py	Tests verify zero virtual-center outputs, nonzero real-center outputs, preserved array metadata, and behavior across both radial modes.

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

Possibly related PRs

• deepmodeling/deepmd-kit#5834: Both changes handle negative virtual atom types in dpmodel.
• deepmodeling/deepmd-kit#5854: Both changes avoid invalid type-index lookups for virtual atoms.
• deepmodeling/deepmd-kit#5856: Both changes handle negative virtual atom types with type remapping or masking.

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 clearly and concisely describes the primary change: masking virtual centers in dpmodel EnvMat.
Linked Issues check	✅ Passed	The changes address issue #5628 by using safe normalization indices, neutral virtual-center values, zero outputs, and regression tests.
Out of Scope Changes check	✅ Passed	The implementation, documentation, and regression tests directly support the linked issue and stated pull request objectives.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.

✨ 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.}

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 79.24%. Comparing base (cc908a8) to head (8fe00ba).
⚠️ Report is 34 commits behind head on master.

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #5833      +/-   ##
==========================================
+ Coverage   79.03%   79.24%   +0.20%     
==========================================
  Files        1055     1073      +18     
  Lines      122233   125306    +3073     
  Branches     4401     4569     +168     
==========================================
+ Hits        96607    99293    +2686     
- Misses      24061    24376     +315     
- Partials     1565     1637      +72     

☔ 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 — 6 commits on changed files; 2 reviews on exact changed files (deepmd/dpmodel/utils/env_mat.py, source/tests/common/dpmodel/test_env_mat.py).
• @OutisLi — 14 reviews on exact changed files (deepmd/dpmodel/utils/env_mat.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 fix is correct and array-API-clean: std -> 1 is applied before the divide (NaN prevented, not just masked), the reshape broadcasts correctly for both radial and full modes, and it's a provable no-op for all-real centers. jax and pt_expt inherit it. Two inline notes; one cross-backend parity note here:

pt and pd carry the same unmasked-center pattern (parity / defense-in-depth, likely a follow-up). deepmd/pt/model/descriptor/env_mat.py (t_avg = mean[atype]; t_std = stddev[atype]) and deepmd/pd/model/descriptor/env_mat.py do the identical unguarded mean[atype] this PR fixes for dpmodel, and aren't touched here. In practice it's not an end-to-end bug — base_atomic_model.forward_common_atomic pre-clamps atype<0 to 0 before the descriptor runs (and zeroes outputs after), so the standard model path never feeds a negative center to mean[-1]; the dpmodel fix (and these tests) are really hardening the direct-call path. So this is a low-priority backend-parity cleanup rather than a correctness fix, but worth tracking so pt/pd don't diverge from dpmodel/the C++ CPU reference.

[njzjz]

Rebased onto master and narrowed the fix in 27f2a5a0f after chasing the CI failure.

What was failing. source/tests/pt_expt/test_training.py::TestCompiledVaryingNatoms::test_compiled_matches_uncompiled_varying_natoms_dpa2 — 100% of force elements differed between the compiled and uncompiled DPA2 models at step 0 (greatest abs diff 0.66, tolerance 1e-10). Reproduced locally and bisected:

variant	result
safe_atype + davg/dstd neutralization only	pass
+ em = xp.where(center_mask, em, 0)	fail
same but with a mask that is always True	fail
em * xp.astype(center_mask, dtype) instead of where	fail
em = em * 1.0 (scalar, no mask tensor)	pass

So it is not the semantics and not the where op — it is that masking makes em/diff/sw depend on atype_ext, and the compiled forward_lower for DPA2 miscompiles the force graph once that dependency exists. That looks like a genuine pt_expt compile bug worth its own issue; it is not something this PR should carry.

What the fix is now. _make_env_mat already zeroes em, diff and sw wherever nlist < 0, so a virtual center — whose neighbor row is empty by the neighbor-list contract — is already zero on the way in. The only thing that pushed it off zero was normalization: xp.take(davg, atype) with atype = -1 silently selects the last real type's row. Gathering with a safe index and neutralizing the offset (0) and scale (1) is therefore the complete fix, and no extra output masking is needed.

Tests. test_normalization_keeps_virtual_centers_zero and test_virtual_center_uses_safe_normalization_indices are unchanged. test_mixed_centers_mask_all_virtual_outputs gave the virtual center real neighbors, deliberately violating the contract, which is exactly the case that no longer holds; it is replaced by test_mixed_centers_keep_virtual_rows_at_zero, which keeps the contract and puts the nonzero davg/dstd in the last row — the one an unguarded take selects for -1 — so borrowing it is still caught. All three new tests fail on master and pass here, and TestCompiledVaryingNatoms is green again.

[wanghan-iapcm]

Both points are resolved, and I want to note that the second one landed somewhere better than what I proposed.

I had suggested either extending the center mask to diff/sw or softening the comment. 6d67a1257 did the former, and 27f2a5a0f then reverted it in favour of neutralizing the normalization instead -- zero offset and unit scale for virtual centers. That is the right call given what the new comment discloses, and the comment now grounds the guarantee in the neighbor-list contract rather than claiming protection against arbitrary malformed lists, which is what I was actually objecting to. The tests were updated to match rather than left behind: test_mixed_centers_keep_virtual_rows_at_zero uses a contract-honoring empty neighbor row and says so, instead of the malformed-list scenario an earlier round advertised. Code and tests agreeing on the narrower claim is worth more than the broader claim was.

The strict Array API gap is closed too -- the mixed case proves the real row stays nonzero while the virtual row stays at zero, so it now discriminates in both directions rather than only checking that nothing crashes. Using dstd = zeros in the all-virtual case is a nice touch: it fails loudly if masking ever moves after the division.

I ran these against a master baseline instead of reading them. Both new cases fail unpatched -- the all-virtual one with RuntimeWarning: divide by zero encountered in divide, which is exactly the hidden NaN the fix is preventing -- and all five pass after. On why nothing caught this before: xp.take(davg, -1) on NumPy silently selects the last real type's row, so a virtual center was getting a plausible-looking normalization rather than an error, which is invisible without an explicit assertion that virtual rows stay at zero.

One request, and it is about something the comment reveals rather than about this diff.

Approving.