feat(pt/pt_expt): O(N) vesin neighbor list for Python/ASE inference

deepmd/calculator.py

@@ -49,6 +49,13 @@ class DP(Calculator):
         will infer this information from model, by default None
     neighbor_list : ase.neighborlist.NeighborList, optional
         The neighbor list object. If None, then build the native neighbor list.
+    nlist_backend : str, default: "auto"
+        Which algorithm builds the neighbor list. ``"auto"`` uses the optional
+        O(N) ``vesin`` cell list when applicable (much faster for large systems)
+        and silently falls back to the native O(N^2) builder otherwise.
+        ``"vesin"`` strictly requires the vesin path (raises if it is missing or
+        the model is spin / hessian); ``"native"`` forces the built-in
+        all-pairs builder. Ignored when an explicit ``neighbor_list`` is given.
     head : Union[str, None], optional
         a specific model branch choosing from pretrained model, by default None
 
@@ -90,13 +97,15 @@ def __init__(
         label: str = "DP",
         type_dict: dict[str, int] | None = None,
         neighbor_list: Optional["NeighborList"] = None,
+        nlist_backend: str = "auto",
         head: str | None = None,
         **kwargs: Any,
     ) -> None:
         Calculator.__init__(self, label=label, **kwargs)
         self.dp = DeepPot(
             str(Path(model).resolve()),
             neighbor_list=neighbor_list,
+            nlist_backend=nlist_backend,
             head=head,
         )
         if type_dict:

deepmd/pt/infer/deep_eval.py

@@ -46,6 +46,9 @@
 from deepmd.pt.model.model import (
     get_model,
 )
+from deepmd.pt.model.model.transform_output import (
+    communicate_extended_output,
+)
 from deepmd.pt.model.network.network import (
     TypeEmbedNetConsistent,
 )
@@ -67,6 +70,10 @@
     to_numpy_array,
     to_torch_tensor,
 )
+from deepmd.pt_expt.utils.nlist import (
+    build_neighbor_list_vesin_torch,
+    is_vesin_torch_available,
+)
 from deepmd.utils.batch_size import (
     RetrySignal,
 )
@@ -116,6 +123,14 @@ class DeepEval(DeepEvalBackend):
     neighbor_list : ase.neighborlist.NewPrimitiveNeighborList, optional
         The ASE neighbor list class to produce the neighbor list. If None, the
         neighbor list will be built natively in the model.
