Neighbor List Backend Integration#

Overview#

DeepMD-kit's pt_expt module (experimental PyTorch extension) implements a pluggable neighbor-list architecture that replaces the historical O(N²) all-pairs builder with optional O(N) cell-list backends. This addresses a key front-end bottleneck in Python/ASE inference for large systems (identified in DPA-4 manuscript §2.4) 1.

There are two distinct subsystems serving different model-export formats:

• NeighborList — for the forward_common Python/ASE inference path (nlist-form models)
• NeighborGraph — for graph-form .pt2 / AOTI / C++ inference (e.g., DPA graph models)

The O(N²) Native Baseline#

The historical builder 2:

1. Extends local coords into ~27 periodic-image buffer regions via extend_coord_with_ghosts
2. Builds a dense O(N²) [N, 27N] distance matrix via build_neighbor_list
3. Materializes ≈27×N ghost atoms regardless of actual neighbor density

Native baseline files:

• deepmd/dpmodel/utils/nlist.py — extend_coord_with_ghosts, build_neighbor_list
• deepmd/pt/utils/nlist.py — PyTorch version using torch.topk

Pluggable NeighborList (nlist-form inference)#

Introduced in PR #5491 3, the strategy pattern is injected at forward_common/call_common. The forward_common_lower export entry (.pt2/AOTI/C++) is left untouched — zero export risk 4.

VesinNeighborList builds an (i, j, S) edge list and materializes only real-neighbor ghosts (coord[j] + S@box), emitting the same (extended_coord, extended_atype, nlist, mapping) quartet. Forces, virials, and atomic virials flow through existing autograd and communicate_extended_output routines unchanged 5.

Backend Selection: nlist_backend#

Configured on pt_expt.DeepEval and the ASE DP calculator 6:

Pluggable NeighborGraph (graph-form .pt2 inference)#

Introduced in PR #5714 7 as PR-C of the NeighborGraph series (after foundation PRs #5581, #5583, #5604):

Both builders follow the pattern: search → per-frame (i, j, S) → neighbor_graph_from_ijs(...), recomputing edge_vec differentiably from the original grad-carrying coords 8.

Backend Selection: neighbor_graph_method#

Wired into pt_expt make_model dispatch and DeepEval .pt2 inference 9:

There is no "auto" selector for the graph path — explicit strings only 10.

Architecture Layers#

dpmodel (torch-free core)
  ├── NeighborList base class, DefaultNeighborList
  └── Graph dispatch: fail-fast on vesin/nv (no torch dependency)

pt_expt (PyTorch experimental)
  ├── VesinNeighborList → nlist-form inference
  ├── vesin_graph_builder → graph-form inference (device-following)
  ├── nv_graph_builder → graph-form inference (CUDA-only, batched)
  └── DeepEval dispatch → nlist_backend / neighbor_graph_method params

The dpmodel layer stays torch-free. dpmodel/jax calls with "vesin" or "nv" fail fast with an explicit error message 11. vesin/nv are not listed in pyproject.toml — they are optional, lazy-imported, guarded by is_vesin_torch_available() / is_nv_available() with ImportError + install hints at initialization time 12.

Error Handling and Known Gaps#

Graceful degradation#

• nlist_backend="auto" falls back silently to native if vesin is unavailable, the model is spin, or an explicit ASE neighbor_list is set 6.
• Strict modes ("vesin", "nv") fail fast at initialization — not at call time.

Known gaps#

• No "auto" for graph-form path: users must specify neighbor_graph_method explicitly; there is no probe-and-fallback 10.
• No real-GPU CI: CUDA paths are manually validated (Tesla T4) but not in automated CI — CUDA compilation failures will not be caught pre-merge 13 14.
• nv capacity heuristic: nvalchemiops starts at max(64, nloc) capacity, grows at 1.25× with a .item() host-sync per overflow. Overflows are re-tried silently, not reported 15.
• vesin per-frame Python loop: vesin.torch.compute is single-system only; multi-frame calls loop in Python. Acceptable for nf=1 inference; prefer "nv" for batched training 16.
• Spin/hessian/dipole/polar models: vesin is gated off for spin; dipole, polar, dos, hessian, and multi-task models are not covered by vesin equivalence tests 17.

Import-time crash fix (PR #5542) 18#

PR #5491's import chain caused plain dp test (pt backend) to crash when the C++ custom-op library (libdeepmd_op_pt.so) was absent. The pt_expt/utils/__init__.py was eagerly importing tabulate_ops, which monkeypatched torch.ops.deepmd.* — then register_fake() failed with "operator does not exist." Fix: removed the eager import and introduced _op_exists(name) to check for a real torch._ops.OpOverloadPacket before calling register_fake().

Dependency packaging (PR #5501) 19#

vesin[torch] was in core dependencies, breaking conda-forge (which packages vesin without torch bindings). Moved to the torch extra in pyproject.toml and backend/find_pytorch.py. Only pip install deepmd-kit[torch] pulls vesin-torch; conda-forge users with nlist_backend="auto" degrade gracefully to native.

Numerics#

All backends produce the same neighbor set (carry-all; sel is normalization-only). Parity is confirmed across 8 descriptor families 20 21:

• CPU: ≤ 1e-12 (fp round-off; difference is ghost-enumeration order only)
• CUDA: ≤ 1e-10

Key References#