fix(dpmodel): exclude virtual atoms from neighbor stats - #5834
Conversation
Mask virtual centers and neighbors before distance and count reductions. Add periodic, mixed-type, and strict Array API regression coverage. Coding-Agent: Codex Codex-Version: codex-cli 0.144.4 Model: gpt-5.6-sol Reasoning-Effort: xhigh
📝 WalkthroughWalkthrough
ChangesVirtual atom neighbor statistics
Estimated code review effort: 3 (Moderate) | ~25 minutes Possibly related PRs
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 |
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## master #5834 +/- ##
==========================================
+ Coverage 78.58% 79.38% +0.80%
==========================================
Files 1050 1085 +35
Lines 120637 126576 +5939
Branches 4356 4592 +236
==========================================
+ Hits 94801 100487 +5686
- Misses 24278 24434 +156
- Partials 1558 1655 +97 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
|
Possible reviewers based on changed lines, exact file history, and exact-file review history:
No review request was made automatically. Coding agent: Codex |
njzjz-bot
left a comment
njzjz-bot
left a comment
There was a problem hiding this comment.
Reviewed head 7ddcaae1f677abd65abf7a11b82c55d8f125d726 with three independent full subagent reviews. No actionable findings were identified in this dpmodel-scoped diff. The pair mask correctly excludes virtual centers and neighbors, preserves periodic self-image neighbors, and keeps strict Array API count reductions numeric. One reviewer noted that Paddle has an analogous independent pre-existing implementation; that is outside this PR's declared dpmodel scope and should be tracked separately rather than block this fix. All current checks pass. Because this PR was opened by the active njzjz-bot account, this is a comment-only review rather than a self-approval event.
The quota is about to reset, so I am concentrating the remaining token budget on these reviews.
Coding agent: Codex
Codex version: codex-cli 0.144.6
Model: gpt-5.6-sol
Reasoning effort: xhigh
wanghan-iapcm
left a comment
wanghan-iapcm
left a comment
There was a problem hiding this comment.
The fix itself is right, and the regression coverage is real: I ran source/tests/common/dpmodel/test_neighbor_stat.py against the pre-fix neighbor_stat.py and it fails, then against this branch and it passes, so the test genuinely exposes the bug rather than merely accompanying the fix. Requiring both ends of a pair to be real also converges dpmodel onto what the TF backend has done since #3275 (virtual_type_mask_i/virtual_type_mask_j), so this is an alignment rather than a new convention.
One gap, in the inline comment: the Paddle backend keeps its own copy of this operator and does not get the fix.
A second, related point that is why nothing in CI will notice. source/tests/consistent/test_neighbor_stat.py is the only test that compares tf/pt/dp/jax/pd against one another, and its gen_sys() builds np.repeat(0, 27), so no negative type ever reaches NeighborStatOP. Both tests this PR adds call the dpmodel operator directly. That is precisely how the original gap survived unnoticed, and it will keep the Paddle divergence invisible too. Adding the virtual-atom geometry to run_neighbor_stat there, where it runs across all five backends, would close both at once.
| real_center = atype >= 0 | ||
| real_neighbor = extend_atype >= 0 |
There was a problem hiding this comment.
real_center and real_neighbor are the fix, and both are correct. The same operator is duplicated in Paddle, which this PR does not touch: deepmd/pd/utils/neighbor_stat.py#L88-L112. It is a hand-written paddle.nn.Layer, not an import of this class, so it does not inherit the change. deepmd/pt, deepmd/jax and deepmd/pt_expt all import NeighborStatOP from here and do pick it up, which leaves Paddle as the only backend behind.
That copy is line-for-line the pre-fix logic: it masks only the eye(nloc, nall) self pair before min_rr2, and its non-mixed branch never excludes virtual centers at all. On this PR's own fixture (coord = [(0,0,0), (0,0,0), (1,0,0), (3,0,0)], atype = [0,-1,0,1], rcut = 1.1) Paddle still returns min_rr2 = [0, 0, 1, 4] and max_nnei = [[2]], against [1, inf, 1, 4] and [[1]] everywhere else. So dp --pd neighbor-stat on a mixed-type dataset whose padding slot lands on a real atom still aborts with RuntimeError: Some atoms are overlapping, and where it does not abort it picks a strictly larger auto-sel than the other backends for the same data. Paddle's own build_neighbor_list already relocates virtual atoms, so that extra sel is pure waste.
CLAUDE.md asks for this explicitly: "Before editing shared code (deepmd/dpmodel/, base classes), find every importer and run all affected backends' tests: backends routinely re-export or subclass generic classes, so one change ripples across tf/pt/pt_expt/dpmodel/jax/pd." The same point came up on #5854 and #5856 and was handled in a follow-up commit each time; the cleanest resolution here is the same.
There was a problem hiding this comment.
Fixed in 8b594d7. The Paddle implementation now excludes self pairs, virtual centers, and virtual neighbors before both the minimum-distance and neighbor-count reductions, matching the generic operator. I also added direct Paddle coverage for periodic/nonperiodic inputs and mixed/type-separated counts using this fixture.
Validation:
pytest source/tests/pd/test_neighbor_stat.py -q: 2 passed, 12 subtests passed- generic neighbor-stat regressions: 2 passed, 6 subtests passed
ruff format .andruff check .: passed
Coding agent: Codex
Codex version: codex-cli 0.144.6
Model: gpt-5.6-sol
Reasoning effort: xhigh
Apply the same real-center and real-neighbor masking as the generic neighbor-stat operator, and cover periodic, nonperiodic, mixed, and type-separated Paddle cases. Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh
![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
🤖 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 `@source/tests/pd/test_neighbor_stat.py`:
- Around line 99-113: Extend the NeighborStatOP parameterized cell cases in the
test around NeighborStatOP to include a periodic cell with at least one axis
shorter than rcut. Add expected assertions for that case verifying the original
self-pair remains excluded while an in-cutoff periodic self-image increases
min_rr2 and max_nnei, without changing the existing non-periodic and large-cell
expectations.
🪄 Autofix
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 Plus
Run ID: 807eb0a0-ff77-4621-a5e2-69e23fa0dcad
📒 Files selected for processing (2)
deepmd/pd/utils/neighbor_stat.pysource/tests/pd/test_neighbor_stat.py
| for cell in ( | ||
| None, | ||
| 10.0 * paddle.eye(3, dtype=paddle.float64).reshape([1, 9]).to(DEVICE), | ||
| ): | ||
| for mixed_types in (False, True): | ||
| with self.subTest(cell=cell is not None, mixed_types=mixed_types): | ||
| min_rr2, max_nnei = NeighborStatOP( | ||
| ntypes=2, | ||
| rcut=1.1, | ||
| mixed_types=mixed_types, | ||
| )(coord, atype, cell) | ||
|
|
||
| np.testing.assert_allclose(min_rr2.numpy(), expected_min_rr2) | ||
| expected_max_nnei = [[1]] if mixed_types else [[1, 0]] | ||
| np.testing.assert_array_equal(max_nnei.numpy(), expected_max_nnei) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Add a periodic self-image regression case.
Line 101 uses cell edges of 10.0 with rcut=1.1. Each periodic self-image is outside the cutoff. This test executes the periodic path, but it cannot verify that the implementation excludes only the original self-pair and retains periodic self-images.
Add a case with at least one cell axis shorter than rcut. Assert that the original self-pair is excluded and that an in-cutoff periodic self-image contributes to min_rr2 and max_nnei.
🤖 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/pd/test_neighbor_stat.py` around lines 99 - 113, Extend the
NeighborStatOP parameterized cell cases in the test around NeighborStatOP to
include a periodic cell with at least one axis shorter than rcut. Add expected
assertions for that case verifying the original self-pair remains excluded while
an in-cutoff periodic self-image increases min_rr2 and max_nnei, without
changing the existing non-periodic and large-cell expectations.
3 participants
Summary
Why existing tests missed this
The existing neighbor-stat tests build datasets containing only type-0 real atoms. They validate aggregate CLI statistics across cutoffs and backends, so neither a negative center type nor a negative neighbor type ever reaches NeighborStatOP. The new geometry contains an overlapping virtual neighbor that previously forced the minimum distance to zero and a virtual center that previously inflated the type-0 maximum neighbor count from one to two.
Validation
Closes #5630
Coding agent: Codex
Codex version: codex-cli 0.144.4
Model: gpt-5.6-sol
Reasoning effort: xhigh
Summary by CodeRabbit
Bug Fixes
Tests