feat(dpmodel): NeighborGraph 3-body angle machinery (PR-E)#5717
Merged
OutisLi merged 14 commits intoJul 9, 2026
Conversation
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:
📝 WalkthroughWalkthroughThis PR adds angle-graph support to ChangesAngle graph construction and aggregation
Estimated code review effort: 3 (Moderate) | ~30 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 |
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
![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: 1
🧹 Nitpick comments (1)
source/tests/common/dpmodel/test_angle_builder.py (1)
326-364: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd a padded-
angle_indexcase to the aggregation tests.Both aggregation tests here only use angle data where every slot is real (no padding). Given the padding-mask precondition gap noted in
angles.py(angle_to_edge_sum/angle_to_node_sum), a test with a paddedangle_indexand non-zero paddingdatavalues would catch regressions if that precondition is ever violated.🤖 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/common/dpmodel/test_angle_builder.py` around lines 326 - 364, The aggregation tests in test_angle_builder only cover fully real angle slots, so add a new padded-angle case for angle_to_edge_sum and angle_to_node_sum using a padded angle_index with non-zero data in the padded positions. Reuse the existing test_angle_aggregation and test_angle_aggregation_torch_namespace patterns, but assert that padding is ignored in both the numpy and torch-namespace paths so regressions around the padding-mask precondition are caught.
🤖 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/dpmodel/utils/neighbor_graph/segment.py`:
- Around line 59-89: The masked softmax in segment_softmax can overflow because
shifted is computed from raw data instead of data_for_max, allowing masked
entries to produce inf and then nan when multiplied by the zero mask. Update
segment_softmax to base the subtraction on data_for_max so masked elements
remain -inf before exp, keeping segment_sum and the denom_e guard from being
poisoned; use the existing segment_max, segment_sum, and mask handling paths to
locate the fix.
---
Nitpick comments:
In `@source/tests/common/dpmodel/test_angle_builder.py`:
- Around line 326-364: The aggregation tests in test_angle_builder only cover
fully real angle slots, so add a new padded-angle case for angle_to_edge_sum and
angle_to_node_sum using a padded angle_index with non-zero data in the padded
positions. Reuse the existing test_angle_aggregation and
test_angle_aggregation_torch_namespace patterns, but assert that padding is
ignored in both the numpy and torch-namespace paths so regressions around the
padding-mask precondition are caught.
🪄 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: b038075f-6809-4579-8ce1-5cf5d1d41c90
📒 Files selected for processing (17)
deepmd/dpmodel/descriptor/dpa1.pydeepmd/dpmodel/utils/neighbor_graph/__init__.pydeepmd/dpmodel/utils/neighbor_graph/angles.pydeepmd/dpmodel/utils/neighbor_graph/env.pydeepmd/dpmodel/utils/neighbor_graph/graph.pydeepmd/dpmodel/utils/neighbor_graph/pairs.pydeepmd/dpmodel/utils/neighbor_graph/segment.pysource/tests/common/dpmodel/test_angle_builder.pysource/tests/common/dpmodel/test_center_edge_pairs.pysource/tests/common/dpmodel/test_dpa1_call_graph_block.pysource/tests/common/dpmodel/test_dpa1_graph_attention_parity.pysource/tests/common/dpmodel/test_graph_angle_cos_parity.pysource/tests/common/dpmodel/test_segment_softmax.pysource/tests/pt_expt/descriptor/test_dpa1.pysource/tests/pt_expt/model/test_dpa1_graph_lower.pysource/tests/pt_expt/model/test_linear_model.pysource/tests/pt_expt/utils/test_neighbor_list.py
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## master #5717 +/- ##
==========================================
- Coverage 79.61% 79.50% -0.11%
==========================================
Files 1014 1015 +1
Lines 115359 115483 +124
Branches 4272 4274 +2
==========================================
- Hits 91839 91813 -26
- Misses 21978 22127 +149
- Partials 1542 1543 +1 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
wanghan-iapcm
force-pushed
the
feat-graph-angles-prE
branch
from
ad73de0 to
dd58c02
Compare
- Remove unused locals flagged by CodeQL (ai_def/ai_full in the ordered-superset test, ev in the angle-cos parity test); strengthen the ordered+include_self test to assert the default pairs are an actual subset of the full pair set (using the previously-dead ai_def/ai_full). - Add masked-padding aggregation tests (CodeRabbit nitpick): pin the precondition that angle_to_edge_sum/angle_to_node_sum ignore guard angles only once data is zeroed via angle_mask, in numpy and torch.
iProzd
reviewed
Jul 7, 2026
iProzd
reviewed
Jul 7, 2026
…nd mask precondition - Document that build_angle_index's strict '<' a_rcut gate intentionally mirrors dpa3's dense angle gate (repflows.py:598 a_dist_mask), unlike the edge channel's '<= rcut' (builder.py:284). Answers iProzd's review question: yes, boundary divergence is intentional dpa3 parity. - Document the angle_to_edge_sum/angle_to_node_sum 'caller must zero data at padded slots first' contract directly in their docstrings (previously only pinned by a test, per iProzd's second review comment).
OutisLi
requested changes
Jul 8, 2026
OutisLi
left a comment
OutisLi
left a comment
Collaborator
There was a problem hiding this comment.
Thanks for the careful machinery and the thorough oracle/parity tests. The forward parity is solid; my one change request is about autograd safety of the norm used in graph_angle_cos (details inline). A couple of minor, non-blocking notes will follow separately.
OutisLi
reviewed
Jul 8, 2026
OutisLi
left a comment
OutisLi
left a comment
Collaborator
There was a problem hiding this comment.
One non-blocking performance follow-up on the unordered angle enumeration (details inline). The change itself lands in the shared center_edge_pairs primitive.
OutisLi
reviewed
Jul 8, 2026
OutisLi
left a comment
OutisLi
left a comment
Collaborator
There was a problem hiding this comment.
One more non-blocking readability note on the pair_mask fold (inline).
Address OutisLi's blocking review on PR deepmodeling#5717. graph_angle_cos normalized with xp.linalg.vector_norm, but the dpa3 dense channel it mirrors uses safe_for_vector_norm for the same step (repflows.py:642-643), and the a_rcut gate mirrors repflows.py:598 which also uses safe_for_vector_norm. The two are value-identical for non-zero vectors (fp64 dense/se_t parity still passes at 1e-12), but differ in gradient at ||v||==0: plain vector_norm back-props NaN under jax, whereas safe_for_vector_norm yields a 0 gradient. edge_vec is the sole autograd leaf, so this was a latent NaN-gradient landmine on the geometry path (not yet triggered: padding angles point at a real, non-zero edge, and no gradients flow until the dpa3 graph descriptor lands in PR-G). Switching both call sites removes the landmine and makes the 'mirrors dpa3 exactly' docstring claims literally true. Add test_graph_angle_cos_zero_edge_grad_is_finite: a jax-autograd test with a zero-length edge_vec asserting the geometry-path gradient stays finite, plus a discrimination check that plain vector_norm DOES produce a NaN gradient under jax on the same construction (torch's vector_norm is already safe at zero, so jax is the discriminating backend). Guarded by pytest.importorskip('jax').
iProzd
approved these changes
Jul 8, 2026
OutisLi
approved these changes
Jul 9, 2026
HydrogenSulfate
pushed a commit
to HydrogenSulfate/deepmd-kit
that referenced
this pull request
Jul 10, 2026
deepmodeling#5728) ## Summary `source/tests/pt_expt/test_plugin.py` (from deepmodeling#5559) pops `deepmd.pt_expt` from `sys.modules` without restoring it, leaving the package's cached submodules bound to a dead parent. Any later import of a cached submodule (e.g. `deepmd.pt_expt.infer.deep_eval`) re-creates a BARE parent package whose `utils`/`infer` attributes are never rebound, and `mock.patch("deepmd.pt_expt.utils...")` in `test_deep_eval_serialize_api.py` then fails with `AttributeError: module 'deepmd.pt_expt' has no attribute 'utils'` under py3.10's mock target resolution (py3.13 tolerates it). The failure is shard-order dependent: it needs (1) something to import `deepmd.pt_expt.infer.deep_eval` before `test_plugin` runs, and (2) the serialize-API test to run after — so it appears/disappears as PRs add test files and reshuffle the duration-based shards. It currently fails `Test Python (4, 3.10)` on deepmodeling#5714 and `Test Python (8, 3.10)` on deepmodeling#5717 while master stays green by ordering luck. ## Fix Snapshot the whole `deepmd.pt_expt.*` module tree before the re-import and restore it (including the `deepmd` parent-package attribute binding) in the `finally`. Verified: fixed test passes together with the serialize-API tests; an inline emulation of the worst-case ordering (child cached → pop/reimport/pop → `mock.patch` target resolution) resolves cleanly. Same one-file fix is cherry-picked on deepmodeling#5714 (as 6422007) and deepmodeling#5717; whichever lands first, the others resolve trivially. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Tests** * Improved import handling in the plugin loading test to better isolate module state between test runs. * Reduced the chance of leftover cached imports affecting later tests in the same session. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: Han Wang <wang_han@iapcm.ac.cn>
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
4 participants
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
NeighborGraph PR-E: the optional 3-body angle extension of the edge-graph neighbor-list contract (design discussion wanghan-iapcm#4). An angle is a pair of edges sharing a center (
dst(edge_a) == dst(edge_b)), stored asangle_index (2, A)into[0, E)+angle_mask (A,)—edge_vecstays the ONLY geometry leaf, so force/virial assembly is untouched (proven by an invariance test).What's added (all in
deepmd/dpmodel/utils/neighbor_graph/)pad_and_guard_angles(graph.py) — angle-axis padder mirroringpad_and_guard_edges(dynamic guard append / staticangle_capacitywith overflow ValueError).angles.py(new):build_angle_index(edge_index, edge_vec, edge_mask, n_total, a_rcut, *, ordered=False, include_self=False, layout=None)— unordered, no-self pairs of edges sharing a center where BOTH edges are withina_rcut; built on PR-D'scenter_edge_pairs;pair_maskfolded intoangle_mask(never discarded).attach_angles(graph, a_rcut, ...)— post-hoc: edge graph in, graph with angle fields out (dataclasses.replace); default builders keep anglesNone.angle_to_edge_sum/angle_to_node_sum— segment-sum aggregation to the query edge / shared center.graph_angle_cos(angle_index, edge_vec, eps=1e-6)— per-angle cos θ mirroring dpa3 repflowscosine_ijeps placement exactly (+epsin norm denominators,*(1-eps)on the product).angle_padding_fraction(graph)— mask-derived padding-waste report for static capacities.Semantics / decisions
a_sel x a_selsquare including thej==kdiagonal; the graph set keeps one entry per unordered{j,k}pair and moves the degenerate diagonal to the (a_rcut-filtered) edge channel. Dense parity is therefore asserted against the OFF-DIAGONALcosine_ij[j,k](j != k) at rtol/atol 1e-12 (same-math fp64), at non-bindinga_sel. The ordered+self full square stays available via flags.a_sel= normalization-only (carry-all withina_rcut), consistent with the edge-seldecision.sw == 1regime (rtol 1e-12).Tests
source/tests/common/dpmodel/test_angle_builder.py(21) +test_graph_angle_cos_parity.py(6): brute-force triplet oracle (all flag combinations, multi-center, static layout,node_capacitybranch), dpa3 dense-parity + no-self-angle assertion, se_t coordinate oracle, force/virial bit-exact invariance with/without angles, padding-fraction (incl.total==0), torch-namespace smoke tests for every new function. Full neighbor-graph suite: 54 passed.Known limitations
mixed_types=False), used as oracle only.nonzeroincenter_edge_pairs) even when a staticlayoutis passed —angle_capacityfixes the output shape only; shape-static enumeration for export is deferred to PR-G.angle_maskbefore summing (padding angles point at edge 0).A ~ sum(deg^2)capacity overhead mitigated bya_rcut < rcutand reported byangle_padding_fraction.Summary by CodeRabbit
New Features
Tests