diff --git a/deepmd/calculator.py b/deepmd/calculator.py index 356bfeb9ce..8c679acf3c 100644 --- a/deepmd/calculator.py +++ b/deepmd/calculator.py @@ -91,6 +91,7 @@ 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) @@ -98,6 +99,7 @@ def __init__( str(Path(model).resolve()), neighbor_list=neighbor_list, head=head, + nlist_backend=nlist_backend, ) if type_dict: self.type_dict = type_dict diff --git a/deepmd/dpmodel/model/ener_model.py b/deepmd/dpmodel/model/ener_model.py index c8c75d3cca..a8280dbebf 100644 --- a/deepmd/dpmodel/model/ener_model.py +++ b/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"] diff --git a/deepmd/dpmodel/model/make_model.py b/deepmd/dpmodel/model/make_model.py index ebcc671f62..28e1118718 100644 --- a/deepmd/dpmodel/model/make_model.py +++ b/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 diff --git a/deepmd/dpmodel/utils/__init__.py b/deepmd/dpmodel/utils/__init__.py index 588c8ea1ae..a9af7a50e5 100644 --- a/deepmd/dpmodel/utils/__init__.py +++ b/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", diff --git a/deepmd/dpmodel/utils/default_neighbor_list.py b/deepmd/dpmodel/utils/default_neighbor_list.py new file mode 100644 index 0000000000..03b4289795 --- /dev/null +++ b/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 diff --git a/deepmd/dpmodel/utils/neighbor_list.py b/deepmd/dpmodel/utils/neighbor_list.py new file mode 100644 index 0000000000..49504cc7c3 --- /dev/null +++ b/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 diff --git a/deepmd/pt/infer/deep_eval.py b/deepmd/pt/infer/deep_eval.py index 0887ceb9df..e398ac9210 100644 --- a/deepmd/pt/infer/deep_eval.py +++ b/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.vesin_neighbor_list import ( + VesinNeighborList, + is_vesin_torch_available, +) from deepmd.utils.batch_size import ( RetrySignal, ) @@ -129,10 +136,14 @@ def __init__( neighbor_list: Optional["ase.neighborlist.NewPrimitiveNeighborList"] = None, head: str | int | None = None, no_jit: bool = False, + nlist_backend: str = "auto", **kwargs: Any, ) -> None: self.output_def = output_def self.model_path = model_file + self.neighbor_list = neighbor_list + # data modifier, populated only for frozen .pth models that carry one + self.modifier = None if str(self.model_path).endswith(".pt"): state_dict = torch.load( model_file, map_location=env.DEVICE, weights_only=True @@ -239,6 +250,68 @@ 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 neighbor-list construction strategy from a user choice. + + ``"native"`` uses the dense all-pairs builder; ``"vesin"`` forces the + O(N) ``vesin.torch`` cell list (raising if it is unavailable or the + model/inputs are unsupported); ``"auto"`` uses vesin when applicable and + silently falls back to the native builder otherwise. Results are + unchanged either way -- only the neighbor-search cost differs. + """ + if nlist_backend not in ("auto", "vesin", "native"): + raise ValueError( + f"Unknown nlist_backend '{nlist_backend}'; " + "expected 'auto', 'vesin', or 'native'." + ) + # reason vesin cannot be used (None means it can) + unsupported = None + if self._has_spin: + unsupported = "spin models" + elif self._has_hessian: + unsupported = "hessian models" + elif self.modifier is not None: + # the vesin path runs forward_common_lower directly, bypassing + # ModelWrapper.forward (which applies the data modifier); fall back + # to the native path so the modifier is still applied. + unsupported = "models with a data modifier" + elif "energy" not in self.dp.model["Default"].model_output_type(): + # _eval_lower_vesin reconstructs the backend output from the + # forward_common_lower / communicate keys via _OUTDEF_DP2BACKEND, + # which matches the model's own translation only for the energy + # model (e.g. the polar fitting key is "polarizability" but the + # backend output is "polar"). Restrict vesin to energy models -- + # the large-system inference target -- and fall back to native + # for the other fitting types. + unsupported = "non-energy models" + ase_provided = self.neighbor_list is not None + if nlist_backend == "native": + self._use_vesin = False + elif nlist_backend == "vesin": + if not is_vesin_torch_available(): + raise ImportError( + "nlist_backend='vesin' was requested but 'vesin.torch' 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}; " + "use nlist_backend='native' (or 'auto')." + ) + if ase_provided: + raise ValueError( + "nlist_backend='vesin' conflicts with an explicitly " + "supplied ASE neighbor_list; pass only one." + ) + self._use_vesin = True + 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 + ) + self._nlist_builder = VesinNeighborList() if self._use_vesin else None def get_rcut(self) -> float: """Get the cutoff radius of this model.""" @@ -586,17 +659,28 @@ def _eval_model( do_atomic_virial = any( x.category == OutputVariableCategory.DERV_C for x in request_defs ) - batch_output = model( - coord_input, - type_input, - box=box_input, - do_atomic_virial=do_atomic_virial, - fparam=fparam_input, - aparam=aparam_input, - charge_spin=charge_spin_input, - ) - if isinstance(batch_output, tuple): - batch_output = batch_output[0] + if self._use_vesin: + batch_output = self._eval_lower_vesin( + coord_input, + type_input, + box_input, + fparam_input, + aparam_input, + charge_spin_input, + do_atomic_virial, + ) + else: + batch_output = model( + coord_input, + type_input, + box=box_input, + do_atomic_virial=do_atomic_virial, + fparam=fparam_input, + aparam=aparam_input, + charge_spin=charge_spin_input, + ) + if isinstance(batch_output, tuple): + batch_output = batch_output[0] results = [] for odef in request_defs: @@ -612,6 +696,52 @@ def _eval_model( ) # this is kinda hacky return tuple(results) + def _eval_lower_vesin( + self, + coord: torch.Tensor, + atype: torch.Tensor, + box: torch.Tensor | None, + fparam: torch.Tensor | None, + aparam: torch.Tensor | None, + charge_spin: torch.Tensor | None, + do_atomic_virial: bool, + ) -> dict[str, torch.Tensor]: + """Evaluate via the O(N) vesin-built ``(i,j,S)`` extended neighbor list. + + Builds the extended representation with the vesin cell list, runs the + model's ``forward_common_lower``, and maps the extended outputs back to + local atoms with ``communicate_extended_output``. Returns a dict keyed + by backend names, matching the normal ``model()`` output so the caller's + extraction is unchanged. ``forward_common_atomic`` sets + ``requires_grad`` on the extended coordinates internally, exactly as on + the native path, so forces/virials are produced identically. + """ + inner = self.dp.model["Default"] + ext_coord, ext_atype, nlist, mapping = self._nlist_builder.build( + coord, atype, box, self.rcut, list(inner.get_sel()) + ) + model_lower = inner.forward_common_lower( + ext_coord, + ext_atype, + nlist, + mapping, + fparam=fparam, + aparam=aparam, + do_atomic_virial=do_atomic_virial, + charge_spin=charge_spin, + ) + predict = communicate_extended_output( + model_lower, + inner.model_output_def(), + mapping, + do_atomic_virial=do_atomic_virial, + ) + return { + backend: predict[internal] + for internal, backend in self._OUTDEF_DP2BACKEND.items() + if predict.get(internal) is not None + } + def _eval_model_spin( self, coords: np.ndarray, diff --git a/deepmd/pt_expt/infer/deep_eval.py b/deepmd/pt_expt/infer/deep_eval.py index 18ca3e4038..42e4181d65 100644 --- a/deepmd/pt_expt/infer/deep_eval.py +++ b/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.vesin_neighbor_list import ( + VesinNeighborList, + 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 @@ -123,6 +128,8 @@ def __init__( "`.pt` (training checkpoint)." ) + self._setup_nlist_backend(nlist_backend) + if isinstance(auto_batch_size, bool): if auto_batch_size: self.auto_batch_size = AutoBatchSize() @@ -135,6 +142,50 @@ 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 neighbor-list construction strategy from a user choice. + + ``"native"`` uses the dense all-pairs builder; ``"vesin"`` forces the + O(N) ``vesin.torch`` cell list (raising if it is unavailable or the + model/inputs are unsupported); ``"auto"`` uses vesin when applicable and + silently falls back to the native builder otherwise. Results are + unchanged either way -- only the neighbor-search cost differs. + """ + if nlist_backend not in ("auto", "vesin", "native"): + raise ValueError( + f"Unknown nlist_backend '{nlist_backend}'; " + "expected 'auto', 'vesin', or 'native'." + ) + is_spin = bool(getattr(self, "_is_spin", False)) + ase_provided = self.neighbor_list is not None + # reason vesin cannot be used (None means it can) + unsupported = "spin models" if is_spin else None + if nlist_backend == "native": + self._use_vesin = False + elif nlist_backend == "vesin": + if not is_vesin_torch_available(): + raise ImportError( + "nlist_backend='vesin' was requested but 'vesin.torch' 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}; " + "use nlist_backend='native' (or 'auto')." + ) + if ase_provided: + raise ValueError( + "nlist_backend='vesin' conflicts with an explicitly " + "supplied ASE neighbor_list; pass only one." + ) + self._use_vesin = True + 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 + ) + self._nlist_builder = VesinNeighborList() if self._use_vesin else None + 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 ( @@ -831,6 +882,21 @@ def _build_nlist_native( sel = self._sel mixed_types = self._mixed_types + if self._nlist_builder is not None: + # O(N) cell-list strategy (e.g. vesin): builds the same extended + # representation. Match the native builder's type handling + # (``distinguish_types=not mixed_types``) so consumers that bypass + # ``forward_common_lower``'s ``format_nlist`` -- e.g. + # ``eval_descriptor`` calling the descriptor directly -- receive the + # type-distinguished nlist a non-mixed-type descriptor expects. The + # main eval path is unaffected (its ``format_nlist`` re-formats). + extended_coord, extended_atype, nlist, mapping = self._nlist_builder.build( + coords, atom_types, cells, rcut, sel + ) + if not mixed_types: + nlist = nlist_distinguish_types(nlist, extended_atype, sel) + return extended_coord, extended_atype, nlist, mapping + if cells is not None: box_input = cells.reshape(nframes, 3, 3) coord_normalized = normalize_coord(coords, box_input) diff --git a/deepmd/pt_expt/model/ener_model.py b/deepmd/pt_expt/model/ener_model.py index 1fdef5eaad..4f868043b6 100644 --- a/deepmd/pt_expt/model/ener_model.py +++ b/deepmd/pt_expt/model/ener_model.py @@ -18,6 +18,9 @@ from deepmd.dpmodel.model.make_hessian_model import ( make_hessian_model, ) +from deepmd.dpmodel.utils.neighbor_list import ( + NeighborList, +) from .make_model import ( make_model, @@ -59,7 +62,22 @@ def forward( aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, charge_spin: torch.Tensor | None = None, + neighbor_list: NeighborList | None = None, ) -> dict[str, torch.Tensor]: + """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.default_neighbor_list.DefaultNeighborList`), + reproducing the historical behavior; an alternative strategy (e.g. + the ``vesin`` O(N) cell list) may be injected to accelerate + neighbor-list construction without changing the model outputs. + """ model_ret = self.call_common( coord, atype, @@ -68,6 +86,7 @@ def forward( 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"] diff --git a/deepmd/pt_expt/utils/vesin_neighbor_list.py b/deepmd/pt_expt/utils/vesin_neighbor_list.py new file mode 100644 index 0000000000..1d55548c56 --- /dev/null +++ b/deepmd/pt_expt/utils/vesin_neighbor_list.py @@ -0,0 +1,232 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Device-resident O(N) neighbor-list strategy backed by ``vesin.torch``. + +This provides a :class:`~deepmd.dpmodel.utils.neighbor_list.NeighborList` +strategy that replaces the historical all-pairs ghost expansion (~27*N images + +a dense ``[N, 27N]`` distance matrix) with a cell list. ``vesin.torch`` returns +an ``(i, j, S)`` edge list (local neighbor index ``j`` plus integer periodic +image ``S``); we materialize only the *real-neighbor* ghosts ``coord[j] + S@box`` +and hand back the standard extended quartet ``(extended_coord, extended_atype, +nlist, mapping)``, so the rest of the model is unchanged. + +The neighbor search runs on the device of the input coordinates (CPU or CUDA), +so on a GPU the whole build stays on the GPU. When the inputs are numpy arrays +(the array-API ``dpmodel`` backend) the build is bridged through a CPU torch +tensor and the result converted back to numpy. The search itself is +non-differentiable -- it runs on detached coordinates -- while the returned +``extended_coord`` is rebuilt from the (possibly grad-carrying) input +coordinates so autograd for forces/virials flows through unchanged. +""" + +from typing import ( + Any, +) + +import torch + +from deepmd.dpmodel.utils.neighbor_list import ( + NeighborList, +) + + +def is_vesin_torch_available() -> bool: + """Whether the device-capable ``vesin.torch`` neighbor list is importable.""" + try: + import vesin.torch # noqa: F401 + except ImportError: + return False + return True + + +class VesinNeighborList(NeighborList): + """O(N) neighbor-list strategy using the ``vesin.torch`` cell list. + + Implements the :class:`~deepmd.dpmodel.utils.neighbor_list.NeighborList` + interface. Works on torch tensors (on their own device) and on numpy arrays + (bridged through a CPU torch tensor); the returned quartet matches the + namespace and device of the input coordinates. + """ + + def build( + self, + coord: Any, + atype: Any, + box: Any, + rcut: float, + sel: list[int], + ) -> tuple[Any, Any, Any, Any]: + """Build the extended system + candidate neighbor list with vesin. + + See :meth:`deepmd.dpmodel.utils.neighbor_list.NeighborList.build`. The + returned ``nlist`` is distance-sorted and truncated to ``sum(sel)`` + (matching the default builder); the lower interface still re-formats / + type-splits it. + """ + is_numpy = not isinstance(coord, torch.Tensor) + # vesin runs on the device of the inputs: numpy (the dpmodel backend) is + # bridged through CPU torch; torch tensors stay on their own device. Pin + # the ambient default device (cf. the ``with torch.device(...)`` guard + # around ``nl.compute`` below) so ``as_tensor`` is not affected by a + # placeholder default device -- e.g. tests set a CUDA default, under + # which a device-less ``as_tensor`` triggers CUDA init even for CPU input. + device = torch.device("cpu") if is_numpy else coord.device + with torch.device(device): + coord_t = torch.as_tensor(coord) + atype_t = torch.as_tensor(atype).to(torch.int64) + box_t = None if box is None else torch.as_tensor(box, dtype=coord_t.dtype) + + nframes = atype_t.shape[0] + nloc = atype_t.shape[1] + coord_t = coord_t.reshape(nframes, nloc, 3) + if box_t is not None: + box_t = box_t.reshape(nframes, 3, 3) + + frame_results = [ + _build_single( + coord_t[ff], + box_t[ff] if box_t is not None else None, + atype_t[ff], + rcut, + sel, + ) + for ff in range(nframes) + ] + max_nall = max(ec.shape[0] for ec, _, _, _ in frame_results) + device = coord_t.device + ext_coords, ext_atypes, nlists, mappings = [], [], [], [] + for ec, ea, nl, mp in frame_results: + pad = max_nall - ec.shape[0] + if pad > 0: + ec = torch.cat( + [ec, torch.zeros((pad, 3), dtype=ec.dtype, device=device)], dim=0 + ) + ea = torch.cat( + [ea, torch.full((pad,), -1, dtype=ea.dtype, device=device)], dim=0 + ) + mp = torch.cat( + [mp, torch.zeros((pad,), dtype=mp.dtype, device=device)], dim=0 + ) + ext_coords.append(ec) + ext_atypes.append(ea) + nlists.append(nl) + mappings.append(mp) + extended_coord = torch.stack(ext_coords, dim=0) + extended_atype = torch.stack(ext_atypes, dim=0) + nlist = torch.stack(nlists, dim=0) + mapping = torch.stack(mappings, dim=0) + + if is_numpy: + return ( + extended_coord.detach().cpu().numpy(), + extended_atype.cpu().numpy(), + nlist.cpu().numpy(), + mapping.cpu().numpy(), + ) + return extended_coord, extended_atype, nlist, mapping + + +def _build_single( + positions: torch.Tensor, + cell: torch.Tensor | None, + atype: torch.Tensor, + rcut: float, + sel: list[int], +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Single-frame ``(i,j,S)`` -> minimal-extended conversion. + + The cell list runs on detached coordinates (the search is + non-differentiable); the returned ``extended_coord`` is rebuilt from + ``positions`` so gradients flow to the local atoms and box. + """ + import vesin.torch + + device = positions.device + nsel = sum(sel) + nloc = positions.shape[0] + + # Empty system: vesin rejects an empty `points` array ("NULL pointer"). + # Return an empty extended representation directly, matching the native + # builder's handling of a zero-atom frame. + if nloc == 0: + return ( + positions, + atype, + torch.full((0, nsel), -1, dtype=torch.int64, device=device), + torch.zeros((0,), dtype=torch.int64, device=device), + ) + + periodic = cell is not None + box = ( + cell if periodic else torch.zeros((3, 3), dtype=positions.dtype, device=device) + ) + + # Pin the default device to the input's device: vesin.torch allocates some + # internal tensors on the ambient default device, which may be a fake/other + # device in some contexts (e.g. tests set a placeholder CUDA default). The + # search runs on detached inputs -- it is non-differentiable. + nl = vesin.torch.NeighborList(cutoff=rcut, full_list=True) + with torch.device(device): + ii, jj, ss = nl.compute( + points=positions.detach(), + box=box.detach(), + periodic=periodic, + quantities="ijS", + ) + ii = ii.to(torch.int64) + jj = jj.to(torch.int64) + ss = ss.to(positions.dtype) + + # ghost atoms: neighbors reached through a non-zero periodic shift. Rebuild + # their coordinates from the grad-carrying `positions`/`box` so autograd for + # forces/virials flows through the extended coordinates unchanged. + out_mask = torch.any(ss != 0, dim=1) + out_idx = jj[out_mask] + out_coords = positions[out_idx] + ss[out_mask] @ box + nghost = int(out_idx.shape[0]) + + extended_coord = torch.cat([positions, out_coords], dim=0) + extended_atype = torch.cat([atype, atype[out_idx]], dim=0) + mapping = torch.cat( + [torch.arange(nloc, dtype=torch.int64, device=device), out_idx], dim=0 + ) + + # remap neighbor column indices: ghosts -> [nloc, nloc + nghost) + neigh_idx = jj.clone() + neigh_idx[out_mask] = torch.arange( + nloc, nloc + nghost, dtype=torch.int64, device=device + ) + + # group pairs by center atom (vesin does not guarantee CSR ordering) + counts = torch.bincount(ii, minlength=nloc) + max_nn = int(counts.max()) if counts.numel() > 0 else 0 + order = torch.argsort(ii, stable=True) + rows = ii[order] + cols = torch.arange(ii.shape[0], dtype=torch.int64, device=device) - ( + torch.repeat_interleave(torch.cumsum(counts, 0) - counts, counts) + ) + dense_idx = torch.full((nloc, max_nn), -1, dtype=torch.int64, device=device) + if ii.shape[0] > 0: + dense_idx[rows, cols] = neigh_idx[order] + + # sort candidates by distance, keep the nsel nearest within rcut, pad with -1 + valid = dense_idx >= 0 + lookup = torch.where(valid, dense_idx, torch.zeros_like(dense_idx)) + dists = torch.linalg.norm( + extended_coord.detach()[lookup] - positions.detach()[:, None, :], dim=-1 + ) + valid &= dists <= rcut + dists = torch.where(valid, dists, torch.full_like(dists, float("inf"))) + sort_order = torch.argsort(dists, dim=-1) + sorted_idx = torch.take_along_dim(dense_idx, sort_order, dim=-1) + sorted_valid = torch.take_along_dim(valid, sort_order, dim=-1) + + nlist = torch.full((nloc, nsel), -1, dtype=torch.int64, device=device) + keep = min(nsel, max_nn) + if keep > 0: + nlist[:, :keep] = torch.where( + sorted_valid[:, :keep], + sorted_idx[:, :keep], + torch.full_like(sorted_idx[:, :keep], -1), + ) + + return extended_coord, extended_atype, nlist, mapping diff --git a/pyproject.toml b/pyproject.toml index 35fc0fdb18..f2a89155fa 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -58,6 +58,8 @@ dependencies = [ 'array-api-compat', 'lmdb', 'msgpack', + # O(N) cell-list neighbor list (vesin.torch) for fast Python/ASE inference + 'vesin[torch]', ] requires-python = ">=3.10" keywords = ["deepmd"] diff --git a/source/tests/pt/model/test_nlist_backend.py b/source/tests/pt/model/test_nlist_backend.py new file mode 100644 index 0000000000..ffc4b3e64f --- /dev/null +++ b/source/tests/pt/model/test_nlist_backend.py @@ -0,0 +1,156 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""``nlist_backend`` dispatch + vesin/native equivalence for the pt backend. + +The pt model is reconstructed eagerly in ``DeepEval`` and evaluated via +``forward_common_lower`` when the O(N) vesin neighbor list is selected (the +exported TorchScript graph is untouched). native and vesin must give identical +results, and the ``nlist_backend`` choice must dispatch / validate correctly. +""" + +import copy + +import numpy as np +import pytest +import torch + +from deepmd.infer.deep_pot import ( + DeepPot, +) +from deepmd.pt.model.model import ( + get_model, +) +from deepmd.pt.train.wrapper import ( + ModelWrapper, +) +from deepmd.pt_expt.utils.vesin_neighbor_list import ( + is_vesin_torch_available, +) + +pytestmark = pytest.mark.skipif( + not is_vesin_torch_available(), reason="vesin.torch is not installed" +) + +TYPE_MAP = ["O", "H", "B"] + +model_se_e2_a = { + "type_map": TYPE_MAP, + "descriptor": { + "type": "se_e2_a", + "sel": [20, 20, 8], + "rcut_smth": 0.5, + "rcut": 4.0, + "neuron": [6, 12], + "resnet_dt": False, + "axis_neuron": 4, + "seed": 1, + }, + "fitting_net": {"neuron": [8, 8], "resnet_dt": True, "seed": 1}, +} + +model_dpa1 = { + "type_map": TYPE_MAP, + "descriptor": { + "type": "se_atten", + "sel": 40, + "rcut_smth": 0.5, + "rcut": 4.0, + "neuron": [6, 12, 24], + "axis_neuron": 4, + "attn": 16, + "attn_layer": 2, + "attn_dotr": True, + "attn_mask": False, + "activation_function": "tanh", + "set_davg_zero": True, + "type_one_side": True, + "seed": 1, + }, + "fitting_net": {"neuron": [8, 8], "resnet_dt": True, "seed": 1}, +} + +ALL_MODELS = {"se_e2_a": model_se_e2_a, "dpa1": model_dpa1} + + +def _save_pt(md_dict: dict, path: str) -> None: + """Write a minimal loadable .pt (state_dict + model_params) for DeepPot.""" + model = get_model(copy.deepcopy(md_dict)).to(torch.float64) + wrapper = ModelWrapper(model, model_params=copy.deepcopy(md_dict)) + torch.save(wrapper.state_dict(), path) + + +def _system(): + rng = np.random.default_rng(20240604) + coords = (rng.random((1, 8, 3)) * 6.0).astype(np.float64) + atype = np.array([0, 0, 1, 1, 2, 0, 1, 2], dtype=np.int64) + box = (np.eye(3) * 6.0).reshape(1, 9).astype(np.float64) + return coords, atype, box + + +def _multiframe_system(nframes: int = 3): + """Frames with different box sizes -> different per-frame ghost counts, + exercising the vesin builder's pad-to-common-nall + stack path. + """ + rng = np.random.default_rng(20240604) + atype = np.array([0, 0, 1, 1, 2, 0, 1, 2], dtype=np.int64) + coords, boxes = [], [] + for ff in range(nframes): + box_len = 6.0 + 1.5 * ff + coords.append((rng.random((len(atype), 3)) * box_len).astype(np.float64)) + boxes.append((np.eye(3) * box_len).reshape(9).astype(np.float64)) + return np.stack(coords, axis=0), atype, np.stack(boxes, axis=0) + + +@pytest.fixture(scope="module") +def pt_files(tmp_path_factory): + d = tmp_path_factory.mktemp("nlist_backend") + files = {} + for name, md in ALL_MODELS.items(): + p = str(d / f"{name}.pt") + _save_pt(md, p) + files[name] = p + return files + + +def test_default_is_auto(pt_files) -> None: + # vesin is available (module skip guard), non-spin/non-hessian -> auto picks it + assert DeepPot(pt_files["se_e2_a"]).deep_eval._use_vesin is True + + +def test_native_disables_vesin(pt_files) -> None: + dp = DeepPot(pt_files["se_e2_a"], nlist_backend="native") + assert dp.deep_eval._use_vesin is False + + +def test_invalid_raises(pt_files) -> None: + with pytest.raises(ValueError): + DeepPot(pt_files["se_e2_a"], nlist_backend="bogus") + + +@pytest.mark.parametrize("name", list(ALL_MODELS)) # descriptor family +@pytest.mark.parametrize("periodic", [False, True]) # non-PBC vs PBC +def test_vesin_matches_native(pt_files, name: str, periodic: bool) -> None: + """Vesin and native give identical energy/force/virial/atomic-virial.""" + coords, atype, box = _system() + cells = box if periodic else None + dp_native = DeepPot(pt_files[name], nlist_backend="native") + dp_vesin = DeepPot(pt_files[name], nlist_backend="vesin") + ref = dp_native.eval(coords, cells, atype, atomic=True) + out = dp_vesin.eval(coords, cells, atype, atomic=True) + for a, b, label in zip(ref, out, ["e", "f", "v", "ae", "av"], strict=True): + np.testing.assert_allclose( + a, b, rtol=1e-9, atol=1e-9, err_msg=f"{name} {label}" + ) + + +@pytest.mark.parametrize("name", list(ALL_MODELS)) # descriptor family +def test_vesin_matches_native_multiframe(pt_files, name: str) -> None: + """Multi-frame eval (frames with differing ghost counts) matches native.""" + coords, atype, box = _multiframe_system() + dp_native = DeepPot(pt_files[name], nlist_backend="native") + dp_vesin = DeepPot(pt_files[name], nlist_backend="vesin") + ref = dp_native.eval(coords, box, atype, atomic=True) + out = dp_vesin.eval(coords, box, atype, atomic=True) + for a, b, label in zip(ref, out, ["e", "f", "v", "ae", "av"], strict=True): + np.testing.assert_allclose( + a, b, rtol=1e-9, atol=1e-9, err_msg=f"{name} {label}" + ) diff --git a/source/tests/pt_expt/infer/test_deep_eval.py b/source/tests/pt_expt/infer/test_deep_eval.py index 7537575f1a..a921a94a7f 100644 --- a/source/tests/pt_expt/infer/test_deep_eval.py +++ b/source/tests/pt_expt/infer/test_deep_eval.py @@ -2203,5 +2203,99 @@ def test_eval_fitting_last_layer_ase_vs_native(self) -> None: np.testing.assert_allclose(f_native, f_ase, rtol=1e-10, atol=1e-10) +class TestDeepEvalNlistBackend(unittest.TestCase): + """Dispatch + equivalence of the ``nlist_backend`` selection (.pte path). + + The vesin O(N) neighbor-list strategy must, through the compiled + ``forward_common_lower``, give identical results to the dense native + builder; the dispatch must validate the choice and fall back / raise + according to ``auto`` / ``native`` / ``vesin``. + """ + + @classmethod + def setUpClass(cls) -> None: + cls.rcut = 4.0 + cls.rcut_smth = 0.5 + cls.sel = [12, 10] + cls.nt = 2 + cls.type_map = ["foo", "bar"] + ds = DescrptSeA(cls.rcut, cls.rcut_smth, cls.sel) + ft = EnergyFittingNet( + cls.nt, ds.get_dim_out(), mixed_types=ds.mixed_types(), seed=GLOBAL_SEED + ) + cls.model = EnergyModel(ds, ft, type_map=cls.type_map).to(torch.float64) + cls.model.eval() + cls.tmpfile = tempfile.NamedTemporaryFile(suffix=".pte", delete=False) + cls.tmpfile.close() + deserialize_to_file( + cls.tmpfile.name, {"model": cls.model.serialize()}, do_atomic_virial=True + ) + + @classmethod + def tearDownClass(cls) -> None: + os.unlink(cls.tmpfile.name) + + def _vesin_available(self) -> bool: + from deepmd.pt_expt.utils.vesin_neighbor_list import ( + is_vesin_torch_available, + ) + + return is_vesin_torch_available() + + def test_default_is_auto(self) -> None: + dp = DeepPot(self.tmpfile.name) + self.assertEqual(dp.deep_eval._use_vesin, self._vesin_available()) + + def test_native_disables_vesin(self) -> None: + dp = DeepPot(self.tmpfile.name, nlist_backend="native") + self.assertFalse(dp.deep_eval._use_vesin) + + def test_invalid_raises(self) -> None: + with self.assertRaises(ValueError): + DeepPot(self.tmpfile.name, nlist_backend="bogus") + + def test_vesin_requested_but_unavailable_raises(self) -> None: + if self._vesin_available(): + self.skipTest("vesin.torch is installed; cannot test the missing path") + with self.assertRaises(ImportError): + DeepPot(self.tmpfile.name, nlist_backend="vesin") + + def test_vesin_matches_native(self) -> None: + if not self._vesin_available(): + self.skipTest("vesin.torch is not installed") + rng = np.random.default_rng(GLOBAL_SEED + 21) + natoms = 6 + coords = rng.random((1, natoms, 3)) * 8.0 + atom_types = np.array([i % self.nt for i in range(natoms)], dtype=np.int32) + dp_native = DeepPot(self.tmpfile.name, nlist_backend="native") + dp_vesin = DeepPot(self.tmpfile.name, nlist_backend="vesin") + for cells in (np.eye(3).reshape(1, 9) * 10.0, None): # PBC and non-PBC + ref = dp_native.eval(coords, cells, atom_types, atomic=True) + out = dp_vesin.eval(coords, cells, atom_types, atomic=True) + for a, b, name in zip(ref, out, ["e", "f", "v", "ae", "av"], strict=True): + np.testing.assert_allclose( + a, b, rtol=1e-10, atol=1e-10, err_msg=f"vesin vs native: {name}" + ) + + def test_eval_descriptor_vesin_matches_native(self) -> None: + """eval_descriptor bypasses forward_common_lower's format_nlist, so the + vesin builder must apply the same type-distinguishing as the native + builder for a non-mixed-type model (regression for the eval_descriptor + mismatch). + """ + if not self._vesin_available(): + self.skipTest("vesin.torch is not installed") + rng = np.random.default_rng(GLOBAL_SEED + 31) + natoms = 6 + coords = rng.random((1, natoms, 3)) * 8.0 + atom_types = np.array([i % self.nt for i in range(natoms)], dtype=np.int32) + dp_native = DeepPot(self.tmpfile.name, nlist_backend="native") + dp_vesin = DeepPot(self.tmpfile.name, nlist_backend="vesin") + for cells in (np.eye(3).reshape(1, 9) * 10.0, None): # PBC and non-PBC + d_native = dp_native.deep_eval.eval_descriptor(coords, cells, atom_types) + d_vesin = dp_vesin.deep_eval.eval_descriptor(coords, cells, atom_types) + np.testing.assert_allclose(d_native, d_vesin, rtol=1e-10, atol=1e-10) + + if __name__ == "__main__": unittest.main() diff --git a/source/tests/pt_expt/utils/test_neighbor_list.py b/source/tests/pt_expt/utils/test_neighbor_list.py new file mode 100644 index 0000000000..ca959eabe8 --- /dev/null +++ b/source/tests/pt_expt/utils/test_neighbor_list.py @@ -0,0 +1,499 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Equivalence of the vesin O(N) ``NeighborList`` strategy with the default builder. + +A ``NeighborList`` strategy is injected at ``forward_common``/``call_common``, +replacing the dense all-pairs ghost expansion (~27*N images + an O(N^2) distance +matrix) with vesin's O(N) cell list. Both strategies hand the *same* extended +representation to the downstream model, so every model output (energy, force, +virial, atomic virial) must match the default builder to fp round-off. + +Two layers are covered: + +* ``test_builder_*`` -- the builder in isolation, asserting the per-atom + neighbor *distance multisets* match the default, for the numpy (dpmodel) and + torch (pt/pt_expt) namespaces, periodic and non-periodic, and that the + returned tensors live on the input device. +* ``test_dpmodel_equivalence`` / ``test_pt_expt_equivalence`` / + ``test_default_fallback`` -- full model equivalence across descriptor families + (non-mixed, attention/mixed-types, message-passing with single and multiple + cutoffs, repflows, hybrid), for dpmodel (energy/atomic energy) and pt_expt + (energy/force/virial/atomic virial), periodic and non-periodic, including the + ``neighbor_list=None`` default falling back to the dense builder + byte-identically. +""" + +import copy + +import numpy as np +import pytest +import torch + +from deepmd.dpmodel.model.model import get_model as get_model_dp +from deepmd.dpmodel.utils import ( + DefaultNeighborList, + NeighborList, +) +from deepmd.pt_expt.model import ( + get_model, +) +from deepmd.pt_expt.utils import ( + env, +) +from deepmd.pt_expt.utils.vesin_neighbor_list import ( + VesinNeighborList, + is_vesin_torch_available, +) + +from ...seed import ( + GLOBAL_SEED, +) + +# --- compact model configs (3-type type_map), reduced layers for test speed --- +TYPE_MAP = ["O", "H", "B"] + +model_se_e2_a = { + "type_map": TYPE_MAP, + "descriptor": { + "type": "se_e2_a", + "sel": [20, 20, 8], + "rcut_smth": 0.5, + "rcut": 4.0, + "neuron": [6, 12], + "resnet_dt": False, + "axis_neuron": 4, + "seed": 1, + }, + "fitting_net": {"neuron": [8, 8], "resnet_dt": True, "seed": 1}, +} + +model_se_r = { + "type_map": TYPE_MAP, + "descriptor": { + "type": "se_e2_r", + "sel": [20, 20, 8], + "rcut_smth": 0.5, + "rcut": 4.0, + "neuron": [6, 12], + "resnet_dt": False, + "seed": 1, + }, + "fitting_net": {"neuron": [8, 8], "resnet_dt": True, "seed": 1}, +} + +model_se_e3 = { + "type_map": TYPE_MAP, + "descriptor": { + "type": "se_e3", + "sel": [12, 12, 4], + "rcut_smth": 0.5, + "rcut": 4.0, + "neuron": [6, 12], + "resnet_dt": False, + "seed": 1, + }, + "fitting_net": {"neuron": [8, 8], "resnet_dt": True, "seed": 1}, +} + +model_dpa1 = { + "type_map": TYPE_MAP, + "descriptor": { + "type": "se_atten", + "sel": 40, + "rcut_smth": 0.5, + "rcut": 4.0, + "neuron": [6, 12, 24], + "axis_neuron": 4, + "attn": 16, + "attn_layer": 2, + "attn_dotr": True, + "attn_mask": False, + "activation_function": "tanh", + "scaling_factor": 1.0, + "normalize": False, + "temperature": 1.0, + "set_davg_zero": True, + "type_one_side": True, + "seed": 1, + }, + "fitting_net": {"neuron": [8, 8], "resnet_dt": True, "seed": 1}, +} + +model_se_atten_v2 = { + "type_map": TYPE_MAP, + "descriptor": { + "type": "se_atten_v2", + "sel": 40, + "rcut_smth": 0.5, + "rcut": 4.0, + "neuron": [6, 12, 24], + "axis_neuron": 4, + "attn": 16, + "attn_layer": 2, + "attn_dotr": True, + "attn_mask": False, + "activation_function": "tanh", + "set_davg_zero": False, + "seed": 1, + }, + "fitting_net": {"neuron": [8, 8], "resnet_dt": True, "seed": 1}, +} + +model_dpa2 = { + "type_map": TYPE_MAP, + "descriptor": { + "type": "dpa2", + "repinit": { + "rcut": 6.0, + "rcut_smth": 2.0, + "nsel": 30, + "neuron": [2, 4, 8], + "axis_neuron": 4, + "activation_function": "tanh", + }, + "repformer": { + "rcut": 4.0, + "rcut_smth": 0.5, + "nsel": 20, + "nlayers": 2, + "g1_dim": 8, + "g2_dim": 5, + "attn2_hidden": 3, + "attn2_nhead": 1, + "attn1_hidden": 5, + "attn1_nhead": 1, + "axis_neuron": 4, + "update_h2": False, + "update_g1_has_conv": True, + "update_g1_has_grrg": True, + "update_g1_has_drrd": True, + "update_g1_has_attn": True, + "update_g2_has_g1g1": True, + "update_g2_has_attn": True, + "attn2_has_gate": True, + }, + "seed": 1, + "add_tebd_to_repinit_out": False, + }, + "fitting_net": {"neuron": [8, 8], "resnet_dt": True, "seed": 1}, +} + +model_dpa3 = { + "type_map": TYPE_MAP, + "descriptor": { + "type": "dpa3", + "repflow": { + "n_dim": 12, + "e_dim": 8, + "a_dim": 6, + "nlayers": 2, + "e_rcut": 6.0, + "e_rcut_smth": 3.0, + "e_sel": 20, + "a_rcut": 4.0, + "a_rcut_smth": 2.0, + "a_sel": 10, + "axis_neuron": 4, + "update_angle": True, + "update_style": "res_residual", + "update_residual": 0.1, + "update_residual_init": "const", + }, + "activation_function": "tanh", + "use_tebd_bias": False, + "concat_output_tebd": False, + "seed": 1, + }, + "fitting_net": {"neuron": [8, 8], "resnet_dt": True, "seed": 1}, +} + +model_hybrid = { + "type_map": TYPE_MAP, + "descriptor": { + "type": "hybrid", + "list": [ + { + "type": "se_e2_a", + "sel": [20, 20, 8], + "rcut_smth": 0.5, + "rcut": 4.0, + "neuron": [6, 12], + "resnet_dt": False, + "axis_neuron": 4, + "seed": 1, + }, + { + "type": "se_e2_r", + "sel": [20, 20, 8], + "rcut_smth": 0.5, + "rcut": 4.0, + "neuron": [6, 12], + "resnet_dt": False, + "seed": 1, + }, + ], + }, + "fitting_net": {"neuron": [8, 8], "resnet_dt": True, "seed": 1}, +} + +ALL_MODELS = { + "se_e2_a": model_se_e2_a, + "se_r": model_se_r, + "se_e3": model_se_e3, + "dpa1": model_dpa1, + "se_atten_v2": model_se_atten_v2, + "dpa2": model_dpa2, + "dpa3": model_dpa3, + "hybrid": model_hybrid, +} + + +def _system(natoms: int = 6, box_len: float = 10.0, seed: int = GLOBAL_SEED): + """A small 3-type periodic system; returns numpy (coord, atype, box).""" + rng = np.random.default_rng(seed) + coord = (rng.random((1, natoms, 3)) * box_len).astype(np.float64) + atype = np.array([[0, 0, 1, 1, 2, 0]], dtype=np.int64)[:, :natoms] + box = (np.eye(3) * box_len).reshape(1, 9).astype(np.float64) + return coord, atype, box + + +def _multiframe_system(nframes: int = 3, natoms: int = 6, seed: int = GLOBAL_SEED): + """Multi-frame 3-type system whose frames have *different* geometries (and + box sizes), so the per-frame ghost counts differ and the builder's + pad-to-common-nall + stack path is exercised. + """ + rng = np.random.default_rng(seed) + coords, boxes = [], [] + for ff in range(nframes): + box_len = 6.0 + 1.5 * ff # vary box -> vary ghost count per frame + coords.append((rng.random((natoms, 3)) * box_len).astype(np.float64)) + boxes.append((np.eye(3) * box_len).reshape(9).astype(np.float64)) + coord = np.stack(coords, axis=0) + atype = np.tile( + np.array([[0, 0, 1, 1, 2, 0]], dtype=np.int64)[:, :natoms], (nframes, 1) + ) + box = np.stack(boxes, axis=0) + return coord, atype, box + + +def _per_atom_neighbor_dists(ext_coord, nlist, coord): + """Sorted, rounded valid-neighbor distances for each local atom.""" + ext_coord = np.asarray(ext_coord).reshape(-1, 3) + coord = np.asarray(coord).reshape(-1, 3) + out = [] + for i in range(coord.shape[0]): + ds = [ + round(float(np.linalg.norm(ext_coord[j] - coord[i])), 6) + for j in np.asarray(nlist)[i] + if j >= 0 + ] + out.append(sorted(ds)) + return out + + +pytestmark = pytest.mark.skipif( + not is_vesin_torch_available(), reason="vesin.torch is not installed" +) + + +def _tol(model_dict: dict) -> dict: + """Equivalence tolerance; loosened only for float32 models.""" + prec = str(model_dict["descriptor"].get("precision", "float64")) + if "32" in prec: + return {"rtol": 1e-5, "atol": 1e-5} + return {"rtol": 1e-9, "atol": 1e-9} + + +def test_builder_isinstance() -> None: + assert isinstance(VesinNeighborList(), NeighborList) + assert isinstance(DefaultNeighborList(), NeighborList) + + +@pytest.mark.parametrize("periodic", [False, True]) # non-PBC vs PBC +def test_builder_matches_default(periodic: bool) -> None: + """The vesin builder produces the same per-atom neighbor distance multisets + as the default all-pairs builder, in both the numpy (dpmodel) and torch + (pt/pt_expt) namespaces. + """ + coord_np, atype_np, box_np = _system() + box_np = box_np if periodic else None + rcut, sel = 4.0, [20, 20, 8] + ec_d, _, nl_d, _ = DefaultNeighborList().build( + coord_np, atype_np, box_np, rcut, sel + ) + ref = _per_atom_neighbor_dists(ec_d, nl_d[0], coord_np[0]) + # numpy (dpmodel) namespace + ec_v, _, nl_v, _ = VesinNeighborList().build(coord_np, atype_np, box_np, rcut, sel) + assert _per_atom_neighbor_dists(ec_v, nl_v[0], coord_np[0]) == ref + # torch (pt/pt_expt) namespace + coord_t = torch.tensor(coord_np, dtype=torch.float64) + atype_t = torch.tensor(atype_np, dtype=torch.int64) + box_t = None if box_np is None else torch.tensor(box_np, dtype=torch.float64) + ec_vt, _, nl_vt, _ = VesinNeighborList().build(coord_t, atype_t, box_t, rcut, sel) + assert torch.is_tensor(ec_vt) + assert ( + _per_atom_neighbor_dists( + ec_vt[0].cpu().numpy(), nl_vt[0].cpu().numpy(), coord_np[0] + ) + == ref + ) + + +@pytest.mark.parametrize("periodic", [False, True]) # non-PBC vs PBC +def test_builder_empty_system(periodic: bool) -> None: + """A zero-atom frame must not crash vesin (which rejects empty points); the + builder returns an empty extended representation, matching the native path. + """ + coord = np.zeros((1, 0, 3), dtype=np.float64) + atype = np.zeros((1, 0), dtype=np.int64) + box = (np.eye(3) * 10.0).reshape(1, 9).astype(np.float64) if periodic else None + sel = [20, 20, 8] + ec, ea, nl, mp = VesinNeighborList().build(coord, atype, box, 4.0, sel) + assert ec.shape == (1, 0, 3) + assert ea.shape == (1, 0) + assert nl.shape == (1, 0, sum(sel)) + assert mp.shape == (1, 0) + + +def test_builder_outputs_on_input_device() -> None: + coord_np, atype_np, box_np = _system() + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + coord_t = torch.tensor(coord_np, dtype=torch.float64, device=device) + atype_t = torch.tensor(atype_np, dtype=torch.int64, device=device) + box_t = torch.tensor(box_np, dtype=torch.float64, device=device) + for t in VesinNeighborList().build(coord_t, atype_t, box_t, 4.0, [20, 20, 8]): + assert t.device.type == device.type + + +@pytest.mark.parametrize("periodic", [False, True]) # non-PBC vs PBC +def test_builder_multiframe_matches_default(periodic: bool) -> None: + """Multi-frame build (frames with differing ghost counts) exercises the + pad-to-common-nall + stack path; every frame's neighbor multiset must still + match the default builder, in numpy and torch namespaces. + """ + coord_np, atype_np, box_np = _multiframe_system() + box_np = box_np if periodic else None + rcut, sel = 4.0, [20, 20, 8] + ec_d, _, nl_d, _ = DefaultNeighborList().build( + coord_np, atype_np, box_np, rcut, sel + ) + ec_v, _, nl_v, _ = VesinNeighborList().build(coord_np, atype_np, box_np, rcut, sel) + coord_t = torch.tensor(coord_np, dtype=torch.float64) + atype_t = torch.tensor(atype_np, dtype=torch.int64) + box_t = None if box_np is None else torch.tensor(box_np, dtype=torch.float64) + ec_vt, _, nl_vt, _ = VesinNeighborList().build(coord_t, atype_t, box_t, rcut, sel) + for ff in range(coord_np.shape[0]): + ref = _per_atom_neighbor_dists(ec_d[ff], nl_d[ff], coord_np[ff]) + assert _per_atom_neighbor_dists(ec_v[ff], nl_v[ff], coord_np[ff]) == ref + assert ( + _per_atom_neighbor_dists( + ec_vt[ff].cpu().numpy(), nl_vt[ff].cpu().numpy(), coord_np[ff] + ) + == ref + ) + + +@pytest.mark.parametrize("name", list(ALL_MODELS)) # descriptor family +@pytest.mark.parametrize("periodic", [False, True]) # non-PBC vs PBC +def test_dpmodel_equivalence(name: str, periodic: bool) -> None: + """Energy / atomic energy (dpmodel) are invariant to the nlist strategy.""" + coord_np, atype_np, box_np = _system() + box_np = box_np if periodic else None + model_dict = ALL_MODELS[name] + md = get_model_dp(copy.deepcopy(model_dict)) + r0 = md.call(coord_np, atype_np, box=box_np, neighbor_list=DefaultNeighborList()) + r1 = md.call(coord_np, atype_np, box=box_np, neighbor_list=VesinNeighborList()) + for k in ("energy", "atom_energy"): + np.testing.assert_allclose( + r0[k], r1[k], err_msg=f"{name} {k}", **_tol(model_dict) + ) + + +@pytest.mark.parametrize("name", list(ALL_MODELS)) # descriptor family +@pytest.mark.parametrize("periodic", [False, True]) # non-PBC vs PBC +def test_pt_expt_equivalence(name: str, periodic: bool) -> None: + """pt_expt energy / force / virial / atomic virial are invariant to the + nlist strategy (force/virial come from the existing autograd routines). + """ + coord_np, atype_np, box_np = _system() + box_np = box_np if periodic else None + model_dict = ALL_MODELS[name] + md = get_model(copy.deepcopy(model_dict)).to(env.DEVICE) + md.eval() + box_t = ( + None + if box_np is None + else torch.tensor(box_np, dtype=torch.float64, device=env.DEVICE) + ) + atype_t = torch.tensor(atype_np, dtype=torch.int64, device=env.DEVICE) + results = {} + for tag, nl in (("def", DefaultNeighborList()), ("ves", VesinNeighborList())): + coord_t = torch.tensor( + coord_np, dtype=torch.float64, device=env.DEVICE + ).requires_grad_(True) + results[tag] = md.forward( + coord_t, atype_t, box=box_t, do_atomic_virial=True, neighbor_list=nl + ) + for k in ("energy", "atom_energy", "force", "virial", "atom_virial"): + a, b = results["def"].get(k), results["ves"].get(k) + if a is None or b is None: + continue + np.testing.assert_allclose( + a.detach().cpu().numpy(), + b.detach().cpu().numpy(), + err_msg=f"{name} {k}", + **_tol(model_dict), + ) + + +@pytest.mark.parametrize("name", ["se_e2_a", "dpa1"]) # non-mixed + attention +def test_pt_expt_multiframe_equivalence(name: str) -> None: + """Multi-frame (frames with differing ghost counts) pt_expt outputs are + invariant to the nlist strategy -- exercises the builder's per-frame pad + + stack feeding the batched model forward. + """ + coord_np, atype_np, box_np = _multiframe_system() + model_dict = ALL_MODELS[name] + md = get_model(copy.deepcopy(model_dict)).to(env.DEVICE) + md.eval() + atype_t = torch.tensor(atype_np, dtype=torch.int64, device=env.DEVICE) + box_t = torch.tensor(box_np, dtype=torch.float64, device=env.DEVICE) + results = {} + for tag, nl in (("def", DefaultNeighborList()), ("ves", VesinNeighborList())): + coord_t = torch.tensor( + coord_np, dtype=torch.float64, device=env.DEVICE + ).requires_grad_(True) + results[tag] = md.forward( + coord_t, atype_t, box=box_t, do_atomic_virial=True, neighbor_list=nl + ) + for k in ("energy", "force", "virial", "atom_virial"): + np.testing.assert_allclose( + results["def"][k].detach().cpu().numpy(), + results["ves"][k].detach().cpu().numpy(), + err_msg=f"{name} {k}", + **_tol(model_dict), + ) + + +@pytest.mark.parametrize("name", list(ALL_MODELS)) # descriptor family +def test_default_fallback(name: str) -> None: + """``neighbor_list=None`` equals an explicit DefaultNeighborList byte-for-byte.""" + coord_np, atype_np, box_np = _system() + md = get_model(copy.deepcopy(ALL_MODELS[name])).to(env.DEVICE) + md.eval() + box_t = torch.tensor(box_np, dtype=torch.float64, device=env.DEVICE) + atype_t = torch.tensor(atype_np, dtype=torch.int64, device=env.DEVICE) + outs = {} + for tag, kw in ( + ("none", {}), + ("explicit", {"neighbor_list": DefaultNeighborList()}), + ): + coord_t = torch.tensor( + coord_np, dtype=torch.float64, device=env.DEVICE + ).requires_grad_(True) + outs[tag] = md.forward(coord_t, atype_t, box=box_t, do_atomic_virial=True, **kw) + for k in ("energy", "force", "virial"): + np.testing.assert_array_equal( + outs["none"][k].detach().cpu().numpy(), + outs["explicit"][k].detach().cpu().numpy(), + err_msg=f"{name} {k}", + )