feat(dpmodel): NeighborGraph foundation — contract, segment toolkit, numpy builder, edge force/virial#5581
Merged
wanghan-iapcm merged 24 commits intoJun 25, 2026
Conversation
added 8 commits
June 24, 2026 16:00
Addresses two-stage review: RUF022 isort ordering of __all__ (CI blocker) and array-API-pure mask multiply (edge_force_virial) per CLAUDE.md guideline.
for more information, see https://pre-commit.ci
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:
📝 WalkthroughWalkthroughAdds a new ChangesBackend-agnostic neighbor graph utilities
Sequence Diagram(s)sequenceDiagram
participant Caller
participant build_neighbor_graph
participant extend_input_and_build_neighbor_list
participant neighbor_graph_from_extended
participant pad_and_guard_edges
participant edge_force_virial
participant segment_sum
Caller->>build_neighbor_graph: coord, atype, box, rcut, sel, layout
build_neighbor_graph->>extend_input_and_build_neighbor_list: coord, atype, rcut, sel, box
extend_input_and_build_neighbor_list-->>build_neighbor_graph: extended_coord, nlist, mapping
build_neighbor_graph->>neighbor_graph_from_extended: extended_coord, nlist, mapping, layout
neighbor_graph_from_extended->>pad_and_guard_edges: edge_index, edge_vec, capacity
pad_and_guard_edges-->>neighbor_graph_from_extended: padded edge_index, edge_vec, edge_mask
neighbor_graph_from_extended-->>build_neighbor_graph: NeighborGraph
Caller->>edge_force_virial: g_e, edge_vec, edge_index, edge_mask, n_node
edge_force_virial->>segment_sum: masked aggregation
segment_sum-->>edge_force_virial: per-node and per-frame reductions
edge_force_virial-->>Caller: force, atom virial, frame virial
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Suggested labels
Suggested reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 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 |
![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: 341a4081e4
ℹ️ 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".
![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)
deepmd/dpmodel/utils/numpy_neighbor_list.py (1)
73-79: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winNormalize or validate
boxbefore frame indexing.Line 78 slices
boxwith[f]unconditionally. If a caller passes a single-frame(3, 3)box,bxbecomes a(3,)row and_frame_edgesreceives the wrong shape. Accept(3, 3)fornf == 1or raise a clear shape error before the loop.Proposed fix
if layout is None: layout = GraphLayout() coord = np.asarray(coord, dtype=np.float64) + if coord.ndim != 3 or coord.shape[-1] != 3: + raise ValueError("coord must have shape (nf, nloc, 3)") nf, nloc = coord.shape[0], coord.shape[1] + box_arr = None + if box is not None: + box_arr = np.asarray(box, dtype=np.float64) + if box_arr.shape == (3, 3): + if nf != 1: + raise ValueError("a single (3, 3) box is only valid for nf == 1") + box_arr = box_arr[None, ...] + elif box_arr.shape != (nf, 3, 3): + raise ValueError("box must have shape (nf, 3, 3)") n_node = np.full((nf,), nloc, dtype=np.int64) src_all, dst_all, vec_all = [], [], [] for f in range(nf): - bx = None if box is None else np.asarray(box, dtype=np.float64)[f] + bx = None if box_arr is None else box_arr[f]🤖 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 `@deepmd/dpmodel/utils/numpy_neighbor_list.py` around lines 73 - 79, The `box` handling in the numpy neighbor list builder is too strict for single-frame inputs because `box[f]` is applied unconditionally before calling `_frame_edges`. Normalize `box` up front so a `(3, 3)` box is accepted when `nf == 1`, or validate and raise a clear shape error before the frame loop if the shape does not match the number of frames. Keep the existing `_frame_edges` call path unchanged once `bx` has the correct per-frame shape.
🤖 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/common/dpmodel/test_numpy_neighbor_list.py`:
- Around line 86-94: The test in test_static_capacity_padding currently uses a
tautological check on ng.edge_mask.sum() versus the same sliced sum, so it does
not verify padding behavior. Update this assertion to compare the number of true
entries against the expected real-edge count for the NumpyNeighborList build
with GraphLayout(edge_capacity=64), and also verify the masked tail is compactly
padded after the last real edge. Keep the existing NumpyNeighborList, build, and
edge_mask/edge_vec checks, but make the mask assertion explicitly validate that
only the real edges are marked and the remaining suffix is masked out.
---
Nitpick comments:
In `@deepmd/dpmodel/utils/numpy_neighbor_list.py`:
- Around line 73-79: The `box` handling in the numpy neighbor list builder is
too strict for single-frame inputs because `box[f]` is applied unconditionally
before calling `_frame_edges`. Normalize `box` up front so a `(3, 3)` box is
accepted when `nf == 1`, or validate and raise a clear shape error before the
frame loop if the shape does not match the number of frames. Keep the existing
`_frame_edges` call path unchanged once `bx` has the correct per-frame shape.
🪄 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: 8cbbb3f7-390c-412c-932a-8d5fbeddb3cd
📒 Files selected for processing (9)
deepmd/dpmodel/utils/__init__.pydeepmd/dpmodel/utils/edge_derivatives.pydeepmd/dpmodel/utils/neighbor_graph.pydeepmd/dpmodel/utils/numpy_neighbor_list.pydeepmd/dpmodel/utils/segment.pysource/tests/common/dpmodel/test_edge_force_virial.pysource/tests/common/dpmodel/test_neighbor_graph.pysource/tests/common/dpmodel/test_numpy_neighbor_list.pysource/tests/common/dpmodel/test_segment.py
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## master #5581 +/- ##
==========================================
+ Coverage 82.27% 82.30% +0.03%
==========================================
Files 882 887 +5
Lines 100009 100161 +152
Branches 4045 4043 -2
==========================================
+ Hits 82283 82440 +157
+ Misses 16267 16263 -4
+ Partials 1459 1458 -1 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
added 2 commits
June 24, 2026 19:18
…yNeighborList Add neighbor_graph_from_extended (quartet -> ghost-free NeighborGraph, src = mapping[neighbor] local owner) and build_neighbor_graph (reuses the tested extend_input_and_build_neighbor_list). Drop the all-pairs NumpyNeighborList from the library; the independent brute-force oracle now lives in the test file (test_neighbor_graph_builder.py) and cross-validates the adapter.
Contributor
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 `@deepmd/dpmodel/utils/neighbor_graph.py`:
- Around line 197-215: The public signature of the neighbor-graph builder is too
broad for the upstream helper it delegates to: the function currently advertises
`sel` as `int | list[int]` but passes it unchanged into
`extend_input_and_build_neighbor_list`, which expects `list[int]`. Update the
`neighbor_graph` entrypoint to either normalize an integer `sel` into the list
form before calling `extend_input_and_build_neighbor_list`, or narrow the `sel`
annotation and doc contract to match the helper’s `list[int]` expectation so the
API and implementation stay aligned.
🪄 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: 26df6060-a292-4698-b218-9fa775e3cf1c
📒 Files selected for processing (4)
deepmd/dpmodel/utils/__init__.pydeepmd/dpmodel/utils/neighbor_graph.pysource/tests/common/dpmodel/test_neighbor_graph.pysource/tests/common/dpmodel/test_neighbor_graph_builder.py
✅ Files skipped from review due to trivial changes (1)
- source/tests/common/dpmodel/test_neighbor_graph.py
added 2 commits
June 24, 2026 19:26
Move the flat utils modules (neighbor_graph.py, segment.py, edge_derivatives.py) into deepmd/dpmodel/utils/neighbor_graph/ as a cohesive subpackage: graph.py (contract + node-validity + edge padding), builder.py (neighbor_graph_from_extended, build_neighbor_graph), segment.py, derivatives.py. Public API re-exported from the package __init__ and from utils/__init__, so import paths (deepmd.dpmodel.utils.*) are unchanged.
…ng test Address CodeRabbit review on deepmodeling#5581: - build_neighbor_graph normalizes an int sel to list form before calling extend_input_and_build_neighbor_list (annotated list[int]), so the public int | list[int] contract is honored. No behavior change -- the underlying build_neighbor_list already accepts int -- but the call is now type-consistent. Add a test that int sel reproduces the list-sel real-edge environment. - test_static_capacity_padding: replace the (already-removed) tautological mask check's intent with explicit assertions on the exact real-edge count and the compact-prefix padding (real edges first, masked tail after).
iProzd
requested changes
Jun 24, 2026
added 5 commits
June 25, 2026 00:15
…el contract Address PR deepmodeling#5581 review (iProzd + chatgpt-codex): - edge_force_virial now returns PER-FRAME virial (nf,3,3) instead of a single (3,3) that collapsed all frames; takes n_node (per-frame counts), assigns each edge to its dst's frame via searchsorted(cumsum(n_node)). Adds a multi-frame regression test (the bug was invisible to the single-frame UT). - builder.py: document the LEGACY-COMPATIBLE / sel-truncating contract of neighbor_graph_from_extended / build_neighbor_graph explicitly; the carry-all graph-native builder is separate (from_ijs via ASE/vesin).
…ding) Address review follow-up: the per-frame-virial bug slipped through because the only test was single-frame. Generalize coverage: - edge_force_virial: add node_capacity for a padded node axis; new test exercises the MOST GENERAL layout (ragged n_node=[3,5], uneven per-frame edges, masked guard edges, padded node axis) + per-frame virial/atom-virial closure. - builder: strengthen multi-frame test to TWO DIFFERENT frames checked per-frame vs the brute-force oracle (non-periodic + periodic); add multi-frame adapter test. (Builder is rectangular-node by construction; ragged-node building is a future ragged builder.)
…ce_virial tests The ragged+padded general test subsumes the square multi-frame and the single-frame per-frame/edge-padding cases. Keep two minimal convention-pinning tests the general test cannot assert: force dst-src sign, and full-to-src attribution (decision deepmodeling#5 — the general closure is attribution-agnostic).
n_node=[3,0,5] with an empty middle frame: verifies searchsorted frame assignment on duplicate cumsum boundaries [3,3,8] skips the zero-width block (node 3 -> frame 2) and the empty frame's virial is exactly zero. Combined with a padded node axis (node_capacity 9 > sum 8) to also exercise node padding at the primitive level (the only node_capacity consumer in this PR; a model-level caller arrives with the graph lower in PR-A/jax in PR-F).
…virial Single- and multi-frame systems with nodes but NO real edges (only masked guard edges): force/atom_virial/per-frame virial must all be exactly zero with correct shapes. Covers isolated-atom / rcut-below-all-distances geometries.
added 6 commits
June 25, 2026 08:47
…onverter to from_dense_quartet build_neighbor_graph now SEARCHES neighbors directly from coordinates and keeps every neighbor within rcut (sel-free), instead of routing through the sel-truncating dense nlist. The legacy quartet adapter is renamed neighbor_graph_from_extended -> from_dense_quartet and documented as the backward-compat converter (inherits sel truncation), separating the 'compute from geometry' dispatcher from the 'adapt an existing list' converters (spec decision deepmodeling#17).
…r exclusion in build_neighbor_graph
…tet over frames Replace the per-frame Python loop in both the carry-all search and the quartet converter with flat (frame, center, neighbor) index grids and cross-frame gathers via (frame * nall + idx) flat indices. No behavior change (edges stay frame-major); tests unchanged and green.
A pair that is a neighbor ONLY across the boundary (direct distance > rcut, minimum image < rcut): build_neighbor_graph must find the image edge with the right vector, find nothing without the box, and agree with the brute-force oracle. Existing periodic test had the pair as a direct neighbor too, so it did not isolate periodic-image handling.
Run build_neighbor_graph and from_dense_quartet on torch.from_numpy inputs (periodic + non-periodic) and assert the neighbor environment matches the numpy result and the output stays in the torch namespace. Guards the vectorized broadcasting/gather against numpy-only ops that break torch/jax.
Remove the shallow 'smoke' test; torch-namespace behavior of the graph builders is covered properly by PR-A's cross-backend consistency harness (full graph forward under pt_expt/torch), not a standalone smoke check.
OutisLi
reviewed
Jun 25, 2026
OutisLi
left a comment
OutisLi
left a comment
Collaborator
There was a problem hiding this comment.
Thanks @wanghan-iapcm — clean, self-contained foundation, and the test coverage is excellent (force sign, the full-to-src ↔ per-frame virial closure, empty/ragged frames, padding/guard edges, and the cross-PBC-image case). I re-derived the force/virial assembly independently and it checks out. A few non-blocking points inline. The only one I'd suggest addressing before merge is the dangling memory/spec_unified_edge_nlist.md references, since they'll be confusing for other maintainers.
github-merge-queue Bot
pushed a commit
that referenced
this pull request
Jun 29, 2026
…5583) ## Summary Adds the graph-native forward path for `dpa1(attn_layer=0)` (the factorizable, mixed-types case), built on the `NeighborGraph` foundation from #5581. Geometry enters the descriptor only through per-edge `edge_vec`; the neighbor-axis reduction becomes a `segment_sum` over edge centers. For `pt_expt` this becomes the **default** forward (force/virial via a single autograd backward through `edge_vec`). ## What it adds - **dpmodel**: `edge_env_mat` (per-edge env-mat 4-vector), `DescrptBlockSeAtten._call_graph` + `DescrptDPA1.call_graph`, model `call_lower_graph` (energy), `neighbor_graph_from_ijs` + an optional **ASE** O(N) carry-all builder. - **pt_expt**: `edge_energy_deriv` (autograd `grad(E, edge_vec)` → `edge_force_virial`) + `forward_common_lower_graph` (energy + force + virial + atom_virial). - The dense `DescrptDPA1.call` becomes a thin adapter (`from_dense_quartet → call_graph`) preserving the 5-tuple ABI; a **shape-static** converter keeps it `jax.jit` / `torch.export`-traceable. ## Default behavior - **pt_expt** defaults graph-eligible `dpa1(attn_layer=0, concat tebd, no exclude_types)` models to the carry-all graph (it has the autograd force/virial path). - **dpmodel/jax** keep the dense default (they compute force/virial analytically; the graph lower is energy-only), and **agree with pt_expt at non-binding `sel`**. - Ineligible configs (attention, strip tebd, `exclude_types`, linear/ZBL) fall back to the dense path unchanged. `neighbor_graph_method="legacy"` forces dense; `"dense"`/`"ase"` force the graph. ## Parity (graph vs legacy dense lower, fp64 CPU) | | energy | force | virial | atom_virial | |---|---|---|---|---| | max abs diff | 0 | ~1e-19 | ~1e-18 | **~1e-18** | atom_virial matches the canonical TF==pt-legacy full-to-src convention. dpa1 descriptor + model consistency suites green across dp/jax/pt_expt. ## Known limitations - Default-flip is **pt_expt-only**; full carry-all default for dp/jax needs analytical/jax graph force (follow-up). - `make_fx` (forward + grad) traces; **full `.pt2` AOTI export is a follow-up** (PR-B). The carry-all builders (`build_neighbor_graph`/`from_ijs`) still use `nonzero` (eager-only); their static variants land with the export PR. - Single-rank only; CUDA unvalidated (CPU box); ASE is opt-in O(N) (vesin O(N) is a follow-up); no jax graph force / dpa2-3 message-passing yet. Also folds in three follow-up fixes to the #5581 foundation from @OutisLi's review (dangling spec refs → design discussion, `edge_force_virial` jax int-sum short-circuit, `Array` typing). <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added graph-native “lowering” for DPA1 when compatible, including graph-native descriptor/forward execution and graph-native descriptor→model output conversion. * Introduced opt-in `neighbor_graph_method` routing for energy/force/virial, with carry-all neighbor graphs and graph-output fitting/post-processing. * Added new neighbor-graph utilities (including ASE-based carry-all building, `(i,j,S)` conversion, and per-edge environment-matrix computation), exported as part of the public API. * **Bug Fixes** * Improved stability for masked/padded edges, virtual atom handling, and parameter protection consistency; refined traced virial assembly when node-capacity is used. * **Tests** * Expanded parity/regression suites for graph lowering, energy/force/virial, conversion correctness, ragged graphs, and FX tracing. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: Han Wang <wang_han@iapcm.ac.cn> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
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
3 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
First, foundational slice of the unified
NeighborGraphneighbor-list design (design discussion: wanghan-iapcm#4). This PR adds the backend-agnosticdeepmd/dpmodelcore only — pure numpy / array-API, no torch/jax/model dependencies — so it is self-contained and not yet wired into any model.What this adds (
deepmd/dpmodel/utils/)neighbor_graph.py—NeighborGraphdataclass (the edge-graph contract:n_node,edge_index (2,E)SoA[src,dst],edge_vec (E,3) = r_src − r_dst,edge_mask, optionaln_local/angle_index/angle_mask) andGraphLayout(edge/angle/node/frame capacities +min_edges). Node validity is derived (arange(N) < sum(n_node)), not a stored field;node_validity_maskandpad_and_guard_edges(compact-prefix layout, capacity/guard, overflow) helpers.segment.py— mask-aware, backend-dispatchedsegment_sum/segment_mean(onxp_add_at).numpy_neighbor_list.py—NumpyNeighborList, an all-pairs reference builder (coords →NeighborGraph, periodic + non-periodic, frame-offset).edge_derivatives.py—edge_force_virial: assembles per-node force, atom virial (attributed full-to-src, the canonical TF==PT convention), and global virial from a per-edge gradientg_e.Tests
source/tests/common/dpmodel/: 22 tests, all passing — dataclasses,node_validity_mask, segment ops,pad_and_guard_edges, builder correctness (brute-force-verified, periodic + non-periodic + multi-frame + capacity padding), and force/virial assembly (force formula, global virial, full-to-srcatom virial summing to global, masked-edge no-op).Scope / deferred
Numpy-only foundation. Deferred to follow-ups: torch/jax backend dispatch + O(N) vesin/nvalchemiops builders; the
forward_common_lower_graphmodel seam +edge_energy_derivautograd; descriptor ports; attentionsegment_softmax+ the 3-body angle list; jax static-E_maxcapacity/bin-packing; C/C++ ABI + LAMMPS + multi-rank.Known limitations
NumpyNeighborListis an O(N²·shells) reference builder with an orthorhombic-shell heuristic — reference/test use; production O(N) builders come later.edge_force_virialconsumes a pre-computedg_e; the autograd that produces it is a follow-up.segment_max/segment_softmaxnot yet implemented (only needed by attention descriptors later).Summary by CodeRabbit