feat(pt_expt): O(N) on-device NeighborGraph builders — vesin & nvalchemiops (PR-C)

[wanghan-iapcm]

NeighborGraph PR-C — O(N) on-device graph builders (vesin / nv)

Third PR of the NeighborGraph series (after #5581 foundation, #5583 PR-A dpa1 graph forward, #5604 PR-B .pt2/C++). Adds two O(N) carry-all NeighborGraph builders behind the World-2 neighbor_graph_method dispatcher, replacing PR-A's O(N²) dense search / per-frame ASE stopgap with on-device cell lists.

What

• build_neighbor_graph_vesin (deepmd/pt_expt/utils/vesin_graph_builder.py) — vesin.torch cell list. Device-following: runs the search on the input tensor's device (CUDA kernel on CUDA input, CPU cell list otherwise).
• build_neighbor_graph_nv (deepmd/pt_expt/utils/nv_graph_builder.py) — nvalchemiops GPU cell list, frame-batched (batch_idx/batch_ptr, one kernel for all frames — no Python loop). CUDA-only.
• Both structurally clone build_neighbor_graph_ase: search → per-frame local (i, j, S) → neighbor_graph_from_ijs(...), which recomputes edge_vec differentiably from the original grad-carrying coords.
• Wired into the pt_expt make_model graph dispatch (neighbor_graph_method ∈ {"legacy","dense","ase","vesin","nv"}) and DeepEval .pt2 graph inference (new neighbor_graph_method kwarg, default "dense" → existing inference byte-identical). dpmodel/jax fail-fast on vesin/nv (torch/CUDA-only).

Perf-only: all builders emit the SAME neighbor set as dense (carry-all, sel=normalization-only), proven by exact set-equality; energy/force/virial are unchanged (parity 1e-12 CPU / 1e-10 CUDA).

Layering

• dpmodel stays torch-free: vesin/nv builders live in pt_expt; the dpmodel dispatch only carries a fail-fast message.
• vesin/nv are optional deps, NOT in pyproject — lazy-imported, guarded by is_vesin_torch_available() / is_nv_available(), ImportError with an install hint on absence.
• nv decode is a faithful transcription of the tested deepmd/pt/utils/nv_nlist.py:_matrix_to_extended_inputs Step-1 extraction.

Testing

• Local (CPU): 13 passed (vesin builder 5, vesin/reject dispatch 2, DeepEval graph 6), nv self-skips.
• Remote GPU (Tesla T4), commit 78f6c24:
  • test_nv_graph_builder.py: 4 passed (set-equality vs dense periodic+non-periodic, frame-batch, differentiable edge_vec).
  • test_graph_builder_dispatch.py: 3 passed on CUDA (vesin parity 1e-10, dpmodel reject vesin+nv, nv parity 1e-10).
  • test_graph_deepeval.py: 6 passed on CUDA (.pt2 graph dense parity + vesin, AOTI compile + torch.as_tensor extraction).

Known limitations

• vesin per-frame Python loop — vesin.torch.compute is single-system (no batch dim), so multi-frame vesin loops over frames (each an O(N) search; nf=1 inference has zero loop cost). nv batches natively — the loop-free path for batched training.
• vesin/nv not in pyproject; no "auto" selector (explicit strings only).
• nv search_capacity = max(64, nloc) initial heuristic + 1.25× grow loop (.item() host-sync per grow).
• nv passes the normalized coord to from_ijs (differentiable lattice translation, identity gradient; edge_vec is lattice-invariant) — consistent with the pt nv-nlist path.
• jax O(N) graph builders (matscipy/jax-md) = PR-F; attention/angles/MP = PR-D/E/G.

Implements plan_neighbor_graph_prC_implementation; design spec: discussion wanghan-iapcm#4.

Summary by CodeRabbit

• New Features
  • Added experimental neighbor-graph builders supporting NVIDIA CUDA (NV) and Vesin (vesin), including periodic handling and frame-aware edge construction.
  • Enabled configurable neighbor_graph_method for graph-form .pt2 inference and tightened dispatch support to dense/ase/vesin/nv.
• Bug Fixes
  • Excluded virtual atoms (type < 0) consistently from neighbor-graph edges for ASE, Vesin, and NV paths.
  • Improved robustness for empty/zero-neighbor cases and ensured edge vectors remain differentiable.
• Tests
  • Added backend parity, NV decode regression, gradient/device checks, virtual-atom coverage, and strengthened plugin import validation.

[coderabbitai]

📝 Walkthrough

Walkthrough

The PR adds NV and vesin neighbor-graph builders for PT experimental graph flows, routes inference and model dispatch through method चयन, filters virtual atoms in ASE graph construction, and expands regression coverage for decoding, parity, and import behavior.

Changes

Neighbor-graph builders and dispatch

Layer / File(s)	Summary
NV builder deepmd/pt_expt/utils/nv_graph_builder.py, deepmd/pt/utils/nv_nlist.py, source/tests/pt_expt/utils/test_nv_matrix_decode.py, source/tests/pt_expt/utils/test_nv_graph_builder.py	A CUDA-backed builder decodes dense neighbor matrices, grows capacity, filters virtual atoms, and is reused by the PT NV neighbor-list path and its decode/parity tests.
Vesin builder deepmd/pt_expt/utils/vesin_graph_builder.py, deepmd/pt_expt/utils/vesin_neighbor_list.py, source/tests/pt_expt/utils/test_vesin_graph_builder.py	A vesin-backed builder computes per-frame graph edges, is shared with the PT vesin neighbor-list path, and is covered by parity, device, gradient, empty-input, and virtual-atom tests.
Graph dispatch deepmd/pt_expt/infer/deep_eval.py, deepmd/pt_expt/model/make_model.py, source/tests/pt_expt/model/test_graph_builder_dispatch.py, source/tests/pt_expt/utils/test_graph_pt2_metadata.py	PT experimental inference and model code now dispatch neighbor-graph construction by method and validate artifact eligibility, with tests for parity and rejection cases.
Virtual atom filtering deepmd/dpmodel/utils/neighbor_graph/ase_builder.py, source/tests/common/dpmodel/test_from_ijs.py	ASE neighbor-graph construction now excludes edges touching negative atype values, with matching regression coverage.
Plugin import state restoration source/tests/pt_expt/test_plugin.py	The plugin-loading test now snapshots and restores the deepmd.pt_expt module subtree during re-import.

Estimated code review effort: 4 (Complex) | ~60 minutes

Possibly related issues

• [Code scan] Filter virtual atoms in ASE neighbor-graph construction #5664: The ASE builder’s atype < 0 filtering and matching regression test address the same virtual-atom edge exclusion behavior.

Possibly related PRs

• deepmodeling/deepmd-kit#5284: Shares the PT-EXPT inference path changes around deepmd/pt_expt/infer/deep_eval.py and deepmd/pt_expt/model/make_model.py.
• deepmodeling/deepmd-kit#5491: Touches the same vesin neighbor-search plumbing via vesin_search_ijs.
• deepmodeling/deepmd-kit#5604: Also modifies the PT-EXPT graph evaluation path in deepmd/pt_expt/infer/deep_eval.py.

Suggested reviewers: OutisLi, njzjz

🚥 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 captures the main change: adding on-device NeighborGraph builders backed by vesin and nvalchemiops.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ 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

❌ Patch coverage is 60.11561% with 69 lines in your changes missing coverage. Please review.
✅ Project coverage is 81.13%. Comparing base (9bd1f16) to head (d2cf920).
⚠️ Report is 6 commits behind head on master.

Files with missing lines	Patch %	Lines
deepmd/pt_expt/utils/nv_graph_builder.py	29.72%	52 Missing ⚠️
deepmd/pt/utils/nv_nlist.py	0.00%	8 Missing ⚠️
deepmd/pt_expt/infer/deep_eval.py	80.76%	5 Missing ⚠️
deepmd/pt_expt/model/make_model.py	70.00%	3 Missing ⚠️
deepmd/pt_expt/utils/vesin_graph_builder.py	97.87%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #5714      +/-   ##
==========================================
- Coverage   81.24%   81.13%   -0.12%     
==========================================
  Files         981      990       +9     
  Lines      109860   111007    +1147     
  Branches     4234     4232       -2     
==========================================
+ Hits        89257    90063     +806     
- Misses      19080    19420     +340     
- Partials     1523     1524       +1     

☔ 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.

[iProzd]

Agree with @OutisLi, and I think this is a correctness gap for the opt-in O(N) builders, not just a docstring scope note. Confirmed in code:

• build_neighbor_graph (dense) drops virtual-atom edges via keep_mask = ... & ~(extended_atype<0) & ~(atype<0).
• vesin/nv/ase build geometric edges from coords only and never look at atype.
• forward_common_atomic_graph masks edges only for pair_excl; for atype<0 it only clamps the type and zeroes the virtual atom's own output — it does not drop edges from a virtual neighbor into a real center.

So for inputs with atype<0 (e.g. mixed-natoms padding), a real center that has a virtual atom within rcut gets an extra type-0 edge under vesin/nv/ase but not under dense, polluting the real atom's descriptor/energy/force. The "SAME neighbor set / parity" claim then holds only for all-real-atom inputs, which is likely all the current set-equality fixtures cover.

Rather than only documenting the caller's responsibility, could the O(N) builders filter edges touching atype<0 (drop where center or neighbor-owner is virtual) to match dense, so the parity claim stays literally true? A set-equality test with a virtual-atom fixture would lock it in. Default users are unaffected (default resolves to dense), so this is opt-in-only severity, but it undermines the stated dense parity.

[OutisLi]

Please rework the nv graph path to reuse the existing NvNeighborList implementation rather than adding a second nvalchemiops search stack. The new graph builder duplicates the kernel setup, capacity loop, device handling, and matrix decode that NvNeighborList already owns, which increases maintenance cost and has already diverged in behavior. This can be fixed by extending NvNeighborList with a graph-oriented return mode or a thin build_graph method while preserving the existing extended/edges modes so current SeZM users are not affected.

[iProzd]

Others LGTM.

[OutisLi]

The same design concern applies to both nv and vesin: graph-form NeighborGraph construction should extend the existing NeighborList strategies instead of adding separate parallel builders. VesinNeighborList and NvNeighborList already provide the O(N) search backends used by SeZM through extended/edges return modes; the graph path should be another output contract on the same strategy layer, not a second copy of each backend. Please unify this design while preserving the existing extended and edges modes for current SeZM users.

[OutisLi]

Please also restore the graph-lower eligibility check in the pt_expt graph dispatch. Explicit neighbor_graph_method values should not bypass the same mixed_types/uses_graph_lower contract enforced by dpmodel and by graph export.

[OutisLi]

Please revisit the public neighbor-backend API design. The PR now exposes two knobs, nlist_backend and neighbor_graph_method, whose scopes are hard to infer from their names and whose supported values differ. From a user perspective both choose the neighbor construction backend, so the API should be made more natural and unambiguous. In particular, if nv is a supported O(N) backend, it should not work only for graph-form .pt2 while the nlist path has no nv option, especially since NvNeighborList already exists.

[coderabbitai]

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

deepmd/pt_expt/utils/nv_graph_builder.py (1)

242-244: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick win

Handle the documented 2D coord form here. Any (nf, nloc*3) tensor is reshaped with nf = 1, so a flattened multi-frame input is merged into one neighbor graph; periodic inputs will also fail once box is reshaped with the wrong frame count. Derive nf from atype (or reject 2D inputs explicitly).

🤖 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/pt_expt/utils/nv_graph_builder.py` around lines 242 - 244, The coord
reshaping logic in nv_graph_builder’s neighbor-graph setup does not correctly
handle the documented 2D coord form, since it forces nf=1 for any non-3D tensor
and can merge flattened multi-frame inputs into one graph. Update the
frame-count derivation in this coord/box preprocessing path to use atype (or
explicitly reject unsupported 2D inputs), and make sure the same nf is applied
consistently before reshaping coord and box.

🧹 Nitpick comments (2)

deepmd/pt_expt/infer/deep_eval.py (1)

149-208: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick win

Consider validating vesin/nv availability at construction, not first eval().

The new fail-fast check (Lines 181-194) only validates that neighbor_graph_method is used on a graph-form artifact; it does not validate that "vesin" (needs vesin.torch) or "nv" (CUDA-only per PR objectives) are actually usable. _setup_nlist_backend (Lines 209-251) eagerly validates its equivalent vesin option (raises ImportError/ValueError at construction), but neighbor_graph_method="vesin"/"nv" will only fail once _build_eval_graph is reached inside the first eval() call — a less consistent, later failure point for the same class of misconfiguration.

💡 Possible early-validation addition

         if (
             neighbor_graph_method != "dense"
             and getattr(self, "metadata", {}).get("lower_input_kind") != "graph"
         ):
             raise ValueError(
                 f"neighbor_graph_method={neighbor_graph_method!r} only applies to "
                 "graph-form .pt2 artifacts (lower_input_kind == 'graph'); this "
                 f"model is not graph-form. Use nlist_backend to select the "
                 "neighbor-list builder for the nlist path."
             )
+        if neighbor_graph_method == "vesin" and not is_vesin_torch_available():
+            raise ImportError(
+                "neighbor_graph_method='vesin' requires 'vesin.torch'; "
+                "install it (`pip install vesin[torch]`) or choose another method."
+            )
+        if neighbor_graph_method == "nv" and not torch.cuda.is_available():
+            raise ValueError(
+                "neighbor_graph_method='nv' requires CUDA; use 'vesin' or 'dense' "
+                "on CPU."
+            )

🤖 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/pt_expt/infer/deep_eval.py` around lines 149 - 208, Validate
neighbor_graph_method availability in __init__ alongside the existing graph-form
check, so misconfigurations fail before first eval(). In deep_eval.py, extend
the constructor logic around neighbor_graph_method to eagerly reject "vesin"
when vesin.torch is unavailable and "nv" when CUDA is unavailable or
unsupported, mirroring the early validation already done by
_setup_nlist_backend. Use the same error style as _setup_nlist_backend and keep
the check tied to _neighbor_graph_method / _is_pt2 /
metadata["lower_input_kind"] so graph-form .pt2 is the only path affected.

deepmd/pt_expt/model/make_model.py (1)

459-469: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Duplicate eligibility-check logic with _resolve_graph_method.

This new guard repeats the exact descriptor = getattr(...); uses_graph_lower = getattr(...) pattern from _resolve_graph_method (Lines 405-409). Extracting a shared helper (e.g. _graph_lower_eligible()) would avoid the two copies drifting if the eligibility rule ever changes.

♻️ Suggested extraction

+        def _graph_lower_eligible(self) -> bool:
+            descriptor = getattr(self.atomic_model, "descriptor", None)
+            uses_graph_lower = getattr(descriptor, "uses_graph_lower", lambda: False)
+            return self.mixed_types() and uses_graph_lower()
+
         def _resolve_graph_method(
             self, neighbor_graph_method: str | None
         ) -> str | None:
             ...
-            descriptor = getattr(self.atomic_model, "descriptor", None)
-            uses_graph_lower = getattr(descriptor, "uses_graph_lower", lambda: False)
-            if self.mixed_types() and uses_graph_lower():
+            if self._graph_lower_eligible():
                 return "dense"
             return None

         def _call_common_graph(...):
             ...
-            descriptor = getattr(self.atomic_model, "descriptor", None)
-            uses_graph_lower = getattr(descriptor, "uses_graph_lower", lambda: False)
-            if not (self.mixed_types() and uses_graph_lower()):
+            if not self._graph_lower_eligible():
                 raise NotImplementedError(
                     "neighbor_graph_method requires a mixed_types descriptor with a "
                     "graph lower (e.g. dpa1 attn_layer=0)"
                 )

🤖 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/pt_expt/model/make_model.py` around lines 459 - 469, The new
neighbor_graph_method guard in make_model repeats the same
descriptor/uses_graph_lower eligibility check already implemented in
_resolve_graph_method, so extract that shared logic into a helper such as
_graph_lower_eligible() and use it in both places. Update the eligibility check
inside the model-building path around atomic_model/descriptor handling so the
rule lives in one place and cannot drift between _resolve_graph_method and the
new guard.

🤖 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.

Outside diff comments:
In `@deepmd/pt_expt/utils/nv_graph_builder.py`:
- Around line 242-244: The coord reshaping logic in nv_graph_builder’s
neighbor-graph setup does not correctly handle the documented 2D coord form,
since it forces nf=1 for any non-3D tensor and can merge flattened multi-frame
inputs into one graph. Update the frame-count derivation in this coord/box
preprocessing path to use atype (or explicitly reject unsupported 2D inputs),
and make sure the same nf is applied consistently before reshaping coord and
box.

---

Nitpick comments:
In `@deepmd/pt_expt/infer/deep_eval.py`:
- Around line 149-208: Validate neighbor_graph_method availability in __init__
alongside the existing graph-form check, so misconfigurations fail before first
eval(). In deep_eval.py, extend the constructor logic around
neighbor_graph_method to eagerly reject "vesin" when vesin.torch is unavailable
and "nv" when CUDA is unavailable or unsupported, mirroring the early validation
already done by _setup_nlist_backend. Use the same error style as
_setup_nlist_backend and keep the check tied to _neighbor_graph_method / _is_pt2
/ metadata["lower_input_kind"] so graph-form .pt2 is the only path affected.

In `@deepmd/pt_expt/model/make_model.py`:
- Around line 459-469: The new neighbor_graph_method guard in make_model repeats
the same descriptor/uses_graph_lower eligibility check already implemented in
_resolve_graph_method, so extract that shared logic into a helper such as
_graph_lower_eligible() and use it in both places. Update the eligibility check
inside the model-building path around atomic_model/descriptor handling so the
rule lives in one place and cannot drift between _resolve_graph_method and the
new guard.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Pro

Run ID: 2d7badff-46b5-40f8-8717-74ed796a6485

📥 Commits

Reviewing files that changed from the base of the PR and between 6396068 and 0134100.

📒 Files selected for processing (8)

• deepmd/pt/utils/nv_nlist.py
• deepmd/pt_expt/infer/deep_eval.py
• deepmd/pt_expt/model/make_model.py
• deepmd/pt_expt/utils/nv_graph_builder.py
• deepmd/pt_expt/utils/vesin_graph_builder.py
• deepmd/pt_expt/utils/vesin_neighbor_list.py
• source/tests/pt_expt/model/test_graph_builder_dispatch.py
• source/tests/pt_expt/utils/test_graph_pt2_metadata.py

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

• source/tests/pt_expt/model/test_graph_builder_dispatch.py
• deepmd/pt_expt/utils/vesin_graph_builder.py