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

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