+    nlist_backend : str, default: "auto"
+        Which algorithm builds the neighbor list on the Python inference path.
+        ``"auto"`` uses the O(N) ``vesin`` cell list when it is applicable
+        (``vesin`` installed and a non-spin, non-hessian model) and silently
+        falls back to the native all-pairs builder otherwise. ``"vesin"``
+        strictly requires the vesin path and raises ``ValueError`` if it cannot
+        be used (vesin missing, or a spin / hessian model). ``"native"`` forces
+        the built-in all-pairs O(N^2) builder.
     **kwargs : dict
         Keyword arguments.
     """
@@ -127,6 +142,7 @@ def __init__(
         *args: Any,
         auto_batch_size: bool | int | AutoBatchSize = True,
         neighbor_list: Optional["ase.neighborlist.NewPrimitiveNeighborList"] = None,
+        nlist_backend: str = "auto",
         head: str | int | None = None,
         no_jit: bool = False,
         **kwargs: Any,
@@ -239,6 +255,48 @@ def __init__(
         if callable(self._has_spin):
             self._has_spin = self._has_spin()
         self._has_hessian = self.model_def_script.get("hessian_mode", False)
+        self._setup_nlist_backend(nlist_backend)
+
+    def _setup_nlist_backend(self, nlist_backend: str) -> None:
+        """Resolve the requested neighbor-list backend for the inference path.
+
+        ``"auto"`` uses the O(N) vesin cell list when applicable and silently
+        falls back to native otherwise; ``"vesin"`` is strict and raises if it
+        cannot be honored; ``"native"`` forces the native builder.
+        """
+        if nlist_backend not in ("auto", "vesin", "native"):
+            raise ValueError(
+                f"Unknown nlist_backend {nlist_backend!r}; "
+                "expected 'auto', 'vesin', or 'native'."
+            )
+        self.nlist_backend = nlist_backend
+        # The vesin O(N) builder only handles the standard (non-spin,
+        # non-hessian) energy path; report why it is unavailable, if so.
+        unsupported = None
+        if self._has_spin:
+            unsupported = "spin"
+        elif self._has_hessian:
+            unsupported = "hessian"
+        if nlist_backend == "vesin":
+            # explicit request: fail loudly if it cannot be honored
+            if not is_vesin_torch_available():
+                raise ValueError(
+                    "nlist_backend='vesin' was requested but the 'vesin.torch' "
+                    "package is not installed; install it (`pip install "
+                    "vesin[torch]`) or use nlist_backend='native' (or 'auto')."
+                )
+            if unsupported is not None:
+                raise ValueError(
+                    f"nlist_backend='vesin' is not supported for {unsupported} "
+                    "models; use nlist_backend='native' (or 'auto')."
+                )
+            self._use_vesin = True
+        elif nlist_backend == "native":
+            self._use_vesin = False
+        else:  # auto: use vesin when possible, otherwise fall back silently
+            self._use_vesin = is_vesin_torch_available() and unsupported is None
+        if self._use_vesin:
+            self._nsel = self.dp.model["Default"].get_nsel()
 
     def get_rcut(self) -> float:
         """Get the cutoff radius of this model."""
@@ -539,6 +597,16 @@ def _eval_model(
         request_defs: list[OutputVariableDef],
         charge_spin: np.ndarray | None,
     ) -> tuple[np.ndarray, ...]:
+        if self._use_vesin:
+            return self._eval_model_vesin(
+                coords,
+                cells,
+                atom_types,
+                fparam,
+                aparam,
+                request_defs,
+                charge_spin,
+            )
         model = self.dp.to(DEVICE)
         prec = NP_PRECISION_DICT[RESERVED_PRECISION_DICT[GLOBAL_PT_FLOAT_PRECISION]]
 
@@ -612,6 +680,111 @@ def _eval_model(
                 )  # this is kinda hacky
         return tuple(results)
 
+    def _eval_model_vesin(
+        self,
+        coords: np.ndarray,
+        cells: np.ndarray | None,
+        atom_types: np.ndarray,
+        fparam: np.ndarray | None,
+        aparam: np.ndarray | None,
+        request_defs: list[OutputVariableDef],
+        charge_spin: np.ndarray | None,
+    ) -> tuple[np.ndarray, ...]:
+        """Evaluate using an O(N) ``vesin`` neighbor list built on the host.
+
+        The neighbor list is built outside the model graph and fed to the
+        exported ``forward_common_lower`` interface; the extended-region output
+        is mapped back to local atoms with :func:`communicate_extended_output`,
+        exactly mirroring the native ``forward_common`` path.  This avoids the
+        native all-pairs O(N^2) neighbor-list build, which dominates runtime for
+        large systems on the Python / ASE inference path.
+        """
+        model = self.dp.model["Default"].to(DEVICE)
+        prec = NP_PRECISION_DICT[RESERVED_PRECISION_DICT[GLOBAL_PT_FLOAT_PRECISION]]
+
+        nframes = coords.shape[0]
+        if len(atom_types.shape) == 1:
+            natoms = len(atom_types)
+            atom_types = np.tile(atom_types, nframes).reshape(nframes, -1)
+        else:
+            natoms = len(atom_types[0])
+
+        coords = coords.reshape([nframes, natoms, 3])
+        # Build the neighbor list on the model's device with the vesin.torch
+        # cell list (CPU or CUDA).  The lower interface re-formats (distance-sort,
+        # truncate, type-split) the candidate list, so a single mixed list of the
+        # nearest sum(sel) neighbors is sufficient here (matches forward_common,
+        # which builds the nlist with mixed_types=True and distinguishes in the
+        # lower interface).
+        coord_t = torch.tensor(
+            coords.astype(prec), dtype=GLOBAL_PT_FLOAT_PRECISION, device=DEVICE
+        )
+        atype_t = torch.tensor(atom_types, dtype=torch.long, device=DEVICE)
+        box_t = (
+            torch.tensor(
+                cells.reshape([nframes, 3, 3]).astype(prec),
+                dtype=GLOBAL_PT_FLOAT_PRECISION,
+                device=DEVICE,
+            )
+            if cells is not None
+            else None
+        )
+        ext_coord_input, ext_atype_input, nlist_input, mapping_input = (
+            build_neighbor_list_vesin_torch(
+                coord_t, box_t, atype_t, self.rcut, [self._nsel], False
+            )
+        )
+
+        if fparam is not None:
+            fparam_input = to_torch_tensor(
+                fparam.reshape(nframes, self.get_dim_fparam())
+            )
+        else:
+            fparam_input = None
+        if aparam is not None:
+            aparam_input = to_torch_tensor(
+                aparam.reshape(nframes, natoms, self.get_dim_aparam())
+            )
+        else:
+            aparam_input = None
+        if charge_spin is not None:
+            charge_spin_input = to_torch_tensor(charge_spin.reshape(nframes, 2))
+        else:
+            charge_spin_input = None
+        do_atomic_virial = any(
+            x.category == OutputVariableCategory.DERV_C for x in request_defs
+        )
+
+        model_ret = model.forward_common_lower(
+            ext_coord_input,
+            ext_atype_input,
+            nlist_input,
+            mapping_input,
+            fparam=fparam_input,
+            aparam=aparam_input,
+            do_atomic_virial=do_atomic_virial,
+            charge_spin=charge_spin_input,
+        )
+        batch_output = communicate_extended_output(
+            model_ret,
+            self.output_def,
+            mapping_input,
+            do_atomic_virial=do_atomic_virial,
+        )
+
+        results = []
+        for odef in request_defs:
+            # communicate_extended_output keys are the internal output names,
+            # which match odef.name (e.g. "energy_redu", "energy_derv_r").
+            if odef.name in batch_output and batch_output[odef.name] is not None:
+                shape = self._get_output_shape(odef, nframes, natoms)
+                out = batch_output[odef.name].reshape(shape).detach().cpu().numpy()
+                results.append(out)
+            else:
+                shape = self._get_output_shape(odef, nframes, natoms)
+                results.append(np.full(np.abs(shape), np.nan, dtype=prec))
+        return tuple(results)
+
     def _eval_model_spin(
         self,
         coords: np.ndarray,

deepmd/pt_expt/infer/deep_eval.py

@@ -54,6 +54,10 @@
 from deepmd.pt.utils.auto_batch_size import (
     AutoBatchSize,
 )
+from deepmd.pt_expt.utils.nlist import (
+    build_neighbor_list_vesin_torch,
+    is_vesin_torch_available,
+)
 
 if TYPE_CHECKING:
     import ase.neighborlist
@@ -103,6 +107,7 @@ def __init__(
         *args: Any,
         auto_batch_size: bool | int | AutoBatchSize = True,
         neighbor_list: Optional["ase.neighborlist.NewPrimitiveNeighborList"] = None,
+        nlist_backend: str = "auto",
         **kwargs: Any,
     ) -> None:
         self.output_def = output_def
@@ -122,6 +127,7 @@ def __init__(
                 "backend: expected `.pt2` / `.pte` (deployable archives) or "
                 "`.pt` (training checkpoint)."
             )
+        self._setup_nlist_backend(nlist_backend)
 
         if isinstance(auto_batch_size, bool):
             if auto_batch_size:
@@ -135,6 +141,48 @@ def __init__(
         else:
             raise TypeError("auto_batch_size should be bool, int, or AutoBatchSize")
 
+    def _setup_nlist_backend(self, nlist_backend: str) -> None:
+        """Resolve the requested neighbor-list backend for the inference path.
+
+        ``"auto"`` uses the O(N) vesin cell list when applicable and silently
+        falls back to native otherwise; ``"vesin"`` is strict and raises if it
+        cannot be honored; ``"native"`` forces the native builder. An explicitly
+        supplied ASE ``neighbor_list`` takes precedence over the vesin path.
+        """
+        if nlist_backend not in ("auto", "vesin", "native"):
+            raise ValueError(
+                f"Unknown nlist_backend {nlist_backend!r}; "
+                "expected 'auto', 'vesin', or 'native'."
+            )
+        self.nlist_backend = nlist_backend
+        ase_provided = self.neighbor_list is not None
+        unsupported = "spin" if self._is_spin else None
+        if nlist_backend == "vesin":
+            # explicit request: fail loudly if it cannot be honored
+            if not is_vesin_torch_available():
+                raise ValueError(
+                    "nlist_backend='vesin' was requested but the 'vesin.torch' "
+                    "package is not installed; install it (`pip install "
+                    "vesin[torch]`) or use nlist_backend='native' (or 'auto')."
+                )
+            if unsupported is not None:
+                raise ValueError(
+                    f"nlist_backend='vesin' is not supported for {unsupported} "
+                    "models; use nlist_backend='native' (or 'auto')."
+                )
+            if ase_provided:
+                raise ValueError(
+                    "nlist_backend='vesin' conflicts with an explicitly supplied "
+                    "neighbor_list; pass only one."
+                )
+            self._use_vesin = True
+        elif nlist_backend == "native":
+            self._use_vesin = False
+        else:  # auto: use vesin when possible, otherwise fall back silently
+            self._use_vesin = (
+                is_vesin_torch_available() and unsupported is None and not ase_provided
+            )
+
     def _init_from_model_json(self, model_json_str: str) -> None:
         """Deserialize model.json and derive model API from the dpmodel instance."""
         from deepmd.pt_expt.model.model import (
@@ -1053,7 +1101,30 @@ def _prepare_inputs(
         )
 
         coord_input = coords.reshape(nframes, natoms, 3)
-        if self.neighbor_list is not None:
+        if self._use_vesin:
+            # device-resident vesin.torch build: on GPU the neighbor search stays
+            # on the GPU.  forward_common_lower re-formats the candidate list, so
+            # the mixed/distinguished choice here only mirrors the ASE path.
+            coord_t = torch.tensor(coord_input, dtype=torch.float64, device=DEVICE)
+            atype_t = torch.tensor(atom_types, dtype=torch.int64, device=DEVICE)
+            box_t = (
+                torch.tensor(
+                    cells.reshape(nframes, 3, 3), dtype=torch.float64, device=DEVICE
+                )
+                if cells is not None
+                else None
+            )
+            ext_coord_t, ext_atype_t, nlist_t, mapping_t = (
+                build_neighbor_list_vesin_torch(
+                    coord_t,
+                    box_t,
+                    atype_t,
+                    self._rcut,
+                    self._sel,
+                    distinguish_types=not self._mixed_types,
+                )
+            )
+        elif self.neighbor_list is not None:
             # ASE path: builds nlist in numpy, then convert to tensors
             extended_coord, extended_atype, nlist, mapping = self._build_nlist_ase(
                 coord_input,