feat(pt_expt): pluggable NeighborList strategy + O(N) vesin neighbor list for Python/ASE inference (deepmodeling#5491)

## Motivation

When deepmd consumes a neighbor list, `forward_common`/`call_common`
extends the local region into ~26 periodic-image buffer regions —
`extend_coord_with_ghosts` then a dense `build_neighbor_list` —
materializing ≈27×N ghost atoms and an O(N²) `[N, 27N]` distance matrix.
This is the Python/ASE front-end bottleneck at large N (DPA4 manuscript
§2.4).

## What this PR does

Makes neighbor-list construction **pluggable** via an optional
`NeighborList` strategy injected at `forward_common`/`call_common` (the
layer where the system is extended). The exported `forward_common_lower`
(the `.pt2`/AOTI/C++ entry) is left untouched, so there is **zero export
risk**.

- **dpmodel (torch-free core)** — `NeighborList` base +
`DefaultNeighborList` (the historical dense extend+build).
`neighbor_list=None` reproduces today's behavior **byte-identically**.
- **pt_expt** — `VesinNeighborList`, a device-aware
[`vesin.torch`](https://github.com/Luthaf/vesin) O(N) cell list: it runs
on the input tensor's device (CPU or CUDA) for torch, and is CPU-bridged
for numpy/dpmodel. It builds an `(i, j, S)` edge list, materializes only
the real-neighbor ghosts `coord[j] + S@box`, and emits the same extended
quartet `(extended_coord, extended_atype, nlist, mapping)`. Because the
representation is identical, force / global-virial / **atomic-virial**
all come out of the existing autograd + `communicate_extended_output`
routines unchanged.
- **inference** — `nlist_backend="auto" | "vesin" | "native"` on the
pt_expt `DeepEval` and the ASE `DP` calculator. `auto` uses vesin when
available/applicable and silently falls back to native otherwise;
`vesin` is strict (raises if unavailable, or for spin /
ASE-`neighbor_list` conflicts); `native` forces the dense builder.
- **pyproject** — depends on `vesin[torch]`.

## Verification

native vs vesin agree to fp round-off (energy `0.0`,
force/virial/atomic-virial ≤ ~1e-18, the only difference being
ghost-enumeration order):

- `source/tests/pt_expt/utils/test_neighbor_list.py` — builder
equivalence (numpy + torch namespaces, PBC/non-PBC, input-device
placement) and full **model** equivalence across 8 descriptor families
(se_e2_a, se_r, se_e3, dpa1, se_atten_v2, dpa2, dpa3, hybrid) for
dpmodel (energy/atomic-energy) and pt_expt
(energy/force/virial/atomic-virial), plus the `neighbor_list=None`
byte-identical fallback.
-
`source/tests/pt_expt/infer/test_deep_eval.py::TestDeepEvalNlistBackend`
— `nlist_backend` dispatch validation and vesin-vs-native equality
through the compiled `.pte`.

## Known limitations

- **Python `forward_common` path only.** This is the path where deepmd
builds the nlist (DeepPot / ASE). C/C++/LAMMPS enter at the exported
`forward_lower` with an externally-supplied list — accepting `(i,j,S)`
there is a planned follow-up.
- **Energy model validated;** spin is gated off vesin;
dipole/polar/dos/hessian/multi-task and fparam/aparam are not yet
covered by vesin equivalence tests (the seam is model-agnostic and the
dense default path is regression-tested).
- **Still materializes O(surface) ghost coords** (the minimal extended
array); a truly buffer-free / sparse-edge-list consumption is deferred.
In dense descriptors the env-mat per-edge tensors dominate memory
anyway.
- **No real-GPU CI** (the device test runs on CPU); ghosts are not
deduped (one per `S≠0` edge — correctness preserved via `mapping`
summation).
- Multi-frame vesin build and the neighbor-truncation (`sel`-exceeded)
path are not yet directly tested.


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* Pluggable neighbor-list backend for calculators and evaluators: "auto"
(default), "native", or "vesin"; preserves historical all‑pairs behavior
when None.
* Adds a Vesin-based O(N) neighbor-list option for faster neighbor
construction and an explicit Default (all‑pairs) builder.

* **Tests**
* Comprehensive test suites validating backend selection, error
handling, device/shape robustness, and numeric equivalence (native vs
vesin), including multi-frame cases.

* **Chores**
  * Adds vesin[torch] runtime entry for the new backend.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Co-authored-by: Han Wang <wang_han@iapcm.ac.cn>

Author: wanghan-iapcm

deepmd/calculator.py

@@ -91,13 +91,15 @@ def __init__(
         type_dict: dict[str, int] | None = None,
         neighbor_list: Optional["NeighborList"] = None,
         head: str | None = None,
+        nlist_backend: str = "auto",
         **kwargs: Any,
     ) -> None:
         Calculator.__init__(self, label=label, **kwargs)
         self.dp = DeepPot(
             str(Path(model).resolve()),
             neighbor_list=neighbor_list,
             head=head,
+            nlist_backend=nlist_backend,
         )
         if type_dict:
             self.type_dict = type_dict

deepmd/dpmodel/model/ener_model.py

@@ -21,6 +21,9 @@
 from deepmd.dpmodel.output_def import (
     FittingOutputDef,
 )
+from deepmd.dpmodel.utils.neighbor_list import (
+    NeighborList,
+)
 
 from .dp_model import (
     DPModelCommon,
@@ -88,7 +91,23 @@ def call(
         aparam: Array | None = None,
         do_atomic_virial: bool = False,
         charge_spin: Array | None = None,
+        neighbor_list: NeighborList | None = None,
     ) -> dict[str, Array]:
+        """Evaluate the energy model.
+
+        Most arguments share the meaning of :meth:`call_common`.
+
+        Parameters
+        ----------
+        neighbor_list
+            The neighbor-list construction strategy forwarded to
+            :meth:`call_common`.  ``None`` uses the default all-pairs builder
+            (:class:`~deepmd.dpmodel.utils.neighbor_list.NeighborList`
+            subclass :class:`~deepmd.dpmodel.utils.default_neighbor_list.DefaultNeighborList`),
+            reproducing the historical behavior; an alternative strategy may be
+            injected to accelerate neighbor-list construction without changing
+            the model outputs.
+        """
         model_ret = self.call_common(
             coord,
             atype,
@@ -97,6 +116,7 @@ def call(
             aparam=aparam,
             charge_spin=charge_spin,
             do_atomic_virial=do_atomic_virial,
+            neighbor_list=neighbor_list,
         )
         model_predict = {}
         model_predict["atom_energy"] = model_ret["energy"]

deepmd/dpmodel/model/make_model.py

@@ -38,10 +38,9 @@
     check_operation_applied,
 )
 from deepmd.dpmodel.utils import (
-    build_neighbor_list,
-    extend_coord_with_ghosts,
+    DefaultNeighborList,
+    NeighborList,
     nlist_distinguish_types,
-    normalize_coord,
 )
 from deepmd.utils.path import (
     DPPath,
@@ -78,6 +77,7 @@ def model_call_from_call_lower(
     do_atomic_virial: bool = False,
     coord_corr_for_virial: Array | None = None,
     charge_spin: Array | None = None,
+    neighbor_list: NeighborList | None = None,
 ) -> dict[str, Array]:
     """Return model prediction from lower interface.
 
@@ -96,6 +96,12 @@ def model_call_from_call_lower(
         atomic parameter. nf x nloc x nda
     do_atomic_virial
         If calculate the atomic virial.
+    neighbor_list
+        The neighbor-list construction strategy.  ``None`` uses the default
+        all-pairs builder (:class:`DefaultNeighborList`), reproducing the
+        historical behavior.  An alternative strategy (e.g. an O(N) cell list)
+        may be injected to speed up neighbor-list construction; it returns the
+        same extended representation, so model outputs are unchanged.
 
     Returns
     -------
@@ -107,26 +113,9 @@ def model_call_from_call_lower(
     nframes, nloc = atype.shape[:2]
     cc, bb, fp, ap = coord, box, fparam, aparam
     del coord, box, fparam, aparam
-    if bb is not None:
-        coord_normalized = normalize_coord(
-            cc.reshape(nframes, nloc, 3),
-            bb.reshape(nframes, 3, 3),
-        )
-    else:
-        xp = array_api_compat.array_namespace(cc)
-        coord_normalized = xp.reshape(cc, (nframes, nloc, 3))
-    extended_coord, extended_atype, mapping = extend_coord_with_ghosts(
-        coord_normalized, atype, bb, rcut
-    )
-    nlist = build_neighbor_list(
-        extended_coord,
-        extended_atype,
-        nloc,
-        rcut,
-        sel,
-        # types will be distinguished in the lower interface,
-        # so it doesn't need to be distinguished here
-        distinguish_types=False,
+    builder = neighbor_list if neighbor_list is not None else DefaultNeighborList()
+    extended_coord, extended_atype, nlist, mapping = builder.build(
+        cc, atype, bb, rcut, sel
     )
     extended_coord = extended_coord.reshape(nframes, -1, 3)
     if coord_corr_for_virial is not None:
@@ -269,6 +258,7 @@ def call_common(
             do_atomic_virial: bool = False,
             coord_corr_for_virial: Array | None = None,
             charge_spin: Array | None = None,
+            neighbor_list: NeighborList | None = None,
         ) -> dict[str, Array]:
             """Return model prediction.
 
@@ -290,6 +280,11 @@ def call_common(
             coord_corr_for_virial
                 The coordinates correction for virial.
                 shape: nf x (nloc x 3)
+            neighbor_list
+                The neighbor-list construction strategy.  ``None`` uses the
+                default all-pairs builder; an alternative strategy (e.g. an O(N)
+                cell list) may be injected to speed up neighbor-list construction
+                without changing model outputs.
 
             Returns
             -------
@@ -316,6 +311,7 @@ def call_common(
                 do_atomic_virial=do_atomic_virial,
                 coord_corr_for_virial=coord_corr_for_virial,
                 charge_spin=cs,
+                neighbor_list=neighbor_list,
             )
             model_predict = self._output_type_cast(model_predict, input_prec)
             return model_predict

deepmd/dpmodel/utils/__init__.py

@@ -1,4 +1,7 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+from .default_neighbor_list import (
+    DefaultNeighborList,
+)
 from .env_mat import (
     EnvMat,
 )
@@ -15,6 +18,9 @@
     is_lmdb,
     make_neighbor_stat_data,
 )
+from .neighbor_list import (
+    NeighborList,
+)
 from .network import (
     EmbeddingNet,
     FittingNet,
@@ -53,6 +59,7 @@
 
 __all__ = [
     "AtomExcludeMask",
+    "DefaultNeighborList",
     "DistributedSameNlocBatchSampler",
     "EmbeddingNet",
     "EnvMat",
@@ -62,6 +69,7 @@
     "LmdbTestDataNlocView",
     "NativeLayer",
     "NativeNet",
+    "NeighborList",
     "NetworkCollection",
     "PairExcludeMask",
     "SameNlocBatchSampler",

deepmd/dpmodel/utils/default_neighbor_list.py

@@ -0,0 +1,59 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Default all-pairs neighbor-list builder (historical deepmd behavior)."""
+
+import array_api_compat
+
+from deepmd.dpmodel.array_api import (
+    Array,
+)
+
+from .neighbor_list import (
+    NeighborList,
+)
+from .nlist import (
+    build_neighbor_list,
+    extend_coord_with_ghosts,
+)
+from .region import (
+    normalize_coord,
+)
+
+
+class DefaultNeighborList(NeighborList):
+    """All-pairs builder: replicate the cell into periodic images and rank by
+    distance (:func:`~deepmd.dpmodel.utils.nlist.extend_coord_with_ghosts` +
+    :func:`~deepmd.dpmodel.utils.nlist.build_neighbor_list`).  This is the
+    default when no strategy is supplied, so results are unchanged.
+    """
+
+    def build(
+        self,
+        coord: Array,
+        atype: Array,
+        box: Array | None,
+        rcut: float,
+        sel: list[int],
+    ) -> tuple[Array, Array, Array, Array]:
+        xp = array_api_compat.array_namespace(coord, atype)
+        nframes, nloc = atype.shape[:2]
+        if box is not None:
+            coord_normalized = normalize_coord(
+                xp.reshape(coord, (nframes, nloc, 3)),
+                xp.reshape(box, (nframes, 3, 3)),
+            )
+        else:
+            coord_normalized = xp.reshape(coord, (nframes, nloc, 3))
+        extended_coord, extended_atype, mapping = extend_coord_with_ghosts(
+            coord_normalized, atype, box, rcut
+        )
+        # types are distinguished in the lower interface, so keep them merged here
+        nlist = build_neighbor_list(
+            extended_coord,
+            extended_atype,
+            nloc,
+            rcut,
+            sel,
+            distinguish_types=False,
+        )
+        extended_coord = xp.reshape(extended_coord, (nframes, -1, 3))
+        return extended_coord, extended_atype, nlist, mapping

deepmd/dpmodel/utils/neighbor_list.py

@@ -0,0 +1,64 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Pluggable neighbor-list construction strategies.
+
+A :class:`NeighborList` turns local coordinates (and an optional cell) into the
+*extended* representation consumed by the model's lower interface.  The default
+all-pairs builder lives in :mod:`deepmd.dpmodel.utils.default_neighbor_list`;
+backend-specific O(N) builders (e.g. the ``vesin``-based one in
+``deepmd.pt_expt.utils.vesin_neighbor_list``) subclass :class:`NeighborList`
+and are injected into the model, so the rest of the model is agnostic to how the
+neighbor list was built.
+"""
+
+from deepmd.dpmodel.array_api import (
+    Array,
+)
+
+
+class NeighborList:
+    """Strategy that builds the extended neighbor environment from local atoms.
+
+    Implementations turn local coordinates into the extended representation: the
+    coordinates and atom types of local-plus-ghost (periodic-image) atoms, a
+    candidate neighbor list indexing the extended atoms, and a mapping from each
+    extended atom to its local owner.  Implementations are stateless --
+    ``rcut``/``sel`` are supplied by the model at call time.
+    """
+
+    def build(
+        self,
+        coord: Array,
+        atype: Array,
+        box: Array | None,
+        rcut: float,
+        sel: list[int],
+    ) -> tuple[Array, Array, Array, Array]:
+        """Build the extended system and a candidate neighbor list.
+
+        Parameters
+        ----------
+        coord
+            local coordinates, shape (nf, nloc, 3) or (nf, nloc*3).
+        atype
+            local atom types, shape (nf, nloc).
+        box
+            simulation cell, shape (nf, 3, 3) or (nf, 9); ``None`` for non-periodic.
+        rcut
+            cutoff radius.
+        sel
+            number of selected neighbors per type.
+
+        Returns
+        -------
+        extended_coord
+            shape (nf, nall, 3).
+        extended_atype
+            shape (nf, nall).
+        nlist
+            shape (nf, nloc, nnei), type-undistinguished candidate neighbors
+            indexing the extended atoms (the lower interface re-formats it:
+            distance sort, truncate to ``sel``, split by type).
+        mapping
+            shape (nf, nall), mapping each extended atom to its local owner.
+        """
+        raise NotImplementedError